From d3547f9d03b519abbc2cbed89590244b526ab2cd Mon Sep 17 00:00:00 2001 From: Abhay Shukla <45638652+AbhayShuklaIIT@users.noreply.github.com> Date: Mon, 5 Oct 2026 19:13:49 +0530 Subject: [PATCH 001/952] [Test] Re-enable Voxtral HF reference test on Transformers v5 (#59771) Signed-off-by: AbhayShuklaIIT Co-authored-by: Claude Opus 5.5 --- tests/models/multimodal/generation/test_voxtral.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral.py b/tests/models/multimodal/generation/test_voxtral.py index 544c06020f66..3c9d2ace8cca 100644 --- a/tests/models/multimodal/generation/test_voxtral.py +++ b/tests/models/multimodal/generation/test_voxtral.py @@ -146,10 +146,6 @@ def _asset_to_openai_chunk(asset): ) -@pytest.mark.skip( - reason="VoxtralProcessor.apply_chat_template() in transformers v5 " - "doesn't resolve chat_template=None to the default template" -) def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets): """Compare vLLM Mistral-format output against HF Transformers reference. @@ -157,9 +153,6 @@ def test_hf_reference(hf_runner, vllm_runner, audio_assets: AudioTestAssets): attention backends), we compare per-token logprobs using the standard check_logprobs_close helper: when tokens diverge at a position, each runner's chosen token must appear in the other's top-k logprobs. - - Marked xfail(strict=False) so remaining edge-case mismatches - don't block CI. """ question = f"What's happening in these {len(audio_assets)} audio clips?" max_tokens = 10 From f9c9e8ac248cd19c1e1826693761e69b0eec1822 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Mon, 5 Oct 2026 22:01:56 +0800 Subject: [PATCH 002/952] [LoRA] Remove tensorizer (#60024) Signed-off-by: Jee Jee Li --- docs/models/extensions/tensorizer.md | 20 +-- examples/features/tensorize_vllm_model.py | 77 +--------- .../completion/test_tensorizer_entrypoint.py | 10 +- tests/lora/test_llama_tp.py | 133 ++---------------- vllm/lora/lora_model.py | 23 +-- vllm/lora/peft_helper.py | 28 +--- vllm/lora/request.py | 1 - vllm/lora/worker_manager.py | 5 +- .../model_executor/model_loader/tensorizer.py | 108 +------------- 9 files changed, 30 insertions(+), 375 deletions(-) diff --git a/docs/models/extensions/tensorizer.md b/docs/models/extensions/tensorizer.md index 0f3ee1ede436..1e0d5da81e6e 100644 --- a/docs/models/extensions/tensorizer.md +++ b/docs/models/extensions/tensorizer.md @@ -32,25 +32,15 @@ python examples/features/tensorize_vllm_model.py \ --suffix v1 ``` -This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`. If you intend on applying a LoRA adapter to your tensorized model, you can pass the HF id of the LoRA adapter in the above command, and the artifacts will be saved there too: - -```bash -python examples/features/tensorize_vllm_model.py \ - --model facebook/opt-125m \ - --lora-path \ - serialize \ - --serialized-directory s3://my-bucket \ - --suffix v1 -``` +This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`. ## Serving the model using Tensorizer -Once the model is serialized where you want it, you can load the model using `vllm serve` or the `LLM` entrypoint. You can pass the directory where you saved the model to the `model` argument for `LLM()` and `vllm serve`. For example, to serve the tensorized model saved previously with the LoRA adapter, you'd do: +Once the model is serialized where you want it, you can load the model using `vllm serve` or the `LLM` entrypoint. You can pass the directory where you saved the model to the `model` argument for `LLM()` and `vllm serve`. For example, to serve the tensorized model saved previously, you'd do: ```bash vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \ - --load-format tensorizer \ - --enable-lora + --load-format tensorizer ``` Or, with `LLM()`: @@ -60,7 +50,6 @@ from vllm import LLM llm = LLM( "s3://my-bucket/vllm/facebook/opt-125m/v1", load_format="tensorizer", - enable_lora=True, ) ``` @@ -73,7 +62,6 @@ As an example, CPU concurrency can be limited when serializing with `tensorizer` ```bash python examples/features/tensorize_vllm_model.py \ --model facebook/opt-125m \ - --lora-path \ serialize \ --serialized-directory s3://my-bucket \ --serialization-kwargs '{"limit_cpu_concurrency": 2}' \ @@ -85,7 +73,6 @@ As an example when customizing the loading process via `TensorDeserializer`, you ```bash vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \ --load-format tensorizer \ - --enable-lora \ --model-loader-extra-config '{"deserialization_kwargs": {"num_readers": 2}}' ``` @@ -96,7 +83,6 @@ from vllm import LLM llm = LLM( "s3://my-bucket/vllm/facebook/opt-125m/v1", load_format="tensorizer", - enable_lora=True, model_loader_extra_config={"deserialization_kwargs": {"num_readers": 2}}, ) ``` diff --git a/examples/features/tensorize_vllm_model.py b/examples/features/tensorize_vllm_model.py index a89b1781264d..4f4a40233c85 100644 --- a/examples/features/tensorize_vllm_model.py +++ b/examples/features/tensorize_vllm_model.py @@ -6,13 +6,11 @@ import os import uuid -from vllm import LLM, SamplingParams +from vllm import LLM from vllm.engine.arg_utils import EngineArgs -from vllm.lora.request import LoRARequest from vllm.model_executor.model_loader.tensorizer import ( TensorizerArgs, TensorizerConfig, - tensorize_lora_adapter, tensorize_vllm_model, tensorizer_kwargs_arg, ) @@ -105,21 +103,6 @@ under the `tensorizer options` section. These can also be used for deserialization in this example script, although `--tensorizer-uri` and `--path-to-tensors` are functionally the same in this case. - -Tensorizer can also be used to save and load LoRA adapters. A LoRA adapter -can be serialized directly with the path to the LoRA adapter on HF Hub and -a TensorizerConfig object. In this script, passing a HF id to a LoRA adapter -will serialize the LoRA adapter artifacts to `--serialized-directory`. - -You can then use the LoRA adapter with `vllm serve`, for instance, by ensuring -the LoRA artifacts are in your model artifacts directory and specifying -`--enable-lora`. For instance: - -``` -vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \ - --load-format tensorizer \ - --enable-lora -``` """ @@ -134,18 +117,6 @@ def get_parser(): ) parser = EngineArgs.add_cli_args(parser) - parser.add_argument( - "--lora-path", - type=str, - required=False, - help="Path to a LoRA adapter to " - "serialize along with model tensors. This can then be deserialized " - "along with the model by instantiating a TensorizerConfig object, " - "creating a dict from it with TensorizerConfig.to_serializable(), " - "and passing it to LoRARequest's initializer with the kwarg " - "tensorizer_config_dict.", - ) - subparsers = parser.add_subparsers(dest="command", required=True) serialize_parser = subparsers.add_parser( @@ -263,42 +234,12 @@ def merge_extra_config_with_tensorizer_config(extra_cfg: dict, cfg: TensorizerCo def deserialize(args, tensorizer_config): - if args.lora_path: - tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir - llm = LLM( - model=args.model, - load_format="tensorizer", - tensor_parallel_size=args.tensor_parallel_size, - model_loader_extra_config=tensorizer_config, - enable_lora=True, - ) - sampling_params = SamplingParams( - temperature=0, max_tokens=256, stop=["[/assistant]"] - ) - - # Truncating this as the extra text isn't necessary - prompts = ["[user] Write a SQL query to answer the question based on ..."] - - # Test LoRA load - print( - llm.generate( - prompts, - sampling_params, - lora_request=LoRARequest( - "sql-lora", - 1, - args.lora_path, - tensorizer_config_dict=tensorizer_config.to_serializable(), - ), - ) - ) - else: - llm = LLM( - model=args.model, - load_format="tensorizer", - tensor_parallel_size=args.tensor_parallel_size, - model_loader_extra_config=tensorizer_config, - ) + llm = LLM( + model=args.model, + load_format="tensorizer", + tensor_parallel_size=args.tensor_parallel_size, + model_loader_extra_config=tensorizer_config, + ) return llm @@ -366,10 +307,6 @@ def main(): **credentials, ) - if args.lora_path: - tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir - tensorize_lora_adapter(args.lora_path, tensorizer_config) - merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config) tensorize_vllm_model(engine_args, tensorizer_config) diff --git a/tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py b/tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py index 29c0c2dc8f97..813b64b2c527 100644 --- a/tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py +++ b/tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py @@ -13,13 +13,11 @@ from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.model_loader.tensorizer import ( TensorizerConfig, - tensorize_lora_adapter, tensorize_vllm_model, ) from vllm.platforms import current_platform MODEL_NAME = "unsloth/llama-3.2-1b-Instruct" -LORA_PATH = "davzoku/finqa_adapter_1b" def _cleanup(): @@ -44,11 +42,10 @@ def model_uri(tmp_dir): @pytest.fixture(scope="module") -def tensorize_model_and_lora(tmp_dir, model_uri): - tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri, lora_dir=tmp_dir) +def tensorize_model(model_uri): + tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri) args = EngineArgs(model=MODEL_NAME) - tensorize_lora_adapter(LORA_PATH, tensorizer_config) tensorize_vllm_model(args, tensorizer_config) # Manually invoke a _cleanup() here, as the cleanup() @@ -59,7 +56,7 @@ def tensorize_model_and_lora(tmp_dir, model_uri): @pytest.fixture(scope="module") -def server(model_uri, tensorize_model_and_lora): +def server(model_uri, tensorize_model): # In this case, model_uri is a directory with a model.tensors # file and all necessary model artifacts, particularly a # HF `config.json` file. In this case, Tensorizer can infer the @@ -72,7 +69,6 @@ def server(model_uri, tensorize_model_and_lora): "tensorizer", "--served-model-name", MODEL_NAME, - "--enable-lora", ] if current_platform.is_rocm(): args += ["--attention-backend", "TRITON_ATTN"] diff --git a/tests/lora/test_llama_tp.py b/tests/lora/test_llama_tp.py index 42f6ddc2f690..9ccb2aec257a 100644 --- a/tests/lora/test_llama_tp.py +++ b/tests/lora/test_llama_tp.py @@ -1,18 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import subprocess -import sys import pytest import vllm import vllm.config -from vllm import LLM from vllm.lora.request import LoRARequest -from vllm.model_executor.model_loader.tensorizer import TensorizerConfig from vllm.platforms import current_platform -from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test +from ..utils import create_new_process_for_each_test, multi_gpu_test PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|> I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request. @@ -40,7 +36,6 @@ def do_sample( llm: vllm.LLM, lora_path: str, lora_id: int, - tensorizer_config_dict: dict | None = None, ) -> list[str]: prompts = [ PROMPT_TEMPLATE.format(context="How many candidates are there?"), @@ -56,27 +51,11 @@ def do_sample( sampling_params = vllm.SamplingParams( temperature=0, max_tokens=64, stop=["<|im_end|>"] ) - if tensorizer_config_dict is not None: - outputs = llm.generate( - prompts, - sampling_params, - lora_request=LoRARequest( - str(lora_id), - lora_id, - lora_path, - tensorizer_config_dict=tensorizer_config_dict, - ) - if lora_id - else None, - ) - else: - outputs = llm.generate( - prompts, - sampling_params, - lora_request=LoRARequest(str(lora_id), lora_id, lora_path) - if lora_id - else None, - ) + outputs = llm.generate( + prompts, + sampling_params, + lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None, + ) lora_request = LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None generated_texts: list[str] = [] for output in outputs: @@ -94,31 +73,13 @@ def do_sample( return generated_texts -def generate_and_test( - llm, llama32_lora_files, tensorizer_config_dict: dict | None = None -): +def generate_and_test(llm, llama32_lora_files): print("lora adapter created") print("lora 1") - assert ( - do_sample( - llm, - llama32_lora_files, - tensorizer_config_dict=tensorizer_config_dict, - lora_id=1, - ) - == EXPECTED_LORA_OUTPUT - ) + assert do_sample(llm, llama32_lora_files, lora_id=1) == EXPECTED_LORA_OUTPUT print("lora 2") - assert ( - do_sample( - llm, - llama32_lora_files, - tensorizer_config_dict=tensorizer_config_dict, - lora_id=2, - ) - == EXPECTED_LORA_OUTPUT - ) + assert do_sample(llm, llama32_lora_files, lora_id=2) == EXPECTED_LORA_OUTPUT print("removing lora") @@ -168,79 +129,3 @@ def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files): fully_sharded_loras=True, ) generate_and_test(llm, llama32_lora_files) - - -@multi_gpu_test(num_gpus=2) -def test_tp2_serialize_and_deserialize_lora( - tmp_path, - llama32_lora_files, -): - # Run the tensorizing of the LoRA adapter and the model in a subprocess - # to guarantee cleanup - - tp_size = 2 - model_name = "model-rank-%03d.tensors" - - model_ref = MODEL_PATH - lora_path = llama32_lora_files - suffix = "test" - try: - result = subprocess.run( - [ - sys.executable, - f"{VLLM_PATH}/examples/features/tensorize_vllm_model.py", - "--model", - MODEL_PATH, - "--lora-path", - lora_path, - "--tensor-parallel-size", - str(tp_size), - "serialize", - "--serialized-directory", - str(tmp_path), - "--suffix", - suffix, - "--serialization-kwargs", - '{"limit_cpu_concurrency": 4}', - ], - check=True, - capture_output=True, - text=True, - ) - except subprocess.CalledProcessError as e: - print("Tensorizing failed.") - print("STDOUT:\n", e.stdout) - print("STDERR:\n", e.stderr) - raise - - print("STDOUT:\n", result.stdout) - - model_uri = tmp_path / "vllm" / model_ref / suffix / model_name - tensorizer_config = TensorizerConfig(tensorizer_uri=str(model_uri)) - - loaded_llm = LLM( - model=model_ref, - load_format="tensorizer", - enable_lora=True, - enforce_eager=True, - model_loader_extra_config=tensorizer_config, - max_num_seqs=7, - max_model_len=1024, - tensor_parallel_size=2, - max_loras=2, - # Leave headroom for LoRA adapter loading and Triton JIT - # compilation, which can allocate GPU memory concurrently - # during the first inference step. - gpu_memory_utilization=0.85, - ) - - tc_as_dict = tensorizer_config.to_serializable() - - print("lora adapter created") - print("lora 1") - assert ( - do_sample( - loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1 - ) - == EXPECTED_LORA_OUTPUT - ) diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 593f85993604..196445332443 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -15,7 +15,6 @@ is_base_embedding_weights, parse_fine_tuned_lora_name, ) -from vllm.model_executor.model_loader.tensorizer import TensorizerConfig from vllm.model_executor.models.utils import WeightsMapper from vllm.utils.torch_utils import PIN_MEMORY @@ -204,7 +203,6 @@ def from_local_checkpoint( dtype: torch.dtype | None = None, model_vocab_size: int | None = None, weights_mapper: WeightsMapper | None = None, - tensorizer_config_dict: dict | None = None, skip_prefixes: list[str] | None = None, moe_ep_spec: MoEEPLoadSpec | None = None, ) -> "LoRAModel": @@ -223,8 +221,6 @@ def from_local_checkpoint( embedding deltas. weights_mapper: Optional mapper rewriting checkpoint weight names to vLLM names. - tensorizer_config_dict: Optional tensorizer config used to load - the checkpoint via tensorizer instead of from disk. skip_prefixes: List of module name prefixes to skip during loading. Models can define this to skip modules not used in inference (e.g., MTP layers). Format: ["mtp."] @@ -280,24 +276,7 @@ def check_unexpected_modules(modules: dict): f" Please verify that the loaded LoRA module is correct" ) - if tensorizer_config_dict: - from tensorizer import TensorDeserializer - - tensorizer_config = TensorizerConfig(**tensorizer_config_dict) - tensorizer_dir = tensorizer_config.tensorizer_dir - if tensorizer_dir is None: - raise ValueError("tensorizer_dir must be set in tensorizer config.") - lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors") - tensorizer_args = tensorizer_config._construct_tensorizer_args() - tensors = TensorDeserializer( - lora_tensor_path, - dtype=tensorizer_config.dtype, - device=device, - **tensorizer_args.deserialization_kwargs, - ) - check_unexpected_modules(tensors) - - elif os.path.isfile(lora_tensor_path): + if os.path.isfile(lora_tensor_path): # Find unexpected modules. # Use safetensor key as a source of truth to find expected modules. # in peft if you have target_modules A, B, C and C does not exist diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index baffe1698762..d84e56bece21 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -11,7 +11,6 @@ from vllm.config.lora import LoRAConfig from vllm.logger import init_logger -from vllm.model_executor.model_loader.tensorizer import TensorizerConfig logger = init_logger(__name__) @@ -89,33 +88,10 @@ def from_local_dir( cls, lora_path: str, max_position_embeddings: int | None, - tensorizer_config_dict: dict | None = None, ) -> "PEFTHelper": lora_config_path = os.path.join(lora_path, "adapter_config.json") - - if tensorizer_config_dict: - tensorizer_config = TensorizerConfig(**tensorizer_config_dict) - tensorizer_args = tensorizer_config._construct_tensorizer_args() - from tensorizer.stream_io import open_stream - - tensorizer_dir = tensorizer_config.tensorizer_dir - if tensorizer_dir is None: - raise ValueError("tensorizer_dir must be set in tensorizer config.") - - lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json") - with open_stream( - lora_config_path, mode="rb", **tensorizer_args.stream_kwargs - ) as f: - config = json.load(f) - - logger.info( - "Successfully deserialized LoRA config from %s", - tensorizer_config.tensorizer_dir, - ) - - else: - with open(lora_config_path) as f: - config = json.load(f) + with open(lora_config_path) as f: + config = json.load(f) config["vllm_max_position_embeddings"] = max_position_embeddings return cls.from_dict(config) diff --git a/vllm/lora/request.py b/vllm/lora/request.py index cd3df54d3b12..36c6881270cc 100644 --- a/vllm/lora/request.py +++ b/vllm/lora/request.py @@ -25,7 +25,6 @@ class LoRARequest( lora_int_id: int lora_path: str = "" base_model_name: str | None = msgspec.field(default=None) - tensorizer_config_dict: dict | None = None load_inplace: bool = False is_3d_lora_weight: bool = False """Whether this adapter's MoE weights are stored in the 3D fused diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 93624523e2f8..fe266879a859 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -119,9 +119,7 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: lora_path = get_adapter_absolute_path(lora_request.lora_path) peft_helper = PEFTHelper.from_local_dir( - lora_path, - self.max_position_embeddings, - lora_request.tensorizer_config_dict, + lora_path, self.max_position_embeddings ) # Validates the LoRA configuration against requirements before @@ -147,7 +145,6 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: device="cpu", dtype=self.lora_config.lora_dtype, model_vocab_size=self.vocab_size, - tensorizer_config_dict=lora_request.tensorizer_config_dict, weights_mapper=hf_to_vllm_mapper, skip_prefixes=lora_skip_prefixes, moe_ep_spec=self._adapter_manager.moe_ep_load_spec, diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index bb48c5dc08af..3159a519fa72 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -23,7 +23,6 @@ import vllm.envs as envs from vllm.config import ModelConfig, ParallelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger -from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.platforms import current_platform from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -167,7 +166,6 @@ class TensorizerConfig(MutableMapping): s3_access_key_id: str | None = None s3_secret_access_key: str | None = None s3_endpoint: str | None = None - lora_dir: str | None = None stream_kwargs: dict[str, Any] | None = None serialization_kwargs: dict[str, Any] | None = None deserialization_kwargs: dict[str, Any] | None = None @@ -185,11 +183,8 @@ class TensorizerConfig(MutableMapping): Attributes: tensorizer_uri: Path to serialized model tensors. Can be a local file - path or a S3 URI. This is a required field unless lora_dir is - provided and the config is meant to be used for the - `tensorize_lora_adapter` function. Unless a `tensorizer_dir` or - `lora_dir` is passed to this object's initializer, this is - a required argument. + path or a S3 URI. Unless a `tensorizer_dir` is passed to this + object's initializer, this is a required argument. tensorizer_dir: Path to a directory containing serialized model tensors, and all other potential model artifacts to load the model, such as configs and tokenizer files. Can be passed instead of @@ -219,10 +214,6 @@ class TensorizerConfig(MutableMapping): be set via the S3_SECRET_ACCESS_KEY environment variable. s3_endpoint: The endpoint for the S3 bucket. Can also be set via the S3_ENDPOINT_URL environment variable. - lora_dir: Path to a directory containing LoRA adapter artifacts for - serialization or deserialization. When serializing LoRA adapters - this is the only necessary parameter to pass to this object's - initializer. """ def __post_init__(self): @@ -232,12 +223,6 @@ def __post_init__(self): and re.search(r"%0\dd", self.tensorizer_uri) is not None ) - if self.tensorizer_dir and self.lora_dir: - raise ValueError( - "Only one of tensorizer_dir or lora_dir may be specified. " - "Use lora_dir exclusively when serializing LoRA adapters, " - "and tensorizer_dir or tensorizer_uri otherwise." - ) if self.tensorizer_dir and self.tensorizer_uri: logger.warning_once( "Provided both tensorizer_dir and tensorizer_uri. " @@ -246,17 +231,12 @@ def __post_init__(self): ) self.tensorizer_dir = os.path.dirname(self.tensorizer_uri) if not self.tensorizer_uri: - if self.lora_dir: - self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors" - elif self.tensorizer_dir: + if self.tensorizer_dir: self.tensorizer_uri = f"{self.tensorizer_dir}/model.tensors" else: raise ValueError( "Unable to resolve tensorizer_uri. " - "A valid tensorizer_uri or tensorizer_dir " - "must be provided for deserialization, and a " - "valid tensorizer_uri, tensorizer_uri, or " - "lora_dir for serialization." + "A valid tensorizer_uri or tensorizer_dir must be provided." ) else: self.tensorizer_dir = os.path.dirname(self.tensorizer_uri) @@ -288,9 +268,6 @@ def to_serializable(self) -> dict[str, Any]: if "tensorizer_uri" in raw_tc_dict and "tensorizer_dir" in raw_tc_dict: blacklisted.append("tensorizer_dir") - if "tensorizer_dir" in raw_tc_dict and "lora_dir" in raw_tc_dict: - blacklisted.append("tensorizer_dir") - tc_dict = {} for k, v in raw_tc_dict.items(): if ( @@ -495,25 +472,6 @@ def _check_tensors_on_meta_device(model: nn.Module) -> None: ) -def _resize_lora_embeddings(model: nn.Module): - """Modify LoRA embedding layers to use bigger tensors - to allow for adapter added tokens.""" - for child in model.modules(): - if ( - isinstance(child, VocabParallelEmbedding) - and child.weight.shape[0] < child.num_embeddings_per_partition - ): - new_weight = torch.empty( - child.num_embeddings_per_partition, - child.embedding_dim, - dtype=child.weight.dtype, - device=child.weight.device, - ) - new_weight[: child.weight.shape[0]].copy_(child.weight.data) - new_weight[child.weight.shape[0] :].fill_(0) - child.weight.data = new_weight - - def init_tensorizer_model( tensorizer_config: TensorizerConfig, vllm_config: VllmConfig ) -> nn.Module: @@ -567,7 +525,6 @@ def deserialize_tensorizer_model( logger.info("Memory usage after: %s", after_mem) _check_tensors_on_meta_device(model) - _resize_lora_embeddings(model) del model.vllm_tensorized_marker @@ -769,60 +726,3 @@ def shutdown_engine_core() -> None: if error is not None: raise error - - -def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig): - """Uses tensorizer to serialize a LoRA adapter. Assumes that the files - needed to load a LoRA adapter are a safetensors-format file called - adapter_model.safetensors and a json config file called adapter_config.json. - - Serializes the files in the tensorizer_config.tensorizer_dir - """ - import safetensors - - from vllm.lora.utils import get_adapter_absolute_path - - lora_dir = get_adapter_absolute_path(lora_path) - - tensor_path = config_path = "" - - for file in os.listdir(lora_dir): - if file.startswith("adapter_model"): - tensor_path = lora_dir + "/" + file - if file.startswith("adapter_config"): - config_path = lora_dir + "/" + file - if tensor_path and config_path: - break - - if tensor_path.endswith(".safetensors"): - tensors = safetensors.torch.load_file(tensor_path) - elif tensor_path.endswith(".bin"): - tensors = torch.load(tensor_path, weights_only=True) - else: - raise ValueError( - f"Unsupported adapter model file: {tensor_path}. " - f"Must be a .safetensors or .bin file." - ) - - with open(config_path) as f: - config = json.load(f) - - tensorizer_args = tensorizer_config._construct_tensorizer_args() - - with open_stream( - f"{tensorizer_config.tensorizer_dir}/adapter_config.json", - mode="wb+", - **tensorizer_args.stream_kwargs, - ) as f: - f.write(json.dumps(config).encode("utf-8")) - - lora_uri = f"{tensorizer_config.tensorizer_dir}/adapter_model.tensors" - with open_stream(lora_uri, mode="wb+", **tensorizer_args.stream_kwargs) as f: - serializer = TensorSerializer(f) - serializer.write_state_dict(tensors) - serializer.close() - - logger.info( - "Successfully serialized LoRA files to %s", - str(tensorizer_config.tensorizer_dir), - ) From 30d4032363d4a1f6412d9a235049673b41309f98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 5 Oct 2026 16:44:14 +0200 Subject: [PATCH 003/952] [Model][GLM-5.3-Flash] Support DCP for the kpool sparse indexer (#59211) Signed-off-by: NickLucche --- .../test_indexer_deepseek_v4_slot_mapping.py | 43 +++++++ tests/v1/core/test_kv_cache_utils.py | 21 ++++ .../attention/dsa/dcp_indexer_cutedsl.py | 21 +++- vllm/models/glm5next/amd/sparse_indexer.py | 4 + vllm/models/glm5next/common/attention.py | 3 + vllm/models/glm5next/nvidia/sparse_indexer.py | 113 +++++++++++++----- .../backends/mla/compressor_utils.py | 28 ++++- vllm/v1/attention/backends/mla/indexer.py | 40 ++++++- vllm/v1/attention/backends/utils.py | 6 +- 9 files changed, 242 insertions(+), 37 deletions(-) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 8acfc0a33481..052088a621c9 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -315,6 +315,9 @@ def test_compressed_slot_mapping_warmup_includes_index_kpool(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=256), model_config=SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=32)), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, cp_kv_cache_interleave_size=1 + ), ) keys = CompressedSlotMappingKernel().get_warmup_keys(config) @@ -343,6 +346,41 @@ def test_compressed_slot_mapping_inherits_padded_token_slots(): assert compressed.tolist() == [-1, -1, -1, -1, -1, 3 * 4 + 2, -1, 3 * 4 + 3] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compressed_slot_mapping_dcp_uses_rank_local_state_positions(): + """Under DCP each rank writes only the states it owns, at the state's + rank-local position.""" + device = torch.device("cuda") + dcp, interleave, ratio, block_size = 2, 4, 2, 4 + num_tokens = 32 + query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + block_table = torch.tensor([[5, 7]], dtype=torch.int32, device=device) + pos = torch.arange(num_tokens) + for rank in range(dcp): + owned = (pos // interleave) % dcp == rank + slot_mapping = torch.where(owned, 0, -1).to(torch.int64).to(device) + compressed = get_compressed_slot_mapping( + num_tokens, + slot_mapping, + query_start_loc, + seq_lens, + block_table, + block_size=block_size, + compress_ratio=ratio, + dcp_world_size=dcp, + cp_interleave=interleave, + ) + local_pos = pos // (interleave * dcp) * interleave + pos % interleave + state = local_pos // ratio + expected = torch.where( + owned & ((pos + 1) % ratio == 0), + block_table[0].cpu()[state // block_size] * block_size + state % block_size, + -1, + ) + assert compressed.tolist() == expected.tolist() + + def test_index_conversion_warmup_uses_physical_block_stride(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=64), @@ -370,6 +408,8 @@ def test_zero_token_pcp_rank_participates_in_compressed_mapping_gather(monkeypat builder.use_pcp = True builder.pcp_world_size = 4 builder.pcp_rank = 0 + builder.dcp_world_size = 1 + builder.cp_kv_cache_interleave_size = 1 builder.kernel_block_size = None builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64) builder.compressed_slot_mapping_buffer = torch.zeros(8, dtype=torch.int64) @@ -383,8 +423,11 @@ def fake_get_compressed_slot_mapping( block_size, compress_ratio, out, + dcp_world_size, + cp_interleave, ): assert num_tokens == 0 + assert dcp_world_size == 1 assert block_size == 64 assert compress_ratio == 4 out.fill_(-1) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 7923e34d0fcd..08db9a286b24 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1873,6 +1873,27 @@ def test_dcp_target_allocates_replicated_draft_independently( assert [len(group) for group in cached_blocks.blocks] == [1, 4] +def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch): + """A replicated kpool tail ring is per-request state, not a draft group.""" + from vllm.v1.attention.backends.utils import resolve_kv_cache_layout + + monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False) + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.parallel_config.decode_context_parallel_size = 4 + config.cache_config.kv_cache_layout = None + tail = KpoolTailSpec( + block_size=4, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=4, + dcp_sharded=False, + ) + layout = resolve_kv_cache_layout(config, [["LBHNC"]], [new_mla_spec(), tail]) + assert layout == KVCacheLayout.LBHNC + + @pytest.mark.parametrize("use_mla", [False, True]) def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py index 34ce32856ad0..8b5d6ab1a5cb 100644 --- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py +++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py @@ -25,6 +25,22 @@ from vllm.triton_utils import tl, triton +def _dcp_merge_topk_and_interleave(vllm_config: Any) -> tuple[int, int]: + """(topk, cp_interleave) the DCP merge runs with. + + Kpool indexers (GLM-5.x) merge pool ids: topk and the interleave are both + in units of ``index_kpool`` tokens. + """ + hf_text_config = vllm_config.model_config.hf_text_config + topk = hf_text_config.index_topk + cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size + kpool = getattr(hf_text_config, "index_kpool", None) or 1 + if kpool > 1: + topk //= kpool + cp_interleave //= kpool + return topk, cp_interleave + + def stable_topk_from_gathered_candidates_cutedsl( gathered: torch.Tensor, topk: int, @@ -171,8 +187,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size if dcp_world_size <= 1: return [] - cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size - topk = vllm_config.model_config.hf_text_config.index_topk + topk, cp_interleave = _dcp_merge_topk_and_interleave(vllm_config) if topk <= 0: return [] @@ -545,7 +560,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size if dcp_world_size <= 1: return [] - topk = vllm_config.model_config.hf_text_config.index_topk + topk, _ = _dcp_merge_topk_and_interleave(vllm_config) if topk <= 0: return [] return self._trace_dispatch(self.dispatch)( diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index c357608c4763..13e67de04bf7 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -733,6 +733,10 @@ def __init__( self.topk_backend = ( cfg.kernel_config.sparse_indexer_topk_backend if cfg is not None else "auto" ) + if cfg is not None and cfg.parallel_config.decode_context_parallel_size > 1: + raise NotImplementedError( + "SparseAttnIndexerKpool on ROCm does not support DCP." + ) def forward_hip( self, diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index d7bf70af0727..a264717f0d78 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -209,6 +209,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig): head_size_v=0, dtype=torch.bfloat16, sliding_window=ring, + # Tiny per-request ring; every DCP rank keeps the full tail so the + # rank owning a pool can compress it. + dcp_sharded=False, ) def get_attn_backend(self): diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 00f51cc6a796..806b633c5836 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -7,10 +7,12 @@ import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config_or_none +from vllm.distributed import get_dcp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.indexer_topk import get_indexer_topk +from vllm.model_executor.layers.sparse_attn_indexer import _merge_dcp_topk_global from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -121,6 +123,9 @@ def sparse_attn_indexer_kpool( tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, topk_backend: str = "auto", + dcp_rank: int = 0, + dcp_world_size: int = 1, + cp_kv_cache_interleave_size: int = 1, ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -170,6 +175,12 @@ def sparse_attn_indexer_kpool( has_decode = attn_metadata_narrowed.num_decodes > 0 has_prefill = attn_metadata_narrowed.num_prefills > 0 num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + # Pools are the DCP sharding unit of the index K cache. + pool_cp_interleave = ( + cp_kv_cache_interleave_size // max(index_kpool, 1) + if dcp_world_size > 1 + else cp_kv_cache_interleave_size + ) # q_scale is required iff the FP4 cache path is enabled; the FP8 path # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. @@ -284,16 +295,18 @@ def sparse_attn_indexer_kpool( scales_spec, ) for chunk in prefill_metadata.chunks if not short_prefill else (): - k_quant = k_quant_full[: chunk.total_seq_lens] - k_scale = k_scale_full[: chunk.total_seq_lens] + # Under DCP this rank only gathers the pools it owns. + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] - if not chunk.skip_kv_gather: + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: ops.cp_gather_indexer_k_quant_cache( kv_cache, k_quant, k_scale, chunk.block_table, - chunk.cu_seq_lens, + chunk.local_cu_seq_lens, ) q_slice = q_quant[chunk.token_start : chunk.token_end] @@ -314,15 +327,18 @@ def sparse_attn_indexer_kpool( k_scale_cast = k_scale.view(torch.float32).squeeze(-1) from vllm.utils.deep_gemm import fp8_fp4_mqa_logits - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - clean_logits=False, - ) - num_rows = logits.shape[0] + num_rows = q_slice.shape[0] + if chunk.local_total_seq_lens == 0: + logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) + else: + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) # kpool: logits are pool-granular (compress_ratio == index_kpool), # so topk selects pools. We pick topk_tokens // kpool pools then @@ -338,15 +354,28 @@ def sparse_attn_indexer_kpool( chunk.token_start : chunk.token_end, :topk_tokens ] - torch.ops._C.top_k_per_row_prefill( + if logits.shape[1] == 0: + topk_dst.fill_(-1) + else: + torch.ops._C.top_k_per_row_prefill( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_dst, + num_rows, + logits.stride(0), + logits.stride(1), + select_k, + ) + # Local pool ids -> request-relative global pool ids. + _merge_dcp_topk_global( logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), select_k, + dcp_rank, + dcp_world_size, + pool_cp_interleave, + row_starts=chunk.cu_seqlen_ks, ) if index_kpool > 1: @@ -584,6 +613,14 @@ def sparse_attn_indexer_kpool( select_k, attn_metadata_narrowed.max_seq_len, ) + _merge_dcp_topk_global( + logits, + topk_dst, + select_k, + dcp_rank, + dcp_world_size, + pool_cp_interleave, + ) # Resolve to token-level indices in the output buffer. if index_kpool > 1: @@ -669,14 +706,33 @@ def __init__( else "auto" ) _parallel = _cfg.parallel_config if _cfg is not None else None - if ( - _parallel is not None - and _parallel.prefill_context_parallel_size > 1 - and _parallel.decode_context_parallel_size > 1 - ): - raise NotImplementedError( - "SparseAttnIndexerKpool does not support PCP+DCP." - ) + self._parallel_config = _parallel + self._cp_kv_cache_interleave_size: int | None = None + self.dcp_world_size = 1 + self.dcp_rank = 0 + if _parallel is not None and _parallel.decode_context_parallel_size > 1: + if _parallel.prefill_context_parallel_size > 1: + raise NotImplementedError( + "SparseAttnIndexerKpool does not support PCP+DCP." + ) + self.dcp_world_size = _parallel.decode_context_parallel_size + self.dcp_rank = get_dcp_group().rank_in_group + + @property + def cp_kv_cache_interleave_size(self) -> int: + """Resolved lazily like ``SparseAttnIndexer``: NIXL P/D can adjust it + after the model is built, and the indexer metadata builder sees the + adjusted value. + """ + if self.dcp_world_size == 1: + return 1 + if self._cp_kv_cache_interleave_size is None: + assert self._parallel_config is not None + value = self._parallel_config.cp_kv_cache_interleave_size + if isinstance(get_forward_context().attn_metadata, dict): + self._cp_kv_cache_interleave_size = value + return value + return self._cp_kv_cache_interleave_size def forward_native( self, @@ -743,4 +799,7 @@ def forward_cuda( self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, self.topk_backend, + self.dcp_rank, + self.dcp_world_size, + self.cp_kv_cache_interleave_size, ) diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 25e24f347c0e..41d973f2963b 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -37,6 +37,8 @@ class CompileKey: compress_ratio: int triton_block_size: int block_size: int + dcp_world_size: int + cp_interleave: int @staticmethod @triton.jit(do_not_specialize=["block_table_stride"]) @@ -53,6 +55,8 @@ def kernel( block_table_ptr, block_table_stride, block_size, + dcp_world_size, + cp_interleave, COMPRESS_RATIO: tl.constexpr, PAD_ID: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, @@ -77,7 +81,14 @@ def kernel( slot_mapping_ptr + query_start + offset, mask=mask, other=PAD_ID ) is_valid = ((pos + 1) % COMPRESS_RATIO == 0) & (slot != PAD_ID) - pos_after_compress = pos // COMPRESS_RATIO + # Under DCP a PAD slot also marks tokens owned by other ranks; the + # interleave is a multiple of COMPRESS_RATIO, so a state lives + # entirely on the rank owning its last token. + local_pos = ( + pos // (cp_interleave * dcp_world_size) * cp_interleave + + pos % cp_interleave + ) + pos_after_compress = local_pos // COMPRESS_RATIO block_ids = pos_after_compress // block_size block_numbers = tl.load( @@ -97,11 +108,15 @@ def dispatch( # type: ignore[override] *, compress_ratio: int, block_size: int, + dcp_world_size: int, + cp_interleave: int, ) -> CompileKey: return self.CompileKey( compress_ratio=compress_ratio, triton_block_size=self.TRITON_BLOCK_SIZE, block_size=triton_scalar_specialization_rep(block_size), + dcp_world_size=triton_scalar_specialization_rep(dcp_world_size), + cp_interleave=triton_scalar_specialization_rep(cp_interleave), ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -115,12 +130,15 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: ) if not compress_ratios: return [] + parallel_config = vllm_config.parallel_config return self._trace_dispatch(self.dispatch)( zip_inputs( *( dict( compress_ratio=ratio, block_size=vllm_config.cache_config.block_size // ratio, + dcp_world_size=parallel_config.decode_context_parallel_size, + cp_interleave=parallel_config.cp_kv_cache_interleave_size, ) for ratio in compress_ratios ) @@ -137,6 +155,8 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: block_table=int32_ptr, block_size=compile_key.block_size, compress_ratio=compile_key.compress_ratio, + dcp_world_size=compile_key.dcp_world_size, + cp_interleave=compile_key.cp_interleave, ) @kernel_launcher @@ -149,6 +169,8 @@ def __call__( block_table: torch.Tensor, block_size: int, compress_ratio: int, + dcp_world_size: int, + cp_interleave: int, ) -> LaunchSpec: return (block_table.shape[0],), dict( block_table_stride=block_table.stride(0), @@ -167,6 +189,8 @@ def get_compressed_slot_mapping( block_size: int, compress_ratio: int, out: torch.Tensor | None = None, + dcp_world_size: int = 1, + cp_interleave: int = 1, ) -> torch.Tensor: """Slot mapping for writing the compressed states of ``num_tokens`` tokens. @@ -196,6 +220,8 @@ def get_compressed_slot_mapping( block_table, block_size, compress_ratio, + dcp_world_size, + cp_interleave, ) return compressed_slot_mapping diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 5cfaf6ca5f2f..9d3823b880b9 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -554,7 +554,15 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: # Cover Triton's divisible, exact-one, and generic i32 classes. query_slice_start=(0, 1, 2), query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), - DCP_INTERLEAVE=dcp_interleave, + # Compressed indexers localize in units of compressed states. + DCP_INTERLEAVE=list( + dict.fromkeys( + dcp_interleave // ratio + for ratio in compress_ratios + if dcp_interleave % ratio == 0 + ) + ) + or [dcp_interleave], BLOCK_SIZE=self.BLOCK_SIZE, COMPRESS_RATIO=list(compress_ratios), # PCP's global cumulative lengths are the second row of one packed @@ -1010,11 +1018,31 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: (self.num_sms + 1, 2), dtype=torch.int32, device=self.device ) - if self.dcp_world_size > 1 and self.compress_ratio > 1: + # Each compressed state must live on a single rank, so the DCP + # interleave has to cover whole states. + if ( + self.dcp_world_size > 1 + and self.compress_ratio > 1 + and ( + ( + isinstance(self.kv_cache_spec, MLAAttentionSpec) + and self.kv_cache_spec.model_version == "deepseek_v4" + ) + or self.cp_kv_cache_interleave_size % self.compress_ratio != 0 + ) + ): raise NotImplementedError( - "DCP is not supported with sparse indexer KV compression " - f"(compress_ratio={self.compress_ratio})." + "DCP with sparse indexer KV compression " + f"(compress_ratio={self.compress_ratio}) requires " + "--cp-kv-cache-interleave-size to be a multiple of it; " + "DeepSeek-V4 is not supported." ) + # DCP interleave in units of compressed states. + self.compressed_cp_interleave_size = ( + self.cp_kv_cache_interleave_size // self.compress_ratio + if self.dcp_world_size > 1 and self.compress_ratio > 1 + else self.cp_kv_cache_interleave_size + ) # Pre-allocate buffers for CUDA graph compatibility when if self.compress_ratio > 1: @@ -1348,6 +1376,8 @@ def build( self.kv_cache_spec.num_states, self.compress_ratio, out=self.compressed_slot_mapping_buffer, + dcp_world_size=self.dcp_world_size, + cp_interleave=self.cp_kv_cache_interleave_size, ) if self.pcp_world_size > 1: compressed_slot_mapping = get_pcp_group().all_gather( @@ -1467,7 +1497,7 @@ def build( skip_kv_gather=query_slice.start > 0, dcp_rank=self.dcp_rank, dcp_world_size=self.dcp_world_size, - cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + cp_kv_cache_interleave_size=self.compressed_cp_interleave_size, pcp_plan=pcp_plan, ) # Skip when total_seq_lens is 0 (i.e., no compressed token). diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 529d816c0665..f93156d53fe5 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -307,8 +307,12 @@ def resolve_kv_cache_layout( f"none is in every supported set: {supported_layouts}." ) + # Self-addressed per-request rings (e.g. the kpool tail) are replicated + # state like Mamba and don't need a separate draft group. dcp_sharding = { - spec.dcp_sharded for spec in kv_cache_specs if isinstance(spec, AttentionSpec) + spec.dcp_sharded + for spec in kv_cache_specs + if isinstance(spec, AttentionSpec) and spec.uses_slot_mapping } page_sizes = {spec.page_size_bytes for spec in kv_cache_specs} if len(dcp_sharding) > 1 and len(page_sizes) > 1: From 60932a6401522039650377d5004a7680459910d1 Mon Sep 17 00:00:00 2001 From: Abhay Joshi <105213625+abhayjoshi201@users.noreply.github.com> Date: Mon, 5 Oct 2026 15:12:40 +0000 Subject: [PATCH 004/952] [Bugfix][Structured Output] Preserve literal values in Guidance disable_additional_properties (#58709) Signed-off-by: Abhay Joshi <105213625+abhayjoshi201@users.noreply.github.com> Co-authored-by: Antigravity Co-authored-by: Artem Perevedentsev --- .../test_backend_guidance.py | 85 ++++++++++++++++++- vllm/v1/structured_output/backend_guidance.py | 49 ++++++++++- 2 files changed, 129 insertions(+), 5 deletions(-) diff --git a/tests/v1/structured_output/test_backend_guidance.py b/tests/v1/structured_output/test_backend_guidance.py index 055ac9c364c4..59d1abf468ff 100644 --- a/tests/v1/structured_output/test_backend_guidance.py +++ b/tests/v1/structured_output/test_backend_guidance.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import time from concurrent.futures import Future @@ -14,7 +15,10 @@ from vllm.tokenizers import get_tokenizer from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_guidance import GuidanceBackend +from vllm.v1.structured_output.backend_guidance import ( + GuidanceBackend, + process_for_additional_properties, +) from vllm.v1.structured_output.backend_types import StructuredOutputOptions TOKENIZER = "openai-community/gpt2" @@ -236,3 +240,82 @@ def test_mistral_tokenizer_compile_grammar( grammar = backend.compile_grammar(request_type, grammar_spec) assert grammar is not None assert not grammar.is_terminated() + + +def test_process_for_additional_properties_preserves_literals() -> None: + """Literal values under const/enum/default/examples must not be rewritten + when they contain a 'properties' or 'patternProperties' key (#58695).""" + feature = { + "type": "Feature", + "geometry": None, + "properties": {"name": "HQ"}, + "patternProperties": {"^x": "y"}, + } + schema = { + "type": "object", + "properties": { + "const_field": {"const": feature}, + "enum_field": {"enum": [feature]}, + "default_field": {"type": "object", "default": feature}, + "examples_field": {"type": "object", "examples": [feature]}, + "properties": {"type": "string"}, + "const": { + "type": "object", + "properties": {"id": {"type": "integer"}}, + }, + }, + "$defs": { + "item": { + "type": "object", + "properties": {"label": {"type": "string"}}, + } + }, + } + + processed = process_for_additional_properties(schema) + + assert processed["additionalProperties"] is False + assert processed["$defs"]["item"]["additionalProperties"] is False + assert processed["properties"]["const"]["additionalProperties"] is False + assert "additionalProperties" not in processed["properties"] + assert processed["properties"]["const_field"]["const"] == feature + assert processed["properties"]["enum_field"]["enum"] == [feature] + assert processed["properties"]["default_field"]["default"] == feature + assert processed["properties"]["examples_field"]["examples"] == [feature] + + +@pytest.mark.parametrize("keyword", ["const", "enum"]) +def test_disable_additional_properties_accepts_literal_with_properties_key( + keyword: str, +) -> None: + """GuidanceBackend with disable_additional_properties=True must accept the + original literal value under const/enum and reject injected + additionalProperties (#58695).""" + feature = {"type": "Feature", "geometry": None, "properties": {"name": "HQ"}} + rewritten = {**feature, "additionalProperties": False} + schema = {keyword: feature if keyword == "const" else [feature]} + + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig( + backend="guidance", + disable_additional_properties=True, + ) + ) + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) + backend = GuidanceBackend( + vllm_config, + tokenizer=tokenizer, + vocab_size=50257, + ) + grammar = backend.compile_grammar( + StructuredOutputOptions.JSON, + json.dumps(schema), + ) + + valid_tokens = tokenizer.encode(json.dumps(feature)) + [tokenizer.eos_token_id] + rewritten_tokens = tokenizer.encode(json.dumps(rewritten)) + [ + tokenizer.eos_token_id + ] + + assert grammar.validate_tokens(valid_tokens) == valid_tokens + assert grammar.validate_tokens(rewritten_tokens) != rewritten_tokens diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 5f43c5a890fc..6dc2dc91ff76 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -34,17 +34,58 @@ logger = init_logger(__name__) +_SCHEMA_MAP_KEYWORDS = ( + "properties", + "patternProperties", + "$defs", + "definitions", + "dependentSchemas", + "dependencies", +) + +_SUBSCHEMA_KEYWORDS = ( + "additionalProperties", + "unevaluatedProperties", + "propertyNames", + "contains", + "additionalItems", + "unevaluatedItems", + "not", + "if", + "then", + "else", + "contentSchema", + "items", + "prefixItems", + "allOf", + "anyOf", + "oneOf", +) + + def _walk_json_for_additional_properties(data: object): if isinstance(data, dict): - for value in data.values(): - _walk_json_for_additional_properties(value) + for key in _SCHEMA_MAP_KEYWORDS: + value = data.get(key) + if isinstance(value, dict): + for subschema in value.values(): + if isinstance(subschema, dict): + _walk_json_for_additional_properties(subschema) + + for key in _SUBSCHEMA_KEYWORDS: + value = data.get(key) + if isinstance(value, (dict, list)): + _walk_json_for_additional_properties(value) + if "additionalProperties" not in data and ( - "properties" in data or "patternProperties" in data + isinstance(data.get("properties"), dict) + or isinstance(data.get("patternProperties"), dict) ): data["additionalProperties"] = False elif isinstance(data, list): for item in data: - _walk_json_for_additional_properties(item) + if isinstance(item, dict): + _walk_json_for_additional_properties(item) def has_guidance_unsupported_json_features(schema: dict[str, Any]) -> bool: From 9e931a0c6c374cd36bd1163c1e31f05480ac0fa7 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 5 Oct 2026 11:18:29 -0400 Subject: [PATCH 005/952] [CI] Add CODEOWNERS for HiSparse (#60060) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .github/CODEOWNERS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b483e1f8b653..ac64cf1225d4 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -4,6 +4,7 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium @wzhao18 +/vllm/distributed/kv_transfer/kv_connector/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 @@ -77,6 +78,7 @@ /vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium +/vllm/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/v1/engine @njhill /vllm/v1/executor @njhill /vllm/v1/worker @njhill From f2d8fbf4590daed7245800840a48298f4c961a9f Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Mon, 5 Oct 2026 17:32:33 +0200 Subject: [PATCH 006/952] [CI/Build] Run the prefill token scoring test under batch invariance (#60056) Signed-off-by: Stefan Koncarevic --- tests/v1/sample/test_logprobs.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 317ac10d4161..4152662fbd50 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1407,6 +1407,10 @@ def test_prompt_logprobs_with_chunking_and_preemption(): def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): """Per-row scores stay row-aligned across chunked prefill and preemption.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + # The ragged check below compares a second generate call against the first + # one. Default kernel selection varies between processes, which moves bf16 + # logprob tails by up to 0.2; under this the two runs are bit-identical. + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") prompts = [ "The following numbers of the sequence " From edde9d2ce4b1cb284679fc4cb2500a57bbc5bbe3 Mon Sep 17 00:00:00 2001 From: charlesli640 <37811657+charlesli640@users.noreply.github.com> Date: Mon, 5 Oct 2026 08:45:37 -0700 Subject: [PATCH 007/952] Remove redundant dependency requirement for TPU (#59977) Signed-off-by: Charles Li --- requirements/tpu.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements/tpu.txt b/requirements/tpu.txt index d8906d8fc1a2..01dc1f37ac98 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -9,7 +9,6 @@ wheel jinja2>=3.1.6 ray[default] ray[data] -setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 tpu-inference==0.30.0 From 528772a4bf589fd52eec40bc8ca79c84d99206b5 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:52:46 +0800 Subject: [PATCH 008/952] [Bugfix][Core] Retain encoder cache references for repeated multimodal inputs (#59942) Signed-off-by: bzsuni Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- tests/v1/core/test_scheduler.py | 27 ++++++++++++++ vllm/v1/core/sched/scheduler.py | 63 ++++++++++++++++++++------------- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 2451ea5ec006..66368a299da5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -688,6 +688,33 @@ def test_throttle_capacity_bound_guard_admits(): assert "b" in output.num_scheduled_tokens +def test_same_step_duplicate_encoder_input_stays_cached(): + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + max_num_batched_tokens=1024, + max_model_len=2048, + ) + request = create_requests( + 1, + num_tokens=2000, + req_ids=["repeated"], + mm_hashes_list=[["image", "image", "image"]], + mm_positions=[ + [PlaceholderRange(offset=offset, length=576) for offset in (0, 600, 1300)] + ], + )[0] + scheduler.add_request(request) + + output = scheduler.schedule() + assert output.scheduled_encoder_inputs == {request.request_id: [0]} + _model_output(scheduler, output, [[]]) + + # The second occurrence is partially consumed; the third is not scheduled. + cache = scheduler.encoder_cache_manager + assert cache.get_cached_input_ids(request) == {1} + assert "image" not in cache.freeable + + def test_no_mm_input_chunking(): # Disable multimodal input chunking. scheduler = create_scheduler( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b390f1aadb57..e25b5adda412 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -709,6 +709,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Schedule encoder inputs. encoder_inputs_to_schedule = None external_load_encoder_input: list[int] = [] + duplicate_encoder_inputs: list[int] = [] new_encoder_compute_budget = encoder_compute_budget if request.has_encoder_inputs: ( @@ -716,6 +717,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) = self._try_schedule_encoder_inputs( request, request.num_computed_tokens, @@ -856,17 +858,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Encoder-related. if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule - # Allocate the encoder cache. - for i in encoder_inputs_to_schedule: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget - if external_load_encoder_input: - for i in external_load_encoder_input: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) + if encoder_inputs_to_schedule or external_load_encoder_input: + self._allocate_encoder_inputs( + request, + (encoder_inputs_to_schedule or []) + external_load_encoder_input, + duplicate_encoder_inputs, + ) # Record the LoRAs in scheduled_running_reqs scheduled_loras: set[int] = set() @@ -1065,6 +1063,7 @@ def skip_request(from_queue: RequestQueue) -> None: encoder_inputs_to_schedule = None external_load_encoder_input = [] + duplicate_encoder_inputs = [] new_encoder_compute_budget = encoder_compute_budget pad_spec_decode = False @@ -1166,6 +1165,7 @@ def skip_request(from_queue: RequestQueue) -> None: num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) = self._try_schedule_encoder_inputs( request, num_computed_tokens, @@ -1335,18 +1335,14 @@ def skip_request(from_queue: RequestQueue) -> None: # Encoder-related. if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule - # Allocate the encoder cache. - for i in encoder_inputs_to_schedule: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget - # Allocate for external load encoder cache - if external_load_encoder_input: - for i in external_load_encoder_input: - self.encoder_cache_manager.allocate(request, i) - if self.ec_connector is not None: - self.ec_connector.update_state_after_alloc(request, i) + if encoder_inputs_to_schedule or external_load_encoder_input: + self._allocate_encoder_inputs( + request, + (encoder_inputs_to_schedule or []) + + external_load_encoder_input, + duplicate_encoder_inputs, + ) # re-queue requests skipped in this pass ahead of older skipped items. if step_skipped_kv_holding: @@ -1742,6 +1738,20 @@ def _reject_on_encoder_cache_embed_mismatch( self.encoder_cache_mismatch_reqs.add(request.request_id) return True + def _allocate_encoder_inputs( + self, + request: Request, + input_ids: list[int], + duplicate_input_ids: list[int], + ) -> None: + for input_id in input_ids: + self.encoder_cache_manager.allocate(request, input_id) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, input_id) + + for input_id in duplicate_input_ids: + self.encoder_cache_manager.check_and_update_cache(request, input_id) + def _try_schedule_encoder_inputs( self, request: Request, @@ -1749,7 +1759,7 @@ def _try_schedule_encoder_inputs( num_new_tokens: int, encoder_compute_budget: int, shift_computed_tokens: int = 0, - ) -> tuple[list[int], int, int, list[int]]: + ) -> tuple[list[int], int, int, list[int], list[int]]: """Determine which encoder inputs need to be scheduled in the current step, and update `num_new_tokens` and encoder token budget accordingly. @@ -1770,12 +1780,13 @@ def _try_schedule_encoder_inputs( blocks and externally cached blocks (via KVConnector). """ if num_new_tokens == 0 or not request.has_encoder_inputs: - return [], num_new_tokens, encoder_compute_budget, [] + return [], num_new_tokens, encoder_compute_budget, [], [] encoder_inputs_to_schedule: list[int] = [] mm_features = request.mm_features assert mm_features is not None assert len(mm_features) > 0 external_load_encoder_input = [] + duplicate_encoder_inputs = [] # NOTE: since scheduler operates on the request level (possibly with # multiple encoder inputs per request), we need to create temporary @@ -1832,9 +1843,10 @@ def _try_schedule_encoder_inputs( scheduled_num_encoder_embeds, num_encoder_embeds, ): - return [], 0, encoder_compute_budget, [] + return [], 0, encoder_compute_budget, [], [] # The same encoder input has already been scheduled in the # current step. + duplicate_encoder_inputs.append(i) continue cached_num_encoder_embeds = ( @@ -1849,7 +1861,7 @@ def _try_schedule_encoder_inputs( num_encoder_embeds, ) ): - return [], 0, encoder_compute_budget, [] + return [], 0, encoder_compute_budget, [], [] if self.encoder_cache_manager.check_and_update_cache(request, i): # The encoder input is already computed and cached from a @@ -1925,6 +1937,7 @@ def _try_schedule_encoder_inputs( num_new_tokens, encoder_compute_budget, external_load_encoder_input, + duplicate_encoder_inputs, ) def _make_scheduled_encoder_input_stats( From d23af12cb1d4071d9ccb4f77012285b62a319f5b Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Mon, 5 Oct 2026 19:22:29 +0300 Subject: [PATCH 009/952] [ROCm][Perf] Reach the fused QSA pre-indexer from the AMD path (#57947) Signed-off-by: mjkvaak-amd Co-authored-by: Cursor --- .buildkite/test_areas/models_basic.yaml | 2 +- tests/models/qwen4_exp/test_qsa_amd.py | 150 +++++ .../qwen4_exp/test_qsa_pre_indexer_amd.py | 344 +++++++++++ vllm/models/qwen4_exp/amd/indexer_qsa.py | 96 ++- .../qwen4_exp/amd/ops/qsa_pre_indexer.py | 558 ++++++++++++++++++ 5 files changed, 1136 insertions(+), 14 deletions(-) create mode 100644 tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py create mode 100644 vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index c2b2b054a30f..865156f59449 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -206,7 +206,7 @@ steps: depends_on: - image-build-amd commands: - - pytest -v -s models/qwen4_exp/test_qsa_amd.py + - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_pre_indexer_amd.py # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py index e8e39a7666dd..15c490223f2d 100644 --- a/tests/models/qwen4_exp/test_qsa_amd.py +++ b/tests/models/qwen4_exp/test_qsa_amd.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect import math from types import SimpleNamespace from typing import Any @@ -9,15 +10,21 @@ import torch from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.models.qwen4_exp.amd import indexer_qsa from vllm.models.qwen4_exp.amd import ( model as _qwen4_exp_model, # noqa: F401 ) from vllm.models.qwen4_exp.amd import ple_layer as ple_layer_module from vllm.models.qwen4_exp.amd.indexer_qsa import ( + QSAIndexer, apply_qsa_rmsnorm, apply_qsa_rope, ) from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops +from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( + qsa_pre_indexer, + supports_fused_pre_indexer, +) from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON @@ -120,6 +127,149 @@ def test_qsa_rmsnorm_uses_portable_implementation(default_vllm_config) -> None: torch.testing.assert_close(output, norm.forward_native(tensor)) +def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: + """A QSAIndexer whose unfused path runs for real, with no projection weights. + + Both branches have to be reachable: a fused-path test is only meaningful if + the unfused one it replaces would otherwise have executed. + """ + heads, head_dim = 2, 4 + indexer = QSAIndexer.__new__(QSAIndexer) + torch.nn.Module.__init__(indexer) + indexer.index_n_heads = heads + indexer.index_kv_heads = 1 + indexer.index_head_dim = head_dim + indexer.compress_ratio = 2 + indexer.use_fused_pre_indexer = use_fused + indexer.skip_topk = True + indexer.index_qk_proj = lambda hidden_states: ( + torch.zeros(hidden_states.shape[0], (heads + 1) * head_dim), + None, + ) + indexer.q_layernorm = GemmaRMSNorm(head_dim, eps=1e-6) + indexer.k_layernorm = GemmaRMSNorm(head_dim, eps=1e-6) + indexer.rotary_emb = SimpleNamespace( + cos_sin_cache=torch.zeros(16, head_dim), + rotary_dim=head_dim // 2, + apply_rotary_emb=lambda tensor, cos, sin: tensor, + _match_cos_sin_cache_dtype=lambda _: torch.zeros(16, head_dim), + ) + # The raw cache's full width carries the MRoPE position tail that the fused + # kernel writes; key_cache is only the leading key columns. + indexer.raw_key_cache = SimpleNamespace( + kv_cache=torch.zeros(1, 4, 1, head_dim + 12), + key_cache=torch.zeros(1, 4, 1, head_dim), + cache_rope_positions=True, + rope_position_offset=head_dim, + ) + indexer.compressed_key_cache = SimpleNamespace( + kv_cache=torch.zeros(1, 2, 1, head_dim) + ) + raw_metadata = SimpleNamespace( + num_actual_tokens=num_tokens, + slot_mapping=torch.arange(num_tokens), + block_table=torch.zeros(1, 1, dtype=torch.int32), + query_start_loc=torch.tensor([0, num_tokens], dtype=torch.int32), + logical_positions=torch.arange(num_tokens), + token_to_req=torch.zeros(num_tokens, dtype=torch.int32), + ) + compressed_metadata = SimpleNamespace( + num_actual_tokens=num_tokens, + slot_mapping=torch.arange(num_tokens), + k_work_metadata=torch.zeros(1, 2, dtype=torch.int32), + ) + indexer._metadata = lambda: (raw_metadata, compressed_metadata) + indexer.test_metadata = (raw_metadata, compressed_metadata) + return indexer + + +@pytest.mark.usefixtures("default_vllm_config") +def test_qsa_fused_pre_indexer_replaces_the_unfused_chain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + indexer = _make_indexer(use_fused=True) + calls = [] + monkeypatch.setattr( + indexer_qsa, + "qsa_pre_indexer", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + monkeypatch.setattr( + QSAIndexer, + "_update_and_compress", + lambda *args, **kwargs: pytest.fail("fused path ran the unfused chain"), + ) + rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) + + actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + + assert actual is rows + assert len(calls) == 1 + args, kwargs = calls[0] + passed = inspect.signature(qsa_pre_indexer).bind(*args, **kwargs).arguments + _, compressed_metadata = indexer.test_metadata + # The kernel addresses the raw cache at its full width and writes the MRoPE + # position tail itself, so it needs kv_cache, not the key_cache view. + assert passed["state_cache"] is indexer.raw_key_cache.kv_cache + assert passed["compressed_cache"] is indexer.compressed_key_cache.kv_cache + assert passed["k_work_metadata"] is compressed_metadata.k_work_metadata + assert passed["compress_ratio"] == indexer.compress_ratio + assert passed["rope_pos_offset"] == indexer.index_head_dim + + +@pytest.mark.usefixtures("default_vllm_config") +def test_qsa_unsupported_config_keeps_the_unfused_chain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + indexer = _make_indexer(use_fused=False) + monkeypatch.setattr( + indexer_qsa, + "qsa_pre_indexer", + lambda *args, **kwargs: pytest.fail("unfused path ran the fused kernel"), + ) + updates = [] + monkeypatch.setattr( + QSAIndexer, + "_update_and_compress", + lambda *args, **kwargs: updates.append(args), + ) + rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) + + actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + + assert actual is rows + assert len(updates) == 1 + + +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + pytest.param({}, True, id="qwen3_8_flash_next"), + pytest.param({"is_neox_style": False}, False, id="not-neox"), + pytest.param({"rotary_dim": 128}, False, id="rotary-dim"), + pytest.param({"mrope_interleaved": False}, False, id="mrope-not-interleaved"), + pytest.param({"mrope_section": (11, 11, 10, 0)}, False, id="mrope-section-len"), + ], +) +def test_fused_pre_indexer_gate_is_config_shaped(overrides, expected) -> None: + # Qwen3.8-Flash-Next's indexer geometry, which both vendors satisfy. + rotary = SimpleNamespace( + rotary_dim=64, + is_neox_style=True, + mrope_section=(11, 11, 10), + mrope_interleaved=True, + ) + for name, value in overrides.items(): + setattr(rotary, name, value) + + assert ( + supports_fused_pre_indexer( + rotary, head_dim=128, num_kv_heads=1, compress_ratio=4 + ) + is expected + ) + + def test_qsa_selection_uses_portable_topk_on_rocm( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py new file mode 100644 index 000000000000..8a39504e0067 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the AMD fused QSA pre-indexer.""" + +import pytest +import torch + +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.models.qwen4_exp.amd.indexer_qsa import apply_qsa_rope +from vllm.models.qwen4_exp.amd.ops.qsa import ( + qsa_compress_groups_with_ratio, + qsa_store_cache_rows, +) +from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( + qsa_pre_indexer, +) +from vllm.models.qwen4_exp.common.qsa_cache import ( + canonical_qsa_rope_positions, + circular_qsa_slot_mapping, + compressed_qsa_slot_mapping, +) +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +requires_qsa_kernels = pytest.mark.skipif( + not current_platform.is_rocm() or not HAS_TRITON, + reason="the AMD QSA kernels require ROCm and Triton", +) + + +def _reference_gemma_rmsnorm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """The unfused Gemma RMSNorm the AMD reference path actually runs.""" + norm = GemmaRMSNorm(weight.shape[-1], eps=eps).to( + device=weight.device, dtype=weight.dtype + ) + with torch.no_grad(): + norm.weight.copy_(weight) + return norm.forward_native(x) + + +HQ, D = 4, 128 +CR = 4 +MROPE_SECTION = (11, 11, 10) +EPS = 1e-6 +BLOCK_SIZE = 16 +COMP_PAGE = BLOCK_SIZE // CR +ROPE_POS_OFFSET = D +RTOL = 1.6e-2 +ATOL = 1e-2 +MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8]) + + +def _make_block_table(block_counts): + num_blocks = sum(block_counts) + table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32) + physical_blocks = torch.randperm(num_blocks) + offset = 0 + for request, count in enumerate(block_counts): + table[request, :count] = physical_blocks[offset : offset + count] + offset += count + return table, num_blocks + + +def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None: + # e4m3 is sign-magnitude, so within a sign the uint8 code order matches the + # value order and one ulp is one code step. The two paths' intermediates + # differ in the pooling accumulation order, which at denormal magnitudes + # (absolute grid step 2^-9) shows up as up to 2 code steps. + code_diff = ( + actual.view(torch.uint8).int() - expected.view(torch.uint8).int() + ).abs() + abs_diff = (actual.float() - expected.float()).abs() + assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all()) + + +@requires_qsa_kernels +@pytest.mark.usefixtures("default_vllm_config") +@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens", + [ + pytest.param(True, True, True, 4, *MIXED_BATCH, id="mrope"), + pytest.param(False, False, False, 4, *MIXED_BATCH, id="text"), + pytest.param(True, False, True, 8, *MIXED_BATCH, id="mrope-model-1d"), + pytest.param(True, True, False, 4, *MIXED_BATCH, id="mrope-no-position-cache"), + pytest.param( + True, + False, + False, + 8, + *MIXED_BATCH, + id="mrope-model-1d-no-position-cache", + ), + pytest.param( + False, False, True, 4, *MIXED_BATCH, id="text-with-position-cache" + ), + pytest.param(True, True, True, 4, [37], [37], [0], id="fresh"), + pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"), + ], +) +def test_qsa_fused_pre_indexer_matches_unfused( + indexer_dtype, + mrope, + is_2d_positions, + cache_rope_positions, + state_size, + seq_lens, + query_lens, + history_lens, +) -> None: + from vllm.model_executor.layers.rotary_embedding import get_rope + + if indexer_dtype == torch.float8_e4m3fn: + pytest.skip("the AMD QSA indexer cache is bf16-only") + + # The FP8 comparison below is a one-ulp bound on data this test generates, + # so an unseeded run decides for itself whether it reproduces a rounding + # difference. Seed it, or the kernel's bit-exactness is only sampled. + torch.manual_seed(0) + + device = "cuda" + rope_params = { + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default", + } + if mrope: + rope_params["mrope_interleaved"] = True + rope_params["mrope_section"] = list(MROPE_SECTION) + with torch.device(device): + rope = get_rope( + head_size=256, + max_position=32768, + rope_parameters=rope_params, + dtype=torch.bfloat16, + ) + + token_to_req = torch.cat( + [ + torch.full((length,), request, dtype=torch.int32) + for request, length in enumerate(query_lens) + ] + ).to(device) + logical_positions = torch.cat( + [ + torch.arange(seq_len - query_len, seq_len, dtype=torch.int64) + for seq_len, query_len in zip(seq_lens, query_lens) + ] + ).to(device) + num_tokens = logical_positions.numel() + query_start_loc = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32 + ).to(device) + positions = ( + torch.stack( + [ + logical_positions, + logical_positions // 7 + 3, + logical_positions // 13 + 11, + ] + ) + if is_2d_positions + else logical_positions + ) + position_rows = ( + canonical_qsa_rope_positions(positions) + if cache_rope_positions + else logical_positions.view(-1, 1, 1).expand(-1, 1, 3) + ) + + compressed_block_counts = [ + (seq_len // CR + COMP_PAGE - 1) // COMP_PAGE for seq_len in seq_lens + ] + raw_block_table, num_raw_blocks = _make_block_table([1] * len(seq_lens)) + compressed_block_table, num_compressed_blocks = _make_block_table( + compressed_block_counts + ) + raw_block_table = raw_block_table.to(device) + raw_slots = circular_qsa_slot_mapping( + raw_block_table, + token_to_req, + logical_positions, + state_size, + query_start_loc, + ) + compressed_slots = compressed_qsa_slot_mapping( + compressed_block_table.to(device), + token_to_req, + logical_positions, + COMP_PAGE, + CR, + ) + group_counts = torch.tensor( + [ + seq_len // CR - (seq_len - query_len) // CR + for seq_len, query_len in zip(seq_lens, query_lens) + ], + dtype=torch.int32, + device=device, + ) + k_work_counts = torch.maximum(group_counts, torch.ones_like(group_counts)) + k_start_loc = torch.cat([k_work_counts.new_zeros(1), k_work_counts.cumsum(0)]) + work_requests = torch.repeat_interleave( + torch.arange(len(query_lens), dtype=torch.int32, device=device), + k_work_counts, + ) + local_work = torch.arange( + int(k_start_loc[-1]), dtype=torch.int32, device=device + ) - torch.repeat_interleave(k_start_loc[:-1], k_work_counts) + max_k_work = (num_tokens + (CR - 1) * len(query_lens)) // CR + k_work_metadata = torch.full((max_k_work, 2), -1, dtype=torch.int32, device=device) + k_work_metadata[: work_requests.numel()] = torch.stack( + (work_requests, local_work), dim=1 + ) + + raw_width = D + 12 if cache_rope_positions else D + # Match vLLM's padded-page cache layout: rows are contiguous, while physical + # blocks have a larger stride than their logical contents. + raw_page_elements = state_size * raw_width + fused_raw_storage = torch.zeros( + num_raw_blocks, + raw_page_elements + 16, + dtype=torch.bfloat16, + device=device, + ) + fused_raw = torch.as_strided( + fused_raw_storage, + (num_raw_blocks, state_size, 1, raw_width), + (raw_page_elements + 16, raw_width, raw_width, 1), + ) + for request, history_len in enumerate(history_lens): + history_end = seq_lens[request] - query_lens[request] + for position in range(history_end - history_len, history_end): + block = int(raw_block_table[request, 0]) + row = fused_raw[block, position % state_size, 0] + row[:D] = torch.randn(D, dtype=torch.bfloat16, device=device) + if cache_rope_positions: + row[ROPE_POS_OFFSET:].view(torch.int64).copy_( + torch.tensor( + [position, position // 7 + 3, position // 13 + 11], + dtype=torch.int64, + device=device, + ) + ) + unfused_raw = fused_raw.clone() + compressed_page_elements = COMP_PAGE * D + fused_compressed_storage = torch.zeros( + num_compressed_blocks, + compressed_page_elements + 16, + dtype=indexer_dtype, + device=device, + ) + fused_compressed = torch.as_strided( + fused_compressed_storage, + (num_compressed_blocks, COMP_PAGE, 1, D), + (compressed_page_elements + 16, D, D, 1), + ) + unfused_compressed = fused_compressed.clone() + + projected_qk = torch.randn( + num_tokens, (HQ + 1) * D, dtype=torch.bfloat16, device=device + ) + q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 + k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 + + fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) + qsa_pre_indexer( + projected_qk[:, : HQ * D], + projected_qk[:, HQ * D :], + positions, + rope.cos_sin_cache, + q_weight, + k_weight, + EPS, + fused_query, + fused_raw, + raw_slots, + raw_block_table, + query_start_loc, + logical_positions, + fused_compressed, + compressed_slots, + k_work_metadata, + compress_ratio=CR, + mrope_section=MROPE_SECTION if mrope else None, + rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None, + ) + + unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D) + unfused_query = _reference_gemma_rmsnorm( + unfused_query.reshape(-1, D), q_weight, EPS + ).reshape_as(unfused_query) + unfused_query = apply_qsa_rope(rope, positions, unfused_query) + + raw_keys = unfused_raw[..., :D] + rope_positions = ( + unfused_raw[..., ROPE_POS_OFFSET:].view(torch.int64) + if cache_rope_positions + else None + ) + pooled, first_positions = qsa_compress_groups_with_ratio( + projected_qk[:, HQ * D :].reshape(-1, 1, D), + position_rows, + raw_keys, + raw_block_table, + token_to_req, + query_start_loc, + logical_positions, + compressed_slots, + CR, + rope_positions, + ) + compressed_rows = _reference_gemma_rmsnorm( + pooled.reshape(-1, D), k_weight, EPS + ).reshape(-1, 1, D) + group_positions = ( + first_positions.transpose(0, 1) if mrope else first_positions[:, 0] + ) + compressed_rows = apply_qsa_rope(rope, group_positions, compressed_rows) + qsa_store_cache_rows(unfused_compressed, compressed_slots, compressed_rows) + qsa_store_cache_rows(raw_keys, raw_slots, projected_qk[:, HQ * D :]) + if rope_positions is not None: + qsa_store_cache_rows(rope_positions, raw_slots, position_rows) + + if indexer_dtype == torch.float8_e4m3fn: + # Both paths round the same ~bf16 intermediates to e4m3. Bitwise + # equality (the dsv4 indexer test's bar) does not hold here: the 1D + # reference rope (forward_cuda) differs from _norm_rope by up to 1 + # bf16 ulp, and even the MRoPE pair flips a code occasionally. + unfused_query = unfused_query.to(indexer_dtype) + assert_fp8_within_one_ulp(fused_query, unfused_query) + else: + torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL) + assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16)) + if indexer_dtype == torch.float8_e4m3fn: + assert_fp8_within_one_ulp(fused_compressed, unfused_compressed) + else: + torch.testing.assert_close( + fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL + ) diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 4345763fd41c..162914f39e08 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -23,6 +23,7 @@ QSAKeyStateCache, canonical_qsa_rope_positions, ) +from .ops.qsa_pre_indexer import qsa_pre_indexer, supports_fused_pre_indexer def apply_qsa_rope( @@ -98,6 +99,12 @@ def __init__( self.token_topk = int(config.indexer_budget) self.compress_ratio = int(config.indexer_compress_ratio) self.rotary_emb = rotary_emb + self.use_fused_pre_indexer = supports_fused_pre_indexer( + rotary_emb, + self.index_head_dim, + self.index_kv_heads, + self.compress_ratio, + ) self.prefix = prefix # MTP step 0 selects the target-aligned rows; later steps reuse them # while continuing to update the QSA side cache. @@ -143,20 +150,27 @@ def __init__( def output_width(self) -> int: return self.token_topk + self.compress_ratio - 1 - def project_qk( + def project( self, hidden_states: torch.Tensor, - positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Project replicated Q/K, normalize+rotate Q, and preserve raw K.""" + """Project replicated Q/K and split, leaving both operands unrotated.""" qk, _ = self.index_qk_proj(hidden_states) - q_raw, token_k = qk.split( + return qk.split( ( self.index_n_heads * self.index_head_dim, self.index_kv_heads * self.index_head_dim, ), dim=-1, ) + + def project_qk( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Project replicated Q/K, normalize+rotate Q, and preserve raw K.""" + q_raw, token_k = self.project(hidden_states) q = q_raw.reshape(-1, self.index_n_heads, self.index_head_dim) q = apply_qsa_rmsnorm( self.q_layernorm, @@ -251,6 +265,52 @@ def _update_and_compress( position_rows, ) + def _fused_pre_indexer( + self, + projected_q: torch.Tensor, + raw_keys: torch.Tensor, + positions: torch.Tensor, + raw_metadata: QSAForwardMetadata, + compressed_metadata: QSAForwardMetadata, + ) -> torch.Tensor: + """Normalize and rotate Q, compress K, and write both caches in one launch. + + Replaces ``project_qk``'s norm/RoPE half and the whole of + ``_update_and_compress``; the two must stay numerically interchangeable. + """ + raw_key_cache = self.raw_key_cache + q = projected_q.new_empty( + raw_metadata.num_actual_tokens, + self.index_n_heads, + self.index_head_dim, + ) + qsa_pre_indexer( + projected_q, + raw_keys, + positions, + self.rotary_emb.cos_sin_cache, + self.q_layernorm.weight, + self.k_layernorm.weight, + self.q_layernorm.variance_epsilon, + q, + raw_key_cache.kv_cache, + raw_metadata.slot_mapping, + raw_metadata.block_table, + raw_metadata.query_start_loc, + raw_metadata.logical_positions, + self.compressed_key_cache.kv_cache, + compressed_metadata.slot_mapping, + compressed_metadata.k_work_metadata, + compress_ratio=self.compress_ratio, + mrope_section=getattr(self.rotary_emb, "mrope_section", None), + rope_pos_offset=( + raw_key_cache.rope_position_offset + if raw_key_cache.cache_rope_positions + else None + ), + ) + return q + def _select( self, q: torch.Tensor, @@ -295,15 +355,25 @@ def forward( return result raw_metadata, compressed_metadata = metadata num_tokens = raw_metadata.num_actual_tokens - q, token_k = self.project_qk( - hidden_states[:num_tokens], positions[..., :num_tokens] - ) - self._update_and_compress( - token_k, - positions[..., :num_tokens], - raw_metadata, - compressed_metadata, - ) + hidden_states = hidden_states[:num_tokens] + positions = positions[..., :num_tokens] + if self.use_fused_pre_indexer: + projected_q, raw_keys = self.project(hidden_states) + q = self._fused_pre_indexer( + projected_q, + raw_keys, + positions, + raw_metadata, + compressed_metadata, + ) + else: + q, token_k = self.project_qk(hidden_states, positions) + self._update_and_compress( + token_k, + positions, + raw_metadata, + compressed_metadata, + ) if self.skip_topk: if out is None: raise RuntimeError("QSA top-k reuse requires an output buffer") diff --git a/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py new file mode 100644 index 000000000000..6a4873882f2c --- /dev/null +++ b/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py @@ -0,0 +1,558 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused QSA pre-indexer kernel for the AMD Qwen4Exp path. + +Started as a copy of the NVIDIA kernel (``nvidia/ops/qsa_prepare.py``) and is +kept separate so either vendor can change its own copy without re-validating +the other. The state source select is two masked loads rather than a pointer +select, which ROCm Triton rejects. +""" + +import torch +from torch import nn + +from vllm.triton_utils import tl, triton + + +def supports_fused_pre_indexer( + rotary_emb: nn.Module, + head_dim: int, + num_kv_heads: int, + compress_ratio: int, +) -> bool: + """Report whether this indexer's shapes match the fused kernel's assumptions. + + The kernel hard-codes the rotary layout and the single-KV-head group + compression it was written for; everything it rejects has a working unfused + path. + """ + rotary_dim = int(rotary_emb.rotary_dim) + mrope_section = getattr(rotary_emb, "mrope_section", None) + return ( + bool(getattr(rotary_emb, "is_neox_style", False)) + and ( + not mrope_section + or ( + len(mrope_section) == 3 + and sum(mrope_section) == rotary_dim // 2 + and bool(getattr(rotary_emb, "mrope_interleaved", False)) + ) + ) + and head_dim == 128 + and rotary_dim == 64 + and num_kv_heads == 1 + and compress_ratio > 1 + and compress_ratio & (compress_ratio - 1) == 0 + ) + + +@triton.jit +def _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + cos_sin_stride, + norm_weight_ptr, + eps, + IS_MROPE: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows.""" + TILE_T: tl.constexpr = x.shape[0] + TILE_H: tl.constexpr = x.shape[1] + D: tl.constexpr = x.shape[2] + ROWS: tl.constexpr = TILE_T * TILE_H + HALF: tl.constexpr = D // 2 + QUARTER: tl.constexpr = D // 4 + pairs = tl.arange(0, QUARTER) + if IS_MROPE: + # Qwen interleaves temporal, height, and width rotary pairs. Each axis + # still indexes the same position-major cos/sin table. + h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_H) + w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_W) + t_mask = ~(h_mask | w_mask) + base = cos_sin_ptr + pairs[None, :] + pos_rows = (pos_t, pos_h, pos_w) + axis_masks = (t_mask, h_mask, w_mask) + cos = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + sin = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + for axis in tl.static_range(3): + cos += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride, + mask=axis_masks[axis][None, :], + other=0, + ) + sin += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride + QUARTER, + mask=axis_masks[axis][None, :], + other=0, + ) + else: + cos = tl.load(cos_sin_ptr + pos_t[:, None] * cos_sin_stride + pairs[None, :]) + sin = tl.load( + cos_sin_ptr + pos_t[:, None] * cos_sin_stride + QUARTER + pairs[None, :] + ) + + cos = tl.reshape( + tl.broadcast_to(cos[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + sin = tl.reshape( + tl.broadcast_to(sin[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + x = tl.reshape(x, (ROWS, D)).to(tl.float32) + weight = tl.load(norm_weight_ptr + tl.arange(0, D)).to(tl.float32) + 1.0 + rrms = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = (x * rrms[:, None] * weight[None, :]).to(cos.dtype) + rotated, passthrough = tl.split( + tl.permute(tl.reshape(y, (ROWS, 2, HALF)), (0, 2, 1)) + ) + r0, r1 = tl.split(tl.permute(tl.reshape(rotated, (ROWS, 2, QUARTER)), (0, 2, 1))) + out0 = r0 * cos - r1 * sin + out1 = r1 * cos + r0 * sin + rotated = tl.reshape(tl.permute(tl.join(out0, out1), (0, 2, 1)), (ROWS, HALF)) + result = tl.reshape(tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), (ROWS, D)) + return tl.reshape(result, (TILE_T, TILE_H, D)) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_state_blocks", + "num_compressed_blocks", + "num_k_work", + ] +) +def _qsa_pre_indexer_kernel( + q_ptr, + q_stride_token, + k_ptr, + k_stride_token, + pos_ptr, + pos_stride_axis, + pos_stride_token, + cos_sin_ptr, + q_norm_weight_ptr, + k_norm_weight_ptr, + eps, + q_out_ptr, + q_out_stride_token, + q_out_stride_head, + state_cache_ptr, + state_cache_stride_block, + state_cache_stride_token, + state_slots_ptr, + state_table_ptr, + state_table_stride_req, + query_start_loc_ptr, + logical_positions_ptr, + compressed_slots_ptr, + k_work_metadata_ptr, + compressed_cache_ptr, + compressed_cache_stride_block, + compressed_cache_stride_token, + num_tokens, + num_state_blocks, + num_compressed_blocks, + num_k_work, + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, + TILE_H_Q: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + STATE_SIZE: tl.constexpr, + COMP_PAGE_SIZE: tl.constexpr, + IS_2D_POSITIONS: tl.constexpr, + IS_K_MROPE: tl.constexpr, + CACHE_HAS_ROPE_POS: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + pid = tl.program_id(0) + # K work occupies the first programs; the remaining programs tile Q. This + # keeps both paths in one launch while leaving their register shapes + # independent. + if pid >= num_k_work: + q_pid = pid - num_k_work + num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q) + token_tile = q_pid // num_head_tiles + head_tile = q_pid % num_head_tiles + tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q) + heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q) + valid_tokens = tokens < num_tokens + valid_heads = heads < HQ + dims = tl.arange(0, D) + mask = valid_tokens[:, None, None] & valid_heads[None, :, None] + x = tl.load( + q_ptr + + tokens[:, None, None] * q_stride_token + + heads[None, :, None] * D + + dims[None, None, :], + mask=mask, + other=0.0, + ) + pos_t = tl.load(pos_ptr + tokens * pos_stride_token, mask=valid_tokens, other=0) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + pos_w = tl.load( + pos_ptr + 2 * pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + y = _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + D // 2, + q_norm_weight_ptr, + eps, + IS_2D_POSITIONS, + MROPE_H, + MROPE_W, + ) + tl.store( + q_out_ptr + + tokens[:, None, None] * q_out_stride_token + + heads[None, :, None] * q_out_stride_head + + dims[None, None, :], + y, + mask=mask, + ) + + if pid < num_k_work: + # One work item owns one completed compression group. Work item zero + # additionally commits the request's current raw-K suffix below. + work_metadata = tl.load(k_work_metadata_ptr + pid * 2 + tl.arange(0, 2)) + request, work_in_request = tl.split(work_metadata) + if request < 0: + return + + query_start = tl.load(query_start_loc_ptr + request) + query_end = tl.load(query_start_loc_ptr + request + 1) + query_len = query_end - query_start + chunk_end = tl.load(logical_positions_ptr + query_end - 1) + chunk_start = chunk_end - query_len + 1 + num_groups = (chunk_end + 1) // COMPRESS_RATIO - chunk_start // COMPRESS_RATIO + dims = tl.arange(0, D) + + if work_in_request < num_groups: + first_boundary = ( + (chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO + ) * COMPRESS_RATIO - 1 + end_position = first_boundary + work_in_request * COMPRESS_RATIO + boundary_token = query_start + end_position - chunk_start + valid_token = ( + (boundary_token >= query_start) + & (boundary_token < query_end) + & (boundary_token < num_tokens) + ) + compressed_slot = tl.load( + compressed_slots_ptr + boundary_token, + mask=valid_token, + other=-1, + ) + valid = ( + valid_token + & (compressed_slot >= 0) + & (compressed_slot < num_compressed_blocks * COMP_PAGE_SIZE) + ) + state_block = tl.load(state_table_ptr + request * state_table_stride_req) + state_block_valid = (state_block >= 0) & (state_block < num_state_blocks) + safe_state_block = tl.maximum(state_block, 0).to(tl.int64) + group_offsets = tl.arange(0, COMPRESS_RATIO) + source_positions = end_position - (COMPRESS_RATIO - 1) + group_offsets + source_in_chunk = source_positions >= chunk_start + source_in_cache = source_positions < chunk_start + source_tokens = query_start + source_positions - chunk_start + source_tokens_valid = ( + (source_tokens >= query_start) + & (source_tokens < query_end) + & (source_tokens < num_tokens) + ) + current_base = ( + k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token + ) + cached_base = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token + ) + # Only the first completed group can cross the chunk boundary, so + # every row comes from exactly one of the two sources. Selecting + # between the base pointers would need a pointer-typed tl.where, + # which ROCm's Triton rejects ("'tt.addptr' op failed to verify + # that result type matches ptr type"), so select the loaded values + # instead. The two masks are complementary, so no lane fetches from + # both and memory traffic is unchanged. + # + # Both bases still need the alignment stated explicitly. Without it + # Triton vectorizes these loads differently, which reorders the + # pooling reduction below; float addition is not associative, so the + # compressed rows stop matching the unfused path bit for bit and an + # FP8 cache rounds a near-zero pooled value to the wrong code. + current_base = tl.multiple_of(current_base, (8, 8)) + cached_base = tl.multiple_of(cached_base, (8, 8)) + from_chunk = valid & source_in_chunk[:, None] & source_tokens_valid[:, None] + from_cache = valid & source_in_cache[:, None] & state_block_valid + current = tl.load(current_base + dims[None, :], mask=from_chunk, other=0.0) + cached = tl.load(cached_base + dims[None, :], mask=from_cache, other=0.0) + source = tl.where(source_in_chunk[:, None], current, cached).to(tl.float32) + # Match the unfused path's BF16 pooled tensor before RMSNorm. + pooled = ( + (tl.sum(source, axis=0) / COMPRESS_RATIO).to(tl.bfloat16).to(tl.float32) + ) + + first_position = end_position - (COMPRESS_RATIO - 1) + if CACHE_HAS_ROPE_POS: + # RoPE uses the first token in the pooled group. Its exact MRoPE + # coordinates may live in this chunk or the raw-state ring. + first_in_chunk = first_position >= chunk_start + first_token = query_start + first_position - chunk_start + first_token_valid = ( + (first_token >= query_start) + & (first_token < query_end) + & (first_token < num_tokens) + ) + safe_first_token = tl.maximum(first_token, 0) + load_current_position = first_in_chunk & first_token_valid + first_pos_t = tl.load( + pos_ptr + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + if IS_2D_POSITIONS: + first_pos_h = tl.load( + pos_ptr + pos_stride_axis + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + first_pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + else: + first_pos_h = first_pos_t + first_pos_w = first_pos_t + tail = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (first_position % STATE_SIZE) * state_cache_stride_token + + D + ).to(tl.pointer_type(tl.int64)) + load_cached_position = ~first_in_chunk & state_block_valid + cached_pos_t = tl.load(tail, mask=load_cached_position, other=0) + cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0) + cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0) + pos_t = tl.where(first_in_chunk, first_pos_t.to(tl.int64), cached_pos_t) + pos_h = tl.where(first_in_chunk, first_pos_h.to(tl.int64), cached_pos_h) + pos_w = tl.where(first_in_chunk, first_pos_w.to(tl.int64), cached_pos_w) + else: + pos_t = first_position + pos_h = first_position + pos_w = first_position + y = _norm_rope( + tl.reshape(pooled, (1, 1, D)), + pos_t + tl.arange(0, 1), + pos_h + tl.arange(0, 1), + pos_w + tl.arange(0, 1), + cos_sin_ptr, + D // 2, + k_norm_weight_ptr, + eps, + IS_K_MROPE, + MROPE_H, + MROPE_W, + ) + compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64) + compressed_row = compressed_slot % COMP_PAGE_SIZE + tl.store( + compressed_cache_ptr + + compressed_block * compressed_cache_stride_block + + compressed_row * compressed_cache_stride_token + + dims, + tl.reshape(y, (D,)), + mask=valid, + ) + + if work_in_request == 0: + # This CTA may have just read historical rows from the circular buffer. + # Keep every lane past those loads before overwriting the same ring. + tl.debug_barrier() + # One CTA per request commits only the suffix retained by the ring. + num_state_rows = tl.minimum(query_len, STATE_SIZE) + for state_offset in tl.range(0, num_state_rows): + token = query_end - num_state_rows + state_offset + valid_token = ( + (token >= query_start) & (token < query_end) & (token < num_tokens) + ) + slot = tl.load(state_slots_ptr + token, mask=valid_token, other=-1) + valid_slot = ( + valid_token & (slot >= 0) & (slot < num_state_blocks * STATE_SIZE) + ) + safe_slot = tl.maximum(slot, 0) + state_row = ( + state_cache_ptr + + (safe_slot // STATE_SIZE).to(tl.int64) * state_cache_stride_block + + (safe_slot % STATE_SIZE) * state_cache_stride_token + ) + k = tl.load( + k_ptr + tl.maximum(token, 0) * k_stride_token + dims, + mask=valid_slot, + other=0.0, + ) + tl.store(state_row + dims, k, mask=valid_slot) + if CACHE_HAS_ROPE_POS: + pos_t = tl.load( + pos_ptr + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + + pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + tail = (state_row + D).to(tl.pointer_type(tl.int64)) + tl.store(tail, pos_t.to(tl.int64), mask=valid_slot) + tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot) + tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) + + +def qsa_pre_indexer( + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + eps: float, + q_out: torch.Tensor, + state_cache: torch.Tensor, + state_slots: torch.Tensor, + state_block_table: torch.Tensor, + query_start_loc: torch.Tensor, + logical_positions: torch.Tensor, + compressed_cache: torch.Tensor, + compressed_slots: torch.Tensor, + k_work_metadata: torch.Tensor, + *, + compress_ratio: int, + mrope_section: tuple[int, int, int] | None, + rope_pos_offset: int | None, +) -> None: + """Normalize Q, compress K, then update the circular raw state.""" + num_tokens = q.shape[0] + if num_tokens == 0: + return + num_q_heads, head_dim = q_out.shape[1:] + assert cos_sin_cache.shape[-1] * 2 == head_dim + assert q.shape == (num_tokens, num_q_heads * head_dim) + assert k.shape == (num_tokens, head_dim) + assert q.stride(-1) == 1 + assert k.stride(-1) == 1 + assert q_out.stride(-1) == 1 + assert cos_sin_cache.is_contiguous() + assert state_cache.stride(-1) == 1 + assert compressed_cache.stride(-1) == 1 + assert k_work_metadata.ndim == 2 and k_work_metadata.shape[1] == 2 + is_2d_positions = positions.ndim == 2 + is_k_mrope = bool(mrope_section) + cache_has_rope_pos = rope_pos_offset is not None + assert rope_pos_offset is None or rope_pos_offset == head_dim + if is_2d_positions: + assert positions.shape == (3, num_tokens) + assert is_k_mrope + pos_stride_axis, pos_stride_token = positions.stride() + else: + assert positions.shape == (num_tokens,) + pos_stride_axis, pos_stride_token = 0, positions.stride(0) + section = mrope_section if mrope_section is not None else (0, 0, 0) + assert len(section) == 3 + + # Swept on MI355X (64-lane wavefronts) over TILE_T_Q 1-16, TILE_H_Q 1-4 and + # num_warps 1/2/4, from 4 to 16384 tokens: no config beat these by more + # than noise, and 2 or 4 warps were slower everywhere. + if num_tokens <= 4096: + TILE_T_Q, TILE_H_Q = 2, 2 + else: + TILE_T_Q, TILE_H_Q = 2, 4 + num_k_work = k_work_metadata.shape[0] + num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q) + _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)]( + q, + q.stride(0), + k, + k.stride(0), + positions, + pos_stride_axis, + pos_stride_token, + cos_sin_cache, + q_norm_weight, + k_norm_weight, + eps, + q_out, + q_out.stride(0), + q_out.stride(1), + state_cache, + state_cache.stride(0), + state_cache.stride(1), + state_slots, + state_block_table, + state_block_table.stride(0), + query_start_loc, + logical_positions, + compressed_slots, + k_work_metadata, + compressed_cache, + compressed_cache.stride(0), + compressed_cache.stride(1), + num_tokens, + state_cache.shape[0], + compressed_cache.shape[0], + num_k_work, + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, + TILE_H_Q=TILE_H_Q, + COMPRESS_RATIO=compress_ratio, + STATE_SIZE=state_cache.shape[1], + COMP_PAGE_SIZE=compressed_cache.shape[1], + IS_2D_POSITIONS=is_2d_positions, + IS_K_MROPE=is_k_mrope, + CACHE_HAS_ROPE_POS=cache_has_rope_pos, + MROPE_H=section[1], + MROPE_W=section[2], + num_warps=1, + ) + + +__all__ = ["qsa_pre_indexer", "supports_fused_pre_indexer"] From 87954c5b042391db2fde6bbccd50760e7b7eb8f0 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Mon, 5 Oct 2026 18:25:42 +0200 Subject: [PATCH 010/952] [CI/Build] Declare the video modality on the CohereCompass reference model (#60063) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas --- .../generation/test_transformers_video.py | 4 ++++ .../multimodal/generation/vlm_utils/model_utils.py | 13 +++++++++++++ 2 files changed, 17 insertions(+) diff --git a/tests/models/multimodal/generation/test_transformers_video.py b/tests/models/multimodal/generation/test_transformers_video.py index bbb31b7be8e5..c6c47f9ac0e1 100644 --- a/tests/models/multimodal/generation/test_transformers_video.py +++ b/tests/models/multimodal/generation/test_transformers_video.py @@ -36,6 +36,7 @@ "<|VISION_START|><|VIDEO_PAD|><|VISION_END|>Describe this video." "<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>" ), + "patch_hf_runner": model_utils.cohere_compass_patch_hf_runner, }, "Qwen/Qwen2.5-VL-3B-Instruct": { "prompt": ( @@ -88,6 +89,9 @@ def test_transformers_video_generation( model_id, dtype="bfloat16", auto_cls=AutoModelForImageTextToText ) as hf_model: hf_model = model_utils.qwen3_vl_patch_hf_runner(hf_model) + patch_hf_runner = VIDEO_MODEL_SETTINGS[model_id].get("patch_hf_runner") + if patch_hf_runner is not None: + hf_model = patch_hf_runner(hf_model) hf_outputs = hf_model.generate_greedy_logprobs_limit( [prompt], 128, num_logprobs=10, videos=[video] ) diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py index bc472c400f8c..9acbe2c0d6b3 100644 --- a/tests/models/multimodal/generation/vlm_utils/model_utils.py +++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py @@ -1111,6 +1111,19 @@ def processor(*args, videos=None, **kwargs): return hf_model +def cohere_compass_patch_hf_runner(hf_model: HfRunner) -> HfRunner: + """Declare the video modality that CohereCompass omits upstream. + + Without it ``generate`` leaves ``pixel_values_videos`` unconsumed and the + forward pass rejects it, see huggingface/transformers#49320. Drop this once + the pinned transformers declares the modality itself. + """ + model = hf_model.model + if "video" not in model.input_modalities: + model.input_modalities = (*model.input_modalities, "video") + return hf_model + + def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner": """Patch HfRunner for Voxtral's conversation-based processor. From 64cb68384241f046d07ba3797e1b14dbc82a0b2d Mon Sep 17 00:00:00 2001 From: Sohaib Ahmed <113179991+Sohaib-Ahmed21@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:28:09 +0500 Subject: [PATCH 011/952] [Model] Add Nemotron 3.5 ASR transcription support (#59827) Signed-off-by: Sohaib-Ahmed21 --- docs/models/supported_models.md | 1 + docs/usage/v1_guide.md | 1 + vllm/model_executor/models/registry.py | 1 + 3 files changed, 3 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a55559b0d6d2..221ab877eaa8 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -199,6 +199,7 @@ Some model architectures are supported via vLLM plugins. These plugins extend vL | ------------ | ------ | ----------------- | | `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) | | `Florence2ForConditionalGeneration` | Florence-2 | [bart-plugin](https://github.com/vllm-project/bart-plugin) | +| `Nemotron3_5AsrForRNNT` | NVIDIA Nemotron 3.5 ASR | [vllm-nemotron-asr-plugin](https://github.com/Sohaib-Ahmed21/vllm-nemotron-asr-plugin) | For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system. diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 4ef2b0b5b6a7..d16db468955f 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -138,6 +138,7 @@ Whisper is supported natively. Other encoder-decoder models are supported via th - **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin). - **Florence-2**: `Florence2ForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin). +- **Nemotron 3.5 ASR**: `Nemotron3_5AsrForRNNT` is supported via [vllm-nemotron-asr-plugin](https://github.com/Sohaib-Ahmed21/vllm-nemotron-asr-plugin). For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md). diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3ff37a256290..48b831bbdf0a 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -849,6 +849,7 @@ "BartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", "Florence2ForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", "MBartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", + "Nemotron3_5AsrForRNNT": "https://github.com/Sohaib-Ahmed21/vllm-nemotron-asr-plugin", } From 54d93af9fb89d0bb8392bdb72db42bb6a17cb4ca Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 5 Oct 2026 12:42:27 -0400 Subject: [PATCH 012/952] [Bugfix][HiSparse] Reject cudagraph_mode=FULL at startup (#59688) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- tests/test_config.py | 11 +++++++++++ vllm/config/vllm.py | 6 ++++++ 2 files changed, 17 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index a2e75b10725d..b1348480cc83 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -502,6 +502,17 @@ def test_hisparse_rejects_pipeline_parallelism(monkeypatch): ) +def test_hisparse_rejects_full_cudagraph_mode(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "support_static_graph_mode", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + with pytest.raises(ValueError, match="does not support cudagraph_mode=FULL"): + VllmConfig( + attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL), + ) + + def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch): monkeypatch.setattr(current_platform, "is_cuda", lambda: True) monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 32b49f42725e..a5e1ad08b4fc 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1836,6 +1836,12 @@ def has_blocked_weights(): raise ValueError( "HiSparse does not support decode context parallelism." ) + if self.compilation_config.cudagraph_mode == CUDAGraphMode.FULL: + raise ValueError( + "HiSparse does not support cudagraph_mode=FULL; use " + "FULL_AND_PIECEWISE (the default), which captures FULL graphs " + "for decode batches." + ) if not self.scheduler_config.scheduler_reserve_full_isl: # Without it, async loads admitted against free host blocks can # each wait on host pages the others hold, and waiting requests From 182247257fb030fe0b41d342ae9e2a4486ae787e Mon Sep 17 00:00:00 2001 From: akii96 Date: Mon, 5 Oct 2026 20:28:07 +0300 Subject: [PATCH 013/952] [ROCm][Perf] Use a stacked hipBLASLt GEMM for the bf16x3 router (#52668) Signed-off-by: Aakif Nawaz --- .buildkite/test_areas/kernels.yaml | 2 + tests/kernels/test_bf16x3_router_gemm_rocm.py | 201 ++++++++++++++++++ .../kernels/test_gate_linear_rocm_dispatch.py | 90 ++++++++ .../router/bf16x3_router_gemm_rocm.py | 163 ++++++++++++++ .../layers/fused_moe/router/gate_linear.py | 142 ++++++++++++- 5 files changed, 594 insertions(+), 4 deletions(-) create mode 100644 tests/kernels/test_bf16x3_router_gemm_rocm.py create mode 100644 vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 515c3966f590..54e7215967cd 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -255,6 +255,7 @@ steps: --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py --ignore=kernels/test_top_k_per_row.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py + --ignore=kernels/test_bf16x3_router_gemm_rocm.py --ignore=kernels/test_ll_bf16_gemm.py --ignore=kernels/test_shuffle_rows.py # BROKEN on main, pending kernel fixes (B200): @@ -280,6 +281,7 @@ steps: kernels/test_apply_repetition_penalties.py kernels/test_awq_int4_to_int8.py kernels/test_bf16_skinny_gemm.py + kernels/test_bf16x3_router_gemm_rocm.py kernels/test_cache_kernels.py kernels/test_compressor_kv_cache.py kernels/test_cp_gather_fp8.py diff --git a/tests/kernels/test_bf16x3_router_gemm_rocm.py b/tests/kernels/test_bf16x3_router_gemm_rocm.py new file mode 100644 index 000000000000..69aa195b55f3 --- /dev/null +++ b/tests/kernels/test_bf16x3_router_gemm_rocm.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the ROCm BF16x3 router GEMM. + +The weight split and the eligibility policy are pure functions, so most of +this runs device-free. Only the GEMM correctness tests need a gfx950 runner. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm as rocm_bf16x3, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") + +# (hidden_size, num_experts) for the models that ship fp32 router weights. +# (6144, 128) -> MiniMax-M3, (3072, 256) -> MiniMax-M2/M2.5 +ROUTER_SHAPES = [(6144, 128), (3072, 256)] +ROUTER_WEIGHT_SCALE = 0.053 + + +def _requires_gfx950(): + if not rocm_bf16x3.platform_supported(): + pytest.skip("ROCm bf16x3 router GEMM requires gfx950") + + +def _inputs(num_tokens, hidden_size, num_experts, seed=42, device="cuda"): + torch.manual_seed(seed) + x = torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) + w = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device) + return x, w * ROUTER_WEIGHT_SCALE + + +def _top_set(logits, k=4): + """Sorted top-k indices; topk's order among near-ties is not meaningful.""" + return logits.topk(k, dim=-1).indices.sort(dim=-1).values + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +def test_split_reconstructs_exactly(hidden_size: int, num_experts: int): + torch.manual_seed(42) + w = torch.randn(num_experts, hidden_size) * ROUTER_WEIGHT_SCALE + split = rocm_bf16x3.split_bf16x3(w) + + assert split.shape == (rocm_bf16x3.BF16X3_TERMS, num_experts, hidden_size) + assert split.dtype == torch.bfloat16 + # router_gemm reinterprets this as (TERMS * E, K) without a copy. + assert split.is_contiguous() + assert torch.equal(split.float().sum(0), w) + + +def test_split_is_exact_across_normal_exponents(): + exps = torch.arange(-100, 100, dtype=torch.float32) + mantissas = torch.linspace(1.0, 1.999, 64, dtype=torch.float32) + w = (mantissas[:, None] * torch.pow(2.0, exps)[None, :]).contiguous() + w = torch.cat([w, -w], dim=0) + assert torch.equal(rocm_bf16x3.split_bf16x3(w).float().sum(0), w) + + +def test_split_rejects_non_finite(): + """ValueError rather than assert: the caller downgrades to the fp32 path + instead of failing model load, and this must survive ``python -O``.""" + for bad in (float("inf"), float("-inf"), float("nan")): + with pytest.raises(ValueError): + rocm_bf16x3.split_bf16x3(torch.full((2, 4), bad, dtype=torch.float32)) + + +def test_split_rejects_non_fp32(): + with pytest.raises(ValueError): + rocm_bf16x3.split_bf16x3(torch.zeros(2, 4, dtype=torch.bfloat16)) + + +def test_split_does_not_mutate_input(): + w = torch.randn(8, 32, dtype=torch.float32) + original = w.clone() + rocm_bf16x3.split_bf16x3(w) + assert torch.equal(w, original) + + +def test_is_supported_rejects_ineligible_inputs(): + meta = torch.device("meta") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device=meta) + w = torch.zeros(128, 6144, dtype=torch.float32, device=meta) + assert rocm_bf16x3.is_supported(x, w) + + assert not rocm_bf16x3.is_supported(x[: rocm_bf16x3.MIN_TOKENS - 1], w) + assert not rocm_bf16x3.is_supported(x.float(), w) + assert not rocm_bf16x3.is_supported(x, w.bfloat16()) + assert not rocm_bf16x3.is_supported(x, w[:, :128]) + assert not rocm_bf16x3.is_supported(x.T.contiguous().T, w) + + +def test_is_supported_rejects_mismatched_devices(): + """The weight can be offloaded while the activations are not.""" + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta") + w = torch.zeros(128, 6144, dtype=torch.float32) + assert not rocm_bf16x3.is_supported(x, w) + + +def test_is_supported_accepts_odd_expert_counts(): + """Unlike a tl.arange-indexed kernel, torch.mm needs no power-of-two E.""" + meta = torch.device("meta") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device=meta) + for num_experts in (17, 384, 512): + w = torch.zeros(num_experts, 6144, dtype=torch.float32, device=meta) + assert rocm_bf16x3.is_supported(x, w) + + +# Measured on gfx950: this path lands at 3.4e-7 to 4.9e-7 and varies <0.2% +# across token counts. Fixed bounds rather than a live comparison, because the +# fp32 fallback swings 5.7e-7 to 1.4e-6 as torch picks different hipBLASLt +# kernels at different M. _FP32_FALLBACK_BEST is the best the fallback was ever +# observed to do, so beating it proves the claim without racing a moving target. +_MAX_REL_L2 = 2e-6 +_FP32_FALLBACK_BEST = 5.7e-7 + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +@pytest.mark.parametrize("num_tokens", [2048, 8192, 16385]) +def test_matches_fp64_reference(hidden_size: int, num_experts: int, num_tokens: int): + _requires_gfx950() + x, w = _inputs(num_tokens, hidden_size, num_experts) + out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w)) + ref = x.double() @ w.double().T + + assert out.shape == (num_tokens, num_experts) + assert out.dtype == torch.float32 + assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2 + + +@pytest.mark.parametrize("num_experts", [17, 384]) +def test_matches_fp64_for_non_power_of_two_experts(num_experts: int): + """Exercises the reduction's N mask, which the model shapes never hit.""" + _requires_gfx950() + x, w = _inputs(8192, 6144, num_experts) + out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w)) + ref = x.double() @ w.double().T + + assert out.shape == (8192, num_experts) + assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2 + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +def test_more_accurate_than_fp32_fallback(hidden_size: int, num_experts: int): + """The claim is not just accuracy but beating the fp32 GEMM it replaces.""" + _requires_gfx950() + x, w = _inputs(16384, hidden_size, num_experts) + ref = x.double() @ w.double().T + out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w)) + assert (out.double() - ref).norm() / ref.norm() < _FP32_FALLBACK_BEST + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +def test_expert_selection_matches_fp64(hidden_size: int, num_experts: int): + """Relative L2 is a proxy; the expert set is what reaches the model. + + Measures zero misroutes on both shapes, where the fp32 fallback misroutes a + nonzero number at some token counts. + """ + _requires_gfx950() + num_tokens = 16384 + x, w = _inputs(num_tokens, hidden_size, num_experts) + ref_set = _top_set(x.double() @ w.double().T) + out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w)) + + misrouted = int((_top_set(out) != ref_set).any(-1).sum()) + assert misrouted / num_tokens < 1e-4 + + +def test_writes_into_supplied_out(): + _requires_gfx950() + x, w = _inputs(8192, 6144, 128) + split = rocm_bf16x3.split_bf16x3(w) + out = torch.empty(8192, 128, dtype=torch.float32, device="cuda") + + returned = rocm_bf16x3.bf16x3_router_gemm(x, split, out=out) + assert returned is out + assert torch.equal(out, rocm_bf16x3.bf16x3_router_gemm(x, split)) + + +def test_validates_layout_assumptions(): + """The reduction assumes a contiguous (TERMS, E, K) split and a contiguous + out; nothing in the kernel enforces either.""" + _requires_gfx950() + x, w = _inputs(8192, 6144, 128) + split = rocm_bf16x3.split_bf16x3(w) + + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split.transpose(1, 2)) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x.float(), split) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split[:2]) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x[:, :512], split) + wide = torch.empty(8192, 256, dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split, out=wide[:, :128]) diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index afa2e57f6aec..06d69fc98569 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -11,6 +11,8 @@ term; a biased gate must fall back so the bias is not silently dropped. """ +import inspect + import pytest import torch @@ -168,3 +170,91 @@ def test_rocm_fp32_router_gemm_rejects_unsupported_configs( on_gfx950=on_gfx950, ) assert not gate.allow_fp32_router_gemm + + +# --------------------------------------------------------------------------- +# ROCm BF16x3 tier (fp32 router weights) +# --------------------------------------------------------------------------- +def _make_bf16x3_gate(monkeypatch, *, supported=True, **kwargs): + """Build a gate with the ROCm bf16x3 platform check mocked.""" + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm as rocm_bf16x3, + ) + + monkeypatch.setattr(rocm_bf16x3, "platform_supported", lambda: supported) + kwargs.setdefault("params_dtype", torch.float32) + return _make_gate(monkeypatch, is_rocm=True, **kwargs) + + +def test_rocm_bf16x3_enabled_for_fp32_weight(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch) + assert gate.allow_rocm_bf16x3_router_gemm + # Built by process_weights_after_loading, not lazily in forward(). + # Registered up front but still None, so named_buffers() skips it. + assert gate._bf16x3_weight is None + assert "_bf16x3_weight" in gate._buffers + + +def test_rocm_bf16x3_disabled_by_bias(monkeypatch): + # torch.mm has no bias term. + gate = _make_bf16x3_gate(monkeypatch, bias=True) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_for_bf16_weight(monkeypatch): + # A bf16 weight belongs to the cuBLAS tier; there is nothing to split. + gate = _make_bf16x3_gate(monkeypatch, params_dtype=torch.bfloat16) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_for_non_fp32_out_dtype(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_on_unsupported_arch(monkeypatch): + # Only gfx950 is benchmarked; gfx942 has a different fp32:bf16 MFMA ratio. + gate = _make_bf16x3_gate(monkeypatch, supported=False) + assert not gate.allow_rocm_bf16x3_router_gemm + assert not gate._rocm_bf16x3_weight_eligible + + +def test_rocm_bf16x3_no_split_built_when_tier_disabled(monkeypatch): + """An ineligible out_dtype must not allocate a split that is never read; + it is 1.5x the router weight on every layer.""" + gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16) + gate.quant_method.process_weights_after_loading(gate) + assert gate._bf16x3_weight is None + + +def test_rocm_bf16x3_set_out_dtype_enables_tier(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch, out_dtype=None) + assert not gate.allow_rocm_bf16x3_router_gemm + gate.set_out_dtype(torch.float32) + assert gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_dispatch_is_a_custom_op(): + """The num_tokens branch must stay opaque to Dynamo; see the op's impl.""" + assert hasattr(torch.ops.vllm, "rocm_bf16x3_router_gemm_dispatch") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta") + w = torch.zeros(128, 6144, dtype=torch.float32, device="meta") + split = torch.zeros(3, 128, 6144, dtype=torch.bfloat16, device="meta") + out = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch(x, w, split) + assert out.shape == (8192, 128) + assert out.dtype == torch.float32 + + +def test_rocm_bf16x3_is_reachable_for_low_m_shapes(monkeypatch): + """The gfx950 low-M fp32 tier claims these shapes too and returns + unconditionally, so bf16x3 is only reachable if checked first.""" + gate = _make_bf16x3_gate( + monkeypatch, input_size=6144, output_size=128, on_gfx950=True + ) + assert gate.allow_fp32_router_gemm + assert gate.allow_rocm_bf16x3_router_gemm + + source = inspect.getsource(type(gate).forward) + assert source.index("allow_rocm_bf16x3_router_gemm") < source.index( + "allow_fp32_router_gemm" + ) diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py new file mode 100644 index 000000000000..8d1fa84c8622 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm BF16x3 MoE router GEMM for fp32 gate weights. + +Gates whose weights ship in fp32 (MiniMax-M2/M3, HunYuan-V3) otherwise land in +the ``F.linear`` fallback, which casts the bf16 activations up to fp32 and runs +the GEMM on the fp32 MFMA pipe. Splitting the weight into three bf16 terms +whose sum reproduces it exactly moves the GEMM onto the bf16 pipe and is more +accurate than the path it replaces, which rounds the activations on the way to +fp32 while this one leaves them untouched. + +The terms are stacked along the expert dimension into one ``(TERMS * E, K)`` +weight, so the router costs a single hipBLASLt GEMM plus a reduction over the +``TERMS`` column blocks of its output. +""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +__all__ = [ + "BF16X3_TERMS", + "MIN_TOKENS", + "bf16x3_router_gemm", + "is_supported", + "platform_supported", + "split_bf16x3", +] + +BF16X3_TERMS = 3 + +# Below this the fp32 fallback is faster: the GEMM plus reduction has a ~32us +# floor from launch and allocation that small batches cannot amortize. +# Measured on gfx950; the crossover is 1536 at (K=6144, E=128) and (K=3072, +# E=256), and this leaves margin for untested shapes. +MIN_TOKENS = 2048 + +_MAX_INT32 = (1 << 31) - 1 + + +def split_bf16x3(weight: torch.Tensor) -> torch.Tensor: + """Split an fp32 weight into bf16 terms whose sum reproduces it exactly. + + Returns a contiguous ``(BF16X3_TERMS, E, K)`` tensor. Three + round-to-nearest bf16 terms carry 24 mantissa bits, matching fp32's 24, so + the reconstruction is bit-exact for finite normals down to ``2**-108``. + + Raises ValueError on non-finite input: ``inf`` rounds to bf16 ``inf`` and + ``inf - inf`` is NaN, which would poison every logit rather than just the + affected experts. + """ + if weight.dtype != torch.float32: + raise ValueError(f"expected fp32 weight, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError("router weight contains inf/NaN") + terms = [] + resid = weight.detach().clone() + for _ in range(BF16X3_TERMS): + hi = resid.to(torch.bfloat16) + terms.append(hi) + resid -= hi.to(torch.float32) + return torch.stack(terms).contiguous() + + +def platform_supported() -> bool: + """Whether this platform has a benchmarked bf16x3 router path. + + gfx942 has a different fp32:bf16 MFMA ratio and is not yet evaluated. + """ + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + +def is_supported(x: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether this path can serve ``x @ weight.T`` -> fp32 profitably.""" + if x.dim() != 2 or weight.dim() != 2: + return False + n, k = weight.shape + m = x.shape[0] + return ( + m >= MIN_TOKENS + and x.dtype == torch.bfloat16 + and weight.dtype == torch.float32 + and x.shape[1] == k + # The weight may be offloaded while the activations are not; the + # fallback this guards would then fail on a device mismatch. + and x.device == weight.device + # torch.mm accepts any stride, but a non-K-contiguous activation costs + # a hidden copy that gives back the win over the fallback. + and x.stride(1) == 1 + and 0 < m * BF16X3_TERMS * n <= _MAX_INT32 + ) + + +@triton.jit +def _reduce_terms_kernel( + y_ptr, + out_ptr, + M, + stride_ym, + stride_om, + N: tl.constexpr, + TERMS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + mask = (offs_m < M)[:, None] & (offs_n < N)[None, :] + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for t in tl.static_range(TERMS): + acc += tl.load( + y_ptr + offs_m[:, None] * stride_ym + t * N + offs_n[None, :], + mask=mask, + other=0.0, + ) + tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :], acc, mask=mask) + + +def bf16x3_router_gemm( + x: torch.Tensor, + weight_split: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Router logits in fp32 from bf16 activations and a bf16x3 split weight.""" + m, k = x.shape + terms, n, k_w = weight_split.shape + assert terms == BF16X3_TERMS, f"expected {BF16X3_TERMS} terms, got {terms}" + assert k == k_w, f"K mismatch: x has {k}, weight_split has {k_w}" + assert weight_split.dtype == torch.bfloat16, weight_split.dtype + assert weight_split.is_contiguous(), "weight_split must be contiguous" + assert x.dtype == torch.bfloat16, x.dtype + if out is None: + out = torch.empty((m, n), dtype=torch.float32, device=x.device) + else: + assert out.shape == (m, n) and out.is_contiguous(), "bad out layout" + + # Rows [t*E, (t+1)*E) of the (TERMS * E, K) view hold term t, so the GEMM + # emits one column block per term and the reduction sums them. + y = torch.empty((m, terms * n), dtype=torch.float32, device=x.device) + torch.mm(x, weight_split.view(terms * n, k).T, out_dtype=torch.float32, out=y) + + # 64 rows/program keeps the accumulator small enough to stay in registers + # at the largest supported expert count. + block_m = 64 + _reduce_terms_kernel[(triton.cdiv(m, block_m),)]( + y, + out, + m, + y.stride(0), + out.stride(0), + N=n, + TERMS=terms, + BLOCK_M=block_m, + BLOCK_N=triton.next_power_of_2(n), + num_warps=4, + ) + return out diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index f91e9c449ef5..8c4961f269b8 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -4,6 +4,7 @@ from torch.nn.parameter import Parameter import vllm._custom_ops as ops +from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ( ReplicatedLinear, @@ -13,6 +14,8 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op +logger = init_logger(__name__) + @PluggableLayer.register("gate_linear") class GateLinear(ReplicatedLinear): @@ -23,8 +26,10 @@ class GateLinear(ReplicatedLinear): 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, M<=32, model-specific shapes) 3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) - 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 5. F.linear via ReplicatedLinear (ultimate fallback) + 4. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, + M>=2048) + 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 6. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -140,6 +145,33 @@ def __init__( and self.out_dtype == torch.float32 ) + # ROCm bf16x3 router GEMM eligibility. out_dtype may still be None here + # and is folded in by set_out_dtype; the split itself is only built + # once the tier is known to be on (see _GateLinearMethod). + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm, + ) + + self._rocm_bf16x3_weight_eligible = ( + self._router_gemm_no_bias + and self.weight.dtype == torch.float32 + and bf16x3_router_gemm_rocm.platform_supported() + ) + self.allow_rocm_bf16x3_router_gemm = ( + self._rocm_bf16x3_weight_eligible and self.out_dtype == torch.float32 + ) + # Filled in by process_weights_after_loading. Not persistent: it is + # derived from `weight` rather than loaded from a checkpoint. + self.register_buffer("_bf16x3_weight", None, persistent=False) + if self._rocm_bf16x3_weight_eligible: + # Safe to swap after create_weights() because + # UnquantizedLinearMethod holds no per-layer state. Asserted + # because the swap would silently drop quantization if a + # quant_config were ever plumbed into this layer. + assert isinstance(self.quant_method, UnquantizedLinearMethod) + self.quant_method = _GateLinearMethod() + logger.info_once("Enabled ROCm BF16x3 router GEMM.") + # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc: # 1. PDL support. Both dot-product and split-K kernels. # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction. @@ -173,6 +205,10 @@ def set_out_dtype(self, out_dtype: torch.dtype) -> None: self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16 # out_dtype may start as None -> recompute eligibility here + self.allow_rocm_bf16x3_router_gemm = ( + self._rocm_bf16x3_weight_eligible and out_dtype == torch.float32 + ) + if self.allow_specialized_router_gemm: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( is_available, @@ -205,6 +241,23 @@ def forward( output = ll_bf16_gemm(x, self.weight) return self._return(output) + # Tier 4: ROCm bf16x3 router GEMM. Checked before tier 2 because both + # tiers claim the same gfx950 shapes and tier 2 returns + # unconditionally; the custom op below picks between them on the + # runtime num_tokens. + if self.allow_rocm_bf16x3_router_gemm and x.dtype == torch.bfloat16: + if self._bf16x3_weight is None: + logger.warning_once( + "ROCm BF16x3 router GEMM is enabled for %s but its weight " + "split is missing; falling back to fp32.", + self.prefix, + ) + else: + output = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch( + x, self.weight, self._bf16x3_weight + ) + return self._return(output) + # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32) # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph # capture does not freeze the runtime num_tokens branch. @@ -226,12 +279,12 @@ def forward( output = bf16x3_router_gemm(x, self.weight) return self._return(output) - # Tier 4: cuBLAS bf16→fp32 + # Tier 5: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return self._return(output) - # Tier 5: F.linear (ReplicatedLinear) + # Tier 6: F.linear (ReplicatedLinear) if ( self.out_dtype is not None and self.is_unquantized @@ -249,9 +302,90 @@ def forward( return output if not self.return_bias else (output, output_bias) +class _GateLinearMethod(UnquantizedLinearMethod): + """UnquantizedLinearMethod plus the ROCm bf16x3 weight split. + + Building the split here rather than lazily in forward() charges it to the + weights memory pool ahead of KV cache profiling, and rebuilds it on weight + reload. The latter matters because an in-place ``param.data.copy_()`` + leaves both ``data_ptr()`` and ``_version`` untouched, so a stale split + cannot be detected from inside forward(). + + ``allow_rocm_bf16x3_router_gemm`` is final by the time this runs: the only + caller of ``set_out_dtype`` does so during model construction. + """ + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # The base CPU branch may delete layer.weight, so it must run first; + # nothing breaks today only because platform_supported() is False on + # CPU. + super().process_weights_after_loading(layer) + if not layer.allow_rocm_bf16x3_router_gemm: + return + try: + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm, + ) + + layer._bf16x3_weight = bf16x3_router_gemm_rocm.split_bf16x3(layer.weight) + except ValueError as exc: + layer.allow_rocm_bf16x3_router_gemm = False + logger.warning_once( + "Disabling ROCm BF16x3 router GEMM for %s: %s", layer.prefix, exc + ) + + _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS +def rocm_bf16x3_router_gemm_dispatch_impl( + x: torch.Tensor, + weight: torch.Tensor, + weight_split: torch.Tensor, +) -> torch.Tensor: + """Run the ROCm bf16x3 router GEMM, or fall back for small batches. + + This must be a custom op because our torch.compile integration does not + support runtime dispatching on num_tokens: models carrying fp32 router + weights (MiniMax-M2, HunYuan-V3) call the gate from + ``@support_torch_compile`` model code and vLLM drops all Dynamo guards, so + a plain Python branch on ``x.shape[0]`` would be frozen at first trace. + """ + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm, + ) + + if bf16x3_router_gemm_rocm.is_supported(x, weight): + return bf16x3_router_gemm_rocm.bf16x3_router_gemm(x, weight_split) + # Below MIN_TOKENS, hand back to the low-M gfx950 kernel this tier is + # checked ahead of. + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: + from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 + can_use_rocm_fp32_router_gemm, + rocm_fp32_router_gemm, + ) + + x = x.contiguous() + if can_use_rocm_fp32_router_gemm(x, weight): + return rocm_fp32_router_gemm(x, weight) + return torch.nn.functional.linear(x.float(), weight) + + +def rocm_bf16x3_router_gemm_dispatch_fake( + x: torch.Tensor, + weight: torch.Tensor, + weight_split: torch.Tensor, +) -> torch.Tensor: + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32) + + +direct_register_custom_op( + op_name="rocm_bf16x3_router_gemm_dispatch", + op_func=rocm_bf16x3_router_gemm_dispatch_impl, + fake_impl=rocm_bf16x3_router_gemm_dispatch_fake, +) + + def fp32_router_gemm_dispatch_impl( x: torch.Tensor, weight: torch.Tensor, From 877ddcc45639d2363b9081917ef489619e3f29de Mon Sep 17 00:00:00 2001 From: Evgeny Savinov Date: Mon, 5 Oct 2026 19:35:51 +0200 Subject: [PATCH 014/952] [Bugfix][Spec Decode] Reserve the bonus KV slot for fill-in DSpark (#59105) Signed-off-by: Evgeny Savinov Co-authored-by: Codex Co-authored-by: Tomas Ruiz --- tests/v1/worker/test_gpu_warmup_blocks.py | 38 ++++++++++++++--------- vllm/config/vllm.py | 12 ++++--- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 04179ae760b0..6fadcbb35eea 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -314,24 +314,28 @@ def _step(num_new_tokens: int) -> None: @pytest.mark.parametrize( - ("method", "expected"), + ("method", "draft_hf_config", "expected"), [ - ("eagle", NUM_SPEC_STEPS), - ("eagle3", NUM_SPEC_STEPS), - ("mtp", NUM_SPEC_STEPS), - ("dspark", NUM_SPEC_STEPS), - ("draft_model", NUM_SPEC_STEPS), + ("eagle", None, NUM_SPEC_STEPS), + ("eagle3", None, NUM_SPEC_STEPS), + ("mtp", None, NUM_SPEC_STEPS), + ("dspark", None, NUM_SPEC_STEPS), + ("dspark", {"sample_from_anchor": True}, NUM_SPEC_STEPS), + ("dspark", {"sample_from_anchor": False}, NUM_SPEC_STEPS + 1), + ("draft_model", None, NUM_SPEC_STEPS), # DFlash's in-fill decoding adds a query for the last sampled token. - ("dflash", NUM_SPEC_STEPS + 1), - ("ngram", 0), - ("ngram_gpu", 0), - ("medusa", 0), - ("mlp_speculator", 0), - ("suffix", 0), - ("extract_hidden_states", 0), + ("dflash", None, NUM_SPEC_STEPS + 1), + ("ngram", None, 0), + ("ngram_gpu", None, 0), + ("medusa", None, 0), + ("mlp_speculator", None, 0), + ("suffix", None, 0), + ("extract_hidden_states", None, 0), ], ) -def test_num_lookahead_tokens_per_method(method: str, expected: int): +def test_num_lookahead_tokens_per_method( + method: str, draft_hf_config: dict | None, expected: int +): """`VllmConfig.num_lookahead_tokens` is the single source of the reservation. Both the scheduler and the warmup read it, so a wrong answer here silently @@ -350,6 +354,12 @@ class _Config: speculative_config = object.__new__(SpeculativeConfig) object.__setattr__(speculative_config, "method", method) object.__setattr__(speculative_config, "num_speculative_tokens", NUM_SPEC_STEPS) + hf_config = SimpleNamespace(**(draft_hf_config or {})) + object.__setattr__( + speculative_config, + "draft_model_config", + SimpleNamespace(hf_config=hf_config), + ) config = _Config() config.speculative_config = speculative_config diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a5e1ad08b4fc..53d4102697c5 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -640,13 +640,15 @@ def num_lookahead_tokens(self) -> int: speculative_config = self.speculative_config if speculative_config is None: return 0 - if speculative_config.use_dflash(): - # DFlash requires an extra lookahead slot since it uses in-fill-style - # decoding instead of standard next-token sampling, so it has a query - # for the last sampled token plus queries for each draft token. + dspark_fill_in = speculative_config.use_dspark() and not getattr( + speculative_config.draft_model_config.hf_config, "sample_from_anchor", True + ) + if speculative_config.use_dflash() or dspark_fill_in: + # Fill-in drafting uses a bonus query plus one query per draft token. + # DSpark's anchor-sampling layout does not need the extra slot. return self.num_speculative_tokens + 1 if speculative_config.use_eagle() or speculative_config.uses_draft_model(): - # DSpark (covered by use_eagle) drafts a block of num_speculative_tokens + # Anchor-sampling DSpark drafts a block of num_speculative_tokens # query tokens in which the anchor itself is the first prediction # position (no separate bonus query), so it needs exactly # num_speculative_tokens lookahead slots. From d09ff7763747bbe7cd153edd8017479804f47be7 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 6 Oct 2026 01:38:18 +0800 Subject: [PATCH 015/952] [Bugfix] Give each data-parallel engine its own global RNG streams (#59788) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/utils_/test_torch_utils.py | 15 +++++++++++++++ vllm/utils/torch_utils.py | 6 +++++- vllm/v1/worker/cpu_worker.py | 4 +++- vllm/v1/worker/gpu_worker.py | 4 +++- 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 79dd3e320ae6..72942534623f 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -14,6 +14,7 @@ is_lossless_cast, is_quantized_kv_cache, set_default_torch_dtype, + set_random_seed, set_torch_threads_for_runtime, startup_omp_num_threads, ) @@ -217,3 +218,17 @@ def test_async_tensor_h2d_staging(device): async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32), torch.tensor([1, 2, 3], dtype=torch.int32, device=device), ) + + +def test_set_random_seed_differs_per_data_parallel_index(): + """Unseeded requests draw from the global RNGs; DP engines sharing the engine + seed must get distinct, reproducible streams, and index 0 keeps the seed.""" + + def draws(*args): + set_random_seed(*args) + return np.random.randint(2**31, size=4).tolist(), torch.rand(4).tolist() + + assert draws(42, 0) == draws(42) + assert draws(42, 1) == draws(42, 1) + assert len({str(draws(42, i)) for i in range(4)}) == 4 + assert draws(42, 1) != draws(43, 1) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b4a13419d5dc..025941eeb678 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -530,8 +530,12 @@ def kv_cache_dtype_str_to_dtype( return STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype] -def set_random_seed(seed: int | None) -> None: +def set_random_seed(seed: int | None, data_parallel_index: int = 0) -> None: if seed is not None: + if data_parallel_index: + # DP engines share the seed but must not sample unseeded requests alike. + ss = np.random.SeedSequence((seed, data_parallel_index)) + seed = int(ss.generate_state(1)[0]) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index e36b15300b25..b103215d65de 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -268,7 +268,9 @@ def compile_or_warm_up_model(self) -> CompilationTimes: self.model_runner.warming_up_model() # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. - set_random_seed(self.model_config.seed) + set_random_seed( + self.model_config.seed, self.parallel_config.data_parallel_index + ) return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3f416b94c789..f9b58831e449 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1020,7 +1020,9 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. - set_random_seed(self.model_config.seed) + set_random_seed( + self.model_config.seed, self.parallel_config.data_parallel_index + ) # Eagerly trigger inductor's once-per-process lazy inits during # warmup (rather than on a later compile cache-miss at runtime). From 88f3c601b336871ab19c742a67915da99b64232c Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 5 Oct 2026 11:35:19 -0700 Subject: [PATCH 016/952] [Docs] Add reviewer area of interest for TheEpicDolphin (#60092) Signed-off-by: Giancarlo Delfin --- docs/community/reviewers.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/community/reviewers.md b/docs/community/reviewers.md index fe85e1af493d..60d37f0d3b4d 100644 --- a/docs/community/reviewers.md +++ b/docs/community/reviewers.md @@ -30,7 +30,7 @@ Sorted alphabetically by GitHub handle: - [@Rohan138](https://github.com/Rohan138): ROCm performance, CI, torch.compile fusions - [@simondanielsson](https://github.com/simondanielsson): ROCm performance (CDNA), MoRI-IO. - [@taneem-ibrahim](https://github.com/taneem-ibrahim): Pooling models -- [@TheEpicDolphin](https://github.com/TheEpicDolphin) +- [@TheEpicDolphin](https://github.com/TheEpicDolphin): Spec decode - [@varun-sundar-rabindranath](https://github.com/varun-sundar-rabindranath): KV cache offloading, LoRA - [@vllmellm](https://github.com/vllmellm) - [@wangxiyuan](https://github.com/wangxiyuan): Platform, KV cache, Mooncake From c32495d58ff580b518cf4c5e5d6b69adc4d38351 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Rialland?= <36076211+TQCB@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:37:13 +0200 Subject: [PATCH 017/952] [watermark] compatibility validation (#56801) Signed-off-by: Raphael Rialland Signed-off-by: Simon Veitner Co-authored-by: Simon Veitner --- .buildkite/test_areas/lm_eval.yaml | 4 + docs/features/watermarking.md | 53 ++- .../test_anthropic_messages_conversion.py | 27 ++ .../cohere/test_serving_conversion.py | 21 ++ .../test_generative_scoring.py | 5 + .../test_completion_with_function_calling.py | 37 +- tests/entrypoints/openai/test_watermarking.py | 231 +++++++++++- .../realtime/test_realtime_validation.py | 42 +++ .../test_speech_to_text_cancellation.py | 2 + tests/samplers/test_beam_search_online.py | 173 ++++++++- tests/test_config.py | 54 ++- tests/test_sampling_params.py | 1 + .../test_input_processor_trace_replay.py | 7 +- tests/v1/engine/test_kv_transfer_rejection.py | 27 ++ tests/v1/sample/test_logprobs.py | 1 + tests/v1/worker/test_mixed_warmup_gate.py | 34 ++ tests/watermarking/test_compatibility.py | 353 ++++++++++++++++++ .../watermarking/test_engine_compatibility.py | 204 ++++++++++ tests/watermarking/test_gumbel.py | 25 +- .../test_parallel_compatibility.py | 106 ++++++ tests/watermarking/test_watermarking.py | 180 +++++++-- vllm/benchmarks/latency.py | 1 + vllm/benchmarks/mm_processor.py | 4 +- vllm/benchmarks/throughput.py | 1 + vllm/config/vllm.py | 54 ++- vllm/entrypoints/anthropic/protocol.py | 6 + vllm/entrypoints/anthropic/serving.py | 1 + vllm/entrypoints/cohere/protocol.py | 6 + vllm/entrypoints/cohere/serving.py | 1 + .../generate/beam_search/offline.py | 4 +- .../generate/beam_search/online.py | 2 + .../generate/generative_scoring/serving.py | 1 + .../openai/chat_completion/protocol.py | 4 +- .../entrypoints/openai/completion/protocol.py | 3 +- vllm/entrypoints/openai/responses/protocol.py | 2 +- .../speech_to_text/base/serving.py | 1 + .../speech_to_text/transcription/protocol.py | 5 + .../speech_to_text/translation/protocol.py | 5 + vllm/sampling_params.py | 10 +- vllm/v1/engine/async_llm.py | 8 + vllm/v1/engine/input_processor.py | 46 ++- vllm/v1/engine/llm_engine.py | 7 + vllm/v1/watermarking/gumbel.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/warmup.py | 6 +- 45 files changed, 1669 insertions(+), 100 deletions(-) create mode 100644 tests/v1/engine/test_kv_transfer_rejection.py create mode 100644 tests/watermarking/test_compatibility.py create mode 100644 tests/watermarking/test_engine_compatibility.py create mode 100644 tests/watermarking/test_parallel_compatibility.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 04337d03d262..bcb111084843 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -270,6 +270,9 @@ steps: source_file_dependencies: - .buildkite/test_areas/lm_eval.yaml - tests/evals/gsm8k/ + - tests/watermarking/test_compatibility.py + - tests/watermarking/test_engine_compatibility.py + - tests/watermarking/test_parallel_compatibility.py - vllm/config/speculative.py - vllm/config/vllm.py - vllm/config/watermarking.py @@ -282,6 +285,7 @@ steps: autorun_on_main: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt + - pytest -s -v watermarking/test_compatibility.py watermarking/test_engine_compatibility.py watermarking/test_parallel_compatibility.py mirror: amd: label: ":amd: (MI355) LM Eval Watermark Feature Combination" diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 7ab02d336341..66b6465eb851 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -18,8 +18,16 @@ vllm serve MODEL \ ``` Watermarking is disabled when `--watermark-config` is omitted. Gumbel-max (see -`gumbel`) is the default algorithm within an enabled `WatermarkConfig`. When -watermarking is configured, it is enabled for requests by default. +`gumbel`) is the default algorithm within an enabled `WatermarkConfig`. + +The request-level `watermarking` option has two effective states: + +- Omitted (`None` in Python) or `true` requests watermarking. +- `false` explicitly opts out. + +When a request asks for watermarking but the engine or generation mode cannot +apply it, vLLM emits a server-side warning once and runs the request without a +watermark. Requests can opt out without changing the engine-level algorithm or key: @@ -30,14 +38,12 @@ sampling_params = SamplingParams(watermarking=False) ``` The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the -same `watermarking: false` request field. The Rust token API accepts it in -`sampling_params`, and the Rust gRPC `GenerateRequest` accepts -the optional `watermarking` field. Omitting the field keeps watermarking enabled -when the engine has a watermark configuration. +same optional `watermarking` request field. The Rust token API accepts it in +`sampling_params`, and the Rust gRPC `GenerateRequest` accepts it directly. +Omitting the field requests watermarking by default. -Deployments that require watermarking must restrict this field to trusted -callers, or strip and validate it at the ingress boundary, so untrusted clients -cannot opt out. +Deployments that require watermarking must restrict explicit opt-out to trusted +callers, or strip and validate the field at the ingress boundary. `context_width` controls how many prior tokens seed each watermark decision and defaults to 4. Larger values make the watermark less robust to @@ -63,6 +69,31 @@ the reference algorithms. `WatermarkDetector` consumes token IDs, so callers remain responsible for using the tokenizer and watermark profile that match generation. +## Interactions with generation features + +- Greedy decoding (`temperature=0`) cannot apply a Gumbel watermark. The first + such request emits a server-side warning; it and subsequent greedy requests + use ordinary greedy sampling without a watermark. This includes + transcription, translation, and realtime transcription, which default to + greedy decoding. Transcription and translation can use a nonzero temperature + to apply the configured watermark. +- Beam search cannot apply a watermark. The first watermarked beam request emits + a server-side warning and runs without watermarking. +- Trace replay replaces the sampled token and therefore cannot apply a + watermark. The first watermarked trace replay request emits a server-side + warning and runs without watermarking. +- Parallel sampling (`n > 1`) is supported, but plain Gumbel candidates can be + identical while their contexts remain identical. +- Structured outputs and tool grammars are applied before watermark sampling. + Restrictive constraints can reduce the statistical evidence available to a + detector. +- Automatic language detection and generative scoring use unwatermarked + auxiliary generations. Language-detection tokens are not part of the returned + transcription, and generative scoring returns scores rather than generated + text. +- Speculative decoding has additional requirements and partial-watermarking + modes described in [Speculative decoding](#speculative-decoding). + ## Speculative decoding Watermarking requires speculative decoding to use probabilistic draft sampling, @@ -245,7 +276,7 @@ watermarked output or to modify watermarked text so it is no longer detected. - Watermarking is currently available only with Model Runner V2. - Not all watermarking algorithms have native speculative-decoding support. -- Beam search expands candidates from model log probabilities and does not apply - Gumbel-max watermarking. - Models that replace the vLLM sampler with a custom sampler cannot use configured watermarking. +- Global custom logits processors are unavailable because Model Runner V2 does + not support them. diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index fd0d0dc2a993..1bccee46016e 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -2117,3 +2117,30 @@ def test_tool_change_outside_system_message_is_rejected( ValidationError, match='only allowed in messages with role "system"' ): request_cls(model="test-model", max_tokens=128, **fields) + + +class TestWatermarking: + def test_defaults_to_unspecified(self): + request = _make_request([{"role": "user", "content": "hi"}]) + + assert _convert(request).watermarking is None + + def test_forwards_explicit_enable(self): + request = _make_request([{"role": "user", "content": "hi"}], watermarking=True) + + assert _convert(request).watermarking is True + + def test_forwards_the_opt_out(self): + request = _make_request([{"role": "user", "content": "hi"}], watermarking=False) + + assert _convert(request).watermarking is False + + @pytest.mark.parametrize("watermarking", [None, True, False]) + def test_reaches_sampling_params(self, watermarking): + request = _make_request( + [{"role": "user", "content": "hi"}], watermarking=watermarking + ) + + params = _convert(request).to_sampling_params(16, {}) + + assert params.watermarking is watermarking diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index 5cbe13df0200..58fa7fe7932b 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -2510,3 +2510,24 @@ def test_envelope_uses_400(self): assert err.error.message == "oops" assert err.error.code == 400 assert err.error.type == "bad_request" + + +def test_watermarking_defaults_to_unspecified(): + assert _convert(_make_request()).watermarking is None + + +def test_watermarking_explicit_enable_is_forwarded(): + assert _convert(_make_request(watermarking=True)).watermarking is True + + +def test_watermarking_opt_out_is_forwarded(): + assert _convert(_make_request(watermarking=False)).watermarking is False + + +@pytest.mark.parametrize("watermarking", [None, True, False]) +def test_watermarking_reaches_sampling_params(watermarking): + request = _make_request(watermarking=watermarking) + + params = _convert(request).to_sampling_params(16, {}) + + assert params.watermarking is watermarking diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index f413869c9009..69be10b19d70 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -12,6 +12,7 @@ import math from dataclasses import dataclass, field +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock @@ -302,8 +303,10 @@ async def test_successful_generation(self): mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0} mock_output = _create_mock_request_output(mock_logprobs) + state = SimpleNamespace(sampling_params=None) async def mock_generate(*args, **kwargs): + state.sampling_params = args[1] yield mock_output mock_engine.generate = mock_generate @@ -317,6 +320,8 @@ async def mock_generate(*args, **kwargs): result = await serving.create_generative_scoring(request, None) assert isinstance(result, GenerativeScoringResponse) + assert state.sampling_params is not None + assert state.sampling_params.watermarking is False assert len(result.data) == 2 for item_result in result.data: assert 0.0 <= item_result.score <= 1.0 diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 61fd2f1576bc..72442670a834 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -154,6 +154,40 @@ async def client(server): yield async_client +@pytest.fixture(scope="module") +def watermarked_server(): + args = [ + "--dtype", + "half", + "--enable-auto-tool-choice", + "--structured-outputs-config.backend", + "xgrammar", + "--tool-call-parser", + "hermes", + "--reasoning-parser", + "qwen3", + "--gpu-memory-utilization", + "0.4", + "--enforce-eager", + "--watermark-config", + '{"algorithm":"gumbel","key":42}', + ] + ROCM_EXTRA_ARGS + + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest.fixture(params=["server", "watermarked_server"], ids=["plain", "watermarked"]) +def tool_use_server(request: pytest.FixtureRequest): + return request.getfixturevalue(request.param) + + +@pytest_asyncio.fixture +async def tool_use_client(tool_use_server): + async with tool_use_server.get_async_client() as async_client: + yield async_client + + @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) @pytest.mark.parametrize("stream", [True, False]) @@ -167,12 +201,13 @@ async def client(server): ) @pytest.mark.parametrize("enable_thinking", [True, False]) async def test_function_tool_use( - client: openai.AsyncOpenAI, + tool_use_client: openai.AsyncOpenAI, model_name: str, stream: bool, tool_choice: str | dict, enable_thinking: bool, ): + client = tool_use_client if not stream: # Non-streaming test chat_completion = await client.chat.completions.create( diff --git a/tests/entrypoints/openai/test_watermarking.py b/tests/entrypoints/openai/test_watermarking.py index 0425b86a028b..e6a750eeb74b 100644 --- a/tests/entrypoints/openai/test_watermarking.py +++ b/tests/entrypoints/openai/test_watermarking.py @@ -1,32 +1,249 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from io import BytesIO + +import pytest +from fastapi import UploadFile + +from vllm.entrypoints.openai.chat_completion.protocol import ( + BatchChatCompletionRequest, + ChatCompletionRequest, +) from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.speech_to_text.transcription.protocol import ( + TranscriptionRequest, +) +from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest +from vllm.sampling_params import StructuredOutputsParams -def test_chat_request_disables_watermarking(): +def test_chat_request_preserves_watermarking_opt_out(): request = ChatCompletionRequest( messages=[{"role": "user", "content": "hello"}], watermarking=False ) params = request.to_sampling_params(max_tokens=1, default_sampling_params={}) - assert not params.watermarking + assert params.watermarking is False -def test_completion_request_disables_watermarking(): +def test_completion_request_preserves_watermarking_opt_out(): request = CompletionRequest(prompt="hello", watermarking=False) params = request.to_sampling_params(max_tokens=1) - assert not params.watermarking + assert params.watermarking is False -def test_responses_request_disables_watermarking(): +def test_responses_request_preserves_watermarking_opt_out(): request = ResponsesRequest(input="hello", watermarking=False) params = request.to_sampling_params(default_max_tokens=1) - assert not params.watermarking + assert params.watermarking is False + + +def test_batch_chat_request_preserves_watermarking_opt_out(): + request = BatchChatCompletionRequest( + messages=[[{"role": "user", "content": "hello"}]], watermarking=False + ) + + converted = request.to_chat_completion_request(request.messages[0]) + params = converted.to_sampling_params(max_tokens=1, default_sampling_params={}) + + assert params.watermarking is False + + +@pytest.mark.parametrize( + "api_request,conversion,kwargs", + [ + ( + ChatCompletionRequest( + messages=[{"role": "user", "content": "hello"}], + use_beam_search=True, + watermarking=False, + ), + "to_beam_search_params", + {"max_tokens": 1, "default_sampling_params": {}}, + ), + ( + CompletionRequest(prompt="hello", use_beam_search=True, watermarking=False), + "to_beam_search_params", + {"max_tokens": 1, "default_sampling_params": {}}, + ), + ( + TranscriptionRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=False, + ), + "to_sampling_params", + {"default_max_tokens": 1}, + ), + ( + TranscriptionRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=False, + ), + "to_beam_search_params", + {"default_max_tokens": 1}, + ), + ( + TranslationRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=False, + ), + "to_sampling_params", + {"default_max_tokens": 1}, + ), + ( + TranslationRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=False, + ), + "to_beam_search_params", + {"default_max_tokens": 1}, + ), + ], +) +def test_request_preserves_watermarking_opt_out(api_request, conversion, kwargs): + params = getattr(api_request, conversion)(**kwargs) + + assert params.watermarking is False + + +@pytest.mark.parametrize("request_cls", [TranscriptionRequest, TranslationRequest]) +def test_speech_requests_leave_watermarking_unspecified(request_cls): + request = request_cls( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + temperature=0.8, + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1) + beam_params = request.to_beam_search_params(default_max_tokens=1) + + assert request.watermarking is None + assert sampling_params.watermarking is None + assert beam_params.watermarking is None + + +@pytest.mark.parametrize( + "api_request,kwargs", + [ + ( + ChatCompletionRequest( + messages=[{"role": "user", "content": "choose A or B"}], + temperature=0.8, + structured_outputs=StructuredOutputsParams(choice=["A", "B"]), + ), + {"max_tokens": 8, "default_sampling_params": {}}, + ), + ( + CompletionRequest( + prompt="choose A or B", + temperature=0.8, + structured_outputs=StructuredOutputsParams(choice=["A", "B"]), + ), + {"max_tokens": 8}, + ), + ( + ResponsesRequest( + input="choose A or B", + temperature=0.8, + structured_outputs=StructuredOutputsParams(choice=["A", "B"]), + ), + {"default_max_tokens": 8}, + ), + ], +) +def test_request_preserves_unspecified_watermarking_with_structured_outputs( + api_request, kwargs +): + params = api_request.to_sampling_params(**kwargs) + + assert params.watermarking is None + assert params.structured_outputs is not None + assert params.structured_outputs.choice == ["A", "B"] + + +@pytest.mark.parametrize( + "api_request,conversion,kwargs", + [ + ( + ChatCompletionRequest( + messages=[{"role": "user", "content": "hello"}], watermarking=True + ), + "to_sampling_params", + {"max_tokens": 1, "default_sampling_params": {}}, + ), + ( + ChatCompletionRequest( + messages=[{"role": "user", "content": "hello"}], + use_beam_search=True, + watermarking=True, + ), + "to_beam_search_params", + {"max_tokens": 1, "default_sampling_params": {}}, + ), + ( + CompletionRequest(prompt="hello", watermarking=True), + "to_sampling_params", + {"max_tokens": 1}, + ), + ( + CompletionRequest(prompt="hello", use_beam_search=True, watermarking=True), + "to_beam_search_params", + {"max_tokens": 1}, + ), + ( + ResponsesRequest(input="hello", watermarking=True), + "to_sampling_params", + {"default_max_tokens": 1}, + ), + ( + TranscriptionRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=True, + ), + "to_sampling_params", + {"default_max_tokens": 1}, + ), + ( + TranslationRequest( + file=UploadFile(file=BytesIO(), filename="audio.wav"), + watermarking=True, + ), + "to_sampling_params", + {"default_max_tokens": 1}, + ), + ], +) +def test_request_preserves_explicit_watermarking(api_request, conversion, kwargs): + params = getattr(api_request, conversion)(**kwargs) + + assert params.watermarking is True + + +def test_batch_chat_request_preserves_explicit_watermarking(): + request = BatchChatCompletionRequest( + messages=[[{"role": "user", "content": "hello"}]], watermarking=True + ) + + converted = request.to_chat_completion_request(request.messages[0]) + params = converted.to_sampling_params(max_tokens=1, default_sampling_params={}) + + assert converted.watermarking is True + assert params.watermarking is True + + +def test_batch_chat_request_inherits_watermarking(): + request = BatchChatCompletionRequest( + messages=[[{"role": "user", "content": "hello"}]] + ) + + converted = request.to_chat_completion_request(request.messages[0]) + params = converted.to_sampling_params(max_tokens=1, default_sampling_params={}) + + assert converted.watermarking is None + assert params.watermarking is None diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index c74246b2f87a..a8c4f35c0409 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -4,6 +4,7 @@ import asyncio import json import warnings +from types import SimpleNamespace import numpy as np import pybase64 as base64 @@ -12,6 +13,7 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer from vllm.assets.audio import AudioAsset +from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection from vllm.multimodal.media.audio import load_audio MISTRAL_FORMAT_ARGS = [ @@ -26,6 +28,46 @@ MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" +def test_realtime_generation_preserves_watermarking(): + transcription_outputs = [] + + class EngineClient: + def generate(self, **kwargs): + async def outputs(): + output = SimpleNamespace( + outputs=[], + prompt_token_ids=[], + is_watermarked=kwargs["sampling_params"].watermarking, + ) + transcription_outputs.append(output) + yield output + + return outputs() + + async def streaming_input(): + chunk: np.ndarray + for chunk in (): + yield chunk + + async def send(_event): + return None + + connection = object.__new__(RealtimeConnection) + connection.connection_id = "test" + connection.serving = SimpleNamespace( + model_cls=SimpleNamespace(realtime_max_tokens=1), + engine_client=EngineClient(), + ) + connection.audio_queue = asyncio.Queue() + connection._is_connected = True + connection.send = send + + asyncio.run(connection._run_generation(streaming_input(), asyncio.Queue())) + + assert len(transcription_outputs) == 1 + assert transcription_outputs[0].is_watermarked is None + + def _get_websocket_url(server: RemoteOpenAIServer) -> str: """Convert HTTP URL to WebSocket URL for realtime endpoint.""" http_url = server.url_root diff --git a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py index 0ff8adb6b5e2..8ed31e4cdb1e 100644 --- a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py +++ b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py @@ -198,4 +198,6 @@ async def test_language_detection_cancel_aborts_engine_request(): with pytest.raises(asyncio.CancelledError): await task + sampling_params = engine_client.generate.call_args.args[1] + assert sampling_params.watermarking is False engine_client.abort.assert_awaited_once_with(request_id) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index cbed58363ad2..22afce2d78eb 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -1,13 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable, Sequence +from typing import Any, cast +from unittest.mock import MagicMock + import pytest +import torch -from vllm import CompletionOutput, RequestOutput +from vllm import CompletionOutput, PromptType, RequestOutput, SamplingParams +from vllm import logger as vllm_logger +from vllm.config import VllmConfig +from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin -from vllm.inputs import TokensInput +from vllm.entrypoints.generate.beam_search.utils import ( + BeamSearchInstance, + BeamSearchSequence, +) +from vllm.inputs import EngineInput, TokensInput from vllm.logprobs import Logprob, SampleLogprobs from vllm.sampling_params import BeamSearchParams +from vllm.v1.structured_output.backend_types import StructuredOutputBackend + + +@pytest.fixture +def reset_warning_once(): + vllm_logger._print_warning_once.cache_clear() + yield + vllm_logger._print_warning_once.cache_clear() class _Tokenizer: @@ -27,8 +47,26 @@ def get_tokenizer(self) -> _Tokenizer: return _Tokenizer() +class _VllmConfig: + watermark_config = object() + speculative_config = None + + def _check_supports_watermarking(self, params: BeamSearchParams) -> bool: + return VllmConfig._check_supports_watermarking(cast(VllmConfig, self), params) + + +class _InputProcessor: + vllm_config = _VllmConfig() + + def resolve_watermarking(self, params): + return self.vllm_config._check_supports_watermarking(params) + + class _EngineClient: + input_processor = _InputProcessor() + async def generate(self, prompt, *args, **kwargs): + assert args[0].watermarking is False yield RequestOutput( request_id=kwargs.get("request_id", "test-request"), prompt=prompt.get("prompt"), @@ -56,11 +94,42 @@ async def generate(self, prompt, *args, **kwargs): ) -class _Serving(BeamSearchOnlineMixin): +class _AsyncServing(BeamSearchOnlineMixin): renderer = _Renderer() engine_client = _EngineClient() +class _OfflineServing(BeamSearchOfflineMixin): + renderer = _Renderer() + llm_engine = _EngineClient() + + def _preprocess_cmpl( + self, + prompts: Sequence[PromptType], + tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, + ) -> Sequence[EngineInput]: + return cast(Sequence[EngineInput], prompts) + + def _lora_request_to_seq(self, lora_request, num_requests): + return [None] * num_requests + + def _beam_search_step( + self, + instances_batch: list[BeamSearchInstance], + base_sampling_params: SamplingParams, + eos_token_id: int | None, + ignore_eos: bool, + beam_width: int, + sort_beams_key: Callable[..., Any], + structured_output_backend: StructuredOutputBackend | None, + structured_output_key: tuple[Any, ...] | None, + structured_output_bitmask: torch.Tensor | None, + ) -> bool: + assert base_sampling_params.watermarking is False + return True + + @pytest.mark.asyncio async def test_beam_search_handles_extra_logprob_candidates() -> None: prompt: TokensInput = { @@ -68,10 +137,11 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None: "prompt": "prompt", "prompt_token_ids": [1], } - params = BeamSearchParams(beam_width=2, max_tokens=1) + params = BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False) outputs = [ - output async for output in _Serving().beam_search(prompt, "request", params) + output + async for output in _AsyncServing().beam_search(prompt, "request", params) ] assert len(outputs) == 1 @@ -101,10 +171,12 @@ async def test_beam_search_respects_skip_special_tokens( max_tokens=1, ignore_eos=True, skip_special_tokens=skip_special_tokens, + watermarking=False, ) outputs = [ - output async for output in _Serving().beam_search(prompt, "request", params) + output + async for output in _AsyncServing().beam_search(prompt, "request", params) ] assert outputs[0].outputs[0].text == expected_text @@ -139,7 +211,7 @@ async def generate(prompt, *args, **kwargs): output.finish_reason = "abort" yield result - serving = _Serving() + serving = _AsyncServing() monkeypatch.setattr(serving.engine_client, "generate", generate) prompt: TokensInput = { "type": "token", @@ -149,7 +221,9 @@ async def generate(prompt, *args, **kwargs): outputs = [ output async for output in serving.beam_search( - prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4) + prompt, + "request", + BeamSearchParams(beam_width=2, max_tokens=4, watermarking=False), ) ] @@ -165,3 +239,86 @@ async def generate(prompt, *args, **kwargs): assert output.logprobs is not None assert len(output.logprobs) == abort_after assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after) + + +@pytest.mark.asyncio +async def test_beam_search_warns_and_disables_watermarking( + caplog_vllm, reset_warning_once +) -> None: + prompt = { + "type": "token", + "prompt": "prompt", + "prompt_token_ids": [1], + } + params = BeamSearchParams(beam_width=2, max_tokens=1) + + with caplog_vllm.at_level("WARNING"): + output = await anext(_AsyncServing().beam_search(prompt, "request", params)) + + assert "beam search requests will run without watermarking" in caplog_vllm.text + # The caller's params object is never rewritten by admission. + assert params.watermarking is None + assert output.finished + + +def test_offline_beam_search_warns_and_disables_watermarking( + caplog_vllm, reset_warning_once +) -> None: + params = BeamSearchParams(beam_width=2, max_tokens=1) + + with caplog_vllm.at_level("WARNING"): + outputs = _OfflineServing().beam_search( + cast(list[PromptType], [{"type": "token", "prompt_token_ids": [1]}]), + params, + ) + + assert "beam search requests will run without watermarking" in caplog_vllm.text + # The caller's params object is never rewritten by admission. + assert params.watermarking is None + assert len(outputs) == 1 + + +def test_offline_beam_search_disables_internal_watermarking() -> None: + outputs = _OfflineServing().beam_search( + cast(list[PromptType], [{"type": "token", "prompt_token_ids": [1]}]), + BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False), + ) + + assert len(outputs) == 1 + + +def test_offline_structured_beam_search_disables_internal_watermarking() -> None: + grammar = MagicMock() + grammar.is_terminated.return_value = False + grammar.fill_bitmask.side_effect = lambda bitmask, index: bitmask[index].fill_(1) + backend = MagicMock() + backend.compile_grammar.return_value = grammar + serving = _OfflineServing() + serving.model_config = MagicMock() + serving.model_config.get_vocab_size.return_value = 32 + beam = BeamSearchSequence( + orig_prompt={"type": "token", "prompt_token_ids": [1]}, + tokens=[1], + logprobs=[], + ) + + base_params = SamplingParams( + logprobs=2, + watermarking=False, + skip_clone=True, + ) + entries = serving._build_beam_sampling_params( + [beam], + base_params, + backend, + ("regex", ".*"), + torch.zeros((1, 1), dtype=torch.int32), + ) + + assert entries[0] is not None + beam_params = entries[0][0] + assert beam_params is not base_params + assert beam_params.logprobs == base_params.logprobs + assert beam_params.watermarking is False + assert beam_params.allowed_token_ids == [0] + assert base_params.allowed_token_ids is None diff --git a/tests/test_config.py b/tests/test_config.py index b1348480cc83..9d4567874b4f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -47,6 +47,7 @@ from vllm.config.utils import get_field from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel from vllm.platforms import current_platform +from vllm.sampling_params import BeamSearchParams from vllm.transformers_utils.config import ( _patch_hf_transformers_nested_rope_validation, get_pooling_config, @@ -3630,7 +3631,7 @@ def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log ) with caplog_vllm.at_level(logging.WARNING): - config._check_watermarking_unsupported() + config._check_supports_watermarking() assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text @@ -3651,7 +3652,7 @@ def test_speculative_watermarking_without_context_dedup_does_not_warn( caplog_vllm.clear() with caplog_vllm.at_level(logging.WARNING): - config._check_watermarking_unsupported() + config._check_supports_watermarking() assert "dedup" not in caplog_vllm.text.lower() @@ -3672,7 +3673,7 @@ def test_speculative_context_dedup_logs_no_unsupported_warning( caplog_vllm.clear() with caplog_vllm.at_level(logging.WARNING): - config._check_watermarking_unsupported() + config._check_supports_watermarking() assert "dedup" not in caplog_vllm.text.lower() @@ -3690,7 +3691,7 @@ def test_gumbel_rejects_speculative_decoding_without_target_only(): ) with pytest.raises(ValueError, match="'gumbel'.*allow_target_only_watermarking"): - config._check_watermarking_unsupported() + config._check_supports_watermarking() def test_dual_key_gumbel_warns_that_configured_alpha_is_unused( @@ -3708,7 +3709,7 @@ def test_dual_key_gumbel_warns_that_configured_alpha_is_unused( ) with caplog_vllm.at_level(logging.WARNING): - config._check_watermarking_unsupported() + config._check_supports_watermarking() assert "The configured alpha=0.25 is not used" in caplog_vllm.text @@ -3734,7 +3735,7 @@ def test_dual_key_gumbel_alpha_warning_is_not_emitted( ) with caplog_vllm.at_level(logging.WARNING): - config._check_watermarking_unsupported() + config._check_supports_watermarking() assert "is not used" not in caplog_vllm.text @@ -3750,7 +3751,7 @@ def test_dual_key_gumbel_requires_probabilistic_drafting(): ) with pytest.raises(ValueError, match="draft_sample_method='probabilistic'"): - config._check_watermarking_unsupported() + config._check_supports_watermarking() @pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"]) @@ -3764,7 +3765,7 @@ def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method): parallel_drafting=False, ) - config._check_watermarking_unsupported() + config._check_supports_watermarking() def test_dual_key_gumbel_supports_dspark(): @@ -3777,7 +3778,7 @@ def test_dual_key_gumbel_supports_dspark(): parallel_drafting=True, ) - config._check_watermarking_unsupported() + config._check_supports_watermarking() def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): @@ -3791,7 +3792,7 @@ def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): ) with pytest.raises(ValueError, match="autoregressive model-based"): - config._check_watermarking_unsupported() + config._check_supports_watermarking() @pytest.mark.parametrize( @@ -3815,17 +3816,40 @@ def test_dual_key_gumbel_rejects_incompatible_speculative_modes(overrides, match config.speculative_config = SimpleNamespace(**values) with pytest.raises(ValueError, match=match): - config._check_watermarking_unsupported() + config._check_supports_watermarking() -def test_gumbel_watermark_rejects_beam_search(): - with pytest.raises(ValueError, match="Beam search is not supported"): - _watermarked_vllm_config()._check_watermarking_unsupported(beam_search=True) +def test_gumbel_watermark_disables_beam_search_watermarking(caplog_vllm): + enabled = _watermarked_vllm_config()._check_supports_watermarking( + BeamSearchParams(beam_width=2, max_tokens=1) + ) + + assert not enabled + assert "beam search requests will run without watermarking" in caplog_vllm.text + + +def test_gumbel_watermark_allows_unwatermarked_beam_search(): + _watermarked_vllm_config()._check_supports_watermarking( + BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False) + ) + + +def test_unwatermarked_beam_does_not_bypass_engine_incompatibilities(): + config = _watermarked_vllm_config() + config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + + with pytest.raises(ValueError, match="Speculative decoding with watermarking"): + config._check_supports_watermarking( + BeamSearchParams(beam_width=2, max_tokens=1, watermarking=False) + ) def test_gumbel_watermark_rejects_custom_sampler(): with pytest.raises(ValueError, match="custom samplers are not supported"): - _watermarked_vllm_config()._check_watermarking_unsupported(custom_sampler=True) + _watermarked_vllm_config()._check_supports_watermarking(custom_sampler=True) def test_watermark_key_must_fit_in_64_bits(): diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index a6e057ad7a41..f32a3d4c367f 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -112,6 +112,7 @@ def test_diffusion_canvas_width_is_accepted_with_either_scheduler( diffusion_config=DiffusionConfig(canvas_length=8), tokenizer=None, validate_logits_processors_params=lambda params: None, + resolve_watermarking=lambda params: False, ) params = SamplingParams(extra_args=extra_args) InputProcessor._validate_params(processor, params, ("generate",)) diff --git a/tests/v1/engine/test_input_processor_trace_replay.py b/tests/v1/engine/test_input_processor_trace_replay.py index 957f70b71858..043de170afad 100644 --- a/tests/v1/engine/test_input_processor_trace_replay.py +++ b/tests/v1/engine/test_input_processor_trace_replay.py @@ -82,7 +82,12 @@ def _validate(enable_trace_replay: bool) -> None: enable_trace_replay=enable_trace_replay, is_diffusion=False, ), - vllm_config=SimpleNamespace(reasoning_config=None), + vllm_config=SimpleNamespace( + reasoning_config=None, + # Engine has no watermark configuration. + _check_supports_watermarking=lambda params: False, + ), + resolve_watermarking=lambda params: False, speculative_config=None, structured_outputs_config=None, diffusion_config=None, diff --git a/tests/v1/engine/test_kv_transfer_rejection.py b/tests/v1/engine/test_kv_transfer_rejection.py new file mode 100644 index 000000000000..a21351900ca2 --- /dev/null +++ b/tests/v1/engine/test_kv_transfer_rejection.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from vllm.v1.engine.async_llm import AsyncLLM + + +@pytest.mark.asyncio +async def test_kv_transfer_rejection_stub_disables_watermarking(): + engine = object.__new__(AsyncLLM) + add_request = AsyncMock() + engine.engine_core = SimpleNamespace( + add_request_async=add_request, + shutdown=lambda **kwargs: None, + ) + + await engine.notify_kv_transfer_request_rejected( + "request", + {"do_remote_prefill": True}, + ) + + request = add_request.call_args.args[0] + assert request.sampling_params.watermarking is False diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 4152662fbd50..813bea0bf3aa 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -518,6 +518,7 @@ def test_prompt_logprob_token_ids_require_v2_model_runner(): structured_outputs_config=None, tokenizer=None, validate_logits_processors_params=lambda params: None, + resolve_watermarking=lambda params: False, ) params = SamplingParams(prompt_logprob_token_ids=[[1, 2]]) with patch.object(SamplingParams, "verify"): diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 04f456a7c9cd..84ddf692a131 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -6,6 +6,7 @@ import pytest +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.worker.gpu import warmup from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup @@ -31,6 +32,39 @@ def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): ) +def test_mixed_warmup_disables_watermarking(): + outputs: list[SchedulerOutput] = [] + runner = SimpleNamespace( + is_pooling_model=False, + max_num_reqs=2, + kv_cache_config=SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16)) + ], + num_blocks=128, + ), + vllm_config=SimpleNamespace(num_lookahead_tokens=0), + max_model_len=128, + model_state=SimpleNamespace(max_encoder_len=0), + kv_connector=SimpleNamespace(set_disabled=lambda disabled: None), + ) + + assert run_mixed_prefill_decode_warmup( + runner, + worker_execute_model=outputs.append, + worker_sample_tokens=lambda grammar_output: None, + num_tokens=32, + ) + + sampling_params = [ + request.sampling_params + for output in outputs + for request in output.scheduled_new_reqs + ] + assert len(sampling_params) == 2 + assert all(params.watermarking is False for params in sampling_params) + + @pytest.mark.parametrize("fail_warmup", [False, True]) def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_warmup): """Startup must warm fixed drafts before calibration and retain its manager.""" diff --git a/tests/watermarking/test_compatibility.py b/tests/watermarking/test_compatibility.py new file mode 100644 index 000000000000..748223d94b56 --- /dev/null +++ b/tests/watermarking/test_compatibility.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from vllm import SamplingParams +from vllm import logger as vllm_logger +from vllm.config import VllmConfig, WatermarkConfig +from vllm.renderers import BaseRenderer +from vllm.sampling_params import BeamSearchParams +from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.engine.llm_engine import LLMEngine + + +def _input_processor(server_uses_watermarking: bool = True) -> InputProcessor: + config = VllmConfig() + config.model_config = SimpleNamespace( + try_get_generation_config=lambda: {}, + logits_processors=None, + return_sampling_mask=False, + enable_trace_replay=True, + is_multimodal_model=False, + supports_multimodal_inputs=False, + runner_type="generate", + is_diffusion=False, + ) + config.watermark_config = ( + WatermarkConfig(key=42) if server_uses_watermarking else None + ) + renderer = MagicMock(spec=BaseRenderer) + renderer.tokenizer = None + renderer._executor = None + return InputProcessor(config, renderer) + + +def _validate( + params: SamplingParams, server_uses_watermarking: bool = True +) -> bool | None: + processor = _input_processor(server_uses_watermarking) + with patch.object(SamplingParams, "verify"): + return processor._validate_params(params, ("generate",)) + + +def _engine_core_request(params: SamplingParams) -> EngineCoreRequest: + return EngineCoreRequest( + request_id="request", + prompt_token_ids=[1], + mm_features=None, + sampling_params=params, + pooling_params=None, + arrival_time=0.0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + ) + + +@pytest.fixture +def reset_warning_once(): + vllm_logger._print_warning_once.cache_clear() + yield + vllm_logger._print_warning_once.cache_clear() + + +def test_greedy_watermarked_requests_warn_once(caplog_vllm, reset_warning_once): + first = SamplingParams(temperature=0) + second = SamplingParams(temperature=0) + + with caplog_vllm.at_level("WARNING"): + first_resolved = _validate(first) + second_resolved = _validate(second) + + message = "subsequent greedy requests will use ordinary greedy sampling" + assert caplog_vllm.text.count(message) == 1 + assert first_resolved is False + assert second_resolved is False + # The caller's params objects are left untouched. + assert first.watermarking is None + assert second.watermarking is None + + +def test_default_watermarking_resolves_when_configured(): + enabled = SamplingParams() + + assert _validate(enabled) is True + # The caller's params object keeps inheriting. + assert enabled.watermarking is None + + +@pytest.mark.parametrize("watermarking", [None, True]) +def test_requested_watermarking_without_engine_config_warns_and_disables( + watermarking, caplog_vllm, reset_warning_once +): + params = SamplingParams(watermarking=watermarking) + + with caplog_vllm.at_level("WARNING"): + resolved = _validate(params, server_uses_watermarking=False) + + assert "engine has no watermark configuration" in caplog_vllm.text + assert resolved is False + assert params.watermarking is watermarking + + +@pytest.mark.parametrize("watermarking", [None, True]) +def test_explicit_and_default_greedy_warn_and_disable( + watermarking, caplog_vllm, reset_warning_once +): + params = SamplingParams(temperature=0, watermarking=watermarking) + + with caplog_vllm.at_level("WARNING"): + resolved = _validate(params) + + assert "ordinary greedy sampling" in caplog_vllm.text + assert resolved is False + assert params.watermarking is watermarking + + +def test_trace_replay_with_watermarking_warns_and_disables( + caplog_vllm, reset_warning_once +): + params = SamplingParams(trace_decode_token_ids=[1]) + + with caplog_vllm.at_level("WARNING"): + resolved = _validate(params) + + assert "trace replay requests will run without watermarking" in caplog_vllm.text + assert resolved is False + assert params.watermarking is None + + +@pytest.mark.parametrize( + "params", + [ + SamplingParams(temperature=0, watermarking=False), + SamplingParams(trace_decode_token_ids=[1], watermarking=False), + ], +) +def test_incompatible_modes_are_allowed_when_watermarking_is_disabled(params): + _validate(params) + + +@pytest.mark.parametrize( + "params", + [ + SamplingParams(temperature=0), + SamplingParams(trace_decode_token_ids=[1]), + ], +) +def test_requests_without_engine_config_are_unwatermarked( + params, caplog_vllm, reset_warning_once +): + with caplog_vllm.at_level("WARNING"): + resolved = _validate(params, server_uses_watermarking=False) + + assert "engine has no watermark configuration" in caplog_vllm.text + assert resolved is False + assert params.watermarking is None + + +def test_trace_replay_takes_precedence_over_greedy_warning( + caplog_vllm, reset_warning_once +): + params = SamplingParams(temperature=0, trace_decode_token_ids=[1]) + + with caplog_vllm.at_level("WARNING"): + _validate(params) + + assert "trace replay requests will run without watermarking" in caplog_vllm.text + assert "ordinary greedy sampling" not in caplog_vllm.text + + +@pytest.mark.parametrize( + "params", + [ + SamplingParams(seed=42), + SamplingParams(n=2), + SamplingParams(n=2, seed=42), + ], +) +def test_seeded_and_parallel_sampling_are_allowed(params): + _validate(params) + + +def test_direct_engine_trace_replay_warns_and_disables( + caplog_vllm, monkeypatch, reset_warning_once +): + engine = object.__new__(LLMEngine) + engine.input_processor = _input_processor() + engine.vllm_config = engine.input_processor.vllm_config + assign_request_id = MagicMock() + monkeypatch.setattr(engine.input_processor, "assign_request_id", assign_request_id) + engine.output_processor = MagicMock() + engine.engine_core = MagicMock() + request = _engine_core_request(SamplingParams(trace_decode_token_ids=[1])) + + with caplog_vllm.at_level("WARNING"): + engine.add_request( + request.request_id, + request, + SamplingParams(watermarking=False), + ) + + assert "trace replay requests will run without watermarking" in caplog_vllm.text + assert request.sampling_params.watermarking is False + assert request.sampling_params._watermarking_skipped is True + assign_request_id.assert_called_once_with(request) + + +def test_direct_engine_request_uses_embedded_opt_out(monkeypatch): + engine = object.__new__(LLMEngine) + engine.input_processor = _input_processor() + engine.vllm_config = engine.input_processor.vllm_config + assign_request_id = MagicMock() + monkeypatch.setattr(engine.input_processor, "assign_request_id", assign_request_id) + engine.output_processor = MagicMock() + engine.engine_core = MagicMock() + request = _engine_core_request(SamplingParams(temperature=0, watermarking=False)) + + engine.add_request( + request.request_id, + request, + SamplingParams(temperature=0), + ) + + assert request.sampling_params._watermarking_skipped is False + assign_request_id.assert_called_once_with(request) + + +def test_direct_async_engine_trace_replay_warns_and_disables( + caplog_vllm, monkeypatch, reset_warning_once +): + engine = object.__new__(AsyncLLM) + engine.engine_core = SimpleNamespace( + resources=SimpleNamespace(engine_dead=False), + shutdown=lambda **kwargs: None, + add_request_async=AsyncMock(), + ) + engine.output_handler = None + engine.input_processor = _input_processor() + engine.vllm_config = engine.input_processor.vllm_config + engine.output_processor = MagicMock() + engine.log_requests = False + monkeypatch.setattr(engine, "_run_output_handler", MagicMock()) + request = _engine_core_request(SamplingParams(trace_decode_token_ids=[1])) + + async def add_request() -> None: + await engine.add_request( + request.request_id, + request, + SamplingParams(watermarking=False), + ) + + with caplog_vllm.at_level("WARNING"): + asyncio.run(add_request()) + + assert "trace replay requests will run without watermarking" in caplog_vllm.text + assert request.sampling_params.watermarking is False + assert request.sampling_params._watermarking_skipped is True + + +def test_direct_async_engine_greedy_request_warns_and_disables_watermarking( + caplog_vllm, monkeypatch, reset_warning_once +): + engine = object.__new__(AsyncLLM) + engine.engine_core = SimpleNamespace( + resources=SimpleNamespace(engine_dead=False), + shutdown=lambda **kwargs: None, + add_request_async=AsyncMock(), + ) + engine.output_handler = None + engine.input_processor = _input_processor() + engine.vllm_config = engine.input_processor.vllm_config + engine.output_processor = MagicMock() + engine.log_requests = False + monkeypatch.setattr(engine, "_run_output_handler", MagicMock()) + request = _engine_core_request(SamplingParams(temperature=0)) + + async def add_request() -> None: + await engine.add_request( + request.request_id, + request, + SamplingParams(watermarking=False), + ) + + with caplog_vllm.at_level("WARNING"): + asyncio.run(add_request()) + + assert "subsequent greedy requests will use ordinary greedy sampling" in ( + caplog_vllm.text + ) + assert request.sampling_params.watermarking is False + assert request.sampling_params._watermarking_skipped is True + + +@pytest.mark.parametrize( + ("kwargs", "server_uses_watermarking", "skipped"), + [ + ({}, True, False), + ({"temperature": 0}, True, True), + ({"temperature": 0, "watermarking": False}, True, False), + ({"temperature": 0}, False, False), + ], +) +def test_apply_watermarking_records_skipped_requests( + kwargs, server_uses_watermarking, skipped, reset_warning_once +): + params = SamplingParams(**kwargs) + processor = _input_processor(server_uses_watermarking) + processor.apply_watermarking(params, processor.resolve_watermarking(params)) + assert params._watermarking_skipped is skipped + + # Resolving an already resolved request keeps the flag. + processor.apply_watermarking(params, processor.resolve_watermarking(params)) + assert params._watermarking_skipped is skipped + + +def test_shared_sampling_params_are_not_rewritten_by_admission(reset_warning_once): + """A params object reused across prompts/engines must keep inheriting.""" + params = SamplingParams(temperature=0.8) + + assert _validate(params, server_uses_watermarking=False) is False + assert params.watermarking is None + # Same object submitted to a watermarked engine must still be watermarked. + assert _validate(params, server_uses_watermarking=True) is True + + greedy = SamplingParams(temperature=0.0) + assert _validate(greedy) is False + greedy.temperature = 0.8 + assert _validate(greedy) is True + + +def test_shared_beam_search_params_are_not_rewritten_by_admission( + caplog_vllm, reset_warning_once +): + """Beam search must keep resolving per engine after reuse.""" + params = BeamSearchParams(beam_width=2, max_tokens=4, temperature=0.8) + + # A plain engine leaves the tri-state unresolved on the caller's object. + assert _input_processor(False).resolve_watermarking(params) is False + assert params.watermarking is None + + # The same object on a watermarked engine is warned about and still + # resolves to False, without the first decision being written back. + with caplog_vllm.at_level("WARNING"): + assert _input_processor(True).resolve_watermarking(params) is False + assert "beam search requests will run without watermarking" in caplog_vllm.text + assert params.watermarking is None diff --git a/tests/watermarking/test_engine_compatibility.py b/tests/watermarking/test_engine_compatibility.py new file mode 100644 index 000000000000..1320306cb399 --- /dev/null +++ b/tests/watermarking/test_engine_compatibility.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Generator + +import pytest + +from vllm import LLM, SamplingParams +from vllm.platforms import current_platform +from vllm.sampling_params import StructuredOutputsParams + +MODEL = "hmellor/tiny-random-LlamaForCausalLM" +MAX_NUM_BATCHED_TOKENS = 32 + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) + + +@pytest.fixture(scope="module") +def watermarked_llm(vllm_runner) -> Generator[LLM, None, None]: + with vllm_runner( + MODEL, + dtype="half", + enforce_eager=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS, + max_model_len=128, + max_num_seqs=4, + disable_log_stats=False, + watermark_config={"algorithm": "gumbel", "key": 42}, + ) as runner: + yield runner.llm + + +def test_watermark_with_combined_sampling_controls(watermarked_llm: LLM): + params = SamplingParams( + temperature=0.8, + top_k=8, + top_p=0.9, + min_p=0.01, + presence_penalty=0.2, + frequency_penalty=0.2, + repetition_penalty=1.1, + logit_bias={10: 100.0}, + allowed_token_ids=[10], + max_tokens=8, + ignore_eos=True, + ) + + output = watermarked_llm.generate("Sampling smoke test", params)[0] + + assert list(output.outputs[0].token_ids) == [10] * 8 + + # The assertion above holds with or without watermarking, because + # allowed_token_ids leaves a single candidate. Repeat with a wider allowed + # set, where only watermarked sampling is reproducible without a seed. + wide = SamplingParams( + temperature=0.8, + top_k=8, + top_p=0.9, + min_p=0.01, + presence_penalty=0.2, + frequency_penalty=0.2, + repetition_penalty=1.1, + allowed_token_ids=list(range(10, 210)), + max_tokens=8, + ignore_eos=True, + ) + + first = watermarked_llm.generate("Sampling smoke test", wide)[0] + second = watermarked_llm.generate("Sampling smoke test", wide)[0] + + assert set(first.outputs[0].token_ids) <= set(range(10, 210)) + assert first.outputs[0].token_ids == second.outputs[0].token_ids + + +def test_watermark_with_penalties_and_nucleus_sampling(watermarked_llm: LLM): + params = SamplingParams( + temperature=0.8, + top_k=8, + top_p=0.9, + min_p=0.01, + presence_penalty=0.2, + frequency_penalty=0.2, + repetition_penalty=1.1, + max_tokens=8, + ignore_eos=True, + ) + + first = watermarked_llm.generate("Penalty smoke test", params)[0] + second = watermarked_llm.generate("Penalty smoke test", params)[0] + + assert len(first.outputs[0].token_ids) == 8 + assert first.outputs[0].token_ids == second.outputs[0].token_ids + + +def test_plain_gumbel_seeded_parallel_sampling_is_deterministic( + watermarked_llm: LLM, +): + output = watermarked_llm.generate( + "Seeded parallel sampling smoke test", + SamplingParams( + n=2, + seed=42, + temperature=0.8, + max_tokens=8, + ignore_eos=True, + ), + )[0] + + assert len(output.outputs) == 2 + assert all(len(candidate.token_ids) == 8 for candidate in output.outputs) + assert output.outputs[0].token_ids == output.outputs[1].token_ids + + # Plain Gumbel keys on the token context, so the request seed does not + # change a watermarked draw. Pin that down rather than implying the seed + # is what makes the candidates match. + unseeded = watermarked_llm.generate( + "Seeded parallel sampling smoke test", + SamplingParams( + n=2, + temperature=0.8, + max_tokens=8, + ignore_eos=True, + ), + )[0] + + assert unseeded.outputs[0].token_ids == output.outputs[0].token_ids + + +@pytest.mark.parametrize( + "structured_outputs", + [ + StructuredOutputsParams(choice=["A", "B"]), + StructuredOutputsParams(grammar='root ::= "A" | "B"'), + ], + ids=["choice", "grammar"], +) +def test_watermark_with_structured_output( + watermarked_llm: LLM, structured_outputs: StructuredOutputsParams +): + params = SamplingParams( + temperature=0.8, + max_tokens=8, + structured_outputs=structured_outputs, + ) + + output = watermarked_llm.generate("Choose A or B:", params)[0] + + assert output.outputs[0].text.strip() in ("A", "B") + + +def test_watermark_with_wide_grammar_is_deterministic(watermarked_llm: LLM): + # A wider grammar makes deterministic watermarking observable. + params = SamplingParams( + temperature=0.8, + max_tokens=8, + structured_outputs=StructuredOutputsParams( + grammar="root ::= [A-Z] [A-Z] [A-Z] [A-Z] [A-Z] [A-Z]" + ), + ) + + first = watermarked_llm.generate("Emit six letters:", params)[0] + second = watermarked_llm.generate("Emit six letters:", params)[0] + + assert first.outputs[0].text.strip().isupper() + assert first.outputs[0].token_ids == second.outputs[0].token_ids + + +def test_watermark_with_prefix_cache_and_chunked_prefill(watermarked_llm: LLM): + prompt = {"prompt_token_ids": [101] * 48} + params = SamplingParams(temperature=0.8, max_tokens=8, ignore_eos=True) + + assert len(prompt["prompt_token_ids"]) > MAX_NUM_BATCHED_TOKENS + + cold = watermarked_llm.generate(prompt, params)[0] + cached = watermarked_llm.generate(prompt, params)[0] + + assert cached.num_cached_tokens > 0 + assert cold.outputs[0].token_ids == cached.outputs[0].token_ids + + +def test_mixed_watermarked_batch(watermarked_llm: LLM): + prompts = [ + {"prompt_token_ids": [101] * 32 + [102]}, + {"prompt_token_ids": [201] * 32 + [202]}, + ] + marked = SamplingParams(temperature=0.8, max_tokens=8, ignore_eos=True) + ordinary = SamplingParams( + temperature=0.8, + seed=123, + max_tokens=8, + ignore_eos=True, + watermarking=False, + ) + + batched = watermarked_llm.generate(prompts, [marked, ordinary]) + marked_alone = watermarked_llm.generate(prompts[0], marked)[0] + ordinary_alone = watermarked_llm.generate(prompts[1], ordinary)[0] + + assert batched[0].outputs[0].token_ids == marked_alone.outputs[0].token_ids + assert batched[1].outputs[0].token_ids == ordinary_alone.outputs[0].token_ids diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index d789692a0e0b..aa0ffa279904 100644 --- a/tests/watermarking/test_gumbel.py +++ b/tests/watermarking/test_gumbel.py @@ -4,9 +4,11 @@ import pytest import torch +from vllm.config.watermarking import WatermarkConfig from vllm.platforms import current_platform from vllm.v1.watermarking import ( DualKeyGumbelWatermarkDetector, + DualKeyGumbelWatermarker, GumbelWatermarkDetector, GumbelWatermarker, derive_watermark_key, @@ -42,6 +44,21 @@ def test_detector_deduplicates_context_even_when_target_differs(): assert detection.num_scored_tokens == 3 +def test_watermarker_respects_filtered_token_support(): + logits = torch.full((2, 8), -torch.inf) + logits[0, 3] = 0 + logits[1, 6] = 0 + contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]]) + + token_ids = ( + GumbelWatermarker(key=42) + .sample(logits, contexts, lambda values: None) + .token_ids + ) + + assert torch.equal(token_ids, torch.tensor([3, 6])) + + def test_dual_key_detector_scores_each_token_against_both_keys(): token_ids = [1, 2, 3, 4, 5] key_a_detector = GumbelWatermarkDetector( @@ -99,8 +116,12 @@ def test_dual_key_detector_rejects_invalid_alpha(): DualKeyGumbelWatermarkDetector(key=42, alpha=1.1) -def test_dual_key_detector_default_alpha(): - assert DualKeyGumbelWatermarkDetector(key=42).alpha == 0.2 +def test_dual_key_alpha_defaults_match(): + config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) + watermarker = DualKeyGumbelWatermarker(key=42) + detector = DualKeyGumbelWatermarkDetector(key=42) + + assert detector.alpha == watermarker.alpha == config.alpha @pytest.mark.skipif( diff --git a/tests/watermarking/test_parallel_compatibility.py b/tests/watermarking/test_parallel_compatibility.py new file mode 100644 index 000000000000..6e47a65b6f4b --- /dev/null +++ b/tests/watermarking/test_parallel_compatibility.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio + +import pytest + +from tests.utils import multi_gpu_test +from vllm import SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.v1.engine.async_llm import AsyncLLM + +MODEL = "hmellor/tiny-random-LlamaForCausalLM" + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "parallel_kwargs", + [ + {"tensor_parallel_size": 2, "enable_batch_sharded_sampling": True}, + {"pipeline_parallel_size": 2}, + ], + ids=["tp2-sharded-sampling", "pp2"], +) +def test_watermarked_generation_with_model_parallelism(vllm_runner, parallel_kwargs): + prompts = ["Parallel watermark smoke"] * 2 + + with vllm_runner( + MODEL, + dtype="half", + enforce_eager=True, + max_model_len=128, + max_num_seqs=4, + watermark_config={"algorithm": "gumbel", "key": 42}, + **parallel_kwargs, + ) as runner: + first = runner.llm.generate( + prompts, + SamplingParams( + temperature=0.8, + seed=1, + max_tokens=8, + ignore_eos=True, + ), + ) + second = runner.llm.generate( + prompts, + SamplingParams( + temperature=0.8, + seed=2, + max_tokens=8, + ignore_eos=True, + ), + ) + + first_tokens = [output.outputs[0].token_ids for output in first] + second_tokens = [output.outputs[0].token_ids for output in second] + assert all(len(token_ids) == 8 for token_ids in first_tokens) + assert first_tokens[0] == first_tokens[1] + assert first_tokens == second_tokens + + +@multi_gpu_test(num_gpus=2) +def test_watermarked_generation_with_data_parallelism(): + async def run() -> None: + engine = AsyncLLM.from_engine_args( + AsyncEngineArgs( + model=MODEL, + dtype="half", + data_parallel_size=2, + enforce_eager=True, + max_model_len=128, + max_num_seqs=4, + watermark_config={"algorithm": "gumbel", "key": 42}, + ) + ) + + async def generate(index: int) -> tuple[int, ...]: + final_output = None + async for output in engine.generate( + "Data-parallel watermark smoke", + SamplingParams( + temperature=0.8, + seed=index, + max_tokens=8, + ignore_eos=True, + ), + request_id=f"watermark-dp-{index}", + data_parallel_rank=index, + ): + final_output = output + assert final_output is not None + return tuple(final_output.outputs[0].token_ids) + + try: + token_ids = await asyncio.gather(generate(0), generate(1)) + assert len(token_ids[0]) == 8 + assert token_ids[0] == token_ids[1] + finally: + engine.shutdown() + + try: + asyncio.run(run()) + finally: + cleanup_dist_env_and_memory() diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 4de6022e34f2..384cc30d570b 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -8,8 +8,11 @@ import torch from vllm import SamplingParams +from vllm import logger as vllm_logger +from vllm.config import VllmConfig from vllm.config.watermarking import WatermarkConfig from vllm.platforms import current_platform +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch from vllm.v1.watermarking import ( DualKeyGumbelWatermarker, SupportsSpeculativeDecoding, @@ -26,8 +29,8 @@ speculative_target_watermark_key, ) from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample +from vllm.v1.worker.gpu import buffer_utils from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor -from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.watermark import ( draft_watermarking_mask, philox_gumbel_sample, @@ -48,6 +51,44 @@ def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor: return logits.argmax(dim=-1) +@pytest.fixture +def make_gpu_watermark_sampler(monkeypatch): + class CPUUvaBuffer: + def __init__(self, size, dtype): + self.cpu = torch.zeros(size, dtype=dtype) + self.np = self.cpu.numpy() + + def uva(self, n=None): + return self.cpu[:n] if n is not None else self.cpu + + monkeypatch.setattr(buffer_utils, "UvaBuffer", CPUUvaBuffer) + + def make(watermarker, *, max_num_reqs=2, vocab_size=8, **kwargs): + req_states = SimpleNamespace( + max_num_reqs=max_num_reqs, + vocab_size=vocab_size, + device=torch.device("cpu"), + all_token_ids=SimpleNamespace(gpu=torch.zeros(max_num_reqs, 1)), + prompt_len=SimpleNamespace( + np=np.zeros(max_num_reqs, dtype=np.int32), + gpu=torch.zeros(max_num_reqs, dtype=torch.int32), + ), + prefill_len=SimpleNamespace(np=np.zeros(max_num_reqs, dtype=np.int32)), + total_len=SimpleNamespace(gpu=torch.zeros(max_num_reqs, dtype=torch.int32)), + ) + return GPUWatermarkSampler( + watermarker, + vllm_config=VllmConfig(), + max_num_reqs=max_num_reqs, + vocab_size=vocab_size, + device=torch.device("cpu"), + req_states=req_states, + **kwargs, + ) + + return make + + @pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"]) def test_watermarker_contract(algorithm: str): watermarker = create_watermarker( @@ -271,49 +312,47 @@ def test_dual_key_derivation_is_stable(): def test_sampling_params_can_disable_watermarking(): - assert SamplingParams().watermarking - assert not SamplingParams.from_optional(watermarking=False).watermarking + assert SamplingParams().watermarking is None + assert SamplingParams.from_optional(watermarking=False).watermarking is False -def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch): - sampler = object.__new__(GPUWatermarkSampler) - sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool)) +def test_gpu_sampler_warns_about_unexpected_greedy_watermarking( + monkeypatch, make_gpu_watermark_sampler +): + vllm_logger._print_warning_once.cache_clear() + sampler = make_gpu_watermark_sampler(StubWatermarker(), max_num_reqs=1) messages: list[str] = [] - monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None) monkeypatch.setattr( - "vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append + "vllm.v1.watermarking.gpu_sampler.logger.warning", + lambda message, *_args, **_kwargs: messages.append(message), ) - sampler.add_request(0, SamplingParams(temperature=0)) - sampler.add_request(0, SamplingParams(temperature=1)) + sampler.add_request(0, SamplingParams(temperature=0, watermarking=True)) + assert sampler.watermarking.np[0] + sampler.add_request(0, SamplingParams(temperature=0, watermarking=True)) + assert sampler.watermarking.np[0] sampler.add_request(0, SamplingParams(temperature=0, watermarking=False)) - + assert not sampler.watermarking.np[0] assert messages == [ ( - "Watermarking is enabled, but greedy decoding (temperature=0) cannot be " - "watermarked. This request will use ordinary greedy sampling." + "Watermarking is enabled, but greedy decoding " + "(temperature=0) cannot be watermarked. This request will use " + "ordinary greedy sampling." ) ] + vllm_logger._print_warning_once.cache_clear() -def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch): - sampler = object.__new__(GPUWatermarkSampler) - sampler.watermarker = StubWatermarker() - sampler.deduplicate_contexts = "single_turn" - sampler.watermarking = SimpleNamespace( - np=np.array([True, False]), gpu=torch.tensor([True, False]) - ) - sampler.sampling_states = SimpleNamespace( - temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)), - seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)), - ) - sampler.use_fp64_gumbel = False +def test_gpu_sampler_respects_mixed_request_watermarking( + monkeypatch, make_gpu_watermark_sampler +): + sampler = make_gpu_watermark_sampler(StubWatermarker(), deduplicate_contexts="none") + sampler.add_request(0, SamplingParams(watermarking=True)) + sampler.add_request(1, SamplingParams(watermarking=False)) + sampler.apply_staged_writes() sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( 2, 1, dtype=torch.int64 ) - sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros( - 2, dtype=torch.bool - ) monkeypatch.setattr( "vllm.v1.watermarking.watermarker.gumbel_sample", lambda *args, **kwargs: torch.tensor([3, 4]), @@ -322,17 +361,92 @@ def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch): sampled, output_logits = sampler._sample_random( logits, - torch.tensor([0, 1]), - np.array([0, 1]), + torch.tensor([1, 0]), + np.array([1, 0]), torch.zeros(2, dtype=torch.int64), None, None, False, ) - assert torch.equal(sampled, torch.tensor([7, 4])) - assert torch.equal(output_logits[0], torch.full((8,), 10.0)) - assert torch.equal(output_logits[1], logits[1]) + assert torch.equal(sampled, torch.tensor([3, 7])) + assert torch.equal(output_logits[0], logits[0]) + assert torch.equal(output_logits[1], torch.full((8,), 10.0)) + + +def test_gpu_sampler_filters_top_k_top_p_before_watermarking( + monkeypatch, make_gpu_watermark_sampler +): + # ``apply_top_k_top_p`` dispatches to a Triton kernel whenever Triton is + # installed, which cannot read this test's CPU tensors. Force the torch + # reference path; the assertion is about the call order, not the kernel. + monkeypatch.setattr( + "vllm.v1.watermarking.gpu_sampler.apply_top_k_top_p", + lambda logits, k, p: apply_top_k_top_p_pytorch( + logits, k, p, allow_cpu_sync=True + ), + ) + + class CapturingWatermarker: + context_width = 1 + captured_logits = None + + def sample(self, logits, contexts, random_sampler=None, skip_mask=None): + self.captured_logits = logits + return WatermarkSample(logits.argmax(dim=-1), logits) + + watermarker = CapturingWatermarker() + sampler = make_gpu_watermark_sampler( + watermarker, max_num_reqs=1, vocab_size=4, deduplicate_contexts="none" + ) + sampler.add_request(0, SamplingParams(top_k=2, top_p=0.8, watermarking=True)) + sampler.apply_staged_writes() + sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros( + 1, 1, dtype=torch.int64 + ) + logits = torch.tensor([[5.0, 4.0, 3.0, 2.0]]) + + sampled, _ = sampler._sample_random( + logits, + torch.tensor([0]), + np.array([0]), + torch.zeros(1, dtype=torch.int64), + torch.tensor([2]), + torch.tensor([0.8]), + False, + ) + + assert watermarker.captured_logits is not None + assert torch.isneginf(watermarker.captured_logits[0, 2:]).all() + assert sampled.item() in (0, 1) + + +def test_watermark_context_ignores_prefix_cache_bookkeeping(): + sampler = object.__new__(GPUWatermarkSampler) + sampler.watermarker = SimpleNamespace(context_width=3) + sampler.req_states = SimpleNamespace( + total_len=SimpleNamespace(gpu=torch.tensor([4, 0, 5])), + prompt_len=SimpleNamespace(gpu=torch.tensor([3, 0, 2])), + all_token_ids=SimpleNamespace( + gpu=torch.tensor( + [ + [10, 11, 12, 40, 0, 0], + [0, 0, 0, 0, 0, 0], + [20, 21, 50, 51, 52, 0], + ] + ) + ), + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([0, 0, 0])), + ) + request_slots = torch.tensor([2, 0]) + + cold_contexts = sampler._get_contexts(request_slots) + sampler.req_states.num_computed_tokens.gpu[:] = torch.tensor([3, 0, 2]) + cached_contexts = sampler._get_contexts(request_slots) + + expected = torch.tensor([[50, 51, 52], [-1, -1, 40]]) + assert torch.equal(cold_contexts, expected) + assert torch.equal(cached_contexts, expected) def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch): diff --git a/vllm/benchmarks/latency.py b/vllm/benchmarks/latency.py index acb3efef153a..0c75c04563b4 100644 --- a/vllm/benchmarks/latency.py +++ b/vllm/benchmarks/latency.py @@ -119,6 +119,7 @@ def llm_generate(): beam_width=args.n, max_tokens=args.output_len, ignore_eos=True, + watermarking=False, ), ) diff --git a/vllm/benchmarks/mm_processor.py b/vllm/benchmarks/mm_processor.py index 40b3494151d3..7a52faf29de3 100644 --- a/vllm/benchmarks/mm_processor.py +++ b/vllm/benchmarks/mm_processor.py @@ -241,6 +241,7 @@ def benchmark_multimodal_processor( temperature=0.0, max_tokens=output_len, detokenize=True, + watermarking=False, ) for output_len in expected_output_lens ] @@ -262,7 +263,8 @@ def benchmark_multimodal_processor( warmup_prompts = [req.prompt for req in warmup_requests] warmup_output_lens = [req.expected_output_len for req in warmup_requests] warmup_sampling_params = [ - SamplingParams(max_tokens=output_len) for output_len in warmup_output_lens + SamplingParams(max_tokens=output_len, watermarking=False) + for output_len in warmup_output_lens ] llm.chat( warmup_prompts, diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 620ce2b08f13..d7c3c0145409 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -178,6 +178,7 @@ def _run_vllm_requests( beam_width=n, max_tokens=output_len, ignore_eos=True, + watermarking=False, ), ) if do_profile: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53d4102697c5..dec1f27fe1c9 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -60,6 +60,7 @@ from transformers import PreTrainedConfig from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.v1.kv_cache_interface import KVCacheConfig else: PreTrainedConfig = Any @@ -1314,15 +1315,22 @@ def _verify_trace_replay_config(self) -> None: if not self.use_v2_model_runner: raise ValueError("trace replay requires Model Runner V2") - def _check_watermarking_unsupported( + def _check_supports_watermarking( self, + config: "SamplingParams | BeamSearchParams | None" = None, *, - beam_search: bool = False, custom_sampler: bool = False, - ) -> None: + ) -> bool: watermark_config = getattr(self, "watermark_config", None) if watermark_config is None: - return + if config is not None and config.watermarking is not False: + logger.warning_once( + "Watermarking is enabled for this request, but the engine has no " + "watermark configuration. This and subsequent requests will run " + "without watermarking.", + scope="global", + ) + return False if self.speculative_config is not None: speculative_config = self.speculative_config if speculative_config.draft_sample_method != "probabilistic": @@ -1377,12 +1385,44 @@ def _check_watermarking_unsupported( watermark_config.alpha, scope="global", ) - if beam_search: - raise ValueError("Beam search is not supported with watermarking.") if custom_sampler: raise ValueError( "Model-specific custom samplers are not supported with watermarking." ) + if config is None: + return True + if config.watermarking is False: + return False + + from vllm.sampling_params import BeamSearchParams, SamplingParams + + if isinstance(config, BeamSearchParams): + logger.warning_once( + "Watermarking is enabled, but beam search cannot be watermarked. " + "This and subsequent beam search requests will run without " + "watermarking.", + scope="global", + ) + return False + if not isinstance(config, SamplingParams): + raise TypeError(f"Unsupported watermarking config: {type(config).__name__}") + if config.trace_decode_token_ids is not None: + logger.warning_once( + "Watermarking is enabled, but trace replay cannot be watermarked. " + "This and subsequent trace replay requests will run without " + "watermarking.", + scope="global", + ) + return False + if config.temperature == 0: + logger.warning_once( + "Watermarking is enabled, but greedy decoding " + "(temperature=0) cannot be watermarked. This and subsequent " + "greedy requests will use ordinary greedy sampling.", + scope="global", + ) + return False + return True def _resolve_and_verify_engram_config(self) -> None: """Resolve defaults and validate n-gram embedding settings.""" @@ -1446,7 +1486,7 @@ def __post_init__(self): self.try_verify_and_update_config() self._resolve_and_verify_engram_config() - self._check_watermarking_unsupported() + self._check_supports_watermarking() # Models may have supplied their own DCP defaults above; anything still # unset falls back to the stock ones. self.parallel_config.set_dcp_defaults() diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 5abfaffed33b..2a14c6293327 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -242,6 +242,12 @@ class AnthropicMessagesRequest(BaseModel): "to 256 bit)." ), ) + watermarking: bool | None = Field( + default=None, + description=( + "Whether to apply the engine's configured watermark to this request." + ), + ) kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters used for disaggregated serving.", diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index cf5fd9ecc163..05c1e24f2a48 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -613,6 +613,7 @@ def _build_base_request( temperature=anthropic_request.temperature, top_p=anthropic_request.top_p, top_k=anthropic_request.top_k, + watermarking=anthropic_request.watermarking, cache_salt=anthropic_request.cache_salt, kv_transfer_params=anthropic_request.kv_transfer_params, ec_transfer_params=anthropic_request.ec_transfer_params, diff --git a/vllm/entrypoints/cohere/protocol.py b/vllm/entrypoints/cohere/protocol.py index 1432e6eb1185..45e0d310dc93 100644 --- a/vllm/entrypoints/cohere/protocol.py +++ b/vllm/entrypoints/cohere/protocol.py @@ -175,6 +175,12 @@ class CohereChatV2Request(BaseModel): # vLLM-specific extensions (not in Cohere spec). These mirror what the # Anthropic and OpenAI surfaces already expose so V2 callers can reach # the same engine knobs when needed. + watermarking: bool | None = Field( + default=None, + description=( + "Whether to apply the engine's configured watermark to this request." + ), + ) kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters used for disaggregated serving.", diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 63d00f0c576a..bdaa098659a7 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -557,6 +557,7 @@ def _build_base_chat_completion( temperature=request.temperature, top_p=request.p, top_k=request.k, + watermarking=request.watermarking, seed=request.seed, frequency_penalty=request.frequency_penalty, presence_penalty=request.presence_penalty, diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index b7095131af31..1ad451ac13ff 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -82,7 +82,7 @@ def beam_search( temperature = params.temperature ignore_eos = params.ignore_eos length_penalty = params.length_penalty - self.llm_engine.vllm_config._check_watermarking_unsupported(beam_search=True) + self.llm_engine.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() eos_token_id = tokenizer.eos_token_id @@ -120,6 +120,7 @@ def beam_search( logprobs=2 * beam_width, max_tokens=1, temperature=temperature, + watermarking=False, detokenize=False, skip_clone=True, # Internal beam search, safe to skip clone ) @@ -457,6 +458,7 @@ def _build_beam_sampling_params( logprobs=base_params.logprobs, max_tokens=1, temperature=base_params.temperature, + watermarking=base_params.watermarking, detokenize=False, allowed_token_ids=( allowed_ids diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index ac7f0f884f5b..5a95f40d3198 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -40,6 +40,7 @@ async def beam_search( temperature = params.temperature length_penalty = params.length_penalty include_stop_str_in_output = params.include_stop_str_in_output + self.engine_client.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() eos_token_id = tokenizer.eos_token_id @@ -62,6 +63,7 @@ async def beam_search( logprobs=logprobs_num, max_tokens=1, temperature=temperature, + watermarking=False, detokenize=False, ) all_beams = [ diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index cb0c4995bf75..4d5a15c2f54c 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -261,6 +261,7 @@ async def create_generative_scoring( logprobs=len(request.label_token_ids), logprob_token_ids=request.label_token_ids, n=1, + watermarking=False, ) # Get trace headers diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 228472ab6b5f..6a70624be4fe 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -268,7 +268,7 @@ class ChatCompletionRequest(OpenAIBaseModel): top_k: int | None = None min_p: float | None = None repetition_penalty: float | None = None - watermarking: bool = True + watermarking: bool | None = None length_penalty: float = 1.0 stop_token_ids: list[int] | None = [] include_stop_str_in_output: bool = False @@ -647,6 +647,7 @@ def to_beam_search_params( max_tokens=max_tokens, ignore_eos=self.ignore_eos, temperature=temperature, + watermarking=self.watermarking, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, skip_special_tokens=self.skip_special_tokens, @@ -1100,6 +1101,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel): top_k: int | None = None min_p: float | None = None repetition_penalty: float | None = None + watermarking: bool | None = None length_penalty: float | None = 1.0 early_stopping: bool = False structured_outputs: StructuredOutputsParams | None = None diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index a0974cb91c69..78722b45ec7b 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -75,7 +75,7 @@ class CompletionRequest(OpenAIBaseModel): top_k: int | None = None min_p: float | None = None repetition_penalty: float | None = None - watermarking: bool = True + watermarking: bool | None = None length_penalty: float = 1.0 stop_token_ids: list[int] | None = [] include_stop_str_in_output: bool = False @@ -303,6 +303,7 @@ def to_beam_search_params( max_tokens=max_tokens, ignore_eos=self.ignore_eos, temperature=temperature, + watermarking=self.watermarking, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, skip_special_tokens=self.skip_special_tokens, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 9c44a04b3a14..8fa96ecdb05c 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -245,7 +245,7 @@ class ResponsesRequest(OpenAIBaseModel): ) # --8<-- [start:responses-extra-params] - watermarking: bool = True + watermarking: bool | None = None request_id: str = Field( default_factory=lambda: f"resp_{random_uuid()}", description=( diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 34b7e9c24928..871f69666f51 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -229,6 +229,7 @@ async def _detect_language( max_tokens=1, temperature=0.0, allowed_token_ids=allowed_token_ids, + watermarking=False, ) result_generator = self.engine_client.generate( diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index e9570fb926d3..7e66c512c720 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -151,6 +151,9 @@ class TranscriptionRequest(OpenAIBaseModel): to automatically increase the temperature until certain thresholds are hit. """ + watermarking: bool | None = None + """Whether to apply the engine's configured watermark to this request.""" + top_p: float | None = None """Enables nucleus (top-p) sampling, where tokens are selected from the smallest possible set whose cumulative probability exceeds `p`. @@ -226,6 +229,7 @@ def to_beam_search_params( beam_width=n, max_tokens=max_tokens, temperature=temperature, + watermarking=self.watermarking, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, ) @@ -264,6 +268,7 @@ def to_sampling_params( return SamplingParams.from_optional( temperature=temperature, + watermarking=self.watermarking, max_tokens=max_tokens, seed=self.seed, top_p=top_p, diff --git a/vllm/entrypoints/speech_to_text/translation/protocol.py b/vllm/entrypoints/speech_to_text/translation/protocol.py index ac8c69fdc284..91ec56e2fba8 100644 --- a/vllm/entrypoints/speech_to_text/translation/protocol.py +++ b/vllm/entrypoints/speech_to_text/translation/protocol.py @@ -102,6 +102,9 @@ class TranslationRequest(OpenAIBaseModel): to automatically increase the temperature until certain thresholds are hit. """ + watermarking: bool | None = None + """Whether to apply the engine's configured watermark to this request.""" + top_p: float | None = None """Enables nucleus (top-p) sampling, where tokens are selected from the smallest possible set whose cumulative probability exceeds `p`. @@ -215,6 +218,7 @@ def to_beam_search_params( beam_width=n, max_tokens=max_tokens, temperature=temperature, + watermarking=self.watermarking, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, ) @@ -251,6 +255,7 @@ def to_sampling_params( return SamplingParams.from_optional( temperature=temperature, + watermarking=self.watermarking, max_tokens=max_tokens, seed=self.seed, top_p=top_p, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 246e2ce90c77..ab9f2711142c 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -251,8 +251,9 @@ class SamplingParams( """Controls the randomness of the sampling. Lower values make the model more deterministic, while higher values make the model more random. Zero means greedy sampling.""" - watermarking: bool = True - """Whether to apply the engine's configured watermark to this request.""" + watermarking: bool | None = None + """Whether to apply the engine's configured watermark to this request. + ``None`` defaults to enabled.""" top_p: float = 1.0 """Controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to 1 to consider all tokens.""" @@ -342,6 +343,7 @@ class SamplingParams( output_text_buffer_length: int = 0 _eos_token_id: int | None = None _all_stop_token_ids: set[int] = msgspec.field(default_factory=set) + _watermarking_skipped: bool = False # Fields used to construct logits processors structured_outputs: StructuredOutputsParams | None = None @@ -395,7 +397,7 @@ def from_optional( frequency_penalty: float | None = 0.0, repetition_penalty: float | None = 1.0, temperature: float | None = 1.0, - watermarking: bool = True, + watermarking: bool | None = None, top_p: float | None = 1.0, top_k: int = 0, min_p: float = 0.0, @@ -1411,6 +1413,7 @@ def for_sampler_warmup() -> "SamplingParams": _bad_words_token_ids=[[0], [1, 2]], logprobs=5, prompt_logprobs=1, + watermarking=True, ) @@ -1430,6 +1433,7 @@ class BeamSearchParams( include_stop_str_in_output: bool = False structured_outputs: StructuredOutputsParams | None = None skip_special_tokens: bool = True + watermarking: bool | None = None def __post_init__(self) -> None: _verify_num_sequences(self.beam_width, "beam_width") diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 4a7a91cb8e07..37025a863230 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -443,6 +443,13 @@ async def add_request( "does not match the EngineCoreRequest.request_id attribute. The " "latter will be used, and the former will be ignored." ) + request_params = request.params + if isinstance(request_params, SamplingParams): + # This request object is owned by the engine from here on. + self.input_processor.apply_watermarking( + request_params, + self.input_processor.resolve_watermarking(request_params), + ) else: if isinstance(prompt, dict) and "type" in prompt: # Rendered EngineInput; no blocking preprocessing needed. @@ -910,6 +917,7 @@ async def notify_kv_transfer_request_rejected( sampling_params=SamplingParams( max_tokens=1, extra_args={"kv_transfer_params": dict(kv_transfer_params)}, + watermarking=False, ), pooling_params=None, arrival_time=time.time(), diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 8b5a69bcf4c0..17e8af5caeef 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -25,7 +25,7 @@ from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer, renderer_from_config from vllm.renderers.inputs.preprocess import parse_model_prompt -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tasks import GENERATION_TASKS, POOLING_TASKS, SupportedTask from vllm.tokenizers import TokenizerLike from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid @@ -112,12 +112,41 @@ def _build_logits_processors_params_validator( return partial(validate_logits_processors_parameters, custom_logitsprocs) + def resolve_watermarking(self, params: SamplingParams | BeamSearchParams) -> bool: + """Return whether watermarking will be applied to this request. + + ``None`` and ``True`` request watermarking, while ``False`` opts out. + If watermarking is unavailable or incompatible with the request, emit + a warning and return ``False``. + """ + return self.vllm_config._check_supports_watermarking(params) + + def apply_watermarking(self, params: SamplingParams, watermarking: bool) -> None: + """Store the resolved `watermarking` on request-local params. + + Also records whether watermarking was requested on a watermarking + engine but could not be applied, as opposed to an explicit opt-out. + """ + # Sticky, so resolving an already resolved request keeps the flag. + params._watermarking_skipped = params._watermarking_skipped or ( + self.vllm_config.watermark_config is not None + and params.watermarking is not False + and not watermarking + ) + params.watermarking = watermarking + def _validate_params( self, params: SamplingParams | PoolingParams, supported_tasks: tuple[SupportedTask, ...], - ) -> None: - """Raise `ValueError` if SamplingParams or PoolingParams is not valid.""" + ) -> bool | None: + """Raise `ValueError` if SamplingParams or PoolingParams is not valid. + + Returns the resolved `watermarking` value for `SamplingParams` (`None` + for `PoolingParams`). The caller applies it to its private copy of the + params so that a caller-owned params object is never mutated. + """ + watermarking: bool | None = None if isinstance(params, SamplingParams): supported_generation_tasks = [ task for task in supported_tasks if task in GENERATION_TASKS @@ -189,6 +218,9 @@ def _validate_params( "enabled. Start the engine with --enable-trace-replay " "to use it." ) + # Resolved last so that feature gates above report the missing + # feature rather than its incompatibility with watermarking. + watermarking = self.resolve_watermarking(params) elif isinstance(params, PoolingParams): supported_pooling_tasks = [ task for task in supported_tasks if task in POOLING_TASKS @@ -217,6 +249,8 @@ def _validate_params( f"but got {type(params).__name__}" ) + return watermarking + def _normalize_trace_replay_params( self, sampling_params: SamplingParams, prompt_len: int ) -> None: @@ -352,7 +386,7 @@ def process_inputs( session_id: str | None = None, kv_hints: KvHintsEnvelope | None = None, ) -> EngineCoreRequest: - self._validate_params(params, supported_tasks) + watermarking = self._validate_params(params, supported_tasks) self._validate_lora(lora_request) parallel_config = self.vllm_config.parallel_config @@ -413,6 +447,10 @@ def process_inputs( if isinstance(params, SamplingParams): # TODO: can we avoid cloning here in multiproc case? sampling_params = params.clone() + # Resolve on the request-local copy: `params` may be shared across + # prompts, requests and even engines by the caller. + assert watermarking is not None + self.apply_watermarking(sampling_params, watermarking) prompt_len = length_from_prompt_token_ids_or_embeds( prompt_token_ids, prompt_embeds ) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index c3fb3efe7e9b..509d46703dd1 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -254,6 +254,13 @@ def add_request( "does not match the EngineCoreRequest.request_id attribute. The " "latter will be used, and the former will be ignored." ) + request_params = request.params + if isinstance(request_params, SamplingParams): + # This request object is owned by the engine from here on. + self.input_processor.apply_watermarking( + request_params, + self.input_processor.resolve_watermarking(request_params), + ) else: request = self.input_processor.process_inputs( request_id, diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 655267d1e4b0..2f953dd38965 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/vllm/v1/watermarking/gumbel.py @@ -221,7 +221,7 @@ def __init__( p_value_threshold: float = 0.01, prf: WatermarkPRFName = "philox", deduplicate_contexts: bool = True, - alpha: float = 0.2, + alpha: float = 0.1, ) -> None: if not 0 <= alpha <= 1: raise ValueError("alpha must be between 0 and 1") diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1b9d0608e26e..b9880f43f242 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -482,7 +482,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: custom = self.model_state.custom_sampler(self.sampler) if custom: - self.vllm_config._check_watermarking_unsupported(custom_sampler=True) + self.vllm_config._check_supports_watermarking(custom_sampler=True) self.sampler, self.rejection_sampler = custom elif self.speculative_config is not None: self.rejection_sampler = RejectionSampler( diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 2af97207ce8d..c29d9958fd15 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -139,7 +139,11 @@ def _alloc_blocks(num_blocks: int) -> list[int]: next_block_id += num_blocks return block_ids - sampling_params = SamplingParams(max_tokens=2, temperature=0.0) + sampling_params = SamplingParams( + max_tokens=2, + temperature=0.0, + watermarking=False, + ) decode_prefill_output = SchedulerOutput.make_empty() decode_prefill_output.scheduled_new_reqs = [ From 74c5cbcd7b6fa07fa5568ba7f7dfbfedc502964e Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 5 Oct 2026 14:54:11 -0400 Subject: [PATCH 018/952] [Bugfix][MRV2][Spec Decode] Honor dynamic K in autoregressive speculators, 1.2~1.3x kernel perf improvement (#57053) Signed-off-by: yewentao256 --- tests/v1/cudagraph/test_cudagraph_manager.py | 80 ++++++++++++++++++- .../test_gpu_autoregressive_speculator.py | 12 ++- tests/v1/worker/test_gpu_model_runner_v2.py | 1 + .../worker/test_gpu_model_runner_v2_eplb.py | 1 + tests/v1/worker/test_gpu_warmup_blocks.py | 9 ++- vllm/v1/worker/gpu/cudagraph_utils.py | 1 + vllm/v1/worker/gpu/model_runner.py | 8 +- .../autoregressive/cudagraph_utils.py | 45 +++++++++++ .../spec_decode/autoregressive/speculator.py | 40 +++++++--- .../gpu/spec_decode/dflash/speculator.py | 1 + .../gpu/spec_decode/draft_model/speculator.py | 1 + .../gpu/spec_decode/extract_hidden_states.py | 1 + .../multi_module_mtp/speculator.py | 1 + vllm/v1/worker/gpu/spec_decode/speculator.py | 1 + vllm/v1/worker/gpu/warmup.py | 6 ++ 15 files changed, 190 insertions(+), 18 deletions(-) diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index 154116b80f13..668fbdf61b4a 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -20,8 +20,12 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBuffers from vllm.v1.worker.gpu.pcp_manager import PCPManager +from vllm.v1.worker.gpu.spec_decode.autoregressive import ( + cudagraph_utils as spec_cudagraph_utils, +) from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( SpeculatorCudaGraphManager, ) @@ -226,8 +230,10 @@ def test_speculator_capture_preserves_decode_query_bounds( ) manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager) manager.max_num_reqs = num_reqs + manager.decode_query_len = uniform_token_count or max_query_len or num_tokens manager.vllm_config = SimpleNamespace( - parallel_config=SimpleNamespace(data_parallel_size=1) + parallel_config=SimpleNamespace(data_parallel_size=1), + num_speculative_tokens=3, ) buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu")) block_tables = MagicMock() @@ -488,6 +494,78 @@ def test_dynamic_spec_decode_shared_token_count_stays_reachable(monkeypatch): ), desc +@pytest.mark.parametrize("specialize", [False, True]) +def test_speculator_replays_graph_for_runtime_width(monkeypatch, specialize): + """Equal batch shapes must replay the graph captured for the selected K.""" + monkeypatch.setattr( + gpu_cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr( + gpu_cudagraph_utils.current_platform, "get_global_graph_pool", lambda: object() + ) + monkeypatch.setattr(gpu_cudagraph_utils, "get_offloader", MagicMock) + monkeypatch.setattr( + spec_cudagraph_utils, "prepare_inputs_to_capture", lambda *a, **kw: ({}, {}) + ) + forward = MagicMock() + replayed_widths: list[int] = [] + + def capture(manager, create_forward_fn, progress_bar_desc): + for desc in manager._capture_descs[CUDAGraphMode.FULL]: + create_forward_fn(desc, warmup=False)(CUDAGraphMode.NONE) + k = forward.call_args.kwargs["num_speculative_steps"] + graph = MagicMock() + graph.replay.side_effect = lambda k=k: replayed_widths.append(k) + manager.graphs[desc] = graph + manager._graphs_captured = True + + monkeypatch.setattr(gpu_cudagraph_utils.CudaGraphManager, "capture", capture) + manager = spec_cudagraph_utils.SpeculatorCudaGraphManager( + _create_decode_vllm_config( + capture_sizes=[4], + num_speculative_tokens=4, + dynamic_spec_schedule=[(1, 2, 4), (3, 4, 2), (5, 6, 1), (7, 8, 0)], + ), + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + decode_query_len=1, + ) + manager.capture( + forward, + MagicMock(), + MagicMock(), + MagicMock(), + [], + MagicMock(), + specialize_spec_tokens=specialize, + ) + captured_widths = [4, 2] if specialize else [4] + assert [c.kwargs["num_speculative_steps"] for c in forward.call_args_list] == ( + captured_widths + ) + + for k in (2, 4, 2, 3): + desc, _ = dispatch_cg_and_sync_dp( + manager, + num_reqs=3, + num_tokens=3, + uniform_token_count=1, + dp_size=1, + dp_rank=0, + ) + desc = manager.specialize_spec_tokens(desc, k) + captured_k = k if specialize else 4 + if captured_k in captured_widths: + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_tokens == 4 + manager.run_fullgraph(desc) + assert replayed_widths[-1] == captured_k + else: + assert desc.cg_mode == CUDAGraphMode.NONE + + @pytest.mark.parametrize( "decode_query_len,capture_sizes", [(1, [1, 2, 4, 8]), (2, [2, 4, 8, 16]), (8, [8, 16, 32, 64])], diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index a87c64725609..0a8cfd1a130e 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -374,19 +374,22 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch): ( "method_name", "cg_mode", + "num_speculative_steps", "expected_eager_calls", "expected_graph_replays", ), [ - ("_multi_step_decode", CUDAGraphMode.NONE, 3, 0), - ("_multi_step_decode", CUDAGraphMode.FULL, 0, 3), - ("_fused_multi_step_decode", CUDAGraphMode.NONE, 3, 0), - ("_fused_multi_step_decode", CUDAGraphMode.FULL, 0, 1), + ("_multi_step_decode", CUDAGraphMode.NONE, 4, 3, 0), + ("_multi_step_decode", CUDAGraphMode.FULL, 4, 0, 3), + ("_fused_multi_step_decode", CUDAGraphMode.NONE, 4, 3, 0), + ("_fused_multi_step_decode", CUDAGraphMode.FULL, 4, 0, 1), + ("_multi_step_decode", CUDAGraphMode.FULL, 3, 0, 2), ], ) def test_multi_step_decode_replays_captured_graph_as_expected( method_name, cg_mode, + num_speculative_steps, expected_eager_calls, expected_graph_replays, ): @@ -416,6 +419,7 @@ def test_multi_step_decode_replays_captured_graph_as_expected( batch_desc=batch_desc, seq_lens_cpu_upper_bound=None, num_tokens_across_dp=None, + num_speculative_steps=num_speculative_steps, ) assert generate_draft.call_count == expected_eager_calls diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 2e0c7271ce62..7586245a3543 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -46,6 +46,7 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(): finished_req_ids=set(), ec_connector_output=None, cudagraph_stats=None, + num_spec_tokens_to_schedule=0, ) runner.sample_tokens(None) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index bbde3295ec62..d8090de2a359 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -190,6 +190,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): ec_connector_output=None, routed_experts=None, cudagraph_stats=None, + num_spec_tokens_to_schedule=0, ) runner.req_states = SimpleNamespace() diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 6fadcbb35eea..e0a99d5e8f8e 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -116,9 +116,14 @@ class _StepRecorder: def __init__(self) -> None: # (blocks held per group, num_computed_tokens, num_scheduled_tokens) self.steps: list[tuple[list[int], int, int]] = [] + self.num_spec_tokens_to_schedule: list[int] = [] self._held: dict[str, list[int]] = {} def execute_model(self, scheduler_output) -> None: + if scheduler_output.total_num_scheduled_tokens: + self.num_spec_tokens_to_schedule.append( + scheduler_output.num_spec_tokens_to_schedule + ) for new_req in scheduler_output.scheduled_new_reqs: self._held[new_req.req_id] = [len(ids) for ids in new_req.block_ids] self._record(new_req.req_id, new_req.num_computed_tokens, scheduler_output) @@ -161,7 +166,7 @@ def _assert_covers_lookahead( # 0 covers eagle / MTP / draft models, 1 covers DFlash's extra in-fill query. @pytest.mark.parametrize("extra_lookahead", [0, 1]) -@pytest.mark.parametrize("num_spec_steps", [2, 3, 5, 7]) +@pytest.mark.parametrize("num_spec_steps", [1, 2, 3, 5, 7]) def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahead): num_lookahead_tokens = num_spec_steps + extra_lookahead recorder = _StepRecorder() @@ -173,6 +178,7 @@ def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahea ) _assert_covers_lookahead(recorder.steps, num_lookahead_tokens) + assert set(recorder.num_spec_tokens_to_schedule) == {num_spec_steps} def test_mixed_warmup_reserves_lookahead_blocks(): @@ -187,6 +193,7 @@ def test_mixed_warmup_reserves_lookahead_blocks(): ) _assert_covers_lookahead(recorder.steps, num_lookahead_tokens) + assert set(recorder.num_spec_tokens_to_schedule) == {NUM_SPEC_STEPS} @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 05fcf5f98f4c..73be62319a64 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -272,6 +272,7 @@ def _init_candidates(self) -> None: if ( speculative_config and speculative_config.uses_dynamic_speculative_decoding() + and self.decode_query_len > self.vllm_config.num_speculative_tokens ): # decode_query_len = num_speculative_steps + num_new_sampled_tokens # _per_step. Recover num_new_sampled_tokens_per_step diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b9880f43f242..b42289ef1111 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2057,6 +2057,7 @@ def execute_model( finished_req_ids=finished_req_ids, ec_connector_output=ec_connector_output, cudagraph_stats=cudagraph_stats, + num_spec_tokens_to_schedule=scheduler_output.num_spec_tokens_to_schedule, ) if not self.is_last_pp_rank: @@ -2086,6 +2087,7 @@ def sample_tokens( finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output cudagraph_stats = self.execute_model_state.cudagraph_stats + num_spec_tokens = self.execute_model_state.num_spec_tokens_to_schedule self.execute_model_state = None if not self.is_last_pp_rank: @@ -2237,9 +2239,12 @@ def sample_tokens( self.req_states.next_prefill_tokens, self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, + num_speculative_tokens=num_spec_tokens, dp_sync=dp_sync, mm_inputs=mm_inputs, ) + if num_spec_tokens < self.num_speculative_steps: + draft_tokens[:, num_spec_tokens:] = -1 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens if self.adaptive_verification is not None: self.adaptive_verification.record_confidences( @@ -2251,7 +2256,7 @@ def sample_tokens( # not have a speculator (i.e. self.speculator is None) self.draft_tokens_handler.set_draft_tokens( input_batch, - self.req_states.draft_tokens[input_batch.idx_mapping], + self.req_states.draft_tokens[input_batch.idx_mapping, :num_spec_tokens], ) if self.pp_handler is not None: self.pp_handler.broadcast_drafts( @@ -2396,6 +2401,7 @@ class ExecuteModelState(NamedTuple): finished_req_ids: set[str] ec_connector_output: ECConnectorOutput | None cudagraph_stats: CUDAGraphStat | None + num_spec_tokens_to_schedule: int class BatchReqState(NamedTuple): diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py index fe3bdef269ea..a7d3eb52b549 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from dataclasses import dataclass, replace import torch from vllm.config.compilation import CUDAGraphMode from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -17,6 +19,11 @@ from vllm.v1.worker.utils import AttentionGroup +@dataclass(frozen=True) +class SpeculatorBatchDescriptor(BatchExecutionDescriptor): + num_speculative_tokens: int = 0 + + class SpeculatorCudaGraphManager(CudaGraphManager): """CudaGraphManager for draft prefill and decode. @@ -27,6 +34,18 @@ class SpeculatorCudaGraphManager(CudaGraphManager): earlier capture would execute kernels with stale buffer contents. """ + def specialize_spec_tokens( + self, desc: BatchExecutionDescriptor, num_speculative_tokens: int + ) -> BatchExecutionDescriptor: + if desc.cg_mode != CUDAGraphMode.FULL or desc in self.graphs: + return desc + desc = SpeculatorBatchDescriptor( + **vars(desc), num_speculative_tokens=num_speculative_tokens + ) + return ( + desc if desc in self.graphs else replace(desc, cg_mode=CUDAGraphMode.NONE) + ) + def capture( self, forward_fn: Callable, @@ -36,7 +55,25 @@ def capture( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, progress_bar_desc: str = "Capturing CUDA graphs", + specialize_spec_tokens: bool = False, ) -> None: + if specialize_spec_tokens: + speculative_config = self.vllm_config.speculative_config + assert speculative_config is not None + dense_schedule = build_dynamic_sd_schedule_lookup( + speculative_config.num_speculative_tokens_per_batch_size, + vllm_max_batch_size=self.max_num_reqs, + vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens, + ) + self._capture_descs = { + mode: [ + SpeculatorBatchDescriptor(**vars(desc), num_speculative_tokens=k) + for desc in descs + for k in sorted(set(dense_schedule) - {0, 1}, reverse=True) + ] + for mode, descs in self._capture_descs.items() + } + def create_forward_fn( desc: BatchExecutionDescriptor, warmup: bool, @@ -60,6 +97,13 @@ def create_forward_fn( max_query_len=desc.max_query_len or desc.uniform_token_count, ) + forward_kwargs = {} + if self.decode_query_len == 1: + forward_kwargs["num_speculative_steps"] = ( + desc.num_speculative_tokens + if isinstance(desc, SpeculatorBatchDescriptor) + else self.vllm_config.num_speculative_tokens + ) return lambda cg_mode: forward_fn( num_reqs, num_tokens, @@ -67,6 +111,7 @@ def create_forward_fn( slot_mappings, num_tokens_across_dp, cg_mode, + **forward_kwargs, ) super().capture(create_forward_fn, progress_bar_desc) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 897e7af8de36..7242a840d249 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -190,6 +190,8 @@ def capture(self) -> None: self.attn_groups, self.kv_cache_config, progress_bar_desc="Capturing decode CUDA graphs", + specialize_spec_tokens=self.use_fused_multi_step_decode + and self.speculative_config.uses_dynamic_speculative_decoding(), ) self.on_multi_step_decode_end(self.max_num_reqs) @@ -220,14 +222,17 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + if num_speculative_tokens is None: + num_speculative_tokens = self.num_speculative_steps num_tokens = input_batch.num_tokens num_tokens_padded = input_batch.num_tokens_after_padding - num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() self.draft_max_seq_len = min( - max_seq_len + self.num_speculative_steps, self.max_model_len + max_seq_len + num_speculative_tokens, self.max_model_len ) # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the @@ -322,9 +327,10 @@ def propose( ) self.on_prefill_end(num_reqs) - if self.num_speculative_steps == 1: - # Early exit. - return self.draft_tokens[:num_reqs, :1] + if num_speculative_tokens <= 1: + if num_speculative_tokens == 0: + self.draft_tokens[:num_reqs].fill_(-1) + return self.draft_tokens[:num_reqs] if self.pcp_manager is not None and not dummy_run: self.block_tables.gather_block_tables( @@ -360,6 +366,10 @@ def propose( need_eager=is_profile, dp_sync=decode_batch_sync, ) + if self.decode_cudagraph_manager is not None: + decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens( + decode_batch_desc, num_speculative_tokens + ) num_tokens_across_dp = ( decode_batch_sync.num_tokens_across_dp if decode_batch_sync is not None @@ -379,6 +389,7 @@ def propose( decode_batch_desc, num_tokens_across_dp, input_batch.seq_lens_cpu_upper_bound, + num_speculative_tokens, ) self.on_multi_step_decode_end(num_reqs) @@ -505,6 +516,7 @@ def _multi_step_decode( batch_desc: BatchExecutionDescriptor, num_tokens_across_dp: torch.Tensor | None, seq_lens_cpu_upper_bound: torch.Tensor, + num_speculative_steps: int, ) -> None: positions = self.input_buffers.positions[:num_reqs] query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] @@ -512,7 +524,7 @@ def _multi_step_decode( attn_metadata = None slot_mappings_by_layer = None - for step in range(1, self.num_speculative_steps): + for step in range(1, num_speculative_steps): # Rebuild every step when positions advance, or just once # on the first step when positions are constant (Gemma4 MTP). if not skip_attn and (self.advance_draft_positions or step == 1): @@ -546,6 +558,7 @@ def _multi_step_decode( slot_mappings_by_layer, num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=batch_desc.cg_mode, + num_speculative_steps=num_speculative_steps, ) def _fused_multi_step_decode( @@ -555,6 +568,7 @@ def _fused_multi_step_decode( batch_desc: BatchExecutionDescriptor, num_tokens_across_dp: torch.Tensor | None, seq_lens_cpu_upper_bound: torch.Tensor, + num_speculative_steps: int, ) -> None: positions = self.input_buffers.positions[:num_reqs] query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] @@ -593,6 +607,7 @@ def _fused_multi_step_decode( slot_mappings_by_layer, num_tokens_across_dp, batch_desc.cg_mode, + num_speculative_steps, ) def _generate_fused_drafts( @@ -602,7 +617,8 @@ def _generate_fused_drafts( attn_metadata: dict[str, Any] | None, slot_mappings: dict[str, torch.Tensor] | None, num_tokens_across_dp: torch.Tensor | None, - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + cudagraph_runtime_mode: CUDAGraphMode, + num_speculative_steps: int, ) -> None: idx_mapping = self.idx_mapping[:num_reqs] positions = self.input_buffers.positions[:num_reqs] @@ -613,7 +629,7 @@ def _generate_fused_drafts( else [] ) - for step in range(1, self.num_speculative_steps): + for step in range(1, num_speculative_steps): self.current_draft_step.fill_(step) self._generate_draft( num_reqs, @@ -622,9 +638,10 @@ def _generate_fused_drafts( slot_mappings, num_tokens_across_dp, cudagraph_runtime_mode, + num_speculative_steps, ) if ( - step < self.num_speculative_steps - 1 + step < num_speculative_steps - 1 and attn_metadata is not None and self.advance_draft_positions ): @@ -644,7 +661,8 @@ def _generate_draft( attn_metadata: dict[str, Any] | None, slot_mappings: dict[str, torch.Tensor] | None, num_tokens_across_dp: torch.Tensor | None, - cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + cudagraph_runtime_mode: CUDAGraphMode, + num_speculative_steps: int, ) -> None: self._prepare_eplb_forward(num_reqs) @@ -682,7 +700,7 @@ def _generate_draft( self.sample_src_positions, num_reqs, self.max_model_len, - self.num_speculative_steps, + num_speculative_steps, advance_draft_positions=self.advance_draft_positions, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 57110e0b05f1..d59d87ae7e64 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -360,6 +360,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: num_reqs = input_batch.num_reqs num_target_tokens = input_batch.num_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py b/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py index 05bc997f7a34..5be7b0c86ddd 100644 --- a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py @@ -355,6 +355,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: assert self.model is not None diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index c411eeae550c..5afbe640b2f2 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -103,6 +103,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: del ( last_hidden_states, diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index c7026acdb88b..71f7f89416be 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -155,6 +155,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: num_reqs = input_batch.num_reqs seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 77a9ece3603e..111066328846 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -97,6 +97,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: pass diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index c29d9958fd15..894e2e5b6bd9 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -146,6 +146,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: ) decode_prefill_output = SchedulerOutput.make_empty() + decode_prefill_output.num_spec_tokens_to_schedule = ( + model_runner.num_speculative_steps + ) decode_prefill_output.scheduled_new_reqs = [ NewRequestData( req_id=decode_req_id, @@ -175,6 +178,7 @@ def _alloc_blocks(num_blocks: int) -> list[int]: ] mixed_output = SchedulerOutput.make_empty() + mixed_output.num_spec_tokens_to_schedule = model_runner.num_speculative_steps mixed_output.scheduled_cached_reqs = cached_decode_req mixed_output.scheduled_new_reqs = [ NewRequestData( @@ -353,6 +357,7 @@ def _alloc_blocks(num_blocks: int) -> list[int]: ] prefill_output = SchedulerOutput.make_empty() + prefill_output.num_spec_tokens_to_schedule = num_spec_steps prefill_output.scheduled_new_reqs = new_reqs prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs @@ -412,6 +417,7 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: step_spec_tokens[req_ids[i]] = [0] * num_spec_steps decode_output = SchedulerOutput.make_empty() + decode_output.num_spec_tokens_to_schedule = num_spec_steps decode_output.scheduled_cached_reqs = cached_req_data decode_output.num_scheduled_tokens = step_num_scheduled_tokens decode_output.scheduled_spec_decode_tokens = step_spec_tokens From 18c8a65eddebb1ac39021d87f059173656e4a77b Mon Sep 17 00:00:00 2001 From: Yida Weng <84164537+YidaWeng@users.noreply.github.com> Date: Tue, 6 Oct 2026 04:39:25 +0800 Subject: [PATCH 019/952] =?UTF-8?q?[BugFix][Frontend]=20Pass=20reasoning?= =?UTF-8?q?=5Fended=20through=20render=20=E2=86=92=20generate=20(#60059)?= =?UTF-8?q?=20(#60062)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Yida Weng <84164537+YidaWeng@users.noreply.github.com> Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/entrypoints/cohere/test_api_router.py | 34 +++++++- .../scale_out/render/test_render.py | 32 ++++++++ .../token_in_token_out/test_protocol.py | 8 ++ .../token_in_token_out/test_serving_tokens.py | 77 +++++++++++++++++++ vllm/entrypoints/cohere/api_router.py | 5 ++ .../openai/chat_completion/protocol.py | 24 +++++- .../openai/chat_completion/serving.py | 14 +--- vllm/entrypoints/scale_out/render/serving.py | 30 ++++++++ .../scale_out/token_in_token_out/protocol.py | 20 +++++ .../scale_out/token_in_token_out/serving.py | 6 ++ vllm/renderers/online_renderer.py | 21 +++-- 11 files changed, 249 insertions(+), 22 deletions(-) diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index 94324dbed76e..1bb2151621c0 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -32,8 +32,14 @@ AssistantMessageResponse, CohereChatV2Response, ) -from vllm.entrypoints.cohere.serving import CohereServingChatV2 -from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest +from vllm.entrypoints.cohere.serving import ( + _API_SERVER_ONLY_TEMPLATE_KWARGS, + CohereServingChatV2, +) +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + GenerateRequest, + ReasoningParserKwargs, +) from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse from vllm.entrypoints.serve.exception_handling.handlers.http import ( http_exception_handler, @@ -97,6 +103,12 @@ class _RenderChatHandler: def to_chat_completion_request(self, request): return CohereServingChatV2._convert_v2_to_chat_completion(request) + def _engine_chat_template_kwargs(self, chat_template_kwargs): + return CohereServingChatV2._engine_chat_template_kwargs( + self, # type: ignore[arg-type] + chat_template_kwargs, + ) + class _RenderHandler: """Minimal stand-in for :class:`ServingRender`. @@ -459,6 +471,24 @@ def test_cohere_request_surface_survives_conversion(self): # The conversion lower-cases safety_mode for the renderer. assert kwargs["safety_mode"] == "strict" + def test_api_server_only_kwargs_not_forwarded(self): + """As on ``/cohere/v2/chat``, the engine's reasoning parser must not + receive the template kwargs that only the API server reads.""" + result = _generate_request() + result.reasoning_parser_kwargs = ReasoningParserKwargs( + chat_template_kwargs={ + "safety_mode": "strict", + **dict.fromkeys(_API_SERVER_ONLY_TEMPLATE_KWARGS, "x"), + } + ) + app = _build_render_app(_RenderChatHandler(), _RenderHandler(result)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat/render", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.OK + assert r.json()["reasoning_parser_kwargs"] == { + "chat_template_kwargs": {"safety_mode": "strict"} + } + def test_501_when_chat_handler_missing(self): app = _build_render_app(None, _RenderHandler(_generate_request())) with TestClient(app) as client: diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 0daf4d405e5d..6c11a518d764 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -38,6 +38,7 @@ def _build_responses_serving_render() -> ServingRender: OnlineRenderer.validate_chat_template.__get__(serving.online_renderer) ) serving.online_renderer.trust_request_chat_template = True + serving.online_renderer.parser = None serving._check_model = AsyncMock(return_value=None) return serving @@ -135,6 +136,37 @@ async def test_render_responses_rejects_empty_token_ids(): assert response.error.message == "No token_ids rendered" +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_render_responses_carries_reasoning_parser_kwargs(): + """/v1/responses gives the engine's reasoning parser the effective template + kwargs, so the rendered request must carry them to generate.""" + serving = _build_responses_serving_render() + online = serving.online_renderer + online.parser = MagicMock(reasoning_parser_cls=object) + online.chat_template = None + online.chat_template_content_format = "auto" + online.default_chat_template_kwargs = {"enable_thinking": True} + online.effective_chat_template_kwargs = ( + OnlineRenderer.effective_chat_template_kwargs.__get__(online) + ) + online.render_responses = AsyncMock( + return_value=MagicMock(messages=[], engine_input={"prompt_token_ids": [7]}) + ) + request = ResponsesRequest( + model=MODEL_NAME, + input="Test prompt", + chat_template_kwargs={"enable_thinking": False}, + ) + + response = await serving.render_responses_request(request) + + assert not isinstance(response, ErrorResponse) + assert response.reasoning_parser_kwargs is not None + kwargs = response.reasoning_parser_kwargs.chat_template_kwargs + assert kwargs["enable_thinking"] is False + + @pytest.fixture(scope="module") def server(): args: list[str] = ["--trust-request-chat-template"] diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index 9894d0918059..6a57e3b2b475 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -124,6 +124,14 @@ def test_generate_request_rejects_placeholder_outside_prompt(): ) +def test_reasoning_parser_kwargs_rejects_unknown_keys(): + """The engine splats these into its reasoning parser's constructor, so a + client must not be able to pass arbitrary keyword arguments.""" + payload = {**_base_payload(), "reasoning_parser_kwargs": {"tokenizer": None}} + with pytest.raises(ValidationError, match="tokenizer"): + GenerateRequest.model_validate(payload) + + def test_output_mode_defaults_to_tokens(): assert GenerateRequest.model_validate(_base_payload()).output_mode == "tokens" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 2ee8f76e0f2e..1027dcef035c 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -589,6 +589,83 @@ async def test_generate_with_lora_adapter(client, tokenizer, messages): assert generate_res == completions_res +def _structured_chat_body(messages, **overrides) -> dict: + return { + "model": MODEL_NAME, + "messages": messages, + "max_tokens": 64, + "temperature": 0.0, + "chat_template_kwargs": {"enable_thinking": True}, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": { + "type": "object", + "properties": {"count": {"type": "integer"}}, + "required": ["count"], + }, + }, + }, + **overrides, + } + + +async def _render_then_generate(client, body) -> tuple[dict, list[int]]: + render_resp = await client.post("/v1/chat/completions/render", json=body) + render_resp.raise_for_status() + generate_request = render_resp.json() + generate_resp = await client.post(GEN_ENDPOINT, json=generate_request) + generate_resp.raise_for_status() + return generate_request, generate_resp.json()["choices"][0]["token_ids"] + + +# deepseek_v3 reads `enable_thinking` from the template kwargs and uses Qwen3's +# tokens, so the engine gates structured outputs differently when +# either reasoning field is dropped between render and generate. +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="the Rust frontend does not forward reasoning fields to the engine", +) +@pytest.mark.parametrize( + "server", [["--reasoning-parser", "deepseek_v3"]], indirect=True +) +async def test_render_then_generate_forwards_reasoning_parser_kwargs(client, messages): + """With thinking on, structured outputs must wait for the reasoning to + end on generate, the same as on chat.""" + body = _structured_chat_body(messages, return_token_ids=True) + chat_resp = await client.post("/v1/chat/completions", json=body) + chat_resp.raise_for_status() + + generate_request, token_ids = await _render_then_generate(client, body) + + assert generate_request["reasoning_ended"] is False + assert token_ids == chat_resp.json()["choices"][0]["token_ids"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="the Rust frontend does not forward reasoning fields to the engine", +) +@pytest.mark.parametrize( + "server", [["--reasoning-parser", "deepseek_v3"]], indirect=True +) +async def test_render_then_generate_forwards_reasoning_ended( + client, tokenizer, messages +): + """`include_reasoning=false` constrains from the first token on chat, so + generate must too, even though the engine's parser expects thinking.""" + body = _structured_chat_body(messages, include_reasoning=False) + + generate_request, token_ids = await _render_then_generate(client, body) + + assert generate_request["reasoning_ended"] is True + output = json.loads(tokenizer.decode(token_ids, skip_special_tokens=True)) + assert isinstance(output["count"], int) + + def _chat_prompt_token_ids(tokenizer, messages) -> list[int]: return tokenizer.apply_chat_template( messages, diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index eaac36f1336f..db6366df7121 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -207,6 +207,11 @@ async def chat_v2_render(request: CohereChatV2Request, raw_request: Request): if isinstance(result, ErrorResponse): return _error_response(result, raw_request) + if (kwargs := result.reasoning_parser_kwargs) is not None: + kwargs.chat_template_kwargs = handler._engine_chat_template_kwargs( + kwargs.chat_template_kwargs + ) + return JSONResponse(content=result.model_dump()) class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware): diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 6a70624be4fe..d2d1cbab93a6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -4,7 +4,7 @@ # Adapted from # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py import time -from typing import Annotated, Any, ClassVar, Literal +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal from openai.types.chat.chat_completion_audio import ( ChatCompletionAudio as OpenAIChatCompletionAudio, @@ -53,6 +53,9 @@ ) from vllm.utils import random_uuid +if TYPE_CHECKING: + from vllm.parser.abstract_parser import Parser + logger = init_logger(__name__) @@ -566,6 +569,25 @@ def _materialize_tool_calls_after(self) -> "ChatCompletionRequest": _grammar_from_parser: bool = PrivateAttr(default=False) """CAUTION: Should only be set by the parser-engine adapter's adjust_request.""" + def resolve_reasoning_ended( + self, parser: "Parser | None", prompt_token_ids: list[int] + ) -> bool | None: + """Resolve the engine's `reasoning_ended` for structured outputs. + + Call after `adjust_request`, which may set `_grammar_from_parser`. + `True` constrains from the first generated token; `None` lets the + engine check the prompt with its own reasoning parser. + """ + if not self.include_reasoning: + return True + if self._grammar_from_parser: + # The Mistral grammar already includes an optional `think?` + # rule that handles both reasoning and non-reasoning outputs. + return True + if parser is not None and parser.reasoning_parser is not None: + return parser.is_reasoning_end(prompt_token_ids) + return None + def build_chat_params( self, default_template: str | None, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 400089c96364..361cc52f72d2 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -364,17 +364,9 @@ async def _create_chat_completion( session_id=session_id, ) else: - if not request.include_reasoning: - reasoning_ended = True - elif request._grammar_from_parser: - # The Mistral grammar already includes an optional - # `think?` rule that handles both reasoning and - # non-reasoning outputs. - reasoning_ended = True - elif parser is not None and parser.reasoning_parser is not None: - reasoning_ended = parser.is_reasoning_end(prompt_token_ids or []) - else: - reasoning_ended = None + reasoning_ended = request.resolve_reasoning_ended( + parser, prompt_token_ids or [] + ) generator = self.engine_client.generate( engine_input, diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index 117ede0f9b36..bddbe90f286f 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -17,6 +17,7 @@ from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, MultiModalFeatures, + ReasoningParserKwargs, ) from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.engine.serving import BaseServing @@ -124,6 +125,20 @@ async def render_chat_request( ) params = request.to_sampling_params(max_tokens, self.default_sampling_params) + # Resolve here what OpenAIServingChat passes to the engine: + # `_grammar_from_parser` (set by `adjust_request`) is not serialized. + reasoning_parser_kwargs = self._reasoning_parser_kwargs(request) + parser = None + if reasoning_parser_kwargs is not None: + assert self.online_renderer.parser is not None + parser = self.online_renderer.parser( + self.online_renderer.renderer.get_tokenizer(), + request.tools, + chat_template_kwargs=reasoning_parser_kwargs.chat_template_kwargs, + model_config=self.model_config, + ) + reasoning_ended = request.resolve_reasoning_ended(parser, token_ids) + request_id = f"chatcmpl-{random_uuid()}" return GenerateRequest( @@ -132,6 +147,8 @@ async def render_chat_request( features=self._extract_mm_features(engine_input), sampling_params=params, model=request.model, + reasoning_ended=reasoning_ended, + reasoning_parser_kwargs=reasoning_parser_kwargs, stream=bool(request.stream), stream_options=(request.stream_options if request.stream else None), cache_salt=request.cache_salt, @@ -263,6 +280,7 @@ async def render_responses_request( features=self._extract_mm_features(engine_input), sampling_params=params, model=request.model, + reasoning_parser_kwargs=self._reasoning_parser_kwargs(request), stream=bool(request.stream), cache_salt=request.cache_salt, priority=request.priority, @@ -271,6 +289,18 @@ async def render_responses_request( token_offsets=engine_input.get("prompt_token_offsets"), ) + def _reasoning_parser_kwargs( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ReasoningParserKwargs | None: + parser_cls = self.online_renderer.parser + if parser_cls is None or parser_cls.reasoning_parser_cls is None: + return None + return ReasoningParserKwargs( + chat_template_kwargs=self.online_renderer.effective_chat_template_kwargs( + request + ) + ) + def _placeholder_metadata_fields(self, modality: str) -> set[str]: """Fields the EC consumer still requires after embeddings transfer.""" if self._placeholder_metadata_parser_failed: diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 31bd7e2eaef9..4e29f7e59dfa 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -4,6 +4,7 @@ from pydantic import ( BaseModel, + ConfigDict, Discriminator, Field, NonNegativeInt, @@ -147,6 +148,16 @@ def _validate_parallel_fields(self) -> "MultiModalFeatures": return self +class ReasoningParserKwargs(BaseModel): + """Kwargs for the engine-side reasoning parser that gates structured + outputs. Typed so clients cannot pass arbitrary constructor kwargs.""" + + model_config = ConfigDict(extra="forbid") + + chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) + """The effective chat template kwargs the prompt was rendered with.""" + + class GenerateRequest(BaseModel): request_id: str = Field( default_factory=lambda: f"{random_uuid()}", @@ -212,6 +223,15 @@ def _require_ec_for_metadata_only(self) -> "GenerateRequest": model: str | None = None + reasoning_ended: bool | None = None + """Whether reasoning has ended before the first generated token, as + resolved by /render. `True` applies structured outputs from the first + token; `None` lets the engine check the prompt with its reasoning parser.""" + + reasoning_parser_kwargs: ReasoningParserKwargs | None = None + """Set by /render when it has a reasoning parser, so the engine-side + parser agrees with the frontend on flags such as `enable_thinking`.""" + return_token_ids: bool | None = Field( default=None, description=( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 6b0195b10426..e7afa0401b29 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -349,6 +349,12 @@ async def serve_tokens( priority=request.priority, data_parallel_rank=data_parallel_rank, session_id=session_id, + reasoning_ended=request.reasoning_ended, + reasoning_parser_kwargs=( + request.reasoning_parser_kwargs.model_dump() + if request.reasoning_parser_kwargs is not None + else None + ), ) assert result_generator is not None diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index 23352ef1cf69..e8220938b26a 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -209,6 +209,18 @@ def warmup(self) -> None: ) ) + def effective_chat_template_kwargs( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> dict[str, Any]: + return ( + request.build_chat_params( + self.chat_template, + self.chat_template_content_format, + ) + .with_defaults(self.default_chat_template_kwargs) + .chat_template_kwargs + ) + async def render_chat( self, request: ChatCompletionRequest, @@ -356,14 +368,7 @@ async def render_responses( else None ), ) - chat_template_kwargs = ( - request.build_chat_params( - self.chat_template, - self.chat_template_content_format, - ) - .with_defaults(self.default_chat_template_kwargs) - .chat_template_kwargs - ) + chat_template_kwargs = self.effective_chat_template_kwargs(request) _, engine_inputs = await self.preprocess_chat( request, messages, From 8542c3eff5809a79d6d825958f39187c54007851 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 5 Oct 2026 17:17:00 -0400 Subject: [PATCH 020/952] [CI] Fix test_mixed_warmup_gate after #57053 (#60116) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- tests/v1/worker/test_mixed_warmup_gate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 84ddf692a131..ba532c2c06b4 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -45,6 +45,7 @@ def test_mixed_warmup_disables_watermarking(): ), vllm_config=SimpleNamespace(num_lookahead_tokens=0), max_model_len=128, + num_speculative_steps=0, model_state=SimpleNamespace(max_encoder_len=0), kv_connector=SimpleNamespace(set_disabled=lambda disabled: None), ) From 2e3154aa18b3d2df2253c3e55f8b32fe9af5cfe8 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:20:32 -0400 Subject: [PATCH 021/952] [HiSparse] Log steady-state max concurrency and expose host-tier utilization gauges (#58949) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 Co-authored-by: Matthew Bonanni --- docs/design/hisparse.md | 8 ++ docs/mkdocs/gen_files/generate_metrics.py | 4 + docs/usage/metrics.md | 9 ++ tests/v1/core/test_kv_cache_utils.py | 43 ++++++++ .../unit/test_hisparse_connector.py | 36 +++++++ .../kv_connector/unit/test_hisparse_stats.py | 81 ++++++++++++++- tests/v1/worker/test_utils.py | 2 + .../kv_connector/v1/hisparse/connector.py | 33 ++++++- .../kv_connector/v1/hisparse/stats.py | 98 +++++++++++++++---- vllm/v1/hisparse/coordinator.py | 7 +- vllm/v1/hisparse/layout.py | 32 ++++++ vllm/v1/hisparse/types.py | 4 + 12 files changed, 323 insertions(+), 34 deletions(-) diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md index 8651b3a06dfd..f5fab3c0f5d2 100644 --- a/docs/design/hisparse.md +++ b/docs/design/hisparse.md @@ -37,6 +37,14 @@ memory consumption is therefore topology- and implementation-dependent. The realized capacity may be slightly smaller because the budget is rounded down to complete host blocks. +Startup logs report two concurrency bounds at `max_model_len`. The generic +`Maximum concurrency` line charges each request its full admission footprint, +including the in-flight window of every resident group. The `HiSparse +steady-state maximum concurrency` line charges running requests that read from +host only their active tail pages, plus one request being admitted at its full +footprint. Host-pool metrics are listed in +[Metrics](../usage/metrics.md#hisparse-kv-connector-metrics). + ## Ownership | Thing | Owner | What “owner” means | diff --git a/docs/mkdocs/gen_files/generate_metrics.py b/docs/mkdocs/gen_files/generate_metrics.py index 4d4fa93cc580..32bbefa88c29 100644 --- a/docs/mkdocs/gen_files/generate_metrics.py +++ b/docs/mkdocs/gen_files/generate_metrics.py @@ -23,6 +23,10 @@ "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py", "key": "metrics-nixl", }, + { + "path": "vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py", + "key": "metrics-hisparse", + }, { "path": "vllm/v1/simple_kv_offload/metrics.py", "key": "metrics-simple-kv-offload", diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 5cb5eebc5582..b87ff1c28fee 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -45,6 +45,15 @@ The following metrics are exposed: --8<-- "gen:metrics-nixl" +## HiSparse KV Connector Metrics + +These metrics are exposed when the `HiSparseConnector` KV connector is +configured. The host-pool gauges are reported by the scheduler once per engine +step, for the host pool shared by all tensor-parallel ranks; evictable cached +blocks count as free. + +--8<-- "gen:metrics-hisparse" + ## Simple CPU Offload Connector Metrics These metrics are exposed when the `SimpleCPUOffloadConnector` KV connector diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 08db9a286b24..7a34d07c07c6 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -60,6 +60,7 @@ from vllm.v1.hisparse.layout import ( create_hisparse_layout, get_hisparse_kv_cache_groups, + get_hisparse_steady_state_concurrency, ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, @@ -198,6 +199,48 @@ def test_hisparse_hma_uses_resolved_gpu_block_size( assert scheduler_block_size == hash_block_size == gpu_block_size +def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): + """Running requests pin only their resident tail; the newest still needs its + full in-flight window to be admitted.""" + block_size = 16 + config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=10 * block_size), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + max_in_flight_tokens=10 * block_size, + ) + attn_spec = FullAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=64, dtype=torch.float16 + ) + page_size = attn_spec.page_size_bytes + kv_cache_config = KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["source"], attn_spec, host_resident=True), + KVCacheGroupSpec(["indexer"], attn_spec), + KVCacheGroupSpec( + ["resident"], + HiSparseResidentSpec(block_size=block_size, page_size=page_size), + ), + KVCacheGroupSpec( + ["hot"], + HiSparseHotSpec( + block_size=block_size, page_size=page_size, blocks_per_request=4 + ), + ), + ], + hisparse_host_num_blocks=100, + ) + + # Admission takes indexer 10 + resident 10 + hot 4 = 24 blocks and steady + # state 16, so 3 fit (2 x 16 + 24 <= 64) but not 4 (3 x 16 + 24 > 64). + assert get_hisparse_steady_state_concurrency(config, kv_cache_config) == 3.5 + # The worst-case bound charges all 10 resident pages. + assert get_max_concurrency_for_kv_cache_config( + config, kv_cache_config + ) == pytest.approx(64 / 24) + + @pytest.mark.parametrize("extra_blocks,ok", [(1, False), (2, True)]) def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok): """Every page needs a host block, so one max_model_len request must fit.""" diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index a2eeb8cd63b1..f08982a5e924 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -64,6 +64,42 @@ def test_hisparse_requires_block_outermost_device_layout(): assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC" +def test_scheduler_stats_report_host_pool_usage(): + """The scheduler-side stats sample the coordinator's host pool level.""" + from tests.v1.core.test_prefix_caching import ( + HISPARSE_BLOCK_SIZE, + make_hisparse_kv_cache_manager, + make_request, + sha256, + ) + from vllm.v1.core.kv_cache_utils import init_none_hash + + init_none_hash(sha256) + + manager = make_hisparse_kv_cache_manager(32, host_num_blocks=8) + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(get_hisparse_coordinator(manager)) + coordinator = scheduler.coordinator + + stats = scheduler.get_kv_connector_stats() + assert stats.data["host_cache_usage_perc"] == [0.0] + assert stats.data["pending_page_transfers"] == [0] + + request = make_request( + "request", list(range(4 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + assert manager.allocate_slots(request, num_new_tokens=64) is not None + host_blocks = coordinator.host_manager.req_to_blocks[request.request_id] + num_used = sum(not block.is_null for block in host_blocks) + + # The pool reserves the null block, so 8 configured blocks yield 7 usable. + expected = pytest.approx(num_used / 7) + assert 0 < num_used <= 7 + for _ in range(2): + stats = scheduler.get_kv_connector_stats() + assert stats.data["host_cache_usage_perc"] == [expected] + + def test_no_forward_enqueues_deferred_hisparse_transfers(): """A zero-token step must still enqueue deferred post-forward transfers.""" connector = object.__new__(ActiveKVConnector) diff --git a/tests/v1/kv_connector/unit/test_hisparse_stats.py b/tests/v1/kv_connector/unit/test_hisparse_stats.py index 09904d0a3caf..91d00b1e2260 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_stats.py +++ b/tests/v1/kv_connector/unit/test_hisparse_stats.py @@ -26,6 +26,28 @@ def test_record_snapshot_and_reduce(): assert reduced["HiSparse host-to-device bytes"] == 64 +def test_record_host_usage_reports_latest_level(): + stats = HiSparseKVConnectorStats() + stats.record_host_usage(usage=0.4, pending_page_transfers=2) + stats.record_host_usage(usage=0.8, pending_page_transfers=0) + + reduced = stats.reduce() + assert reduced["HiSparse host KV cache usage %"] == 80.0 + + +def test_aggregate_keeps_latest_host_usage_level(): + first = HiSparseKVConnectorStats() + first.record_host_usage(usage=0.4, pending_page_transfers=2) + second = HiSparseKVConnectorStats() + second.record_host_usage(usage=0.8, pending_page_transfers=0) + second.record_host_usage(usage=0.6, pending_page_transfers=1) + + first.aggregate(second) + + assert first.data["host_cache_usage_perc"] == [0.4, 0.6] + assert first.data["pending_page_transfers"] == [2, 1] + + def test_aggregate_extends_snapshot_deltas(): first = HiSparseKVConnectorStats() first.record_snapshot(hits=7, misses=3, host_to_device_bytes=48) @@ -34,11 +56,9 @@ def test_aggregate_extends_snapshot_deltas(): first.aggregate(second) - assert first.data == { - "cache_hits": [7, 5], - "cache_misses": [3, 1], - "host_to_device_bytes": [48, 16], - } + assert first.data["cache_hits"] == [7, 5] + assert first.data["cache_misses"] == [3, 1] + assert first.data["host_to_device_bytes"] == [48, 16] def test_aggregate_skips_empty_stats(): @@ -60,3 +80,54 @@ def test_build_kv_connector_stats_round_trip(): assert rebuilt is not None assert isinstance(rebuilt, HiSparseKVConnectorStats) assert rebuilt.reduce() == stats.reduce() + + +def test_prom_metrics_observe_host_usage_gauges(): + from types import SimpleNamespace + from typing import Any + + from prometheus_client import Counter, Gauge, Histogram + + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( + HiSparsePromMetrics, + ) + + class _FakeMetric: + def __init__(self, **kwargs: Any): + self.kwargs = kwargs + self.increments: list[int | float] = [] + self.set_values: list[int | float] = [] + + def labels(self, *labelvalues: object) -> "_FakeMetric": + return self + + def inc(self, value: int | float) -> None: + self.increments.append(value) + + def set(self, value: int | float) -> None: + self.set_values.append(value) + + created: dict[str, _FakeMetric] = {} + + class _NamedFake(_FakeMetric): + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + created[kwargs["name"]] = self + + prom = HiSparsePromMetrics( + vllm_config=SimpleNamespace(kv_transfer_config=None), + metric_types={Gauge: _NamedFake, Counter: _NamedFake, Histogram: _NamedFake}, + labelnames=["model_name"], + per_engine_labelvalues={0: ["model"]}, + ) + + stats = HiSparseKVConnectorStats() + stats.record_snapshot(hits=3, misses=2, host_to_device_bytes=32) + stats.record_host_usage(usage=0.5, pending_page_transfers=1) + stats.record_host_usage(usage=0.75, pending_page_transfers=3) + prom.observe(stats.to_dict()) + + assert created["vllm:hisparse_cache_hits"].increments == [3] + assert created["vllm:hisparse_host_to_device_bytes"].increments == [32] + assert created["vllm:hisparse_host_cache_usage_perc"].set_values == [0.75] + assert created["vllm:hisparse_pending_page_transfers"].set_values == [3] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index abcbdb24f2f6..f6f7227954bf 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -92,6 +92,8 @@ def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeyp "cache_hits": [12], "cache_misses": [4], "host_to_device_bytes": [64], + "host_cache_usage_perc": [], + "pending_page_transfers": [], } diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index dabf9095dbd0..7e8960f644a9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -30,10 +30,12 @@ PromMetric, PromMetricT, ) +from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.hisparse.coordinator import get_hisparse_coordinator +from vllm.v1.hisparse.layout import get_hisparse_steady_state_concurrency from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror from vllm.v1.outputs import KVConnectorOutput @@ -44,6 +46,8 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request +logger = init_logger(__name__) + @dataclass class HiSparseConnectorMetadata(KVConnectorMetadata): @@ -99,6 +103,17 @@ def bind_coordinator(self, coordinator: HiSparseCoordinator) -> None: assert self.coordinator is None self.coordinator = coordinator + def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None: + if self.coordinator is None: + return None + host_pool = self.coordinator.get_host_block_pool() + assert host_pool is not None + stats = HiSparseKVConnectorStats() + stats.record_host_usage( + host_pool.get_usage(), len(self.coordinator.pending_spills) + ) + return stats + def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> HiSparseConnectorMetadata: @@ -223,6 +238,16 @@ def __init__( ), draft_kv_lookahead=vllm_config.num_lookahead_tokens, ) + max_model_len = vllm_config.model_config.max_model_len + steady_concurrency = get_hisparse_steady_state_concurrency( + vllm_config, kv_cache_config + ) + logger.info_once( + "HiSparse steady-state maximum concurrency for %s tokens per " + "request: %.2fx (running requests reading from host).", + f"{max_model_len:,}", + steady_concurrency, + ) elif role == KVConnectorRole.WORKER: self.connector_worker = HiSparseConnectorWorker( vllm_config, kv_cache_config @@ -263,9 +288,11 @@ def reset_capture_state(self) -> None: self.connector_worker.reset_hot_state() def get_kv_connector_stats(self) -> KVConnectorStats | None: - if self.connector_worker is None: - return None - return self.connector_worker.get_kv_connector_stats() + if self.connector_worker is not None: + return self.connector_worker.get_kv_connector_stats() + if self.connector_scheduler is not None: + return self.connector_scheduler.get_kv_connector_stats() + return None @classmethod def build_kv_connector_stats( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py index f2abf51fc4e9..6f6cfced9ede 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py @@ -14,14 +14,7 @@ ) from vllm.v1.metrics.utils import create_metric_per_engine -_HISPARSE_COUNTERS: tuple[tuple[str, str], ...] = ( - ("cache_hits", "Number of HiSparse device hot-buffer hits."), - ("cache_misses", "Number of HiSparse device hot-buffer misses."), - ( - "host_to_device_bytes", - "Bytes transferred from host KV storage to HiSparse hot buffers.", - ), -) +_HISPARSE_LEVEL_KEYS = frozenset({"host_cache_usage_perc", "pending_page_transfers"}) @dataclass @@ -30,6 +23,8 @@ class HiSparseKVConnectorStats(KVConnectorStats): Each list entry is the delta recorded over one device counter snapshot interval, so a list can hold multiple snapshots per logging interval. + The host-tier gauges are level values, not deltas: aggregation keeps + only the most recent observation. """ def __post_init__(self): @@ -39,10 +34,12 @@ def __post_init__(self): def reset(self): # Must be serializable - self.data: dict[str, list[int]] = { + self.data: dict[str, list[int | float]] = { "cache_hits": [], "cache_misses": [], "host_to_device_bytes": [], + "host_cache_usage_perc": [], + "pending_page_transfers": [], } def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): @@ -50,21 +47,32 @@ def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): self.data["cache_misses"].append(misses) self.data["host_to_device_bytes"].append(host_to_device_bytes) + def record_host_usage(self, usage: float, pending_page_transfers: int) -> None: + self.data["host_cache_usage_perc"].append(usage) + self.data["pending_page_transfers"].append(pending_page_transfers) + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: if not other.is_empty(): for k, v in other.data.items(): accumulator = self.data[k] assert isinstance(accumulator, list) - accumulator.extend(v) + if k in _HISPARSE_LEVEL_KEYS: + accumulator.extend(v[-1:]) + else: + accumulator.extend(v) return self def reduce(self) -> dict[str, int | float]: # Compute compact representative stats suitable for CLI logging. - return { + reduced: dict[str, int | float] = { "HiSparse hot-buffer hits": sum(self.data["cache_hits"]), "HiSparse hot-buffer misses": sum(self.data["cache_misses"]), "HiSparse host-to-device bytes": sum(self.data["host_to_device_bytes"]), } + if self.data["host_cache_usage_perc"]: + usage = self.data["host_cache_usage_perc"][-1] + reduced["HiSparse host KV cache usage %"] = round(100 * usage, 1) + return reduced def is_empty(self) -> bool: return all(len(values) == 0 for values in self.data.values()) @@ -80,18 +88,66 @@ def __init__( ): super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) - self.hisparse_counters: dict[str, Any] = {} - for name, documentation in _HISPARSE_COUNTERS: - counter = self._counter_cls( - name=f"vllm:hisparse_{name}", - documentation=documentation, - labelnames=labelnames, - ) - self.hisparse_counters[name] = create_metric_per_engine( - counter, self.per_engine_labelvalues - ) + counter_cache_hits = self._counter_cls( + name="vllm:hisparse_cache_hits", + documentation="Number of HiSparse device hot-buffer hits.", + labelnames=labelnames, + ) + counter_cache_misses = self._counter_cls( + name="vllm:hisparse_cache_misses", + documentation="Number of HiSparse device hot-buffer misses.", + labelnames=labelnames, + ) + counter_host_to_device_bytes = self._counter_cls( + name="vllm:hisparse_host_to_device_bytes", + documentation=( + "Bytes transferred from host KV storage to HiSparse hot buffers." + ), + labelnames=labelnames, + ) + self.hisparse_counters: dict[str, Any] = { + "cache_hits": create_metric_per_engine( + counter_cache_hits, self.per_engine_labelvalues + ), + "cache_misses": create_metric_per_engine( + counter_cache_misses, self.per_engine_labelvalues + ), + "host_to_device_bytes": create_metric_per_engine( + counter_host_to_device_bytes, self.per_engine_labelvalues + ), + } + + gauge_host_cache_usage = self._gauge_cls( + name="vllm:hisparse_host_cache_usage_perc", + documentation=( + "HiSparse host KV-cache usage. 1 means 100 percent usage. " + "Evictable cached blocks count as free." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + gauge_pending_page_transfers = self._gauge_cls( + name="vllm:hisparse_pending_page_transfers", + documentation=( + "HiSparse page transfers (write-backs and restores) not yet completed." + ), + multiprocess_mode="mostrecent", + labelnames=labelnames, + ) + self.hisparse_gauges: dict[str, Any] = { + "host_cache_usage_perc": create_metric_per_engine( + gauge_host_cache_usage, self.per_engine_labelvalues + ), + "pending_page_transfers": create_metric_per_engine( + gauge_pending_page_transfers, self.per_engine_labelvalues + ), + } def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): for name, counter in self.hisparse_counters.items(): for value in transfer_stats_data.get(name, []): counter[engine_idx].inc(value) + for name, gauge in self.hisparse_gauges.items(): + values = transfer_stats_data.get(name, []) + if values: + gauge[engine_idx].set(values[-1]) diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index bc9bff1a8b00..766c926b12b7 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -19,6 +19,7 @@ SingleTypeKVCacheManager, ) from vllm.v1.hisparse.types import ( + ACTIVE_TAIL_PAGES, SparseKVOffloadCommand, SparseKVPageTransfer, SparseKVRowMirror, @@ -32,10 +33,6 @@ if TYPE_CHECKING: from vllm.v1.core.kv_cache_manager import KVCacheManager -# Sealed pages this many positions behind the block-table tail stay pinned so -# a page written by an in-flight step is never handed out under it. -_ACTIVE_TAIL_PAGES = 2 - @dataclass class _PendingPublication: @@ -343,7 +340,7 @@ def _resident_page_blocks( req_blocks = manager.req_to_blocks.get(request_id) if ( req_blocks is None - or page_idx >= len(req_blocks) - _ACTIVE_TAIL_PAGES + or page_idx >= len(req_blocks) - ACTIVE_TAIL_PAGES or req_blocks[page_idx].is_null ): return None diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 8ebc63c0edbb..98636866f8ac 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -14,6 +14,7 @@ get_hisparse_host_block_stride, use_shared_hisparse_host_pool, ) +from vllm.v1.hisparse.types import ACTIVE_TAIL_PAGES from vllm.v1.kv_cache_interface import ( HiSparseHotSpec, HiSparseResidentSpec, @@ -373,3 +374,34 @@ def get_hisparse_kv_cache_config( vllm_config.cache_config.prefix_cache_retention_interval ), ) + + +def get_hisparse_steady_state_concurrency( + vllm_config: VllmConfig, kv_cache_config: KVCacheConfig +) -> float: + """Max concurrency at max_model_len once running requests read from host. + + A request reading from host pins only its active tail of each resident + group, but admitting one still takes its full in-flight window, so the + bound is all-but-one requests at steady state plus one being admitted. + """ + admission_blocks = 0 + steady_blocks = 0 + host_blocks = 0 + for group in kv_cache_config.kv_cache_groups: + spec = group.kv_cache_spec + required = cdiv(spec.max_memory_usage_bytes(vllm_config), spec.page_size_bytes) + if group.host_resident: + host_blocks += required + continue + admission_blocks += required + if isinstance(spec, HiSparseResidentSpec): + required = min(required, ACTIVE_TAIL_PAGES) + steady_blocks += required + assert kv_cache_config.hisparse_host_num_blocks is not None + num_blocks = kv_cache_config.num_blocks + if num_blocks < admission_blocks: + gpu_concurrency = num_blocks / admission_blocks + else: + gpu_concurrency = 1 + (num_blocks - admission_blocks) / steady_blocks + return min(gpu_concurrency, kv_cache_config.hisparse_host_num_blocks / host_blocks) diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py index 269793ea8521..db0ee06e9399 100644 --- a/vllm/v1/hisparse/types.py +++ b/vllm/v1/hisparse/types.py @@ -3,6 +3,10 @@ from dataclasses import dataclass +# Sealed pages this many positions behind the block-table tail stay pinned so +# a page written by an in-flight step is never handed out under it. +ACTIVE_TAIL_PAGES = 2 + @dataclass(frozen=True) class SparseKVPageTransfer: From 3fb1cd796bfd318d3261844237a88c371704829d Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 5 Oct 2026 22:57:45 +0100 Subject: [PATCH 022/952] [Frontend] Extract parse and assembly out of batch chat derender (#60073) Signed-off-by: Martin Hickey Co-authored-by: Flora Feng <4florafeng@gmail.com> --- vllm/renderers/online_derenderer.py | 184 +++++++++++++++------------- 1 file changed, 101 insertions(+), 83 deletions(-) diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index d221b32a8b32..1c6875159a9f 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -151,6 +151,12 @@ def _derender_chat( skip_special_tokens=skip_special, ) + chat_template_kwargs = ( + self._resolve_chat_template_kwargs(chat_request) + if has_parser and chat_request is not None + else {} + ) + for choice in generate_response.choices: if not choice.token_ids: raise ValueError(f"choice {choice.index} has empty or null token_ids") @@ -161,83 +167,29 @@ def _derender_chat( else None ) - if has_parser: - assert self.parser is not None and chat_request is not None - # Parser path: decode with special tokens preserved - # so the parser can see markers like , - # , or Harmony channel tokens. - decoded_text, _ = self._detokenize_delta( - tokenizer, - choice.token_ids, - seed_state, - skip_special_tokens=skip_special, - spaces_between_special_tokens=spaces_between, - ) - - chat_template_kwargs: dict[str, Any] = {} - if not self.use_harmony: - chat_template_kwargs = ( - chat_request.build_chat_params( - self.chat_template, - self.chat_template_content_format, - ) - .with_defaults(self.default_chat_template_kwargs) - .chat_template_kwargs - ) + # With a parser, special tokens are preserved so it can see + # markers like , , or Harmony channel tokens. + # Without one, the request's skip_special_tokens is honoured + # (default True when no request was given). + decoded_text, _ = self._detokenize_delta( + tokenizer, + choice.token_ids, + seed_state, + skip_special_tokens=skip_special, + spaces_between_special_tokens=spaces_between, + ) - parser = self.parser( + if has_parser: + assert chat_request is not None + message = self._parse_and_assemble( tokenizer, - chat_request.tools, - chat_template_kwargs=chat_template_kwargs, - model_config=self.model_config, - ) - if generate_response.prompt_token_ids is not None: - parser.set_prompt_token_ids(generate_response.prompt_token_ids) - reasoning, content, tool_calls = parser.parse( decoded_text, + choice.token_ids, chat_request, - enable_auto_tools=self.enable_auto_tools, - model_output_token_ids=choice.token_ids, - ) - - if not getattr(chat_request, "include_reasoning", True): - reasoning = None - - tc_items = ( - [ - ToolCall( - id=random_uuid(), - function=tc, - ) - for tc in tool_calls - ] - if tool_calls - else [] - ) - - is_named_tool_choice = ( - type(chat_request.tool_choice) is ChatCompletionNamedToolChoiceParam - ) - is_required_tool_choice = chat_request.tool_choice == "required" - if is_named_tool_choice or is_required_tool_choice: - content = content or "" - - message = ChatMessage( - role="assistant", - reasoning=reasoning, - content=content, - tool_calls=tc_items, + chat_template_kwargs, + generate_response.prompt_token_ids, ) else: - # No parser: plain detokenization honouring the request's - # skip_special_tokens (default True when no request was given). - decoded_text, _ = self._detokenize_delta( - tokenizer, - choice.token_ids, - seed_state, - skip_special_tokens=skip_special, - spaces_between_special_tokens=spaces_between, - ) message = ChatMessage(role="assistant", content=decoded_text) choices.append( @@ -251,6 +203,83 @@ def _derender_chat( return choices + def _resolve_chat_template_kwargs( + self, chat_request: ChatCompletionRequest + ) -> dict[str, Any]: + if self.use_harmony: + return {} + return ( + chat_request.build_chat_params( + self.chat_template, + self.chat_template_content_format, + ) + .with_defaults(self.default_chat_template_kwargs) + .chat_template_kwargs + ) + + def _parse_and_assemble( + self, + tokenizer: TokenizerLike, + text: str, + token_ids: Sequence[int], + chat_request: ChatCompletionRequest, + chat_template_kwargs: dict[str, Any], + prompt_token_ids: list[int] | None, + ) -> ChatMessage: + """Parse decoded output text into an assistant `ChatMessage`. + + Args: + tokenizer: Tokenizer handed to the parser. + text: Decoded output text to parse. + token_ids: Output token IDs the text was decoded from. + chat_request: Request supplying tools, `tool_choice` and + `include_reasoning`. + chat_template_kwargs: Already resolved template kwargs. + prompt_token_ids: Prompt token IDs, if known. + + Returns: + The assembled assistant message. + + """ + assert self.parser is not None + parser = self.parser( + tokenizer, + chat_request.tools, + chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, + ) + if prompt_token_ids is not None: + parser.set_prompt_token_ids(prompt_token_ids) + reasoning, content, tool_calls = parser.parse( + text, + chat_request, + enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=token_ids, + ) + + if not getattr(chat_request, "include_reasoning", True): + reasoning = None + + tc_items = ( + [ToolCall(id=random_uuid(), function=tc) for tc in tool_calls] + if tool_calls + else [] + ) + + is_named_tool_choice = ( + type(chat_request.tool_choice) is ChatCompletionNamedToolChoiceParam + ) + is_required_tool_choice = chat_request.tool_choice == "required" + if is_named_tool_choice or is_required_tool_choice: + content = content or "" + + return ChatMessage( + role="assistant", + reasoning=reasoning, + content=content, + tool_calls=tc_items, + ) + def _detokenize_delta( self, tokenizer: TokenizerLike, @@ -497,21 +526,10 @@ def _derender_chat_stream_parsed( """ tokenizer = self.renderer.get_tokenizer() - chat_template_kwargs: dict[str, Any] = {} - if not self.use_harmony: - chat_template_kwargs = ( - chat_request.build_chat_params( - self.chat_template, - self.chat_template_content_format, - ) - .with_defaults(self.default_chat_template_kwargs) - .chat_template_kwargs - ) - parser = parser_cls( tokenizer, chat_request.tools, - chat_template_kwargs=chat_template_kwargs, + chat_template_kwargs=self._resolve_chat_template_kwargs(chat_request), model_config=self.model_config, ) From 2e06395fd384f3bd6755faded96d90a19d87e31e Mon Sep 17 00:00:00 2001 From: RomeJay Date: Tue, 6 Oct 2026 06:01:03 +0800 Subject: [PATCH 023/952] [Bugfix][Frontend] Build tool-call grammars from the tools the prompt renders (#59879) Signed-off-by: Rongjie Jiang Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structural_tag_registry.py | 127 ++++++++++++++++++ tests/tool_use/test_tool_choice_required.py | 6 + vllm/parser/abstract_parser.py | 7 +- vllm/parser/harmony.py | 2 +- vllm/tool_parsers/structural_tag_registry.py | 46 ++++++- vllm/tool_parsers/utils.py | 102 ++++++++------ 6 files changed, 247 insertions(+), 43 deletions(-) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index d5ab3e4f1ae6..70f1921e459a 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -21,6 +21,8 @@ ChatCompletionToolsParam, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.utils import construct_tool_dicts +from vllm.exceptions import VLLMValidationError from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.plamo3 import ( BEGIN_TOOL_ARGUMENTS, @@ -203,6 +205,131 @@ def test_get_model_structural_tag_supports_all_xgrammar_builtins( assert isinstance(tag, StructuralTag) +_MCP_TOOL = {"type": "mcp", "server_label": "docs", "server_url": "http://mcp"} + + +def _namespace_tool(strict: bool | None = None) -> dict[str, Any]: + lookup = { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}}, + } + if strict is not None: + lookup["strict"] = strict + return {"type": "namespace", "name": "crm", "description": "", "tools": [lookup]} + + +def _hermes_call(name: str) -> str: + return ( + f'\n{{"name": "{name}", "arguments": {{"id": "1"}}}}\n' + ) + + +@pytest.mark.parametrize("model", sorted(SUPPORTED_STRUCTURAL_TAG_MODELS)) +def test_no_structural_tag_without_function_tools(model: str): + """Builtin tools never reach the prompt, so there is nothing to constrain.""" + request = ResponsesRequest.model_validate( + {"input": "hi", "tools": [_MCP_TOOL, {"type": "web_search"}]} + ) + + tag = get_model_structural_tag( + model=model, + tools=request.tools, + tool_choice="auto", + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + + assert tag is None + with pytest.raises(VLLMValidationError, match="no function tool"): + get_model_structural_tag( + model=model, + tools=request.tools, + tool_choice="required", + reasoning=False, + ) + + +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +def test_structural_tag_allows_exactly_the_prompt_tools(tool_choice: str): + """The grammar lets the model call every tool its prompt lists, and no other.""" + request = ResponsesRequest.model_validate( + { + "input": "hi", + "tools": [ + {"type": "function", "name": "get_weather", "parameters": {}}, + _namespace_tool(), + _MCP_TOOL, + ], + } + ) + prompt_names = [ + tool["function"]["name"] + for tool in construct_tool_dicts(request.tools, tool_choice) + ] + + tag = get_model_structural_tag( + model="hermes", + tools=request.tools, + tool_choice=tool_choice, + reasoning=False, + strict_level=ToolStrictLevel.FUNCTION, + ) + + assert prompt_names == ["get_weather", "crm__lookup"] + assert tag is not None + grammar = Grammar.from_structural_tag(json.dumps(tag.model_dump())) + for name in prompt_names: + assert _is_grammar_accept_string(grammar, _hermes_call(name)) + assert not _is_grammar_accept_string(grammar, _hermes_call("docs")) + + +def test_strict_namespace_function_enables_auto_structural_tag(): + request = ResponsesRequest.model_validate( + {"input": "hi", "tools": [_namespace_tool(strict=True)]} + ) + + tag = get_model_structural_tag( + model="hermes", tools=request.tools, tool_choice="auto", reasoning=False + ) + + assert tag is not None + grammar = Grammar.from_structural_tag(json.dumps(tag.model_dump())) + assert _is_grammar_accept_string(grammar, _hermes_call("crm__lookup")) + + +def test_harmony_structural_tag_keeps_builtin_tools(): + """Harmony can call builtin tools, so they still shape its grammar.""" + request = ResponsesRequest.model_validate( + {"input": "hi", "tools": [{"type": "web_search"}]} + ) + + tag = get_model_structural_tag( + model="harmony", + tools=request.tools, + tool_choice="required", + reasoning=False, + ) + + assert tag is not None + + +def test_glm47_default_level_allows_builtin_only_tools( + monkeypatch: pytest.MonkeyPatch, +): + """GLM-4.7's default FUNCTION level must not reject builtin-only tools.""" + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", True) + + class TestParser(DelegatingParser): + tool_parser_cls = Glm47MoeModelToolParser + + request = ResponsesRequest.model_validate({"input": "hi", "tools": [_MCP_TOOL]}) + parser = TestParser(MagicMock(), tools=None) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) + + assert parser.adjust_request(request).structured_outputs is None + + def test_get_model_structural_tag_supports_vllm_hermes( sample_tools_strict: list[ChatCompletionToolsParam], ): diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index a859659417b8..fc868bbcab6e 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -13,6 +13,7 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, ) +from vllm.exceptions import VLLMValidationError from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.tool_parsers.utils import ( find_tool_properties, @@ -403,6 +404,11 @@ def test_get_json_schema_with_mixed_tools(self): assert len(any_of) == 1 assert any_of[0]["properties"]["name"]["enum"] == ["get_weather"] + def test_get_json_schema_rejects_only_non_function_tools(self): + # An empty anyOf would compile to a grammar no output can satisfy. + with pytest.raises(VLLMValidationError, match="no function tool"): + get_json_schema_from_tools(tools=[WEB_SEARCH_TOOL], tool_choice="required") + class TestParallelToolCallsConstraint: """`parallel_tool_calls=false` must be enforced by the decoding grammar. diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index bee090e7f6d1..f6b009efde07 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -50,6 +50,7 @@ extract_required_tool_call_streaming, ) from vllm.tool_parsers.structural_tag_registry import ( + get_structural_tag_tools, limit_to_single_tool_call, resolve_tool_strictness, ) @@ -520,7 +521,11 @@ def _apply_structural_tag( resolved_tools = None if tool_parser.structural_tag_model is not None: resolved_tools = resolve_tool_strictness( - request.tools, + get_structural_tag_tools( + tool_parser.structural_tag_model, + request.tools, + request.tool_choice, + ), request.tool_choice, strict_level, ) diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 6a1399db66cf..68b50250defb 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -472,7 +472,7 @@ def _assemble_tag( return StructuralTag(format=SequenceFormat(elements=tags)) -@register_vllm_structural_tag("harmony") +@register_vllm_structural_tag("harmony", builtin_tools=True) def get_harmony_structural_tag( tools: list[FunctionToolParam], builtin_tools: list[BuiltinToolParam], diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index a1ee263dc2a2..1d3a7c4a4c7f 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -39,6 +39,7 @@ ChatCompletionToolsParam, ) from vllm.tool_parsers.tool_strict_level import ToolStrictLevel +from vllm.tool_parsers.utils import get_function_tools, require_function_tools ToolChoice: TypeAlias = ( Literal["none", "auto", "required"] @@ -85,20 +86,60 @@ ) _VLLM_STRUCTURAL_TAG_REGISTRY: dict[str, StructuralTagBuilder] = {} +_BUILTIN_TOOL_STRUCTURAL_TAG_MODELS: set[str] = set() def register_vllm_structural_tag( model: str, + *, + builtin_tools: bool = False, ) -> Callable[[StructuralTagBuilder], StructuralTagBuilder]: - """Register a vLLM-owned structural tag builder.""" + """Register a vLLM-owned structural tag builder. + + Args: + model: Structural tag model key. + builtin_tools: Whether the model can call non-function tools, so its + builder receives every request tool instead of only the function + tools from ``get_function_tools``. + + """ def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder: _VLLM_STRUCTURAL_TAG_REGISTRY[model] = func + if builtin_tools: + _BUILTIN_TOOL_STRUCTURAL_TAG_MODELS.add(model) return func return decorator +def get_structural_tag_tools( + model: str, + tools: Sequence[ChatCompletionToolsParam | ResponsesTool], + tool_choice: ToolChoice, +) -> Sequence[ChatCompletionToolsParam | ResponsesTool]: + """Return the tools a structural tag for ``model`` lets the model call. + + Most models are only shown function tools, so their grammar must cover + exactly those: a tool missing from the prompt cannot be called, and one + missing from the grammar is shown but blocked. + + Raises: + VLLMValidationError: ``tool_choice`` requires a call, but the model + has no tool it can call. + + """ + if model in _BUILTIN_TOOL_STRUCTURAL_TAG_MODELS: + return tools + if ( + tool_choice is None + or tool_choice == "auto" + or getattr(tool_choice, "mode", None) == "auto" + ): + return get_function_tools(tools) + return require_function_tools(tools) + + def _tool_is_strict(tool: ChatCompletionToolsParam | ResponsesTool) -> bool: if isinstance(tool, FunctionTool): return tool.strict is True @@ -164,6 +205,9 @@ def get_model_structural_tag( if not tools or tool_choice == "none": return None + tools = get_structural_tag_tools(model, tools, tool_choice) + if not tools: + return None tools = resolve_tool_strictness(tools, tool_choice, strict_level) if tools is None: return None diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index a2992261daaf..03bcf1b652cc 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -6,6 +6,7 @@ import keyword as _python_keyword import math import warnings +from collections.abc import Sequence from dataclasses import dataclass from json import JSONDecodeError, JSONDecoder from typing import Any, TypeAlias @@ -29,6 +30,7 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, ) +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger Tool: TypeAlias = ChatCompletionToolsParam | ResponsesTool @@ -179,43 +181,74 @@ def flat_namespace_tool_name(namespace: str, name: str) -> str: return f"{namespace}{_NAMESPACE_TOOL_SEPARATOR}{name}" +def get_function_tools( + tools: Sequence[Tool], +) -> list[ChatCompletionToolsParam | FunctionTool]: + """Return the function tools the model is shown, in request order. + + This is the single source for both the chat template's tool list and the + tools a tool-call grammar constrains, so the two cannot drift apart. + Namespace functions are flattened to ``__``; other + Responses tool types are never rendered for the model and are dropped. + """ + function_tools: list[ChatCompletionToolsParam | FunctionTool] = [] + for tool in tools: + if isinstance(tool, NamespaceTool): + function_tools.extend( + FunctionTool.model_validate( + { + **namespaced_tool.model_dump(by_alias=True), + "name": flat_namespace_tool_name( + tool.name, namespaced_tool.name + ), + } + ) + for namespaced_tool in tool.tools + if namespaced_tool.type == "function" + ) + elif isinstance(tool, (FunctionTool, ChatCompletionToolsParam)): + function_tools.append(tool) + return function_tools + + +def require_function_tools( + tools: Sequence[Tool], +) -> list[ChatCompletionToolsParam | FunctionTool]: + """Return ``get_function_tools(tools)``, rejecting a request that has none. + + Raises: + VLLMValidationError: ``tools`` has no tool the model can call, so a + tool choice that requires a call can never be satisfied. + + """ + function_tools = get_function_tools(tools) + if not function_tools: + raise VLLMValidationError( + "tool_choice requires a function call, but tools contains no " + "function tool the model can call.", + parameter="tool_choice", + ) + return function_tools + + def iter_response_function_tool_info( tool: ResponsesTool, ) -> list[tuple[str, dict[str, Any] | None]]: - if isinstance(tool, FunctionTool): - return [(tool.name, tool.parameters)] - if not isinstance(tool, NamespaceTool): - return [] - - namespace = tool.name return [ - ( - flat_namespace_tool_name(namespace, namespaced_tool.name), - namespaced_tool.parameters, - ) - for namespaced_tool in tool.tools - if namespaced_tool.type == "function" + (function_tool.name, function_tool.parameters) + for function_tool in get_function_tools([tool]) + if isinstance(function_tool, FunctionTool) ] def iter_response_function_tool_dicts( tools: list[ResponsesTool], ) -> list[dict[str, Any]]: - function_tools: list[dict[str, Any]] = [] - for tool in tools: - if isinstance(tool, NamespaceTool): - namespace = tool.name - for namespaced_tool in tool.tools: - if namespaced_tool.type != "function": - continue - tool_dict = namespaced_tool.model_dump() - tool_dict["name"] = flat_namespace_tool_name( - namespace, namespaced_tool.name - ) - function_tools.append(tool_dict) - elif isinstance(tool, FunctionTool): - function_tools.append(tool.model_dump()) - return function_tools + return [ + tool.model_dump() + for tool in get_function_tools(tools) + if isinstance(tool, FunctionTool) + ] def build_responses_tool_call_name_map( @@ -353,19 +386,8 @@ def _get_json_schema_from_tools( tools: list[Tool], parallel_tool_calls: bool | None = None, ) -> dict: - fn_tool_schemas: list[dict[str, Any]] = [] - fn_tools: list[Tool] = [] - for tool in tools: - if isinstance(tool, (FunctionTool, NamespaceTool)): - fn_tool_schemas.extend( - _get_tool_schema_from_name_and_params(name, params) - for name, params in iter_response_function_tool_info(tool) - ) - if isinstance(tool, FunctionTool): - fn_tools.append(tool) - elif _is_function_tool(tool): - fn_tool_schemas.append(_get_tool_schema_from_tool(tool)) - fn_tools.append(tool) + fn_tools: list[Tool] = list(require_function_tools(tools)) + fn_tool_schemas = [_get_tool_schema_from_tool(tool) for tool in fn_tools] json_schema: dict[str, Any] = { "type": "array", "minItems": 1, From 89a00995c944d934c4c83ef853670cdd838c8bcb Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 5 Oct 2026 15:25:11 -0700 Subject: [PATCH 024/952] [CI] Read each profiler round's trace as it stops in test_gpu_profiler (#59769) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- tests/v1/worker/test_gpu_profiler.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index 5b407e80e903..e140803ef0d3 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import os from contextlib import nullcontext from types import SimpleNamespace @@ -98,7 +99,14 @@ def test_torch_profiler_rebuilds_one_shot_profiler_each_round(tmp_path): def test_torch_profiler_records_each_profile_round(tmp_path): - traces: list[torch.profiler.profile] = [] + trace_paths: list[str] = [] + + def export_trace(prof: torch.profiler.profile) -> None: + # Export as each round stops instead of parsing every round at the end. + path = str(tmp_path / f"round_{len(trace_paths)}.json") + prof.export_chrome_trace(path) + trace_paths.append(path) + wrapper = TorchProfilerWrapper( ProfilerConfig( profiler="torch", @@ -108,7 +116,7 @@ def test_torch_profiler_records_each_profile_round(tmp_path): worker_name="worker", local_rank=1, activities=["CPU"], - on_trace_ready=traces.append, + on_trace_ready=export_trace, ) for run in range(2): @@ -117,10 +125,16 @@ def test_torch_profiler_records_each_profile_round(tmp_path): pass wrapper.stop() - assert len(traces) == 2 - for run, trace in enumerate(traces): + assert len(trace_paths) == 2 + for run, path in enumerate(trace_paths): + # The CPU profiler records every thread in the process, and an event + # name from elsewhere in a long session can hold invalid UTF-8. + with open(path, encoding="utf-8", errors="replace") as f: + events = json.load(f)["traceEvents"] assert { - event.name for event in trace.events() if event.name.startswith("run_") + event["name"] + for event in events + if event.get("name", "").startswith("run_") } == {f"run_{run}"} From 4a705c7b54ac775e6436eae6e1ee573f8d4c7775 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Tue, 6 Oct 2026 00:43:31 +0200 Subject: [PATCH 025/952] [Bugfix][Frontend] Preserve logprobs when top_logprobs is null (#47838) Signed-off-by: Ting Sun Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../openai/chat_completion/test_chat.py | 17 +++++++++++++++++ .../openai/responses/test_sampling_params.py | 14 ++++++++++++++ vllm/entrypoints/generate/base/protocol.py | 7 +++++++ .../openai/chat_completion/batch_serving.py | 4 +--- .../openai/chat_completion/protocol.py | 5 +++-- .../openai/chat_completion/serving.py | 8 ++------ vllm/entrypoints/openai/responses/protocol.py | 3 ++- 7 files changed, 46 insertions(+), 12 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_chat.py b/tests/entrypoints/openai/chat_completion/test_chat.py index 5541b605d99e..c246482d8fcd 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_chat.py @@ -16,6 +16,7 @@ from tests.utils import RemoteOpenAIServer from vllm.entrypoints.openai.chat_completion.protocol import ( + BatchChatCompletionRequest, ChatCompletionRequest, ) from vllm.exceptions import VLLMValidationError @@ -1012,6 +1013,22 @@ def test_chat_completion_request_n_parameter_default(): assert sampling_params.n == 1, f"Expected n=1 (default), got n={sampling_params.n}" +@pytest.mark.parametrize( + ("request_cls", "messages"), + [ + (ChatCompletionRequest, [{"role": "user", "content": "Hello"}]), + (BatchChatCompletionRequest, [[{"role": "user", "content": "Hello"}]]), + ], +) +def test_null_top_logprobs_is_same_as_omitted(request_cls, messages): + """`top_logprobs: null` must not switch off `logprobs: true`.""" + request = request_cls.model_validate( + {"messages": messages, "logprobs": True, "top_logprobs": None} + ) + + assert request.top_logprobs == 0 + + def test_chat_completion_request_accepts_model_specific_reasoning_effort(): request = ChatCompletionRequest( model="test-model", diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 6ede3f1f7f28..877e8724e8b2 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -84,6 +84,20 @@ def test_default_values(self): assert sampling_params.stop == [] # Empty list assert sampling_params.extra_args == {} # Empty dict + def test_null_top_logprobs_is_same_as_omitted(self): + """`top_logprobs: null` must not drop the requested output logprobs.""" + request = ResponsesRequest.model_validate( + { + "input": "test input", + "include": ["message.output_text.logprobs"], + "top_logprobs": None, + } + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.logprobs == 0 + def test_seed_bounds_validation(self): """Test that seed values outside torch.long bounds are rejected.""" # Test seed below minimum diff --git a/vllm/entrypoints/generate/base/protocol.py b/vllm/entrypoints/generate/base/protocol.py index c1834981431a..b858ea19a5c5 100644 --- a/vllm/entrypoints/generate/base/protocol.py +++ b/vllm/entrypoints/generate/base/protocol.py @@ -8,6 +8,7 @@ from pydantic import ( BaseModel, + BeforeValidator, Field, model_serializer, ) @@ -26,6 +27,12 @@ str | Annotated[list[str], Field(max_length=envs.VLLM_MAX_STOP_STRINGS)] | None ) +# `top_logprobs` is nullable in the OpenAI spec; null means the same as omitted. +TopLogprobsParam: TypeAlias = Annotated[ + int, + BeforeValidator(lambda v: 0 if v is None else v, json_schema_input_type=int | None), +] + _CACHE_SALT_FORBIDDEN_CHARS = frozenset("@/\\\x00") _MAX_CACHE_SALT_LENGTH = 128 diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index 8c7ae3764778..e630b97f182a 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -247,9 +247,7 @@ async def chat_completion_full_generator_batch( for output in final_res.outputs: self._raise_if_error(output.finish_reason, request_id) - if request.logprobs and ( - request.top_logprobs is not None or request.logprob_token_ids - ): + if request.logprobs: assert output.logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=output.token_ids, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index d2d1cbab93a6..aee1c8a7a63a 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -33,6 +33,7 @@ StopParam, StreamOptions, ToolCall, + TopLogprobsParam, structured_outputs_from_response_format, validate_cache_salt, validate_structural_tag_response_format, @@ -222,7 +223,7 @@ class ChatCompletionRequest(OpenAIBaseModel): frequency_penalty: float | None = None logit_bias: dict[str, float] | None = None logprobs: bool | None = False - top_logprobs: int | None = 0 + top_logprobs: TopLogprobsParam = 0 max_tokens: int | None = Field( default=None, deprecated="max_tokens is deprecated in favor of " @@ -1095,7 +1096,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel): frequency_penalty: float | None = 0.0 logit_bias: dict[str, float] | None = None logprobs: bool | None = False - top_logprobs: int | None = 0 + top_logprobs: TopLogprobsParam = 0 logprob_token_ids: list[int] | None = Field( default=None, description=( diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 361cc52f72d2..9b0737c25d47 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -625,9 +625,7 @@ async def chat_completion_stream_generator( self._raise_if_error(output.finish_reason, request_id) - if request.logprobs and ( - request.top_logprobs is not None or request.logprob_token_ids - ): + if request.logprobs: assert output.logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=output.token_ids, @@ -970,9 +968,7 @@ async def chat_completion_full_generator( token_ids = output.token_ids out_logprobs = output.logprobs - if request.logprobs and ( - request.top_logprobs is not None or request.logprob_token_ids - ): + if request.logprobs: assert out_logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=token_ids, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 8fa96ecdb05c..5d4a7411ac21 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -61,6 +61,7 @@ from vllm.entrypoints.generate.base.protocol import ( PerRequestMetrics, StopParam, + TopLogprobsParam, validate_cache_salt, ) from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel @@ -210,7 +211,7 @@ class ResponsesRequest(OpenAIBaseModel): text: ResponseTextConfig | None = None tool_choice: ToolChoice = "auto" tools: list[Tool] = Field(default_factory=list) - top_logprobs: int | None = 0 + top_logprobs: TopLogprobsParam = 0 top_p: float | None = None top_k: int | None = None truncation: Literal["auto", "disabled"] | None = "disabled" From 0eac15270710ddb7f5296d7244fcd0496c2daa0f Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 6 Oct 2026 03:03:44 +0400 Subject: [PATCH 026/952] [Bugfix] Route every speculation-capable row through the speculative path so recurrent state stays consistent (#56531) Signed-off-by: Vadim Gimpelson Co-authored-by: Cursor Agent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/models/kimi_k3/test_kda_metadata.py | 2 +- .../test_gdn_zero_draft_spec_rows.py | 211 +++++++++++++++ ...odel_runner_prefill_tail_classification.py | 249 ++++++++++++++++++ .../worker/test_mamba_hybrid_model_state.py | 171 ++++++++++++ vllm/models/kimi_k3/nvidia/kda_metadata.py | 18 +- vllm/v1/attention/backends/gdn_attn.py | 9 +- .../worker/gpu/model_states/mamba_hybrid.py | 21 +- vllm/v1/worker/gpu_model_runner.py | 22 +- 8 files changed, 672 insertions(+), 31 deletions(-) create mode 100644 tests/v1/attention/test_gdn_zero_draft_spec_rows.py create mode 100644 tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 6b3c71bdc674..74414279552c 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -447,7 +447,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset(): 2, False, [False, False], - None, + 1, id="no-scheduled-draft-tokens", ), pytest.param( diff --git a/tests/v1/attention/test_gdn_zero_draft_spec_rows.py b/tests/v1/attention/test_gdn_zero_draft_spec_rows.py new file mode 100644 index 000000000000..55955886ce5d --- /dev/null +++ b/tests/v1/attention/test_gdn_zero_draft_spec_rows.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which path GDN metadata takes when a decode row carries no draft tokens. + +`build_for_cudagraph_capture` derives the draft count from the query length +(`num_decode_draft_tokens = query_len - 1`), so a capture shape of one token +per request describes rows with zero drafts. `GDNAttentionMetadataBuilder` +routes every row with `num_decode_draft_tokens_cpu >= 0` through the +speculative path, zero drafts included: that path is the only one that +applies each row's accepted-token offset to the recurrent state, so a +zero-draft decode row must not fall back to the plain decode path. +""" + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +NUM_SPEC = 3 + + +def _builder( + use_spec_decode: bool, + mamba_cache_mode: str = "none", + shapes: tuple[tuple[int, ...], ...] = ((16, 64),), + dtypes: tuple[torch.dtype, ...] = (torch.float16,), +) -> GDNAttentionMetadataBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + vllm_config.speculative_config = ( + SpeculativeConfig(method="ngram", num_speculative_tokens=NUM_SPEC) + if use_spec_decode + else None + ) + # "none" keeps the block table as-is; the align gather needs columns for + # the speculative blocks, which `_decode_metadata` appends on request. + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + mamba_spec = MambaSpec( + block_size=BLOCK_SIZE, + shapes=shapes, + dtypes=dtypes, + num_speculative_blocks=NUM_SPEC, + ) + + return GDNAttentionMetadataBuilder( + kv_cache_spec=mamba_spec, + layer_names=["layer.0"], + vllm_config=vllm_config, + device=DEVICE, + ) + + +def _capture_metadata(builder, query_len: int, batch_size: int = 2): + """Metadata as produced during CUDA-graph capture of a decode shape.""" + batch = BatchSpec( + seq_lens=[64] * batch_size, + query_lens=[query_len] * batch_size, + ) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE) + return builder.build_for_cudagraph_capture(common) + + +def _decode_metadata(builder, num_accepted: list[int], seq_len: int = 64): + """Metadata for a live decode batch where no row was given drafts. + + Every row schedules a single token, so `num_decode_draft_tokens_cpu` is all + zeros -- the shape the runner produces once `zero_draft_decode_mask` has + promoted zero-draft decode rows out of the `-1` placeholder. + `num_accepted` is what the previous step accepted (bonus token included, + so one is the minimum), which is what decides the state slot. + """ + batch = BatchSpec( + seq_lens=[seq_len] * len(num_accepted), query_lens=[1] * len(num_accepted) + ) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE, arange_block_indices=True + ) + if builder.vllm_config.cache_config.mamba_cache_mode == "align": + # Align gathers the last 1 + num_speculative_blocks columns starting at + # the request's current block, so the table has to reach past the end of + # the sequence -- at runtime it is sized by max_model_len, while the + # test helper sizes it by seq_len. + rows, cols = common.block_table_tensor.shape + extra = ( + torch.arange( + rows * NUM_SPEC, dtype=common.block_table_tensor.dtype, device=DEVICE + ).view(rows, NUM_SPEC) + + rows * cols + ) + common.block_table_tensor = torch.cat([common.block_table_tensor, extra], dim=1) + accepted = torch.tensor(num_accepted, dtype=torch.int32, device=DEVICE) + num_decode_draft_tokens_cpu = torch.zeros(len(num_accepted), dtype=torch.int32) + return builder.build(0, common, accepted, num_decode_draft_tokens_cpu) + + +def test_zero_draft_capture_shape_routes_through_the_speculative_path(): + """The one-token capture shape is what breaks backends whose decode kernels + assume a full num_spec + 1 query window: `num_decode_draft_tokens_cpu = + query_len - 1 = 0` for every row, which is `>= 0` and must still take the + speculative path (not the plain decode path that drops the state offset), + with a state tensor sized for the full window even though every row only + supplies a single query token. + + Only `query_len=1` (zero drafts) discriminates this from the plain-decode + path before the fix: for any `query_len > 1`, `num_decode_draft_tokens_cpu` + is already `> 0` and both the pre-fix and post-fix code agree, so this + does not need parametrizing over larger draft counts. + """ + meta = _capture_metadata(_builder(use_spec_decode=True), query_len=1) + + assert meta.num_spec_decodes == 2 + assert meta.num_decodes == 0 + slots = meta.spec_state_indices_tensor + assert slots is not None + # The kernels index columns 0..num_spec of this tensor. + assert slots.shape[1] == NUM_SPEC + 1 + assert int(slots.min()) >= 0 + query_lens = torch.diff(meta.spec_query_start_loc) + assert torch.all(query_lens == 1) + + +def test_plain_decode_without_spec_decode_skips_the_speculative_path(): + """Without a speculative config the builder never builds spec metadata, + so a one-token-per-row decode capture takes the plain decode path. + + Not a regression check for this patch (this branch is untouched by it), + but it is what keeps the test above from passing vacuously: it would also + pass if the builder routed every row through the speculative path + unconditionally, regardless of whether speculative decoding is even + configured. + """ + meta = _capture_metadata(_builder(use_spec_decode=False), query_len=1) + + assert meta.num_spec_decodes == 0 + assert meta.num_decodes == 2 + + +@pytest.mark.parametrize( + "mamba_cache_mode, shapes, dtypes", + [ + ("none", ((16, 64),), (torch.float16,)), + # Kimi-Linear's KDA layers subclass GatedDeltaNetAttention and consume + # this same metadata, and prefix caching (which Kimi enables by + # default) puts the cache in align mode -- so the guarantee has to hold + # for that state shape and mode too. + ( + "align", + MambaStateShapeCalculator.kda_state_shape( + tp_world_size=1, num_heads=4, head_dim=32, num_spec=NUM_SPEC + ), + (torch.bfloat16, torch.bfloat16), + ), + ], + ids=["gdn_none", "kimi_kda_align"], +) +def test_zero_draft_batch_still_applies_the_accepted_token_offset( + mamba_cache_mode, shapes, dtypes +): + """The numeric difference the fix makes, on the batch shape that triggers it. + + The recurrent kernels read their state at column `num_accepted_tokens - 1` + of `spec_state_indices_tensor`. Before the fix, a batch in which *every* + speculative row drafted nothing was pushed onto the non-speculative path, + which passes `non_spec_state_indices_tensor` -- always column 0 -- and + drops `num_accepted_tokens` entirely. A row that accepted drafts on the + previous step wrote its fresh state further along, so column 0 is stale. + """ + accepted = [NUM_SPEC, 1] + meta = _decode_metadata( + _builder( + use_spec_decode=True, + mamba_cache_mode=mamba_cache_mode, + shapes=shapes, + dtypes=dtypes, + ), + accepted, + ) + + assert meta.num_spec_decodes == len(accepted) + assert meta.num_decodes == 0 + # Dropping these is what silently sent the kernels to the wrong slot. + assert meta.num_accepted_tokens is not None + assert meta.num_accepted_tokens.tolist()[: len(accepted)] == accepted + assert meta.non_spec_state_indices_tensor is None + + slots = meta.spec_state_indices_tensor + assert slots is not None + for row, num_accepted in enumerate(accepted): + read = int(slots[row, num_accepted - 1]) + fallback = int(slots[row, 0]) + # Only the row that accepted a draft moved off column 0; the row that + # accepted nothing but its bonus token legitimately stays there. This + # also confirms the columns are genuinely distinct blocks (via + # `arange_block_indices=True` in `_decode_metadata`), not just + # distinct labels for the same one -- otherwise `read != fallback` + # here would mean nothing. + assert (read != fallback) == (num_accepted > 1) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py new file mode 100644 index 000000000000..9375c477a224 --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""`GPUModelRunner._prepare_inputs` must tell a chunked-prefill continuation +row apart from a genuine zero-draft decode row. + +Both kinds of row are absent from `scheduled_spec_decode_tokens` and both +schedule exactly one token, so `num_decode_draft_tokens` starts at the +placeholder `-1` for either. The `zero_draft_decode_mask` in +`gpu_model_runner.py` is what tells them apart afterwards, using +`num_computed_tokens >= num_prompt_tokens`: a prefill continuation must keep +`-1` (it carries no drafts and no accepted-token offset to apply), while a +zero-draft decode row must become `0` (it still has to run the speculative +path so the recurrent backends apply that offset). Getting this wrong in +either direction was the root cause behind issue #40875: Kimi-Linear crashed +with an illegal memory access when prefill tails were misclassified as +zero-draft decodes. + +Each batch below also includes one row with a real draft token, so +`scheduled_spec_decode_tokens` is non-empty even under the pre-fix +`use_spec_decode = len(scheduled_spec_decode_tokens) > 0` predicate. That +keeps the assertions targeted at the mask fix alone, not the unrelated +`use_spec_decode` predicate fix (`self.speculative_config is not None`) that +landed in the same patch. +""" + +import numpy as np +import pytest +import torch + +from vllm.config import ( + CacheConfig, + ModelConfig, + ParallelConfig, + SchedulerConfig, + SpeculativeConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.model_executor.layers.attention import Attention +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, +) +from vllm.v1.worker.gpu_input_batch import InputBatch +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +BLOCK_SIZE = 16 +NUM_BLOCKS = 10 +DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu" + + +def _get_vllm_config() -> VllmConfig: + model_config = ModelConfig(model="facebook/opt-125m", dtype="float16", seed=42) + scheduler_config = SchedulerConfig( + max_num_seqs=10, + max_num_batched_tokens=512, + max_model_len=512, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + cache_config = CacheConfig( + block_size=BLOCK_SIZE, gpu_memory_utilization=0.9, cache_dtype="auto" + ) + cache_config.kv_cache_layout = "LBNHC" + return VllmConfig( + model_config=model_config, + cache_config=cache_config, + scheduler_config=scheduler_config, + parallel_config=ParallelConfig(), + # ngram needs no draft model weights, so the fixture stays as cheap as + # the speculative-decoding-free one in test_gpu_model_runner.py. + speculative_config=SpeculativeConfig(method="ngram", num_speculative_tokens=3), + ) + + +def _initialize_kv_cache(runner: GPUModelRunner) -> None: + attn_spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config), + head_size=runner.model_config.get_head_size(), + dtype=runner.kv_cache_dtype, + ) + tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor( + size=tensor_size, + layers=["layer.0"], + layer_stride=tensor_size, + block_stride=attn_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec) + ], + ) + runner.kv_cache_config = kv_cache_config + runner.input_batch = InputBatch( + max_num_reqs=runner.max_num_reqs, + max_model_len=runner.max_model_len, + max_num_batched_tokens=runner.max_num_tokens, + device=runner.device, + vocab_size=runner.model_config.get_vocab_size(), + block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], + kernel_block_sizes=[ + kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size + ], + max_num_blocks_per_req=[NUM_BLOCKS], + ) + runner.initialize_attn_backend(kv_cache_config) + + +@pytest.fixture +def spec_decode_model_runner(dist_init): + vllm_config = _get_vllm_config() + with set_current_vllm_config(vllm_config): + model_config = vllm_config.model_config + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config) + head_size = model_config.get_head_size() + vllm_config.compilation_config.static_forward_context["layer.0"] = Attention( + num_heads, head_size, 0.1 + ) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) + _initialize_kv_cache(runner) + yield runner + + +def _new_request( + req_id: str, prompt_len: int, num_computed_tokens: int +) -> tuple[NewRequestData, int]: + """A request with `prompt_len` prompt tokens, `num_computed_tokens` of + which are already computed (e.g. a prefix-cache hit), scheduled for the + single remaining/next token this step.""" + return ( + NewRequestData( + req_id=req_id, + prompt_token_ids=list(range(prompt_len)), + mm_features=[], + sampling_params=SamplingParams(), + pooling_params=None, + block_ids=([0],), + num_computed_tokens=num_computed_tokens, + lora_request=None, + ), + 1, + ) + + +def test_prefill_continuation_keeps_placeholder_decode_zero_draft_becomes_zero( + spec_decode_model_runner, +): + """One request has one token left to prefill (`num_computed_tokens == + num_prompt_tokens - 1`); another has already finished its prompt and is + now decoding with no draft tokens proposed this step; a third is decoding + with a real draft. `decode_with_draft` gives the batch a non-empty + `scheduled_spec_decode_tokens`, so this also exercises the `-1`-vs-`0` + distinction on the pre-fix definition of `use_spec_decode` (`len(sched + uled_spec_decode_tokens) > 0`), isolating the mask fix from the unrelated + `use_spec_decode` predicate fix that landed in the same patch.""" + runner = spec_decode_model_runner + + prefill_req, prefill_sched = _new_request( + "prefill_tail", prompt_len=11, num_computed_tokens=10 + ) + decode_req, decode_sched = _new_request( + "decode_zero_draft", prompt_len=3, num_computed_tokens=3 + ) + drafted_req, _ = _new_request( + "decode_with_draft", prompt_len=3, num_computed_tokens=3 + ) + num_scheduled_tokens = { + "prefill_tail": prefill_sched, + "decode_zero_draft": decode_sched, + "decode_with_draft": 2, # 1 draft token + 1 bonus token + } + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[prefill_req, decode_req, drafted_req], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), + scheduled_spec_decode_tokens={"decode_with_draft": [99]}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + runner._update_states(scheduler_output) + runner._prepare_inputs( + scheduler_output, + np.array( + [num_scheduled_tokens[r] for r in runner.input_batch.req_ids], + dtype=np.int32, + ), + ) + + prefill_idx = runner.input_batch.req_id_to_index["prefill_tail"] + decode_idx = runner.input_batch.req_id_to_index["decode_zero_draft"] + drafted_idx = runner.input_batch.req_id_to_index["decode_with_draft"] + assert runner.num_decode_draft_tokens.np[prefill_idx] == -1 + assert runner.num_decode_draft_tokens.np[decode_idx] == 0 + assert runner.num_decode_draft_tokens.np[drafted_idx] == 1 + + +def test_prefill_tail_of_two_or_three_tokens_also_keeps_the_placeholder( + spec_decode_model_runner, +): + """A chunked-prefill tail can be 1, 2, or 3 tokens (`long_prefill_token_ + threshold` rounding); the mask only fires for `num_scheduled_tokens == 1`, + so wider tails must keep `-1` regardless of how much prompt is left. A + third, drafted decode row keeps `scheduled_spec_decode_tokens` non-empty + so the assertions hold on the pre-fix `use_spec_decode` predicate too.""" + runner = spec_decode_model_runner + + tail_2_req, _ = _new_request("tail_2", prompt_len=12, num_computed_tokens=10) + tail_3_req, _ = _new_request("tail_3", prompt_len=13, num_computed_tokens=10) + drafted_req, _ = _new_request( + "decode_with_draft", prompt_len=3, num_computed_tokens=3 + ) + num_scheduled_tokens = {"tail_2": 2, "tail_3": 3, "decode_with_draft": 2} + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[tail_2_req, tail_3_req, drafted_req], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), + scheduled_spec_decode_tokens={"decode_with_draft": [99]}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + runner._update_states(scheduler_output) + runner._prepare_inputs( + scheduler_output, + np.array( + [num_scheduled_tokens[r] for r in runner.input_batch.req_ids], + dtype=np.int32, + ), + ) + + tail_2_idx = runner.input_batch.req_id_to_index["tail_2"] + tail_3_idx = runner.input_batch.req_id_to_index["tail_3"] + assert runner.num_decode_draft_tokens.np[tail_2_idx] == -1 + assert runner.num_decode_draft_tokens.np[tail_3_idx] == -1 diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index dba5444150a6..f87201e48de2 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -140,6 +140,177 @@ def test_padded_prompt_tail_builds_as_spec_decode( assert meta.num_prefill_tokens == k + 1 +def _mamba_hybrid_state(num_speculative_tokens: int) -> MambaHybridModelState: + state = object.__new__(MambaHybridModelState) + state.vllm_config = SimpleNamespace(num_speculative_tokens=num_speculative_tokens) + state.model_config = SimpleNamespace(max_model_len=8192) + state._align_mode = False + state.recoverssm = None + # A request that accepted 3 drafts (plus the bonus token) on the previous + # step; the others have never had a draft accepted. + state.num_accepted_tokens_gpu = torch.tensor([4, 1, 1], dtype=torch.int32) + return state + + +def _input_batch( + num_scheduled_tokens: list[int], + num_draft_tokens_per_req: list[int] | None, + is_prefilling: list[bool], + num_reqs_after_padding: int | None = None, +) -> SimpleNamespace: + """`num_draft_tokens_per_req=None` is a batch where nobody proposed a + draft this step, e.g. because ngram found no match for anyone.""" + num_reqs = len(num_scheduled_tokens) + query_start_loc = np.zeros(num_reqs + 1, dtype=np.int32) + query_start_loc[1:] = np.cumsum(num_scheduled_tokens) + num_tokens = int(query_start_loc[-1]) + seq_lens = torch.full((num_reqs,), 10, dtype=torch.int32) + return SimpleNamespace( + num_reqs=num_reqs, + num_reqs_after_padding=num_reqs_after_padding or num_reqs, + num_tokens=num_tokens, + num_tokens_after_padding=num_tokens, + query_start_loc_np=query_start_loc, + num_scheduled_tokens=np.array(num_scheduled_tokens, dtype=np.int32), + max_query_len=None, + seq_lens_cpu_upper_bound=seq_lens, + is_prefilling_np=np.array(is_prefilling), + prefill_runs_as_decode_np=None, + num_draft_tokens_per_req=None + if num_draft_tokens_per_req is None + else np.array(num_draft_tokens_per_req, dtype=np.int32), + idx_mapping=torch.arange(num_reqs), + query_start_loc=torch.from_numpy(query_start_loc), + seq_lens=seq_lens, + dcp_local_seq_lens=None, + positions=torch.zeros(num_tokens, dtype=torch.int64), + prompt_lens=None, + ) + + +def _prepare_attn_metadata( + state: MambaHybridModelState, + monkeypatch, + input_batch: SimpleNamespace | None = None, + cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, +) -> SimpleNamespace: + captured = {} + + def fake_build_attn_metadata(**kwargs): + captured.update(kwargs) + return Mock() + + monkeypatch.setattr( + "vllm.v1.worker.gpu.model_states.mamba_hybrid.build_attn_metadata", + fake_build_attn_metadata, + ) + + if input_batch is None: + input_batch = _input_batch([1, 1], None, [False, False]) + state.prepare_attn( + input_batch=input_batch, + cudagraph_mode=cudagraph_mode, + block_tables=(), + slot_mappings={}, + attn_groups=[], + kv_cache_config=Mock(), + ) + return captured["model_specific_attn_metadata"] + + +def test_prepare_attn_keeps_the_accepted_offset_when_the_batch_drafted_nothing( + monkeypatch, +): + """`num_accepted_tokens` must reach the attention builders whenever the + speculative config is on, even when *this* batch scheduled no draft + tokens at all. Before the fix, `use_spec_decode` was gated on this step's + scheduled drafts (`len(scheduled_spec_decode_tokens) > 0` in the V1 + runner) rather than on the speculative config, so a request that accepted + drafts on a previous step would silently lose that offset the moment + ngram found no match for anyone in the batch. + """ + state = _mamba_hybrid_state(num_speculative_tokens=3) + + metadata = _prepare_attn_metadata(state, monkeypatch) + + assert metadata.num_accepted_tokens is not None + assert metadata.num_accepted_tokens.tolist() == [4, 1] + assert metadata.num_decode_draft_tokens_cpu is not None + + +def test_prepare_attn_omits_the_offset_without_a_speculative_config(monkeypatch): + """Mirrors the check above: with speculative decoding off entirely, there + is no accepted-token offset to carry, and the metadata must say so + explicitly (`None`) rather than a stale or default value.""" + state = _mamba_hybrid_state(num_speculative_tokens=0) + + metadata = _prepare_attn_metadata(state, monkeypatch) + + assert metadata.num_accepted_tokens is None + assert metadata.num_decode_draft_tokens_cpu is None + + +def test_zero_draft_decode_row_is_marked_as_a_speculative_row(monkeypatch): + """A decode row that got no drafts still has to reach the speculative path. + + The recurrent backends only apply a row's accepted-token offset on that + path, and a row that accepted drafts on the previous step keeps that offset + even when the current step proposes nothing (the scheduler drops drafts + whenever the token budget truncates the step). Marking such a row with the + -1 sentinel sends it to the plain decode path, which reads the state at + column 0 and ignores the offset. + """ + state = _mamba_hybrid_state(num_speculative_tokens=3) + input_batch = _input_batch([3, 1, 1], [2, 0, 0], [False, False, True]) + + metadata = _prepare_attn_metadata(state, monkeypatch, input_batch) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [ + 2, # decode row with drafts + 0, # decode row with no drafts: speculative path, offset applied + -1, # one-token prefill tail: no offset to apply, stays a plain row + ] + + +def test_batch_without_any_drafts_still_marks_its_decode_rows(monkeypatch): + """`num_draft_tokens_per_req` is None whenever no row in the batch drafted + anything. Those rows are ordinary decode rows carrying an offset from the + previous step, so the absence of drafts in this step must not be read as + speculative decoding being off.""" + state = _mamba_hybrid_state(num_speculative_tokens=3) + input_batch = _input_batch([1, 1], None, [False, True]) + + metadata = _prepare_attn_metadata(state, monkeypatch, input_batch) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, -1] + + +def test_padded_rows_keep_the_sentinel(monkeypatch): + """Full CUDA-graph capture pads the batch; padded rows describe no request + and must not be picked up as speculative rows.""" + state = _mamba_hybrid_state(num_speculative_tokens=3) + input_batch = _input_batch([1, 1], [0, 0], [False, False], num_reqs_after_padding=4) + + metadata = _prepare_attn_metadata( + state, monkeypatch, input_batch, cudagraph_mode=CUDAGraphMode.FULL + ) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, 0, -1, -1] + + +def test_chunked_prefill_tail_of_two_or_three_tokens_keeps_the_sentinel( + monkeypatch, +): + """A prefill continuation is excluded by `is_prefilling`, not by its token + count, so tails wider than one token are covered by the same rule.""" + state = _mamba_hybrid_state(num_speculative_tokens=3) + input_batch = _input_batch([2, 3], [0, 0], [True, True]) + + metadata = _prepare_attn_metadata(state, monkeypatch, input_batch) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [-1, -1] + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)]) def test_postprocess_state_scalar_with_int32_mapping( diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index e5a63f13a343..58c74e625818 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -409,20 +409,12 @@ def build( # type: ignore[override] spec_sequence_masks_cpu &= ( query_start_loc_cpu.diff() <= self.num_spec + 1 ) - # Native KDA can use its regular decode path when no draft token - # was scheduled. RecoverSSM must preserve its extended conv window. - if ( - not self.use_recoverssm - and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() - == 0 - ): + # A batch whose rows all drafted nothing still has to run the + # speculative path: that is the only path that applies each row's + # accepted-token offset to the recurrent state. + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + if num_spec_decodes == 0: spec_sequence_masks_cpu = None - num_spec_decodes = 0 - active_non_spec_mask_cpu = None - else: - num_spec_decodes = spec_sequence_masks_cpu.sum().item() - if num_spec_decodes == 0: - spec_sequence_masks_cpu = None spec_request_indices = None spec_token_start = None diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 5fece32a1ba3..95f6a808a597 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -260,11 +260,10 @@ def build( # type: ignore[override] else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 num_spec_decodes = spec_sequence_masks_cpu.sum().item() - if ( - num_spec_decodes == 0 - or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() - == 0 - ): + # A batch whose rows all drafted nothing still has to run the + # speculative path: that is the only path that applies each row's + # accepted-token offset to the recurrent state. + if num_spec_decodes == 0: num_spec_decodes = 0 spec_sequence_masks = None spec_sequence_masks_cpu = None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 32924555b123..f3f3c666e0b9 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -277,18 +277,19 @@ def prepare_attn( ] # GDN uses >= 0 to select spec-decode rows, so non-decode rows - # need the -1 sentinel rather than a raw zero draft count. + # need the -1 sentinel. Decode rows without drafts stay spec rows: + # only that path applies the previous step's accepted offset. num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req - if num_draft_tokens_per_req is not None: - # Test request state, not num_scheduled_tokens == draft_count+1: - # adaptive rewrites num_scheduled_tokens to an even split, so that - # equality rarely holds and would demote every verify row to decode. - is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0) - spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode - num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( - spec_decode_mask, num_draft_tokens_per_req, -1 - ) + if num_draft_tokens_per_req is None: + num_draft_tokens_per_req = np.zeros(input_batch.num_reqs, np.int32) + # Test request state, not num_scheduled_tokens == draft_count+1: + # adaptive rewrites num_scheduled_tokens to an even split, so that + # equality rarely holds and would demote every verify row to decode. + is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0) + num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( + is_decode, num_draft_tokens_per_req, -1 + ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) if self._align_mode: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index df7dc96516ee..76299308f866 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -840,6 +840,12 @@ def __init__( self.num_accepted_tokens = self._make_buffer( self.max_num_reqs, dtype=torch.int32 ) + # The count includes the bonus token, so one is the minimum, and the + # recurrent kernels pick their state slot as `accepted - 1`. Dummy runs + # skip `_prepare_inputs`, so the zeros from the allocation would reach + # those kernels as index -1 and fault on the out-of-bounds read. + self.num_accepted_tokens.np.fill(1) + self.num_accepted_tokens.copy_to_gpu() # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: @@ -2199,7 +2205,7 @@ def _prepare_inputs( ) self.mrope_positions.gpu[:, :total_num_scheduled_tokens] += drift - use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + use_spec_decode = self.speculative_config is not None if not use_spec_decode: # NOTE(woosuk): Due to chunked prefills, the batch may contain # partial requests. While we should not sample any token @@ -2226,6 +2232,18 @@ def _prepare_inputs( num_draft_tokens[req_idx] = draft_len if num_scheduled_tokens[req_idx] == draft_len + 1: num_decode_draft_tokens[req_idx] = draft_len + # Some recurrent backends may read the previous step's state at an + # offset that only the speculative path applies. So a decode step + # that got no drafts must still count as a speculative row with + # zero drafts, not as a plain decode. Prompt chunks are left out: + # their tokens are prompt tokens and carry no such offset. + zero_draft_decode_mask = ( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + >= self.input_batch.num_prompt_tokens[:num_reqs] + ) & (num_scheduled_tokens[:num_reqs] == 1) + num_decode_draft_tokens[ + (num_decode_draft_tokens < 0) & zero_draft_decode_mask + ] = 0 spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens ) @@ -4312,7 +4330,7 @@ def execute_model( self.mamba_state_idx, ) - use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + use_spec_decode = self.speculative_config is not None ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices slot_mappings_by_group, slot_mappings = self._get_slot_mappings( From 6dd813e0b7ca8314fed28c13a73cfb5fc5310163 Mon Sep 17 00:00:00 2001 From: Mason <150880684+usberkeley@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:24:14 +0800 Subject: [PATCH 027/952] [Bugfix][MoE] Allow MiniMax2 routing in TRTLLM BF16 monolithic backend (#59571) Signed-off-by: Mason Pitt Co-authored-by: Yongye Zhu --- .../kernels/moe/test_unquantized_backend_selection.py | 11 +++++++++-- .../layers/fused_moe/experts/trtllm_bf16_moe.py | 1 + 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 9d0f4d6900c4..25e6b7dec618 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -446,8 +446,15 @@ def test_select_cuda_deepep_ht_falls_back_from_trtllm( "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", return_value=True, ) +@pytest.mark.parametrize( + "routing_method", + [RoutingMethodType.Renormalize, RoutingMethodType.MiniMax2], +) +@pytest.mark.parametrize("moe_backend", ["auto", "flashinfer_trtllm"]) def test_select_cuda_flashinfer_trtllm_ag_rs_uses_monolithic( mock_supports_current_device, + routing_method, + moe_backend, ): """Test AG/RS stays on BF16 TRTLLM monolithic when TRTLLM is supported.""" with ( @@ -462,8 +469,8 @@ def test_select_cuda_flashinfer_trtllm_ag_rs_uses_monolithic( ), ): moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2) - moe_config.moe_backend = "flashinfer_trtllm" - moe_config.routing_method = RoutingMethodType.Renormalize + moe_config.moe_backend = moe_backend + moe_config.routing_method = routing_method moe_config.moe_parallel_config.use_ep = True moe_config.moe_parallel_config.dp_size = 2 moe_config.moe_parallel_config.ep_size = 2 diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 782d6bddd5cf..760be97e15d5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -254,6 +254,7 @@ def _supports_routing_method( RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.SigmoidRenorm, + RoutingMethodType.MiniMax2, RoutingMethodType.Sigmoid, ] From fa423b614c5328913ffcb87eaf28817f5acd89f2 Mon Sep 17 00:00:00 2001 From: Aarushi Jain Date: Mon, 5 Oct 2026 18:39:59 -0500 Subject: [PATCH 028/952] [CI][ROCm] Wait for engine teardown between LM Eval models (#60100) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- .../test_lm_eval_correctness.py | 68 ++++++++++++++++++- 1 file changed, 66 insertions(+), 2 deletions(-) diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index 5576c84b05cb..e4c99adb7b40 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -8,16 +8,28 @@ --tp-size=1 """ +import gc import os +import time from contextlib import contextmanager import lm_eval import pytest +import torch import yaml +from lm_eval.api.registry import get_model +from lm_eval.utils import simple_parse_args_string +from vllm import LLM +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.logger import init_logger from vllm.platforms import current_platform +logger = init_logger(__name__) + DEFAULT_RTOL = 0.08 +ROCM_ENGINE_SHUTDOWN_TIMEOUT_S = 60.0 +MEMORY_RELEASE_TIMEOUT_S = 240.0 @contextmanager @@ -46,6 +58,48 @@ def scoped_env_vars(new_env: dict[str, str]): os.environ.pop(key, None) +def _wait_for_memory_release(gpu_memory_utilization: float) -> None: + devices = range(torch.accelerator.device_count()) + deadline = time.monotonic() + MEMORY_RELEASE_TIMEOUT_S + while True: + memory = [torch.accelerator.get_memory_info(d) for d in devices] + if all(free >= total * gpu_memory_utilization for free, total in memory): + return + if time.monotonic() >= deadline: + usage = ", ".join( + f"{d}: {free / 2**30:.2f}/{total / 2**30:.2f} GiB free" + for d, (free, total) in zip(devices, memory) + ) + logger.warning( + "GPU memory not released after %ss (%s); the next model needs %.2f", + MEMORY_RELEASE_TIMEOUT_S, + usage, + gpu_memory_utilization, + ) + return + time.sleep(1) + + +def _shutdown_lm(lm) -> None: + if not current_platform.is_rocm(): + return + llm = getattr(lm, "model", None) + if not isinstance(llm, LLM): + return + gpu_memory_utilization = ( + llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization + ) + try: + llm.llm_engine.engine_core.shutdown(timeout=ROCM_ENGINE_SHUTDOWN_TIMEOUT_S) + except Exception: + logger.exception("Engine core shutdown raised; GPU memory may leak") + del llm + lm.model = None + gc.collect() + cleanup_dist_env_and_memory() + _wait_for_memory_release(gpu_memory_utilization) + + def launch_lm_eval(eval_config, tp_size): trust_remote_code = eval_config.get("trust_remote_code", False) max_model_len = eval_config.get("max_model_len", 4096) @@ -82,9 +136,17 @@ def launch_lm_eval(eval_config, tp_size): env_vars = eval_config.get("env_vars", None) with scoped_env_vars(env_vars): + if current_platform.is_rocm(): + model = get_model(backend).create_from_arg_string( + model_args, + {"batch_size": batch_size, "max_batch_size": None, "device": None}, + ) + model_kwargs = {"metadata": simple_parse_args_string(model_args)} + else: + model = backend + model_kwargs = {"model_args": model_args} results = lm_eval.simple_evaluate( - model=backend, - model_args=model_args, + model=model, tasks=[task["name"] for task in eval_config["tasks"]], num_fewshot=eval_config["num_fewshot"], limit=eval_config["limit"], @@ -98,7 +160,9 @@ def launch_lm_eval(eval_config, tp_size): # Forward decoding and early-stop controls (e.g., max_gen_toks, until=...) gen_kwargs=eval_config.get("gen_kwargs"), batch_size=batch_size, + **model_kwargs, ) + _shutdown_lm(model) return results From 5cd8a2626b7a1378930e27ffcafcd5f109193d4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mai=20Ho=C3=A0ng=20=C4=90=E1=BB=A9c?= <80498645+Mazukiri@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:06:59 +0700 Subject: [PATCH 029/952] [Bugfix][Structured Output] Cap JSON schema nesting to prevent 500s and API-server stalls (#60036) Signed-off-by: Mazukiri Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../openai/chat_completion/test_chat_error.py | 30 +++++- tests/parser/cohere/test_structural_tags.py | 18 ++++ tests/parser/mistral/test_tool_calls.py | 14 +++ tests/parser/test_abstract_parser.py | 18 +++- tests/v1/structured_output/test_validation.py | 101 +++++++++++++++++- vllm/entrypoints/generate/base/protocol.py | 9 +- vllm/parser/abstract_parser.py | 8 +- vllm/parser/cohere_command.py | 4 +- vllm/parser/mistral.py | 3 +- vllm/sampling_params.py | 58 ++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 4 +- 11 files changed, 259 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index baa8bcec59ea..e467d1bcc5b0 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from dataclasses import dataclass, field from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -642,7 +643,15 @@ def test_batch_structural_tag_response_format_invalid(format_value): ) -@pytest.mark.parametrize("structural_tag", ["not json", ""]) +@pytest.mark.parametrize( + "structural_tag", + [ + "not json", + "", + # json.loads raises a plain ValueError, not JSONDecodeError + '{"type": "structural_tag", "x": ' + "1" * 5000 + "}", + ], +) def test_structured_outputs_structural_tag_invalid(structural_tag): """Malformed direct structured_outputs structural tags should be rejected.""" with pytest.raises( @@ -656,6 +665,25 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): ) +@pytest.mark.parametrize("field", ["response_format", "structured_outputs"]) +def test_deeply_nested_structural_tag_rejected(field): + """Structural tags are compiled during request parsing, so a deeply nested + one must be rejected there instead of blocking the server while converting, + and reported as too deep rather than malformed.""" + schema = '{"type": "array", "items": ' * 200 + "{}" + "}" * 200 + tag = ( + '{"type": "structural_tag", "format": {"type": "json_schema", ' + f'"json_schema": {schema}}}}}' + ) + value = json.loads(tag) if field == "response_format" else {"structural_tag": tag} + with pytest.raises(VLLMValidationError, match="nested too deeply"): + ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hello"}], + **{field: value}, + ) + + @pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"]) def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a diff --git a/tests/parser/cohere/test_structural_tags.py b/tests/parser/cohere/test_structural_tags.py index 893cfc80f710..149c5229a7f7 100644 --- a/tests/parser/cohere/test_structural_tags.py +++ b/tests/parser/cohere/test_structural_tags.py @@ -19,6 +19,7 @@ ) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.exceptions import VLLMValidationError from vllm.parser.cohere_command import ( CohereCommandParser, CohereNormalizedTool, @@ -339,6 +340,14 @@ def test_structured_outputs_json_string_raises( _make_chat_request(structured_outputs={"json": json_value}), ) + def test_deeply_nested_json_string_rejected(self, parser) -> None: + """json.loads raises RecursionError on a string nested this deeply.""" + schema = '{"type": "array", "items": ' * 20_000 + "{}" + "}" * 20_000 + with pytest.raises(VLLMValidationError, match="nested too deeply"): + parser.adjust_request( + _make_chat_request(structured_outputs={"json": schema}), + ) + @pytest.mark.parametrize( "construct", [ @@ -450,6 +459,15 @@ def test_draft07_definitions(self) -> None: self._tool_calls_json("get_weather", {"location": {"city": 1}}), ) + def test_deeply_nested_parameters_rejected(self) -> None: + """Tool schemas are compiled here, before request validation, and end up + inside an EBNF grammar that the nesting check cannot see later.""" + parameters = json.loads('{"type": "array", "items": ' * 200 + "{}" + "}" * 200) + with pytest.raises(VLLMValidationError, match="nested too deeply"): + collect_tool_schema( + [CohereNormalizedTool(name="deep", parameters=parameters)] + ) + def test_input_parameters_not_mutated(self) -> None: params = json.loads(json.dumps(NESTED_DEFS_PARAMETERS)) collect_tool_schema( diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py index 25b41f44c2d2..a3359dccc68e 100644 --- a/tests/parser/mistral/test_tool_calls.py +++ b/tests/parser/mistral/test_tool_calls.py @@ -35,6 +35,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) +from vllm.exceptions import VLLMValidationError from vllm.parser.engine.events import EventType from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine from vllm.parser.mistral import _DEFAULT_JSON_SCHEMA, MistralParser, mistral_config @@ -1890,6 +1891,19 @@ def test_adjust_request_structured_outputs_generates_grammar( assert len(result.structured_outputs.grammar) > 0 +def test_adjust_request_deeply_nested_json_string_rejected( + mistral_tool_parser: MistralToolParser, +) -> None: + """The schema string is parsed here, before request validation, and + json.loads raises RecursionError on a string nested this deeply.""" + schema = '{"type": "array", "items": ' * 20_000 + "{}" + "}" * 20_000 + request = _make_request( + structured_outputs=StructuredOutputsParams(json=schema), + ) + with pytest.raises(VLLMValidationError, match="nested too deeply"): + mistral_tool_parser.adjust_request(request) + + @pytest.mark.parametrize( "response_format_kwargs,expected_json_schema", [ diff --git a/tests/parser/test_abstract_parser.py b/tests/parser/test_abstract_parser.py index 56d20417d15f..7952ff4114e5 100644 --- a/tests/parser/test_abstract_parser.py +++ b/tests/parser/test_abstract_parser.py @@ -13,7 +13,8 @@ ChatCompletionRequest, ChatCompletionToolsParam, ) -from vllm.parser.abstract_parser import DelegatingParser +from vllm.exceptions import VLLMValidationError +from vllm.parser.abstract_parser import DelegatingParser, structured_outputs_to_format from vllm.sampling_params import StructuredOutputsParams from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser @@ -303,3 +304,18 @@ def test_required(self): assert _is_grammar_accept_string(grammar, self._qwen_tool_call()) assert not _is_grammar_accept_string(grammar, '{"text": "hi"}') mock_warn.assert_called_once() + + +@pytest.mark.parametrize("field", ["json", "structural_tag"]) +def test_structured_outputs_to_format_rejects_deeply_nested_string(field): + """Parsers convert the constraint before request validation, and json.loads + raises RecursionError on a string nested this deeply.""" + schema = '{"type": "array", "items": ' * 20_000 + "{}" + "}" * 20_000 + value = schema + if field == "structural_tag": + value = ( + '{"type": "structural_tag", "format": {"type": "json_schema", ' + f'"json_schema": {schema}}}}}' + ) + with pytest.raises(VLLMValidationError, match="nested too deeply"): + structured_outputs_to_format(StructuredOutputsParams(**{field: value})) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 767153cac7cc..cd5526c7b760 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -2,11 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Request-time validation of structured output requests.""" +import json + import pytest from vllm.config import StructuredOutputsConfig from vllm.exceptions import VLLMClientError, VLLMValidationError -from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.sampling_params import ( + MAX_STRUCTURED_OUTPUT_JSON_NESTING, + STRUCTURAL_TAG_WRAPPER_NESTING, + SamplingParams, + StructuredOutputsParams, +) pytestmark = pytest.mark.cpu_test @@ -110,6 +117,98 @@ def test_regex_with_nul_byte_rejected(regex): validate_xgrammar_grammar(params) +def _nested_array_schema(levels: int) -> str: + # Built as text: json.dumps on a deeply nested dict would itself recurse. + wrappers = levels - 1 + return ( + '{"type": "array", "items": ' * wrappers + '{"type": "string"}' + "}" * wrappers + ) + + +def _structural_tag(schema: str) -> str: + return ( + '{"type": "structural_tag", "format": {"type": "json_schema", ' + f'"json_schema": {schema}}}}}' + ) + + +@pytest.mark.parametrize( + "structured_outputs", + [ + pytest.param( + {"json": _nested_array_schema(MAX_STRUCTURED_OUTPUT_JSON_NESTING + 1)}, + id="json-str-over-limit", + ), + pytest.param( + {"json": json.loads(_nested_array_schema(2_000))}, + id="json-dict-past-recursion-limit", + ), + pytest.param( + {"json": _nested_array_schema(200_000)}, + id="json-str-past-json-loads-recursion-limit", + ), + pytest.param( + { + "structural_tag": _structural_tag( + _nested_array_schema( + MAX_STRUCTURED_OUTPUT_JSON_NESTING + + STRUCTURAL_TAG_WRAPPER_NESTING + ) + ) + }, + id="structural-tag-over-limit", + ), + ], +) +@pytest.mark.parametrize( + "backend", ["auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"] +) +def test_deeply_nested_schema_rejected(structured_outputs, backend): + """Converting a schema to a grammar gets steeply more expensive with nesting + and runs on the API server; past the recursion limit it used to escape as a + RecursionError (HTTP 500 with an empty message). Deep nesting must be a + validation error before any backend touches the schema.""" + params = SamplingParams( + structured_outputs=StructuredOutputsParams(**structured_outputs) + ) + with pytest.raises(VLLMValidationError, match="nested too deeply"): + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend=backend), + tokenizer=object(), + ) + + +@pytest.mark.parametrize("field", ["json", "structural_tag"]) +def test_xgrammar_validation_rejects_deeply_nested_schema(field): + """Request parsing and tool parsers call validate_xgrammar_grammar directly, + before _validate_structured_outputs, so it must check nesting itself.""" + from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar + + schema = _nested_array_schema(2_000) + value = _structural_tag(schema) if field == "structural_tag" else schema + params = SamplingParams( + structured_outputs=StructuredOutputsParams(**{field: value}) + ) + with pytest.raises(VLLMValidationError, match="nested too deeply"): + validate_xgrammar_grammar(params) + + +@pytest.mark.parametrize("field", ["json", "structural_tag"]) +def test_schema_at_nesting_limit_accepted(field): + """A structural tag wraps its schema in a few more levels (more for the + tags tool parsers generate), which must not count against the schema.""" + schema = _nested_array_schema(MAX_STRUCTURED_OUTPUT_JSON_NESTING) + value = _structural_tag(schema) if field == "structural_tag" else schema + structured_outputs = StructuredOutputsParams(**{field: value}) + SamplingParams(structured_outputs=structured_outputs)._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(), + tokenizer=object(), + ) + assert structured_outputs._backend == "xgrammar" + + INVALID_JSON_SCHEMA = {"type": "object", "properties": {"name": {"type": "str"}}} diff --git a/vllm/entrypoints/generate/base/protocol.py b/vllm/entrypoints/generate/base/protocol.py index b858ea19a5c5..89f6816c41cb 100644 --- a/vllm/entrypoints/generate/base/protocol.py +++ b/vllm/entrypoints/generate/base/protocol.py @@ -230,7 +230,11 @@ def validate_structural_tag_response_format( def validate_structural_tag_payload(payload: Any, *, parameter: str) -> None: - from vllm.sampling_params import SamplingParams, StructuredOutputsParams + from vllm.sampling_params import ( + SamplingParams, + StructuredOutputsParams, + check_json_nesting, + ) from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar if isinstance(payload, str) and not payload: @@ -239,6 +243,9 @@ def validate_structural_tag_payload(payload: Any, *, parameter: str) -> None: parameter=parameter, ) + if isinstance(payload, str): + # Raised here so the error is not reported as a malformed tag below + check_json_nesting(payload, structural_tag=True) try: validate_xgrammar_grammar( SamplingParams( diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f6b009efde07..c28563fabf65 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -42,7 +42,11 @@ from vllm.parser.metrics import record_tool_parser_invocation from vllm.parser.utils import count_history_tool_calls from vllm.reasoning.abs_reasoning_parsers import ReasoningParser -from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.sampling_params import ( + SamplingParams, + StructuredOutputsParams, + check_json_nesting, +) from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.streaming import ( @@ -274,6 +278,7 @@ def structured_outputs_to_format(params: StructuredOutputsParams) -> Format | No if params.json is not None: schema = params.json if isinstance(schema, str): + check_json_nesting(schema) schema = json.loads(schema) return JSONSchemaFormat(json_schema=schema) if params.regex is not None: @@ -295,6 +300,7 @@ def structured_outputs_to_format(params: StructuredOutputsParams) -> Format | No raise VLLMValidationError("Invalid grammar specification.") from e return GrammarFormat(grammar=grammar) if params.structural_tag is not None: + check_json_nesting(params.structural_tag, structural_tag=True) s_tag = json.loads(params.structural_tag) if "structures" in s_tag: # LegacyStructuralTagResponseFormat diff --git a/vllm/parser/cohere_command.py b/vllm/parser/cohere_command.py index a1df4914011d..97a9d3d6f3da 100644 --- a/vllm/parser/cohere_command.py +++ b/vllm/parser/cohere_command.py @@ -45,7 +45,7 @@ BaseCohereCommandReasoningParser, ) from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY -from vllm.sampling_params import StructuredOutputsParams +from vllm.sampling_params import StructuredOutputsParams, check_json_nesting from vllm.tool_parsers.cohere_command_tool_parser import BaseCohereCommandToolParser if TYPE_CHECKING: @@ -127,6 +127,7 @@ def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: tool_dictionary: dict[str, str] = {} for tool in tool_schema: tool_name = tool["name"] + check_json_nesting(tool["parameters"]) tool_parameters = dict(tool["parameters"]) json_schema: dict[str, Any] = { "type": "object", @@ -404,6 +405,7 @@ def _schema_dict_from_structured_outputs( if isinstance(raw, str): if not raw.strip(): raise ValueError("structured_outputs.json cannot be empty.") + check_json_nesting(raw) try: raw = json.loads(raw) except json.JSONDecodeError as e: diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index 247ef2803fbe..61f133598c8c 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -48,7 +48,7 @@ ParserState, Transition, ) -from vllm.sampling_params import StructuredOutputsParams +from vllm.sampling_params import StructuredOutputsParams, check_json_nesting from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: @@ -381,6 +381,7 @@ def adjust_request( json_schema = _DEFAULT_JSON_SCHEMA elif request.structured_outputs.json is not None: if isinstance(request.structured_outputs.json, str): + check_json_nesting(request.structured_outputs.json) json_schema = json.loads(request.structured_outputs.json) else: json_schema = request.structured_outputs.json diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index ab9f2711142c..8a37c5a3fd17 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -33,6 +33,15 @@ """Upper bound on `SamplingParams.logprob_token_ids` list length. Must match the per-request row width allocated by the sampler's `LogprobTokenIdsState`.""" +MAX_STRUCTURED_OUTPUT_JSON_NESTING = 64 +"""Upper bound on nested objects and arrays in a structured-output JSON schema +or structural tag. Grammar conversion runs on the API server and its cost grows +steeply with nesting. OpenAI allows 10 levels of schema nesting.""" + +STRUCTURAL_TAG_WRAPPER_NESTING = 16 +"""Extra nesting allowed for a structural tag, which wraps its JSON schemas in +tag structure (up to 15 levels for tool-call tags generated by tool parsers).""" + def _verify_num_sequences(value: int, parameter_name: str) -> None: if not isinstance(value, int): @@ -202,6 +211,52 @@ class RequestOutputKind(Enum): FINAL_ONLY = 2 +def _exceeds_json_nesting(spec: str | dict[str, Any], limit: int) -> bool: + """Whether a JSON document nests objects and arrays deeper than `limit`.""" + if isinstance(spec, str): + try: + spec = json_mod.loads(spec) + except RecursionError: + return True + except ValueError: + # Malformed JSON is reported by the backend validation. + return False + + level: list[Any] = [spec] + for _ in range(limit): + level = [ + child + for node in level + if isinstance(node, (dict, list)) + for child in (node.values() if isinstance(node, dict) else node) + if isinstance(child, (dict, list)) + ] + if not level: + return False + return True + + +def check_json_nesting( + spec: str | dict[str, Any] | None, *, structural_tag: bool = False +) -> None: + """Reject a JSON schema nested more than `MAX_STRUCTURED_OUTPUT_JSON_NESTING` + levels deep, or a structural tag nested deeper than that plus + `STRUCTURAL_TAG_WRAPPER_NESTING`. + + Raises: + VLLMValidationError: If `spec` is nested too deeply. + + """ + limit = MAX_STRUCTURED_OUTPUT_JSON_NESTING + if structural_tag: + limit += STRUCTURAL_TAG_WRAPPER_NESTING + if spec and _exceeds_json_nesting(spec, limit): + raise VLLMValidationError( + "The JSON schema is nested too deeply: objects and arrays may nest at " + f"most {MAX_STRUCTURED_OUTPUT_JSON_NESTING} levels" + ) + + def _is_non_tekken_mistral(tokenizer: TokenizerLike) -> bool: return is_mistral_tokenizer(tokenizer) and not tokenizer.is_tekken @@ -1230,6 +1285,9 @@ def _validate_structured_outputs( raise VLLMValidationError( "structured_outputs.json cannot be an empty string" ) + # Reject deep nesting before any backend walks or converts the schema + check_json_nesting(self.structured_outputs.json) + check_json_nesting(self.structured_outputs.structural_tag, structural_tag=True) # Reject json_object=False early to avoid engine-side crashes if self.structured_outputs.json_object is False: raise VLLMValidationError( diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 337b2285f0c8..4697b00487fe 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -10,7 +10,7 @@ import vllm.envs from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import SamplingParams, check_json_nesting from vllm.utils.import_utils import LazyLoader from vllm.utils.mistral import is_mistral_tokenizer from vllm.v1.structured_output.backend_types import ( @@ -392,6 +392,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: return if so_params.json: + check_json_nesting(so_params.json) if isinstance(so_params.json, str): try: schema = json.loads(so_params.json) @@ -426,6 +427,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: return if so_params.structural_tag: + check_json_nesting(so_params.structural_tag, structural_tag=True) try: s_tag = json.loads(so_params.structural_tag) From b2863c0de953204b64e4e3d184f96eac26bd5b46 Mon Sep 17 00:00:00 2001 From: Brandon Zhang <31413216+brandonfzhang@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:13:55 -0700 Subject: [PATCH 030/952] [Kimi-K3] Use the fused K/V pack kernel on the decode-context-parallel prefill path (#56881) Signed-off-by: Brandon Zhang <31413216+brandonfzhang@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../kimi_k3/test_mla_prefill_context.py | 114 ++++++++++++++++++ .../layers/attention/mla_attention.py | 21 +++- vllm/models/kimi_k3/nvidia/mla.py | 14 ++- 3 files changed, 142 insertions(+), 7 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index 6c6e4bf20e87..e8cfb014f63c 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -13,8 +13,10 @@ import pytest import torch +import vllm.model_executor.layers.attention.mla_attention as mla_attention from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBaseImpl, + MLACommonImpl, MLACommonPrefillMetadata, build_mla_chunked_context_metadata, ) @@ -317,3 +319,115 @@ def test_fused_context_rejects_an_unquantized_query() -> None: ) with pytest.raises(AssertionError, match="new-token epilogue"): layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) + + +# -------------------------------------------------------------------------- +# Decode context parallelism: the generic impl keeps its loop (allgather + +# reorg) and only the per-chunk K/V pack is handed to the layer's fused kernel. +# -------------------------------------------------------------------------- +_DCP = 2 +_DCP_TOKENS = 48 + + +def _run_dcp_context_loop(monkeypatch, *, q_data_type, fused_mla_kv_concat_fn): + """Drive ``MLACommonImpl._context_parallel_compute_prefill_context`` on one + chunk with the collective pieces (gather, allgather, reorg) stubbed out.""" + torch.manual_seed(0) + device = torch.device("cuda") + kv_b_proj = _KVBProj(device, weight_dtype=torch.bfloat16) + workspace = torch.randn( + (_DCP_TOKENS * (_DCP + 1), _ENTRY), dtype=torch.bfloat16, device=device + ) + impl = SimpleNamespace( + kv_cache_dtype="auto", + kv_lora_rank=_KV_LORA_RANK, + qk_rope_head_dim=_QK_ROPE, + qk_nope_head_dim=_QK_NOPE, + v_head_dim=_V_HEAD_DIM, + num_heads=_NUM_HEADS, + kv_b_proj=kv_b_proj, + _use_flashinfer_concat_mla_k=False, + ) + impl._concat_k_nope_k_pe = MLACommonImpl._concat_k_nope_k_pe.__get__(impl) + + monkeypatch.setattr(mla_attention.ops, "cp_gather_cache", lambda **_: None) + monkeypatch.setattr(mla_attention, "_get_kv_b_proj_input_dtype", lambda *_: None) + monkeypatch.setattr( + mla_attention, + "reorg_kvcache", + lambda kv_c, k_pe, **_: (kv_c[:_DCP_TOKENS].squeeze(1), k_pe[:_DCP_TOKENS]), + ) + chunk = SimpleNamespace( + index=0, + padded_local_seq_lens=[_DCP_TOKENS], + local_context_lens_allranks=[[_DCP_TOKENS // _DCP] * _DCP], + padded_local_cu_seq_lens=torch.tensor([0, _DCP_TOKENS], device=device), + padded_local_token_to_seq=torch.zeros( + _DCP_TOKENS, dtype=torch.int32, device=device + ), + local_starts=[0], + num_local_context_tokens=_DCP_TOKENS, + request_slice=slice(0, 1), + num_requests=1, + starts=torch.tensor([0], device=device), + num_context_tokens=_DCP_TOKENS, + max_seq_len=_DCP_TOKENS, + token_slice=slice(0, 4), + ) + backend = _RecordingPrefillBackend() + attn_metadata = SimpleNamespace( + prefill=SimpleNamespace( + prefill_backend=backend, + chunked_context=SimpleNamespace( + workspace=workspace, + chunks=[chunk], + empty_token_slices=False, + dcp_manager=SimpleNamespace(kv_gather=lambda dst, src: None), + ), + q_data_type=q_data_type, + block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), + ) + ) + q = torch.randn( + (4, _NUM_HEADS, _QK_NOPE + _QK_ROPE), dtype=torch.bfloat16, device=device + ) + out, lse = MLACommonImpl._context_parallel_compute_prefill_context( + impl, + q, + torch.empty(0, device=device), + attn_metadata, + k_scale=torch.ones(1, device=device), + dcp_world_size=_DCP, + fused_mla_kv_concat_fn=fused_mla_kv_concat_fn, + ) + return backend, out, lse + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@torch.inference_mode() +def test_dcp_context_fused_pack_matches_generic_tail( + monkeypatch, kv_cache_dtype: str +) -> None: + """Under DCP the layer's fused K/V pack must feed the prefill backend exactly + what the generic impl's cast + split + concat tail feeds it.""" + q_data_type = ( + current_platform.fp8_dtype() if kv_cache_dtype == "fp8" else torch.bfloat16 + ) + layer = SimpleNamespace(qk_nope_head_dim=_QK_NOPE, v_head_dim=_V_HEAD_DIM) + fused_fn = MultiHeadLatentAttention._fused_mla_kv_concat.__get__(layer) + + backend_ref, ref_out, ref_lse = _run_dcp_context_loop( + monkeypatch, q_data_type=q_data_type, fused_mla_kv_concat_fn=None + ) + backend_fused, fused_out, fused_lse = _run_dcp_context_loop( + monkeypatch, q_data_type=q_data_type, fused_mla_kv_concat_fn=fused_fn + ) + assert len(backend_ref.calls) == len(backend_fused.calls) == 1 + for name, fused_t, ref_t in zip( + ("q", "k", "v"), backend_fused.calls[0], backend_ref.calls[0], strict=True + ): + torch.testing.assert_close( + fused_t, ref_t, atol=0, rtol=0, msg=lambda m, n=name: f"{n} differs: {m}" + ) + torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0) + torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b952a36821e0..00e70d6fd0ae 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -3033,7 +3033,13 @@ def _context_parallel_compute_prefill_context( attn_metadata: MLACommonMetadata, k_scale: torch.Tensor, dcp_world_size: int, + fused_mla_kv_concat_fn: Callable[ + [torch.Tensor, torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor] + ] + | None = None, ): + # fused_mla_kv_concat_fn(kv_nope, k_pe, use_fp8_prefill) -> (k, v) replaces + # the per-chunk cast + split + concat below with one fused kernel. assert attn_metadata.prefill is not None prefill_metadata = attn_metadata.prefill assert prefill_metadata.prefill_backend is not None @@ -3127,11 +3133,16 @@ def _context_parallel_compute_prefill_context( kv_nope = self.kv_b_proj(kv_c_normed)[0].view( -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim ) - if use_fp8_prefill: - kv_nope = kv_nope.to(prefill_metadata.q_data_type) - k_pe = k_pe.to(prefill_metadata.q_data_type) - k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k = self._concat_k_nope_k_pe(k_nope, k_pe) + if fused_mla_kv_concat_fn is not None: + k, v = fused_mla_kv_concat_fn(kv_nope, k_pe, use_fp8_prefill) + else: + if use_fp8_prefill: + kv_nope = kv_nope.to(prefill_metadata.q_data_type) + k_pe = k_pe.to(prefill_metadata.q_data_type) + k_nope, v = kv_nope.split( + [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + k = self._concat_k_nope_k_pe(k_nope, k_pe) attn_output, attn_softmax_lse = ( prefill_metadata.prefill_backend.run_prefill_context_chunk( diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 188f72e66016..bf7c53e9e515 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -785,6 +785,15 @@ def _decode_concat_cache( cos_sin_cache=cos_sin_cache, ) + def _fused_mla_kv_concat( + self, kv_nope: torch.Tensor, k_pe: torch.Tensor, use_fp8_prefill: bool + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pack one DCP context chunk's ``(k, v)`` with the fused kernels.""" + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + if use_fp8_prefill: + return fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) + return fused_mla_kv_concat(k_nope, k_pe), v + def _compute_prefill_context( self, q: torch.Tensor, @@ -814,8 +823,8 @@ def _compute_prefill_context( copied per chunk. Decode context parallelism keeps using - ``impl._context_parallel_compute_prefill_context``; its extra allgather - and reorg are not fused here. + ``impl._context_parallel_compute_prefill_context`` for its allgather and + reorg, with this layer's ``_fused_mla_kv_concat`` as its per-chunk pack. """ prefill = attn_metadata.prefill assert prefill is not None @@ -1074,6 +1083,7 @@ def _forward_prefill_fused( attn_metadata, k_scale=self._k_scale, dcp_world_size=self.dcp_world_size, + fused_mla_kv_concat_fn=self._fused_mla_kv_concat, ) ) else: From 97bd8c74ebf12c3d847fb59f24966a83bf6f6920 Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Mon, 5 Oct 2026 17:57:46 -0700 Subject: [PATCH 031/952] [ROCm][DSv4.1][Perf] One-pass dequantize and a gather-sized grid for the compressed K cache (#56720) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/test_compressor_kv_cache.py | 74 +++++++++++ vllm/models/deepseek_v41/amd/rocm.py | 5 + .../deepseek_v41/common/ops/cache_utils.py | 118 +++++++++++++++++- 3 files changed, 195 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 020dd607b6af..cc4f4c72a931 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -405,6 +405,80 @@ def test_v41_rope_insert_mxfp8_record(compress_ratio: int): torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="needs a CUDA or ROCm device" +) +@pytest.mark.parametrize("gather_len", [None, 70, 200]) +@pytest.mark.parametrize("one_pass", [False, True]) +def test_v41_fp8_gather_is_grid_width_invariant(gather_len, one_pass, monkeypatch): + """The gather grid only partitions tokens, so its width cannot move the + result, and the rows it writes must match a host-side dequantization. + + Covers the fp8_ds_mla record (584 B); the MXFP8 one is covered below. + Both dequantization forms are checked against the same host reference on + whatever device is present: the wrapper picks the one-pass form on ROCm + alone, so otherwise each platform would only ever see one of them. + """ + monkeypatch.setattr(current_platform, "is_rocm", lambda: one_pass) + from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + dequantize_and_gather_k_cache_triton, + ) + + torch.manual_seed(7) + device = "cuda" + block_size, num_blocks = 64, 6 + num_tokens = block_size * num_blocks + + k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + # insert takes [num_blocks, block_bytes]; gather takes the record view. + cache = torch.zeros(num_blocks, block_size * 584, dtype=torch.uint8, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + quantize_and_insert_k_cache(k, cache, slot_mapping, block_size=block_size) + cache_view = cache.view(num_blocks, block_size, 584) + + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gls = ( + None + if gather_len is None + else torch.tensor([gather_len], dtype=torch.int32, device=device) + ) + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(1, -1) + + def gather(bound): + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache_triton( + out, + cache_view, + seq_lens, + gls, + block_table, + block_size, + 0, + max_gather_len=bound, + ) + return out + + # None keeps the historical 128; the rest make the wrapper pick other widths + # for the same work, and the output must not notice. + baseline = gather(None) + for bound in (8, num_tokens, 1 << 20): + torch.testing.assert_close(gather(bound), baseline, rtol=0, atol=0) + + n = num_tokens if gather_len is None else gather_len + start = num_tokens - n + for i in (0, 1, n // 2, n - 1): + values, scales = _ue8m0_reference(k[start + i, :448], 64, 448.0) + nope = (values.to(torch.float32).view(7, 64) * scales.view(7, 1)).reshape(448) + torch.testing.assert_close( + baseline[0, i, :448], nope.to(torch.bfloat16), rtol=4e-3, atol=1e-6 + ) + # The RoPE dims are stored verbatim, so they come back exactly. + torch.testing.assert_close( + baseline[0, i, 448:], k[start + i, 448:], rtol=0, atol=0 + ) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") def test_v41_mxfp8_cache_round_trip(): """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record. diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 99fcfa6702b5..64b4a9d3d1f2 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -1387,6 +1387,9 @@ def _forward_prefill( gather_lens = swa_metadata.prefill_gather_lens assert seq_lens is not None assert gather_lens is not None + # CPU copy, so sizing the gather grid from it does not sync. + seq_lens_cpu = swa_metadata.prefill_seq_lens_cpu + assert seq_lens_cpu is not None query_start_loc_cpu = swa_metadata.query_start_loc_cpu query_start_loc = swa_metadata.query_start_loc @@ -1438,6 +1441,8 @@ def _forward_prefill( block_size=attn_metadata.block_size // self.compress_ratio, offset=0, use_fnuz=False, + max_gather_len=int(seq_lens_cpu[chunk_start:chunk_end].max()) + // self.compress_ratio, ) swa_block_table = swa_metadata.block_table[num_decodes:] diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index ad03c3294524..b7f07d8d84c0 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -33,7 +33,7 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl -from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.math_utils import cdiv, next_power_of_2 # Per-token byte width of the two paged fp8 records. Both put a page's whole # data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the @@ -436,6 +436,95 @@ def _dequantize_and_gather_k_kernel( tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) +@triton.jit +def _dequantize_and_gather_k_one_pass_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + # Constants + max_blocks_per_seq: tl.constexpr, + fp8_dim: tl.constexpr, # 448 + bf16_dim: tl.constexpr, # 64 + scale_dim: tl.constexpr, # 8 + quant_block: tl.constexpr, # 64 (quantization block size) + cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size) + token_data_size: tl.constexpr, # 576 bytes per token data + block_stride: tl.constexpr, # total bytes per block (padded) int32 + output_dim: tl.constexpr, # 512 + fp8_max: tl.constexpr, + n_quant_blocks: tl.constexpr, # 7 real blocks; the per-block form only + use_fnuz: tl.constexpr = False, +): + """``_dequantize_and_gather_k_kernel`` with the token read in one tile. + + The seven quantization blocks tile the fp8 region exactly, so the per-block + form's bounds never bite and its twenty-two memory operations per token + reduce to four. Selected on ROCm; CUDA keeps the per-block kernel. + """ + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + + token_fp8_ptr = cache_block_ptr + pos_in_block * token_data_size + token_bf16_ptr = token_fp8_ptr + fp8_dim + token_scale_ptr = ( + cache_block_ptr + + cache_block_size * token_data_size + + pos_in_block * scale_dim + ) + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + + # The tile is scale_dim * quant_block = 512 against an fp8 region of + # 448: the tail reads the token's own bf16 bytes, is scaled by a scale + # nothing reads, and is dropped by the store mask. + tile_dim: tl.constexpr = scale_dim * quant_block + tl.static_assert(tile_dim <= token_data_size, "tile outruns the token") + fp8_offsets = tl.arange(0, tile_dim) + x_uint8 = tl.load(token_fp8_ptr + fp8_offsets) + if use_fnuz: + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + + encoded_scale = tl.load(token_scale_ptr + tl.arange(0, scale_dim)) + scale = tl.exp2(encoded_scale.to(tl.float32) - 127.0) # UE8M0 + tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) + x_dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (tile_dim,)) + tl.store( + output_row_ptr + fp8_offsets, + x_dequant.to(tl.bfloat16), + mask=fp8_offsets < fp8_dim, + ) + + bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) + bf16_offsets = tl.arange(0, bf16_dim) + tl.store( + output_row_ptr + fp8_dim + bf16_offsets, + tl.load(bf16_cache_ptr + bf16_offsets), + ) + + @triton.jit def _dequantize_and_gather_k_mxfp8_kernel( out_ptr, @@ -566,6 +655,22 @@ def _dequantize_and_gather_k_nvfp4_kernel( tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) +# The per-token loop is a chain of dependent loads, so the grid has to hold +# enough chains in flight to cover the latency. The cap binds above 64k tokens, +# where a worker takes more than the eight steps this aims for. +_GATHER_STEPS_PER_WORKER = 8 +_GATHER_MIN_WORKERS = 128 +_GATHER_MAX_WORKERS = 8192 + + +def gather_num_workers(max_gather_len: int | None) -> int: + """Grid width for a gather of at most ``max_gather_len`` tokens.""" + if max_gather_len is None: + return _GATHER_MIN_WORKERS + workers = cdiv(max_gather_len, _GATHER_STEPS_PER_WORKER) + return max(_GATHER_MIN_WORKERS, min(_GATHER_MAX_WORKERS, workers)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -580,6 +685,7 @@ def dequantize_and_gather_k_cache_triton( block_size: int, offset: int, use_fnuz: bool = False, + max_gather_len: int | None = None, ) -> None: num_reqs = seq_lens.shape[0] NUM_WORKERS = 128 @@ -633,7 +739,13 @@ def dequantize_and_gather_k_cache_triton( FP8_MAX = 448.0 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( + gather_workers = gather_num_workers(max_gather_len) + kernel = ( + _dequantize_and_gather_k_one_pass_kernel + if current_platform.is_rocm() + else _dequantize_and_gather_k_kernel + ) + kernel[(num_reqs, gather_workers)]( out, out.stride(0), out.stride(1), @@ -671,6 +783,7 @@ def dequantize_and_gather_k_cache( block_size: int, offset: int, use_fnuz: bool = False, + max_gather_len: int | None = None, ) -> None: """Dequantize and gather a paged DSv4 K cache. @@ -711,6 +824,7 @@ def dequantize_and_gather_k_cache( block_table, block_size, offset, + max_gather_len=max_gather_len, use_fnuz=use_fnuz, ) From 4e198cef835edc6c6ce429be5a8f19d5aae308da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Tue, 6 Oct 2026 03:19:01 +0200 Subject: [PATCH 032/952] [TEST][XPU][CI] Fix ci uva kernel test instability (#60041) Signed-off-by: Your Name Co-authored-by: Your Name Co-authored-by: Kunshang Ji --- tests/kernels/core/test_uva.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index b7c2984d6a23..0a183207c859 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -34,6 +34,7 @@ def test_cpu_write(device): cpu_tensor[4, 5] = -1 gpu_view.mul_(2) + torch.accelerator.synchronize(gpu_view.device) assert gpu_view[0, 0] == 2 assert gpu_view[2, 3] == 4 assert gpu_view[4, 5] == -2 @@ -55,6 +56,7 @@ def test_gpu_write(device): gpu_view[2, 3] = 2 gpu_view[4, 5] = -1 gpu_view.mul_(2) + torch.accelerator.synchronize(gpu_view.device) assert cpu_tensor[0, 0] == 2 assert cpu_tensor[2, 3] == 4 From 0e8f8ecba838eb96fc4cb036429dab13108c6c37 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 5 Oct 2026 21:23:20 -0400 Subject: [PATCH 033/952] [Bugfix][KVConnector][NIXL] Return NumPy descriptors from mixed-memory local registration (#60108) Signed-off-by: Lucas Wilkinson Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Matthew Bonanni --- tests/v1/kv_connector/unit/test_nixl_connector.py | 4 +++- .../kv_transfer/kv_connector/v1/nixl/base_worker.py | 11 +++-------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d7ab2da55221..43a3c7030a83 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2100,7 +2100,9 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert handle == 22 assert worker._dram_src_handles_by_block_size[worker.block_size] == 11 - assert [block[2] for block in blocks] == [0, 0, 3, 3] + # Must stay an array: P_TP > D_TP splits call .tolist() on it. + assert isinstance(blocks, np.ndarray) + assert blocks[:, 2].tolist() == [0, 0, 3, 3] memory_types = [ call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index dffd70643e52..a9a0d1980a13 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2174,14 +2174,9 @@ def register_local_xfer_handler( self._desc_pos_by_block_size[block_size] = desc_pos # DRAM descriptors are registered under CPU device 0. - blocks_data = [ - (addr, length, 0) if is_dram else (addr, length, dev) - for (addr, length, dev), is_dram in zip( - blocks_data, desc_is_dram, strict=True - ) - ] - dram_blocks = [blocks_data[i] for i in dram_idx] - vram_blocks = [blocks_data[i] for i in vram_idx] + blocks_data[desc_is_dram, 2] = 0 + dram_blocks = blocks_data[dram_idx] + vram_blocks = blocks_data[vram_idx] dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") self._dram_src_handles_by_block_size[block_size] = ( self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) From 2537629eb22a5e3134cb8b09d367e07efbbfdf00 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 5 Oct 2026 21:27:50 -0400 Subject: [PATCH 034/952] [Bugfix][KVConnector][NIXL] Size per-region replicate flags by each region's block count (#60107) Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Matthew Bonanni --- .../kv_connector/v1/nixl/base_worker.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index a9a0d1980a13..dade5056eeb0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -340,7 +340,7 @@ def _build_local_splits_from_plan( ) # Per-FA-descriptor replicate flag, in _build_fa_local emission order. - fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs, block_size_ratio) sharded_desc_end = len(src_blocks_data) - ( self._logical_num_blocks if self._ple_region_index is not None else 0 ) @@ -383,19 +383,22 @@ def _needs_split_local_xfer_handles(self, tp_ratio: int, plan: TPMapping) -> boo """ return tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1) - def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: - """Per-FA-descriptor replicate flag, in _build_fa_local emission order - (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``. + def _fa_desc_replicated( + self, num_fa_descs: int, block_size_ratio: int = 1 + ) -> list[bool]: + """Per-FA-descriptor replicate flag, in _build_fa_local emission order: + region-major, ``region_num_blocks[i] * block_size_ratio`` descs per + region. Length ``num_fa_descs``. """ assert self.transfer_topo is not None n_regions = len(self.block_len_per_layer) if n_regions == 0 or self.num_regions == 0: return [False] * num_fa_descs - nblk = num_fa_descs // self.num_regions + region_indices = self._transfer_layer_region_indices or range(n_regions) flags: list[bool] = [] - for i in range(n_regions): - replicated = self._is_region_replicated(i) - flags.extend([replicated] * nblk) + for i in region_indices: + num_descs = self.region_num_blocks[i] * block_size_ratio + flags.extend([self._is_region_replicated(i)] * num_descs) assert len(flags) == num_fa_descs, ( f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}" ) From efd001756bff47791113857ef262d4683f7f7b55 Mon Sep 17 00:00:00 2001 From: amd-xavierwang Date: Mon, 5 Oct 2026 21:37:18 -0400 Subject: [PATCH 035/952] [ROCm] Disable AITER attention query quantization on RDNA3 (#57750) Signed-off-by: xawang Co-authored-by: TJian --- vllm/v1/attention/backends/rocm_aiter_unified_attn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index d57b420cd7da..0d4fd5519484 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -184,8 +184,11 @@ def __init__( ) from aiter.ops.triton.unified_attention import unified_attention + from vllm.platforms.rocm import on_gfx11 + self.unified_attention = unified_attention - self.supports_quant_query_input = True + # RDNA3 has no low precision MMA, preserve query precision + self.supports_quant_query_input = not on_gfx11() def _split_kv_cache( self, kv_cache: torch.Tensor From 9bf4e510303d8d4cbdec504490b50f5f68762eca Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:37:20 -0500 Subject: [PATCH 036/952] [Bugfix][Frontend] Malformed `$defs` in a tool schema should be a 400, not a 500 (#54850) Signed-off-by: lzhan011 <35493221+lzhan011@users.noreply.github.com> Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/tool_use/test_tool_choice_required.py | 95 +++++++++++++++++++++ vllm/tool_parsers/utils.py | 13 ++- 2 files changed, 105 insertions(+), 3 deletions(-) diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index fc868bbcab6e..26978fbb7c0e 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -410,6 +410,101 @@ def test_get_json_schema_rejects_only_non_function_tools(self): get_json_schema_from_tools(tools=[WEB_SEARCH_TOOL], tool_choice="required") +class TestMalformedToolSchemaDefs: + """Malformed `$defs` in caller-supplied tool parameters is a 400, not a 500.""" + + @staticmethod + def _tool(params: dict) -> ChatCompletionToolsParam: + return TypeAdapter(ChatCompletionToolsParam).validate_python( + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": params, + }, + } + ) + + @pytest.mark.parametrize( + "defs", + [ + pytest.param(None, id="null"), + pytest.param([], id="list"), + pytest.param("nope", id="string"), + pytest.param(1, id="int"), + ], + ) + def test_non_object_defs_is_a_client_error(self, defs): + tools = [ + self._tool( + { + "type": "object", + "properties": {"a": {"type": "string"}}, + "$defs": defs, + } + ) + ] + + with pytest.raises(VLLMValidationError) as exc_info: + get_json_schema_from_tools(tools=tools, tool_choice="required") + + assert exc_info.value.parameter == "tools" + assert "$defs" in str(exc_info.value) + assert "get_weather" in str(exc_info.value) + + def test_duplicate_def_names_is_a_client_error(self): + """Also caller-caused, and it used to raise a plain `ValueError`.""" + tools = [ + self._tool( + { + "type": "object", + "properties": {"a": {"$ref": "#/$defs/D"}}, + "$defs": {"D": {"type": "string"}}, + } + ), + self._tool( + { + "type": "object", + "properties": {"b": {"$ref": "#/$defs/D"}}, + "$defs": {"D": {"type": "integer"}}, + } + ), + ] + + with pytest.raises(VLLMValidationError) as exc_info: + get_json_schema_from_tools(tools=tools, tool_choice="required") + + assert exc_info.value.parameter == "tools" + assert "multiple schemas" in str(exc_info.value) + + def test_well_formed_defs_still_hoisted(self): + """The type check must not reject the shape it is guarding.""" + tools = [ + self._tool( + { + "type": "object", + "properties": {"a": {"$ref": "#/$defs/D"}}, + "$defs": {"D": {"type": "string"}}, + } + ) + ] + + schema = get_json_schema_from_tools(tools=tools, tool_choice="required") + + assert isinstance(schema, dict) + assert schema["$defs"] == {"D": {"type": "string"}} + + def test_absent_defs_still_works(self): + tools = [ + self._tool({"type": "object", "properties": {"a": {"type": "string"}}}) + ] + + schema = get_json_schema_from_tools(tools=tools, tool_choice="required") + + assert isinstance(schema, dict) + + class TestParallelToolCallsConstraint: """`parallel_tool_calls=false` must be enforced by the decoding grammar. diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 03bcf1b652cc..86d68a9b5130 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -368,15 +368,22 @@ def _get_tool_schema_defs( ) -> dict: all_defs: dict[str, dict[str, Any]] = {} for tool in tools: - _, params = _extract_tool_info(tool) + name, params = _extract_tool_info(tool) if params is None: continue defs = params.pop("$defs", {}) + if not isinstance(defs, dict): + raise VLLMValidationError( + f"`$defs` in the parameters of tool '{name}' must be an " + f"object, got {type(defs).__name__}.", + parameter="tools", + ) for def_name, def_schema in defs.items(): if def_name in all_defs and all_defs[def_name] != def_schema: - raise ValueError( + raise VLLMValidationError( f"Tool definition '{def_name}' has multiple schemas, " - "which is not supported." + "which is not supported.", + parameter="tools", ) all_defs[def_name] = def_schema return all_defs From 9465e1c783ca3ead11af4bdf2d930a921cb35436 Mon Sep 17 00:00:00 2001 From: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com> Date: Tue, 6 Oct 2026 09:46:53 +0800 Subject: [PATCH 037/952] [Bugfix][NIXL] Recover locally invalidated pull peer metadata (#55471) Signed-off-by: xijiade.aihemaiti <3146335281@qq.com> Signed-off-by: Nick Hill Co-authored-by: OpenAI Codex Co-authored-by: Nick Hill Co-authored-by: Yannik Hinteregger <37209495+YannikHinteregger@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../kv_connector/unit/test_nixl_connector.py | 222 +++++++++++++++++- .../unit/test_nixl_connector_hma.py | 4 + .../unit/test_nixl_push_connector.py | 2 + .../kv_connector/v1/nixl/base_worker.py | 92 +++++++- .../kv_connector/v1/nixl/pull_worker.py | 26 +- 5 files changed, 314 insertions(+), 32 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 43a3c7030a83..35ac7d887ea4 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -14,7 +14,7 @@ from concurrent.futures import Future from types import SimpleNamespace from typing import Any, cast -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import msgspec import numpy as np @@ -144,6 +144,7 @@ class FakeNixlWrapper: def __init__(self, agent_name: str, *args, **kwargs): self._cycles_before_xfer_done = 0 self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(lambda: 0) + self.invalid_remote_agents: set[str] = set() def get_reg_descs(self, caches_data, memory_type: str) -> list: return [str(uuid.uuid4()) for _ in caches_data] @@ -164,6 +165,7 @@ def get_agent_metadata(self) -> bytes: return self.AGENT_METADATA def add_remote_agent(self, agent_metadata: bytes) -> str: + self.invalid_remote_agents.discard(self.REMOTE_AGENT_NAME) return self.REMOTE_AGENT_NAME def get_new_notifs(self) -> dict[str, list[bytes]]: @@ -185,6 +187,9 @@ def release_dlist_handle(self, handle: int) -> None: def remove_remote_agent(self, agent: str) -> None: pass + def check_remote_metadata(self, agent: str) -> bool: + return agent not in self.invalid_remote_agents + def send_notif(self, agent_name: str, notif_msg: bytes) -> None: pass @@ -2127,6 +2132,8 @@ def recv_worker(): worker._handshake_futures = {} worker._remote_agents = {} worker._engine_by_address = {} + worker._failed_remote_engines = set() + worker._invalid_remote_engines = set() worker._replicated_pcp_done_sending = set() worker._invalid_block_ids = queue.Queue() worker._pending_recv_notifs = {} @@ -2809,9 +2816,9 @@ def _connect(self, engine_id="new", host="localhost", port=1234): self.nixl.REMOTE_AGENT_NAME = engine_id future.set_result(self.worker._nixl_handshake(host, port, 1, engine_id)) - def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True): + def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True, req_id=None): metadata = NixlConnectorMetadata() - metadata.reqs_to_recv[f"{engine_id}-req"] = ReqMeta( + metadata.reqs_to_recv[req_id or f"{engine_id}-req"] = ReqMeta( local_block_ids=(blocks,), local_physical_block_ids=(blocks,), tp_size=1, @@ -2822,6 +2829,11 @@ def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True): ) return metadata + def _fail_old_peer(self): + self.worker.start_load_kv(self._request("old")) + self.nixl.invalid_remote_agents.add("old") + self.transport.check_xfer_state.return_value = "ERR" + def test_cleans_old_peer_after_confirmed_handshake(self): self._connect("healthy", host="other-host") old_handle = self.worker.dst_xfer_side_handles["old"][0] @@ -2871,7 +2883,12 @@ def test_preserves_old_peer_when_handshake_fails(self): ], ids=["healthy", "failed-sibling-completes", "failed-sibling-fails"], ) - def test_waits_for_reads_before_cleanup_and_failure_reporting(self, states, failed): + @pytest.mark.parametrize("invalidated", [False, True]) + def test_waits_for_reads_before_cleanup_and_failure_reporting( + self, states, failed, invalidated + ): + if invalidated: + self.nixl.invalid_remote_agents.add("old") self.worker._recving_metadata.update(self._request("old").reqs_to_recv) self.worker._recving_transfers["old-req"] = [1, 2] if failed else [1] self.transport.check_xfer_state.side_effect = states @@ -2890,6 +2907,7 @@ def test_waits_for_reads_before_cleanup_and_failure_reporting(self, states, fail ) assert self.transport.release_xfer_handle.call_count == (2 if failed else 1) self.transport.remove_remote_agent.assert_called_once_with("old") + assert not self.worker._invalid_remote_engines def test_waits_for_queued_notification(self): metadata = self._request("old", blocks=()) @@ -2910,22 +2928,31 @@ def test_waits_for_queued_notification(self): assert result.failed_recving == set() self.transport.remove_remote_agent.assert_called_once_with("old") - def test_waits_for_other_handshakes(self): + @pytest.mark.parametrize("invalidated", [False, True]) + def test_waits_for_other_handshakes(self, invalidated): + if invalidated: + self._fail_old_peer() + else: + self._connect() self.worker._ensure_handshake("other", "other-host", 1234, 1) - self._connect() self.worker.get_transfer_results() self.transport.remove_remote_agent.assert_not_called() + self.transport.check_remote_metadata.assert_not_called() self._connect("other", host="other-host") self.worker.get_transfer_results() self.transport.remove_remote_agent.assert_called_once_with("old") - def test_rehandshakes_queued_request_for_released_engine(self): + @pytest.mark.parametrize("invalidated", [False, True]) + def test_rehandshakes_queued_request_for_released_engine(self, invalidated): # A request queued after its engine's handshake, but drained only after # that engine was replaced and released, must not read from it. + if invalidated: + self._fail_old_peer() metadata = self._request("old", blocks=(), awaiting_kvs=False) self.worker._ready_requests.put(("old-req", metadata.reqs_to_recv["old-req"])) - self._connect() + if not invalidated: + self._connect() self.worker.get_transfer_results() self.transport.remove_remote_agent.assert_called_once_with("old") @@ -2933,6 +2960,185 @@ def test_rehandshakes_queued_request_for_released_engine(self): assert "old" in self.worker._handshake_futures self.transport.send_notif.assert_not_called() + @pytest.mark.parametrize("missing", [False, True, RuntimeError("probe failed")]) + def test_recovery_requires_confirmed_missing_metadata(self, missing): + self._fail_old_peer() + # A missing rank suffices even if another rank still has metadata. + self.worker._remote_agents["old"][(0, 1)] = "old-other-rank" + self.nixl.invalid_remote_agents.clear() + if missing is True: + self.nixl.invalid_remote_agents.add("old-other-rank") + elif isinstance(missing, Exception): + self.transport.check_remote_metadata.side_effect = missing + + result = self.worker.get_transfer_results() + assert result.finished_recving == result.failed_recving == {"old-req"} + assert self.worker.get_block_ids_with_load_errors() == {1} + assert ("old" not in self.worker._remote_agents) is (missing is True) + assert self.transport.release_dlist_handle.call_count == (missing is True) + + @pytest.mark.parametrize("handshake_fails", [False, True]) + def test_recovery_rehandshakes_same_id_with_fresh_descriptors( + self, handshake_fails + ): + old_handle = self.worker.dst_xfer_side_handles["old"][0] + self._fail_old_peer() + self.worker.get_transfer_results() + self.transport.release_dlist_handle.assert_called_once_with(old_handle) + assert not self.worker._engine_by_address + with pytest.raises(KeyError): + self.worker.transfer_topo.get_engine_info("old") + self.transport.make_prepped_xfer.reset_mock() + + self.worker.start_load_kv(self._request("old", req_id="next")) + self.transport.make_prepped_xfer.assert_not_called() + if handshake_fails: + self.worker._handshake_futures["old"].set_exception( + RuntimeError("handshake failed") + ) + result = self.worker.get_transfer_results() + assert result.failed_recving == result.finished_recving == {"next"} + assert "old" not in self.worker._remote_agents + assert self.worker._ready_requests.empty() + return + + self._connect("old") + fresh_handle = self.worker.dst_xfer_side_handles["old"][0] + assert fresh_handle != old_handle + assert self.worker._engine_by_address == {("localhost", 1234): "old"} + self.transport.check_xfer_state.return_value = "DONE" + self.worker.start_load_kv(NixlConnectorMetadata()) + assert self.transport.make_prepped_xfer.call_args.args[3] == fresh_handle + result = self.worker.get_transfer_results() + assert result.finished_recving == {"next"} and not result.failed_recving + self.transport.release_dlist_handle.assert_called_once_with(old_handle) + + def test_recovery_drains_reads_and_rejects_new_ones_without_harming_other_peers( + self, + ): + self._connect("healthy", host="other-host") + self.worker._recving_metadata.update(self._request("old").reqs_to_recv) + self.worker._recving_metadata.update(self._request("healthy").reqs_to_recv) + self.worker._recving_transfers.update( + {"old-req": [10, 11], "healthy-req": [12]} + ) + self.nixl.invalid_remote_agents.add("old") + self.transport.check_xfer_state.side_effect = ( + lambda h: "ERR" if h == 10 else "PROC" + ) + old_handle = self.worker.dst_xfer_side_handles["old"][0] + + result = self.worker.get_transfer_results() + assert not result.finished_recving and not result.failed_recving + self.transport.release_dlist_handle.assert_not_called() + self.worker.start_load_kv(self._request("old", req_id="blocked")) + self.transport.make_prepped_xfer.assert_not_called() + + self.transport.check_xfer_state.side_effect = ( + lambda h: "DONE" if h == 11 else "PROC" + ) + result = self.worker.get_transfer_results() + assert ( + result.finished_recving == result.failed_recving == {"old-req", "blocked"} + ) + assert set(self.worker._remote_agents) == {"healthy"} + assert self.worker._recving_transfers == {"healthy-req": [12]} + self.transport.remove_remote_agent.assert_called_once_with("old") + calls = self.transport.mock_calls + assert calls.index(call.release_xfer_handle(11)) < calls.index( + call.release_dlist_handle(old_handle) + ) + + def test_repeated_invalidation_backs_off_before_rehandshake(self): + # A peer that is invalidated again soon after recovery (eg a persistent + # RDMA fault) must not re-handshake on every request. + self._fail_old_peer() + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_called_once_with("old") + + self._connect("old") + self._fail_old_peer() + assert self.worker.get_transfer_results().failed_recving == {"old-req"} + assert "old" in self.worker._remote_agents + + self.transport.make_prepped_xfer.reset_mock() + self.worker.start_load_kv(self._request("old", req_id="next")) + assert self.worker.get_transfer_results().failed_recving == {"next"} + self.transport.make_prepped_xfer.assert_not_called() + assert "old" not in self.worker._handshake_futures + self.transport.remove_remote_agent.assert_called_once() + + backoff, release_at = self.worker._invalid_engine_backoff["old"] + assert backoff == 1.0 + self.worker._invalid_engine_backoff["old"] = (backoff, release_at - backoff) + self.worker.get_transfer_results() + assert self.transport.remove_remote_agent.call_count == 2 + assert "old" not in self.worker._remote_agents + + def test_recovery_waits_until_failed_handle_can_be_released(self): + self._fail_old_peer() + with patch.object( + self.transport, "release_xfer_handle", side_effect=RuntimeError("in use") + ): + result = self.worker.get_transfer_results() + assert not result.finished_recving and not result.failed_recving + assert self.worker._recving_transfers["old-req"] + self.transport.release_dlist_handle.assert_not_called() + + result = self.worker.get_transfer_results() + assert result.finished_recving == result.failed_recving == {"old-req"} + self.transport.remove_remote_agent.assert_called_once_with("old") + + @pytest.mark.parametrize("expired", [False, True]) + def test_success_and_expiry_do_not_probe_metadata(self, expired): + metadata = self._request("old") + if expired: + metadata.reqs_to_recv["old-req"].remote.blocks_expiry_time = 0.0 + self.worker._bidirectional_kv_xfer_enabled = True + with ( + patch.object(self.worker.xfer_stats, "record_failed_transfer") as failure, + patch.object(self.worker.xfer_stats, "record_kv_expired_req") as expiry, + ): + self.worker.start_load_kv(metadata) + result = self.worker.get_transfer_results() + assert result.finished_recving == {"old-req"} + assert result.failed_recving == ({"old-req"} if expired else set()) + assert expiry.call_count == int(expired) + failure.assert_not_called() + self.transport.check_remote_metadata.assert_not_called() + self.transport.remove_remote_agent.assert_not_called() + + @pytest.mark.parametrize("awaiting_kvs", [False, True]) + def test_invalid_peer_empty_recv_is_reported_only_when_awaited(self, awaiting_kvs): + self._fail_old_peer() + self.worker._recving_transfers["old-req"].append(11) + self.transport.check_xfer_state.side_effect = ( + lambda h: "PROC" if h == 11 else "ERR" + ) + self.worker.get_transfer_results() + self.transport.make_prepped_xfer.reset_mock() + + self.worker.start_load_kv( + self._request("old", blocks=(), awaiting_kvs=awaiting_kvs, req_id="empty") + ) + result = self.worker.get_transfer_results() + expected = {"empty"} if awaiting_kvs else set() + assert result.finished_recving == result.failed_recving == expected + self.transport.make_prepped_xfer.assert_not_called() + self.transport.send_notif.assert_not_called() + self.transport.remove_remote_agent.assert_not_called() + + def test_eviction_clears_pending_recovery(self): + self.worker._failed_remote_engines.add("old") + self.worker._invalid_remote_engines.add("old") + self.worker._engine_ttl = 1 + self.worker._engine_last_active["old"] = time.perf_counter() - 10 + self.worker._evict_stale_engines() + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_called_once_with("old") + assert not self.worker._failed_remote_engines + assert not self.worker._invalid_remote_engines + def test_transfer_topology_unregister(): """TransferTopology.unregister_remote_engine removes the engine.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 2a1ca5a87717..f96595572162 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -60,6 +60,8 @@ def region_pull_worker(): worker._transfer_layer_group_ids = () worker._bidirectional_kv_xfer_enabled = False worker._recving_transfers = {} + worker._failed_remote_engines = set() + worker._invalid_remote_engines = set() worker.use_mla, worker._has_mamba = True, False worker.dcp_size = 1 worker.dcp_rank = 0 @@ -690,6 +692,8 @@ def test_read_blocks_for_req_expands_remote_ids( worker = object.__new__(NixlConnectorWorker) worker._physical_blocks_per_logical_kv_block = local_physical_per_logical + worker._invalid_remote_engines = set() + worker._remote_agents = {"remote-engine": {}} worker._engine_last_active = {} worker._recving_transfers = {} worker._bidirectional_kv_xfer_enabled = False diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 47779c3611a5..a36c7e5b2a71 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -379,6 +379,8 @@ def fresh(cls) -> _StubWriterWorker: w.pp_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} + w._failed_remote_engines = set() + w._invalid_remote_engines = set() w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index dade5056eeb0..fa29d41a825e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -101,6 +101,11 @@ _SHARED_REGION_GROUP_ID = -1 +# Re-handshake delay bounds for a pull peer whose local metadata keeps getting +# invalidated, eg due to a persistent RDMA fault. +_INVALID_PEER_MIN_BACKOFF_S = 1.0 +_INVALID_PEER_MAX_BACKOFF_S = 60.0 + def _region_sort_key(layer_name: str) -> tuple[tuple[int, int | str], ...]: """Sort transfer regions in model-layer order, then by cache name.""" @@ -685,6 +690,12 @@ def __init__( self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict( dict ) + # Failed reads need a local metadata probe; invalid peers wait for drain. + self._failed_remote_engines: set[EngineId] = set() + self._invalid_remote_engines: set[EngineId] = set() + # Invalid engine -> (backoff, earliest release time). Kept across release + # so that repeat invalidations back off. + self._invalid_engine_backoff: dict[EngineId, tuple[float, float]] = {} # Map of engine_id -> clock offset. self._engine_clock_offset: dict[EngineId, float] = {} @@ -3029,7 +3040,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults: # Handle timeout to avoid stranding blocks on remote. self._reap_expired_send_leases(done_sending) - self._cleanup_replaced_remote_engines() + self._cleanup_remote_engines() return KVConnectorTransferResults( finished_sending=done_sending, finished_recving=done_recving, @@ -3195,6 +3206,12 @@ def _handle_failed_transfer( """ if record_failed_transfer: self.xfer_stats.record_failed_transfer() + if ( + self._TRANSFER_MODE == "pull" + and (meta := self._recving_metadata.get(req_id)) is not None + and meta.remote is not None + ): + self._failed_remote_engines.add(meta.remote.engine_id) if failed_req_ids is not None: failed_req_ids.add(req_id) return handle is None or self._try_release_xfer_handle(req_id, handle) @@ -3642,15 +3659,38 @@ def _track_remote_engine_replacement( engine_id, ) - def _cleanup_replaced_remote_engines(self) -> None: - """Release replaced pull peers once requests and handshakes have drained.""" + def _cleanup_remote_engines(self) -> None: + """Release replaced or locally invalid pull peers after their reads drain.""" + # TODO: Also handle push mode, which handshakes in both directions. if self._TRANSFER_MODE != "pull": return + # A background handshake may be loading native metadata with the GIL + # released. Defer both probing and cleanup until it has finished. with self._handshake_lock: if self._handshake_futures: return replaced = self._remote_agents.keys() - self._engine_by_address.values() - if not replaced: + now = time.perf_counter() + for engine_id in self._failed_remote_engines - self._invalid_remote_engines: + agents = self._remote_agents.get(engine_id, {}).values() + try: + # No descriptors: query local existence, not peer health. + if not all(map(self.nixl_wrapper.check_remote_metadata, agents)): + self._mark_remote_engine_invalid(engine_id, now) + except Exception: + logger.warning( + "Could not check local NIXL metadata for engine %s.", + engine_id, + exc_info=True, + ) + self._failed_remote_engines.clear() + # Invalid peers keep rejecting new reads until their backoff has elapsed. + invalid = { + engine_id + for engine_id in self._invalid_remote_engines + if self._invalid_engine_backoff.get(engine_id, (0.0, 0.0))[1] <= now + } + if not replaced and not invalid: return # _recving_metadata is only accessed from this (main) thread. busy = { @@ -3658,9 +3698,42 @@ def _cleanup_replaced_remote_engines(self) -> None: for meta in self._recving_metadata.values() if meta.remote is not None } - for engine_id in replaced - busy: + # Replaced peers can still serve queued requests. Invalid peers reject + # new reads, but must retain descriptors until outstanding handles end. + invalid -= self._engines_with_inflight_transfers() + for engine_id in (replaced - busy) | invalid: self._cleanup_remote_engine(engine_id, log_eviction=False) - logger.info("Released NIXL state for replaced remote engine %s.", engine_id) + logger.info( + "Released NIXL state for replaced or locally invalid remote engine %s.", + engine_id, + ) + + def _mark_remote_engine_invalid(self, engine_id: EngineId, now: float) -> None: + """Schedule release of an engine whose local NIXL metadata is gone. + + The first invalidation is released as soon as reads drain. Repeats + before the peer has stayed healthy for the max backoff double the + delay, during which reads fail fast instead of re-handshaking. + """ + # Forget engines that have stayed healthy long enough. + for eid, (_, release_at) in list(self._invalid_engine_backoff.items()): + if now - release_at > _INVALID_PEER_MAX_BACKOFF_S: + del self._invalid_engine_backoff[eid] + prev = self._invalid_engine_backoff.get(engine_id) + backoff = 0.0 + if prev is not None: + backoff = min( + max(2 * prev[0], _INVALID_PEER_MIN_BACKOFF_S), + _INVALID_PEER_MAX_BACKOFF_S, + ) + self._invalid_engine_backoff[engine_id] = (backoff, now + backoff) + self._invalid_remote_engines.add(engine_id) + logger.warning( + "Local NIXL metadata for remote engine %s was invalidated; " + "re-handshaking after reads drain%s.", + engine_id, + f" and a {backoff:.0f}s backoff" if backoff else "", + ) def _cleanup_remote_engine( self, engine_id: EngineId, *, log_eviction: bool = True @@ -3668,8 +3741,8 @@ def _cleanup_remote_engine( """Remove all state for a single remote engine. Releases NIXL resources (dlist handles, remote agents) and clears - all per-engine data structures. Used by TTL eviction, replaced peer - cleanup and shutdown. + all per-engine data structures. Used by TTL eviction, replaced/invalid + peer cleanup and shutdown. """ assert engine_id in self._remote_agents @@ -3685,6 +3758,9 @@ def _cleanup_remote_engine( for agent_name in agents.values(): self.nixl_wrapper.remove_remote_agent(agent_name) + self._failed_remote_engines.discard(engine_id) + self._invalid_remote_engines.discard(engine_id) + self.kv_caches_base_addr.pop(engine_id, None) self.dst_num_blocks.pop(engine_id, None) self.dst_region_num_blocks.pop(engine_id, None) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 4a18632eb4fe..5bcfdb9560a0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -72,26 +72,11 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): # Aborted cleanup requests have already been removed from the scheduler. if meta.awaiting_kvs or any(meta.local_block_ids): self._recving_metadata[req_id] = meta - if remote_engine_id not in self._remote_agents: - # Initiate handshake with remote engine to exchange metadata. - with self._handshake_lock: - if remote_engine_id not in self._remote_agents: - self._background_nixl_handshake(req_id, remote_engine_id, meta) - continue - - # Handshake already completed, start async read xfer. self._read_blocks_for_req(req_id, meta) # Start transfers for requests whose handshakes have now finished. while not self._ready_requests.empty(): - req_id, meta = self._ready_requests.get_nowait() - assert meta.remote is not None - if meta.remote.engine_id not in self._remote_agents: - # The engine was released after its handshake completed, so - # handshake again. This fails if the engine is gone. - self._background_nixl_handshake(req_id, meta.remote.engine_id, meta) - continue - self._read_blocks_for_req(req_id, meta) + self._read_blocks_for_req(*self._ready_requests.get_nowait()) if self.pcp_rank > 0 and not self.pcp_dcp_sharded: # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank @@ -151,6 +136,15 @@ def _is_turn2_read_expired(self, meta: ReqMeta) -> bool: def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): assert meta.remote is not None and self.transfer_topo is not None engine_id = meta.remote.engine_id + if engine_id in self._invalid_remote_engines: + self._handle_failed_transfer( + req_id, None, self._recv_failures, record_failed_transfer=False + ) + return + if engine_id not in self._remote_agents: + # Also re-handshakes queued requests that outlived invalid-peer cleanup. + self._background_nixl_handshake(req_id, engine_id, meta) + return # Update last activity from this remote. Mind that cleanup is done on main # thread (this one), so we don't race on this structure. self._engine_last_active[engine_id] = time.perf_counter() From b0b23a48d466712d18d390e73b62b169842f65e2 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:56:44 +0400 Subject: [PATCH 038/952] [Bugfix] Fix Qwen4Exp FP8 PLE pinned lookup kernel failing to compile below SM89 (#60142) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Thien Tran --- .../qwen4_exp/common/ngram_embedding.py | 27 +++++++++++-------- vllm/triton_utils/importing.py | 2 ++ 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 24eac4f8768d..6dad6b7b5afe 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -373,29 +373,33 @@ def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: @triton.jit def _lookup_ple_embedding_from_pinned_kernel( - weight_ptr, + weight_ptr: tl.pointer_type(tl.uint8), # type: ignore[valid-type] ids_ptr, - output_ptr, - embedding_dim, + output_ptr: tl.pointer_type(tl.uint8), # type: ignore[valid-type] + row_bytes, tp_vocab_start, tp_vocab_end, BLOCK_D: tl.constexpr, ): - """Look up TP-owned PLE rows through a CUDA view of pinned host memory.""" - row_id = tl.program_id(0) + """Copy TP-owned PLE rows as raw bytes from a CUDA view of pinned host memory. + + Byte pointers keep the storage dtype out of the kernel signature, so + dtypes Triton cannot lower on the GPU (FP8 E4M3FN before SM89) still work. + """ + row_id = tl.program_id(0).to(tl.int64) global_idx = tl.load(ids_ptr + row_id) in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end) local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0) offsets = tl.arange(0, BLOCK_D) - store_mask = offsets < embedding_dim + store_mask = offsets < row_bytes load_mask = store_mask & in_range values = tl.load( - weight_ptr + local_idx * embedding_dim + offsets, + weight_ptr + local_idx * row_bytes + offsets, mask=load_mask, - other=0.0, + other=0, ) tl.store( - output_ptr + row_id * embedding_dim + offsets, + output_ptr + row_id * row_bytes + offsets, values, mask=store_mask, ) @@ -433,7 +437,8 @@ def __init__( data_parallel_rank=data_parallel_rank, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) - self._block_d = triton.next_power_of_2(self.embedding_dim) + self._row_bytes = self.embedding_dim * self.weight.element_size() + self._block_d = triton.next_power_of_2(self._row_bytes) self._prefetch_stream: torch.cuda.Stream | None = None self._prefetch_buffer: torch.Tensor | None = None self._prefetch_alloc_lock = threading.Lock() @@ -485,7 +490,7 @@ def _lookup( self._uva_weight, flat_ids, output, - self.embedding_dim, + self._row_bytes, self.shard_indices.org_vocab_start_index, self.shard_indices.org_vocab_end_index, BLOCK_D=self._block_d, diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index e312564c0bd4..12f19cc52b8f 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -128,6 +128,8 @@ def __init__(self): self.dtype = None self.int64 = None self.int32 = None + self.uint8 = None + self.pointer_type = lambda element_ty: None self.tensor = None self.exp = None self.exp2 = None From 2b32d9ac6a910faf786a2904b87d5abed044c228 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Mon, 5 Oct 2026 19:51:03 -0700 Subject: [PATCH 039/952] [Bugfix][Core] Drop stale block hashes when a session truncates tokens (#59103) Signed-off-by: David Cheung Co-authored-by: Brian Woods Co-authored-by: Change72 Co-authored-by: Claude Opus 5 (1M context) --- .../test_scheduler_streaming.py | 49 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 1 + 2 files changed, 50 insertions(+) diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index bfef3dd33c8d..1f97d8a1c0a2 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -7,6 +7,7 @@ import numpy as np import torch +import vllm.v1.core.kv_cache_utils as kv_cache_utils from vllm.config import DeviceConfig, VllmConfig from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -14,6 +15,12 @@ PlaceholderRange, ) from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + hash_block_tokens, + init_none_hash, +) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.engine import FinishReason from vllm.v1.kv_cache_interface import ( @@ -365,6 +372,48 @@ def test_update_request_as_session_with_output_tokens(self): num_new_tokens = session.num_tokens - session.num_computed_tokens assert num_new_tokens == 2 + def test_update_request_as_session_drops_stale_block_hashes(self): + """A discarded sampled token must not stay fingerprinted in + block_hashes (#49377).""" + init_none_hash(sha256) + scheduler = create_scheduler() + hash_block_size = scheduler.hash_block_size + + # hash_block_size - 1 prompt tokens, so the discarded sampled token + # completes hash block 0. + prompt = list(range(1, hash_block_size)) + discarded_token = 999 + session = Request( + request_id="session", + prompt_token_ids=list(prompt), + sampling_params=SamplingParams(max_tokens=16), + pooling_params=None, + block_hasher=get_request_block_hasher(hash_block_size, sha256), + resumable=True, + ) + session.append_output_token_ids(discarded_token) + assert len(session.block_hashes) == 1 + stale_hash = session.block_hashes[0] + session.num_computed_tokens = len(prompt) + + next_token = 42 + update = StreamingUpdate.from_request( + DummyRequest(request_id="session", prompt_token_ids=[next_token]) + ) + scheduler._update_request_as_session(session, update) + + assert list(session.all_token_ids) == prompt + [next_token] + assert len(session.block_hashes) == 1 + expected = hash_block_tokens( + sha256, kv_cache_utils.NONE_HASH, tuple(prompt + [next_token]) + ) + assert session.block_hashes[0] == expected, ( + "block_hashes[0] does not identify the current token sequence" + ) + assert session.block_hashes[0] != stale_hash, ( + "block_hashes[0] still fingerprints the discarded sampled token" + ) + def test_streaming_e2e_lifecycle(self): """Comprehensive integration test covering complete streaming request lifecycle including scheduler state management and aliasing bug prevention. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index e25b5adda412..0f2425312c60 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1628,6 +1628,7 @@ def _update_request_as_session( session.num_prompt_tokens : num_computed_tokens ] del session._all_token_ids[num_computed_tokens:] + del session.block_hashes[num_computed_tokens // self.hash_block_size :] session._output_token_ids.clear() assert session.prompt_token_ids is not None # Extend prompt with kept output tokens. From 6fa67afed999d7cbab6a9a7492b2538dd8c2aff0 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Tue, 6 Oct 2026 05:01:19 +0200 Subject: [PATCH 040/952] [Bugfix][Core] Retire stale Mamba spec blocks in prefill checkpoint steps (#59759) Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../core/test_single_type_kv_cache_manager.py | 70 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 12 ++-- 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index dcb7ba923e51..415609aefb8d 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -6,10 +6,14 @@ import pytest import torch +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + get_request_block_hasher, + init_none_hash, make_block_hash_with_group_id, ) from vllm.v1.core.single_type_kv_cache_manager import ( @@ -30,6 +34,7 @@ RSWASpec, SlidingWindowSpec, ) +from vllm.v1.request import Request pytestmark = pytest.mark.cpu_test @@ -238,6 +243,71 @@ def estimate(num_tokens, total_computed_tokens, apply_admission_cap): assert admission_estimate == allocation_estimate == allocated +@pytest.mark.parametrize("num_speculative_blocks", [1, 2, 3]) +@pytest.mark.parametrize("first_chunk_blocks", [2, 4]) +@pytest.mark.parametrize("retention_interval", [0, None]) +def test_mamba_checkpoint_hash_maps_only_to_checkpoint_block( + num_speculative_blocks, first_chunk_blocks, retention_interval +): + """A checkpoint step must not leave a never-written speculative block + cached.""" + block_size = 128 + spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + num_prefill_checkpoint_blocks=1, + prefill_checkpoint_alignment=16, + ) + pool = BlockPool(num_gpu_blocks=64, enable_caching=True, hash_block_size=128) + manager = MambaManager( + spec, + block_pool=pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=block_size, + ) + # MTP/EAGLE moves the checkpoint one block below the last full block. + manager.drop_eagle_checkpoint_block = True + init_none_hash(sha256) + prompt_tokens = 6 * block_size + 50 + request = Request( + request_id="r", + prompt_token_ids=list(range(prompt_tokens)), + sampling_params=SamplingParams(max_tokens=1), + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, sha256), + ) + + def step(start, end): + estimate = manager.get_num_blocks_to_allocate("r", end, [], start, start, end) + free_before = pool.get_num_free_blocks() + manager.allocate_new_blocks("r", end, end) + assert free_before - pool.get_num_free_blocks() == estimate + manager.cache_blocks( + request, + end, + retention_interval=retention_interval, + replay_boundaries=[5 * block_size], + ) + request.num_computed_tokens = end + + step(0, first_chunk_blocks * block_size) + speculative_blocks = manager.req_to_blocks["r"][-num_speculative_blocks:] + manager.new_step_starts() + step(first_chunk_blocks * block_size, prompt_tokens) + + checkpoint_position, checkpoint_idx = manager._checkpoints["r"] + checkpoint_block = manager.req_to_blocks["r"][checkpoint_idx] + assert all( + b.block_hash is None or b is checkpoint_block for b in speculative_blocks + ) + block_hash = request.block_hashes[checkpoint_position // block_size - 1] + assert pool.get_cached_block(block_hash, [0]) == [checkpoint_block] + + def get_sliding_window_manager( sliding_window_spec, block_pool, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 0ff805956223..5940a1cd0a3f 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1812,7 +1812,7 @@ def get_num_blocks_to_allocate( if num_new_blocks > 0: blocks_allocated = request_id in self._allocated_block_reqs physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block - if not blocks_allocated or checkpoint_block: + if not blocks_allocated: physical_block_cap += self.num_speculative_blocks num_new_blocks = min(num_new_blocks, physical_block_cap) @@ -1885,12 +1885,14 @@ def allocate_new_blocks( [self._null_block for _ in range(prev_block_len, null_end)] ) - if blocks_allocated and not checkpoint_block: - # Relocate exclusively owned speculative scratch blocks. + if blocks_allocated: + # Relocate exclusively owned scratch blocks this step leaves + # behind. In a checkpoint step, one at the checkpoint index + # stays: the worker exports the checkpoint into that column. for block_idx in range( prev_block_len - self.num_speculative_blocks, prev_block_len ): - if block_idx < num_skipped_blocks: + if block_idx < num_skipped_blocks - checkpoint_block: self._relocate_speculative_block(req_blocks, block_idx) else: break @@ -1898,7 +1900,7 @@ def allocate_new_blocks( if has_partial_hit: num_new_blocks = max(num_new_blocks, 0) + 1 max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block - if not blocks_allocated or checkpoint_block: + if not blocks_allocated: max_new_blocks += self.num_speculative_blocks assert num_new_blocks <= max_new_blocks new_blocks = self.block_pool.get_new_blocks(num_new_blocks) From 48eb44b71972df0f1da88c0ab4672ac952db13f0 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 6 Oct 2026 11:19:41 +0800 Subject: [PATCH 041/952] [Bugfix][KVConnector][NIXL] Fix TP handshake test after per-region FA desc flags (#60157) Signed-off-by: Tianyu Guo --- tests/v1/kv_connector/unit/test_nixl_connector.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 35ac7d887ea4..18097d220e79 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -901,11 +901,16 @@ def test_prefill_tp_size_greater_than_decode_tp_size( vllm_config = create_vllm_config() - connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) - ) + # num_blocks=1 keeps the fake handshake's region_num_blocks consistent + # with the minimal single-block registration below (num_descs == 1); + # _fa_desc_replicated cross-checks the two. + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=1) + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 + vllm_config, + connector.engine_id, + hand_shake_latency=0, + kv_cache_config=kv_cache_config, ) worker = connector.connector_worker From 8c417a2dae6476c84dcf6503fca52d8f86ff711c Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 6 Oct 2026 11:35:53 +0800 Subject: [PATCH 042/952] [Bugfix][Model] Fix Qwen3-Omni interleaved M-RoPE boundary positions (#59842) Signed-off-by: Tianyu Guo Co-authored-by: OpenAI Codex --- tests/model_executor/test_qwen3_omni_mrope.py | 48 +++++++++++++------ .../models/qwen3_omni_moe_thinker.py | 5 +- 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/tests/model_executor/test_qwen3_omni_mrope.py b/tests/model_executor/test_qwen3_omni_mrope.py index bb8de63dc5f5..2c96c328df66 100644 --- a/tests/model_executor/test_qwen3_omni_mrope.py +++ b/tests/model_executor/test_qwen3_omni_mrope.py @@ -145,27 +145,45 @@ def test_mrope_positions_video_at_prompt_end(): assert positions.shape == (3, len(input_tokens)) -def test_mrope_positions_audio_in_video_at_prompt_end(): +@pytest.mark.parametrize("num_videos", [1, 2]) +@pytest.mark.parametrize("trailing_tokens", [0, 3]) +def test_mrope_positions_audio_in_video_boundaries(num_videos, trailing_tokens): + """Interleaved BOS/EOS positions and decode delta match Qwen3-Omni HF.""" model = make_model() - # video grid (1, 2, 2) -> 4 video tokens; audio_feature_length=16 -> 2 - # audio tokens, interleaved after the video tokens. - # [T, T, <|vision_start|>, <|audio_start|>, V x 4, A x 2, - # <|audio_end|>, <|vision_end|>] - input_tokens = ( - [TEXT_TOKEN_ID] * 2 - + [VISION_START_TOKEN_ID, AUDIO_START_TOKEN_ID] + video_tokens = ( + [VISION_START_TOKEN_ID, AUDIO_START_TOKEN_ID] + [VIDEO_TOKEN_ID] * 4 + [AUDIO_TOKEN_ID] * 2 + [AUDIO_END_TOKEN_ID, VISION_END_TOKEN_ID] ) - mm_features = [ - make_video_feature(offset=3, grid_thw=(1, 2, 2), use_audio_in_video=True), - make_audio_feature(offset=3, audio_feature_length=16), - ] - - positions, _ = model.get_mrope_input_positions(input_tokens, mm_features) + input_tokens = [TEXT_TOKEN_ID] * 2 + mm_features = [] + expected = [torch.tensor([[0, 1]]).expand(3, -1)] + video_positions = torch.tensor( + [ + [2, 3, 4, 4, 4, 4, 4, 5, 6, 7], + [2, 3, 4, 4, 5, 5, 4, 5, 6, 7], + [2, 3, 4, 5, 4, 5, 4, 5, 6, 7], + ] + ) + for index in range(num_videos): + offset = len(input_tokens) + 1 + mm_features.extend( + [ + make_video_feature(offset, (1, 2, 2), use_audio_in_video=True), + make_audio_feature(offset, audio_feature_length=16), + ] + ) + input_tokens.extend(video_tokens) + expected.append(video_positions + index * 6) + + input_tokens.extend([TEXT_TOKEN_ID] * trailing_tokens) + next_position = 2 + num_videos * 6 + expected.append((torch.arange(trailing_tokens) + next_position).expand(3, -1)) + positions, delta = model.get_mrope_input_positions(input_tokens, mm_features) - assert positions.shape == (3, len(input_tokens)) + torch.testing.assert_close(positions, torch.cat(expected, dim=1)) + assert delta == -4 * num_videos def test_mrope_positions_image_then_text(): diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index d3510f3e647b..4ae80a4a36e8 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -2311,8 +2311,9 @@ def get_mrope_input_positions( llm_pos_ids_list.append(eos_pos) st = offset + video_len + 1 else: - audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1)) + audio_bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(audio_bos_pos) + st_idx += 1 pos_ids, _ = self._compute_interleaved_positions(st_idx, data) llm_pos_ids_list.append(pos_ids) @@ -2320,7 +2321,7 @@ def get_mrope_input_positions( eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - llm_pos_ids_list.append(eos_pos) + llm_pos_ids_list.append(eos_pos + 1) video_len = grid_t * grid_h * grid_w audio_len = self._compute_audio_token_count( From a3e0243b1c3ea313b902a43b0f6559fb4148ff81 Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:37:30 -0500 Subject: [PATCH 043/952] [Bugfix][Parser] Mistral pre-v11: don't fail requests on unexpected tool call JSON (#54844) Signed-off-by: lzhan011 <35493221+lzhan011@users.noreply.github.com> Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/mistral/test_tool_calls.py | 139 ++++++++++++++++++++++-- vllm/parser/mistral.py | 125 +++++++++++++-------- 2 files changed, 208 insertions(+), 56 deletions(-) diff --git a/tests/parser/mistral/test_tool_calls.py b/tests/parser/mistral/test_tool_calls.py index a3359dccc68e..c17f618ba8dc 100644 --- a/tests/parser/mistral/test_tool_calls.py +++ b/tests/parser/mistral/test_tool_calls.py @@ -475,17 +475,22 @@ def test_extract_tool_calls_pre_v11_tokenizer( assert extracted_tool_calls.content == expected_content -def test_extract_tool_calls_pre_v11_multiple_bot_tokens_raises( +def test_extract_tool_calls_pre_v11_multiple_bot_tokens( mistral_pre_v11_tool_parser, ): + """A second BOT token must not fail the request; like streaming, only the + calls after the first BOT token are extracted.""" model_output = ( '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1}}]' '[TOOL_CALLS] [{"name": "sub", "arguments":{"b": 2}}]' ) - with pytest.raises(ValueError, match="Only one BOT token"): - mistral_pre_v11_tool_parser.extract_tool_calls( - model_output, request=_DUMMY_REQUEST - ) + result = mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert [(tc.function.name, tc.function.arguments) for tc in result.tool_calls] == [ + ("add", '{"a": 1}') + ] def test_extract_tool_calls_pre_v11_regex_fallback( @@ -506,15 +511,126 @@ def test_extract_tool_calls_pre_v11_regex_fallback( assert result.tool_calls[0].function.arguments == json.dumps({"a": 1, "b": 2}) -def test_extract_tool_calls_pre_v11_regex_fallback_fails( - mistral_pre_v11_tool_parser, +@pytest.mark.parametrize( + "stringified_tool_calls", + ["not json at all", "42", '["add"]', '{"foo": 1}'], + ids=["invalid_json", "scalar", "list_of_str", "not_a_call"], +) +def test_extract_tool_calls_pre_v11_no_tool_call_returns_content( + mistral_pre_v11_tool_parser, stringified_tool_calls ): - model_output = "[TOOL_CALLS] not json at all" + """Output that holds no tool call is returned as content instead of failing + the request.""" result = mistral_pre_v11_tool_parser.extract_tool_calls( - model_output, request=_DUMMY_REQUEST + f"[TOOL_CALLS] {stringified_tool_calls}", request=_DUMMY_REQUEST ) assert result == ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content="not json at all" + tools_called=False, tool_calls=[], content=stringified_tool_calls + ) + + +# (model_output, expected (name, arguments), streamed content). Non-streaming +# drops the text after the calls; streaming returns it as content. +_PRE_V11_UNUSUAL_CALLS = [ + pytest.param( + '[TOOL_CALLS] {"name": "add", "arguments": {"a": 1}}', + [("add", '{"a": 1}')], + "", + id="bare_object", + ), + # The next three end the bare object and start the trailing text in the + # same token, in each state the object can close from. + pytest.param( + '[TOOL_CALLS] {"name": "add", "arguments": {"a": 1}}, ' + '{"name": "sub", "arguments": {"b": 2}}', + [("add", '{"a": 1}')], + ', {"name": "sub", "arguments": {"b": 2}}', + id="bare_object_then_text", + ), + pytest.param( + '[TOOL_CALLS] {"arguments": {"a": 1}, "name": "add"}, ' + '{"name": "sub", "arguments": {"b": 2}}', + [("add", '{"a": 1}')], + ', {"name": "sub", "arguments": {"b": 2}}', + id="bare_object_name_last_then_text", + ), + pytest.param( + '[TOOL_CALLS] {"name": "add", "arguments": "{\\"a\\": 1}"}, ' + '{"name": "sub", "arguments": {"b": 2}}', + [("add", '{"a": 1}')], + ', {"name": "sub", "arguments": {"b": 2}}', + id="bare_object_string_arguments_then_text", + ), + pytest.param( + '[TOOL_CALLS] [{"arguments": {"a": 1}}]', + [("", '{"a": 1}')], + "", + id="missing_name", + ), + pytest.param( + '[TOOL_CALLS] [{"name": "add", "arguments": "{\\"a\\": 1}"}, ' + '{"name": "sub", "arguments": {"b": 2}}]', + [("add", '{"a": 1}'), ("sub", '{"b": 2}')], + "", + id="string_arguments", + ), + pytest.param( + '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 1}}, "junk"]', + [("add", '{"a": 1}')], + "", + id="entry_not_a_call", + ), + # Mistral-7B-Instruct-v0.3 emits repeated BOT tokens at temperature 0. + pytest.param( + '[TOOL_CALLS] [TOOL_CALLS] [{"name": "add", "arguments": {"a": 1}}]', + [("add", '{"a": 1}')], + "", + id="consecutive_bot_tokens", + ), + pytest.param( + '[TOOL_CALLS] [TOOL_CALLS] {"name": "add", "arguments": {"a": 1}}', + [("add", '{"a": 1}')], + "", + id="consecutive_bot_tokens_bare_object", + ), +] + + +@pytest.mark.parametrize("model_output,expected,_", _PRE_V11_UNUSUAL_CALLS) +def test_extract_tool_calls_pre_v11_unusual_calls( + mistral_pre_v11_tool_parser, model_output, expected, _ +): + result = mistral_pre_v11_tool_parser.extract_tool_calls( + model_output, request=_DUMMY_REQUEST + ) + assert result.tools_called + assert [ + (tc.function.name, tc.function.arguments) for tc in result.tool_calls + ] == expected + assert result.content is None + + +@pytest.mark.parametrize( + "model_output,expected,expected_content", _PRE_V11_UNUSUAL_CALLS +) +def test_extract_tool_calls_streaming_pre_v11_unusual_calls( + mistral_pre_v11_tool_parser, + mistral_pre_v11_tokenizer, + model_output, + expected, + expected_content, +): + """Streaming extracts the same calls as the non-streaming path.""" + _test_extract_tool_calls_streaming( + mistral_pre_v11_tool_parser, + mistral_pre_v11_tokenizer, + model_output, + None, + [ + ToolCall(function=FunctionCall(name=name, arguments=arguments)) + for name, arguments in expected + ], + expected_content, ) @@ -882,6 +998,7 @@ def _test_extract_tool_calls_streaming( # if a new tool is being called, set up empty arguments if tool_call.index != tool_call_idx: tool_call_idx = tool_call.index + function_names.append("") function_args_strs.append("") tool_call_ids.append(None) @@ -895,7 +1012,7 @@ def _test_extract_tool_calls_streaming( # IN ENTIRETY, exactly one time. if tool_call.function.name: assert isinstance(tool_call.function.name, str) - function_names.append(tool_call.function.name) + function_names[tool_call.index] = tool_call.function.name if tool_call.function.arguments: # make sure they're a string and then add them to the list diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index 61f133598c8c..f6d3435bd7c5 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -117,6 +117,37 @@ def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool: return _ARGS not in vocab +def _pre_v11_function_calls(parsed: Any) -> list[FunctionCall]: + """Decode the JSON after a pre-v11 ``[TOOL_CALLS]`` marker. + + Mirrors the streaming path: a bare object is a single call, entries that + are not call objects are skipped, a missing or non-string name becomes + ``""`` (as in the v11+ path) and string arguments are kept as-is. + """ + if isinstance(parsed, dict): + parsed = [parsed] + if not isinstance(parsed, list): + return [] + function_calls: list[FunctionCall] = [] + for tool_call in parsed: + if not isinstance(tool_call, dict) or not ( + "name" in tool_call or "arguments" in tool_call + ): + continue + name = tool_call.get("name") + arguments = tool_call.get("arguments") + if arguments is None: + arguments = "{}" + elif not isinstance(arguments, str): + arguments = json.dumps(arguments, ensure_ascii=False) + function_calls.append( + FunctionCall( + name=name if isinstance(name, str) else "", arguments=arguments + ) + ) + return function_calls + + @functools.cache def mistral_config( *, @@ -301,6 +332,7 @@ def __init__( self.starting_new_tool: bool = False self.streamed_args_for_tool: list[str] = [] self._is_pre_v11: bool = _is_pre_v11_tokeniser(tokenizer) + self.current_tool_string_arguments: str | None = None self.parse_coro = None if self._is_pre_v11: self.parse_coro = ijson.parse_coro( @@ -551,16 +583,13 @@ def _legacy_extract_tool_calls( content: str | None = None if self.bot_token in model_output: - content_and_raw_tool_calls = model_output.split(self.bot_token) - content = content_and_raw_tool_calls[0] - raw_tool_calls = content_and_raw_tool_calls[1:] # pre-v11: content[BOT] [{tool_call1},{tool_call2}] - if len(raw_tool_calls) != 1: - raise ValueError( - "Only one BOT token should have been outputted, " - f"but got {model_output}." - ) - stringified_tool_calls = raw_tool_calls[0].strip() + # Repeated BOT tokens with nothing between them count as one, and + # output after a later BOT token is trailing text, as in streaming. + content, *segments = model_output.split(self.bot_token) + stringified_tool_calls = next( + (segment.strip() for segment in segments if segment.strip()), "" + ) elif tool_choice == "required" or isinstance( tool_choice, ChatCompletionNamedToolChoiceParam ): @@ -575,49 +604,31 @@ def _legacy_extract_tool_calls( # Use raw_decode to parse the first valid JSON value, # ignoring trailing tokens the model may emit after # the tool call array. - tool_calls, _ = json.JSONDecoder().raw_decode(stringified_tool_calls) + parsed, _ = json.JSONDecoder().raw_decode(stringified_tool_calls) except json.JSONDecodeError: try: raw_tool_call = self.tool_call_regex.findall(stringified_tool_calls)[0] - tool_calls = json.loads(raw_tool_call) - tool_calls = [ - { - "name": tool_call["name"], - "arguments": json.dumps( - tool_call.get("arguments", {}), - ensure_ascii=False, - ), - } - for tool_call in tool_calls - ] + parsed = json.loads(raw_tool_call) except (IndexError, json.JSONDecodeError): logger.exception("Error in extracting tool call from response.") - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=stringified_tool_calls, + parsed = None + + function_calls = _pre_v11_function_calls(parsed) + if not function_calls: + if parsed is not None: + logger.warning( + "No tool call found in the JSON after %s; returning it as content.", + self.bot_token, ) - else: - tool_calls = [ - { - "name": tool_call["name"], - "arguments": json.dumps( - tool_call.get("arguments", {}), - ensure_ascii=False, - ), - } - for tool_call in tool_calls - ] + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=stringified_tool_calls, + ) mistral_tool_calls: list[MistralToolCall] = [ - MistralToolCall( - type="function", - function=FunctionCall( - name=tool_call["name"], - arguments=tool_call.get("arguments", "{}"), - ), - ) - for tool_call in tool_calls + MistralToolCall(type="function", function=function_call) + for function_call in function_calls ] return ExtractedToolCallInformation( @@ -669,9 +680,16 @@ def extract_tool_calls_streaming( @ijson.coroutine def update_stream_state_pre_v11_tokenizer(self): + bare_object = False while True: (prefix, event, value) = yield + # A bare top-level object is a single call: parse it as an item. + if prefix == "" and event == "start_map": + bare_object = True + if bare_object: + prefix = f"item.{prefix}" if prefix else "item" + if prefix == "item" and event == "start_map": self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY self.starting_new_tool = True @@ -686,8 +704,14 @@ def update_stream_state_pre_v11_tokenizer(self): self.streaming_state = StreamingState.PARSING_ARGUMENTS if prefix == "item.arguments" and event == "end_map": self.streaming_state = StreamingState.PARSING_ARGUMENTS_COMPLETED + if prefix == "item.arguments" and event == "string": + self.current_tool_string_arguments = value if prefix == "item" and event == "end_map": - self.streaming_state = StreamingState.TOOL_COMPLETE + self.streaming_state = ( + StreamingState.ALL_TOOLS_COMPLETE + if bare_object + else StreamingState.TOOL_COMPLETE + ) if prefix == "" and event == "end_array": self.streaming_state = StreamingState.ALL_TOOLS_COMPLETE @@ -723,22 +747,27 @@ def _extract_tool_calls_streaming_pre_v11_tokenizer( delta_text=delta_text, stop_after_opening_curly_braces=1, ) + # Every state in which the call object can close stops after "}", + # so text following a bare object is never sent to ijson. elif self.streaming_state == StreamingState.WAITING_FOR_TOOL_KEY: delta_to_be_parsed, delta_text = self._split_delta( delta_text=delta_text, stop_after_colon=1, stop_after_opening_curly_braces=1, + stop_after_closing_curly_braces=1, ) elif self.streaming_state == StreamingState.PARSING_NAME: delta_to_be_parsed, delta_text = self._split_delta( delta_text=delta_text, stop_after_comma=1, stop_after_closing_brackets=1, + stop_after_closing_curly_braces=1, ) elif self.streaming_state == StreamingState.WAITING_FOR_ARGUMENTS_START: delta_to_be_parsed, delta_text = self._split_delta( delta_text=delta_text, stop_after_opening_curly_braces=1, + stop_after_closing_curly_braces=1, ) elif self.streaming_state == StreamingState.PARSING_ARGUMENTS: delta_to_be_parsed, delta_text = self._split_delta( @@ -799,6 +828,12 @@ def _extract_tool_calls_streaming_pre_v11_tokenizer( self.current_tool_name ) self.current_tool_name = None + if self.current_tool_string_arguments is not None: + current_tool_call_modified = True + current_tool_call.function.arguments = ( + self.current_tool_string_arguments + ) + self.current_tool_string_arguments = None if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED: self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY if self.streaming_state in [ From 30ab428235e4a444dbb26f092d51c051480b8198 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 6 Oct 2026 12:13:23 +0800 Subject: [PATCH 044/952] [Bugfix][Attention] Restrict FlashInfer sparse MLA FULL graphs to decode (#59751) Signed-off-by: Tianyu Guo --- .../test_dspark_noncausal_sparse_mla.py | 30 +++++++++---- tests/v1/worker/test_attn_utils.py | 43 ++++++++++++++++++- .../backends/mla/flashinfer_mla_sparse.py | 14 +++++- 3 files changed, 76 insertions(+), 11 deletions(-) diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 3d4af1741362..7cd7863b0659 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -55,7 +55,6 @@ ) from vllm.utils.math_utils import cdiv -from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseTRTLLMBackend, ) @@ -122,6 +121,7 @@ def _run_sparse_backend_vs_sdpa( qk_nope_head_dim: int = 128, v_head_dim: int = 128, stale_cpu_query_lens: list[int] | None = None, + capture_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Run a sparse-MLA backend with the given per-token indices and compute a dense per-token SDPA reference over the SAME indices. @@ -394,9 +394,25 @@ def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa): metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device ) with torch.inference_mode(): + impl.record_logical_topk_ready() backend_output = mock_layer.forward_impl( query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, out_buffer ) + if capture_graph: + device_boundaries = common_attn_metadata.query_start_loc.clone() + common_attn_metadata.query_start_loc.copy_( + common_attn_metadata.query_start_loc_cpu + ) + metadata = builder.build(0, common_attn_metadata) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + impl.record_logical_topk_ready() + backend_output = mock_layer.forward_impl( + query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, out_buffer + ) + common_attn_metadata.query_start_loc.copy_(device_boundaries) + builder.build(0, common_attn_metadata) + graph.replay() return backend_output, sdpa_reference, causal_reference @@ -423,22 +439,19 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability") -def test_flashinfer_sparse_mla_adaptive_varlen_matches_sdpa( +@pytest.mark.parametrize("query_lens", [[1, 7, 3, 5], [1, 7, 8, 0]]) +def test_flashinfer_sparse_mla_adaptive_varlen_graph_matches_sdpa( default_vllm_config, dist_init, workspace_init, + query_lens, ): - """Adaptive request boundaries must drive SM100 sparse index conversion.""" + """Replay uses device boundaries, including empty padding, despite stale CPU.""" backend_cls = FlashInferMLASparseTRTLLMBackend _skip_if_backend_unavailable(backend_cls, "fp8", 64) - assert ( - backend_cls.get_builder_cls().get_cudagraph_support(None, None) - == AttentionCGSupport.ALWAYS - ) device = torch.device(DEVICE_TYPE) seq_lens = [257, 270, 265, 276] - query_lens = [1, 7, 3, 5] sparse_indices = _build_dspark_noncausal_indices( seq_lens, query_lens, @@ -459,6 +472,7 @@ def test_flashinfer_sparse_mla_adaptive_varlen_matches_sdpa( qk_nope_head_dim=192, v_head_dim=256, stale_cpu_query_lens=[4, 4, 4, 4], + capture_graph=True, ) torch.testing.assert_close( diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 0e6089d529a6..41148070528a 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -16,7 +16,7 @@ import vllm.v1.hisparse.binding as attn_utils_module from tests.v1.attention.utils import dense_kv_cache_views -from vllm.config.compilation import CUDAGraphMode +from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches @@ -228,6 +228,47 @@ def group( assert unsupported([[target, never]], config, 1) == ("_DraftBackend", None) +@pytest.mark.parametrize("num_speculative_tokens", [0, 7]) +def test_flashinfer_sparse_full_graphs_exclude_prefill_keep_varlen_decode( + num_speculative_tokens, +): + """Prefill must fall back without disabling adaptive verification graphs.""" + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseTRTLLMBackend, + ) + + backend = FlashInferMLASparseTRTLLMBackend + config: Any = SimpleNamespace( + speculative_config=SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, parallel_drafting=False + ), + use_v2_model_runner=True, + ) + spec = MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=torch.bfloat16 + ) + group = AttentionGroup(backend, ["target"], spec, 0) + group.metadata_builders = [object.__new__(backend.get_builder_cls())] + groups = [[group]] + support = get_attn_cg_support(groups, config) + compilation = CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL) + mode = compilation.resolve_cudagraph_mode_and_sizes( + support.min_cg_support, + support.min_cg_attn_backend, + uniform_decode_query_len=1 + num_speculative_tokens, + use_v2_model_runner=True, + ) + assert mode.mixed_mode() != CUDAGraphMode.FULL + assert mode.decode_mode() == CUDAGraphMode.FULL + assert get_query_lens_mismatch_unsupported_backend(groups) is None + unsupported = attn_utils.get_varlen_cudagraph_unsupported_backend + assert unsupported(groups, config, 1 + num_speculative_tokens) is None + assert unsupported(groups, config, 2 + num_speculative_tokens) == ( + backend.__name__, + 1 + num_speculative_tokens, + ) + + def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyPatch): """Fast prefill applies to the contiguous suffix of KV-sharing layers. diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 05af2de79710..4ebda50178cb 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -26,6 +26,7 @@ AttentionType, MLAAttentionImpl, MultipleOf, + max_decode_query_len, ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( @@ -34,7 +35,7 @@ triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec if TYPE_CHECKING: from vllm.model_executor.models.deepseek_v2 import Indexer @@ -272,7 +273,16 @@ def __init__( class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilder): """Metadata builder for the SM100 TRT-LLM sparse MLA kernel.""" - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + @classmethod + def get_varlen_cudagraph_max_query_len( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + # Decode uses device request boundaries; prefill metadata is not graph-safe. + return max_decode_query_len(vllm_config) def __init__( self, From 21d93d0d8c0e9627900020382bfce4730e61cab7 Mon Sep 17 00:00:00 2001 From: Luca Motz Date: Tue, 6 Oct 2026 06:18:29 +0200 Subject: [PATCH 045/952] [Attention][Model] FP8 KV cache for Triton DiffKV and MiMo-V2.6-Flash (#58128) Signed-off-by: Luca Motz Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 Co-authored-by: Lucas Wilkinson --- .../test_triton_unified_attention_diffkv.py | 67 +++++++++++++++++-- vllm/model_executor/models/mimo_v2.py | 6 ++ vllm/model_executor/models/mimo_v2_mtp.py | 7 +- .../attention/backends/triton_attn_diffkv.py | 24 ++++--- .../ops/triton_unified_attention_diffkv.py | 22 ++++++ 5 files changed, 111 insertions(+), 15 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 213d896c033c..3b6cf524e2fc 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -35,6 +35,8 @@ HEAD_SIZES = [(128, 128), (192, 128)] BLOCK_SIZES = [16] DTYPES = [torch.bfloat16] +# None keeps the KV cache in the query dtype. +KV_CACHE_DTYPES = [None, current_platform.fp8_dtype()] NUM_BLOCKS = 2048 @@ -80,6 +82,7 @@ def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v @pytest.mark.parametrize("sliding_window", [None, 128]) @pytest.mark.parametrize("soft_cap", [None, 50.0]) @pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES) @pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES) @torch.inference_mode() def test_triton_unified_attn_diffkv_vs_reference( @@ -89,13 +92,17 @@ def test_triton_unified_attn_diffkv_vs_reference( sliding_window: int | None, soft_cap: float | None, dtype: torch.dtype, + kv_cache_dtype: torch.dtype | None, block_size: int, seq_threshold_3D: int, + quant_query: bool = False, ) -> None: head_size_qk, head_size_v = head_sizes - # Keep the FA3/FA4 comparison on NVIDIA; ROCm uses the PyTorch oracle. - if not current_platform.is_rocm(): + # FA3/FA4 is the reference on NVIDIA; ROCm and the FP8 KV cache (FA would + # also quantize Q) use the PyTorch oracle. + use_torch_ref = current_platform.is_rocm() or kv_cache_dtype is not None + if not use_torch_ref: fa_version = get_flash_attn_version( head_size=head_size_qk, head_size_v=head_size_v ) @@ -126,6 +133,20 @@ def test_triton_unified_attn_diffkv_vs_reference( ) key_cache = kv_cache[..., :head_size_qk] value_cache = kv_cache[..., head_size_qk:] + q_descale = k_descale = v_descale = None + if kv_cache_dtype is not None: + # Non-power-of-two scales exercise descale rounding. + k_descale = torch.tensor(0.3, dtype=torch.float32) + v_descale = torch.tensor(0.7, dtype=torch.float32) + kv_cache = torch.cat( + [key_cache / k_descale, value_cache / v_descale], dim=-1 + ).to(kv_cache_dtype) + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk:] + + if quant_query: + q_descale = torch.tensor(0.2, dtype=torch.float32) + query = (query / q_descale).to(kv_cache_dtype) cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( dim=0, dtype=torch.int32 @@ -137,12 +158,20 @@ def test_triton_unified_attn_diffkv_vs_reference( 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 ) - if current_platform.is_rocm(): + if use_torch_ref: + ref_key_cache = key_cache.float() + ref_value_cache = value_cache.float() + if kv_cache_dtype is not None: + ref_key_cache *= k_descale + ref_value_cache *= v_descale # FP32 also keeps the helper's in-place scaling off the kernel input. + ref_query = query.float() + if quant_query: + ref_query *= q_descale ref_out = ref_paged_attn( - query.float(), - key_cache.float(), - value_cache.float(), + ref_query, + ref_key_cache, + ref_value_cache, query_lens, kv_lens, block_tables, @@ -200,6 +229,32 @@ def test_triton_unified_attn_diffkv_vs_reference( softmax_segm_output=segm_output, softmax_segm_max=segm_max, softmax_segm_expsum=segm_expsum, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, ) torch.testing.assert_close(triton_out, ref_out, atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(89), + reason="FP8 query requires CUDA SM89+", +) +@pytest.mark.parametrize( + "seq_lens,seq_threshold_3D", + [([(129, 463)], 0), ([(1, 2011)], 0), ([(1, 2011)], 8)], +) +def test_triton_unified_attn_diffkv_fp8_query(seq_lens, seq_threshold_3D): + test_triton_unified_attn_diffkv_vs_reference( + seq_lens=seq_lens, + num_heads=(8, 2), + head_sizes=(192, 128), + sliding_window=None, + soft_cap=None, + dtype=torch.bfloat16, + kv_cache_dtype=current_platform.fp8_dtype(), + block_size=16, + seq_threshold_3D=seq_threshold_3D, + quant_query=True, + ) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 6267a37dc040..a6464c7c84ee 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -372,6 +372,10 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: max_position_embeddings = getattr(config, "max_position_embeddings", 32768) v_scale = getattr(config, "attention_value_scale", None) + # cache_config.sliding_window holds MiMo's SWA window; global layers + # would otherwise fall back to it in Attention. + cache_config = vllm_config.cache_config + cache_config.sliding_window = None if self.is_compressed_softmax_layer(): self.self_attn = MiMoV2Attention( @@ -389,6 +393,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: layer_id=layer_id, rope_theta=getattr(config, "swa_rope_theta", rope_theta), max_position_embeddings=max_position_embeddings, + cache_config=cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", @@ -406,6 +411,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: layer_id=layer_id, rope_theta=rope_theta, max_position_embeddings=max_position_embeddings, + cache_config=cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index c413b65ba927..fe6d64552d9f 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -24,7 +24,7 @@ import torch.nn as nn from transformers import PreTrainedConfig -from vllm.config import VllmConfig +from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -64,6 +64,7 @@ def __init__( self, config: PreTrainedConfig, prefix: str, + cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() @@ -100,6 +101,7 @@ def __init__( layer_id=0, rope_theta=swa_rope_theta, max_position_embeddings=getattr(config, "max_position_embeddings", 32768), + cache_config=cache_config, quant_config=quant_config, partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), prefix=f"{prefix}.self_attn", @@ -147,6 +149,7 @@ def __init__( self, config: PreTrainedConfig, num_mtp_layers: int, + cache_config: CacheConfig | None, quant_config: QuantizationConfig | None, prefix: str, ) -> None: @@ -156,6 +159,7 @@ def __init__( str(i): MiMoV2MTPLayer( config=config, prefix=f"{prefix}.{i}", + cache_config=cache_config, quant_config=quant_config, ) for i in range(num_mtp_layers) @@ -182,6 +186,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.mtp = _MiMoV2MTPLayers( config=config, num_mtp_layers=num_mtp_layers, + cache_config=vllm_config.cache_config, quant_config=vllm_config.quant_config, prefix=maybe_prefix(prefix, "mtp.layers"), ) diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 4b42838a1503..932f483af37b 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -70,10 +70,11 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend): # V head dim — set per layer via ``set_head_size_v`` before instantiation. head_size_v: int = 128 - # No FP8 / int8 KV cache for the DiffKV path yet; require fp16/bf16/fp32. supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "bfloat16", + "fp8", + "fp8_e4m3", ] @classmethod @@ -110,16 +111,15 @@ class TritonAttentionDiffKVImpl(TritonAttentionImpl): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) - if is_quantized_kv_cache(self.kv_cache_dtype): + if ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype + not in TritonAttentionDiffKVBackend.supported_kv_cache_dtypes + ): raise NotImplementedError( - "TritonAttentionDiffKVBackend does not yet support quantized " + "TritonAttentionDiffKVBackend supports only fp8/fp8_e4m3 quantized " f"KV cache (got kv_cache_dtype={self.kv_cache_dtype!r})." ) - if self._is_per_token_head_quant: - raise NotImplementedError( - "TritonAttentionDiffKVBackend does not support per-token-head " - "quantization." - ) if self.chunk_lookback > -1: raise NotImplementedError( "TritonAttentionDiffKVBackend does not support chunked " @@ -191,6 +191,9 @@ def forward( # Triton DiffKV kernels consume (B, N, H, D) cache views. kv_cache = kv_cache.transpose(1, 2) + quantized = is_quantized_kv_cache(self.kv_cache_dtype) + if quantized: + kv_cache = kv_cache.view(self.fp8_dtype) key_cache = kv_cache[..., :head_size_qk] value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v] @@ -202,6 +205,11 @@ def forward( cu_seqlens_q=attn_metadata.query_start_loc, seqused_k=attn_metadata.seq_lens, softmax_scale=self.scale, + q_descale=( + layer._q_scale if quantized and query.dtype == self.fp8_dtype else None + ), + k_descale=layer._k_scale if quantized else None, + v_descale=layer._v_scale if quantized else None, causal=True, alibi_slopes=self.alibi_slopes, use_alibi_sqrt=self.use_alibi_sqrt, diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py index eaf62b6bce66..8480754256b2 100644 --- a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -60,6 +60,9 @@ def kernel_unified_attention_diffkv( query_ptr, key_cache_ptr, # view of packed cache: [..., :head_size_qk] value_cache_ptr, # view of packed cache: [..., head_size_qk:hqk+hv] + q_descale_ptr, + k_descale_ptr, + v_descale_ptr, sink_ptr, block_tables_ptr, seq_lens_ptr, @@ -83,6 +86,8 @@ def kernel_unified_attention_diffkv( USE_ALIBI_SQRT: tl.constexpr, USE_SOFTCAP: tl.constexpr, USE_SINKS: tl.constexpr, + USE_Q_SCALE: tl.constexpr, + USE_KV_SCALES: tl.constexpr, SLIDING_WINDOW: tl.constexpr, # Strides for both cache views (they share the same packed buffer, so # dims 0/1/2 strides match; only the per-head extent differs). @@ -154,6 +159,12 @@ def kernel_unified_attention_diffkv( other=0.0, ) + if USE_Q_SCALE: + scale *= tl.load(q_descale_ptr) + if USE_KV_SCALES: + scale *= tl.load(k_descale_ptr) + v_descale = tl.load(v_descale_ptr) + block_table_offset = seq_idx * block_table_stride M = init_softmax_M( @@ -262,6 +273,9 @@ def kernel_unified_attention_diffkv( acc += tl.dot(P.to(V.dtype), V) # ---- Epilogue -------------------------------------------------------- + if USE_KV_SCALES: + acc *= v_descale + if IS_3D: # Store per-segment partials; finalized by reduce_segments_diffkv. segm_output_offset = ( @@ -403,6 +417,9 @@ def unified_attention_diffkv( softmax_segm_output: torch.Tensor | None = None, softmax_segm_max: torch.Tensor | None = None, softmax_segm_expsum: torch.Tensor | None = None, + k_descale: torch.Tensor | None = None, + v_descale: torch.Tensor | None = None, + q_descale: torch.Tensor | None = None, ): assert causal, "Only causal attention is supported" @@ -470,6 +487,9 @@ def unified_attention_diffkv( query_ptr=q, key_cache_ptr=k, value_cache_ptr=v, + q_descale_ptr=q_descale, + k_descale_ptr=k_descale, + v_descale_ptr=v_descale, sink_ptr=sinks, block_tables_ptr=block_table, seq_lens_ptr=seqused_k, @@ -493,6 +513,8 @@ def unified_attention_diffkv( USE_ALIBI_SQRT=use_alibi_sqrt, USE_SOFTCAP=(softcap > 0), USE_SINKS=(sinks is not None), + USE_Q_SCALE=q_descale is not None, + USE_KV_SCALES=k_descale is not None, SLIDING_WINDOW=sliding_window_val, stride_k_cache_0=k.stride(0), stride_k_cache_1=k.stride(1), From eb2901b7356dd6782b7418ac071d3cf56ce018f7 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Mon, 5 Oct 2026 22:11:45 -0700 Subject: [PATCH 046/952] [Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs (#59532) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- ...test_deepseek_v41_decoder_replay_layers.py | 68 ++++++- .../models/test_deepseek_v41_replay_batch.py | 45 +++-- tests/v1/worker/test_gpu_model_runner_v2.py | 4 +- vllm/config/compilation.py | 23 +++ .../deepseek_v41/decoder_replay_layers.py | 77 ++++++-- .../nvidia/decoder_replay_cudagraph.py | 102 +++++++++++ vllm/models/deepseek_v41/nvidia/model.py | 10 ++ .../models/deepseek_v41/nvidia/model_state.py | 168 ++++++++++++++---- vllm/v1/worker/gpu/model_runner.py | 6 + vllm/v1/worker/gpu/model_states/interface.py | 13 +- 10 files changed, 445 insertions(+), 71 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 4af797a1e6c8..9ba2df5c8f4c 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -2,15 +2,26 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DecoderReplayLayers runs the replay layers on the forward's replay batch.""" +from unittest.mock import MagicMock + import pytest import torch +from vllm.config import CompilationConfig, CUDAGraphMode from vllm.forward_context import ( + BatchDescriptor, ForwardContext, get_forward_context, override_forward_context, ) -from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers +from vllm.models.deepseek_v41.decoder_replay_layers import ( + DecoderReplayLayers, + ReplayBatch, +) +from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( + DecoderReplayCudaGraphManager, +) +from vllm.v1.worker.gpu import cudagraph_utils pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -32,9 +43,9 @@ def _context(metadata): def _set_replay_batch(layers, rows, metadata): - layers.rows = torch.tensor(rows, device=DEVICE) - layers.forward_context = ForwardContext( - no_compile_layers={}, attn_metadata=metadata, slot_mapping={} + layers.replay_batch = ReplayBatch( + torch.tensor(rows, device=DEVICE), + ForwardContext(no_compile_layers={}, attn_metadata=metadata, slot_mapping={}), ) @@ -69,7 +80,7 @@ def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual seen["is_padding"] = replay_context.is_padding return residual, pre_mix - layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates]) + layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates], set()) states = _states(NUM_TOKENS) hidden, residual = states[0], states[-1] full, sub = object(), object() @@ -96,9 +107,54 @@ def run_layers(hidden_states, *rest): seen["attn_metadata"] = get_forward_context().attn_metadata return (hidden_states,) - layers = DecoderReplayLayers(WINDOW, run_layers, []) + layers = DecoderReplayLayers(WINDOW, run_layers, [], set()) states = _states(NUM_TOKENS) full = object() with override_forward_context(_context(full)): (output,) = layers(*states) assert output is states[0] and seen["attn_metadata"] is full + + +def _graph_context(num_tokens): + return ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + batch_descriptor=BatchDescriptor(num_tokens=num_tokens), + ) + + +def test_replay_graph_matches_eager(monkeypatch): + """The replay graph of the next captured size pads and runs the rows.""" + monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) + capture = torch.cuda.stream(torch.cuda.Stream()) + monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) + compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4]) + cfg = MagicMock(compilation_config=compilation, speculative_config=None) + cfg.scheduler_config = MagicMock(max_num_seqs=2, max_num_batched_tokens=64) + cfg.cache_config.use_kda_recoverssm = False + cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) + cfg.model_config.dtype = torch.float32 + + def run_layers(hidden, positions, ids, pre, post, mix, residual): + return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] + + layers = DecoderReplayLayers(4, run_layers, [], set()) + manager = DecoderReplayCudaGraphManager(cfg, DEVICE, layers) + + def prepare(desc): + rows = torch.arange(desc.num_tokens) + layers.replay_batch = ReplayBatch(rows, _graph_context(desc.num_tokens)) + + manager.capture_replay_graphs(prepare) + states = [ + torch.randn(12, *b.shape[1:], device=DEVICE).to(b.dtype) + for b in manager.input_buffers + ] + rows = torch.tensor([11, 2, 7], device=DEVICE) + desc = manager.dispatch(2, 3, None, 0) + with override_forward_context(_graph_context(desc.num_tokens)): + actual = manager.run(rows, states) + for a, e in zip(actual, run_layers(*(t[rows] for t in states))): + torch.testing.assert_close(a, e, rtol=0, atol=0) diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index f6efe5c51608..810a1dbd9a02 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -12,6 +12,7 @@ from vllm.config import CUDAGraphMode from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.default import DefaultModelState @@ -57,7 +58,8 @@ def state(monkeypatch): cfg.scheduler_config.max_num_batched_tokens = 1024 cfg.parallel_config.data_parallel_size = 1 cfg.compilation_config.fast_moe_cold_start = False - layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None) + cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) builds: list = [] @@ -127,13 +129,8 @@ def _prepare(state, batch, cg_mode): [], GROUPS, ) - layers = state.decoder_replay_layers - if layers.rows is None: - assert layers.forward_context is None - return None, None - return SimpleNamespace( - rows=layers.rows, forward_context=layers.forward_context - ), state.builds[-1] + replay = state.decoder_replay_layers.replay_batch + return (None, None) if replay is None else (replay, state.builds[-1]) def test_replay_batch_keeps_each_request_window(state): @@ -191,9 +188,9 @@ def test_replay_batch_keeps_adaptive_verification_query_bound(state): assert replay is not None and build.batch.max_query_len == 200 -def test_only_eager_steps_trim(state): - """A CUDA graph keeps the layers on its whole batch; an eager step trims a - padded batch to its real rows.""" +def test_graph_steps_trim_only_piecewise_at_threshold(state): + """A CUDA graph keeps the layers on its whole batch, unless it is a PIECEWISE + one at the trim threshold or above; an eager step trims a batch to its rows.""" batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512) for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL): assert _prepare(state, batch, cg_mode) == (None, None) @@ -202,6 +199,10 @@ def test_only_eager_steps_trim(state): assert replay.rows.tolist() == REPLAY_ROWS assert build.cg_mode == CUDAGraphMode.NONE assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS) + state._trim_threshold = 256 + replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE) + assert replay is not None and replay.rows.tolist() == REPLAY_ROWS + assert _prepare(state, batch, CUDAGraphMode.FULL) == (None, None) def test_whole_batch_forwards_get_no_replay_batch(state): @@ -251,9 +252,25 @@ def test_dp_ranks_replay_together(dp_state): def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state): - dummy = _input_batch([512], [512], [False]) + """An idle rank's dummy keeps one row, so peers do not wait on its replay.""" + dummy = _input_batch([256, 256], [256, 256], [False, False]) dp_state.other = (True, 129) replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE) - assert replay is not None and replay.rows.shape[0] == WINDOW + assert replay is not None and replay.rows.shape[0] == 1 dp_metadata = replay.forward_context.dp_metadata - assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129] + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129] + + +def test_dp_ranks_pad_to_one_replay_graph(dp_state): + """Every rank pads its replay batch to the graph fitting the largest rank's.""" + dp_state.replay_cudagraphs = MagicMock() + dp_state.replay_cudagraphs.dispatch.return_value = BatchExecutionDescriptor( + CUDAGraphMode.PIECEWISE, 512, None + ) + dp_state.other = (True, 300) + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING) + replay, _ = _prepare(dp_state, batch, CUDAGraphMode.NONE) + assert dp_state.replay_cudagraphs.dispatch.call_args.args[:2] == (3, 300) + assert replay.forward_context.batch_descriptor.num_tokens == 512 + dp_metadata = replay.forward_context.dp_metadata + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [512, 512] diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 7586245a3543..0caa132d7e45 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -259,7 +259,9 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner: """Minimal V2 runner for capture_model: fakes everything except the cudagraph_manager's needs_capture decision.""" runner = GPUModelRunner.__new__(GPUModelRunner) - runner.model_state = SimpleNamespace(supports_mm_inputs=False) + runner.model_state = SimpleNamespace( + supports_mm_inputs=False, capture_inner_cudagraphs=lambda *args: None + ) runner.cudagraph_manager = SimpleNamespace( needs_capture=lambda: captured, capture=lambda *args, **kwargs: None, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index eefcf1f49e24..afe576daa9d3 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -570,6 +570,24 @@ class CompilationConfig: If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).""" + cudagraph_decoder_replay: bool = True + """Run the decoder replay layers of YOCO models (e.g. DeepSeek-V4.1) in CUDA + graphs of their own, and trim them in PIECEWISE graph steps. Requires + breakable PIECEWISE graphs; off under LoRA, prompt embeddings and non-first + PP ranks.""" + + decoder_replay_cudagraph_capture_sizes: list[int] = field(default_factory=list) + """Decoder replay CUDA graph sizes for YOCO models (e.g. DeepSeek-V4.1), at + most max_num_batched_tokens. If empty: multiples of sliding_window up to + min(max_cudagraph_capture_size, max_num_seqs * sliding_window), coarser past + 16 * sliding_window. Larger replay batches run eagerly.""" + + decoder_replay_trim_threshold: int = Field(default=768, ge=0) + """For YOCO models (e.g. DeepSeek-V4.1), PIECEWISE graphs with at least this + many padded tokens trim the decoder replay batch; must exceed the replay + window (128), below which nothing trims. Independent of the replay graph + capture sizes; eager steps still trim.""" + # Inductor capture compile_sizes: list[int | str] | None = None """Sizes to compile for inductor. In addition @@ -1033,6 +1051,11 @@ def __post_init__(self) -> None: f"Invalid backend for piecewise compilation: {self.backend}" ) + if any(s <= 0 for s in self.decoder_replay_cudagraph_capture_sizes): + raise ValueError( + "All decoder_replay_cudagraph_capture_sizes must be positive" + ) + # Validate encoder CUDA graph configuration if ( self.cudagraph_mm_encoder diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index 869029aa7484..d2d7ffd15364 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -7,25 +7,43 @@ ``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention metadata and a forward context of its own, like a microbatch; ``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers -under that context and scatters the outputs back to batch rows. Steps that run -in a CUDA graph keep the layers on the whole batch, inside the graph. +under that context and scatters the outputs back to batch rows. FULL graphs and +small PIECEWISE ones keep the layers on the whole batch; larger PIECEWISE graphs +break out to the replay batch, which may run in a graph of its own +(decoder_replay_cudagraph.py). """ from collections.abc import Callable +from dataclasses import dataclass import torch -from vllm.forward_context import ForwardContext, override_forward_context +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.forward_context import ( + ForwardContext, + in_piecewise_cudagraph, + override_forward_context, +) from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput +@dataclass +class ReplayBatch: + """One forward's replay batch; ``run_graph`` is set when it runs in a graph.""" + + rows: torch.Tensor + forward_context: ForwardContext + run_graph: Callable[..., tuple[torch.Tensor, ...]] | None = None + + class DecoderReplayLayers: """Runs the replay layers on the step's replay batch. ``run_layers`` takes a batch's layer inputs and returns its per-row outputs. ``row_buffers`` hold per-row results the source layer's indexer publishes for the layers after it; they are compacted to the replay rows - in place. + in place. ``metadata_prefixes`` are the attention metadata keys the layers + read, so the replay batch builds only those. """ def __init__( @@ -33,32 +51,57 @@ def __init__( window: int, run_layers: Callable[..., tuple[torch.Tensor, ...]], row_buffers: list[torch.Tensor], + metadata_prefixes: set[str], ) -> None: self.window = window self.run_layers = run_layers self.row_buffers = row_buffers - # The replay batch, set by the model state every step: its rows of the - # batch and its forward context. None runs the layers on the batch. - self.rows: torch.Tensor | None = None - self.forward_context: ForwardContext | None = None + self.metadata_prefixes = metadata_prefixes + # Set by the model state every step; None runs the layers on the batch. + self.replay_batch: ReplayBatch | None = None + self._break_outputs: list[torch.Tensor] | None = None def __call__( self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None ) -> tuple[torch.Tensor, ...]: - rows = self.rows - if rows is None: + if self.replay_batch is not None and in_piecewise_cudagraph(): + return self._run_in_graph_break(hidden_states, *states) + return self._run(hidden_states, *states) + + @eager_break_during_capture + def _run_in_graph_break( + self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None + ) -> tuple[torch.Tensor, ...]: + """A graph break: the outputs land in fixed buffers the next segment reads.""" + outputs = self._run(hidden_states, *states) + if self._break_outputs is None: + # Sized by the first call: the largest model graph is captured first. + self._break_outputs = [torch.empty_like(out) for out in outputs] + return tuple( + buf[: out.shape[0]].copy_(out) + for buf, out in zip(self._break_outputs, outputs) + ) + + def _run( + self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None + ) -> tuple[torch.Tensor, ...]: + replay_batch = self.replay_batch + if replay_batch is None: return self.run_layers(hidden_states, *states) - # A trimming step holds a prefill longer than the window, more tokens - # than any step whose MoE leaves its finalize to the next layer. + # Replay batch steps hold more than `window` tokens: no deferred MoE finalize. assert isinstance(hidden_states, torch.Tensor) + rows = replay_batch.rows num_rows = rows.shape[0] for buf in self.row_buffers: buf[:num_rows].copy_(buf.index_select(0, rows)) - with override_forward_context(self.forward_context): - row_outputs = self.run_layers( - hidden_states.index_select(0, rows), - *(None if t is None else t.index_select(0, rows) for t in states), - ) + inputs = (hidden_states, *states) + with override_forward_context(replay_batch.forward_context): + if replay_batch.run_graph is not None: + row_outputs = replay_batch.run_graph(rows, inputs) + else: + row_outputs = self.run_layers( + *(None if t is None else t.index_select(0, rows) for t in inputs) + ) # The trimmed rows' outputs stay zero; nothing reads them. num_tokens = hidden_states.shape[0] return tuple( diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py new file mode 100644 index 000000000000..871f535584c6 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA graphs of the decoder replay layers on a step's replay batch.""" + +from collections.abc import Callable +from copy import copy +from dataclasses import replace + +import torch + +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.forward_context import get_forward_context, override_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) + + +class DecoderReplayCudaGraphManager(CudaGraphManager): + """Breakable PIECEWISE graphs of the replay layers keyed by the replay batch's + padded size. They read the layer inputs from static buffers.""" + + def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers): + compilation = copy(vllm_config.compilation_config) + scheduler = vllm_config.scheduler_config + sizes = compilation.decoder_replay_cudagraph_capture_sizes + if not sizes: + bound = min( + compilation.max_cudagraph_capture_size, + scheduler.max_num_seqs * layers.window, + ) + # Replay batches hold at least `window` rows: step by window, or by 1/8 + # of the power-of-two interval once wider (256 in [2048, 4096)). + sizes, s = [bound], layers.window + while s < bound: + sizes.append(s) + s += max(layers.window, 1 << max(s.bit_length() - 4, 0)) + sizes = sorted(set(sizes)) + if max(sizes) > scheduler.max_num_batched_tokens: + raise ValueError( + "decoder_replay_cudagraph_capture_sizes must not exceed " + "max_num_batched_tokens" + ) + compilation.cudagraph_capture_sizes = sizes + compilation.max_cudagraph_capture_size = max(sizes) + replay_config = copy(vllm_config) + replay_config.compilation_config = compilation + super().__init__(replay_config, device, CUDAGraphMode.PIECEWISE, 1) + self.layers = layers + self.breakable_cg_runner: BreakableCUDAGraphWrapper = BreakableCUDAGraphWrapper( + layers.run_layers, replay_config + ) + config, act = vllm_config.model_config.hf_config, vllm_config.model_config.dtype + assert isinstance(act, torch.dtype) + hc, h, f32 = config.hc_mult, config.hidden_size, torch.float32 + mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + ids = torch.int64 if mega_moe else torch.int32 + # hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual + shapes = [(h,), (), (), (hc,), (hc, 1), (hc, hc), (hc, h)] + dtypes = [act, torch.int64, ids, f32, f32, f32, act] + n = compilation.max_cudagraph_capture_size + self.input_buffers = [ + torch.zeros(n, *s, dtype=d, device=device) for s, d in zip(shapes, dtypes) + ] + + def capture_replay_graphs( + self, prepare: Callable[[BatchExecutionDescriptor], None] + ) -> None: + """Capture each size on the replay batch ``prepare`` sets on the layers.""" + + def create_forward_fn(desc: BatchExecutionDescriptor, warmup: bool): + prepare(desc) + assert self.layers.replay_batch is not None + context = self.layers.replay_batch.forward_context + inputs = [buf[: desc.num_tokens] for buf in self.input_buffers] + + def forward_fn(cg_mode: CUDAGraphMode) -> None: + with override_forward_context( + replace(context, cudagraph_runtime_mode=cg_mode) + ): + self.breakable_cg_runner(*inputs) + + return forward_fn + + self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs") + self.layers.replay_batch = None + + def run(self, rows: torch.Tensor, inputs: tuple) -> tuple[torch.Tensor, ...]: + """Replay the forward context's graph on ``rows`` of ``inputs``.""" + num_rows = rows.shape[0] + for buf, t in zip(self.input_buffers, inputs): + torch.index_select(t, 0, rows, out=buf[:num_rows]) + batch_descriptor = get_forward_context().batch_descriptor + assert batch_descriptor is not None + num_tokens = batch_descriptor.num_tokens + outputs = self.breakable_cg_runner( + *(buf[:num_tokens] for buf in self.input_buffers) + ) + return tuple(out[:num_rows] for out in outputs) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index bd91e8f12b4f..9c3bfa7a927c 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -766,6 +766,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): and self.layers[cut].attn.swa_cache_layer.bounded_replay ): self.decoder_replay_start = cut + 1 + # The attention metadata keys the replay layers read. + metadata_prefixes = set() + for layer in islice(self.layers, cut + 1, self.end_layer): + attn = typing.cast(DeepseekV4DecoderLayer, layer).attn + metadata_prefixes.add(attn.swa_cache_layer.prefix) + if attn.compressed_cache_prefix is not None: + metadata_prefixes.add(attn.compressed_cache_prefix) + if attn.indexer is not None: + metadata_prefixes.add(attn.indexer.k_cache.prefix) self.decoder_replay_layers = DecoderReplayLayers( config.sliding_window, self._run_replay_layers, @@ -774,6 +783,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for buf in (self.topk_indices_buffer, self.candidate_block_buffer) if buf is not None ], + metadata_prefixes, ) logger.info_once( "Decoder SWA bounded replay: in eager prefill steps, layers " diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 30b287db22ee..c09e4582ea34 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -8,10 +8,17 @@ import torch.distributed as dist import torch.nn as nn +from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled from vllm.config import CUDAGraphMode, VllmConfig -from vllm.distributed.parallel_state import get_dp_group -from vllm.forward_context import DPMetadata, create_forward_context -from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers +from vllm.distributed.parallel_state import get_dp_group, get_pp_group +from vllm.forward_context import BatchDescriptor, DPMetadata, create_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import ( + DecoderReplayLayers, + ReplayBatch, +) +from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( + DecoderReplayCudaGraphManager, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -19,8 +26,10 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata @@ -230,6 +239,31 @@ def __init__( self.max_num_reqs, dtype=torch.int32, device=device ) self._replay_attn_groups: list[list[AttentionGroup]] | None = None + layers = self.decoder_replay_layers + compilation_config = vllm_config.compilation_config + self.replay_cudagraphs: DecoderReplayCudaGraphManager | None = None + # PIECEWISE graphs of this many tokens or more break out to the replay batch. + self._trim_threshold: int | None = None + if ( + layers is not None + and compilation_config.cudagraph_decoder_replay + and compilation_config.cudagraph_mode.has_piecewise_cudagraphs() + and is_breakable_cudagraph_enabled() + and get_pp_group().is_first_rank + # The graphs' fixed inputs hold no LoRA mapping and no None input ids. + and vllm_config.lora_config is None + and not vllm_config.model_config.enable_prompt_embeds + ): + threshold = compilation_config.decoder_replay_trim_threshold + if threshold <= layers.window: + raise ValueError( + "decoder_replay_trim_threshold must exceed the replay window " + f"of {layers.window} tokens: smaller graphs have nothing to trim" + ) + self.replay_cudagraphs = DecoderReplayCudaGraphManager( + vllm_config, device, layers + ) + self._trim_threshold = threshold def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) @@ -409,15 +443,21 @@ def _prepare_replay_batch( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, replay_start: torch.Tensor, + capture_desc: BatchExecutionDescriptor | None = None, ) -> None: - """Set the replay layers' batch for this forward, or none when nothing - trims. Only eager steps trim: a CUDA graph keeps the layers on the whole - batch, and ranks share the graph mode. Under data parallelism every rank - replays if any does.""" + """Set the replay layers' batch for this forward, or none. Eager steps trim; + captures and PIECEWISE graphs at the trim threshold or above break out to a + replay batch, trimmed or not; FULL and smaller PIECEWISE graphs keep the + whole batch. Under DP every rank replays if any does.""" layers = self.decoder_replay_layers assert layers is not None - layers.rows = layers.forward_context = None - if cudagraph_mode != CUDAGraphMode.NONE: + layers.replay_batch = None + force_replay_batch = capture_desc is not None or ( + cudagraph_mode == CUDAGraphMode.PIECEWISE + and self._trim_threshold is not None + and input_batch.num_tokens_after_padding >= self._trim_threshold + ) + if cudagraph_mode != CUDAGraphMode.NONE and not force_replay_batch: return num_reqs = input_batch.num_reqs @@ -427,20 +467,47 @@ def _prepare_replay_batch( query_lens, np.minimum(query_lens, layers.window), ) + # A batch of context-free non-prefill rows is a dummy (an idle DP rank's): + # nothing reads it and its idx_mapping hits stale slots, so keep one row. + seq_lens = input_batch.seq_lens_cpu_upper_bound[:num_reqs].numpy() + is_dummy = not input_batch.is_prefilling_np[:num_reqs].any() and bool( + (seq_lens == query_lens).all() + ) + if is_dummy and capture_desc is None: + kept_lens = np.zeros_like(kept_lens) + kept_lens[0] = 1 # Dummy batches (captures, an idle DP rank's) are not prefills: they # keep their rows unless a rank trims. - trims = bool( + needs_replay_batch = force_replay_batch or bool( (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any() ) num_tokens = int(kept_lens.sum()) - dp_metadata = None + dp_tokens = None if self.vllm_config.parallel_config.data_parallel_size > 1: - trims, dp_metadata = self._agree_across_dp(trims, num_tokens) - if not trims: + needs_replay_batch, dp_tokens = self._agree_across_dp( + needs_replay_batch, num_tokens + ) + if not needs_replay_batch: return + cudagraphs, run_graph = self.replay_cudagraphs, None + num_padded, cg_mode = num_tokens, CUDAGraphMode.NONE + if cudagraphs is not None: + # Every rank runs the graph that fits the largest rank's batch. + dp_max_tokens = num_tokens if dp_tokens is None else int(dp_tokens.max()) + desc = capture_desc or cudagraphs.dispatch(num_reqs, dp_max_tokens, None, 0) + if desc.cg_mode == CUDAGraphMode.PIECEWISE: + num_padded, cg_mode = desc.num_tokens, desc.cg_mode + run_graph = cudagraphs.run + if dp_tokens is not None: + dp_tokens.fill_(num_padded) + dp_metadata = None + if dp_tokens is not None: + dp_metadata = DPMetadata.make( + self.vllm_config.parallel_config, num_padded, dp_tokens + ) kept_batch, kept_slot_mappings = self._kept_input_batch( - input_batch, slot_mappings, replay_start, kept_lens + input_batch, slot_mappings, replay_start, kept_lens, num_padded ) attn_metadata = super().prepare_attn( kept_batch, @@ -453,22 +520,55 @@ def _prepare_replay_batch( self._kept_kv_start[:num_reqs] ), ) - layers.rows = self._kept_rows[:num_tokens] - layers.forward_context = create_forward_context( - attn_metadata, - self.vllm_config, - dp_metadata=dp_metadata, - slot_mapping=build_slot_mappings_by_layer( - kept_slot_mappings, kv_cache_config + layers.replay_batch = ReplayBatch( + self._kept_rows[:num_tokens], + create_forward_context( + attn_metadata, + self.vllm_config, + dp_metadata=dp_metadata, + cudagraph_runtime_mode=cg_mode, + batch_descriptor=BatchDescriptor(num_tokens=num_padded), + slot_mapping=build_slot_mappings_by_layer( + kept_slot_mappings, kv_cache_config + ), ), + run_graph, ) + def capture_inner_cudagraphs( + self, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + ) -> None: + def prepare(desc: BatchExecutionDescriptor) -> None: + # Capture reuses the runtime path on a dummy batch, so the graph reads + # the buffers a real step fills. + num_reqs = min(desc.num_tokens, self.max_num_reqs) + self._prepare_replay_batch( + InputBatch.make_dummy(num_reqs, desc.num_tokens, input_buffers), + CUDAGraphMode.NONE, + block_tables.get_dummy_block_tables(num_reqs), + block_tables.get_dummy_slot_mappings(desc.num_tokens), + attn_groups, + kv_cache_config, + self._replay_start[:num_reqs].zero_(), + capture_desc=desc, + ) + + if self.replay_cudagraphs is not None: + # profile_only captures use pre-KV-cache groups: drop the cache. + self._replay_attn_groups = None + self.replay_cudagraphs.capture_replay_graphs(prepare) + def _kept_input_batch( self, input_batch: InputBatch, slot_mappings: torch.Tensor, replay_start: torch.Tensor, kept_lens: np.ndarray, + num_padded: int, ) -> tuple[InputBatch, torch.Tensor]: """Build the sub-``InputBatch`` of each request's last ``kept_lens`` rows, and its slot mappings. @@ -486,7 +586,8 @@ def _kept_input_batch( input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np ) assert self._kept_slot_mappings is not None - kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens] + kept_slot_mappings = self._kept_slot_mappings[:, :num_padded] + kept_slot_mappings[:, num_tokens:].fill_(PAD_SLOT_ID) _gather_replay_batch_kernel[(num_reqs,)]( input_batch.query_start_loc, dropped_before, @@ -511,42 +612,45 @@ def _kept_input_batch( kept_batch = replace( input_batch, num_tokens=num_tokens, - num_tokens_after_padding=num_tokens, + num_tokens_after_padding=num_padded, query_start_loc=self._kept_query_start_loc[: num_reqs + 1], query_start_loc_np=kept_query_start_loc_np, max_query_len=max_query_len, - positions=self._kept_positions[:num_tokens], + positions=self._kept_positions[:num_padded], fast_prefill=None, ) return kept_batch, kept_slot_mappings def _agree_across_dp( - self, trims: bool, num_tokens: int - ) -> tuple[bool, DPMetadata | None]: - """Whether any rank trims and, then, the replay layers' DP metadata from - every rank's replay token count.""" + self, needs_replay_batch: bool, num_tokens: int + ) -> tuple[bool, torch.Tensor | None]: + """Whether any rank needs a replay batch and, if so, each rank's token count.""" parallel_config = self.vllm_config.parallel_config agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32) agreed[:, parallel_config.data_parallel_rank] = torch.tensor( - [trims, num_tokens] + [needs_replay_batch, num_tokens] ) if not should_skip_dp_coordination(): dist.all_reduce(agreed, group=get_dp_group().cpu_group) if not agreed[0].any(): return False, None - return True, DPMetadata.make(parallel_config, num_tokens, agreed[1]) + return True, agreed[1] def _replay_groups( self, attn_groups: list[list[AttentionGroup]] ) -> list[list[AttentionGroup]]: """The attention groups with metadata builders of their own, like each microbatch's: a builder keeps the metadata it built, and the runner's - hold the batch's.""" + hold the batch's. Only the groups the replay layers read are built.""" if self._replay_attn_groups is None: + assert self.decoder_replay_layers is not None + prefixes = self.decoder_replay_layers.metadata_prefixes self._replay_attn_groups = [] for groups in attn_groups: replay_groups = [] for group in groups: + if prefixes.isdisjoint(group.layer_names): + continue replay_group = replace(group, metadata_builders=[]) replay_group.create_metadata_builders( self.vllm_config, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b42289ef1111..e42917918afe 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1077,6 +1077,12 @@ def capture_model(self, *, profile_only: bool = False) -> int: self.lora_config, self ), ) + self.model_state.capture_inner_cudagraphs( + input_buffers, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + ) if self.speculator is not None: with use_workspace_lane(self._draft_workspace_lane): self.speculator.capture() diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index ee3e519a6e74..f03caf4dba83 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -17,7 +17,8 @@ from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.states import RequestState @@ -126,6 +127,16 @@ def remove_request(self, req_id: str) -> None: def apply_staged_writes(self) -> None: return None + def capture_inner_cudagraphs( + self, + input_buffers: InputBuffers, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + ) -> None: + """Capture the CUDA graphs this state runs inside the model's forward.""" + return None + def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: """Cudagraph support of attention groups this ModelState builds outside ``init_attn_backend`` (e.g. encoder-only layers). From 533d81b4b063a1ce167a191f97780021dae24401 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Tue, 6 Oct 2026 10:42:19 +0530 Subject: [PATCH 047/952] [ZenCPU][Test] Add zentorch/torch version pin consistency test (#52888) Signed-off-by: Chinmay Kulkarni Co-authored-by: Cursor Agent --- .buildkite/hardware_tests/cpu.yaml | 2 + docker/Dockerfile.cpu | 5 ++ docs/getting_started/installation/cpu.md | 3 + tests/test_zentorch_version_pin.py | 79 ++++++++++++++++++++++++ 4 files changed, 89 insertions(+) create mode 100644 tests/test_zentorch_version_pin.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index d54197f83678..fbd43d80ddb2 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -20,6 +20,7 @@ steps: - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/moe/test_zen_cpu_int8_moe.py + - tests/test_zentorch_version_pin.py - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py - tests/kernels/test_onednn.py - tests/kernels/test_awq_int4_to_int8.py @@ -45,6 +46,7 @@ steps: tests/kernels/moe/test_cpu_fused_moe.py \ tests/kernels/moe/test_cpu_quant_fused_moe.py \ tests/kernels/moe/test_zen_cpu_int8_moe.py \ + tests/test_zentorch_version_pin.py \ tests/kernels/mamba/test_cpu_short_conv.py \ tests/kernels/test_onednn.py \ tests/kernels/test_awq_int4_to_int8.py \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index f61752475d25..655a956ba7cd 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -441,6 +441,11 @@ ADD ./cmake/ ./cmake/ # (e.g. tests/tools/test_check_test_tethering.py imports # tools.pre_commit.check_test_tethering). ADD ./tools/ ./tools/ +# tests/test_zentorch_version_pin.py string-matches the zentorch extra in +# setup.py against the torch pin in requirements/cpu.txt. The wheel does not +# ship either file, and this stage does not inherit the vllm-src tree. +ADD ./setup.py ./setup.py +ADD ./requirements/cpu.txt ./requirements/cpu.txt # install development dependencies (for testing) RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 20834529b19b..fbd20f334a6c 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -241,6 +241,9 @@ export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/rel uv pip install "vllm[zen]" --extra-index-url https://wheels.vllm.ai/${VLLM_VERSION}/cpu --index-strategy first-index --torch-backend cpu ``` +!!! warning + `zentorch` is not supported with vLLM 0.27.0 and 0.27.1. + vLLM auto-detects the platform and routes linear layers through ZenDNN-optimized kernels - no flag needed. To verify it is engaged, look for the platform-selection line in the server's startup logs: ```bash diff --git a/tests/test_zentorch_version_pin.py b/tests/test_zentorch_version_pin.py new file mode 100644 index 000000000000..91125863beac --- /dev/null +++ b/tests/test_zentorch_version_pin.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Static consistency check between the pinned zentorch and torch versions. + +zentorch (the AMD Zen CPU optimization extra) tracks the torch release that +vLLM depends on, differing only by an extra trailing "patch level" segment. +Concretely, torch ``X.Y.Z`` corresponds to zentorch ``X.Y.Z.W``. + +This test parses the raw text of ``setup.py`` and ``requirements/cpu.txt`` and +string-matches the versions. +""" + +from pathlib import Path + +import pytest +import regex as re + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + +REPO_ROOT = Path(__file__).resolve().parent.parent +SETUP_PY = REPO_ROOT / "setup.py" +# zentorch is the AMD Zen *CPU* extra, so the CPU requirement file is the +# authoritative source for the torch version it must track. +CPU_REQUIREMENTS = REPO_ROOT / "requirements" / "cpu.txt" + + +def _extract_zentorch_version() -> str: + """Pull the pinned zentorch version out of setup.py's extras_require.""" + text = SETUP_PY.read_text() + matches = re.findall(r"zentorch==([0-9][0-9A-Za-z.\-+]*)", text) + assert matches, f"No 'zentorch==' pin found in {SETUP_PY}" + assert len(set(matches)) == 1, ( + f"Multiple conflicting zentorch pins found in {SETUP_PY}: {matches}" + ) + return matches[0] + + +def _extract_torch_version() -> str: + """Pull the pinned torch version out of requirements/cpu.txt. + + The local version label (e.g. the ``+cpu`` in ``2.13.0+cpu``) is stripped + so only the upstream release ``X.Y.Z`` is compared. + """ + text = CPU_REQUIREMENTS.read_text() + # Match lines like: torch==2.13.0+cpu; platform_machine == "x86_64" ... + # but not torchaudio / torchvision / open-clip-torch etc. + matches = re.findall(r"(?m)^\s*torch==([0-9][0-9A-Za-z.\-+]*)", text) + assert matches, f"No 'torch==' pin found in {CPU_REQUIREMENTS}" + # Strip local version labels (+cpu, +cu130, ...) and keep the release only. + releases = {m.split("+", 1)[0] for m in matches} + assert len(releases) == 1, ( + f"Multiple conflicting torch release pins found in {CPU_REQUIREMENTS}: " + f"{sorted(releases)}" + ) + return releases.pop() + + +def test_zentorch_pin_tracks_torch_release(): + zentorch_version = _extract_zentorch_version() + torch_version = _extract_torch_version() + + zentorch_parts = zentorch_version.split(".") + torch_parts = torch_version.split(".") + + # torch is X.Y.Z; zentorch is X.Y.Z.W (one extra trailing patch-level + # segment). Enforce both the shape and the shared prefix. + assert len(torch_parts) == 3, ( + f"Expected torch release of the form 'X.Y.Z', got {torch_version!r}" + ) + assert len(zentorch_parts) == 4, ( + f"Expected zentorch pin of the form 'X.Y.Z.W', got {zentorch_version!r}" + ) + assert zentorch_parts[:3] == torch_parts, ( + "zentorch version is out of sync with torch. " + f"setup.py pins zentorch=={zentorch_version} (base " + f"{'.'.join(zentorch_parts[:3])}) but requirements/cpu.txt pins " + f"torch=={torch_version}. When bumping torch, bump the first three " + "components of the zentorch pin in setup.py to match." + ) From 9a1f6fd4cbe4ed8fd26215bc32a3f285b6a92d21 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 6 Oct 2026 13:20:09 +0800 Subject: [PATCH 048/952] [Model][DeepSeek-V4.1] Permute mega-attention wq_b / wo_a while loading (#60064) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/kernels/test_dsv41_mega_attn_layouts.py | 42 +++++++++++-------- .../deepseek_v41/common/ops/fused_layout.py | 36 ++++++---------- vllm/models/deepseek_v41/nvidia/dspark.py | 14 ------- .../nvidia/flash_mla_mega_attn.py | 38 +++++------------ vllm/models/deepseek_v41/nvidia/model.py | 12 ------ 5 files changed, 48 insertions(+), 94 deletions(-) diff --git a/tests/kernels/test_dsv41_mega_attn_layouts.py b/tests/kernels/test_dsv41_mega_attn_layouts.py index 6f9ff56bd236..3d454183605d 100644 --- a/tests/kernels/test_dsv41_mega_attn_layouts.py +++ b/tests/kernels/test_dsv41_mega_attn_layouts.py @@ -3,25 +3,34 @@ """Layout plumbing for DeepSeek V4.1 mega attention. The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and -``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs -agree with a reference that permutes the activation instead. +``wo_a`` columns in their weight loaders, so what has to hold is that the +permuted GEMMs agree with a reference that permutes the activation instead. """ import torch +from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.models.deepseek_v41.common.ops.fused_layout import ( WV_GROUP_SIZE, o_fused_chunk_permutation, o_fused_permutation, + permute_on_load, permute_q_to_fused, - permute_wo_a_, - permute_wq_b_, q_fused_permutation, ) HEAD_DIM = 512 +def _load(checkpoint: torch.Tensor, perm: torch.Tensor, dim: int) -> torch.Tensor: + """Load ``checkpoint`` into a fresh parameter through a permuting loader.""" + param = torch.nn.Parameter(torch.empty_like(checkpoint), requires_grad=False) + param.weight_loader = default_weight_loader # type: ignore[attr-defined] + permute_on_load(param, perm, dim) + param.weight_loader(param, checkpoint) # type: ignore[attr-defined] + return param.data + + def test_permuted_wq_b_gemm_matches_permuted_activation(): """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly.""" torch.manual_seed(0) @@ -32,18 +41,17 @@ def test_permuted_wq_b_gemm_matches_permuted_activation(): 0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8 ) qr = torch.randn(num_tokens, q_lora_rank) - orig_weight, orig_scale = weight.clone(), scale.clone() standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) expected = permute_q_to_fused(standard) - permute_wq_b_(weight, scale, num_heads) - fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) + perm = q_fused_permutation(num_heads, HEAD_DIM) + loaded_weight = _load(weight, perm, dim=0) + fused = (qr @ loaded_weight.T).view(num_tokens, num_heads, HEAD_DIM) torch.testing.assert_close(fused, expected) # The scale has to follow its row, or dequant reads another row's scale. - perm = q_fused_permutation(num_heads, HEAD_DIM) - torch.testing.assert_close(weight, orig_weight[perm]) - torch.testing.assert_close(scale, orig_scale[perm]) + torch.testing.assert_close(loaded_weight, weight[perm]) + torch.testing.assert_close(_load(scale, perm, dim=0), scale[perm]) def test_permuted_wo_a_consumes_fused_output(): @@ -54,18 +62,18 @@ def test_permuted_wo_a_consumes_fused_output(): weight = torch.randn(out_features, in_features) scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8) standard_o = torch.randn(num_tokens, in_features) - orig_weight, orig_scale = weight.clone(), scale.clone() expected = standard_o @ weight.T # The kernel emits the same values in the fused chunk order. perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM) fused_o = standard_o[:, perm] - permute_wo_a_(weight, scale, WV_GROUP_SIZE) - torch.testing.assert_close(weight, orig_weight[:, perm]) - torch.testing.assert_close( - scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)] - ) + loaded_weight = _load(weight, perm, dim=1) + torch.testing.assert_close(loaded_weight, weight[:, perm]) + chunk_perm = o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM) + torch.testing.assert_close(_load(scale, chunk_perm, dim=1), scale[:, chunk_perm]) # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the # summation order changes and the result differs in the last few ulps. - torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3) + torch.testing.assert_close( + fused_o @ loaded_weight.T, expected, rtol=1e-3, atol=1e-3 + ) diff --git a/vllm/models/deepseek_v41/common/ops/fused_layout.py b/vllm/models/deepseek_v41/common/ops/fused_layout.py index 071afde6bca4..50df8de3514a 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_layout.py +++ b/vllm/models/deepseek_v41/common/ops/fused_layout.py @@ -12,6 +12,8 @@ import torch +from vllm.model_executor.model_loader.weight_utils import composed_weight_loader + HEAD_DIM = 512 Q_CHUNK = 16 O_CHUNK = 32 @@ -60,27 +62,13 @@ def _bytes_view(t: torch.Tensor) -> torch.Tensor: return t.view(torch.uint8) if t.element_size() == 1 else t -def permute_wq_b_( - weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int -) -> None: - """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place.""" - head_dim = weight.shape[0] // num_local_heads - perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device) - for t in (weight, weight_scale): - b = _bytes_view(t) - b.copy_(b[perm]) - - -def permute_wo_a_( - weight: torch.Tensor, - weight_scale: torch.Tensor, - heads_per_group: int = WV_GROUP_SIZE, -) -> None: - """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place.""" - head_dim = weight.shape[1] // heads_per_group - perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device) - w = _bytes_view(weight) - w.copy_(w[:, perm]) - chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim) - s = _bytes_view(weight_scale) - s.copy_(s[:, chunk_perm.to(weight.device)]) +def permute_on_load(param: torch.Tensor, perm: torch.Tensor, dim: int) -> None: + """Make ``param``'s weight loader gather the loaded shard by ``perm``.""" + + def gather(t: torch.Tensor) -> torch.Tensor: + return _bytes_view(t).index_select(dim, perm.to(t.device)).view(t.dtype) + + param.weight_loader = composed_weight_loader( # type: ignore[attr-defined] + param.weight_loader, # type: ignore[attr-defined] + gather, + ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index d1534d52483f..a09c3159c235 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -500,22 +500,8 @@ def _finalize_moe(self) -> None: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() - def _finalize_attn(self) -> None: - """Run the attention backend's post-load weight step on the draft layers. - - They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide - attention backend and owe it whatever the target model owes it -- mega - attention permutes wq_b / wo_a here and refuses to run without it. - Idempotent, like the target's. - """ - for layer in self.model.layers: - finalize = getattr(layer.attn, "finalize_loaded_weights", None) - if finalize is not None: - finalize() - def process_weights_after_loading(self) -> None: self._finalize_moe() - self._finalize_attn() def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 8efc987739b1..7a7f4f8c9016 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -32,8 +32,10 @@ ) from vllm.models.deepseek_v41.common.ops.fused_layout import ( WV_GROUP_SIZE, - permute_wo_a_, - permute_wq_b_, + o_fused_chunk_permutation, + o_fused_permutation, + permute_on_load, + q_fused_permutation, ) from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.models.deepseek_v41.sparse_mla import ( @@ -182,7 +184,13 @@ def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: "wo_a group." ) self.n_wv_group = self.padded_heads // WV_GROUP_SIZE - self._fused_layouts_ready = False + q_perm = q_fused_permutation(self.n_local_heads, self.head_dim) + permute_on_load(self.wq_b.weight, q_perm, dim=0) + permute_on_load(self.wq_b.weight_scale, q_perm, dim=0) + o_perm = o_fused_permutation(WV_GROUP_SIZE, self.head_dim) + permute_on_load(self.wo_a.weight, o_perm, dim=1) + o_chunk_perm = o_fused_chunk_permutation(WV_GROUP_SIZE, self.head_dim) + permute_on_load(self.wo_a.weight_scale, o_chunk_perm, dim=1) # ---- interface contract ------------------------------------------------ @@ -227,25 +235,6 @@ def _o_proj( ) return self._wo_b_proj(z.flatten(1)) - # ---- weights ----------------------------------------------------------- - - def finalize_loaded_weights(self) -> None: - """Permute wq_b rows / wo_a columns into the kernel's layouts. - - Idempotent: a second post-load pass must not permute twice. - """ - if self._fused_layouts_ready: - return - permute_wq_b_( - self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads - ) - permute_wo_a_( - self.wo_a.weight.data, - self.wo_a.weight_scale.data, - self.n_local_heads // self.n_local_groups, - ) - self._fused_layouts_ready = True - # ---- forward ----------------------------------------------------------- def forward_mqa( @@ -255,11 +244,6 @@ def forward_mqa( positions: torch.Tensor, output: QuantizedActivation, ) -> None: - if not self._fused_layouts_ready: - raise RuntimeError( - f"{self.prefix}: wq_b / wo_a were never permuted for the mega " - "attention kernel; refusing to run with mismatched layouts." - ) attn_metadata = get_forward_context().attn_metadata if attn_metadata is None: # Warmup dummy run: reserve the prefill workspace, produce zeros. diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 9c3bfa7a927c..94fa0b4b8b67 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -1364,17 +1364,6 @@ def finalize_mega_moe_weights(self) -> None: for layer in islice(self.layers, self.start_layer, self.end_layer): layer.ffn.finalize_mega_moe_weights() - def finalize_mega_attn_weights(self) -> None: - """Permute wq_b / wo_a into FlashMLA's mega-attention layouts. - - A no-op for every other attention layer, and idempotent, so a second - post-load pass cannot permute twice. - """ - for layer in islice(self.layers, self.start_layer, self.end_layer): - finalize = getattr(layer.attn, "finalize_loaded_weights", None) - if finalize is not None: - finalize() - def finalize_mhc_broadcast_weights(self) -> None: if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer: return @@ -1647,7 +1636,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() - self.model.finalize_mega_attn_weights() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() From 6fddc9b21e763bfc947cc01d3ce815179aac1dc6 Mon Sep 17 00:00:00 2001 From: czhu-cohere Date: Mon, 5 Oct 2026 22:21:19 -0700 Subject: [PATCH 049/952] [Profiler] Add per-session profiling controls to Python and Rust frontends (#57875) Signed-off-by: czhu Signed-off-by: Conway Zhu Co-authored-by: aoshen02 --- docs/contributing/profiling.md | 37 +++ rust/src/engine-core-client/src/client.rs | 34 +- rust/src/engine-core-client/src/error.rs | 2 + rust/src/server/src/error.rs | 9 + rust/src/server/src/routes/profile.rs | 77 ++++- rust/src/server/src/routes/tests.rs | 290 +++++++++++++++-- tests/benchmarks/test_audio_dataset.py | 47 +++ .../serve/test_profile_api_router.py | 86 ++++++ tests/v1/engine/test_async_llm.py | 292 +++++++++++++++++- tests/v1/engine/test_llm_engine.py | 33 ++ tests/v1/worker/test_gpu_profiler.py | 184 ++++++++++- vllm/benchmarks/lib/endpoint_request_func.py | 30 +- vllm/benchmarks/serve.py | 33 +- vllm/config/profiler.py | 46 ++- vllm/engine/protocol.py | 10 +- vllm/entrypoints/llm.py | 20 +- vllm/entrypoints/serve/profile/api_router.py | 23 +- vllm/exceptions.py | 11 + vllm/profiler/wrapper.py | 130 ++++++-- vllm/v1/engine/async_llm.py | 91 ++++-- vllm/v1/engine/core.py | 15 +- vllm/v1/engine/core_client.py | 59 +++- vllm/v1/engine/llm_engine.py | 32 +- vllm/v1/executor/abstract.py | 39 ++- vllm/v1/worker/gpu_worker.py | 21 +- 25 files changed, 1492 insertions(+), 159 deletions(-) create mode 100644 tests/entrypoints/serve/test_profile_api_router.py diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 86bd68acef0c..5e5e27026372 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -31,6 +31,18 @@ When using `vllm bench serve`, you can enable profiling by passing the `--profil Traces can be visualized using . +Each worker writes its trace and summary table under `torch_profiler_dir`, +named after its rank, for example `rank0..pt.trace.json.gz` and +`rank0.profiler_out.txt`. Depending on the parallel configuration, the rank +name can also include the parallel ranks, such as +`dp0_pp0_tp0_dcp0_ep0_rank0`. The AsyncLLM CPU trace uses +`_.async_llm` instead of the rank name. When a session sets +`profile_prefix`, it is prepended as `_`, so summaries from +different sessions do not overwrite each other. + +!!! note + The summary table was previously written to `profiler_out_.txt`. + !!! tip You can directly call bench module without installing vLLM using `python -m vllm.entrypoints.cli.main bench`. @@ -88,6 +100,31 @@ curl -X POST http://localhost:8000/v1/chat/completions \ $ curl -X POST http://localhost:8000/stop_profile ``` +`/start_profile` also accepts an optional JSON body with settings for that +profiling session. The Python and Rust API frontends accept the same fields: + +```shell +curl -X POST http://localhost:8000/start_profile \ + -H "Content-Type: application/json" \ + -d '{ + "profile_prefix": "sharegpt_run_1", + "delay_iterations": 5000, + "max_iterations": 20 + }' +``` + +`profile_prefix` is prepended to the generated trace and summary file names. It +must be 1-128 ASCII characters, start with a letter or number, and contain only +letters, numbers, `.`, `_`, or `-`. `delay_iterations` skips worker iterations +before collection starts, and `max_iterations` limits the number of collected +worker iterations (`0` means no limit). Both must be non-negative JSON integers; +values of the wrong type, such as `"5"` or `5.0`, are rejected rather than +converted. Unknown fields are ignored. These iteration +bounds do not apply to the AsyncLLM frontend CPU profiler, which records from +`/start_profile` until `/stop_profile`. Always call `/stop_profile`, including +after the worker reaches `max_iterations`; starting another session before +stopping the current one returns HTTP 409. + ## Profile with Triton Proton [Proton](https://github.com/triton-lang/triton/tree/main/third_party/proton) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 0514292b902e..36a4fbffa5e1 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -9,7 +9,7 @@ use futures::future::join_all; use itertools::Itertools; use serde::Serialize; use serde_json::Value as JsonValue; -use tokio::sync::mpsc; +use tokio::sync::{Mutex, mpsc}; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info, trace}; @@ -283,6 +283,9 @@ pub struct EngineCoreClient { inner: Arc, coordinator: Option, abort_tx: mpsc::UnboundedSender, + /// Whether a profiling session is awaiting an explicit stop. Held across each + /// start/stop utility call so a stale stop reply cannot clear a newer session. + profile_active: Mutex, /// Runtime used to send messages to the engine and drive all background tasks. runtime: BackgroundShutdownRuntime, @@ -431,6 +434,7 @@ impl EngineCoreClient { inner, coordinator, abort_tx, + profile_active: Mutex::new(false), runtime, output_task, dispatcher_task, @@ -963,14 +967,38 @@ impl EngineCoreClient { } /// Start profiling the engine. - pub async fn start_profile(&self, profile_prefix: Option<&str>) -> Result<()> { - self.call_utility::<(), _>("profile", (true, profile_prefix)).await?; + pub async fn start_profile( + &self, + profile_prefix: Option<&str>, + delay_iterations: Option, + max_iterations: Option, + ) -> Result<()> { + let mut active = self.profile_active.lock().await; + if *active { + return Err(Error::ProfileAlreadyActive); + } + + // A cancelled start may still reach the engine, so only a reported + // failure clears the session; otherwise a stop is needed. + *active = true; + if let Err(error) = self + .call_utility::<(), _>( + "profile", + (true, profile_prefix, delay_iterations, max_iterations), + ) + .await + { + *active = false; + return Err(error); + } Ok(()) } /// Stop profiling the engine. pub async fn stop_profile(&self, profile_prefix: Option<&str>) -> Result<()> { + let mut active = self.profile_active.lock().await; self.call_utility::<(), _>("profile", (false, profile_prefix)).await?; + *active = false; Ok(()) } diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index d43952f0239f..d55f2cbbcb9e 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -83,6 +83,8 @@ pub enum Error { ClientClosed { message: String }, #[error("request output stream for `{request_id}` closed unexpectedly")] RequestStreamClosed { request_id: String }, + #[error("profiling session is already active")] + ProfileAlreadyActive, #[error("utility call `{method}` failed (call_id={call_id}): {message}")] UtilityCallFailed { method: String, diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 5f7adc8fbb90..77f7295b561d 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -22,6 +22,8 @@ pub enum ApiError { ModelNotFound { model: String }, /// The request body could not be parsed as valid JSON. JsonParseError { message: String }, + /// The requested operation conflicts with current server state. + Conflict { message: String }, /// An unexpected internal failure happened before streaming started. ServerError { message: String }, } @@ -32,6 +34,7 @@ impl ApiError { match self { Self::InvalidRequest { .. } => StatusCode::BAD_REQUEST, Self::ModelNotFound { .. } => StatusCode::NOT_FOUND, + Self::Conflict { .. } => StatusCode::CONFLICT, Self::ServerError { .. } => StatusCode::INTERNAL_SERVER_ERROR, Self::JsonParseError { .. } => StatusCode::BAD_REQUEST, } @@ -65,6 +68,12 @@ impl ApiError { param: None, code: Some("json_parse_error".to_string()), }, + Self::Conflict { message } => ErrorDetail { + message: message.clone(), + error_type: "conflict_error".to_string(), + param: None, + code: Some("conflict_error".to_string()), + }, }; ErrorResponse { error } diff --git a/rust/src/server/src/routes/profile.rs b/rust/src/server/src/routes/profile.rs index de793a7b3b00..763a4f8c1b8b 100644 --- a/rust/src/server/src/routes/profile.rs +++ b/rust/src/server/src/routes/profile.rs @@ -3,22 +3,75 @@ use std::sync::Arc; +use axum::Json; +use axum::body::Bytes; use axum::extract::State; use axum::http::StatusCode; +use serde::Deserialize; use tracing::info; +use vllm_engine_core_client::Error as EngineCoreClientError; use crate::error::ApiError; use crate::state::AppState; use crate::utils::utility_call_error; +#[derive(Debug, Default, Deserialize)] +pub(crate) struct StartProfileRequest { + profile_prefix: Option, + delay_iterations: Option, + max_iterations: Option, +} + +fn valid_profile_prefix(prefix: &str) -> bool { + let bytes = prefix.as_bytes(); + (1..=128).contains(&bytes.len()) + && bytes[0].is_ascii_alphanumeric() + && bytes[1..] + .iter() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-')) +} + /// Start profiling the engine. -pub async fn start_profile(State(state): State>) -> Result { +pub async fn start_profile( + State(state): State>, + body: Bytes, +) -> Result { + let body = if body.is_empty() { + StartProfileRequest::default() + } else { + Json::::from_bytes(&body)?.0 + }; + if body + .profile_prefix + .as_deref() + .is_some_and(|prefix| !valid_profile_prefix(prefix)) + { + return Err(ApiError::invalid_request( + "profile_prefix must be 1-128 ASCII characters, start with a letter or \ + number, and contain only letters, numbers, '.', '_', or '-'" + .to_string(), + Some("profile_prefix"), + )); + } info!("starting profiler"); - state + if let Err(error) = state .engine_core_client() - .start_profile(None) + .start_profile( + body.profile_prefix.as_deref(), + body.delay_iterations, + body.max_iterations, + ) .await - .map_err(|error| utility_call_error("start_profile", error))?; + { + if matches!(error, EngineCoreClientError::ProfileAlreadyActive) { + return Err(ApiError::conflict( + "A profiling session is already active. Call /stop_profile before \ + starting another session." + .to_string(), + )); + } + return Err(utility_call_error("start_profile", error)); + } info!("profiler started"); Ok(StatusCode::OK) } @@ -34,3 +87,19 @@ pub async fn stop_profile(State(state): State>) -> Result) -> u64 { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + assert_eq!(array[2], Value::from("profile")); + assert_eq!(array[3], Value::Array(expected_args)); + array[1].as_u64().expect("call id") +} - assert_eq!(array[2], Value::from("profile")); - assert_eq!(array[3], Value::Array(vec![Value::from(false), Value::Nil])); +/// A start issued while a stop is in flight waits for that stop instead of +/// racing it, so a stale stop reply cannot clear the new session. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn start_profile_waits_for_inflight_stop() { + let (stop_received_tx, stop_received_rx) = tokio::sync::oneshot::channel(); + let (release_stop_tx, release_stop_rx) = tokio::sync::oneshot::channel::<()>(); + let (app, engine_task) = test_profiling_app_with_engine_script(move |dealer, push| { + boxed_test_future(async move { + let start = || vec![Value::from(true), Value::Nil, Value::Nil, Value::Nil]; + let call_id = recv_profile_call(dealer, start()).await; + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + + let call_id = recv_profile_call(dealer, vec![Value::from(false), Value::Nil]).await; + stop_received_tx.send(()).expect("notify stop received"); + release_stop_rx.await.expect("release stop"); + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + + let call_id = recv_profile_call(dealer, start()).await; + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let post = |uri: &'static str| { + let mut app = app.clone(); + async move { + app.call( + Request::builder() + .method("POST") + .uri(uri) + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app") + .status() + } + }; + + let timeout = Duration::from_secs(2); + assert_eq!(post("/start_profile").await, StatusCode::OK); + let stop = tokio::spawn(post("/stop_profile")); + tokio::time::timeout(timeout, stop_received_rx) + .await + .expect("stop dispatch timeout") + .expect("stop received"); + + // Nothing before the profile lock awaits, so one poll reaches it; an + // unserialized client would instead reject against the stale session. + let mut restart = std::pin::pin!(post("/start_profile")); + assert!(futures::poll!(restart.as_mut()).is_pending()); + + release_stop_tx.send(()).expect("release stop"); + let stop = tokio::time::timeout(timeout, stop).await.expect("stop timeout"); + assert_eq!(stop.expect("stop task"), StatusCode::OK); + let restart = tokio::time::timeout(timeout, restart).await.expect("restart timeout"); + assert_eq!(restart, StatusCode::OK); + tokio::time::timeout(timeout, engine_task) + .await + .expect("mock engine timeout") + .expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn start_profile_route_rejects_invalid_prefix() { + let (app, engine_task) = + test_profiling_app_with_engine_script(|_, _| boxed_test_future(async move {})).await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/start_profile") + .header("content-type", "application/json") + .body(Body::from( + json!({"profile_prefix": "../trace"}).to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + engine_task.await.expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stop_profile_route_sends_expected_utility_call() { + let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let stop = vec![Value::from(false), Value::Nil]; + let call_id = recv_profile_call(dealer, stop).await; send_outputs(push, utility_outputs(call_id, utility_none_result())).await; }) }) diff --git a/tests/benchmarks/test_audio_dataset.py b/tests/benchmarks/test_audio_dataset.py index e81c0cd1737c..54a165dc61fd 100644 --- a/tests/benchmarks/test_audio_dataset.py +++ b/tests/benchmarks/test_audio_dataset.py @@ -98,6 +98,17 @@ def post(self, *, url: str, data: _FakeFormData, headers: dict[str, str]): return _FakeResponse() +class _FakeProfileSession: + def __init__(self) -> None: + self.url: str | None = None + self.headers: dict[str, str] | None = None + + def post(self, *, url: str, headers: dict[str, str]): + self.url = url + self.headers = headers + return _FakeResponse() + + def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None: audio_path = tmp_path / "earnings.wav" _write_wav(audio_path, duration_s=0.1) @@ -222,6 +233,42 @@ def test_async_request_openai_audio_handles_decoded_audio_arrays( assert output.generated_text == "hello" +def test_audio_request_rejects_profile_endpoint() -> None: + request_input = RequestFuncInput( + prompt="", + api_url="http://localhost:8000/start_profile", + prompt_len=1, + output_len=32, + model="openai/whisper-large-v3", + ) + + with pytest.raises(ValueError, match="OpenAI Audio API"): + asyncio.run( + request_func_module.async_request_openai_audio( + request_input, _FakeProfileSession() + ) + ) + + +@pytest.mark.parametrize("endpoint", ["start_profile", "stop_profile"]) +def test_audio_benchmark_profile_request_has_no_inference_payload( + endpoint: str, +) -> None: + session = _FakeProfileSession() + + output = asyncio.run( + request_func_module.async_request_profile( + f"http://localhost:8000/{endpoint}", + session, + extra_headers={"X-Test-Header": "profile"}, + ) + ) + + assert output.success is True + assert session.url == f"http://localhost:8000/{endpoint}" + assert session.headers == {"X-Test-Header": "profile"} + + _COHERE_ASR_PROMPT = ( "<|startofcontext|><|startoftranscript|>" "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>" diff --git a/tests/entrypoints/serve/test_profile_api_router.py b/tests/entrypoints/serve/test_profile_api_router.py new file mode 100644 index 000000000000..c5ceddd1bc44 --- /dev/null +++ b/tests/entrypoints/serve/test_profile_api_router.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm.entrypoints.serve.profile.api_router import router + + +@pytest.fixture +def engine_client() -> AsyncMock: + return AsyncMock() + + +@pytest.fixture +def client(engine_client: AsyncMock) -> TestClient: + app = FastAPI() + app.state.engine_client = engine_client + app.include_router(router) + return TestClient(app) + + +def test_start_profile_without_body_is_backwards_compatible( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post("/start_profile") + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with() + + +def test_start_profile_ignores_completion_payload( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post( + "/start_profile", + json={ + "model": "Qwen/Qwen3-0.6B", + "prompt": "profile this request", + "max_tokens": 16, + "stream": False, + }, + ) + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with() + + +def test_start_profile_forwards_session_overrides( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post( + "/start_profile", + json={ + "profile_prefix": "sharegpt_run-1", + "delay_iterations": 5000, + "max_iterations": 20, + }, + ) + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with( + profile_prefix="sharegpt_run-1", + delay_iterations=5000, + max_iterations=20, + ) + + +@pytest.mark.parametrize( + "body", + [ + {"profile_prefix": 1}, + {"delay_iterations": "1"}, + {"max_iterations": "20"}, + ], +) +def test_start_profile_rejects_wrongly_typed_overrides( + client: TestClient, engine_client: AsyncMock, body: dict +) -> None: + response = client.post("/start_profile", json=body) + + assert response.status_code == 422 + engine_client.start_profile.assert_not_awaited() diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index dc3cf1aa001d..3ac687c918f1 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import threading import time from contextlib import ExitStack +from http import HTTPStatus from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -11,7 +13,7 @@ import vllm.v1.engine.async_llm as async_llm_module from vllm import SamplingParams from vllm.assets.image import ImageAsset -from vllm.config import LoggingConfig, VllmConfig +from vllm.config import LoggingConfig, ProfilerConfig, VllmConfig from vllm.engine.arg_utils import AsyncEngineArgs from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, @@ -20,10 +22,11 @@ from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import ProfilerAlreadyActiveError, VLLMValidationError from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform +from vllm.profiler.wrapper import WorkerProfiler from vllm.sampling_params import RequestOutputKind from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine.async_llm import AsyncLLM @@ -87,13 +90,19 @@ def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch): return engine_core -def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): +def _mock_vllm_config() -> MagicMock: vllm_config = MagicMock() vllm_config.logging_config = LoggingConfig() vllm_config.observability_config.otlp_traces_endpoint = None vllm_config.scheduler_config.stream_interval = 1 + return vllm_config + + +def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): + vllm_config = _mock_vllm_config() vllm_config.profiler_config.profiler = "cuda" vllm_config.profiler_config.ignore_frontend = False + vllm_config.profiler_config.should_profile_frontend = False engine_core = _mock_async_llm_dependencies(monkeypatch) engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) @@ -105,22 +114,22 @@ async def profile(): asyncio.run(profile()) assert engine.profiler is None - engine_core.profile_async.assert_has_awaits([call(True, None), call(False)]) + engine_core.profile_async.assert_has_awaits( + [call(True, None, None, None), call(False)] + ) def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( monkeypatch: pytest.MonkeyPatch, ): - vllm_config = MagicMock() - vllm_config.logging_config = LoggingConfig() - vllm_config.observability_config.otlp_traces_endpoint = None - vllm_config.scheduler_config.stream_interval = 1 + vllm_config = _mock_vllm_config() vllm_config.profiler_config.profiler = "torch" vllm_config.profiler_config.ignore_frontend = False vllm_config.profiler_config.torch_profiler_activities = ["CUDA"] + vllm_config.profiler_config.should_profile_frontend = False _mock_async_llm_dependencies(monkeypatch) profiler = MagicMock() - monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler) + monkeypatch.setattr(async_llm_module, "create_frontend_profiler", profiler) engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) @@ -128,6 +137,271 @@ def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( profiler.assert_not_called() +def test_profile_forwards_overrides_rejects_duplicate_and_allows_restart( + monkeypatch: pytest.MonkeyPatch, +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config.should_profile_frontend = False + engine_core = _mock_async_llm_dependencies(monkeypatch) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + async def profile(): + for kwargs in ( + {"profile_prefix": "../trace"}, + {"delay_iterations": -1}, + {"max_iterations": -1}, + ): + with pytest.raises(ValueError): + await engine.start_profile(**kwargs) + await engine.start_profile( + "session", + delay_iterations=5, + max_iterations=2, + ) + with pytest.raises(ProfilerAlreadyActiveError) as exc_info: + await engine.start_profile("duplicate") + assert exc_info.value.http_status == HTTPStatus.CONFLICT + await engine.stop_profile() + await engine.start_profile("next-session") + await engine.stop_profile() + + asyncio.run(profile()) + + engine_core.profile_async.assert_has_awaits( + [ + call(True, "session", 5, 2), + call(False), + call(True, "next-session", None, None), + call(False), + ] + ) + + +def test_stop_profile_waits_for_inflight_start(monkeypatch: pytest.MonkeyPatch): + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + events: list[str] = [] + frontend_profiler = MagicMock() + frontend_profiler.start.side_effect = lambda: events.append("frontend_start") + frontend_profiler.stop.side_effect = lambda: events.append("frontend_stop") + + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + start_rpc_entered = asyncio.Event() + release_start_rpc = asyncio.Event() + + async def profile_async(is_start, *args): + if is_start: + start_rpc_entered.set() + await release_start_rpc.wait() + events.append("engine_start_done") + else: + events.append("engine_stop") + + engine_core.profile_async.side_effect = profile_async + + start_task = asyncio.create_task(engine.start_profile()) + await start_rpc_entered.wait() + + stop_task = asyncio.create_task(engine.stop_profile()) + await asyncio.sleep(0) + release_start_rpc.set() + await asyncio.wait_for(asyncio.gather(start_task, stop_task), timeout=5) + + asyncio.run(profile()) + + assert events[:2] == ["engine_start_done", "frontend_start"] + assert sorted(events[2:]) == ["engine_stop", "frontend_stop"] + assert engine._profile_session_active is False + + +@pytest.mark.parametrize("failure", ["error", "cancel"]) +def test_interrupted_engine_profile_start_skips_frontend_profiler( + monkeypatch: pytest.MonkeyPatch, failure: str +): + """The frontend trace starts only after the engine start succeeds. A + cancelled start may still reach the engine, so it keeps the session active + until /stop_profile.""" + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + frontend_profiler = MagicMock() + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + start_rpc_entered = asyncio.Event() + + async def profile_async(is_start, *args): + if not is_start: + return + start_rpc_entered.set() + if failure == "error": + raise RuntimeError("engine profiler failed to start") + await asyncio.Event().wait() + + engine_core.profile_async.side_effect = profile_async + + start_task = asyncio.create_task(engine.start_profile()) + await start_rpc_entered.wait() + if failure == "cancel": + start_task.cancel() + expected = RuntimeError if failure == "error" else asyncio.CancelledError + with pytest.raises(expected): + await start_task + + assert engine._profile_session_active is (failure == "cancel") + await engine.stop_profile() + assert engine._profile_session_active is False + + asyncio.run(profile()) + + frontend_profiler.start.assert_not_called() + + +class _BlockingStartProfiler(WorkerProfiler): + def __init__(self) -> None: + self.in_start = threading.Event() + self.release_start = threading.Event() + self.stop_call_count = 0 + super().__init__( + ProfilerConfig(profiler="torch", torch_profiler_dir="/tmp/mock") + ) + + def _start(self) -> None: + self.in_start.set() + if not self.release_start.wait(timeout=5): + raise TimeoutError("start was never released") + + def _stop(self) -> None: + self.stop_call_count += 1 + + +def test_stop_after_cancelled_frontend_start_stops_frontend_profiler( + monkeypatch: pytest.MonkeyPatch, +): + """Cancelling start_profile while the frontend profiler is starting keeps + the session active, and the next stop_profile still stops the backend.""" + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + frontend_profiler = _BlockingStartProfiler() + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + try: + start_task = asyncio.create_task(engine.start_profile()) + assert await asyncio.to_thread(frontend_profiler.in_start.wait, 5) + start_task.cancel() + with pytest.raises(asyncio.CancelledError): + await start_task + assert engine._profile_session_active is True + + stop_task = asyncio.create_task(engine.stop_profile()) + done, _ = await asyncio.wait({stop_task}, timeout=0.1) + assert not done, "stop_profile did not wait for the in-flight start" + frontend_profiler.release_start.set() + await asyncio.wait_for(stop_task, timeout=5) + finally: + frontend_profiler.release_start.set() + + asyncio.run(profile()) + + assert frontend_profiler.stop_call_count == 1 + assert frontend_profiler._active is False + assert frontend_profiler._running is False + assert engine._profile_session_active is False + assert engine_core.profile_async.await_args_list == [ + call(True, None, None, None), + call(False), + ] + + +def test_multi_client_profile_uses_idempotent_engine_requests( + monkeypatch: pytest.MonkeyPatch, tmp_path +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_activities=["CPU"], + ) + engine_core = _mock_async_llm_dependencies(monkeypatch) + profiler = MagicMock() + monkeypatch.setattr(async_llm_module, "create_frontend_profiler", profiler) + engines = [ + AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + client_count=2, + client_index=client_index, + ) + for client_index in range(2) + ] + + assert all(engine.profiler is None for engine in engines) + profiler.assert_not_called() + + async def profile(): + await engines[0].start_profile() + await engines[1].start_profile() + await engines[1].stop_profile() + await engines[0].start_profile() + + asyncio.run(profile()) + + engine_core.profile_async.assert_has_awaits( + [ + call(True, None, None, None), + call(True, None, None, None), + call(False), + call(True, None, None, None), + ] + ) + + +def test_frontend_profiler_ignores_worker_iteration_bounds( + monkeypatch: pytest.MonkeyPatch, tmp_path +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_activities=["CPU", "CUDA"], + delay_iterations=5, + max_iterations=2, + wait_iterations=3, + warmup_iterations=1, + ) + _mock_async_llm_dependencies(monkeypatch) + profiler = MagicMock() + monkeypatch.setattr("vllm.profiler.wrapper.TorchProfilerWrapper", profiler) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + frontend_config = profiler.call_args.args[0] + assert frontend_config.delay_iterations == 0 + assert frontend_config.max_iterations == 0 + assert frontend_config.wait_iterations == 0 + assert frontend_config.warmup_iterations == 0 + assert engine.profiler is profiler.return_value + + async def generate( engine: AsyncLLM, request_id: str, diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index b09f9405d7b5..69241889dc44 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -7,6 +7,7 @@ import pytest from vllm import LLM +from vllm.exceptions import ProfilerAlreadyActiveError from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.engine.llm_engine import LLMEngine from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector @@ -52,6 +53,38 @@ def test_step_invalidates_mm_processor_cache_misses(): ] +def test_duplicate_profile_start_rejected_before_engine_core_dispatch(): + engine = object.__new__(LLMEngine) + engine.engine_core = MagicMock() + engine._profile_session_active = False + + engine.start_profile("first") + + with pytest.raises(ProfilerAlreadyActiveError): + engine.start_profile("duplicate") + + engine.engine_core.profile.assert_called_once_with(True, "first", None, None) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"profile_prefix": "../trace"}, + {"delay_iterations": -1}, + {"max_iterations": -1}, + ], +) +def test_invalid_profile_controls_rejected_before_engine_core_dispatch(kwargs): + engine = object.__new__(LLMEngine) + engine.engine_core = MagicMock() + engine._profile_session_active = False + + with pytest.raises(ValueError): + engine.start_profile(**kwargs) + + engine.engine_core.profile.assert_not_called() + + def _vllm_model( apc: bool, vllm_runner: type[VllmRunner], diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index e140803ef0d3..e272ed5106e0 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os +import threading from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, Mock, call, patch @@ -27,6 +28,7 @@ validate_worker_profiler_config, ) from vllm.v1.core.sched.output import CachedRequestData +from vllm.v1.executor.abstract import Executor from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.xpu_worker import XPUWorker @@ -159,13 +161,33 @@ def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device wrapper.start() wrapper.stop() - summary = tmp_path / "profiler_out_0.txt" + summary = tmp_path / "worker.profiler_out.txt" assert summary.exists() == dump_device_time assert ("device times" in capsys.readouterr().out) == dump_device_time if dump_device_time: assert summary.read_text() == "device times\n" +def test_torch_profiler_summary_follows_output_name(tmp_path): + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + ) + with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile: + table = profile.return_value.key_averages.return_value.table + table.side_effect = ["first run", "second run"] + wrapper = TorchProfilerWrapper( + config, worker_name="initial", local_rank=1, activities=["CUDA"] + ) + for output_name in ["first_rank1", "second_rank1"]: + wrapper.set_output_name(output_name) + wrapper.start() + wrapper.stop() + + assert (tmp_path / "first_rank1.profiler_out.txt").read_text() == "first run\n" + assert (tmp_path / "second_rank1.profiler_out.txt").read_text() == "second run\n" + + @pytest.mark.parametrize( "activities", [[], ["CPU", "CPU"], ["INVALID"]], @@ -184,6 +206,30 @@ def test_torch_profiler_activities_require_torch_profiler(): ProfilerConfig(torch_profiler_activities=["CPU"]) +@pytest.mark.parametrize( + ("profiler", "activities", "ignore_frontend", "expected"), + [ + ("torch", None, False, True), + ("torch", ["CPU"], False, True), + ("torch", ["CPU", "CUDA"], False, True), + ("torch", ["CUDA"], False, False), + ("torch", ["CPU"], True, False), + ("cuda", None, False, False), + ], +) +def test_should_profile_frontend( + profiler, activities, ignore_frontend, expected, tmp_path +): + config = ProfilerConfig( + profiler=profiler, + torch_profiler_dir=str(tmp_path) if profiler == "torch" else "", + torch_profiler_activities=activities, + ignore_frontend=ignore_frontend, + ) + + assert config.should_profile_frontend is expected + + @pytest.mark.parametrize( ("device_type", "activities"), [ @@ -270,16 +316,81 @@ def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"), patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, ): - worker.profile() + worker.profile( + profile_prefix="first", + delay_iterations=5, + max_iterations=2, + ) worker.profile(is_start=False) - worker.profile() + worker.profile(profile_prefix="second") assert worker.profiler is wrapper.return_value wrapper.assert_called_once() - assert wrapper.return_value.start.call_count == 2 + assert wrapper.return_value.set_output_name.call_args_list == [ + call("first_rank0"), + call("second_rank0"), + ] + assert wrapper.return_value.start.call_args_list == [ + call(delay_iterations=5, max_iterations=2), + call(delay_iterations=None, max_iterations=None), + ] wrapper.return_value.stop.assert_called_once_with() +def test_executor_profile_only_forwards_supplied_override_kwargs(): + executor = MagicMock() + executor.vllm_config.parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" + + Executor.profile(executor, True, "default") + Executor.profile( + executor, + True, + "bounded", + delay_iterations=5, + max_iterations=2, + ) + + assert executor.collective_rpc.call_args_list == [ + call("profile", args=(True, "default"), kwargs=None), + call( + "profile", + args=(True, "bounded"), + kwargs={"delay_iterations": 5, "max_iterations": 2}, + ), + ] + + +class _LegacyProfileWorker: + def profile(self, is_start: bool = True, profile_prefix: str | None = None): + pass + + +def test_executor_rejects_overrides_unsupported_by_worker(): + """Out-of-tree workers with the old ``profile`` signature are rejected + before dispatch, so a failed collective_rpc cannot leave replies queued.""" + executor = MagicMock() + executor.vllm_config.parallel_config.worker_cls = f"{__name__}._LegacyProfileWorker" + + with pytest.raises(ValueError, match="max_iterations"): + Executor.profile(executor, True, None, max_iterations=2) + Executor.profile(executor, True, "legacy") + + executor.collective_rpc.assert_called_once_with( + "profile", args=(True, "legacy"), kwargs=None + ) + + +def test_worker_rejects_invalid_profile_prefix(): + worker = object.__new__(Worker) + worker.profiler = None + worker.profiler_config = ProfilerConfig( + profiler="torch", torch_profiler_dir="/tmp/mock" + ) + + with pytest.raises(ValueError, match="profile_prefix"): + worker.profile(profile_prefix="../trace") + + def test_immediate_start_stop(default_profiler_config): """Test standard start without delay.""" profiler = ConcreteWorkerProfiler(default_profiler_config) @@ -294,6 +405,33 @@ def test_immediate_start_stop(default_profiler_config): assert profiler.stop_call_count == 1 +def test_stop_waits_for_inflight_start(default_profiler_config): + """A stop() from another thread during a slow start() must still stop the + backend instead of leaving it running with the session marked inactive.""" + profiler = ConcreteWorkerProfiler(default_profiler_config) + in_start = threading.Event() + release_start = threading.Event() + + def slow_start(): + in_start.set() + release_start.wait() + + with patch.object(profiler, "_start", side_effect=slow_start): + start_thread = threading.Thread(target=profiler.start) + start_thread.start() + assert in_start.wait(timeout=5) + stop_thread = threading.Thread(target=profiler.stop) + stop_thread.start() + stop_thread.join(timeout=0.1) + release_start.set() + start_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert profiler._active is False + assert profiler._running is False + assert profiler.stop_call_count == 1 + + def test_delayed_start(default_profiler_config): """Test that profiler waits for N steps before actually starting.""" default_profiler_config.delay_iterations = 2 @@ -400,19 +538,37 @@ def test_delayed_start_and_max_iters(default_profiler_config): assert profiler.stop_call_count == 1 -def test_idempotency(default_profiler_config): - """Test that calling start/stop multiple times doesn't break logic.""" +def test_duplicate_start_is_ignored(default_profiler_config): profiler = ConcreteWorkerProfiler(default_profiler_config) - # Double Start profiler.start() profiler.start() - assert profiler.start_call_count == 1 # Should only start once + assert profiler.start_call_count == 1 + + profiler.stop() + profiler.stop() + assert profiler.stop_call_count == 1 + + +def test_session_bounds_override_and_reset_to_defaults(): + profiler = ConcreteWorkerProfiler( + ProfilerConfig( + profiler="torch", + torch_profiler_dir="/tmp/mock", + delay_iterations=3, + max_iterations=4, + ) + ) - # Double Stop + profiler.start(delay_iterations=1, max_iterations=2) + assert profiler._delay_iters == 1 + assert profiler._max_iters == 2 profiler.stop() + + profiler.start() + assert profiler._delay_iters == 3 + assert profiler._max_iters == 4 profiler.stop() - assert profiler.stop_call_count == 1 # Should only stop once def test_step_inactive(default_profiler_config): @@ -1095,7 +1251,9 @@ def test_gpu_worker_creates_proton_profiler(): Worker.profile(worker) wrapper.assert_called_once_with(worker.profiler_config, worker_name="rank1") - worker.profiler.start.assert_called_once_with() + worker.profiler.start.assert_called_once_with( + delay_iterations=None, max_iterations=None + ) @_requires_cuda_for_proton @@ -1141,7 +1299,9 @@ def test_gpu_worker_reuses_cuda_graph_proton_session(): Worker.profile(worker, is_start=False) worker.profiler.set_output_name.assert_called_once_with("first_rank1") - worker.profiler.start.assert_called_once_with() + worker.profiler.start.assert_called_once_with( + delay_iterations=None, max_iterations=None + ) worker.profiler.stop.assert_called_once_with() diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 14ff6c1ab132..778a4e50d334 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -122,8 +122,6 @@ def _validate_api_url( if isinstance(expected_suffixes, str): expected_suffixes = {expected_suffixes} - expected_suffixes = {*expected_suffixes, "profile"} - if not api_url.endswith(tuple(expected_suffixes)): raise ValueError(f"{api_name} URL must end with one of: {expected_suffixes}.") @@ -158,6 +156,34 @@ def _get_headers(content_type: str | None = None) -> dict[str, str]: return headers +async def async_request_profile( + api_url: str, + session: aiohttp.ClientSession, + extra_headers: dict | None = None, +) -> RequestFuncOutput: + _validate_api_url( + api_url, + "Profiling API", + {"start_profile", "stop_profile"}, + ) + + headers = _get_headers() + if extra_headers: + headers |= extra_headers + + output = RequestFuncOutput() + try: + async with session.post(url=api_url, headers=headers) as response: + output.success = response.status == 200 + if not output.success: + output.error = response.reason or "" + except Exception: + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + async def async_request_openai_completions( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 19130716f4f0..b551361f55d9 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -49,6 +49,7 @@ POOLING_BACKENDS, RequestFuncInput, RequestFuncOutput, + async_request_profile, ) from vllm.benchmarks.lib.ready_checker import wait_for_endpoint from vllm.benchmarks.lib.utils import ( @@ -936,22 +937,10 @@ async def warmup_limited_request_func(): if profile: print("Starting profiler...") - profile_input = RequestFuncInput( - model=model_id, - model_name=model_name, - prompt=test_prompt, - api_url=base_url + "/start_profile", - prompt_len=test_prompt_len, - output_len=test_output_len, - logprobs=logprobs, - multi_modal_content=test_mm_content, - ignore_eos=ignore_eos, + profile_output = await async_request_profile( + base_url + "/start_profile", + session, extra_headers=extra_headers, - extra_body=test_extra_body, - chat_messages=test_chat_messages, - ) - profile_output = await request_func( - request_func_input=profile_input, session=session ) if profile_output.success: print("Profiler started") @@ -1469,16 +1458,10 @@ def process_one_metric( if profile: print("Stopping profiler...") - profile_input = RequestFuncInput( - model=model_id, - prompt=test_prompt, - api_url=base_url + "/stop_profile", - prompt_len=test_prompt_len, - output_len=test_output_len, - logprobs=logprobs, - ) - profile_output = await request_func( - request_func_input=profile_input, session=session + profile_output = await async_request_profile( + base_url + "/stop_profile", + session, + extra_headers=extra_headers, ) if profile_output.success: print("Profiler stopped") diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 3ade4bceed87..7547546d140d 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -22,6 +22,34 @@ ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"] +def validate_profile_prefix(prefix: str | None) -> str | None: + """Validate a trace filename prefix supplied for one profiling session.""" + if prefix is None: + return None + if ( + not 1 <= len(prefix) <= 128 + or not prefix.isascii() + or not prefix[0].isalnum() + or any(not (char.isalnum() or char in "._-") for char in prefix[1:]) + ): + raise ValueError( + "profile_prefix must be 1-128 ASCII characters, start with a letter " + "or number, and contain only letters, numbers, '.', '_', or '-'" + ) + return prefix + + +def validate_profile_iteration_bounds( + delay_iterations: int | None, + max_iterations: int | None, +) -> None: + """Validate iteration bounds supplied for one profiling session.""" + if delay_iterations is not None and delay_iterations < 0: + raise ValueError("delay_iterations must be greater than or equal to 0") + if max_iterations is not None and max_iterations < 0: + raise ValueError("max_iterations must be greater than or equal to 0") + + def _is_uri_path(path: str) -> bool: """Check if path is a URI (scheme://...), excluding Windows drive letters. @@ -136,6 +164,16 @@ class ProfilerConfig: Defaults to 0, meaning no limit. """ + @property + def should_profile_frontend(self) -> bool: + """Whether the frontend CPU profiler should be enabled.""" + activities = self.torch_profiler_activities + return ( + self.profiler == "torch" + and not self.ignore_frontend + and (activities is None or "CPU" in activities) + ) + warmup_iterations: int = Field(default=0, ge=0) """Number of warmup iterations for PyTorch profiler schedule. During warmup, the profiler runs but data is discarded. This helps reduce @@ -188,13 +226,7 @@ def _validate_profiler_config(self) -> Self: raise ValueError( "torch_profiler_activities must not contain duplicates" ) - records_cpu_activity = activities is None or "CPU" in activities - if ( - self.profiler == "torch" - and has_delay_or_limit - and not self.ignore_frontend - and records_cpu_activity - ): + if has_delay_or_limit and self.should_profile_frontend: logger.warning_once( "Using 'torch' profiler with delay_iterations or max_iterations " "while ignore_frontend is False may result in high overhead." diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index a34981d651fd..2bb128ba073e 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -162,8 +162,14 @@ async def check_health(self) -> None: ... @abstractmethod - async def start_profile(self) -> None: - """Start profiling the engine.""" + async def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + """Start profiling with optional per-session overrides.""" ... @abstractmethod diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 3de7bab1339a..45b6aede583c 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -785,16 +785,30 @@ def enqueue_chat( mm_processor_kwargs=mm_processor_kwargs, ) - def start_profile(self, profile_prefix: str | None = None) -> None: - """Start profiling with optional custom trace prefix. + def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + """Start profiling with optional per-session overrides. Args: profile_prefix: Optional prefix for the trace file names. If provided, trace files will be named as "_dp_pp_tp". If not provided, default naming will be used. + delay_iterations: Optional number of worker iterations to skip before + profiling starts. + max_iterations: Optional maximum number of worker iterations to profile. + Zero means no limit. """ - self.llm_engine.start_profile(profile_prefix) + self.llm_engine.start_profile( + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def stop_profile(self) -> None: self.llm_engine.stop_profile() diff --git a/vllm/entrypoints/serve/profile/api_router.py b/vllm/entrypoints/serve/profile/api_router.py index eeed6b45ef4e..2051da2b1d2a 100644 --- a/vllm/entrypoints/serve/profile/api_router.py +++ b/vllm/entrypoints/serve/profile/api_router.py @@ -4,6 +4,7 @@ from fastapi import APIRouter, FastAPI, Request from fastapi.responses import Response +from pydantic import BaseModel, ConfigDict from vllm.config import ProfilerConfig from vllm.engine.protocol import EngineClient @@ -14,14 +15,32 @@ router = APIRouter() +class StartProfileRequest(BaseModel): + """Optional overrides for one profiling session. + + Values are validated by the engine client. + """ + + model_config = ConfigDict(extra="ignore", strict=True) + + profile_prefix: str | None = None + delay_iterations: int | None = None + max_iterations: int | None = None + + def engine_client(request: Request) -> EngineClient: return request.app.state.engine_client @router.post("/start_profile") -async def start_profile(raw_request: Request): +async def start_profile( + raw_request: Request, profile_request: StartProfileRequest | None = None +): logger.info("Starting profiler...") - await engine_client(raw_request).start_profile() + # Only pass supplied overrides so clients without them keep the + # zero-argument call. + overrides = profile_request.model_dump(exclude_none=True) if profile_request else {} + await engine_client(raw_request).start_profile(**overrides) logger.info("Profiler started.") return Response(status_code=200) diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 9bfd4634bcfa..11ca13fe0bc3 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -181,3 +181,14 @@ def __init__(self): "Please try again later or on a different instance.", HTTPStatus.SERVICE_UNAVAILABLE, ) + + +class ProfilerAlreadyActiveError(GracefulHTTPError): + """Raised when a profiling session is already active.""" + + def __init__(self): + super().__init__( + "A profiling session is already active. Call /stop_profile before " + "starting another session.", + HTTPStatus.CONFLICT, + ) diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index bd9b6c231db7..ebe7d90c3bf3 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -5,6 +5,7 @@ import inspect import json import os +import threading from abc import ABC, abstractmethod from collections.abc import Callable, Iterator, Sequence from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress @@ -17,7 +18,13 @@ import vllm.version from vllm.config import ProfilerConfig -from vllm.config.profiler import ProfilerKind, TorchProfilerActivity, _is_uri_path +from vllm.config.profiler import ( + ProfilerKind, + TorchProfilerActivity, + _is_uri_path, + validate_profile_iteration_bounds, +) +from vllm.config.utils import replace from vllm.logger import init_logger from vllm.platforms import current_platform @@ -28,14 +35,16 @@ class WorkerProfiler(ABC): def __init__(self, profiler_config: ProfilerConfig) -> None: - self._delay_iters = profiler_config.delay_iterations + self._default_delay_iters = profiler_config.delay_iterations + self._delay_iters = self._default_delay_iters if self._delay_iters > 0: logger.info_once( "GPU profiling will start " f"{self._delay_iters} steps after start_profile." ) - self._max_iters = profiler_config.max_iterations + self._default_max_iters = profiler_config.max_iterations + self._max_iters = self._default_max_iters if self._max_iters > 0: logger.info_once( "GPU profiling will stop " @@ -43,6 +52,9 @@ def __init__(self, profiler_config: ProfilerConfig) -> None: "or when stop_profile is received." ) + # Serializes start()/stop(), which AsyncLLM may call from different threads. + self._lifecycle_lock = threading.Lock() + # Track when the profiler gets triggered by start_profile self._active_iteration_count = 0 self._active = False @@ -88,17 +100,32 @@ def _call_stop(self) -> None: logger.warning("Failed to stop profiler: %s", e) self._running = False # Always mark as not running, assume stop worked - def start(self) -> None: + def start( + self, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: """Attempt to start the profiler, accounting for delayed starts.""" - if self._active: - logger.debug( - "start_profile received when profiler is already active. " - "Ignoring request." + with self._lifecycle_lock: + if self._active: + logger.debug( + "start_profile received when profiler is already active. " + "Ignoring request." + ) + return + validate_profile_iteration_bounds(delay_iterations, max_iterations) + self._delay_iters = ( + self._default_delay_iters + if delay_iterations is None + else delay_iterations ) - return - self._active = True - if self._delay_iters == 0: - self._call_start() + self._max_iters = ( + self._default_max_iters if max_iterations is None else max_iterations + ) + self._active = True + if self._delay_iters == 0: + self._call_start() def step(self) -> None: """Update the profiler state at each worker step, @@ -148,17 +175,19 @@ def _profiler_step(self) -> bool: def stop(self) -> None: """Attempt to stop the profiler, accounting for overlapped calls.""" - if not self._active: - logger.debug( - "stop_profile received when profiler is not active. Ignoring request." - ) - return - self._active = False - self._active_iteration_count = 0 - self._profiling_for_iters = 0 + with self._lifecycle_lock: + if not self._active: + logger.debug( + "stop_profile received when profiler is not active. " + "Ignoring request." + ) + return + self._active = False + self._active_iteration_count = 0 + self._profiling_for_iters = 0 - if self._running: - self._call_stop() + if self._running: + self._call_stop() def shutdown(self) -> None: """Ensure profiler is stopped when shutting down.""" @@ -175,6 +204,10 @@ def capture_cuda_graphs(self) -> AbstractContextManager[None]: """Observe graph creation for backends that attribute replay activity.""" return nullcontext() + def set_output_name(self, worker_name: str) -> None: + """Set the output name for the next profiling session.""" + return None + def annotate_context_manager(self, name: str): """Return a context manager to annotate profiler traces.""" return nullcontext() @@ -202,6 +235,10 @@ def __init__( self.local_rank = local_rank self.profiler_config = profiler_config torch_profiler_trace_dir = profiler_config.torch_profiler_dir + self._torch_profiler_trace_dir = torch_profiler_trace_dir + self._torch_profiler_use_gzip = profiler_config.torch_profiler_use_gzip + self._worker_name = worker_name + self._custom_trace_handler = on_trace_ready is not None if local_rank in (None, 0): logger.info_once( "Torch profiling enabled. Traces will be saved to: %s", @@ -276,6 +313,21 @@ def __init__( self._warmup_steps_remaining = self._initial_warmup_steps_remaining self._version_metadata_added = False + @override + def set_output_name(self, worker_name: str) -> None: + if self._active: + return + self._worker_name = worker_name + if self._custom_trace_handler: + return + self._profiler_kwargs["on_trace_ready"] = ( + torch.profiler.tensorboard_trace_handler( + self._torch_profiler_trace_dir, + worker_name=worker_name, + use_gzip=self._torch_profiler_use_gzip, + ) + ) + def _build_profiler_table( self, sort_key: str, @@ -293,13 +345,15 @@ def _build_profiler_table( row_limit=row_limit, ) - def _write_profiler_table(self, rank: int, table: str) -> None: + def _write_profiler_table(self, table: str) -> None: profiler_dir = self.profiler_config.torch_profiler_dir # Skip file write for URI paths (gs://, s3://, etc.) # as standard file I/O doesn't work with URI schemes if not _is_uri_path(profiler_dir): - profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt" + profiler_out_file = os.path.join( + profiler_dir, f"{self._worker_name}.profiler_out.txt" + ) with open(profiler_out_file, "w") as f: print(table, file=f) @@ -344,7 +398,7 @@ def _stop(self) -> None: rank = self.local_rank if self.dump_device_time_total: table = self._build_profiler_table(sort_key="self_device_time_total") - self._write_profiler_table(rank, table) + self._write_profiler_table(table) # only print profiler results on rank 0 if rank == 0: @@ -354,7 +408,7 @@ def _stop(self) -> None: table = self._build_profiler_table( sort_key="self_cpu_time_total", row_limit=50 ) - self._write_profiler_table(rank, table) + self._write_profiler_table(table) # only print profiler results on rank 0 if rank == 0: @@ -494,6 +548,7 @@ def _create_session(self, output_path: str) -> int: def has_cuda_graph_session(self) -> bool: return self._graph_session + @override def set_output_name(self, worker_name: str) -> None: """Set the next run's output name after startup graph capture.""" if self._active: @@ -701,6 +756,29 @@ def create_worker_profiler( return ProtonProfilerWrapper(profiler_config, worker_name=worker_name) +def create_frontend_profiler( + profiler_config: ProfilerConfig, *, worker_name: str +) -> TorchProfilerWrapper: + """Create the CPU-only frontend profiler, which runs for the whole session + rather than following engine iterations.""" + logger.info( + "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", + profiler_config.torch_profiler_dir, + ) + return TorchProfilerWrapper( + replace( + profiler_config, + delay_iterations=0, + max_iterations=0, + wait_iterations=0, + warmup_iterations=0, + ), + worker_name=worker_name, + local_rank=0, + activities=["CPU"], + ) + + def create_graph_capture_profiler( profiler_config: ProfilerConfig, global_rank: int ) -> WorkerProfiler | None: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 37025a863230..7e50aa045123 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -13,6 +13,10 @@ from vllm import TokensPrompt from vllm.config import VllmConfig from vllm.config.kv_events import KVEventsConfig +from vllm.config.profiler import ( + validate_profile_iteration_bounds, + validate_profile_prefix, +) from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, WeightTransferUpdateRequest, @@ -23,6 +27,7 @@ from vllm.exceptions import ( GracefulHTTPError, MaxQueuedTokensError, + ProfilerAlreadyActiveError, QueueOverflowError, VLLMClientError, VLLMValidationError, @@ -33,7 +38,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput from vllm.pooling_params import PoolingParams -from vllm.profiler.wrapper import TorchProfilerWrapper +from vllm.profiler.wrapper import TorchProfilerWrapper, create_frontend_profiler from vllm.renderers import renderer_from_config from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.sampling_params import RequestOutputKind, SamplingParams @@ -211,24 +216,27 @@ def __init__( pass self.profiler = profiler - configured_activities = vllm_config.profiler_config.torch_profiler_activities + self._frontend_profiler_injected = profiler is not None + self._profile_lock = asyncio.Lock() + self._profile_session_guard_enabled = client_count == 1 + self._profile_session_active = False if ( - vllm_config.profiler_config.profiler == "torch" - and not vllm_config.profiler_config.ignore_frontend - and (configured_activities is None or "CPU" in configured_activities) + not self._frontend_profiler_injected + and vllm_config.profiler_config.should_profile_frontend ): - profiler_dir = vllm_config.profiler_config.torch_profiler_dir - logger.info( - "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501 - profiler_dir, - ) - worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm" - self.profiler = TorchProfilerWrapper( - vllm_config.profiler_config, - worker_name=worker_name, - local_rank=0, - activities=["CPU"], - ) + if self._profile_session_guard_enabled: + self._frontend_profiler_worker_name = ( + f"{socket.gethostname()}_{os.getpid()}.async_llm" + ) + self.profiler = create_frontend_profiler( + vllm_config.profiler_config, + worker_name=self._frontend_profiler_worker_name, + ) + else: + logger.warning( + "Frontend CPU profiling is disabled when multiple API server " + "processes share an engine." + ) @classmethod def from_vllm_config( @@ -1081,17 +1089,48 @@ async def check_health(self) -> None: if self.errored: raise self.dead_error - async def start_profile(self, profile_prefix: str | None = None) -> None: - coros = [self.engine_core.profile_async(True, profile_prefix)] - if self.profiler is not None: - coros.append(asyncio.to_thread(self.profiler.start)) - await asyncio.gather(*coros) + async def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + async with self._profile_lock: + if self._profile_session_guard_enabled and self._profile_session_active: + raise ProfilerAlreadyActiveError() + validate_profile_prefix(profile_prefix) + validate_profile_iteration_bounds(delay_iterations, max_iterations) + + if self._profile_session_guard_enabled: + self._profile_session_active = True + # A cancelled start may still take effect, so only a reported + # engine failure clears the session; otherwise /stop_profile is needed. + try: + await self.engine_core.profile_async( + True, + profile_prefix, + delay_iterations, + max_iterations, + ) + except Exception: + self._profile_session_active = False + raise + if self.profiler is not None: + if not self._frontend_profiler_injected: + worker_name = self._frontend_profiler_worker_name + if profile_prefix is not None: + worker_name = f"{profile_prefix}_{worker_name}" + self.profiler.set_output_name(worker_name) + await asyncio.to_thread(self.profiler.start) async def stop_profile(self) -> None: - coros = [self.engine_core.profile_async(False)] - if self.profiler is not None: - coros.append(asyncio.to_thread(self.profiler.stop)) - await asyncio.gather(*coros) + async with self._profile_lock: + coros = [self.engine_core.profile_async(False)] + if self.profiler is not None: + coros.append(asyncio.to_thread(self.profiler.stop)) + await asyncio.gather(*coros) + self._profile_session_active = False async def reset_mm_cache(self) -> None: # Join the background MM warmup first: the mm_processor_cache is not diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index cc0a8bedddf3..e0a003af6254 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -810,8 +810,19 @@ def shutdown(self): cleanup_dist_env_and_memory() logger.debug_once("[shutdown] EngineCore: local resource teardown complete") - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - self.model_executor.profile(is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + self.model_executor.profile( + is_start, + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def reset_mm_cache(self): # NOTE: Since this is mainly for debugging, we don't attempt to diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index cb6705966ac2..88fdc1d1e7f1 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -181,7 +181,13 @@ def get_supported_tasks(self) -> tuple[SupportedTask, ...]: def add_request(self, request: EngineCoreRequest) -> None: raise NotImplementedError - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: raise NotImplementedError def reset_mm_cache(self) -> None: @@ -275,7 +281,11 @@ async def add_request_async(self, request: EngineCoreRequest) -> None: raise NotImplementedError async def profile_async( - self, is_start: bool = True, profile_prefix: str | None = None + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, ) -> None: raise NotImplementedError @@ -392,8 +402,19 @@ def shutdown(self, timeout: float | None = None) -> None: if self._finalizer.detach() is not None: self.engine_core.shutdown() - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: - self.engine_core.profile(is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + self.engine_core.profile( + is_start, + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def reset_mm_cache(self) -> None: self.engine_core.reset_mm_cache() @@ -1019,8 +1040,20 @@ def abort_requests(self, request_ids: list[str]) -> None: if request_ids and not self.resources.engine_dead: self._send_input(EngineCoreRequestType.ABORT, request_ids) - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: - self.call_utility("profile", is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + self.call_utility( + "profile", + is_start, + profile_prefix, + delay_iterations, + max_iterations, + ) def reset_mm_cache(self) -> None: self.call_utility("reset_mm_cache") @@ -1294,9 +1327,19 @@ async def is_scheduler_paused_async(self) -> bool: return await self.call_utility_consensus_async("is_scheduler_paused") async def profile_async( - self, is_start: bool = True, profile_prefix: str | None = None + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, ) -> None: - await self.call_utility_async("profile", is_start, profile_prefix) + await self.call_utility_async( + "profile", + is_start, + profile_prefix, + delay_iterations, + max_iterations, + ) async def reset_mm_cache_async(self) -> None: await self.call_utility_async("reset_mm_cache") diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 509d46703dd1..2dd70fc17917 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -12,9 +12,14 @@ import vllm.envs as envs from vllm.config import ParallelConfig, VllmConfig +from vllm.config.profiler import ( + validate_profile_iteration_bounds, + validate_profile_prefix, +) from vllm.distributed import stateless_destroy_torch_distributed_process_group from vllm.distributed.parallel_state import get_dp_group from vllm.engine.arg_utils import EngineArgs +from vllm.exceptions import ProfilerAlreadyActiveError from vllm.inputs import EngineInput, PromptType from vllm.logger import configure_logging_if_needed, init_logger from vllm.lora.request import LoRARequest @@ -116,6 +121,7 @@ def __init__( log_stats=self.log_stats, renderer=renderer, ) + self._profile_session_active = False self.logger_manager: StatLoggerManager | None = None if self.log_stats: @@ -358,11 +364,33 @@ def step(self) -> list[RequestOutput | PoolingRequestOutput]: return processed_outputs.request_outputs - def start_profile(self, profile_prefix: str | None = None): - self.engine_core.profile(True, profile_prefix) + def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + if self._profile_session_active: + raise ProfilerAlreadyActiveError() + validate_profile_prefix(profile_prefix) + validate_profile_iteration_bounds(delay_iterations, max_iterations) + + self._profile_session_active = True + try: + self.engine_core.profile( + True, + profile_prefix, + delay_iterations, + max_iterations, + ) + except BaseException: + self._profile_session_active = False + raise def stop_profile(self): self.engine_core.profile(False) + self._profile_session_active = False def reset_mm_cache(self): # Join the background MM warmup first: the mm_processor_cache is not diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index f234e943b4fd..0af9508662cb 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect import time from abc import ABC, abstractmethod from collections.abc import Callable @@ -38,6 +39,18 @@ FailureCallback = Callable[[], None] +def _check_worker_profile_kwargs(worker_cls: str, kwargs: dict[str, int]) -> None: + """Raise if the worker's ``profile`` cannot accept the given overrides.""" + params = inspect.signature(resolve_obj_by_qualname(worker_cls).profile).parameters + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return + if unsupported := [name for name in kwargs if name not in params]: + raise ValueError( + f"Worker class {worker_cls} does not support per-session profiling " + f"overrides: {', '.join(unsupported)}" + ) + + class Executor(ABC): """Abstract base class for vLLM executors. @@ -269,8 +282,30 @@ def take_draft_token_ids(self) -> DraftTokenIds | None: output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") return output[0] - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - self.collective_rpc("profile", args=(is_start, profile_prefix)) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + kwargs: dict[str, int] = {} + if delay_iterations is not None: + kwargs["delay_iterations"] = delay_iterations + if max_iterations is not None: + kwargs["max_iterations"] = max_iterations + if kwargs: + # Reject before dispatch: a failing collective_rpc may leave + # unread worker replies queued. + _check_worker_profile_kwargs( + self.vllm_config.parallel_config.worker_cls, kwargs + ) + self.collective_rpc( + "profile", + args=(is_start, profile_prefix), + kwargs=kwargs or None, + ) def save_sharded_state( self, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index f9b58831e449..d4c64c554273 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -19,6 +19,7 @@ import vllm.envs as envs from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config from vllm.config.compilation import CompilationMode +from vllm.config.profiler import validate_profile_prefix from vllm.device_allocator import get_mem_allocator_instance from vllm.distributed import ( ensure_model_parallel_initialized, @@ -1359,7 +1360,14 @@ def execute_model( def take_draft_token_ids(self) -> DraftTokenIds | None: return self.model_runner.take_draft_token_ids() - def profile(self, is_start: bool = True, profile_prefix: str | None = None): + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): # Check if profiling is enabled if self.profiler_config is None or self.profiler_config.profiler is None: raise RuntimeError( @@ -1370,6 +1378,8 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): ) if is_start: + validate_profile_prefix(profile_prefix) + # Generate the trace name by combining prefix with comprehensive rank suffix from vllm.distributed.utils import get_worker_rank_suffix @@ -1381,9 +1391,6 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): else: trace_name = rank_suffix - if self.profiler_config.profiler == "proton" and self.profiler is not None: - self.profiler.set_output_name(trace_name) - # Create the profiler wrapper only on the first start call if self.profiler is None: self.profiler = create_worker_profiler( @@ -1392,7 +1399,11 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None): local_rank=self.local_rank, ) - self.profiler.start() + self.profiler.set_output_name(trace_name) + self.profiler.start( + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) else: if self.profiler is None: logger.warning("Profiler was not started, nothing to stop.") From 5145e35e50ca7fa2bf381eaacdb30e234d774d1c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 6 Oct 2026 13:46:27 +0800 Subject: [PATCH 050/952] [Frontend] Port the MiniMax M3 tool parser to the parser engine (#59743) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- tests/parser/engine/test_minimax_m3.py | 253 +++++++++++++++ tests/parser/engine/test_replay.py | 9 +- tests/parser/engine/trace_builder.py | 62 ++++ .../test_minimax_m3_tool_parser.py | 57 ++-- vllm/parser/engine/registered_adapters.py | 6 + vllm/parser/minimax_m3.py | 293 ++++++++++++++++++ vllm/tool_parsers/minimax_m3_tool_parser.py | 19 +- 7 files changed, 658 insertions(+), 41 deletions(-) create mode 100644 tests/parser/engine/test_minimax_m3.py create mode 100644 vllm/parser/minimax_m3.py diff --git a/tests/parser/engine/test_minimax_m3.py b/tests/parser/engine/test_minimax_m3.py new file mode 100644 index 000000000000..48c65eeadb04 --- /dev/null +++ b/tests/parser/engine/test_minimax_m3.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from tests.parser.engine.streaming_helpers import ( + collect_content, + collect_function_name, + collect_tool_arguments, + simulate_tool_streaming, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionToolsParam, + FunctionDefinition, +) +from vllm.parser.minimax_m3 import ( + INVOKE_END, + MAX_ELEMENT_DEPTH, + NAMESPACE, + TOOL_CALL_END, + TOOL_CALL_START, + MinimaxM3Parser, +) + +NS = NAMESPACE + +TOOLS = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="create_order", + parameters={ + "type": "object", + "properties": { + "user_id": {"type": "integer"}, + "note": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "shipping": { + "type": "object", + "properties": {"zip": {"type": "integer"}}, + }, + "either": { + "anyOf": [ + {"type": "array", "items": {"type": "integer"}}, + {"type": "string"}, + ] + }, + "free": {"type": "object"}, + }, + }, + ), + ) +] + + +@pytest.fixture +def parser(): + return MinimaxM3Parser(make_mock_tokenizer({}), tools=TOOLS) + + +def element(name: str, body: str) -> str: + return f"{NS}<{name}>{body}{NS}" + + +def tool_block(*invokes: tuple[str, str]) -> str: + body = "\n".join( + f'{NS}{params}{INVOKE_END}' for name, params in invokes + ) + return f"{TOOL_CALL_START}\n{body}\n{TOOL_CALL_END}" + + +def parse_arguments(parser, request, params: str, name: str = "create_order"): + result = parser.extract_tool_calls(tool_block((name, params)), request) + assert [tc.function.name for tc in result.tool_calls] == [name] + return json.loads(result.tool_calls[0].function.arguments) + + +class TestArguments: + def test_nested_elements_follow_schema(self, parser, mock_request): + params = ( + element("tags", element("item", "a") + element("item", "b")) + + element("shipping", element("zip", "018956")) + + element("either", element("item", "1") + element("item", "2")) + + element("free", element("t", "x") + element("t", "y")) + ) + assert parse_arguments(parser, mock_request, params) == { + "tags": ["a", "b"], + "shipping": {"zip": 18956}, + "either": [1, 2], + "free": {"t": ["x", "y"]}, + } + + def test_empty_elements_follow_schema(self, parser, mock_request): + # The chat template renders [] and {} as empty elements. + params = ( + element("tags", "") + + element("shipping", "") + + element("note", "") + + element("either", "") + ) + assert parse_arguments(parser, mock_request, params) == { + "tags": [], + "shipping": {}, + "note": "", + "either": "", + } + + def test_unknown_tool_keeps_text_and_nests_objects(self, parser, mock_request): + params = element("n", "1") + element( + "list", element("item", "a") + element("item", "b") + ) + assert parse_arguments(parser, mock_request, params, name="other") == { + "n": "1", + "list": {"item": ["a", "b"]}, + } + + def test_mixed_text_uses_reserved_field(self, parser, mock_request): + params = element("free", "hi " + element("k", "v") + " there") + element( + "shipping", "x" + element("$text", "child") + ) + assert parse_arguments(parser, mock_request, params) == { + "free": {"k": "v", "$text": "hi there"}, + "shipping": {"$text": "child", "$$text": "x"}, + } + + def test_leaf_text_is_verbatim_for_strings(self, parser, mock_request): + params = element("note", "\n a & c\n") + element("user_id", " 7 ") + assert parse_arguments(parser, mock_request, params) == { + "note": "\n a & c\n", + "user_id": 7, + } + + def test_block_marker_named_element_is_argument(self, parser, mock_request): + params = element("free", element("tool_call", "v")) + assert parse_arguments(parser, mock_request, params) == { + "free": {"tool_call": "v"}, + } + + @pytest.mark.parametrize( + "tail", + [ + "junk" + element("note", "dropped"), + f"{NS}a{NS}", + f"{NS}unterminated", + ], + ) + def test_malformed_element_keeps_preceding_params(self, parser, mock_request, tail): + params = element("user_id", "1") + tail + assert parse_arguments(parser, mock_request, params) == {"user_id": 1} + + def test_depth_limit(self, parser, mock_request): + def nested(depth: int) -> str: + body = "v" + for _ in range(depth - 1): + body = element("n", body) + return element("free", body) + + accepted = parse_arguments(parser, mock_request, nested(MAX_ELEMENT_DEPTH)) + assert "free" in accepted + rejected = parse_arguments(parser, mock_request, nested(MAX_ELEMENT_DEPTH + 1)) + assert rejected == {} + + +class TestToolBlock: + def test_multiple_invokes_and_prefix_content(self, parser, mock_request): + result = parser.extract_tool_calls( + "Checking.\n" + + tool_block( + ("create_order", element("user_id", "1")), + ("create_order", element("user_id", "2")), + ), + mock_request, + ) + assert result.content == "Checking.\n" + assert [json.loads(tc.function.arguments) for tc in result.tool_calls] == [ + {"user_id": 1}, + {"user_id": 2}, + ] + + def test_text_after_tool_block_is_dropped(self, parser, mock_request): + result = parser.extract_tool_calls( + "pre" + tool_block(("create_order", "")) + "post", mock_request + ) + assert result.content == "pre" + assert len(result.tool_calls) == 1 + + def test_invoke_outside_tool_block_is_content(self, parser, mock_request): + text = f'{NS}{INVOKE_END}' + result = parser.extract_tool_calls(text, mock_request) + assert not result.tools_called + assert result.content == text + + +class TestStreaming: + def test_arguments_arrive_with_invoke_end(self, parser, mock_request): + block = tool_block( + ("create_order", element("user_id", "42") + element("tags", "")) + ) + invoke_end = block.index(INVOKE_END) + results = simulate_tool_streaming( + parser, + mock_request, + ["Sure. ", block[:invoke_end], block[invoke_end:]], + ) + + before_end, at_end = results[1][0], results[2][0] + assert before_end is not None and before_end.tool_calls + assert before_end.tool_calls[0].function.name == "create_order" + assert not before_end.tool_calls[0].function.arguments + assert collect_content(results) == "Sure. " + assert json.loads(collect_tool_arguments([(at_end, "")])) == { + "user_id": 42, + "tags": [], + } + + def test_char_chunks_match_complete_parse(self, parser, mock_request): + text = "a" + tool_block( + ("create_order", element("shipping", element("zip", "1"))), + ("create_order", element("note", "x")), + ) + results = simulate_tool_streaming(parser, mock_request, list(text)) + + assert collect_content(results) == "a" + assert collect_function_name(results) == "create_order" + streamed = [ + tc.function.arguments + for delta, _ in results + if delta and delta.tool_calls + for tc in delta.tool_calls + if tc.function and tc.function.arguments + ] + assert [json.loads(args) for args in streamed] == [ + {"shipping": {"zip": 1}}, + {"note": "x"}, + ] + + def test_truncated_invoke_keeps_completed_params(self, parser, mock_request): + block = tool_block( + ("create_order", element("user_id", "1") + element("note", "x")) + ) + truncated = block[: block.index(f"{NS}")] + results = simulate_tool_streaming(parser, mock_request, [truncated]) + results.append((parser.finish_streaming(), truncated)) + + assert collect_function_name(results) == "create_order" + assert json.loads(collect_tool_arguments(results)) == {"user_id": 1} + + +def test_adjust_request_keeps_special_tokens_skipped(parser, mock_request): + mock_request.skip_special_tokens = True + assert parser.adjust_request(mock_request).skip_special_tokens is True diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py index 89791ba29bfe..e6339f895998 100644 --- a/tests/parser/engine/test_replay.py +++ b/tests/parser/engine/test_replay.py @@ -74,10 +74,11 @@ def _discover_parsers() -> list[_ParserInfo]: if cfg.name not in _BUILDERS: missing_builders.append(f"{obj.__name__} (config.name={cfg.name!r})") continue - if cfg.name in ("inkling", "plamo3"): - # Inkling uses typed structural blocks. PLaMo markers can span - # several token IDs. Neither fits this TOOL_END-token harness; - # their direct and delegating coverage lives in dedicated tests. + if cfg.name in ("inkling", "plamo3", "minimax_m3"): + # Inkling uses typed structural blocks. PLaMo and MiniMax M3 + # markers can span several token IDs. None fits this TOOL_END-token + # harness; their direct and delegating coverage lives in dedicated + # tests. continue if cfg.name == "granite": # Granite has a JSON-array tool body with no TOOL_END terminal, so diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 3fa6e05078d4..284ff7048daf 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -42,6 +42,7 @@ KimiK2Parser, MiMoParser, MinimaxM2Parser, + MinimaxM3Parser, NemotronV3Parser, Plamo3Parser, Qwen3Parser, @@ -539,6 +540,66 @@ def _build_minimax_m2(scenario: Scenario, validate: bool = True) -> Sample: return sample +# ── MiniMax M3 (namespaced XML elements, no reasoning) ─────────────── + +_MINIMAX_M3_NS = "]<]minimax[>[" + + +def _minimax_m3_value(value: Any) -> str: + """Render a value the way the M3 chat template's ``to_xml`` does.""" + ns = _MINIMAX_M3_NS + if isinstance(value, dict): + return "".join( + f"{ns}<{key}>{_minimax_m3_value(item)}{ns}" + for key, item in value.items() + ) + if isinstance(value, list): + return "".join( + f"{ns}{_minimax_m3_value(item)}{ns}" for item in value + ) + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + +def _minimax_m3_segments(scenario: Scenario) -> list[tuple[str, bool]]: + ns = _MINIMAX_M3_NS + segs: list[tuple[str, bool]] = [] + if scenario.content: + # The tool parser receives reasoning-stripped content. + segs.append((scenario.content, False)) + if scenario.tool_calls: + invokes = "".join( + f'{ns}' + f"{_minimax_m3_value(tc.arguments)}{ns}\n" + for tc in scenario.tool_calls + ) + segs.append((f"{ns}\n{invokes}{ns}", False)) + return segs + + +def _minimax_m3_expected_content(scenario: Scenario) -> str | None: + if scenario.tool_calls and not (scenario.content or "").strip(): + return None + return scenario.content + + +def _build_minimax_m3(scenario: Scenario, validate: bool = True) -> Sample: + sample = _make_sample( + sample_id=f"minimax_m3-{scenario.id}", + description=scenario.description, + vocab={}, + segments=_minimax_m3_segments(scenario), + expected_reasoning=None, + expected_content=_minimax_m3_expected_content(scenario), + expected_tool_calls=_expected_tc(scenario), + tools=_expected_tools(scenario), + ) + if validate: + _validate_sample(sample, MinimaxM3Parser) + return sample + + # ── Gemma4 (channel reasoning, custom arg format) ──────────────────── _GEMMA4_VOCAB: dict[str, int] = { @@ -1275,6 +1336,7 @@ def _build_granite(scenario: Scenario, validate: bool = True) -> Sample: "gemma4": _build_gemma4, "granite": _build_granite, "minimax_m2": _build_minimax_m2, + "minimax_m3": _build_minimax_m3, "nemotron_v3": _build_nemotron_v3, "granite_thinking_parser": _build_granite_thinking, "seed_oss": _build_seed_oss, diff --git a/tests/tool_parsers/test_minimax_m3_tool_parser.py b/tests/tool_parsers/test_minimax_m3_tool_parser.py index 575892217a75..c01ab89d1e2c 100644 --- a/tests/tool_parsers/test_minimax_m3_tool_parser.py +++ b/tests/tool_parsers/test_minimax_m3_tool_parser.py @@ -8,16 +8,13 @@ from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, ChatCompletionToolsParam, FunctionDefinition, ) from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.minimax_m3_tool_parser import MinimaxM3ToolParser -# MinimaxM3ToolParser extends RustToolParser; skip when the PyO3 extension -# is absent (mirrors the guard in test_rust_tool_parser.py). -pytest.importorskip("vllm._rust_tool_parser") - pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] NS = "]<]minimax[>[" @@ -108,6 +105,23 @@ def build_order_call() -> str: ) +ORDER_ARGUMENTS = { + "user_id": 42, + "urgent": True, + "note": "Please leave at front desk.", + "shipping": {"city": "Singapore", "zip": 18956}, + "items": [ + {"sku": "book-001", "qty": 2}, + {"sku": "pen-007", "qty": 5}, + ], + "metadata": { + "source": "mobile", + "campaign": "may-launch", + }, + "duplicate_demo": {"tag": ["a", "b"]}, +} + + def build_order_invocation(user_id: int) -> str: return ( f'{NS}' @@ -181,6 +195,20 @@ def test_minimax_m3_parser_registered(): assert ToolParserManager.get_tool_parser("minimax_m3") is MinimaxM3ToolParser +def test_required_tool_choice_keeps_native_syntax(parser): + request = ChatCompletionRequest( + model="minimax-m3", + messages=[{"role": "user", "content": "order"}], + tools=sample_tools(), + tool_choice="required", + ) + + adjusted = parser.adjust_request(request) + + assert adjusted.structured_outputs is None + assert adjusted.skip_special_tokens + + def test_non_streaming_nested_tool_call(parser): result = parser.extract_tool_calls( "I will create it.\n" + build_order_call(), @@ -192,21 +220,7 @@ def test_non_streaming_nested_tool_call(parser): assert len(result.tool_calls) == 1 tool_call = result.tool_calls[0] assert tool_call.function.name == "create_order" - assert json.loads(tool_call.function.arguments) == { - "user_id": 42, - "urgent": True, - "note": "Please leave at front desk.", - "shipping": {"city": "Singapore", "zip": 18956}, - "items": [ - {"sku": "book-001", "qty": 2}, - {"sku": "pen-007", "qty": 5}, - ], - "metadata": { - "source": "mobile", - "campaign": "may-launch", - }, - "duplicate_demo": {"tag": ["a", "b"]}, - } + assert json.loads(tool_call.function.arguments) == ORDER_ARGUMENTS def test_non_streaming_without_tool_call_keeps_content(parser): @@ -258,8 +272,5 @@ def test_streaming_nested_tool_call(parser): assert len(tool_calls) == 1 assert tool_calls[0]["name"] == "create_order" assert tool_calls[0]["id"] is not None - assert json.loads(tool_calls[0]["arguments"]) == json.loads( - parser.streamed_args_for_tool[0] - ) - assert json.loads(parser.prev_tool_call_arr[0]["arguments"])["items"][1]["qty"] == 5 + assert json.loads(tool_calls[0]["arguments"]) == ORDER_ARGUMENTS assert results[-1].content is None diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index 31295869c5a0..6d610ce3b5b4 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -19,6 +19,7 @@ from vllm.parser.kimi_k2 import KimiK2Parser from vllm.parser.mimo import MiMoParser from vllm.parser.minimax_m2 import MinimaxM2Parser +from vllm.parser.minimax_m3 import MinimaxM3Parser from vllm.parser.mistral import MistralParser from vllm.parser.nemotron_v3 import NemotronV3Parser from vllm.parser.plamo3 import Plamo3Parser @@ -46,6 +47,11 @@ MinimaxM2ParserToolAdapter, ) = make_adapters(MinimaxM2Parser) +( + MinimaxM3ParserReasoningAdapter, + MinimaxM3ParserToolAdapter, +) = make_adapters(MinimaxM3Parser) + ( Gemma4ParserReasoningAdapter, Gemma4ParserToolAdapter, diff --git a/vllm/parser/minimax_m3.py b/vllm/parser/minimax_m3.py new file mode 100644 index 000000000000..dafe1457e982 --- /dev/null +++ b/vllm/parser/minimax_m3.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax M3 parser for namespace-delimited XML-style tool calls. + +MiniMax M3 prefixes every structural tag with the ``]<]minimax[>[`` marker:: + + ]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[42]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[]<]minimax[>[book-001]<]minimax[>[]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[ + +Each ```` becomes one tool call. Parameters are elements named after +the argument and may nest: an element holding child elements becomes an +object, or an array when its schema says so (the chat template renders list +items as ````). Arguments are emitted once the whole ```` is +parsed. Reasoning is handled by ``MiniMaxM3ReasoningParser``. +""" + +from __future__ import annotations + +import functools +import json +from typing import TYPE_CHECKING, Any + +from vllm.parser.engine.events import EventType +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) +from vllm.tool_parsers.utils import ( + coerce_to_schema_type, + extract_types_from_schema, + find_tool_properties, +) + +if TYPE_CHECKING: + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + +NAMESPACE = "]<]minimax[>[" +TOOL_CALL_START = f"{NAMESPACE}" +TOOL_CALL_END = f"{NAMESPACE}" +INVOKE_PREFIX_DQ = f'{NAMESPACE}' +NAME_END_SQ = "'>" +NAME_END_UNQUOTED = ">" +ELEMENT_START = f"{NAMESPACE}<" +ELEMENT_END_START = f"{NAMESPACE} tuple[str, _Element, int] | None: + """Parse the element starting at ``pos``; ``None`` if malformed.""" + if depth > MAX_ELEMENT_DEPTH: + return None + name_start = pos + len(ELEMENT_START) + name_end = text.find(">", name_start) + if name_end <= name_start: + return None + name = text[name_start:name_end] + if name.startswith("/") or not name.strip(): + return None + + close_tag = f"{ELEMENT_END_START}{name}>" + pos = name_end + 1 + text_parts: list[str] = [] + children: list[tuple[str, _Element]] = [] + while True: + marker = text.find(NAMESPACE, pos) + if marker < 0: + return None + text_parts.append(text[pos:marker]) + if text.startswith(close_tag, marker): + pos = marker + len(close_tag) + break + if not text.startswith(ELEMENT_START, marker): + return None + child = _parse_element(text, marker, depth + 1) + if child is None: + return None + child_name, child_value, pos = child + children.append((child_name, child_value)) + + body_text = "".join(text_parts) + if not children: + return name, body_text, pos + if body_text.strip(): + # Keep mixed text under a reserved field that no child uses. + field = MIXED_TEXT_FIELD + while any(child_name == field for child_name, _ in children): + field = "$" + field + children.append((field, body_text)) + return name, children, pos + + +def _parse_invoke_params(raw_args: str) -> list[tuple[str, _Element]]: + """Parse the parameter elements of one ```` body. + + Parsing stops at the first malformed or unterminated element, keeping the + parameters before it. + """ + params: list[tuple[str, _Element]] = [] + pos = 0 + while True: + while pos < len(raw_args) and raw_args[pos].isspace(): + pos += 1 + if not raw_args.startswith(ELEMENT_START, pos): + return params + element = _parse_element(raw_args, pos, depth=1) + if element is None: + return params + name, value, pos = element + params.append((name, value)) + + +def _container_schema(schema: dict[str, Any]) -> tuple[str, dict[str, Any]] | None: + """Return the first ``array`` or ``object`` alternative of ``schema``.""" + types = schema.get("type") + for kind in ("array", "object"): + if types == kind or (isinstance(types, list) and kind in types): + return kind, schema + if types is None: + if "items" in schema: + return "array", schema + if "properties" in schema or "additionalProperties" in schema: + return "object", schema + for choice_field in ("anyOf", "oneOf"): + for choice in schema.get(choice_field) or (): + if isinstance(choice, dict) and (found := _container_schema(choice)): + return found + return None + + +def _convert_element(value: _Element, schema: Any) -> Any: + if not isinstance(schema, dict): + schema = None + container = _container_schema(schema) if schema is not None else None + + if isinstance(value, str): + if schema is None: + return value + types = extract_types_from_schema(schema) + # The template renders an empty list or dict as an empty element. + if container is not None and "string" not in types and not value.strip(): + return [] if container[0] == "array" else {} + return coerce_to_schema_type(value, types) + + if container is not None and container[0] == "array": + items = container[1].get("items") + return [_convert_element(child, items) for _, child in value] + return _convert_object(value, container[1] if container is not None else None) + + +def _convert_object( + elements: list[tuple[str, _Element]], schema: dict[str, Any] | None +) -> dict[str, Any]: + """Convert named elements to an object; repeated names collect into a list.""" + properties = (schema or {}).get("properties") or {} + additional = (schema or {}).get("additionalProperties") + result: dict[str, Any] = {} + repeated: set[str] = set() + for name, element in elements: + converted = _convert_element(element, properties.get(name, additional)) + if name not in result: + result[name] = converted + elif name in repeated: + result[name].append(converted) + else: + result[name] = [result[name], converted] + repeated.add(name) + return result + + +def _minimax_m3_arg_converter( + raw_args: str, + partial: bool, + properties: dict[str, Any] | None = None, +) -> str: + params = _parse_invoke_params(raw_args) + arguments = _convert_object(params, {"properties": properties or {}}) + return json.dumps(arguments, ensure_ascii=False) + + +@functools.cache +def minimax_m3_config() -> ParserEngineConfig: + return ParserEngineConfig( + name="minimax_m3", + initial_state=ParserState.CONTENT, + terminals={ + "TOOL_START": TOOL_CALL_START, + "TOOL_END": TOOL_CALL_END, + "INVOKE_PREFIX_DQ": INVOKE_PREFIX_DQ, + "INVOKE_PREFIX_SQ": INVOKE_PREFIX_SQ, + "INVOKE_PREFIX_UNQUOTED": INVOKE_PREFIX_UNQUOTED, + "INVOKE_END": INVOKE_END, + "NAME_END_DQ": NAME_END_DQ, + "NAME_END_SQ": NAME_END_SQ, + "NAME_END_UNQUOTED": NAME_END_UNQUOTED, + }, + transitions={ + (ParserState.CONTENT, "TOOL_START"): Transition( + ParserState.TOOL_PREAMBLE, + (), + ), + **{ + (state, terminal): Transition( + ParserState.TOOL_NAME, + (EventType.TOOL_CALL_START,), + ) + for state in (ParserState.TOOL_PREAMBLE, ParserState.TOOL_BETWEEN) + for terminal in ( + "INVOKE_PREFIX_DQ", + "INVOKE_PREFIX_SQ", + "INVOKE_PREFIX_UNQUOTED", + ) + }, + **{ + (ParserState.TOOL_NAME, terminal): Transition( + ParserState.TOOL_ARGS, + (), + ) + for terminal in ("NAME_END_DQ", "NAME_END_SQ", "NAME_END_UNQUOTED") + }, + # Nested parameter elements are argument text; only the invoke + # closer ends the call. + (ParserState.TOOL_ARGS, "INVOKE_END"): Transition( + ParserState.TOOL_BETWEEN, + (EventType.TOOL_CALL_END,), + ), + # The tool block ends the message: stay in TOOL_BETWEEN so any + # text after it is dropped. + (ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition( + ParserState.TOOL_BETWEEN, + (), + ), + (ParserState.TOOL_BETWEEN, "TOOL_END"): Transition( + ParserState.TOOL_BETWEEN, + (), + ), + }, + arg_converter=_minimax_m3_arg_converter, + # Emit arguments once per ``: converted values are not + # prefix-stable while elements arrive, e.g. a repeated parameter name + # turns an already streamed scalar into a list. + stream_arg_deltas=False, + tool_args_json=False, + strip_content_whitespace_with_tools=False, + ) + + +class MinimaxM3Parser(ParserEngine): + """MiniMax M3 tool-call parser backed by the declarative parser engine.""" + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + kwargs.setdefault("parser_engine_config", minimax_m3_config()) + super().__init__(tokenizer, tools, **kwargs) + self._arg_converter = self._convert_args + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + # The M3 markers are non-special added tokens, so they decode with + # special tokens skipped. + return request + + def _convert_args(self, raw_args: str, partial: bool) -> str: + func_name = next( + (slot.name for slot in self._tool_slots if slot.args == raw_args), "" + ) + properties = find_tool_properties(self._tools, func_name.strip()) + return _minimax_m3_arg_converter(raw_args, partial, properties) diff --git a/vllm/tool_parsers/minimax_m3_tool_parser.py b/vllm/tool_parsers/minimax_m3_tool_parser.py index a8628448c448..6be06872738c 100644 --- a/vllm/tool_parsers/minimax_m3_tool_parser.py +++ b/vllm/tool_parsers/minimax_m3_tool_parser.py @@ -1,19 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.tool_parsers.rust_tool_parser import RustToolParser +from vllm.parser.engine.registered_adapters import MinimaxM3ParserToolAdapter -class MinimaxM3ToolParser(RustToolParser): - """Adapter from the Rust MiniMax M3 parser to vLLM ToolParser. - - The real M3 grammar lives in the Rust tool-parser crate. This class only - configures the generic Rust bridge with the MiniMax M3 parser name. - - M3 is not M2 with renamed tags: it prefixes each structural tag with the - MiniMax namespace marker, allows multiple ```` tags in one wrapper, - and represents nested arguments with parameter-name XML tags. - """ - - rust_parser_name = "MinimaxM3ToolParser" - tool_call_start_token = "]<]minimax[>[" +class MinimaxM3ToolParser(MinimaxM3ParserToolAdapter): # type: ignore[valid-type, misc] + # Required and named tool choices parse the native M3 syntax instead of + # forcing JSON output. + supports_required_and_named = False From b3082d304469437cc38a20efbdb20c3559c8a902 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 6 Oct 2026 13:46:27 +0800 Subject: [PATCH 051/952] [Rust Frontend] Remove the PyO3 tool-parser bridge (#59744) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../scripts/run-rust-frontend-cargo-ci.sh | 32 +- rust/Cargo.lock | 85 ---- rust/Cargo.toml | 3 - rust/src/parser/python/Cargo.toml | 19 - rust/src/parser/python/src/lib.rs | 398 ------------------ setup.py | 2 + tests/tool_parsers/test_rust_tool_parser.py | 328 --------------- tools/build_rust.py | 17 +- vllm/tool_parsers/rust_tool_parser.py | 322 -------------- 9 files changed, 13 insertions(+), 1193 deletions(-) delete mode 100644 rust/src/parser/python/Cargo.toml delete mode 100644 rust/src/parser/python/src/lib.rs delete mode 100644 tests/tool_parsers/test_rust_tool_parser.py delete mode 100644 vllm/tool_parsers/rust_tool_parser.py diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index 62a70fd66450..fa1cddb95f57 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -23,7 +23,7 @@ export PATH="$CARGO_HOME/bin:$PATH" CARGO_BINSTALL_VERSION="${CARGO_BINSTALL_VERSION:-1.20.1}" UV_VERSION="${UV_VERSION:-0.11.28}" -PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}" +UV_SCRIPT_PYTHON_VERSION="${UV_SCRIPT_PYTHON_VERSION:-3.12}" CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}" CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}" @@ -83,34 +83,7 @@ install_uv() { "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/uv-installer.sh" \ | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh uv --version -} - -setup_pyo3_python() { - log_section "Installing Python ${PYO3_PYTHON_VERSION} for PyO3 tests" - uv python install "$PYO3_PYTHON_VERSION" - PYO3_PYTHON="$(uv python find \ - --managed-python \ - --no-project \ - --resolve-links \ - "$PYO3_PYTHON_VERSION")" - export PYO3_PYTHON - - local python_libdir - python_libdir="$("$PYO3_PYTHON" - <<'PY' -import pathlib -import sysconfig - -libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR")) -ldlibrary = sysconfig.get_config_var("LDLIBRARY") -assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1 -assert ldlibrary -assert (libdir / ldlibrary).exists(), libdir / ldlibrary -print(libdir) -PY -)" - - export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}" - export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}" + uv python install "$UV_SCRIPT_PYTHON_VERSION" } run_style_clippy() { @@ -144,7 +117,6 @@ run_style_clippy() { run_tests() { install_uv - setup_pyo3_python install_cargo_binstall install_cargo_nextest diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 2ba9832b6bcc..9dcdb0bdbc2c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3403,74 +3403,6 @@ version = "0.1.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f" -[[package]] -name = "pyo3" -version = "0.29.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" -dependencies = [ - "libc", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", -] - -[[package]] -name = "pyo3-build-config" -version = "0.29.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9" -dependencies = [ - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.29.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-macros" -version = "0.29.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.29.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "pythonize" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89" -dependencies = [ - "pyo3", - "serde", - "serde_json", -] - [[package]] name = "quick-error" version = "2.0.1" @@ -4792,12 +4724,6 @@ dependencies = [ "libc", ] -[[package]] -name = "target-lexicon" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" - [[package]] name = "task-local" version = "0.1.1" @@ -6113,17 +6039,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "vllm-tool-parser-py" -version = "0.1.0" -dependencies = [ - "pyo3", - "pythonize", - "serde_json", - "thiserror-ext", - "vllm-parser", -] - [[package]] name = "vllm-tracing" version = "0.1.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index e9935d9e2832..38436467f862 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -11,7 +11,6 @@ members = [ "src/metrics", "src/mock-engine", "src/parser", - "src/parser/python", "src/server", "src/text", "src/tokenizer", @@ -79,8 +78,6 @@ prometheus-client-derive-encode = "0.5.0" prost = "0.14.3" prost-types = "0.14.3" protox = "0.9.1" -pyo3 = "0.29.0" -pythonize = "0.29.0" rand = "0.9.3" rand_distr = "0.5.1" rayon = "1.11.0" diff --git a/rust/src/parser/python/Cargo.toml b/rust/src/parser/python/Cargo.toml deleted file mode 100644 index aadae5638f90..000000000000 --- a/rust/src/parser/python/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "vllm-tool-parser-py" -version.workspace = true -edition.workspace = true -license.workspace = true - -[lib] -name = "_rust_tool_parser" -crate-type = ["cdylib", "rlib"] - -[dependencies] -pyo3.workspace = true -pythonize = { workspace = true, features = ["serde_json"] } -serde_json.workspace = true -thiserror-ext.workspace = true -vllm-parser.workspace = true - -[lints] -workspace = true diff --git a/rust/src/parser/python/src/lib.rs b/rust/src/parser/python/src/lib.rs deleted file mode 100644 index 38b5f1d461a7..000000000000 --- a/rust/src/parser/python/src/lib.rs +++ /dev/null @@ -1,398 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -//! Thin PyO3 bindings for `vllm_parser::tool`. -//! -//! This crate exposes the Rust tool parser trait and data shapes to Python -//! while keeping parser state, grammar, and schema-aware argument conversion in -//! Rust. Python callers should use this module as a typed bridge and keep any -//! vLLM protocol adaptation outside the binding. - -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use pyo3::types::{PyAny, PyModule}; -use pythonize::{depythonize, pythonize}; -use serde_json::Value; -use thiserror_ext::AsReport as _; -use vllm_parser::tool::{Tool, ToolCallDelta, ToolParser, ToolParserOutput}; - -macro_rules! tool_parser_factory { - ($($parser:ident),+ $(,)?) => { - fn create_tool_parser( - name: &str, - tools: &[Tool], - ) -> PyResult> { - match name { - $( - stringify!($parser) => { - ::create(tools) - } - )+ - _ => { - return Err(PyValueError::new_err(format!( - "unsupported tool parser `{name}`" - ))); - } - } - .map_err(|error| PyValueError::new_err(error.to_report_string())) - } - }; -} - -// Export a tool parser to Python by registering it here. -tool_parser_factory! { - MinimaxM3ToolParser, - - // Below are the parsers just for testing purposes on Python side. - DeepSeekV4ToolParser, - KimiK2ToolParser, -} - -#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)] -#[derive(Clone)] -struct PyTool(Tool); - -#[pymethods] -impl PyTool { - #[new] - #[pyo3(signature = (name, description, parameters, strict=None))] - fn new( - name: String, - description: Option, - parameters: &Bound<'_, PyAny>, - strict: Option, - ) -> PyResult { - let parameters = depythonize::(parameters).map_err(|error| { - PyValueError::new_err(format!( - "failed to convert tool parameters from Python to JSON: {error}" - )) - })?; - Ok(Self(Tool { - name, - description, - parameters, - strict, - })) - } - - #[getter] - fn name(&self) -> &str { - &self.0.name - } - - #[getter] - fn description(&self) -> Option<&str> { - self.0.description.as_deref() - } - - #[getter] - fn parameters(&self, py: Python<'_>) -> PyResult> { - pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| { - PyValueError::new_err(format!( - "failed to convert tool parameters from JSON to Python: {error}" - )) - }) - } - - #[getter] - fn strict(&self) -> Option { - self.0.strict - } -} - -#[pyclass( - name = "ToolCallDelta", - module = "vllm._rust_tool_parser", - skip_from_py_object -)] -#[derive(Clone)] -struct PyToolCallDelta(ToolCallDelta); - -#[pymethods] -impl PyToolCallDelta { - #[new] - #[pyo3(signature = (tool_index, name, arguments))] - fn new(tool_index: usize, name: Option, arguments: String) -> Self { - Self(ToolCallDelta { - tool_index, - name, - arguments, - }) - } - - #[getter] - fn tool_index(&self) -> usize { - self.0.tool_index - } - - #[getter] - fn name(&self) -> Option<&str> { - self.0.name.as_deref() - } - - #[getter] - fn arguments(&self) -> &str { - &self.0.arguments - } -} - -#[pyclass( - name = "ToolParserOutput", - module = "vllm._rust_tool_parser", - skip_from_py_object -)] -#[derive(Clone)] -struct PyToolParserOutput(ToolParserOutput); - -#[pymethods] -impl PyToolParserOutput { - #[new] - #[pyo3(signature = (normal_text="", calls=None))] - fn new(py: Python<'_>, normal_text: &str, calls: Option>>) -> Self { - let mut output = ToolParserOutput::default(); - output.push_text(normal_text); - for call in calls.unwrap_or_default() { - output.push_call(call.borrow(py).0.clone()); - } - Self(output) - } - - #[getter] - fn normal_text(&self) -> String { - self.0.normal_text() - } - - #[getter] - fn calls(&self) -> Vec { - self.0.calls().into_iter().cloned().map(PyToolCallDelta).collect() - } - - fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) { - self.0.append(other.0.clone()); - } - - fn coalesce(&self) -> Self { - Self(self.0.clone().coalesce()) - } -} - -#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)] -struct PyToolParser(Box); - -impl PyToolParser { - fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> { - self.0 - .parse_into(chunk, &mut output.0) - .map_err(|error| PyValueError::new_err(error.to_report_string())) - } -} - -#[pymethods] -impl PyToolParser { - #[new] - fn new(py: Python<'_>, parser_name: &str, tools: Vec>) -> PyResult { - let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::>(); - create_tool_parser(parser_name, &tools).map(Self) - } - - fn parse_into( - &mut self, - chunk: &str, - mut output: PyRefMut<'_, PyToolParserOutput>, - ) -> PyResult<()> { - self.parse_into_output(chunk, &mut output) - } - - fn finish(&mut self) -> PyResult { - self.0 - .finish() - .map(PyToolParserOutput) - .map_err(|error| PyValueError::new_err(error.to_report_string())) - } - - fn reset(&mut self) -> String { - self.0.reset() - } - - fn preserve_special_tokens(&self) -> bool { - self.0.preserve_special_tokens() - } - - fn tool_call_id(&self, tool_index: usize) -> Option<&str> { - self.0.tool_call_id(tool_index) - } -} - -#[pymodule] -fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - Ok(()) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - fn with_python(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R { - Python::initialize(); - Python::attach(f) - } - - fn tool_schema() -> Value { - json!({ - "type": "object", - "properties": { - "user_id": {"type": "integer"}, - "shipping": { - "type": "object", - "properties": { - "city": {"type": "string"}, - "zip": {"type": "integer"} - } - } - } - }) - } - - fn build_call() -> String { - r#"<｜DSML｜tool_calls> -<｜DSML｜invoke name="create_order"> -<｜DSML｜parameter name="user_id" string="false">42 -<｜DSML｜parameter name="shipping" string="false">{"city":"Singapore","zip":18956} - -"# - .to_owned() - } - - fn make_py_tool(py: Python<'_>) -> PyResult> { - let parameters = pythonize(py, &tool_schema()).map_err(|error| { - PyValueError::new_err(format!( - "failed to convert test schema from JSON to Python: {error}" - )) - })?; - Py::new( - py, - PyTool::new( - "create_order".to_owned(), - Some("Create an order".to_owned()), - ¶meters, - None, - )?, - ) - } - - #[test] - fn tool_round_trips_typed_fields() { - with_python(|py| { - let tool = make_py_tool(py)?; - let borrowed = tool.borrow(py); - assert_eq!(borrowed.name(), "create_order"); - assert_eq!(borrowed.description(), Some("Create an order")); - assert_eq!(borrowed.strict(), None); - - let parameters = borrowed.parameters(py)?; - let parameters = depythonize::(parameters.bind(py))?; - assert_eq!(parameters, tool_schema()); - PyResult::Ok(()) - }) - .unwrap(); - } - - #[test] - fn output_append_and_coalesce() { - with_python(|py| { - let first = Py::new( - py, - PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()), - )?; - let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?; - let mut output = PyToolParserOutput::new(py, "text", Some(vec![first])); - let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?; - output.append(other.borrow(py)); - - let coalesced = output.coalesce(); - assert_eq!(coalesced.normal_text(), "text"); - let calls = coalesced.calls(); - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].tool_index(), 0); - assert_eq!(calls[0].name(), Some("create_order")); - assert_eq!(calls[0].arguments(), "{\"a\":1}"); - PyResult::Ok(()) - }) - .unwrap(); - } - - #[test] - fn parser_parse_finish_and_preserve_special_tokens() { - with_python(|py| { - let tool = make_py_tool(py)?; - let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?; - assert!(parser.preserve_special_tokens()); - - let mut output = PyToolParserOutput::new(py, "", None); - parser.parse_into_output(&build_call(), &mut output)?; - let finish = Py::new(py, parser.finish()?)?; - output.append(finish.borrow(py)); - let output = output.coalesce(); - - assert_eq!(output.normal_text(), ""); - let calls = output.calls(); - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].name(), Some("create_order")); - assert_eq!( - serde_json::from_str::(calls[0].arguments()).unwrap(), - json!({ - "user_id": 42, - "shipping": { - "city": "Singapore", - "zip": 18956 - } - }) - ); - - assert_eq!(parser.reset(), ""); - PyResult::Ok(()) - }) - .unwrap(); - } - - #[test] - fn parser_exposes_model_emitted_tool_call_ids() { - with_python(|py| { - let tool = make_py_tool(py)?; - let mut parser = PyToolParser::new(py, "KimiK2ToolParser", vec![tool])?; - - let input = "<|tool_calls_section_begin|>\ - <|tool_call_begin|>functions.create_order:0<|tool_call_argument_begin|>\ - {\"user_id\":42}<|tool_call_end|>\ - <|tool_calls_section_end|>"; - let mut output = PyToolParserOutput::new(py, "", None); - parser.parse_into_output(input, &mut output)?; - - assert_eq!(parser.tool_call_id(0), Some("functions.create_order:0")); - assert_eq!(parser.tool_call_id(1), None); - PyResult::Ok(()) - }) - .unwrap(); - } - - #[test] - fn parser_errors_for_unknown_name() { - with_python(|py| { - let tool = make_py_tool(py)?; - let error = match PyToolParser::new(py, "missing", vec![tool]) { - Ok(_) => panic!("missing parser name unexpectedly succeeded"), - Err(error) => error, - }; - let message = format!("{error}"); - assert!(message.contains("unsupported tool parser `missing`")); - PyResult::Ok(()) - }) - .unwrap(); - } -} diff --git a/setup.py b/setup.py index 67f192cdd393..190d65d9fb62 100644 --- a/setup.py +++ b/setup.py @@ -87,6 +87,8 @@ def get_missing_precompiled_rust_extension_modules() -> list[str]: def has_precompiled_rust_extensions() -> bool: + if not rust_build.rust_py_extension_module_names(): + return False return not get_missing_precompiled_rust_extension_modules() diff --git a/tests/tool_parsers/test_rust_tool_parser.py b/tests/tool_parsers/test_rust_tool_parser.py deleted file mode 100644 index 2349d4d292ab..000000000000 --- a/tests/tool_parsers/test_rust_tool_parser.py +++ /dev/null @@ -1,328 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -from collections.abc import Sequence -from unittest.mock import MagicMock - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionToolsParam, -) -from vllm.tool_parsers.rust_tool_parser import RustToolParser - -# The PyO3 extension is an optional build artifact; skip when absent. -_rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser") - -MOCK_TOKENIZER = MagicMock() -MOCK_TOKENIZER.get_vocab.return_value = {} - -TC_START = "<｜DSML｜tool_calls>" -TC_END = "" -INV_START = '<｜DSML｜invoke name="' -INV_END = "" -PARAM_START = '<｜DSML｜parameter name="' -PARAM_END = "" - - -class DeepSeekV4RustToolParser(RustToolParser): - rust_parser_name = "DeepSeekV4ToolParser" - tool_call_start_token = TC_START - - -class KimiK2RustToolParser(RustToolParser): - rust_parser_name = "KimiK2ToolParser" - tool_call_start_token = "<|tool_calls_section_begin|>" - - -def sample_tools() -> list[ChatCompletionToolsParam]: - return [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "get_weather", - "description": "Get weather for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string"}, - "date": {"type": "string"}, - }, - }, - }, - ), - ChatCompletionToolsParam( - type="function", - function={ - "name": "add", - "description": "Add two integers", - "parameters": { - "type": "object", - "properties": { - "x": {"type": "integer"}, - "y": {"type": "integer"}, - }, - }, - }, - ), - ] - - -EXPECTED_CALLS = [ - ("get_weather", {"location": "SF", "date": "2024-01-16"}), - ("add", {"x": 3, "y": 5}), -] - - -def build_invoke( - function_name: str, - params: Sequence[tuple[str, str, bool]], -) -> str: - param_text = "\n".join( - f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}' - for name, value, is_string in params - ) - return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n' - - -def build_tool_call() -> str: - weather = build_invoke( - "get_weather", - [ - ("location", "SF", True), - ("date", "2024-01-16", True), - ], - ) - add = build_invoke( - "add", - [ - ("x", "3", False), - ("y", "5", False), - ], - ) - return f"{TC_START}\n{weather}{add}{TC_END}" - - -def parse_streaming( - parser: DeepSeekV4RustToolParser, - text: str, - chunk_size: int, -) -> list: - deltas = [] - previous_text = "" - for start in range(0, len(text), chunk_size): - delta_text = text[start : start + chunk_size] - current_text = previous_text + delta_text - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[1], - request=MagicMock(), - ) - previous_text = current_text - if delta is not None: - deltas.append(delta) - - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=previous_text, - delta_text="", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[2], - request=MagicMock(), - ) - if delta is not None: - deltas.append(delta) - - return deltas - - -def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str: - return "".join( - tool_call.function.arguments - for delta in deltas - for tool_call in delta.tool_calls or [] - if ( - tool_call.index == tool_index - and tool_call.function is not None - and tool_call.function.arguments is not None - ) - ) - - -def test_rust_tool_parser_extension_typed_api() -> None: - tools = [ - _rust_tool_parser.Tool( - tool.function.name, - tool.function.description, - tool.function.parameters, - None, - ) - for tool in sample_tools() - ] - parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools) - output = _rust_tool_parser.ToolParserOutput() - - parser.parse_into(build_tool_call(), output) - output.append(parser.finish()) - output = output.coalesce() - - assert parser.preserve_special_tokens() - assert output.normal_text == "" - assert len(output.calls) == 2 - for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS): - assert call.name == name - assert json.loads(call.arguments) == arguments - - -def test_rust_tool_parser_adapter_extracts_complete_output() -> None: - tools = sample_tools() - parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) - - result = parser.extract_tool_calls( - "Let me create it. " + build_tool_call(), - ChatCompletionRequest(messages=[], model="m", tools=tools), - ) - - assert result.tools_called - assert result.content == "Let me create it. " - assert len(result.tool_calls) == 2 - for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): - assert tool_call.function.name == name - assert json.loads(tool_call.function.arguments) == arguments - - -def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None: - parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) - - deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) - - names = [ - tool_call.function.name - for delta in deltas - for tool_call in delta.tool_calls or [] - if tool_call.function is not None and tool_call.function.name is not None - ] - assert names == [name for name, _ in EXPECTED_CALLS] - for index, (_, arguments) in enumerate(EXPECTED_CALLS): - assert json.loads(collect_streamed_arguments(deltas, index)) == arguments - - -def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None: - parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) - text = build_tool_call() - split_at = len(TC_START) + 8 - deltas = [] - previous_text = "" - - for delta_text in (text[:split_at], "", text[split_at:], ""): - current_text = previous_text + delta_text - delta = parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[1], - request=MagicMock(), - ) - previous_text = current_text - if delta is not None: - deltas.append(delta) - - names = [ - tool_call.function.name - for delta in deltas - for tool_call in delta.tool_calls or [] - if tool_call.function is not None and tool_call.function.name is not None - ] - assert names == [name for name, _ in EXPECTED_CALLS] - for index, (_, arguments) in enumerate(EXPECTED_CALLS): - assert json.loads(collect_streamed_arguments(deltas, index)) == arguments - - -KIMI_EXPECTED_IDS = ["functions.get_weather:0", "functions.add:1"] - - -def build_kimi_tool_call() -> str: - return ( - "<|tool_calls_section_begin|>" - "<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>" - '{"location": "SF", "date": "2024-01-16"}<|tool_call_end|>' - "<|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>" - '{"x": 3, "y": 5}<|tool_call_end|>' - "<|tool_calls_section_end|>" - ) - - -def test_rust_tool_parser_adapter_complete_prefers_model_tool_call_ids() -> None: - tools = sample_tools() - parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=tools) - - result = parser.extract_tool_calls( - "Let me check. " + build_kimi_tool_call(), - ChatCompletionRequest(messages=[], model="m", tools=tools), - ) - - assert result.tools_called - assert [tool_call.id for tool_call in result.tool_calls] == KIMI_EXPECTED_IDS - for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): - assert tool_call.function.name == name - assert json.loads(tool_call.function.arguments) == arguments - - -def test_rust_tool_parser_adapter_streaming_prefers_model_tool_call_ids() -> None: - parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) - - deltas = parse_streaming(parser, build_kimi_tool_call(), chunk_size=5) - - ids = [ - tool_call.id - for delta in deltas - for tool_call in delta.tool_calls or [] - if tool_call.id is not None - ] - assert ids == KIMI_EXPECTED_IDS - for index, (_, arguments) in enumerate(EXPECTED_CALLS): - assert json.loads(collect_streamed_arguments(deltas, index)) == arguments - - -def test_rust_tool_parser_adapter_streaming_generates_ids_as_fallback() -> None: - # DeepSeekV4 never emits model tool call IDs, so the bridge mints them. - parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) - - deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) - - ids = [ - tool_call.id - for delta in deltas - for tool_call in delta.tool_calls or [] - if tool_call.function is not None and tool_call.function.name is not None - ] - assert len(ids) == len(EXPECTED_CALLS) - assert all(ids) - assert len(set(ids)) == len(ids) - - -def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None: - tools = sample_tools() - parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) - request = ChatCompletionRequest( - messages=[], - model="m", - tools=tools, - tool_choice="required", - skip_special_tokens=True, - ) - - adjusted = parser.adjust_request(request) - - assert adjusted is request - assert adjusted.skip_special_tokens is False - assert adjusted.structured_outputs is None diff --git a/tools/build_rust.py b/tools/build_rust.py index 320cf75e3db4..b5c3945d63f9 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -55,14 +55,15 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]: binding=Binding.Exec, optional=optional, ), - RustExtension( - target="vllm._rust_tool_parser", - path="rust/src/parser/python/Cargo.toml", - features=["pyo3/abi3-py38"], - binding=Binding.PyO3, - optional=optional, - py_limited_api=True, - ), + # Example PyO3 module (the removed tool-parser bridge): + # RustExtension( + # target="vllm._rust_tool_parser", + # path="rust/src/parser/python/Cargo.toml", + # features=["pyo3/abi3-py38"], + # binding=Binding.PyO3, + # optional=optional, + # py_limited_api=True, + # ), ] diff --git a/vllm/tool_parsers/rust_tool_parser.py b/vllm/tool_parsers/rust_tool_parser.py deleted file mode 100644 index da85825d4cc9..000000000000 --- a/vllm/tool_parsers/rust_tool_parser.py +++ /dev/null @@ -1,322 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import importlib -from collections.abc import Sequence -from typing import Any - -from openai.types.responses.function_tool import FunctionTool - -from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.generate.base.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionToolsParam, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser - -logger = init_logger(__name__) - - -def _rust_tool_parser_module() -> Any: - try: - return importlib.import_module("vllm._rust_tool_parser") - except ImportError as exc: - raise RuntimeError( - "Rust tool parsing requires the vllm._rust_tool_parser PyO3 " - "extension. Rebuild vLLM with Rust frontend/extensions enabled." - ) from exc - - -class RustToolParser(ToolParser): - """Adapter from an opaque Rust parser to the vLLM ToolParser API. - - Subclasses provide only model-specific configuration: the exact Rust parser - name and an optional tool-call start marker for fast complete-output - rejection. - - This class keeps the vLLM-specific bridge work: - - convert vLLM tool definitions into the Rust ``Tool`` shape; - - translate typed Rust parser outputs into vLLM protocol objects; and - - maintain vLLM streaming bookkeeping used by finish-reason handling. - - The parser grammar and incremental parser state stay in Rust. - """ - - # Rust-backed parsers are opaque to Python by default. Do not use vLLM's - # standard JSON required/named handling; let the Rust parser consume the - # model's native tool-call syntax. - supports_required_and_named = False - - rust_parser_name: str - tool_call_start_token: str | None = None - - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - self._parser: Any | None = None - self._error: Exception | None = None - - if not self.model_tokenizer: - raise ValueError( - "The model tokenizer must be passed to the ToolParser " - "constructor during construction." - ) - - logger.debug( - "vLLM successfully imported tool parser %s", self.__class__.__name__ - ) - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - """Adjust request options without installing Python-side constraints. - - Rust-backed parsers are treated as source-of-truth opaque parsers. The - bridge intentionally avoids ``super().adjust_request()`` so Python does - not install JSON schema guidance or structural-tag constraints that may - conflict with the Rust parser's native grammar. - """ - if self._get_parser().preserve_special_tokens(): - request.skip_special_tokens = False - return request - - def _rust_tools(self) -> list[Any]: - """Build Rust ``Tool`` objects from vLLM tool definitions.""" - if not self.tools: - return [] - - tools: list[Any] = [] - for tool in self.tools: - if isinstance(tool, FunctionTool): - name = tool.name - description = tool.description - parameters = tool.parameters or {} - strict = getattr(tool, "strict", None) - elif isinstance(tool, ChatCompletionToolsParam): - name = tool.function.name - description = tool.function.description - parameters = tool.function.parameters or {} - strict = getattr(tool.function, "strict", None) - else: - continue - tools.append( - _rust_tool_parser_module().Tool(name, description, parameters, strict) - ) - return tools - - def _new_parser(self) -> Any: - """Create a fresh Rust parser with the current tool schemas.""" - return _rust_tool_parser_module().ToolParser( - self.rust_parser_name, self._rust_tools() - ) - - def _get_parser(self) -> Any: - if self._parser is None: - self._parser = self._new_parser() - return self._parser - - def _reset_streaming_state(self) -> None: - """Reset parser state for a new request on a reused parser instance.""" - self._parser = self._new_parser() - self._error = None - self.prev_tool_call_arr.clear() - self.streamed_args_for_tool.clear() - self.current_tool_id = -1 - self.current_tool_name_sent = False - - def _ensure_tool_state(self, index: int) -> None: - """Grow vLLM streaming state arrays to contain ``index``.""" - while len(self.prev_tool_call_arr) <= index: - self.prev_tool_call_arr.append({}) - while len(self.streamed_args_for_tool) <= index: - self.streamed_args_for_tool.append("") - - def _record_delta( - self, index: int, name: str | None, arguments: str | None - ) -> str | None: - """Mirror a Rust parser delta into vLLM streaming bookkeeping. - - ``prev_tool_call_arr`` and ``streamed_args_for_tool`` are read later by - the chat serving layer to decide the final ``tool_calls`` finish reason - and to flush any remaining argument bytes. - """ - tool_call_id = None - self._ensure_tool_state(index) - - if name is not None: - # Prefer the model-emitted ID surfaced by the Rust parser (e.g. - # Kimi K2) over a randomly generated one. - tool_call_id = self._get_parser().tool_call_id(index) or make_tool_call_id() - self.prev_tool_call_arr[index] = {"name": name, "arguments": {}} - self.current_tool_name_sent = True - - if arguments is not None: - self.streamed_args_for_tool[index] += arguments - self.prev_tool_call_arr[index]["arguments"] = self.streamed_args_for_tool[ - index - ] - self.current_tool_id = index - - return tool_call_id - - def _delta_message_from_parser_output( - self, parser_output: Any | None - ) -> DeltaMessage | None: - """Translate one Rust parser output into a vLLM ``DeltaMessage``.""" - if parser_output is None: - return None - - normal_text = parser_output.normal_text or None - tool_calls: list[DeltaToolCall] = [] - for tool_call in parser_output.calls: - index = tool_call.tool_index - name = tool_call.name - arguments: str | None = tool_call.arguments - if name is None and arguments is None: - continue - - tool_call_id = self._record_delta(index, name, arguments) - tool_calls.append( - DeltaToolCall( - index=index, - id=tool_call_id, - type="function" if name is not None else None, - function=DeltaFunctionCall( - name=name, - arguments=arguments, - ), - ) - ) - - if normal_text is None and not tool_calls: - return None - return DeltaMessage(content=normal_text, tool_calls=tool_calls) - - def _parse_complete(self, model_output: str) -> tuple[Any, dict[int, str]] | None: - """Parse complete model output with a throwaway Rust parser instance. - - Returns the coalesced parser output along with any model-emitted tool - call IDs keyed by tool index. - """ - parser = self._new_parser() - output = _rust_tool_parser_module().ToolParserOutput() - try: - parser.parse_into(model_output, output) - # finish() clears parser state, so snapshot model-emitted IDs first. - tool_call_ids = { - call.tool_index: tool_call_id - for call in output.calls - if (tool_call_id := parser.tool_call_id(call.tool_index)) is not None - } - output.append(parser.finish()) - except Exception: - logger.exception( - "Error parsing %s tool call output.", self.rust_parser_name - ) - return None - return output.coalesce(), tool_call_ids - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - """Extract tool calls from complete model output (non-streaming).""" - if ( - self.tool_call_start_token is not None - and self.tool_call_start_token not in model_output - ): - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=model_output, - ) - - parse_result = self._parse_complete(model_output) - if parse_result is None: - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=model_output, - ) - parsed, tool_call_ids = parse_result - - tool_calls: list[ToolCall] = [] - self.prev_tool_call_arr.clear() - for parsed_tool_call in parsed.calls: - name = parsed_tool_call.name - arguments = parsed_tool_call.arguments or "{}" - if name is None: - continue - tool_calls.append( - ToolCall( - id=tool_call_ids.get(parsed_tool_call.tool_index) - or make_tool_call_id(), - type="function", - function=FunctionCall(name=name, arguments=arguments), - ) - ) - self.prev_tool_call_arr.append({"name": name, "arguments": arguments}) - - if not tool_calls: - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=model_output, - ) - - content = parsed.normal_text or None - return ExtractedToolCallInformation( - tools_called=True, - tool_calls=tool_calls, - content=content, - ) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], # pylint: disable=unused-argument - current_token_ids: Sequence[int], # pylint: disable=unused-argument - delta_token_ids: Sequence[int], # pylint: disable=unused-argument - request: ChatCompletionRequest, # pylint: disable=unused-argument - ) -> DeltaMessage | None: - """Extract tool calls from streaming model output. - - The Rust parser owns the incremental buffer, so this adapter feeds only - the newest text delta and lets the serving layer handle final empty - chunks. - """ - # TODO: Add a final-chunk hook if streaming needs to call Rust finish(). - if not previous_text: - self._reset_streaming_state() - - if self._error is not None: - return None - - parser_output = _rust_tool_parser_module().ToolParserOutput() - try: - self._get_parser().parse_into(delta_text, parser_output) - except Exception as error: - self._error = error - logger.exception( - "Error parsing %s streaming tool call output.", - self.rust_parser_name, - ) - - delta_message = self._delta_message_from_parser_output(parser_output) - if delta_message is not None: - return delta_message - - return None From 40c86ae1ce84900876f9413d175c820bdea20934 Mon Sep 17 00:00:00 2001 From: lazymio Date: Tue, 6 Oct 2026 13:49:13 +0800 Subject: [PATCH 052/952] [Kernel] Enable OAI Triton MXFP4 MoE on SM12x (#58877) Signed-off-by: lazymio Co-authored-by: Claude Fable 5.1 Co-authored-by: Shengqi Chen --- .../moe/test_gpt_oss_triton_kernels.py | 41 ++++++++++++++++- .../experts/gpt_oss_triton_kernels_moe.py | 44 ++++++++++++++++--- .../layers/fused_moe/oracle/mxfp4.py | 9 ++++ .../layers/quantization/utils/mxfp4_utils.py | 9 +++- 4 files changed, 96 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index c1d83d8349bf..406e525bdfab 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -32,8 +32,11 @@ from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( + _sm12x_block_m, triton_kernel_moe_forward, ) +from vllm.model_executor.layers.fused_moe.oracle import mxfp4 as mxfp4_oracle +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs from vllm.utils.math_utils import round_up @@ -314,9 +317,13 @@ def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init): pc2, ) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8) - if current_platform.is_device_capability_family(100): + if current_platform.is_device_capability_family( + 100 + ) or current_platform.is_device_capability_family(120): + # Mirrors `_swizzle_mxfp4`; SM12x also needs epilogue_subtile=1. constraints = { "is_persistent": True, + "epilogue_subtile": 1, } opt_flags.update_opt_flags_constraints(constraints) @@ -356,6 +363,38 @@ def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init): assert_close(ref=out_ref, tri=out_triton_monolithic, maxtol=0.025, rmstol=0.005) +@pytest.mark.parametrize( + ("num_rows", "num_experts", "expected"), + [ + (4, 32, 16), # bs=1 decode: below the 16-row floor + (32 * 20, 32, 32), # 20 tokens/expert rounds up to 32 + (32 * 33, 32, 32), # heuristic would pick 64; capped for 99KB smem + (12000, 32, 32), # prefill: heuristic would pick 128 + ], +) +def test_sm12x_block_m(num_rows, num_experts, expected): + assert _sm12x_block_m(num_rows, num_experts) == expected + + +@pytest.mark.parametrize("sm12x", [True, False]) +def test_gpt_oss_priority_prefers_marlin_on_sm12x( + monkeypatch: pytest.MonkeyPatch, sm12x: bool +): + """SM12x runs the OAI Triton kernels at small tiles only, so MARLIN must + stay the default there while TRITON remains selectable explicitly.""" + monkeypatch.setattr(mxfp4_oracle.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + mxfp4_oracle.current_platform, + "is_device_capability_family", + lambda family: sm12x and family == 120, + ) + order = mxfp4_oracle._get_priority_backends_for_gpt_oss() + triton_first = order.index(Mxfp4MoeBackend.TRITON) < order.index( + Mxfp4MoeBackend.MARLIN + ) + assert triton_first is not sm12x + + def test_unit_shuffle(): N = ModelConfig.intermediate_size K = ModelConfig.hidden_size diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index d23edaa01ab4..e3b6995d693c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -41,11 +43,11 @@ def _triton_kernel_moe_supports_current_device() -> bool: p = current_platform if p.is_cuda(): cap = p.get_device_capability() - # Keep the original `(9, 0) <= cap < (11, 0)` window on - # CUDA (covers Hopper SM90 and Blackwell SM100, excludes - # SM120) — this PR is ROCm-scoped and the broader CUDA - # range was not validated. - return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) + # Hopper SM90, datacenter Blackwell SM100/103 and consumer Blackwell + # SM120/121. SM12x needs the opt-flag constraints applied by + # `_swizzle_mxfp4` and `_constrain_sm12x_block_m` to fit its 99KB of + # shared memory; it is still ordered after MARLIN by the selector. + return cap is not None and (9, 0) <= (cap.major, cap.minor) < (13, 0) if p.is_rocm(): from vllm.platforms.rocm import on_gfx1x, on_gfx9 @@ -57,6 +59,36 @@ def _triton_kernel_moe_supports_current_device() -> bool: return False +@functools.cache +def _is_sm12x() -> bool: + return current_platform.is_cuda() and current_platform.is_device_capability_family( + 120 + ) + + +# Consumer Blackwell has 99KB of opt-in shared memory: with block_m >= 64 the +# persistent mxfp4 kernel fits a single pipeline stage and runs 2-3x slower +# than with 32, but triton_kernels' block_m heuristic still tops out at 128. +_SM12X_MAX_BLOCK_M = 32 + + +def _sm12x_block_m(num_rows: int, num_experts: int) -> int: + """triton_kernels' tokens-per-expert block_m heuristic, capped at 32.""" + tokens_per_expt = max(1, num_rows // num_experts) + return max(16, min(triton.next_power_of_2(tokens_per_expt), _SM12X_MAX_BLOCK_M)) + + +def _constrain_sm12x_block_m(num_rows: int, num_experts: int) -> None: + """Pin block_m for the following matmul_ogs calls on SM12x.""" + if not _is_sm12x(): + return + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags + + opt_flags.update_opt_flags_constraints( + {"block_m": _sm12x_block_m(num_rows, num_experts)} + ) + + def _patch_make_bitmatrix_metadata() -> None: """Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k. @@ -784,6 +816,7 @@ def triton_kernel_fused_experts( ) ) gammas = routing_data.gate_scal if routing_data else None + _constrain_sm12x_block_m(M * topk, E) matmul_ogs( hidden_states, @@ -1403,6 +1436,7 @@ def apply( intermediate_cache2 = _resize_cache(workspace13, (M * topk, activation_out_dim)) gammas = routing_data.gate_scal if routing_data else None + _constrain_sm12x_block_m(M * topk, E) matmul_ogs( hidden_states, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index a3c65dc4a5d5..691d849228f8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -386,6 +386,15 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.CPU, Mxfp4MoeBackend.EMULATION, ] + if current_platform.is_cuda() and current_platform.is_device_capability_family(120): + # Consumer Blackwell: the OAI Triton kernels run but their 99KB shared + # memory budget caps them at small tiles, so MARLIN prefills ~10% + # faster at equal decode throughput. Keep TRITON as an opt-in. + _AVAILABLE_BACKENDS.remove(Mxfp4MoeBackend.TRITON) + _AVAILABLE_BACKENDS.insert( + _AVAILABLE_BACKENDS.index(Mxfp4MoeBackend.BATCHED_MARLIN) + 1, + Mxfp4MoeBackend.TRITON, + ) return _AVAILABLE_BACKENDS diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 8cc626f3866a..45562b12b1da 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -120,7 +120,14 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): K = scale.shape[-1] pad_k = -K % 4 scale = torch.nn.functional.pad(scale, (0, pad_k)) - elif current_platform.is_device_capability_family(100): + elif current_platform.is_device_capability_family( + 100 + ) or current_platform.is_device_capability_family(120): + # Native MXFP (cap >= 10.0) requires the persistent kernel. On + # SM12x, epilogue_subtile=1 also keeps the heuristic away from a + # triton_kernels 3.5.1 codegen bug (SWAP_XW with a subtiled + # epilogue); block_m is pinned per call in + # gpt_oss_triton_kernels_moe to fit the 99KB shared memory. constraints = { "is_persistent": True, "epilogue_subtile": 1, From 1e5d0ea888c035a7240f2ed0eb77fd9857635be9 Mon Sep 17 00:00:00 2001 From: sungbin1015 Date: Tue, 6 Oct 2026 14:58:47 +0900 Subject: [PATCH 053/952] [CI/Build] Fix incorrect CLI arg validation tests in test_cli_args.py (#46779) Signed-off-by: sungbin1015 Co-authored-by: Claude --- tests/entrypoints/launchers/test_cli_args.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index 6b9d496822c3..84279e155582 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -242,7 +242,7 @@ def test_rust_grpc_port_requires_serve_subcommand(monkeypatch): validate_parsed_serve_args(args) -def test_enable_auto_choice_passes_without_tool_call_parser(serve_parser): +def test_enable_auto_choice_fails_without_tool_call_parser(serve_parser): """Ensure validation fails if tool choice is enabled with no call parser.""" # If we enable-auto-tool-choice, explode with no tool-call-parser args = serve_parser.parse_args(args=["--enable-auto-tool-choice"]) @@ -262,17 +262,23 @@ def test_enable_auto_choice_passes_with_tool_call_parser(serve_parser): validate_parsed_serve_args(args) -def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser): - """Ensure validation fails if reasoning is enabled with auto tool choice.""" +def test_auto_tool_choice_with_reasoning_parser_passes(serve_parser): + """Auto tool choice combined with a reasoning parser is supported. + + See docs/features/tool_calling.md, which documents using + ``--tool-call-parser`` together with ``--reasoning-parser``. This + combination must pass validation as long as a tool call parser is set. + """ args = serve_parser.parse_args( args=[ "--enable-auto-tool-choice", + "--tool-call-parser", + "mistral", "--reasoning-parser", "deepseek_r1", ] ) - with pytest.raises(TypeError): - validate_parsed_serve_args(args) + validate_parsed_serve_args(args) def test_passes_with_reasoning_parser(serve_parser): From 6ad2e06c4a5bdd878977dfccc7443831e4fdc889 Mon Sep 17 00:00:00 2001 From: Tarun Kumar Date: Tue, 6 Oct 2026 12:01:13 +0530 Subject: [PATCH 054/952] fix: exclude flashinfer-jit-cache from wheel install_requires (#60052) Signed-off-by: Tarun Kumar --- setup.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 190d65d9fb62..0ac53a45b5b1 100644 --- a/setup.py +++ b/setup.py @@ -1341,10 +1341,10 @@ def _read_requirements(filename: str) -> list[str]: # vllm-flash-attn is built only for CUDA 12.x. # Skip for other versions. continue - if "flashinfer-cubin" in req: - # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so - # it cannot be a wheel dependency; flashinfer falls back to - # fetching cubins at runtime when the package is absent. + if "flashinfer-cubin" in req or "flashinfer-jit-cache" in req: + # Not on PyPI (only https://flashinfer.ai/whl), so they + # cannot be wheel dependencies; flashinfer handles the + # absence of pre-compiled cubins/jit-cache at runtime. continue if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": # [cu13] extra is the default; strip it on CUDA 12 builds. From 15c1cc12ac50c1fb36fd92978c14ad2e2f865055 Mon Sep 17 00:00:00 2001 From: LostFox11 <140548160+LostFox11@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:44:57 +0800 Subject: [PATCH 055/952] [Bugfix][EPLB] Use group-local ranks for torch P2P transfers (#55804) Signed-off-by: LostFox11 Co-authored-by: LostFox11 Co-authored-by: Codex Co-authored-by: YangShuai52 <178856220+YangShuai52@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_eplb_execute.py | 45 +++++++++++++++++++++- vllm/distributed/eplb/eplb_communicator.py | 16 ++++---- 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 2037535e3ddf..acdb9f39fcac 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -36,7 +36,11 @@ def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch): gsd._install_copy_checkers() monkeypatch.setattr(eplb_comm, "is_local_first_rank", lambda: False) - monkeypatch.setattr(eplb_comm, "P2POp", lambda op, tensor, peer, group: tensor) + monkeypatch.setattr( + eplb_comm, + "P2POp", + lambda op, tensor, peer=None, group=None, tag=0, group_peer=None: tensor, + ) def receive(tensors): for tensor in tensors: @@ -849,3 +853,42 @@ def test_rearrange_expert_weights_profile_mode(world_size): _test_rearrange_expert_weights_profile_mode, world_size, ) + + +def _test_eplb_communicator_pipeline_parallel(env, world_size, backend): + set_env_vars_and_device(env) + config = VllmConfig() + config.parallel_config.tensor_parallel_size = 2 + config.parallel_config.pipeline_parallel_size = world_size // 2 + + with set_current_vllm_config(config): + ensure_model_parallel_initialized( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=world_size // 2, + ) + ep_group = get_tp_group() + peer_rank = 1 - ep_group.rank_in_group + rank = torch.distributed.get_rank() + send_tensor = torch.full((16,), rank, device="cuda", dtype=torch.int32) + recv_tensor = torch.empty_like(send_tensor) + communicator = create_eplb_communicator_or_raise( + group_coordinator=ep_group, + backend=backend, + expert_weights=[[send_tensor]], + expert_buffer=[recv_tensor], + ) + communicator.add_send([send_tensor], peer_rank, expert_id=0) + communicator.add_recv([recv_tensor], peer_rank, expert_id=0) + communicator.execute() + + expected = torch.full_like(recv_tensor, ep_group.ranks[peer_rank]) + torch.testing.assert_close(recv_tensor, expected) + + +@pytest.mark.parametrize("world_size", [2, 4]) +@pytest.mark.parametrize("backend", ["torch_nccl", "torch_gloo"]) +def test_eplb_communicator_pipeline_parallel(world_size, backend): + """Transfer weights within each PP stage, including nonzero global ranks.""" + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Need at least {world_size} GPUs to run the test") + distributed_run(_test_eplb_communicator_pipeline_parallel, world_size, backend) diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 730c01a3882e..473058d5258e 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -120,8 +120,8 @@ def add_send( P2POp( torch.distributed.isend, tensor, - dst_rank, - self._ep_group, + group=self._ep_group, + group_peer=dst_rank, ) ) @@ -136,8 +136,8 @@ def add_recv( P2POp( torch.distributed.irecv, tensor, - src_rank, - self._ep_group, + group=self._ep_group, + group_peer=src_rank, ) ) @@ -199,8 +199,8 @@ def build_ops() -> None: P2POp( torch.distributed.isend, cpu_tensor, - peer_rank, - self._cpu_group, + group=self._cpu_group, + group_peer=peer_rank, ) ) continue @@ -211,8 +211,8 @@ def build_ops() -> None: P2POp( torch.distributed.irecv, cpu_tensor, - peer_rank, - self._cpu_group, + group=self._cpu_group, + group_peer=peer_rank, ) ) recv_staging.append((tensor, cpu_tensor)) From edfccc946925a8439e27abd76da8d1a640d1a1b7 Mon Sep 17 00:00:00 2001 From: big_yellow_duck <83417790+big-yellow-duck@users.noreply.github.com> Date: Tue, 6 Oct 2026 15:53:47 +0800 Subject: [PATCH 056/952] [ROCM][Perf] Add ROCM_SEGMENTED_ATTN attention-backend works across RDNA3/3.5/4 (#59132) Signed-off-by: root Signed-off-by: big-yellow-duck Co-authored-by: Tan Pin Siang <1716735+tanpinsiang@users.noreply.github.com> Co-authored-by: vllmellm <190700713+vllmellm@users.noreply.github.com> Co-authored-by: Codex --- .../attention/test_segmented_attention.py | 1408 +++++++++++++++++ .../test_rocm_attention_backends_selection.py | 210 +++ vllm/config/kernel.py | 4 + vllm/model_executor/models/qwen3_dflash.py | 19 + vllm/model_executor/warmup/kernel_warmup.py | 7 + .../rocm_segmented_attn_autotune_warmup.py | 1335 ++++++++++++++++ vllm/v1/attention/backends/registry.py | 3 + .../attention/backends/rocm_segmented_attn.py | 371 +++++ vllm/v1/attention/ops/segmented_attention.py | 1074 +++++++++++++ 9 files changed, 4431 insertions(+) create mode 100644 tests/kernels/attention/test_segmented_attention.py create mode 100644 vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py create mode 100644 vllm/v1/attention/backends/rocm_segmented_attn.py create mode 100644 vllm/v1/attention/ops/segmented_attention.py diff --git a/tests/kernels/attention/test_segmented_attention.py b/tests/kernels/attention/test_segmented_attention.py new file mode 100644 index 000000000000..0fbf5fe8944e --- /dev/null +++ b/tests/kernels/attention/test_segmented_attention.py @@ -0,0 +1,1408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Segmented attention accuracy, cache semantics, and startup tuning contracts.""" + +import math + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import triton + +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + +pytestmark = [ + pytest.mark.skip_global_cleanup, + pytest.mark.skipif( + not (current_platform.is_rocm() and on_gfx1x()), + reason="Segmented attention requires RDNA GPUs", + ), +] + + +def _make_paged_attention_case( + query_lens: list[int], + context_lens: list[int], + *, + num_heads: int, + num_kv_heads: int, + head_size: int, + block_size: int, + fp8: bool, + causal: bool = True, + sliding_window: int = 0, + softcap: float = 0.0, + sinks: torch.Tensor | None = None, +): + device = torch.device("cuda:0") + generator = torch.Generator(device=device).manual_seed(419) + total_queries = sum(query_lens) + qkv = ( + torch.randn( + total_queries, + num_heads + 2 * num_kv_heads, + head_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + * 0.25 + ) + query, key, value = qkv.split((num_heads, num_kv_heads, num_kv_heads), dim=1) + cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16 + k_scale = torch.tensor(0.125 if fp8 else 1.0, device=device) + v_scale = torch.tensor(0.25 if fp8 else 1.0, device=device) + max_seq_len = max(q + c for q, c in zip(query_lens, context_lens)) + blocks_per_seq = triton.cdiv(max_seq_len, block_size) + num_blocks = len(query_lens) * blocks_per_seq + key_cache = ( + torch.randn( + num_blocks, + block_size, + num_kv_heads, + head_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 0.25 + / k_scale + ).to(cache_dtype) + value_cache = ( + torch.randn( + num_blocks, + block_size, + num_kv_heads, + head_size, + device=device, + dtype=torch.float32, + generator=generator, + ) + * 0.25 + / v_scale + ).to(cache_dtype) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32).view( + len(query_lens), blocks_per_seq + ) + starts = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + seq_lens = torch.tensor( + [q + c for q, c in zip(query_lens, context_lens)], + device=device, + dtype=torch.int32, + ) + + first = 0 + for seq, (query_len, context_len) in enumerate(zip(query_lens, context_lens)): + for query_pos in range(query_len): + cache_pos = context_len + query_pos + block = block_table[seq, cache_pos // block_size] + offset = cache_pos % block_size + key_cache[block, offset] = (key[first + query_pos].float() / k_scale).to( + cache_dtype + ) + value_cache[block, offset] = ( + value[first + query_pos].float() / v_scale + ).to(cache_dtype) + first += query_len + + references = [] + first = 0 + for seq, (query_len, context_len) in enumerate(zip(query_lens, context_lens)): + physical = block_table[seq].long() + seq_len = context_len + query_len + full_key = ( + key_cache[physical].flatten(0, 1)[:seq_len].float() * k_scale + ).repeat_interleave(num_heads // num_kv_heads, dim=1) + full_value = ( + value_cache[physical].flatten(0, 1)[:seq_len].float() * v_scale + ).repeat_interleave(num_heads // num_kv_heads, dim=1) + scores = torch.einsum( + "qhd,khd->hqk", query[first : first + query_len].float(), full_key + ) / math.sqrt(head_size) + if softcap > 0: + scores = softcap * torch.tanh(scores / softcap) + query_positions = context_len + torch.arange(query_len, device=device) + key_positions = torch.arange(seq_len, device=device) + causal_mask = key_positions[None, :] > query_positions[:, None] + if causal: + scores.masked_fill_(causal_mask[None], -float("inf")) + if sliding_window > 0: + window_mask = key_positions[None, :] < ( + query_positions[:, None] - sliding_window + 1 + ) + scores.masked_fill_(window_mask[None], -float("inf")) + if sinks is not None: + sink_scores = sinks.float()[:, None, None].expand(-1, query_len, 1) + probabilities = torch.cat((scores, sink_scores), dim=-1).softmax(-1) + probabilities = probabilities[..., :-1] + else: + probabilities = scores.softmax(-1) + references.append(torch.einsum("hqk,khd->qhd", probabilities, full_value)) + first += query_len + + return { + "query": query, + "key": key, + "value": value, + "key_cache": key_cache, + "value_cache": value_cache, + "block_table": block_table, + "starts": starts, + "seq_lens": seq_lens, + "k_scale": k_scale, + "v_scale": v_scale, + "reference": torch.cat(references), + "max_seq_len": max_seq_len, + } + + +def _run_segmented_case(case, output, max_query_len, *, force_splits=None, **kwargs): + from vllm.v1.attention.ops.segmented_attention import ( + run_segmented_attention, + select_segmented_config, + ) + + config = None + heads, dim = case["query"].shape[1:] + if force_splits is not None: + config = dict( + select_segmented_config( + case["seq_lens"].numel(), + max_query_len, + case["max_seq_len"], + heads, + case["key_cache"].shape[2], + dim, + case["key_cache"].element_size() == 1, + ), + splits=force_splits, + ) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + max_query_len, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + dim**-0.5, + config=config, + **kwargs, + ) + + +def test_segmented_tuning_candidates_preserve_workspace_bound(): + """D128/D256 candidates retain their existing split scratch bound.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + from vllm.v1.attention.ops import segmented_attention as segmented + + for batch, query_len, seq_len, heads, kv_heads, dim, fp8 in ( + (1, 1, 128, 8, 8, 128, True), + (32, 1, 262144, 16, 1, 128, True), + (4, 1024, 8192, 16, 4, 256, False), + (1, 1024, 131072, 6, 1, 256, True), + ): + default = segmented.select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + candidates = tuning._candidate_configs(default, batch, query_len, heads, dim) + assert candidates + assert ( + tuning._normalized_config(default, batch, query_len, heads, dim) + in candidates + ) + assert all(config["splits"] <= default["splits"] for config in candidates) + assert all(dim % config["bk"] == 0 for config in candidates) + + +@pytest.mark.parametrize("dim", (128, 256)) +def test_segmented_fp8_long_extend_split_workspace_bound(dim): + from vllm.v1.attention.ops.segmented_attention import ( + MAX_LONG_EXTEND_WORKSPACE_BYTES, + segmented_query_capacity, + select_segmented_config, + ) + + for batch, query_len, seq_len, heads in ( + (1, 256, 8192, 6), + (1, 256, 32768, 6), + (1, 256, 131072, 6), + (1, 1024, 32768, 1), + (1, 1024, 131072, 1), + (1, 1024, 131072, 6), + (1, 4096, 131072, 1), + (4, 1024, 131072, 6), + (5, 257, 131072, 6), + (21, 256, 131072, 6), + ): + config = select_segmented_config(batch, query_len, seq_len, heads, 1, dim, True) + if seq_len >= 32768: + scratch = ( + batch + * segmented_query_capacity(query_len) + * heads + * config["splits"] + * (dim + 1) + * 4 + ) + assert scratch <= MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + + for batch in (1, 4, 16): + for query_len in (257, 1025, 4095, 8192): + for ratio in (1, 6, 16): + for kv_heads in (1, 4): + for seq_len in (32768, 131072): + heads = ratio * kv_heads + config = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, True + ) + if config["splits"] == 1: + continue + scratch = ( + batch + * segmented_query_capacity(query_len) + * heads + * config["splits"] + * (dim + 1) + * 4 + ) + assert scratch <= MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + + +@torch.inference_mode() +def test_segmented_bf16_d256_long_extend_matches_dense_reference(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x segmented prefill") + + case = _make_paged_attention_case( + [1024], + [7168], + num_heads=6, + num_kv_heads=1, + head_size=256, + block_size=1568, + fp8=False, + ) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + 1024, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + 256**-0.5, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize( + "query_len,heads,qk_amplitude,dim", + [ + (256, 1, 1, 256), + (256, 6, 1, 256), + (512, 6, 1, 256), + (256, 6, 4, 256), + (256, 6, 16, 256), + (256, 6, 1, 128), + (512, 6, 4, 128), + (256, 6, 16, 128), + (8192, 1, 1, 128), + (8192, 1, 1, 256), + ], +) +@torch.inference_mode() +def test_segmented_fp8_long_extend_matches_dense_reference( + query_len, heads, qk_amplitude, dim +): + from vllm.platforms.rocm import on_gfx12x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx12x(): + pytest.skip("FP8 KV requires gfx12") + case = _make_paged_attention_case( + [query_len], + [8192 - query_len], + num_heads=heads, + num_kv_heads=1, + head_size=dim, + block_size=1568, + fp8=True, + ) + if qk_amplitude != 1: + # Check peaked logits too: one-term FP8 Q quantization passed only the + # nearly uniform fixture and deviated by 4-70% at these amplitudes. + case["query"].mul_(qk_amplitude) + case["k_scale"].mul_(qk_amplitude) + full_key = ( + case["key_cache"][case["block_table"][0].long()] + .flatten(0, 1)[:8192] + .float() + * case["k_scale"] + ).repeat_interleave(heads, dim=1) + full_value = ( + case["value_cache"][case["block_table"][0].long()] + .flatten(0, 1)[:8192] + .float() + * case["v_scale"] + ).repeat_interleave(heads, dim=1) + scores = torch.einsum( + "qhd,khd->hqk", case["query"].float(), full_key + ) / math.sqrt(dim) + future = torch.arange(8192, device=case["query"].device)[None, :] > ( + 8192 + - query_len + + torch.arange(query_len, device=case["query"].device)[:, None] + ) + scores.masked_fill_(future[None], -float("inf")) + case["reference"] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), full_value) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + query_len, + case["max_seq_len"], + case["k_scale"], + case["v_scale"], + dim**-0.5, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +def test_segmented_tuning_protects_static_incumbent(): + """Noise-sized gains must not replace the static configuration.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + incumbent = {"bm": 16} + challenger = {"bm": 32} + samples = { + tuning._config_key(incumbent): [100.0, 101.0, 99.0, 100.0, 100.0], + tuning._config_key(challenger): [99.0, 100.0, 98.0, 99.0, 99.0], + } + winner, comparisons = tuning._select_tuned_winner( + incumbent, [incumbent, challenger], samples + ) + assert winner["config"] == incumbent + assert comparisons[1]["paired_speedup_vs_default"] < 1.02 + + +def test_segmented_tuning_promotes_verified_challenger(): + """A finalist with a stable material gain should replace the incumbent.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + incumbent = {"bm": 16} + challenger = {"bm": 32} + samples = { + tuning._config_key(incumbent): [100.0, 102.0, 98.0, 101.0, 99.0], + tuning._config_key(challenger): [94.0, 96.0, 92.0, 95.0, 93.0], + } + winner, _ = tuning._select_tuned_winner(incumbent, [incumbent, challenger], samples) + assert winner["config"] == challenger + assert winner["paired_speedup_vs_default"] > 1.02 + + +def test_segmented_tuning_balances_tp_workloads_without_overlap(): + """Every missing bucket belongs to exactly one reasonably balanced rank.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + workloads = list(tuning._workloads(8192, 262144, 32)) + shards = tuning._shard_workloads( + workloads, + 4, + 8192, + 12, + 2, + 256, + True, + ) + flattened = [workload for shard in shards for workload in shard] + assert sorted(flattened) == sorted(workloads) + assert len(flattened) == len(set(flattened)) + owners = { + workload[:2]: rank for rank, shard in enumerate(shards) for workload in shard + } + assert all( + owners[workload[:2]] == rank + for rank, shard in enumerate(shards) + for workload in shard + ) + + weights = [ + sum( + tuning._workload_weight(workload, 8192, 12, 2, 256, True) + for workload in shard + ) + for shard in shards + ] + largest_group = max( + sum( + tuning._workload_weight(workload, 8192, 12, 2, 256, True) + for workload in workloads + if workload[:2] == group + ) + for group in {workload[:2] for workload in workloads} + ) + assert max(weights) - min(weights) <= largest_group + + +def test_segmented_tuning_prunes_scheduler_and_kv_limits(): + """Generated buckets must be reachable under scheduler and cache limits.""" + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + limits = (8192, 262144, 32) + raw = set( + tuning._workloads( + *limits, + dtype=torch.bfloat16, + kv_dtype=torch.float8_e4m3fn, + heads=16, + kv_heads=1, + dim=128, + page=16, + ) + ) + assert {batch for batch, _, _ in raw} == {1, 2, 4, 8, 16, 32} + assert {query for _, query, _ in raw} == set(tuning._QUERY_BUCKETS) + assert (32, 1, 262144) in raw + assert (1, 8192, 262144) in raw + assert (1, 8192, 131072) in raw + assert (1, 8, 131072) in raw + assert all( + sum(tuning._query_lengths(batch, query, limits[0])) <= limits[0] + and batch <= limits[2] + and query <= min(limits[0], tuning.MAX_QUERY_LEN) + and query <= seq_len <= limits[1] + for batch, query, seq_len in raw + ) + + layouts = ((16, 16 * 128 * 2),) + bounded = set( + tuning._workloads( + *limits, + memory_budget_bytes=2**50, + dtype=torch.bfloat16, + kv_dtype=torch.float8_e4m3fn, + heads=16, + kv_heads=1, + dim=128, + page=16, + cache_layouts=layouts, + cache_budget_bytes=2**20, + ) + ) + assert bounded < raw + assert all( + batch * sum(((seq_len + block - 1) // block) * size for block, size in layouts) + <= 2**20 + for batch, _, seq_len in bounded + ) + + +@pytest.fixture +def segmented_tuner(tmp_path, monkeypatch): + """Exercise real cache files and lookup while replacing GPU timing only.""" + from types import SimpleNamespace + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + from vllm.v1.attention.ops import segmented_attention as segmented + + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _: SimpleNamespace( + name="test", gcnArchName="gfx1201", multi_processor_count=32 + ), + ) + monkeypatch.setattr(tuning, "_TABLES", {}) + calls = [] + + def tune(*args, **kwargs): + heads, kv_heads, dim = args[3:6] + batch, query_len, seq_len = args[-1] + default = segmented.select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, False + ) + best = tuning._candidate_configs(default, batch, query_len, heads, dim)[0] + calls.append(args[-1]) + return { + "workload": list(args[-1]), + "query_lengths": tuning._query_lengths(batch, query_len, args[-2]), + "default": default, + "best": best, + "results": [], + } + + monkeypatch.setattr(tuning, "_tune_workload", tune) + geometry = (torch.device("cuda:0"), torch.bfloat16, 4, 2, 128, 16, 128**-0.5) + return SimpleNamespace( + module=tuning, + warmup_args=(*geometry, 4, 8, 2), + lookup_args=(*geometry[:2], torch.bfloat16, *geometry[2:], 1, 1, 8), + calls=calls, + cache_dir=tmp_path, + ) + + +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("device", ["cuda", "cuda:0"]) +def test_segmented_tuning_releases_only_new_kernel_modules(monkeypatch, fail, device): + """Tuning releases candidate modules even on failure without invalidating graphs.""" + import weakref + from types import SimpleNamespace + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + class Module: + pass + + existing = Module() + other_device = Module() + caches = [{"existing": existing} for _ in range(2)] + for name, cache in zip( + ("_segmented_attention_stage", "_segmented_attention_reduce"), caches + ): + monkeypatch.setattr( + tuning, + name, + SimpleNamespace(device_caches={0: (cache,), 1: ({"other": other_device},)}), + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda _: None) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + references = [] + try: + with tuning._temporary_tuning_kernels(torch.device(device)): + for cache in caches: + cache["candidate"] = Module() + references.append(weakref.ref(cache["candidate"])) + if fail: + raise RuntimeError("tuning failed") + except RuntimeError: + assert fail + assert all(reference() is None for reference in references) + assert all(cache == {"existing": existing} for cache in caches) + for kernel in ( + tuning._segmented_attention_stage, + tuning._segmented_attention_reduce, + ): + assert kernel.device_caches[1][0] == {"other": other_device} + + +def test_segmented_tuning_cleanup_preserves_graphs_and_reloads_kernels(): + """Existing graphs survive cleanup, and released variants still run correctly.""" + import weakref + + from vllm.model_executor.warmup import ( + rocm_segmented_attn_autotune_warmup as tuning, + ) + + case = _make_paged_attention_case( + [3], + [11], + num_heads=10, + num_kv_heads=2, + head_size=128, + fp8=False, + block_size=16, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, 3, force_splits=1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_segmented_case(case, output, 3, force_splits=1) + caches = [ + kernel.device_caches[0][0] + for kernel in ( + tuning._segmented_attention_stage, + tuning._segmented_attention_reduce, + ) + ] + keys = [set(cache) for cache in caches] + with tuning._temporary_tuning_kernels(torch.device("cuda:0")): + _run_segmented_case(case, output, 3, force_splits=2) + references = [ + weakref.ref(cache[key]) + for cache, existing in zip(caches, keys) + for key in cache.keys() - existing + ] + assert references + assert all(reference() is None for reference in references) + output.fill_(float("nan")) + graph.replay() + torch.testing.assert_close(output.float(), case["reference"], atol=0.01, rtol=0.01) + output.fill_(float("nan")) + _run_segmented_case(case, output, 3, force_splits=2) + torch.testing.assert_close(output.float(), case["reference"], atol=0.01, rtol=0.01) + + +@pytest.mark.parametrize( + "mode", + [ + pytest.param({}, id="causal"), + pytest.param({"causal": False}, id="noncausal"), + pytest.param({"sliding_window": 7}, id="sliding-window"), + pytest.param({"sliding_window": 7, "causal": False}, id="draft-window"), + pytest.param({"has_sinks": True}, id="sinks"), + ], +) +def test_segmented_tuning_persists_without_retuning(segmented_tuner, monkeypatch, mode): + """A fresh in-memory table reloads the same winner without timing or writing.""" + state = segmented_tuner + tuning = state.module + tuning.warmup_segmented_attention( + *state.warmup_args, memory_budget_bytes=2**50, **mode + ) + assert state.calls + expected = tuning.get_segmented_config(*state.lookup_args, **mode) + assert expected is not None + (cache,) = state.cache_dir.rglob("*.json") + saved = cache.read_bytes(), cache.stat().st_mtime_ns + tuning._TABLES.clear() + + def forbidden(*args, **kwargs): + pytest.fail("Persistent segmented startup cache attempted to retune") + + monkeypatch.setattr(tuning, "_tune_workload", forbidden) + tuning.warmup_segmented_attention( + *state.warmup_args, memory_budget_bytes=2**50, **mode + ) + assert tuning.get_segmented_config(*state.lookup_args, **mode) == expected + assert (cache.read_bytes(), cache.stat().st_mtime_ns) == saved + + +@pytest.mark.parametrize( + "max_tokens,mode", + [ + pytest.param(3, {}, id="scheduler-token-budget"), + pytest.param(4, {"causal": False}, id="noncausal"), + pytest.param(4, {"sliding_window": 7}, id="sliding-window"), + pytest.param(4, {"sliding_window": 7, "causal": False}, id="draft-window"), + pytest.param(4, {"has_sinks": True}, id="sinks"), + ], +) +def test_segmented_tuning_does_not_reuse_incompatible_winners( + segmented_tuner, max_tokens, mode +): + """Changed attention semantics or query mix require their own measurements.""" + state = segmented_tuner + tuning = state.module + tuning.warmup_segmented_attention(*state.warmup_args, memory_budget_bytes=2**50) + previous_calls = len(state.calls) + if mode: + assert tuning.get_segmented_config(*state.lookup_args, **mode) is None + args = (*state.warmup_args[:7], max_tokens, *state.warmup_args[8:]) + tuning.warmup_segmented_attention(*args, memory_budget_bytes=2**50, **mode) + assert len(state.calls) > previous_calls + assert len(list(state.cache_dir.rglob("*.json"))) == 2 + assert tuning.get_segmented_config(*state.lookup_args, **mode) is not None + + +@pytest.mark.parametrize( + "dim,page,hq,hk,dtype,kv_dtype,byte_cache,force_splits", + [ + (128, 16, 4, 4, torch.bfloat16, torch.bfloat16, False, None), + (128, 32, 8, 2, torch.float16, torch.float16, False, None), + (256, 784, 12, 2, torch.bfloat16, torch.bfloat16, False, None), + (256, 1568, 12, 2, torch.bfloat16, torch.float8_e4m3fn, False, None), + (128, 32, 16, 1, torch.bfloat16, torch.float8_e4m3fn, True, None), + (256, 32, 4, 2, torch.float16, torch.float8_e4m3fn, False, None), + (128, 32, 8, 2, torch.bfloat16, torch.bfloat16, False, 32), + ], +) +@torch.inference_mode() +def test_segmented_attention_ragged_unified_graph_replay( + monkeypatch, dim, page, hq, hk, dtype, kv_dtype, byte_cache, force_splits +): + """Unified-cache segmented prefill follows metadata on graph replay.""" + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not ( + on_gfx12x() if kv_dtype.itemsize == 1 else on_gfx1x() + ): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + if force_splits is not None: + original = dispatcher.select_segmented_config + monkeypatch.setattr( + dispatcher, + "select_segmented_config", + lambda *args: dict(original(*args), splits=force_splits), + ) + query_lengths = [0, 1, 2, 7, 33, 129] + contexts = [0, 67, 0, page - 1, page + 1, 33] + device = torch.device("cuda:0") + total = sum(query_lengths) + q = torch.randn(total, hq + 1, dim, device=device, dtype=dtype)[:, :hq] + k = torch.randn(total, hk + 1, dim, device=device, dtype=dtype)[:, :hk] + v = torch.randn_like(k) + blocks_per_seq = triton.cdiv(max(map(sum, zip(query_lengths, contexts))), page) + num_blocks = len(contexts) * blocks_per_seq + page_elements = page * hk * dim + backing = torch.empty(num_blocks * 2 * page_elements, device=device, dtype=kv_dtype) + cache_shape = (num_blocks, page, hk, dim) + strides = (2 * page_elements, hk * dim, dim, 1) + kc = torch.as_strided(backing, cache_shape, strides) + vc = torch.as_strided(backing, cache_shape, strides, page_elements) + ks = torch.tensor(0.13 if kv_dtype.itemsize == 1 else 1.0, device=device) + vs = torch.tensor(0.27 if kv_dtype.itemsize == 1 else 1.0, device=device) + kc.copy_((torch.randn(cache_shape, device=device).float() * 0.25 / ks).to(kv_dtype)) + vc.copy_((torch.randn(cache_shape, device=device).float() * 0.25 / vs).to(kv_dtype)) + table = torch.randperm(num_blocks, device=device, dtype=torch.int32).view( + len(contexts), blocks_per_seq + ) + host_table = table.cpu().tolist() + starts = torch.empty(len(contexts) + 1, device=device, dtype=torch.int32) + lengths = torch.empty(len(contexts), device=device, dtype=torch.int32) + output = torch.full((total, hq + 1, dim), 17.0, device=device, dtype=dtype)[:, :hq] + + def metadata(qlens): + starts.copy_( + torch.tensor( + [0, *torch.tensor(qlens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + ) + lengths.copy_( + torch.tensor( + [c + n for c, n in zip(contexts, qlens)], + device=device, + dtype=torch.int32, + ) + ) + first = 0 + for seq, (nq, context) in enumerate(zip(qlens, contexts)): + for local in range(nq): + position = context + local + physical = host_table[seq][position // page] + offset = position % page + kc[physical, offset].copy_((k[first + local].float() / ks).to(kv_dtype)) + vc[physical, offset].copy_((v[first + local].float() / vs).to(kv_dtype)) + first += nq + + def run(): + dispatcher.segmented_attention( + query=q, + key=k, + value=v, + output=output, + kv_cache_dtype="fp8" if kv_dtype.itemsize == 1 else "auto", + key_cache=kc.view(torch.uint8) if byte_cache else kc, + value_cache=vc.view(torch.uint8) if byte_cache else vc, + block_table=table, + query_start_loc=starts, + seq_lens=lengths, + max_seq_len=max(contexts) + 129, + max_query_len=129, + k_scale=ks, + v_scale=vs, + sm_scale=dim**-0.5, + ) + + def check(qlens): + first = 0 + for seq, (nq, context) in enumerate(zip(qlens, contexts)): + if nq == 0: + first += nq + continue + physical = table[seq].long() + full_k = ( + kc[physical].reshape(-1, hk, dim)[: context + nq].float() * ks + ).repeat_interleave(hq // hk, 1) + full_v = ( + vc[physical].reshape(-1, hk, dim)[: context + nq].float() * vs + ).repeat_interleave(hq // hk, 1) + logits = torch.einsum( + "qhd,khd->hqk", q[first : first + nq].float(), full_k + ) / math.sqrt(dim) + mask = ( + torch.arange(context + nq, device=device)[None, :] + > context + torch.arange(nq, device=device)[:, None] + ) + logits.masked_fill_(mask[None], -float("inf")) + ref = torch.einsum("hqk,khd->qhd", logits.softmax(-1), full_v) + actual = output[first : first + nq].float() + assert torch.isfinite(actual).all() + relative = (actual - ref).norm(dim=-1) / ref.norm(dim=-1).clamp_min(1e-6) + assert relative.max().item() < 0.01 + first += nq + + metadata(query_lengths) + run() + torch.accelerator.synchronize() + check(query_lengths) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + replay_lengths = [1, 0, 7, 2, 34, 128] + q.mul_(0.75) + k.add_(0.25) + v.add_(0.5) + metadata(replay_lengths) + output.fill_(17.0) + graph.replay() + torch.accelerator.synchronize() + check(replay_lengths) + + +@pytest.mark.parametrize( + "fp8,force_splits", + [(False, None), (True, None), (False, 4), (True, 4)], +) +@torch.inference_mode() +def test_segmented_ragged_prefill_matches_dense_reference(fp8, force_splits): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + case = _make_paged_attention_case( + [1, 2, 7, 129], + [31, 64, 97, 128], + num_heads=12, + num_kv_heads=2, + head_size=256, + block_size=32, + fp8=fp8, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, max_query_len=129, force_splits=force_splits) + + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@torch.inference_mode() +def test_rocm_segmented_attn_layout_cache_update(fp8): + from types import SimpleNamespace + + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + cache_dtype = torch.uint8 if fp8 else torch.bfloat16 + kv_cache_dtype = "fp8" if fp8 else "auto" + impl = RocmSegmentedAttentionImpl( + 12, + 256, + 256**-0.5, + 2, + None, + None, + kv_cache_dtype, + attn_type=AttentionType.DECODER, + ) + cache = torch.zeros(3, 2, 32, 512, device="cuda", dtype=cache_dtype) + key = torch.randn(5, 2, 256, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + slots = torch.tensor([0, 7, 32, 65, 95], device="cuda", dtype=torch.int64) + k_scale = torch.tensor(0.125 if fp8 else 1.0, device="cuda") + v_scale = torch.tensor(0.25 if fp8 else 1.0, device="cuda") + layer = SimpleNamespace(_k_scale=k_scale, _v_scale=v_scale) + + impl.do_kv_cache_update(layer, key, value, cache, slots) + key_cache, value_cache = impl._split_kv_cache(cache) + if fp8: + key_cache = key_cache.view(impl.fp8_dtype) + value_cache = value_cache.view(impl.fp8_dtype) + cached_key = torch.stack( + [key_cache[int(slot) // 32, int(slot) % 32] for slot in slots] + ) + cached_value = torch.stack( + [value_cache[int(slot) // 32, int(slot) % 32] for slot in slots] + ) + expected_key = (key.float() / k_scale).to(key_cache.dtype) + expected_value = (value.float() / v_scale).to(value_cache.dtype) + torch.testing.assert_close(cached_key, expected_key, rtol=0, atol=0) + torch.testing.assert_close(cached_value, expected_value, rtol=0, atol=0) + + +@torch.inference_mode() +def test_segmented_mixed_causality_matches_dense_reference(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x unified attention per-request causal fallback") + + query_lens = [2, 3] + context_lens = [5, 6] + args = dict(num_heads=8, num_kv_heads=2, head_size=128, block_size=32) + case = _make_paged_attention_case( + query_lens, context_lens, fp8=False, causal=True, **args + ) + noncausal_case = _make_paged_attention_case( + query_lens, context_lens, fp8=False, causal=False, **args + ) + reference = torch.cat( + ( + case["reference"][: query_lens[0]], + noncausal_case["reference"][query_lens[0] :], + ) + ) + causal = torch.tensor([True, False], device=case["query"].device) + output = torch.empty_like(case["query"]) + dispatcher.segmented_attention( + query=case["query"], + key=case["key"], + value=case["value"], + output=output, + kv_cache_dtype="auto", + key_cache=case["key_cache"], + value_cache=case["value_cache"], + block_table=case["block_table"], + query_start_loc=case["starts"], + seq_lens=case["seq_lens"], + max_seq_len=case["max_seq_len"], + max_query_len=max(query_lens), + k_scale=case["k_scale"], + v_scale=case["v_scale"], + sm_scale=128**-0.5, + causal=causal, + ) + + relative = (output.float() - reference).norm(dim=-1) / reference.norm( + dim=-1 + ).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("window_size", [1, 2048]) +@torch.inference_mode() +def test_segmented_sliding_window_matches_dense_reference(fp8, causal, window_size): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + case = _make_paged_attention_case( + [8], + [4096], + num_heads=32, + num_kv_heads=8, + head_size=128, + block_size=32, + fp8=fp8, + causal=causal, + sliding_window=window_size, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case( + case, output, max_query_len=8, sliding_window=window_size - 1, causal=causal + ) + + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("window_size", [0, 128]) +@pytest.mark.parametrize("head_size", [64, 128]) +@torch.inference_mode() +def test_segmented_sinks_match_dense_reference(fp8, window_size, head_size): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + sinks = torch.linspace( + -1.0, + 8.0, + 16, + device="cuda:0", + dtype=torch.bfloat16 if head_size == 64 else torch.float32, + ) + case = _make_paged_attention_case( + [1, 7], + [127, 130], + num_heads=16, + num_kv_heads=2, + head_size=head_size, + block_size=32, + fp8=fp8, + sliding_window=window_size, + sinks=sinks, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case( + case, + output, + max_query_len=7, + sliding_window=window_size - 1 if window_size else -1, + sinks=sinks, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("window_size", [0, 128]) +@pytest.mark.parametrize("query_len", [1, 8]) +@torch.inference_mode() +def test_segmented_sink_split_reduce_matches_reference(fp8, window_size, query_len): + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + from vllm.v1.attention.ops.segmented_attention import ( + run_segmented_attention, + select_segmented_config, + ) + + if not current_platform.is_rocm() or not (on_gfx12x() if fp8 else on_gfx1x()): + pytest.skip("gfx1x segmented prefill (FP8 requires gfx12)") + sinks = torch.linspace(-1.0, 8.0, 16, device="cuda:0") + case = _make_paged_attention_case( + [query_len], + [8192 - query_len], + num_heads=16, + num_kv_heads=2, + head_size=64, + block_size=64, + fp8=fp8, + sliding_window=window_size, + sinks=sinks, + ) + attention_span = min(8192, window_size + query_len) if window_size else 8192 + config = dict( + select_segmented_config(1, query_len, attention_span, 16, 2, 64, fp8), + splits=4, + ) + output = torch.empty_like(case["query"]) + run_segmented_attention( + case["query"], + output, + case["key_cache"], + case["value_cache"], + case["block_table"], + case["starts"], + case["seq_lens"], + query_len, + 8192, + case["k_scale"], + case["v_scale"], + 64**-0.5, + sliding_window=window_size - 1 if window_size else -1, + sinks=sinks, + config=config, + ) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize("query_len,context_len", [(1, 8191), (8, 32760), (1, 131071)]) +@torch.inference_mode() +def test_segmented_long_context_fp8_sinks_match_dense_reference(query_len, context_len): + from vllm.platforms.rocm import on_gfx12x + + if not current_platform.is_rocm() or not on_gfx12x(): + pytest.skip("FP8 KV requires gfx12") + sinks = torch.linspace(-1.0, 1.0, 16, device="cuda:0", dtype=torch.bfloat16) + case = _make_paged_attention_case( + [query_len], + [context_len], + num_heads=16, + num_kv_heads=2, + head_size=64, + block_size=64, + fp8=True, + sinks=sinks, + ) + output = torch.empty_like(case["query"]) + _run_segmented_case(case, output, max_query_len=query_len, sinks=sinks) + relative = (output.float() - case["reference"]).norm(dim=-1) / case[ + "reference" + ].norm(dim=-1).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@pytest.mark.parametrize( + "feature", + [ + "softcap", + "sliding_softcap", + "output_scale", + ], +) +@torch.inference_mode() +def test_segmented_softcap_and_output_scaling_match_dense_reference(feature): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops import segmented_attention as dispatcher + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x unified attention feature fallback") + + num_heads = 8 + configured_sliding_window = 16 if feature.startswith("sliding") else 0 + sliding_window = configured_sliding_window - 1 if configured_sliding_window else -1 + softcap = 5.0 if "softcap" in feature else 0.0 + case = _make_paged_attention_case( + [7, 2], + [40, 35], + num_heads=num_heads, + num_kv_heads=2, + head_size=128, + block_size=32, + fp8=False, + sliding_window=configured_sliding_window, + softcap=softcap, + sinks=None, + ) + output_scale = None + if feature == "output_scale": + output_scale = torch.tensor(0.5, device="cuda:0") + output = torch.empty_like(case["query"], dtype=current_platform.fp8_dtype()) + else: + output = torch.empty_like(case["query"]) + + dispatcher.segmented_attention( + query=case["query"], + key=case["key"], + value=case["value"], + output=output, + kv_cache_dtype="auto", + key_cache=case["key_cache"], + value_cache=case["value_cache"], + block_table=case["block_table"], + query_start_loc=case["starts"], + seq_lens=case["seq_lens"], + max_seq_len=case["max_seq_len"], + max_query_len=7, + k_scale=case["k_scale"], + v_scale=case["v_scale"], + sm_scale=128**-0.5, + sliding_window=sliding_window, + softcap=softcap, + output_scale=output_scale, + sinks=None, + ) + + actual = output.float() + if output_scale is not None: + actual *= output_scale + torch.testing.assert_close(actual, case["reference"], atol=0.2, rtol=0.2) + else: + relative = (actual - case["reference"]).norm(dim=-1) / case["reference"].norm( + dim=-1 + ).clamp_min(1e-6) + assert torch.isfinite(output).all() and relative.max().item() < 0.01 + + +@torch.inference_mode() +def test_segmented_attention_cache_offsets_cross_int32_boundary(): + """A small logical prefix can live beyond 2**31 elements in a strided cache.""" + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.ops.segmented_attention import run_segmented_attention + + if not current_platform.is_rocm() or not on_gfx1x(): + pytest.skip("gfx1x segmented prefill") + torch.accelerator.empty_cache() + free, _ = torch.accelerator.get_memory_info() + if free < 10 * 2**30: + pytest.skip("The address-boundary fixture needs 10 GiB free VRAM") + device = torch.device("cuda:0") + stride = 2**30 + 4096 + kc = torch.empty_strided( + (3, 32, 1, 128), + (stride, 128, 128, 1), + device=device, + dtype=torch.bfloat16, + ) + vc = torch.empty_strided( + (3, 32, 1, 128), + (stride, 128, 128, 1), + device=device, + dtype=torch.bfloat16, + ) + dense_k = torch.randn(2, 32, 1, 128, device=device, dtype=torch.bfloat16) + dense_v = torch.randn_like(dense_k) + for logical, physical in enumerate((2, 1)): + kc[physical].copy_(dense_k[logical]) + vc[physical].copy_(dense_v[logical]) + q = torch.randn(2, 4, 128, device=device, dtype=torch.bfloat16) + out = torch.empty_like(q) + table = torch.tensor([[2, 1]], device=device, dtype=torch.int32) + starts = torch.tensor([0, 2], device=device, dtype=torch.int32) + lengths = torch.tensor([35], device=device, dtype=torch.int32) + one = torch.ones((), device=device) + run_segmented_attention( + q, out, kc, vc, table, starts, lengths, 2, 35, one, one, 128**-0.5 + ) + full_k = dense_k.flatten(0, 1)[:35].float().repeat_interleave(4, 1) + full_v = dense_v.flatten(0, 1)[:35].float().repeat_interleave(4, 1) + scores = torch.einsum("qhd,khd->hqk", q.float(), full_k) / math.sqrt(128) + scores[:, 0, -1] = -float("inf") + expected = torch.einsum("hqk,khd->qhd", scores.softmax(-1), full_v) + error = ((out.float() - expected).norm(dim=-1) / expected.norm(dim=-1)).max() + assert torch.isfinite(out).all() and error < 0.01 + + +@pytest.mark.parametrize("dim,hq,hk", [(128, 16, 1), (256, 12, 2)]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_segmented_attention_reservation_covers_ragged_query_caps( + monkeypatch, dim, hq, hk, fp8 +): + """Every supported query bucket fits the startup buffer after locking.""" + from vllm.v1.attention.ops import segmented_attention as segmented + + sizes = segmented.segmented_attention_workspace_size( + 32, + hq, + hk, + dim, + 65536, + fp8=fp8, + ) + workspace = tuple(torch.empty(size) for size in sizes) + query_lengths = { + query_len + for capacity in segmented._query_capacity_buckets() + for query_len in (max(1, capacity // 2 + 1), capacity) + } + for batch in range(1, 33): + for query_len in query_lengths: + qcap = segmented.segmented_query_capacity(query_len) + cfg = segmented.select_segmented_config( + batch, query_len, 65536, hq, hk, dim, fp8 + ) + shapes = segmented.segmented_workspace_shapes( + batch, qcap, hq, hk, dim, cfg["splits"] + ) + if shapes is not None: + assert math.prod(shapes[0]) <= workspace[0].numel() + assert math.prod(shapes[1]) <= workspace[1].numel() + + +def test_segmented_attention_reservation_respects_scheduler_token_limit( + monkeypatch, +) -> None: + """Workspace planning excludes batch/query pairs the scheduler cannot form.""" + from vllm.v1.attention.ops import segmented_attention as segmented + + calls = [] + + def record_config(batch, query_len, *_args): + calls.append((batch, query_len)) + return {"splits": 1} + + monkeypatch.setattr(segmented, "select_segmented_config", record_config) + segmented.segmented_attention_workspace_size.cache_clear() + segmented.segmented_attention_workspace_size( + 32, + 16, + 2, + 128, + 65536, + max_tokens=8, + ) + + assert calls + assert all(batch + query_len - 1 <= 8 for batch, query_len in calls) + assert (8, 1) in calls + assert (1, 8) in calls + + +def test_segmented_attention_query_capacity_buckets() -> None: + from vllm.v1.attention.ops.segmented_attention import ( + MAX_QUERY_LEN, + segmented_query_capacity, + ) + + assert [ + segmented_query_capacity(query_len) + for query_len in (1, 2, 3, 33, 129, 1025, 2049, MAX_QUERY_LEN) + ] == [1, 2, 4, 64, 256, 2048, 4096, 8192] + with pytest.raises(ValueError, match="Query length"): + segmented_query_capacity(MAX_QUERY_LEN + 1) + + +def test_segmented_scratch_is_stable_and_isolated_across_runtime_lanes(monkeypatch): + """Draft/DBO scratch never aliases target scratch or generic workspace.""" + import vllm.v1.worker.workspace as workspace + from vllm.v1.attention.ops import segmented_attention as segmented + + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + monkeypatch.setattr(segmented, "is_workspace_manager_initialized", lambda: True) + monkeypatch.setattr(segmented, "current_workspace_manager", lambda: manager) + buffers = {} + for ubatch in range(2): + active_ubatch[0] = ubatch + for lane in range(2): + with workspace.use_workspace_lane(lane): + buffers[ubatch, lane] = segmented.get_segmented_attention_workspace( + torch.device("cpu"), (1024, 16) + ) + manager.get_simultaneous(((2048,), torch.float32)) + manager.lock() + pointers: set[int] = set() + for (ubatch, lane), expected in buffers.items(): + active_ubatch[0] = ubatch + with workspace.use_workspace_lane(lane): + actual = segmented.get_segmented_attention_workspace( + torch.device("cpu"), (1024, 16) + ) + assert actual is expected + pointers.update(t.data_ptr() for t in actual) + assert ( + manager.get_simultaneous(((1,), torch.float32))[0].data_ptr() + not in pointers + ) + assert len(pointers) == 8 diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index b99a33f9f7a9..067d1dbc8bbd 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -61,6 +61,216 @@ def test_aiter_unified_attention_uses_dedicated_metadata_builder(): ) +def test_segmented_attention_uses_dedicated_backend_components(): + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionBackend, + RocmSegmentedAttentionImpl, + RocmSegmentedAttentionMetadataBuilder, + ) + + assert ( + RocmSegmentedAttentionBackend.get_builder_cls() + is RocmSegmentedAttentionMetadataBuilder + ) + assert RocmSegmentedAttentionBackend.get_impl_cls() is RocmSegmentedAttentionImpl + assert RocmSegmentedAttentionBackend.get_name() == "ROCM_SEGMENTED_ATTN" + assert RocmSegmentedAttentionBackend.supports_non_causal() + assert RocmSegmentedAttentionBackend.supports_sliding_window() + assert RocmSegmentedAttentionBackend.supports_sink() + assert not RocmSegmentedAttentionBackend.supports_mm_prefix() + + +@pytest.mark.parametrize( + "sliding_window,logits_soft_cap,has_sinks", + [(4096, None, False), (None, 50.0, False), (None, None, True)], +) +def test_segmented_attention_admits_unified_fallback_features( + monkeypatch, sliding_window, logits_soft_cap, has_sinks +): + from vllm.platforms import rocm + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setattr(rocm, "on_gfx12x", lambda: True) + sinks = torch.zeros(8) if has_sinks else None + + impl = RocmSegmentedAttentionImpl( + num_heads=8, + head_size=128, + scale=128**-0.5, + num_kv_heads=2, + alibi_slopes=None, + sliding_window=sliding_window, + kv_cache_dtype="auto", + logits_soft_cap=logits_soft_cap, + attn_type=AttentionType.DECODER, + sinks=sinks, + ) + + assert isinstance(impl, RocmSegmentedAttentionImpl) + + +@pytest.mark.parametrize( + "platform,rdna", + [("cuda", False), ("cpu", False), ("xpu", False), ("rocm", False), ("rocm", True)], +) +def test_segmented_attention_requires_rocm_rdna(monkeypatch, platform, rdna): + """Reject unsupported platforms before probing ROCm or constructing kernels.""" + from vllm.platforms import rocm + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backends import rocm_segmented_attn as segmented + + monkeypatch.setattr( + segmented, + "current_platform", + SimpleNamespace(is_rocm=lambda: platform == "rocm"), + ) + + def on_rdna(): + assert platform == "rocm", "Must not probe ROCm on another platform" + return rdna + + monkeypatch.setattr(rocm, "on_gfx1x", on_rdna) + reason = segmented.RocmSegmentedAttentionBackend.supports_combination( + 128, + torch.bfloat16, + "auto", + 32, + False, + False, + False, + False, + DeviceCapability(12, 0), + ) + if rdna: + assert reason is None + segmented.RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + else: + assert reason == "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm" + with pytest.raises(ValueError, match="requires AMD RDNA GPUs on ROCm"): + segmented.RocmSegmentedAttentionImpl( + 8, 128, 128**-0.5, 2, None, None, "auto" + ) + + +@pytest.mark.parametrize("explicit_cache", [False, True]) +def test_segmented_triton_cache_survives_standalone_compilation( + monkeypatch, tmp_path, explicit_cache +): + """Backend initialization preserves one cache across model compilations.""" + import os + + from vllm.compilation.compiler_interface import InductorStandaloneAdaptor + from vllm.platforms import rocm + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) + monkeypatch.delenv("TORCHINDUCTOR_CACHE_DIR", raising=False) + expected = tmp_path / "rocm_segmented_attention" / "triton" + if explicit_cache: + expected = tmp_path / "custom" + monkeypatch.setenv("TRITON_CACHE_DIR", str(expected)) + for model in ("target", "draft", "restart_target"): + RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + assert os.environ["TRITON_CACHE_DIR"] == str(expected) + assert expected.is_dir() + InductorStandaloneAdaptor("binary").initialize_cache(str(tmp_path / model)) + assert os.environ["TRITON_CACHE_DIR"] == str(expected) + + +def test_segmented_attention_is_opt_in(monkeypatch): + from vllm.platforms import rocm + from vllm.platforms.rocm import RocmPlatform, _get_backend_priorities + + assert AttentionBackendEnum.ROCM_SEGMENTED_ATTN not in _get_backend_priorities( + use_mla=False, use_sparse=False + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + monkeypatch.setattr(rocm, "on_gfx12x", lambda: True) + config = AttentionSelectorConfig( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=32, + ) + + path = RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN, + attn_selector_config=config, + ) + + assert path == AttentionBackendEnum.ROCM_SEGMENTED_ATTN.get_path() + + +def test_segmented_attention_rejects_mm_prefix_at_selection(monkeypatch): + from vllm.platforms import rocm + from vllm.platforms.rocm import RocmPlatform + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + config = AttentionSelectorConfig( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=32, + use_mm_prefix=True, + ) + + with pytest.raises(ValueError, match="multimodal token full attention"): + RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN, + attn_selector_config=config, + ) + + +def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch): + from vllm.config import KernelConfig + from vllm.platforms import rocm + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx1x", lambda: True) + impl = RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto") + impl._segmented_attention_config = SimpleNamespace( + kernel_config=KernelConfig(), + scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=3), + model_config=SimpleNamespace(max_model_len=16), + ) + query = torch.empty(3, 8, 128, dtype=torch.bfloat16) + key = torch.empty(3, 2, 128, dtype=torch.bfloat16) + value = torch.empty_like(key) + output = torch.empty_like(query) + kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16) + metadata = SimpleNamespace( + use_cascade=False, + num_actual_tokens=3, + block_table=torch.zeros(1, 1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 3], dtype=torch.int32), + seq_lens=torch.tensor([3], dtype=torch.int32), + max_seq_len=3, + max_query_len=3, + causal=True, + ) + layer = SimpleNamespace(_k_scale=torch.ones(()), _v_scale=torch.ones(())) + + with patch( + "vllm.v1.attention.backends.rocm_segmented_attn.segmented_attention" + ) as dispatch: + result = impl.forward(layer, query, key, value, kv_cache, metadata, output) + + assert result is output + dispatch.assert_called_once() + assert dispatch.call_args.kwargs["key_cache"].shape == (2, 16, 2, 128) + + def test_aiter_unified_attention_capture_preserves_query_start_locations(): from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( RocmAiterUnifiedAttentionMetadataBuilder, diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index df69f171493c..b8bfca082a4a 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -278,6 +278,9 @@ class KernelConfig: enable_flashinfer_autotune: bool = None # type: ignore[assignment] """If True, run FlashInfer autotuning during kernel warmup.""" + enable_rocm_segmented_attn_autotune: bool = False + """If True, autotune ROCm segmented attention during kernel warmup on RDNA GPUs.""" + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated # to the shared JIT warmup infrastructure. # https://github.com/vllm-project/vllm/pull/47451 @@ -406,6 +409,7 @@ def compute_hash(self) -> str: "enable_cutedsl_warmup", "enable_jit_warmup", "enable_flashinfer_autotune", + "enable_rocm_segmented_attn_autotune", "ir_op_priority", # handled separately below } if self.linear_backend_per_quant is None: diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 3161cf22616f..44b7865daaa6 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -34,6 +34,7 @@ VocabParallelEmbedding, ) from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform from vllm.transformers_utils.config import set_default_rope_theta from vllm.transformers_utils.repo_utils import get_hf_file_bytes from vllm.v1.attention.backend import AttentionType @@ -269,6 +270,24 @@ def __init__( sinks=self.attention_sink_bias, ) self.causal = causal + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + + if on_gfx1x() and isinstance(self.attn.impl, RocmSegmentedAttentionImpl): + # Startup segmented tuning needs the draft's actual attention mode. + # Causality is normally supplied later through attention metadata. + self.attn.segmented_causal = causal + speculative_config = get_current_vllm_config().speculative_config + if ( + speculative_config is not None + and speculative_config.num_speculative_tokens is not None + ): + self.attn.segmented_query_limit = ( + speculative_config.num_speculative_tokens + 1 + ) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 228ce73803fa..55d5ea7a7df7 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -220,6 +220,13 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): if process_local_only: return + if current_platform.is_rocm(): + from vllm.model_executor.warmup.rocm_segmented_attn_autotune_warmup import ( + rocm_segmented_attn_autotune_warmup, + ) + + rocm_segmented_attn_autotune_warmup(worker) + flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) diff --git a/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py b/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py new file mode 100644 index 000000000000..0ebcc58c3cce --- /dev/null +++ b/vllm/model_executor/warmup/rocm_segmented_attn_autotune_warmup.py @@ -0,0 +1,1335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup-only, persistent launch tuning for segmented ROCm attention.""" + +from __future__ import annotations + +import gc +import hashlib +import itertools +import json +import math +import os +import statistics +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +from filelock import FileLock + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.v1.attention.ops import segmented_attention +from vllm.v1.attention.ops.segmented_attention import ( + MAX_QUERY_LEN, + _segmented_attention_reduce, + _segmented_attention_stage, + compile_segmented_attention, + run_segmented_attention, + segmented_query_capacity, + segmented_workspace_shapes, + select_segmented_config, +) + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_TABLES: dict[tuple, dict] = {} +_QUERY_BUCKETS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192) +_BATCH_BUCKETS = (1, 2, 4, 8, 16, 32) +_SEQUENCE_BUCKETS = (128, 512, 2048, 8192, 32768, 131072, 262144) +_COMPILE_WORKERS = 2 +_CONDITION_ROUNDS = 2 +_SCREEN_ROUNDS = 5 +_FINALIST_COUNT = 3 +_FINAL_ROUNDS = 7 +_MIN_PROMOTION_SPEEDUP = 1.02 + + +def _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window=-1, + causal=True, + has_sinks=False, +): + return ( + torch.device(device).index, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + + +def _normalized_config(config, batch, query_len, heads, dim): + qcap = segmented_query_capacity(query_len) + result = dict(config) + result.setdefault("waves_per_eu", 2) + result.setdefault("reduce_d", 64 if batch * qcap * heads < 64 else dim) + result.setdefault("reduce_warps", 4) + return result + + +def _candidate_configs(default, batch, query_len, heads, dim, seq_len=None): + """Search launch tiles and split counts within a bounded workspace.""" + base = _normalized_config(default, batch, query_len, heads, dim) + max_splits = base["splits"] + long_d64_decode = ( + dim == 64 + and batch == 1 + and query_len == 1 + and seq_len is not None + and seq_len >= 131072 + ) + if long_d64_decode: + max_splits = 512 + elif dim == 64 and query_len <= 512 and seq_len is not None and seq_len >= 4096: + max_splits = min(64, 2 * max_splits) + candidates: list[dict] = [] + seen = set() + + def add(**updates): + config = {**base, **updates} + if ( + config["splits"] > max_splits + or dim % config["bk"] + or config["reduce_d"] > dim + ): + return + identity = tuple(sorted(config.items())) + if identity not in seen: + seen.add(identity) + candidates.append(config) + + if long_d64_decode: + add() + for splits in (128, 256, 512, 64): + for bn, warps, stages in ( + (32, 1, 1), + (32, 2, 1), + (32, 2, 2), + (64, 1, 1), + (64, 2, 1), + (64, 2, 2), + ): + add( + bm=16, + bn=bn, + bk=64, + splits=splits, + warps=warps, + stages=stages, + ) + return candidates + + split_choices = [base["splits"]] + if max_splits > base["splits"]: + split_choices.append(max_splits) + for divisor in (2, 4): + split_choices.append(max(1, base["splits"] // divisor)) + split_choices.append(1) + for splits in split_choices: + add(splits=splits) + + tile_variants: tuple[tuple[int, int, int, int, int, int], ...] = ( + (16, 32, dim, 4, 1, 2), + (16, 64, min(128, dim), 4, 1, 2), + (32, 32, 64, 4, 2, 2), + (32, 64, 64, 4, 1, 2), + (64, 32, min(128, dim), 4, 1, 6), + (128, 32, min(128, dim), 8, 1, 6), + ) + if dim == 64: + tile_variants += ( + (16, 16, 64, 4, 1, 2), + (32, 16, 64, 4, 2, 2), + (64, 32, 64, 4, 1, 2), + (64, 64, 64, 4, 1, 2), + (128, 32, 64, 4, 1, 2), + (128, 64, 64, 4, 1, 2), + ) + for bm, bn, bk, warps, stages, waves in tile_variants: + add( + bm=bm, + bn=bn, + bk=bk, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + half_splits = max(1, base["splits"] // 2) + for bm, bn, bk, warps, stages, waves in tile_variants[:3]: + add( + bm=bm, + bn=bn, + bk=bk, + splits=half_splits, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + if dim == 64: + for splits in (max_splits, half_splits): + for bm, bn, bk, warps, stages, waves in tile_variants[2:]: + add( + bm=bm, + bn=bn, + bk=bk, + splits=splits, + warps=warps, + stages=stages, + waves_per_eu=waves, + ) + + add(waves_per_eu=6) + add(stages=2 if base["stages"] == 1 else 1) + add(reduce_d=64 if base["reduce_d"] != 64 else dim) + return candidates + + +def _identity( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + sliding_window, + causal, + max_query_len, + has_sinks=False, + physical_max_len=None, +): + properties = torch.cuda.get_device_properties(device) + from triton._C.libtriton import get_cache_invalidating_env_vars + + source = Path(__file__) + kernel = Path(segmented_attention.__file__) + return { + "schema": 5, + "gpu": properties.name, + "arch": properties.gcnArchName, + "compute_units": properties.multi_processor_count, + "torch": torch.__version__, + "hip": torch.version.hip, + "compiler_env": get_cache_invalidating_env_vars(), + "kernel_sha256": hashlib.sha256(kernel.read_bytes()).hexdigest(), + "tuner_sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + "dtype": str(dtype), + "kv_dtype": str(kv_dtype), + "heads": heads, + "kv_heads": kv_heads, + "dim": dim, + "page": page, + "scale": scale, + "max_tokens": max_tokens, + "max_len": max_len, + "max_seqs": max_seqs, + "sliding_window": sliding_window, + "causal": causal, + "max_query_len": max_query_len, + "has_sinks": has_sinks, + "physical_max_len": physical_max_len, + "query_buckets": list(_QUERY_BUCKETS), + "batch_buckets": list(_BATCH_BUCKETS), + "sequence_buckets": list(_SEQUENCE_BUCKETS), + "compile_workers": _COMPILE_WORKERS, + "condition_rounds": _CONDITION_ROUNDS, + "screen_rounds": _SCREEN_ROUNDS, + "finalist_count": _FINALIST_COUNT, + "final_rounds": _FINAL_ROUNDS, + "min_promotion_speedup": _MIN_PROMOTION_SPEEDUP, + } + + +def _memory_budget(device): + free, total = torch.accelerator.get_memory_info(device) + available = int(min(free // 2, total // 4)) + return 1 << (available.bit_length() - 1) if available else 0 + + +def _query_lengths(batch, query_len, max_tokens): + if batch * query_len <= max_tokens: + return [query_len] * batch + return [min(query_len, max_tokens - batch + 1)] + [1] * (batch - 1) + + +def _scratch_bytes( + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_max_len=None, +): + batch, query_len, seq_len = workload + queries = _query_lengths(batch, query_len, max_tokens) + tokens = sum(queries) + scratch_seq_len = ( + max(seq_len, physical_max_len) + if physical_max_len is not None and batch == 1 and query_len <= 8 + else seq_len + ) + blocks = batch * math.ceil(scratch_seq_len / page) + cache = 2 * blocks * page * kv_heads * dim * kv_dtype.itemsize + tensors = tokens * dim * (4 * heads + 2 * kv_heads) * dtype.itemsize + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + candidates = _candidate_configs( + default, batch, query_len, heads, dim, seq_len=seq_len + ) + shapes = segmented_workspace_shapes( + batch, + segmented_query_capacity(query_len), + heads, + kv_heads, + dim, + max(config["splits"] for config in candidates), + ) + workspace = 0 if shapes is None else sum(math.prod(shape) * 4 for shape in shapes) + return cache + tensors + workspace + 384 * 1024**2 + + +def _workloads( + max_tokens, + max_len, + max_seqs, + *, + memory_budget_bytes=None, + dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + heads=8, + kv_heads=1, + dim=128, + page=16, + cache_layouts=(), + cache_budget_bytes=None, + max_query_len=MAX_QUERY_LEN, + physical_max_len=None, +): + query_limit = min(max_tokens, max_len, max_query_len, MAX_QUERY_LEN) + queries = {q for q in _QUERY_BUCKETS if q <= query_limit} + if query_limit: + queries.add(query_limit) + for query_len in sorted(queries): + batch_limit = min(max_seqs, max_tokens - query_len + 1) + if batch_limit < 1: + continue + if query_len <= 2: + batches = {min(batch, batch_limit) for batch in _BATCH_BUCKETS} + else: + batches = {1, min(4, batch_limit), batch_limit} + if query_len <= 2: + sequences = { + seq for seq in _SEQUENCE_BUCKETS if query_len <= seq <= max_len + } + sequences.update((query_len, max_len)) + else: + sequences = { + query_len, + max(query_len, min(max_len, 2048)), + max(query_len, min(max_len, 8192)), + max_len, + } + # Keep long-context verification and prefill buckets even when + # the model's maximum length is pruned by the KV memory budget. + sequences.update( + seq for seq in (32768, 131072) if query_len <= seq <= max_len + ) + for batch, seq_len in itertools.product(sorted(batches), sorted(sequences)): + workload = (batch, query_len, seq_len) + if ( + memory_budget_bytes is not None + and _scratch_bytes( + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_max_len, + ) + > memory_budget_bytes + ): + continue + if cache_budget_bytes is not None: + cache_bytes = batch * sum( + math.ceil(seq_len / block) * size if block else size + for block, size in cache_layouts + ) + if cache_bytes > cache_budget_bytes: + continue + yield workload + + +def _make_inputs( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_seq_len=None, +): + batch, query_len, seq_len = workload + physical_seq_len = max(seq_len, physical_seq_len or seq_len) + generator = torch.Generator(device=device).manual_seed(1234) + + def randn(*shape): + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + queries = _query_lengths(batch, query_len, max_tokens) + tokens = sum(queries) + q = randn(tokens, heads, dim) + k = randn(tokens, kv_heads, dim) + v = randn(tokens, kv_heads, dim) + blocks_per_seq = math.ceil(seq_len / page) + physical_blocks_per_seq = math.ceil(physical_seq_len / page) + cache_blocks_per_seq = ( + physical_blocks_per_seq if batch == 1 and query_len <= 8 else blocks_per_seq + ) + cache_shape = (batch * cache_blocks_per_seq, 2, page, kv_heads, dim) + scale_value = 0.125 if kv_dtype.itemsize == 1 else 1.0 + cache = (randn(*cache_shape).float() / scale_value).to(kv_dtype) + kc, vc = cache[:, 0], cache[:, 1] + table = ( + torch.arange(physical_blocks_per_seq, device=device, dtype=torch.int32)[None, :] + .remainder(cache_blocks_per_seq) + .expand(batch, -1) + .contiguous() + ) + table += ( + torch.arange(batch, device=device, dtype=torch.int32)[:, None] + * cache_blocks_per_seq + ) + starts = torch.tensor( + [0, *itertools.accumulate(queries)], device=device, dtype=torch.int32 + ) + lengths = torch.full((batch,), physical_seq_len, device=device, dtype=torch.int32) + scale = torch.full((), scale_value, device=device, dtype=torch.float32) + return q, k, v, kc, vc, table, starts, lengths, scale + + +def _config_key(config): + return tuple(sorted(config.items())) + + +def _rotated(configs, offset): + offset %= len(configs) + return configs[offset:] + configs[:offset] + + +def _measure_configs(configs, run, device, eviction, rounds, graph_calls=0): + """Interleave candidates so clock and thermal drift affect them evenly.""" + stream = torch.cuda.current_stream(device) + samples: dict[tuple, list[float]] = {_config_key(config): [] for config in configs} + graphs = {} + if graph_calls: + for config in configs: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(graph_calls): + run(config) + graphs[_config_key(config)] = graph + for graph in graphs.values(): + graph.replay() + torch.accelerator.synchronize(device) + for round_index in range(rounds): + events = [] + order = _rotated(configs, round_index) + if round_index % 2: + order = list(reversed(order)) + for config in order: + if not graph_calls: + eviction.zero_() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + if graph_calls: + graphs[_config_key(config)].replay() + else: + run(config) + end.record(stream) + events.append((_config_key(config), start, end)) + events[-1][2].synchronize() + for key, start, end in events: + samples[key].append(start.elapsed_time(end) * 1000 / max(1, graph_calls)) + return samples + + +def _condition_configs(configs, run, device): + for round_index in range(_CONDITION_ROUNDS): + for config in _rotated(configs, round_index): + run(config) + torch.accelerator.synchronize(device) + + +def _select_tuned_winner( + default, finalists, samples, min_promotion_speedup=_MIN_PROMOTION_SPEEDUP +): + """Return a finalist only when it reliably beats the static incumbent.""" + default_key = _config_key(default) + default_samples = samples[default_key] + comparisons = [] + for config in finalists: + key = _config_key(config) + config_samples = samples[key] + paired_speedup = statistics.median( + baseline / candidate + for baseline, candidate in zip(default_samples, config_samples, strict=True) + ) + comparisons.append( + { + "config": config, + "us": statistics.median(config_samples), + "paired_speedup_vs_default": paired_speedup, + } + ) + challenger = max( + comparisons, key=lambda result: result["paired_speedup_vs_default"] + ) + if challenger["paired_speedup_vs_default"] >= min_promotion_speedup: + return challenger, comparisons + return ( + next(result for result in comparisons if result["config"] == default), + comparisons, + ) + + +def _precompile_configs( + configs, + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + scale, + workspace, + sliding_window, + causal, + sinks=None, +): + workers = min(_COMPILE_WORKERS, len(configs)) + with ( + ThreadPoolExecutor(max_workers=workers) as executor, + triton.AsyncCompileMode(executor, ignore_errors=True), + ): + for config in configs: + compile_segmented_attention( + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + kv_scale, + scale, + config, + workspace, + sliding_window=sliding_window, + causal=causal, + sinks=sinks, + ) + + +def _tune_workload( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + workload, + *, + sliding_window=-1, + causal=True, + has_sinks=False, + physical_seq_len=None, +): + batch, query_len, seq_len = workload + q, _, _, kc, vc, table, starts, lengths, kv_scale = _make_inputs( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + max_tokens, + workload, + physical_seq_len, + ) + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + configs = _candidate_configs(default, batch, query_len, heads, dim, seq_len=seq_len) + sinks = ( + torch.linspace(-1, 1, heads, dtype=torch.float32, device=device) + if has_sinks + else None + ) + qcap = segmented_query_capacity(query_len) + shapes = segmented_workspace_shapes( + batch, qcap, heads, kv_heads, dim, max(c["splits"] for c in configs) + ) + workspace = ( + None + if shapes is None + else tuple( + torch.empty(shape, dtype=torch.float32, device=device) for shape in shapes + ) + ) + reference = torch.empty_like(q) + output = torch.empty_like(q) + + def run(config, out=output): + run_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + query_len, + seq_len, + kv_scale, + kv_scale, + scale, + sliding_window=sliding_window, + causal=causal, + config=config, + workspace=workspace, + sinks=sinks, + ) + + _precompile_configs( + configs, + q, + output, + kc, + vc, + table, + starts, + lengths, + query_len, + kv_scale, + scale, + workspace, + sliding_window, + causal, + sinks, + ) + incumbent = configs[0] + reference.fill_(float("nan")) + run(incumbent, reference) + torch.accelerator.synchronize(device) + if not torch.isfinite(reference).all(): + raise AssertionError("Default segmented attention output is not finite") + reference_norm = reference.float().norm().clamp_min(1e-6) + eviction = torch.empty(256 * 1024**2, dtype=torch.int8, device=device) + results = [] + valid_configs = [] + for config in configs: + try: + output.fill_(float("nan")) + run(config) + torch.accelerator.synchronize(device) + torch.testing.assert_close(output, reference, atol=0.01, rtol=0.01) + relative_l2 = ( + (output.float() - reference.float()).norm() / reference_norm + ).item() + if relative_l2 > 0.005: + raise AssertionError(f"Relative L2 error {relative_l2} exceeds 0.5%") + results.append({"config": config, "relative_l2": relative_l2}) + valid_configs.append(config) + except ( + torch.OutOfMemoryError, + triton.OutOfResources, + triton.CompilationError, + AssertionError, + RuntimeError, + ) as error: + logger.warning("Rejected segmented attention config %s: %s", config, error) + if not results: + raise RuntimeError("No valid segmented attention launch configurations") + if incumbent not in valid_configs: + raise RuntimeError("Static segmented attention configuration was rejected") + + _condition_configs(valid_configs, run, device) + graph_calls = 5 if dim == 64 and has_sinks else 0 + screen_samples = _measure_configs( + valid_configs, run, device, eviction, _SCREEN_ROUNDS, graph_calls + ) + by_key = {_config_key(result["config"]): result for result in results} + for key, samples in screen_samples.items(): + by_key[key]["screen_us"] = statistics.median(samples) + by_key[key]["screen_samples_us"] = samples + + challengers = sorted( + (config for config in valid_configs if config != incumbent), + key=lambda config: statistics.median(screen_samples[_config_key(config)]), + ) + finalists = [incumbent, *challengers[: _FINALIST_COUNT - 1]] + _condition_configs(finalists, run, device) + final_samples = _measure_configs( + finalists, run, device, eviction, _FINAL_ROUNDS, graph_calls + ) + min_promotion_speedup = ( + 1.10 + if dim == 64 and has_sinks and sliding_window >= 0 + else _MIN_PROMOTION_SPEEDUP + ) + winner, comparisons = _select_tuned_winner( + incumbent, finalists, final_samples, min_promotion_speedup + ) + for comparison in comparisons: + result = by_key[_config_key(comparison["config"])] + result["final_us"] = comparison["us"] + result["final_samples_us"] = final_samples[_config_key(comparison["config"])] + result["paired_speedup_vs_default"] = comparison["paired_speedup_vs_default"] + winner_result = by_key[_config_key(winner["config"])] + logger.info( + "ROCm segmented attention autotune: B=%d Q=%d seq=%d best=%s " + "%.2f us speedup=%.3fx%s", + batch, + query_len, + seq_len, + winner["config"], + winner["us"], + winner["paired_speedup_vs_default"], + " (static retained)" if winner["config"] == incumbent else "", + ) + return { + "workload": list(workload), + "query_lengths": _query_lengths(batch, query_len, max_tokens), + "default": default, + "best": winner["config"], + "selection": { + "incumbent": incumbent, + "finalists": finalists, + "min_promotion_speedup": min_promotion_speedup, + "winner_us": winner_result["final_us"], + "paired_speedup_vs_default": winner["paired_speedup_vs_default"], + }, + "results": results, + } + + +def _save(path, data): + with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as file: + temporary = Path(file.name) + try: + json.dump(data, file, indent=2) + file.flush() + os.fsync(file.fileno()) + except BaseException: + temporary.unlink(missing_ok=True) + raise + try: + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _load_records(path, identity, heads, kv_heads, dim, fp8): + """Load a canonical table and any crash-recovery TP shards.""" + records: dict[tuple, dict] = {} + merged_shards = [] + sources = [path, *sorted(path.parent.glob(f"{path.name}.tp*.part"))] + for source in sources: + if not source.exists(): + continue + try: + saved = json.loads(source.read_text()) + if saved["identity"] != identity or not all( + _valid_record(record, heads, kv_heads, dim, fp8) + for record in saved["records"] + ): + raise ValueError("identity or record validation failed") + records.update( + (tuple(record["workload"]), record) for record in saved["records"] + ) + if source != path: + merged_shards.append(source) + except (ValueError, KeyError, TypeError, AssertionError): + logger.warning( + "Ignoring invalid segmented attention tuning cache %s", source + ) + return {"identity": identity, "records": list(records.values())}, merged_shards + + +def _tp_context(): + """Return the initialized TP coordinator, or local-only tuning state.""" + from vllm.distributed.parallel_state import ( + get_tp_group, + model_parallel_is_initialized, + ) + + if not model_parallel_is_initialized(): + return None, 0, 1 + group = get_tp_group() + return group, group.rank_in_group, group.world_size + + +def _gather_tp(group, value): + if group is None or group.world_size == 1: + return [value] + gathered = [None] * group.world_size + torch.distributed.all_gather_object(gathered, value, group=group.cpu_group) + return gathered + + +def _workload_weight(workload, max_tokens, heads, kv_heads, dim, fp8): + """Estimate compile plus execution cost for TP workload balancing.""" + batch, query_len, seq_len = workload + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + candidates = len( + _candidate_configs(default, batch, query_len, heads, dim, seq_len=seq_len) + ) + query_tokens = sum(_query_lengths(batch, query_len, max_tokens)) + attention_work = query_tokens * seq_len * heads * dim + return candidates * (1 << 40) + attention_work + + +def _shard_workloads(workloads, world_size, max_tokens, heads, kv_heads, dim, fp8): + """Balance compile-affine ``(batch, query)`` groups across TP ranks.""" + assignments: list[list[tuple[int, tuple]]] = [[] for _ in range(world_size)] + loads = [0] * world_size + groups: dict[tuple, dict] = {} + for index, workload in enumerate(workloads): + compile_group = groups.setdefault(workload[:2], {"weight": 0, "items": []}) + compile_group["weight"] += _workload_weight( + workload, max_tokens, heads, kv_heads, dim, fp8 + ) + compile_group["items"].append((index, workload)) + for compile_group in sorted( + groups.values(), key=lambda item: item["weight"], reverse=True + ): + rank = min( + range(world_size), key=lambda candidate: (loads[candidate], candidate) + ) + assignments[rank].extend(compile_group["items"]) + loads[rank] += compile_group["weight"] + return [ + [workload for _, workload in sorted(assignment)] for assignment in assignments + ] + + +def _valid_record(record, heads, kv_heads, dim, fp8): + workload = record["workload"] + if len(workload) != 3 or any( + type(value) is not int or value < 1 for value in workload + ): + return False + batch, query_len, seq_len = workload + default = select_segmented_config( + batch, query_len, seq_len, heads, kv_heads, dim, fp8 + ) + return record["best"] in _candidate_configs( + default, batch, query_len, heads, dim, seq_len=seq_len + ) + + +@contextmanager +def _temporary_tuning_kernels(device): + """Release tuning-only HIP modules without invalidating existing graphs.""" + device_index = torch.device(device).index + if device_index is None: + device_index = torch.accelerator.current_device_index() + caches = [] + for kernel in (_segmented_attention_stage, _segmented_attention_reduce): + cache = kernel.device_caches[device_index][0] + caches.append((cache, set(cache))) + try: + yield + finally: + torch.accelerator.synchronize(device) + for cache, existing_keys in caches: + for key in cache.keys() - existing_keys: + del cache[key] + # CompiledKernel destruction unloads the HIP module. Keep pre-existing + # kernels alive, since an earlier graph capture may reference them. + gc.collect() + torch.accelerator.empty_cache() + + +@torch.inference_mode() +def warmup_segmented_attention( + device, + dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + *, + memory_budget_bytes=None, + cache_layouts=(), + cache_budget_bytes=None, + kv_dtype=torch.bfloat16, + sliding_window=-1, + causal=True, + max_query_len=MAX_QUERY_LEN, + has_sinks=False, + physical_max_len=None, +): + """Tune reachable segmented buckets before graph capture.""" + if ( + dtype not in (torch.bfloat16, torch.float16) + or kv_dtype + not in ( + torch.bfloat16, + torch.float16, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + or dim not in (64, 128, 256) + or kv_heads < 1 + or heads % kv_heads + or not 1 <= heads // kv_heads <= 16 + ): + return + key = _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + identity = _identity( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + max_len, + max_seqs, + sliding_window, + causal, + max_query_len, + has_sinks, + physical_max_len, + ) + if memory_budget_bytes is None: + memory_budget_bytes = _memory_budget(device) + workload_args = dict( + memory_budget_bytes=memory_budget_bytes, + dtype=dtype, + kv_dtype=kv_dtype, + heads=heads, + kv_heads=kv_heads, + dim=dim, + page=page, + cache_layouts=cache_layouts, + cache_budget_bytes=cache_budget_bytes, + max_query_len=max_query_len, + physical_max_len=( + physical_max_len if has_sinks and sliding_window >= 0 else None + ), + ) + workloads = list(_workloads(max_tokens, max_len, max_seqs, **workload_args)) + active = _TABLES.get(key, {}) + warmed = ( + {tuple(record["workload"]) for record in active.get("records", [])} + if active.get("identity") == identity + else set() + ) + if all(workload in warmed for workload in workloads): + return + digest = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + path = Path(envs.VLLM_CACHE_ROOT) / "rocm_segmented_attention" / f"{digest}.json" + path.parent.mkdir(parents=True, exist_ok=True) + group, tp_rank, tp_size = _tp_context() + if group is not None: + digests = _gather_tp(group, digest) + if any(peer != digest for peer in digests): + logger.warning( + "ROCm segmented attention TP ranks have different tuning " + "identities; tuning this rank independently" + ) + group, tp_rank, tp_size = None, 0, 1 + + unpruned = len( + list(_workloads(max_tokens, max_len, max_seqs, max_query_len=max_query_len)) + ) + start = time.monotonic() + lock = FileLock(str(path) + ".lock") if tp_rank == 0 else None + if lock is not None: + lock.acquire() + try: + if tp_rank == 0: + data, merged_shards = _load_records( + path, identity, heads, kv_heads, dim, kv_dtype.itemsize == 1 + ) + else: + data, merged_shards = None, [] + if group is not None: + data = group.broadcast_object(data, src=0) + assert data is not None + records = {tuple(record["workload"]): record for record in data["records"]} + missing = [workload for workload in workloads if workload not in records] + loaded = len(workloads) - len(missing) + assignments = _shard_workloads( + missing, + tp_size, + max_tokens, + heads, + kv_heads, + dim, + kv_dtype.itemsize == 1, + ) + local_workloads = assignments[tp_rank] + logger.info( + "ROCm segmented attention tuning plan: workloads=%d missing=%d " + "pruned=%d tp_rank=%d/%d assigned=%d scratch_budget=%d MiB " + "cache_budget=%s", + len(workloads), + len(missing), + unpruned - len(workloads), + tp_rank, + tp_size, + len(local_workloads), + memory_budget_bytes // 2**20, + cache_budget_bytes, + ) + + if not missing: + if tp_rank == 0 and merged_shards: + _save(path, data) + for shard in merged_shards: + shard.unlink(missing_ok=True) + _TABLES[key] = data + logger.info( + "ROCm segmented attention autotune ready: tp_rank=%d/%d " + "local_tuned=0 tuned=0 loaded=%d failed=0 elapsed=%.2fs cache=%s", + tp_rank, + tp_size, + loaded, + time.monotonic() - start, + path, + ) + return + + shard_path = path.with_suffix(path.suffix + f".tp{tp_rank}.part") + local_records = {} + local_failed = 0 + fatal = None + try: + with _temporary_tuning_kernels(device): + for workload in local_workloads: + try: + record = _tune_workload( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + max_tokens, + workload, + sliding_window=sliding_window, + causal=causal, + has_sinks=has_sinks, + physical_seq_len=( + physical_max_len + if has_sinks and sliding_window >= 0 + else None + ), + ) + except ( + torch.OutOfMemoryError, + RuntimeError, + AssertionError, + ) as error: + local_failed += 1 + logger.warning( + "Skipping segmented attention tuning workload %s: %s", + workload, + error, + ) + else: + local_records[workload] = record + _save( + shard_path, + { + "identity": identity, + "records": list(local_records.values()), + }, + ) + finally: + gc.collect() + torch.accelerator.empty_cache() + except BaseException as error: + fatal = f"{type(error).__name__}: {error}" + + gathered = _gather_tp( + group, + { + "records": list(local_records.values()), + "failed": local_failed, + "fatal": fatal, + }, + ) + if tp_rank == 0: + fatal_errors = [] + for rank, contribution in enumerate(gathered): + records.update( + (tuple(record["workload"]), record) + for record in contribution["records"] + ) + if contribution["fatal"] is not None: + fatal_errors.append(f"TP rank {rank}: {contribution['fatal']}") + data = {"identity": identity, "records": list(records.values())} + _save(path, data) + for shard in path.parent.glob(f"{path.name}.tp*.part"): + shard.unlink(missing_ok=True) + result = { + "data": data, + "tuned": sum(len(item["records"]) for item in gathered), + "failed": sum(item["failed"] for item in gathered), + "fatal": fatal_errors, + } + else: + result = None + if group is not None: + result = group.broadcast_object(result, src=0) + assert result is not None + data = result["data"] + finally: + if lock is not None: + lock.release() + + _TABLES[key] = data + logger.info( + "ROCm segmented attention autotune ready: tp_rank=%d/%d " + "local_tuned=%d tuned=%d loaded=%d failed=%d elapsed=%.2fs cache=%s", + tp_rank, + tp_size, + len(local_records), + result["tuned"], + loaded, + result["failed"], + time.monotonic() - start, + path, + ) + if result["fatal"]: + raise RuntimeError( + "ROCm segmented attention distributed tuning failed: " + + "; ".join(result["fatal"]) + ) + + +def get_segmented_config( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + batch, + query_len, + seq_len, + sliding_window=-1, + causal=True, + has_sinks=False, +): + """Return a warmed ceiling bucket, or None for the static fallback.""" + data = _TABLES.get( + _key( + device, + dtype, + kv_dtype, + heads, + kv_heads, + dim, + page, + scale, + sliding_window, + causal, + has_sinks, + ) + ) + if data is None: + return None + records = data["records"] + eligible = [ + record + for record in records + if record["workload"][0] >= batch + and record["workload"][1] >= query_len + and record["workload"][2] >= seq_len + ] + if not eligible: + return None + winner = min( + eligible, + key=lambda record: ( + record["workload"][1], + record["workload"][0], + record["workload"][2], + ), + ) + return dict(winner["best"]) + + +def _warmup_segmented_attention(layer, device, **limits) -> None: + impl = layer.impl + config = impl._segmented_attention_config + assert config is not None + spec = layer.get_kv_cache_spec(config) + assert spec is not None + sliding_window = impl.sliding_window[0] + query_limit = getattr(layer, "segmented_query_limit", None) + if not isinstance(query_limit, int): + query_limit = MAX_QUERY_LEN + causal = getattr(layer, "segmented_causal", True) + if not isinstance(causal, bool): + causal = True + tuning_max_len = config.model_config.max_model_len + if sliding_window >= 0: + tuning_max_len = min(tuning_max_len, sliding_window + query_limit) + warmup_segmented_attention( + device, + config.model_config.dtype, + impl.num_heads, + impl.num_kv_heads, + impl.head_size, + spec.block_size, + impl.scale, + config.scheduler_config.max_num_batched_tokens, + tuning_max_len, + config.scheduler_config.max_num_seqs, + kv_dtype=spec.dtype if spec.dtype != torch.uint8 else impl.fp8_dtype, + sliding_window=sliding_window, + causal=causal, + has_sinks=impl.sinks is not None, + physical_max_len=config.model_config.max_model_len, + max_query_len=query_limit, + **limits, + ) + impl._segmented_attention_warmed_up = True + + +def rocm_segmented_attn_autotune_warmup(worker: Worker) -> None: + """Autotune RDNA segmented attention after KV allocation and before capture.""" + if ( + not current_platform.is_rocm() + or not worker.vllm_config.kernel_config.enable_rocm_segmented_attn_autotune + ): + return + + from vllm.platforms.rocm import on_gfx1x + from vllm.v1.attention.backends.rocm_segmented_attn import ( + RocmSegmentedAttentionImpl, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec + + if not on_gfx1x(): + return + layers = [] + for layer in worker.vllm_config.compilation_config.static_forward_context.values(): + impl = getattr(layer, "impl", None) + if not isinstance(impl, RocmSegmentedAttentionImpl): + continue + config = impl._segmented_attention_config + if ( + config is not None + and config.kernel_config.enable_rocm_segmented_attn_autotune + and not impl._segmented_attention_warmed_up + and impl.alibi_slopes is None + and not impl.logits_soft_cap + ): + layers.append(layer) + if not layers: + return + + specs = worker.get_kv_cache_spec() + layouts = tuple( + ( + (spec.block_size, spec.page_size_bytes) + if isinstance(spec, FullAttentionSpec) + else (0, spec.max_memory_usage_bytes(worker.vllm_config)) + ) + for spec in specs.values() + ) + cache_budget = sum( + tensor.size + for tensor in worker.model_runner.kv_cache_config.kv_cache_tensors + if not tensor.host_resident + ) + budget = _memory_budget(worker.device) + for layer in layers: + _warmup_segmented_attention( + layer, + worker.device, + memory_budget_bytes=budget, + cache_layouts=layouts, + cache_budget_bytes=cache_budget, + ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 32e2fabeb567..f28a331854b6 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -53,6 +53,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend" ) ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" + ROCM_SEGMENTED_ATTN = ( + "vllm.v1.attention.backends.rocm_segmented_attn.RocmSegmentedAttentionBackend" + ) ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" ROCM_AITER_TRITON_MLA = ( "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend" diff --git a/vllm/v1/attention/backends/rocm_segmented_attn.py b/vllm/v1/attention/backends/rocm_segmented_attn.py new file mode 100644 index 000000000000..c541eed45166 --- /dev/null +++ b/vllm/v1/attention/backends/rocm_segmented_attn.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm attention backend for token-major segmented Triton attention.""" + +import os +from pathlib import Path +from typing import TYPE_CHECKING, ClassVar + +import torch + +import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.backend import AttentionLayer, AttentionType +from vllm.v1.attention.backends.rocm_attn import ( + RocmAttentionBackend, + RocmAttentionImpl, + RocmAttentionMetadata, + RocmAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import get_num_attention_heads_from_layers +from vllm.v1.attention.ops.segmented_attention import ( + get_segmented_attention_workspace, + segmented_attention, + segmented_attention_workspace_size, +) +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +if TYPE_CHECKING: + from vllm.platforms.interface import DeviceCapability + +logger = init_logger(__name__) + + +def is_rdna() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx1x + + return on_gfx1x() + + +class RocmSegmentedAttentionMetadataBuilder(RocmAttentionMetadataBuilder): + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + model_config = vllm_config.model_config + self.num_heads_q = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or model_config.get_num_attention_heads(vllm_config.parallel_config) + self.num_heads_kv = kv_cache_spec.num_kv_heads + self.headdim = kv_cache_spec.head_size + + +class RocmSegmentedAttentionBackend(RocmAttentionBackend): + """Explicit token-major ROCm backend optimized for segmented prefill.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + @staticmethod + def get_name() -> str: + return "ROCM_SEGMENTED_ATTN" + + @staticmethod + def get_impl_cls() -> type["RocmSegmentedAttentionImpl"]: + return RocmSegmentedAttentionImpl + + @staticmethod + def get_builder_cls() -> type["RocmSegmentedAttentionMetadataBuilder"]: + return RocmSegmentedAttentionMetadataBuilder + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [64, 128, 256] + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + return attn_type == AttentionType.DECODER + + @classmethod + def supports_non_causal(cls) -> bool: + return True + + @classmethod + def supports_sliding_window(cls) -> bool: + return True + + @classmethod + def supports_sink(cls) -> bool: + return True + + @classmethod + def supports_mm_prefix(cls) -> bool: + return False + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: "DeviceCapability", + ) -> str | None: + del ( + head_size, + dtype, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ) + if not is_rdna(): + return "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm" + + from vllm.platforms.rocm import on_gfx12x + + if kv_cache_dtype in ("fp8", "fp8_e4m3") and not on_gfx12x(): + return "FP8 segmented attention requires gfx12" + return None + + +class RocmSegmentedAttentionImpl(RocmAttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: int | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + if not is_rdna(): + raise ValueError("ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm") + + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + sinks, + ) + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + fp8 = kv_cache_dtype in ("fp8", "fp8_e4m3") + arch_supported = on_gfx12x() if fp8 else on_gfx1x() + gqa_ratio_supported = ( + num_kv_heads > 0 + and num_heads % num_kv_heads == 0 + and 1 <= self.num_queries_per_kv <= 16 + ) + if not ( + arch_supported + and attn_type == AttentionType.DECODER + and head_size in (64, 128, 256) + and gqa_ratio_supported + and kv_cache_dtype in ("auto", "float16", "bfloat16", "fp8", "fp8_e4m3") + and alibi_slopes is None + ): + raise ValueError( + "ROCM_SEGMENTED_ATTN requires decoder attention on gfx1x " + "(gfx12 for FP8), head size 64, 128 or 256, GQA ratio 1-16, and " + "does not support ALiBi." + ) + # Startup tuning and target/draft compilation must share a Triton cache. + cache_dir = os.environ.setdefault( + "TRITON_CACHE_DIR", + str(Path(envs.VLLM_CACHE_ROOT) / "rocm_segmented_attention" / "triton"), + ) + Path(cache_dir).mkdir(parents=True, exist_ok=True) + logger.info_once("Using token-major ROCm segmented Triton attention") + self._segmented_attention_warmed_up = False + self._segmented_attention_config: VllmConfig | None = None + + def process_weights_after_loading(self, act_dtype: torch.dtype): + super().process_weights_after_loading(act_dtype) + config = get_current_vllm_config() + self._segmented_attention_config = config + + def _autotune_enabled(self) -> bool: + config = self._segmented_attention_config + return ( + config is not None + and config.kernel_config.enable_rocm_segmented_attn_autotune + ) + + def _get_workspace(self, device): + config = self._segmented_attention_config + assert config is not None + sizes = segmented_attention_workspace_size( + config.scheduler_config.max_num_seqs, + self.num_heads, + self.num_kv_heads, + self.head_size, + config.model_config.max_model_len, + max_tokens=config.scheduler_config.max_num_batched_tokens, + fp8=self.kv_cache_dtype in ("fp8", "fp8_e4m3"), + autotune=self._autotune_enabled(), + ) + return get_segmented_attention_workspace(device, sizes) + + def _split_kv_cache( + self, + kv_cache: torch.Tensor, + num_kv_heads: int | None = None, + head_size: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_kv_heads = num_kv_heads or self.num_kv_heads + head_size = head_size or self.head_size + return ( + kv_cache[:, 0].unflatten(-1, (num_kv_heads, head_size)), + kv_cache[:, 1].unflatten(-1, (num_kv_heads, head_size)), + ) + + def forward( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: RocmAttentionMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + if output_block_scale is not None: + raise NotImplementedError( + "fused block_scale output quantization is not supported" + ) + # Profile runs allocate the same fixed scratch used during graph replay. + workspace = self._get_workspace(query.device) + if attn_metadata is None: + return output.fill_(0) + assert not attn_metadata.use_cascade + + num_actual_tokens = attn_metadata.num_actual_tokens + key_cache, value_cache = self._split_kv_cache(kv_cache) + if is_quantized_kv_cache(self.kv_cache_dtype): + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0: + raise NotImplementedError( + "A non-1.0 q_scale with an FP8 query is not supported" + ) + + segmented_attention( + query=query[:num_actual_tokens], + key=key[:num_actual_tokens], + value=value[:num_actual_tokens], + output=output[:num_actual_tokens], + kv_cache_dtype=self.kv_cache_dtype, + key_cache=key_cache, + value_cache=value_cache, + block_table=attn_metadata.block_table, + query_start_loc=attn_metadata.query_start_loc, + seq_lens=attn_metadata.seq_lens, + max_seq_len=attn_metadata.max_seq_len, + max_query_len=attn_metadata.max_query_len, + k_scale=layer._k_scale, + v_scale=layer._v_scale, + sm_scale=self.scale, + sliding_window=self.sliding_window[0], + output_scale=output_scale, + sinks=self.sinks, + causal=attn_metadata.causal, + softcap=self.logits_soft_cap, + workspace=workspace, + use_tuned_config=self._autotune_enabled(), + ) + return output + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ): + key_cache, value_cache = self._split_kv_cache(kv_cache) + if is_quantized_kv_cache(self.kv_cache_dtype): + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + triton_reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + def do_rope_and_kv_cache_update( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + key_cache, value_cache = self._split_kv_cache( + kv_cache, + layer.num_kv_heads, # type: ignore[attr-defined] + layer.head_size, # type: ignore[attr-defined] + ) + is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) + if is_fp8_kv_cache: + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + + rocm_aiter_ops.triton_rope_and_cache( + query, + key, + value, + positions, + cos_sin_cache, + is_neox, + key_cache, + value_cache, + layer_slot_mapping, + layer._k_scale, + layer._v_scale, + True, + is_fp8_kv_cache, + ) diff --git a/vllm/v1/attention/ops/segmented_attention.py b/vllm/v1/attention/ops/segmented_attention.py new file mode 100644 index 000000000000..f98b07df2a3b --- /dev/null +++ b/vllm/v1/attention/ops/segmented_attention.py @@ -0,0 +1,1074 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ROCm token-major segmented attention kernels and dispatch.""" + +from __future__ import annotations + +from functools import lru_cache + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.kv_cache_interface import get_kv_quant_mode +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +from .triton_unified_attention import unified_attention + +MAX_QUERY_LEN = 8192 +MAX_SPLITS = 64 +# D128 scratch cap; wider outputs scale this budget with head dimension. +MAX_LONG_EXTEND_WORKSPACE_BYTES = 32 * 1024 * 1024 + + +def segmented_query_capacity(max_query_len: int) -> int: + """Round a supported query length to its compiled power-of-two capacity.""" + if not 1 <= max_query_len <= MAX_QUERY_LEN: + raise ValueError(f"Query length must be in [1, {MAX_QUERY_LEN}]") + return 1 << (max_query_len - 1).bit_length() + + +def _query_capacity_buckets() -> tuple[int, ...]: + return tuple(1 << exponent for exponent in range(MAX_QUERY_LEN.bit_length())) + + +def _long_extend_splits( + base_splits: int, + batch: int, + query_len: int, + seq_len: int, + hq: int, + hk: int, + dim: int, + bm: int, +) -> int: + """Add long-prefix parallelism while bounding the extra split workspace.""" + if seq_len < 32768: + return base_splits + groups = batch * hk * triton.cdiv(query_len * (hq // hk), bm) + split_cap = 16 if seq_len < 131072 else 32 + scratch_per_split = batch * segmented_query_capacity(query_len) * hq * (dim + 1) * 4 + scratch_limit = MAX_LONG_EXTEND_WORKSPACE_BYTES * dim // 128 + scratch_splits = scratch_limit // scratch_per_split + scratch_cap = 1 << (scratch_splits.bit_length() - 1) if scratch_splits else 1 + occupancy_splits = triton.next_power_of_2(triton.cdiv(512, groups)) + return max(base_splits, min(split_cap, scratch_cap, occupancy_splits)) + + +@triton.jit +def _segmented_attention_stage( + Q, + K, + V, + TABLE, + STARTS, + LENS, + KS, + VS, + SINKS, + PART, + LSE, + OUT, + Q0: tl.constexpr, + Q1: tl.constexpr, + O0: tl.constexpr, + O1: tl.constexpr, + K0: tl.constexpr, + K1: tl.constexpr, + K2: tl.constexpr, + K3: tl.constexpr, + V0: tl.constexpr, + V1: tl.constexpr, + V2: tl.constexpr, + V3: tl.constexpr, + T0: tl.constexpr, + T1: tl.constexpr, + PAGE: tl.constexpr, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + QCAP: tl.constexpr, + LIMIT_Q: tl.constexpr, + SCALE: tl.constexpr, + FP8: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + SPLITS: tl.constexpr, + WINDOW: tl.constexpr, + CAUSAL: tl.constexpr, + HAS_SINKS: tl.constexpr, + PREFIX_FAST: tl.constexpr, + PV_SPLIT: tl.constexpr, + QK_PIPELINE: tl.constexpr, +): + item, mb, split = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq, kh = item // HK, item % HK + first = tl.load(STARTS + seq) + nq = tl.load(STARTS + seq + 1) - first + if nq < 1 or nq > LIMIT_Q: + return + G: tl.constexpr = HQ // HK + ROWS: tl.constexpr = QCAP * G + if mb * BM >= nq * G: + return + sequence = tl.load(LENS + seq) + prefix = sequence - nq + m = mb * BM + tl.arange(0, BM) + qpos, qh = m // G, kh * G + m % G + n = tl.arange(0, BN) + d = tl.arange(0, D) + if BK == D: + q_full = tl.load( + Q + (first + qpos[:, None]) * Q0 + qh[:, None] * Q1 + d[None, :], + mask=(m < nq * G)[:, None], + other=0.0, + ) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denom = tl.zeros((BM,), tl.float32) + if HAS_SINKS: + sink = tl.load(SINKS + qh, mask=m < nq * G, other=0).to(tl.float32) + maximum = tl.where(split == 0, sink * 1.4426950408889634, maximum) + denom = tl.where(split == 0, 1.0, denom) + if PV_SPLIT: + half = tl.arange(0, D // 2) + acc_low = tl.zeros((BM, D // 2), tl.float32) + acc_high = tl.zeros((BM, D // 2), tl.float32) + else: + acc = tl.zeros((BM, D), tl.float32) + tile_rows = tl.minimum((mb + 1) * BM, nq * G) + range_begin = 0 + if WINDOW >= 0: + first_qpos = (mb * BM) // G + range_begin = tl.maximum(0, prefix + first_qpos - WINDOW) + range_begin = (range_begin // BN) * BN + range_end = sequence + if CAUSAL: + range_end = tl.minimum(range_end, prefix + tl.cdiv(tile_rows, G)) + span = tl.maximum(0, range_end - range_begin) + segment = tl.cdiv(span, SPLITS * BN) * BN + begin = range_begin + split * segment + end = tl.minimum(begin + segment, range_end) + k_scale = 1.0 + v_scale = 1.0 + if FP8: + k_scale = tl.load(KS) + v_scale = tl.load(VS) + for tail in tl.static_range(2): + if tail: + loop_begin = tl.maximum(begin, (end // BN) * BN) + loop_end = end + else: + loop_begin = begin + loop_end = (end // BN) * BN + for start in tl.range(loop_begin, loop_end, BN): + ns = start + n + valid = tl.full((BN,), True, tl.int1) + if tail: + valid = ns < end + if FP8 and PAGE >= BN: + page = start // PAGE + inside = start - page * PAGE + n + crossed = inside >= PAGE + page += crossed + inside = tl.where(crossed, inside - PAGE, inside) + else: + page = ns // PAGE + inside = ns % PAGE + block = tl.load(TABLE + seq * T0 + page * T1, mask=valid, other=0).to( + tl.int64 + ) + scores = tl.zeros((BM, BN), tl.float32) + for ki in tl.range(D // BK, num_stages=QK_PIPELINE): + kd = ki * BK + tl.arange(0, BK) + if BK == D: + q = q_full + else: + q = tl.load( + Q + + (first + qpos[:, None]) * Q0 + + qh[:, None] * Q1 + + kd[None, :], + mask=(m < nq * G)[:, None], + other=0.0, + ) + k = tl.load( + K + + block[None, :] * K0 + + inside[None, :] * K1 + + kh * K2 + + kd[:, None] * K3, + mask=valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k.to(q.dtype), scores) + scores *= SCALE * 1.4426950408889634 * k_scale + # A complete prefix tile is visible to every valid causal row when + # sliding-window masking is disabled. + if not ( + PREFIX_FAST + and CAUSAL + and WINDOW < 0 + and start + BN <= prefix + and mb * BM + BM <= nq * G + ): + score_mask = valid[None, :] & (m[:, None] < nq * G) + if CAUSAL: + score_mask = score_mask & (ns[None, :] <= prefix + qpos[:, None]) + if WINDOW >= 0: + score_mask = score_mask & ( + ns[None, :] >= prefix + qpos[:, None] - WINDOW + ) + scores = tl.where(score_mask, scores, -float("inf")) + new_max = tl.maximum(maximum, tl.max(scores, 1)) + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + alpha = tl.exp2(maximum - safe_max) + p = tl.exp2(scores - safe_max[:, None]) + if PV_SPLIT: + acc_low *= alpha[:, None] + acc_high *= alpha[:, None] + v_low = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + half[None, :] * V3, + mask=valid[:, None], + other=0.0, + ) + v_high = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + (half[None, :] + D // 2) * V3, + mask=valid[:, None], + other=0.0, + ) + prob = p.to(Q.dtype.element_ty) + acc_low = tl.dot(prob, v_low.to(prob.dtype), acc_low) + acc_high = tl.dot(prob, v_high.to(prob.dtype), acc_high) + else: + acc *= alpha[:, None] + v = tl.load( + V + + block[:, None] * V0 + + inside[:, None] * V1 + + kh * V2 + + d[None, :] * V3, + mask=valid[:, None], + other=0.0, + ) + if FP8 and D == 256 and BM >= 32: + # Defer the scalar V dequantization to shorten the D-wide + # accumulator live range for the register-heavy 256-D tiles. + acc = tl.dot( + p.to(Q.dtype.element_ty), v.to(Q.dtype.element_ty), acc + ) + elif FP8: + acc += ( + tl.dot(p.to(Q.dtype.element_ty), v.to(Q.dtype.element_ty)) + * v_scale + ) + else: + acc = tl.dot(p.to(Q.dtype.element_ty), v, acc) + denom = denom * alpha + tl.sum(p, 1) + maximum = safe_max + norm = tl.where(denom > 0.0, denom, 1.0) + if PV_SPLIT: + result = tl.cat(acc_low, acc_high, dim=1) / norm[:, None] + else: + result = acc / norm[:, None] + if FP8 and D == 256 and BM >= 32: + result *= v_scale + if SPLITS == 1: + tl.store( + OUT + (first + qpos[:, None]) * O0 + qh[:, None] * O1 + d[None, :], + result, + mask=(m < nq * G)[:, None], + ) + else: + row = ((seq * HK + kh) * ROWS + m) * SPLITS + split + tl.store( + PART + row[:, None] * D + d[None, :], + result, + mask=(m < nq * G)[:, None], + ) + tl.store( + LSE + row, + maximum + tl.log2(denom), + mask=m < nq * G, + ) + + +@triton.jit +def _segmented_attention_reduce( + PART, + LSE, + OUT, + STARTS, + O0: tl.constexpr, + O1: tl.constexpr, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + QCAP: tl.constexpr, + LIMIT_Q: tl.constexpr, + SPLITS: tl.constexpr, + REDUCE_D: tl.constexpr, +): + m, item, db = tl.program_id(0), tl.program_id(1), tl.program_id(2) + seq, kh = item // HK, item % HK + first = tl.load(STARTS + seq) + nq = tl.load(STARTS + seq + 1) - first + G: tl.constexpr = HQ // HK + ROWS: tl.constexpr = QCAP * G + if nq < 1 or nq > LIMIT_Q or m >= nq * G: + return + s = tl.arange(0, SPLITS) + d = db * REDUCE_D + tl.arange(0, REDUCE_D) + row = ((seq * HK + kh) * ROWS + m) * SPLITS + lse = tl.load(LSE + row + s) + weights = tl.exp2(lse - tl.max(lse, 0)) + part = tl.load( + PART + (row + s)[:, None] * D + d[None, :], + mask=d[None, :] < D, + other=0.0, + ) + result = tl.sum(part * weights[:, None], 0) / tl.sum(weights, 0) + tl.store( + OUT + (first + m // G) * O0 + (kh * G + m % G) * O1 + d, + result, + mask=d < D, + ) + + +@lru_cache(maxsize=512) +def select_segmented_config(batch, max_query_len, max_seq_len, hq, hk, dim, fp8): + """Select the token-major cache configuration for segmented attention.""" + qcap = min(max_query_len, MAX_QUERY_LEN) + if batch < 1 or qcap < 1: + raise ValueError("Batch and query capacity must be positive") + + if qcap <= 2: + if dim == 64 and qcap == 1 and batch == 1 and max_seq_len >= 131072: + return dict( + bm=16, + bn=32, + bk=64, + splits=256, + warps=1, + stages=1, + ) + if dim == 64 and max_seq_len <= 256: + return dict( + bm=16, + bn=32 if fp8 else 64, + bk=64, + splits=1, + warps=4, + stages=1, + ) + if dim == 64 and max_seq_len >= 32768: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 16) + return dict( + bm=16, + bn=32, + bk=64, + splits=min(64, triton.next_power_of_2(triton.cdiv(128, groups))), + warps=4, + stages=1, + ) + if batch * hk == 1: + cfg = dict( + bm=16, + bn=64, + bk=min(128, dim), + splits=32, + warps=4, + stages=1, + ) + else: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 16) + cfg = dict( + bm=16, + bn=32, + bk=dim, + splits=min( + MAX_SPLITS, + triton.next_power_of_2(triton.cdiv(64, groups)), + ), + warps=4, + stages=1, + ) + while cfg["splits"] > 1 and max_seq_len < cfg["splits"] * cfg["bn"] * 2: + cfg["splits"] //= 2 + return cfg + + if qcap <= 8: + bm, bn = 16, 64 + bk, stages = 64, 1 + elif qcap <= 32: + bm, bk = 32, 64 + bn, stages = (32, 2) if fp8 else (64, 1) + else: + bm, bn, bk = 32, 64, 64 + stages = 1 if fp8 else 2 + groups = batch * hk * triton.cdiv(qcap * (hq // hk), bm) + target = 96 + if fp8 and qcap > 32: + target = 192 + splits = min(32, triton.next_power_of_2(triton.cdiv(target, groups))) + while splits > 1 and max_seq_len < splits * 128: + splits //= 2 + cfg = dict( + bm=bm, + bn=bn, + bk=bk, + splits=splits, + warps=4, + stages=stages, + ) + if dim == 64 and qcap >= 64 and qcap < max_seq_len <= qcap + 256: + cfg.update( + bm=64 if fp8 or qcap <= 256 else 128, + bn=32, + bk=64, + splits=1, + warps=4, + stages=1, + ) + elif dim == 64 and fp8 and max_seq_len <= 256: + cfg.update(bn=32, splits=1) + elif dim == 64 and qcap == 8 and max_seq_len >= 32768: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 32) + cfg.update( + bm=32, + bn=32 if fp8 else 64, + splits=min(64, triton.next_power_of_2(triton.cdiv(256, groups))), + ) + elif dim == 64 and fp8 and qcap == 8 and max_seq_len >= 4096: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 32) + cfg.update( + bm=32, + bn=64, + splits=min(32, triton.next_power_of_2(triton.cdiv(128, groups))), + ) + elif dim == 64 and fp8 and qcap >= 64 and max_seq_len >= 4096: + groups = batch * hk * triton.cdiv(qcap * (hq // hk), 64) + cfg.update( + bm=64, + bn=32, + splits=min(8, triton.next_power_of_2(triton.cdiv(1024, groups))), + stages=1, + ) + elif dim == 128 and fp8 and qcap <= 256 and max_seq_len <= qcap: + cfg.update(splits=1) + elif ( + dim == 128 + and fp8 + and qcap >= 256 + and max_seq_len >= 8192 + and ( + hq // hk >= 4 + or (max_seq_len >= 32768 and (hq // hk >= 2 or qcap >= 512)) + or max_seq_len >= 131072 + ) + ): + # Shorter prefixes need sufficient grouped-query reuse; at GQA1/Q256 + # the wider tile pays off only once the prefix reaches 131K tokens. + cfg.update( + bm=64, + bn=128, + bk=64, + warps=8, + stages=3, + waves_per_eu=6, + prefix_fast=True, + ) + if qcap >= 512 or max_seq_len >= 32768: + cfg["qk_pipeline"] = 3 + cfg["splits"] = _long_extend_splits( + cfg["splits"], batch, qcap, max_seq_len, hq, hk, dim, cfg["bm"] + ) + elif dim == 256 and fp8 and qcap >= 256 and max_seq_len >= 8192: + # Wider tiles amortize each FP8 KV load over twice as many queries and + # keys. Six waves is the best gfx1201 pressure/occupancy tradeoff. + cfg.update( + bm=64, + bn=128, + bk=64, + warps=8, + stages=1, + waves_per_eu=6, + ) + cfg["splits"] = _long_extend_splits( + cfg["splits"], batch, qcap, max_seq_len, hq, hk, dim, cfg["bm"] + ) + elif dim == 256 and not fp8 and qcap >= 1024 and max_seq_len >= 8192: + # Two half-width PV accumulators avoid spilling this BF16 D256 tile. + # Keep Q256 on the incumbent: the wider tile loses at that capacity. + cfg.update( + bm=64, + bn=64, + bk=64, + warps=8, + stages=3, + waves_per_eu=6, + pv_split=True, + qk_pipeline=3, + ) + elif dim == 128 and not fp8 and qcap > 128: + if max_seq_len >= 8192: + cfg.update( + bm=128, + bn=32, + bk=128, + warps=8, + stages=1, + waves_per_eu=6, + ) + groups = batch * hk * triton.cdiv(qcap * (hq // hk), cfg["bm"]) + cfg["splits"] = min(16, triton.next_power_of_2(triton.cdiv(192, groups))) + elif qcap >= 4096: + cfg.update( + bm=128, + bn=32, + bk=128, + splits=1, + warps=8, + stages=1, + waves_per_eu=6, + ) + elif qcap > 256 or hq > 10: + cfg.update( + bm=64, + bn=32, + bk=128, + splits=1, + warps=4, + stages=1, + waves_per_eu=6, + ) + else: + cfg.update(bm=32, bn=64, bk=64, splits=1, warps=4, stages=2) + return cfg + + +def segmented_workspace_shapes(batch, query_cap, hq, hk, dim, splits): + if splits == 1: + return None + partial = (batch, hk, query_cap * (hq // hk), splits, dim) + return partial, partial[:-1] + + +@lru_cache(maxsize=128) +def segmented_attention_workspace_size( + max_batch, + hq, + hk, + dim, + max_seq_len, + *, + max_tokens=None, + fp8=False, + autotune=False, +): + """Size the largest selected attention workspace before graph capture. + + ``max_tokens`` bounds reachable ``(batch, query_len)`` pairs using one + longest query and one token for each remaining sequence. Omitting it keeps + the legacy reservation behavior for callers without scheduler limits. + ``autotune`` reserves additional splits for tuned D64 candidates. + """ + largest = 0 + previous_capacity = 0 + for query_capacity in _query_capacity_buckets(): + query_lengths = {previous_capacity + 1, query_capacity} + for query_len in query_lengths: + if max_tokens is not None and query_len > max_tokens: + continue + batch_limit = max_batch + if max_tokens is not None: + batch_limit = min(batch_limit, max_tokens - query_len + 1) + for batch in range(1, batch_limit + 1): + cfg = select_segmented_config( + batch, query_len, max_seq_len, hq, hk, dim, fp8 + ) + splits = cfg["splits"] + if dim == 64 and query_len <= 512 and max_seq_len >= 4096 and autotune: + split_limit = ( + 512 + if batch == 1 and query_len == 1 and max_seq_len >= 131072 + else MAX_SPLITS + ) + splits = min(split_limit, 2 * splits) + if splits > 1: + largest = max( + largest, + splits * batch * query_capacity * hq, + ) + previous_capacity = query_capacity + return largest * dim, largest + + +def get_segmented_attention_workspace(device, sizes): + """Keep scratch storage stable and separate for each runtime lane/ubatch.""" + + def allocate(): + return tuple( + torch.empty(size, device=device, dtype=torch.float32) for size in sizes + ) + + if is_workspace_manager_initialized(): + return current_workspace_manager().get_persistent_resource( + ("rocm_segmented_attention", device, sizes), allocate + ) + return allocate() + + +def can_use_segmented_attention( + q, k, v, out, kc, vc, table, starts, lengths, k_scale, v_scale +): + """Check tensor metadata; unsupported feature checks live in the caller.""" + if ( + k is None + or v is None + or q.ndim != 3 + or kc.ndim != 4 + or vc.ndim != 4 + or q.dtype not in (torch.bfloat16, torch.float16) + or out.dtype != q.dtype + or q.shape != out.shape + or q.shape[-1] not in (64, 128, 256) + or kc.dtype != vc.dtype + or k.dtype != q.dtype + or v.dtype != q.dtype + or kc.dtype not in (q.dtype, torch.float8_e4m3fn, torch.float8_e4m3fnuz) + or table.ndim != 2 + or starts.ndim != 1 + or lengths.ndim != 1 + or table.shape[0] != lengths.numel() + or starts.numel() != lengths.numel() + 1 + or table.dtype != torch.int32 + or starts.dtype != torch.int32 + or lengths.dtype != torch.int32 + or starts.stride(0) != 1 + or lengths.stride(0) != 1 + or not q.is_cuda + ): + return False + hkv = kc.shape[2] + dim = q.shape[-1] + page = kc.shape[1] + if hkv < 1 or q.shape[1] % hkv or not 1 <= q.shape[1] // hkv <= 16: + return False + if ( + k.shape != (q.shape[0], hkv, dim) + or v.shape != k.shape + or kc.shape != (kc.shape[0], page, hkv, dim) + or vc.shape != kc.shape + or any( + t.device != q.device for t in (k, v, out, kc, vc, table, starts, lengths) + ) + or any(t.stride(-1) != 1 for t in (q, k, v, out, kc, vc)) + ): + return False + if kc.element_size() == 1: + return all( + isinstance(s, torch.Tensor) + and s.numel() == 1 + and s.dtype == torch.float32 + and s.device == q.device + for s in (k_scale, v_scale) + ) + return True + + +def run_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + max_query_len, + max_seq_len, + k_scale, + v_scale, + scale, + *, + sliding_window=-1, + causal=True, + sinks=None, + config=None, + workspace=None, +): + """Write eligible unified-cache attention rows; leave other requests untouched.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + if batch == 0 or max_query_len == 0 or q.numel() == 0: + return out + query_len = min(max_query_len, MAX_QUERY_LEN) + qcap = segmented_query_capacity(query_len) + fp8 = kc.element_size() == 1 + hk = kc.shape[2] + attention_span = max_seq_len + if sliding_window >= 0: + attention_span = min(max_seq_len, sliding_window + max_query_len) + cfg = config or select_segmented_config( + batch, query_len, attention_span, hq, hk, dim, fp8 + ) + splits = cfg["splits"] + shapes = segmented_workspace_shapes(batch, qcap, hq, hk, dim, splits) + if shapes is None: + partial = lse = out + elif workspace is not None: + partial, lse = ( + buffer.view(-1)[: shape.numel()].view(shape) + for buffer, shape in zip(workspace, map(torch.Size, shapes)) + ) + elif is_workspace_manager_initialized(): + partial, lse = current_workspace_manager().get_simultaneous( + (shapes[0], torch.float32), (shapes[1], torch.float32) + ) + else: + partial = torch.empty(shapes[0], device=q.device, dtype=torch.float32) + lse = torch.empty(shapes[1], device=q.device, dtype=torch.float32) + _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + cfg, + partial, + lse, + qcap, + sliding_window, + causal, + compile_only=False, + ) + return out + + +def compile_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + max_query_len, + k_scale, + v_scale, + scale, + config, + workspace, + *, + sliding_window=-1, + causal=True, + sinks=None, +): + """Compile one segmented attention configuration without launching it.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + qcap = segmented_query_capacity(min(max_query_len, MAX_QUERY_LEN)) + hk = kc.shape[2] + shapes = segmented_workspace_shapes(batch, qcap, hq, hk, dim, config["splits"]) + if shapes is None: + partial = lse = out + else: + partial, lse = workspace + _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + config, + partial, + lse, + qcap, + sliding_window, + causal, + compile_only=True, + ) + + +def _launch_segmented_attention( + q, + out, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + scale, + cfg, + partial, + lse, + qcap, + sliding_window, + causal, + *, + compile_only, +): + """Launch or compile the exact stage and reduction specializations.""" + batch, hq, dim = lengths.numel(), q.shape[1], q.shape[2] + hk = kc.shape[2] + fp8 = kc.element_size() == 1 + splits = cfg["splits"] + stage_grid = ( + batch * hk, + triton.cdiv(qcap * (hq // hk), cfg["bm"]), + splits, + ) + stage_args = ( + q, + kc, + vc, + table, + starts, + lengths, + k_scale, + v_scale, + sinks, + partial, + lse, + out, + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + *kc.stride(), + *vc.stride(), + *table.stride(), + kc.shape[1], + hq, + hk, + dim, + qcap, + MAX_QUERY_LEN, + scale, + fp8, + cfg["bm"], + cfg["bn"], + cfg["bk"], + splits, + sliding_window, + causal, + sinks is not None, + cfg.get("prefix_fast", False), + cfg.get("pv_split", False), + cfg.get("qk_pipeline"), + ) + stage_options = dict( + num_warps=cfg["warps"], + num_stages=cfg["stages"], + waves_per_eu=cfg.get("waves_per_eu", 2), + ) + if compile_only: + _segmented_attention_stage.warmup(*stage_args, grid=stage_grid, **stage_options) + else: + _segmented_attention_stage[stage_grid](*stage_args, **stage_options) + if splits > 1: + reduce_d = cfg.get("reduce_d", 64 if batch * qcap * hq < 64 else dim) + reduce_grid = ( + qcap * (hq // hk), + batch * hk, + triton.cdiv(dim, reduce_d), + ) + reduce_args = ( + partial, + lse, + out, + starts, + out.stride(0), + out.stride(1), + hq, + hk, + dim, + qcap, + MAX_QUERY_LEN, + splits, + reduce_d, + ) + reduce_options = {"num_warps": cfg.get("reduce_warps", 4)} + if compile_only: + _segmented_attention_reduce.warmup( + *reduce_args, grid=reduce_grid, **reduce_options + ) + else: + _segmented_attention_reduce[reduce_grid](*reduce_args, **reduce_options) + + +def get_segmented_config(*args, **kwargs): + """Load the tuner on demand to avoid a kernel/tuner import cycle.""" + from vllm.model_executor.warmup.rocm_segmented_attn_autotune_warmup import ( + get_segmented_config as tuned_config, + ) + + return tuned_config(*args, **kwargs) + + +def segmented_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + kv_cache_dtype: str, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + max_query_len: int, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + *, + sliding_window: int = -1, + output_scale: torch.Tensor | None = None, + sinks: torch.Tensor | None = None, + causal: bool | torch.Tensor = True, + softcap: float = 0.0, + workspace=None, + use_tuned_config: bool = True, +) -> None: + """Run segmented attention when eligible, otherwise use unified Triton.""" + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + fp8_dtype = current_platform.fp8_dtype() + if key_cache.dtype == torch.uint8: + key_cache = key_cache.view(fp8_dtype) + value_cache = value_cache.view(fp8_dtype) + + from vllm.platforms.rocm import on_gfx1x, on_gfx12x + + segmented_pattern = ( + isinstance(causal, bool) + and (on_gfx12x() if key_cache.element_size() == 1 else on_gfx1x()) + and not softcap + and output_scale is None + and ( + sinks is None + or ( + sinks.ndim == 1 + and sinks.shape[0] == query.shape[1] + and sinks.device == query.device + and sinks.dtype in (torch.float16, torch.bfloat16, torch.float32) + ) + ) + ) + if ( + segmented_pattern + and 0 < max_query_len <= MAX_QUERY_LEN + and ( + can_use_segmented_attention( + query, + key, + value, + output, + key_cache, + value_cache, + block_table, + query_start_loc, + seq_lens, + k_scale, + v_scale, + ) + ) + ): + attention_span = max_seq_len + if sliding_window >= 0: + attention_span = min(max_seq_len, sliding_window + max_query_len) + config = None + if use_tuned_config: + config = get_segmented_config( + query.device, + query.dtype, + key_cache.dtype, + query.shape[1], + key_cache.shape[2], + query.shape[2], + key_cache.shape[1], + sm_scale, + len(seq_lens), + max_query_len, + attention_span, + sliding_window, + causal, + has_sinks=sinks is not None, + ) + if config is None: + config = select_segmented_config( + len(seq_lens), + max_query_len, + attention_span, + query.shape[1], + key_cache.shape[2], + query.shape[2], + key_cache.element_size() == 1, + ) + run_segmented_attention( + query, + output, + key_cache, + value_cache, + block_table, + query_start_loc, + seq_lens, + max_query_len, + max_seq_len, + k_scale, + v_scale, + sm_scale, + sliding_window=sliding_window, + causal=causal, + sinks=sinks, + config=config, + workspace=workspace, + ) + return + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=query_start_loc, + max_seqlen_q=max_query_len, + seqused_k=seq_lens, + max_seqlen_k=max_seq_len, + softmax_scale=sm_scale, + causal=causal, + window_size=(-1, -1) if sliding_window < 0 else (sliding_window, 0), + block_table=block_table, + softcap=softcap, + q_descale=None, + k_descale=k_scale, + v_descale=v_scale, + output_scale=output_scale, + sinks=sinks, + kv_quant_mode=get_kv_quant_mode(kv_cache_dtype), + ) From e11962fc1ce542c725dacda6f641daed8d54337f Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:53:51 +0300 Subject: [PATCH 057/952] [Feat][Mamba2] Enable internal prefill checkpoints (#57329) Signed-off-by: Roi Koren Signed-off-by: zjy0516 Co-authored-by: zjy0516 --- .../tool_parsers/test_granite4_tool_parser.py | 4 + .../mamba/test_mamba2_checkpoint_state.py | 163 ++++++++++++++++++ .../test_mamba2_prefill_checkpoints.py | 163 ++++++++++++++++++ .../model_executor/layers/mamba/checkpoint.py | 4 + .../layers/mamba/mamba_mixer2.py | 41 ++++- .../layers/mamba/ops/ssd_checkpoint.py | 115 ++++++++++++ vllm/v1/attention/backends/mamba2_attn.py | 59 ++++++- vllm/v1/attention/backends/mamba_attn.py | 60 +++++-- 8 files changed, 586 insertions(+), 23 deletions(-) create mode 100644 tests/kernels/mamba/test_mamba2_checkpoint_state.py create mode 100644 tests/v1/attention/test_mamba2_prefill_checkpoints.py create mode 100644 vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py diff --git a/tests/entrypoints/tool_parsers/test_granite4_tool_parser.py b/tests/entrypoints/tool_parsers/test_granite4_tool_parser.py index 71fe7637bf4a..67826d956475 100644 --- a/tests/entrypoints/tool_parsers/test_granite4_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_granite4_tool_parser.py @@ -24,6 +24,10 @@ def server(): "4096", "--max-num-seqs", "2", + # Streaming and non-streaming runs must match bitwise; a prefix-cache + # hit resumes the second from a mid-prefill Mamba checkpoint, whose + # numerics differ slightly from the cold single-pass prefill. + "--no-enable-prefix-caching", ] with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server: yield server diff --git a/tests/kernels/mamba/test_mamba2_checkpoint_state.py b/tests/kernels/mamba/test_mamba2_checkpoint_state.py new file mode 100644 index 000000000000..54fcf80205ce --- /dev/null +++ b/tests/kernels/mamba/test_mamba2_checkpoint_state.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""An SSD intermediate state equals an independent run over the same prefix. + +This is the central claim of internal prefill checkpoints: the state the SSD +scan already materializes at the checkpoint's chunk boundary is exactly the +state a fresh prefill over `[0, checkpoint)` would end with, so it can be +cached and later resumed from. The last test covers the other half: moving +that state, and the conv window, into the paged cache. +""" + +import pytest +import torch + +from tests.kernels.mamba.test_mamba_ssm_ssd import generate_random_inputs +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointMetadata, +) +from vllm.model_executor.layers.mamba.ops.ssd_checkpoint import ( + store_prefill_checkpoint, +) +from vllm.model_executor.layers.mamba.ops.ssd_combined import ( + mamba_chunk_scan_combined_varlen, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +DEVICE = current_platform.device_type +compute_chunk_metadata = BaseMambaAttentionMetadataBuilder._compute_chunk_metadata + + +def _run(inputs, seqlen, chunk_size, checkpoint_offset): + """Run one prefill of `seqlen` tokens, optionally exporting a checkpoint. + + Returns (states, checkpoint_chunk_index). `states` holds every chunk's + state when a checkpoint was requested, else the per-sequence final state. + """ + A, dt, X, B, C = inputs + offsets = [checkpoint_offset] if checkpoint_offset else None + cu, seq_idx, last, ckpt = compute_chunk_metadata( + chunk_size, + 1, + torch.tensor([0]), + torch.tensor([0, seqlen]), + checkpoint_offsets_p=offsets, + ) + + x = X[:seqlen] + out = torch.empty_like(x) + states = mamba_chunk_scan_combined_varlen( + x, + dt[:seqlen], + A, + B[:seqlen], + C[:seqlen], + chunk_size, + cu_seqlens=torch.tensor([0, seqlen], dtype=torch.int32, device=DEVICE), + cu_chunk_seqlens=torch.tensor(cu, dtype=torch.int32, device=DEVICE), + last_chunk_indices=torch.tensor(last, dtype=torch.int32, device=DEVICE), + seq_idx=torch.tensor(seq_idx, dtype=torch.int32, device=DEVICE), + out=out, + D=None, + return_intermediate_states=checkpoint_offset > 0, + state_dtype=torch.float32, + ) + return states, ckpt[0] + + +@pytest.mark.parametrize("checkpoint_offset", [32, 48, 96]) +def test_checkpoint_state_matches_an_independent_prefix_run(checkpoint_offset): + """states[checkpoint_chunk_idx] == the final state of a prefix-only run. + + 48 is deliberately not a multiple of chunk_size: align-mode block sizes + carry no chunk_size factor, so off-grid checkpoints are the common case + and are what the chunk split exists to support. + """ + seqlen, chunk_size, n_heads, d_head = 128, 32, 4, 32 + inputs = generate_random_inputs(1, seqlen, n_heads, d_head, torch.float32) + inputs = tuple(t.squeeze(0) if t.dim() > 1 else t for t in inputs) + + states, ckpt_idx = _run(inputs, seqlen, chunk_size, checkpoint_offset) + assert ckpt_idx >= 0, "the checkpoint did not land on a chunk boundary" + checkpoint_state = states[ckpt_idx] + + prefix_final, _ = _run(inputs, checkpoint_offset, chunk_size, 0) + + torch.testing.assert_close( + checkpoint_state.to(torch.float32), + prefix_final[0].to(torch.float32), + atol=8e-3, + rtol=5e-3, + ) + + +def test_checkpoint_does_not_perturb_the_final_state(): + """Exporting a checkpoint must not change what the prefill produces. + + The chunk split is an extra boundary, not a different computation, so the + sequence's final state has to match a run that took no checkpoint. + """ + seqlen, chunk_size, n_heads, d_head = 128, 32, 4, 32 + inputs = generate_random_inputs(1, seqlen, n_heads, d_head, torch.float32) + inputs = tuple(t.squeeze(0) if t.dim() > 1 else t for t in inputs) + + with_ckpt, ckpt_idx = _run(inputs, seqlen, chunk_size, 48) + without_ckpt, _ = _run(inputs, seqlen, chunk_size, 0) + + torch.testing.assert_close( + with_ckpt[-1].to(torch.float32), + without_ckpt[0].to(torch.float32), + atol=8e-3, + rtol=5e-3, + ) + + +@pytest.mark.parametrize("dim_first", [True, False]) +def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first): + """The store reads the checkpoint row out of every chunk state. + + Row 0 checkpoints 32 tokens in, its state being chunk 3 rather than the + program's own row, so the store only passes if ``chunk_idx`` is used. + Row 1 declines and must write nothing. The conv input is a column slice of + a wider projection, the SD conv layout is a transposed view, and SSM rows + are padded pages, matching the strides the mixer passes. + """ + dim, state_len, heads, head_dim, d_state = 16, 4, 2, 4, 8 + projection = torch.randn(64, dim + 8, device=DEVICE) + conv_input = projection[:, :dim] + if dim_first: + conv_state = torch.zeros(3, dim, state_len, device=DEVICE) + else: + conv_state = torch.zeros(3, state_len, dim, device=DEVICE).transpose(-1, -2) + varlen_states = torch.randn(5, heads, head_dim, d_state, device=DEVICE) + # Blocks are pages holding more than the SSM state, so rows are padded. + row = heads * head_dim * d_state + ssm_storage = torch.zeros(3, row + 8, device=DEVICE) + ssm_state = ssm_storage[:, :row].view(3, heads, head_dim, d_state) + checkpoint = MambaPrefillCheckpointMetadata( + checkpoint_offsets=torch.tensor([32, 0], dtype=torch.int32, device=DEVICE), + state_indices=torch.tensor( + [2, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE + ), + offsets=[32, 0], + ) + + store_prefill_checkpoint( + checkpoint, + torch.tensor([3, 0], device=DEVICE), + conv_input.T, + conv_state, + varlen_states, + ssm_state, + torch.tensor([0, 40, 64], dtype=torch.int32, device=DEVICE), + ) + + exact = dict(rtol=0, atol=0) + torch.testing.assert_close( + conv_state[2], conv_input[32 - state_len : 32].T, **exact + ) + torch.testing.assert_close(ssm_state[2], varlen_states[3], **exact) + assert not conv_state[:2].any() and not ssm_storage[:2].any() + assert not ssm_storage[:, row:].any(), "wrote into the page padding" diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py new file mode 100644 index 000000000000..e5276786e23c --- /dev/null +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Internal prefill checkpoints for Mamba2 align-mode prefix caching.""" + +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.v1.attention.backends.mamba2_attn import ( + Mamba2AttentionMetadataBuilder, +) +from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import MambaSpec + +compute_chunk_metadata = BaseMambaAttentionMetadataBuilder._compute_chunk_metadata + + +def test_no_checkpoint_offsets_preserves_legacy_chunking(): + """Passing no offsets must reproduce the pre-feature chunk layout.""" + num_computed = torch.tensor([0, 300]) + qsl = torch.tensor([0, 700, 1000]) + + cu, seq_idx, last, ckpt = compute_chunk_metadata(256, 2, num_computed, qsl) + + # Row 0: 700 fresh tokens from sequence position 0 -> 256, 256, 188. + # Row 1: 300 new tokens resuming at 300. 300 % 256 == 44, so it finishes + # its partial chunk with 256 - 44 == 212 tokens, then 88. + assert cu == [0, 256, 512, 700, 912, 1000] + assert seq_idx == [0, 0, 0, 1, 1] + assert last == [2, 4] + assert ckpt == [-1, -1] + + +def test_chunking_realigns_to_the_grid_after_a_checkpoint(): + """After the forced break, chunks resume breaking on chunk_size. + + 100 -> then 156 to reach the 256 grid, then full chunks. Without the + realignment the second chunk would span 100..356 and cross a boundary. + """ + num_computed = torch.tensor([0]) + qsl = torch.tensor([0, 700]) + + cu, _, _, _ = compute_chunk_metadata( + 256, 1, num_computed, qsl, checkpoint_offsets_p=[100] + ) + + assert cu == [0, 100, 256, 512, 700] + + +def test_checkpoint_on_a_resumed_request_is_relative_to_the_query(): + """Offsets count from the query start, not the sequence start.""" + num_computed = torch.tensor([300]) + qsl = torch.tensor([0, 500]) + + cu, _, _, ckpt = compute_chunk_metadata( + 256, 1, num_computed, qsl, checkpoint_offsets_p=[212] + ) + + # Row resumes at sequence position 300. 300 % 256 == 44, so the partial + # chunk is 212 tokens and already ends exactly at the checkpoint (query + # offset 212 == sequence position 512) — no extra split needed. + assert cu == [0, 212, 468, 500] + assert ckpt == [0] + + +MAMBA_BLOCK_SIZE = 256 +CHUNK_SIZE = 256 +DEVICE = torch.device("cpu") + + +def _create_mamba2_builder( + mamba_cache_mode: str = "align", + num_prefill_checkpoint_blocks: int = 1, + prefill_checkpoint_alignment: int | None = 1, +) -> Mamba2AttentionMetadataBuilder: + vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + # get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`. + vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE + spec = MambaSpec( + block_size=MAMBA_BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + mamba_cache_mode=mamba_cache_mode, + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, + prefill_checkpoint_alignment=prefill_checkpoint_alignment, + ) + return Mamba2AttentionMetadataBuilder( + kv_cache_spec=spec, + layer_names=["layer.0"], + vllm_config=vllm_config, + device=DEVICE, + ) + + +def _build(builder, seq_lens, query_lens): + common = create_common_attn_metadata( + BatchSpec(seq_lens=seq_lens, query_lens=query_lens), + MAMBA_BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ) + # create_common_attn_metadata leaves is_prefilling unset; the builder + # asserts on it. Every row here is a full prefill. + common = common.replace(is_prefilling=torch.ones(len(seq_lens), dtype=torch.bool)) + return builder.build(common_prefix_len=0, common_attn_metadata=common) + + +def test_builder_emits_one_checkpoint_entry_per_prefill_row(): + """Declined rows stay in place, masked, and a chunk ends on the checkpoint.""" + builder = _create_mamba2_builder() + # Row 0 (seq_len 900) checkpoints at 768; row 1 (seq_len 100) is too short + # for a checkpoint at all. + meta = _build(builder, seq_lens=[900, 100], query_lens=[900, 100]) + + assert meta.checkpoint_chunk_idx is not None + assert meta.checkpoint_meta is not None + assert meta.checkpoint_chunk_idx.numel() == 2 + # The exporter masks row 1 off by its zero offset and null block, but + # still loads its chunk index first, so the placeholder must be in range. + assert meta.checkpoint_meta.checkpoint_offsets.tolist() == [768, 0] + assert meta.checkpoint_meta.state_indices[1].item() == NULL_BLOCK_ID + assert meta.checkpoint_chunk_idx[1].item() == 0 + # The checkpoint's token offset is the end of its chunk. Asserting it here + # also proves the chunk split actually landed on the checkpoint. + ckpt_chunk = meta.checkpoint_chunk_idx[0] + assert meta.cu_chunk_seqlen_p[ckpt_chunk + 1].item() == 768 + + +def test_update_block_table_regathers_checkpoint_blocks(): + """Groups sharing a spec reuse one group's metadata via update_block_table. + + Checkpoint destinations are block-table entries, so each group needs its + own: layers of different groups share physical tensors, and reusing the + first group's indices writes every group's checkpoint into its blocks. + """ + builder = _create_mamba2_builder() + meta = _build(builder, seq_lens=[900, 100], query_lens=[900, 100]) + assert meta.checkpoint_meta is not None + + # A second group's table: same shape as the first, disjoint block ids. + other_table = torch.arange(1000, 1016, dtype=torch.int32).reshape(2, 8) + updated = builder.update_block_table( + meta, other_table, torch.zeros(1000, dtype=torch.int64) + ) + + # Row 0 checkpoints in column cdiv(900, 256) - 2 == 2; row 1 declined. + assert updated.checkpoint_meta is not None + assert updated.checkpoint_meta.state_indices.tolist() == [ + other_table[0, 2].item(), + NULL_BLOCK_ID, + ] + + +def test_builder_emits_nothing_outside_align_mode(): + builder = _create_mamba2_builder(mamba_cache_mode="none") + meta = _build(builder, seq_lens=[900], query_lens=[900]) + + assert meta.checkpoint_chunk_idx is None diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index feef0afa05ff..45888113e152 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -70,6 +70,9 @@ def _gather_checkpoint_state_indices( class MambaPrefillCheckpointMetadata: checkpoint_offsets: torch.Tensor state_indices: torch.Tensor + # Host copy of ``checkpoint_offsets``, for backends that lay out the + # checkpoint before any tensor exists (Mamba2's SSD chunk split). + offsets: list[int] | None = None # Gather indices into the block table, kept so that another KV cache group # with the same spec can re-derive its own ``state_indices``. request_rows: torch.Tensor | None = None @@ -139,6 +142,7 @@ def build( _gather_checkpoint_state_indices( m.block_table_tensor, request_rows_tensor, checkpoint_cols_tensor ), + offsets=checkpoint_offsets, request_rows=request_rows_tensor, block_cols=checkpoint_cols_tensor, ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a67480ebc26d..3753153fa3f2 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -3,11 +3,12 @@ from collections.abc import Sequence +from dataclasses import replace import torch from torch import nn -from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.config.mamba import MambaBackendEnum from vllm.distributed import ( divide, @@ -39,6 +40,7 @@ from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 selective_state_update_replayssm_output_only, ) +from vllm.model_executor.layers.mamba.ops.ssd_checkpoint import store_prefill_checkpoint from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -66,7 +68,7 @@ from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum -from vllm.v1.kv_cache_interface import KVCacheGroupSpec +from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec, MambaSpec logger = init_logger(__name__) @@ -766,6 +768,8 @@ def conv_ssm_forward( query_start_loc_d = attn_metadata.query_start_loc_d num_decodes = attn_metadata.num_decodes num_decode_tokens = attn_metadata.num_decode_tokens + checkpoint_chunk_idx = attn_metadata.checkpoint_chunk_idx + checkpoint_meta = attn_metadata.checkpoint_meta if attn_metadata is None: # V1 profile run -- warm up SSD kernels so that autotuning @@ -842,6 +846,7 @@ def conv_ssm_forward( # NOTE: final output is an in-place update of out tensor assert preallocated_ssm_out_p is not None + has_checkpoints = checkpoint_chunk_idx is not None varlen_states = mamba_chunk_scan_combined_varlen( hidden_states_p.view( num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim @@ -859,6 +864,7 @@ def conv_ssm_forward( cu_chunk_seqlens=cu_chunk_seqlen_p, last_chunk_indices=last_chunk_indices_p, initial_states=initial_states, + return_intermediate_states=has_checkpoints, dt_softplus=True, dt_limit=(0.0, float("inf")), out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim), @@ -869,7 +875,24 @@ def conv_ssm_forward( # - varlen state is a (num_prefills, nheads, headdim, dstate) # tensor assert state_indices_tensor_p is not None - ssm_state[state_indices_tensor_p] = varlen_states + final_states = ( + varlen_states[last_chunk_indices_p] + if has_checkpoints + else varlen_states + ) + if has_checkpoints: + assert checkpoint_meta is not None + assert query_start_loc_p is not None + store_prefill_checkpoint( + checkpoint_meta, + checkpoint_chunk_idx, + x, + conv_state, + varlen_states, + ssm_state, + query_start_loc_p, + ) + ssm_state[state_indices_tensor_p] = final_states if ring_start is not None and self._updates_replayssm_trackers: assert prev_num_accepted is not None reset_replayssm_ring_trackers( @@ -1070,6 +1093,18 @@ def conv_ssm_forward( is_blackwell=self.is_blackwell, ) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + spec = super().get_kv_cache_spec(vllm_config) + if not isinstance(spec, MambaSpec): + return spec + # The block-boundary state can be exported from inside one forward + # pass, so the scheduler need not split the prefill; every reader + # acts on this only in align mode. The metadata builder puts a chunk + # end exactly on the checkpoint, so any token position works. + return replace( + spec, num_prefill_checkpoint_blocks=1, prefill_checkpoint_alignment=1 + ) + def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None diff --git a/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py new file mode 100644 index 000000000000..25bfcedbdc88 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointMetadata, +) +from vllm.triton_utils import tl, triton + + +@triton.jit +def _store_prefill_checkpoint_kernel( + x_ptr, + conv_state_ptr, + chunk_states_ptr, + ssm_state_ptr, + query_start_loc_ptr, + checkpoint_offsets_ptr, + state_indices_ptr, + chunk_idx_ptr, + stride_x_dim, + stride_x_token, + stride_conv_block, + stride_conv_dim, + stride_conv_token, + stride_chunk_states, + stride_ssm_block, + DIM: tl.constexpr, + STATE_LEN: tl.constexpr, + SSM_ROW_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + checkpoint_offset = tl.load(checkpoint_offsets_ptr + row) + valid = checkpoint_offset > 0 + dst = tl.load(state_indices_ptr + row).to(tl.int64) + + conv_mask = valid & (cols < DIM * STATE_LEN) + dim_idx = cols // STATE_LEN + history_idx = cols % STATE_LEN + token_idx = ( + tl.load(query_start_loc_ptr + row) + checkpoint_offset - STATE_LEN + history_idx + ).to(tl.int64) + conv = tl.load( + x_ptr + dim_idx * stride_x_dim + token_idx * stride_x_token, mask=conv_mask + ) + tl.store( + conv_state_ptr + + dst * stride_conv_block + + dim_idx * stride_conv_dim + + history_idx * stride_conv_token, + conv, + mask=conv_mask, + ) + + ssm_mask = valid & (cols < SSM_ROW_SIZE) + src = tl.load(chunk_idx_ptr + row).to(tl.int64) + ssm = tl.load(chunk_states_ptr + src * stride_chunk_states + cols, mask=ssm_mask) + tl.store(ssm_state_ptr + dst * stride_ssm_block + cols, ssm, mask=ssm_mask) + + +def store_prefill_checkpoint( + checkpoint: MambaPrefillCheckpointMetadata, + chunk_idx: torch.Tensor, + x: torch.Tensor, + conv_state: torch.Tensor, + chunk_states: torch.Tensor, + ssm_state: torch.Tensor, + query_start_loc: torch.Tensor, +) -> None: + """Write each prefill row's mid-prefill checkpoint into its paged block. + + Rows without a checkpoint (offset 0) write nothing. + + Args: + checkpoint: Per-row offsets and destination blocks. + chunk_idx: Per-row index into ``chunk_states`` of the checkpoint. + x: Pre-conv activations, ``(dim, num_prefill_tokens)``. + conv_state: Paged conv state, ``(num_blocks, dim, state_len)``. + chunk_states: SSD states at every chunk end. + ssm_state: Paged SSM state. + query_start_loc: Prefill query start locations. + + """ + dim = x.size(0) + state_len = conv_state.size(-1) + ssm_row_size = chunk_states[0].numel() + block_size = 1024 + grid = ( + chunk_idx.numel(), + triton.cdiv(max(dim * state_len, ssm_row_size), block_size), + ) + _store_prefill_checkpoint_kernel[grid]( + x, + conv_state, + chunk_states, + ssm_state, + query_start_loc, + checkpoint.checkpoint_offsets, + checkpoint.state_indices, + chunk_idx, + x.stride(0), + x.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + chunk_states.stride(0), + ssm_state.stride(0), + DIM=dim, + STATE_LEN=state_len, + SSM_ROW_SIZE=ssm_row_size, + BLOCK_SIZE=block_size, + ) diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index e5e3827bcd67..fc81cf680ea0 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -7,6 +7,10 @@ import torch from vllm.config import VllmConfig +from vllm.model_executor.layers.mamba.checkpoint import ( + MambaPrefillCheckpointBuilder, + MambaPrefillCheckpointMetadata, +) from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, @@ -109,6 +113,11 @@ class Mamba2AttentionMetadata(BaseMambaAttentionMetadata): # Chunk-related metadata (only for prefill) seq_idx_p: torch.Tensor | None = None + # Internal prefill checkpoints, one entry per prefill row. The chunk + # index selects the varlen_states row holding the checkpoint state. + checkpoint_chunk_idx: torch.Tensor | None = None + checkpoint_meta: MambaPrefillCheckpointMetadata | None = None + class Mamba2AttentionMetadataBuilder( BaseMambaAttentionMetadataBuilder[Mamba2AttentionMetadata] @@ -128,6 +137,9 @@ def __init__( "chunk_size needs to be set in the model config for Mamba2 models" ) self.chunk_size: int = chunk_size + self.checkpoint_builder = MambaPrefillCheckpointBuilder( + vllm_config, kv_cache_spec + ) def build( self, @@ -145,6 +157,8 @@ def build( seq_idx_p = None cu_chunk_seqlen_p = None last_chunk_indices_p = None + checkpoint_chunk_idx = None + checkpoint_meta = None prep_initial_states = False # Compute seq_idx for prefill only @@ -162,12 +176,27 @@ def build( ) prep_initial_states = bool((num_computed_tokens_p_cpu > 0).any()) - cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p = ( - self._build_chunk_metadata_tensors( - self.chunk_size, - common, - common_attn_metadata, - ) + checkpoint_offsets_p = None + first = common.num_reqs - common.num_prefills + checkpoint = self.checkpoint_builder.build( + common_attn_metadata, + list(range(first, common.num_reqs)), + ) + if checkpoint is not None: + # The host offsets place a chunk boundary on the checkpoint. + checkpoint_offsets_p = checkpoint.offsets + checkpoint_meta = checkpoint + + ( + cu_chunk_seqlen_p, + seq_idx_p, + last_chunk_indices_p, + checkpoint_chunk_idx, + ) = self._build_chunk_metadata_tensors( + self.chunk_size, + common, + common_attn_metadata, + checkpoint_offsets_p, ) return replace( @@ -177,4 +206,22 @@ def build( seq_idx_p=seq_idx_p, cu_chunk_seqlen_p=cu_chunk_seqlen_p, last_chunk_indices_p=last_chunk_indices_p, + checkpoint_chunk_idx=checkpoint_chunk_idx, + checkpoint_meta=checkpoint_meta, + ) + + def update_block_table( + self, + metadata: Mamba2AttentionMetadata, + blk_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> Mamba2AttentionMetadata: + new_metadata = super().update_block_table(metadata, blk_table, slot_mapping) + if metadata.checkpoint_meta is None: + return new_metadata + # Checkpoint destinations are block-table entries, so each group + # re-gathers its own rather than writing into the source group's. + return replace( + new_metadata, + checkpoint_meta=metadata.checkpoint_meta.regather_state_indices(blk_table), ) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 9953461500a2..3e7ae905a30d 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -233,13 +233,14 @@ def build( num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, ) + @staticmethod def _compute_chunk_metadata( - self, chunk_size: int, num_prefills: int, num_computed_tokens_p_cpu: torch.Tensor, query_start_loc_p_cpu: torch.Tensor, - ) -> tuple[list[int], list[int], list[int]]: + checkpoint_offsets_p: list[int] | None = None, + ) -> tuple[list[int], list[int], list[int], list[int]]: """Compute chunk-specific metadata for Mamba models. The code below carefully constructs the chunks such that: @@ -251,20 +252,34 @@ def _compute_chunk_metadata( of prefix caching for mamba (wip). We need to take care of the interaction with chunked prefill in order to satisfy constraint (2). + + `checkpoint_offsets_p[i]` forces an extra chunk end that many tokens + into row i's query (0 = none), because the SSD scan only materializes + states at chunk ends. Returns its chunk index per row, or -1. """ # TODO (tdoublep): This code could probably be optimized. cu_chunk_seqlen = [] seq_idx = [] last_chunk_indices = [] + checkpoint_chunk_indices = [] seqlen_pos = 0 + # A checkpoint splits its request into two segments so that a chunk + # ends exactly on it; each segment then chunks by the rules below. + segments = [] for req_idx in range(num_prefills): - this_num_computed = num_computed_tokens_p_cpu[req_idx].item() - this_new_tokens = ( + computed = num_computed_tokens_p_cpu[req_idx].item() + tokens = ( query_start_loc_p_cpu[req_idx + 1].item() - query_start_loc_p_cpu[req_idx].item() ) + ckpt = checkpoint_offsets_p[req_idx] if checkpoint_offsets_p else 0 + if ckpt: + segments.append((req_idx, computed, ckpt, True)) + segments.append((req_idx, computed + ckpt, tokens - ckpt, False)) + checkpoint_chunk_indices = [-1] * num_prefills + for req_idx, this_num_computed, this_new_tokens, is_ckpt_end in segments: # if computed tokens are not chunk-aligned, use the first # chunk to finish it off if this_num_computed % chunk_size != 0: @@ -288,11 +303,14 @@ def _compute_chunk_metadata( this_new_tokens -= chunk_len assert this_new_tokens == 0 - last_chunk_indices.append(len(cu_chunk_seqlen) - 1) + if is_ckpt_end: + checkpoint_chunk_indices[req_idx] = len(cu_chunk_seqlen) - 1 + else: + last_chunk_indices.append(len(cu_chunk_seqlen) - 1) cu_chunk_seqlen.append(seqlen_pos) - return cu_chunk_seqlen, seq_idx, last_chunk_indices + return cu_chunk_seqlen, seq_idx, last_chunk_indices, checkpoint_chunk_indices def _prefill_cpu_metadata( self, @@ -326,9 +344,13 @@ def _build_chunk_metadata_tensors( chunk_size: int, common: M, common_attn_metadata: CommonAttentionMetadata, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + checkpoint_offsets_p: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """Compute chunk metadata and return as device tensors. - Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p). + Returns (cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p, + checkpoint_chunk_idx_p). The last is None unless a row checkpoints; + it has one entry per prefill row, holding the logical chunk that ends + on the checkpoint, and 0 for rows that decline. """ num_prefills = common.num_prefills @@ -339,11 +361,14 @@ def _build_chunk_metadata_tensors( common.num_decode_tokens, ) - cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata( - chunk_size, - num_prefills, - num_computed_tokens_p_cpu, - query_start_loc_p_cpu, + cu_chunk_seqlen, seq_idx, last_chunk_indices, ckpt_idxs = ( + self._compute_chunk_metadata( + chunk_size, + num_prefills, + num_computed_tokens_p_cpu, + query_start_loc_p_cpu, + checkpoint_offsets_p, + ) ) device = common_attn_metadata.query_start_loc.device @@ -356,7 +381,14 @@ def _build_chunk_metadata_tensors( last_chunk_indices_p = async_tensor_h2d( last_chunk_indices, dtype=torch.int32, device=device ) - return cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p + ckpt_idx_p = ( + async_tensor_h2d( + [max(idx, 0) for idx in ckpt_idxs], dtype=torch.int64, device=device + ) + if any(idx >= 0 for idx in ckpt_idxs) + else None + ) + return cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p, ckpt_idx_p def _compute_common_metadata( self, From 31e2443c90542a33a4a4a293ea7186fba2796c67 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 6 Oct 2026 16:47:02 +0800 Subject: [PATCH 058/952] [Model] LongCat-Flash: scale MLA norms while loading and drop the post-load sweep (#60080) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/longcat_flash.py | 87 ++++++------------- .../models/longcat_flash_mtp.py | 47 +--------- 2 files changed, 30 insertions(+), 104 deletions(-) diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index 50962e972877..4ddd1ea9bf3f 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -58,13 +58,17 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.utils.int8_utils import block_dequant 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.model_loader.weight_utils import ( + LoaderFunction, + composed_weight_loader, + default_weight_loader, +) from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention +from vllm.model_executor.utils import set_weight_attrs from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP @@ -80,6 +84,23 @@ logger = init_logger(__name__) +def _scaled_weight_loader(scale: float) -> LoaderFunction: + return composed_weight_loader(default_weight_loader, lambda x: x * scale) + + +def scale_mla_lora_norms_on_load(self_attn: nn.Module, config) -> None: + """Fold LongCat's MLA LoRA scaling into the q/kv_a_layernorm weight loaders.""" + for norm, rank, enabled in ( + (self_attn.q_a_layernorm, config.q_lora_rank, config.mla_scale_q_lora), + (self_attn.kv_a_layernorm, config.kv_lora_rank, config.mla_scale_kv_lora), + ): + if enabled: + scale = (config.hidden_size / rank) ** 0.5 + set_weight_attrs( + norm.weight, {"weight_loader": _scaled_weight_loader(scale)} + ) + + class FlashMLP(nn.Module): """Flash MLP layer.""" @@ -265,6 +286,8 @@ def __init__( for i in range(2) ] ) + for self_attn in self.self_attn: + scale_mla_lora_norms_on_load(self_attn, config) self.input_layernorm = nn.ModuleList( [RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for i in range(2)] ) @@ -450,6 +473,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue + # The MTP layers are loaded by LongCatFlashMTP. + if name.startswith("mtp."): + continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue @@ -461,9 +487,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name.endswith(".bias") or name.endswith("_bias") ) and name not in params_dict: continue - # Skip mtp - if ".mtp." in name: - continue if is_pp_missing_parameter(name, self): continue param = params_dict[name] @@ -478,9 +501,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue is_expert_weight = True name_mapped = name.replace(weight_name, param_name) - # Skip mtp - if ".mtp." in name_mapped: - continue if ( name_mapped.endswith(".bias") or name_mapped.endswith("_bias") ) and name not in params_dict: @@ -515,9 +535,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Skip loading kv_scale from ckpts towards new design. if name.endswith(".kv_scale") and name not in params_dict: continue - # Skip mtp - if ".mtp." in name: - continue if name is None: continue if is_pp_missing_parameter(name, self): @@ -528,54 +545,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) weight_loader(param, loaded_weight) loaded_params.add(name) - for layer_id in range(self.config.num_layers): - for i in range(2): - if isinstance(self.layers[layer_id], PPMissingLayer): - continue - self_attn = self.layers[layer_id].self_attn[i] - if ( - self.quant_config is not None - and hasattr(self.quant_config, "weight_block_size") - and self_attn.kv_b_proj.weight.dtype - in ( - torch.float8_e4m3fn, - torch.float8_e4m3fnuz, - ) - ): - weight_block_size = self.quant_config.weight_block_size - if weight_block_size is not None: - assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") - dtype = torch.get_default_dtype() - w = block_dequant( - self_attn.kv_b_proj.weight, - self_attn.kv_b_proj.weight_scale_inv, - weight_block_size, - ).to(dtype) - else: - w = self_attn.kv_b_proj.weight - - w_kc, w_vc = w.unflatten( - 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) - ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) - self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) - self_attn.w_vc = w_vc.contiguous().transpose(1, 2) - # Guard against compounding on incremental load_weights calls: - # the in-place ``*=`` would otherwise re-apply the MLA LoRA - # scaling to the layernorm weights on each pass. - if self.config.mla_scale_q_lora and not getattr( - self_attn, "_mla_q_lora_scaled", False - ): - self_attn.q_a_layernorm.weight.data *= ( - self.config.hidden_size / self.config.q_lora_rank - ) ** 0.5 - self_attn._mla_q_lora_scaled = True - if self.config.mla_scale_kv_lora and not getattr( - self_attn, "_mla_kv_lora_scaled", False - ): - self_attn.kv_a_layernorm.weight.data *= ( - self.config.hidden_size / self.config.kv_lora_rank - ) ** 0.5 - self_attn._mla_kv_lora_scaled = True return loaded_params diff --git a/vllm/model_executor/models/longcat_flash_mtp.py b/vllm/model_executor/models/longcat_flash_mtp.py index 506f7a7451ac..ae380f92b6bb 100644 --- a/vllm/model_executor/models/longcat_flash_mtp.py +++ b/vllm/model_executor/models/longcat_flash_mtp.py @@ -14,7 +14,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.utils.int8_utils import block_dequant from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -23,6 +22,7 @@ from vllm.sequence import IntermediateTensors from .deepseek_v2 import DeepseekV2DecoderLayer +from .longcat_flash import scale_mla_lora_norms_on_load from .utils import maybe_prefix @@ -45,6 +45,7 @@ def __init__( prefix="eh_proj", ) self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix) + scale_mla_lora_norms_on_load(self.mtp_block.self_attn, config) self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( @@ -263,50 +264,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) - spec_layer_id = self.model.mtp_start_layer_idx - self_attn = self.model.layers[str(spec_layer_id)].mtp_block.self_attn - if ( - self.quant_config is not None - and hasattr(self.quant_config, "weight_block_size") - and self_attn.kv_b_proj.weight.dtype - in ( - torch.float8_e4m3fn, - torch.float8_e4m3fnuz, - ) - ): - weight_block_size = self.quant_config.weight_block_size - if weight_block_size is not None: - dtype = torch.get_default_dtype() - w = block_dequant( - self_attn.kv_b_proj.weight, - self_attn.kv_b_proj.weight_scale_inv, - weight_block_size, - ).to(dtype) - else: - w = self_attn.kv_b_proj.weight - else: - w = self_attn.kv_b_proj.weight - w_kc, w_vc = w.unflatten( - 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) - ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) - self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) - self_attn.w_vc = w_vc.contiguous().transpose(1, 2) - # Guard against compounding on incremental load_weights calls (the - # in-place *= would otherwise double-apply the LoRA scaling). - if self.config.mla_scale_q_lora and not getattr( - self_attn, "_mla_q_lora_scaled", False - ): - self_attn.q_a_layernorm.weight.data *= ( - self.config.hidden_size / self.config.q_lora_rank - ) ** 0.5 - self_attn._mla_q_lora_scaled = True - if self.config.mla_scale_kv_lora and not getattr( - self_attn, "_mla_kv_lora_scaled", False - ): - self_attn.kv_a_layernorm.weight.data *= ( - self.config.hidden_size / self.config.kv_lora_rank - ) ** 0.5 - self_attn._mla_kv_lora_scaled = True return loaded_params def _rewrite_spec_layer_name( From 6bbad6acdffd58b63037829f9fe27b92c6650071 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 6 Oct 2026 18:00:15 +0800 Subject: [PATCH 059/952] [Rust Frontend] Pass tool defer_loading through to chat templates (#60203) Co-authored-by: Claude Opus 5.5 Signed-off-by: Bugen Zhao --- rust/src/chat/src/output/default/mod.rs | 2 + rust/src/chat/src/output/default/unified.rs | 1 + rust/src/chat/src/output/harmony/tests.rs | 1 + .../chat/src/renderer/deepseek_v32/tests.rs | 1 + .../chat/src/renderer/deepseek_v4/tests.rs | 1 + rust/src/chat/src/renderer/hf/mod.rs | 18 ++++++- rust/src/chat/src/renderer/kimi_k3/tests.rs | 2 + rust/src/chat/src/renderer/test_utils.rs | 1 + rust/src/chat/src/request.rs | 2 + rust/src/chat/tests/chat.rs | 1 + rust/src/chat/tests/roundtrip.rs | 2 + rust/src/parser/src/output_grammar.rs | 1 + .../tool/deepseek_dsml/deepseek_v41/tests.rs | 1 + rust/src/parser/src/tool/mimo.rs | 1 + rust/src/parser/src/tool/minimax_m3.rs | 1 + rust/src/parser/src/tool/mod.rs | 3 ++ rust/src/parser/src/tool/parameters.rs | 1 + rust/src/parser/src/tool/test_utils.rs | 5 ++ rust/src/parser/src/unified/combined.rs | 1 + rust/src/parser/src/unified/gemma4.rs | 8 +++ rust/src/parser/src/unified/hf/tests.rs | 2 + rust/src/parser/src/unified/hy/hy_v3.rs | 1 + rust/src/parser/src/unified/hy/hy_v4.rs | 2 + rust/src/parser/src/unified/inkling.rs | 1 + .../routes/openai/chat_completions/convert.rs | 51 +++++++++++++++++++ .../openai/chat_completions/validate.rs | 6 +++ .../server/src/routes/openai/utils/types.rs | 7 +++ rust/src/server/src/routes/tokenize/types.rs | 1 + 28 files changed, 124 insertions(+), 1 deletion(-) diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 3f75aed63d4e..b31c5e9e4cbe 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -311,6 +311,7 @@ mod tests { "required": ["query"] }), strict, + defer_loading: None, }]; let mut request = ChatRequest { tool_context: ResolvedToolContext::new(&[], tools, Some(choice), true).unwrap(), @@ -346,6 +347,7 @@ mod tests { description: None, parameters: serde_json::json!({"type": "object"}), strict: Some(true), + defer_loading: None, }]; let mut request = ChatRequest { tool_context: ResolvedToolContext::new( diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index b4a2dcbe5d30..9f872d97945b 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -884,6 +884,7 @@ mod tests { description: None, parameters: serde_json::json!({ "type": "object" }), strict: None, + defer_loading: None, }]; let parser = Gemma4UnifiedParser::new(&tools, Arc::new(tokenizer)).unwrap(); let events = vec![ diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 035aef3cd4aa..5b89e53bb9aa 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -99,6 +99,7 @@ fn request_with_tools() -> ChatRequest { "required": ["city"] }), strict: None, + defer_loading: None, }]; ChatRequest { tool_context: ResolvedToolContext::new(&[], tools, Some(ChatToolChoice::Auto), true) diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index dd9ec13ddc97..804e4c2bc9fa 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -152,6 +152,7 @@ fn request_level_tools_are_lowered_as_synthetic_leading_system_message() { "required": ["query"] }), strict: None, + defer_loading: None, }]; let mut request = ChatRequest { request_id: "deepseek-v32-tools".to_string(), diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 8702e8d73c08..78e9044a3be0 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -183,6 +183,7 @@ fn tools_keep_historical_developer() { description: None, parameters: Value::Object(Default::default()), strict: None, + defer_loading: None, }; let request = thinking_request(vec![ ChatMessage::developer("old instruction", Some(vec![tool])), diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 5cb0262191bf..06b58e27ee3b 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -324,6 +324,8 @@ struct TemplateToolDefinition { parameters: JsonValue, #[serde(skip_serializing_if = "Option::is_none")] strict: Option, + #[serde(skip_serializing_if = "Option::is_none")] + defer_loading: Option, } /// Convert chat messages into the JSON shape expected by Jinja chat templates. @@ -604,6 +606,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec { description: tool.description.clone(), parameters: tool.parameters.clone(), strict: tool.strict, + defer_loading: tool.defer_loading, }, }) .collect() @@ -649,6 +652,7 @@ mod tests { description: None, parameters: serde_json::json!({"type": "object"}), strict: None, + defer_loading: None, }]), ), ChatMessage::user("hello"), @@ -1020,6 +1024,7 @@ mod tests { "required": ["city"], }), strict: Some(true), + defer_loading: None, }]), )]); @@ -1419,6 +1424,7 @@ mod tests { "required": ["city"], }), strict: None, + defer_loading: None, }]; request.tool_context = crate::request::ResolvedToolContext::new( &request.messages, @@ -1445,6 +1451,7 @@ mod tests { description: Some("Get weather".to_string()), parameters: serde_json::json!({"type": "object"}), strict: None, + defer_loading: None, }]; request.tool_context = crate::request::ResolvedToolContext::new( &request.messages, @@ -1472,12 +1479,21 @@ mod tests { description: None, parameters: Value::Null, strict: None, + defer_loading: None, }, ChatTool { name: "with_strict".to_string(), description: Some("description".to_string()), parameters: serde_json::json!({"type": "object"}), strict: Some(false), + defer_loading: None, + }, + ChatTool { + name: "deferred".to_string(), + description: None, + parameters: Value::Null, + strict: None, + defer_loading: Some(true), }, ]; request.tool_context = crate::request::ResolvedToolContext::new( @@ -1498,7 +1514,7 @@ mod tests { assert_eq!( rendered, - "name=\"without_strict\"|description=null|parameters=null|;name=\"with_strict\"|description=\"description\"|parameters={\"type\": \"object\"}|strict=false|;" + "name=\"without_strict\"|description=null|parameters=null|;name=\"with_strict\"|description=\"description\"|parameters={\"type\": \"object\"}|strict=false|;name=\"deferred\"|description=null|parameters=null|defer_loading=true|;" ); } diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 5aadb3d403c6..29a6ecb47349 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -178,6 +178,7 @@ fn tool_declare_omits_absent_fields_and_preserves_parameter_nulls() { } }), strict: None, + defer_loading: None, }]; request.tool_context = crate::request::ResolvedToolContext::new(&request.messages, tools, None, true) @@ -200,6 +201,7 @@ fn tool_declare_preserves_present_optional_fields() { description: Some("Look up a record".to_string()), parameters: json!({"type": "object"}), strict: Some(false), + defer_loading: None, }]; request.tool_context = crate::request::ResolvedToolContext::new(&request.messages, tools, None, true) diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 2c59ed7a22c4..e42dba76a518 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -280,6 +280,7 @@ fn to_chat_tools(tools: &[FixtureTool]) -> Vec { description: tool.function.description.clone(), parameters: tool.function.parameters.clone(), strict: tool.function.strict, + defer_loading: None, }) .collect() } diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index f7ddd9b12843..b52e9fbc0b81 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -801,6 +801,7 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: Some(true), + defer_loading: None, }]), ); @@ -815,6 +816,7 @@ mod tests { description: None, parameters: json!({"type": "object"}), strict: None, + defer_loading: None, } } diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ab366b8827a4..b920b5f0ed49 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -297,6 +297,7 @@ fn sample_tool_request(request_id: &str) -> ChatRequest { "required": ["city"], }), strict: None, + defer_loading: None, }]; request.tool_context = vllm_chat::ResolvedToolContext::new( &request.messages, diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index b1a6505cf88c..079ce3e1a7e7 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -1289,6 +1289,7 @@ fn test_tools() -> Vec { "required": ["location"] }), strict: None, + defer_loading: None, }, ChatTool { name: "add".to_string(), @@ -1306,6 +1307,7 @@ fn test_tools() -> Vec { "required": ["y", "x", "items"] }), strict: None, + defer_loading: None, }, ] } diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index 9058b4a58b8d..53c52538ee1d 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -255,6 +255,7 @@ mod tests { "required": ["query"] }), strict, + defer_loading: None, } } diff --git a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs index 208d88ea33a8..58043cc95c63 100644 --- a/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v41/tests.rs @@ -28,6 +28,7 @@ fn tools() -> Vec { } }), strict: None, + defer_loading: None, }] } diff --git a/rust/src/parser/src/tool/mimo.rs b/rust/src/parser/src/tool/mimo.rs index f6491e9caacc..f7ba06fb06dd 100644 --- a/rust/src/parser/src/tool/mimo.rs +++ b/rust/src/parser/src/tool/mimo.rs @@ -65,6 +65,7 @@ mod tests { name: "convert".into(), description: None, strict: None, + defer_loading: None, parameters: json!({"type": "object", "properties": { "flag": {"type": "boolean"}, "empty": {"type": ["string", "null"]}, "text": {"type": "string"}, "literal": {"type": "string"}, "payload": {"type": "object"} diff --git a/rust/src/parser/src/tool/minimax_m3.rs b/rust/src/parser/src/tool/minimax_m3.rs index b0643dd1a14e..b6e179917ed2 100644 --- a/rust/src/parser/src/tool/minimax_m3.rs +++ b/rust/src/parser/src/tool/minimax_m3.rs @@ -459,6 +459,7 @@ mod tests { } }), strict: None, + defer_loading: None, }); tools } diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index 77897526261d..e4afbda3e902 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -52,6 +52,9 @@ pub struct Tool { pub description: Option, pub parameters: Value, pub strict: Option, + /// Whether the tool's schema is loaded on demand (tool search). Only chat + /// templates consume it; the tool stays callable for parsing and grammars. + pub defer_loading: Option, } /// One tool-call update emitted while parsing assistant text. diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index 19d700830bc9..ef59ce08759b 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -502,6 +502,7 @@ mod tests { description: None, parameters, strict: None, + defer_loading: None, } } diff --git a/rust/src/parser/src/tool/test_utils.rs b/rust/src/parser/src/tool/test_utils.rs index ef3143183fe8..ec45a31ec1ec 100644 --- a/rust/src/parser/src/tool/test_utils.rs +++ b/rust/src/parser/src/tool/test_utils.rs @@ -24,6 +24,7 @@ pub fn test_tools() -> Vec { } }), strict: None, + defer_loading: None, }, Tool { name: "add".to_string(), @@ -36,6 +37,7 @@ pub fn test_tools() -> Vec { } }), strict: None, + defer_loading: None, }, Tool { name: "convert".to_string(), @@ -51,6 +53,7 @@ pub fn test_tools() -> Vec { } }), strict: None, + defer_loading: None, }, Tool { name: "calculate_area".to_string(), @@ -64,6 +67,7 @@ pub fn test_tools() -> Vec { } }), strict: None, + defer_loading: None, }, Tool { name: "update_record".to_string(), @@ -80,6 +84,7 @@ pub fn test_tools() -> Vec { } }), strict: None, + defer_loading: None, }, ] } diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index 6f6849ec95b9..ba44e7fb9613 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -197,6 +197,7 @@ mod tests { }, }), strict: None, + defer_loading: None, }] } diff --git a/rust/src/parser/src/unified/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs index 2f90555b26ba..1c65a7171cd5 100644 --- a/rust/src/parser/src/unified/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -617,48 +617,56 @@ mod tests { description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "get_time".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "write_file".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "Edit".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "search".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "set".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "get_status".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, Tool { name: "todowrite".to_string(), description: None, parameters: json!({ "type": "object" }), strict: None, + defer_loading: None, }, ] } diff --git a/rust/src/parser/src/unified/hf/tests.rs b/rust/src/parser/src/unified/hf/tests.rs index 7e27a42805f9..ed2efa6f39f5 100644 --- a/rust/src/parser/src/unified/hf/tests.rs +++ b/rust/src/parser/src/unified/hf/tests.rs @@ -686,6 +686,7 @@ fn tools_convert_string_arguments() { }, }), strict: None, + defer_loading: None, }]; let text = "\n\n\n7\n\n\ntrue\n\n\nwake up\n\n12\n\n"; assert_eq!( @@ -888,6 +889,7 @@ fn differential_against_transformers() { description: None, parameters: tool["function"]["parameters"].clone(), strict: None, + defer_loading: None, }) .collect(); let expected = match DIVERGENCES.iter().find(|(divergent, _)| *divergent == text) { diff --git a/rust/src/parser/src/unified/hy/hy_v3.rs b/rust/src/parser/src/unified/hy/hy_v3.rs index 30e021b5bce5..3599ce2cb9a0 100644 --- a/rust/src/parser/src/unified/hy/hy_v3.rs +++ b/rust/src/parser/src/unified/hy/hy_v3.rs @@ -114,6 +114,7 @@ mod tests { "properties": { "city": { "type": "string" } }, }), strict: None, + defer_loading: None, }] } diff --git a/rust/src/parser/src/unified/hy/hy_v4.rs b/rust/src/parser/src/unified/hy/hy_v4.rs index 777688bcaa05..a5d81e0f2ae7 100644 --- a/rust/src/parser/src/unified/hy/hy_v4.rs +++ b/rust/src/parser/src/unified/hy/hy_v4.rs @@ -114,12 +114,14 @@ mod tests { "properties": { "city": { "type": "string" } }, }), strict: None, + defer_loading: None, }, Tool { name: "get_current_date".to_string(), description: None, parameters: json!({ "type": "object", "properties": {} }), strict: None, + defer_loading: None, }, ] } diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs index 4372cafb24e3..7d819b06408b 100644 --- a/rust/src/parser/src/unified/inkling.rs +++ b/rust/src/parser/src/unified/inkling.rs @@ -557,6 +557,7 @@ mod tests { }, }), strict: None, + defer_loading: None, }] } diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 418ab23ce805..998ca4c939ab 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -409,6 +409,7 @@ pub(crate) fn convert_tools(tools: Option>) -> Result, A description: tool.function.description, parameters: tool.function.parameters, strict: tool.function.strict, + defer_loading: tool.function.defer_loading.or(tool.defer_loading), }) }) .collect() @@ -712,10 +713,50 @@ mod tests { description: None, parameters: serde_json::Value::Null, strict: None, + defer_loading: None, }] ); } + #[test] + fn prepare_chat_request_maps_tool_defer_loading() { + let request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + {"function": {"name": "tool_level"}, "defer_loading": true}, + {"function": {"name": "function_level", "defer_loading": true}}, + { + "function": {"name": "function_wins", "defer_loading": false}, + "defer_loading": true, + }, + ], + })) + .expect("parse defer_loading"); + + let prepared = prepare_chat_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("prepare defer_loading"); + + let defer_loading = prepared + .chat_request + .tools() + .iter() + .map(|tool| (tool.name.as_str(), tool.defer_loading)) + .collect::>(); + assert_eq!( + defer_loading, + [ + ("tool_level", Some(true)), + ("function_level", Some(true)), + ("function_wins", Some(false)), + ] + ); + } + #[test] fn chat_http_request_rejects_empty_cache_salt() { let request: ChatCompletionRequest = serde_json::from_value(json!({ @@ -1022,7 +1063,9 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: Some(true), + defer_loading: None, }, + defer_loading: None, }]), name: None, }], @@ -1048,6 +1091,7 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: Some(true), + defer_loading: None, }]), )] ); @@ -1378,7 +1422,9 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: None, + defer_loading: None, }, + defer_loading: None, }]), tool_choice: Some(ToolChoice::Value(ToolChoiceValue::None)), ..base_request() @@ -1413,6 +1459,7 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: None, + defer_loading: None, }] ); assert_eq!(prepared.chat_request.tool_choice(), &ChatToolChoice::None); @@ -1431,7 +1478,9 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: None, + defer_loading: None, }, + defer_loading: None, }]), tool_choice: Some(ToolChoice::Value(ToolChoiceValue::Required)), ..base_request() @@ -1464,7 +1513,9 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: None, + defer_loading: None, }, + defer_loading: None, }]), tool_choice: Some(ToolChoice::Function { tool_type: "function".to_string(), diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 1bfc1d9a091d..579b815ba104 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -195,7 +195,9 @@ mod tests { description: None, parameters: json!({}), strict: None, + defer_loading: None, }, + defer_loading: None, }]), ..base_request() }; @@ -212,7 +214,9 @@ mod tests { description: None, parameters: json!({}), strict: None, + defer_loading: None, }, + defer_loading: None, }]), name: None, }], @@ -234,7 +238,9 @@ mod tests { description: None, parameters: json!({}), strict: None, + defer_loading: None, }, + defer_loading: None, }]), name: None, }], diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs index 85a891341c8e..e2187bae673c 100644 --- a/rust/src/server/src/routes/openai/utils/types.rs +++ b/rust/src/server/src/routes/openai/utils/types.rs @@ -271,6 +271,10 @@ pub struct Tool { #[serde(default = "default_function_tool_type")] pub tool_type: String, pub function: Function, + /// Tool-level form of [`Function::defer_loading`]; the function-level value + /// takes precedence. + #[serde(skip_serializing_if = "Option::is_none")] + pub defer_loading: Option, } fn default_function_tool_type() -> String { @@ -286,6 +290,9 @@ pub struct Function { pub parameters: Value, /// Whether to enable strict schema adherence (OpenAI structured outputs). pub strict: Option, + /// Whether the schema is loaded on demand via tool search (vLLM extension, + /// mirroring the OpenAI Responses and Anthropic Messages field). + pub defer_loading: Option, } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index d907073c87ae..ed3fd11d1d00 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -215,6 +215,7 @@ mod tests { "properties": {"city": {"type": "string"}}, }), strict: None, + defer_loading: None, }] ); } From 888074b42ded93c0d3aac083bbf281daeec963e4 Mon Sep 17 00:00:00 2001 From: Louie Tsai Date: Tue, 6 Oct 2026 03:15:40 -0700 Subject: [PATCH 060/952] [CPU][Recipes] Detect model head constraints for automatic tensor parallel selection in vLLM Recipes Tool (#60170) Signed-off-by: louie-tsai --- tools/recipes/RUNTIME_TUNING.md | 31 +++++++- tools/recipes/model_constraints.py | 111 +++++++++++++++++++++++++++++ tools/recipes/runtime_tuning.py | 42 +++++++++-- 3 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 tools/recipes/model_constraints.py diff --git a/tools/recipes/RUNTIME_TUNING.md b/tools/recipes/RUNTIME_TUNING.md index d3a359446c93..d0b8ee929974 100644 --- a/tools/recipes/RUNTIME_TUNING.md +++ b/tools/recipes/RUNTIME_TUNING.md @@ -79,7 +79,7 @@ on a validation environment that differs from the user's target deployment. | Runtime parameter | Why it may need deployment-time refinement | Main decision input | Current policy | | --- | --- | --- | --- | -| `tensor-parallel-size` | The effective CPU/NUMA topology available to a container or pod can differ from the system used to validate the recipe. | Hardware topology | Use the largest power-of-two TP value that does not exceed the effective NUMA-node count. | +| `tensor-parallel-size` | Available topology and model head counts constrain TP. | Hardware topology + model configuration | Use the largest head-compatible power-of-two TP within the effective NUMA-node count, for DP=1 and PP=1. | | `gpu-memory-utilization` | Available memory can differ by machine size, container limits, and other memory use. The vLLM option name is also used by the CPU backend. | Hardware memory + recipe baseline | Calculate a conservative fraction from the most constrained NUMA node. | | `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active at the same time. | Workload concurrency | Set `max-num-seqs` to `--concurrency` when supplied. | | `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Calculate decode budget + expected prefill demand, with vLLM scheduler constraints as floors. | @@ -99,6 +99,35 @@ For example, 2 effective NUMA nodes produce TP=2, while 6 nodes currently produce TP=4. This is a topology-based starting point and avoids automatically selecting unusual non-power-of-two TP sizes before they are validated. +The policy then reads the model configuration through vLLM's config loader, +without loading weights. It honors the model path/ID, revision, code revision, +`trust-remote-code`, config format and JSON `hf-overrides`. No model-specific +numeric head-count table is maintained. + +Each candidate must divide the attention-head count. KV heads must divide +across TP ranks when KV heads >= TP; otherwise TP must be a multiple of the KV +head count so the heads can be replicated. An incompatible candidate is halved +until it passes every detected constraint. Nested text configuration, global KV +heads, Whisper encoder/decoder heads and hybrid linear-attention heads are +checked. Vision head divisibility is checked conservatively unless +`mm-encoder-tp-mode: data` is configured. + +For Phi-4-reasoning on four NUMA nodes, the detected 40 attention heads and 10 +KV heads reduce the initial TP=4 suggestion to TP=2. The printed tuning notes +show the detected counts and the reduction; `config.yml` receives TP=2. + +If configuration loading or head extraction fails, the tool keeps the recipe +TP (vLLM's default TP=1 when absent) and prints why detection failed. That +fallback does not certify an existing recipe TP as compatible. Plain conversion +without hardware detection does not fetch the model configuration. + +Automatic TP selection currently requires recipe DP=1 and PP=1; otherwise the +recipe's parallelism is retained. Automatic DP tuning remains disabled pending +CPU DP binding and capacity validation. No unused NUMA nodes are converted into +DP replicas. These checks establish head compatibility, not memory fit or +complete model/backend support. Explicit `vllm serve` CLI overrides still take +precedence over the generated YAML and are validated by vLLM at startup. + ### `gpu-memory-utilization` For every effective NUMA memory node, the policy calculates: diff --git a/tools/recipes/model_constraints.py b/tools/recipes/model_constraints.py new file mode 100644 index 000000000000..c25deb6f57e9 --- /dev/null +++ b/tools/recipes/model_constraints.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Read model head constraints without loading model weights.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class HeadConstraint: + name: str + attention_heads: int + kv_heads: int + + def supports(self, tp: int) -> bool: + if self.attention_heads % tp: + return False + if self.kv_heads >= tp: + return self.kv_heads % tp == 0 + return tp % self.kv_heads == 0 + + +def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"Cannot determine a positive integer for {name}: {value!r}") + return value + + +def extract_head_constraints( + hf_config: Any, *, include_vision: bool = True +) -> tuple[HeadConstraint, ...]: + """Extract decoder, global-attention, hybrid and encoder head constraints.""" + constraints = [] + if getattr(hf_config, "model_type", None) == "whisper": + for component in ("encoder", "decoder"): + heads = _positive_int( + getattr(hf_config, f"{component}_attention_heads", None), component + ) + constraints.append(HeadConstraint(component, heads, heads)) + return tuple(constraints) + + text = hf_config.get_text_config() + heads = _positive_int(getattr(text, "num_attention_heads", None), "attention heads") + kv = getattr(text, "num_key_value_heads", None) + if kv is None: + kv = heads + kv = _positive_int(kv, "KV heads") + constraints.append(HeadConstraint("text", heads, kv)) + + global_kv = getattr(text, "num_global_key_value_heads", None) + if global_kv is not None: + constraints.append( + HeadConstraint( + "global attention", heads, _positive_int(global_kv, "global KV") + ) + ) + + # Hybrid layers have separate head counts; use conservative divisibility + # rather than assuming they support standard attention's KV replication. + for field in ("linear_num_key_heads", "linear_num_value_heads"): + count = getattr(text, field, None) + if count is not None: + count = _positive_int(count, field) + constraints.append(HeadConstraint(field, count, count)) + + vision = getattr(hf_config, "vision_config", None) + if include_vision and vision is not None: + vision_heads = getattr(vision, "num_attention_heads", None) + if vision_heads is None: + vision_heads = getattr(vision, "num_heads", None) + vision_heads = _positive_int(vision_heads, "vision heads") + constraints.append(HeadConstraint("vision", vision_heads, vision_heads)) + return tuple(constraints) + + +def load_head_constraints(config: dict[str, Any]) -> tuple[HeadConstraint, ...]: + """Load the effective serving configuration through vLLM's config loader.""" + from vllm.transformers_utils.config import get_config + + hf_overrides = config.get("hf-overrides") + if isinstance(hf_overrides, str): + hf_overrides = json.loads(hf_overrides) + if hf_overrides is not None and not isinstance(hf_overrides, dict): + raise ValueError("Head detection requires hf-overrides to be a JSON object") + hf_config = get_config( + config["model"], + trust_remote_code=config.get("trust-remote-code", False), + revision=config.get("revision"), + code_revision=config.get("code-revision"), + config_format=config.get("config-format", "auto"), + hf_overrides_kw=hf_overrides, + ) + return extract_head_constraints( + hf_config, include_vision=config.get("mm-encoder-tp-mode") != "data" + ) + + +def suggest_tensor_parallel_size( + candidate: int, constraints: tuple[HeadConstraint, ...] +) -> int: + """Reduce an automatic power-of-two candidate until all head checks pass.""" + _positive_int(candidate, "TP candidate") + if candidate & (candidate - 1) or not constraints: + raise ValueError("Expected a power-of-two TP candidate and model constraints") + while not all(constraint.supports(candidate) for constraint in constraints): + candidate //= 2 + return candidate diff --git a/tools/recipes/runtime_tuning.py b/tools/recipes/runtime_tuning.py index e251de89d72e..b6b2d298d24d 100644 --- a/tools/recipes/runtime_tuning.py +++ b/tools/recipes/runtime_tuning.py @@ -19,6 +19,7 @@ from typing import Any from hardware_detection import HardwareInfo +from model_constraints import load_head_constraints, suggest_tensor_parallel_size from vllm.config.scheduler import SchedulerConfig @@ -119,19 +120,50 @@ def _resolve_tensor_parallel_size( workload: WorkloadHints, result: TuningResult, ) -> None: - del config, workload + del workload if hardware is None or hardware.numa_node_count < 1: return + if any( + config.get(key, 1) != 1 + for key in ("data-parallel-size", "pipeline-parallel-size") + ): + result.notes.append( + "tensor-parallel-size: kept recipe value; automatic TP selection " + "currently requires DP=1 and PP=1" + ) + return + # Draft policy: use the largest power-of-two TP that does not exceed the # effective NUMA-node count. This deliberately avoids hard-coding unusual # non-power-of-two topologies into the converter. candidate = 1 << (hardware.numa_node_count.bit_length() - 1) + try: + constraints = load_head_constraints(config) + except Exception as exc: + result.notes.append( + "tensor-parallel-size: kept recipe value (vLLM default TP=1 if absent); " + f"model head detection failed: {exc}" + ) + return + + for constraint in constraints: + result.notes.append( + f"detected {constraint.name}: " + f"attention_heads={constraint.attention_heads}, " + f"kv_heads={constraint.kv_heads}" + ) + selected = suggest_tensor_parallel_size(candidate, constraints) + if selected != candidate: + result.notes.append( + f"tensor-parallel-size: reduced NUMA-derived TP={candidate} to " + f"TP={selected} for model head partitioning/replication" + ) _record_override( result, "tensor-parallel-size", - candidate, - "derived from the effective NUMA-node count", + selected, + "derived from effective NUMA nodes and detected model head constraints", ) @@ -256,8 +288,8 @@ def _resolve_data_parallel_size( return result.notes.append( - "data-parallel-size: target QPS supplied; kept recipe value because " - "per-replica capacity is not known yet" + "data-parallel-size: automatic DP tuning is disabled; kept recipe value " + "because CPU DP auto-binding and per-replica capacity need validation" ) From 68088ed3927e91bce8918db78f2efd9e644a4134 Mon Sep 17 00:00:00 2001 From: zhangchengzhucufe-dev Date: Tue, 6 Oct 2026 18:19:54 +0800 Subject: [PATCH 061/952] =?UTF-8?q?Fix=20grammar=20in=20`BlockPool.reset?= =?UTF-8?q?=5Fprefix=5Fcache`=20docstring:=20"to=20invalid"=20=E2=86=92=20?= =?UTF-8?q?"to=20invalidate"=20(#59886)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: zhangchengzhucufe-dev --- vllm/v1/core/block_pool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 664804d50de4..75bdb40c24be 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -829,7 +829,7 @@ def evict_blocks(self, block_ids: set[int]) -> None: def reset_prefix_cache(self) -> bool: """Reset prefix cache. This function may be used in RLHF - flows to invalid prefix caching after the weights are updated, + flows to invalidate prefix caching after the weights are updated, or used for resetting prefix caching status for benchmarking. Returns: From 09e0ce3d82f404e842fdfc9ab34ba27e76f95a16 Mon Sep 17 00:00:00 2001 From: Shimi Bandiel Date: Tue, 6 Oct 2026 03:46:24 -0700 Subject: [PATCH 062/952] [Bugfix][Frontend] Apply model-default reasoning parser in GPU-less render server (#54835) Signed-off-by: Shimi Bandiel Co-authored-by: Claude Co-authored-by: Cyrus Leung --- .../entrypoints/launchers/render/__init__.py | 2 + .../launchers/render/test_app_state.py | 84 +++++++++++++++++++ .../scale_out/derender/test_derender.py | 66 +++++++++++++++ vllm/engine/arg_utils.py | 27 ++++-- .../entrypoints/launchers/render/app_state.py | 14 +++- vllm/entrypoints/launchers/render/entry.py | 5 ++ 6 files changed, 189 insertions(+), 9 deletions(-) create mode 100644 tests/entrypoints/launchers/render/__init__.py create mode 100644 tests/entrypoints/launchers/render/test_app_state.py diff --git a/tests/entrypoints/launchers/render/__init__.py b/tests/entrypoints/launchers/render/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/entrypoints/launchers/render/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/entrypoints/launchers/render/test_app_state.py b/tests/entrypoints/launchers/render/test_app_state.py new file mode 100644 index 000000000000..f0f54d9bafd3 --- /dev/null +++ b/tests/entrypoints/launchers/render/test_app_state.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace + +import pytest +from starlette.datastructures import State + +import vllm.entrypoints.launchers.render.app_state as app_state_mod +from vllm.config import ModelConfig, VllmConfig +from vllm.entrypoints.launchers.cli_args import make_arg_parser +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class _CaptureKwargs: + """Stands in for OnlineRenderer/OnlineDerenderer; records init kwargs.""" + + captured: list[dict] + + def __init__(self, **kwargs): + type(self).captured.append(kwargs) + + def warmup(self): + pass + + +def _render_cli_args(*argv: str) -> Namespace: + """Parse ``argv`` the way ``vllm launch render`` does, so new serve flags + carry their defaults instead of needing a hand-maintained ``Namespace``.""" + args = make_arg_parser(FlexibleArgumentParser()).parse_args(list(argv)) + if args.model_tag is not None: + args.model = args.model_tag + return args + + +@pytest.mark.asyncio +async def test_render_app_state_uses_config_resolved_reasoning_parser(monkeypatch): + """gpt_oss defaults its reasoning parser through verify_and_update_config; + the render server must pick up that resolved value like the main API + server does, otherwise harmony markup leaks unparsed into derender output. + """ + model_config = ModelConfig("openai/gpt-oss-20b") + vllm_config = VllmConfig(model_config=model_config) + assert vllm_config.structured_outputs_config.reasoning_parser == "openai_gptoss" + + class _Renderer(_CaptureKwargs): + captured = [] + + class _Derenderer(_CaptureKwargs): + captured = [] + + async def _noop_async(*args, **kwargs): + return None + + monkeypatch.setattr(app_state_mod, "renderer_from_config", lambda cfg: object()) + monkeypatch.setattr(app_state_mod, "OnlineRenderer", _Renderer) + monkeypatch.setattr(app_state_mod, "OnlineDerenderer", _Derenderer) + monkeypatch.setattr(app_state_mod, "ServingTokenization", lambda *a, **kw: object()) + monkeypatch.setattr(app_state_mod, "init_render_state", lambda *a, **kw: None) + monkeypatch.setattr(app_state_mod, "init_endpoint_plugins_state", _noop_async) + + args = _render_cli_args("openai/gpt-oss-20b") + + await app_state_mod.init_render_app_state(vllm_config, State(), args) + + assert _Renderer.captured[0]["reasoning_parser"] == "openai_gptoss" + assert _Derenderer.captured[0]["reasoning_parser"] == "openai_gptoss" + + +def test_explicit_reasoning_parser_flag_wins_over_model_default(): + """An explicit --reasoning-parser must survive VllmConfig construction + for models that define their own default (gpt_oss); the default only + applies when the resolved value is empty.""" + from vllm.engine.arg_utils import AsyncEngineArgs + + engine_args = AsyncEngineArgs( + model="openai/gpt-oss-20b", reasoning_parser="deepseek_r1" + ) + model_config = engine_args.create_model_config() + vllm_config = VllmConfig( + model_config=model_config, + structured_outputs_config=engine_args.create_structured_outputs_config(), + ) + + assert vllm_config.structured_outputs_config.reasoning_parser == "deepseek_r1" diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 659aaed1aabe..40b051304e22 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -1298,6 +1298,72 @@ async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer): assert "Four" in content +@pytest.fixture(scope="module") +def harmony_default_server(): + """gpt-oss render server with NO parser flags. + + Regression fixture: the model-default reasoning parser + ("openai_gptoss", applied by verify_and_update_config) must be + resolved by the render server without any --reasoning-parser flag, + otherwise harmony markup leaks unparsed into derender output. + """ + _ensure_harmony_vocab() + with RemoteLaunchRenderServer( + HARMONY_MODEL, ["--trust-remote-code"] + ) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def harmony_default_client(harmony_default_server): + async with httpx.AsyncClient( + base_url=harmony_default_server.url_for(""), timeout=60.0 + ) as http_client: + yield http_client + + +@pytest.mark.asyncio +async def test_e2e_harmony_reasoning_default_parser( + harmony_default_client, harmony_tokenizer +): + """GPT-OSS reasoning parses on a bare launch (no parser flags).""" + messages = [{"role": "user", "content": "Add 2 and 3."}] + gen_req = await _e2e_render_chat(harmony_default_client, HARMONY_MODEL, messages) + + reasoning_text = "The user wants 2 plus 3." + answer_text = "The answer is 5." + assistant_msg = { + "role": "assistant", + "thinking": reasoning_text, + "content": answer_text, + } + output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg) + + decoded = harmony_tokenizer.decode(output_ids) + if reasoning_text not in decoded: + pytest.skip("Harmony template did not render thinking") + + resp = await harmony_default_client.post( + "/v1/chat/completions/derender", + json={ + "model": HARMONY_MODEL, + "generate_response": _e2e_generate_response(output_ids), + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": HARMONY_MODEL, + "messages": messages, + "include_reasoning": True, + }, + }, + ) + assert resp.status_code == 200, resp.text + msg = resp.json()["choices"][0]["message"] + assert msg["reasoning"] is not None + assert reasoning_text in msg["reasoning"] + assert answer_text in (msg["content"] or "") + assert "<|channel|>" not in (msg["content"] or "") + + @pytest.mark.asyncio async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer): """GPT-OSS reasoning: analysis channel extracted.""" diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 4f1f41728948..98a0b7cd88fd 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2136,6 +2136,25 @@ def create_watermark_config(self) -> WatermarkConfig | None: cfg = json.loads(cfg) return WatermarkConfig(**cfg) + def create_structured_outputs_config(self) -> StructuredOutputsConfig: + """Merge frontend parser flags into the structured outputs config. + + Mutates `self.structured_outputs_config` in place and returns it + (not a copy). Model-specific defaults (e.g. gpt_oss -> + "openai_gptoss") are applied later by `verify_and_update_config` + only when the resolved value is still empty, so explicit CLI flags + take precedence. + """ + if self.reasoning_parser: + self.structured_outputs_config.reasoning_parser = self.reasoning_parser + + if self.reasoning_parser_plugin: + self.structured_outputs_config.reasoning_parser_plugin = ( + self.reasoning_parser_plugin + ) + + return self.structured_outputs_config + def create_observability_config(self) -> ObservabilityConfig: return ObservabilityConfig( show_hidden_metrics_for_version=self.show_hidden_metrics_for_version, @@ -2728,13 +2747,7 @@ def create_engine_config( load_config = self.create_load_config() # Pass reasoning_parser into StructuredOutputsConfig - if self.reasoning_parser: - self.structured_outputs_config.reasoning_parser = self.reasoning_parser - - if self.reasoning_parser_plugin: - self.structured_outputs_config.reasoning_parser_plugin = ( - self.reasoning_parser_plugin - ) + self.create_structured_outputs_config() observability_config = self.create_observability_config() diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py index bbd9a6127521..33aebb9bc39e 100644 --- a/vllm/entrypoints/launchers/render/app_state.py +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -50,6 +50,16 @@ async def init_render_app_state( default_chat_template_kwargs = resolve_default_chat_template_kwargs(args) state.tool_server = await init_tool_server(args) + # The config-resolved reasoning parser carries both the CLI flag (merged + # in the entrypoint via `create_structured_outputs_config`) and any + # model-specific default applied by `verify_and_update_config` + # (e.g. "openai_gptoss" for gpt_oss), matching the main API server. + # Keep `args.reasoning_parser` as the first source so callers that build + # a VllmConfig without that merge still honor an explicit flag. + reasoning_parser = ( + args.reasoning_parser or vllm_config.structured_outputs_config.reasoning_parser + ) + state.online_renderer = OnlineRenderer( model_config=vllm_config.model_config, renderer=renderer, @@ -62,7 +72,7 @@ async def init_render_app_state( exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, tool_strict_level=args.tool_strict_level, - reasoning_parser=args.reasoning_parser, + reasoning_parser=reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, ) @@ -79,7 +89,7 @@ async def init_render_app_state( exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, tool_parser=args.tool_call_parser, tool_strict_level=args.tool_strict_level, - reasoning_parser=args.reasoning_parser, + reasoning_parser=reasoning_parser, default_chat_template_kwargs=default_chat_template_kwargs, log_error_stack=args.log_error_stack, ) diff --git a/vllm/entrypoints/launchers/render/entry.py b/vllm/entrypoints/launchers/render/entry.py index b42e1f591171..92c29ea320df 100644 --- a/vllm/entrypoints/launchers/render/entry.py +++ b/vllm/entrypoints/launchers/render/entry.py @@ -87,9 +87,14 @@ def _interrupt_init(*_) -> None: # cache space warning from CpuPlatform.check_and_update_config. envs.VLLM_CPU_KVCACHE_SPACE = 0 + # Merge frontend parser flags (e.g. --reasoning-parser) into the + # structured outputs config, mirroring `create_engine_config`, which this + # GPU-less entrypoint bypasses. Model-specific defaults are then applied + # by `verify_and_update_config` during VllmConfig construction. vllm_config = VllmConfig( model_config=model_config, logging_config=engine_args.create_logging_config(), + structured_outputs_config=engine_args.create_structured_outputs_config(), ) shutdown_task = await build_and_serve_renderer( vllm_config, listen_address, sock, args From 3313825983dad6fd4291986e7f9ba647c059d1ff Mon Sep 17 00:00:00 2001 From: Mia Ojeda Date: Tue, 6 Oct 2026 05:25:51 -0600 Subject: [PATCH 063/952] [Bugfix][DiffusionGemma] Read the causal mask as int32 and compile the sample step (#59992) Signed-off-by: ojeda-e --- csrc/cpu/cpu_attn.cpp | 10 ++++++---- csrc/cpu/cpu_attn_impl.hpp | 13 +++++++------ tests/kernels/attention/test_cpu_attn.py | 2 +- tests/v1/sample/test_diffusion_gemma_reads.py | 4 ++++ vllm/model_executor/models/diffusion_gemma.py | 11 ++++++----- 5 files changed, 24 insertions(+), 16 deletions(-) diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 13e430fb4192..afdc0e0ecae8 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -78,8 +78,9 @@ torch::Tensor get_scheduler_metadata( input.causal = causal; input.isa = isa; input.enable_kv_split = enable_kv_split; - input.dynamic_causal = - dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + input.dynamic_causal = dynamic_causal.has_value() + ? dynamic_causal->data_ptr() + : nullptr; const int64_t kv_cache_idx = static_cast(parse_fp8_kv_dtype(kv_cache_dtype)); @@ -256,8 +257,9 @@ void cpu_attention_with_kv_cache( input.s_aux = nullptr; input.s_aux_is_bf16 = false; } - input.dynamic_causal = - dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + input.dynamic_causal = dynamic_causal.has_value() + ? dynamic_causal->data_ptr() + : nullptr; input.scale = scale; input.causal = causal; input.sliding_window_size = sliding_window; diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index d914e150636c..0666ee413f13 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -397,7 +397,7 @@ class AttentionScheduler { int32_t max_num_q_per_iter; // max Q head num can be hold in registers int32_t kv_block_alignment; // context length alignment requirement bool enable_kv_split; - bool* dynamic_causal; + int32_t* dynamic_causal; }; static constexpr int32_t MaxQTileIterNum = 128; @@ -450,7 +450,7 @@ class AttentionScheduler { const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; const bool req_causal = - is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + is_dynamic_causal ? input.dynamic_causal[req_id] != 0 : causal; const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; @@ -492,7 +492,7 @@ class AttentionScheduler { const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; const bool req_causal = - is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + is_dynamic_causal ? input.dynamic_causal[req_id] != 0 : causal; const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; @@ -884,7 +884,7 @@ struct AttentionInput { // provided as fp32 and executed in full float precision. const void* s_aux; bool s_aux_is_bf16; - bool* dynamic_causal; + int32_t* dynamic_causal; float scale; bool causal; int32_t sliding_window_size; @@ -1571,7 +1571,7 @@ class AttentionMainLoop { const float* alibi_slopes = input->alibi_slopes; const void* s_aux = input->s_aux; const bool s_aux_is_bf16 = input->s_aux_is_bf16; - const bool* dynamic_causal = input->dynamic_causal; + const int32_t* dynamic_causal = input->dynamic_causal; const bool is_dynamic_causal = dynamic_causal != nullptr; const bool causal = input->causal; @@ -1658,7 +1658,8 @@ class AttentionMainLoop { const int32_t current_group_idx = current_workitem_group->req_id; const int32_t current_group_causal = - is_dynamic_causal ? dynamic_causal[current_group_idx] : causal; + is_dynamic_causal ? dynamic_causal[current_group_idx] != 0 + : causal; auto [sliding_window_left, sliding_window_right] = AttentionScheduler::calcu_sliding_window_size( sliding_window_size, current_group_causal); diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index f79b1b15d2d9..7f88f76d0b79 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -472,7 +472,7 @@ def varlen_with_paged_kv( scale = head_size**-0.5 token_num = sum(query_lens) dynamic_causal_tensor = ( - torch.tensor(dynamic_causal, dtype=torch.bool) + torch.tensor(dynamic_causal, dtype=torch.int32) if dynamic_causal is not None else None ) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 7bd04c0379a6..4abadc254d0a 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -227,6 +227,10 @@ def test_single_step_tile_skips_self_conditioning(): assert not states.self_conditioning_embeds[0].any() +def test_sample_step_is_compiled(): + assert hasattr(_compiled_sample_step, "_torchdynamo_orig_callable") + + @pytest.mark.parametrize("stance", ["default", "force_eager"]) def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance): # The model's embeddings are bf16 while the buffer is fp32. Compiled code diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 404c154b5468..1b27df68abb0 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -445,11 +445,6 @@ def _mask_rows_to_allowed( return logits if out is None else out -# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, -# on compute_sc and on sizes that coincide with the state buffers. That set -# is small but passes Dynamo's default of 8, after which every step would run -# eager: about twice as slow for a self-conditioned step. -@torch._dynamo.config.patch(recompile_limit=64) @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _denoise_temperature( step_tensor: torch.Tensor, @@ -464,6 +459,12 @@ def _denoise_temperature( return t_min + (t_max - t_min) * (remaining / max_denoising_steps) +# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, +# on compute_sc and on sizes that coincide with the state buffers. That set +# is small but passes Dynamo's default of 8, after which every step would run +# eager: about twice as slow for a self-conditioned step. +@torch._dynamo.config.patch(recompile_limit=64) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _compiled_sample_step( # Per-position statistics of the temperature-scaled logits, from # sample_row_stats: [num_decode, CL] each, and the softmax From 3e182185aa5d143b0e69c44607f58a0cc55f3971 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 6 Oct 2026 08:15:12 -0400 Subject: [PATCH 064/952] [Deprecation] Deprecate sonet dataset as scheduled (#60129) Signed-off-by: yewentao256 --- benchmarks/sonnet.txt | 518 ------------------ docs/benchmarking/cli.md | 10 +- pyproject.toml | 1 - rust/src/bench/AGENTS.md | 1 - rust/src/bench/README.md | 24 +- rust/src/bench/src/benchmark.rs | 23 - rust/src/bench/src/cli.rs | 15 - rust/src/bench/src/config.rs | 9 +- rust/src/bench/src/datasets/mod.rs | 1 - rust/src/bench/src/datasets/sonnet.rs | 131 ----- rust/src/bench/src/datasets/sonnet.txt | 518 ------------------ rust/src/bench/src/multi_turn.rs | 5 - rust/src/bench/tests/python_serve_flags.txt | 3 - vllm/benchmarks/datasets/__init__.py | 2 - .../datasets/create_txt_slices_dataset.py | 6 +- vllm/benchmarks/datasets/datasets.py | 165 +----- vllm/benchmarks/serve.py | 11 +- vllm/benchmarks/throughput.py | 11 +- 18 files changed, 22 insertions(+), 1432 deletions(-) delete mode 100644 benchmarks/sonnet.txt delete mode 100644 rust/src/bench/src/datasets/sonnet.rs delete mode 100644 rust/src/bench/src/datasets/sonnet.txt diff --git a/benchmarks/sonnet.txt b/benchmarks/sonnet.txt deleted file mode 100644 index 34c444e8ce8e..000000000000 --- a/benchmarks/sonnet.txt +++ /dev/null @@ -1,518 +0,0 @@ -FROM fairest creatures we desire increase, -That thereby beauty's rose might never die, -But as the riper should by time decease, -His tender heir might bear his memory: -But thou, contracted to thine own bright eyes, -Feed'st thy light'st flame with self-substantial fuel, -Making a famine where abundance lies, -Thyself thy foe, to thy sweet self too cruel. -Thou that art now the world's fresh ornament -And only herald to the gaudy spring, -Within thine own bud buriest thy content -And, tender churl, makest waste in niggarding. -Pity the world, or else this glutton be, -To eat the world's due, by the grave and thee. -When forty winters shall beseige thy brow, -And dig deep trenches in thy beauty's field, -Thy youth's proud livery, so gazed on now, -Will be a tatter'd weed, of small worth held: -Then being ask'd where all thy beauty lies, -Where all the treasure of thy lusty days, -To say, within thine own deep-sunken eyes, -Were an all-eating shame and thriftless praise. -How much more praise deserved thy beauty's use, -If thou couldst answer 'This fair child of mine -Shall sum my count and make my old excuse,' -Proving his beauty by succession thine! -This were to be new made when thou art old, -And see thy blood warm when thou feel'st it cold. -Look in thy glass, and tell the face thou viewest -Now is the time that face should form another; -Whose fresh repair if now thou not renewest, -Thou dost beguile the world, unbless some mother. -For where is she so fair whose unear'd womb -Disdains the tillage of thy husbandry? -Or who is he so fond will be the tomb -Of his self-love, to stop posterity? -Thou art thy mother's glass, and she in thee -Calls back the lovely April of her prime: -So thou through windows of thine age shall see -Despite of wrinkles this thy golden time. -But if thou live, remember'd not to be, -Die single, and thine image dies with thee. -Unthrifty loveliness, why dost thou spend -Upon thyself thy beauty's legacy? -Nature's bequest gives nothing but doth lend, -And being frank she lends to those are free. -Then, beauteous niggard, why dost thou abuse -The bounteous largess given thee to give? -Profitless usurer, why dost thou use -So great a sum of sums, yet canst not live? -For having traffic with thyself alone, -Thou of thyself thy sweet self dost deceive. -Then how, when nature calls thee to be gone, -What acceptable audit canst thou leave? -Thy unused beauty must be tomb'd with thee, -Which, used, lives th' executor to be. -Those hours, that with gentle work did frame -The lovely gaze where every eye doth dwell, -Will play the tyrants to the very same -And that unfair which fairly doth excel: -For never-resting time leads summer on -To hideous winter and confounds him there; -Sap cheque'd with frost and lusty leaves quite gone, -Beauty o'ersnow'd and bareness every where: -Then, were not summer's distillation left, -A liquid prisoner pent in walls of glass, -Beauty's effect with beauty were bereft, -Nor it nor no remembrance what it was: -But flowers distill'd though they with winter meet, -Leese but their show; their substance still lives sweet. -Then let not winter's ragged hand deface -In thee thy summer, ere thou be distill'd: -Make sweet some vial; treasure thou some place -With beauty's treasure, ere it be self-kill'd. -That use is not forbidden usury, -Which happies those that pay the willing loan; -That's for thyself to breed another thee, -Or ten times happier, be it ten for one; -Ten times thyself were happier than thou art, -If ten of thine ten times refigured thee: -Then what could death do, if thou shouldst depart, -Leaving thee living in posterity? -Be not self-will'd, for thou art much too fair -To be death's conquest and make worms thine heir. -Lo! in the orient when the gracious light -Lifts up his burning head, each under eye -Doth homage to his new-appearing sight, -Serving with looks his sacred majesty; -And having climb'd the steep-up heavenly hill, -Resembling strong youth in his middle age, -yet mortal looks adore his beauty still, -Attending on his golden pilgrimage; -But when from highmost pitch, with weary car, -Like feeble age, he reeleth from the day, -The eyes, 'fore duteous, now converted are -From his low tract and look another way: -So thou, thyself out-going in thy noon, -Unlook'd on diest, unless thou get a son. -Music to hear, why hear'st thou music sadly? -Sweets with sweets war not, joy delights in joy. -Why lovest thou that which thou receivest not gladly, -Or else receivest with pleasure thine annoy? -If the true concord of well-tuned sounds, -By unions married, do offend thine ear, -They do but sweetly chide thee, who confounds -In singleness the parts that thou shouldst bear. -Mark how one string, sweet husband to another, -Strikes each in each by mutual ordering, -Resembling sire and child and happy mother -Who all in one, one pleasing note do sing: -Whose speechless song, being many, seeming one, -Sings this to thee: 'thou single wilt prove none.' -Is it for fear to wet a widow's eye -That thou consumest thyself in single life? -Ah! if thou issueless shalt hap to die. -The world will wail thee, like a makeless wife; -The world will be thy widow and still weep -That thou no form of thee hast left behind, -When every private widow well may keep -By children's eyes her husband's shape in mind. -Look, what an unthrift in the world doth spend -Shifts but his place, for still the world enjoys it; -But beauty's waste hath in the world an end, -And kept unused, the user so destroys it. -No love toward others in that bosom sits -That on himself such murderous shame commits. -For shame! deny that thou bear'st love to any, -Who for thyself art so unprovident. -Grant, if thou wilt, thou art beloved of many, -But that thou none lovest is most evident; -For thou art so possess'd with murderous hate -That 'gainst thyself thou stick'st not to conspire. -Seeking that beauteous roof to ruinate -Which to repair should be thy chief desire. -O, change thy thought, that I may change my mind! -Shall hate be fairer lodged than gentle love? -Be, as thy presence is, gracious and kind, -Or to thyself at least kind-hearted prove: -Make thee another self, for love of me, -That beauty still may live in thine or thee. -As fast as thou shalt wane, so fast thou growest -In one of thine, from that which thou departest; -And that fresh blood which youngly thou bestowest -Thou mayst call thine when thou from youth convertest. -Herein lives wisdom, beauty and increase: -Without this, folly, age and cold decay: -If all were minded so, the times should cease -And threescore year would make the world away. -Let those whom Nature hath not made for store, -Harsh featureless and rude, barrenly perish: -Look, whom she best endow'd she gave the more; -Which bounteous gift thou shouldst in bounty cherish: -She carved thee for her seal, and meant thereby -Thou shouldst print more, not let that copy die. -When I do count the clock that tells the time, -And see the brave day sunk in hideous night; -When I behold the violet past prime, -And sable curls all silver'd o'er with white; -When lofty trees I see barren of leaves -Which erst from heat did canopy the herd, -And summer's green all girded up in sheaves -Borne on the bier with white and bristly beard, -Then of thy beauty do I question make, -That thou among the wastes of time must go, -Since sweets and beauties do themselves forsake -And die as fast as they see others grow; -And nothing 'gainst Time's scythe can make defence -Save breed, to brave him when he takes thee hence. -O, that you were yourself! but, love, you are -No longer yours than you yourself here live: -Against this coming end you should prepare, -And your sweet semblance to some other give. -So should that beauty which you hold in lease -Find no determination: then you were -Yourself again after yourself's decease, -When your sweet issue your sweet form should bear. -Who lets so fair a house fall to decay, -Which husbandry in honour might uphold -Against the stormy gusts of winter's day -And barren rage of death's eternal cold? -O, none but unthrifts! Dear my love, you know -You had a father: let your son say so. -Not from the stars do I my judgment pluck; -And yet methinks I have astronomy, -But not to tell of good or evil luck, -Of plagues, of dearths, or seasons' quality; -Nor can I fortune to brief minutes tell, -Pointing to each his thunder, rain and wind, -Or say with princes if it shall go well, -By oft predict that I in heaven find: -But from thine eyes my knowledge I derive, -And, constant stars, in them I read such art -As truth and beauty shall together thrive, -If from thyself to store thou wouldst convert; -Or else of thee this I prognosticate: -Thy end is truth's and beauty's doom and date. -When I consider every thing that grows -Holds in perfection but a little moment, -That this huge stage presenteth nought but shows -Whereon the stars in secret influence comment; -When I perceive that men as plants increase, -Cheered and cheque'd even by the self-same sky, -Vaunt in their youthful sap, at height decrease, -And wear their brave state out of memory; -Then the conceit of this inconstant stay -Sets you most rich in youth before my sight, -Where wasteful Time debateth with Decay, -To change your day of youth to sullied night; -And all in war with Time for love of you, -As he takes from you, I engraft you new. -But wherefore do not you a mightier way -Make war upon this bloody tyrant, Time? -And fortify yourself in your decay -With means more blessed than my barren rhyme? -Now stand you on the top of happy hours, -And many maiden gardens yet unset -With virtuous wish would bear your living flowers, -Much liker than your painted counterfeit: -So should the lines of life that life repair, -Which this, Time's pencil, or my pupil pen, -Neither in inward worth nor outward fair, -Can make you live yourself in eyes of men. -To give away yourself keeps yourself still, -And you must live, drawn by your own sweet skill. -Who will believe my verse in time to come, -If it were fill'd with your most high deserts? -Though yet, heaven knows, it is but as a tomb -Which hides your life and shows not half your parts. -If I could write the beauty of your eyes -And in fresh numbers number all your graces, -The age to come would say 'This poet lies: -Such heavenly touches ne'er touch'd earthly faces.' -So should my papers yellow'd with their age -Be scorn'd like old men of less truth than tongue, -And your true rights be term'd a poet's rage -And stretched metre of an antique song: -But were some child of yours alive that time, -You should live twice; in it and in my rhyme. -Shall I compare thee to a summer's day? -Thou art more lovely and more temperate: -Rough winds do shake the darling buds of May, -And summer's lease hath all too short a date: -Sometime too hot the eye of heaven shines, -And often is his gold complexion dimm'd; -And every fair from fair sometime declines, -By chance or nature's changing course untrimm'd; -But thy eternal summer shall not fade -Nor lose possession of that fair thou owest; -Nor shall Death brag thou wander'st in his shade, -When in eternal lines to time thou growest: -So long as men can breathe or eyes can see, -So long lives this and this gives life to thee. -Devouring Time, blunt thou the lion's paws, -And make the earth devour her own sweet brood; -Pluck the keen teeth from the fierce tiger's jaws, -And burn the long-lived phoenix in her blood; -Make glad and sorry seasons as thou fleets, -And do whate'er thou wilt, swift-footed Time, -To the wide world and all her fading sweets; -But I forbid thee one most heinous crime: -O, carve not with thy hours my love's fair brow, -Nor draw no lines there with thine antique pen; -Him in thy course untainted do allow -For beauty's pattern to succeeding men. -Yet, do thy worst, old Time: despite thy wrong, -My love shall in my verse ever live young. -A woman's face with Nature's own hand painted -Hast thou, the master-mistress of my passion; -A woman's gentle heart, but not acquainted -With shifting change, as is false women's fashion; -An eye more bright than theirs, less false in rolling, -Gilding the object whereupon it gazeth; -A man in hue, all 'hues' in his controlling, -Much steals men's eyes and women's souls amazeth. -And for a woman wert thou first created; -Till Nature, as she wrought thee, fell a-doting, -And by addition me of thee defeated, -By adding one thing to my purpose nothing. -But since she prick'd thee out for women's pleasure, -Mine be thy love and thy love's use their treasure. -So is it not with me as with that Muse -Stirr'd by a painted beauty to his verse, -Who heaven itself for ornament doth use -And every fair with his fair doth rehearse -Making a couplement of proud compare, -With sun and moon, with earth and sea's rich gems, -With April's first-born flowers, and all things rare -That heaven's air in this huge rondure hems. -O' let me, true in love, but truly write, -And then believe me, my love is as fair -As any mother's child, though not so bright -As those gold candles fix'd in heaven's air: -Let them say more than like of hearsay well; -I will not praise that purpose not to sell. -My glass shall not persuade me I am old, -So long as youth and thou are of one date; -But when in thee time's furrows I behold, -Then look I death my days should expiate. -For all that beauty that doth cover thee -Is but the seemly raiment of my heart, -Which in thy breast doth live, as thine in me: -How can I then be elder than thou art? -O, therefore, love, be of thyself so wary -As I, not for myself, but for thee will; -Bearing thy heart, which I will keep so chary -As tender nurse her babe from faring ill. -Presume not on thy heart when mine is slain; -Thou gavest me thine, not to give back again. -As an unperfect actor on the stage -Who with his fear is put besides his part, -Or some fierce thing replete with too much rage, -Whose strength's abundance weakens his own heart. -So I, for fear of trust, forget to say -The perfect ceremony of love's rite, -And in mine own love's strength seem to decay, -O'ercharged with burden of mine own love's might. -O, let my books be then the eloquence -And dumb presagers of my speaking breast, -Who plead for love and look for recompense -More than that tongue that more hath more express'd. -O, learn to read what silent love hath writ: -To hear with eyes belongs to love's fine wit. -Mine eye hath play'd the painter and hath stell'd -Thy beauty's form in table of my heart; -My body is the frame wherein 'tis held, -And perspective it is the painter's art. -For through the painter must you see his skill, -To find where your true image pictured lies; -Which in my bosom's shop is hanging still, -That hath his windows glazed with thine eyes. -Now see what good turns eyes for eyes have done: -Mine eyes have drawn thy shape, and thine for me -Are windows to my breast, where-through the sun -Delights to peep, to gaze therein on thee; -Yet eyes this cunning want to grace their art; -They draw but what they see, know not the heart. -Let those who are in favour with their stars -Of public honour and proud titles boast, -Whilst I, whom fortune of such triumph bars, -Unlook'd for joy in that I honour most. -Great princes' favourites their fair leaves spread -But as the marigold at the sun's eye, -And in themselves their pride lies buried, -For at a frown they in their glory die. -The painful warrior famoused for fight, -After a thousand victories once foil'd, -Is from the book of honour razed quite, -And all the rest forgot for which he toil'd: -Then happy I, that love and am beloved -Where I may not remove nor be removed. -Lord of my love, to whom in vassalage -Thy merit hath my duty strongly knit, -To thee I send this written embassage, -To witness duty, not to show my wit: -Duty so great, which wit so poor as mine -May make seem bare, in wanting words to show it, -But that I hope some good conceit of thine -In thy soul's thought, all naked, will bestow it; -Till whatsoever star that guides my moving -Points on me graciously with fair aspect -And puts apparel on my tatter'd loving, -To show me worthy of thy sweet respect: -Then may I dare to boast how I do love thee; -Till then not show my head where thou mayst prove me. -Weary with toil, I haste me to my bed, -The dear repose for limbs with travel tired; -But then begins a journey in my head, -To work my mind, when body's work's expired: -For then my thoughts, from far where I abide, -Intend a zealous pilgrimage to thee, -And keep my drooping eyelids open wide, -Looking on darkness which the blind do see -Save that my soul's imaginary sight -Presents thy shadow to my sightless view, -Which, like a jewel hung in ghastly night, -Makes black night beauteous and her old face new. -Lo! thus, by day my limbs, by night my mind, -For thee and for myself no quiet find. -How can I then return in happy plight, -That am debarr'd the benefit of rest? -When day's oppression is not eased by night, -But day by night, and night by day, oppress'd? -And each, though enemies to either's reign, -Do in consent shake hands to torture me; -The one by toil, the other to complain -How far I toil, still farther off from thee. -I tell the day, to please them thou art bright -And dost him grace when clouds do blot the heaven: -So flatter I the swart-complexion'd night, -When sparkling stars twire not thou gild'st the even. -But day doth daily draw my sorrows longer -And night doth nightly make grief's strength seem stronger. -When, in disgrace with fortune and men's eyes, -I all alone beweep my outcast state -And trouble deal heaven with my bootless cries -And look upon myself and curse my fate, -Wishing me like to one more rich in hope, -Featured like him, like him with friends possess'd, -Desiring this man's art and that man's scope, -With what I most enjoy contented least; -Yet in these thoughts myself almost despising, -Haply I think on thee, and then my state, -Like to the lark at break of day arising -From sullen earth, sings hymns at heaven's gate; -For thy sweet love remember'd such wealth brings -That then I scorn to change my state with kings. -When to the sessions of sweet silent thought -I summon up remembrance of things past, -I sigh the lack of many a thing I sought, -And with old woes new wail my dear time's waste: -Then can I drown an eye, unused to flow, -For precious friends hid in death's dateless night, -And weep afresh love's long since cancell'd woe, -And moan the expense of many a vanish'd sight: -Then can I grieve at grievances foregone, -And heavily from woe to woe tell o'er -The sad account of fore-bemoaned moan, -Which I new pay as if not paid before. -But if the while I think on thee, dear friend, -All losses are restored and sorrows end. -Thy bosom is endeared with all hearts, -Which I by lacking have supposed dead, -And there reigns love and all love's loving parts, -And all those friends which I thought buried. -How many a holy and obsequious tear -Hath dear religious love stol'n from mine eye -As interest of the dead, which now appear -But things removed that hidden in thee lie! -Thou art the grave where buried love doth live, -Hung with the trophies of my lovers gone, -Who all their parts of me to thee did give; -That due of many now is thine alone: -Their images I loved I view in thee, -And thou, all they, hast all the all of me. -If thou survive my well-contented day, -When that churl Death my bones with dust shall cover, -And shalt by fortune once more re-survey -These poor rude lines of thy deceased lover, -Compare them with the bettering of the time, -And though they be outstripp'd by every pen, -Reserve them for my love, not for their rhyme, -Exceeded by the height of happier men. -O, then vouchsafe me but this loving thought: -'Had my friend's Muse grown with this growing age, -A dearer birth than this his love had brought, -To march in ranks of better equipage: -But since he died and poets better prove, -Theirs for their style I'll read, his for his love.' -Full many a glorious morning have I seen -Flatter the mountain-tops with sovereign eye, -Kissing with golden face the meadows green, -Gilding pale streams with heavenly alchemy; -Anon permit the basest clouds to ride -With ugly rack on his celestial face, -And from the forlorn world his visage hide, -Stealing unseen to west with this disgrace: -Even so my sun one early morn did shine -With all triumphant splendor on my brow; -But out, alack! he was but one hour mine; -The region cloud hath mask'd him from me now. -Yet him for this my love no whit disdaineth; -Suns of the world may stain when heaven's sun staineth. -Why didst thou promise such a beauteous day, -And make me travel forth without my cloak, -To let base clouds o'ertake me in my way, -Hiding thy bravery in their rotten smoke? -'Tis not enough that through the cloud thou break, -To dry the rain on my storm-beaten face, -For no man well of such a salve can speak -That heals the wound and cures not the disgrace: -Nor can thy shame give physic to my grief; -Though thou repent, yet I have still the loss: -The offender's sorrow lends but weak relief -To him that bears the strong offence's cross. -Ah! but those tears are pearl which thy love sheds, -And they are rich and ransom all ill deeds. -No more be grieved at that which thou hast done: -Roses have thorns, and silver fountains mud; -Clouds and eclipses stain both moon and sun, -And loathsome canker lives in sweetest bud. -All men make faults, and even I in this, -Authorizing thy trespass with compare, -Myself corrupting, salving thy amiss, -Excusing thy sins more than thy sins are; -For to thy sensual fault I bring in sense-- -Thy adverse party is thy advocate-- -And 'gainst myself a lawful plea commence: -Such civil war is in my love and hate -That I an accessary needs must be -To that sweet thief which sourly robs from me. -Let me confess that we two must be twain, -Although our undivided loves are one: -So shall those blots that do with me remain -Without thy help by me be borne alone. -In our two loves there is but one respect, -Though in our lives a separable spite, -Which though it alter not love's sole effect, -Yet doth it steal sweet hours from love's delight. -I may not evermore acknowledge thee, -Lest my bewailed guilt should do thee shame, -Nor thou with public kindness honour me, -Unless thou take that honour from thy name: -But do not so; I love thee in such sort -As, thou being mine, mine is thy good report. -As a decrepit father takes delight -To see his active child do deeds of youth, -So I, made lame by fortune's dearest spite, -Take all my comfort of thy worth and truth. -For whether beauty, birth, or wealth, or wit, -Or any of these all, or all, or more, -Entitled in thy parts do crowned sit, -I make my love engrafted to this store: -So then I am not lame, poor, nor despised, -Whilst that this shadow doth such substance give -That I in thy abundance am sufficed -And by a part of all thy glory live. -Look, what is best, that best I wish in thee: -This wish I have; then ten times happy me! \ No newline at end of file diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index d05f0c233e84..a1001149aa1f 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -23,7 +23,6 @@ th { | ShareGPT4V (Image) | ✅ | ✅ | `wget https://huggingface.co/datasets/Lin-Chen/ShareGPT4V/resolve/main/sharegpt4v_instruct_gpt4-vision_cap100k.json`
Note that the images need to be downloaded separately. For example, to download COCO's 2017 Train images:
`wget http://images.cocodataset.org/zips/train2017.zip` | | ShareGPT4Video (Video) | ✅ | ✅ | `git clone https://huggingface.co/datasets/ShareGPT4Video/ShareGPT4Video` | | BurstGPT | ✅ | ✅ | `wget https://github.com/HPMLL/BurstGPT/releases/download/v1.1/BurstGPT_without_fails_2.csv` | -| Sonnet (deprecated) | ✅ | ✅ | Local file: `benchmarks/sonnet.txt` | | Random | ✅ | ✅ | `synthetic` | | RandomMultiModal (Image/Video) | ✅ | ✅ | `synthetic` | | RandomForReranking | ✅ | ✅ | `synthetic` | @@ -773,16 +772,17 @@ Using KV cache metrics for load pattern configuration: ```bash vllm bench throughput \ --model NousResearch/Hermes-3-Llama-3.1-8B \ - --dataset-name sonnet \ - --dataset-path vllm/benchmarks/sonnet.txt \ + --dataset-name random \ + --random-input-len 500 \ + --random-output-len 150 \ --num-prompts 10 ``` If successful, you will see the following output ```text -Throughput: 7.15 requests/s, 4656.00 total tokens/s, 1072.15 output tokens/s -Total num prompt tokens: 5014 +Throughput: 7.15 requests/s, 4647.50 total tokens/s, 1072.50 output tokens/s +Total num prompt tokens: 5000 Total num output tokens: 1500 ``` diff --git a/pyproject.toml b/pyproject.toml index 890dd6bdb7ce..53ce6dfbe06d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -192,7 +192,6 @@ python = "./.venv" [tool.typos.files] # these files may be written in non english words extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*", - "benchmarks/sonnet.txt", "rust/src/bench/src/datasets/sonnet.txt", "tests/lora/data/*", "build/*", "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md index d9a17f63f29b..5cf33e085897 100644 --- a/rust/src/bench/AGENTS.md +++ b/rust/src/bench/AGENTS.md @@ -47,7 +47,6 @@ cargo test -p vllm-bench -- --ignored - `src/datasets/random.rs` — Random dataset generation with rayon parallelism - `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`) - `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching -- `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`) - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) - `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (parquet shards via hf-hub into the standard HF hub cache; datasets-server for config/split discovery; native main-branch parquet/JSON fallback for private datasets; column auto-detection) - `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template) diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index fb1e5b529dc8..043605d9c827 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -13,7 +13,7 @@ vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \ - **Fast** — ~7 ms startup, single ~7 MB static binary, no Python imports. - **Scales** — `Arc` prompt sharing + mimalloc keep memory <100 MB at 1400+ concurrency. -- **Many datasets** — `random`, `random-mm` (VLM), `sharegpt`, `sonnet`, `speed-bench`, and any HuggingFace dataset. +- **Many datasets** — `random`, `random-mm` (VLM), `sharegpt`, `speed-bench`, and any HuggingFace dataset. - **Many backends** — completions, chat, embeddings, pooling, and rerank. - **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**. - **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain. @@ -127,7 +127,7 @@ vllm-bench \
-Datasets (ShareGPT / Sonnet / HuggingFace / SPEED-Bench) +Datasets (ShareGPT / HuggingFace / SPEED-Bench) ```bash # ShareGPT (auto-downloads from HuggingFace on first run, cached afterwards) @@ -141,14 +141,6 @@ vllm-bench \ --dataset-name sharegpt --dataset-path /path/to/ShareGPT_V3.json \ --num-prompts 500 --save-result -# Sonnet — built-in Shakespeare sonnets, no dataset file needed. -# Generates prompts of a controllable token length with a shared prefix. -vllm-bench \ - --backend openai-chat --base-url http://127.0.0.1:8000 --model \ - --dataset-name sonnet \ - --sonnet-input-len 550 --sonnet-output-len 150 --sonnet-prefix-len 200 \ - --num-prompts 500 - # Any HuggingFace dataset (downloads parquet shards into the standard HF hub # cache — shared with `hf download` and Python tooling — auto-detects columns) vllm-bench \ @@ -449,7 +441,7 @@ to measure a concurrent serving-style workload. | `vllm-pooling` | `/pooling` | vLLM native pooling endpoint | | `vllm-rerank` | `/v1/rerank` | vLLM reranking (query from prompt, documents via `--extra-body`) | -Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Use `--dataset-name sharegpt`, `sonnet`, or `hf` for text-based embedding/rerank benchmarks, or `random` for token-ID-based embedding benchmarks. +Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Use `--dataset-name sharegpt` or `hf` for text-based embedding/rerank benchmarks, or `random` for token-ID-based embedding benchmarks. ## Supported Datasets @@ -458,7 +450,6 @@ Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Us | `random` | Synthetic prompts with exact token-length matching (default) | | `random-mm` | Synthetic multimodal prompts with random JPEG images for VLM benchmarking (requires `openai-chat`) | | `sharegpt` | Real conversations from ShareGPT (auto-downloads from HuggingFace, or use `--dataset-path`) | -| `sonnet` | Built-in Shakespeare sonnets; controllable token length + shared prefix, no dataset file needed | | `speed-bench` | NVIDIA SPEED-Bench for speculative decoding evaluation (auto-downloads, 11 categories) | | `hf` | Any HuggingFace dataset with parquet data (downloads parquet shards via hf-hub into the standard HF hub cache, auto-detects chat/text columns) | @@ -523,8 +514,8 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below. | Flag | Default | Description | | ------ | --------- | ------------- | -| `--dataset-name` | `random` | Dataset type (`random`, `random-mm`, `sharegpt`, `sonnet`, `speed-bench`, `hf`) | -| `--dataset-path` | — | Path to dataset file (optional for `sharegpt`/`sonnet`, which auto-source) | +| `--dataset-name` | `random` | Dataset type (`random`, `random-mm`, `sharegpt`, `speed-bench`, `hf`) | +| `--dataset-path` | — | Path to dataset file (`sharegpt` auto-downloads when omitted) | | `--num-prompts` | `1000` | Number of prompts to generate (conversations in multi-turn mode) | | `--max-model-len` | — | Filter out requests where `prompt_len + output_len` exceeds this context length | | `--input-len` | — | Override input length (general) | @@ -545,10 +536,6 @@ Run `vllm-bench --help` for the authoritative list. Grouped reference below. | `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | `(height,width,frames)` → probability (Python tuple syntax; frames=1 = image) | | **ShareGPT** | | | | `--sharegpt-output-len` | — | Override output length | -| **Sonnet** | | | -| `--sonnet-input-len` | `550` | Input tokens per request | -| `--sonnet-output-len` | `150` | Output tokens per request | -| `--sonnet-prefix-len` | `200` | Prefix tokens shared across requests | | **SPEED-Bench** | | | | `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`); alias: `--speed-bench-dataset-subset` (Python name) | | `--speed-bench-output-len` | `4096` | Output tokens per request (matches Python default) | @@ -817,7 +804,6 @@ src/ │ ├── random_mm.rs # Random multimodal dataset (JPEG generation, bucket sampling) │ ├── multi_turn.rs # Multi-turn synthetic + ShareGPT conversation generators │ ├── sharegpt.rs # ShareGPT JSON dataset loader -│ ├── sonnet.rs # Sonnet dataset (built-in Shakespeare sonnets) │ ├── speed_bench.rs # NVIDIA SPEED-Bench loader (auto-download + cache) │ └── hf_dataset.rs # Generic HuggingFace dataset (auto-download, column detection) ├── metrics/ diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 54b6e8701f27..161646fe65f6 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -370,14 +370,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.num_prompts, config.dataset_path.as_deref().unwrap_or("auto-download") ), - DatasetName::Sonnet => format!( - "{} prompts from Sonnet ({}, isl={}, osl={}, prefix={})", - config.num_prompts, - config.dataset_path.as_deref().unwrap_or("built-in"), - config.sonnet_input_len, - config.sonnet_output_len, - config.sonnet_prefix_len, - ), DatasetName::SpeedBench => { let truncate_info = config .speed_bench_max_input_len @@ -484,21 +476,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.disable_shuffle, )? } - DatasetName::Sonnet => { - let tok = tokenizer - .as_ref() - .ok_or_else(|| BenchError::Config("Sonnet dataset requires a tokenizer".into()))?; - crate::datasets::sonnet::load_sonnet_dataset( - tok, - config.dataset_path.as_deref(), - config.num_prompts, - config.sonnet_input_len, - config.sonnet_output_len, - config.sonnet_prefix_len, - config.seed, - &config.request_id_prefix, - )? - } DatasetName::SpeedBench => { let tok = tokenizer.as_ref().ok_or_else(|| { BenchError::Config("SPEED-Bench dataset requires a tokenizer".into()) diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 302dbb479231..4782818ccb11 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -73,8 +73,6 @@ pub enum DatasetName { RandomMm, #[value(name = "sharegpt")] ShareGpt, - #[value(name = "sonnet")] - Sonnet, #[value(name = "speed-bench", alias = "speed_bench")] SpeedBench, #[value(name = "hf")] @@ -521,19 +519,6 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = false)] pub disable_shuffle: bool, - // --- Sonnet dataset --- - /// Number of input tokens per request (sonnet dataset). - #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_INPUT_LEN)] - pub sonnet_input_len: usize, - - /// Number of output tokens per request (sonnet dataset). - #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_OUTPUT_LEN)] - pub sonnet_output_len: usize, - - /// Number of prefix tokens shared across requests (sonnet dataset). - #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_PREFIX_LEN)] - pub sonnet_prefix_len: usize, - /// SPEED-Bench config/split (qualitative, throughput_1k, throughput_2k, throughput_8k, /// throughput_16k, throughput_32k). #[arg( diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 84f55ad427df..197cf08e7dd7 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -145,9 +145,6 @@ pub struct BenchConfig { pub prefix_repetition_num_prefixes: usize, pub prefix_repetition_output_len: usize, pub sharegpt_output_len: Option, - pub sonnet_input_len: usize, - pub sonnet_output_len: usize, - pub sonnet_prefix_len: usize, pub no_oversample: bool, pub disable_shuffle: bool, pub num_prompts: usize, @@ -447,8 +444,7 @@ impl BenchConfig { (MmLimitPerPrompt::default(), Vec::new()) }; - // Note: --dataset-path is optional for sharegpt (auto-downloads) and - // sonnet (uses built-in Shakespeare's sonnets). + // --dataset-path is optional for sharegpt (auto-downloads). // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1)) let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?; @@ -689,9 +685,6 @@ impl BenchConfig { random_cache_hit_fraction: args.random_cache_hit_fraction, random_cache_ratio: args.random_cache_ratio, sharegpt_output_len: args.sharegpt_output_len, - sonnet_input_len: args.sonnet_input_len, - sonnet_output_len: args.sonnet_output_len, - sonnet_prefix_len: args.sonnet_prefix_len, no_oversample: args.no_oversample, disable_shuffle: args.disable_shuffle, num_prompts: args.num_prompts, diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs index 8f81f8d3c75d..4848ec530362 100644 --- a/rust/src/bench/src/datasets/mod.rs +++ b/rust/src/bench/src/datasets/mod.rs @@ -10,7 +10,6 @@ pub mod random; pub mod random_mm; pub mod random_rerank; pub mod sharegpt; -pub mod sonnet; pub mod speed_bench; use std::sync::Arc; diff --git a/rust/src/bench/src/datasets/sonnet.rs b/rust/src/bench/src/datasets/sonnet.rs deleted file mode 100644 index 6fc7a75c0b5b..000000000000 --- a/rust/src/bench/src/datasets/sonnet.rs +++ /dev/null @@ -1,131 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -use std::sync::Arc; - -use rand::SeedableRng; -use rand::rngs::StdRng; -use rand::seq::IndexedRandom; - -use super::SampleRequest; -use crate::error::{BenchError, Result}; -use crate::tokenizer::TokenizerKind; - -/// Default values mirror Python's SonnetDataset defaults (datasets.py). -pub const DEFAULT_PREFIX_LEN: usize = 200; -pub const DEFAULT_INPUT_LEN: usize = 550; -pub const DEFAULT_OUTPUT_LEN: usize = 150; - -const BASE_PROMPT: &str = "Pick as many lines as you can from these poem lines:\n"; - -/// Shakespeare's sonnets, public domain. Bundled so `--dataset-name sonnet` works -/// out of the box without `--dataset-path`. Source: -/// https://raw.githubusercontent.com/vllm-project/vllm/main/benchmarks/sonnet.txt -const BUILTIN_SONNET: &str = include_str!("sonnet.txt"); - -/// Load the sonnet dataset and generate `num_requests` prompts targeting `input_len` -/// total prompt tokens. -/// -/// Mirrors Python's `SonnetDataset.sample()` from `vllm/benchmarks/datasets/datasets.py`. -/// The Rust port skips `apply_chat_template` (no Jinja runtime here): `base_offset` -/// and `prompt_len` are computed from the raw text. The resulting prompts are slightly -/// shorter than Python's chat-template-formatted version (off by the chat scaffolding -/// tokens, typically <20). -pub fn load_sonnet_dataset( - tokenizer: &TokenizerKind, - dataset_path: Option<&str>, - num_requests: usize, - input_len: usize, - output_len: usize, - prefix_len: usize, - seed: u64, - request_id_prefix: &str, -) -> Result> { - let content = match dataset_path { - Some(path) => std::fs::read_to_string(path) - .map_err(|e| BenchError::Config(format!("Failed to read sonnet file '{path}': {e}")))?, - None => BUILTIN_SONNET.to_string(), - }; - - // Match Python's f.readlines(): keep trailing newlines so that joining lines - // reconstructs the original text without inserting extra separators. - let lines: Vec = content.split_inclusive('\n').map(|s| s.to_string()).collect(); - - if lines.is_empty() { - let src = dataset_path.unwrap_or(""); - return Err(BenchError::Config(format!("Sonnet file '{src}' is empty"))); - } - - // Average tokens per line (used to estimate how many lines to draw). - let mut total_tokens: usize = 0; - for line in &lines { - let ids = tokenizer.encode(line, false)?; - total_tokens += ids.len(); - } - let avg_len = total_tokens as f64 / lines.len() as f64; - if avg_len <= 0.0 { - return Err(BenchError::Config( - "Sonnet lines tokenized to zero tokens on average".into(), - )); - } - - let base_ids = tokenizer.encode(BASE_PROMPT, false)?; - let base_offset = base_ids.len(); - if input_len <= base_offset { - return Err(BenchError::Config(format!( - "--sonnet-input-len ({input_len}) must be larger than the base prompt length ({base_offset})" - ))); - } - - let num_input_lines = ((input_len - base_offset) as f64 / avg_len).round() as i64; - let num_prefix_lines = - (((prefix_len as i64 - base_offset as i64) as f64) / avg_len).round().max(0.0) as i64; - let num_input_lines = num_input_lines.max(0) as usize; - let num_prefix_lines = (num_prefix_lines as usize).min(lines.len()); - let num_input_lines = num_input_lines.max(num_prefix_lines); - - let prefix_lines: &[String] = &lines[..num_prefix_lines]; - let extras_per_request = num_input_lines - num_prefix_lines; - - let mut rng = StdRng::seed_from_u64(seed); - let mut samples = Vec::with_capacity(num_requests); - let mut ind = 0usize; - let mut attempts = 0usize; - let max_attempts = num_requests.saturating_mul(20).max(1000); - - while samples.len() < num_requests { - if attempts >= max_attempts { - return Err(BenchError::Config(format!( - "Could not assemble {num_requests} sonnet prompts under input_len={input_len} \ - after {attempts} attempts. Try increasing --sonnet-input-len." - ))); - } - attempts += 1; - - let mut prompt = String::with_capacity(BASE_PROMPT.len() + 256 * num_input_lines); - prompt.push_str(BASE_PROMPT); - for line in prefix_lines { - prompt.push_str(line); - } - for _ in 0..extras_per_request { - // random.choices with replacement — duplicates are allowed. - let line = lines.choose(&mut rng).unwrap(); - prompt.push_str(line); - } - - let prompt_ids = tokenizer.encode(&prompt, false)?; - let prompt_len = prompt_ids.len(); - if prompt_len <= input_len { - samples.push(SampleRequest { - prompt: Arc::from(prompt), - prompt_len, - expected_output_len: output_len, - request_id: Some(format!("{request_id_prefix}{ind}")), - ..Default::default() - }); - ind += 1; - } - } - - Ok(samples) -} diff --git a/rust/src/bench/src/datasets/sonnet.txt b/rust/src/bench/src/datasets/sonnet.txt deleted file mode 100644 index 34c444e8ce8e..000000000000 --- a/rust/src/bench/src/datasets/sonnet.txt +++ /dev/null @@ -1,518 +0,0 @@ -FROM fairest creatures we desire increase, -That thereby beauty's rose might never die, -But as the riper should by time decease, -His tender heir might bear his memory: -But thou, contracted to thine own bright eyes, -Feed'st thy light'st flame with self-substantial fuel, -Making a famine where abundance lies, -Thyself thy foe, to thy sweet self too cruel. -Thou that art now the world's fresh ornament -And only herald to the gaudy spring, -Within thine own bud buriest thy content -And, tender churl, makest waste in niggarding. -Pity the world, or else this glutton be, -To eat the world's due, by the grave and thee. -When forty winters shall beseige thy brow, -And dig deep trenches in thy beauty's field, -Thy youth's proud livery, so gazed on now, -Will be a tatter'd weed, of small worth held: -Then being ask'd where all thy beauty lies, -Where all the treasure of thy lusty days, -To say, within thine own deep-sunken eyes, -Were an all-eating shame and thriftless praise. -How much more praise deserved thy beauty's use, -If thou couldst answer 'This fair child of mine -Shall sum my count and make my old excuse,' -Proving his beauty by succession thine! -This were to be new made when thou art old, -And see thy blood warm when thou feel'st it cold. -Look in thy glass, and tell the face thou viewest -Now is the time that face should form another; -Whose fresh repair if now thou not renewest, -Thou dost beguile the world, unbless some mother. -For where is she so fair whose unear'd womb -Disdains the tillage of thy husbandry? -Or who is he so fond will be the tomb -Of his self-love, to stop posterity? -Thou art thy mother's glass, and she in thee -Calls back the lovely April of her prime: -So thou through windows of thine age shall see -Despite of wrinkles this thy golden time. -But if thou live, remember'd not to be, -Die single, and thine image dies with thee. -Unthrifty loveliness, why dost thou spend -Upon thyself thy beauty's legacy? -Nature's bequest gives nothing but doth lend, -And being frank she lends to those are free. -Then, beauteous niggard, why dost thou abuse -The bounteous largess given thee to give? -Profitless usurer, why dost thou use -So great a sum of sums, yet canst not live? -For having traffic with thyself alone, -Thou of thyself thy sweet self dost deceive. -Then how, when nature calls thee to be gone, -What acceptable audit canst thou leave? -Thy unused beauty must be tomb'd with thee, -Which, used, lives th' executor to be. -Those hours, that with gentle work did frame -The lovely gaze where every eye doth dwell, -Will play the tyrants to the very same -And that unfair which fairly doth excel: -For never-resting time leads summer on -To hideous winter and confounds him there; -Sap cheque'd with frost and lusty leaves quite gone, -Beauty o'ersnow'd and bareness every where: -Then, were not summer's distillation left, -A liquid prisoner pent in walls of glass, -Beauty's effect with beauty were bereft, -Nor it nor no remembrance what it was: -But flowers distill'd though they with winter meet, -Leese but their show; their substance still lives sweet. -Then let not winter's ragged hand deface -In thee thy summer, ere thou be distill'd: -Make sweet some vial; treasure thou some place -With beauty's treasure, ere it be self-kill'd. -That use is not forbidden usury, -Which happies those that pay the willing loan; -That's for thyself to breed another thee, -Or ten times happier, be it ten for one; -Ten times thyself were happier than thou art, -If ten of thine ten times refigured thee: -Then what could death do, if thou shouldst depart, -Leaving thee living in posterity? -Be not self-will'd, for thou art much too fair -To be death's conquest and make worms thine heir. -Lo! in the orient when the gracious light -Lifts up his burning head, each under eye -Doth homage to his new-appearing sight, -Serving with looks his sacred majesty; -And having climb'd the steep-up heavenly hill, -Resembling strong youth in his middle age, -yet mortal looks adore his beauty still, -Attending on his golden pilgrimage; -But when from highmost pitch, with weary car, -Like feeble age, he reeleth from the day, -The eyes, 'fore duteous, now converted are -From his low tract and look another way: -So thou, thyself out-going in thy noon, -Unlook'd on diest, unless thou get a son. -Music to hear, why hear'st thou music sadly? -Sweets with sweets war not, joy delights in joy. -Why lovest thou that which thou receivest not gladly, -Or else receivest with pleasure thine annoy? -If the true concord of well-tuned sounds, -By unions married, do offend thine ear, -They do but sweetly chide thee, who confounds -In singleness the parts that thou shouldst bear. -Mark how one string, sweet husband to another, -Strikes each in each by mutual ordering, -Resembling sire and child and happy mother -Who all in one, one pleasing note do sing: -Whose speechless song, being many, seeming one, -Sings this to thee: 'thou single wilt prove none.' -Is it for fear to wet a widow's eye -That thou consumest thyself in single life? -Ah! if thou issueless shalt hap to die. -The world will wail thee, like a makeless wife; -The world will be thy widow and still weep -That thou no form of thee hast left behind, -When every private widow well may keep -By children's eyes her husband's shape in mind. -Look, what an unthrift in the world doth spend -Shifts but his place, for still the world enjoys it; -But beauty's waste hath in the world an end, -And kept unused, the user so destroys it. -No love toward others in that bosom sits -That on himself such murderous shame commits. -For shame! deny that thou bear'st love to any, -Who for thyself art so unprovident. -Grant, if thou wilt, thou art beloved of many, -But that thou none lovest is most evident; -For thou art so possess'd with murderous hate -That 'gainst thyself thou stick'st not to conspire. -Seeking that beauteous roof to ruinate -Which to repair should be thy chief desire. -O, change thy thought, that I may change my mind! -Shall hate be fairer lodged than gentle love? -Be, as thy presence is, gracious and kind, -Or to thyself at least kind-hearted prove: -Make thee another self, for love of me, -That beauty still may live in thine or thee. -As fast as thou shalt wane, so fast thou growest -In one of thine, from that which thou departest; -And that fresh blood which youngly thou bestowest -Thou mayst call thine when thou from youth convertest. -Herein lives wisdom, beauty and increase: -Without this, folly, age and cold decay: -If all were minded so, the times should cease -And threescore year would make the world away. -Let those whom Nature hath not made for store, -Harsh featureless and rude, barrenly perish: -Look, whom she best endow'd she gave the more; -Which bounteous gift thou shouldst in bounty cherish: -She carved thee for her seal, and meant thereby -Thou shouldst print more, not let that copy die. -When I do count the clock that tells the time, -And see the brave day sunk in hideous night; -When I behold the violet past prime, -And sable curls all silver'd o'er with white; -When lofty trees I see barren of leaves -Which erst from heat did canopy the herd, -And summer's green all girded up in sheaves -Borne on the bier with white and bristly beard, -Then of thy beauty do I question make, -That thou among the wastes of time must go, -Since sweets and beauties do themselves forsake -And die as fast as they see others grow; -And nothing 'gainst Time's scythe can make defence -Save breed, to brave him when he takes thee hence. -O, that you were yourself! but, love, you are -No longer yours than you yourself here live: -Against this coming end you should prepare, -And your sweet semblance to some other give. -So should that beauty which you hold in lease -Find no determination: then you were -Yourself again after yourself's decease, -When your sweet issue your sweet form should bear. -Who lets so fair a house fall to decay, -Which husbandry in honour might uphold -Against the stormy gusts of winter's day -And barren rage of death's eternal cold? -O, none but unthrifts! Dear my love, you know -You had a father: let your son say so. -Not from the stars do I my judgment pluck; -And yet methinks I have astronomy, -But not to tell of good or evil luck, -Of plagues, of dearths, or seasons' quality; -Nor can I fortune to brief minutes tell, -Pointing to each his thunder, rain and wind, -Or say with princes if it shall go well, -By oft predict that I in heaven find: -But from thine eyes my knowledge I derive, -And, constant stars, in them I read such art -As truth and beauty shall together thrive, -If from thyself to store thou wouldst convert; -Or else of thee this I prognosticate: -Thy end is truth's and beauty's doom and date. -When I consider every thing that grows -Holds in perfection but a little moment, -That this huge stage presenteth nought but shows -Whereon the stars in secret influence comment; -When I perceive that men as plants increase, -Cheered and cheque'd even by the self-same sky, -Vaunt in their youthful sap, at height decrease, -And wear their brave state out of memory; -Then the conceit of this inconstant stay -Sets you most rich in youth before my sight, -Where wasteful Time debateth with Decay, -To change your day of youth to sullied night; -And all in war with Time for love of you, -As he takes from you, I engraft you new. -But wherefore do not you a mightier way -Make war upon this bloody tyrant, Time? -And fortify yourself in your decay -With means more blessed than my barren rhyme? -Now stand you on the top of happy hours, -And many maiden gardens yet unset -With virtuous wish would bear your living flowers, -Much liker than your painted counterfeit: -So should the lines of life that life repair, -Which this, Time's pencil, or my pupil pen, -Neither in inward worth nor outward fair, -Can make you live yourself in eyes of men. -To give away yourself keeps yourself still, -And you must live, drawn by your own sweet skill. -Who will believe my verse in time to come, -If it were fill'd with your most high deserts? -Though yet, heaven knows, it is but as a tomb -Which hides your life and shows not half your parts. -If I could write the beauty of your eyes -And in fresh numbers number all your graces, -The age to come would say 'This poet lies: -Such heavenly touches ne'er touch'd earthly faces.' -So should my papers yellow'd with their age -Be scorn'd like old men of less truth than tongue, -And your true rights be term'd a poet's rage -And stretched metre of an antique song: -But were some child of yours alive that time, -You should live twice; in it and in my rhyme. -Shall I compare thee to a summer's day? -Thou art more lovely and more temperate: -Rough winds do shake the darling buds of May, -And summer's lease hath all too short a date: -Sometime too hot the eye of heaven shines, -And often is his gold complexion dimm'd; -And every fair from fair sometime declines, -By chance or nature's changing course untrimm'd; -But thy eternal summer shall not fade -Nor lose possession of that fair thou owest; -Nor shall Death brag thou wander'st in his shade, -When in eternal lines to time thou growest: -So long as men can breathe or eyes can see, -So long lives this and this gives life to thee. -Devouring Time, blunt thou the lion's paws, -And make the earth devour her own sweet brood; -Pluck the keen teeth from the fierce tiger's jaws, -And burn the long-lived phoenix in her blood; -Make glad and sorry seasons as thou fleets, -And do whate'er thou wilt, swift-footed Time, -To the wide world and all her fading sweets; -But I forbid thee one most heinous crime: -O, carve not with thy hours my love's fair brow, -Nor draw no lines there with thine antique pen; -Him in thy course untainted do allow -For beauty's pattern to succeeding men. -Yet, do thy worst, old Time: despite thy wrong, -My love shall in my verse ever live young. -A woman's face with Nature's own hand painted -Hast thou, the master-mistress of my passion; -A woman's gentle heart, but not acquainted -With shifting change, as is false women's fashion; -An eye more bright than theirs, less false in rolling, -Gilding the object whereupon it gazeth; -A man in hue, all 'hues' in his controlling, -Much steals men's eyes and women's souls amazeth. -And for a woman wert thou first created; -Till Nature, as she wrought thee, fell a-doting, -And by addition me of thee defeated, -By adding one thing to my purpose nothing. -But since she prick'd thee out for women's pleasure, -Mine be thy love and thy love's use their treasure. -So is it not with me as with that Muse -Stirr'd by a painted beauty to his verse, -Who heaven itself for ornament doth use -And every fair with his fair doth rehearse -Making a couplement of proud compare, -With sun and moon, with earth and sea's rich gems, -With April's first-born flowers, and all things rare -That heaven's air in this huge rondure hems. -O' let me, true in love, but truly write, -And then believe me, my love is as fair -As any mother's child, though not so bright -As those gold candles fix'd in heaven's air: -Let them say more than like of hearsay well; -I will not praise that purpose not to sell. -My glass shall not persuade me I am old, -So long as youth and thou are of one date; -But when in thee time's furrows I behold, -Then look I death my days should expiate. -For all that beauty that doth cover thee -Is but the seemly raiment of my heart, -Which in thy breast doth live, as thine in me: -How can I then be elder than thou art? -O, therefore, love, be of thyself so wary -As I, not for myself, but for thee will; -Bearing thy heart, which I will keep so chary -As tender nurse her babe from faring ill. -Presume not on thy heart when mine is slain; -Thou gavest me thine, not to give back again. -As an unperfect actor on the stage -Who with his fear is put besides his part, -Or some fierce thing replete with too much rage, -Whose strength's abundance weakens his own heart. -So I, for fear of trust, forget to say -The perfect ceremony of love's rite, -And in mine own love's strength seem to decay, -O'ercharged with burden of mine own love's might. -O, let my books be then the eloquence -And dumb presagers of my speaking breast, -Who plead for love and look for recompense -More than that tongue that more hath more express'd. -O, learn to read what silent love hath writ: -To hear with eyes belongs to love's fine wit. -Mine eye hath play'd the painter and hath stell'd -Thy beauty's form in table of my heart; -My body is the frame wherein 'tis held, -And perspective it is the painter's art. -For through the painter must you see his skill, -To find where your true image pictured lies; -Which in my bosom's shop is hanging still, -That hath his windows glazed with thine eyes. -Now see what good turns eyes for eyes have done: -Mine eyes have drawn thy shape, and thine for me -Are windows to my breast, where-through the sun -Delights to peep, to gaze therein on thee; -Yet eyes this cunning want to grace their art; -They draw but what they see, know not the heart. -Let those who are in favour with their stars -Of public honour and proud titles boast, -Whilst I, whom fortune of such triumph bars, -Unlook'd for joy in that I honour most. -Great princes' favourites their fair leaves spread -But as the marigold at the sun's eye, -And in themselves their pride lies buried, -For at a frown they in their glory die. -The painful warrior famoused for fight, -After a thousand victories once foil'd, -Is from the book of honour razed quite, -And all the rest forgot for which he toil'd: -Then happy I, that love and am beloved -Where I may not remove nor be removed. -Lord of my love, to whom in vassalage -Thy merit hath my duty strongly knit, -To thee I send this written embassage, -To witness duty, not to show my wit: -Duty so great, which wit so poor as mine -May make seem bare, in wanting words to show it, -But that I hope some good conceit of thine -In thy soul's thought, all naked, will bestow it; -Till whatsoever star that guides my moving -Points on me graciously with fair aspect -And puts apparel on my tatter'd loving, -To show me worthy of thy sweet respect: -Then may I dare to boast how I do love thee; -Till then not show my head where thou mayst prove me. -Weary with toil, I haste me to my bed, -The dear repose for limbs with travel tired; -But then begins a journey in my head, -To work my mind, when body's work's expired: -For then my thoughts, from far where I abide, -Intend a zealous pilgrimage to thee, -And keep my drooping eyelids open wide, -Looking on darkness which the blind do see -Save that my soul's imaginary sight -Presents thy shadow to my sightless view, -Which, like a jewel hung in ghastly night, -Makes black night beauteous and her old face new. -Lo! thus, by day my limbs, by night my mind, -For thee and for myself no quiet find. -How can I then return in happy plight, -That am debarr'd the benefit of rest? -When day's oppression is not eased by night, -But day by night, and night by day, oppress'd? -And each, though enemies to either's reign, -Do in consent shake hands to torture me; -The one by toil, the other to complain -How far I toil, still farther off from thee. -I tell the day, to please them thou art bright -And dost him grace when clouds do blot the heaven: -So flatter I the swart-complexion'd night, -When sparkling stars twire not thou gild'st the even. -But day doth daily draw my sorrows longer -And night doth nightly make grief's strength seem stronger. -When, in disgrace with fortune and men's eyes, -I all alone beweep my outcast state -And trouble deal heaven with my bootless cries -And look upon myself and curse my fate, -Wishing me like to one more rich in hope, -Featured like him, like him with friends possess'd, -Desiring this man's art and that man's scope, -With what I most enjoy contented least; -Yet in these thoughts myself almost despising, -Haply I think on thee, and then my state, -Like to the lark at break of day arising -From sullen earth, sings hymns at heaven's gate; -For thy sweet love remember'd such wealth brings -That then I scorn to change my state with kings. -When to the sessions of sweet silent thought -I summon up remembrance of things past, -I sigh the lack of many a thing I sought, -And with old woes new wail my dear time's waste: -Then can I drown an eye, unused to flow, -For precious friends hid in death's dateless night, -And weep afresh love's long since cancell'd woe, -And moan the expense of many a vanish'd sight: -Then can I grieve at grievances foregone, -And heavily from woe to woe tell o'er -The sad account of fore-bemoaned moan, -Which I new pay as if not paid before. -But if the while I think on thee, dear friend, -All losses are restored and sorrows end. -Thy bosom is endeared with all hearts, -Which I by lacking have supposed dead, -And there reigns love and all love's loving parts, -And all those friends which I thought buried. -How many a holy and obsequious tear -Hath dear religious love stol'n from mine eye -As interest of the dead, which now appear -But things removed that hidden in thee lie! -Thou art the grave where buried love doth live, -Hung with the trophies of my lovers gone, -Who all their parts of me to thee did give; -That due of many now is thine alone: -Their images I loved I view in thee, -And thou, all they, hast all the all of me. -If thou survive my well-contented day, -When that churl Death my bones with dust shall cover, -And shalt by fortune once more re-survey -These poor rude lines of thy deceased lover, -Compare them with the bettering of the time, -And though they be outstripp'd by every pen, -Reserve them for my love, not for their rhyme, -Exceeded by the height of happier men. -O, then vouchsafe me but this loving thought: -'Had my friend's Muse grown with this growing age, -A dearer birth than this his love had brought, -To march in ranks of better equipage: -But since he died and poets better prove, -Theirs for their style I'll read, his for his love.' -Full many a glorious morning have I seen -Flatter the mountain-tops with sovereign eye, -Kissing with golden face the meadows green, -Gilding pale streams with heavenly alchemy; -Anon permit the basest clouds to ride -With ugly rack on his celestial face, -And from the forlorn world his visage hide, -Stealing unseen to west with this disgrace: -Even so my sun one early morn did shine -With all triumphant splendor on my brow; -But out, alack! he was but one hour mine; -The region cloud hath mask'd him from me now. -Yet him for this my love no whit disdaineth; -Suns of the world may stain when heaven's sun staineth. -Why didst thou promise such a beauteous day, -And make me travel forth without my cloak, -To let base clouds o'ertake me in my way, -Hiding thy bravery in their rotten smoke? -'Tis not enough that through the cloud thou break, -To dry the rain on my storm-beaten face, -For no man well of such a salve can speak -That heals the wound and cures not the disgrace: -Nor can thy shame give physic to my grief; -Though thou repent, yet I have still the loss: -The offender's sorrow lends but weak relief -To him that bears the strong offence's cross. -Ah! but those tears are pearl which thy love sheds, -And they are rich and ransom all ill deeds. -No more be grieved at that which thou hast done: -Roses have thorns, and silver fountains mud; -Clouds and eclipses stain both moon and sun, -And loathsome canker lives in sweetest bud. -All men make faults, and even I in this, -Authorizing thy trespass with compare, -Myself corrupting, salving thy amiss, -Excusing thy sins more than thy sins are; -For to thy sensual fault I bring in sense-- -Thy adverse party is thy advocate-- -And 'gainst myself a lawful plea commence: -Such civil war is in my love and hate -That I an accessary needs must be -To that sweet thief which sourly robs from me. -Let me confess that we two must be twain, -Although our undivided loves are one: -So shall those blots that do with me remain -Without thy help by me be borne alone. -In our two loves there is but one respect, -Though in our lives a separable spite, -Which though it alter not love's sole effect, -Yet doth it steal sweet hours from love's delight. -I may not evermore acknowledge thee, -Lest my bewailed guilt should do thee shame, -Nor thou with public kindness honour me, -Unless thou take that honour from thy name: -But do not so; I love thee in such sort -As, thou being mine, mine is thy good report. -As a decrepit father takes delight -To see his active child do deeds of youth, -So I, made lame by fortune's dearest spite, -Take all my comfort of thy worth and truth. -For whether beauty, birth, or wealth, or wit, -Or any of these all, or all, or more, -Entitled in thy parts do crowned sit, -I make my love engrafted to this store: -So then I am not lame, poor, nor despised, -Whilst that this shadow doth such substance give -That I in thy abundance am sufficed -And by a part of all thy glory live. -Look, what is best, that best I wish in thee: -This wish I have; then ten times happy me! \ No newline at end of file diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 366fd4f6fbc1..b2b39b14cf19 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -170,11 +170,6 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { - return Err(BenchError::Config( - "Sonnet multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), - )); - } DatasetName::SpeedBench => { return Err(BenchError::Config( "SPEED-Bench multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), diff --git a/rust/src/bench/tests/python_serve_flags.txt b/rust/src/bench/tests/python_serve_flags.txt index d200e5bf245a..320619617804 100644 --- a/rust/src/bench/tests/python_serve_flags.txt +++ b/rust/src/bench/tests/python_serve_flags.txt @@ -86,9 +86,6 @@ --sharegpt-output-len --skip-chat-template --skip-tokenizer-init ---sonnet-input-len ---sonnet-output-len ---sonnet-prefix-len --spec-bench-category --spec-bench-output-len --speed-bench-category diff --git a/vllm/benchmarks/datasets/__init__.py b/vllm/benchmarks/datasets/__init__.py index b003ee4c059a..f1f58954b0be 100644 --- a/vllm/benchmarks/datasets/__init__.py +++ b/vllm/benchmarks/datasets/__init__.py @@ -27,7 +27,6 @@ RandomMultiModalDataset, SampleRequest, ShareGPTDataset, - SonnetDataset, SpecBench, VisionArenaDataset, add_dataset_parser, @@ -71,7 +70,6 @@ "RandomMultiModalDataset", "SampleRequest", "ShareGPTDataset", - "SonnetDataset", "SpecBench", "VisionArenaDataset", "add_dataset_parser", diff --git a/vllm/benchmarks/datasets/create_txt_slices_dataset.py b/vllm/benchmarks/datasets/create_txt_slices_dataset.py index 1bcd93c8ae11..ed3523a65311 100644 --- a/vllm/benchmarks/datasets/create_txt_slices_dataset.py +++ b/vllm/benchmarks/datasets/create_txt_slices_dataset.py @@ -12,8 +12,8 @@ :: python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\ - --input sonnet.txt \\ - --output sonnet_dataset.jsonl \\ + --input corpus.txt \\ + --output dataset.jsonl \\ --tokenizer gpt2 \\ --num-prompts 1000 \\ --input-len 1024 \\ @@ -23,7 +23,7 @@ python -m vllm.benchmarks.serve \\ --dataset-name custom \\ - --dataset-path sonnet_dataset.jsonl \\ + --dataset-path dataset.jsonl \\ ... """ diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index abaf299a1992..3746ac727e0f 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -5,7 +5,6 @@ generation. Supported dataset types include: - ShareGPT - Random (synthetic) - - Sonnet - BurstGPT - HuggingFace - VisionArena @@ -31,7 +30,6 @@ import numpy as np import pybase64 as base64 from PIL import Image -from typing_extensions import deprecated from vllm.benchmarks.datasets.utils import ( RangeRatio, @@ -1598,7 +1596,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser): choices=[ "sharegpt", "burstgpt", - "sonnet", "random", "random-mm", "random-rerank", @@ -1622,8 +1619,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "--dataset-path", type=str, default=None, - help="Path to the sharegpt/sonnet dataset or the HF dataset ID if " - "using HF dataset.", + help="Path to the ShareGPT dataset or the HF dataset ID if using HF dataset.", ) parser.add_argument( "--no-oversample", @@ -1681,26 +1677,6 @@ def add_dataset_parser(parser: FlexibleArgumentParser): help="Category for spec bench dataset. If None, use all categories.", ) - sonnet_group = parser.add_argument_group("sonnet dataset options") - sonnet_group.add_argument( - "--sonnet-input-len", - type=int, - default=550, - help="Number of input tokens per request, used only for sonnet dataset.", - ) - sonnet_group.add_argument( - "--sonnet-output-len", - type=int, - default=150, - help="Number of output tokens per request, used only for sonnet dataset.", - ) - sonnet_group.add_argument( - "--sonnet-prefix-len", - type=int, - default=200, - help="Number of prefix tokens per request, used only for sonnet dataset.", - ) - sharegpt_group = parser.add_argument_group("sharegpt dataset options") sharegpt_group.add_argument( "--sharegpt-output-len", @@ -2134,45 +2110,6 @@ def get_samples( no_oversample=args.no_oversample, ) - elif args.dataset_name == "sonnet": - assert tokenizer is not None, ( - "Tokenizer must be initialized for the 'sonnet' dataset." - ) - sonnet_dataset = SonnetDataset( - dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle - ) - # For the "sonnet" dataset, formatting depends on the backend. - # Chat-style backends leave templating to the server; completions - # backends need the prompt rendered client side. - if args.backend in ("openai-chat", "openai-responses"): - input_requests = sonnet_dataset.sample( - num_requests=args.num_prompts, - input_len=args.sonnet_input_len, - output_len=args.sonnet_output_len, - prefix_len=args.sonnet_prefix_len, - tokenizer=tokenizer, - return_prompt_formatted=False, - request_id_prefix=args.request_id_prefix, - no_oversample=args.no_oversample, - ) - else: - assert ( - hasattr(tokenizer, "chat_template") and tokenizer.chat_template - ) or ( - hasattr(tokenizer, "default_chat_template") - and tokenizer.default_chat_template - ), "Tokenizer/model must have chat template for sonnet dataset." - input_requests = sonnet_dataset.sample( - num_requests=args.num_prompts, - input_len=args.sonnet_input_len, - output_len=args.sonnet_output_len, - prefix_len=args.sonnet_prefix_len, - tokenizer=tokenizer, - return_prompt_formatted=True, - request_id_prefix=args.request_id_prefix, - no_oversample=args.no_oversample, - ) - elif args.dataset_name == "hf": # all following datasets are implemented from the # HuggingFaceDataset base class @@ -3070,106 +3007,6 @@ def sample( ) -# ----------------------------------------------------------------------------- -# Sonnet Dataset Implementation -# ----------------------------------------------------------------------------- - - -@deprecated( - "SonnetDataset is deprecated and will be removed in a future version.", -) -class SonnetDataset(BenchmarkDataset): - """Simplified implementation of the Sonnet dataset. Loads poem lines from a - text file and generates sample requests. Default values here copied from - `benchmark_serving.py` for the sonnet dataset. - """ - - DEFAULT_PREFIX_LEN = 200 - DEFAULT_INPUT_LEN = 550 - DEFAULT_OUTPUT_LEN = 150 - - def __init__( - self, - **kwargs, - ) -> None: - super().__init__(**kwargs) - self.load_data() - - def load_data(self) -> None: - if not self.dataset_path: - raise ValueError("dataset_path must be provided.") - with open(self.dataset_path, encoding="utf-8") as f: - self.data = f.readlines() - - def sample( - self, - tokenizer: TokenizerLike, - num_requests: int, - request_id_prefix: str = "", - no_oversample: bool = False, - prefix_len: int = DEFAULT_PREFIX_LEN, - input_len: int = DEFAULT_INPUT_LEN, - output_len: int = DEFAULT_OUTPUT_LEN, - return_prompt_formatted: bool = False, - **kwargs, - ) -> list[SampleRequest]: - poem_lines = self.data - assert poem_lines is not None - # Calculate average token length for a poem line. - tokenized_lines = [tokenizer(line).input_ids for line in poem_lines] - avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines) - - # Build the base prompt. - base_prompt = "Pick as many lines as you can from these poem lines:\n" - base_msg = [{"role": "user", "content": base_prompt}] - base_fmt = tokenizer.apply_chat_template( - base_msg, # type: ignore[arg-type] - add_generation_prompt=True, - tokenize=False, - ) - assert isinstance(base_fmt, str) - base_offset = len(tokenizer(base_fmt).input_ids) - if input_len <= base_offset: - raise ValueError( - f"'input_len' must be higher than the base prompt length " - f"({base_offset})." - ) - - # Determine how many poem lines to use. - num_input_lines = max(round((input_len - base_offset) / avg_len), 1) - num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0) - prefix_lines = poem_lines[:num_prefix_lines] - - samples: list[SampleRequest] = [] - ind = 0 - while len(samples) < num_requests: - extra_lines = random.choices( - poem_lines, k=num_input_lines - num_prefix_lines - ) - prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}" - msg = [{"role": "user", "content": prompt}] - prompt_formatted = tokenizer.apply_chat_template( - msg, # type: ignore[arg-type] - add_generation_prompt=True, - tokenize=False, - ) - assert isinstance(prompt_formatted, str) - prompt_len = len(tokenizer(prompt_formatted).input_ids) - if prompt_len <= input_len: - samples.append( - SampleRequest( - prompt=( - prompt_formatted if return_prompt_formatted else prompt - ), - prompt_len=prompt_len, - expected_output_len=output_len, - request_id=request_id_prefix + str(ind), - ) - ) - ind += 1 - return samples - - # ----------------------------------------------------------------------------- # BurstGPT Dataset Implementation # ----------------------------------------------------------------------------- diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index b551361f55d9..9e3c0b57e45f 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -1653,8 +1653,7 @@ def add_cli_args(parser: FlexibleArgumentParser): "--input-len", type=int, default=None, - help="General input length for datasets. Maps to dataset-specific " - "input length arguments (e.g., --random-input-len, --sonnet-input-len). " + help="General input length for datasets. Maps to --random-input-len. " "If not specified, uses dataset defaults.", ) parser.add_argument( @@ -1662,7 +1661,7 @@ def add_cli_args(parser: FlexibleArgumentParser): type=int, default=None, help="General output length for datasets. Maps to dataset-specific " - "output length arguments (e.g., --random-output-len, --sonnet-output-len). " + "output length arguments (e.g., --random-output-len, --hf-output-len). " "If not specified, uses dataset defaults.", ) parser.add_argument( @@ -2115,18 +2114,16 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: raise ValueError( f"Cannot use '{args.dataset_name}' dataset with --dataset-path. " "Please specify the appropriate --dataset-name (e.g., " - "'sharegpt', 'custom', 'sonnet') for your dataset file: " + "'sharegpt', 'custom', 'hf') for your dataset file: " f"{args.dataset_path}" ) - # Map general --input-len and --output-len to all dataset-specific arguments + # Map general length options to dataset-specific arguments. if args.input_len is not None: args.random_input_len = args.input_len - args.sonnet_input_len = args.input_len if args.output_len is not None: args.random_output_len = args.output_len - args.sonnet_output_len = args.output_len args.sharegpt_output_len = args.output_len args.custom_output_len = args.output_len args.hf_output_len = args.output_len diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index d7c3c0145409..6245b20e86c8 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -554,10 +554,6 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace: # matching prior throughput behaviour of omitting output_len when unset. d["hf_output_len"] = output_len d["sharegpt_output_len"] = output_len - # sonnet reads dedicated attrs; fall back to SonnetDataset's own defaults. - d["sonnet_input_len"] = input_len if input_len is not None else 550 - d["sonnet_output_len"] = output_len if output_len is not None else 150 - d["sonnet_prefix_len"] = prefix_len # Explicit --enable-multimodal-chat wins; otherwise auto-enable for the # multimodal chat backend (preserves today's vllm-chat handling). Callers # without a --backend flag (e.g. bench mm-processor) drive every request @@ -710,14 +706,14 @@ def validate_args(args): ) # --prefix-len: only used when dataset_name is 'random', 'random-mm', - # 'sonnet', or not set. + # or not set. if ( - args.dataset_name not in {"random", "random-mm", "sonnet", None} + args.dataset_name not in {"random", "random-mm", None} and args.prefix_len is not None ): warnings.warn( "--prefix-len will be ignored since --dataset-name\ - is not 'random', 'random-mm', 'sonnet', or not set.", + is not 'random', 'random-mm', or not set.", stacklevel=2, ) @@ -803,7 +799,6 @@ def add_cli_args(parser: FlexibleArgumentParser): choices=[ "sharegpt", "random", - "sonnet", "burstgpt", "hf", "prefix_repetition", From 25385102fe99708f831eedafeb1b03f51287c95a Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Tue, 6 Oct 2026 13:25:08 +0100 Subject: [PATCH 065/952] [Bugfix][KV Offload] Honor speculative cacheability in SimpleCPU (#60071) Signed-off-by: Sherif Waly --- tests/v1/simple_kv_offload/test_scheduler.py | 185 ++++++++++++++++++- vllm/v1/core/kv_cache_coordinator.py | 62 +++---- vllm/v1/simple_kv_offload/manager.py | 20 +- 3 files changed, 214 insertions(+), 53 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 5e299d27889f..dee5446c614b 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -6,6 +6,7 @@ import logging from dataclasses import dataclass +from unittest.mock import Mock import pytest import torch @@ -17,6 +18,7 @@ KVTransferConfig, ModelConfig, SchedulerConfig, + SpeculativeConfig, VllmConfig, ) from vllm.config.cache import MambaCacheMode @@ -26,7 +28,7 @@ ) from vllm.utils.hashing import sha256 from vllm.v1.core.block_pool import BlockPool -from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_utils import ( get_request_block_hasher, init_none_hash, @@ -3089,3 +3091,184 @@ def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None: without_ring, max_batched ) assert target_with == target_without + + +@pytest.mark.parametrize("retention_interval", [None, 0]) +@pytest.mark.parametrize("chunk_size,prefill_lookahead", [(64, 1), (1024, 1), (1, 3)]) +@pytest.mark.parametrize( + "draft_sliding_window,draft_sharded", + [(True, False), (False, False), (False, True)], + ids=["swa", "replicated_full", "sharded_full"], +) +def test_eager_store_matches_gpu_cacheable_prefix( + retention_interval: int | None, + chunk_size: int, + prefill_lookahead: int, + draft_sliding_window: bool, + draft_sharded: bool, +) -> None: + """Preserve GPU hits for replicated SWA and DCP-sharded full-attention drafts.""" + block_size = 64 + scheduler_block_size = 8 * block_size + draft_block_size = scheduler_block_size if draft_sharded else block_size + prompt_len = 16 * scheduler_block_size + draft_block_size + prefill_lookahead + kv_cache_config = _make_kv_cache_config(256, num_groups=2) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + group.kv_cache_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + dcp_sharded=group_id == 0 or draft_sharded, + ) + for tensor in kv_cache_config.kv_cache_tensors: + tensor.size *= block_size // BLOCK_SIZE + tensor.layer_stride *= block_size // BLOCK_SIZE + tensor.block_stride *= block_size // BLOCK_SIZE + if draft_sliding_window: + kv_cache_config.kv_cache_groups[1].kv_cache_spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + sliding_window=scheduler_block_size, + extra_retained_tokens=prefill_lookahead - 1, + dcp_sharded=False, + ) + kv_cache_config.kv_cache_groups[1].is_eagle_group = True + kv_cache_config.prefix_cache_retention_interval = retention_interval + vllm_config = _make_cp_vllm_config(dcp_world_size=8) + vllm_config.model_config.max_model_len = 16384 + speculative_config = Mock(spec=SpeculativeConfig) + speculative_config.num_speculative_tokens = 3 + speculative_config.use_eagle.return_value = True + speculative_config.use_eagle_block_drop.return_value = True + speculative_config.use_multi_module_mtp.return_value = prefill_lookahead > 1 + vllm_config.speculative_config = speculative_config + scheduler = SimpleCPUOffloadScheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + cpu_capacity_bytes=kv_cache_config.kv_cache_tensors[0].size, + scheduler_block_size=scheduler_block_size, + hash_block_size=block_size, + ) + gpu_manager = KVCacheManager( + kv_cache_config, + max_model_len=16384, + max_in_flight_tokens=chunk_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=block_size, + dcp_world_size=8, + use_eagle=True, + num_prefill_lookahead=prefill_lookahead, + ) + scheduler.bind_gpu_block_pool(gpu_manager.block_pool) + request = Request( + request_id="dcp-draft-source", + prompt_token_ids=list(range(prompt_len)), + sampling_params=SamplingParams(max_tokens=2), + pooling_params=None, + mm_features=None, + block_hasher=get_request_block_hasher(block_size, sha256), + ) + # Follow allocate -> connector metadata -> compute completion, including + # one decode step so the final prefill blocks can be offloaded. + while request.num_computed_tokens <= prompt_len: + if request.num_computed_tokens == prompt_len: + request.append_output_token_ids([prompt_len]) + num_new_tokens = min( + chunk_size, request.num_tokens - request.num_computed_tokens + ) + new_blocks = gpu_manager.allocate_slots( + request, num_new_tokens, num_lookahead_tokens=3 + ) + assert new_blocks is not None + if request.num_computed_tokens == 0: + scheduler.update_state_after_alloc( + request, gpu_manager.get_blocks(request.request_id), 0 + ) + output = make_scheduler_output( + {request.request_id: num_new_tokens}, + new_reqs={request.request_id: new_blocks.get_block_ids()}, + ) + else: + output = make_scheduler_output( + {request.request_id: num_new_tokens}, + cached_req_new_blocks={request.request_id: new_blocks.get_block_ids()}, + ) + store_meta = scheduler.build_connector_meta(output) + request.num_computed_tokens += num_new_tokens + if store_meta.store_event >= 0: + simulate_store_completion(scheduler, store_meta.store_event) + + matching_request = Request( + request_id="dcp-draft-load", + prompt_token_ids=request.prompt_token_ids, + sampling_params=request.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=request._block_hasher, + ) + expected_hit = ( + 15 if draft_sliding_window and retention_interval is not None else 16 + ) * scheduler_block_size + assert gpu_manager.get_computed_blocks(matching_request)[1] == expected_hit + scheduler.request_finished(request, []) + gpu_manager.free(request) + assert gpu_manager.block_pool.reset_prefix_cache() + + hit_tokens, is_async = scheduler.get_num_new_matched_tokens( + matching_request, num_computed_tokens=0 + ) + assert hit_tokens is not None + assert (hit_tokens, is_async) == (expected_hit, True) + + assert ( + gpu_manager.allocate_slots( + matching_request, + num_new_tokens=0, + num_external_computed_tokens=hit_tokens, + delay_cache_blocks=True, + ) + is not None + ) + load_blocks = gpu_manager.get_blocks(matching_request.request_id) + scheduler.update_state_after_alloc(matching_request, load_blocks, hit_tokens) + load_meta = scheduler.build_connector_meta( + make_scheduler_output( + {matching_request.request_id: 0}, + new_reqs={matching_request.request_id: load_blocks.get_block_ids()}, + ) + ) + # SWA loads only the eight-block tail after null padding; a sharded full + # draft must load every prefix block, with the same DCP geometry as the target. + num_draft_blocks = 8 if draft_sliding_window else hit_tokens // draft_block_size + assert load_meta.load_gpu_blocks == [ + block.block_id + for block in ( + *load_blocks.blocks[0], + *load_blocks.blocks[1][-num_draft_blocks:], + ) + ] + assert ( + len(load_meta.load_gpu_blocks) + == hit_tokens // scheduler_block_size + num_draft_blocks + ) + target_hashes = [ + make_block_hash_with_group_id(matching_request.block_hashes[(i + 1) * 8 - 1], 0) + for i in range(hit_tokens // scheduler_block_size) + ] + draft_hashes = [ + make_block_hash_with_group_id( + matching_request.block_hashes[(i + 1) * draft_block_size // block_size - 1], + 1, + ) + for i in range( + hit_tokens // draft_block_size - num_draft_blocks, + hit_tokens // draft_block_size, + ) + ] + assert [ + scheduler.cpu_block_pool.blocks[bid].block_hash + for bid in load_meta.load_cpu_blocks + ] == target_hashes + draft_hashes diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 9b91bb192068..deb7f89967b2 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -341,6 +341,14 @@ def get_replay_boundaries(self, request: Request) -> tuple[int, ...]: extension = request.num_prompt_tokens // block * block return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) + def get_num_cacheable_tokens( + self, num_computed_tokens: int, kv_cache_group_id: int + ) -> int: + """Return the prefix eligible for hashing in this group.""" + # Only cache tokens with finalized KV. The last num_reprefillable_tokens + # tokens can be re-prefilled during multi-module MTP. + return max(0, num_computed_tokens - self.num_reprefillable_tokens) + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """Cache the blocks for the request. @@ -352,17 +360,12 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """ boundaries = self.get_replay_boundaries(request) - for manager in self.single_type_managers: + for group_id, manager in enumerate(self.single_type_managers): if not manager.enable_caching: continue - # Only cache tokens with finalized KV. The last num_reprefillable_tokens - # tokens can be re-prefilled during multi-module MTP. - num_tokens_to_cache = max( - 0, num_computed_tokens - self.num_reprefillable_tokens - ) manager.cache_blocks( request, - num_tokens_to_cache, + self.get_num_cacheable_tokens(num_computed_tokens, group_id), retention_interval=self.retention_interval, replay_boundaries=boundaries, ) @@ -809,39 +812,22 @@ def _align_cacheable(self, num_tokens: int) -> int: return num_tokens return round_down(num_tokens, self.scheduler_block_size) - def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + def get_num_cacheable_tokens( + self, num_computed_tokens: int, kv_cache_group_id: int + ) -> int: cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) - boundaries = self.get_replay_boundaries(request) - for manager in self.single_type_managers: - if not manager.enable_caching: - continue - num_tokens_to_cache = cached_num_computed_tokens - # EAGLE groups match one block past each aligned boundary and drop - # it, so make that lookahead block eligible to be cached. - if manager.use_eagle and cached_num_computed_tokens > 0: - # Only cache tokens with finalized KV. The last - # num_reprefillable_tokens tokens can be re-prefilled during - # multi-module MTP. - num_finalized_computed_tokens = max( - 0, num_computed_tokens - self.num_reprefillable_tokens - ) - cached_num_finalized_computed_tokens = self._align_cacheable( - num_finalized_computed_tokens - ) - num_tokens_to_cache = min( - num_finalized_computed_tokens, - cached_num_finalized_computed_tokens + manager.block_size, - ) - # The manager already knows the fine hit granularity - # (``scheduler_block_size``); retention is passed separately so it - # can keep both the coarse segment tails and the fine replay - # boundary (which needs the fine value). - manager.cache_blocks( - request, - num_tokens_to_cache, - retention_interval=self.retention_interval, - replay_boundaries=boundaries, + manager = self.single_type_managers[kv_cache_group_id] + # EAGLE matches one block past the aligned boundary, then drops it. + if manager.use_eagle and cached_num_computed_tokens > 0: + num_finalized_computed_tokens = max( + 0, num_computed_tokens - self.num_reprefillable_tokens + ) + return min( + num_finalized_computed_tokens, + self._align_cacheable(num_finalized_computed_tokens) + + manager.block_size, ) + return cached_num_computed_tokens def find_longest_cache_hit( self, diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 9d95940a11f9..dfc24e420a5d 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -216,6 +216,7 @@ def __init__( scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, allow_partial_hash_hits=not lazy_offload, + num_prefill_lookahead=vllm_config.num_prefill_lookahead_tokens, ) self.group_block_sizes = self.cpu_coordinator.group_block_sizes # FA group's own resolved block_size; divides scheduler_block_size (the @@ -837,19 +838,6 @@ def _select_eager_blocks_to_store( ) request = state.request confirmed_tokens = request.num_computed_tokens - request.num_output_placeholders - # Truncate to the granularity a lookup can actually land on. With - # fine-grained hits that is hash_block_size; otherwise hits only land on - # the scheduler block (the group LCM). Using the LCM unconditionally - # would drop whole blocks that sit between the last LCM boundary and a - # reachable fine-grained boundary, so the other groups would hold that - # boundary and this one would not, and the joint hybrid lookup would - # reconcile to zero. - store_alignment = ( - self.hash_block_size - if self.cpu_coordinator.enable_partial_hash_hits - else self.block_size - ) - aligned_tokens = confirmed_tokens // store_alignment * store_alignment num_free = self.cpu_block_pool.get_num_free_blocks() for g, group_gpu_ids in enumerate(block_ids_by_group): @@ -864,7 +852,11 @@ def _select_eager_blocks_to_store( # num_stored_blocks can be stale and omit evicted blocks in # the middle of the request. group_size = self.group_block_sizes[g] - ready = min(len(group_gpu_ids), aligned_tokens // group_size) + # Mimic the GPU KV cache manager's cacheable prefix. + cacheable_tokens = self.cpu_coordinator.get_num_cacheable_tokens( + confirmed_tokens, g + ) + ready = min(len(group_gpu_ids), cacheable_tokens // group_size) resolved_hashes = resolve_block_hashes( request.block_hashes, self.hash_block_size, group_size ) From 3faf2140516d52dbaf143ab82402bf4681abc7c3 Mon Sep 17 00:00:00 2001 From: Rodrigo Garcia <32329949+roG0d@users.noreply.github.com> Date: Tue, 6 Oct 2026 14:39:35 +0200 Subject: [PATCH 066/952] [vllm-bench, feature] Added support for Mooncake-style, timed-traces replay. (#55937) --- rust/src/bench/AGENTS.md | 3 +- rust/src/bench/src/benchmark.rs | 117 +++++- rust/src/bench/src/cli.rs | 39 ++ rust/src/bench/src/config.rs | 183 +++++++++ rust/src/bench/src/datasets/custom.rs | 50 +-- rust/src/bench/src/datasets/mod.rs | 64 +++ .../bench/src/datasets/prefix_repetition.rs | 9 +- rust/src/bench/src/datasets/random_rerank.rs | 9 +- rust/src/bench/src/datasets/timed_trace.rs | 376 ++++++++++++++++++ rust/src/bench/src/multi_turn.rs | 5 +- rust/src/bench/src/rate_control.rs | 20 + rust/src/bench/tests/cli_parity.rs | 9 - 12 files changed, 794 insertions(+), 90 deletions(-) create mode 100644 rust/src/bench/src/datasets/timed_trace.rs diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md index 5cf33e085897..62fd963bbf9d 100644 --- a/rust/src/bench/AGENTS.md +++ b/rust/src/bench/AGENTS.md @@ -50,12 +50,13 @@ cargo test -p vllm-bench -- --ignored - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) - `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (parquet shards via hf-hub into the standard HF hub cache; datasets-server for config/split discovery; native main-branch parquet/JSON fallback for private datasets; column auto-detection) - `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template) +- `src/datasets/timed_trace.rs` — Mooncake-style replay trace (`{timestamp, input_length, output_length, hash_ids}` JSONL; `hash_ids` expand to memoized token chunks so prefix sharing reproduces; `--self-timed` fires each request at its recorded offset instead of `--request-rate`; mirrors Python `TimedTrace`) - `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`) - `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`) - `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`) - `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs - `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics` -- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop) +- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` and off under `--self-timed` (closed-loop only) - `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown - `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics` diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 161646fe65f6..5833a3bdcea5 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -17,9 +17,13 @@ use crate::metrics::calculator::{calculate_embedding_metrics, calculate_metrics} use crate::metrics::steady_state; use crate::output::console::print_results; use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result}; -use crate::rate_control::compute_schedule; +use crate::rate_control::{compute_schedule, trace_schedule}; use crate::ready_checker::{get_first_model, wait_for_endpoint}; +/// Past this first-arrival offset a self-timed run is almost certainly a +/// wrongly scaled or epoch-anchored trace rather than a real idle head. +const TRACE_START_WARN_SECONDS: f64 = 3600.0; + /// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the /// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This /// avoids repeated DNS lookups under high concurrency which can cause @@ -404,6 +408,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { "{} random rerank requests (batch={}, reranker={})", config.num_prompts, config.random_batch_size, config.is_reranker, ), + DatasetName::TimedTrace => format!( + "{} requests from timed trace ({}, chunk_hash_size={})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("unknown"), + config.timed_trace_chunk_hash_size, + ), }; tracing::info!( dataset = ?config.dataset_name, @@ -581,6 +591,30 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.is_reranker, )? } + DatasetName::TimedTrace => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("Timed trace dataset requires a tokenizer".into()) + })?; + let path = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config( + "--dataset-path is required for --dataset-name timed_trace".into(), + ) + })?; + crate::datasets::timed_trace::load_timed_trace_dataset( + tok, + path, + config.num_prompts, + &config.request_id_prefix, + &crate::datasets::timed_trace::TimedTraceOptions { + chunk_size: config.timed_trace_chunk_hash_size, + sec_multiplier: config.timed_trace_sec_multiplier, + label_timestamp: &config.timed_trace_label_timestamp, + label_input_length: &config.timed_trace_label_input_length, + label_output_length: &config.timed_trace_label_output_length, + label_hash_ids: &config.timed_trace_label_hash_ids, + }, + )? + } }; let gen_elapsed = gen_start.elapsed(); @@ -800,20 +834,30 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { tracing::info!("detected speculative decoding; collecting metrics"); } - // Main benchmark - let distribution = if config.burstiness == 1.0 { - "Poisson process" + // Main benchmark. Self-timed runs ignore --request-rate/--burstiness + // entirely (config.rs rejects them), so don't report a rate that isn't used. + if config.self_timed { + tracing::info!( + timing = "trace", + max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), + prompts = config.num_prompts, + "starting main benchmark run; firing requests at their trace timestamps" + ); } else { - "Gamma distribution" - }; - tracing::info!( - request_rate = config.request_rate, - burstiness = config.burstiness, - distribution, - max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), - prompts = config.num_prompts, - "starting main benchmark run" - ); + let distribution = if config.burstiness == 1.0 { + "Poisson process" + } else { + "Gamma distribution" + }; + tracing::info!( + request_rate = config.request_rate, + burstiness = config.burstiness, + distribution, + max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), + prompts = config.num_prompts, + "starting main benchmark run" + ); + } // Pre-assign LoRA adapters to each request (None when --lora-modules not set). let lora_assignments = assign_lora_modules( @@ -832,14 +876,40 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { ); } - // Compute request schedule - let schedule = compute_schedule( - input_requests.len(), - config.request_rate, - config.burstiness, - config.seed, - config.ramp_up.as_ref(), - ); + // Compute request schedule: the trace's recorded arrival offsets when + // self-timed, synthetic gamma inter-arrivals otherwise. Either way the + // dispatch loop below sleeps to benchmark_start + delay, an absolute + // deadline, so trace replay does not accumulate drift. + let schedule = if config.self_timed { + let schedule = trace_schedule(input_requests.iter().map(|r| r.timestamp)); + let first = schedule.delays.first().copied().unwrap_or(0.0); + let last = schedule.delays.iter().copied().fold(0.0, f64::max); + tracing::info!( + first_arrival_seconds = first, + last_arrival_seconds = last, + "replaying trace arrival schedule" + ); + // Timestamps are absolute offsets from now, never normalized (Python + // does the same). An unscaled millisecond trace or an epoch-anchored + // one parks every request past any plausible run length, so say so + // rather than sitting at 0/N. + if first > TRACE_START_WARN_SECONDS { + tracing::warn!( + first_arrival_seconds = first, + "first trace arrival is far in the future; check \ + --timed-trace-sec-multiplier or re-anchor the trace timestamps" + ); + } + schedule + } else { + compute_schedule( + input_requests.len(), + config.request_rate, + config.burstiness, + config.seed, + config.ramp_up.as_ref(), + ) + }; // Progress bar let pb = if config.disable_tqdm { @@ -1058,7 +1128,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { }; // Attach steady-state metrics when the closed-loop scope gate passes. + // Trace replay is open-loop: an in-flight plateau there reflects a burst + // in the trace, not a saturated server, so it is out of scope. let scope_ok = !config.no_steady_state + && !config.self_timed && config.max_concurrency.is_some() && config.request_rate.is_infinite(); if scope_ok { diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 4782818ccb11..93aa0d9c053b 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -83,6 +83,8 @@ pub enum DatasetName { PrefixRepetition, #[value(name = "random-rerank")] RandomRerank, + #[value(name = "timed_trace", alias = "timed-trace")] + TimedTrace, } /// Ramp-up strategy for request rate. @@ -321,6 +323,33 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = 128)] pub prefix_repetition_output_len: usize, + // --- Timed-trace dataset --- + /// How many tokens each prefix hash in the trace represents + /// (e.g. 512 for Moonshot traces, 16 for Qwen/Alibaba). + #[arg(long, default_value_t = 16)] + pub timed_trace_chunk_hash_size: usize, + + /// Multiplier converting trace timestamps to seconds + /// (e.g. 0.001 if timestamps are in milliseconds). + #[arg(long, default_value_t = 1.0)] + pub timed_trace_sec_multiplier: f64, + + /// JSON key of the timestamp field in the trace. + #[arg(long, default_value = "timestamp")] + pub timed_trace_label_timestamp: String, + + /// JSON key of the input length field in the trace. + #[arg(long, default_value = "input_length")] + pub timed_trace_label_input_length: String, + + /// JSON key of the output length field in the trace. + #[arg(long, default_value = "output_length")] + pub timed_trace_label_output_length: String, + + /// JSON key of the hash ids field in the trace. + #[arg(long, default_value = "hash_ids")] + pub timed_trace_label_hash_ids: String, + /// Number of prompts to generate. #[arg(long, default_value_t = 1000)] pub num_prompts: usize, @@ -333,6 +362,16 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = 1.0)] pub burstiness: f64, + /// Schedule requests at the timestamps recorded in the trace instead of + /// --request-rate. Defaults to on for --dataset-name timed_trace (the only + /// dataset that carries timestamps); an error for any other dataset. + #[arg(long, overrides_with = "no_self_timed")] + pub self_timed: bool, + + /// Force trace-driven timing off for timed_trace and use --request-rate. + #[arg(long, overrides_with = "self_timed")] + pub no_self_timed: bool, + /// Maximum number of concurrent requests. #[arg(long)] pub max_concurrency: Option, diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 197cf08e7dd7..3e180cfb3085 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -198,6 +198,15 @@ pub struct BenchConfig { pub hf_subset: Option, pub hf_output_len: Option, pub hf_text_column: Option, + /// Fire each request at its trace-recorded timestamp instead of a + /// synthetic --request-rate schedule. Only timed_trace can set this. + pub self_timed: bool, + pub timed_trace_chunk_hash_size: usize, + pub timed_trace_sec_multiplier: f64, + pub timed_trace_label_timestamp: String, + pub timed_trace_label_input_length: String, + pub timed_trace_label_output_length: String, + pub timed_trace_label_hash_ids: String, pub reset_prefix_cache: bool, pub prompt_token_ids: bool, // --- Random multimodal dataset --- @@ -376,6 +385,79 @@ impl BenchConfig { let mut multi_turn_min_turns = args.multi_turn_num_turns; let mut multi_turn_max_turns = args.multi_turn_num_turns; + let self_timed = if args.dataset_name == DatasetName::TimedTrace { + if !matches!(args.backend, BackendKind::Vllm | BackendKind::Openai) { + return Err(BenchError::Config( + "timed_trace passes pre-tokenized prompts (list[int]) and requires \ + a completions backend ('vllm' or 'openai')" + .into(), + )); + } + if args.dataset_path.is_none() { + return Err(BenchError::Config( + "--dataset-path is required for --dataset-name timed_trace".into(), + )); + } + // A zero, negative or non-finite scale turns trace timestamps into + // offsets the scheduler cannot sleep on. + if !(args.timed_trace_sec_multiplier.is_finite() + && args.timed_trace_sec_multiplier > 0.0) + { + return Err(BenchError::Config(format!( + "--timed-trace-sec-multiplier must be finite and positive, got {}", + args.timed_trace_sec_multiplier + ))); + } + let self_timed = !args.no_self_timed; + if self_timed { + // Flags that shape a synthetic schedule are unreachable once + // the trace supplies arrival times. Accepting them silently + // would still stamp the result JSON and its filename with a + // request rate that never ran. + let mut unused: Vec<&str> = Vec::new(); + if args.request_rate.is_finite() { + unused.push("--request-rate"); + } + if args.burstiness != 1.0 { + unused.push("--burstiness"); + } + if args.ramp_up_strategy.is_some() { + unused.push("--ramp-up-strategy"); + } + if args.sweep_request_rate.is_some() { + unused.push("--sweep-request-rate"); + } + // A per-point num_prompts truncates the trace to a different + // number of rows, so each sweep point replays a different span. + if args.sweep_max_concurrency.is_some() && args.sweep_num_prompts_factor.is_some() { + return Err(BenchError::Config( + "--sweep-num-prompts-factor truncates the trace to a different \ + number of rows per sweep point, so each point replays a \ + different span; drop it to sweep concurrency over the whole trace" + .into(), + )); + } + if !unused.is_empty() { + return Err(BenchError::Config(format!( + "{} {} no effect under trace-driven timing; pass --no-self-timed \ + to schedule timed_trace by --request-rate instead", + unused.join(" / "), + if unused.len() == 1 { "has" } else { "have" }, + ))); + } + } + self_timed + } else { + if args.self_timed || args.no_self_timed { + return Err(BenchError::Config( + "--self-timed/--no-self-timed is only supported with \ + --dataset-name timed_trace" + .into(), + )); + } + false + }; + // For random datasets with openai-compatible backends, default to ignore_eos. // Exception: multi-turn mode, where ignore_eos causes unbounded context growth // across turns. Multi-turn uses min_tokens instead for output length control. @@ -384,6 +466,8 @@ impl BenchConfig { false } else { args.ignore_eos + // timed_trace: generation must run to the trace's output_length + || args.dataset_name == DatasetName::TimedTrace || ((args.dataset_name == DatasetName::Random || args.dataset_name == DatasetName::RandomMm) && args.backend.is_openai_compatible() @@ -745,6 +829,13 @@ impl BenchConfig { hf_subset: args.hf_subset.clone(), hf_output_len: args.hf_output_len, hf_text_column: args.hf_text_column.clone(), + self_timed, + timed_trace_chunk_hash_size: args.timed_trace_chunk_hash_size, + timed_trace_sec_multiplier: args.timed_trace_sec_multiplier, + timed_trace_label_timestamp: args.timed_trace_label_timestamp.clone(), + timed_trace_label_input_length: args.timed_trace_label_input_length.clone(), + timed_trace_label_output_length: args.timed_trace_label_output_length.clone(), + timed_trace_label_hash_ids: args.timed_trace_label_hash_ids.clone(), reset_prefix_cache: args.reset_prefix_cache, prompt_token_ids: args.prompt_token_ids, random_mm_base_items_per_request: args.random_mm_base_items_per_request, @@ -899,6 +990,98 @@ mod tests { TestCli::parse_from(args).args } + fn base_timed_trace_args() -> Vec<&'static str> { + vec![ + "vllm-bench", + "--model", + "test-model", + "--dataset-name", + "timed_trace", + "--dataset-path", + "trace.jsonl", + ] + } + + #[test] + fn test_timed_trace_defaults_to_self_timed() { + let config = BenchConfig::from_args(&parse_args(base_timed_trace_args())).unwrap(); + assert!(config.self_timed); + assert!(config.ignore_eos); + } + + #[test] + fn test_no_self_timed_restores_rate_scheduling() { + let mut args = base_timed_trace_args(); + args.extend(["--no-self-timed", "--request-rate", "10"]); + let config = BenchConfig::from_args(&parse_args(args)).unwrap(); + assert!(!config.self_timed); + assert_eq!(config.request_rate, 10.0); + } + + /// Rate/burstiness/ramp-up shape a synthetic schedule the trace replaces; + /// accepting them would mislabel the result JSON with a rate that never ran. + #[test] + fn test_self_timed_rejects_synthetic_schedule_flags() { + for flag in [ + vec!["--request-rate", "10"], + vec!["--burstiness", "0.5"], + vec!["--sweep-request-rate", "1,5,10"], + ] { + let mut args = base_timed_trace_args(); + args.extend(flag.iter().copied()); + let msg = BenchConfig::from_args(&parse_args(args)) + .expect_err("self-timed replay should reject synthetic schedule flags") + .to_string(); + assert!( + msg.contains(flag[0]) && msg.contains("--no-self-timed"), + "{flag:?}: {msg}" + ); + } + } + + /// A per-point num_prompts truncates the trace differently at every sweep + /// point, so the replayed span stops being comparable across them. + #[test] + fn test_self_timed_rejects_sweep_num_prompts_factor() { + let mut args = base_timed_trace_args(); + args.extend([ + "--sweep-max-concurrency", + "8,16", + "--sweep-num-prompts-factor", + "4", + ]); + let err = BenchConfig::from_args(&parse_args(args)).expect_err("should reject the factor"); + let msg = err.to_string(); + assert!( + msg.contains("--sweep-num-prompts-factor") && msg.contains("different span"), + "{msg}" + ); + } + + #[test] + fn test_self_timed_allows_concurrency_sweep_over_full_trace() { + let mut args = base_timed_trace_args(); + args.extend(["--sweep-max-concurrency", "8,16"]); + assert!(BenchConfig::from_args(&parse_args(args)).is_ok()); + } + + #[test] + fn test_non_positive_sec_multiplier_fails() { + // `=` form: clap reads a bare "-0.001" as a flag. + for value in [ + "--timed-trace-sec-multiplier=0", + "--timed-trace-sec-multiplier=-0.001", + "--timed-trace-sec-multiplier=inf", + "--timed-trace-sec-multiplier=nan", + ] { + let mut args = base_timed_trace_args(); + args.push(value); + let err = BenchConfig::from_args(&parse_args(args)) + .expect_err("should reject a multiplier that cannot scale timestamps"); + assert!(err.to_string().contains("--timed-trace-sec-multiplier")); + } + } + fn base_multi_turn_args() -> Vec<&'static str> { vec![ "vllm-bench", diff --git a/rust/src/bench/src/datasets/custom.rs b/rust/src/bench/src/datasets/custom.rs index 11ddee02b17d..a9fcba335498 100644 --- a/rust/src/bench/src/datasets/custom.rs +++ b/rust/src/bench/src/datasets/custom.rs @@ -43,29 +43,10 @@ pub fn load_custom_dataset( no_oversample: bool, disable_shuffle: bool, ) -> Result> { - let content = std::fs::read_to_string(path) - .map_err(|e| BenchError::Config(format!("Failed to read custom dataset '{path}': {e}")))?; - - let mut lines: Vec = Vec::new(); - for (lineno, line) in content.lines().enumerate() { - let line = line.trim(); - if line.is_empty() { - continue; - } - let parsed: CustomLine = serde_json::from_str(line).map_err(|e| { - BenchError::Config(format!( - "Invalid JSONL at {path}:{}: {e} (each line must be an object \ - with a 'prompt' field)", - lineno + 1 - )) - })?; - lines.push(parsed); - } - if lines.is_empty() { - return Err(BenchError::Config(format!( - "Custom dataset '{path}' contains no entries" - ))); - } + let mut lines: Vec = super::read_jsonl(path, "custom dataset", None)? + .into_iter() + .map(|(_, line)| line) + .collect(); // Python shuffles the loaded data (seeded) before taking num_requests. if !disable_shuffle { @@ -118,26 +99,13 @@ pub fn load_custom_dataset( #[cfg(test)] mod tests { + use super::super::test_support::{test_tokenizer, write_temp_jsonl}; use super::*; - fn write_temp_jsonl(name: &str, content: &str) -> String { - let path = std::env::temp_dir().join(format!("vllm-bench-custom-{name}.jsonl")); - std::fs::write(&path, content).unwrap(); - path.to_string_lossy().into_owned() - } - - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - #[test] fn test_load_custom_dataset_basic() { let path = write_temp_jsonl( - "basic", + "custom-basic", r#"{"prompt": "hello world", "output_tokens": 10} {"prompt": "foo bar baz", "output_tokens": 20} "#, @@ -154,7 +122,7 @@ mod tests { #[test] fn test_load_custom_dataset_per_line_output_tokens() { let path = write_temp_jsonl( - "perline", + "custom-perline", r#"{"prompt": "hello", "output_tokens": 10} {"prompt": "world", "output_tokens": 20} "#, @@ -167,7 +135,7 @@ mod tests { #[test] fn test_load_custom_dataset_missing_output_tokens_errors() { - let path = write_temp_jsonl("missing", r#"{"prompt": "hello"}"#); + let path = write_temp_jsonl("custom-missing", r#"{"prompt": "hello"}"#); let err = load_custom_dataset(&test_tokenizer(), &path, 1, -1, 0, "t-", true, true) .expect_err("should fail without output_tokens"); assert!(err.to_string().contains("output_tokens")); @@ -175,7 +143,7 @@ mod tests { #[test] fn test_load_custom_dataset_missing_prompt_errors() { - let path = write_temp_jsonl("noprompt", r#"{"text": "hello"}"#); + let path = write_temp_jsonl("custom-noprompt", r#"{"text": "hello"}"#); assert!( load_custom_dataset(&test_tokenizer(), &path, 1, 256, 0, "t-", true, true).is_err() ); diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs index 4848ec530362..bfdc4006a695 100644 --- a/rust/src/bench/src/datasets/mod.rs +++ b/rust/src/bench/src/datasets/mod.rs @@ -11,6 +11,7 @@ pub mod random_mm; pub mod random_rerank; pub mod sharegpt; pub mod speed_bench; +pub mod timed_trace; use std::sync::Arc; @@ -54,6 +55,10 @@ pub struct SampleRequest { /// rerank sends `[0]` as the query and `[1..]` as documents (random-rerank). /// Mirrors Python's list-valued `SampleRequest.prompt`. pub prompt_list: Option]>>, + /// Arrival time in seconds relative to benchmark start (timed_trace only). + /// With --self-timed the scheduler fires the request at this offset + /// instead of a synthetic --request-rate schedule. + pub timestamp: Option, } impl Default for SampleRequest { @@ -69,6 +74,7 @@ impl Default for SampleRequest { multi_modal_content: None, chat_messages_json: None, prompt_list: None, + timestamp: None, } } } @@ -137,6 +143,64 @@ pub fn batch_requests( .collect() } +/// Read a JSONL file into `T` values, skipping blank lines. +/// +/// `what` names the dataset in error messages ("custom dataset", "timed +/// trace"); `max_entries` stops the read early. Returns `(line_number, value)` +/// pairs with 0-based line numbers so callers can index or report by row. +pub fn read_jsonl( + path: &str, + what: &str, + max_entries: Option, +) -> crate::error::Result> { + use crate::error::BenchError; + + let content = std::fs::read_to_string(path) + .map_err(|e| BenchError::Config(format!("Failed to read {what} '{path}': {e}")))?; + + let mut entries: Vec<(usize, T)> = Vec::new(); + for (lineno, line) in content.lines().enumerate() { + if max_entries.is_some_and(|max| entries.len() >= max) { + break; + } + let line = line.trim(); + if line.is_empty() { + continue; + } + let value = serde_json::from_str(line).map_err(|e| { + BenchError::Config(format!("Invalid JSONL at {path}:{}: {e}", lineno + 1)) + })?; + entries.push((lineno, value)); + } + + if entries.is_empty() { + return Err(BenchError::Config(format!("No entries in {what} '{path}'"))); + } + Ok(entries) +} + +/// Fixtures shared by the dataset test modules. +#[cfg(test)] +pub(crate) mod test_support { + use crate::tokenizer::TokenizerKind; + + /// gpt2 via built-in tiktoken encoding — loads without network access. + pub(crate) fn test_tokenizer() -> TokenizerKind { + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) + } + + /// Write a JSONL fixture to the temp dir. `name` must be unique across the + /// whole test suite — all fixtures share one directory. + pub(crate) fn write_temp_jsonl(name: &str, content: &str) -> String { + let path = std::env::temp_dir().join(format!("vllm-bench-{name}.jsonl")); + std::fs::write(&path, content).unwrap(); + path.to_string_lossy().into_owned() + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs index e4dfcd36a3a9..7eb5fac415bf 100644 --- a/rust/src/bench/src/datasets/prefix_repetition.rs +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -107,16 +107,9 @@ pub fn generate_prefix_repetition_dataset( #[cfg(test)] mod tests { + use super::super::test_support::test_tokenizer; use super::*; - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - #[test] fn test_prefix_repetition_structure() { let tok = test_tokenizer(); diff --git a/rust/src/bench/src/datasets/random_rerank.rs b/rust/src/bench/src/datasets/random_rerank.rs index ad3686cf3f21..76a7eb342788 100644 --- a/rust/src/bench/src/datasets/random_rerank.rs +++ b/rust/src/bench/src/datasets/random_rerank.rs @@ -135,16 +135,9 @@ pub fn generate_random_rerank_dataset( #[cfg(test)] mod tests { + use super::super::test_support::test_tokenizer; use super::*; - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - fn fixed_ratio() -> RangeRatio { RangeRatio::parse("0.0").unwrap() } diff --git a/rust/src/bench/src/datasets/timed_trace.rs b/rust/src/bench/src/datasets/timed_trace.rs new file mode 100644 index 000000000000..f90a94450a87 --- /dev/null +++ b/rust/src/bench/src/datasets/timed_trace.rs @@ -0,0 +1,376 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Timed-trace dataset: JSONL replay trace with per-request arrival timestamps. +//! +//! ```jsonl +//! {"timestamp": 12.4, "input_length": 42, "output_length": 8, "hash_ids": [981, 4412, 7]} +//! ``` +//! +//! Mirrors Python's `TimedTrace`. Prompts carry no text: each `hash_id` expands +//! to a pseudo-random token chunk of `chunk_size` tokens (the last chunk may be +//! partial so the prompt length equals `input_length` exactly). Chunks are +//! memoized per (hash_id, size), so requests sharing a hash prefix get +//! byte-identical token prefixes and prefix-cache hits reproduce. Field names +//! are remappable via `--timed-trace-label-*`; timestamps are scaled by +//! `--timed-trace-sec-multiplier`. +//! +//! Deviation from Python: chunk contents are seeded from (hash_id, size) with a +//! fixed mixing function, so they are bit-reproducible across runs. Python +//! seeds via `hash(f"{h}:{size}")`, which is randomized per process unless +//! PYTHONHASHSEED is pinned. + +use std::collections::HashMap; +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use serde_json::Value; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +pub struct TimedTraceOptions<'a> { + pub chunk_size: usize, + pub sec_multiplier: f64, + pub label_timestamp: &'a str, + pub label_input_length: &'a str, + pub label_output_length: &'a str, + pub label_hash_ids: &'a str, +} + +/// Load a timed-trace JSONL dataset. Row order (and thus the timestamp +/// schedule) is preserved: no shuffle, no oversampling. +pub fn load_timed_trace_dataset( + tokenizer: &TokenizerKind, + path: &str, + num_requests: usize, + request_id_prefix: &str, + opts: &TimedTraceOptions<'_>, +) -> Result> { + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + + let entries: Vec<(usize, Value)> = super::read_jsonl(path, "timed trace", Some(num_requests))?; + + let mut chunk_cache: HashMap<(u64, usize), Arc<[u32]>> = HashMap::new(); + let mut requests: Vec = Vec::with_capacity(entries.len()); + + for (lineno, entry) in &entries { + let input_length = require_u64( + entry, + opts.label_input_length, + "--timed-trace-label-input-length", + )? as usize; + let output_length = require_u64( + entry, + opts.label_output_length, + "--timed-trace-label-output-length", + )? as usize; + let timestamp = require_f64(entry, opts.label_timestamp, "--timed-trace-label-timestamp")? + * opts.sec_multiplier; + // A non-finite offset panics in `Duration::from_secs_f64` once the + // scheduler consumes it; reject here, naming the offending row. + if !timestamp.is_finite() { + return Err(BenchError::Config(format!( + "Field '{}' at {path}:{} scales to a non-finite offset ({timestamp}); \ + check --timed-trace-sec-multiplier", + opts.label_timestamp, + lineno + 1, + ))); + } + let hash_ids = parse_hash_ids(entry, opts.label_hash_ids)?; + + let prompt_ids = expand_prompt( + &hash_ids, + input_length, + opts.chunk_size, + &allowed_tokens, + &mut chunk_cache, + ); + + requests.push(SampleRequest { + prompt_len: prompt_ids.len(), + prompt_token_ids: Some(Arc::from(prompt_ids)), + expected_output_len: output_length, + request_id: Some(format!("{request_id_prefix}{lineno}")), + timestamp: Some(timestamp), + ..Default::default() + }); + } + + Ok(requests) +} + +/// Expand prefix-hash ids into token ids, consuming `chunk_size` tokens per +/// hash until `target_len` is reached (the final chunk may be partial). A +/// repeated (hash_id, size) yields the identical chunk via the cache. If the +/// trace supplies fewer hashes than `target_len / chunk_size`, the prompt +/// comes out shorter — same as Python. +fn expand_prompt( + hash_ids: &[u64], + target_len: usize, + chunk_size: usize, + allowed_tokens: &[u32], + cache: &mut HashMap<(u64, usize), Arc<[u32]>>, +) -> Vec { + // `target_len` is raw trace input, so reserve only what the hashes can + // actually produce: a corrupt row, or one in the wrong unit, would + // otherwise ask the allocator for terabytes before any length filter runs. + let capacity = target_len.min(hash_ids.len().saturating_mul(chunk_size)); + let mut prompt = Vec::with_capacity(capacity); + let mut remaining = target_len; + for &h in hash_ids { + let size = remaining.min(chunk_size); + let chunk = cache.entry((h, size)).or_insert_with(|| { + let mut rng = StdRng::seed_from_u64(chunk_seed(h, size)); + (0..size) + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())]) + .collect() + }); + prompt.extend_from_slice(chunk); + remaining -= size; + if remaining == 0 { + break; + } + } + prompt +} + +/// Stable seed for a chunk. The golden-ratio multiply is bijective on u64, so +/// distinct `hash_id`s never collide at a fixed `size`. The trailing xor is not +/// injective across sizes: two unrelated hashes seen at different sizes can +/// alias and share a token prefix, at ~2^-64 odds per pair. +fn chunk_seed(hash_id: u64, size: usize) -> u64 { + hash_id.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (size as u64) +} + +fn require_field<'v>(entry: &'v Value, label: &str, flag: &str) -> Result<&'v Value> { + entry.get(label).ok_or_else(|| { + let available: Vec<&str> = entry + .as_object() + .map(|o| o.keys().map(String::as_str).collect()) + .unwrap_or_default(); + BenchError::Config(format!( + "Field '{label}' not found in trace entry. Available fields: \ + {available:?}. Use {flag} to specify the correct field name." + )) + }) +} + +/// Float-encoded integers (`512.0`) are accepted, matching Python's `int()`. +fn require_u64(entry: &Value, label: &str, flag: &str) -> Result { + let v = require_field(entry, label, flag)?; + v.as_u64() + .or_else(|| { + v.as_f64() + .filter(|f| f.fract() == 0.0 && (0.0..u64::MAX as f64).contains(f)) + .map(|f| f as u64) + }) + .ok_or_else(|| { + BenchError::Config(format!( + "Field '{label}' must be a non-negative integer, got: {v}" + )) + }) +} + +fn require_f64(entry: &Value, label: &str, flag: &str) -> Result { + let v = require_field(entry, label, flag)?; + v.as_f64() + .ok_or_else(|| BenchError::Config(format!("Field '{label}' must be a number, got: {v}"))) +} + +/// Missing hash_ids field means an empty prompt (mirrors Python's +/// `entry.get(label, [])`); a present field must be an array of integers. +/// +/// Ids are kept as `u64` because traces carry raw unsigned 64-bit block hashes; +/// negative ids keep the two's-complement value `chunk_seed` used before. +fn parse_hash_ids(entry: &Value, label: &str) -> Result> { + let Some(v) = entry.get(label) else { + return Ok(Vec::new()); + }; + let arr = v + .as_array() + .ok_or_else(|| BenchError::Config(format!("Field '{label}' must be an array, got: {v}")))?; + arr.iter() + .map(|x| { + x.as_u64().or_else(|| x.as_i64().map(|i| i as u64)).ok_or_else(|| { + BenchError::Config(format!( + "Field '{label}' must contain only integers, got: {x}" + )) + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::super::test_support::{test_tokenizer, write_temp_jsonl}; + use super::*; + + fn opts(chunk_size: usize, sec_multiplier: f64) -> TimedTraceOptions<'static> { + TimedTraceOptions { + chunk_size, + sec_multiplier, + label_timestamp: "timestamp", + label_input_length: "input_length", + label_output_length: "output_length", + label_hash_ids: "hash_ids", + } + } + + #[test] + fn test_prompt_len_matches_input_length_and_fields_map() { + let path = write_temp_jsonl( + "timed-trace-basic", + r#"{"timestamp": 1000, "input_length": 40, "output_length": 7, "hash_ids": [1, 2, 3]} +{"timestamp": 2500, "input_length": 32, "output_length": 9, "hash_ids": [1, 4]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 0.001)) + .expect("load should succeed"); + assert_eq!(reqs.len(), 2); + // 40 = 16 + 16 + partial 8 from three hashes + assert_eq!(reqs[0].prompt_len, 40); + assert_eq!(reqs[0].prompt_token_ids.as_ref().unwrap().len(), 40); + assert_eq!(reqs[0].expected_output_len, 7); + assert_eq!(reqs[0].timestamp, Some(1.0)); + assert_eq!(reqs[1].timestamp, Some(2.5)); + assert_eq!(reqs[0].request_id.as_deref(), Some("t-0")); + } + + #[test] + fn test_shared_hash_prefix_gives_identical_tokens() { + let path = write_temp_jsonl( + "timed-trace-prefix", + r#"{"timestamp": 0, "input_length": 32, "output_length": 1, "hash_ids": [7, 8]} +{"timestamp": 1, "input_length": 32, "output_length": 1, "hash_ids": [7, 9]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + let a = reqs[0].prompt_token_ids.as_ref().unwrap(); + let b = reqs[1].prompt_token_ids.as_ref().unwrap(); + assert_eq!(a[..16], b[..16], "shared hash 7 must expand identically"); + assert_ne!(a[16..], b[16..], "distinct hashes 8/9 must differ"); + } + + #[test] + fn test_reload_is_bit_identical() { + let path = write_temp_jsonl( + "timed-trace-determinism", + r#"{"timestamp": 0, "input_length": 40, "output_length": 1, "hash_ids": [11, 12, 13]}"#, + ); + let load = || { + load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed") + }; + assert_eq!(load()[0].prompt_token_ids, load()[0].prompt_token_ids); + } + + #[test] + fn test_too_few_hashes_yields_short_prompt() { + let path = write_temp_jsonl( + "timed-trace-short", + r#"{"timestamp": 0, "input_length": 100, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16, "one hash covers only one chunk"); + } + + /// `input_length` is untrusted: the reservation must follow what the hashes + /// can produce, not the raw field, or the allocator aborts the process. + #[test] + fn test_absurd_input_length_does_not_overallocate() { + let path = write_temp_jsonl( + "timed-trace-absurd-len", + r#"{"timestamp": 0, "input_length": 10000000000000, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + } + + #[test] + fn test_num_requests_caps_rows() { + let path = write_temp_jsonl( + "timed-trace-cap", + r#"{"timestamp": 0, "input_length": 8, "output_length": 1, "hash_ids": [1]} +{"timestamp": 1, "input_length": 8, "output_length": 1, "hash_ids": [2]} +{"timestamp": 2, "input_length": 8, "output_length": 1, "hash_ids": [3]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 2, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs.len(), 2); + } + + #[test] + fn test_missing_field_error_names_label_and_flag() { + let path = write_temp_jsonl( + "timed-trace-missing", + r#"{"timestamp": 0, "input_length": 8, "output_length": 1, "hash_ids": [1]}"#, + ); + let mut o = opts(16, 1.0); + o.label_input_length = "input_len"; + let err = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &o) + .expect_err("should fail on missing field"); + let msg = err.to_string(); + assert!(msg.contains("input_len") && msg.contains("--timed-trace-label-input-length")); + } + + /// Hashes above `i64::MAX` are ordinary 64-bit block hashes, and Python + /// takes them as unbounded ints. + #[test] + fn test_accepts_unsigned_64bit_hash_ids() { + let path = write_temp_jsonl( + "timed-trace-u64-hash", + r#"{"timestamp": 0, "input_length": 16, "output_length": 1, "hash_ids": [18446744073709551615]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + } + + /// Exporters that round-trip lengths through floats emit `512.0`; Python's + /// `int()` accepts it. + #[test] + fn test_accepts_float_encoded_lengths() { + let path = write_temp_jsonl( + "timed-trace-float-len", + r#"{"timestamp": 0, "input_length": 16.0, "output_length": 4.0, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + assert_eq!(reqs[0].expected_output_len, 4); + } + + #[test] + fn test_non_finite_scaled_timestamp_is_rejected() { + let path = write_temp_jsonl( + "timed-trace-overflow-ts", + r#"{"timestamp": 1e308, "input_length": 16, "output_length": 1, "hash_ids": [1]}"#, + ); + let err = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1e10)) + .expect_err("should reject an offset that overflows to inf"); + assert!(err.to_string().contains("non-finite offset")); + } + + /// Negative offsets load unchanged; the scheduler clamps them to fire + /// immediately (Python's `get_request` does the same). + #[test] + fn test_negative_timestamp_loads() { + let path = write_temp_jsonl( + "timed-trace-negative-ts", + r#"{"timestamp": -5.0, "input_length": 16, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].timestamp, Some(-5.0)); + } +} diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index b2b39b14cf19..085684d5de75 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -175,7 +175,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { + DatasetName::Custom + | DatasetName::PrefixRepetition + | DatasetName::RandomRerank + | DatasetName::TimedTrace => { return Err(BenchError::Config( "This dataset does not support multi-turn. Use 'random' or 'sharegpt' with --multi-turn.".into(), )); diff --git a/rust/src/bench/src/rate_control.rs b/rust/src/bench/src/rate_control.rs index aa05f8f8a8df..1311cb9354be 100644 --- a/rust/src/bench/src/rate_control.rs +++ b/rust/src/bench/src/rate_control.rs @@ -104,10 +104,30 @@ pub fn compute_schedule( } } +/// Schedule built from trace-recorded arrival offsets (`--self-timed`). +/// +/// Offsets are clamped at zero so a row anchored before the trace start fires +/// immediately, as Python's `get_request` does; `Duration::from_secs_f64` in +/// the dispatch loop would otherwise panic. `rates` carries no meaning here. +pub fn trace_schedule(timestamps: impl ExactSizeIterator>) -> RequestSchedule { + let rates = vec![0.0; timestamps.len()]; + RequestSchedule { + delays: timestamps.map(|t| t.unwrap_or(0.0).max(0.0)).collect(), + rates, + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_trace_schedule_uses_offsets_and_clamps_negatives() { + let sched = trace_schedule([Some(-3.0), Some(0.0), Some(2.5), None].into_iter()); + assert_eq!(sched.delays, vec![0.0, 0.0, 2.5, 0.0]); + assert_eq!(sched.rates.len(), 4); + } + #[test] fn test_infinite_rate_all_zero_delays() { let sched = compute_schedule(100, f64::INFINITY, 1.0, 42, None); diff --git a/rust/src/bench/tests/cli_parity.rs b/rust/src/bench/tests/cli_parity.rs index b0fd09ab9e36..b0bb707aa111 100644 --- a/rust/src/bench/tests/cli_parity.rs +++ b/rust/src/bench/tests/cli_parity.rs @@ -36,13 +36,6 @@ const PYTHON_ONLY: &[&str] = &[ // spec_bench dataset "--spec-bench-category", "--spec-bench-output-len", - // timed-trace dataset - "--timed-trace-chunk-hash-size", - "--timed-trace-label-hash-ids", - "--timed-trace-label-input-length", - "--timed-trace-label-output-length", - "--timed-trace-label-timestamp", - "--timed-trace-sec-multiplier", // client-side chat templating / request shaping "--chat-template-kwargs", "--custom-ensure-client-side-data", @@ -53,8 +46,6 @@ const PYTHON_ONLY: &[&str] = &[ "--timeline-itl-thresholds", // misc python-side controls "--hf-name", - "--no-self-timed", - "--self-timed", "--no-stream", "--probe-request-rate", ]; From 1a86733c7b02f1cd5b035e64b9af5ce23e726484 Mon Sep 17 00:00:00 2001 From: Roy Wang Date: Tue, 6 Oct 2026 20:40:10 +0800 Subject: [PATCH 067/952] [Docs] Add model recipes skill (#57786) --- .agents/skills/model-recipes/SKILL.md | 151 ++++++++++++++++++ .../skills/model-recipes/agents/openai.yaml | 4 + .../model-recipes/scripts/read_recipe.sh | 66 ++++++++ .claude/skills/model-recipes | 1 + 4 files changed, 222 insertions(+) create mode 100644 .agents/skills/model-recipes/SKILL.md create mode 100644 .agents/skills/model-recipes/agents/openai.yaml create mode 100755 .agents/skills/model-recipes/scripts/read_recipe.sh create mode 120000 .claude/skills/model-recipes diff --git a/.agents/skills/model-recipes/SKILL.md b/.agents/skills/model-recipes/SKILL.md new file mode 100644 index 000000000000..c462b85a8a0a --- /dev/null +++ b/.agents/skills/model-recipes/SKILL.md @@ -0,0 +1,151 @@ +--- +name: model-recipes +description: Find vLLM deployment recipes and configs for LLMs and Hugging Face models by model and hardware. Use for recipe lookup, creation, and updates. +--- + +# Model Recipes + +Model serving recipes are maintained in the +[models/ directory of vllm-project/recipes](https://github.com/vllm-project/recipes/tree/main/models). +The published catalog is at +[recipes.vllm.ai](https://recipes.vllm.ai). + +## Usage notes + +1. Recipe lookup returns commands and deployment guidance without launching + `vllm serve`. +2. Use published configurations or the recipes repository's command generator + for supported selections. Return commands as published, including their + flags and parallelism settings. +3. Recipes are starting points for deployment and may not be tuned for your + workload. Check each recipe for its verified configurations, and use + benchmarks of your workload to guide further tuning. +4. To create or update a recipe, follow the + [authoring workflow below](#create-or-update-a-recipe). + +## Find a recipe + +Read recipes from the published JSON on recipes.vllm.ai. It is generated from +the source YAML with variant, hardware, and strategy settings already +resolved; the YAML contains comments and conditional overrides, so do not use +it for lookup. When presenting results, link the website page. Read the source +YAML only when creating or updating a recipe. + +Each recipe has a page at `https://recipes.vllm.ai//` and JSON at +the same path plus `.json`. Paths use the full Hugging Face ID, preserving +case; do not substitute `meta.slug` or a display title. For example, +`Qwen/Qwen3.8-27B` maps to +[recipes.vllm.ai/Qwen/Qwen3.8-27B](https://recipes.vllm.ai/Qwen/Qwen3.8-27B). + +Recipe JSON is often over 20 KB, mostly the `guide`. Instead of printing it, +read only what the request needs: each step below returns the paths for the +next, so stop once the request is answered. Commands need `curl` and `jq`. +[scripts/read_recipe.sh](scripts/read_recipe.sh) is in this skill's directory; +the examples use its path from the vLLM repository root. + +### Search the catalog + +Skip this step if the exact HF ID is known. Otherwise search by model ID, +family, or provider using case-insensitive keywords. Each +whitespace-separated word must occur in the ID, title, or provider: + +```bash +curl -fsSL https://recipes.vllm.ai/models.json | + jq --arg q 'qwen 27b' ' + ($q | ascii_downcase | [scan("\\S+")]) as $words | .[] | + ([.hf_id, .title, .provider] | join(" ") | ascii_downcase) as $text | + select(all($words[]; . as $word | $text | contains($word))) | + {hf_id, json, url, derived_from}' +``` + +`json` is the recipe data path; `url` is its page path. Prefix both with +`https://recipes.vllm.ai`. A `derived_from` entry is a variant of a parent +recipe: fetch its own `json` to get that checkpoint's command. Its page may +belong to the parent, so link the returned `url` rather than assuming every +checkpoint has a separate page. + +If there is no exact match, search the family and inspect `variants` before +concluding the recipe is missing. + +### Summarize the recipe + +Save the recipe JSON to a file and summarize it: + +```bash +read_recipe=.agents/skills/model-recipes/scripts/read_recipe.sh +recipe=$(mktemp) +curl -fsSL https://recipes.vllm.ai/Qwen/Qwen3.8-27B.json -o "$recipe" +"$read_recipe" "$recipe" +``` + +`default` is the recommended hardware, strategy, and variant; `hardware` maps +each hardware ID to its command path; `recipe_verified_hardware` lists hardware +where the recipe as a whole was verified end to end. For capability +requirements, such as tool calling or text-to-video, check `features` and +`omni_tasks`, then confirm in the guide. +Some recipes, including Omni, lack `hardware`; skip hardware selection and +use the guide's launch and task-specific sections. + +### Select hardware and strategy + +Match a hardware ID to the user's GPU and fetch its path from `hardware`, +prefixed with `https://recipes.vllm.ai`, rather than constructing a URL. That +response is the command object itself: single-node objects carry `command`, +and multi-node objects carry `head_command` and `worker_commands` instead, so +list its keys before selecting fields. +Check `hardware_profile` to confirm that the selected configuration matches +the user's GPU model and specifications. Its `gpu_count` describes the profile, +not necessarily the number of GPUs used by the command. Check the command's +parallelism settings, `node_count`, and all `worker_commands` against the +user's available topology, and compare the variant's `vram_minimum_gb`, an +estimate, with the total memory of the GPUs the command uses; a configuration +below that minimum is not a working option unless the guide says it was +tested. Ask when the user's hardware is unclear. If no +published command or guide command fits, report the mismatch and return the +closest one unchanged; list any adjustments the user's setup would need +separately and mark them unverified instead of editing the command. + +For another strategy, follow `alternatives` in the same response. These are +static files; query parameters do not change their configuration. For other +feature/topology selections, use the recipes repository's command synthesis. + +### Read deployment guidance + +Read the guide sections that apply to the selection, such as prerequisites, +launch notes for the selected hardware or variant, troubleshooting, and known +limitations. Pass entries from `guide_sections` or prefixes of their text. +The script prints every matching section with its subsections and fails when +a section matches nothing: + +```bash +"$read_recipe" "$recipe" prerequisites troubleshooting +``` + +Query other fields with `jq` on the saved file, such as `features.`, +`hardware_overrides`, or `omni.tasks` for request examples. + +### Return the result + +Return the website link, selected checkpoint and hardware, the requested +command, and relevant prerequisites and caveats from the guide. Pick one +command and copy it, including its Docker form, exactly as the command object +or guide section gives it; do not combine commands from different sections or +add, drop, or reorder flags. List optional changes separately. + +For each command, state its source, either a generated hardware or strategy +configuration or a named guide section, and whether that source says it was +tested. A `recipe_verified_hardware` entry covers the recipe on that +hardware, not every variant, strategy, or guide command, so it does not make a +command verified. Generated configurations are not automatically verified. + +## Create or update a recipe + +Recipes cover official or widely used checkpoints, with configurations verified +end to end on real hardware. Do not add recipes for arbitrary community +uploads, untested hardware claims, or non-standard flags that only work for one +setup. Before opening a pull request, confirm that the user ran the +configuration and can share the results. + +Follow the recipes repository's +[add-recipe skill](https://github.com/vllm-project/recipes/blob/main/.agents/skills/add-recipe/SKILL.md) +in a checkout of `vllm-project/recipes` or the user's intended fork. diff --git a/.agents/skills/model-recipes/agents/openai.yaml b/.agents/skills/model-recipes/agents/openai.yaml new file mode 100644 index 000000000000..6ad0f4d4db41 --- /dev/null +++ b/.agents/skills/model-recipes/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Model Recipes" + short_description: "Find vLLM deployment recipes and configs for LLMs and Hugging Face models by model and hardware." + default_prompt: "Use $model-recipes to find a recipe for my model and hardware, and return its commands and deployment caveats." diff --git a/.agents/skills/model-recipes/scripts/read_recipe.sh b/.agents/skills/model-recipes/scripts/read_recipe.sh new file mode 100755 index 000000000000..28e880537e4e --- /dev/null +++ b/.agents/skills/model-recipes/scripts/read_recipe.sh @@ -0,0 +1,66 @@ +#!/bin/bash +# Usage: read_recipe.sh [guide-section...] +# Without sections, print a compact summary of a recipes.vllm.ai recipe JSON. +# With sections, print each guide section whose heading starts with one of +# them (case-insensitive), including its subsections. + +set -euo pipefail + +if [ $# -lt 1 ]; then + echo "Usage: $0 [guide-section...]" >&2 + exit 1 +fi + +# Print "\t\t" for guide lines, marking headings +# outside code fences with their level and other lines with 0. +guide_lines() { + jq -r '.guide // ""' "$1" | awk ' + /^[[:space:]]*(```|~~~)/ { fence = !fence } + !fence && match($0, /^#+ /) { + print RLENGTH - 1 "\t" substr($0, RLENGTH + 1) "\t" $0 + next + } + { print "0\t\t" $0 }' +} + +if [ $# -gt 1 ]; then + recipe=$1 + shift + status=0 + for section in "$@"; do + guide_lines "$recipe" | awk -F '\t' -v h="$section" ' + BEGIN { h = tolower(h); sub(/^#+[[:space:]]*/, "", h) } + $1 > 0 && lvl && $1 <= lvl { lvl = 0 } + $1 > 1 && !lvl && index(tolower($2), h) == 1 { lvl = $1; found = 1 } + lvl { sub(/^[^\t]*\t[^\t]*\t/, ""); print } + END { + if (!found) { + print "No guide section starts with: " h > "/dev/stderr" + exit 1 + } + }' || status=1 + done + exit $status +fi + +sections=$(guide_lines "$1" | awk -F '\t' '$1 == 2 || $1 == 3 { print $3 }' | + jq -R . | jq -s .) + +jq --argjson sections "$sections" '{ + hf_id, + title: .meta.title, + tasks: .meta.tasks, + recipe_verified_hardware: .meta.hardware, + omni_tasks: (.omni.tasks // null | + if . then map(objects |= {id, label, endpoint, model_id, description}) + else null end), + variants: (.variants // {} | map_values( + {model_id, precision, vram_minimum_gb, supported_hardware} | + with_entries(select(.value != null)))), + features: (.features // {} | keys), + opt_in_features, + default: (.recommended_command // null | + if . then {hardware, strategy, variant} else null end), + hardware: .recommended_command.by_hardware, + guide_sections: $sections +} | with_entries(select(.value != null))' "$1" diff --git a/.claude/skills/model-recipes b/.claude/skills/model-recipes new file mode 120000 index 000000000000..5231457b9c8a --- /dev/null +++ b/.claude/skills/model-recipes @@ -0,0 +1 @@ +../../.agents/skills/model-recipes \ No newline at end of file From 049507aa76200090f8bc6ca4e48ab0e09f2f0d29 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Tue, 6 Oct 2026 16:01:04 +0300 Subject: [PATCH 068/952] [Bugfix][Structured Output] Flag multi-branch allOf as unsupported for xgrammar (#59061) Signed-off-by: Artem Perevedentsev Co-authored-by: Shashank Varma <324153016+shashankvarma499@users.noreply.github.com> --- ...test_utils.py => test_backend_xgrammar.py} | 179 ++++++++++++++++-- vllm/v1/structured_output/backend_xgrammar.py | 30 ++- 2 files changed, 184 insertions(+), 25 deletions(-) rename tests/v1/structured_output/{test_utils.py => test_backend_xgrammar.py} (72%) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_backend_xgrammar.py similarity index 72% rename from tests/v1/structured_output/test_utils.py rename to tests/v1/structured_output/test_backend_xgrammar.py index d3187c428cc0..339f0de17f03 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_backend_xgrammar.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import pytest from xgrammar import Grammar from xgrammar.testing import _is_grammar_accept_string @@ -19,6 +20,11 @@ def grammar_accepts(schema: dict, text: str) -> bool: return _is_grammar_accept_string(Grammar.from_json_schema(schema), text) +# ================================================ +# Unsupported schemas +# ================================================ + + @pytest.fixture def unsupported_string_schemas(): return [ @@ -35,16 +41,12 @@ def unsupported_string_schemas(): @pytest.fixture -def unsupported_integer_schemas(): - return [ - {"type": "integer", "multipleOf": 120}, - ] - - -@pytest.fixture -def unsupported_number_schemas(): +def unsupported_multipleOf_schemas(): return [ {"type": "number", "multipleOf": 120}, + {"Even": {"type": "number", "multipleOf": 2}}, + {"type": "integer", "multipleOf": 120}, + {"Even": {"type": "integer", "multipleOf": 2}}, ] @@ -140,6 +142,81 @@ def unsupported_pattern_properties_combinations(): ] +@pytest.fixture +def unsupported_multibranch_allof(): + """Regression tests for issue #56556 + [Bug]: JSON schema with multiple allOf branches is silently ignored + by the xgrammar structured-output backend #56556 + + Xgrammar currently does NOT support multi-branch allOf. + Bug reported to xgrammar team for tracking progress on resolving: + https://github.com/mlc-ai/xgrammar/issues/937 + """ + return [ + { + "allOf": [ + {"type": "string"}, + {"enum": ["yes", "no"]}, + ] + }, + # Multi-branch test case with > 2 branches: + { + "allOf": [ + {"type": "string"}, + {"minLength": 2}, + {"maxLength": 10}, + ] + }, + { + "allOf": [ + {"type": "integer"}, + {"minimum": 0}, + ] + }, + # Non root multi-branch allOf: + { + "type": "object", + "properties": { + "is_this_sparta": { + "allOf": [ + {"type": "string"}, + {"enum": ["yes", "no"]}, + ], + }, + }, + }, + ] + + +@pytest.fixture +def unsupported_vllm_issue_56556_schema(): + """Bug repro from vLLM issue #56556""" + return [ + { + "$defs": { + "Base": { + "type": "object", + "properties": {"x": {"type": "integer", "minimum": 10}}, + "required": ["x"], + } + }, + "allOf": [ + {"$ref": "#/$defs/Base"}, + { + "type": "object", + "properties": {"y": {"type": "string"}}, + "required": ["y"], + }, + ], + }, + ] + + +# ================================================ +# Supported schemas +# ================================================ + + @pytest.fixture def supported_frankenstein_schema(): # IMPORTANT(arpera): @@ -177,11 +254,6 @@ def supported_frankenstein_schema(): } -@pytest.fixture -def property_names_schema(): - return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} - - @pytest.fixture def pattern_properties_schema(): return { @@ -190,19 +262,80 @@ def pattern_properties_schema(): } +@pytest.fixture +def property_names_schema(): + return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} + + +@pytest.fixture +def supported_allof_anyof_and_oneof(): + """Xgrammar DO support some variants of anyOf, allOf, oneOf + This test case was made during working on covering + unsupported multi-branch allOf to make sure the bug + is NOT present in other combinators. + See for more context in this file test `unsupported_multibranch_allof` + """ + return [ + # single-branch allOf: + { + "allOf": [ + {"type": "string"}, + ], + }, + # Corner cases: empty allOf, this is a valid schema + {"allOf": []}, + # multi-branch allOf: + # is NOT supported yet, see vLLM issue #56556 + # single-branch anyOf: + { + "anyOf": [], + }, + # multi-branch anyOf: + { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + ], + }, + # single-branch oneOf: + { + "oneOf": [ + {"type": "integer", "minimum": 0, "maximum": 100}, + ], + }, + # multi-branch oneOf: + { + "oneOf": [ + {"type": "integer", "minimum": 0, "maximum": 100}, + {"type": "string", "enum": ["auto", "none"]}, + ], + }, + # Corner case: + # "allOf" is a property name, not allOf combinator keyword + {"type": "object", "properties": {"allOf": {"type": "string"}}}, + ] + + +# ================================================ +# Test has_xgrammar_unsupported_json_features functionality +# ================================================ + + class TestHasXGrammarUnsupportedJsonFeatures: @pytest.mark.parametrize( "schema_type", [ "unsupported_string_schemas", - "unsupported_integer_schemas", - "unsupported_number_schemas", + "unsupported_multipleOf_schemas", "unsupported_array_schemas", "unsupported_property_names_combinations", "unsupported_pattern_properties_combinations", + # vLLM issue #56556 + "unsupported_multibranch_allof", + "unsupported_vllm_issue_56556_schema", ], ) - def test_unsupported_json_features_by_type(self, schema_type, request): + def test_unsupported_json_features(self, schema_type, request): schemas = request.getfixturevalue(schema_type) for schema in schemas: assert has_xgrammar_unsupported_json_features(schema), ( @@ -215,13 +348,19 @@ def test_unsupported_json_features_by_type(self, schema_type, request): "supported_frankenstein_schema", "pattern_properties_schema", "property_names_schema", + # Additional test cases implemented during work on + # vLLM issue #56556 + "supported_allof_anyof_and_oneof", ], ) def test_supported_json_features(self, schema_type, request): - schema = request.getfixturevalue(schema_type) - assert not has_xgrammar_unsupported_json_features(schema), ( - f"Schema should be supported: {schema}" - ) + schemas = request.getfixturevalue(schema_type) + if not isinstance(schemas, list): + schemas = [schemas] + for schema in schemas: + assert not has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be supported: {schema}" + ) class TestPR48416Regressions: @pytest.mark.parametrize( diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 4697b00487fe..733ed7db092b 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -335,6 +335,24 @@ def check_object(obj: dict[str, Any]) -> bool: ): return True + # Note(arpera): + # Xgrammar lacks support of multi-branch allOf + # For instance, this schema: + # { + # "allOf": [ + # { "type": "string" }, + # { "enum": ["yes", "no"] } + # ] + # } + # would accept any kind of json, such as + # "maybe", "", 42, {}, [], {"a": 1}, etc. + # which is NOT what is expected. + # Reported this issue to xgrammar team to track progress on resolving: + # https://github.com/mlc-ai/xgrammar/issues/937 + allof = obj.get("allOf") + if isinstance(allof, list) and len(allof) >= 2: + return True + # Recursively check all nested objects and arrays for value in obj.values(): if isinstance(value, dict): @@ -401,13 +419,15 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: else: schema = so_params.json - if has_xgrammar_unsupported_json_features(schema): - raise VLLMValidationError( - "The provided JSON schema contains features not supported by xgrammar." - ) - try: + if has_xgrammar_unsupported_json_features(schema): + raise VLLMValidationError( + "The provided JSON schema contains features not supported " + "by xgrammar." + ) xgr.Grammar.from_json_schema(schema) + except VLLMValidationError: + raise except Exception as err: raise VLLMValidationError( f"Failed to transform json schema into a grammar: {err}" From b267fe488102eb26555d553a005127b4a523ff4b Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Tue, 6 Oct 2026 15:22:33 +0200 Subject: [PATCH 069/952] [ROCm][Perf] W4A16: pad gfx11 weight and activation strides (#56301) Signed-off-by: Matthias Gehre --- .../benchmark_rdna_hybrid_w4a16_gemm.py | 37 ++-- csrc/rocm/skinny_gemms_int4.cu | 61 +++--- .../quantization/test_rdna_hybrid_w4a16.py | 179 ++++++++++++++++++ .../mixed_precision/rdna_hybrid_w4a16.py | 85 ++++++++- 4 files changed, 316 insertions(+), 46 deletions(-) diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 2d5ec80c6aae..7ce49e2e1132 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -15,6 +15,9 @@ import torch +from vllm.model_executor.kernels.linear.mixed_precision.rdna_hybrid_w4a16 import ( + pack_skinny_int4, +) from vllm.triton_utils import triton # --------------------------------------------------------------------------- @@ -56,20 +59,25 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"): Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives its int32 view from w_q_skinny, so no separate int32 buffer is returned. + + Weights go through ``pack_skinny_int4``, the packer the layer uses, so the + benchmark sees the production row stride rather than an always-dense one. """ num_groups = K // group_size - # Random packed weights — actual values don't matter for throughput - w_q_skinny_i32 = torch.randint( - 0, 2**31, (N, K // 8), dtype=torch.int32, device=device - ) - w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous() + # Values don't matter for throughput, but the layout does. + unpacked = torch.randint(0, 16, (N, K), dtype=torch.int32, device=device) + w_q_skinny = pack_skinny_int4(unpacked) + del unpacked w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01 - # Raw per-group zero-points for asymmetric benchmarks - w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to( - torch.float16 - ) + # Per-group zero-points for asymmetric benchmarks, in the packed layout the + # kernels read: [N//8, num_groups] int32, row n's nibble at bits 4*(n%8). + zp_raw = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device) + shifts = (torch.arange(8, device=device, dtype=torch.int32) * 4)[:, None] + w_zp = torch.sum( + (zp_raw.view(N // 8, 8, num_groups) & 0xF) << shifts, dim=1, dtype=torch.int32 + ).contiguous() # FP16 baseline for F.linear w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01 @@ -97,6 +105,13 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"): ) ) def benchmark(batch_size, provider, N, K, group_size, weights): + """Time one GEMM. + + ``do_bench`` rather than ``do_bench_cudagraph``: it zeroes a 256 MB buffer + between reps, so the weight is read from DRAM as it would be in a model. + With a cudagraph replay the weight stays resident in the 32 MB MALL, which + flatters wide-BLOCK_N tiles and hides row-stride effects entirely. + """ M = batch_size device = "cuda" dtype = torch.float16 @@ -106,7 +121,7 @@ def benchmark(batch_size, provider, N, K, group_size, weights): if provider == "torch-fp16": w_fp16 = weights["w_fp16"] - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = triton.testing.do_bench( lambda: torch.nn.functional.linear(a, w_fp16), quantiles=quantiles, ) @@ -133,7 +148,7 @@ def run(): group_size, ) - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = triton.testing.do_bench( run, quantiles=quantiles, ) diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index a20ca41d413e..4588b216fe01 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -8,6 +8,7 @@ #include #include +#include // std::in_range #include "../cuda_compat.h" #include "dispatch_utils.h" @@ -121,6 +122,7 @@ __device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed, // GROUP_SIZE: 0 = per-channel scale [M], >0 = per-group scale [M, // K/GROUP_SIZE]. // Requires GROUP_SIZE % A_CHUNK == 0 when GROUP_SIZE > 0. +// B_row_stride_bytes is the per-row byte stride of the packed weights. #if defined(__HIP__GFX1X__) template @@ -130,9 +132,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) const scalar_t* __restrict__ A, const scalar_t* scale, const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { + const int _WvPrGrp, const int CuCount, + const int B_row_stride_bytes) { constexpr int max_lds_len = LDS_SIZE / 2; - const int K_packed = K / 2; union bigTypeA { scalar_t h[A_CHUNK]; @@ -181,9 +183,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t k_ = k + threadIdx.x * A_CHUNK; if (k_ >= K) break; - const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2]; + const uint8_t* B_ = &B_packed[(m + 0) * B_row_stride_bytes + k_ / 2]; for (int y = 0; y < YTILE; y++) { - const float* src = (const float*)(&B_[y * K_packed]); + const float* src = (const float*)(&B_[y * B_row_stride_bytes]); #pragma unroll for (int i = 0; i < A_CHUNK / 8; i++) bigB[y][k2].f[i] = loadnt((float*)&src[i]); @@ -344,14 +346,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #else // !defined(__HIP__GFX1X__) template -__global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, - const int By, const uint8_t* B_packed, - const scalar_t* __restrict__ A, - const scalar_t* scale, - const uint32_t* zero_points, - const scalar_t* __restrict__ BIAS, - scalar_t* C, const int _WvPrGrp, - const int CuCount) { +__global__ void wvSplitK_int4_hf_sml_( + const int K, const int M, const int Bx, const int By, + const uint8_t* B_packed, const scalar_t* __restrict__ A, + const scalar_t* scale, const uint32_t* zero_points, + const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, + const int CuCount, const int B_row_stride_bytes) { UNREACHABLE_CODE } #endif // defined(__HIP__GFX1X__) @@ -367,9 +367,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { + const int _WvPrGrp, const int CuCount, + const int B_row_stride_bytes) { constexpr int max_lds_len = LDS_SIZE / 2; - const int K_packed = K / 2; union bigTypeA { scalar_t h[A_CHUNK]; @@ -430,9 +430,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) uint32_t k_ = k + threadIdx.x * A_CHUNK; if (k_ >= K) break; - const uint8_t* B_ = &B_packed[(m + 0) * K_packed + k_ / 2]; + const uint8_t* B_ = &B_packed[(m + 0) * B_row_stride_bytes + k_ / 2]; for (int y = 0; y < YTILE; y++) { - const float* src = (const float*)(&B_[y * K_packed]); + const float* src = (const float*)(&B_[y * B_row_stride_bytes]); #pragma unroll for (int i = 0; i < A_CHUNK / 8; i++) bigB[y][k2].f[i] = loadnt((float*)&src[i]); @@ -606,14 +606,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) #else // !defined(__HIP__GFX1X__) template -__global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx, - const int By, const uint8_t* B_packed, - const scalar_t* __restrict__ A, - const scalar_t* scale, - const uint32_t* zero_points, - const scalar_t* __restrict__ BIAS, - scalar_t* C, const int _WvPrGrp, - const int CuCount) { +__global__ void wvSplitK_int4_hf_( + const int K, const int M, const int Bx, const int By, + const uint8_t* B_packed, const scalar_t* __restrict__ A, + const scalar_t* scale, const uint32_t* zero_points, + const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, + const int CuCount, const int B_row_stride_bytes) { UNREACHABLE_CODE } #endif // defined(__HIP__GFX1X__) @@ -658,10 +656,17 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, ? in_bias->size(0) : 1; - int64_t expected_weight_bytes = M_in * K_in / 2; - int64_t actual_weight_bytes = in_a.numel() * in_a.element_size(); + const int64_t expected_weight_bytes = M_in * K_in / 2; + const int64_t actual_weight_bytes = in_a.numel() * in_a.element_size(); TORCH_CHECK(actual_weight_bytes == expected_weight_bytes, "Weight tensor must contain M*K/2 bytes for int4 packing"); + const int64_t b_row_stride_bytes = in_a.stride(0) * in_a.element_size(); + TORCH_CHECK(b_row_stride_bytes >= K_in / 2, "Weight row stride (", + b_row_stride_bytes, " bytes) must hold at least K/2=", K_in / 2, + " bytes per row"); + TORCH_CHECK(std::in_range(b_row_stride_bytes), "Weight row stride (", + b_row_stride_bytes, " bytes) exceeds int range"); + const int b_row_stride_bytes_i32 = static_cast(b_row_stride_bytes); TORCH_CHECK( in_b.dtype() == torch::kFloat16 || in_b.dtype() == torch::kBFloat16, "Activation must be float16 or bfloat16"); @@ -722,12 +727,12 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, wvSplitK_int4_hf_sml_<<>>( \ K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ - __wvPrGrp, CuCount); \ + __wvPrGrp, CuCount, b_row_stride_bytes_i32); \ else \ wvSplitK_int4_hf_<<>>( \ K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, cptr, \ - __wvPrGrp, CuCount); \ + __wvPrGrp, CuCount, b_row_stride_bytes_i32); \ } #define WVSPLITK_INT4G(_YTILE, _UNRL, _N, _GS, _HAS_ZP) \ diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index fe3db8dab360..d6045a791c1f 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -563,3 +563,182 @@ def test_wvsplitk_int4_g_rejects_unpacked_zero_points(): with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"): ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None) + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.skipif( + not hasattr(torch.ops, "_rocm_C") + or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"), + reason="wvSplitK_int4_g not built", +) +def test_wvsplitk_int4_g_rejects_undersized_weight_view(): + """A sufficient row stride does not prove the row holds K/2 packed bytes. + + Row padding makes stride(0) > K/2 legitimate, but the logical size still + has to be checked: this view has a valid stride and half the required + bytes, and without the size check the kernel reads past it. + """ + import vllm._custom_ops as ops + from vllm.utils.platform_utils import num_compute_units + + K, N, G, M = 512, 256, 128, 1 + a = torch.randn((M, K), device=device, dtype=torch.float16) + # Full-size backing allocation, so a missing check misreads rather than + # running off the end of the buffer. + backing = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view( + torch.int8 + ) + undersized = backing[:, : K // 4] + assert undersized.stride(0) == K // 2 + scales = torch.rand((N, K // G), device=device, dtype=torch.float16) + + with pytest.raises(RuntimeError, match=r"must contain M\*K/2 bytes"): + ops.wvSplitK_int4_g(undersized, a, scales, num_compute_units(), G, None, None) + + +# --------------------------------------------------------------------------- +# gfx11 weight row-stride padding +# --------------------------------------------------------------------------- + +_weight_pad_bytes = hybrid_module._weight_pad_bytes +_act_pad_bytes = hybrid_module._act_pad_bytes +_WEIGHT_CLIFF_BYTES = hybrid_module._WEIGHT_CLIFF_BYTES +_ACT_CLIFF_BYTES = hybrid_module._ACT_CLIFF_BYTES +_STRIDE_PAD_BYTES = hybrid_module._STRIDE_PAD_BYTES + + +def test_weight_pad_bytes_only_moves_strides_on_the_cliff(): + """Pad 1024 B multiples, leave everything else dense.""" + # On the cliff (K % 2048 == 0): 1024 B multiples. + for row_bytes in (1024, 2048, 4096, 5120, 6144, 7168, 8192): + assert _weight_pad_bytes(row_bytes) == _STRIDE_PAD_BYTES + + # Off it, including strides that are 512 B multiples but not 1024 B ones: + # padding those measured as a real loss on gfx1151 (-14% at M=1 on a + # 4864 B row). + for row_bytes in (1280, 1536, 2560, 4864, 9472, 12800): + assert _weight_pad_bytes(row_bytes) == 0 + + # The padded stride is never back on the cliff. + for row_bytes in range(16, 16384, 16): + padded = row_bytes + _weight_pad_bytes(row_bytes) + assert padded % _WEIGHT_CLIFF_BYTES != 0 or row_bytes % _WEIGHT_CLIFF_BYTES + + +def test_act_pad_bytes_only_moves_strides_on_the_cliff(): + """Activations sit on a wider cliff than the packed weight: 2048 B.""" + # On the cliff (K % 1024 == 0 for a 2-byte dtype): 2048 B multiples. + for row_bytes in (2048, 4096, 8192, 16384, 20480, 24576, 43008): + assert _act_pad_bytes(row_bytes) == _STRIDE_PAD_BYTES + + # Off it, including 1024 B multiples that are not 2048 B ones -- the + # weight cliff must not be applied to activations. + for row_bytes in (1024, 3072, 5120, 10752, 19456, 37888): + assert _act_pad_bytes(row_bytes) == 0 + + # The padded stride is never back on the cliff. + for row_bytes in range(16, 65536, 16): + padded = row_bytes + _act_pad_bytes(row_bytes) + assert padded % _ACT_CLIFF_BYTES != 0 or row_bytes % _ACT_CLIFF_BYTES + + +def _row_padded_copy(t: torch.Tensor, pad_cols: int) -> torch.Tensor: + """Copy of ``t`` whose stride(0) is ``t.shape[1] + pad_cols``.""" + rows, cols = t.shape + buf = torch.empty((rows, cols + pad_cols), dtype=t.dtype, device=t.device) + buf[:, :cols].copy_(t) + return buf[:, :cols] + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True]) +@pytest.mark.parametrize("M", [1, MAX_SKINNY_BATCH_SIZE + 1]) +def test_weight_row_padding_does_not_change_results(dtype, has_zp, M): + """Padding the weight rows is a pure layout change. + + Both the HIP skinny kernel (M small) and the Triton prefill kernel take the + weight row stride from the tensor, so a padded layout must produce exactly + the same numbers as the dense one. + """ + from vllm.utils.platform_utils import num_compute_units + + set_random_seed(0) + K, N, G = 512, 256, 128 + + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) + w_q = pack_int4_exllama_shuffle(w_int4_nk).view(torch.int8) + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( + dtype + ) + zp = ( + _pack_zp_rows_for_kernel( + torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32) + ) + if has_zp + else None + ) + + # 32 int32 columns = 128 B of pad, the production pad size. + w_q_pad = _row_padded_copy(w_q.view(torch.int32), 32).view(torch.int8) + assert w_q_pad.stride(0) == w_q.stride(0) + _STRIDE_PAD_BYTES + + cu_count = num_compute_units() + dense = torch.ops.vllm.rdna_hybrid_w4a16_apply( + a, w_q, scales, zp, None, cu_count, G + ) + padded = torch.ops.vllm.rdna_hybrid_w4a16_apply( + a, w_q_pad, scales, zp, None, cu_count, G + ) + torch.testing.assert_close(padded, dense, rtol=0, atol=0) + + +@pytest.mark.skipif( + not hybrid_module._on_gfx1151(), + reason="row-stride padding is only enabled on gfx1151", +) +@pytest.mark.parametrize( + "K,expected_weight_stride", + [ + (512, 256), # 256 B row: off the cliff, stored dense + (8192, 4096 + 128), # 4096 B row: on the cliff, padded + ], +) +def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init): + """The stored weight row stride follows the padding rule.""" + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, + ) + from vllm.scalar_type import scalar_types + + set_random_seed(0) + N, G = 128, 128 + + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + layer = _build_dummy_layer( + _pack_int4_along_k_to_ckpt(w_int4_kn), + 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16), + zeros_ckpt=None, + ) + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=G, + zero_points=False, + ) + RDNAHybridW4A16LinearKernel( + config, + w_q_param_name="weight_packed", + w_s_param_name="weight_scale", + w_zp_param_name=None, + ).process_weights_after_loading(layer) + + # Weight rows: stride is in int8 elements, i.e. bytes. + assert layer.weight_packed.stride(0) == expected_weight_stride + # The int32 view the Triton path uses survives the padded stride. + assert tuple(layer.weight_packed.view(torch.int32).shape) == (N, K // 8) + # Metadata is untouched by this change. + assert layer.weight_scale.is_contiguous() diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 86e029b5f88d..9d5303469945 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -84,6 +84,8 @@ def _triton_w4a16_skinny_fmt_kernel( K, K8, # K // 8 num_groups, # K // group_size + stride_bn, # b_ptr row stride; >= K8, the rows may be padded + stride_am, # a_ptr row stride in elements; >= K, rows may be padded # Quantization parameters group_size, ZP_BIAS: tl.constexpr, @@ -127,12 +129,12 @@ def _triton_w4a16_skinny_fmt_kernel( offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K) mask_k = offs_k < K - a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :] + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] mask_a = (offs_m[:, None] < M) & mask_k[None, :] a = tl.load(a_ptrs, mask=mask_a, other=0.0) offs_k8 = k_start * (BLOCK_K // 8) + tl.arange(0, BLOCK_K // 8) - b_ptrs = b_ptr + offs_n[:, None] * K8 + offs_k8[None, :] + b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k8[None, :] mask_b = (offs_n[:, None] < N) & (offs_k8[None, :] < K8) b_packed = tl.load(b_ptrs, mask=mask_b, other=0) @@ -210,14 +212,16 @@ def triton_w4a16_skinny_fmt_gemm( Output matrix [M, N], same dtype as a. """ - assert a.is_contiguous(), "Activation matrix must be contiguous" - assert b_q.is_contiguous(), "Weight matrix must be contiguous" + assert a.stride(1) == 1, "Activation rows must be contiguous" + assert b_q.stride(1) == 1, "Weight rows must be contiguous" assert scales.is_contiguous(), "Scales must be contiguous" M, K = a.shape N = b_q.shape[0] K8 = K // 8 num_groups = K // group_size + stride_bn = b_q.stride(0) + stride_am = a.stride(0) assert b_q.shape == (N, K8), f"b_q shape mismatch: {b_q.shape} vs ({N}, {K8})" assert scales.shape == (N, num_groups), ( @@ -331,6 +335,8 @@ def triton_w4a16_skinny_fmt_gemm( K, K8, num_groups, + stride_bn, + stride_am, group_size=group_size, ZP_BIAS=zp_bias, HAS_ZP=has_zp, @@ -368,11 +374,77 @@ def pack_int4_exllama_shuffle(w_uint4: torch.Tensor) -> torch.Tensor: ) +# gfx11 row strides alias in the vector caches when they land on a power-of-two +# multiple. +_WEIGHT_CLIFF_BYTES = 1024 # packed weight row, K/2 B +_ACT_CLIFF_BYTES = 2048 # activation row, K*2 B +_STRIDE_PAD_BYTES = 128 # one cache line + + +def _weight_pad_bytes(row_bytes: int) -> int: + """Bytes to add to a packed weight row stride to move it off the cliff. + + Only strides that are a multiple of ``_WEIGHT_CLIFF_BYTES`` are moved, and + only by one cache line; padding a stride that is already off the cliff + costs memory and measures slower. + """ + if row_bytes % _WEIGHT_CLIFF_BYTES: + return 0 + return _STRIDE_PAD_BYTES + + +def _act_pad_bytes(row_bytes: int) -> int: + """Bytes to add to an activation row stride to move it off the cliff.""" + if row_bytes % _ACT_CLIFF_BYTES: + return 0 + return _STRIDE_PAD_BYTES + + +def pack_skinny_int4(unpacked: torch.Tensor) -> torch.Tensor: + """Pack [N, K] uint4 into the skinny weight layout the kernels consume. + + ExLlama shuffle to [N, K//8] int32, viewed as int8 [N, K//2]. On gfx1151 + the row stride is nudged off the cliff (see ``_weight_pad_bytes``). + """ + shuffled = pack_int4_exllama_shuffle(unpacked) + n_rows, k8 = shuffled.shape + pad_int32 = _weight_pad_bytes(k8 * 4) // 4 + if not (pad_int32 and _on_gfx1151()): + return shuffled.contiguous().view(torch.int8) + padded = torch.empty( + (n_rows, k8 + pad_int32), dtype=torch.int32, device=shuffled.device + ) + padded[:, :k8].copy_(shuffled) + # The int8 view keeps stride(0) = 4 * (k8 + pad_int32) bytes, and + # .view(torch.int32) at apply time recovers the [N, K//8] int32 view. + return padded.view(torch.int8)[:, : k8 * 4] + + # --------------------------------------------------------------------------- # Hybrid dispatch logic # --------------------------------------------------------------------------- +def _pad_activation_rows(x_2d: torch.Tensor) -> torch.Tensor: + """Copy ``x_2d`` into a row-padded buffer when its row stride is on the cliff. + + The packed weight is padded once at load time, but activations are produced + fresh every step, so this materialises a padded copy. + """ + pad_bytes = _act_pad_bytes(x_2d.shape[1] * x_2d.element_size()) + if not (pad_bytes and _on_gfx1151()): + return x_2d + pad_elems = pad_bytes // x_2d.element_size() + buf = torch.empty( + (x_2d.shape[0], x_2d.shape[1] + pad_elems), + dtype=x_2d.dtype, + device=x_2d.device, + ) + padded = buf[:, : x_2d.shape[1]] + padded.copy_(x_2d) + return padded + + def _rdna_hybrid_w4a16_apply_impl( x_2d: torch.Tensor, w_q: torch.Tensor, @@ -410,7 +482,7 @@ def _rdna_hybrid_w4a16_apply_impl( ) with ctx: output = triton_w4a16_skinny_fmt_gemm( - a=x_2d, + a=_pad_activation_rows(x_2d), b_q=w_q.view(torch.int32), scales=w_s, group_size=group_size, @@ -513,9 +585,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if getattr(w_q_raw, "output_dim", 0) != 0: unpacked = unpacked.t().contiguous() - shuffled = pack_int4_exllama_shuffle(unpacked) # Store as int8; Triton reinterprets via .view(torch.int32) at apply time. - w_q_skinny = shuffled.contiguous().view(torch.int8) + w_q_skinny = pack_skinny_int4(unpacked) permute_param_layout_(w_s_raw, input_dim=1, output_dim=0) w_s_skinny = w_s_raw.data.contiguous() From bf049ddb8f4f305cdb01be63f89743f83b397228 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 6 Oct 2026 22:12:18 +0800 Subject: [PATCH 070/952] [Docs] Run pr-checklist before opening a PR (#60235) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 --- .agents/skills/pr-checklist/SKILL.md | 2 +- AGENTS.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.agents/skills/pr-checklist/SKILL.md b/.agents/skills/pr-checklist/SKILL.md index 886e6211a72c..4d73be6adf3e 100644 --- a/.agents/skills/pr-checklist/SKILL.md +++ b/.agents/skills/pr-checklist/SKILL.md @@ -1,6 +1,6 @@ --- name: pr-checklist -description: Prepare a vLLM change for human PR review or re-review, or analyze an open pull request. Use to check design fit, behavioral coverage, performance evidence, diff quality, and closure of previous feedback before requesting maintainer attention. +description: Prepare a vLLM change for human PR review or re-review, or analyze an open pull request. Use before opening a PR, drafts included (e.g. `gh pr create`), or pushing review fixes, to check design fit, behavioral coverage, performance evidence, diff quality, and closure of previous feedback. --- # PR Checklist diff --git a/AGENTS.md b/AGENTS.md index b3ce2603834c..933fa05cfcab 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -31,7 +31,7 @@ Do not open one-off PRs for tiny edits (single typo, isolated style change, one - Test commands run and results. - Model evaluation results when the change affects output, accuracy, or serving. - Clear statement that AI assistance was used. -- Use vLLM's [`/pr-checklist`](.agents/skills/pr-checklist/SKILL.md) skill before submitting a PR to ensure quality and adherence to the contribution guidelines. +- Before opening a PR (drafts included) or requesting re-review, run the [`pr-checklist`](.agents/skills/pr-checklist/SKILL.md) skill and address its findings. ### Fail-closed behavior From 4a99ca40d647a0b5618315fe44fcb85554ddcb53 Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Tue, 6 Oct 2026 16:37:01 +0200 Subject: [PATCH 071/952] [Perf][MoE] Eliminate staging from NCCL symmetric reduce-scatter (#49194) Signed-off-by: samnordmann Co-authored-by: Misha Goin --- tests/distributed/test_mnnvl_alltoall.py | 209 +++++++++++++++++- tests/distributed/test_nccl_symm_mem.py | 154 ++++++++++++- .../device_communicators/all2all.py | 36 +++ .../base_device_communicator.py | 55 +++++ .../device_communicators/cuda_communicator.py | 183 +++++++++++---- vllm/distributed/parallel_state.py | 39 ++++ .../layers/fused_moe/modular_kernel.py | 72 ++++-- .../fused_moe/prepare_finalize/naive_dp_ep.py | 44 +++- 8 files changed, 723 insertions(+), 69 deletions(-) diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index f7dac37940e4..bd7ff4df608a 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -150,7 +150,9 @@ def _init_dp_environment(world_size, rank, port, dp_size, dp_port): @contextmanager -def _make_forward_context(rank, world_size, num_tokens_per_rank): +def _make_forward_context( + rank, world_size, num_tokens_per_rank, num_tokens_across_dp=None +): """Keep the vLLM config and mock-DP forward context active together. Returns a context manager suitable for ``with`` statements. @@ -175,15 +177,15 @@ class _AttnMeta: is_moe_model=True, data_parallel_rank=rank, ) + if num_tokens_across_dp is None: + num_tokens_across_dp = [num_tokens_per_rank] * world_size with ( set_current_vllm_config(vllm_config), set_forward_context( _AttnMeta(), vllm_config, num_tokens=num_tokens_per_rank, - num_tokens_across_dp=torch.tensor( - [num_tokens_per_rank] * world_size, dtype=torch.int - ), + num_tokens_across_dp=torch.tensor(num_tokens_across_dp, dtype=torch.int), ), ): yield @@ -516,6 +518,121 @@ def test_one_sided_manager_workspace_grow(world_size): # --------------------------------------------------------------------------- +class _FakeCombineGroup: + def __init__(self, combine_input): + self.combine_input = combine_input + self.allocation_calls = 0 + self.combined_input = None + + def allocate_combine_input(self, shape, dtype, device, **kwargs): + self.allocation_calls += 1 + assert self.combine_input is not None + assert self.combine_input.shape == shape + assert self.combine_input.dtype == dtype + assert self.combine_input.device == device + return self.combine_input + + def combine_into_output(self, hidden_states, output, **kwargs): + self.combined_input = hidden_states + output.copy_(hidden_states) + return output + + +def test_args_cpu_workspace_allocation(monkeypatch): + from types import SimpleNamespace + + import vllm.model_executor.layers.fused_moe.modular_kernel as mk + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + kernel = object.__new__(mk.FusedMoEKernelModularImpl) + kernel.fused_experts = SimpleNamespace( + workspace_dtype=lambda dtype: dtype, + workspace_shapes=lambda *args: ((0,), (0,), (2, 3)), + finalize_weight_and_reduce_impl=lambda: None, + ) + kernel.prepare_finalize = SimpleNamespace( + allocate_fused_expert_output=lambda *args: None, + ) + monkeypatch.setattr(mk.current_platform, "is_cpu", lambda: True) + + workspace13, workspace2, output, specialized = kernel._allocate_buffers( + torch.float32, + torch.device("cpu"), + 2, + 2, + 3, + 3, + 1, + 1, + 1, + None, + MoEActivation.SILU, + ) + assert workspace13.numel() == workspace2.numel() == 0 + assert output.shape == (2, 3) + assert output.device.type == "cpu" + assert not specialized + + +def test_args_finalize_noop_reuses_fused_expert_output(monkeypatch): + import vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep as pf + from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, + ) + + fused_expert_output = torch.arange(6, dtype=torch.float32).view(2, 3) + output = torch.empty_like(fused_expert_output) + group = _FakeCombineGroup(combine_input=None) + monkeypatch.setattr(pf, "get_ep_group", lambda: group) + + impl = pf.MoEPrepareAndFinalizeNaiveDPEPModular() + impl.finalize( + output, + fused_expert_output, + torch.ones((2, 1)), + torch.zeros((2, 1), dtype=torch.long), + False, + TopKWeightAndReduceNoOP(), + ) + + assert group.allocation_calls == 0 + assert group.combined_input is fused_expert_output + torch.testing.assert_close(output, fused_expert_output) + + +def test_args_finalize_contiguous_writes_combine_input(monkeypatch): + import vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep as pf + import vllm.model_executor.layers.fused_moe.topk_weight_and_reduce as wr + + fused_expert_output = torch.arange(12, dtype=torch.float32).view(2, 2, 3) + original = fused_expert_output.clone() + topk_weights = torch.tensor([[0.25, 0.75], [0.4, 0.6]]) + combine_input = torch.empty((2, 3), dtype=torch.float32) + output = torch.empty_like(combine_input) + group = _FakeCombineGroup(combine_input) + monkeypatch.setattr(pf, "get_ep_group", lambda: group) + monkeypatch.setattr( + wr.ops, + "moe_sum", + lambda hidden_states, out: out.copy_(hidden_states.sum(dim=1)), + ) + + impl = pf.MoEPrepareAndFinalizeNaiveDPEPModular() + impl.finalize( + output, + fused_expert_output, + topk_weights, + torch.zeros((2, 2), dtype=torch.long), + False, + wr.TopKWeightAndReduceContiguous(), + ) + + expected = (original * topk_weights.unsqueeze(-1)).sum(dim=1) + assert group.allocation_calls == 1 + assert group.combined_input is combine_input + torch.testing.assert_close(output, expected) + + def _args_dispatch_combine_worker(rank, world_size): from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager from vllm.forward_context import get_forward_context @@ -604,14 +721,39 @@ def _args_dispatch_combine_worker(rank, world_size): # -- combine (reduce-scatter) -- # Each token i has value i in all columns; after reduce-scatter # each rank gets its slice, summed across ranks. - expert_out = ( + expert_values = ( torch.arange(total_tokens, device=device, dtype=torch.float32) .unsqueeze(1) .expand(total_tokens, hidden_size) .contiguous() ) + expert_out = manager.allocate_combine_input( + expert_values.shape, + expert_values.dtype, + expert_values.device, + is_sequence_parallel=True, + ) + if expert_out is None: + expert_out = expert_values + else: + from vllm.distributed.device_communicators.pynccl_allocator import ( + is_symmetric_memory_tensor, + ) + + assert is_symmetric_memory_tensor(expert_out) + expert_out.copy_(expert_values) - combined = manager.combine(expert_out, is_sequence_parallel=True) + combine_output = torch.empty( + (tokens_per_rank, hidden_size), + device=device, + dtype=expert_out.dtype, + ) + combined = manager.combine_into_output( + expert_out, + combine_output, + is_sequence_parallel=True, + ) + assert combined.data_ptr() == combine_output.data_ptr() assert combined.shape == (tokens_per_rank, hidden_size) for i in range(tokens_per_rank): @@ -623,6 +765,61 @@ def _args_dispatch_combine_worker(rank, world_size): torch.distributed.barrier() + uneven_sizes = [tokens_per_rank + r for r in range(world_size)] + with _make_forward_context( + rank, + world_size, + uneven_sizes[rank], + num_tokens_across_dp=uneven_sizes, + ): + dp_metadata = get_forward_context().dp_metadata + assert dp_metadata is not None + with dp_metadata.sp_local_sizes(sequence_parallel_size=1): + assert ( + manager.allocate_combine_input( + (sum(uneven_sizes), hidden_size), + torch.float32, + device, + is_sequence_parallel=True, + ) + is None + ) + + total_uneven_tokens = sum(uneven_sizes) + uneven_input = ( + torch.arange( + total_uneven_tokens, + dtype=torch.float32, + device=device, + ) + .unsqueeze(1) + .expand(total_uneven_tokens, hidden_size) + .contiguous() + ) + uneven_input.add_(rank) + uneven_output = torch.empty( + (uneven_sizes[rank], hidden_size), + dtype=torch.float32, + device=device, + ) + combined = get_ep_group().reduce_scatterv_into_output( + uneven_input, + uneven_output, + dim=0, + sizes=uneven_sizes, + ) + assert combined.data_ptr() == uneven_output.data_ptr() + + start = sum(uneven_sizes[:rank]) + expected = torch.arange( + start, + start + uneven_sizes[rank], + dtype=torch.float32, + device=device, + ) * world_size + sum(range(world_size)) + expected = expected.unsqueeze(1).expand_as(combined) + torch.testing.assert_close(combined, expected) + @requires_multi_gpu @pytest.mark.parametrize("world_size", [2]) diff --git a/tests/distributed/test_nccl_symm_mem.py b/tests/distributed/test_nccl_symm_mem.py index bd0270fc4849..d3013445cf3b 100644 --- a/tests/distributed/test_nccl_symm_mem.py +++ b/tests/distributed/test_nccl_symm_mem.py @@ -3,6 +3,7 @@ import random import typing +from types import SimpleNamespace import pytest import torch @@ -12,11 +13,15 @@ import vllm.envs as envs from tests.utils import ensure_current_vllm_config from vllm.distributed import cleanup_dist_env_and_memory -from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator +from vllm.distributed.device_communicators.cuda_communicator import ( + NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION, + CudaCommunicator, +) from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops from vllm.distributed.device_communicators.pynccl_allocator import ( get_nccl_mem_pool, is_symmetric_memory_enabled, + is_symmetric_memory_tensor, ) from vllm.distributed.parallel_state import ( get_tp_group, @@ -32,6 +37,53 @@ test_size_elements = 4 * 1024 * 1024 +def test_disabled_pynccl_does_not_allocate_symmetric_buffer(): + communicator = object.__new__(CudaCommunicator) + communicator.pynccl_comm = SimpleNamespace(disabled=True) + + assert ( + communicator.get_symmetric_memory_buffer( + "test", (1,), torch.float32, torch.device("cuda") + ) + is None + ) + + +@pytest.mark.parametrize("use_aiter", [False, True]) +def test_reduce_scatterv_preserves_caller_output(monkeypatch, use_aiter): + import vllm.distributed.device_communicators.cuda_communicator as cc + + communicator = object.__new__(CudaCommunicator) + communicator.world_size = 2 + communicator.rank_in_group = 0 + input_ = torch.arange(12, dtype=torch.float32).view(4, 3) + output = torch.full((2, 3), float("nan")) + + def reduce_scatter(out, inp): + assert out is output + out.copy_(inp.view(2, 2, 3).sum(dim=0)) + + communicator.pynccl_comm = SimpleNamespace( + disabled=False, nccl_version=22902, reduce_scatter=reduce_scatter + ) + communicator._can_use_aiter_ag_rs = lambda sizes: use_aiter and sizes is None + communicator.aiter_ar_comm = SimpleNamespace( + should_custom_rs=lambda inp, dim: True, + custom_reduce_scatter=lambda inp, out, dim: reduce_scatter(out, inp), + ) + monkeypatch.setattr(cc, "should_nccl_symm_mem_ag_rs", lambda: True) + monkeypatch.setattr(cc, "is_symmetric_memory_tensor", lambda tensor: False) + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0") + + # Explicit uniform sizes must not stage an ordinary input into symmetric + # scratch, and AITER must write the supplied output after normalization. + result = communicator.reduce_scatterv_into_output( + input_, output, dim=0, sizes=[2, 2] + ) + assert result.data_ptr() == output.data_ptr() + torch.testing.assert_close(output, input_.view(2, 2, 3).sum(dim=0)) + + def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int): monkeypatch = pytest.MonkeyPatch() with monkeypatch.context() as m: @@ -193,15 +245,107 @@ def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int): pytest.skip("NCCL symmetric memory is disabled.") per_rank_size = test_size_elements // world_size - input_tensor = torch.randint( + ordinary_input = torch.randint( 1, 23, (test_size_elements,), dtype=dtype, device=device ) - input_clone = input_tensor.clone() - output = cuda_communicator.reduce_scatter(input_tensor, dim=0) + ordinary_clone = ordinary_input.clone() + ordinary_output = cuda_communicator.reduce_scatter(ordinary_input, dim=0) group = get_tp_group().device_group - expected = torch.empty(per_rank_size, dtype=dtype, device=device) + ordinary_expected = torch.empty(per_rank_size, dtype=dtype, device=device) + dist.reduce_scatter_tensor(ordinary_expected, ordinary_clone, group=group) + torch.testing.assert_close( + ordinary_output, ordinary_expected, atol=2.5, rtol=0.1 + ) + assert is_symmetric_memory_tensor(ordinary_output) + + # Both calls reuse rs_out; validate the first result before overwriting it. + for sizes in (None, [per_rank_size] * world_size): + ordinary_v_output = cuda_communicator.reduce_scatterv( + ordinary_input, dim=0, sizes=sizes + ) + torch.testing.assert_close( + ordinary_v_output, ordinary_expected, atol=2.5, rtol=0.1 + ) + assert is_symmetric_memory_tensor(ordinary_v_output) + + pynccl_comm = cuda_communicator.pynccl_comm + assert pynccl_comm is not None + if pynccl_comm.nccl_version < NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION: + return + + from vllm.v1.worker import ubatching + + m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 0) + ubatch0 = cuda_communicator._get_symm_scratch("ag_out", (128,), dtype, device) + ubatch0.fill_(1) + m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 1) + ubatch1 = cuda_communicator._get_symm_scratch("ag_out", (128,), dtype, device) + ubatch1.fill_(2) + assert ubatch0.data_ptr() != ubatch1.data_ptr() + assert (ubatch0 == 1).all() + + m.setattr(ubatching, "dbo_current_ubatch_id", lambda: 0) + smaller = cuda_communicator._get_symm_scratch("ag_out", (64,), dtype, device) + assert smaller.shape == (64,) + assert smaller.data_ptr() == ubatch0.data_ptr() + + grown = cuda_communicator._get_symm_scratch("ag_out", (129,), dtype, device) + assert grown.shape == (129,) + assert grown.data_ptr() != ubatch0.data_ptr() + reused = cuda_communicator._get_symm_scratch("ag_out", (200,), dtype, device) + assert reused.data_ptr() == grown.data_ptr() + + input_tensor = cuda_communicator._get_symm_scratch( + "test_graph_rs_input", (test_size_elements,), dtype, device + ) + input_tensor.random_(1, 23) + input_clone = input_tensor.clone() + output = torch.empty(per_rank_size, dtype=dtype, device=device) + assert is_symmetric_memory_tensor(input_tensor) + assert not is_symmetric_memory_tensor(output) + + result = cuda_communicator.reduce_scatterv_into_output( + input_tensor, + output, + dim=0, + sizes=[per_rank_size] * world_size, + ) + assert result.data_ptr() == output.data_ptr() + + expected = torch.empty_like(output) dist.reduce_scatter_tensor(expected, input_clone, group=group) + torch.testing.assert_close(result, expected, atol=2.5, rtol=0.1) + + dist.barrier(group=get_tp_group().cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_result = cuda_communicator.reduce_scatterv_into_output( + input_tensor, + output, + dim=0, + sizes=[per_rank_size] * world_size, + ) + assert graph_result.data_ptr() == output.data_ptr() + + input_tensor.random_(24, 47) + input_clone.copy_(input_tensor) + dist.reduce_scatter_tensor(expected, input_clone, group=group) + captured_input_ptr = input_tensor.data_ptr() + del input_tensor + grown_input = cuda_communicator._get_symm_scratch( + "test_graph_rs_input", (test_size_elements + 1,), dtype, device + ) + assert grown_input.data_ptr() != captured_input_ptr + retired = cuda_communicator.__dict__["_retired_symm_scratch_bufs"] + assert any( + buf.data_ptr() == captured_input_ptr + for buffers in retired.values() + for buf in buffers + ) + output.zero_() + graph.replay() + torch.accelerator.synchronize() torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index d3c48756ca6e..d8790dbf42cb 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -171,6 +171,42 @@ def combine( hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes) return hidden_states + def allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + is_sequence_parallel: bool = False, + ) -> torch.Tensor | None: + dist_group = self._get_comm_group(is_sequence_parallel) + sizes = self._get_sizes(shape[0] // dist_group.world_size, dist_group) + if sum(sizes) != shape[0] or any(size != sizes[0] for size in sizes): + return None + device_communicator = dist_group.device_communicator + if device_communicator is None: + return None + return device_communicator.get_symmetric_memory_buffer( + "moe_ag_rs_combine", shape, dtype, device + ) + + def combine_into_output( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + is_sequence_parallel: bool = False, + ) -> torch.Tensor: + dist_group = self._get_comm_group(is_sequence_parallel) + sizes = self._get_sizes( + hidden_states.shape[0] // dist_group.world_size, + dist_group, + ) + return dist_group.reduce_scatterv_into_output( + hidden_states, + output, + dim=0, + sizes=sizes, + ) + def destroy(self): pass diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index f70bf0921518..588163b84602 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -157,6 +157,24 @@ def checkpoint_restore(self) -> None: def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False): raise NotImplementedError + def allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + is_sequence_parallel: bool = False, + ) -> torch.Tensor | None: + return None + + def combine_into_output( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + is_sequence_parallel: bool = False, + ) -> torch.Tensor: + output.copy_(self.combine(hidden_states, is_sequence_parallel)) + return output + def destroy(self): pass @@ -315,6 +333,25 @@ def reduce_scatterv( ) -> torch.Tensor: raise NotImplementedError + def reduce_scatterv_into_output( + self, + input_: torch.Tensor, + output: torch.Tensor, + dim: int = -1, + sizes: list[int] | None = None, + ) -> torch.Tensor: + output.copy_(self.reduce_scatterv(input_, dim, sizes)) + return output + + def get_symmetric_memory_buffer( + self, + role: str, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor | None: + return None + def gather( self, input_: torch.Tensor, dst: int = 0, dim: int = -1 ) -> torch.Tensor | None: @@ -417,5 +454,23 @@ def combine( """ return hidden_states + def allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + is_sequence_parallel: bool = False, + ) -> torch.Tensor | None: + return None + + def combine_into_output( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + is_sequence_parallel: bool = False, + ) -> torch.Tensor: + output.copy_(self.combine(hidden_states, is_sequence_parallel)) + return output + def batch_isend_irecv(self, p2p_ops: list): raise NotImplementedError diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 99ebcd717ddf..d4d588bf450d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -16,6 +16,7 @@ from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops from vllm.distributed.device_communicators.pynccl_allocator import ( is_symmetric_memory_enabled, + is_symmetric_memory_tensor, ) from vllm.logger import init_logger from vllm.platforms import current_platform @@ -26,6 +27,10 @@ logger = init_logger(__name__) +# Since NCCL 2.29.2, ncclSymkPickKernel accepts registered input with +# non-registered output for ReduceScatter (ncclSymSendRegRecvNonreg). +NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION = 22902 + class CudaCommunicator(DeviceCommunicatorBase): def __init__( @@ -505,13 +510,33 @@ def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): def reduce_scatterv( self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None ): + return self._reduce_scatterv(input_, dim, sizes) + + def reduce_scatterv_into_output( + self, + input_: torch.Tensor, + output: torch.Tensor, + dim: int = -1, + sizes: list[int] | None = None, + ) -> torch.Tensor: + return self._reduce_scatterv(input_, dim, sizes, output) + + def _reduce_scatterv( + self, + input_: torch.Tensor, + dim: int, + sizes: list[int] | None, + output: torch.Tensor | None = None, + ) -> torch.Tensor: world_size = self.world_size pynccl_comm = self.pynccl_comm assert pynccl_comm is not None + assert not pynccl_comm.disabled, "reduce_scatterv requires PyNccl" if dim < 0: # Convert negative dim to positive. dim += input_.dim() + sizes_are_explicit = sizes is not None # 'sizes' is not needed if all inputs in the same group have the same # shape if sizes is not None and all(s == sizes[0] for s in sizes): @@ -529,35 +554,63 @@ def reduce_scatterv( assert input_tensor.shape[0] % world_size == 0 chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] + if output is not None: + assert dim == 0, "preallocated reduce-scatter output requires dim=0" + assert output.shape == output_shape + assert output.dtype == input_tensor.dtype + assert output.device == input_tensor.device + assert output.is_contiguous() if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None if aiter_comm.should_custom_rs(input_tensor, dim=0): - output = torch.empty( - output_shape, dtype=input_tensor.dtype, device=input_tensor.device - ) + if output is None: + output = torch.empty( + output_shape, + dtype=input_tensor.dtype, + device=input_tensor.device, + ) aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0) return output.movedim(0, dim).contiguous() # Symmetric memory is only used when all ranks have uniform sizes. # ncclCommWindowRegister is collective: asymmetric pool allocations # from variable per-rank sizes cause deadlocks. - use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs() + uniform_sizes = sizes is None or all(size == sizes[0] for size in sizes) + direct_output_supported = ( + output is None + or is_symmetric_memory_tensor(output) + or pynccl_comm.nccl_version >= NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION + ) + use_symm_mem = ( + uniform_sizes + and direct_output_supported + and ( + output is None + or not sizes_are_explicit + or is_symmetric_memory_tensor(input_tensor) + ) + and should_nccl_symm_mem_ag_rs() + ) if use_symm_mem: - output = self._reduce_scatter_symm_mem(input_tensor) + output = self._reduce_scatter_symm_mem(input_tensor, output) else: - output = torch.empty( - output_shape, dtype=input_tensor.dtype, device=input_tensor.device - ) + if output is None: + output = torch.empty( + output_shape, dtype=input_tensor.dtype, device=input_tensor.device + ) use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2 if use_deterministic_rs: # Reduce to a fixed root (0) for determinism reduced = torch.empty_like(input_tensor) - sizes = sizes if sizes else [chunk_size] * world_size + scatter_sizes = sizes + if scatter_sizes is None: + scatter_sizes = [chunk_size] * world_size pynccl_comm.reduce(reduced, input_tensor, root=0) - pynccl_comm.scatter(output, reduced, sizes, root=0) - elif sizes is not None and sizes.count(sizes[0]) != len(sizes): + pynccl_comm.scatter(output, reduced, scatter_sizes, root=0) + elif not uniform_sizes: + assert sizes is not None pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes) else: pynccl_comm.reduce_scatter(output, input_tensor) @@ -576,54 +629,64 @@ def _get_symm_scratch( Allocating a fresh symm tensor per collective pays the ``nccl_symm_mem_context`` snapshot + window-registration scan on every - call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we - allocate once per ``(role, shape, dtype)``, register once, and reuse. - - Safe for serial (eager) sequence parallelism: each collective's result - is consumed on the same stream before the next same-role collective - reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both - full-size) get distinct buffers so a reduce-scatter input copy never - clobbers a still-live all-gather output. + call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead, + each ``(role, shape[1:], dtype, device)`` keeps a geometric high-water + allocation and returns a leading slice. Superseded allocations remain + referenced because a captured CUDA graph may still use their pointers; + geometric growth bounds their total capacity to less than twice the + current allocation. + + Safe across serial MoE layers and eager sequence parallelism: the + producer and collective are ordered on the same stream before the next + same-role operation reuses the buffer. DBO microbatches use distinct + cache entries. Any future cross-layer communication overlap must also + use distinct roles. """ from vllm.distributed.device_communicators.pynccl_allocator import ( nccl_symm_mem_context, ) + from vllm.v1.worker.ubatching import dbo_current_ubatch_id pynccl_comm = self.pynccl_comm assert pynccl_comm is not None + assert shape, "symmetric scratch buffers require at least one dimension" cache = self.__dict__.setdefault("_symm_scratch_bufs", {}) - key = (role, tuple(shape), dtype) + key = (role, dbo_current_ubatch_id(), tuple(shape[1:]), dtype, device) buf = cache.get(key) - if buf is None: + requested_rows = shape[0] + if buf is None or buf.shape[0] < requested_rows: + capacity = ( + requested_rows if buf is None else max(requested_rows, 2 * buf.shape[0]) + ) with nccl_symm_mem_context(pynccl_comm): - buf = torch.empty(shape, dtype=dtype, device=device) + new_buf = torch.empty( + (capacity, *shape[1:]), dtype=dtype, device=device + ) + if buf is not None: + retired = self.__dict__.setdefault("_retired_symm_scratch_bufs", {}) + retired.setdefault(key, []).append(buf) + buf = new_buf cache[key] = buf - return buf + return buf[:requested_rows] def _reduce_scatter_symm_mem( self, input_tensor: torch.Tensor, + output: torch.Tensor | None = None, ) -> torch.Tensor: """ReduceScatter using NCCL symmetric memory (NVLS). - Only called for uniform-size reduce_scatter (variable sizes are - guarded out by the caller to avoid asymmetric ncclCommWindowRegister). - Uses persistent pre-registered scratch (see _get_symm_scratch). + The MoE reduce_scatterv path passes explicit uniform sizes and calls + this only with an already-registered input, so that path never stages. + Uniform calls without a caller-owned output retain their existing + opt-in behavior and stage ordinary input into registered scratch. + The output may be ordinary memory. """ - from vllm.distributed.device_communicators.pynccl_allocator import ( - is_symmetric_memory_tensor, - ) - pynccl_comm = self.pynccl_comm assert pynccl_comm is not None - chunk = input_tensor.shape[0] // self.world_size output_shape = (chunk,) + tuple(input_tensor.shape[1:]) - symm_output = self._get_symm_scratch( - "rs_out", output_shape, input_tensor.dtype, input_tensor.device - ) - # NVLS reduce-scatter (LDMC) requires the input in symmetric memory. if is_symmetric_memory_tensor(input_tensor): symm_input = input_tensor else: @@ -635,8 +698,31 @@ def _reduce_scatter_symm_mem( ) symm_input.copy_(input_tensor) - pynccl_comm.reduce_scatter(symm_output, symm_input) - return symm_output + if output is None: + output = self._get_symm_scratch( + "rs_out", output_shape, input_tensor.dtype, input_tensor.device + ) + pynccl_comm.reduce_scatter(output, symm_input) + return output + + def get_symmetric_memory_buffer( + self, + role: str, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor | None: + pynccl_comm = self.pynccl_comm + if ( + pynccl_comm is None + or pynccl_comm.disabled + or pynccl_comm.world_size == 1 + or pynccl_comm.nccl_version < NCCL_DIRECT_SYMM_RS_OUTPUT_MIN_VERSION + or not should_nccl_symm_mem_ag_rs() + ): + return None + output = self._get_symm_scratch(role, shape, dtype, device) + return output if is_symmetric_memory_tensor(output) else None def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: """Sends a tensor to the destination rank in a blocking way.""" @@ -925,6 +1011,29 @@ def combine( is_sequence_parallel, ) + def allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + is_sequence_parallel: bool = False, + ) -> torch.Tensor | None: + assert self.all2all_manager is not None + return self.all2all_manager.allocate_combine_input( + shape, dtype, device, is_sequence_parallel + ) + + def combine_into_output( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + is_sequence_parallel: bool = False, + ) -> torch.Tensor: + assert self.all2all_manager is not None + return self.all2all_manager.combine_into_output( + hidden_states, output, is_sequence_parallel + ) + def batch_isend_irecv(self, p2p_ops: list): pynccl_comm = self.pynccl_comm if pynccl_comm is not None and not pynccl_comm.disabled: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 88bb3d7c5a61..de6680e02745 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -801,6 +801,19 @@ def reduce_scatterv( raise ValueError("No device communicator found") return self.device_communicator.reduce_scatterv(input_, dim, sizes) + def reduce_scatterv_into_output( + self, + input_: torch.Tensor, + output: torch.Tensor, + dim: int = -1, + sizes: list[int] | None = None, + ) -> torch.Tensor: + if self.device_communicator is None: + raise ValueError("No device communicator found") + return self.device_communicator.reduce_scatterv_into_output( + input_, output, dim, sizes + ) + def _reduce_scatter_out_place(self, input_: torch.Tensor, dim: int) -> torch.Tensor: if self.device_communicator is None: raise ValueError("No device communicator found") @@ -1442,6 +1455,32 @@ def combine( else: return hidden_states + def allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + is_sequence_parallel: bool = False, + ) -> torch.Tensor | None: + if self.device_communicator is None: + return None + return self.device_communicator.allocate_combine_input( + shape, dtype, device, is_sequence_parallel + ) + + def combine_into_output( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + is_sequence_parallel: bool = False, + ) -> torch.Tensor: + if self.device_communicator is None: + output.copy_(hidden_states) + return output + return self.device_communicator.combine_into_output( + hidden_states, output, is_sequence_parallel + ) + _WORLD: GroupCoordinator | None = None _INNER_DP_WORLD: GroupCoordinator | None = None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 9bde03795c9a..7a4cbd423746 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -197,6 +197,20 @@ def post_init_setup(self, fused_experts: "FusedMoEExperts"): """ return + def allocate_fused_expert_output( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + weight_and_reduce_impl: TopKWeightAndReduce, + ) -> torch.Tensor | None: + """Optionally provide the exact output buffer for the expert kernel. + + Implementations returning a tensor must preserve its identity through + finalize; replacing it may silently reintroduce a write-back copy. + """ + return None + @property @abstractmethod def activation_format(self) -> FusedMoEActivationFormat: @@ -1054,13 +1068,14 @@ def _allocate_buffers( local_num_experts: int, expert_tokens_meta: ExpertTokensMetadata | None, activation: MoEActivation, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]: """Allocate temporary and output buffers for the fused experts op. Inputs: - out_dtype: output type of workspace and output tensors. - device: the device of the workspace and output tensors. See `workspace_shapes` for a description of the remainder of arguments. - Returns a tuple of (workspace13, workspace2, output) tensors. + Returns the workspace tensors, output tensor, and whether the output + uses a prepare/finalize-provided allocation. """ assert M_full > 0 and M_chunk > 0 @@ -1090,9 +1105,20 @@ def _allocate_buffers( activation, ) + fused_out_alias = self.prepare_finalize.allocate_fused_expert_output( + fused_out_shape, + workspace_dtype, + device, + self.fused_experts.finalize_weight_and_reduce_impl(), + ) + # We can reuse the memory between cache1 and cache3 because by the # time we need cache3, we're done with cache1. # Reuse workspace13 for the output since there is only one chunk. + # Keep the original fused-output reservation even when the output is + # externally aliased. The workspace manager profiles and locks a + # process-wide high-water mark; later graph shapes may need this + # storage for workspace13. max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) if current_platform.is_cpu(): @@ -1110,16 +1136,23 @@ def _allocate_buffers( ) workspace13 = _resize_cache(common_workspace, workspace13_shape) fused_out = _resize_cache(common_workspace, fused_out_shape) - return workspace13, workspace2, fused_out + return workspace13, workspace2, fused_out, False common_workspace, workspace2 = current_workspace_manager().get_simultaneous( ((max_shape_size,), workspace_dtype), (workspace2_shape, workspace_dtype), ) workspace13 = _resize_cache(common_workspace, workspace13_shape) - fused_out = _resize_cache(common_workspace, fused_out_shape) + if fused_out_alias is None: + fused_out = _resize_cache(common_workspace, fused_out_shape) + else: + assert fused_out_alias.shape == fused_out_shape + assert fused_out_alias.dtype == workspace_dtype + assert fused_out_alias.device == device + assert fused_out_alias.is_contiguous() + fused_out = fused_out_alias - return workspace13, workspace2, fused_out + return workspace13, workspace2, fused_out, fused_out_alias is not None def _maybe_apply_shared_experts( self, @@ -1251,22 +1284,25 @@ def _fused_experts( if M_full == 0: return torch.empty_like(a1q, dtype=in_dtype) - workspace13, workspace2, fused_out = self._allocate_buffers( - in_dtype, - a1q.device, - M_full, - M_full, - N, - K, - top_k, - global_num_experts, - local_num_experts, - expert_tokens_meta, - activation, + workspace13, workspace2, fused_out, fused_out_is_specialized = ( + self._allocate_buffers( + in_dtype, + a1q.device, + M_full, + M_full, + N, + K, + top_k, + global_num_experts, + local_num_experts, + expert_tokens_meta, + activation, + ) ) use_output_alias = ( - output_alias is not None + not fused_out_is_specialized + and output_alias is not None and output_alias.shape == fused_out.shape and output_alias.dtype == fused_out.dtype and output_alias.device == fused_out.device diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index 97d4e0d284c6..7087b9016d8f 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -8,6 +8,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate, + TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.utils.flashinfer import nvfp4_block_scale_interleave @@ -108,6 +109,30 @@ def num_dispatchers(self) -> int: def output_is_reduced(self) -> bool: return False + def _allocate_combine_input( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor | None: + return get_ep_group().allocate_combine_input( + shape, + dtype, + device, + is_sequence_parallel=self.is_sequence_parallel, + ) + + def allocate_fused_expert_output( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + weight_and_reduce_impl: mk.TopKWeightAndReduce, + ) -> torch.Tensor | None: + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP): + return self._allocate_combine_input(shape, dtype, device) + return None + def prepare( self, a1: torch.Tensor, @@ -194,16 +219,29 @@ def finalize( if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate): weight_and_reduce_impl = TopKWeightAndReduceContiguous() + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP): + # allocate_fused_expert_output already selected the combine input. + # Reuse it directly; the symmetric allocator may return a new view + # object for the same high-water backing allocation. + combine_input = fused_expert_output + else: + combine_input = self._allocate_combine_input( + (topk_ids.shape[0], fused_expert_output.shape[-1]), + fused_expert_output.dtype, + fused_expert_output.device, + ) out = weight_and_reduce_impl.apply( - output=None, + output=combine_input, fused_expert_output=fused_expert_output, topk_weights=topk_weights, topk_ids=topk_ids, apply_router_weight_on_input=apply_router_weight_on_input, ) - output.copy_( - get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel) + get_ep_group().combine_into_output( + out, + output, + is_sequence_parallel=self.is_sequence_parallel, ) From 8bd7737a03bfcd4a1bca008f9f0e94729f8d67dc Mon Sep 17 00:00:00 2001 From: Sunny Rangnani <1923481+SunnyR@users.noreply.github.com> Date: Tue, 6 Oct 2026 08:39:59 -0600 Subject: [PATCH 072/952] [Docs] Clarify that vLLM does not isolate tenants sharing a server (#59811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: SunnyR Co-authored-by: Claude Co-authored-by: Juan Pérez de Algaba <124347725+jperezdealgaba@users.noreply.github.com> --- docs/usage/security.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index d350b6bdfd9d..3dad19d6cb03 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -550,6 +550,15 @@ ensure that only trusted principals can submit work to the cluster: - Place the Ray cluster on an isolated network segment. - Do not expose the Ray client port or dashboard to untrusted networks. +## Multi-Tenant Deployments + +vLLM does not provide isolation between tenants that share the same server process. Requests from different callers are scheduled together and share the same caches. Options such as `cache_salt` reduce specific cross-tenant risks, but they are not an isolation boundary. + +If tenants must be isolated from each other, enforce it in the deployment architecture: + +- Run a dedicated vLLM instance per tenant. +- Place a gateway in front of vLLM that authenticates callers and scopes client-supplied cache identifiers, such as `cache_salt` and multimodal `uuid` values, per tenant. + ## Prefix Cache Timing Side-Channel Mitigation (Cache Salting) ### Background @@ -633,7 +642,7 @@ For additional cross-tenant isolation, set `cache_salt` on each request (see [Pr - **Multi-tenant deployments**: Always generate cryptographically random UUIDs per media item. Additionally, set `cache_salt` to a per-tenant secret for defense in depth. - **Single-tenant deployments**: Ensure UUIDs are unique per distinct media content. `cache_salt` is unnecessary when there is no cross-tenant threat. -- **Default behavior**: Omitting `uuid` entirely preserves the default content-hash-based identity, which is safe against this class of collision but requires hashing the media bytes on every request. +- **Default behavior**: Omitting `uuid` uses a hash of the media bytes as the cache identity, which requires hashing the media on every request. This avoids accidental collisions between callers, but it does not isolate tenants from each other. See [Multi-Tenant Deployments](#multi-tenant-deployments). ## Reporting Security Vulnerabilities From 0719975ceeeab568cbd3a3e88a5c790a91b56709 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 6 Oct 2026 22:53:17 +0800 Subject: [PATCH 073/952] [CI][Rust Frontend] Pass `--locked` to cargo binstall (#60182) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .buildkite/scripts/run-rust-frontend-cargo-ci.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh index fa1cddb95f57..aa16302da8c4 100755 --- a/.buildkite/scripts/run-rust-frontend-cargo-ci.sh +++ b/.buildkite/scripts/run-rust-frontend-cargo-ci.sh @@ -60,12 +60,12 @@ install_cargo_binstall() { install_cargo_sort() { log_section "Installing cargo-sort ${CARGO_SORT_VERSION_REQ}" - cargo binstall --no-confirm --force "cargo-sort@${CARGO_SORT_VERSION_REQ}" + cargo binstall --no-confirm --force --locked "cargo-sort@${CARGO_SORT_VERSION_REQ}" } install_cargo_deny() { log_section "Installing cargo-deny ${CARGO_DENY_VERSION_REQ}" - cargo binstall --no-confirm --force "cargo-deny@${CARGO_DENY_VERSION_REQ}" + cargo binstall --no-confirm --force --locked "cargo-deny@${CARGO_DENY_VERSION_REQ}" } install_cargo_nextest() { @@ -74,6 +74,7 @@ install_cargo_nextest() { --no-confirm \ --force \ --secure \ + --locked \ "cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}" } From cbe1f9740d0ae7f6eb82bd752a984c7db8ab85fe Mon Sep 17 00:00:00 2001 From: ovidiusm Date: Tue, 6 Oct 2026 16:55:03 +0200 Subject: [PATCH 074/952] [NIXL] Bump NIXL version to 1.5.0 (#56907) Signed-off-by: Ovidiu Mara --- requirements/kv_connectors.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index 9d9d1278818e..f6cacd200b78 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,7 +2,7 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. cupy-cuda13x != 14.1.0 -nixl == 1.4.1 +nixl == 1.5.0 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. mooncake-transfer-engine >= 0.3.12 From cc73cca9c88bc19a368a910a630e9c5caec61ac3 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 6 Oct 2026 11:39:04 -0400 Subject: [PATCH 075/952] [Build] Migrate vendored DeepGEMM from pybind to TORCH_LIBRARY (abi3) (#48962) Signed-off-by: Chris Leonard Co-authored-by: Shengqi Chen --- cmake/external_projects/deepgemm.cmake | 94 ++++++---------- tools/build_deepgemm_C.py | 100 ++++++------------ tools/check_wheel_deepgemm.py | 64 +++++------ tools/install_deepgemm.sh | 4 +- .../model_executor/warmup/deep_gemm_warmup.py | 2 + 5 files changed, 107 insertions(+), 157 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index a17031cc5e0f..1dcee345f845 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -33,7 +33,8 @@ else() # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x # layout header fix from vllm-project/DeepGEMM#12. - set(_DEEPGEMM_UPSTREAM_TAG "e1f418c2a4f20818221f6b0e578b4c2f634d4c3f") + # TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "1e1842a833699298f7afc02eefb2ed168fff6938") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -43,7 +44,7 @@ else() set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build") set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild") - if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp") + if(EXISTS "${_deepgemm_src}/deep_gemm/_C.py") set(deepgemm_SOURCE_DIR "${_deepgemm_src}") set(deepgemm_BINARY_DIR "${_deepgemm_bin}") else() @@ -93,33 +94,17 @@ if(DEEPGEMM_ARCHS) # # DeepGEMM integration notes # -------------------------- - # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a - # `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The - # per-Python build is delegated to tools/build_deepgemm_C.py. + # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle: + # - deep_gemm/_C.py (Python shim over torch.ops.deep_gemm) + # - deep_gemm/_C_extension.abi3.so (single limited-API extension) + # The build is delegated to tools/build_deepgemm_C.py (setup.py build_ext). # - # Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links - # private CPython symbols — a single `_C.abi3.so` is not viable today - # (see #41476 / #41512 for the failed attempt). - # - # TODOs (tracked in vllm-project/vllm#42431): - # - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim - # binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to - # collapse to one `_C.abi3.so`. Needs either an upstream change or - # a maintained binding fork in vLLM. - # - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop - # the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime - # requirement. + # TODO: AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop + # the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime + # requirement. # - # DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths. - # Empty/unset → fall back to the build interpreter (editable installs). - # (Empty-but-set env vars test as DEFINED in cmake — treat as unset.) - if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "") - string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}") - else() - set(_dg_pythons "${Python_EXECUTABLE}") - endif() - message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") + message(STATUS "DeepGEMM extension will be built with: ${Python_EXECUTABLE}") # add_custom_command does no implicit header scanning; glob explicitly so # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild. @@ -132,40 +117,29 @@ if(DEEPGEMM_ARCHS) "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h" "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp") - set(_dg_markers) - set(_dg_seen_soabis) - foreach(_pybin IN LISTS _dg_pythons) - execute_process( - COMMAND "${_pybin}" -c - "import sysconfig; print(sysconfig.get_config_var('SOABI'))" - OUTPUT_VARIABLE _dg_soabi - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY) - # Dedup interpreters that resolve to the same CPython. - if(_dg_soabi IN_LIST _dg_seen_soabis) - continue() - endif() - list(APPEND _dg_seen_soabis "${_dg_soabi}") - set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}") - set(_dg_marker "${_dg_dir}/.built") - add_custom_command( - OUTPUT "${_dg_marker}" - COMMAND "${Python_EXECUTABLE}" - "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" - "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}" - COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}" - DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" - "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp" - ${_dg_headers} - COMMENT "Building DeepGEMM _C for ${_pybin}" - VERBATIM) - list(APPEND _dg_markers "${_dg_marker}") - install(DIRECTORY "${_dg_dir}/" - DESTINATION vllm/third_party/deep_gemm - COMPONENT _deep_gemm_C - FILES_MATCHING PATTERN "_C.cpython-*.so") - endforeach() - add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers}) + set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C") + set(_dg_marker "${_dg_dir}/.built") + add_custom_command( + OUTPUT "${_dg_marker}" + COMMAND "${Python_EXECUTABLE}" + "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" + "${deepgemm_SOURCE_DIR}" "${_dg_dir}" + COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}" + DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp" + "${deepgemm_SOURCE_DIR}/deep_gemm/_C.py" + "${deepgemm_SOURCE_DIR}/setup.py" + ${_dg_headers} + COMMENT "Building DeepGEMM _C_extension (abi3)" + VERBATIM) + add_custom_target(_deep_gemm_C ALL DEPENDS "${_dg_marker}") + + install(DIRECTORY "${_dg_dir}/" + DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C + FILES_MATCHING + PATTERN "_C.py" + PATTERN "_C_extension*.so") # # Vendor DeepGEMM Python package files diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index f4438c67552f..36824318c051 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -1,85 +1,53 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Build DeepGEMM's `_C` pybind11 extension for . +"""Build DeepGEMM's TORCH_LIBRARY extension and copy vendored artifacts. -Driven from cmake/external_projects/deepgemm.cmake. The driver runs against -the build interpreter's torch; is only consulted for INCLUDEPY -and SOABI, so target venvs don't need torch installed. +DeepGEMM now registers ops via TORCH_LIBRARY into ``deep_gemm._C_extension`` +(abi3) and exposes the legacy API through ``deep_gemm/_C.py``. This driver +delegates to DeepGEMM's ``setup.py build_ext --inplace`` and copies the shim +plus extension into the cmake output directory. -Usage: python build_deepgemm_C.py +Usage: python build_deepgemm_C.py """ -import json import os +import shutil import subprocess import sys from pathlib import Path -import torch -from torch.utils import cpp_extension - -if len(sys.argv) != 4: - sys.exit(f"usage: {sys.argv[0]} ") +if len(sys.argv) != 3: + sys.exit(f"usage: {sys.argv[0]} ") src = Path(sys.argv[1]).resolve() out = Path(sys.argv[2]).resolve() -target_py = sys.argv[3] +_pkg = src / "deep_gemm" out.mkdir(parents=True, exist_ok=True) -info = json.loads( - subprocess.check_output( - [ - target_py, - "-c", - "import sysconfig, json; " - "print(json.dumps({k: sysconfig.get_config_var(k) " - "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))", - ] - ).decode() +if not (_pkg / "_C.py").is_file(): + sys.exit( + f"DeepGEMM source at {src} is missing deep_gemm/_C.py; " + "expected TORCH_LIBRARY migration layout" + ) + +# Avoid DeepGEMM's clean-git assertion when vendoring a local dirty tree. +env = os.environ.copy() +env.pop("DG_SKIP_CUDA_BUILD", None) + +print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True) +subprocess.check_call( + [sys.executable, "setup.py", "build_ext", "--inplace"], + cwd=src, + env=env, ) -cuda_home = cpp_extension.CUDA_HOME -if cuda_home is None: - sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C") -# CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py). -includes = [ - info["INCLUDEPY"], - f"{cuda_home}/include", - f"{cuda_home}/include/cccl", - str(src / "csrc"), - str(src / "deep_gemm/include"), - str(src / "third-party/cutlass/include"), - str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/deep_jit/include"), - *cpp_extension.include_paths(device_type="cuda"), -] +shim = _pkg / "_C.py" +shutil.copy2(shim, out / shim.name) -cmd = [ - os.environ.get("CXX", "g++"), - "-shared", - "-fPIC", - "-std=c++20", - "-O3", - "-g0", - "-Wno-psabi", - "-Wno-deprecated-declarations", - "-DTORCH_API_INCLUDE_EXTENSION_H", - "-DTORCH_EXTENSION_NAME=_C", - f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}", - *(f"-I{p}" for p in includes), - str(src / "csrc/python_api.cpp"), - *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")), - f"-L{cuda_home}/lib64", - "-ltorch", - "-ltorch_python", - "-ltorch_cpu", - "-ltorch_cuda", - "-lc10", - "-lc10_cuda", - "-lcudart", - "-lnvrtc", - "-o", - str(out / f"_C{info['EXT_SUFFIX']}"), -] -print("[build_deepgemm_C] " + " ".join(cmd), flush=True) -subprocess.check_call(cmd) +so = _pkg / "_C_extension.abi3.so" +if not so.is_file(): + sys.exit(f"DeepGEMM build did not produce {so}") +for stale in out.glob("_C_extension*.so"): + stale.unlink() +shutil.copy2(so, out / so.name) +print(f"[build_deepgemm_C] installed {so.name} -> {out}", flush=True) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index 6f8a03ffd3dc..fcc35ead58c4 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -1,41 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython -covered by `requires-python`. Fails closed if a Python's `.so` is missing -from the wheel — i.e. the regression that surfaced in #41476/#41512. +"""Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout. -Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm -kernel tests in .buildkite/test_areas/kernels.yaml. +Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under +``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100 +deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. """ import importlib.util -import os import sys from pathlib import Path -import regex as re -import tomllib -SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-") - - -def required_pythons() -> list[str]: - pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml" - spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"] - m = re.match(r">=3\.(\d+),<3\.(\d+)", spec) - if not m: - sys.exit(f"unexpected requires-python format: {spec!r}") - return [f"3.{v}" for v in range(int(m[1]), int(m[2]))] - - -spec = importlib.util.find_spec("vllm.third_party.deep_gemm") -if spec is None or spec.origin is None: - sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?") -pkg_dir = Path(spec.origin).parent - -found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))} -required = required_pythons() -missing = [v for v in required if v not in found] -print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}") -sys.exit(1 if missing else 0) +def main() -> int: + spec = importlib.util.find_spec("vllm.third_party.deep_gemm") + if spec is None or spec.origin is None: + print( + "vllm.third_party.deep_gemm not importable; is vllm installed?", + file=sys.stderr, + ) + return 1 + pkg_dir = Path(spec.origin).parent + + shim = pkg_dir / "_C.py" + so_files = sorted(pkg_dir.glob("_C_extension*.so")) + missing = [] + if not shim.is_file(): + missing.append("_C.py") + if len(so_files) != 1: + missing.append("exactly one _C_extension*.so") + + print( + f"deepgemm vendored binding: shim={shim.is_file()}, " + f"extensions={[p.name for p in so_files]}" + ) + if missing: + print(f"missing: {missing}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 949aa28da533..2fb2351a0ca5 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -10,7 +10,7 @@ DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch, which carries the sm120 # and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE # activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. -DEEPGEMM_GIT_REF="e1f418c2a4f20818221f6b0e578b4c2f634d4c3f" +DEEPGEMM_GIT_REF="1e1842a833699298f7afc02eefb2ed168fff6938" WHEEL_DIR="" # Parse command line arguments @@ -93,6 +93,8 @@ pushd "$INSTALL_DIR/deepgemm" # Checkout the specific reference git checkout "$DEEPGEMM_GIT_REF" +git submodule sync --recursive +git submodule update --init --recursive # Clean previous build artifacts # (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 1266f4ccb265..a4775ce22e45 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -179,6 +179,7 @@ def _deepgemm_fp8_gemm_nt_warmup( device = w.device a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn) + # Must be initialized (UE8M0 packing asserts zero sign/mantissa bits). a1q_scales = torch.zeros( (max_tokens, k // block_m), device=device, dtype=torch.float32 ) @@ -284,6 +285,7 @@ def _warmup(w: torch.Tensor, w_scale: torch.Tensor): ((MAX_M, k // block_m), torch.float32), ((MAX_M, n), torch.bfloat16), ) + # Must be initialized (UE8M0 packing asserts zero sign/mantissa bits). a1q_scales.zero_() for num_tokens, align_used, expert_ids in warmup_cases: From aad4579bbe56ce87a7799d31df57a19756d6f549 Mon Sep 17 00:00:00 2001 From: Jeff-Tseng-dev Date: Tue, 6 Oct 2026 15:44:56 +0000 Subject: [PATCH 076/952] Upgrade tpu-inference to v0.31.0 (#60190) --- requirements/tpu.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/tpu.txt b/requirements/tpu.txt index 01dc1f37ac98..a16281664c39 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -11,4 +11,4 @@ ray[default] ray[data] setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.30.0 +tpu-inference==0.31.0 From 4ea0c28bc4ba28968b8386728ed64235fbe1ac79 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 00:55:27 +0900 Subject: [PATCH 077/952] [Bugfix] Warn only on per-request do_normalize/do_rescale overrides (#60240) Signed-off-by: lesj0610 --- vllm/config/multimodal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 24c534b36429..3bba0434c3e4 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -725,7 +725,7 @@ def merge_mm_processor_kwargs( # The model expects raw pixels and normalises on device; # overriding these per request would silently double-normalise. for key in ("do_normalize", "do_rescale"): - if key in inference_kwargs: + if inference_kwargs.get(key): logger.warning_once( "mm_device_do_normalize is enabled; ignoring " "per-request mm_processor_kwargs[%r].", From fcf53d240c73d44bd9e3fe62996f8a2f1169fee9 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 6 Oct 2026 09:17:30 -0700 Subject: [PATCH 078/952] [Spec Decode] Remove eager metadata rebuild during MTP fused multi-step decode (#58463) Signed-off-by: Giancarlo Delfin --- .../test_indexer_deepseek_v4_slot_mapping.py | 61 +++++++++++++++++++ vllm/v1/attention/backend.py | 29 ++++----- .../backends/mla/flashattn_mla_sparse.py | 15 ++++- vllm/v1/attention/backends/mla/indexer.py | 22 +++++++ vllm/v1/attention/backends/mla/sparse_swa.py | 15 ++++- vllm/v1/attention/ops/metadata.py | 24 ++++++++ .../spec_decode/autoregressive/speculator.py | 50 ++++++--------- vllm/v1/worker/gpu/spec_decode/speculator.py | 16 +++++ 8 files changed, 182 insertions(+), 50 deletions(-) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 052088a621c9..17dea5d03c19 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -122,6 +122,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = 16 + builder.kernel_block_size = None builder.num_sms = num_sms builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) builder.kv_cache_spec = SimpleNamespace(num_states=64) @@ -140,6 +141,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): requires_padding=False, schedule_metadata=schedule_metadata, ), + block_table=block_table, ) builder.update_draft_decode_metadata(metadata) @@ -167,6 +169,65 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): torch.testing.assert_close(metadata.slot_mapping, expected_slot_mapping) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize( + "block_size,kernel_block_size,compress_ratio", [(64, 64, 1), (256, 64, 4)] +) +def test_indexer_draft_decode_refresh_follows_new_batch( + monkeypatch, block_size, kernel_block_size, compress_ratio +): + """Without a prior build(), the refresh must not keep the last batch's state.""" + monkeypatch.setattr(indexer_module, "has_deep_gemm", lambda: False) + device = torch.device("cuda") + num_reqs, width = 3, 8 + factor = block_size // kernel_block_size + + builder = object.__new__(DeepseekV32IndexerMetadataBuilder) + builder.dcp_world_size = 1 + builder.compress_ratio = compress_ratio + builder.kernel_block_size = kernel_block_size + builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) + builder.kv_cache_spec = SimpleNamespace(block_size=block_size, num_states=64) + + def stale(*shape: int) -> torch.Tensor: + return torch.full(shape, -7, dtype=torch.int32, device=device) + + decode = DeepSeekV32IndexerDecodeMetadata( + block_table=stale(num_reqs, width // factor), + seq_lens=stale(num_reqs, 1), + decode_lens=stale(num_reqs), + requires_padding=False, + schedule_metadata=stale(1), + indices=stale(num_reqs), + per_req_decode_lens=stale(num_reqs), + ) + request_block_table = torch.randint( + 0, 1000, (num_reqs, width), dtype=torch.int32, device=device + ) + seq_lens = torch.tensor([17, 33, 65], dtype=torch.int32, device=device) + metadata = DeepseekV32IndexerMetadata( + seq_lens=seq_lens, + max_seq_len=128, + slot_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), + num_decodes=num_reqs, + num_decode_tokens=num_reqs, + num_prefills=0, + num_prefill_tokens=0, + decode=decode, + block_table=request_block_table, + ) + + builder.update_draft_decode_metadata(metadata) + + expected_block_table = request_block_table[:, ::factor] // factor + torch.testing.assert_close(decode.block_table, expected_block_table) + torch.testing.assert_close(decode.indices, builder.arange_buffer[:num_reqs]) + assert torch.all(decode.per_req_decode_lens == 1) + assert torch.all(decode.decode_lens == 1) + expected_seq_lens = torch.div(seq_lens, compress_ratio, rounding_mode="floor") + torch.testing.assert_close(decode.seq_lens.flatten(), expected_seq_lens) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize( "query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0], [0, 0, 0]] diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ffb68582eab0..3387a31794c9 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -506,19 +506,13 @@ def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: # but not the right per-request boundaries. Padding requests have a query # length of zero and are skipped by the device boundary search. num_mapped_tokens = int(self.query_start_loc_cpu[-1]) - from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel + from vllm.v1.attention.ops.metadata import compute_token_to_req_indices num_output_tokens = max(num_mapped_tokens, num_tokens) assert buffer.shape[0] >= num_output_tokens - _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)]( - self.query_start_loc, - buffer, - self.query_start_loc.shape[0] - 1, - num_mapped_tokens, - num_output_tokens, - num_warps=4, + self._token_to_req_indices_cache = compute_token_to_req_indices( + self.query_start_loc, buffer, num_mapped_tokens, num_output_tokens ) - self._token_to_req_indices_cache = buffer[:num_output_tokens] return self._token_to_req_indices_cache[:num_tokens] # TODO(lucas): remove once we have FULL-CG spec-decode support @@ -609,8 +603,11 @@ class AttentionMetadataBuilder(ABC, Generic[M]): supports_update_block_table: bool = False # Whether the builder constructor requires the block-table width. requires_block_table_width: ClassVar[bool] = False - # Whether all step-dependent draft decode metadata can be updated in place, - # allowing one metadata build to be reused across autoregressive draft steps. + # Whether update_draft_decode_metadata() can regenerate all decode metadata + # from persistent device buffers with capture-safe ops only. Lets a + # speculator record the refresh inside its draft decode CUDA graph, so no + # eager build() is needed between steps or between the batches that replay + # it, whether the graph holds one draft forward or a multi-step loop. supports_draft_decode_metadata_update: bool = False @abstractmethod @@ -758,12 +755,12 @@ def build_for_drafting( ) def update_draft_decode_metadata(self, metadata: M) -> None: - """Update step-dependent draft decode metadata in place. + """Update draft decode metadata in place. - The fused draft loop may call this method during full CUDA graph - capture. CUDA graph replay does not run this Python method, so - implementations must emit capture-safe operations and keep replayed - tensor state in persistent storage. + Speculators can record this call inside their draft decode CUDA graphs, + so it runs once at capture and never again in Python. Implementations + must emit only capture-safe operations and write results into persistent + tensors that ``metadata`` already references. """ raise NotImplementedError diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 3ca88992aa28..aabd68224093 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -27,6 +27,7 @@ flat_kv_row_view, triton_convert_req_index_to_global_index, ) +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func @@ -132,8 +133,18 @@ def __init__( self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True) self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 - def update_draft_decode_metadata(self, _metadata: SparseMLACommonMetadata) -> None: - pass + def update_draft_decode_metadata(self, metadata: SparseMLACommonMetadata) -> None: + num_tokens = metadata.num_decode_tokens + if num_tokens == 0: + return + # Everything else is a view of runner buffers; the per-token request + # map is a builder buffer that build() filled for the captured batch. + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.req_id_per_token, + num_tokens, + num_tokens, + ) class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 9d3823b880b9..8b186502f2a9 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -678,6 +678,7 @@ class DeepseekV32IndexerMetadata: decode: DeepSeekV32IndexerDecodeMetadata | None = None prefill: DeepseekV32IndexerPrefillMetadata | None = None positions: torch.Tensor | None = None + block_table: torch.Tensor | None = None @triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) @@ -1713,6 +1714,7 @@ def build( num_prefill_tokens=num_prefill_tokens, prefill=prefill_metadata, decode=decode_metadata, + block_table=common_attn_metadata.block_table_tensor, ) return attn_metadata @@ -1728,8 +1730,28 @@ def update_draft_decode_metadata( assert metadata.num_prefills == 0 assert metadata.num_decodes == metadata.num_decode_tokens assert decode.seq_lens.numel() == metadata.num_decode_tokens + assert metadata.block_table is not None assert self.dcp_world_size == 1 + # One query token per request, so the per-token decode block table is + # the request block table (at the indexer kernel's page size). build() + # may copy it into a builder buffer, so re-derive it for this batch. + block_table = metadata.block_table[: metadata.num_decode_tokens] + kernel_block_size = self.kernel_block_size + if ( + self.compress_ratio > 1 + and kernel_block_size is not None + and self.kv_cache_spec.block_size != kernel_block_size + and self.kv_cache_spec.block_size % kernel_block_size == 0 + ): + factor = self.kv_cache_spec.block_size // kernel_block_size + block_table = block_table[:, ::factor] // factor + decode.block_table.copy_(block_table) + if decode.indices is not None: + decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens]) + if decode.per_req_decode_lens is not None: + decode.per_req_decode_lens.fill_(1) + if self.compress_ratio > 1: get_compressed_slot_mapping( metadata.num_decode_tokens, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 0f5bff41442a..c5cbb4ef2822 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -33,6 +33,7 @@ ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, @@ -829,7 +830,8 @@ def update_draft_decode_metadata( self, metadata: DeepseekSparseSWAMetadata, ) -> None: - if metadata.num_decode_tokens == 0: + num_tokens = metadata.num_decode_tokens + if num_tokens == 0: return assert metadata.query_start_loc is not None assert metadata.seq_lens is not None @@ -839,6 +841,17 @@ def update_draft_decode_metadata( assert metadata.decode_swa_lens is not None assert metadata.replay_start is not None + # Recompute the validity mask and token->request map from the device + # buffers. Their padding differs from the dummy batch the graph was + # captured with. + torch.ge(metadata.slot_mapping, 0, out=metadata.is_valid_token) + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.token_to_req_indices, + num_tokens, + num_tokens, + ) + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( metadata.decode_swa_indices, metadata.decode_swa_lens, diff --git a/vllm/v1/attention/ops/metadata.py b/vllm/v1/attention/ops/metadata.py index f5baa51a9297..1852e77ee28f 100644 --- a/vllm/v1/attention/ops/metadata.py +++ b/vllm/v1/attention/ops/metadata.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Device-side request mapping and sparse indexer metadata.""" +import torch + from vllm.triton_utils import tl, triton @@ -27,6 +29,28 @@ def _token_request_mapping_kernel( tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens) +def compute_token_to_req_indices( + query_start_loc: torch.Tensor, + out: torch.Tensor, + num_mapped_tokens: int, + num_tokens: int, +) -> torch.Tensor: + """Map each of the first ``num_tokens`` tokens to its request index. + + Reads only the device ``query_start_loc``, so it is safe to record in a + CUDA graph. Tokens at or past ``num_mapped_tokens`` are mapped to 0. + """ + _token_request_mapping_kernel[((num_tokens + 255) // 256,)]( + query_start_loc, + out, + query_start_loc.shape[0] - 1, + num_mapped_tokens, + num_tokens, + num_warps=4, + ) + return out[:num_tokens] + + @triton.jit( do_not_specialize=[ "num_reqs", diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 7242a840d249..46ff82681561 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -99,10 +99,6 @@ def _configure_fused_multi_step_decode(self) -> None: self.use_fused_multi_step_decode = False return - if not self.advance_draft_positions: - self.use_fused_multi_step_decode = True - return - unsupported_backends = sorted( { attn_group.backend.get_name() @@ -574,6 +570,11 @@ def _fused_multi_step_decode( query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] idx_mapping = self.idx_mapping[:num_reqs] + if batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.run_fullgraph(batch_desc) + return + attn_metadata = None slot_mappings_by_layer = None if not skip_attn: @@ -583,10 +584,9 @@ def _fused_multi_step_decode( positions, batch_desc.num_tokens, ) - if batch_desc.cg_mode != CUDAGraphMode.FULL: - slot_mappings_by_layer = build_slot_mappings_by_layer( - slot_mappings, self.kv_cache_config - ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, self.kv_cache_config + ) attn_metadata = self._build_uniform_attn_metadata( num_reqs=num_reqs, batch_desc=batch_desc, @@ -595,11 +595,6 @@ def _fused_multi_step_decode( step=1, ) - if batch_desc.cg_mode == CUDAGraphMode.FULL: - assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.run_fullgraph(batch_desc) - return - self._generate_fused_drafts( num_reqs, batch_desc.num_tokens, @@ -623,13 +618,19 @@ def _generate_fused_drafts( idx_mapping = self.idx_mapping[:num_reqs] positions = self.input_buffers.positions[:num_reqs] query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] - attn_groups = ( - [group for groups in self.attn_groups for group in groups] - if attn_metadata is not None - else [] - ) for step in range(1, num_speculative_steps): + if attn_metadata is not None and ( + self.advance_draft_positions or step == 1 + ): + self.block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded, + ) + self._update_draft_decode_metadata(attn_metadata, num_reqs) + self.current_draft_step.fill_(step) self._generate_draft( num_reqs, @@ -640,19 +641,6 @@ def _generate_fused_drafts( cudagraph_runtime_mode, num_speculative_steps, ) - if ( - step < num_speculative_steps - 1 - and attn_metadata is not None - and self.advance_draft_positions - ): - self.block_tables.compute_slot_mappings( - idx_mapping, - query_start_loc, - positions, - num_tokens_padded, - ) - for attn_group in attn_groups: - attn_group.update_draft_decode_metadata(attn_metadata) def _generate_draft( self, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 111066328846..d07c043de915 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -562,3 +562,19 @@ def _build_uniform_attn_metadata( causal=causal, dcp_local_seq_lens=dcp_local_seq_lens, ) + + def _update_draft_decode_metadata( + self, attn_metadata: dict[str, Any], num_reqs: int + ) -> None: + if self.block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + for groups in self.attn_groups: + for group in groups: + group.update_draft_decode_metadata(attn_metadata) From 451689662688bfb3152f5c2218b70c24b12c5852 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 6 Oct 2026 18:23:37 +0200 Subject: [PATCH 079/952] [CI/Build] Warm the DP engines before measuring load balance in test_load (#60241) Signed-off-by: Stefan Koncarevic --- tests/v1/distributed/test_async_llm_dp.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index cb23910566c7..ecdef1f38b0b 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -132,6 +132,26 @@ def log_engine_initialized(self): ) after.callback(engine.shutdown) + # The first request an engine serves pays a few hundred milliseconds of + # one-off setup, and while it does the engine holds the requests already + # routed to it while reporting itself idle, which is enough to skew the + # split asserted on below. Pay it on every engine first, then measure. + await asyncio.gather( + *( + generate( + engine, + f"warmup-{rank}", + prompt, + output_kind, + max_tokens=1, + data_parallel_rank=rank, + ) + for rank in range(DP_SIZE) + ) + ) + for sl in stats_loggers.values(): + sl.finished_req_count = 0 + NUM_REQUESTS = 100 NUM_EXPECTED_TOKENS = 10 From 29f955cce0e8713e22d50d521dc0589f2c4bd319 Mon Sep 17 00:00:00 2001 From: eky-amd Date: Tue, 6 Oct 2026 12:26:38 -0400 Subject: [PATCH 080/952] [ROCm] Fuse MLA dual RMSNorm + FP8 group quant for DeepSeek-R1 (#54857) Signed-off-by: Ethan Ky Co-authored-by: Claude Fable 5 --- .../passes/test_fuse_mla_dual_rms_norm.py | 131 ++++++++++++++++ vllm/_aiter_ops.py | 4 + .../passes/fusion/rocm_aiter_fusion.py | 141 +++++++++++++++++- 3 files changed, 272 insertions(+), 4 deletions(-) diff --git a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py index 0f631b9472ad..75a0eee469b0 100644 --- a/tests/compile/passes/test_fuse_mla_dual_rms_norm.py +++ b/tests/compile/passes/test_fuse_mla_dual_rms_norm.py @@ -285,3 +285,134 @@ def test_fuse_mla_dual_rms_norm_fp8_per_token( backend.check_before_ops(model.ops_in_model_before()) backend.check_after_ops(model.ops_in_model_after()) + + +class MLADualRMSNormFp8GroupTestModel(torch.nn.Module): + """Minimal model reproducing the FP8 MLA attention path with *group* quant + (block-scaled fp8 weights, 128-wide activation groups): + + linear -> split([q_dim, kv_dim]) + +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fp8_group_quant -> dequant + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim]) + +-- kv_c (getitem 0) -> rms_norm (bf16) + +-- k_pe + """ + + GROUP_SIZE = 128 + + def __init__( + self, + hidden_size: int, + q_dim: int = Q_DIM, + kv_c_dim: int = KV_C_DIM, + k_pe_dim: int = K_PE_DIM, + eps: float = EPS, + ): + super().__init__() + self.q_dim = q_dim + self.kv_dim = kv_c_dim + k_pe_dim + self.kv_c_dim = kv_c_dim + self.k_pe_dim = k_pe_dim + self.eps = eps + self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False) + self.q_weight = torch.nn.Parameter(torch.ones(q_dim)) + self.kv_norm = RMSNorm(kv_c_dim, eps=eps) + + def _dequant(self, x_fp8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + # Group: one (M, N / 128) scale per 128-column group of the row. + m = x_fp8.shape[0] + x = x_fp8.to(torch.float32).view(m, -1, self.GROUP_SIZE) + return (x * scale.view(m, -1, 1)).view(m, -1).to(torch.bfloat16) + + def forward(self, x: torch.Tensor): + # Avoid graph input being a direct arg to a matched pattern node + x = torch.relu(x) + projected = self.proj(x) + q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1) + q_fp8, q_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant( + q_c, self.q_weight, self.eps, self.GROUP_SIZE + ) + kv_normed = self.kv_norm(kv_c) + # Also return the group scales so the test can check them directly: + # equal scales pin any fp8 payload difference to rounding alone. + return self._dequant(q_fp8, q_scale), q_scale, kv_normed, k_pe + + def ops_in_model_before(self): + return [ + torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default, + torch.ops.vllm_ir.rms_norm.default, + ] + + def ops_in_model_after(self): + return [torch.ops.vllm.fused_mla_dual_rms_norm_group_quant.default] + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("hidden_size", [7168]) +@pytest.mark.skipif( + not is_aiter_found_and_supported(), + reason="Only test on ROCm with AITER installed and supported", +) +def test_fuse_mla_dual_rms_norm_fp8_group( + dtype: torch.dtype, + hidden_size: int, + monkeypatch: pytest.MonkeyPatch, +): + torch._dynamo.reset() + vllm_config = VllmConfig( + model_config=ModelConfig(dtype=dtype), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + custom_ops=["+rms_norm"], + pass_config=PassConfig( + fuse_mla_dual_rms_norm=True, + eliminate_noops=True, + ), + ), + ) + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: + from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( + MLADualRMSNormFusionPass, + ) + + torch.set_default_device("cuda") + torch.set_default_dtype(dtype) + torch.manual_seed(42) + m.setenv("VLLM_ROCM_USE_AITER", "1") + rocm_aiter_ops.refresh_env_variables() + + fusion_pass = MLADualRMSNormFusionPass(vllm_config) + passes = [ + NoOpEliminationPass(vllm_config), + fusion_pass, + PostCleanupPass(vllm_config), + ] + backend = TestBackend(*passes) + model = MLADualRMSNormFp8GroupTestModel(hidden_size) + + x = torch.randn(64, hidden_size) + torch._dynamo.mark_dynamic(x, 0) + + with torch.inference_mode(): + outputs_unfused = model(x) + model_fused = torch.compile(model, backend=backend) + outputs_fused = model_fused(x) + + q_deq_u, q_scale_u, kv_normed_u, k_pe_u = outputs_unfused + q_deq_f, q_scale_f, kv_normed_f, k_pe_f = outputs_fused + + torch.testing.assert_close(k_pe_u, k_pe_f, atol=0, rtol=0) + torch.testing.assert_close(kv_normed_u, kv_normed_f, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(q_scale_u, q_scale_f, atol=0, rtol=1e-5) + + E4M3_STEP = 0.125 + exact_frac = (q_deq_u == q_deq_f).float().mean().item() + assert exact_frac > 0.99, f"q: only {exact_frac:.4f} of elements bit-exact" + torch.testing.assert_close(q_deq_u, q_deq_f, atol=1e-2, rtol=E4M3_STEP) + + assert fusion_pass.matched_count == 1, ( + f"Expected 1 fused pair, got {fusion_pass.matched_count}" + ) + backend.check_before_ops(model.ops_in_model_before()) + backend.check_after_ops(model.ops_in_model_after()) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index ca98e63b394f..941f66e5b71f 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2994,6 +2994,10 @@ def fused_qk_rmsnorm_group_quant( transpose_scale, ) + @staticmethod + def get_fused_mla_dual_rms_norm_group_quant_op() -> OpOverload: + return torch.ops.vllm.fused_mla_dual_rms_norm_group_quant.default + @staticmethod def rms_norm( x: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index 19bbfc2d8754..849cd5d67d42 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -1074,16 +1074,148 @@ def _replacement( return _replacement +class MLADualRMSGroupQuantPattern( + VllmPatternReplacement[ + ..., + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + ] +): + """Fuse the MLA FP8 *block-scale* attention path -- q-latent RMSNorm + FP8 + group quant plus kv-latent RMSNorm -- into AITER's + ``fused_qk_rmsnorm_group_quant``. + + With a block-scale FP8 ``q_b_proj`` (DeepSeek fp8), the earlier + ``RocmAiterRMSNormQuantFusionPass`` folds the q side into + ``rocm_aiter_rmsnorm_fp8_group_quant`` (the Triton rms+group-quant kernel) + and leaves the kv side a plain ``vllm_ir.rms_norm``. This pattern matches + that asymmetric pair:: + + gemm -> split_with_sizes([q_dim, kv_dim]) + +-- q_c -> rocm_aiter_rmsnorm_fp8_group_quant -> (q_fp8, q_scale) + +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim]) + +-- kv_c -> vllm_ir.rms_norm -> kv_normed (bf16) + +-- k_pe + """ + + GROUP_QUANT_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op() + FUSED_OP = rocm_aiter_ops.get_fused_mla_dual_rms_norm_group_quant_op() + + def __init__(self, epsilon: float) -> None: + self._epsilon = epsilon + + def get_inputs(self) -> list[torch.Tensor]: + q_dim, kv_c_dim, k_pe_dim = 256, 128, 64 + return [ + self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim), + self.empty_bf16(q_dim), + self.empty_bf16(kv_c_dim), + ] + + @property + def pattern( + self, + ) -> Callable[ + ..., + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + ]: + eps = self._epsilon + group_quant_op = self.GROUP_QUANT_OP + + def _pattern( + projected: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + q_dim = q_weight.shape[0] + kv_dim = projected.shape[-1] - q_dim + kv_c_dim = kv_weight.shape[0] + k_pe_dim = kv_dim - kv_c_dim + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) + quant_kwargs = dict( + x=q_c, + weight=q_weight, + variance_epsilon=eps, + group_size=128, + ) + q_quant = group_quant_op(**quant_kwargs) + kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps) + return q_quant[0], q_quant[1], kv_normed, k_pe + + return _pattern + + @property + def replacement( + self, + ) -> Callable[ + ..., + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + ]: + eps = self._epsilon + fused_op = self.FUSED_OP + + def _replacement( + projected: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + q_dim = q_weight.shape[0] + kv_dim = projected.shape[-1] - q_dim + kv_c_dim = kv_weight.shape[0] + k_pe_dim = kv_dim - kv_c_dim + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1) + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1) + at = fused_op( + q_c, + q_weight, + kv_c, + kv_weight, + eps, + eps, + 128, + False, + ) + # q_fp8, q_scale, kv_normed, k_pe + return at[0], at[1], at[2], k_pe + + return _replacement + + class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass): """Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's ``fused_qk_rmsnorm`` HIP kernel. The FP8 attention path is also handled via - :class:`MLADualRMSPerTokenQuantPattern`, which fuses the q-latent RMSNorm + - FP8 per-token quant together with the kv-latent RMSNorm into - ``fused_mla_dual_rms_norm_per_token_quant`` backed by aiter's - ``fused_qk_rmsnorm_per_token_quant`` HIP kernel. + :class:`MLADualRMSPerTokenQuantPattern` or :class:`MLADualRMSGroupQuantPattern`, + which fuse the q-latent RMSNorm + FP8 per-token/group quant together with + the kv-latent RMSNorm into ``fused_mla_dual_rms_norm__quant`` + backed by aiter's ``fused_qk_rmsnorm__quant`` HIP kernel. """ def __init__(self, config: VllmConfig) -> None: @@ -1092,3 +1224,4 @@ def __init__(self, config: VllmConfig) -> None: for epsilon in [1e-5, 1e-6]: self.register(MLADualRMSNormPattern(epsilon)) self.register(MLADualRMSPerTokenQuantPattern(epsilon)) + self.register(MLADualRMSGroupQuantPattern(epsilon)) From f8d92cae410a56822bc5a89ccb96012b44bea222 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Tue, 6 Oct 2026 18:38:18 +0200 Subject: [PATCH 081/952] [Quantization] Layer re-quantization for linear layers through online quantization API (MXFP8 -> FP8 PTPC showcase) (#55684) Signed-off-by: Felix Marty Signed-off-by: Tan Pin Siang Co-authored-by: vllmellm Co-authored-by: Tan Pin Siang Co-authored-by: Misha Goin --- docs/features/quantization/online.md | 29 +++ tests/quantization/test_online.py | 204 +++++++++++++++--- .../layers/quantization/base_config.py | 42 +++- .../layers/quantization/modelopt.py | 12 ++ .../layers/quantization/online/base.py | 1 + .../layers/quantization/online/fp8.py | 56 ++++- .../layers/quantization/online/moe_base.py | 7 + .../layers/quantization/online/mxfp4.py | 3 +- .../layers/quantization/online/mxfp8.py | 3 +- .../model_loader/base_loader.py | 27 ++- .../model_loader/reload/layerwise.py | 33 ++- vllm/model_executor/model_loader/utils.py | 50 ++++- 12 files changed, 422 insertions(+), 45 deletions(-) diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index d3a4edd2710f..a81ea219f718 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -99,6 +99,35 @@ vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8 Combine with `--moe-backend` to pin a specific kernel family. +### Re-quantizing layers to a different precision + +An online target can re-quantize a checkpoint-quantized (prequantized) layer at load time. + +vLLM uses the checkpoint method to load and dequantize the checkpoint weight, and then applies the online quantization method. This requires the checkpoint +method to implement weight dequantization, unsupported source and target +combinations raise an error. + +Currently supported combinations are: + +| Checkpoint layer | Online target | Requirements | +| ---------------- | ------------- | ------------ | +| Serialized ModelOpt MXFP8 linear | any | Target-specific hardware requirements | + +For example, serialized ModelOpt MXFP8 linears can be converted to +per-channel-weight/per-token-activation FP8 at load time: + +```bash +vllm serve MiniMaxAI/MiniMax-M3-MXFP8 \ + --linear-backend auto \ + --quantization-config.linear fp8_per_channel +``` + +!!! info + + This is only available for linear layers at the moment. + + Checkpoint-exclusion policy is not enforced by online quantization/re-quantization. Only exclusion through `--quantization-config.targets`/`--quantization-config.ignore` is enforced. + ### Online quantization on unquantized layers from partially-quantized checkpoints Online quantization can be used on already quantized checkpoints independently of their original `quant_method` (`modelopt`, `compressed-tensors`, `quark`, etc.), for layers that are left unquantized in the original checkpoint. diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index b4dffbb37588..f3b72707f3d6 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -58,6 +58,7 @@ ) from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, + ModelOptLinearMethod, ModelOptMxFp8Config, ) from vllm.model_executor.layers.quantization.online.base import ( @@ -93,12 +94,14 @@ ) from vllm.model_executor.layers.quantization.utils import quant_utils from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( amax_for_moe_weight_quant, amax_for_tp_weight_quant, kMxfp8Dynamic, + kMxfp8Static, weight_amax, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -112,6 +115,7 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe +from vllm.utils.torch_utils import set_default_torch_dtype if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx942, on_gfx950 @@ -307,6 +311,14 @@ def _fully_quantized_modelopt_config() -> ModelOptFp8Config: ) +def _fully_quantized_modelopt_mxfp8_config() -> ModelOptMxFp8Config: + return ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=["lm_head"], + ) + + def _moe_only_compressed_tensors_config() -> CompressedTensorsConfig: return CompressedTensorsConfig( target_scheme_map={"RoutedExperts": {}}, @@ -334,10 +346,19 @@ def _write_minimal_llama_config( @pytest.mark.parametrize( - "checkpoint_config_factory,raises_conflict", + "checkpoint_config_factory,raises_requantization_error", [ pytest.param(_fully_quantized_quark_config, True, id="quark"), - pytest.param(_fully_quantized_modelopt_config, True, id="modelopt"), + pytest.param( + _fully_quantized_modelopt_config, + True, + id="modelopt", + ), + pytest.param( + _fully_quantized_modelopt_mxfp8_config, + False, + id="modelopt_mxfp8", + ), pytest.param( _moe_only_compressed_tensors_config, False, @@ -345,9 +366,13 @@ def _write_minimal_llama_config( ), ], ) +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="MXFP8 requantization requires CUDA or ROCm.", +) def test_online_prequantized_compatibility( checkpoint_config_factory, - raises_conflict: bool, + raises_requantization_error: bool, default_vllm_config, dist_init, ) -> None: @@ -370,35 +395,69 @@ def test_online_prequantized_compatibility( "disable_tp": True, } - if raises_conflict: - with pytest.raises(ValueError, match="pre-quantized layer"): - ColumnParallelLinear(**layer_kwargs) + layer = ColumnParallelLinear(**layer_kwargs) + assert isinstance(layer.quant_method, Mxfp8OnlineLinearMethod) + + if layer.weight.is_meta: + # Weight loading. + layer = layer.to_empty(device=DEVICE) + for parameter in layer.parameters(): + parameter.zero_() + else: + layer = layer.to(device=DEVICE) + + if raises_requantization_error: + with pytest.raises( + NotImplementedError, + match="does not implement dequantize_weight|only supported for MXFP8", + ): + layer.quant_method.process_weights_after_loading(layer) else: - layer = ColumnParallelLinear(**layer_kwargs) + layer.quant_method.process_weights_after_loading(layer) + + if isinstance(checkpoint_config, ModelOptMxFp8Config): assert isinstance(layer.quant_method, Mxfp8OnlineLinearMethod) + assert not layer.quant_method.uses_meta_device + assert isinstance( + layer.quant_method.requantization_source, ModelOptLinearMethod + ) + elif not raises_requantization_error: + assert layer.quant_method.uses_meta_device + assert layer.quant_method.requantization_source is None -def test_online_target_rejects_prequantized_layer( - default_vllm_config, dist_init +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()) + or not is_quant_method_supported("fp8"), + reason="Requires FP8 support on CUDA or ROCm.", +) +def test_online_block_fp8_requantization_removes_source_scales( + default_vllm_config, dist_init, workspace_init ) -> None: - """A targets match participates in checkpoint compatibility checks.""" - default_vllm_config.model_config = ModelConfig() - prefix = "model.layers.0.self_attn.o_proj" - quant_config = _fully_quantized_quark_config() + """Blockwise FP8 replaces MXFP8 scales with its own scale parameter.""" + default_vllm_config.model_config = ModelConfig(dtype="bfloat16") + quant_config = _fully_quantized_modelopt_mxfp8_config() quant_config.online_quantization_config = OnlineQuantizationConfig( - QuantizationConfigArgs(targets={prefix: "mxfp8"}) + QuantizationConfigArgs(linear="fp8_per_block") ) - - with pytest.raises(ValueError, match="pre-quantized layer"): - ColumnParallelLinear( - input_size=32, - output_size=32, + with set_default_torch_dtype(torch.bfloat16), torch.device(DEVICE): + layer = ColumnParallelLinear( + input_size=256, + output_size=256, bias=False, params_dtype=torch.bfloat16, quant_config=quant_config, - prefix=prefix, + prefix="model.layers.0.self_attn.o_proj", disable_tp=True, ) + with torch.no_grad(): + layer.weight.fill_(1) + layer.weight_scale.fill_(127) # E8M0 encoding of scale 1. + + layer.quant_method.process_weights_after_loading(layer) + + assert not hasattr(layer, "weight_scale") + assert layer.weight_scale_inv.numel() > 0 def test_online_ignore_keeps_checkpoint_quantization_linear( @@ -431,17 +490,25 @@ def test_online_ignore_keeps_checkpoint_quantization_linear( def test_online_quantization_rejects_prequantized_moe( - default_vllm_config, dist_init + default_vllm_config, dist_init, monkeypatch ) -> None: - """Online linear and MoE quantization reject a pre-quantized MoE layer.""" + """Reject pre-quantized MoE before constructing an online target backend.""" default_vllm_config.model_config = ModelConfig() prefix = "model.layers.0.mlp.experts" quant_config = _fully_quantized_quark_config() quant_config.online_quantization_config = OnlineQuantizationConfig( quant_config_args(linear="mxfp4", moe="mxfp4") ) + monkeypatch.setattr( + quant_config.online_quantization_config, + "get_quant_method", + lambda *args: pytest.fail("unsupported online backend must not be constructed"), + ) - with pytest.raises(ValueError, match="pre-quantized layer"): + with pytest.raises( + NotImplementedError, + match="Requantizing checkpoint-quantized MoE layers is not supported", + ): FusedMoEFactory( num_experts=4, top_k=2, @@ -591,7 +658,7 @@ def test_log_online_quantization_for_composable_config(monkeypatch) -> None: assert log_args == [ ( - "Quantized %d layers of types: %s", + "Quantizing %d layers of types: %s", 2, "; ".join(online_config.quantized_layer_summaries), ) @@ -836,6 +903,24 @@ def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtyp {}, id="partially_prequantized_checkpoint", ), + pytest.param( + PARTIALLY_PREQUANTIZED_MODEL_NAME, + None, + {"targets": {"model.layers.1.self_attn.o_proj": "mxfp8"}}, + Mxfp8OnlineLinearMethod, + CompressedTensorsMoEMethod, + {}, + id="partially_prequantized_checkpoint", + ), + pytest.param( + "mgoin/Qwen3-0.6B-MXFP8", + None, + {"linear": "fp8_per_channel"}, + Fp8PtpcOnlineLinearMethod, + None, + {}, + id="requantization_mxfp8_ptcp_fp8", + ), ], ) @pytest.mark.parametrize( @@ -888,6 +973,25 @@ def test_online_quantization( moe_backend=runner_kwargs.get("moe_backend", "auto"), ), } + if model_name == "mgoin/Qwen3-0.6B-MXFP8": + original_process = Fp8PtpcOnlineLinearMethod.process_weights_after_loading + + def assert_source_weights_released(method, layer) -> None: + source_parameters = dict(method.requantization_source_parameters) + assert set(source_parameters) == {"weight", "weight_scale"} + + original_process(method, layer) + + assert not method.requantization_source_parameters + assert layer.weight is not source_parameters["weight"] + assert layer.weight_scale is not source_parameters["weight_scale"] + assert layer.weight_scale.shape[-1] == 1 + + monkeypatch.setattr( + Fp8PtpcOnlineLinearMethod, + "process_weights_after_loading", + assert_source_weights_released, + ) if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: model, vllm_config = load_model_without_vllm_runner( @@ -943,11 +1047,17 @@ def test_online_quantization( if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: o_proj = model.model.layers[1].self_attn.o_proj moe = model.model.layers[0].mlp.experts + elif model_name == "mgoin/Qwen3-0.6B-MXFP8": + o_proj = model.model.layers[0].self_attn.o_proj + moe = None else: o_proj = model.model.layers[0].self_attn.o_proj moe = model.model.layers[0].block_sparse_moe.experts + assert isinstance(o_proj.quant_method, expected_linear_cls) - assert isinstance(moe._quant_method, expected_moe_cls) + + if moe is not None: + assert isinstance(moe._quant_method, expected_moe_cls) if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( o_proj.quant_method.kernel, MarlinMxfp8LinearKernel @@ -1242,7 +1352,7 @@ def record_info(message: str, *args: object) -> None: log_online_quantization(default_vllm_config) assert logged_messages == [ - "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: " + "Quantizing 3 layers of types: mlp.down_proj: 2 (from linear: " "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: " "re:.*qkv_proj.*, mxfp4)" ] @@ -1365,6 +1475,48 @@ def test_online_linear_tp_weight_quant_matches_unsharded( assert torch.equal(tp_scale, full_scale) +@pytest.mark.skipif( + not is_quant_method_supported("fp8"), + reason="FP8 is not supported on this GPU type.", +) +def test_modelopt_mxfp8_requantized_ptpc_tp_shards_match_unsharded() -> None: + """MXFP8 source shards convert to slices of the unsharded PTPC weight.""" + torch.manual_seed(0) + serialized_weight = torch.randn(16, 64, device=DEVICE, dtype=torch.bfloat16).to( + MXFP8_VALUE_DTYPE + ) + serialized_scale = torch.randint( + 124, + 131, + (16, 2), + device=DEVICE, + dtype=MXFP8_SCALE_DTYPE, + ) + + source = ModelOptLinearMethod.__new__(ModelOptLinearMethod) + source.wkey = SimpleNamespace(key=kMxfp8Static) + + def dequantize(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + layer = SimpleNamespace(weight=weight, weight_scale=scale) + return source.dequantize_weight(layer) + + full_weight = dequantize(serialized_weight, serialized_scale) + full_scale = _fp8_channel_scale(weight_amax(full_weight, dim=-1, keepdim=True)) + full_qweight = _fp8_quant_per_channel(full_weight, full_scale) + + # Mimic the two K-dimension shards of a row-parallel layer. PTPC's global + # per-output-channel scale is the same on both ranks after its MAX reduce. + for rank in range(2): + start = rank * 32 + shard_weight = serialized_weight[:, start : start + 32].contiguous() + shard_scale = serialized_scale[:, rank : rank + 1].contiguous() + weight = dequantize(shard_weight, shard_scale) + qweight = _fp8_quant_per_channel(weight, full_scale) + + assert torch.equal(weight, full_weight[:, start : start + 32]) + assert torch.equal(qweight, full_qweight[:, start : start + 32]) + + def _quantize_moe(weight, scheme, moe_tp_size): if scheme == "nvfp4": return _quantize_moe_weight_to_nvfp4(weight, moe_tp_size) diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 1b0457854eb1..6c1f93473906 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -84,6 +84,13 @@ def process_weights_after_loading(self, layer: nn.Module) -> None: """ return + def dequantize_weight(self, layer: nn.Module) -> torch.Tensor: + """Materialize a serialized quantized weight for requantization.""" + raise NotImplementedError( + f"The quantization method {type(self)} does not implement " + "dequantize_weight. Please open an issue." + ) + def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: """Not all quant methods have embedding implemented, so we need to check that @@ -294,13 +301,16 @@ def resolve_quant_method( LinearBase, UnquantizedLinearMethod, ) + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase base_quant_method = quant_config.get_quant_method(layer, prefix) if quant_config.online_quantization_config is None: + # No online configuration: retain the checkpoint method. return base_quant_method # Online quantization currently supports only LinearBase and RoutedExperts. # Embeddings and ParallelLMHead retain their checkpoint quantization method. if not isinstance(layer, (LinearBase, RoutedExperts)): + # Online quantization only supports linear and routed MoE layers. return base_quant_method quant_config.online_quantization_config.packed_modules_mapping = ( @@ -312,14 +322,34 @@ def resolve_quant_method( online_target = quant_config.online_quantization_config.resolve_quant_method_cls( layer, prefix ) + if checkpoint_is_quantized: - if online_target is not None: - raise ValueError( - f"Cannot apply requested online quantization {online_target[3]} to " - f"pre-quantized layer {prefix}: {base_quant_method} was already " - "selected by the checkpoint quantization config." + if online_target is None: + # The checkpoint quant method is applied as there is no online override. + return base_quant_method + + if isinstance(layer, RoutedExperts): + raise NotImplementedError( + "Requantizing checkpoint-quantized MoE layers is not supported." ) - return base_quant_method + + online_quant_method = quant_config.online_quantization_config.get_quant_method( + layer, prefix + ) + + assert base_quant_method is not None and not isinstance( + base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) + ) + + assert isinstance(online_quant_method, OnlineLinearBase) + online_quant_method.set_requantization_source(base_quant_method) + + # The online method dequantizes the checkpoint method before requantizing. + return online_quant_method + if online_target is None: + # The layer is unquantized and no online target applies. return base_quant_method + + # Quantize an unquantized layer with the online method. return quant_config.online_quantization_config.get_quant_method(layer, prefix) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 272eb457bbe0..e33226cad3c8 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -76,6 +76,7 @@ MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, + dequant_mxfp8_to_bf16, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( FP4_DTYPE, @@ -2572,6 +2573,17 @@ def process_weights_after_loading(self, layer) -> None: layer._nvfp4_group_size_for_gather = self.ctx.group_size self.kernel.process_weights_after_loading(layer) + def dequantize_weight(self, layer: torch.nn.Module) -> torch.Tensor: + """Reconstruct serialized weights for online requantization.""" + if self.wkey.key is kMxfp8Static: + return dequant_mxfp8_to_bf16( + layer.weight.contiguous(), layer.weight_scale.contiguous() + ) + else: + raise NotImplementedError( + "ModelOpt weight dequantization is only supported for MXFP8." + ) + def apply(self, layer, x, bias=None): return self.fmt.apply( layer, diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 8ed48b2b64f3..3f35b56a31b8 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -156,6 +156,7 @@ def __init__( self.args = args self.ignored_layers: list[str] = args.ignore self.quantized_layers: dict[str, tuple[str, str, str | None]] = {} + self.online_quantization_time = 0.0 @property def quantized_layer_summaries(self) -> list[str]: diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index ed018c81bf0a..23a36c8fbc3b 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -34,6 +34,7 @@ from vllm.model_executor.layers.linear import ( LinearMethodBase, ) +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) @@ -125,6 +126,13 @@ class OnlineLinearBase(LinearMethodBase): def __init__(self): self.out_dtype = torch.get_default_dtype() self.input_dtype = get_current_vllm_config().model_config.dtype + self.requantization_source: QuantizeMethodBase | None = None + self.requantization_source_parameters: dict[str, torch.nn.Parameter] = {} + + def set_requantization_source(self, source_method: QuantizeMethodBase) -> None: + """Configure serialized-weight conversion before online quantization.""" + self.requantization_source = source_method + self.uses_meta_device = False def create_weights( self, @@ -136,6 +144,24 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): + if self.requantization_source is not None: + existing_parameter_names = set(layer._parameters) + self.requantization_source.create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + self.requantization_source_parameters = { + name: parameter + for name, parameter in layer._parameters.items() + if name not in existing_parameter_names and parameter is not None + } + return + output_size_per_partition = sum(output_partition_sizes) weight_loader = extra_weight_attrs.get("weight_loader") layer.logical_widths = output_partition_sizes @@ -159,6 +185,19 @@ def create_weights( initialize_online_processing(layer) + def release_requantization_source_weights(self, layer: torch.nn.Module) -> None: + """Release checkpoint parameters after successful requantization.""" + for name, source_parameter in self.requantization_source_parameters.items(): + if layer._parameters.get(name) is source_parameter: + delattr(layer, name) + self.requantization_source_parameters.clear() + + def get_weight_for_quantization(self, layer: Module) -> torch.Tensor: + """Return checkpoint weights materialized for online quantization.""" + if self.requantization_source is None: + return layer.weight + return self.requantization_source.dequantize_weight(layer) + class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. @@ -213,10 +252,11 @@ def process_weights_after_loading(self, layer: Module) -> None: return layer.input_scale = None - amax = weight_amax(layer.weight).reshape(1) + weight = self.get_weight_for_quantization(layer) + amax = weight_amax(weight).reshape(1) amax = amax_for_tp_weight_quant(amax, _is_tp_sharded(layer)) weight_scale = _fp8_scale(amax) - qweight, _ = ops.scaled_fp8_quant(layer.weight, scale=weight_scale) + qweight, _ = ops.scaled_fp8_quant(weight, scale=weight_scale) # Update layer with new values. replace_parameter(layer, "weight", qweight.t().data) @@ -312,14 +352,18 @@ def process_weights_after_loading(self, layer: Module) -> None: layer.input_scale = None block_size = self.weight_block_size + weight = self.get_weight_for_quantization(layer) qweight, weight_scale_inv = per_block_cast_to_fp8( - layer.weight, block_size=block_size, use_ue8m0=False + weight, block_size=block_size, use_ue8m0=False ) replace_parameter(layer, "weight", qweight.data) replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data) + if self.requantization_source is not None and hasattr(layer, "weight_scale"): + del layer.weight_scale + self.fp8_linear.process_weights_after_loading(layer) # Prevent duplicate processing (e.g., during weight reload) @@ -395,18 +439,20 @@ def process_weights_after_loading(self, layer: Module) -> None: return layer.input_scale = None - amax = weight_amax(layer.weight, dim=-1, keepdim=True) + weight = self.get_weight_for_quantization(layer) + amax = weight_amax(weight, dim=-1, keepdim=True) amax = amax_for_tp_weight_quant( amax, _is_tp_sharded(layer, reduces_output_dim=False) ) weight_scale = _fp8_channel_scale(amax) - qweight = _fp8_quant_per_channel(layer.weight, weight_scale) + qweight = _fp8_quant_per_channel(weight, weight_scale) replace_parameter(layer, "weight", qweight.t()) replace_parameter(layer, "weight_scale", weight_scale) self.fp8_linear.process_weights_after_loading(layer) expose_input_quant_key(layer, self.fp8_linear) + self.release_requantization_source_weights(layer) layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 6e21bc909224..e296eab37389 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -11,6 +11,7 @@ SharedExperts, ) from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.reload.layerwise import ( initialize_online_processing, ) @@ -24,6 +25,12 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): uses_meta_device: bool = True + def set_requantization_source(self, source_method: QuantizeMethodBase) -> None: + """Reject requantization from a checkpoint-quantized MoE method.""" + raise NotImplementedError( + "Requantizing checkpoint-quantized MoE layers is not supported." + ) + def create_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index 72914a9726c8..b949c9b7ff8d 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -119,7 +119,8 @@ def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return - weight_fp4, weight_scale = mxfp4_quantize(layer.weight.contiguous()) + weight = self.get_weight_for_quantization(layer) + weight_fp4, weight_scale = mxfp4_quantize(weight.contiguous()) layer.input_scale = None replace_parameter(layer, "weight", weight_fp4.data) diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 9d5814206c1d..9cddab85c12a 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -75,7 +75,8 @@ def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return - weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous()) + weight = self.get_weight_for_quantization(layer) + weight_fp8, weight_scale = mxfp8_e4m3_quantize(weight.contiguous()) layer.input_scale = None replace_parameter(layer, "weight", weight_fp8.data) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 2b7cb72fffe0..4c774808bf79 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -101,6 +101,7 @@ def load_model( finalize_layerwise_processing(model, model_config) process_weights_after_loading(model, model_config, target_device) + log_online_quantization_time(vllm_config) return model.eval() @@ -115,8 +116,8 @@ def log_model_inspection(model: nn.Module) -> None: logger.info("vLLM model structure:\n%s", format_model_inspection(model)) -def log_online_quantization(vllm_config: VllmConfig) -> None: - """Log the online-quantized layer count and types, when applicable.""" +def _get_online_quantization_config(vllm_config: VllmConfig): + """Return the online quantization config when one is configured.""" from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, ) @@ -128,15 +129,35 @@ def log_online_quantization(vllm_config: VllmConfig) -> None: if isinstance(online_quantization_config, OnlineQuantizationConfig): quant_config = online_quantization_config if not isinstance(quant_config, OnlineQuantizationConfig): + return None + return quant_config + + +def log_online_quantization(vllm_config: VllmConfig) -> None: + """Log the online-quantized layer count and types, when applicable.""" + quant_config = _get_online_quantization_config(vllm_config) + if quant_config is None: return logger.info( - "Quantized %d layers of types: %s", + "Quantizing %d layers of types: %s", len(quant_config.quantized_layers), "; ".join(quant_config.quantized_layer_summaries), ) +def log_online_quantization_time(vllm_config: VllmConfig) -> None: + """Log online quantization processing time, when applicable.""" + quant_config = _get_online_quantization_config(vllm_config) + if quant_config is None: + return + + logger.info( + "Online quantization as part of model loading took %.2f seconds", + quant_config.online_quantization_time, + ) + + def _has_online_quant(model: nn.Module): for module in model.modules(): quant_method = getattr(module, "quant_method", None) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 062bfa60906b..df91522732fe 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import inspect from collections.abc import Callable +from contextlib import nullcontext from functools import wraps from weakref import WeakKeyDictionary, WeakSet @@ -321,7 +322,21 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) - _get_weight_loader(param)(*args.args, **args.kwargs) if quant_method is not None: - quant_method.process_weights_after_loading(layer) + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) + from vllm.model_executor.model_loader.utils import ( + online_quantization_timing_context, + ) + + timing_context = ( + online_quantization_timing_context(layer, quant_method) + if isinstance(quant_method, (OnlineLinearBase, OnlineMoEMethodBase)) + else nullcontext() + ) + with timing_context: + quant_method.process_weights_after_loading(layer) _copy_and_restore_kernel_tensors(layer, info) @@ -356,7 +371,21 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): # Process weights (quantization, repacking, etc.) quant_method = getattr(layer, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): - quant_method.process_weights_after_loading(layer) + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) + from vllm.model_executor.model_loader.utils import ( + online_quantization_timing_context, + ) + + timing_context = ( + online_quantization_timing_context(layer, quant_method) + if isinstance(quant_method, (OnlineLinearBase, OnlineMoEMethodBase)) + else nullcontext() + ) + with timing_context: + quant_method.process_weights_after_loading(layer) # Re-reconcile parameter TP state: process_weights_after_loading may # have re-created Parameters (stamped with the global rank), which would # otherwise break replicated (disable_tp) weights on a subsequent reload. diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 465703f75983..c12b8e76c5f5 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -3,6 +3,7 @@ """Utilities for selecting and loading models.""" import inspect +import time import warnings from contextlib import contextmanager, nullcontext from typing import Any @@ -141,6 +142,11 @@ def process_weights_after_loading( """ # Reclaim memory when an explicit lm_head has been # loaded, but it is identical to the input embeddings. + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) + maybe_retie_word_embeddings(model, model_config) for name, module in model.named_modules(): @@ -166,7 +172,12 @@ def process_weights_after_loading( if quant_method.requires_device_loading else nullcontext() ) - with loading_context: + timing_context = ( + online_quantization_timing_context(module, quant_method) + if isinstance(quant_method, (OnlineLinearBase, OnlineMoEMethodBase)) + else nullcontext() + ) + with loading_context, timing_context: quant_method.process_weights_after_loading(module) # process_weights_after_loading may swap in freshly-created # Parameters (e.g. FP8 requantization), which are stamped with the @@ -245,6 +256,43 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device) p._vllm_is_uva_offloaded = True +@contextmanager +def online_quantization_timing_context( + layer: nn.Module, quant_method: QuantizeMethodBase +): + """Accumulate online quantization processing time for ``quant_method``.""" + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, + ) + + assert isinstance(quant_method, (OnlineLinearBase, OnlineMoEMethodBase)) + quant_config = layer.quant_config + assert quant_config is not None + online_quantization_config = ( + quant_config + if isinstance(quant_config, OnlineQuantizationConfig) + else quant_config.online_quantization_config + ) + # reload/layerwise.py's _layerwise_process calls this for experts_int8, + # whose online MoE method has no OnlineQuantizationConfig. + if online_quantization_config is None: + yield + return + assert isinstance(online_quantization_config, OnlineQuantizationConfig) + + start_time = time.perf_counter() + try: + yield + finally: + online_quantization_config.online_quantization_time += ( + time.perf_counter() - start_time + ) + + _MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]() """Caches the outputs of `_get_model_architecture`.""" From fba31f31a9d14a8fbe52e4246ba5b3fcd1a9ba63 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:54:20 -0400 Subject: [PATCH 082/952] [GLM5.3 Perf] Switch to fp8 kv cache by default for GLM, 2.3%~5.5% E2E Throughput Improvement (#60140) Signed-off-by: yewentao256 --- vllm/model_executor/models/config.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 8831c5f6c406..d80663731d1e 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -43,6 +43,18 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: + from vllm.platforms import current_platform + + cache_config = vllm_config.cache_config + if cache_config.cache_dtype == "auto": + if current_platform.is_xpu(): + cache_config.cache_dtype = "bfloat16" + elif current_platform.is_cuda_alike(): + capability = current_platform.get_device_capability() + cache_config.cache_dtype = ( + "fp8_e4m3" if capability and capability.major >= 10 else "bfloat16" + ) + # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs # in most cases. vllm_config.parallel_config.set_dcp_defaults( From e82b80099acba001693af29361a0255a7b171b8a Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 6 Oct 2026 13:04:55 -0400 Subject: [PATCH 083/952] [Perf][HiSparse] Cache per-request residency instead of rescanning every page each step (#60083) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- vllm/v1/hisparse/coordinator.py | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 766c926b12b7..390e0d27ced2 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -168,6 +168,10 @@ def __init__( # request -> prefix length an external load is still filling in. self._pending_imports: dict[str, int] = {} self.block_table_updates: set[str] = set() + # Requests whose resident pages were all non-null when last scanned. + # Pages only become null at admission or import (before the first + # scan) and when the pool reuses an unpinned page (_lose_page). + self._resident_requests: set[str] = set() self.spills_to_send: list[SparseKVPageTransfer] = [] self.pending_spills: dict[int, _PendingSpill] = {} self.request_states: dict[str, _HiSparseRequestState] = {} @@ -188,6 +192,7 @@ def _get_request_state(self, request_id: str) -> _HiSparseRequestState: def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None: """Account for a prefix hit on host pages the request now references.""" + self._resident_requests.discard(request_id) if not self.resident_managers or self.host_manager is None: return host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) @@ -382,6 +387,7 @@ def _on_block_reused(self, block: KVCacheBlock) -> None: self._lose_page(request_id, page_idx) def _lose_page(self, request_id: str, page_idx: int) -> None: + self._resident_requests.discard(request_id) state = self.request_states.get(request_id) if state is None: return @@ -530,6 +536,7 @@ def _publish_host_blocks_if_ready(self, request_id: str) -> None: def record_pending_host_import(self, request_id: str, num_tokens: int) -> None: """Note a prefix an external load is populating in host pages.""" + self._resident_requests.discard(request_id) self._pending_imports[request_id] = num_tokens def finish_host_import(self, request_id: str, *, failed: bool) -> None: @@ -674,15 +681,22 @@ def all_context_pages_resident( """Return whether every scheduled request can read only resident KV.""" if not self.resident_managers: return False - block_size = self.resident_managers[0].block_size + first = self.resident_managers[0] for request_id, num_computed_tokens, num_scheduled_tokens in requests: - num_pages = cdiv(num_computed_tokens + num_scheduled_tokens, block_size) - for page_idx in range(num_pages): - if any( - manager.get_resident_page(request_id, page_idx) is None - for manager in self.resident_managers - ): - return False + num_pages = cdiv( + num_computed_tokens + num_scheduled_tokens, first.block_size + ) + if len(first.req_to_blocks.get(request_id, ())) < num_pages: + return False + if request_id in self._resident_requests: + continue + if any( + block.is_null + for manager in self.resident_managers + for block in manager.req_to_blocks.get(request_id, ()) + ): + return False + self._resident_requests.add(request_id) return True def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: @@ -806,6 +820,7 @@ def _publish_detached_blocks(self, state: _HiSparseRequestState) -> None: def free(self, request_id: str) -> None: """Detach the request; its clean pages stay readable copies in the pool.""" + self._resident_requests.discard(request_id) self._pending_imports.pop(request_id, None) state = self.request_states.pop(request_id, None) if state is None: From 7b665f7a581b8cfb33974675174a9b0f29ec57e2 Mon Sep 17 00:00:00 2001 From: Yin Li Date: Tue, 6 Oct 2026 21:22:47 +0400 Subject: [PATCH 084/952] [Frontend] Integer token IDs for generate output logprobs (GenerateLogProbs) (#58181) Signed-off-by: Yin Li Co-authored-by: Yin Li Co-authored-by: Claude --- docs/serving/online_serving/derenderer.md | 4 +- docs/serving/online_serving/renderer.md | 54 ++++++ .../online_serving/token_in_token_out.md | 2 +- .../server/src/routes/inference/generate.rs | 146 ++++++++++++--- .../src/routes/inference/generate/convert.rs | 10 +- .../src/routes/inference/generate/types.rs | 32 +++- rust/src/server/src/routes/tests.rs | 20 +- .../scale_out/derender/test_derender.py | 43 +++-- .../derender/test_derender_parity.py | 44 ++++- .../derender/test_derender_stream.py | 38 ++-- .../test_generate_logprobs_conversion.py | 172 ++++++++++++++++++ .../test_generate_stream.py | 3 +- .../token_in_token_out/test_serving_tokens.py | 2 +- .../test_tokens_logprobs.py | 112 +++++++++++- vllm/entrypoints/generate/base/serving.py | 50 ++--- .../scale_out/token_in_token_out/protocol.py | 39 +++- .../scale_out/token_in_token_out/serving.py | 115 ++++++++++-- vllm/renderers/online_derenderer.py | 75 ++++---- 18 files changed, 773 insertions(+), 188 deletions(-) create mode 100644 tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index 9585031312fe..f85433a963ea 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -105,7 +105,7 @@ The server keeps no state between calls. Everything the next call needs is in `s - **Send `prompt_token_ids` on the first chunk, even without a parser.** It seeds detokenization so the first token's leading space matches the coupled endpoint (see [Request format](#request-format)). Later chunks already carry that context in `stream_state`. - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk. -Streaming derender resolves `logprobs` on the plain detokenization path, including `token_id:N` placeholders, same as the non-streaming endpoints (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them. +`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them. ## Streaming cost @@ -129,7 +129,7 @@ Streaming derender is stateless on the server side: all mutable state lives in t - `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings - `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0 -When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. +When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the integer `token_id`s are decoded per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. ## Example diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index fc6579c25807..79d0ade85b4f 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -79,6 +79,60 @@ for the existing authentication boundaries. For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md). +## Generate Output Logprobs + +`/inference/v1/generate` is a token in / token out API, so its output logprobs +identify tokens by integer ID rather than by the OpenAI string token. With +`sampling_params.logprobs` set and `output_mode: "tokens"` (the default), each +choice carries a `GenerateLogProbs`; `output_mode: "text"` returns decoded +`ChatCompletionLogProbs` instead (see [token in, token out](token_in_token_out.md)): + +```json +{ + "logprobs": { + "content": [ + { + "token_id": 262, + "logprob": -0.10, + "rank": 1, + "top_logprobs": [ + {"token_id": 262, "logprob": -0.10, "rank": 1}, + {"token_id": 257, "logprob": -1.20, "rank": 2} + ] + } + ] + } +} +``` + +- `content` has one entry per generated token, in generation order. +- `top_logprobs` is a list, not a dict: JSON turns dict keys into strings and + the ordering would be implicit. It follows the engine's order: the sampled + token first, then the remaining candidates in rank order. With non-greedy + sampling the sampled token can sit outside the top k (for example ranks + `[5, 1, 2]` at `logprobs=2`); it then takes one of the `logprobs` slots and + the rank-k candidate is left out, as on the OpenAI endpoints. +- `rank` is the token's rank in the vocabulary distribution (1 = most likely) + on every entry, the sampled one included; a top-k candidate's rank is its + top-k position. The list is not sorted by it, so sort by `rank` if you need + rank order. It is `null` when the engine could not rank the token (a NaN + logprob, which is sent as `-9999.0`). +- There is no `token` or `bytes` field. The generate server has no tokenizer; + [derender](derenderer.md) fills those in when it converts the response to the + OpenAI shapes. +- `prompt_logprobs` on the same response is unchanged + (`list[dict[int, Logprob] | None]`). + +!!! warning "Changed in this release" + Output logprobs used to be `ChatCompletionLogProbs` with every token written + as a `"token_id:N"` placeholder string, and `bytes` set to the UTF-8 bytes of + that placeholder by the Rust frontend but left unset by the Python one. + Clients that read only `content[i].logprob` are unaffected. Clients that + parsed the placeholder should read `content[i].token_id` instead. + `return_tokens_as_token_ids` on `/v1/chat/completions` and `/v1/completions` + is unchanged: it is a user-facing OpenAI option and still uses the + `token_id:N` format. + ## Multimodal Render Features Multimodal render responses include a `features` object with per-modality diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index b4e0d465bdd7..e22d4906155b 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -43,7 +43,7 @@ print(response["choices"][0]["token_ids"]) ### Logprobs -With `output_mode: "tokens"`, logprob entries carry `token_id:N` placeholders. With `output_mode: "text"`, they carry the decoded token strings and their UTF-8 `bytes` for the sampled token and every entry in `top_logprobs`. A server started with `--return-tokens-as-token-ids` always returns the placeholders, as `/v1/completions` does. +With `output_mode: "tokens"`, logprobs are `GenerateLogProbs` with an integer `token_id` and `rank` per entry and no `token` or `bytes` (see [Generate Output Logprobs](renderer.md#generate-output-logprobs)). With `output_mode: "text"`, they are `ChatCompletionLogProbs` carrying the decoded token strings and their UTF-8 `bytes` for the sampled token and every entry in `top_logprobs`; a server started with `--return-tokens-as-token-ids` returns `token_id:N` placeholders there instead, as `/v1/completions` does. ### Errors diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 0e682f7c45c2..99d7aec1b770 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -5,7 +5,7 @@ mod convert; mod types; mod validate; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::convert::Infallible; use std::result::Result; use std::sync::Arc; @@ -31,16 +31,17 @@ use vllm_llm::{ use self::convert::{ResponseOptions, prepare_generate_request}; use self::types::{ - GenerateLogprob, GenerateResponse, GenerateResponseChoice, GenerateResponseStreamChoice, - GenerateStreamResponse, MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo, - SpeculativeDecodingMetrics, StreamingSpeculativeDecodingMetrics, + GenerateLogProb, GenerateLogProbs, GenerateLogProbsContent, GenerateLogprob, GenerateResponse, + GenerateResponseChoice, GenerateResponseStreamChoice, GenerateStreamResponse, + MultiModalPlaceholders, PerRequestMetrics, PlaceholderRangeInfo, SpeculativeDecodingMetrics, + StreamingSpeculativeDecodingMetrics, }; pub(crate) use self::types::{GenerateRequest, GenerateSamplingParams}; pub(crate) use self::validate::validate_request_compat; use crate::config::ApiServerOptions; use crate::error::{ApiError, bail_server_error, server_error, text_submit_error}; use crate::routes::openai::utils::logprobs::clamp_logprob; -use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage}; +use crate::routes::openai::utils::types::Usage; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::{resolve_request_context, sse_response}; @@ -157,7 +158,7 @@ async fn generate_chunk_stream( ResponseOptions { include_usage, include_continuous_usage, - include_logprobs, + logprobs: output_logprobs, // Ignored: raw generate streaming has no prompt-logprobs wire shape. include_prompt_logprobs: _, return_token_ids, @@ -212,13 +213,15 @@ async fn generate_chunk_stream( continue; } - let logprobs = if include_logprobs && !token_ids.is_empty() { + let logprobs = if let Some(requested) = output_logprobs + && !token_ids.is_empty() + { let logprobs = output.logprobs.as_ref().ok_or_else(|| { server_error!( "raw generate stream requested logprobs but generation returned none" ) })?; - Some(raw_logprobs_to_openai_chat(logprobs)?) + Some(raw_logprobs_to_generate(logprobs, requested)?) } else { None }; @@ -282,19 +285,19 @@ fn collect_generate( include_usage: _, // Ignored: continuous usage is a streaming-only option. include_continuous_usage: _, - include_logprobs, + logprobs: output_logprobs, include_prompt_logprobs, return_token_ids, }: ResponseOptions, mm_placeholders: Option, ) -> Result { - let logprobs = if include_logprobs { + let logprobs = if let Some(requested) = output_logprobs { let logprobs = collected.logprobs.as_ref().ok_or_else(|| { ApiError::server_error( "raw generate response requested logprobs but generation returned none".to_string(), ) })?; - Some(raw_logprobs_to_openai_chat(logprobs)?) + Some(raw_logprobs_to_generate(logprobs, requested)?) } else { None }; @@ -371,14 +374,17 @@ fn extract_mm_placeholders(features: Option<&[MmFeatureSpec]>) -> Option Result { +fn raw_logprobs_to_generate( + logprobs: &Logprobs, + requested: i32, +) -> Result { let content = logprobs .positions .iter() - .map(position_to_chat_logprobs_content) + .map(|position| position_to_generate_logprobs_content(position, requested)) .collect::, _>>()?; - Ok(ChatLogProbs { + Ok(GenerateLogProbs { content: Some(content), }) } @@ -396,35 +402,47 @@ fn raw_prompt_logprobs_to_maps( .collect() } -fn position_to_chat_logprobs_content( +/// `top_logprobs` matches the Python generate server: the engine's entries (the +/// sampled token first, then ranks 1..k) with repeated token ids dropped, cut to +/// `max(requested, 1)` entries, or all of them for `-1`. +fn position_to_generate_logprobs_content( position: &PositionLogprobs, -) -> Result { + requested: i32, +) -> Result { let chosen = position.entries.first().ok_or_else(|| { ApiError::server_error( "raw generate logprobs position unexpectedly had no token candidates".to_string(), ) })?; - let token = format_token_id(chosen.token_id); - Ok(ChatLogProbsContent { - token: token.clone(), + // One pass: the first entry of each token id wins, as in Python's dict. The + // sampled token is the only one the engine can list twice. + let mut seen = HashSet::with_capacity(position.entries.len()); + Ok(GenerateLogProbsContent { + token_id: chosen.token_id, logprob: clamp_logprob(chosen.logprob), - bytes: Some(token.as_bytes().to_vec()), + rank: wire_rank(chosen.rank), top_logprobs: position .entries .iter() - .map(|entry| { - let token = format_token_id(entry.token_id); - TopLogProb { - token: token.clone(), - logprob: clamp_logprob(entry.logprob), - bytes: Some(token.into_bytes()), - } + .filter(|entry| seen.insert(entry.token_id)) + .take(usize::try_from(requested).map_or(usize::MAX, |n| n.max(1))) + .map(|entry| GenerateLogProb { + token_id: entry.token_id, + logprob: clamp_logprob(entry.logprob), + rank: wire_rank(entry.rank), }) .collect(), }) } +/// The engine reports rank 0 for a sampled token whose logprob is NaN (no value +/// compares >= NaN). That is not a rank, so it goes on the wire as `None`, the +/// way `clamp_logprob` handles the logprob itself. +fn wire_rank(rank: u32) -> Option { + (rank > 0).then_some(rank) +} + fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap { position .entries @@ -523,6 +541,80 @@ mod tests { use super::*; + fn position(entries: &[(u32, f32, u32)]) -> PositionLogprobs { + PositionLogprobs { + entries: entries + .iter() + .map(|&(token_id, logprob, rank)| { + vllm_engine_core_client::protocol::logprobs::TokenLogprob { + token_id, + logprob, + rank, + } + }) + .collect(), + } + } + + fn top(position: &PositionLogprobs, requested: i32) -> Vec<(u32, Option)> { + position_to_generate_logprobs_content(position, requested) + .unwrap() + .top_logprobs + .iter() + .map(|entry| (entry.token_id, entry.rank)) + .collect() + } + + #[test] + fn generate_top_logprobs_drop_the_repeated_sampled_token() { + // Sampled token 7 is also top-1: the engine row is [7, 7, 8]. + let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2)]); + assert_eq!(top(&pos, 2), vec![(7, Some(1)), (8, Some(2))]); + } + + #[test] + fn generate_top_logprobs_keep_sampled_first_when_outside_top_k() { + // Sampled token 50 has vocab rank 5; it takes a slot and rank 2 drops + // out, as on the Python generate server and the OpenAI endpoints. + let pos = position(&[(50, -3.0, 5), (10, -0.2, 1), (20, -1.0, 2)]); + assert_eq!(top(&pos, 2), vec![(50, Some(5)), (10, Some(1))]); + } + + #[test] + fn generate_top_logprobs_respect_zero_and_all() { + let pos = position(&[(7, -0.1, 1), (7, -0.1, 1), (8, -2.0, 2), (9, -3.0, 3)]); + assert_eq!(top(&pos, 0), vec![(7, Some(1))]); + assert_eq!( + top(&pos, -1), + vec![(7, Some(1)), (8, Some(2)), (9, Some(3))] + ); + } + + #[test] + fn generate_top_logprobs_full_vocab_is_linear() { + // logprobs = -1 returns the whole vocabulary; deduping it must not be + // quadratic in the number of entries. + const VOCAB: u32 = 200_000; + let mut entries = vec![(5, -0.5, 1)]; + entries.extend((0..VOCAB).map(|t| (t, -1.0 - t as f32, t + 1))); + let pos = position(&entries); + let start = std::time::Instant::now(); + let got = top(&pos, -1); + assert!(start.elapsed() < std::time::Duration::from_secs(5)); + assert_eq!(got.len(), VOCAB as usize); + assert_eq!(got[0], (5, Some(1))); + } + + #[test] + fn generate_rank_zero_from_a_nan_logprob_is_none() { + // The engine reports rank 0 for a sampled token whose logprob is NaN. + let pos = position(&[(7, f32::NAN, 0), (8, -0.2, 1)]); + let content = position_to_generate_logprobs_content(&pos, 1).unwrap(); + assert_eq!(content.rank, None); + assert_eq!(content.logprob, -9999.0); + assert_eq!(top(&pos, 1), vec![(7, None)]); + } + #[tokio::test] async fn generate_chunk_stream_captures_late_prompt_info() { let stream = stream::iter(vec![ diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 4c54619ed4a7..36a497f6ddf3 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -26,8 +26,10 @@ pub(super) struct ResponseOptions { pub include_usage: bool, /// Whether the caller asked for usage on every streamed chunk. pub include_continuous_usage: bool, - /// Whether the caller requested output logprobs on generate choices. - pub include_logprobs: bool, + /// Requested output logprobs (`sampling_params.logprobs`): `None` omits + /// them, `Some(n)` returns `max(n, 1)` candidates per position, `Some(-1)` + /// all of them. + pub logprobs: Option, /// Whether the caller requested top-level prompt logprobs. pub include_prompt_logprobs: bool, /// Whether the caller requested final prompt token metadata. @@ -56,7 +58,7 @@ pub(super) fn prepare_generate_request( .as_ref() .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); - let include_logprobs = request.sampling_params.inner.logprobs.is_some(); + let logprobs = request.sampling_params.inner.logprobs; let include_prompt_logprobs = request.sampling_params.inner.prompt_logprobs.is_some(); let return_token_ids = request.return_token_ids.unwrap_or(false); let mut sampling_params = request.sampling_params.inner; @@ -96,7 +98,7 @@ pub(super) fn prepare_generate_request( options: ResponseOptions { include_usage, include_continuous_usage, - include_logprobs, + logprobs, include_prompt_logprobs, return_token_ids, }, diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index b0b93d7080a6..bd277ef34a0a 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -10,7 +10,7 @@ use validator::Validate; use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_text::SamplingParams; -use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage}; +use crate::routes::openai::utils::types::{Normalizable, StreamOptions, Usage}; /// Sampling parameters for the token-in/token-out generate API. /// @@ -61,7 +61,7 @@ impl Normalizable for GenerateRequest {} #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateResponseChoice { pub index: u32, - pub logprobs: Option, + pub logprobs: Option, pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, @@ -72,7 +72,7 @@ pub(super) struct GenerateResponseChoice { #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateResponseStreamChoice { pub index: u32, - pub logprobs: Option, + pub logprobs: Option, pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, @@ -197,6 +197,32 @@ pub(super) struct PlaceholderRangeInfo { pub length: usize, } +/// Mirrors the Python vLLM `GenerateLogProbs` class: output logprobs for one +/// choice, carrying integer token ids rather than the OpenAI string token. +#[derive(Debug, Clone, Serialize)] +pub(super) struct GenerateLogProbs { + pub content: Option>, +} + +/// Mirrors the Python vLLM `GenerateLogProbsContent` class: the sampled token +/// at one position plus its top-k candidates, in the engine's order: the +/// sampled token first, then the remaining candidates in rank order. +#[derive(Debug, Clone, Serialize)] +pub(super) struct GenerateLogProbsContent { + pub token_id: u32, + pub logprob: f32, + pub rank: Option, + pub top_logprobs: Vec, +} + +/// Mirrors the Python vLLM `GenerateLogProb` class: one candidate token. +#[derive(Debug, Clone, Serialize)] +pub(super) struct GenerateLogProb { + pub token_id: u32, + pub logprob: f32, + pub rank: Option, +} + /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads. #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateLogprob { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 31c328451b15..0ed2ea4d4590 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -4453,13 +4453,15 @@ async fn non_stream_raw_generate_returns_token_output_envelope() { assert_eq!(json["choices"][0]["index"], 0); assert_eq!(json["choices"][0]["token_ids"], json!([33, 44])); assert_eq!(json["choices"][0]["finish_reason"], "stop"); + // Output logprobs carry integer token ids and rank, not the OpenAI + // "token_id:N" placeholder string; `bytes` is gone. + assert_eq!(json["choices"][0]["logprobs"]["content"][0]["token_id"], 33); + assert_eq!(json["choices"][0]["logprobs"]["content"][0]["rank"], 1); + assert!(json["choices"][0]["logprobs"]["content"][0]["token"].is_null()); + assert!(json["choices"][0]["logprobs"]["content"][0]["bytes"].is_null()); assert_eq!( - json["choices"][0]["logprobs"]["content"][0]["token"], - "token_id:33" - ); - assert_eq!( - json["choices"][0]["logprobs"]["content"][1]["top_logprobs"][0]["token"], - "token_id:44" + json["choices"][0]["logprobs"]["content"][1]["top_logprobs"][0]["token_id"], + 44 ); assert_eq!(json["prompt_logprobs"][0], serde_json::Value::Null); assert_eq!( @@ -4581,9 +4583,11 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() { assert_eq!(first["choices"][0]["index"], 0); assert_eq!(first["choices"][0]["token_ids"], json!([33])); assert_eq!( - first["choices"][0]["logprobs"]["content"][0]["token"], - "token_id:33" + first["choices"][0]["logprobs"]["content"][0]["token_id"], + 33 ); + assert_eq!(first["choices"][0]["logprobs"]["content"][0]["rank"], 1); + assert!(first["choices"][0]["logprobs"]["content"][0]["token"].is_null()); assert_eq!(first["usage"]["prompt_tokens"], 2); assert_eq!(first["usage"]["completion_tokens"], 1); diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 40b051304e22..d29aa3f53f73 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -74,14 +74,13 @@ def _make_generate_response( } -def _make_logprobs_with_placeholders(token_id: int = 1234) -> dict: +def _make_generate_logprobs(token_id: int = 1234) -> dict: + """A `GenerateLogProbs` payload: integer token ids, no token/bytes.""" entry = { - "token": f"token_id:{token_id}", + "token_id": token_id, "logprob": -1.0, - "bytes": None, - "top_logprobs": [ - {"token": f"token_id:{token_id + 1}", "logprob": -2.0, "bytes": None} - ], + "rank": 1, + "top_logprobs": [{"token_id": token_id + 1, "logprob": -2.0, "rank": 2}], } return {"content": [entry]} @@ -154,7 +153,7 @@ async def test_derender_chat_usage_default(client): @pytest.mark.asyncio async def test_derender_chat_logprobs(client): - """token_id:N placeholders in content.token are resolved to real strings.""" + """Integer token ids are decoded to real strings by derender.""" gen_req = await _render_chat(client) synthetic_ids = gen_req["token_ids"][:3] token_id = synthetic_ids[0] @@ -165,7 +164,7 @@ async def test_derender_chat_logprobs(client): "model": MODEL_NAME, "generate_response": _make_generate_response( synthetic_ids, - logprobs=_make_logprobs_with_placeholders(token_id), + logprobs=_make_generate_logprobs(token_id), ), }, ) @@ -176,8 +175,8 @@ async def test_derender_chat_logprobs(client): content = logprobs["content"] assert content is not None and len(content) == 1 token_str = content[0]["token"] - assert not token_str.startswith("token_id:"), ( - f"Placeholder was not resolved: {token_str!r}" + assert token_str and not token_str.startswith("token_id:"), ( + f"token id was not decoded: {token_str!r}" ) @@ -194,7 +193,7 @@ async def test_derender_chat_logprobs_bytes(client): "model": MODEL_NAME, "generate_response": _make_generate_response( synthetic_ids, - logprobs=_make_logprobs_with_placeholders(token_id), + logprobs=_make_generate_logprobs(token_id), ), }, ) @@ -208,7 +207,7 @@ async def test_derender_chat_logprobs_bytes(client): @pytest.mark.asyncio async def test_derender_chat_top_logprobs(client): - """top_logprobs entries also have their placeholders resolved.""" + """top_logprobs entries are decoded too.""" gen_req = await _render_chat(client) synthetic_ids = gen_req["token_ids"][:3] token_id = synthetic_ids[0] @@ -219,7 +218,7 @@ async def test_derender_chat_top_logprobs(client): "model": MODEL_NAME, "generate_response": _make_generate_response( synthetic_ids, - logprobs=_make_logprobs_with_placeholders(token_id), + logprobs=_make_generate_logprobs(token_id), ), }, ) @@ -228,7 +227,7 @@ async def test_derender_chat_top_logprobs(client): top = content[0]["top_logprobs"] assert len(top) == 1 assert not top[0]["token"].startswith("token_id:"), ( - f"top_logprobs placeholder not resolved: {top[0]['token']!r}" + f"top_logprobs token id not decoded: {top[0]['token']!r}" ) @@ -633,8 +632,8 @@ async def test_derender_completion_empty_generate_responses(client): @pytest.mark.asyncio async def test_derender_completion_logprobs(client): - """token_id:N placeholders in logprobs are resolved; CompletionLogProbs - flat-list structure is returned with non-empty tokens and text_offsets.""" + """Integer token ids are decoded; CompletionLogProbs flat-list structure is + returned with non-empty tokens and text_offsets.""" gr1 = await _render_completion(client, "Hello world") ids1 = gr1["token_ids"][:3] token_id = ids1[0] @@ -647,7 +646,7 @@ async def test_derender_completion_logprobs(client): _make_completion_generate_response( ids1, gr1["request_id"], - logprobs=_make_logprobs_with_placeholders(token_id), + logprobs=_make_generate_logprobs(token_id), ), ], }, @@ -658,7 +657,7 @@ async def test_derender_completion_logprobs(client): tokens = logprobs["tokens"] assert len(tokens) == 1 assert not tokens[0].startswith("token_id:"), ( - f"Placeholder was not resolved: {tokens[0]!r}" + f"token id was not decoded: {tokens[0]!r}" ) assert len(logprobs["token_logprobs"]) == 1 assert isinstance(logprobs["token_logprobs"][0], float) @@ -794,7 +793,7 @@ async def test_derender_chat_oversized_logprobs_rejected(client): """logprobs.content longer than max_model_len returns 400.""" oversized_logprobs: dict = { "content": [ - {"token": "x", "logprob": -1.0, "bytes": None, "top_logprobs": []} + {"token_id": 42, "logprob": -1.0, "rank": 1, "top_logprobs": []} for _ in range(1_000_000) ] } @@ -825,11 +824,11 @@ async def test_derender_chat_oversized_top_logprobs_rejected(client): oversized_top_logprobs = { "content": [ { - "token": "x", + "token_id": 42, "logprob": -1.0, - "bytes": None, + "rank": 1, "top_logprobs": [ - {"token": f"t{i}", "logprob": -float(i), "bytes": None} + {"token_id": i, "logprob": -float(i), "rank": i + 1} for i in range(25) ], } diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 66d3e8396939..4ea3bc246192 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -237,21 +237,47 @@ async def test_parity_reasoning_and_tool_call(client): assert _tool_sig(disagg["choices"][0]) +def _to_generate_logprobs(openai_logprobs: dict) -> dict: + """Turn a coupled response's `token_id:N` logprobs into `GenerateLogProbs`. + + A real GPU-less worker emits integer ids directly; here the coupled server + with `return_tokens_as_token_ids=True` is the only source of the same ids, + so the placeholders are parsed back out to build the generate-shaped + payload that `/derender` now consumes. + """ + + def _id(token: str) -> int: + assert token.startswith("token_id:"), token + return int(token.removeprefix("token_id:")) + + return { + "content": [ + { + "token_id": _id(entry["token"]), + "logprob": entry["logprob"], + "top_logprobs": [ + {"token_id": _id(top["token"]), "logprob": top["logprob"]} + for top in entry["top_logprobs"] + ], + } + for entry in openai_logprobs["content"] + ] + } + + @pytest.mark.asyncio async def test_parity_logprobs(client): - """token_id:N resolution parity vs. the coupled server's real strings. + """Integer-id logprob decoding parity vs. the coupled server's strings. - A real disaggregated worker only has token IDs so it emits logprobs - with `token_id:N` placeholders (`return_tokens_as_token_ids=True` - reproduces that shape here). `/derender` must resolve those - placeholders to the same token strings/bytes the coupled server - resolves them to directly. + A real disaggregated worker only has token IDs, so it emits + `GenerateLogProbs` with integer `token_id`s. `/derender` must decode those + to the same token strings/bytes the coupled server produces directly. """ messages = [{"role": "user", "content": "What is 2+2?"}] extra = {"logprobs": True, "top_logprobs": 3} - # What a real GPU less worker would hand to /derender is token IDs plus - # logprobs still in token_id:N placeholder form + # The coupled server with placeholders is used only as a source of the + # token ids a GPU-less worker would have had. placeholder = await _coupled( client, messages, return_tokens_as_token_ids=True, **extra ) @@ -263,7 +289,7 @@ async def test_parity_logprobs(client): placeholder["usage"]["prompt_tokens"], ch["finish_reason"], chat_request, - logprobs=ch["logprobs"], + logprobs=_to_generate_logprobs(ch["logprobs"]), ) # The coupled server resolving the same greedy generation diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 0e4b168677ba..0a7d96f3cf90 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -64,6 +64,7 @@ _seed_stream_state, ) from vllm.tokenizers import get_tokenizer +from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens from vllm.utils import random_uuid MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" @@ -212,17 +213,12 @@ def _make_stream_chunk( ) -def _placeholder_logprobs(token_ids: list[int]) -> dict: - """Per-token logprob entries using token_id:N placeholders, as sent by - the generate worker.""" +def _generate_logprobs(token_ids: list[int]) -> dict: + """Per-token `GenerateLogProbs` entries (integer token ids, no token or + bytes), as sent by the generate worker.""" return { "content": [ - { - "token": f"token_id:{tid}", - "logprob": -0.5, - "bytes": None, - "top_logprobs": [], - } + {"token_id": tid, "logprob": -0.5, "rank": 1, "top_logprobs": []} for tid in token_ids ] } @@ -892,12 +888,12 @@ async def test_finish_reason_forwarded(self, derenderer, tokenizer): class TestStreamLogprobs: - """Streaming derender must carry per-chunk logprobs with placeholders - resolved, matching what the generate streaming path emits.""" + """Streaming derender must carry per-chunk logprobs with token ids + decoded, matching what the generate streaming path emits.""" @pytest.mark.asyncio async def test_chat_stream_logprobs_resolved_per_chunk(self, derenderer, tokenizer): - """Each streamed chunk carries logprobs with token_id:N resolved.""" + """Each streamed chunk carries logprobs with its token ids decoded.""" token_ids = tokenizer.encode("hello world")[:6] mid = len(token_ids) // 2 @@ -906,17 +902,15 @@ async def test_chat_stream_logprobs_resolved_per_chunk(self, derenderer, tokeniz chunk, state = await derenderer.derender_chat_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( - part, logprobs=_placeholder_logprobs(part) + part, logprobs=_generate_logprobs(part) ), state=state, ) logprobs = chunk.choices[0].logprobs assert logprobs is not None and logprobs.content is not None - assert len(logprobs.content) == len(part) - for entry in logprobs.content: - assert not entry.token.startswith("token_id:"), ( - f"placeholder not resolved: {entry.token!r}" - ) + assert [entry.token for entry in logprobs.content] == ( + convert_ids_list_to_tokens(tokenizer, part) + ) @pytest.mark.asyncio async def test_chat_stream_logprobs_multibyte_across_chunks( @@ -936,7 +930,7 @@ async def test_chat_stream_logprobs_multibyte_across_chunks( last_chunk, state = await derenderer.derender_chat_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( - [tid], logprobs=_placeholder_logprobs([tid]) + [tid], logprobs=_generate_logprobs([tid]) ), state=state, ) @@ -959,13 +953,13 @@ async def test_completion_stream_logprobs_text_offset_absolute( chunk1, state = await derenderer.derender_completion_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( - token_ids[:mid], logprobs=_placeholder_logprobs(token_ids[:mid]) + token_ids[:mid], logprobs=_generate_logprobs(token_ids[:mid]) ), ) chunk2, _ = await derenderer.derender_completion_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( - token_ids[mid:], logprobs=_placeholder_logprobs(token_ids[mid:]) + token_ids[mid:], logprobs=_generate_logprobs(token_ids[mid:]) ), state=state, ) @@ -982,7 +976,7 @@ async def test_stream_logprob_state_stays_bounded(self, derenderer, tokenizer): _, state = await derenderer.derender_chat_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( - [tid], logprobs=_placeholder_logprobs([tid]) + [tid], logprobs=_generate_logprobs([tid]) ), state=state, ) diff --git a/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py b/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py new file mode 100644 index 000000000000..96d937ad5ae4 --- /dev/null +++ b/tests/entrypoints/scale_out/derender/test_generate_logprobs_conversion.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the generate -> OpenAI logprob conversion in derender. + +These exercise `_resolve_logprobs` directly. A stub tokenizer covers the +byte-fallback correction path; a real SentencePiece tokenizer covers the +leading-space restoration. +""" + +from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + GenerateLogProb, + GenerateLogProbs, + GenerateLogProbsContent, +) +from vllm.renderers.online_derenderer import ( + _convert_chat_logprobs_to_completion_logprobs, + _resolve_logprobs, +) +from vllm.tokenizers import get_tokenizer +from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens + +# A two-byte character split across two byte-fallback tokens: neither decodes +# to anything printable on its own, but the pair decodes to "é". +FIRST_BYTE = 10 +SECOND_BYTE = 11 +PLAIN = 12 + +VOCAB = {FIRST_BYTE: "<0xC3>", SECOND_BYTE: "<0xA9>", PLAIN: "ok"} + + +class _StubTokenizer: + """Decodes byte-fallback tokens only when both halves are present.""" + + def convert_ids_to_tokens(self, ids: list[int]) -> list[str | None]: + return [VOCAB.get(i) for i in ids] + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return self.decode([k for t in tokens for k, v in VOCAB.items() if v == t]) + + def decode(self, ids: list[int], **kwargs) -> str: + out = "" + i = 0 + while i < len(ids): + if ids[i] == FIRST_BYTE and i + 1 < len(ids) and ids[i + 1] == SECOND_BYTE: + out += "é" + i += 2 + elif ids[i] in (FIRST_BYTE, SECOND_BYTE): + out += "�" + i += 1 + else: + out += VOCAB.get(ids[i], "") + i += 1 + return out + + +def test_token_ids_are_decoded_with_bytes_and_order_preserved(): + logprobs = GenerateLogProbs( + content=[ + GenerateLogProbsContent( + token_id=PLAIN, + logprob=-0.5, + rank=1, + top_logprobs=[ + GenerateLogProb(token_id=PLAIN, logprob=-0.5, rank=1), + GenerateLogProb(token_id=FIRST_BYTE, logprob=-2.0, rank=2), + ], + ) + ] + ) + + resolved = _resolve_logprobs(logprobs, _StubTokenizer()) + + assert resolved.content is not None + entry = resolved.content[0] + assert entry.token == "ok" + assert entry.logprob == -0.5 + assert entry.bytes == list(b"ok") + # Rank order survives the conversion (the OpenAI shape has no rank field). + assert [t.logprob for t in entry.top_logprobs] == [-0.5, -2.0] + assert entry.top_logprobs[0].token == "ok" + + +def test_byte_fallback_second_half_uses_preceding_token_as_context(): + """U+FFFD correction needs the preceding sampled ids, which the integer + shape carries directly (no placeholder parsing).""" + logprobs = GenerateLogProbs( + content=[ + GenerateLogProbsContent(token_id=FIRST_BYTE, logprob=-0.1), + GenerateLogProbsContent(token_id=SECOND_BYTE, logprob=-0.2), + ] + ) + + resolved = _resolve_logprobs(logprobs, _StubTokenizer()) + + assert resolved.content is not None + # The first half has no preceding context to repair it with, so it stays + # empty (existing `_correct_decoded_token` behaviour). The second resolves + # to the complete character once the first is used as context. + assert resolved.content[0].token == "" + assert resolved.content[0].bytes == [] + assert resolved.content[1].token == "é" + assert resolved.content[1].bytes == list("é".encode()) + + +def test_content_none_is_not_an_error(): + resolved = _resolve_logprobs(GenerateLogProbs(), _StubTokenizer()) + assert resolved.content is None + + +def test_sentencepiece_leading_space_matches_engine(): + """`▁true` and `true` are distinct tokens and must stay distinct: the + strings match the engine's per-token detokenization (leading space + restored), so `/v1/completions/derender` top_logprobs keys do not collide. + Regression test for #59241.""" + tokenizer = get_tokenizer("hmellor/tiny-random-LlamaForCausalLM") + spaced = tokenizer.convert_tokens_to_ids("▁true") + bare = tokenizer.convert_tokens_to_ids("true") + assert spaced != bare + logprobs = GenerateLogProbs( + content=[ + GenerateLogProbsContent( + token_id=spaced, + logprob=-0.5, + rank=1, + top_logprobs=[ + GenerateLogProb(token_id=spaced, logprob=-0.5, rank=1), + GenerateLogProb(token_id=bare, logprob=-1.5, rank=2), + ], + ) + ] + ) + + resolved = _resolve_logprobs(logprobs, tokenizer) + + assert resolved.content is not None + entry = resolved.content[0] + assert entry.token == " true" + assert entry.bytes == list(b" true") + assert [t.token for t in entry.top_logprobs] == convert_ids_list_to_tokens( + tokenizer, [spaced, bare] + ) + completion = _convert_chat_logprobs_to_completion_logprobs(resolved) + assert completion.top_logprobs == [{" true": -0.5, "true": -1.5}] + + +def test_unknown_top_id_does_not_shift_the_other_candidates(): + """Ids are decoded in one batch per position; an id with no vocab entry + becomes an empty token without shifting its neighbours.""" + unknown = 99 + logprobs = GenerateLogProbs( + content=[ + GenerateLogProbsContent( + token_id=PLAIN, + logprob=-0.5, + rank=1, + top_logprobs=[ + GenerateLogProb(token_id=unknown, logprob=-1.0, rank=1), + GenerateLogProb(token_id=PLAIN, logprob=-0.5, rank=2), + ], + ) + ] + ) + + resolved = _resolve_logprobs(logprobs, _StubTokenizer()) + + assert resolved.content is not None + entry = resolved.content[0] + assert entry.token == "ok" + assert [(t.token, t.bytes) for t in entry.top_logprobs] == [ + ("", None), + ("ok", list(b"ok")), + ] diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 7e7f74de9f12..7f0cd1286041 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -1123,7 +1123,8 @@ async def mock_generate(*args, **kwargs): lp = dc["choices"][0]["logprobs"] assert lp is not None assert len(lp["content"]) == 1 - assert lp["content"][0]["token"].startswith("token_id:") + assert isinstance(lp["content"][0]["token_id"], int) + assert "token" not in lp["content"][0] @pytest.mark.asyncio diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 1027dcef035c..0919431abac6 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -183,7 +183,7 @@ async def test_generate_sampling_mask(client): assert token_id in support # processed_logprobs: exactly the support carries finite probability. for top in entry["top_logprobs"]: - in_support = int(top["token"].removeprefix("token_id:")) in support + in_support = top["token_id"] in support assert in_support == (top["logprob"] > -9999.0) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py index c417a5fc2b57..05dba167fe6e 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py @@ -1,22 +1,62 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from unittest.mock import MagicMock +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateLogProb from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.logprobs import Logprob def test_top_logprobs_alternatives_have_own_token_ids(): - """Each top_logprobs alternative must carry its own token_id placeholder.""" + """Each top_logprobs alternative must carry its own integer token id.""" result = ServingTokens._create_tokens_logprobs( None, token_ids=[262], top_logprobs=[{262: Logprob(-0.1), 257: Logprob(-1.2), 428: Logprob(-2.3)}], num_output_top_logprobs=2, ) - tokens = {e.token for e in result.content[0].top_logprobs} - assert tokens == {"token_id:262", "token_id:257"}, f"got {tokens}" + token_ids = {e.token_id for e in result.content[0].top_logprobs} + assert token_ids == {262, 257}, f"got {token_ids}" + + +def test_sampled_entry_carries_token_id_and_rank(): + """The sampled token is identified by id, and the engine's rank survives.""" + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[262], + top_logprobs=[ + {262: Logprob(-0.1, rank=1), 257: Logprob(-1.2, rank=2)}, + ], + num_output_top_logprobs=2, + ) + entry = result.content[0] + assert entry.token_id == 262 + assert entry.logprob == -0.1 + assert entry.rank == 1 + assert [(t.token_id, t.rank) for t in entry.top_logprobs] == [(262, 1), (257, 2)] + # No tokenizer on the generate server: no string token, no bytes. + assert not hasattr(entry, "token") + assert not hasattr(entry, "bytes") + + +def test_sampled_token_absent_from_topk_uses_sentinel(): + """A sampled token missing from the engine's map keeps the -9999.0 sentinel.""" + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[5], + top_logprobs=[{7: Logprob(-0.9)}], + num_output_top_logprobs=1, + ) + entry = result.content[0] + assert entry.token_id == 5 + assert entry.logprob == -9999.0 + assert entry.rank is None + assert entry.top_logprobs == [] def test_logprobs_zero_emits_sampled_token(): @@ -40,21 +80,23 @@ def test_logprobs_minus_one_emits_all_tokens(): assert len(result.content[0].top_logprobs) == 2 -def test_tokens_level_logprobs_have_no_bytes(): - result = ServingTokens._create_tokens_logprobs( +def test_text_logprobs_without_tokenizer_are_placeholders_without_bytes(): + """--return-tokens-as-token-ids keeps `token_id:N` placeholders in text mode.""" + result = ServingTokens._create_text_logprobs( None, token_ids=[7], top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}], num_output_top_logprobs=2, ) entry = result.content[0] + assert entry.token == "token_id:7" assert entry.bytes is None assert all(t.bytes is None for t in entry.top_logprobs) def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes(): tokenizer = MagicMock() - result = ServingTokens._create_tokens_logprobs( + result = ServingTokens._create_text_logprobs( None, token_ids=[7], top_logprobs=[ @@ -75,7 +117,7 @@ def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes(): def test_resolved_logprobs_decode_tokens_the_engine_did_not(): tokenizer = MagicMock() tokenizer.decode.side_effect = lambda ids: f"<{ids[0]}>" - result = ServingTokens._create_tokens_logprobs( + result = ServingTokens._create_text_logprobs( None, token_ids=[7, 9], top_logprobs=[{7: Logprob(-0.9)}, None], @@ -85,3 +127,59 @@ def test_resolved_logprobs_decode_tokens_the_engine_did_not(): assert [e.token for e in result.content] == ["<7>", "<9>"] assert result.content[0].top_logprobs[0].token == "<7>" assert result.content[1].bytes == list(b"<9>") + + +def test_sampled_token_outside_topk_comes_first(): + """The engine puts the sampled token first, then ranks 1..k. When the + sampled token is outside the top k it takes one of the k slots, so rank k + is left out (same as the OpenAI endpoints).""" + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[50], + top_logprobs=[ + { + 50: Logprob(-3.0, rank=5), + 10: Logprob(-0.2, rank=1), + 20: Logprob(-1.0, rank=2), + } + ], + num_output_top_logprobs=2, + ) + assert [t.rank for t in result.content[0].top_logprobs] == [5, 1] + + +def test_logprob_is_required_on_the_wire(): + """A payload missing logprob is rejected rather than read as -9999.""" + with pytest.raises(ValidationError): + GenerateLogProb.model_validate({"token_id": 1}) + + +def test_rank_zero_from_a_nan_logprob_is_none(): + """The engine reports rank 0 for a sampled token whose logprob is NaN.""" + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[7], + top_logprobs=[{7: Logprob(float("nan"), rank=0), 8: Logprob(-0.2, rank=1)}], + num_output_top_logprobs=1, + ) + entry = result.content[0] + assert entry.rank is None + assert entry.top_logprobs[0].rank is None + # Clamped like the Rust frontend (max(nan, x) is nan in Python). + assert entry.logprob == -9999.0 + assert entry.top_logprobs[0].logprob == -9999.0 + json.dumps(result.model_dump(), allow_nan=False) + + +def test_text_logprobs_clamp_a_nan_logprob(): + result = ServingTokens._create_text_logprobs( + None, + token_ids=[7], + top_logprobs=[{7: Logprob(float("nan"), rank=0, decoded_token="a")}], + num_output_top_logprobs=1, + tokenizer=MagicMock(), + ) + entry = result.content[0] + assert entry.logprob == -9999.0 + assert entry.top_logprobs[0].logprob == -9999.0 + json.dumps(result.model_dump(), allow_nan=False) diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 5142346df661..d1f3e94b006f 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -32,6 +32,7 @@ from vllm.logprobs import Logprob, PromptLogprobs from vllm.lora.request import LoRARequest from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens from vllm.tracing import ( contains_trace_headers, extract_trace_headers, @@ -346,32 +347,33 @@ def format_token_id_placeholder(token_id: int) -> str: return f"token_id:{token_id}" -def resolve_token_id_placeholder( - token: str, tokenizer: TokenizerLike -) -> tuple[str, list[int] | None]: - """Decode a 'token_id:N' placeholder back to a token string and UTF-8 bytes. +def decode_token_ids( + token_ids: list[int], tokenizer: TokenizerLike +) -> list[tuple[str, list[int] | None]]: + """Decode token ids individually to their token strings and UTF-8 bytes. - Returns (token, None) unchanged if token is not a placeholder. - This is the inverse of format_token_id_placeholder / _get_decoded_token - when return_as_token_id=True. + Uses the engine's per-token detokenization, which restores the + SentencePiece leading space that `convert_tokens_to_string` drops, so the + strings match the coupled endpoints. Ids are decoded in one batch (callers + pass a position's sampled id together with its top-k ids). An id with no + vocab entry decodes to ("", None). """ - suffix = token.removeprefix("token_id:") - if suffix == token: - return token, None - try: - token_id = int(suffix) - except ValueError: - return token, None - token_repr = tokenizer.convert_ids_to_tokens([token_id])[0] - if token_repr is None: - logger.warning_once( - "resolve_token_id_placeholder: token_id %d has no vocab entry; " - "substituting empty string", - token_id, - ) - return "", None - token_str = tokenizer.convert_tokens_to_string([token_repr]) - return token_str, list(token_str.encode("utf-8", errors="replace")) + pieces = tokenizer.convert_ids_to_tokens(token_ids) + known = [tid for tid, piece in zip(token_ids, pieces) if piece is not None] + decoded = iter(convert_ids_list_to_tokens(tokenizer, known)) + out: list[tuple[str, list[int] | None]] = [] + for tid, piece in zip(token_ids, pieces): + if piece is None: + logger.warning_once( + "decode_token_ids: token_id %d has no vocab entry; " + "substituting empty string", + tid, + ) + out.append(("", None)) + continue + token_str = next(decoded) + out.append((token_str, list(token_str.encode("utf-8", errors="replace")))) + return out def clamp_prompt_logprobs( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 4e29f7e59dfa..84a6a2c62388 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -360,6 +360,39 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: ) +class GenerateLogProb(BaseModel): + """A single (token, logprob) candidate on the generate wire protocol. + + Unlike the OpenAI logprob shapes this carries the integer token id: the + generate server has no tokenizer, so decoding to a string belongs in + derender (or the coupled chat/completions path), not here. + """ + + token_id: int + logprob: float + rank: int | None = None + + +class GenerateLogProbsContent(GenerateLogProb): + """The sampled token at one position, plus its top-k candidates. + + ``top_logprobs`` is a list, not a dict: JSON turns dict keys into strings + and the order would be implicit. It is in the engine's order: the sampled + token first, then the remaining candidates in rank order. + """ + + top_logprobs: list[GenerateLogProb] = [] + + +class GenerateLogProbs(BaseModel): + """Output logprobs for one choice. + + ``content`` holds one entry per generated token. + """ + + content: list[GenerateLogProbsContent] | None = None + + class GenerateChoiceBase(BaseModel): """Fields shared by every `output_mode` of a non-streaming choice.""" @@ -388,8 +421,8 @@ def validate_token_ids(cls, v: list[int] | None) -> list[int] | None: class GenerateTokensChoice(GenerateChoiceBase): - logprobs: ChatCompletionLogProbs | None = None - """Logprobs whose tokens are `token_id:N` placeholders.""" + logprobs: GenerateLogProbs | None = None + """Logprobs with integer token ids (no tokenizer on the generate server).""" class GenerateTextChoice(GenerateChoiceBase): @@ -413,7 +446,7 @@ class GenerateStreamChoiceBase(BaseModel): class GenerateTokensStreamChoice(GenerateStreamChoiceBase): - logprobs: ChatCompletionLogProbs | None = None + logprobs: GenerateLogProbs | None = None class GenerateTextStreamChoice(GenerateStreamChoiceBase): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index e7afa0401b29..b6d7bf19ba93 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -3,6 +3,7 @@ import asyncio +import math import time from collections.abc import AsyncGenerator from collections.abc import Sequence as GenericSequence @@ -57,6 +58,9 @@ placeholder_ranges_from_engine_input, ) from .protocol import ( + GenerateLogProb, + GenerateLogProbs, + GenerateLogProbsContent, GenerateRequest, GenerateResponse, GenerateTextChoice, @@ -72,6 +76,15 @@ logger = init_logger(__name__) +def _clamp_logprob(logprob: float) -> float: + """The OpenAI shapes' floor for a logprob: ``-inf`` and NaN become ``-9999.0``. + + ``max(nan, -9999.0)`` is NaN in Python, so NaN needs its own check; NaN + would otherwise fail ``JSONResponse`` or reach the client as ``null``. + """ + return -9999.0 if math.isnan(logprob) else max(logprob, -9999.0) + + def _logprob_token( token_id: int, logprob: Logprob | None, tokenizer: TokenizerLike | None ) -> tuple[str, list[int] | None]: @@ -94,7 +107,7 @@ def _top_logprob( ) -> ChatCompletionLogProb: token, token_bytes = _logprob_token(token_id, logprob, tokenizer) return ChatCompletionLogProb( - token=token, logprob=max(logprob.logprob, -9999.0), bytes=token_bytes + token=token, logprob=_clamp_logprob(logprob.logprob), bytes=token_bytes ) @@ -406,12 +419,20 @@ async def serve_tokens_full_generator( # This is top_logprobs in completions API if sampling_params.logprobs is not None: assert out_logprobs is not None, "Did not output logprobs" - logprobs = self._create_tokens_logprobs( - token_ids=token_ids, - top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, - tokenizer=tokenizer, - ) + logprobs: GenerateLogProbs | ChatCompletionLogProbs | None + if text_mode: + logprobs = self._create_text_logprobs( + token_ids=token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=sampling_params.logprobs, + tokenizer=tokenizer, + ) + else: + logprobs = self._create_tokens_logprobs( + token_ids=token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=sampling_params.logprobs, + ) else: logprobs = None @@ -571,12 +592,20 @@ async def serve_tokens_stream_generator( if sampling_params.logprobs is not None: out_logprobs = output.logprobs assert out_logprobs is not None, "Did not output logprobs" - logprobs = self._create_tokens_logprobs( - token_ids=delta_token_ids, - top_logprobs=out_logprobs, - num_output_top_logprobs=sampling_params.logprobs, - tokenizer=tokenizer, - ) + logprobs: GenerateLogProbs | ChatCompletionLogProbs | None + if text_mode: + logprobs = self._create_text_logprobs( + token_ids=delta_token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=sampling_params.logprobs, + tokenizer=tokenizer, + ) + else: + logprobs = self._create_tokens_logprobs( + token_ids=delta_token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=sampling_params.logprobs, + ) else: logprobs = None @@ -691,17 +720,18 @@ def _logprobs_tokenizer(self, text_mode: bool) -> TokenizerLike | None: return None return self.renderer.tokenizer - def _create_tokens_logprobs( + def _create_text_logprobs( self, token_ids: GenericSequence[int], top_logprobs: GenericSequence[dict[int, Logprob] | None], num_output_top_logprobs: int | None = None, tokenizer: TokenizerLike | None = None, ) -> ChatCompletionLogProbs: - """Create OpenAI-style logprobs. + """Create OpenAI-style logprobs for ``output_mode="text"``. - Tokens are ``token_id:N`` placeholders. With a ``tokenizer`` they are - decoded strings and carry ``bytes``. + With a ``tokenizer`` the tokens are decoded strings and carry ``bytes``; + without one (``--return-tokens-as-token-ids``) they are ``token_id:N`` + placeholders. """ logprobs_content: list[ChatCompletionLogProbsContent] = [] @@ -719,7 +749,7 @@ def _create_tokens_logprobs( logprobs_content.append( ChatCompletionLogProbsContent( token=token, - logprob=max(step_token.logprob, -9999.0), + logprob=_clamp_logprob(step_token.logprob), bytes=token_bytes, top_logprobs=[ _top_logprob(top_id, logprob, tokenizer) @@ -736,3 +766,52 @@ def _create_tokens_logprobs( ) return ChatCompletionLogProbs(content=logprobs_content) + + def _create_tokens_logprobs( + self, + token_ids: GenericSequence[int], + top_logprobs: GenericSequence[dict[int, Logprob] | None], + num_output_top_logprobs: int | None = None, + ) -> GenerateLogProbs: + """Create generate-shaped logprobs (integer token ids, no tokenizer). + + The engine reports rank 0 for a sampled token whose logprob is NaN; that + is not a rank, so it is sent as ``None`` (the logprob is clamped). + """ + logprobs_content: list[GenerateLogProbsContent] = [] + + for i, token_id in enumerate(token_ids): + step_top_logprobs = top_logprobs[i] + if step_top_logprobs is None or step_top_logprobs.get(token_id) is None: + # Same sentinel the OpenAI shapes use when the sampled token + # has no entry in the engine's top-k map. + logprobs_content.append( + GenerateLogProbsContent(token_id=token_id, logprob=-9999.0) + ) + else: + step_token = step_top_logprobs[token_id] + + logprobs_content.append( + GenerateLogProbsContent( + token_id=token_id, + logprob=_clamp_logprob(step_token.logprob), + rank=step_token.rank or None, + top_logprobs=[ + GenerateLogProb( + token_id=top_token_id, + logprob=_clamp_logprob(top_logprob.logprob), + rank=top_logprob.rank or None, + ) + for rank_index, (top_token_id, top_logprob) in enumerate( + step_top_logprobs.items() + ) + if num_output_top_logprobs is not None + and ( + num_output_top_logprobs == -1 + or rank_index < max(num_output_top_logprobs, 1) + ) + ], + ) + ) + + return GenerateLogProbs(content=logprobs_content) diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 1c6875159a9f..bd5bf88fbe8c 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -9,9 +9,11 @@ DeltaMessage, ToolCall, ) -from vllm.entrypoints.generate.base.serving import resolve_token_id_placeholder +from vllm.entrypoints.generate.base.serving import decode_token_ids from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionLogProb, ChatCompletionLogProbs, + ChatCompletionLogProbsContent, ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponseChoice, @@ -28,6 +30,7 @@ ) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderStreamState, + GenerateLogProbs, GenerateTokensResponse, GenerateTokensStreamResponse, ) @@ -963,16 +966,6 @@ def _seed_stream_state( ) -def _parse_token_id_placeholder(token: str) -> int | None: - """Extract token ID from a 'token_id:N' placeholder string.""" - if not token.startswith("token_id:"): - return None - try: - return int(token[len("token_id:") :]) - except ValueError: - return None - - def _correct_decoded_token( token_id: int, context_token_ids: list[int], tokenizer: TokenizerLike ) -> str: @@ -1012,53 +1005,63 @@ def _correct_decoded_token( def _resolve_logprobs( - logprobs: ChatCompletionLogProbs, + logprobs: GenerateLogProbs, tokenizer: TokenizerLike, initial_context_token_ids: Sequence[int] = (), ) -> ChatCompletionLogProbs: - """Resolve token_id:N placeholders in a ChatCompletionLogProbs object. - - ``initial_context_token_ids`` seeds the byte-fallback correction context - with sampled IDs from preceding chunks (streaming), so multi-byte - characters split across chunk boundaries still resolve. + """Convert generate's integer-id logprobs to the OpenAI chat shape. + + The generate server has no tokenizer, so `token` and `bytes` are filled + here, with the same U+FFFD byte-fallback correction the coupled path + applies (`_correct_decoded_token`, which needs the preceding sampled token + ids as context). ``initial_context_token_ids`` seeds that context with + sampled IDs from preceding chunks (streaming), so multi-byte characters + split across chunk boundaries still resolve. """ if logprobs.content is None: - return logprobs + return ChatCompletionLogProbs() context_token_ids: list[int] = list(initial_context_token_ids) resolved_content = [] for entry in logprobs.content: - token_str, token_bytes = resolve_token_id_placeholder(entry.token, tokenizer) - sampled_id = _parse_token_id_placeholder(entry.token) + # One batch per position: the sampled id and its top-k ids. + (token_str, token_bytes), *top_decoded = decode_token_ids( + [entry.token_id, *(top.token_id for top in entry.top_logprobs)], + tokenizer, + ) - if token_str.endswith("�") and sampled_id is not None: - token_str = _correct_decoded_token(sampled_id, context_token_ids, tokenizer) + if token_str.endswith("\ufffd"): + token_str = _correct_decoded_token( + entry.token_id, context_token_ids, tokenizer + ) token_bytes = list(token_str.encode("utf-8")) resolved_top = [] - for top in entry.top_logprobs: - top_str, top_bytes = resolve_token_id_placeholder(top.token, tokenizer) - top_id = _parse_token_id_placeholder(top.token) - if top_str.endswith("�") and top_id is not None: - top_str = _correct_decoded_token(top_id, context_token_ids, tokenizer) + for top, (top_str, top_bytes) in zip(entry.top_logprobs, top_decoded): + if top_str.endswith("\ufffd"): + top_str = _correct_decoded_token( + top.token_id, context_token_ids, tokenizer + ) top_bytes = list(top_str.encode("utf-8")) resolved_top.append( - top.model_copy(update={"token": top_str, "bytes": top_bytes}) + ChatCompletionLogProb( + token=top_str, + logprob=top.logprob, + bytes=top_bytes, + ) ) resolved_content.append( - entry.model_copy( - update={ - "token": token_str, - "bytes": token_bytes, - "top_logprobs": resolved_top, - } + ChatCompletionLogProbsContent( + token=token_str, + logprob=entry.logprob, + bytes=token_bytes, + top_logprobs=resolved_top, ) ) - if sampled_id is not None: - context_token_ids.append(sampled_id) + context_token_ids.append(entry.token_id) return ChatCompletionLogProbs(content=resolved_content) From 02b83919aa2eece2fd7efb8d8c906e4db9a02ff4 Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Tue, 6 Oct 2026 14:43:38 -0300 Subject: [PATCH 085/952] [Model] Support EmbeddingGemma2 multimodal pooling architecture (#60254) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins --- .../pooling/test_embedding_gemma2.py | 687 ++++++++++ tests/models/registry.py | 5 + vllm/model_executor/models/config.py | 106 ++ .../model_executor/models/embedding_gemma2.py | 1142 +++++++++++++++++ vllm/model_executor/models/registry.py | 1 + vllm/multimodal/video.py | 178 +++ .../attention/ops/triton_prefill_attention.py | 5 +- 7 files changed, 2122 insertions(+), 2 deletions(-) create mode 100644 tests/models/multimodal/pooling/test_embedding_gemma2.py create mode 100644 vllm/model_executor/models/embedding_gemma2.py diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py new file mode 100644 index 000000000000..701d0914ddbc --- /dev/null +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +import inspect +import os +import textwrap + +import pytest +import torch + +MODEL = os.environ.get("EG2_MODEL_PATH") +needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") +TEXT_COS, MM_COS = 0.999, 0.999 + + +def _cos(a, b): + return torch.nn.functional.cosine_similarity( + torch.tensor(a, dtype=torch.float32), + torch.tensor(b, dtype=torch.float32), + dim=-1, + ) + + +def _ensure_synthetic_video( + path: str, duration: float, fps: float, width: int, height: int +) -> str: + if os.path.exists(path): + os.remove(path) + import cv2 + import numpy as np + + fourcc_fn = getattr(cv2, "VideoWriter_fourcc", cv2.VideoWriter.fourcc) + fourcc = fourcc_fn(*"mp4v") + num_frames = int(duration * fps) + writer = cv2.VideoWriter(path, fourcc, fps, (width, height)) + y_coords, x_coords = np.mgrid[0:height, 0:width] + for i in range(num_frames): + # Spatially varying color pattern that evolves dynamically over frames + frame: np.ndarray = np.zeros((height, width, 3), dtype=np.uint8) + frame[..., 0] = ((x_coords * 255 // max(1, width)) + i * 17) % 256 + frame[..., 1] = ((y_coords * 255 // max(1, height)) + i * 29) % 256 + frame[..., 2] = (x_coords + y_coords + i * 13) % 256 + writer.write(frame) + writer.release() + return path + + +@pytest.fixture(scope="module") +def st_model(): + from sentence_transformers import SentenceTransformer + + return SentenceTransformer( + MODEL, device="cuda", model_kwargs={"torch_dtype": torch.bfloat16} + ) + + +@needs_model +def test_structure(vllm_runner, monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + + def probe(model): + lm = model.language_model + win = [layer.self_attn.attn.impl.sliding_window for layer in lm.layers] + return win, model.config.text_config.sliding_window + + win, cfg_w = vm.apply_model(probe)[0] + assert (cfg_w, cfg_w) in win # (512, 512) == W-1 with W = 513 + + +@needs_model +@pytest.mark.parametrize( + "prompt", + [ + "short query", + "task: search result | query: what is vllm", + "word " * 700, # > 512: crosses sliding window + "word " * 1500, # > 1024 + ], +) +def test_text_parity(vllm_runner, st_model, prompt): + ref = st_model.encode([prompt], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([prompt]) + sim = _cos(out, ref).min().item() + assert sim >= TEXT_COS, ( + f"Text cosine {sim} < {TEXT_COS} for prompt len {len(prompt)}" + ) + + +@needs_model +def test_mm_recipe(vllm_runner, st_model): + from PIL import Image + + img = Image.new("RGB", (224, 224), color=(255, 0, 0)) + prompt = "Describe this image: <|image|>" + + # SentenceTransformer multimodal encode + ref = st_model.encode([{"image": img, "text": prompt}], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([{"prompt": prompt, "multi_modal_data": {"image": img}}]) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Image cosine {sim} < {MM_COS}" + + +@needs_model +def test_natural_image_parity(vllm_runner, st_model): + from PIL import Image + + img_path = "/projects/gemma4-vllm/images/cat.jpg" + if not os.path.exists(img_path): + pytest.skip(f"Test asset {img_path} not found") + img = Image.open(img_path) + prompt = "What is this image? <|image|>" + ref = st_model.encode([{"image": img, "text": prompt}], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([{"prompt": prompt, "multi_modal_data": {"image": img}}]) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Natural image cosine {sim} < {MM_COS}" + + +@needs_model +def test_audio_parity(vllm_runner, st_model): + import torchaudio + + audio_path = "/projects/gemma4-vllm/audio/tellMeAboutTheSun.wav" + if not os.path.exists(audio_path): + pytest.skip(f"Test asset {audio_path} not found") + wav, sr = torchaudio.load(audio_path) + mono_wav = wav.mean(dim=0).numpy() + prompt = "Describe this audio: <|audio|>" + ref = st_model.encode( + [{"audio": audio_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"audio": (mono_wav, sr)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Audio cosine {sim} < {MM_COS}" + + +@needs_model +def test_mixed_parity(vllm_runner, st_model): + import torchaudio + from PIL import Image + + img_path = "/projects/gemma4-vllm/images/cat.jpg" + audio_path = "/projects/gemma4-vllm/audio/tellMeAboutTheSun.wav" + if not os.path.exists(img_path) or not os.path.exists(audio_path): + pytest.skip("Test assets for mixed parity not found") + img = Image.open(img_path) + wav, sr = torchaudio.load(audio_path) + mono_wav = wav.mean(dim=0).numpy() + prompt = "Describe image <|image|> and audio <|audio|>" + ref = st_model.encode( + [{"image": img, "audio": audio_path, "text": prompt}], + normalize_embeddings=True, + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [ + { + "prompt": prompt, + "multi_modal_data": { + "image": img, + "audio": (mono_wav, sr), + }, + } + ] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Mixed cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_parity_uniform_overflow(vllm_runner, st_model): + from transformers.video_utils import load_video + + video_path = _ensure_synthetic_video( + "/tmp/test_40s.mp4", duration=40.0, fps=10.0, width=160, height=120 + ) + frames, meta = load_video(video_path, backend="opencv") + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"40s video cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_parity_portrait(vllm_runner, st_model): + from transformers.video_utils import load_video + + video_path = _ensure_synthetic_video( + "/tmp/test_portrait.mp4", duration=3.0, fps=10.0, width=120, height=160 + ) + frames, meta = load_video(video_path, backend="opencv") + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Portrait video cosine {sim} < {MM_COS}" + + +@needs_model +def test_tower_skip(vllm_runner, monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + def nparams(m): + return sum(p.numel() for p in m.parameters()) / 1e6 + + with vllm_runner( + MODEL, + runner="pooling", + max_model_len=2048, + limit_mm_per_prompt={"image": 0, "audio": 0, "video": 0}, + ) as vm: + p_text = vm.apply_model(nparams)[0] + assert abs(p_text - 271.00) <= 0.05, ( + f"Expected 271.00M text params, got {p_text}M" + ) + + with vllm_runner(MODEL, runner="pooling", max_model_len=2048) as vm: + p_full = vm.apply_model(nparams)[0] + assert abs(p_full - 744.37) <= 0.05, ( + f"Expected 744.37M full params, got {p_full}M" + ) + + +def test_encoder_contract(): + from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model + from vllm.model_executor.models.gemma4_mm import ( + Gemma4ForConditionalGeneration as G, + ) + + borrowed_methods = [ + "_parse_and_validate_multimodal_inputs", + "_parse_and_validate_image_input", + "_parse_and_validate_video_input", + "_parse_and_validate_audio_input", + "_process_image_input", + "_process_video_input", + "_process_audio_input", + "embed_multimodal", + "_encoder_chunk", + ] + + for name in borrowed_methods: + # Check method identity (L-F6 + T-F4) + assert getattr(EmbeddingGemma2Model, name) is getattr(G, name) + + # AST checks: verify no zero-arg super() calls in borrowed methods + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "super" + ): + assert len(node.args) > 0, ( + f"Zero-arg super() call found in borrowed method {name}" + ) + + # AST check: all self. non-callable accesses are assigned in __init__ + init_src = inspect.getsource(EmbeddingGemma2Model.__init__) + for name in borrowed_methods: + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + ): + attr = node.attr + assert hasattr(EmbeddingGemma2Model, attr) or attr in init_src, ( + f"Attribute {attr} accessed on self in {name} is neither " + f"defined on class nor in __init__" + ) + + +@needs_model +def test_instance_attribute_access(vllm_runner): + """Check self. accesses on a real model instance.""" + from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model + + borrowed_methods = [ + "_parse_and_validate_multimodal_inputs", + "_parse_and_validate_image_input", + "_parse_and_validate_video_input", + "_parse_and_validate_audio_input", + "_process_image_input", + "_process_video_input", + "_process_audio_input", + "embed_multimodal", + "_encoder_chunk", + ] + attrs_accessed = set() + for name in borrowed_methods: + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + ): + attrs_accessed.add(node.attr) + + def check_attrs(model): + for attr in attrs_accessed: + assert hasattr(model, attr), ( + f"Real instance of EmbeddingGemma2Model is missing attribute '{attr}' " + f"accessed by borrowed method" + ) + return True + + with vllm_runner(MODEL, runner="pooling", max_model_len=2048) as vm: + res = vm.apply_model(check_attrs) + assert res[0] is True + + +@needs_model +def test_video_loader_real_video(vllm_runner, st_model): + """Decode a real MP4 file through EmbeddingGemma2VideoBackend.load_bytes + + and VideoMediaIO / MediaConnector with default num_frames=32. + """ + from vllm.multimodal.media import MediaConnector + from vllm.multimodal.media.video import VideoMediaIO + from vllm.multimodal.video import VIDEO_LOADER_REGISTRY + + video_path = "/projects/gemma4-vllm/videos/test_intel.mp4" + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + with open(video_path, "rb") as f: + data = f.read() + + # Direct backend load + frames, meta = backend_cls.load_bytes(data) + + # VideoMediaIO with default num_frames=32 (VideoMediaIO passes num_frames=32) + video_io = VideoMediaIO(image_io=None, video_backend="embedding_gemma2") + loaded = video_io.load_bytes(data) + frames_io, meta_io = loaded.media + assert len(frames_io) == len(frames) + assert meta_io.get("do_sample_frames") is False + + # MediaConnector with video_processor="EmbeddingGemma2VideoProcessor" + connector = MediaConnector(allowed_local_media_path="/projects/gemma4-vllm/videos") + fetched = connector.fetch_video( + f"file://{video_path}", + video_processor="EmbeddingGemma2VideoProcessor", + ) + frames_conn, meta_conn = fetched.media + assert len(frames_conn) == 32 + assert meta_conn.get("do_sample_frames") is False + + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [ + { + "prompt": prompt, + "multi_modal_data": {"video": (frames_conn, meta_conn)}, + } + ] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Real video loader cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): + """End-to-end duplicate expansion test with low-fps clip.""" + import numpy as np + + from vllm.multimodal.video import VIDEO_LOADER_REGISTRY + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + # 4s clip at 0.5 fps = 2 unique frames + lowfps_path = "/tmp/test_lowfps_05.mp4" + _ensure_synthetic_video(lowfps_path, duration=4.0, fps=0.5, width=64, height=48) + with open(lowfps_path, "rb") as f: + data = f.read() + # Default fps=1.0 -> 4 frames expanding 2 unique frames to [0, 0, 1, 1] + frames, meta = backend_cls.load_bytes(data) + assert frames.shape[0] == 4 + # Compare returned frames directly, not only indices + assert np.array_equal(frames[0], frames[1]) + assert not np.array_equal(frames[0], frames[2]) + assert np.array_equal(frames[2], frames[3]) + + prompt = "Describe this video: <|video|>" + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + ref = st_model.encode( + [{"video": lowfps_path, "text": prompt}], normalize_embeddings=True + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Duplicate expansion cosine {sim} < {MM_COS}" + assert len(out) == 1 + assert len(out[0]) == 768 + + +def test_video_loader_indices(): + from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + from transformers.video_utils import VideoMetadata + + from vllm.multimodal.video import ( + VIDEO_LOADER_REGISTRY, + VideoSourceMetadata, + VideoTargetMetadata, + ) + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + hf_proc = EmbeddingGemma2VideoProcessor() + + # Case 1: clip shorter than 32s (e.g. 5s @ 24fps = 120 frames) + src_5s = VideoSourceMetadata(total_frames_num=120, original_fps=24.0, duration=5.0) + tgt = VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + vllm_indices = backend_cls.compute_frames_index_to_sample(src_5s, tgt) + hf_meta_5s = VideoMetadata(fps=24.0, duration=5.0, total_num_frames=120) + hf_indices_5s = hf_proc.sample_frames( + hf_meta_5s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_indices == hf_indices_5s + + # Case 2: clip shorter than 1s (e.g. 0.5s @ 24fps = 12 frames) + src_sub1s = VideoSourceMetadata( + total_frames_num=12, original_fps=24.0, duration=0.5 + ) + vllm_sub1s = backend_cls.compute_frames_index_to_sample(src_sub1s, tgt) + hf_meta_sub1s = VideoMetadata(fps=24.0, duration=0.5, total_num_frames=12) + hf_indices_sub1s = hf_proc.sample_frames( + hf_meta_sub1s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_sub1s == hf_indices_sub1s + + # Case 3: clip longer than 32s (e.g. 40s @ 10fps = 400 frames) + # -> uniform overflow cap 32 frames + src_40s = VideoSourceMetadata( + total_frames_num=400, original_fps=10.0, duration=40.0 + ) + vllm_40s = backend_cls.compute_frames_index_to_sample(src_40s, tgt) + hf_meta_40s = VideoMetadata(fps=10.0, duration=40.0, total_num_frames=400) + hf_indices_40s = hf_proc.sample_frames( + hf_meta_40s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_40s == hf_indices_40s + assert len(vllm_40s) == 32 + + # Case 4: default target fps (target.fps <= 0 defaults to 1.0) + tgt_default_fps = VideoTargetMetadata(num_frames=-1, fps=-1.0, max_duration=300) + vllm_def = backend_cls.compute_frames_index_to_sample(src_5s, tgt_default_fps) + assert vllm_def == hf_indices_5s + + # Case 5 (NEW-2): 60 fps / 300 s (18000 frames) vs HF sample_frames + src_300s = VideoSourceMetadata( + total_frames_num=18000, original_fps=60.0, duration=300.0 + ) + v_300s = backend_cls.compute_frames_index_to_sample( + src_300s, VideoTargetMetadata(num_frames=-1, fps=60.0, max_duration=300) + ) + h_300s = hf_proc.sample_frames( + VideoMetadata(fps=60.0, duration=300.0, total_num_frames=18000), + fps=60.0, + max_frames=32, + overflow_strategy="uniform", + ).tolist() + assert v_300s == h_300s, f"60fps/300s mismatch: {v_300s[-3:]} vs {h_300s[-3:]}" + + # Case 6 (NEW-2): 1 fps / 4 h (432000 frames @ 30 native fps) vs HF sample_frames + src_4h = VideoSourceMetadata( + total_frames_num=432000, original_fps=30.0, duration=14400.0 + ) + v_4h = backend_cls.compute_frames_index_to_sample( + src_4h, VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + ) + h_4h = hf_proc.sample_frames( + VideoMetadata(fps=30.0, duration=14400.0, total_num_frames=432000), + fps=1.0, + max_frames=32, + overflow_strategy="uniform", + ).tolist() + assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" + + +def test_video_loader_duplicate_expansion(): + from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + from transformers.video_utils import VideoMetadata + + from vllm.multimodal.video import ( + VIDEO_LOADER_REGISTRY, + VideoSourceMetadata, + VideoTargetMetadata, + ) + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + hf_proc = EmbeddingGemma2VideoProcessor() + + # 4s clip @ 0.5fps = 2 frames total. Requested target fps = 1.0 -> 4 frames. + # Source has fewer frames than requested sampled frames, requiring duplicates. + src_lowfps = VideoSourceMetadata(total_frames_num=2, original_fps=0.5, duration=4.0) + tgt = VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + vllm_indices = backend_cls.compute_frames_index_to_sample(src_lowfps, tgt) + hf_meta = VideoMetadata(fps=0.5, duration=4.0, total_num_frames=2) + hf_indices = hf_proc.sample_frames( + hf_meta, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_indices == hf_indices + assert len(vllm_indices) == 4 + assert vllm_indices == [0, 0, 1, 1] + + +def test_video_media_io_embedding_gemma2(): + """Verify VideoMediaIO with default num_frames=32 loads video bytes for + + EmbeddingGemma2 without error. + """ + import os + from unittest.mock import MagicMock + + from vllm.multimodal.media import MediaConnector + from vllm.multimodal.media.video import VideoMediaIO + + video_path = "/projects/gemma4-vllm/videos/test_intel.mp4" + if not os.path.exists(video_path): + return + + with open(video_path, "rb") as f: + data = f.read() + + # VideoMediaIO defaults to num_frames=32 + video_io = VideoMediaIO(image_io=MagicMock(), video_backend="embedding_gemma2") + loaded = video_io.load_bytes(data) + frames, meta = loaded.media + assert len(frames) == 32 + assert meta.get("do_sample_frames") is False + + # MediaConnector path + connector = MediaConnector(allowed_local_media_path="/projects/gemma4-vllm/videos") + fetched = connector.fetch_video( + f"file://{video_path}", + video_processor="EmbeddingGemma2VideoProcessor", + ) + frames_conn, meta_conn = fetched.media + assert len(frames_conn) == 32 + assert meta_conn.get("do_sample_frames") is False + + +def test_config_capping_and_explicit_limits(caplog): + """Verify EmbeddingGemma2ModelConfig capping and explicit user overrides.""" + import logging + from types import SimpleNamespace + from unittest.mock import patch + + from vllm.model_executor.models.config import ( + EmbeddingGemma2ModelConfig, + Gemma4Config, + ) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + # (a) Defaults: uncapped max_model_len=262144 -> capped to 8192, + # and derived scheduler limits (>= 262144) lowered to 8192. + cfg_default = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=262144, original_max_model_len=None + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=262144, + max_num_encoder_input_tokens=262144, + encoder_cache_size=262144, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_default) + + assert cfg_default.attention_config.backend == AttentionBackendEnum.TRITON_ATTN + assert cfg_default.model_config.max_model_len == 8192 + assert cfg_default.scheduler_config.max_num_batched_tokens == 8192 + assert cfg_default.scheduler_config.max_num_encoder_input_tokens == 8192 + assert cfg_default.scheduler_config.encoder_cache_size == 8192 + + # (b) Reachable explicit limits: explicit max_model_len=8192 and explicit + # max_num_batched_tokens=32768 stays 32768. + # Note: omitting max_model_len while passing max_num_batched_tokens < 262144 + # is rejected earlier by SchedulerConfig validation (32768 < 262144) before + # this hook runs. + cfg_explicit_batched = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=8192, original_max_model_len=8192 + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=32768, + max_num_encoder_input_tokens=32768, + encoder_cache_size=32768, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_batched) + + assert cfg_explicit_batched.model_config.max_model_len == 8192 + assert cfg_explicit_batched.scheduler_config.max_num_batched_tokens == 32768 + assert cfg_explicit_batched.scheduler_config.max_num_encoder_input_tokens == 32768 + assert cfg_explicit_batched.scheduler_config.encoder_cache_size == 32768 + + # (c) Explicit value > uncapped_len (e.g. 300000) with uncapped model len is lowered + # to 8192 and logs a warning because it exceeds uncapped_len. + cfg_explicit_large = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=262144, original_max_model_len=None + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=300000, + max_num_encoder_input_tokens=300000, + encoder_cache_size=300000, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with ( + caplog.at_level(logging.WARNING), + patch.object(Gemma4Config, "verify_and_update_config"), + ): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_large) + + assert cfg_explicit_large.model_config.max_model_len == 8192 + assert cfg_explicit_large.scheduler_config.max_num_batched_tokens == 8192 + assert cfg_explicit_large.scheduler_config.max_num_encoder_input_tokens == 8192 + assert cfg_explicit_large.scheduler_config.encoder_cache_size == 8192 + assert any( + "lowering explicitly set scheduler_config.max_num_batched_tokens" + " from 300000 to 8192" in rec.message + for rec in caplog.records + ) + + # (d) Explicit max_model_len=16384 is respected and not capped to 8192 + cfg_explicit_model_len = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=16384, original_max_model_len=16384 + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=16384, + max_num_encoder_input_tokens=16384, + encoder_cache_size=16384, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_model_len) + + assert cfg_explicit_model_len.model_config.max_model_len == 16384 + assert cfg_explicit_model_len.scheduler_config.max_num_batched_tokens == 16384 diff --git a/tests/models/registry.py b/tests/models/registry.py index b764cad595a8..95b424ede580 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -581,6 +581,11 @@ def check_available_online( trust_remote_code=True, hf_overrides={"model_type": "deepseek_v3", "auto_map": None}, ), + "EmbeddingGemma2Model": _HfExamplesInfo( + "google/embeddinggemma-2", + min_transformers_version="5.18.0", + is_available_online=False, + ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"), "GteModel": _HfExamplesInfo( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d80663731d1e..97b9bfa0d38b 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -282,6 +282,111 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +class EmbeddingGemma2ModelConfig(Gemma4Config): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Gemma4Config.verify_and_update_config(vllm_config) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attn_config = vllm_config.attention_config + if attn_config.backend is None: + attn_config.backend = AttentionBackendEnum.TRITON_ATTN + logger.info( + "EmbeddingGemma2: attention backend not specified; defaulting to " + "TRITON_ATTN (validated backend for heterogeneous head dimensions)." + ) + else: + logger.info( + "EmbeddingGemma2: using attention backend %s", attn_config.backend + ) + + model_config = vllm_config.model_config + orig_len = getattr(model_config, "original_max_model_len", None) + if orig_len in (None, -1): + default_len = 8192 + # Check sentence_bert_config.json if available + try: + import json + from pathlib import Path + + cfg_path = Path(model_config.model) / "sentence_bert_config.json" + if cfg_path.exists(): + data = json.loads(cfg_path.read_text()) + if "max_seq_length" in data and isinstance( + data["max_seq_length"], int + ): + default_len = data["max_seq_length"] + except Exception: + pass + + if model_config.max_model_len > default_len: + uncapped_len = model_config.max_model_len + logger.info( + "EmbeddingGemma2: max_model_len not explicitly set; capping from " + "%d to %d.", + model_config.max_model_len, + default_len, + ) + model_config.max_model_len = default_len + scheduler_config = vllm_config.scheduler_config + if scheduler_config is not None: + new_batched_tokens = max( + model_config.max_model_len, scheduler_config.max_num_seqs + ) + if scheduler_config.max_num_batched_tokens >= uncapped_len: + if scheduler_config.max_num_batched_tokens > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.max_num_batched_tokens from %d to %d " + "to match capped max_model_len.", + scheduler_config.max_num_batched_tokens, + new_batched_tokens, + ) + else: + logger.info( + "EmbeddingGemma2: lowering " + "scheduler_config.max_num_batched_tokens from %d to %d " + "to match capped max_model_len.", + scheduler_config.max_num_batched_tokens, + new_batched_tokens, + ) + scheduler_config.max_num_batched_tokens = new_batched_tokens + if hasattr(vllm_config, "_set_compile_ranges"): + vllm_config._set_compile_ranges() + if scheduler_config.max_num_encoder_input_tokens >= uncapped_len: + if scheduler_config.max_num_encoder_input_tokens > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.max_num_encoder_input_tokens " + "from %d to %d to match capped max_model_len.", + scheduler_config.max_num_encoder_input_tokens, + new_batched_tokens, + ) + scheduler_config.max_num_encoder_input_tokens = ( + new_batched_tokens + ) + if scheduler_config.encoder_cache_size >= uncapped_len: + if scheduler_config.encoder_cache_size > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.encoder_cache_size from %d to %d " + "to match capped max_model_len.", + scheduler_config.encoder_cache_size, + new_batched_tokens, + ) + scheduler_config.encoder_cache_size = new_batched_tokens + if hasattr(scheduler_config, "verify_max_model_len"): + scheduler_config.verify_max_model_len( + model_config.max_model_len + ) + elif model_config.max_model_len > 32768: + logger.warning( + "EmbeddingGemma2: max_model_len=%d is large; " + "consider --max-model-len 8192 to bound encoder memory.", + model_config.max_model_len, + ) + + class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: @@ -1025,6 +1130,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 + "EmbeddingGemma2Model": EmbeddingGemma2ModelConfig, "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 "FalconMambaForCausalLM": MambaModelConfig, "Gemma3TextModel": Gemma3TextModelConfig, diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py new file mode 100644 index 000000000000..2d0012714eab --- /dev/null +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -0,0 +1,1142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EmbeddingGemma2: bidirectional Gemma4-derived multimodal embedding model.""" + +import dataclasses +from collections.abc import Iterable, Mapping, Sequence +from typing import Any, cast + +import torch +from torch import nn +from transformers import AutoModel, BatchFeature, PreTrainedConfig +from transformers.models.embedding_gemma2.configuration_embedding_gemma2 import ( + EmbeddingGemma2Config, + EmbeddingGemma2TextConfig, +) +from transformers.models.embedding_gemma2.processing_embedding_gemma2 import ( + EmbeddingGemma2Processor, +) +from transformers.video_utils import VideoMetadata + +from vllm.config import VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.pooler import DispatchPooler +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.gemma4 import ( + Gemma4MLP, + gemma4_layer_config, +) +from vllm.model_executor.models.gemma4_mm import ( + _SUPPORTED_SOFT_TOKENS, + Gemma4DummyInputsBuilder, + Gemma4ForConditionalGeneration, + Gemma4MultimodalEmbedder, + Gemma4MultiModalProcessor, + Gemma4ProcessingInfo, + _get_max_soft_tokens, +) +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.models.interfaces_base import ( + VllmModelForPooling, + default_pooling_type, +) +from vllm.model_executor.models.transformers.utils import ( + recursive_replace_linear, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalKwargsItems, + VideoItem, +) +from vllm.multimodal.parse import ( + AudioProcessorItems, + ImageProcessorItems, + MultiModalDataItems, +) +from vllm.multimodal.processing import ( + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backend import AttentionType + +logger = init_logger(__name__) + +_WINDOW_OFFSET = 1 # HF masks |q-k| <= W; vLLM ENCODER_ONLY masks <= W-1. +_VIDEO_FRAME_WIDTH = 672 +_VIDEO_FRAME_HEIGHT = 480 +_VIDEO_MAX_FRAMES = 32 +_VIDEO_MAX_SOFT_TOKENS = 140 + + +class EmbeddingGemma2Attention(nn.Module): + """Bidirectional attention with sliding window offset for EmbeddingGemma2.""" + + def __init__( + self, + config: EmbeddingGemma2TextConfig, + layer_cfg, + layer_idx: int, + cache_config, + quant_config: QuantizationConfig | None, + max_position: int, + prefix: str, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + tp = get_tensor_model_parallel_world_size() + self.num_heads = config.num_attention_heads // tp + self.num_kv_heads = max(1, layer_cfg.num_key_value_heads // tp) + self.head_dim = layer_cfg.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + config.num_attention_heads, + layer_cfg.num_key_value_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False) + + assert config.layer_types is not None + layer_type = config.layer_types[layer_idx] + self.is_sliding = layer_type == "sliding_attention" + assert config.rope_parameters is not None + rope_parameters = dict(config.rope_parameters[layer_type]) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position, + rope_parameters=rope_parameters, + is_neox_style=True, + ) + + window = config.sliding_window if self.is_sliding else None + self.attn = Attention( + self.num_heads, + self.head_dim, + scale=1.0, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=(window + _WINDOW_OFFSET) + if window is not None + else None, + attn_type=AttentionType.ENCODER_ONLY, + prefix=f"{prefix}.attn", + ) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + q = self.q_norm(q) + q = q.flatten(-2, -1) + + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + k = self.k_norm(k) + k = k.flatten(-2, -1) + q, k = self.rotary_emb(positions, q, k) + + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)) + v = self.v_norm(v) + v = v.flatten(-2, -1) + + attn_output = self.attn(q, k, v) + out, _ = self.o_proj(attn_output) + return out + + +class EmbeddingGemma2PLEBlock(nn.Module): + """Per-layer-embedding residual block (mirrors HF ple_block).""" + + def __init__( + self, + config: EmbeddingGemma2TextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + self.hidden_size_per_layer_input = config.hidden_size_per_layer_input + + self.per_layer_input_gate = ReplicatedLinear( + self.hidden_size, + self.hidden_size_per_layer_input, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_input_gate", + return_bias=False, + ) + self.per_layer_projection = ReplicatedLinear( + self.hidden_size_per_layer_input, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_projection", + return_bias=False, + ) + self.post_per_layer_input_norm = RMSNorm( + self.hidden_size, eps=config.rms_norm_eps + ) + self.act_fn = get_act_fn(config.hidden_activation) + + def forward( + self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor + ) -> torch.Tensor: + residual = hidden_states + gate = self.per_layer_input_gate(hidden_states) + gate = self.act_fn(gate) + gated = gate * per_layer_input + proj = self.per_layer_projection(gated) + proj = self.post_per_layer_input_norm(proj) + return residual + proj + + +class EmbeddingGemma2DecoderLayer(nn.Module): + """Transformer decoder layer for EmbeddingGemma2 text model with PLE gating.""" + + def __init__( + self, + config: EmbeddingGemma2TextConfig, + idx: int, + cache_config, + quant_config: QuantizationConfig | None, + max_position: int, + prefix: str, + ): + super().__init__() + lc = gemma4_layer_config(config, idx) + self.self_attn = EmbeddingGemma2Attention( + config, + lc, + idx, + cache_config, + quant_config, + max_position, + f"{prefix}.self_attn", + ) + self.mlp = Gemma4MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.ple_block = EmbeddingGemma2PLEBlock( + config, + quant_config=quant_config, + prefix=f"{prefix}.ple_block", + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_feedforward_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_feedforward_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.layer_scalar: torch.Tensor + self.register_buffer("layer_scalar", torch.ones(1), persistent=True) + + def forward( + self, + positions: torch.Tensor, + h: torch.Tensor, + per_layer_input: torch.Tensor, + ) -> torch.Tensor: + r = h + h = ( + self.post_attention_layernorm( + self.self_attn(positions, self.input_layernorm(h)) + ) + + r + ) + r = h + h = ( + self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(h))) + + r + ) + h = self.ple_block(h, per_layer_input) + return h * cast(torch.Tensor, self.layer_scalar) + + +class EmbeddingGemma2TextPLE(nn.Module): + """Per-layer embedding (PLE) projection and normalization module.""" + + def __init__( + self, + config: EmbeddingGemma2TextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size_per_layer_input = config.hidden_size_per_layer_input + self.per_layer_model_projection = ColumnParallelLinear( + config.hidden_size, + config.num_hidden_layers * config.hidden_size_per_layer_input, + bias=False, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_model_projection", + ) + self.per_layer_model_projection_scale = config.hidden_size**-0.5 + self.per_layer_projection_norm = RMSNorm( + config.hidden_size_per_layer_input, eps=config.rms_norm_eps + ) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + proj, _ = self.per_layer_model_projection(inputs_embeds) + proj = proj * self.per_layer_model_projection_scale + proj = proj.reshape( + *inputs_embeds.shape[:-1], + self.num_hidden_layers, + self.hidden_size_per_layer_input, + ) + return self.per_layer_projection_norm(proj) + + +class EmbeddingGemma2TextModel(nn.Module): + """Text embedding backbone for EmbeddingGemma2 with bidirectional attention.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + cfg = vllm_config.model_config.hf_config.text_config + self.config = cfg + max_pos = min( + cfg.max_position_embeddings, vllm_config.model_config.max_model_len + ) + self.embed_tokens = VocabParallelEmbedding( + cfg.vocab_size, + cfg.hidden_size, + quant_config=vllm_config.quant_config, + prefix=f"{prefix}.embed_tokens", + ) + self.layers = nn.ModuleList( + [ + EmbeddingGemma2DecoderLayer( + cfg, + i, + vllm_config.cache_config, + vllm_config.quant_config, + max_pos, + f"{prefix}.layers.{i}", + ) + for i in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + self.ple = EmbeddingGemma2TextPLE( + cfg, + quant_config=vllm_config.quant_config, + prefix=f"{prefix}.ple", + ) + self.embedding_projection = ReplicatedLinear( + cfg.hidden_size, + cfg.embedding_dim, + bias=False, + quant_config=vllm_config.quant_config, + prefix=f"{prefix}.embedding_projection", + return_bias=False, + ) + self.embed_scale: torch.Tensor + self.register_buffer( + "embed_scale", + torch.tensor(cfg.hidden_size**0.5), + persistent=False, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + embed_scale = cast(torch.Tensor, self.embed_scale) + dtype = cast(torch.dtype, self.embed_tokens.weight.dtype) + return self.embed_tokens(input_ids) * embed_scale.to(dtype) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + per_layer = self.ple(inputs_embeds) + h = inputs_embeds + for i, layer in enumerate(self.layers): + h = layer(positions, h, per_layer[..., i, :]) + h = self.norm(h) + h = self.embedding_projection(h) + return h + + +class EmbeddingGemma2ProcessingInfo(Gemma4ProcessingInfo): + """Processing information and token budgeting for EmbeddingGemma2.""" + + def get_hf_config(self) -> EmbeddingGemma2Config: + return self.ctx.get_hf_config(EmbeddingGemma2Config) + + def get_hf_processor( # type: ignore[override] + self, **kwargs: object + ) -> EmbeddingGemma2Processor: + return self.ctx.get_hf_processor(EmbeddingGemma2Processor, **kwargs) + + def get_default_tok_params(self): + # Base-class behavior: add_special_tokens=True (BOS + EOS) + return BaseProcessingInfo.get_default_tok_params(self) + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int] | None: + config = self.get_hf_config() + merged_kwargs = self.ctx.get_merged_mm_kwargs({}) + val, _ = _get_max_soft_tokens(merged_kwargs) + if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: + tokens_per_image = val + 2 + else: + tokens_per_image = ( + getattr(config.vision_config, "default_output_length", 256) + 2 + ) + tokens: dict[str, int] = {"image": tokens_per_image} + if config.audio_config is not None: + processor = self.get_hf_processor() + tokens["audio"] = processor.audio_seq_length + 2 + + processor = self.get_hf_processor() + video_proc = getattr(processor, "video_processor", None) + max_frames = getattr(video_proc, "max_frames", 32) + mm_config = self.ctx.model_config.get_multimodal_config() + + # The loader's max_frames cap is 128 when media_io_kwargs can raise it + video_media_io = ( + getattr(mm_config, "media_io_kwargs", {}).get("video", {}) + if mm_config + else {} + ) + media_io_max_frames = video_media_io.get("max_frames") + if media_io_max_frames is not None: + max_frames = min(int(media_io_max_frames), 128) + elif "max_frames" in merged_kwargs: + max_frames = min(int(merged_kwargs["max_frames"]), 128) + + video_opts = mm_config.limit_per_prompt.get("video") if mm_config else None + if video_opts is not None: + num_frames_opt = getattr(video_opts, "num_frames", None) + if isinstance(num_frames_opt, int): + max_frames = min(max_frames, num_frames_opt) + + video_max_soft = merged_kwargs.get( + "max_soft_tokens", + getattr(video_proc, "max_soft_tokens", _VIDEO_MAX_SOFT_TOKENS), + ) + if video_max_soft is not None and video_max_soft not in _SUPPORTED_SOFT_TOKENS: + raise ValueError( + f"Unsupported max_soft_tokens={video_max_soft} for video. " + f"Valid values are {_SUPPORTED_SOFT_TOKENS}." + ) + if not isinstance(video_max_soft, int): + video_max_soft = _VIDEO_MAX_SOFT_TOKENS + + # Token allocation per video frame verified against HF EmbeddingGemma2Processor: + # - Without timestamps: boi (1) + video_token * num_soft_tokens + eoi (1) + # Total = num_soft_tokens + 2 + # - With timestamps: timestamp prefix string (e.g. ' 01:23 ' = up to 7 tokens) + # Total = num_soft_tokens + 2 + 7 + add_ts = getattr(video_proc, "add_timestamps", False) or merged_kwargs.get( + "add_timestamps", False + ) + ts_slack = 7 if add_ts else 0 + tokens["video"] = max_frames * (video_max_soft + 2 + ts_slack) + return tokens + + def get_image_repl( # type: ignore[override] + self, + *, + image_width: int, + image_height: int, + processor: EmbeddingGemma2Processor | None, + max_soft_tokens: int | None = None, + ) -> PromptUpdateDetails: + if processor is None: + processor = self.get_hf_processor() + + num_soft = self._compute_num_soft_tokens( + image_width, + image_height, + max_soft_tokens=max_soft_tokens, + ) + config = self.get_hf_config() + assert config.boi_token_id is not None + assert processor.image_token_id is not None + assert config.eoi_token_id is not None + token_ids = ( + [config.boi_token_id] + + [processor.image_token_id] * num_soft + + [config.eoi_token_id] + ) + return PromptUpdateDetails.select_token_id(token_ids, processor.image_token_id) + + def get_audio_repl( # type: ignore[override] + self, + *, + audio_len: int, + processor: EmbeddingGemma2Processor | None, + ) -> PromptUpdateDetails: + if processor is None: + processor = self.get_hf_processor() + sampling_rate = processor.feature_extractor.sampling_rate + num_tokens = self._compute_audio_num_tokens( + audio_len, sampling_rate, processor.audio_seq_length + ) + config = self.get_hf_config() + assert config.boa_token_id is not None + assert processor.audio_token_id is not None + assert config.eoa_token_index is not None + token_ids = ( + [config.boa_token_id] + + [processor.audio_token_id] * num_tokens + + [config.eoa_token_index] + ) + return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id) + + def get_video_repl( # type: ignore[override] + self, + *, + timestamps: list[float], + num_soft_tokens_per_frame: list[int], + processor: EmbeddingGemma2Processor, + add_timestamps: bool = False, + ) -> PromptUpdateDetails: + tokenizer = self.ctx.get_tokenizer() + config = self.get_hf_config() + + boi_token_id = config.boi_token_id + eoi_token_id = config.eoi_token_id + video_token_id = processor.video_token_id + assert boi_token_id is not None + assert eoi_token_id is not None + assert video_token_id is not None + + all_token_ids: list[int] = [] + if not add_timestamps: + for n_tokens in num_soft_tokens_per_frame: + all_token_ids.append(boi_token_id) + all_token_ids.extend([video_token_id] * n_tokens) + all_token_ids.append(eoi_token_id) + else: + for i, (ts, n_tokens) in enumerate( + zip(timestamps, num_soft_tokens_per_frame) + ): + minutes = int(ts // 60) + seconds = int(ts % 60) + ts_str = f"{minutes:02d}:{seconds:02d}" + prefix = f" {ts_str} " if i > 0 else f"{ts_str} " + ts_token_ids = tokenizer.encode(prefix, add_special_tokens=False) + all_token_ids.extend(ts_token_ids) + all_token_ids.append(boi_token_id) + all_token_ids.extend([video_token_id] * n_tokens) + all_token_ids.append(eoi_token_id) + + return PromptUpdateDetails.select_token_id(all_token_ids, video_token_id) + + +class EmbeddingGemma2DummyInputsBuilder(Gemma4DummyInputsBuilder): + """Dummy input builder for EmbeddingGemma2 profiling and testing.""" + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + num_audios = mm_counts.get("audio", 0) + num_videos = mm_counts.get("video", 0) + processor = self.info.get_hf_processor() + text = ( + (processor.image_token * num_images) + + (processor.audio_token * num_audios if processor.audio_token else "") + + (processor.video_token * num_videos) + ) + return text + + def _get_dummy_videos( # type: ignore[override] + self, + *, + width: int, + height: int, + num_frames: int, + num_videos: int, + overrides=None, + ) -> list[VideoItem]: + num_frames = _VIDEO_MAX_FRAMES + width = _VIDEO_FRAME_WIDTH + height = _VIDEO_FRAME_HEIGHT + videos = super(Gemma4DummyInputsBuilder, self)._get_dummy_videos( + width=width, + height=height, + num_frames=num_frames, + num_videos=num_videos, + overrides=overrides, + ) + video_items: list[VideoItem] = [] + for video in videos: + video_num_frames = video.shape[0] + video_metadata = { + "fps": 1.0, + "duration": float(video_num_frames), + "total_num_frames": video_num_frames, + "frames_indices": list(range(video_num_frames)), + "video_backend": "opencv", + "do_sample_frames": False, + } + video_items.append((video, video_metadata)) + return video_items + + +class EmbeddingGemma2MultiModalProcessor(Gemma4MultiModalProcessor): + """Multimodal input processor for EmbeddingGemma2.""" + + info: EmbeddingGemma2ProcessingInfo + + def _apply_hf_processor_main( + self, + mm_items: MultiModalDataItems, + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( + mm_items, hf_kwargs + ) + + if not hf_data: + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, BatchFeature() + ) + + prompt_text = hf_data.pop("text") + assert isinstance(prompt_text, str) + + processor = self.info.get_hf_processor() + video_outputs: dict[str, Any] = {} + + if videos := hf_data.pop("videos", []): + assert isinstance(videos, list) + + all_video_pixel_values: list[torch.Tensor] = [] + all_video_position_ids: list[torch.Tensor] = [] + video_num_soft_tokens_per_video: list[list[int]] = [] + video_timestamps_per_video: list[list[float]] = [] + video_frame_counts: list[int] = [] + + for item in videos: + video_array, metadata = item + + if hf_kwargs.get("num_frames") not in (-1, None): + raise ValueError( + "EmbeddingGemma2 does not accept num_frames; " + "use max_frames/fps." + ) + + if isinstance(metadata, dict): + if not metadata.get("do_sample_frames", True) and any( + k in hf_kwargs + for k in ("fps", "max_frames", "overflow_strategy") + ): + logger.warning_once( + "For file/URL-loaded videos, " + "fps/max_frames/overflow_strategy in mm_processor_kwargs " + "have no effect because frames are decoded and sampled by " + "media_io. Please specify these sampling options via " + "media_io_kwargs instead." + ) + valid_meta_fields = { + f.name for f in dataclasses.fields(VideoMetadata) + } + meta_dict = { + k: v for k, v in metadata.items() if k in valid_meta_fields + } + metadata_obj = VideoMetadata(**meta_dict) + else: + if getattr(metadata, "do_sample_frames", True) is False and any( + k in hf_kwargs + for k in ("fps", "max_frames", "overflow_strategy") + ): + logger.warning_once( + "For file/URL-loaded videos, " + "fps/max_frames/overflow_strategy in mm_processor_kwargs " + "have no effect because frames are decoded and sampled by " + "media_io. Please specify these sampling options via " + "media_io_kwargs instead." + ) + metadata_obj = metadata + + video_mm_kwargs = dict(hf_kwargs) + video_mm_kwargs.pop("num_frames", None) + video_mm_kwargs["return_metadata"] = True + if isinstance(metadata, dict) and "do_sample_frames" in metadata: + video_mm_kwargs["do_sample_frames"] = metadata["do_sample_frames"] + elif hasattr(metadata, "do_sample_frames"): + video_mm_kwargs["do_sample_frames"] = metadata.do_sample_frames + + single_video_out = self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**video_mm_kwargs), + dict( + text=processor.video_token, + videos=[video_array], + video_metadata=[metadata_obj], + ), + video_mm_kwargs, + ) + + all_video_pixel_values.append(single_video_out["pixel_values_videos"]) + all_video_position_ids.append(single_video_out["video_position_ids"]) + + num_frames = int(single_video_out["num_frames_per_video"][0]) + video_frame_counts.append(num_frames) + + num_soft = single_video_out.get("num_soft_tokens_per_video") + if num_soft is not None: + soft_val = num_soft[0] + if isinstance(soft_val, torch.Tensor): + soft_val = soft_val.item() + soft_tokens_list = [int(soft_val)] * num_frames + elif "video_position_ids" in single_video_out and num_frames > 0: + pos = single_video_out["video_position_ids"] + pooling_k2 = ( + getattr(processor.video_processor, "pooling_kernel_size", 3) + ** 2 + ) + valid_patches = (pos[0] != -1).all(dim=-1).sum().item() + soft_tokens_list = [valid_patches // pooling_k2] * num_frames + elif "input_ids" in single_video_out and num_frames > 0: + vtok = ( + ( + single_video_out["input_ids"][0] + == processor.tokenizer.video_token_id + ) + .sum() + .item() + ) + soft_tokens_list = [vtok // num_frames] * num_frames + else: + raise ValueError( + f"Unable to determine soft tokens for video frame: " + f"keys={list(single_video_out.keys())}" + ) + + video_num_soft_tokens_per_video.append(soft_tokens_list) + + # Use returned video_metadata if available, else input metadata_obj + ret_meta = single_video_out.get("video_metadata") + meta_for_ts = ( + ret_meta[0] if (ret_meta and len(ret_meta) > 0) else metadata_obj + ) + + add_ts = hf_kwargs.get( + "add_timestamps", + getattr(processor.video_processor, "add_timestamps", False), + ) + if add_ts: + try: + ts = meta_for_ts.timestamps + except ValueError as e: + raise ValueError( + "Asked to build a prompt with frame timestamps, but no " + "valid fps/frames_indices was available in video metadata." + ) from e + video_timestamps_per_video.append([float(t) for t in ts]) + else: + try: + ts = meta_for_ts.timestamps + video_timestamps_per_video.append([float(t) for t in ts]) + except ValueError: + video_timestamps_per_video.append([0.0] * num_frames) + + video_outputs = { + "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0), + "pixel_position_ids_videos": torch.cat(all_video_position_ids, dim=0), + "video_frame_counts": torch.tensor( + video_frame_counts, dtype=torch.long + ), + "video_num_soft_tokens": video_num_soft_tokens_per_video, + "video_timestamps": video_timestamps_per_video, + } + + if hf_data: + remaining_counts: dict[str, int] = {} + if "images" in hf_data: + images_list = hf_data["images"] + remaining_counts["image"] = ( + len(images_list) if isinstance(images_list, list) else 1 + ) + if "audio" in hf_data: + audio_list = hf_data["audio"] + remaining_counts["audio"] = ( + len(audio_list) if isinstance(audio_list, (list, tuple)) else 1 + ) + remaining_prompt_text = self.dummy_inputs.get_dummy_text(remaining_counts) + + call_kwargs = dict(hf_kwargs) + call_kwargs.pop("num_frames", None) + processed_data = self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**call_kwargs), + dict(text=remaining_prompt_text, **hf_data), + call_kwargs, + ) + else: + processed_data = BatchFeature() + + if "image_position_ids" in processed_data: + processed_data["pixel_position_ids"] = processed_data.pop( + "image_position_ids" + ) + + if "input_features" in processed_data: + input_features = processed_data.pop("input_features") + input_features_mask = processed_data.pop("input_features_mask") + if not isinstance(input_features, list): + unpadded_features = [] + unpadded_masks = [] + for i in range(input_features.shape[0]): + mask = input_features_mask[i] + valid_len = int(mask.sum().item()) + unpadded_features.append(input_features[i, :valid_len]) + unpadded_masks.append(mask[:valid_len]) + processed_data["input_features_padded"] = unpadded_features + processed_data["input_features_mask"] = unpadded_masks + else: + processed_data["input_features_padded"] = input_features + processed_data["input_features_mask"] = input_features_mask + + processed_data.update(video_outputs) + processed_data.pop("num_frames_per_video", None) + processed_data.pop("video_position_ids", None) + + return self._finalize_hf_mm_data( + hf_data, hf_kwargs, passthrough_data, processed_data + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + prompt_updates = [] + + if "image" in mm_items: + image_token_id = hf_processor.image_token_id + + def get_replacement_image(item_idx: int): + images = mm_items.get_items("image", ImageProcessorItems) + image_size = images.get_image_size(item_idx) + merged_kwargs = self.info.ctx.get_merged_mm_kwargs( + hf_processor_mm_kwargs, + ) + val, _ = _get_max_soft_tokens(merged_kwargs) + max_soft_tokens = ( + val + if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS + else None + ) + return self.info.get_image_repl( + image_width=image_size.width, + image_height=image_size.height, + processor=hf_processor, + max_soft_tokens=max_soft_tokens, + ) + + prompt_updates.append( + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_replacement_image, + ) + ) + + if "video" in mm_items: + video_token_id = hf_processor.video_token_id + assert video_token_id is not None + + def get_replacement_video(item_idx: int): + out_item = out_mm_kwargs["video"][item_idx] + raw_ts = out_item["video_timestamps"].data + if isinstance(raw_ts, torch.Tensor): + ts_list: list[Any] = raw_ts.tolist() + elif isinstance(raw_ts, (list, tuple)): + ts_list = list(raw_ts) + else: + ts_list = [raw_ts] + timestamps = [float(t) for t in ts_list] + + raw_soft = out_item["video_num_soft_tokens"].data + if isinstance(raw_soft, torch.Tensor): + soft_list: list[Any] = raw_soft.tolist() + elif isinstance(raw_soft, (list, tuple)): + soft_list = list(raw_soft) + else: + soft_list = [raw_soft] + num_soft = [int(n) for n in soft_list] + + add_ts = hf_processor_mm_kwargs.get( + "add_timestamps", + getattr(hf_processor.video_processor, "add_timestamps", False), + ) + return self.info.get_video_repl( + timestamps=timestamps, + num_soft_tokens_per_frame=num_soft, + processor=hf_processor, + add_timestamps=add_ts, + ) + + prompt_updates.append( + PromptReplacement( + modality="video", + target=[video_token_id], + replacement=get_replacement_video, + ) + ) + + if "audio" in mm_items: + audio_token_id = hf_processor.audio_token_id + assert audio_token_id is not None + + def get_replacement_audio(item_idx: int): + audios = mm_items.get_items("audio", AudioProcessorItems) + audio_len = audios.get_audio_length(item_idx) + return self.info.get_audio_repl( + audio_len=audio_len, + processor=hf_processor, + ) + + prompt_updates.append( + PromptReplacement( + modality="audio", + target=[audio_token_id], + replacement=get_replacement_audio, + ) + ) + + return prompt_updates + + +def _embedding_gemma2_weights_mapper( + config_or_layers: PreTrainedConfig | int | None = None, +) -> WeightsMapper: + return WeightsMapper( + orig_to_new_stacked={ + ".self_attn.q_proj.weight": (".self_attn.qkv_proj.weight", "q"), + ".self_attn.k_proj.weight": (".self_attn.qkv_proj.weight", "k"), + ".self_attn.v_proj.weight": (".self_attn.qkv_proj.weight", "v"), + ".self_attn.q_proj.bias": (".self_attn.qkv_proj.bias", "q"), + ".self_attn.k_proj.bias": (".self_attn.qkv_proj.bias", "k"), + ".self_attn.v_proj.bias": (".self_attn.qkv_proj.bias", "v"), + ".mlp.gate_proj.weight": (".mlp.gate_up_proj.weight", 0), + ".mlp.up_proj.weight": (".mlp.gate_up_proj.weight", 1), + ".mlp.gate_proj.bias": (".mlp.gate_up_proj.bias", 0), + ".mlp.up_proj.bias": (".mlp.gate_up_proj.bias", 1), + } + ) + + +@MULTIMODAL_REGISTRY.register_processor( + EmbeddingGemma2MultiModalProcessor, + info=EmbeddingGemma2ProcessingInfo, + dummy_inputs=EmbeddingGemma2DummyInputsBuilder, +) +@default_pooling_type(seq_pooling_type="MEAN", tok_pooling_type="ALL") +class EmbeddingGemma2Model(nn.Module, SupportsMultiModal, VllmModelForPooling): + """Multimodal pooling embedding model supporting text, vision, and audio.""" + + is_pooling_model = True + + hf_to_vllm_mapper = _embedding_gemma2_weights_mapper() + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # TODO(lucianomartins): Extract Gemma4EncoderMixin into a shared module + # once generative model tests are migrated. + # TODO(lucianomartins): Investigate @support_torch_compile integration for encoder. + _parse_and_validate_image_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_image_input + ) + _parse_and_validate_video_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_video_input + ) + _parse_and_validate_audio_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_audio_input + ) + _parse_and_validate_multimodal_inputs = ( + Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs + ) + _process_image_input = Gemma4ForConditionalGeneration._process_image_input + _process_video_input = Gemma4ForConditionalGeneration._process_video_input + _process_audio_input = Gemma4ForConditionalGeneration._process_audio_input + embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal + _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + self.config = config + self.quant_config = quant_config + self.multimodal_config = multimodal_config + self.model_dtype = vllm_config.model_config.dtype + self.vllm_config = vllm_config + self.hf_to_vllm_mapper = _embedding_gemma2_weights_mapper(config) + lora_config = vllm_config.lora_config + self._enable_mm_lora = bool( + lora_config is not None and lora_config.enable_tower_connector_lora + ) + + tower_quant: QuantizationConfig | None + if quant_config and quant_config.get_name() in [ + "bitsandbytes", + "torchao", + "compressed-tensors", + ]: + tower_quant = quant_config + else: + vision_cfg = getattr(config, "vision_config", None) + quantizable = ( + vision_cfg is not None + and vision_cfg.hidden_size % 64 == 0 + and vision_cfg.intermediate_size % 64 == 0 + ) + tower_quant = quant_config if quantizable else None + + # ---- Vision tower (shared by image and video) ---- + self.embed_vision: Gemma4MultimodalEmbedder | None + if getattr(config, "vision_config", None) is not None: + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.embed_vision = Gemma4MultimodalEmbedder( + config.vision_config, + config.text_config, + quant_config=tower_quant, + prefix=maybe_prefix(prefix, "embed_vision"), + ) + recursive_replace_linear( + self.vision_tower, + tower_quant, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + else: + self.vision_tower = None + self.embed_vision = None + + # ---- Audio tower ---- + self.embed_audio: Gemma4MultimodalEmbedder | None + if getattr(config, "audio_config", None) is not None: + with self._mark_tower_model(vllm_config, "audio"): + self.audio_tower = AutoModel.from_config(config=config.audio_config) + self.audio_tower.post_init() + self.embed_audio = Gemma4MultimodalEmbedder( + config.audio_config, + config.text_config, + quant_config=tower_quant, + prefix=maybe_prefix(prefix, "embed_audio"), + ) + recursive_replace_linear( + self.audio_tower, + tower_quant, + prefix=maybe_prefix(prefix, "audio_tower"), + ) + else: + self.audio_tower = None + self.embed_audio = None + + # ---- Language model ---- + with self._mark_language_model(vllm_config): + self.language_model = EmbeddingGemma2TextModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + assert vllm_config.model_config.pooler_config is not None + self.pooler = DispatchPooler.for_embedding( + vllm_config.model_config.pooler_config + ) + + def get_language_model(self) -> nn.Module: # type: ignore[override] + return self.language_model + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor: + return self.language_model(input_ids, positions, inputs_embeds=inputs_embeds) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + ignore_prefixes = [ + "embed_vision.embedding.", + "embed_audio.embedding.", + ] + if self.audio_tower is None: + ignore_prefixes.extend(["audio_tower.", "embed_audio."]) + if self.vision_tower is None: + ignore_prefixes.extend(["vision_tower.", "embed_vision."]) + return AutoWeightsLoader( + self, + ignore_unexpected_prefixes=ignore_prefixes, + ).load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int = 0) -> str | None: + if modality == "image": + return "<|image|>" + if modality == "audio": + return "<|audio|>" + if modality == "video": + return "<|video|>" + raise ValueError(f"Unsupported modality: {modality}") diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 48b831bbdf0a..e019d674f502 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -246,6 +246,7 @@ # [Multimodal] "CLIPModel": ("clip", "CLIPEmbeddingModel"), "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "EmbeddingGemma2Model": ("embedding_gemma2", "EmbeddingGemma2Model"), "LlamaNemotronVLModel": ("nemotron_vl", "LlamaNemotronVLForEmbedding"), "LlavaNextForConditionalGeneration": ( "llava_next", diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 44a8e9452dc3..8f22bffe958c 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -405,6 +405,184 @@ def load_bytes( ) +@VIDEO_LOADER_REGISTRY.register( + "embedding_gemma2", + video_processor="EmbeddingGemma2VideoProcessor", +) +class EmbeddingGemma2VideoBackend(VideoBackend): + """Samples frames with the same indices as HF EmbeddingGemma2VideoProcessor.""" + + PROCESSOR_NAMES = ("EmbeddingGemma2Processor", "EmbeddingGemma2VideoProcessor") + + _MAX_ALLOWED_FRAMES: ClassVar[int] = 128 + _MAX_ALLOWED_FPS: ClassVar[float] = 60.0 + + @classmethod + def _get_processor_defaults(cls, **kwargs) -> tuple[float, int]: + """Read default fps and max_frames from video_processor if available.""" + proc = kwargs.get("video_processor") or kwargs.get("processor") + if proc is not None: + vp = getattr(proc, "video_processor", proc) + fps = getattr(vp, "fps", None) + max_frames = getattr(vp, "max_frames", None) + if fps is not None and max_frames is not None: + return float(fps), int(max_frames) + try: + from transformers.models.embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + + vp = EmbeddingGemma2VideoProcessor() + return float(getattr(vp, "fps", 1.0)), int(getattr(vp, "max_frames", 32)) + except Exception: + return 1.0, 32 + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + # Ignores target.num_frames (EmbeddingGemma2 does not accept num_frames) + default_fps, default_max_frames = cls._get_processor_defaults(**kwargs) + + fps = kwargs.get("fps") + if fps is None: + fps = target.fps if target.fps > 0 else default_fps + if fps is not None: + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") + if fps > cls._MAX_ALLOWED_FPS: + raise ValueError( + f"fps {fps} exceeds maximum allowed {cls._MAX_ALLOWED_FPS}" + ) + + max_frames = kwargs.get("max_frames", default_max_frames) + if max_frames is not None: + if max_frames <= 0: + raise ValueError(f"max_frames must be positive, got {max_frames}") + if max_frames > cls._MAX_ALLOWED_FRAMES: + raise ValueError( + f"max_frames {max_frames} exceeds maximum allowed " + f"{cls._MAX_ALLOWED_FRAMES}" + ) + + overflow_strategy = kwargs.get("overflow_strategy", "uniform") + + total_frames = source.total_frames_num + original_fps = source.original_fps + duration = source.duration + + if fps is not None and (original_fps <= 0 or duration <= 0): + fps = None + + if overflow_strategy is not None: + if max_frames is None: + raise ValueError( + "You must pass max_frames when requesting an " + f"overflow_strategy={overflow_strategy}!" + ) + if overflow_strategy not in ("truncate", "uniform"): + raise ValueError( + f"You passed overflow_strategy={overflow_strategy} but " + "expected one of ['truncate', 'uniform']" + ) + + if fps is None: + if overflow_strategy is None or total_frames <= max_frames: + indices = np.arange(total_frames, dtype=int) + elif overflow_strategy == "truncate": + indices = np.arange(max_frames, dtype=int) + elif overflow_strategy == "uniform": + indices = np.linspace(0, total_frames - 1, max_frames, dtype=int) + else: + step = original_fps / fps + num_sampled = max(1, int(duration * fps)) + + if overflow_strategy is None or num_sampled <= max_frames: + i_arr = np.arange(num_sampled, dtype=float) + indices = np.minimum(total_frames - 1, (i_arr * step).astype(int)) + elif overflow_strategy == "truncate": + i_arr = np.arange(max_frames, dtype=float) + indices = np.minimum(total_frames - 1, (i_arr * step).astype(int)) + elif overflow_strategy == "uniform": + linspace_idx = np.linspace(0, num_sampled - 1, max_frames, dtype=int) + indices = np.minimum( + total_frames - 1, (linspace_idx * step).astype(int) + ) + + return indices.tolist() + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int | float = -1, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: VideoDecoderBackend = "opencv", + **kwargs: Any, + ) -> tuple[DecodedFrames, dict[str, Any]]: + # Accept num_frames as required by VideoLoader interface (VideoMediaIO + # passes default num_frames=32), but ignore it since EmbeddingGemma2 uses + # fps/max_frames. Force num_frames=-1 so compute_frames_index_to_sample + # drives sampling via fps/duration. + kwargs.pop("num_frames", None) + frames, metadata = super().load_bytes( + data, + num_frames=-1, + fps=fps, # type: ignore[arg-type] + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + + source = VideoSourceMetadata( + total_frames_num=metadata["total_num_frames"], + original_fps=metadata["fps"], + duration=metadata["duration"], + ) + target = VideoTargetMetadata(num_frames=-1, fps=fps, max_duration=max_duration) + sampling_kwargs, _ = resolve_video_backend_kwargs(backend, kwargs) + requested_indices = cls.compute_frames_index_to_sample( + source, target, **sampling_kwargs + ) + + # Expand deduplicated frames back to requested_indices if needed + # (e.g., OpenCV deduplicates frame_idx using set(frame_idx)) + valid_frame_indices = metadata.get("frames_indices") + if valid_frame_indices is not None and requested_indices != valid_frame_indices: + idx_to_pos = {idx: pos for pos, idx in enumerate(valid_frame_indices)} + if all(idx in idx_to_pos for idx in requested_indices): + positions = [idx_to_pos[idx] for idx in requested_indices] + if isinstance(frames, np.ndarray): + frames = frames[positions] + else: + frames = frames[torch.tensor(positions, device=frames.device)] + metadata["frames_indices"] = list(requested_indices) + + return frames, metadata + + @classmethod + def create_hf_metadata( + cls, + source: VideoSourceMetadata, + valid_frame_indices: list[int], + video_backend: str, + ) -> dict[str, Any]: + meta = super().create_hf_metadata( + source=source, + valid_frame_indices=valid_frame_indices, + video_backend=video_backend, + ) + meta["do_sample_frames"] = False + return meta + + @VIDEO_LOADER_REGISTRY.register( "qwen2_vl", video_processor="Qwen2VLVideoProcessor", diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 1ed55c3e28b5..9e8482b5ecdd 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -243,9 +243,10 @@ def context_attention_fwd( b_seq_len: [b] out: [b * s, head, head_dim] """ - BLOCK = get_block_size(q.dtype) - Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1] + BLOCK = get_block_size(q.dtype) + if Lk >= 512: + BLOCK = min(BLOCK, 32) sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale # rescale with 1/ln(2) for triton exp2 From 3403e0f176efb90b0d92cc0ff1d350f092ed9e9f Mon Sep 17 00:00:00 2001 From: Madhav Jivrajani <12381034+MadhavJivrajani@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:51:01 -0400 Subject: [PATCH 086/952] [Bugfix] Persist FlashInfer autotune cache per rank to fix the rank-0-only cache-hit deadlock (#57635) Signed-off-by: Madhav Jivrajani Signed-off-by: NickLucche Co-authored-by: NickLucche Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../test_flashinfer_autotune_warmup.py | 91 ++++++++++++++----- vllm/model_executor/warmup/kernel_warmup.py | 63 +++++++++---- 2 files changed, 116 insertions(+), 38 deletions(-) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index 22bb54d71ffd..d65f2a4ad116 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -46,6 +46,7 @@ def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto vllm_config=SimpleNamespace( kernel_config=SimpleNamespace(linear_backend=linear_backend), attention_config=SimpleNamespace(hisparse_config=None), + parallel_config=SimpleNamespace(data_parallel_rank=0), ), get_model=Mock( return_value=SimpleNamespace(modules=Mock(return_value=modules)) @@ -146,12 +147,20 @@ def __init__(self, run, ranks): def record(self, operation): self.run.collectives[self.ranks][self.run.rank].append(operation) - def broadcast_object(self, obj, src=0): - assert src == 0 - self.record(("broadcast", src)) - if self.rank_in_group == src: - self.run.broadcasts[self.ranks] = obj - return self.run.broadcasts[self.ranks] + def all_gather_object(self, out, obj): + # Ranks run sequentially: snapshot every rank's file before any saves. + self.record(("all_gather",)) + if isinstance(obj, bool): + out[:] = [rank not in self.run.rejected_load_ranks for rank in self.ranks] + assert out[self.rank_in_group] == obj + return + out[:] = self.run.gathered.setdefault( + self.ranks, + [ + warmup._autotune_cache_fingerprint(self.run.cache_path(rank)) + for rank in self.ranks + ], + ) def barrier(self): self.record(("barrier",)) @@ -162,10 +171,18 @@ def __init__(self, run): self.run = run self.cache = {} self.loaded = None + self._dirty = False def load_configs(self, path): + if self.run.rank in self.run.rejected_load_ranks: + return False self.loaded = json.loads(Path(path).read_text()) self.cache.update(self.loaded) + return True + + def clear_cache(self): + self.cache.clear() + self._dirty = False def save_configs(self, path): self.run.saves.append((self.run.rank, Path(path), dict(self.cache))) @@ -182,21 +199,27 @@ def profile(self, operation): if group is not None: group.record(("all_reduce", operation, tactic)) self.cache[operation] = self.run.rank // self.run.tp + self._dirty = True class _AutotuneRun: - def __init__(self, pp, tp): + def __init__(self, pp, tp, cache_dir): self.pp, self.tp = pp, tp + self.cache_dir = cache_dir self.rank = 0 self.tuning_group = None self.collectives: dict[tuple[int, ...], dict[int, list[tuple[Any, ...]]]] = ( defaultdict(lambda: defaultdict(list)) ) - self.broadcasts = {} + self.gathered = {} + self.rejected_load_ranks: set[int] = set() self.tuners = {} self.saves = [] self.profile_groups = defaultdict(list) + def cache_path(self, rank): + return self.cache_dir / f"autotune_configs_dp0_rank{rank}.json" + def world(self): return _AutotuneGroup(self, range(self.pp * self.tp)) @@ -235,7 +258,7 @@ def autotune_run(monkeypatch, tmp_path): from vllm.distributed import parallel_state def make_run(*, pp=2, tp=4): - run = _AutotuneRun(pp, tp) + run = _AutotuneRun(pp, tp, tmp_path) autotuner = ModuleType("flashinfer.autotuner") monkeypatch.setattr( autotuner, @@ -251,6 +274,10 @@ def make_run(*, pp=2, tp=4): monkeypatch.setattr(parallel_state, "get_tp_group", run.tensor_group) monkeypatch.setattr(parallel_state, "get_pp_group", run.pipeline_group) monkeypatch.setattr(fi_utils, "autotune", lambda **kwargs: nullcontext()) + monkeypatch.setattr( + "torch.distributed.all_gather_object", + lambda out, obj, group: group.all_gather_object(out, obj), + ) monkeypatch.setattr( warmup, "resolve_flashinfer_autotune_file", @@ -286,19 +313,18 @@ def test_pp_stage_cache_roundtrip_isolated_and_asymmetric_hits_safe( legacy = tmp_path / "autotune_configs.json" legacy.write_text('{"legacy_world_cache": 99}') cold = autotune_run().execute() - assert [rank for rank, _, _ in cold.saves] == [0, 4] - paths = [path for _, path, _ in cold.saves] - assert len(set(paths)) == 2 and legacy not in paths - assert json.loads(legacy.read_text()) == {"legacy_world_cache": 99} - assert cold.saves[0][2] == {"shared_gemm": 0, "pp0_extra_gemm": 0} - assert cold.saves[1][2] == {"shared_gemm": 1} cold.assert_collectives_match() + assert [rank for rank, _, _ in cold.saves] == list(range(8)) + assert [path for _, path, _ in cold.saves] == [cold.cache_path(r) for r in range(8)] + assert json.loads(legacy.read_text()) == {"legacy_world_cache": 99} + stage_caches = [{"shared_gemm": 0, "pp0_extra_gemm": 0}, {"shared_gemm": 1}] + assert all(cache == stage_caches[r // 4] for r, _, cache in cold.saves) warm = autotune_run().execute() warm.assert_collectives_match() - assert not warm.profile_groups + assert not warm.profile_groups and not warm.saves for rank, tuner in warm.tuners.items(): - assert tuner.loaded == cold.saves[rank // 4][2] - paths[1].unlink() + assert tuner.loaded == stage_caches[rank // 4] + cold.cache_path(5).unlink() mixed = autotune_run().execute() mixed.assert_collectives_match() assert set(mixed.profile_groups) == {4, 5, 6, 7} @@ -306,10 +332,33 @@ def test_pp_stage_cache_roundtrip_isolated_and_asymmetric_hits_safe( assert all(mixed.tuners[rank].loaded is None for rank in range(4, 8)) -def test_pp1_retains_world_synchronization_and_existing_cache_name(autotune_run): +def test_pp1_tunes_world_group_and_saves_per_rank(autotune_run): run = autotune_run(pp=1, tp=4).execute() run.assert_collectives_match() - assert [(rank, path.name) for rank, path, _ in run.saves] == [ - (0, "autotune_configs.json") + assert [(rank, path) for rank, path, _ in run.saves] == [ + (rank, run.cache_path(rank)) for rank in range(4) ] assert all(groups == [(0, 1, 2, 3)] for groups in run.profile_groups.values()) + + +def test_mismatched_rank_caches_are_ignored_by_every_rank(autotune_run): + """One rank loading while others tune deadlocks the per-tactic reduce.""" + cold = autotune_run(pp=1, tp=4).execute() + stale = cold.cache_path(2) + stale.write_text(json.dumps({**json.loads(stale.read_text()), "old_gemm": 0})) + rerun = autotune_run(pp=1, tp=4).execute() + rerun.assert_collectives_match() + assert all(tuner.loaded is None for tuner in rerun.tuners.values()) + assert set(rerun.profile_groups) == {0, 1, 2, 3} + + +@pytest.mark.parametrize("pp, tp", [(1, 1), (1, 4), (2, 4)]) +def test_rejected_cache_retunes_every_rank_in_its_tuning_group(autotune_run, pp, tp): + """Matching files can still be incompatible with one rank's runtime.""" + autotune_run(pp=pp, tp=tp).execute() + rerun = autotune_run(pp=pp, tp=tp) + rerun.rejected_load_ranks = {tp - 1} + rerun.execute() + rerun.assert_collectives_match() + assert set(rerun.profile_groups) == set(range(tp)) + assert {rank for rank, _, _ in rerun.saves} == set(range(tp)) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 55d5ea7a7df7..06f04b75d959 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -5,8 +5,10 @@ happen during model execution. """ +import json import sys import time +from pathlib import Path from typing import TYPE_CHECKING import torch @@ -18,7 +20,6 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, - write_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( autotune_hisparse_flashinfer_attention, @@ -408,6 +409,26 @@ def _run_flashinfer_bf16_autotune_dummy_run( ) +def _autotune_cache_fingerprint(path: Path) -> tuple[str, int] | None: + """Identify a saved autotune file by its FlashInfer metadata and size.""" + try: + configs = json.loads(path.read_text()) + metadata = configs.pop("_metadata", None) + except (OSError, ValueError, AttributeError): + return None + return json.dumps(metadata, sort_keys=True), len(configs) + + +def _all_ranks_have_matching_cache(path: Path, group) -> bool: + """True iff every rank in ``group`` has its own, mutually consistent file.""" + fingerprint = _autotune_cache_fingerprint(path) + if group.world_size == 1: + return fingerprint is not None + gathered: list[tuple[str, int] | None] = [None] * group.world_size + torch.distributed.all_gather_object(gathered, fingerprint, group=group.cpu_group) + return fingerprint is not None and all(f == fingerprint for f in gathered) + + def flashinfer_autotune(runner: "GPUModelRunner") -> None: """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, @@ -421,6 +442,12 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: so each stage's TP group tunes separately with its own cache file; otherwise the world group tunes together. Per-tactic timings are averaged over the tuning group so all its ranks select the same tactic. + + Results are persisted per rank: FlashInfer keys MoE entries by tp/ep rank + (``MoERunner.get_cache_key_extras``), so one rank's file only hits on that + rank. A rank with a cache hit skips the per-tactic reduce the others block + in, so ranks keep loaded configs only if every rank in the tuning group + has a matching file and successfully loads it. """ from flashinfer.autotuner import AutoTuner, set_autotune_process_group @@ -447,11 +474,11 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: autotune_kwargs["skip_ops"] = skip_ops cache_path = resolve_flashinfer_autotune_file(runner) - if pp_size > 1: - ranks = "-".join(str(rank) for rank in tune_group.ranks) - cache_path = cache_path.with_name( - f"{cache_path.stem}_tp_{ranks}{cache_path.suffix}" - ) + # The world group only spans DP ranks when vLLM folds DP into it. + dp_rank = runner.vllm_config.parallel_config.data_parallel_rank + cache_path = cache_path.with_name( + f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" + ) if is_leader: logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) @@ -460,16 +487,16 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: # which lead to some EP ranks receiving no tokens and skipping their # MoE kernel entirely, and cause hang due to all-reduce collective # during synchronized autotuning. - # Read cached autotune results and broadcast within the tuning group. - cached_results: bytes | None = None - if is_leader and cache_path.exists(): - with open(cache_path, "rb") as f: - cached_results = f.read() - cached_results = tune_group.broadcast_object(cached_results, src=0) - if cached_results is not None: - write_flashinfer_autotune_cache(cache_path, cached_results) - tune_group.barrier() - tuner.load_configs(str(cache_path)) + if _all_ranks_have_matching_cache(cache_path, tune_group): + loaded = tuner.load_configs(str(cache_path)) + if tune_group.world_size > 1: + loaded_by_rank: list[bool | None] = [None] * tune_group.world_size + torch.distributed.all_gather_object( + loaded_by_rank, loaded, group=tune_group.cpu_group + ) + loaded = all(loaded_by_rank) + if not loaded: + tuner.clear_cache() group = tune_group.cpu_group if tune_group.world_size > 1 else None set_autotune_process_group(group) @@ -497,5 +524,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: if world.world_size > 1: world.barrier() - if is_leader: + # Skip the rewrite when nothing was tuned this start (every entry came from + # the file). FlashInfer gates its own autotune(cache=...) save the same way. + if tuner._dirty: tuner.save_configs(str(cache_path)) From ad0f67a7ec8f0393c8b4ae50a2a0e66672db372f Mon Sep 17 00:00:00 2001 From: Nishant Revur Date: Tue, 6 Oct 2026 11:03:44 -0700 Subject: [PATCH 087/952] [Bugfix][Spec Decode] Keep heterogeneous-vocab draft models on Model Runner V1 (#59541) Signed-off-by: Nishant Revur Co-authored-by: Claude --- tests/test_config.py | 24 ++++++++++++++++++++++++ vllm/config/vllm.py | 5 +++++ 2 files changed, 29 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index 9d4567874b4f..a0c64099dc32 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -936,6 +936,30 @@ def test_v2_model_runner_supports_extract_hidden_states(): assert config._get_v2_model_runner_unsupported_features() == [] +@pytest.mark.parametrize("use_heterogeneous_vocab", [False, True]) +def test_v2_model_runner_heterogeneous_vocab_draft_falls_back_to_v1( + use_heterogeneous_vocab, +): + """The V2 draft-model speculator exchanges token ids with the target + without a vocab mapping, so a heterogeneous-vocab (TLI) draft must stay on + V1, where VocabMapping translates them.""" + config = VllmConfig() + config.speculative_config = cast( + SpeculativeConfig, + SimpleNamespace( + method="draft_model", + parallel_drafting=False, + enable_adaptive_verification=False, + use_heterogeneous_vocab=use_heterogeneous_vocab, + ), + ) + + expected = ["heterogeneous-vocabulary draft models"] + assert config._get_v2_model_runner_unsupported_features() == ( + expected if use_heterogeneous_vocab else [] + ) + + def test_v2_model_runner_supports_custom_logits_processors(): config = VllmConfig() config.model_config = cast( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index dec1f27fe1c9..b847da4320ab 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3159,6 +3159,11 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: ): unsupported.append("parallel drafting for EAGLE speculative decoding") + # The V2 draft-model speculator has no token mapping between the + # draft and target vocabularies (TLI is TBD in #47172). + if getattr(speculative_config, "use_heterogeneous_vocab", False): + unsupported.append("heterogeneous-vocabulary draft models") + if self.parallel_config.use_ubatching: unsupported.extend(self._get_dbo_unsupported_features()) From e37e51dd246cf421c554e7d6f53e178e3fd29085 Mon Sep 17 00:00:00 2001 From: Cameron Quilici Date: Tue, 6 Oct 2026 13:52:05 -0500 Subject: [PATCH 088/952] [Metrics] Expose cached prompt tokens by cache tier (#56318) Signed-off-by: Cam Quilici Signed-off-by: Cam Quilici Co-authored-by: Cam Quilici Co-authored-by: Nick Hill Co-authored-by: Yifan Qiao --- rust/src/engine-core-client/src/metrics.rs | 28 +- .../engine-core-client/src/protocol/stats.rs | 59 ++++ rust/src/metrics/src/scheduler.rs | 16 + tests/v1/core/test_scheduler.py | 62 +++- .../offloading_connector/test_scheduler.py | 288 +++++++++++++++++- .../unit/offloading_connector/utils.py | 2 + .../unit/test_cache_pollution_prevention.py | 4 +- .../unit/test_error_propagation.py | 6 +- .../unit/test_invalid_blocks_correctness.py | 8 +- .../unit/test_kv_load_failure_recovery.py | 12 +- .../unit/test_mooncake_store_connector.py | 24 ++ .../kv_connector/unit/test_multi_connector.py | 13 + .../unit/test_nixl_simple_cpu_offload.py | 7 + tests/v1/kv_connector/unit/utils.py | 10 + tests/v1/kv_offload/tiering/test_factory.py | 22 +- .../tiering/test_tiering_offloading.py | 41 ++- tests/v1/metrics/test_stats.py | 93 ++++++ tests/v1/simple_kv_offload/test_scheduler.py | 33 ++ tests/v1/worker/test_gpu_profiler.py | 1 + .../kv_transfer/kv_connector/v1/base.py | 23 ++ .../v1/mooncake/mooncake_connector.py | 3 + .../v1/moriio/moriio_connector.py | 3 + .../kv_connector/v1/multi_connector.py | 13 + .../kv_connector/v1/nixl/connector.py | 3 + .../kv_connector/v1/offloading/scheduler.py | 78 ++++- .../kv_connector/v1/offloading_connector.py | 11 + .../v1/simple_cpu_offload_connector.py | 4 + vllm/v1/core/sched/scheduler.py | 12 +- vllm/v1/kv_offload/base.py | 13 + vllm/v1/kv_offload/cpu/manager.py | 7 + vllm/v1/kv_offload/tiering/base.py | 12 + vllm/v1/kv_offload/tiering/manager.py | 28 ++ vllm/v1/kv_offload/tiering/p2p/manager.py | 5 +- vllm/v1/metrics/cache_hit_source.py | 32 ++ vllm/v1/metrics/loggers.py | 34 ++- vllm/v1/metrics/stats.py | 39 ++- 36 files changed, 1005 insertions(+), 44 deletions(-) create mode 100644 vllm/v1/metrics/cache_hit_source.py diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index b392d5260185..e8b0741066f1 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -6,14 +6,15 @@ use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ - EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, - LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, - MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, - U64Counter, U64Gauge, WaitingReasonLabels, + CacheHitSourceLabels, EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, + LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, + MooncakeOperationHistogramFamily, MooncakeOperationLabels, RequestMetrics, + SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::{ - KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, SchedulerStats, + CachedTokensBySource, KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, + SchedulerStats, }; use crate::transport::ConnectedEngine; @@ -60,6 +61,8 @@ struct SchedulerStatsHandles { prefix_cache_hits: U64Counter, external_prefix_cache_queries: U64Counter, external_prefix_cache_hits: U64Counter, + /// One handle per `CachedTokensBySource::SOURCES` entry. + prompt_tokens_cached_by_source: [U64Counter; 5], // Speculative decoding counters. spec_decode_num_drafts: U64Counter, @@ -165,6 +168,15 @@ fn resolve_scheduler_stats_handles( .external_prefix_cache_queries .get_or_create_owned(&labels), external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels), + prompt_tokens_cached_by_source: CachedTokensBySource::SOURCES.map(|source| { + metrics + .prompt_tokens_cached_by_source + .get_or_create_owned(&CacheHitSourceLabels { + model_name: model_name.to_string(), + engine, + source, + }) + }), spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels), spec_decode_num_draft_tokens: metrics .spec_decode_num_draft_tokens @@ -220,6 +232,8 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & // Prefix-cache counters, including the connector-backed external cache path. handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries); handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits); + // Local prefix-cache hits are the `device` tier. + handles.prompt_tokens_cached_by_source[0].inc_by(stats.prefix_cache_stats.base.hits); if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats { handles @@ -228,6 +242,10 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & handles .external_prefix_cache_hits .inc_by(connector_prefix_cache_stats.base.hits); + let counts = connector_prefix_cache_stats.hits_by_source.counts(); + for (handle, count) in handles.prompt_tokens_cached_by_source.iter().zip(counts) { + handle.inc_by(count); + } } // Speculative decoding counters. diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs index aa1b9c69ab59..7797ecf6e48f 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -21,6 +21,34 @@ pub struct BaseCacheStats { pub hits: u64, } +/// Cache-hit tokens per source tier. Python sends a sparse +/// `dict[CacheHitSource, int]`, so absent tiers default to zero. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default)] +pub struct CachedTokensBySource { + pub device: u64, + pub host: u64, + pub p2p: u64, + pub disk: u64, + pub external_unspecified: u64, +} + +impl CachedTokensBySource { + /// Label values, in the same order as [`Self::counts`]. + pub const SOURCES: [&'static str; 5] = + ["device", "host", "p2p", "disk", "external_unspecified"]; + + pub fn counts(&self) -> [u64; 5] { + [ + self.device, + self.host, + self.p2p, + self.disk, + self.external_unspecified, + ] + } +} + /// Stores prefix cache hit statistics. /// - `reset`: Whether `reset_prefix_cache` was invoked. /// - `queries`: Refers to the number of tokens that were queried. @@ -38,6 +66,9 @@ pub struct PrefixCacheStats { pub preempted_queries: u64, /// The `hits` number for preempted requests. pub preempted_hits: u64, + /// `hits` split by the cache tier that supplied them (connector stats only). + #[serde(default)] + pub hits_by_source: CachedTokensBySource, } /// Single KV cache block eviction sample. @@ -304,3 +335,31 @@ pub struct SchedulerStats { /// Estimated MFU/performance stats, when enabled. pub perf_stats: Option, } + +#[cfg(test)] +mod tests { + use super::{CachedTokensBySource, PrefixCacheStats}; + + #[test] + fn prefix_cache_stats_hits_by_source_defaults_and_decodes() { + // Engines without per-source attribution omit the field entirely. + let legacy = serde_json::json!({ + "reset": false, "requests": 1, "queries": 8, "hits": 4, + "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0 + }); + let wire = rmp_serde::to_vec_named(&legacy).unwrap(); + let stats: PrefixCacheStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!(stats.base.hits, 4); + assert_eq!(stats.hits_by_source, CachedTokensBySource::default()); + + let payload = serde_json::json!({ + "reset": false, "requests": 1, "queries": 16, "hits": 12, + "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0, + "hits_by_source": {"host": 8, "p2p": 4} + }); + let wire = rmp_serde::to_vec_named(&payload).unwrap(); + let stats: PrefixCacheStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!(stats.hits_by_source.counts(), [0, 8, 4, 0, 0]); + assert_eq!(CachedTokensBySource::SOURCES[1], "host"); + } +} diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index cdc25a244dc8..64ca46541889 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -86,6 +86,13 @@ pub struct EngineLabels { pub engine: u32, } +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct CacheHitSourceLabels { + pub model_name: String, + pub engine: u32, + pub source: &'static str, +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct EnginePositionLabels { pub model_name: String, @@ -236,6 +243,7 @@ pub struct SchedulerMetrics { pub prefix_cache_hits: Family, pub external_prefix_cache_queries: Family, pub external_prefix_cache_hits: Family, + pub prompt_tokens_cached_by_source: Family, // Speculative decoding counters. pub spec_decode_num_drafts: Family, @@ -339,6 +347,13 @@ impl SchedulerMetrics { ); let external_prefix_cache_hits = Family::default(); + + let prompt_tokens_cached_by_source = Family::default(); + registry.register( + "vllm:prompt_tokens_cached_by_source", + "Prefix-cache hit tokens by the cache tier that supplied them: device (local HBM), host (offloaded to DRAM), disk, p2p (transferred from another vLLM instance) or external_unspecified. Counted at admission; sums to vllm:prefix_cache_hits + vllm:external_prefix_cache_hits.", + prompt_tokens_cached_by_source.clone(), + ); registry.register( "vllm:external_prefix_cache_hits", "External prefix cache hits from KV connector cross-instance cache sharing, in terms of number of cached tokens.", @@ -528,6 +543,7 @@ impl SchedulerMetrics { prefix_cache_hits, external_prefix_cache_queries, external_prefix_cache_hits, + prompt_tokens_cached_by_source, spec_decode_num_drafts, spec_decode_num_draft_tokens, spec_decode_num_accepted_tokens, diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 66368a299da5..f78db1917db0 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3,7 +3,7 @@ import dataclasses from concurrent.futures import Future from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import numpy as np import pytest @@ -60,6 +60,7 @@ KVCacheGroupSpec, MambaSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import ( DraftTokenIds, ECConnectorOutput, @@ -2613,6 +2614,65 @@ def test_kv_connector_honors_skip_reading_prefix_cache(): assert output.num_scheduled_tokens[scoring.request_id] == BLOCK_SIZE * 4 +@pytest.mark.parametrize("is_async", [False, True]) +def test_kv_connector_records_external_cache_hit_sources(monkeypatch, is_async): + block_size = 16 + num_matched_tokens = 2 * block_size + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv( + matched_tokens=num_matched_tokens, + is_async=is_async, + ), + block_size=block_size, + ) + request = create_requests( + num_requests=1, + num_tokens=2 * num_matched_tokens, + block_size=block_size, + )[0] + assert scheduler.connector is not None + get_sources = Mock( + return_value={CacheHitSource.P2P: block_size, CacheHitSource.HOST: block_size} + ) + update_state = Mock(wraps=scheduler.connector.update_state_after_alloc) + calls = Mock() + calls.attach_mock(update_state, "allocated") + calls.attach_mock(get_sources, "sources") + monkeypatch.setattr(scheduler.connector, "update_state_after_alloc", update_state) + monkeypatch.setattr( + scheduler.connector, "get_external_cache_hit_sources", get_sources + ) + + scheduler.add_request(request) + # A failed allocation is not an admission: nothing is attributed. + with patch.object(scheduler.kv_cache_manager, "allocate_slots", return_value=None): + scheduler.schedule() + get_sources.assert_not_called() + update_state.assert_not_called() + + scheduler.schedule() + # Attribution runs once, after the connector has built its load plan. + get_sources.assert_called_once_with(request, num_matched_tokens) + assert [c[0] for c in calls.mock_calls] == ["allocated", "sources"] + connector_stats = scheduler.connector_prefix_cache_stats + assert connector_stats is not None + assert connector_stats.hits == num_matched_tokens + assert connector_stats.hits_by_source == { + CacheHitSource.P2P: block_size, + CacheHitSource.HOST: block_size, + } + + if is_async: + # Re-admission after the async load completes is not a new hit. + scheduler.make_stats() + _step_until_kv_transfer_finished(scheduler, [request.request_id]) + get_sources.assert_called_once() + connector_stats = scheduler.connector_prefix_cache_stats + assert connector_stats is not None + assert connector_stats.hits_by_source == {} + + @pytest.mark.parametrize("is_async", [False, True]) def test_kv_connector_basic(is_async: bool): """Test whether Scheduler with KVConnector schedules tokens, allocates diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 5b3bf1c8bebe..e71c3dbf18d4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -77,6 +77,7 @@ make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -197,14 +198,18 @@ def test_swa_offload_window_covers_unaligned_hit( manager.complete_load(job.keys, state.req_context) -def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: +def _make_partial_tail_scheduler( + recurrent_only: bool = False, +) -> OffloadingConnectorScheduler: vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True}) - vllm_config.cache_config.prefix_match_unit = 4 + vllm_config.cache_config.prefix_match_unit = None if recurrent_only else 4 vllm_config.speculative_config = None vllm_config.kv_events_config = KVEventsConfig( enable_kv_cache_events=True, publisher="null" ) kv_cache_config = _make_mamba_hybrid_kv_cache_config() + if recurrent_only: + kv_cache_config.kv_cache_groups = kv_cache_config.kv_cache_groups[1:] spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) @@ -576,6 +581,284 @@ def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): assert 42 in loaded +@pytest.mark.parametrize( + ("recurrent_source", "expected"), + [ + ("host", {CacheHitSource.P2P: 16, CacheHitSource.HOST: 12}), + ("p2p", {CacheHitSource.P2P: 28}), + (None, {CacheHitSource.P2P: 16, CacheHitSource.HOST: 12}), + ], +) +def test_external_cache_hit_sources_preserve_partial_tail_origin( + recurrent_source, expected +): + """Recurrent state covers the whole prefix; the outermost tier wins.""" + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.update_offload_keys() + + scheduler.manager.lookup.return_value = LookupResult.HIT + assert scheduler._lookup(req_status) == 28 + + def get_load_source(key, req_context): + if get_offload_group_idx(key) == 1: + assert recurrent_source is not None + return recurrent_source + block_hash = get_offload_block_hash(key) + return "p2p" if block_hash == b"h3" else "host" + + scheduler.manager.get_load_source.side_effect = get_load_source + + recurrent_block = KVCacheBlock(41) + if recurrent_source is None: + recurrent_block.set_block_hash( + make_block_hash_with_group_id(BlockHash(b"cached"), 1), 12 + ) + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32)], + [KVCacheBlock(0, is_null=True), recurrent_block], + ) + ), + num_external_tokens=28, + ) + assert req_status.partial_tail_boundary is None + assert scheduler.get_external_cache_hit_sources(request, 28) == expected + + +@pytest.mark.parametrize( + ("second_source", "expected"), + [ + ("host", {CacheHitSource.HOST: 4, CacheHitSource.DISK: 4}), + ("p2p", {CacheHitSource.P2P: 4, CacheHitSource.DISK: 4}), + ], +) +def test_external_cache_hit_sources_reconcile_different_group_chunk_sizes( + request_runner, + second_source, + expected, +): + """Groups with different chunk sizes overlap; the outermost tier wins.""" + block_size = 4 + groups = [ + KVCacheGroupSpec( + [f"layer{group_idx}"], + FullAttentionSpec( + block_size=block_size * (group_idx + 1), + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx in range(2) + ] + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=False, + kv_cache_groups=groups, + ) + scheduler = runner.connector_scheduler + request = MagicMock(request_id="mixed") + req_status = RequestOffloadState( + config=scheduler.config, + req=request, + req_context=ReqContext(req_id=request.request_id), + offloading_context=RequestOffloadingContext(policy=OffloadPolicy.CHUNK_LEVEL), + ) + for group_idx, group_state in enumerate(req_status.group_states): + group_state.offload_keys = [ + make_offload_key(f"h{chunk_idx}".encode(), group_idx) + for chunk_idx in range(2 // (group_idx + 1)) + ] + scheduler._req_status[request.request_id] = req_status + runner.manager.get_load_source.side_effect = lambda key, req_context: ( + ("host" if get_offload_block_hash(key) == b"h0" else "disk") + if get_offload_group_idx(key) == 0 + else second_source + ) + scheduler.update_state_after_alloc( + request, + KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(3)])), + 2 * block_size, + ) + + assert scheduler.get_external_cache_hit_sources(request, 2 * block_size) == expected + + +@pytest.mark.parametrize( + ("ranges", "expected"), + [ + pytest.param( + [(8, 12, "host"), (0, 8, "host"), (4, 10, "host")], + {CacheHitSource.HOST: 8, CacheHitSource.EXTERNAL_UNSPECIFIED: 4}, + id="overlapping-same-source-and-adjacent-endpoints", + ), + pytest.param( + [(10, 20, "p2p"), (6, 10, "disk"), (0, 8, "host"), (4, 12, "host")], + {CacheHitSource.HOST: 2, CacheHitSource.DISK: 4, CacheHitSource.P2P: 6}, + id="nested-tiers-and-clipped-ranges", + ), + pytest.param( + [(0, 4, "disk"), (16, 20, "disk"), (8, 8, "disk")], + {CacheHitSource.EXTERNAL_UNSPECIFIED: 12}, + id="no-covering-ranges", + ), + ], +) +def test_external_cache_hit_sources_range_boundaries(ranges, expected): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + state = scheduler._req_status[request.request_id] + state.num_locally_computed_tokens = 4 + for idx, (lo, hi, source) in enumerate(ranges): + key = make_offload_key(str(idx).encode(), 0) + state.load_key_ranges.append((lo, hi, key)) + scheduler.manager.get_load_source.side_effect = lambda key, ctx: ranges[ + int(get_offload_block_hash(key)) + ][2] + + result = scheduler.get_external_cache_hit_sources(request, 12) + assert result == expected + assert sum(result.values()) == 12 + + +@pytest.mark.parametrize("source", ["host", "disk", "p2p", "external_unspecified"]) +def test_external_cache_hit_sources_recurrent_only_state(source): + scheduler = _make_partial_tail_scheduler(recurrent_only=True) + request = _make_partial_tail_request(scheduler) + request.num_prompt_tokens = request.num_tokens = 64 + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(4)] + request.all_token_ids = list(range(64)) + state = scheduler._req_status[request.request_id] + state.update_offload_keys() + scheduler.manager.lookup.return_value = LookupResult.HIT + assert scheduler._lookup(state) == 48 + scheduler.manager.get_load_source.return_value = source + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [ + KVCacheBlock(0, is_null=True), + KVCacheBlock(0, is_null=True), + KVCacheBlock(41), + ], + ) + ), + 48, + ) + assert scheduler.get_external_cache_hit_sources(request, 48) == { + CacheHitSource(source): 48 + } + scheduler.manager.get_load_source.assert_called_once() + + +@pytest.mark.parametrize("sparse_kind", ["sliding", "chunked_local"]) +@pytest.mark.parametrize("blocks_per_chunk", [1, 2]) +@pytest.mark.parametrize("local_tokens", [0, 4]) +@pytest.mark.parametrize("full_attention", [False, True]) +@pytest.mark.parametrize( + "sparse_source", ["host", "disk", "p2p", "external_unspecified", "mixed"] +) +def test_external_cache_hit_sources_use_loaded_attention_ranges( + sparse_kind, + blocks_per_chunk, + local_tokens, + full_attention, + sparse_source, +): + """SWA/local attention affects only loaded positions, not omitted history.""" + block_size = 4 + chunk_size = block_size * blocks_per_chunk + spec_kwargs = dict( + block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + sparse_spec = ( + SlidingWindowSpec(**spec_kwargs, sliding_window=2 * chunk_size) + if sparse_kind == "sliding" + else ChunkedLocalAttentionSpec( + **spec_kwargs, attention_chunk_size=2 * chunk_size + ) + ) + groups = [KVCacheGroupSpec(["sparse"], sparse_spec)] + if full_attention: + groups.insert(0, KVCacheGroupSpec(["full"], FullAttentionSpec(**spec_kwargs))) + vllm_config = _make_vllm_config(extra_config={"block_size": chunk_size}) + vllm_config.cache_config.block_size = block_size + vllm_config.speculative_config = None + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) + scheduler = OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config) + request = MagicMock(request_id="sparse") + req_status = RequestOffloadState( + config=scheduler.config, + req=request, + req_context=ReqContext(req_id=request.request_id), + offloading_context=RequestOffloadingContext(policy=OffloadPolicy.CHUNK_LEVEL), + num_locally_computed_tokens=local_tokens, + ) + scheduler._req_status[request.request_id] = req_status + block_groups = [] + for group_idx, state in enumerate(req_status.group_states): + state.offload_keys = [ + make_offload_key(str(i).encode(), group_idx) for i in range(4) + ] + block_groups.append( + [ + KVCacheBlock( + 0 + if group_idx == len(groups) - 1 and i < 2 * blocks_per_chunk + else 1 + group_idx * 100 + i, + is_null=group_idx == len(groups) - 1 and i < 2 * blocks_per_chunk, + ) + for i in range(4 * blocks_per_chunk) + ] + ) + + def get_load_source(key, req_context): + chunk_idx = int(get_offload_block_hash(key)) + if get_offload_group_idx(key) == len(groups) - 1: + assert chunk_idx >= 2 # Earlier sparse keys were not loaded. + if sparse_source == "mixed": + return "host" if chunk_idx == 2 else "disk" + return sparse_source + return "host" + + scheduler.manager.get_load_source.side_effect = get_load_source + count = 4 * chunk_size - local_tokens + scheduler.update_state_after_alloc( + request, KVCacheBlocks(tuple(block_groups)), count + ) + expected: dict[CacheHitSource, int] = {} + + def expect(source: CacheHitSource, num_tokens: int) -> None: + expected[source] = expected.get(source, 0) + num_tokens + + expect( + CacheHitSource.HOST if full_attention else CacheHitSource.EXTERNAL_UNSPECIFIED, + 2 * chunk_size - local_tokens, + ) + for chunk_idx in range(2, 4): + source = ( + ("host" if chunk_idx == 2 else "disk") + if sparse_source == "mixed" + else sparse_source + ) + expect(CacheHitSource(source), chunk_size) + result = scheduler.get_external_cache_hit_sources(request, count) + assert result == expected + assert sum(result.values()) == count + assert scheduler.get_external_cache_hit_sources(request, 0) == {} + + def test_partial_lookup_requires_every_cache_group(): scheduler = _make_partial_tail_scheduler() _make_partial_tail_request(scheduler) @@ -1400,7 +1683,6 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # are within the window → loads blocks 1,2 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), ) - runner.manager.touch.assert_not_called() # 3 blocks are hit on GPU [0, 1, 2] diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index cb766c6676c8..9621890f3c26 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -57,6 +57,7 @@ make_offload_key, ) from vllm.v1.kv_offload.config import OffloadingConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -129,6 +130,7 @@ def __init__(self, config: OffloadingConfig): self.manager = MagicMock(spec=OffloadingManager) self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.manager.lookup.return_value = LookupResult.MISS + self.manager.get_load_source.return_value = CacheHitSource.EXTERNAL_UNSPECIFIED self.manager.get_stats.return_value = None self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() diff --git a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py index 21b50a13794b..0515994543b6 100644 --- a/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py +++ b/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py @@ -8,7 +8,6 @@ """ from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -16,6 +15,7 @@ from vllm.v1.request import Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -73,7 +73,7 @@ def test_invalid_blocks_evicted_prevents_cache_pollution( } # mock connector indicating sync load - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py index 35a3749e082a..ded110d55b6b 100644 --- a/tests/v1/kv_connector/unit/test_error_propagation.py +++ b/tests/v1/kv_connector/unit/test_error_propagation.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -10,6 +9,7 @@ from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -56,7 +56,7 @@ def test_error_propagation_sync_load(fail_scheduler: Scheduler): request.request_id: num_external_computed_tokens, } - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -110,7 +110,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler): request.request_id: num_external_computed_tokens, } - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 7ea2170584ad..f3283326962f 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -11,7 +11,6 @@ """ from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -19,6 +18,7 @@ from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -83,7 +83,7 @@ def test_sync_recompute_blocks_not_freed_for_running_requests( } # mock connector indicating sync load - recompute_scheduler.connector = Mock() + recompute_scheduler.connector = create_mock_connector() recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -206,7 +206,7 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler): } # mock connector indicating sync load - fail_scheduler.connector = Mock() + fail_scheduler.connector = create_mock_connector() fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False) ) @@ -323,7 +323,7 @@ def test_async_recompute_blocks_not_cached_when_invalid( } # mock connector indicating async load - recompute_scheduler.connector = Mock() + recompute_scheduler.connector = create_mock_connector() recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) ) diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 3dbbfdc6b286..5dce592576df 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock import pytest @@ -10,6 +9,7 @@ from vllm.v1.request import Request, RequestStatus from .utils import ( + create_mock_connector, create_model_runner_output, create_request, create_scheduler, @@ -41,7 +41,7 @@ def test_failed_receive_completion_honors_failure_policy(policy): scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy)) request = create_request(num_tokens=3 * scheduler.block_size) scheduler.add_request(request) - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = [ (2 * scheduler.block_size, True), @@ -103,7 +103,7 @@ def test_async_load_failure( request3.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) @@ -181,7 +181,7 @@ def test_sync_load_failure( request3.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False) ) @@ -262,7 +262,7 @@ def test_sync_load_failure_with_shared_blocks( request1.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False) ) @@ -333,7 +333,7 @@ def test_async_progressive_load_failure( request.request_id: num_external_computed_tokens, } - scheduler.connector = Mock() + scheduler.connector = create_mock_connector() scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index f3ffa55a5439..824f806f487b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -38,6 +38,7 @@ KVCacheTensor, MambaSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from .utils import create_vllm_config @@ -208,6 +209,29 @@ def test_scheduler_role_initializes_store_scheduler_only(): mock_scheduler.return_value.bind_gpu_block_pool.assert_called_once_with(block_pool) +def test_scheduler_reports_mooncake_cache_source(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreScheduler" + ), + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + + # Mooncake Store is an external cache service. Its internal memory and + # local-storage policy is intentionally hidden behind that stable label. + assert connector.get_external_cache_hit_sources( + None, # type: ignore[arg-type] + 32, + ) == {CacheHitSource.EXTERNAL_UNSPECIFIED: 32} + + def test_worker_methods_delegate_to_store_worker(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 67cfdc1ab029..e576e3d0d9e6 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -34,6 +34,7 @@ NixlKVConnectorStats, ) from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, KVConnectorWorkerMetadata MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" @@ -246,6 +247,18 @@ def test_multi_connector_rejects_multiple_mem_pool_contexts(mc: MultiConnector): mc.get_mem_pool_context() +def test_cache_hit_sources_delegate_to_selected_connector(mc: MultiConnector): + request = MagicMock(request_id="request") + mc._requests_to_connector[request.request_id] = 1 + sources = {CacheHitSource.DISK: 32} + mc._connectors[1].get_external_cache_hit_sources.return_value = sources + + assert mc.get_external_cache_hit_sources(request, 32) is sources + mc._connectors[1].get_external_cache_hit_sources.assert_called_once_with( + request, 32 + ) + + # Helper function to compare directories recursively def _compare_directories(dir1: Path, dir2: Path) -> bool: """Compares two directories recursively for identical content.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index bdab91a21ffe..5ae235c24ddb 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -36,6 +36,7 @@ KVCacheTensor, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from vllm.v1.simple_kv_offload.metadata import ( @@ -232,6 +233,9 @@ def test_nixl_wins_load_over_cpu_offload(): assert not sched_out.has_sync_kv_loads assert mc._requests_to_connector[request.request_id] == 0 + assert mc.get_external_cache_hit_sources(request, 2 * BLOCK_SIZE) == { + CacheHitSource.P2P: 2 * BLOCK_SIZE + } meta = sched_out.kv_connector_metadata assert isinstance(meta, MultiKVConnectorMetadata) @@ -302,6 +306,9 @@ def test_cpu_offload_wins_when_nixl_has_no_match(): assert hit_tokens is not None and hit_tokens > 0 assert mc._requests_to_connector[req2.request_id] == 1 assert is_async is True + assert mc.get_external_cache_hit_sources(req2, hit_tokens) == { + CacheHitSource.HOST: hit_tokens + } @pytest.mark.parametrize("swa_enabled", [False, True], ids=["fa_only", "fa_sw"]) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 6f27212281e9..86062fc47848 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from itertools import chain, count from typing import Any, Literal +from unittest.mock import Mock import torch @@ -42,6 +43,7 @@ MambaSpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -49,6 +51,14 @@ EOS_TOKEN_ID = 50256 +def create_mock_connector() -> Mock: + connector = Mock() + connector.get_external_cache_hit_sources.side_effect = lambda request, num_tokens: ( + {CacheHitSource.EXTERNAL_UNSPECIFIED: num_tokens} if num_tokens else {} + ) + return connector + + def assert_scheduler_empty(scheduler: Scheduler): """Confirm the scheduler is "empty" - i.e. no leaks.""" # Scheduler Metadata. diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f91f54bb0d16..262af6cb833f 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -17,6 +17,7 @@ from vllm.v1.kv_offload.tiering.base import SecondaryTierManager from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource # --------------------------------------------------------------------------- # Fixtures @@ -43,12 +44,22 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): """CI sentinel: registered tiers import and yield SecondaryTierManager.""" + # Derived from each tier's ``medium``; P2P has none and sets its own. + expected_sources = { + "example": CacheHitSource.HOST, + "fs": CacheHitSource.DISK, + "obj": CacheHitSource.DISK, + "p2p": CacheHitSource.P2P, + } for tier_type in SecondaryTierFactory._registry: # KVCR is an optional external dependency and may not be installed. if tier_type == "kvcr" and find_spec("kvcr") is None: continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) + assert cls.cache_hit_source == expected_sources.get( + tier_type, CacheHitSource.EXTERNAL_UNSPECIFIED + ) def test_example_tier_registered(): @@ -95,7 +106,8 @@ def test_create_multiple_tiers(): assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers) -def test_register_new_tier_type(): +@pytest.mark.parametrize("tier_type", ["custom_tier", "cpu", "nvme"]) +def test_register_new_tier_type(tier_type): """Verify that new tier types can be registered and created. This is how external projects add custom secondary tiers @@ -103,20 +115,22 @@ def test_register_new_tier_type(): """ # Register a new tier type (reuse example manager for simplicity) SecondaryTierFactory.register_tier( - "custom_tier", + tier_type, "vllm.v1.kv_offload.tiering.example.manager", "ExampleSecondaryTierManager", ) primary_kv_view, offloading_spec = _make_mock_args() tier = SecondaryTierFactory.create_secondary_tier( - {"type": "custom_tier", "custom_param": 99}, + {"type": tier_type, "custom_param": 99}, primary_kv_view, offloading_spec, ) - assert tier.tier_type == "custom_tier" + assert tier.tier_type == tier_type assert isinstance(tier, ExampleSecondaryTierManager) + # Attribution follows the tier's medium, never its registration name. + assert tier.cache_hit_source is CacheHitSource.HOST # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 63b4c2c4d0b5..93b49afb4f44 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -53,6 +53,7 @@ TieringOffloadingManager, ) from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec +from vllm.v1.metrics.cache_hit_source import CacheHitSource _CTX = ReqContext(req_id="test") _MOCK_OFFLOADING_SPEC = MagicMock() @@ -500,8 +501,13 @@ def test_lookup_from_primary(self, manager_setup): # Lookup should find all chunks in primary assert count_hits(self.manager, chunks) == 3 - def test_promotion_from_secondary(self, manager_setup): + @pytest.mark.parametrize( + "source", + [CacheHitSource.EXTERNAL_UNSPECIFIED, CacheHitSource.DISK, CacheHitSource.P2P], + ) + def test_promotion_from_secondary(self, manager_setup, monkeypatch, source): """Test promotion of chunks from secondary to primary tier.""" + monkeypatch.setattr(ExampleSecondaryTierManager, "cache_hit_source", source) chunks = to_keys(range(3)) # Manually add chunks to secondary tier (simulate previous cascade) @@ -525,6 +531,27 @@ def test_promotion_from_secondary(self, manager_setup): # Next lookup should succeed assert count_hits(self.manager, chunks) == 3 + # The request that caused the promotion retains the secondary origin. + assert all( + self.manager.get_load_source(chunk, _CTX) == source for chunk in chunks + ) + + # A later request sees the chunks as ordinary host-memory primary hits. + later_context = ReqContext(req_id="later") + self.manager.on_new_request(later_context) + assert all( + self.manager.lookup(chunk, later_context) is LookupResult.HIT + for chunk in chunks + ) + assert all( + self.manager.get_load_source(chunk, later_context) == "host" + for chunk in chunks + ) + + @pytest.mark.parametrize( + "source", + [CacheHitSource.EXTERNAL_UNSPECIFIED, CacheHitSource.DISK, CacheHitSource.P2P], + ) @pytest.mark.parametrize( ("successful_indices", "expected_results"), [ @@ -540,8 +567,9 @@ def test_promotion_from_secondary(self, manager_setup): ids=["partial", "legacy-full-failure"], ) def test_failed_promotion_keeps_only_successful_chunks( - self, manager_setup, successful_indices, expected_results + self, manager_setup, monkeypatch, source, successful_indices, expected_results ): + monkeypatch.setattr(ExampleSecondaryTierManager, "cache_hit_source", source) chunks = to_keys(range(3)) for chunk in chunks: self.secondary_tier1.chunks[chunk] = True @@ -571,6 +599,14 @@ def submit_partial(job_metadata: TransferJob) -> None: assert [ self.primary_tier.lookup(chunk, _CTX) for chunk in chunks ] == expected_results + expected_source_keys = ( + set() + if successful_indices is None + else {chunks[i] for i in successful_indices} + ) + assert set(self.manager._request_load_sources[_CTX.req_id]) == ( + expected_source_keys + ) def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): """Resolved lookups report one sync delay sample per tier and chunk.""" @@ -1315,6 +1351,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): # Orchestrator state cleared. assert self.manager._jobs == {} assert self.manager._pending_load_submissions == {} + assert self.manager._request_load_sources == {} assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id} assert self.manager._processed_jobs_this_step is False diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 0de74f0faaa4..c51c223d28e7 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from vllm.v1.core.sched.output import ScheduledEncoderInputStats, SchedulerOutput from vllm.v1.engine import EngineCoreOutputs, FinishReason +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.stats import ( IterationStats, PrefillStats, + PrefixCacheStats, PromptTokenStats, RequestStateStats, SchedulerIterationDetails, @@ -45,6 +48,96 @@ def test_scheduler_iteration_details_serialization(): assert decoded.scheduler_stats.iteration_details == iteration_details +def test_prefix_cache_stats_hits_by_source_accumulate(): + stats = PrefixCacheStats() + stats.record( + num_tokens=8, + num_hits=5, + preempted=False, + hits_by_source={CacheHitSource.HOST: 3, CacheHitSource.DISK: 2}, + ) + stats.record( + num_tokens=8, + num_hits=6, + preempted=False, + hits_by_source={CacheHitSource.P2P: 5, CacheHitSource.HOST: 1}, + ) + # Zero counts never create an entry. + stats.record( + num_tokens=8, + num_hits=4, + preempted=False, + hits_by_source={CacheHitSource.DISK: 4, CacheHitSource.P2P: 0}, + ) + + assert stats.hits_by_source == { + CacheHitSource.HOST: 4, + CacheHitSource.DISK: 6, + CacheHitSource.P2P: 5, + } + assert sum(stats.hits_by_source.values()) == stats.hits == 15 + + +def test_prefix_cache_stats_record_hits_by_source(): + stats = PrefixCacheStats() + stats.record(num_tokens=32, num_hits=16, preempted=False) + stats.record( + num_tokens=32, + num_hits=24, + preempted=False, + hits_by_source={CacheHitSource.HOST: 16, CacheHitSource.DISK: 8}, + ) + # Preempted re-admissions never contribute to hits or the split. + stats.record( + num_tokens=32, + num_hits=32, + preempted=True, + hits_by_source={CacheHitSource.P2P: 32}, + ) + # A connector miscount is reported as unspecified, not dropped or raised. + stats.record( + num_tokens=8, + num_hits=8, + preempted=False, + hits_by_source={CacheHitSource.HOST: 4}, + ) + + assert stats.requests == 3 + assert stats.hits == 48 + assert stats.preempted_hits == 32 + assert stats.hits_by_source == { + CacheHitSource.HOST: 16, + CacheHitSource.DISK: 8, + CacheHitSource.EXTERNAL_UNSPECIFIED: 8, + } + assert sum(stats.hits_by_source.values()) == 48 - 16 + + +def test_prefix_cache_stats_hits_by_source_serialization(): + connector_stats = PrefixCacheStats() + connector_stats.record( + num_tokens=16, + num_hits=8, + preempted=False, + hits_by_source={CacheHitSource.P2P: 4, CacheHitSource.HOST: 4}, + ) + outputs = EngineCoreOutputs( + scheduler_stats=SchedulerStats( + connector_prefix_cache_stats=connector_stats, + ) + ) + + encoded = MsgpackEncoder().encode(outputs) + decoded = MsgpackDecoder(EngineCoreOutputs).decode(encoded) + + assert decoded.scheduler_stats is not None + assert decoded.scheduler_stats.connector_prefix_cache_stats is not None + assert decoded.scheduler_stats.connector_prefix_cache_stats.hits_by_source == { + CacheHitSource.P2P: 4, + CacheHitSource.HOST: 4, + } + + def test_compute_iteration_details_includes_encoder_stats(): scheduler_output = SchedulerOutput.make_empty() scheduler_output.scheduled_encoder_input_stats = ScheduledEncoderInputStats( diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index dee5446c614b..08759bc1e34b 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -23,6 +23,7 @@ ) from vllm.config.cache import MambaCacheMode from vllm.config.kv_events import KVEventsConfig +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) @@ -53,6 +54,7 @@ MambaSpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler @@ -213,6 +215,37 @@ class SchedulerFixture: num_groups: int = 1 +@pytest.mark.parametrize( + ("backend", "source"), + [("cpu", CacheHitSource.HOST), ("disk", CacheHitSource.DISK)], +) +def test_connector_reports_configured_cache_source( + tmp_path, backend: str, source: CacheHitSource +): + kv_cache_config = _make_kv_cache_config(num_blocks=4) + vllm_config = _make_vllm_config() + extra_config = { + "cpu_bytes_to_use_per_rank": _BYTES_PER_BLOCK * 4, + "kv_offload_backend": backend, + } + if backend == "disk": + extra_config.update( + disk_path=str(tmp_path / "kv-offload.bin"), + disk_capacity_bytes=_BYTES_PER_BLOCK * 4, + ) + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = extra_config + + connector = SimpleCPUOffloadConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + kv_cache_config, + ) + + assert connector.get_external_cache_hit_sources(None, 32) == {source: 32} # type: ignore[arg-type] + assert connector.get_external_cache_hit_sources(None, 0) == {} # type: ignore[arg-type] + + def make_scheduler( num_cpu_blocks: int = 8, num_gpu_blocks: int = 16, diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index e272ed5106e0..4b5d2178c49c 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -114,6 +114,7 @@ def export_trace(prof: torch.profiler.profile) -> None: profiler="torch", torch_profiler_dir=str(tmp_path), torch_profiler_dump_cuda_time_total=False, + torch_profiler_with_stack=False, ), worker_name="worker", local_rank=1, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 61de1802d5aa..bb9d4597ce3c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -51,6 +51,7 @@ from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -179,6 +180,10 @@ def aggregate( class KVConnectorBase_V1(ABC): """Base class for KV connectors.""" + # Source label for this connector's external hits. Subclasses must set + # HOST/DISK/P2P or override get_external_cache_hit_sources(). + _cache_hit_source = CacheHitSource.EXTERNAL_UNSPECIFIED + @property def supports_divergent_local_hybrid_hits(self) -> bool: """Whether external hits can complete divergent local hybrid hits. @@ -527,6 +532,24 @@ def get_num_new_matched_tokens( """ pass + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + """Split ``num_external_tokens`` by the cache tier that supplied them. + + Called after :meth:`update_state_after_alloc`, so the load plan is + known. Counts must sum to ``num_external_tokens``; a mismatch is + reported as ``external_unspecified``. Only non-zero counts are + included. + + Default: all tokens under ``_cache_hit_source``. + """ + if num_external_tokens == 0: + return {} + return {self._cache_hit_source: num_external_tokens} + @abstractmethod def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 172d41efdb45..5f064259e8b2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -67,6 +67,7 @@ SlidingWindowMLASpec, SlidingWindowSpec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import RequestStatus from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size @@ -700,6 +701,8 @@ def add_new_req( class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): + _cache_hit_source = CacheHitSource.P2P + def __init__( self, vllm_config: VllmConfig, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 1dc6f8459ab3..167c7ef3c299 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -99,6 +99,7 @@ SlidingWindowSpec, is_full_attention_spec, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -222,6 +223,8 @@ def resolve_moriio_transfer_ack( class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): + _cache_hit_source = CacheHitSource.P2P + @property def supports_divergent_local_hybrid_hits(self) -> bool: # The READ path always transfers the recurrent-state slot, including diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 752acb01eccb..1ded62adda1a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -31,6 +31,7 @@ from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -451,6 +452,18 @@ def get_num_new_matched_tokens( to_return = (toks, load_async) return to_return + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + chosen_connector = self._requests_to_connector.get(request.request_id) + if chosen_connector is None: + return super().get_external_cache_hit_sources(request, num_external_tokens) + return self._connectors[chosen_connector].get_external_cache_hit_sources( + request, num_external_tokens + ) + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 27b4d410f043..6e43356c0d7b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -58,6 +58,7 @@ from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -77,6 +78,8 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): """Base connector with common logic shared by pull and push modes.""" + _cache_hit_source = CacheHitSource.P2P + @property def supports_divergent_local_hybrid_hits(self) -> bool: return True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 64b3ad5bb22a..806267dcf3dd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +from collections import Counter from collections.abc import Iterable, Sequence from dataclasses import dataclass, field, replace from itertools import chain, islice @@ -56,6 +57,7 @@ TierMatcher, make_offload_key, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request, RequestStatus @@ -375,6 +377,9 @@ class RequestOffloadState: # Fine-grained token boundary selected beyond the last complete offload # chunk. It is consumed when the corresponding load is scheduled. partial_tail_boundary: int | None = None + # Loaded attention keys cover their token positions. Recurrent state + # supports the whole reused external prefix. + load_key_ranges: list[tuple[int, int, OffloadKey]] = field(default_factory=list) # True once on_request_finished has been signaled to the manager. finished_signaled: bool = False @@ -587,8 +592,9 @@ def _sliding_window_sort_key(i: int) -> int: sliding_window_groups.sort(key=_sliding_window_sort_key, reverse=True) # used by _lookup + self._full_attention_groups: tuple[int, ...] = tuple(full_attention_groups) self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) - self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups + self._lookup_groups = self._full_attention_groups + self._sliding_window_groups self._mamba_align_size: int | None = resolve_mamba_align_size( spec, kv_cache_config ) @@ -1090,6 +1096,53 @@ def get_num_new_matched_tokens( return num_hit_tokens, bool(num_hit_tokens) + def get_external_cache_hit_sources( + self, + request: Request, + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + """Split the accepted external hit by the tier each loaded key came from. + + Attention keys cover only their loaded token ranges; recurrent state + covers the whole reused prefix. Overlaps report the outermost tier. + Uncovered tokens, including omitted sliding-window history, are + ``external_unspecified``. + """ + sources: dict[CacheHitSource, int] = {} + if num_external_tokens == 0: + return sources + req_status = self._req_status[request.request_id] + start = req_status.num_locally_computed_tokens + end = start + num_external_tokens + + events: list[tuple[int, int, CacheHitSource]] = [ + (end, 0, CacheHitSource.EXTERNAL_UNSPECIFIED) + ] + for lo, hi, key in req_status.load_key_ranges: + lo, hi = max(start, lo), min(end, hi) + if lo >= hi: + continue + source = CacheHitSource( + self.manager.get_load_source(key, req_status.req_context) + ) + events.extend(((lo, 1, source), (hi, -1, source))) + + # Sweep endpoints instead of rescanning every range for each interval. + active: Counter[CacheHitSource] = Counter() + position = start + for boundary, delta, source in sorted(events): + if boundary > position: + tiers = [tier for tier, count in active.items() if count > 0] + tier = ( + CacheHitSource.outermost(tiers) + if tiers + else CacheHitSource.EXTERNAL_UNSPECIFIED + ) + sources[tier] = sources.get(tier, 0) + boundary - position + active[source] += delta + position = boundary + return sources + def update_state_after_alloc( self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int ): @@ -1105,6 +1158,7 @@ def update_state_after_alloc( assert partial_tail_boundary == num_cached_tokens keys_to_load: list[OffloadKey] = [] + req_status.load_key_ranges.clear() dst_block_ids: list[int] = [] # per group group_sizes: list[int] = [] @@ -1157,9 +1211,9 @@ def update_state_after_alloc( ) end_chunk_idx = num_chunks - (partial_tail_boundary is not None) assert len(offload_keys) >= end_chunk_idx - keys_to_load.extend(offload_keys[start_chunk_idx:end_chunk_idx]) + group_keys_to_load = offload_keys[start_chunk_idx:end_chunk_idx] if partial_tail_boundary is not None: - keys_to_load.append( + group_keys_to_load.append( self._make_boundary_key( request, group_config.group_idx, @@ -1167,6 +1221,24 @@ def update_state_after_alloc( req_status.req_context, ) ) + keys_to_load.extend(group_keys_to_load) + if isinstance(group_config.kv_cache_spec, MambaSpec): + req_status.load_key_ranges.extend( + (num_locally_computed_tokens, num_cached_tokens, key) + for key in group_keys_to_load + ) + else: + for chunk_idx, key in enumerate( + group_keys_to_load, start_chunk_idx + ): + range_start = max( + load_start_gpu_block_idx * tokens_per_block, + chunk_idx * tokens_per_chunk, + ) + range_end = min( + num_cached_tokens, (chunk_idx + 1) * tokens_per_chunk + ) + req_status.load_key_ranges.append((range_start, range_end, key)) dst_block_ids.extend( block.block_id diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index def6e1e1f3e8..af3e113b5f09 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -45,6 +45,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.kv_offload.factory import OffloadingSpecFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request @@ -179,6 +180,16 @@ def _max_loadable_tokens( bound = min(per_group_hits[group_id] for group_id in self._bounding_group_ids) return max(0, bound - num_computed_tokens) + def get_external_cache_hit_sources( + self, + request: "Request", + num_external_tokens: int, + ) -> dict[CacheHitSource, int]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_external_cache_hit_sources( + request, num_external_tokens + ) + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 9124d9eac14a..0bd8ebe3cd66 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -23,6 +23,7 @@ ) from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput from vllm.v1.simple_kv_offload.manager import ( BoundaryStoreStats, @@ -109,6 +110,9 @@ def __init__( f"expected one of {VALID_KV_OFFLOAD_BACKENDS}" ) disk_mode = kv_offload_backend == "disk" + self._cache_hit_source = ( + CacheHitSource.DISK if disk_mode else CacheHitSource.HOST + ) disk_path = extra_config.get("disk_path", None) or None disk_capacity_bytes = int( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 0f2425312c60..8bdf91da86da 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1261,10 +1261,20 @@ def skip_request(from_queue: RequestQueue) -> None: self.connector_prefix_cache_stats is not None and connector_prefix_cache_queries != 0 ): + preempted = request.num_preemptions > 0 self.connector_prefix_cache_stats.record( num_tokens=connector_prefix_cache_queries, num_hits=connector_prefix_cache_hits, - preempted=request.num_preemptions > 0, + preempted=preempted, + # The connector knows the tier only after the load + # plan above exists. + hits_by_source=( + self.connector.get_external_cache_hit_sources( + request, connector_prefix_cache_hits + ) + if connector_prefix_cache_hits and not preempted + else None + ), ) # Record at admission so unscheduled lookups are not counted. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index f8e92abed604..398bc269e6f2 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -18,6 +18,7 @@ from vllm.v1.kv_hints import KvHintsEnvelope from vllm.v1.kv_offload.config import OffloadingConfig +from vllm.v1.metrics.cache_hit_source import CacheHitSource # `OffloadKey` identifies an offloaded block. It combines a block hash with # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead. @@ -248,6 +249,18 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ pass + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + """Return the cache tier that supplied a successful lookup. + + This is queried only after ``lookup`` resolves to ``HIT``. Composing + managers can override it to retain the origin across asynchronous + promotion into a primary tier. The default keeps out-of-tree managers + compatible while making unknown provenance explicit. + """ + return CacheHitSource.EXTERNAL_UNSPECIFIED + @abstractmethod def prepare_load( self, diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index e42cb47f2530..2c08ac897c6c 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -27,6 +27,7 @@ ) from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +from vllm.v1.metrics.cache_hit_source import CacheHitSource @dataclass(slots=True) @@ -193,6 +194,12 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.HIT_PENDING return LookupResult.HIT + @override + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + return CacheHitSource.HOST + @override def prepare_load( self, diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index a90d33eb1136..b855b6b000a5 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -24,6 +24,7 @@ from vllm.v1.kv_offload.tiering.backpressure import ( BackpressureDetector, ) +from vllm.v1.metrics.cache_hit_source import CacheHitSource if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( @@ -141,6 +142,17 @@ class SecondaryTierManager(ABC): """ medium: ClassVar[Medium | None] = None + # Cache tier reported for hits served from this tier. Derived from + # ``medium`` unless a subclass sets it; ``None`` medium means unknown. + cache_hit_source: ClassVar[CacheHitSource] = CacheHitSource.EXTERNAL_UNSPECIFIED + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if "cache_hit_source" not in cls.__dict__ and cls.medium is not None: + cls.cache_hit_source = { + Medium.CPU: CacheHitSource.HOST, + Medium.STORAGE: CacheHitSource.DISK, + }[cls.medium] def __init__( self, diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ce518042b656..83ff2e2c1280 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -54,6 +54,7 @@ TransferJob, ) from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker +from vllm.v1.metrics.cache_hit_source import CacheHitSource logger = init_logger(__name__) @@ -232,6 +233,11 @@ def __init__( # complete_store(), since complete_store() can still submit cascades. self._req_state: dict[str, RequestState] = {} + # Preserve the original tier for this request's cache-hit metrics, + # even after its KV blocks are promoted into host memory + # (otherwise a disk or P2P hit would be mislabeled as a host hit). + self._request_load_sources: dict[str, dict[OffloadKey, CacheHitSource]] = {} + # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { tier: _SecondaryTierFacingParent(self, tier_idx) @@ -293,6 +299,13 @@ def _complete_promotion( successful_keys = () failed_keys = transfer_job.keys + load_sources = self._request_load_sources.get(transfer_job.req_context.req_id) + if load_sources is not None: + source = self.secondary_tiers[job_metadata.tier_idx].cache_hit_source + for key in failed_keys: + if load_sources.get(key) == source: + del load_sources[key] + if successful_keys: self.primary_tier.complete_write( successful_keys, @@ -464,6 +477,15 @@ def lookup( return LookupResult.RETRY return LookupResult.MISS + @override + def get_load_source( + self, key: OffloadKey, req_context: ReqContext + ) -> CacheHitSource: + load_sources = self._request_load_sources.get(req_context.req_id) + if load_sources is not None and key in load_sources: + return load_sources[key] + return self.primary_tier.get_load_source(key, req_context) + def _initiate_promotion( self, tier_idx: int, @@ -501,6 +523,10 @@ def _initiate_promotion( store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) + load_sources = self._request_load_sources.setdefault(req_context.req_id, {}) + source = self.secondary_tiers[tier_idx].cache_hit_source + for promoted_key in primary_write_result.keys_to_store: + load_sources[promoted_key] = source # Defer submit_load to on_schedule_end(). Group by (tier, request) so # each request's chunks are submitted as one batched job per tier. tier_pending = self._pending_load_submissions.setdefault(tier_idx, {}) @@ -848,6 +874,7 @@ def _maybe_finalize_request( continue tier.on_request_finished(state.req_context) self._metrics.on_request_finished(state.req_context) + self._request_load_sources.pop(req_id, None) del self._req_state[req_id] @override @@ -930,6 +957,7 @@ def reset_cache(self) -> None: # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() + self._request_load_sources.clear() self._metrics.assert_idle() finished_req_ids = [] diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 59b9c320695f..f759a4481e5d 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -11,7 +11,7 @@ import uuid from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar from typing_extensions import override @@ -35,6 +35,7 @@ from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport from vllm.v1.kv_offload.tiering.p2p.data import DataTransport, NixlTransport from vllm.v1.kv_offload.tiering.p2p.session import P2PSession +from vllm.v1.metrics.cache_hit_source import CacheHitSource if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec @@ -206,6 +207,8 @@ class P2PSecondaryTierManager(SecondaryTierManager): and existing sessions are polled even when no requests are scheduled. """ + cache_hit_source: ClassVar[CacheHitSource] = CacheHitSource.P2P + def __init__( self, offloading_spec: OffloadingSpec, diff --git a/vllm/v1/metrics/cache_hit_source.py b/vllm/v1/metrics/cache_hit_source.py new file mode 100644 index 000000000000..b6fc5270e776 --- /dev/null +++ b/vllm/v1/metrics/cache_hit_source.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterable +from enum import Enum + + +class CacheHitSource(str, Enum): + """Where cached prompt tokens' KV came from. + + Members are listed in the order blocks are offloaded: accelerator, then + host memory, then secondary tiers. Per-source token counts are plain + ``dict[CacheHitSource, int]`` mappings holding only non-zero entries. + """ + + DEVICE = "device" + HOST = "host" + P2P = "p2p" + DISK = "disk" + # Fallback: uninstrumented connector or remote store. Last, so mixing + # with a known tier never reports the known one. + EXTERNAL_UNSPECIFIED = "external_unspecified" + + @classmethod + def outermost(cls, sources: Iterable[CacheHitSource]) -> CacheHitSource: + """The tier in ``sources`` farthest from the accelerator.""" + return max(sources, key=_SOURCE_ORDER.__getitem__) + + +_SOURCE_ORDER = {source: rank for rank, source in enumerate(CacheHitSource)} diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 681f80ce2942..7ff1d2993cad 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -23,6 +23,7 @@ from vllm.plugins import STAT_LOGGER_PLUGINS_GROUP, load_plugins_by_group from vllm.v1.engine import FinishReason from vllm.v1.metrics.buckets import histogram_buckets +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.perf import PerfMetricsLogging, PerfMetricsProm from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import ( @@ -752,6 +753,23 @@ def __init__( counter_prompt_tokens_cached, per_engine_labelvalues ) + self.counter_prompt_tokens_cached_by_source = self._counter_cls( + name="vllm:prompt_tokens_cached_by_source", + documentation=( + "Prefix-cache hit tokens by the cache tier that supplied them: " + "device (local HBM), host (offloaded to DRAM), disk, p2p " + "(transferred from another vLLM instance) or " + "external_unspecified. Counted at admission; sums to " + "vllm:prefix_cache_hits + vllm:external_prefix_cache_hits." + ), + labelnames=labelnames + ["source"], + ) + for source in CacheHitSource: + for labelvalues in per_engine_labelvalues.values(): + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, source.value + ) + counter_generation_tokens = self._counter_cls( name="vllm:generation_tokens", documentation="Number of generation tokens processed.", @@ -1111,13 +1129,23 @@ def record( scheduler_stats.prefix_cache_stats.hits ) - if scheduler_stats.connector_prefix_cache_stats is not None: + labelvalues = self.per_engine_labelvalues[engine_idx] + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, CacheHitSource.DEVICE.value + ).inc(scheduler_stats.prefix_cache_stats.hits) + + connector_stats = scheduler_stats.connector_prefix_cache_stats + if connector_stats is not None: self.counter_connector_prefix_cache_queries[engine_idx].inc( - scheduler_stats.connector_prefix_cache_stats.queries + connector_stats.queries ) self.counter_connector_prefix_cache_hits[engine_idx].inc( - scheduler_stats.connector_prefix_cache_stats.hits + connector_stats.hits ) + for tier, num_tokens in connector_stats.hits_by_source.items(): + self.counter_prompt_tokens_cached_by_source.labels( + *labelvalues, tier.value + ).inc(num_tokens) if scheduler_stats.spec_decoding_stats is not None: self.spec_decoding_prom.observe( diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 4db98884fe38..0046dcefa92e 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -8,9 +8,13 @@ import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.logger import init_logger +from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.metrics.perf import PerfStats from vllm.v1.spec_decode.metrics import SpecDecodingStats +logger = init_logger(__name__) + if TYPE_CHECKING: from vllm.v1.engine import EngineCoreEvent, EngineCoreOutput, FinishReason @@ -128,19 +132,42 @@ class PrefixCacheStats(BaseCacheStats): preempted_hits: int = 0 """The `hits` number for preempted requests.""" + hits_by_source: dict[CacheHitSource, int] = field(default_factory=dict) + """`hits` split by the cache tier that supplied them (connector stats only).""" - def record(self, num_tokens: int, num_hits: int, preempted: bool) -> None: + def record( + self, + num_tokens: int, + num_hits: int, + preempted: bool, + hits_by_source: dict[CacheHitSource, int] | None = None, + ) -> None: """Aggregate request information into the stats.""" if preempted: # Previously preempted request self.preempted_requests += 1 self.preempted_queries += num_tokens self.preempted_hits += num_hits - else: - # New request - self.requests += 1 - self.queries += num_tokens - self.hits += num_hits + return + # New request + self.requests += 1 + self.queries += num_tokens + self.hits += num_hits + if hits_by_source is None: + return + if sum(hits_by_source.values()) != num_hits: + logger.warning_once( + "Connector attributed %s for %d cached tokens; reporting them as %s.", + str(hits_by_source), + num_hits, + CacheHitSource.EXTERNAL_UNSPECIFIED.value, + ) + hits_by_source = {CacheHitSource.EXTERNAL_UNSPECIFIED: num_hits} + for source, num_source_hits in hits_by_source.items(): + if num_source_hits: + self.hits_by_source[source] = ( + self.hits_by_source.get(source, 0) + num_source_hits + ) @dataclass From bb87d227d4b964abb2a966cdf9194f3d376d9bbe Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Tue, 6 Oct 2026 15:10:00 -0400 Subject: [PATCH 089/952] [Bugfix][Model] Fix mypy safe-super error in EmbeddingGemma2Model (#60280) Signed-off-by: Roger Wang Co-authored-by: Kimi Code --- vllm/model_executor/models/embedding_gemma2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py index 2d0012714eab..150128505d5e 100644 --- a/vllm/model_executor/models/embedding_gemma2.py +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -1020,7 +1020,7 @@ class EmbeddingGemma2Model(nn.Module, SupportsMultiModal, VllmModelForPooling): _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() + nn.Module.__init__(self) config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config From d1f3d8b87083a95f914c674c477796950e2ba5fe Mon Sep 17 00:00:00 2001 From: Jinzhen Lin Date: Wed, 7 Oct 2026 03:31:35 +0800 Subject: [PATCH 090/952] [Quantization] Prefer Humming before Marlin backends on SM90 (#56997) Signed-off-by: jinzhen.ljz Signed-off-by: mgoin Co-authored-by: Codex Co-authored-by: mgoin Co-authored-by: Claude Opus 5.5 --- .../test_scaled_mm_kernel_selection.py | 21 ++++++++ tests/quantization/test_modelopt.py | 20 ++++++-- tests/quantization/test_online.py | 3 +- .../model_executor/kernels/linear/__init__.py | 14 ++++-- .../kernels/linear/mixed_precision/humming.py | 21 +++++++- .../layers/fused_moe/oracle/fp8.py | 18 +++++++ .../layers/fused_moe/oracle/int_wna16.py | 4 +- .../layers/fused_moe/oracle/mxfp4.py | 9 ++-- .../layers/fused_moe/oracle/mxfp8.py | 4 +- .../layers/fused_moe/oracle/nvfp4.py | 2 + .../quantization/utils/humming/__init__.py | 6 +++ .../quantization/utils/humming/linear.py | 4 +- .../quantization/utils/humming/priority.py | 48 +++++++++++++++++++ 13 files changed, 159 insertions(+), 15 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/utils/humming/priority.py diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index f8380dcecf36..641c458608f4 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -14,6 +14,8 @@ from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.kernels.linear import ( + _POSSIBLE_KERNELS, + _POSSIBLE_NVFP4_KERNELS, AiterInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, HummingFP8ScaledMMLinearKernel, @@ -26,6 +28,7 @@ init_int8_linear_kernel, register_linear_kernel, ) +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) @@ -197,3 +200,21 @@ def test_register_oot_linear_kernel(platform_mock): assert isinstance(kernel, OOTInt8ScaledMMLinearKernel), ( "init_int8_linear_kernel should return an instance of the registered kernel" ) + + +@pytest.mark.parametrize( + "possible_kernels", [_POSSIBLE_KERNELS, _POSSIBLE_NVFP4_KERNELS] +) +def test_humming_priority_only_reorders_humming_and_marlin(possible_kernels): + """Preferring Humming must not demote Marlin below unrelated kernels.""" + kernels = possible_kernels[PlatformEnum.CUDA] + original = list(kernels) + + assert prioritize_humming(kernels, compute_capability=80) == original + + names = [k.__name__ for k in prioritize_humming(kernels, compute_capability=90)] + humming = next(i for i, name in enumerate(names) if "Humming" in name) + assert "Marlin" in names[humming + 1] + del names[humming] + assert names == [k.__name__ for k in original if "Humming" not in k.__name__] + assert kernels == original diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index eea4b6218de8..7b1131a00e39 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -688,7 +688,8 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): @pytest.mark.parametrize( ("linear_backend", "kernel_cls"), [ - ("auto", MarlinNvFp4LinearKernel), + ("auto", None), + ("marlin", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel), ("flashinfer_cutedsl", FlashInferCuteDslNvFp4W4A16LinearKernel), ], @@ -696,7 +697,7 @@ def test_modelopt_linear_method_builder_registry_override(monkeypatch): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``: - ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The + ``use_a16=True`` follows platform priorities, but an explicit backend wins. The generic method routes this through ``select_linear_kernel``.""" from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.quantization.modelopt import ( @@ -704,7 +705,20 @@ def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): select_linear_kernel, ) - if linear_backend != "auto": + if linear_backend == "auto": + capability = current_platform.get_device_capability() + assert capability is not None + cc = capability.to_int() + if ( + cc in (100, 103) + and FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported()[0] + ): + kernel_cls = FlashInferCuteDslNvFp4W4A16LinearKernel + elif cc == 90 and HummingNvFp4LinearKernel.is_supported()[0]: + kernel_cls = HummingNvFp4LinearKernel + else: + kernel_cls = MarlinNvFp4LinearKernel + else: is_supported, reason = kernel_cls.is_supported() if not is_supported: pytest.skip(reason) diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index f3b72707f3d6..3f27757b9017 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -34,6 +34,7 @@ from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp8.humming import HummingMxfp8LinearKernel from vllm.model_executor.kernels.linear.mxfp8.marlin import ( MarlinMxfp8LinearKernel, ) @@ -1060,7 +1061,7 @@ def assert_source_weights_released(method, layer) -> None: assert isinstance(moe._quant_method, expected_moe_cls) if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( - o_proj.quant_method.kernel, MarlinMxfp8LinearKernel + o_proj.quant_method.kernel, (MarlinMxfp8LinearKernel, HummingMxfp8LinearKernel) ): assert o_proj.weight.dtype == torch.int32 elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index e5fed78c40a7..484143ffc14a 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -228,6 +228,10 @@ from vllm.model_executor.kernels.linear.scaled_mm.zentorch import ( ZentorchInt8ScaledMMLinearKernel, ) +from vllm.model_executor.layers.quantization.utils.humming import ( + prefers_humming, + prioritize_humming, +) from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms import PlatformEnum, current_platform @@ -376,8 +380,9 @@ def _resolve_backend_kernels( layer_desc: str, *, quantization: str, + compute_capability: int | None = None, ) -> list[type]: - """Apply --linear-backend filtering to one layer type's kernel list. + """Apply device priorities and --linear-backend filtering to a kernel list. When the requested backend has no kernel for this layer type, fall back to the unfiltered list (with a WARNING log) instead of failing engine @@ -386,6 +391,7 @@ def _resolve_backend_kernels( layer types (e.g. NVFP4 MoE projections next to FP8 attention projections). """ + kernels = prioritize_humming(kernels, compute_capability) linear_backend = _get_linear_backend(quantization=quantization) if linear_backend == "auto": return kernels @@ -682,6 +688,7 @@ def choose_scaled_mm_linear_kernel( platform_kernels, "scaled-mm", quantization=quantization, + compute_capability=compute_capability, ) for kernel in platform_kernels: @@ -845,6 +852,7 @@ def choose_mp_linear_kernel( platform_kernels, "mixed-precision", quantization="mixed_precision", + compute_capability=compute_capability, ) failure_reasons = [] @@ -1121,13 +1129,13 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: _cc = current_platform.get_device_capability() compute_capability = _cc.to_int() if _cc is not None else None # Weight-only: prefer FlashInfer CuTe-DSL W4A16 on SM100/103, - # otherwise Marlin. + # Humming then Marlin where Humming is preferred, and Marlin elsewhere. cutedsl_ok, _ = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported( compute_capability ) if compute_capability in (100, 103) and cutedsl_ok: force_kernel = FlashInferCuteDslNvFp4W4A16LinearKernel - else: + elif not prefers_humming(compute_capability): force_kernel = MarlinNvFp4LinearKernel if force_kernel is not None: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 4655181f78c6..a2d88705d291 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -4,11 +4,21 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + kInt8DynamicTokenSym, +) from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig +# Per-token activation quantization requested through `act_type`. +_ACT_QUANT_KEYS = { + torch.int8: kInt8DynamicTokenSym, + torch.float8_e4m3fn: kFp8DynamicTokenSym, +} + class HummingLinearKernel(MPLinearKernel): @classmethod @@ -21,6 +31,8 @@ def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: return False, "Humming is only supported on CUDA" if not has_humming(): return False, "Humming is not installed" + if c.act_type.itemsize == 1 and c.act_type not in _ACT_QUANT_KEYS: + return False, f"Activation type ({c.act_type}) not supported by Humming" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -28,6 +40,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name} @@ -48,8 +61,14 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) input_quant_config = getattr(layer, "_humming_input_quant_config", None) + input_schema = None + if (act_key := _ACT_QUANT_KEYS.get(self.config.act_type)) is not None: + input_schema = quant_key_to_input_schema(act_key) self.layer_config = prepare_humming_linear_layer_config( - layer, quant_config, input_quant_config=input_quant_config + layer, + quant_config, + input_quant_config=input_quant_config, + input_schema=input_schema, ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 5219fff617d3..446825c59047 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( prepare_fp8_moe_layer_for_deepgemm, ) +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( prepare_fp8_moe_layer_for_marlin, ) @@ -37,6 +38,7 @@ kFp8Dynamic128Sym, kFp8Static128BlockSym, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.utils.math_utils import round_up @@ -98,6 +100,7 @@ def _get_priority_backends( Fp8MoeBackend.CPU, Fp8MoeBackend.HPC, ] + _AVAILABLE_BACKENDS = prioritize_humming(_AVAILABLE_BACKENDS) def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None: backends.insert(0, backends.pop(backends.index(backend))) @@ -607,6 +610,21 @@ def convert_to_fp8_moe_kernel_format( convert_to_humming_moe_kernel_format, ) + # Online quantization has not installed the quantized tensors yet. + scale_name = ( + "weight_scale_inv" + if hasattr(layer, "w13_weight_scale_inv") + else "weight_scale" + ) + for prefix, weight, scale, input_scale in ( + ("w13", w13, w13_scale, w13_input_scale), + ("w2", w2, w2_scale, w2_input_scale), + ): + replace_parameter(layer, f"{prefix}_weight", weight) + replace_parameter(layer, f"{prefix}_{scale_name}", scale) + if input_scale is not None: + replace_parameter(layer, f"{prefix}_input_scale", input_scale) + convert_to_humming_moe_kernel_format( layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale) ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 239fe1d55a85..3e91fd9e041c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -32,6 +32,7 @@ TrtLlmMxint4ExpertsMonolithic, ) from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_moe_marlin_supports_config, marlin_act_int8_process_scales, @@ -129,7 +130,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] - return [ + backends = [ # Native HIP kernel, gated on gfx1100 by _supports_current_device(). WNA16MoEBackend.RDNA3, WNA16MoEBackend.FLASHINFER_TRTLLM, @@ -139,6 +140,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.HUMMING, WNA16MoEBackend.EMULATION, ] + return prioritize_humming(backends) def _backend_incompatibility_reason( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 691d849228f8..5463853caba1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -35,6 +35,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( _swizzle_mxfp4, ) @@ -382,6 +383,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: # TRITON_UNFUSED Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, + Mxfp4MoeBackend.HUMMING, Mxfp4MoeBackend.XPU, Mxfp4MoeBackend.CPU, Mxfp4MoeBackend.EMULATION, @@ -395,12 +397,12 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: _AVAILABLE_BACKENDS.index(Mxfp4MoeBackend.BATCHED_MARLIN) + 1, Mxfp4MoeBackend.TRITON, ) - return _AVAILABLE_BACKENDS + return prioritize_humming(_AVAILABLE_BACKENDS) def _get_priority_backends() -> list[Mxfp4MoeBackend]: """Get available backends in priority order. SM100+ prefers DeepGEMM FP4 / - TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the + TRTLLM MXFP8; SM90 falls through to Humming or Marlin (the backend-level ``is_supported_config`` check filters by device capability). """ if current_platform.is_rocm(): @@ -420,8 +422,9 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: # TRITON_UNFUSED Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, + Mxfp4MoeBackend.HUMMING, ] - return _AVAILABLE_BACKENDS + return prioritize_humming(_AVAILABLE_BACKENDS) def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 49f61764f0fd..c4417cdd76ee 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -8,6 +8,7 @@ Fp8MoeBackend, backend_to_kernel_cls, ) +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp8Dynamic, kMxfp8Static, @@ -129,7 +130,8 @@ def select_mxfp8_moe_backend( return backend, _select_kernel_cls(backend, config) # Auto-select: pick the first supported backend. - for backend in _SUPPORTED_BACKENDS: + backends = prioritize_humming(list(_SUPPORTED_BACKENDS)) + for backend in backends: try: experts_cls = _select_kernel_cls(backend, config) except ValueError: diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 68ba62782c8f..330e7b6aa13a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -26,6 +26,7 @@ prepare_nvfp4_moe_layer_for_fi_or_cutlass, prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, ) +from vllm.model_executor.layers.quantization.utils.humming import prioritize_humming from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_nvfp4_moe_layer_for_marlin, ) @@ -209,6 +210,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.HUMMING, NvFp4MoeBackend.EMULATION, ] + AVAILABLE_BACKENDS = prioritize_humming(AVAILABLE_BACKENDS) NVFP4_BACKENDS_WITH_CLAMP = { NvFp4MoeBackend.B12X, diff --git a/vllm/model_executor/layers/quantization/utils/humming/__init__.py b/vllm/model_executor/layers/quantization/utils/humming/__init__.py index 4303f74221f0..91a5713bbfad 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/__init__.py +++ b/vllm/model_executor/layers/quantization/utils/humming/__init__.py @@ -19,6 +19,10 @@ make_humming_moe_quant_config, select_humming_moe_experts, ) +from vllm.model_executor.layers.quantization.utils.humming.priority import ( + prefers_humming, + prioritize_humming, +) from vllm.model_executor.layers.quantization.utils.humming.schema import ( check_and_fallback_input_schema, humming_is_layer_skipped, @@ -41,6 +45,8 @@ "convert_linear_layer_to_humming_standard", "get_humming_linear_compute_config", "prepare_humming_linear_layer_config", + "prefers_humming", + "prioritize_humming", "make_humming_moe_kernel", "select_humming_moe_experts", "input_schema_to_quant_key", diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py index 0c64645c4395..fcce3dabfef2 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/linear.py +++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py @@ -87,7 +87,7 @@ def prepare_humming_linear_layer_config( shape_n_stacks = layer.output_partition_sizes # Step 1: convert weight and input schemas to humming standard format - source_tensors = dict(layer.named_parameters()) + source_tensors = dict(layer.named_parameters(recurse=False)) weight_schema, tensors = weight_schema.convert_humming( tensors=source_tensors, shape_n_stacks=shape_n_stacks, @@ -143,7 +143,7 @@ def prepare_humming_linear_layer_config( if "bias" in tensors: zero_bias = torch.zeros_like(tensors["bias"]) layer.register_buffer("zero_bias", zero_bias) - for name, _ in list(layer.named_parameters()): + for name, _ in list(layer.named_parameters(recurse=False)): delattr(layer, name) for name, tensor in tensors.items(): param = torch.nn.Parameter(tensor, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/utils/humming/priority.py b/vllm/model_executor/layers/quantization/utils/humming/priority.py new file mode 100644 index 000000000000..4a1452c2a96a --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/priority.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Humming kernel selection priority.""" + +from enum import Enum +from typing import TypeVar + +_KernelT = TypeVar("_KernelT", bound=type | Enum) + +# Compute capabilities where Humming outranks Marlin. +_HUMMING_PREFERRED_CAPABILITIES = frozenset({90}) + + +def prefers_humming(compute_capability: int | None = None) -> bool: + """Whether Humming outranks Marlin. A missing compute capability uses the + current device.""" + if compute_capability is None: + from vllm.platforms import current_platform + + if current_platform.is_cuda(): + cc = current_platform.get_device_capability() + compute_capability = cc.to_int() if cc is not None else None + return compute_capability in _HUMMING_PREFERRED_CAPABILITIES + + +def prioritize_humming( + kernels: list[_KernelT], + compute_capability: int | None = None, +) -> list[_KernelT]: + """Move Humming directly ahead of Marlin where Humming is preferred. + + Every other entry keeps its relative order and the input list is not + modified. Match kernel class names or MoE backend enum names. + """ + if not prefers_humming(compute_capability): + return kernels + + names = [ + (kernel.name if isinstance(kernel, Enum) else kernel.__name__).lower() + for kernel in kernels + ] + humming = next((i for i, name in enumerate(names) if "humming" in name), None) + marlin = next((i for i, name in enumerate(names) if "marlin" in name), None) + if humming is None or marlin is None or humming < marlin: + return kernels + kernels = kernels.copy() + kernels.insert(marlin, kernels.pop(humming)) + return kernels From e43db1f5e2eb7ca16fc284a810c548a0d3d44b14 Mon Sep 17 00:00:00 2001 From: PatrykSaffer Date: Tue, 6 Oct 2026 22:03:05 +0200 Subject: [PATCH 091/952] Pad to 64 not 128 Q heads in flashMLA sparse (#60029) Signed-off-by: Patryk Saffer Co-authored-by: Misha Goin --- vllm/models/hy_v4/nvidia/flashmla_sparse.py | 6 +++--- vllm/v1/attention/backends/mla/flashmla_sparse.py | 12 +++++------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/vllm/models/hy_v4/nvidia/flashmla_sparse.py b/vllm/models/hy_v4/nvidia/flashmla_sparse.py index 1d62f2f556de..52656b06a6f7 100644 --- a/vllm/models/hy_v4/nvidia/flashmla_sparse.py +++ b/vllm/models/hy_v4/nvidia/flashmla_sparse.py @@ -228,9 +228,9 @@ def _bf16_flash_mla_kernel( -1, 1, kv_c_and_k_pe_cache.shape[-1] ) - # NOTE(Chen): kernel requires num_local_head to be a multiple of - # 64 on hopper and 128 on blackwell. Pad from q's head count, not - # self.num_heads: under DCP the heads are all-gathered before this. + # NOTE(Chen): kernel requires num_local_head to be a multiple of 64. + # Pad from q's head count, not self.num_heads: under DCP the heads are + # all-gathered before this. if actual_num_heads is None: actual_num_heads = q.shape[1] padded_num_heads = ( diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index dea9b81116e4..d441934c022a 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -730,10 +730,8 @@ def __init__( **mla_args, ) self.softmax_scale = scale - # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell - self.prefill_padding = ( - 128 if current_platform.is_device_capability_family(100) else 64 - ) + # Prefill BF16 kernel requires 64 heads on Hopper and Blackwell. + self.prefill_padding = 64 self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads) # The kernels read queries as wide as the cache rows. A NoPE model on @@ -1323,9 +1321,9 @@ def _bf16_flash_mla_kernel( -1, 1, kv_c_and_k_pe_cache.shape[-1] ) - # NOTE(Chen): kernel requires num_local_head to be a multiple of - # 64 on hopper and 128 on blackwell. Pad from q's head count, not - # self.num_heads: under DCP the heads are all-gathered before this. + # NOTE(Chen): kernel requires num_local_head to be a multiple of 64. + # Pad from q's head count, not self.num_heads: under DCP the heads are + # all-gathered before this. if actual_num_heads is None: actual_num_heads = q.shape[1] padded_num_heads = ( From dac35c5b66efb0ddabb6f67af5d1817b2c74edc0 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Tue, 6 Oct 2026 16:13:44 -0400 Subject: [PATCH 092/952] [Misc] Add chat-completion options to gsm8k_eval.py CLI (#60114) Signed-off-by: Benjamin Chislett Co-authored-by: Claude Opus 5.5 --- tests/evals/gsm8k/README.md | 16 ++++ tests/evals/gsm8k/gsm8k_eval.py | 143 +++++++++++++++++++++++++++----- 2 files changed, 139 insertions(+), 20 deletions(-) diff --git a/tests/evals/gsm8k/README.md b/tests/evals/gsm8k/README.md index db37d2e22438..ea61b6bf4562 100644 --- a/tests/evals/gsm8k/README.md +++ b/tests/evals/gsm8k/README.md @@ -21,6 +21,22 @@ vllm serve Qwen/Qwen2.5-1.5B-Instruct --port 8000 python tests/evals/gsm8k/gsm8k_eval.py --port 8000 ``` +For chat-only models, send the prompts through `/v1/chat/completions`: + +```bash +python tests/evals/gsm8k/gsm8k_eval.py --port 8000 \ + --use-chat-completions --max-tokens 1024 \ + --chat-template-kwargs '{"enable_thinking": false}' +``` + +The server applies the chat template and uses its served model, so `--model` is +optional. In chat mode, `temperature` is sent only when `--temperature` is given, +so the server's default sampling applies otherwise; completions mode keeps the +default of 0. Chat mode sends no stop strings; completions mode keeps them. +`--reasoning-effort`, `--top-p`, and `--top-k` set the matching request fields +when given. With a reasoning parser, only the final answer `content` is scored. +`--save-results` records the request settings used. + ## Configuration Format Model configs in `configs/` directory use this YAML format: diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 7e2534b55a95..4eec87945b6f 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -76,14 +76,20 @@ def get_answer_value(answer_str: str) -> int: return INVALID +def _optional_params(**params: object) -> dict[str, object]: + return {k: v for k, v in params.items() if v is not None} + + async def call_vllm_api( session: aiohttp.ClientSession, prompt: str, - temperature: float, + temperature: float | None, max_tokens: int, stop: list[str] | None = None, url: str | None = None, seed: int | None = None, + top_p: float | None = None, + top_k: int | None = None, ) -> tuple[str, int]: """Call vLLM's OpenAI-compatible completions endpoint. @@ -93,12 +99,12 @@ async def call_vllm_api( """ data = { "prompt": prompt, - "temperature": temperature, "max_tokens": max_tokens, "stop": stop, } - if seed is not None: - data["seed"] = seed + data.update( + _optional_params(temperature=temperature, seed=seed, top_p=top_p, top_k=top_k) + ) try: async with session.post(f"{url}/v1/completions", json=data) as response: @@ -114,24 +120,45 @@ async def call_vllm_api( async def call_vllm_chat_api( session: aiohttp.ClientSession, - model: str, + model: str | None, prompt: str, - temperature: float, + temperature: float | None, max_tokens: int, stop: list[str] | None = None, url: str | None = None, seed: int | None = None, + top_p: float | None = None, + top_k: int | None = None, + reasoning_effort: str | None = None, + chat_template_kwargs: dict[str, object] | None = None, ) -> tuple[str, int]: - """Call vLLM's OpenAI-compatible chat completions endpoint.""" + """Call vLLM's OpenAI-compatible chat completions endpoint. + + ``model``, ``temperature``, ``stop`` and the other optional fields are + omitted from the request when ``None``, so the server defaults apply (vLLM + uses its served model when ``model`` is omitted). + + Returns: + Tuple of (final answer content, completion_tokens). Reasoning returned + separately by a reasoning parser is not part of the answer. + + """ data = { - "model": model, "messages": [{"role": "user", "content": prompt}], - "temperature": temperature, "max_tokens": max_tokens, - "stop": stop, } - if seed is not None: - data["seed"] = seed + data.update( + _optional_params( + stop=stop, + model=model, + temperature=temperature, + seed=seed, + top_p=top_p, + top_k=top_k, + reasoning_effort=reasoning_effort, + chat_template_kwargs=chat_template_kwargs, + ) + ) try: async with session.post(f"{url}/v1/chat/completions", json=data) as response: @@ -214,16 +241,34 @@ def evaluate_gsm8k( use_chat_completions: bool = False, host: str = "http://127.0.0.1", port: int = 8000, - temperature: float = 0.0, + temperature: float | None = 0.0, seed: int | None = 42, request_timeout_seconds: float = 600, gen_prefix: str = "", max_concurrency: int | None = None, + top_p: float | None = None, + top_k: int | None = None, + reasoning_effort: str | None = None, + chat_template_kwargs: dict[str, object] | None = None, ) -> dict[str, float | int]: """Evaluate GSM8K accuracy using vLLM serve endpoint. + ``temperature``, ``top_p`` and ``top_k`` are sent only when not ``None``; + otherwise the server default applies. ``model`` is used only in chat mode, + where it may be ``None`` (the server uses its served model). + ``reasoning_effort`` and ``chat_template_kwargs`` require + ``use_chat_completions=True``. Stop strings are sent only in completions + mode. + Returns dict with accuracy, invalid_rate, latency, etc. """ + if not use_chat_completions and ( + reasoning_effort is not None or chat_template_kwargs is not None + ): + raise ValueError( + "reasoning_effort and chat_template_kwargs require " + "use_chat_completions=True" + ) base_url = f"{host}:{port}" prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) num_questions = len(prompts) @@ -233,19 +278,21 @@ async def run_async_evaluation(): output_tokens: list[int] = [0] * num_questions async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]: - stop = ["Question", "Assistant:", "<|separator|>"] if use_chat_completions: - if model is None: - raise ValueError("model is required for chat completions") + # No stop strings: the chat template delimits the turn, and they + # would also cut off reasoning. answer, tokens = await call_vllm_chat_api( session=session, model=model, prompt=prompts[i], temperature=temperature, max_tokens=max_tokens, - stop=stop, url=base_url, seed=seed, + top_p=top_p, + top_k=top_k, + reasoning_effort=reasoning_effort, + chat_template_kwargs=chat_template_kwargs, ) else: answer, tokens = await call_vllm_api( @@ -253,9 +300,11 @@ async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]: prompt=prompts[i], temperature=temperature, max_tokens=max_tokens, - stop=stop, + stop=["Question", "Assistant:", "<|separator|>"], url=base_url, seed=seed, + top_p=top_p, + top_k=top_k, ) states[i] = answer output_tokens[i] = tokens @@ -352,7 +401,38 @@ def main() -> None: parser.add_argument("--host", type=str, default="http://127.0.0.1", help="Host URL") parser.add_argument("--port", type=int, default=8000, help="Port number") parser.add_argument( - "--temperature", type=float, default=0.0, help="Temperature for generation" + "--temperature", + type=float, + help="Temperature for generation (default: 0 for completions; " + "server default with --use-chat-completions)", + ) + parser.add_argument( + "--top-p", type=float, help="Top-p for generation (server default if unset)" + ) + parser.add_argument( + "--top-k", type=int, help="Top-k for generation (server default if unset)" + ) + parser.add_argument( + "--use-chat-completions", + action="store_true", + help="Send each few-shot prompt as a user message to /v1/chat/completions " + "instead of /v1/completions (for chat-only models)", + ) + parser.add_argument( + "--model", + type=str, + help="Model name to send with chat requests (default: omitted, so the " + "server uses its served model)", + ) + parser.add_argument( + "--reasoning-effort", + type=str, + help="reasoning_effort for chat completions, e.g. none, low, high", + ) + parser.add_argument( + "--chat-template-kwargs", + type=json.loads, + help="JSON chat_template_kwargs, e.g. '{\"enable_thinking\": false}'", ) parser.add_argument( "--seed", type=int, default=42, help="Random seed for reproducibility" @@ -371,6 +451,9 @@ def main() -> None: parser.add_argument("--save-results", type=str, help="Save results to JSON file") args = parser.parse_args() + temperature = args.temperature + if temperature is None and not args.use_chat_completions: + temperature = 0.0 result = evaluate_gsm8k( num_questions=args.num_questions, @@ -378,10 +461,16 @@ def main() -> None: max_tokens=args.max_tokens, host=args.host, port=args.port, - temperature=args.temperature, + temperature=temperature, seed=args.seed, max_concurrency=args.max_concurrency, request_timeout_seconds=args.request_timeout_seconds, + model=args.model, + use_chat_completions=args.use_chat_completions, + top_p=args.top_p, + top_k=args.top_k, + reasoning_effort=args.reasoning_effort, + chat_template_kwargs=args.chat_template_kwargs, ) # Print results to terminal @@ -395,6 +484,20 @@ def main() -> None: # Optional file saving if args.save_results: + # None means the field was not sent and the server default applied. + result["request_params"] = { + "endpoint": "chat" if args.use_chat_completions else "completions", + "model": args.model, + "temperature": temperature, + "top_p": args.top_p, + "top_k": args.top_k, + "seed": args.seed, + "stop": None + if args.use_chat_completions + else ["Question", "Assistant:", "<|separator|>"], + "reasoning_effort": args.reasoning_effort, + "chat_template_kwargs": args.chat_template_kwargs, + } with open(args.save_results, "w") as f: json.dump(result, f, indent=2) print(f"Results saved to {args.save_results}") From 819852df6bffb873304c421f7091b7f8133fa87b Mon Sep 17 00:00:00 2001 From: PatchyTIS <58251192+PatchouliTIS@users.noreply.github.com> Date: Wed, 7 Oct 2026 04:32:59 +0800 Subject: [PATCH 093/952] [ModelRunner V2] Speculative Decoding NGram GPU Implementations (#40704) Signed-off-by: PatchouliTaisa Signed-off-by: PatchouliTaisa Signed-off-by: Nick Hill Signed-off-by: Xuanan Chen Co-authored-by: PatchouliTaisa Co-authored-by: Nick Hill Co-authored-by: Xuanan Chen Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> --- tests/test_config.py | 5 +- tests/v1/spec_decode/test_max_len.py | 33 +++ tests/v1/spec_decode/test_ngram_gpu.py | 279 ++++++++++++++++++ ...st_gpu_extract_hidden_states_speculator.py | 3 +- tests/v1/worker/test_gpu_model_runner_v2.py | 20 ++ vllm/compilation/backends.py | 5 +- vllm/config/speculative.py | 3 + vllm/config/vllm.py | 3 - .../distributed/elastic_ep/elastic_execute.py | 4 +- vllm/v1/worker/gpu/model_runner.py | 72 +++-- vllm/v1/worker/gpu/spec_decode/__init__.py | 18 +- .../worker/gpu/spec_decode/ngram/__init__.py | 0 .../gpu/spec_decode/ngram/speculator.py | 278 +++++++++++++++++ vllm/v1/worker/gpu/states.py | 7 +- 14 files changed, 687 insertions(+), 43 deletions(-) create mode 100644 tests/v1/spec_decode/test_ngram_gpu.py create mode 100644 vllm/v1/worker/gpu/spec_decode/ngram/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/ngram/speculator.py diff --git a/tests/test_config.py b/tests/test_config.py index a0c64099dc32..0098ecdc3dac 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -922,12 +922,13 @@ def test_dsa_models_select_matching_mtp(model_type, expected_architecture): assert hf_config.architectures == [expected_architecture] -def test_v2_model_runner_supports_extract_hidden_states(): +@pytest.mark.parametrize("method", ["extract_hidden_states", "ngram", "ngram_gpu"]) +def test_v2_model_runner_supports_speculative_method(method): config = VllmConfig() config.speculative_config = cast( SpeculativeConfig, SimpleNamespace( - method="extract_hidden_states", + method=method, parallel_drafting=False, enable_adaptive_verification=False, ), diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 81f842419375..17dd318563b7 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -36,6 +36,39 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): runner.llm.generate(_PROMPTS, sampling_params) +@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) +@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"]) +def test_ngram_gpu_max_len( + method: str, + num_speculative_tokens: int, + vllm_runner, + monkeypatch: pytest.MonkeyPatch, +): + """V2 n-gram decoding stops at max_model_len.""" + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + "facebook/opt-125m", + trust_remote_code=False, + max_model_len=100, + enable_chunked_prefill=None, + enforce_eager=True, # For faster initialization. + speculative_config={ + "method": method, + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": num_speculative_tokens, + }, + ) as runner: + assert runner.llm.llm_engine.vllm_config.use_v2_model_runner + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + outputs = runner.llm.generate(_PROMPTS, sampling_params) + for output in outputs: + assert output.prompt_token_ids is not None + assert ( + len(output.prompt_token_ids) + len(output.outputs[0].token_ids) == 100 + ) + + @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) def test_eagle_max_len( diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py new file mode 100644 index 000000000000..2359f0cdef2e --- /dev/null +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU-accelerated n-gram speculator. + +These tests target the Triton proposer in +``vllm.v1.worker.gpu.spec_decode.ngram.speculator`` and complement the CPU +``NgramProposer`` tests in ``test_ngram.py``. The GPU speculator follows a +slightly different policy than the CPU one: when multiple n-gram matches of +the same length exist, the GPU kernel picks the right-most (most recent) +match inside the active context, whereas the CPU implementation returns the +left-most. The expectations below reflect the GPU behavior. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import ( + ModelConfig, + SchedulerConfig, + SpeculativeConfig, + VllmConfig, +) +from vllm.v1.worker.gpu.spec_decode.ngram.speculator import NgramGPUSpeculator +from vllm.v1.worker.gpu.states import RequestState + +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for NgramGPUSpeculator tests", + allow_module_level=True, + ) + +DEVICE = torch.device("cuda") + + +def _make_vllm_config( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 8, + max_model_len: int = 64, + method: str = "ngram_gpu", +) -> VllmConfig: + model_config = ModelConfig( + model="facebook/opt-125m", + max_model_len=max_model_len, + enforce_eager=True, + ) + scheduler_config = SchedulerConfig.default_factory( + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) + speculative_config = SpeculativeConfig( + method=method, + prompt_lookup_min=min_n, + prompt_lookup_max=max_n, + num_speculative_tokens=k, + ) + return VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + speculative_config=speculative_config, + ) + + +def _make_request_state(cfg: VllmConfig) -> RequestState: + assert cfg.speculative_config is not None + return RequestState( + max_num_reqs=cfg.scheduler_config.max_num_seqs, + max_model_len=cfg.model_config.max_model_len, + max_num_batched_tokens=cfg.scheduler_config.max_num_batched_tokens, + num_speculative_steps=cfg.speculative_config.num_speculative_tokens, + vocab_size=cfg.model_config.get_vocab_size(), + device=DEVICE, + use_dense_all_token_ids=True, + ) + + +def _make_speculator( + min_n: int, + max_n: int, + k: int, + max_num_seqs: int = 8, + max_model_len: int = 32, +) -> NgramGPUSpeculator: + cfg = _make_vllm_config( + min_n=min_n, + max_n=max_n, + k=k, + max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + ) + return NgramGPUSpeculator(cfg, DEVICE, _make_request_state(cfg)) + + +def _propose( + spec: NgramGPUSpeculator, + rows: list[list[int]], + seq_lens: list[int] | None = None, + num_sampled: list[int] | None = None, + last_sampled: list[int] | None = None, + slots: list[int] | None = None, +) -> list[list[int]]: + """Place each batch row at a request slot and run propose(). + + Returns the drafts as python lists in batch order. + """ + B = len(rows) + if seq_lens is None: + seq_lens = [len(r) for r in rows] + if num_sampled is None: + num_sampled = [1] * B + if last_sampled is None: + last_sampled = [0] * B + if slots is None: + slots = list(range(B)) + + max_num_reqs = spec.max_num_reqs + all_token_ids = spec.req_states.all_token_ids.gpu + total_len = spec.req_states.total_len.gpu + all_token_ids.zero_() + total_len.zero_() + last_sampled_t = torch.zeros((max_num_reqs, 1), dtype=torch.int64, device=DEVICE) + for row, slot, seq_len, last in zip(rows, slots, seq_lens, last_sampled): + if row: + all_token_ids[slot, : len(row)] = torch.tensor( + row, dtype=torch.int32, device=DEVICE + ) + total_len[slot] = seq_len + last_sampled_t[slot, 0] = last + + idx_mapping = torch.tensor(slots, dtype=torch.int64, device=DEVICE) + input_batch = SimpleNamespace(num_reqs=B, idx_mapping=idx_mapping) + + drafts = spec.propose( + input_batch=input_batch, + attn_metadata=None, + slot_mappings=None, + last_hidden_states=torch.empty(0, device=DEVICE), + aux_hidden_states=None, + num_sampled=torch.tensor(num_sampled, dtype=torch.int32, device=DEVICE), + num_rejected=torch.zeros(B, dtype=torch.int32, device=DEVICE), + last_sampled=last_sampled_t, + next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), + temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), + seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), + dp_sync=None, + ) + return drafts.cpu().tolist() + + +@pytest.mark.parametrize("max_model_len", [32, 300]) +def test_no_match_clears_previous_proposal(max_model_len): + spec = _make_speculator(min_n=2, max_n=2, k=2, max_model_len=max_model_len) + row = [0] * (max_model_len - 5) + [1, 2, 3, 1, 2] + assert _propose(spec, [row]) == [[3, 1]] + drafts = _propose(spec, [[1, 2, 3, 4, 5]], last_sampled=[42]) + assert drafts == [[42, 42]] + + +def test_no_4gram_match_only(): + """No 4-gram match in [1,2,3,4,1,2,3] → 0 valid drafts.""" + spec = _make_speculator(min_n=4, max_n=4, k=2) + drafts = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]], last_sampled=[7]) + assert drafts == [[7, 7]] + + +def test_falls_back_to_3gram_when_4gram_missing(): + """No 4-gram match but a 3-gram match exists → propose [4, 1].""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts = _propose(spec, [[1, 2, 3, 4, 1, 2, 3]]) + assert drafts == [[4, 1]] + + +def test_prefers_longer_ngram(): + """Prefer a 4-gram match over a more recent 3-gram match.""" + spec = _make_speculator(min_n=3, max_n=4, k=2) + drafts = _propose(spec, [[1, 2, 3, 4, 50, 51, 2, 3, 4, 60, 61, 1, 2, 3, 4]]) + assert drafts == [[50, 51]] + + +def test_picks_longest_match_among_2_3_4_grams(): + """Prefer a 3-gram match over a more recent 2-gram match.""" + spec = _make_speculator(min_n=2, max_n=4, k=2) + drafts = _propose(spec, [[2, 3, 4, 50, 51, 3, 4, 60, 61, 1, 2, 3, 4]]) + assert drafts == [[50, 51]] + + +@pytest.mark.parametrize("max_model_len", [32, 128, 257, 1025]) +def test_picks_rightmost_when_multiple_matches(max_model_len): + """Pick the last valid match across blocks, ignoring trailing tokens.""" + spec = _make_speculator(min_n=3, max_n=3, k=2, max_model_len=max_model_len) + padding = [0] * (spec.block_l - 5) if spec.n_blocks > 1 else [] + row = [1, 2, 3, 100] + padding + [1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3] + drafts = _propose(spec, [row + [1, 2, 3, 999, 1, 2, 3]], seq_lens=[len(row)]) + assert drafts == [[300, 1]] + + +def test_short_context_yields_zero_valid(): + """The only length-2 window overlaps the suffix itself → no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts = _propose(spec, [[5, 6]], last_sampled=[99]) + assert drafts == [[99, 99]] + + +def test_zero_sampled_disables_proposal(): + """num_sampled==0 disables proposals for that request regardless of match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts = _propose(spec, [[1, 2, 3, 1, 2]], num_sampled=[0], last_sampled=[77]) + assert drafts == [[77, 77]] + + +def test_tail_falls_back_when_few_tokens_after_match(): + """Fewer than k tokens after the match → the tail falls back. + + Tokens: [1, 2, 1, 2] (seq_len=4). Suffix (1, 2) matches at position 0 + (the match at position 2 is the suffix itself and is excluded). With + k=3, only 2 slots map to tokens inside the context. + """ + spec = _make_speculator(min_n=2, max_n=2, k=3) + drafts = _propose(spec, [[1, 2, 1, 2]], last_sampled=[55]) + assert drafts == [[1, 2, 55]] + + +def test_multibatch_mixed(): + """Mixed batch: row 0 matches, row 1 has no match.""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts = _propose( + spec, + [[1, 2, 3, 1, 2], [4, 5, 6]], + last_sampled=[10, 20], + ) + assert drafts[0] == [3, 1] + assert drafts[1] == [20, 20] + + +def test_multibatch_independent_choice_of_n(): + """Each row independently picks its longest matched n.""" + spec = _make_speculator(min_n=2, max_n=3, k=2) + drafts = _propose( + spec, + [ + [9, 1, 2, 3, 8, 1, 2, 3], # 3-gram (1,2,3) at idx 1 → [8, 1] + [7, 1, 2, 9, 1, 2], # 2-gram (1,2) at idx 1 → [9, 1] + ], + ) + assert drafts[0] == [8, 1] + assert drafts[1] == [9, 1] + + +def test_min_n_eq_1(): + """min_n=max_n=1 — single-token n-grams always match if context > 1.""" + spec = _make_speculator(min_n=1, max_n=1, k=2) + drafts = _propose(spec, [[1, 2, 3, 4, 1]]) + assert drafts == [[2, 3]] + + +def test_noncontiguous_idx_mapping(): + """propose() reads token rows in place via idx_mapping (non-contiguous).""" + spec = _make_speculator(min_n=2, max_n=2, k=2) + drafts = _propose( + spec, + [[7, 8, 9, 7, 8], [1, 2, 3, 1, 2]], + slots=[3, 0], + ) + assert drafts == [[9, 7], [3, 1]] + + +def test_construction_validates_speculative_config(): + spec = _make_speculator(min_n=2, max_n=3, k=2) + assert spec.min_n == 2 + assert spec.max_n == 3 + assert spec.num_speculative_steps == 2 + # No-op hooks must not raise. + spec.init_cudagraph_manager(None) + spec.capture() diff --git a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py index 74464f542ffc..5fa44d347569 100644 --- a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py +++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py @@ -45,7 +45,8 @@ def fake_speculator(config, target_device): monkeypatch.setattr(spec_module, "ExtractHiddenStatesSpeculator", fake_speculator) - assert init_speculator(vllm_config, device) == (vllm_config, device) + speculator = init_speculator(vllm_config, device, cast(Any, None)) + assert speculator == (vllm_config, device) def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch): diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 0caa132d7e45..a10f54a97c75 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -337,3 +337,23 @@ def test_capture_model_profile_only_skips_lock(monkeypatch): runner.capture_model(profile_only=True) assert lock_calls == [] + + +@pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"]) +def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer): + """Targets allocate the MTP hidden buffer only for hidden-state drafters.""" + runner = GPUModelRunner.__new__(GPUModelRunner) + hidden_states = torch.zeros(4, 8) + buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8) + if target_buffer == "absent": + runner.model = SimpleNamespace() + else: + returned = buffer if target_buffer == "tensor" else None + runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) + + out = runner._get_drafter_hidden_states(hidden_states) + + if target_buffer == "tensor": + assert torch.equal(out, buffer[:4]) + else: + assert out is hidden_states diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 43252ce3d8e3..d985d8de757a 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1088,10 +1088,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE. disable_cache = not is_compile_cache_enabled(self.inductor_config) - # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors. + # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm + # torch compile cache errors. The V2 implementation is pure Triton and + # does not need the cache disabled. is_ngram_gpu_enabled = ( vllm_config.speculative_config is not None and vllm_config.speculative_config.use_ngram_gpu() + and not vllm_config.use_v2_model_runner ) disable_cache = disable_cache or is_ngram_gpu_enabled diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 47ffb3b7f7f9..cd42206f983a 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1970,6 +1970,9 @@ def uses_extract_hidden_states(self) -> bool: def use_ngram_gpu(self) -> bool: return self.method == "ngram_gpu" + def use_ngram(self) -> bool: + return self.method in ("ngram", "ngram_gpu") + def use_multi_module_mtp(self) -> bool: if self.method != "mtp" or self.draft_model_config is None: return False diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b847da4320ab..8fc89f1591cd 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3140,9 +3140,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: if speculative_config is not None: if speculative_config.method in ( - # https://github.com/vllm-project/vllm/pull/40704 - "ngram", - "ngram_gpu", "suffix", "medusa", "mlp_speculator", diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 10b5e88d29a2..182a4b9bb4f3 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -454,8 +454,8 @@ def _release_cuda_graphs(self) -> None: torch.compiler.reset() with set_current_vllm_config(self.worker.vllm_config): reset_compile_wrapper(self.worker.model_runner.get_model()) - if speculator is not None: - reset_compile_wrapper(speculator.model) + if draft_model := getattr(speculator, "model", None): + reset_compile_wrapper(draft_model) gc.collect() torch.accelerator.synchronize() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e42917918afe..c3ac543ff290 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -265,13 +265,32 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.encoder_cache = EncoderCache() self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache) + self.num_speculative_steps = vllm_config.num_speculative_tokens + num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) + use_dense_all_token_ids = ( + self.speculative_config is not None and self.speculative_config.use_ngram() + ) + + # General request states. + self.req_states = RequestState( + max_num_reqs=self.max_num_reqs, + max_model_len=self.max_model_len, + max_num_batched_tokens=self.max_num_tokens, + num_speculative_steps=self.num_speculative_steps, + vocab_size=self.vocab_size, + device=self.device, + num_prefill_lookahead=num_prefill_lookahead, + use_dense_all_token_ids=use_dense_all_token_ids, + ) + # Speculative decoding. self.speculator = None self.use_aux_hidden_state_outputs = False - self.num_speculative_steps = vllm_config.num_speculative_tokens if self.speculative_config is not None: if self.is_last_pp_rank: - self.speculator = init_speculator(self.vllm_config, self.device) + self.speculator = init_speculator( + self.vllm_config, self.device, self.req_states + ) if self.speculative_config.method in ( "eagle3", @@ -291,20 +310,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None - num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) - self.step_timing = StepTimingCollector() - - # General request states. - self.req_states = RequestState( - max_num_reqs=self.max_num_reqs, - max_model_len=self.max_model_len, - max_num_batched_tokens=self.max_num_tokens, - num_speculative_steps=self.num_speculative_steps, - vocab_size=self.vocab_size, - device=self.device, - num_prefill_lookahead=num_prefill_lookahead, - ) self.adaptive_verification: AdaptiveVerificationManager | None = None self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, @@ -909,14 +915,7 @@ def _dummy_run( ) mm_inputs = [], all_false - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = self._get_drafter_hidden_states(hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping @@ -2075,6 +2074,24 @@ def execute_model( ) return None + def _get_drafter_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Hidden states fed to the drafter. + + Targets such as DeepSeek V4 expose the pre-hc_head residual through + get_mtp_target_hidden_states(). The buffer is sized at + max_num_batched_tokens and only allocated for drafters that consume + target hidden states, so None means "use the regular hidden states". + """ + get_target_hidden_states = getattr( + self.model, "get_mtp_target_hidden_states", None + ) + if get_target_hidden_states is None: + return hidden_states + target_hidden_states = get_target_hidden_states() + if target_hidden_states is None: + return hidden_states + return target_hidden_states[: hidden_states.shape[0]] + @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2220,14 +2237,7 @@ def sample_tokens( self.speculator.observe_verification( input_batch.idx_mapping, num_sampled, num_rejected ) - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = draft_hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] + spec_hidden_states = self._get_drafter_hidden_states(draft_hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 0616f12dd57a..902a4f085ca9 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,11 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import torch from vllm.config import VllmConfig +if TYPE_CHECKING: + from vllm.v1.worker.gpu.states import RequestState + -def init_speculator(vllm_config: VllmConfig, device: torch.device): +def init_speculator( + vllm_config: VllmConfig, + device: torch.device, + req_states: "RequestState", +): + """Build the speculator for this config.""" speculative_config = vllm_config.speculative_config assert speculative_config is not None if speculative_config.method == "extract_hidden_states": @@ -66,5 +76,11 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): ) return PlainDraftModelSpeculator(vllm_config, device) + elif speculative_config.use_ngram(): + from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( + NgramGPUSpeculator, + ) + + return NgramGPUSpeculator(vllm_config, device, req_states) else: raise NotImplementedError(f"{speculative_config.method} is not supported yet.") diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py b/vllm/v1/worker/gpu/spec_decode/ngram/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py new file mode 100644 index 000000000000..404f576274e8 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.speculator import BaseSpeculator + +if TYPE_CHECKING: + from vllm.config.compilation import CUDAGraphMode + from vllm.v1.worker.gpu.dp_utils import DPSyncState + from vllm.v1.worker.gpu.states import RequestState + + +@triton.jit +def _ngram_scan_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] batch_idx -> req_state_idx + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + scratch_ptr, # *int64 [B, scratch_stride] (output) + scratch_stride, + L, # int64 scalar (= max_model_len) + MIN_N: tl.constexpr, + MAX_N: tl.constexpr, + MAX_N_PO2: tl.constexpr, + BLOCK_L: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + blk = tl.program_id(1).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + eligible_row = (num_sampled > 0) & (seq_len >= MIN_N) + + scratch_off = b * scratch_stride + blk + + # Ineligible rows, and blocks fully past the last candidate match + # position, write 0 and exit. + if not (eligible_row & (blk * BLOCK_L <= seq_len - MIN_N - 1)): + tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64)) + return + + row_off = req_state_idx * token_ids_stride + + # Load the length-MAX_N suffix once into registers. + suf_iota = tl.arange(0, MAX_N_PO2).to(tl.int64) + suf_pos = seq_len - MAX_N + suf_iota + suf_in_range = (suf_iota < MAX_N) & (suf_pos >= 0) & (suf_pos < seq_len) + suffix = tl.load( + token_ids_ptr + row_off + suf_pos, + mask=suf_in_range, + other=-1, + ).to(tl.int32) + + pos_iota = tl.arange(0, BLOCK_L).to(tl.int64) + pos = blk * BLOCK_L + pos_iota # ascending + + best_score = tl.zeros([BLOCK_L], dtype=tl.int64) + + for n_iter in tl.static_range(MIN_N, MAX_N + 1): + max_pos_n = seq_len - n_iter - 1 + match = (pos >= 0) & (pos <= max_pos_n) + for j in tl.static_range(0, n_iter): + tok = tl.load( + token_ids_ptr + row_off + (pos + j), + mask=match, + other=0, + ).to(tl.int32) + suf_idx = (MAX_N - n_iter) + j + suf_val = tl.sum(tl.where(suf_iota == suf_idx, suffix, 0)) + match = match & (tok == suf_val) + + # Pack (n, pos) so a single max yields longest-n, rightmost-pos. + cand = n_iter * Lp1 + pos + 1 + best_score = tl.where(match, cand, best_score) + + block_best = tl.max(best_score, axis=0) + tl.store(scratch_ptr + scratch_off, block_best) + + +@triton.jit +def _ngram_finalize_kernel( + token_ids_ptr, # *int32 [max_num_reqs, token_ids_stride] + token_ids_stride, + idx_mapping_ptr, # *int64 [B] + total_len_ptr, # *int32 [max_num_reqs] + num_sampled_ptr, # *int32 [B] + last_sampled_ptr, # *int64 [max_num_reqs] + scratch_ptr, # *int64 [B, scratch_stride] + scratch_stride, + drafts_ptr, # *int64 [B, K] (output) + L, + N_BLOCKS, + K: tl.constexpr, + K_PO2: tl.constexpr, + N_BLOCKS_PO2: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + Lp1 = tl.cast(L, tl.int64) + 1 + NB = tl.cast(N_BLOCKS, tl.int64) + + req_state_idx = tl.load(idx_mapping_ptr + b).to(tl.int64) + + nb_iota = tl.arange(0, N_BLOCKS_PO2).to(tl.int64) + nb_in_range = nb_iota < NB + block_scores = tl.load( + scratch_ptr + b * scratch_stride + nb_iota, + mask=nb_in_range, + other=0, + ) + score = tl.max(block_scores, axis=0) + + seq_len = tl.load(total_len_ptr + req_state_idx).to(tl.int64) + num_sampled = tl.load(num_sampled_ptr + b) + last_tok = tl.load(last_sampled_ptr + req_state_idx) + + has_match = score > 0 + s1 = score - 1 + best_n = tl.where(has_match, s1 // Lp1, tl.zeros_like(s1)) + best_pos = tl.where(has_match, s1 - best_n * Lp1, tl.zeros_like(s1)) + draft_start = tl.where(has_match, best_pos + best_n, tl.zeros_like(s1)) + + tokens_avail = tl.maximum(seq_len - draft_start, 0) + write_ok = (num_sampled > 0) & has_match + + row_off = req_state_idx * token_ids_stride + k_iota = tl.arange(0, K_PO2).to(tl.int64) + k_in_range = k_iota < K + gather_idx = tl.minimum(draft_start + k_iota, tl.cast(L, tl.int64) - 1) + slot_valid = (k_iota < tokens_avail) & write_ok & k_in_range + gathered = tl.load( + token_ids_ptr + row_off + gather_idx, + mask=slot_valid, + other=0, + ).to(tl.int64) + # Invalid slots fall back to the last sampled token; they are verified as + # ordinary (rejectable) drafts, so the fill value only affects efficiency. + out = tl.where(slot_valid, gathered, last_tok) + tl.store(drafts_ptr + b * K + k_iota, out, mask=k_in_range) + + +class NgramGPUSpeculator(BaseSpeculator): + """V2-compatible GPU n-gram speculator.""" + + supports_mm_inputs = False + draft_logits = None + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + req_states: RequestState, + ): + if not HAS_TRITON: + raise RuntimeError("ngram_gpu speculative decoding requires Triton.") + spec = vllm_config.speculative_config + assert spec is not None + assert spec.prompt_lookup_min is not None, ( + "prompt_lookup_min must be configured for ngram_gpu" + ) + assert spec.prompt_lookup_max is not None, ( + "prompt_lookup_max must be configured for ngram_gpu" + ) + assert 1 <= spec.prompt_lookup_min <= spec.prompt_lookup_max + + self.vllm_config = vllm_config + self.device = device + self.req_states = req_states + self.speculative_config = spec + self.num_speculative_steps: int = spec.num_speculative_tokens + + self.min_n: int = spec.prompt_lookup_min + self.max_n: int = spec.prompt_lookup_max + + self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs + self.max_model_len: int = vllm_config.model_config.max_model_len + + L = self.max_model_len + if L >= 1024: + self.block_l = 256 + elif L >= 256: + self.block_l = 128 + elif L >= 64: + self.block_l = 64 + else: + self.block_l = max(16, triton.next_power_of_2(max(L, 1))) + self.n_blocks = triton.cdiv(L, self.block_l) + + self.scratch = torch.zeros( + (self.max_num_reqs, self.n_blocks), dtype=torch.int64, device=device + ) + # Batch-ordered draft output, scattered into RequestState.draft_tokens + # by the model runner (same contract as the model-based speculators). + self.drafts = torch.zeros( + (self.max_num_reqs, self.num_speculative_steps), + dtype=torch.int64, + device=device, + ) + + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + del cudagraph_mode + + def capture(self) -> None: + return None + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: Any, + slot_mappings: Any, + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + dp_sync: DPSyncState | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + if dummy_run: + return self.drafts[:num_reqs] + + req_states = self.req_states + token_ids = req_states.all_token_ids.gpu + idx_mapping = input_batch.idx_mapping + + _ngram_scan_kernel[(num_reqs, self.n_blocks)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + self.scratch, + self.scratch.stride(0), + self.max_model_len, + self.min_n, + self.max_n, + max(1, triton.next_power_of_2(self.max_n)), + self.block_l, + num_warps=4, + num_stages=2, + ) + + _ngram_finalize_kernel[(num_reqs,)]( + token_ids, + token_ids.stride(0), + idx_mapping, + req_states.total_len.gpu, + num_sampled, + last_sampled.view(-1), + self.scratch, + self.scratch.stride(0), + self.drafts, + self.max_model_len, + self.n_blocks, + self.num_speculative_steps, + max(1, triton.next_power_of_2(self.num_speculative_steps)), + max(1, triton.next_power_of_2(self.n_blocks)), + num_warps=2, + num_stages=1, + ) + return self.drafts[:num_reqs] diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index db27bf013f73..16260b0ddb9e 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -16,6 +16,7 @@ def __init__( vocab_size: int, device: torch.device, num_prefill_lookahead: int = 1, + use_dense_all_token_ids: bool = False, ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -30,12 +31,14 @@ def __init__( # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs) # depending on the configured max_num_reqs and max_model_len. - # To save GPU memory, we use UVA instead of GPU for this tensor. + # To save GPU memory, we use UVA instead of GPU by default, but + # ngram_gpu benefits from dense device residency because it scans + # active rows repeatedly during proposal. self.all_token_ids = StagedWriteTensor( (self.max_num_reqs, self.max_model_len), dtype=torch.int32, device=device, - uva_instead_of_gpu=True, + uva_instead_of_gpu=not use_dense_all_token_ids, ) # NOTE(woosuk): Distinguish clearly between prompt_len and prefill_len: # - prompt_len: Number of tokens in the user-provided prompt. From 385d86a6cf6f05051899b53d83ec55c63cccb379 Mon Sep 17 00:00:00 2001 From: aditi-amd Date: Tue, 6 Oct 2026 13:42:52 -0700 Subject: [PATCH 094/952] [Attention] Add UltraQuant 4-bit KV cache backend (FlyDSL D=256) (#57057) Signed-off-by: Aditi Rana --- tests/kernels/ultraquant/test_config.py | 36 + .../ultraquant/test_decode_vs_reference.py | 111 ++ .../ultraquant/test_dequant_vs_reference.py | 88 ++ .../ultraquant/test_flydsl_decode_d256.py | 149 +++ .../kernels/ultraquant/test_seam_contract.py | 106 ++ .../ultraquant/test_store_vs_reference.py | 149 +++ vllm/config/cache.py | 1 + vllm/engine/arg_utils.py | 18 +- vllm/platforms/interface.py | 34 +- vllm/platforms/rocm.py | 59 +- vllm/utils/torch_utils.py | 1 + vllm/v1/attention/backends/registry.py | 1 + vllm/v1/attention/backends/ultraquant_attn.py | 609 +++++++++ .../flydsl_kernels/ultraquant_decode_hd256.py | 983 +++++++++++++++ .../attention/ops/flydsl_ultraquant_decode.py | 834 +++++++++++++ vllm/v1/attention/ops/ultraquant/__init__.py | 19 + vllm/v1/attention/ops/ultraquant/format.py | 218 ++++ vllm/v1/attention/ops/ultraquant/reference.py | 349 ++++++ .../ops/ultraquant/triton_dequant.py | 190 +++ .../attention/ops/ultraquant/triton_store.py | 351 ++++++ .../ultraquant/triton_unified_attention.py | 1087 +++++++++++++++++ vllm/v1/kv_cache_interface.py | 8 + 22 files changed, 5377 insertions(+), 24 deletions(-) create mode 100644 tests/kernels/ultraquant/test_config.py create mode 100644 tests/kernels/ultraquant/test_decode_vs_reference.py create mode 100644 tests/kernels/ultraquant/test_dequant_vs_reference.py create mode 100644 tests/kernels/ultraquant/test_flydsl_decode_d256.py create mode 100644 tests/kernels/ultraquant/test_seam_contract.py create mode 100644 tests/kernels/ultraquant/test_store_vs_reference.py create mode 100644 vllm/v1/attention/backends/ultraquant_attn.py create mode 100644 vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py create mode 100644 vllm/v1/attention/ops/flydsl_ultraquant_decode.py create mode 100644 vllm/v1/attention/ops/ultraquant/__init__.py create mode 100644 vllm/v1/attention/ops/ultraquant/format.py create mode 100644 vllm/v1/attention/ops/ultraquant/reference.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_dequant.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_store.py create mode 100644 vllm/v1/attention/ops/ultraquant/triton_unified_attention.py diff --git a/tests/kernels/ultraquant/test_config.py b/tests/kernels/ultraquant/test_config.py new file mode 100644 index 000000000000..14036b44e853 --- /dev/null +++ b/tests/kernels/ultraquant/test_config.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only integration tests for the UltraQuant cache dtype.""" + +from vllm.v1.attention.backends.ultraquant_attn import ( + UltraQuantAttentionBackend, +) +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + ultraquant_flydsl_decode_eligible, +) + + +def test_ultraquant_prefers_block_size_64(): + # FlyDSL D=256 decode geometry requires a 64-token KV block. + assert UltraQuantAttentionBackend.get_preferred_block_size(16) == 64 + assert UltraQuantAttentionBackend.get_preferred_block_size(128) == 64 + + +def test_ultraquant_flydsl_decode_routing(): + """CPU check of FlyDSL vs unified-Triton decode eligibility.""" + eligible = dict( + head_size=256, + num_kv_groups=8, + has_sinks=False, + sliding_window=None, + flydsl_loaded=True, + ) + assert ultraquant_flydsl_decode_eligible(**eligible) + for gqa in (6, 8, 16): + assert ultraquant_flydsl_decode_eligible(**{**eligible, "num_kv_groups": gqa}) + + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "flydsl_loaded": False}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "head_size": 128}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "num_kv_groups": 4}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "has_sinks": True}) + assert not ultraquant_flydsl_decode_eligible(**{**eligible, "sliding_window": 4096}) diff --git a/tests/kernels/ultraquant/test_decode_vs_reference.py b/tests/kernels/ultraquant/test_decode_vs_reference.py new file mode 100644 index 000000000000..9a8d431001f6 --- /dev/null +++ b/tests/kernels/ultraquant/test_decode_vs_reference.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the UltraQuant Triton unified-attention fallback.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import reference_ultraquant_attention +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store +from vllm.v1.attention.ops.ultraquant.triton_unified_attention import ( + ultraquant_unified_attention, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _build_paged_cache(key, value, block_size): + batch_size, seq_len, num_kv_heads, head_size = key.shape + blocks_per_seq = (seq_len + block_size - 1) // block_size + kv_cache = torch.zeros( + batch_size * blocks_per_seq, + block_size, + num_kv_heads, + slot_size(head_size), + dtype=torch.uint8, + device=key.device, + ) + block_table = torch.arange( + batch_size * blocks_per_seq, + dtype=torch.int32, + device=key.device, + ).view(batch_size, blocks_per_seq) + slot_mapping = torch.arange( + batch_size * blocks_per_seq * block_size, + dtype=torch.int64, + device=key.device, + ).view(batch_size, -1)[:, :seq_len] + ultraquant_store( + key.reshape(-1, num_kv_heads, head_size), + value.reshape(-1, num_kv_heads, head_size), + kv_cache, + slot_mapping.reshape(-1), + ) + return kv_cache, block_table + + +def _attention_quality(actual: torch.Tensor, expected: torch.Tensor) -> dict: + a = actual.float().reshape(-1) + b = expected.float().reshape(-1) + diff = (a - b).abs() + return { + "cos_sim": torch.nn.functional.cosine_similarity(a, b, dim=0).item(), + "mean_abs": diff.mean().item(), + "max_abs": diff.max().item(), + "norm_ratio": (a.norm() / b.norm()).item(), + } + + +def _unified_decode(query, kv_cache, block_table, seq_lens, scale, sinks=None): + batch_size = query.shape[0] + return ultraquant_unified_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + query_start_loc=torch.arange( + batch_size + 1, dtype=torch.int32, device=query.device + ), + scale=scale, + max_query_len=1, + max_seq_len=int(seq_lens.max().item()), + sinks=sinks, + ) + + +@pytest.mark.parametrize("query_group_size", [6, 8]) +def test_unified_decode_d256_matches_reference(query_group_size): + """The D=256 fallback covers both target Qwen GQA layouts.""" + device = torch.device("cuda") + set_random_seed(0xD256 + query_group_size) + + B, Hk, N, D = 2, 1, 256, 256 + Hq = query_group_size * Hk + block_size = 32 + dtype = torch.bfloat16 + scale = D**-0.5 + + query = torch.randn(B, Hq, D, dtype=dtype, device=device) + key = torch.randn(B, N, Hk, D, dtype=dtype, device=device) + value = torch.randn_like(key) + kv_cache, block_table = _build_paged_cache(key, value, block_size) + seq_lens = torch.full((B,), N, dtype=torch.int32, device=device) + + actual = _unified_decode(query, kv_cache, block_table, seq_lens, scale) + expected = reference_ultraquant_attention( + query=query, + key=key, + value=value, + scale=scale, + ) + + quality = _attention_quality(actual, expected) + assert quality["cos_sim"] > 0.98, quality + assert quality["mean_abs"] < 0.08, quality + assert quality["max_abs"] < 0.8, quality diff --git a/tests/kernels/ultraquant/test_dequant_vs_reference.py b/tests/kernels/ultraquant/test_dequant_vs_reference.py new file mode 100644 index 000000000000..db56f56e2947 --- /dev/null +++ b/tests/kernels/ultraquant/test_dequant_vs_reference.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the UltraQuant full-KV dequant kernel (continuation prefill). + +The kernel reconstructs cached K/V from FP4 codes plus UE8M0 group scales. +Because dequant is an exact table lookup times a power of two, the result +must match the reference to within the output dtype's rounding. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import ( + ultraquant_dequant, + ultraquant_encode, +) +from vllm.v1.attention.ops.ultraquant.triton_dequant import ( + ultraquant_full_dequant_kv, +) +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("cached_len", [32, 96]) +def test_full_dequant_matches_reference(dtype, head_dim, num_kv_heads, cached_len): + device = torch.device("cuda") + set_random_seed(0xD4) + + block_size = 32 + alloc_len = cached_len + num_blocks = alloc_len // block_size + 1 + + key = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device) + value = torch.randn(alloc_len, num_kv_heads, head_dim, dtype=dtype, device=device) + + kv_cache = torch.zeros( + num_blocks, + block_size, + num_kv_heads, + slot_size(head_dim), + dtype=torch.uint8, + device=device, + ) + # Leave block 0 unused so the block table exercises a non-identity mapping. + slot_mapping = torch.arange( + block_size, block_size + alloc_len, dtype=torch.int64, device=device + ) + ultraquant_store(key, value, kv_cache, slot_mapping) + + block_table = torch.arange( + 1, num_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + + k_out = torch.zeros( + 1, num_kv_heads, alloc_len, head_dim, dtype=dtype, device=device + ) + v_out = torch.zeros_like(k_out) + ultraquant_full_dequant_kv( + kv_cache=kv_cache, + block_table=block_table, + k_out=k_out, + v_out=v_out, + alloc_len=alloc_len, + ) + torch.accelerator.synchronize() + + # K is stored Hadamard-rotated, V unrotated; the kernel returns them as + # stored, so the reference must not inverse-rotate either. + ref_k = ultraquant_dequant(ultraquant_encode(key, rotate=True), head_dim) + ref_v = ultraquant_dequant(ultraquant_encode(value, rotate=False), head_dim) + + got_k = k_out[0].transpose(0, 1).to(torch.float32) + got_v = v_out[0].transpose(0, 1).to(torch.float32) + + # Reference is fp32; the kernel rounds to the output dtype once. + tol = 1e-2 if dtype is torch.bfloat16 else 1e-3 + torch.testing.assert_close(got_k, ref_k.to(torch.float32), rtol=tol, atol=tol) + torch.testing.assert_close(got_v, ref_v.to(torch.float32), rtol=tol, atol=tol) diff --git a/tests/kernels/ultraquant/test_flydsl_decode_d256.py b/tests/kernels/ultraquant/test_flydsl_decode_d256.py new file mode 100644 index 000000000000..33a1a8273f20 --- /dev/null +++ b/tests/kernels/ultraquant/test_flydsl_decode_d256.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the production UltraQuant D=256 FlyDSL decode.""" + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + flydsl_ultraquant_decode_attention, + is_ultraquant_flydsl_available, + ultraquant_flydsl_decode_eligible, +) +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.reference import ( + hadamard_matrix, + reference_ultraquant_attention, +) +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _build_paged_cache(key, value, block_size): + batch_size, seq_len, num_kv_heads, head_size = key.shape + blocks_per_seq = (seq_len + block_size - 1) // block_size + kv_cache = torch.zeros( + batch_size * blocks_per_seq, + block_size, + num_kv_heads, + slot_size(head_size), + dtype=torch.uint8, + device=key.device, + ) + block_table = torch.arange( + batch_size * blocks_per_seq, + dtype=torch.int32, + device=key.device, + ).view(batch_size, blocks_per_seq) + slot_mapping = torch.arange( + batch_size * blocks_per_seq * block_size, + dtype=torch.int64, + device=key.device, + ).view(batch_size, -1)[:, :seq_len] + ultraquant_store( + key.reshape(-1, num_kv_heads, head_size), + value.reshape(-1, num_kv_heads, head_size), + kv_cache, + slot_mapping.reshape(-1), + ) + return kv_cache, block_table + + +@pytest.mark.parametrize( + ("batch_size", "query_group_size", "block_size", "contiguous_query"), + [ + (1, 8, 32, True), + (8, 8, 64, True), + (9, 8, 256, True), + (1, 16, 32, True), + (1, 6, 16, True), + (1, 6, 32, True), + (2, 6, 64, True), + (1, 6, 128, True), + (1, 6, 256, True), + (1, 6, 32, False), + ], +) +def test_flydsl_decode_matches_reference( + batch_size, + query_group_size, + block_size, + contiguous_query, +): + if not ultraquant_flydsl_decode_eligible( + head_size=256, + num_kv_groups=query_group_size, + has_sinks=False, + sliding_window=None, + flydsl_loaded=is_ultraquant_flydsl_available(), + ): + pytest.skip("UltraQuant D=256 FlyDSL kernel is unavailable") + + set_random_seed(0x57A1D + batch_size + block_size + query_group_size) + device = torch.device("cuda") + dtype = torch.bfloat16 + num_kv_heads = 1 + head_size = 256 + seq_len = 2048 + scale = head_size**-0.5 + + key = torch.randn( + batch_size, + seq_len, + num_kv_heads, + head_size, + dtype=dtype, + device=device, + ) + value = torch.randn_like(key) + if contiguous_query: + qkv = torch.randn( + batch_size, + query_group_size + 2 * num_kv_heads, + head_size, + dtype=dtype, + device=device, + ) + query = qkv[:, :query_group_size] + else: + query_storage = torch.randn( + batch_size, + query_group_size, + head_size * 2, + dtype=dtype, + device=device, + ) + query = query_storage[..., ::2] + kv_cache, block_table = _build_paged_cache(key, value, block_size) + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + hadamard = hadamard_matrix(head_size, device) + + actual = flydsl_ultraquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + scale=scale, + PiT=hadamard, + max_seq_len=seq_len, + ) + expected = reference_ultraquant_attention( + query=query, + key=key, + value=value, + scale=scale, + ) + + actual_f32 = actual.float() + expected_f32 = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_f32.flatten(), expected_f32.flatten(), dim=0 + ) + error = (actual_f32 - expected_f32).abs() + assert cosine > 0.995 + assert error.mean() < 0.05 + assert error.max() < 0.8 diff --git a/tests/kernels/ultraquant/test_seam_contract.py b/tests/kernels/ultraquant/test_seam_contract.py new file mode 100644 index 000000000000..a5f211679787 --- /dev/null +++ b/tests/kernels/ultraquant/test_seam_contract.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant inherits TurboQuant's forward/do_kv_cache_update orchestration. + +UltraQuant only overrides the store/decode/prefill seams, so a TurboQuant-side +change to a seam signature, to an inherited helper, or to the ``layer._tq_*`` +attributes ``forward`` reads will break UltraQuant at serve time without +touching any UltraQuant kernel test. These checks pin that contract. +""" + +import ast +import inspect +import textwrap + +from vllm.v1.attention.backends import ultraquant_attn +from vllm.v1.attention.backends.turboquant_attn import TurboQuantAttentionImpl +from vllm.v1.attention.backends.ultraquant_attn import UltraQuantAttentionImpl + +_SEAMS = ("_ensure_on_device", "_store_kv", "_decode_attention", "_prefill_attention") + + +def _method_ast(cls: type, name: str) -> ast.AST: + return ast.parse(textwrap.dedent(inspect.getsource(getattr(cls, name)))) + + +def _tq_attrs_read(cls: type, name: str) -> set[str]: + """``_tq_*`` attribute names read (not assigned) in ``cls.name``.""" + tree = _method_ast(cls, name) + assigned = { + t.attr + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Attribute) + } + read = { + node.attr + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) and node.attr.startswith("_tq_") + } + return read - assigned + + +def _tq_attrs_written(cls: type, name: str) -> set[str]: + tree = _method_ast(cls, name) + return { + t.attr + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + for t in node.targets + if isinstance(t, ast.Attribute) and t.attr.startswith("_tq_") + } + + +def test_ultraquant_matches_turboquant_seam_contract(): + # 1. Each overridden seam accepts every parameter the base declares, so any + # call the inherited orchestration makes stays valid for UltraQuant. + for name in _SEAMS: + base = list( + inspect.signature(getattr(TurboQuantAttentionImpl, name)).parameters + ) + override = list( + inspect.signature(getattr(UltraQuantAttentionImpl, name)).parameters + ) + assert override[: len(base)] == base, ( + f"UltraQuantAttentionImpl.{name} no longer accepts the arguments " + f"TurboQuantAttentionImpl passes: base={base} override={override}" + ) + + # 2. Keyword arguments UltraQuant passes to *inherited* TurboQuant helpers + # must exist on those helpers (e.g. _flash_attn_varlen). + own = vars(UltraQuantAttentionImpl) + module = ast.parse(inspect.getsource(ultraquant_attn)) + for node in ast.walk(module): + if not isinstance(node, ast.Call): + continue + fn = node.func + if not ( + isinstance(fn, ast.Attribute) + and isinstance(fn.value, ast.Name) + and fn.value.id == "self" + ): + continue + target = getattr(TurboQuantAttentionImpl, fn.attr, None) + if target is None or fn.attr in own or not callable(target): + continue + params = inspect.signature(target).parameters + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + continue + for kw in node.keywords: + if kw.arg is None: + continue + assert kw.arg in params, ( + f"UltraQuant calls inherited {fn.attr}(..., {kw.arg}=...) but " + f"TurboQuantAttentionImpl.{fn.attr} accepts {list(params)}" + ) + + # 3. Every layer._tq_* attribute the inherited forward reads is set by + # UltraQuant's _ensure_on_device. + required = _tq_attrs_read(TurboQuantAttentionImpl, "forward") + provided = _tq_attrs_written(UltraQuantAttentionImpl, "_ensure_on_device") + assert required, "expected inherited forward to read layer._tq_* attributes" + assert required <= provided, ( + f"UltraQuant._ensure_on_device does not set {sorted(required - provided)}, " + f"which the inherited TurboQuant forward reads" + ) diff --git a/tests/kernels/ultraquant/test_store_vs_reference.py b/tests/kernels/ultraquant/test_store_vs_reference.py new file mode 100644 index 000000000000..2f02e9578ce3 --- /dev/null +++ b/tests/kernels/ultraquant/test_store_vs_reference.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Triton UltraQuant store kernel. + +The Triton kernel must produce bit-identical codes and scale bytes vs +the PyTorch reference (``reference.ultraquant_encode``). +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops.ultraquant.format import ( + k_codes_offset, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.reference import ultraquant_encode +from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA/HIP device" +) + + +def _make_cache( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_dim: int, + device: torch.device, +) -> torch.Tensor: + slot_b = slot_size(head_dim) + return torch.zeros( + num_blocks, block_size, num_kv_heads, slot_b, dtype=torch.uint8, device=device + ) + + +def _read_codes_and_scales( + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + num_kv_heads: int, + head_dim: int, +): + """Pull (k_codes, k_scale_bytes, v_codes, v_scale_bytes) for each + (token, head) from the AoS slot layout. Scales come back as uint8 + (one E8M0 byte per group).""" + block_size = kv_cache.shape[1] + N = slot_mapping.shape[0] + Gk = n_groups(head_dim) + + k_codes_off = k_codes_offset(head_dim) + k_scales_off = k_scales_offset(head_dim) + v_codes_off = v_codes_offset(head_dim) + v_scales_off = v_scales_offset(head_dim) + codes_bytes = head_dim // 2 + + k_codes = torch.zeros( + N, num_kv_heads, codes_bytes, dtype=torch.uint8, device=kv_cache.device + ) + k_scale_bytes = torch.zeros( + N, num_kv_heads, Gk, dtype=torch.uint8, device=kv_cache.device + ) + v_codes = torch.zeros_like(k_codes) + v_scale_bytes = torch.zeros_like(k_scale_bytes) + + for ti in range(N): + slot = int(slot_mapping[ti].item()) + if slot < 0: + continue + blk = slot // block_size + off = slot % block_size + for h in range(num_kv_heads): + slot_bytes = kv_cache[blk, off, h] # [slot_size_aligned] uint8 + + k_codes[ti, h] = slot_bytes[k_codes_off : k_codes_off + codes_bytes] + k_scale_bytes[ti, h] = slot_bytes[k_scales_off : k_scales_off + Gk] + v_codes[ti, h] = slot_bytes[v_codes_off : v_codes_off + codes_bytes] + v_scale_bytes[ti, h] = slot_bytes[v_scales_off : v_scales_off + Gk] + + return k_codes, k_scale_bytes, v_codes, v_scale_bytes + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("head_dim", [128, 256]) +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("num_tokens", [1, 16, 17]) +def test_store_matches_reference(dtype, head_dim, num_kv_heads, num_tokens): + """Triton store output (codes + E8M0 scale bytes) matches the reference + bit-exactly.""" + device = torch.device("cuda") + set_random_seed(0xF8) + + block_size = 32 + num_blocks = max(2, (num_tokens + block_size - 1) // block_size + 1) + + key = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device) + value = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device) + + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + if num_tokens >= 4: + slot_mapping[1] = -1 # skip-path test + + kv_cache = _make_cache(num_blocks, block_size, num_kv_heads, head_dim, device) + + ultraquant_store(key, value, kv_cache, slot_mapping) + torch.accelerator.synchronize() + + ref_k = ultraquant_encode(key, rotate=True) + ref_v = ultraquant_encode(value, rotate=False) + + k_codes, k_scale_bytes, v_codes, v_scale_bytes = _read_codes_and_scales( + kv_cache, slot_mapping, num_kv_heads, head_dim + ) + + if num_tokens >= 4: + assert (kv_cache[0, 1, :, :] == 0).all(), ( + "skipped slot's physical cache region was unexpectedly written" + ) + + for ti in range(num_tokens): + slot = int(slot_mapping[ti].item()) + if slot < 0: + continue + + assert torch.equal(k_codes[ti], ref_k.codes_packed[ti]), ( + f"K codes mismatch at token {ti}\n" + f"got: {k_codes[ti].cpu().numpy().tolist()[:16]}\n" + f"expected: {ref_k.codes_packed[ti].cpu().numpy().tolist()[:16]}" + ) + assert torch.equal(k_scale_bytes[ti], ref_k.scale_bytes[ti]), ( + f"K scale bytes mismatch at token {ti}\n" + f"got: {k_scale_bytes[ti].cpu().numpy().tolist()}\n" + f"expected: {ref_k.scale_bytes[ti].cpu().numpy().tolist()}" + ) + + assert torch.equal(v_codes[ti], ref_v.codes_packed[ti]), ( + f"V codes mismatch at token {ti}" + ) + assert torch.equal(v_scale_bytes[ti], ref_v.scale_bytes[ti]), ( + f"V scale bytes mismatch at token {ti}\n" + f"got: {v_scale_bytes[ti].cpu().numpy().tolist()}\n" + f"expected: {ref_v.scale_bytes[ti].cpu().numpy().tolist()}" + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 92b984c0faa2..e6436f98f9cb 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -50,6 +50,7 @@ def _layout_from_name(layout_name: str) -> KVCacheLayout: "turboquant_4bit_nc", "turboquant_k3v4_nc", "turboquant_3bit_nc", + "ultraquant_4bit", "int4_per_token_head", "int8_per_token_head", "fp8_per_token_head", diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 98a0b7cd88fd..527b8956d033 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2271,7 +2271,13 @@ def create_engine_config( kv_offloading_backend=self.kv_offloading_backend, ) - if resolved_cache_dtype.startswith("turboquant_"): + # TurboQuant and UltraQuant both keep boundary attention layers at the + # native dtype, and compute those layers the same way. + uses_packed_kv_backend = ( + resolved_cache_dtype.startswith("turboquant_") + or resolved_cache_dtype == "ultraquant_4bit" + ) + if uses_packed_kv_backend: from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, ) @@ -2681,14 +2687,18 @@ def create_engine_config( # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert # FlashAttentionImpl which fails with TurboQuantAttentionImpl. - if resolved_cache_dtype.startswith("turboquant_") and ( + # UltraQuant subclasses that impl, so it inherits the same limit. + if uses_packed_kv_backend and ( attention_config.flash_attn_version is None or attention_config.flash_attn_version >= 3 ): logger.warning( - "TurboQuant is not yet compatible with FlashAttention >= 3. " + "%s is not yet compatible with FlashAttention >= 3. " "Overriding flash_attn_version to 2. To silence this " - "warning, pass --attention-config.flash_attn_version=2" + "warning, pass --attention-config.flash_attn_version=2", + "TurboQuant" + if resolved_cache_dtype.startswith("turboquant_") + else "UltraQuant", ) attention_config.flash_attn_version = 2 diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 57a7a82f219e..f4583daba26f 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -906,15 +906,29 @@ def _align_hybrid_block_size( cache_dtype_str=cache_config.cache_dtype, kv_quant_mode=kv_quant_mode, ).page_size_bytes - elif cache_config.cache_dtype.startswith("turboquant_"): - # TQ has a packed K|V layout; the standard FullAttentionSpec - # formula over-sizes it and trips unify_kv_cache_spec_page_size - # when all attention layers are TQ. With mixed skip+TQ the skip - # layers still use the standard layout — take max so mamba - # padding covers the largest actual page. - from vllm.v1.attention.backends.turboquant_attn import ( - TurboQuantAttentionBackend, - ) + elif ( + cache_config.cache_dtype.startswith("turboquant_") + or cache_config.cache_dtype == "ultraquant_4bit" + ): + # TurboQuant and UltraQuant have packed K|V layouts; the standard + # FullAttentionSpec formula over-sizes them and trips + # unify_kv_cache_spec_page_size when all attention layers are + # packed. With mixed skip+packed the skip layers still use the + # standard layout — take max so mamba padding covers the largest + # actual page. Each dtype packs via its own backend. + _pack_backend: type[AttentionBackend] + if cache_config.cache_dtype == "ultraquant_4bit": + from vllm.v1.attention.backends.ultraquant_attn import ( + UltraQuantAttentionBackend, + ) + + _pack_backend = UltraQuantAttentionBackend + else: + from vllm.v1.attention.backends.turboquant_attn import ( + TurboQuantAttentionBackend, + ) + + _pack_backend = TurboQuantAttentionBackend tq_spec = FullAttentionSpec( block_size=1, @@ -923,7 +937,7 @@ def _align_hybrid_block_size( dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ) - tq_page = TurboQuantAttentionBackend.customize_spec(tq_spec).page_size_bytes + tq_page = _pack_backend.customize_spec(tq_spec).page_size_bytes if cache_config.kv_cache_dtype_skip_layers: skip_page = FullAttentionSpec( block_size=1, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 487cffa67010..f394e56bf640 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -30,6 +30,7 @@ _KV_CACHE_DTYPE_REASON = "kv_cache_dtype not supported" _TURBOQUANT_LAYOUT_REASON = "no KV cache layout in common with TURBOQUANT" +_ULTRAQUANT_LAYOUT_REASON = "no KV cache layout in common with ULTRAQUANT" try: from amdsmi import ( @@ -505,6 +506,7 @@ def _get_backend_priorities( backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) backends.append(AttentionBackendEnum.TRITON_ATTN) backends.append(AttentionBackendEnum.TURBOQUANT) + backends.append(AttentionBackendEnum.ULTRAQUANT) return backends @@ -517,6 +519,14 @@ def _uses_turboquant(vllm_config: "VllmConfig | None") -> bool: ) +def _uses_ultraquant(vllm_config: "VllmConfig | None") -> bool: + """Whether the run's KV cache dtype is the UltraQuant 4-bit format.""" + cache_config = getattr(vllm_config, "cache_config", None) + return ( + cache_config is not None and str(cache_config.cache_dtype) == "ultraquant_4bit" + ) + + def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> bool: """Whether this backend reads a KV cache layout TURBOQUANT reads too. @@ -535,12 +545,29 @@ def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> b return not set(layouts).isdisjoint(turboquant_layouts) +def _shares_layout_with_ultraquant(backend_class: type["AttentionBackend"]) -> bool: + """Whether this backend reads a KV cache layout ULTRAQUANT reads too. + + Same constraint as _shares_layout_with_turboquant: an ultraquant_4bit run + keeps its boundary layers at the native dtype, and one layout has to serve + the whole worker. + """ + layouts = backend_class.supported_kv_cache_layouts() + ultraquant_layouts = ( + AttentionBackendEnum.ULTRAQUANT.get_class().supported_kv_cache_layouts() + ) + if layouts is None or ultraquant_layouts is None: + return True + return not set(layouts).isdisjoint(ultraquant_layouts) + + def _get_invalid_reasons( backend_class: type["AttentionBackend"], device_capability: DeviceCapability, attn_selector_config: "AttentionSelectorConfig", *, is_turboquant_run: bool, + is_ultraquant_run: bool = False, ) -> list[str]: """Why this backend cannot serve the layer, empty when it can.""" invalid_reasons = backend_class.validate_configuration( @@ -553,6 +580,12 @@ def _get_invalid_reasons( and not _shares_layout_with_turboquant(backend_class) ): invalid_reasons = [_TURBOQUANT_LAYOUT_REASON] + if ( + not invalid_reasons + and is_ultraquant_run + and not _shares_layout_with_ultraquant(backend_class) + ): + invalid_reasons = [_ULTRAQUANT_LAYOUT_REASON] return invalid_reasons @@ -683,7 +716,7 @@ def get_valid_backends( @classmethod def get_attn_backend_cls( cls, - selected_backend: "AttentionBackendEnum", + selected_backend: "AttentionBackendEnum | None", attn_selector_config: "AttentionSelectorConfig", num_heads: int | None = None, ) -> str: @@ -695,13 +728,16 @@ def get_attn_backend_cls( # Keep lazy: vllm.config imports current_platform during initialization. from vllm.config import get_current_vllm_config_or_none - is_turboquant_run = _uses_turboquant(get_current_vllm_config_or_none()) + vllm_config = get_current_vllm_config_or_none() + is_turboquant_run = _uses_turboquant(vllm_config) + is_ultraquant_run = _uses_ultraquant(vllm_config) try: sel_invalid_reasons = _get_invalid_reasons( selected_backend.get_class(), device_capability, attn_selector_config, is_turboquant_run=is_turboquant_run, + is_ultraquant_run=is_ultraquant_run, ) except ImportError: sel_invalid_reasons = ["ImportError"] @@ -711,21 +747,23 @@ def get_attn_backend_cls( selected_backend.name, ) return selected_backend.get_path() - # Only tolerate the mismatch when turboquant is in play: boundary - # layers keep the native dtype while every other layer needs - # TURBOQUANT, so no single --attention-backend can serve every layer. - # For any other dtype the selection is genuinely invalid -> fail loud. + # Only tolerate the mismatch when turboquant or ultraquant is in play: + # boundary layers keep the native dtype while every other layer needs + # the packed backend, so no single --attention-backend can serve every + # layer. For any other dtype the selection is genuinely invalid -> fail + # loud. kv_dtype = attn_selector_config.kv_cache_dtype layer_is_turboquant = kv_dtype is not None and str(kv_dtype).startswith( "turboquant" ) - is_turboquant_fallback = ( - is_turboquant_run or layer_is_turboquant + is_packed_kv_fallback = ( + is_turboquant_run or layer_is_turboquant or is_ultraquant_run ) and sel_invalid_reasons in ( [_KV_CACHE_DTYPE_REASON], [_TURBOQUANT_LAYOUT_REASON], + [_ULTRAQUANT_LAYOUT_REASON], ) - if not is_turboquant_fallback: + if not is_packed_kv_fallback: raise ValueError( f"Selected backend {selected_backend} is not valid for " f"this configuration. Reason: {sel_invalid_reasons}" @@ -733,10 +771,11 @@ def get_attn_backend_cls( # NOTE: pass a str (not the list) -- info_once hashes its args. logger.info_once( "Selected backend %s is incompatible with this layer (%s) of " - "the turboquant run; using the auto-selected per-layer backend. " + "the %s run; using the auto-selected per-layer backend. " "Reason: %s", selected_backend.name, attn_selector_config.attn_type, + "ultraquant" if is_ultraquant_run else "turboquant", str(sel_invalid_reasons), ) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 025941eeb678..dce5e84fce25 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -50,6 +50,7 @@ "turboquant_4bit_nc": torch.uint8, "turboquant_k3v4_nc": torch.uint8, "turboquant_3bit_nc": torch.uint8, + "ultraquant_4bit": torch.uint8, "nvfp4": torch.uint8, "nvfp4_4over6": torch.uint8, } diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index f28a331854b6..ad184f091329 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -152,6 +152,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend" AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend" TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" + ULTRAQUANT = "vllm.v1.attention.backends.ultraquant_attn.UltraQuantAttentionBackend" # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None diff --git a/vllm/v1/attention/backends/ultraquant_attn.py b/vllm/v1/attention/backends/ultraquant_attn.py new file mode 100644 index 000000000000..b62b7a683023 --- /dev/null +++ b/vllm/v1/attention/backends/ultraquant_attn.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant 4-bit KV-cache attention backend. + +UltraQuant is a 4-bit KV-cache format (FP4 E2M1 codes + UE8M0 power-of-two +per-group-32 scales) with Hadamard-rotated keys. It reuses the TurboQuant +attention scaffolding (forward orchestration, metadata, cache-layout addressing, +CUDA-graph pre-warm) by subclassing, and overrides only the format-specific +seams: + + * store -> ultraquant_store (rotate K, pack FP4 + UE8M0 scales) + * decode -> FlyDSL D=256 scaled-MFMA kernel, Triton unified fallback + * continuation -> small chunk: Triton unified; large chunk: dequant + FA + +``UltraQuantAttentionBackend`` is a standalone backend with its own identity +(``ULTRAQUANT``) and dtype (``ultraquant_4bit``); the TurboQuant backend/impl +carry no UltraQuant knowledge. +""" + +import contextlib +import math +from dataclasses import replace +from typing import Any, ClassVar + +import torch +import torch.nn.functional as F + +from vllm.config import get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.turboquant_attn import ( + _CONTINUATION_DECODE_THRESHOLD, + _HAS_FLASH_ATTN, + TurboQuantAttentionBackend, + TurboQuantAttentionImpl, + TurboQuantMetadata, + _build_hadamard, +) +from vllm.v1.attention.ops.flydsl_ultraquant_decode import ( + flydsl_ultraquant_decode_attention, + is_ultraquant_flydsl_available, + ultraquant_flydsl_decode_eligible, +) +from vllm.v1.attention.ops.ultraquant.format import slot_size +from vllm.v1.attention.ops.ultraquant.triton_unified_attention import ( + ultraquant_unified_attention, +) +from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) + +logger = init_logger(__name__) + +# CK FlashAttention on ROCm has no compiled kernel above head_dim 256 and no +# sinks argument; those cases route through the Triton unified kernel instead. +_CK_MAX_HEAD_DIM = 256 + + +class UltraQuantAttentionBackend(TurboQuantAttentionBackend): + """Standalone backend for the ``ultraquant_4bit`` KV-cache dtype. + + Subclasses :class:`TurboQuantAttentionBackend` for the shared metadata and + cache-layout plumbing, and overrides identity, spec, impl, block size, and + capabilities so no UltraQuant knowledge lives in the TurboQuant backend. + """ + + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["ultraquant_4bit"] + # FlyDSL D=256 decode loads query as bf16; fp16 is not supported. Selection + # reads this from the backend class, not from the impl. + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """UltraQuant packs K+V into one 4-bit FP4/UE8M0 slot per head.""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_ultraquant: + return spec + return replace(spec, state_content_bytes=slot_size(spec.head_size)) + + @staticmethod + def get_name() -> str: + return "ULTRAQUANT" + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + # FlyDSL D=256 decode geometry requires a 64-token KV block; pin it so + # the KV-cache manager pages stay dense and splittable into kernel blocks. + return 64 + + @staticmethod + def get_impl_cls() -> type["UltraQuantAttentionImpl"]: + return UltraQuantAttentionImpl + + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: + return kv_cache_dtype == "ultraquant_4bit" + + @classmethod + def supports_sink(cls) -> bool: + # Decode and continuation prefill route sinks to the Triton unified + # kernel, but first-chunk prefill goes through flash_attn_varlen, which + # takes no sinks argument. Until that path is covered, reject sink + # models at selection rather than returning silently wrong output. + return False + + +class UltraQuantAttentionImpl(TurboQuantAttentionImpl): + """UltraQuant 4-bit KV-cache attention (FP4 E2M1 + UE8M0, rotated K). + + Overrides the init, on-device setup, store, decode, and prefill/continuation + seams of :class:`TurboQuantAttentionImpl`; the shared ``forward`` / + ``do_kv_cache_update`` orchestration and metadata handling are inherited. + """ + + # FlyDSL D=256 decode loads query as bf16; fp16 is not supported. + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + alibi_slopes: list[float] | None = None, + sliding_window: int | None = None, + kv_cache_dtype: str = "auto", + logits_soft_cap: float | None = None, + attn_type: str = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + **kwargs, + ): + self.num_heads = num_heads + self.head_size = head_size + self.scale = scale + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.num_kv_groups = num_heads // self.num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + + self.fa_version = get_flash_attn_version(head_size=head_size) + + vllm_config = get_current_vllm_config() + self.max_num_kv_splits = ( + vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph + ) + self.sliding_window = sliding_window + self.sinks = kwargs.get("sinks") + self._max_model_len = vllm_config.model_config.max_model_len + + self._use_ultraquant_flydsl = is_ultraquant_flydsl_available() + # FlyDSL decode eligibility is static per layer config; compute once. + self._flydsl_decode_eligible = ultraquant_flydsl_decode_eligible( + head_size=self.head_size, + num_kv_groups=self.num_kv_groups, + has_sinks=self.sinks is not None, + sliding_window=self.sliding_window, + flydsl_loaded=self._use_ultraquant_flydsl, + ) + + def _ensure_on_device(self, layer, device): + """Build the shared Hadamard rotation and pre-warm FlyDSL buffers. + + UltraQuant needs no Lloyd-Max centroids, but the inherited forward reads + ``layer._tq_centroids``, so an empty tensor is kept on the layer. + """ + if self._use_ultraquant_flydsl: + # CUDA-graph capture safety for the FlyDSL decode path on ROCm: + # pre-allocate index buffers and grow the workspace to its max size + # before capture so replay never lands on stale addresses. + _max_len = self._max_model_len + _already_ok = ( + hasattr(self, "_arange_cache") + and self._arange_cache.device.type == str(device).split(":")[0] + and self._arange_cache.shape[0] >= _max_len + 2 + ) + if not _already_ok: + self._arange_cache = torch.arange( + 0, _max_len + 2, device=device, dtype=torch.int32 + ) + if not hasattr(self, "_cu_2") or self._cu_2.device != torch.device(device): + self._cu_2 = torch.zeros(2, device=device, dtype=torch.int32) + if ( + is_workspace_manager_initialized() + and not current_workspace_manager().is_locked() + ): + B_max = self._max_capture_batch_size() + D = self.head_size + Hq = self.num_heads + S = self.max_num_kv_splits + _pre_warm_bytes = ( + B_max * Hq * (S * (D + 1) + D) * 4 # fp32 mid_o + fp32 lse + + B_max * Hq * D * 2 # query-dtype output (bf16 = 2 B) + + 512 # alignment padding + ) + with contextlib.suppress(AssertionError): + current_workspace_manager().get_simultaneous( + ((_pre_warm_bytes,), torch.uint8) + ) + + if not hasattr(layer, "_tq_cached"): + D = self.head_size + # Pure Hadamard: orthonormal + symmetric (H = H^T), enabling + # in-kernel butterfly fusion and trivial inverse for continuation. + H = _build_hadamard(D, str(device)) + layer._tq_PiT = H + layer._tq_Pi = H + layer._tq_Pi_half = H.to(torch.float16) + # Inherited TurboQuantAttentionImpl.forward reads _tq_centroids + # unconditionally; UltraQuant has no centroid table, so an empty + # tensor satisfies the read. + layer._tq_centroids = torch.empty(0, device=device, dtype=torch.float32) + layer._tq_cached = True + + def _store_kv( + self, + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + slot_mapping: torch.Tensor, + layer: Any, + ): + """Rotate K by the Hadamard and pack K+V into FP4/UE8M0 slots.""" + from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + + ultraquant_store( + key, + value, + kv_cache, + slot_mapping, + PiT=layer._tq_PiT, + ) + + def _decode_attention( + self, + query: torch.Tensor, # (B, Hq, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: torch.nn.Module | None = None, + ) -> torch.Tensor: + # Acquire shared decode scratch from the WorkspaceManager (UltraQuant + # only needs the output buffer; the mid_o/lse slots keep the triple + # allocation identical to the base decode for buffer reuse). + B = query.shape[0] + D = self.head_size + S = self.max_num_kv_splits + Hq = self.num_heads + output_buf = None + if is_workspace_manager_initialized(): + _mid_o, output_buf, _lse = current_workspace_manager().get_simultaneous( + ((B, Hq, S, D + 1), torch.float32), + ((B, Hq, D), query.dtype), + ((B, Hq), torch.float32), + ) + + if self._flydsl_decode_eligible: + return flydsl_ultraquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + scale=self.scale, + PiT=PiT, + max_seq_len=attn_metadata.max_seq_len, + output_buf=output_buf, + buf_holder=layer, + max_num_kv_splits=self.max_num_kv_splits, + sinks=self.sinks, + ) + + logger.info_once( + "UltraQuant FlyDSL is ineligible for head_size=%s, GQA=%s, " + "sinks=%s, sliding_window=%s; using Triton fallback.", + self.head_size, + self.num_kv_groups, + self.sinks is not None, + self.sliding_window, + ) + return ultraquant_unified_attention( + query=query, + kv_cache=kv_cache, + block_table=attn_metadata.block_table, + seq_lens=attn_metadata.seq_lens, + query_start_loc=attn_metadata.query_start_loc, + scale=self.scale, + PiT=PiT, + output=output_buf, + max_query_len=1, + max_seq_len=attn_metadata.max_seq_len, + sinks=self.sinks, + sliding_window=self.sliding_window, + ) + + def _prefill_attention( + self, + query: torch.Tensor, # (N, Hq, D) + key: torch.Tensor, # (N, Hk, D) + value: torch.Tensor, # (N, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + attn_metadata: TurboQuantMetadata, + Pi: torch.Tensor, + centroids: torch.Tensor, + PiT: torch.Tensor | None = None, + layer: Any = None, + ) -> torch.Tensor: + N, Hq, D = query.shape + + # Fast path: flash_attn for first-chunk prefills (all K/V in batch). + # CK cannot serve head dims above 256, so those fall through to the + # per-request unified-attention path below. + if ( + _HAS_FLASH_ATTN + and attn_metadata.max_query_len == attn_metadata.max_seq_len + and D <= _CK_MAX_HEAD_DIM + ): + return self._flash_attn_varlen( + q=query, + k=key, + v=value, + cu_seqlens_q=attn_metadata.query_start_loc, + cu_seqlens_k=attn_metadata.query_start_loc, + max_seqlen_q=attn_metadata.max_query_len, + max_seqlen_k=attn_metadata.max_query_len, + ) + + # Continuation or no flash_attn: per-request attention. + Hk = key.shape[1] + use_gqa = Hk < Hq + query_start_loc = attn_metadata.query_start_loc + num_reqs = query_start_loc.shape[0] - 1 + + output = torch.zeros(N, Hq, D, device=query.device, dtype=query.dtype) + + if attn_metadata.query_start_loc_cpu is not None: + qsl = attn_metadata.query_start_loc_cpu.tolist() + else: + qsl = query_start_loc.tolist() + if attn_metadata.seq_lens_cpu is not None: + seq_lens_list = attn_metadata.seq_lens_cpu.tolist() + else: + seq_lens_list = attn_metadata.seq_lens.tolist() + + if not hasattr(self, "_cu_2"): + self._cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32) + _max_seq = attn_metadata.max_seq_len + _ac: torch.Tensor | None = getattr(self, "_arange_cache", None) + if _ac is None or _ac.shape[0] <= _max_seq: + _ac = torch.arange( + 0, _max_seq + 1, device=query.device, dtype=attn_metadata.seq_lens.dtype + ) + self._arange_cache = _ac + _arange_cache: torch.Tensor = _ac + + for i in range(num_reqs): + q_start = qsl[i] + q_end = qsl[i + 1] + q_len = q_end - q_start + if q_len <= 0: + continue + + seq_len = seq_lens_list[i] + q_seq = query[q_start:q_end] # (q_len, Hq, D) + k_seq = key[q_start:q_end] # (q_len, Hk, D) + v_seq = value[q_start:q_end] # (q_len, Hk, D) + + if q_len == seq_len: + # First-chunk prefill: all K/V are in the current batch. + if D > _CK_MAX_HEAD_DIM: + # Head dims CK cannot handle (Gemma 4 global layers, + # D=512). This kernel is generic over head size. + from vllm.v1.attention.ops.triton_unified_attention import ( + unified_attention, + ) + + out = torch.empty_like(q_seq) + cu_single = torch.tensor( + [0, q_len], dtype=torch.int32, device=query.device + ) + seq_lens_single = torch.tensor( + [q_len], dtype=torch.int32, device=query.device + ) + block_table_single = torch.zeros( + (1, 1), dtype=torch.int32, device=query.device + ) + unified_attention( + q=q_seq, + k=k_seq.unsqueeze(0), + v=v_seq.unsqueeze(0), + out=out, + cu_seqlens_q=cu_single, + max_seqlen_q=q_len, + seqused_k=seq_lens_single, + max_seqlen_k=q_len, + softmax_scale=self.scale, + causal=True, + window_size=(-1, -1), + block_table=block_table_single, + softcap=0.0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=None, + ) + elif _HAS_FLASH_ATTN: + self._cu_2[1:2] = q_len + cu = self._cu_2 + out = self._flash_attn_varlen( + q=q_seq, + k=k_seq, + v=v_seq, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=q_len, + max_seqlen_k=q_len, + ) + else: + q_t = q_seq.transpose(0, 1).contiguous() + k_t = k_seq.transpose(0, 1).contiguous() + v_t = v_seq.transpose(0, 1).contiguous() + out = F.scaled_dot_product_attention( + q_t, + k_t, + v_t, + is_causal=True, + scale=self.scale, + enable_gqa=use_gqa, + ).transpose(0, 1) + output[q_start:q_end] = out.to(query.dtype) + else: + # Continuation chunk: prefix already stored (Hadamard-rotated + # FP4) by do_kv_cache_update. Small chunks reuse the Triton + # unified decode-shaped kernel; large chunks dequant the prefix + # once and run a dense flash-attention prefill. + cached_len = seq_len - q_len + if q_len <= _CONTINUATION_DECODE_THRESHOLD: + synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1] + synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1) + out = ultraquant_unified_attention( + query=q_seq, + kv_cache=kv_cache, + block_table=synth_bt, + seq_lens=synth_seq_lens, + query_start_loc=_arange_cache[: q_len + 1], + scale=self.scale, + PiT=PiT, + max_query_len=1, + max_seq_len=seq_len, + sinks=self.sinks, + sliding_window=self.sliding_window, + ) + else: + out = self._ultraquant_continuation_prefill( + layer=layer, + query=q_seq, + key_chunk=k_seq, + val_chunk=v_seq, + kv_cache=kv_cache, + block_table=attn_metadata.block_table[i : i + 1], + cached_len=cached_len, + seq_len=seq_len, + PiT=PiT, + ) + output[q_start:q_end] = out.to(query.dtype) + + return output + + def _ultraquant_continuation_prefill( + self, + *, + layer: Any, + query: torch.Tensor, # (q_len, Hq, D) unrotated + key_chunk: torch.Tensor, # (q_len, Hk, D) unrotated + val_chunk: torch.Tensor, # (q_len, Hk, D) + kv_cache: torch.Tensor, # (num_blocks, block_size, Hk, slot_size) + block_table: torch.Tensor, # (1, max_num_blocks) + cached_len: int, + seq_len: int, + PiT: torch.Tensor, # (D, D) fp32 Hadamard + ) -> torch.Tensor: + """Large continuation chunk: dequant the cached prefix, then flash_attn. + + The cache stores K already Hadamard-rotated, so only the current chunk's + K and the queries need rotating. Hadamard is symmetric orthonormal, so + rotating both sides leaves the attention scores unchanged. + """ + from vllm.v1.attention.ops.ultraquant.triton_dequant import ( + ultraquant_full_dequant_kv, + ) + + q_len, Hq, D = query.shape + Hk = key_chunk.shape[1] + device = query.device + qdtype = query.dtype + block_size = kv_cache.shape[1] + alloc_len = math.ceil(cached_len / block_size) * block_size + + # Dequant target, shared across layers by the workspace manager so long + # context does not pay one allocation per layer. + buf_shape = (1, Hk, alloc_len, D) + k_buf, v_buf = current_workspace_manager().get_simultaneous( + (buf_shape, qdtype), + (buf_shape, qdtype), + ) + ultraquant_full_dequant_kv( + kv_cache=kv_cache, + block_table=block_table, + k_out=k_buf, + v_out=v_buf, + alloc_len=alloc_len, + ) + + # Layer-cached prefix+chunk concatenation, grown once to worst case. + cap = block_table.shape[1] * block_size + k_full_buf = getattr(layer, "_uq_kfull_buf", None) + if ( + k_full_buf is None + or k_full_buf.shape[0] < cap + or k_full_buf.dtype != qdtype + ): + k_full_buf = torch.empty(cap, Hk, D, dtype=qdtype, device=device) + layer._uq_kfull_buf = k_full_buf + v_full_buf = getattr(layer, "_uq_vfull_buf", None) + if ( + v_full_buf is None + or v_full_buf.shape[0] < cap + or v_full_buf.dtype != qdtype + ): + v_full_buf = torch.empty(cap, Hk, D, dtype=qdtype, device=device) + layer._uq_vfull_buf = v_full_buf + k_full = k_full_buf[:seq_len] + v_full = v_full_buf[:seq_len] + + k_full[:cached_len] = k_buf[0, :, :cached_len, :].transpose(0, 1) + v_full[:cached_len] = v_buf[0, :, :cached_len, :].transpose(0, 1) + k_full[cached_len:] = (key_chunk.to(torch.float32) @ PiT).to(qdtype) + v_full[cached_len:] = val_chunk + q_rot = (query.to(torch.float32) @ PiT).to(qdtype) + + # Unequal q/k lengths with causal=True gives lower-right alignment: + # query i (absolute position cached_len + i) sees keys 0..cached_len+i. + # + # FA2 on ROCm has no sinks parameter, and CK cannot serve head dims + # above 256, so those cases route through the Triton unified kernel. + # Q/K are already in Hadamard-rotated space (q_rot, k_full) so the dot + # products match the original-space attention scores. + if self.sinks is not None or D > _CK_MAX_HEAD_DIM: + from vllm.v1.attention.ops.triton_unified_attention import ( + unified_attention, + ) + + out = torch.empty_like(q_rot) + cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=device) + seqused_k = torch.tensor([seq_len], dtype=torch.int32, device=device) + block_table_single = torch.zeros((1, 1), dtype=torch.int32, device=device) + window = ( + (self.sliding_window - 1, 0) + if self.sliding_window and self.sliding_window > 0 + else (-1, -1) + ) + unified_attention( + q=q_rot, + k=k_full.unsqueeze(0), + v=v_full.unsqueeze(0), + out=out, + cu_seqlens_q=cu_q, + max_seqlen_q=q_len, + seqused_k=seqused_k, + max_seqlen_k=seq_len, + softmax_scale=self.scale, + causal=True, + window_size=window, + block_table=block_table_single, + softcap=0.0, + q_descale=None, + k_descale=None, + v_descale=None, + sinks=self.sinks, + ) + return out.to(qdtype) + if _HAS_FLASH_ATTN: + cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=device) + cu_k = torch.tensor([0, seq_len], dtype=torch.int32, device=device) + return self._flash_attn_varlen( + q=q_rot, + k=k_full, + v=v_full, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=q_len, + max_seqlen_k=seq_len, + ) + q_t = q_rot.transpose(0, 1).unsqueeze(0) + k_t = k_full.transpose(0, 1).unsqueeze(0) + v_t = v_full.transpose(0, 1).unsqueeze(0) + if Hq > Hk: + k_t = k_t.expand(1, Hq, -1, -1) + v_t = v_t.expand(1, Hq, -1, -1) + return ( + F.scaled_dot_product_attention( + q_t, k_t, v_t, is_causal=True, scale=self.scale + ) + .squeeze(0) + .transpose(0, 1) + .to(qdtype) + ) diff --git a/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py b/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py new file mode 100644 index 000000000000..6a03f06f458f --- /dev/null +++ b/vllm/v1/attention/ops/flydsl_kernels/ultraquant_decode_hd256.py @@ -0,0 +1,983 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# UltraQuant D=256 decode (FlyDSL) — gfx950 / CDNA4 +# +# Production path only: scaled FP4×E4M3 QK MFMA, native V CVT, HW V-transpose, +# Q hoist, strided tile-groups, and in-kernel Walsh-Hadamard for GQA 6/8/16. +# +# AoS slot (D=256, group_size=32 → 8 groups, 272 B): +# [0:128) K FP4 codes; [128:136) K UE8M0; [136:264) V FP4; [264:272) V UE8M0 +# QK MFMA: A=K, B=Q → C[token, query]; PV MFMA: A=V_T, B=P → C[head_dim, query] + +from __future__ import annotations + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf as _scf +from flydsl.expr import ( + arith, + gpu, + rocdl, +) +from flydsl.expr.primitive import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + + +def _vector_insert(value, dest, *, static_position): + """Insert a DSL scalar using FlyDSL 0.3's raw vector dialect binding.""" + return vector.insert( + vector.as_ir_value(value), + vector.as_ir_value(dest), + dynamic_position=[], + static_position=static_position, + ) + + +# === Constants (256-wide-head ultraquant decode profile) ======================= +HEAD_SIZE = 256 +KV_BLOCK_SIZE = 16 # default; overridable via build(kv_block_size=...) +TILE_SIZE = 16 # MFMA tile = 16 tokens +N_CENTROIDS = 16 # FP4 E2M1 value table (16 entries indexed by nibble) +QUERY_GROUP_SIZE = 16 # default for tests; overridable via build(...) +WARP_SIZE = 64 +BLOCK_THREADS = WARP_SIZE +KV_COMPUTE_BLOCK = 256 # 16 K-tiles × 16 tokens + +# ultraquant AoS slot layout (per (slot, head), D=256, group_size=32 => 8 groups) +FP8_GROUP_SIZE = 32 +N_GROUPS = HEAD_SIZE // FP8_GROUP_SIZE # 8 UE8M0 scale bytes per head +KEY_CODE_BYTES = HEAD_SIZE // 2 # 128 +VAL_CODE_BYTES = HEAD_SIZE // 2 # 128 +K_SCALES_OFFSET = KEY_CODE_BYTES # 128 +V_CODES_OFFSET = KEY_CODE_BYTES + N_GROUPS # 136 +V_SCALES_OFFSET = V_CODES_OFFSET + VAL_CODE_BYTES # 264 +SLOT_CONTENT_BYTES = V_SCALES_OFFSET + N_GROUPS # 272 + +# MFMA +MFMA_N = 16 +PV_N_CHUNKS = HEAD_SIZE // MFMA_N # 16 for HEAD_SIZE=256 + +# --- Per-lane data-movement geometry (generalized from the 128 kernel) ------ +# The K/V dequant lane assignment splits each token's packed code bytes across +# the 4 ``chunk_in_tok`` sub-lanes (lane % 4). Each lane processes +# ``KEY_CODE_BYTES // 4`` code bytes = 32 B at 256. buffer_load maxes at 16 B +# (dwordx4) so a lane issues ``HALVES``=2 loads; each 16-B half covers exactly +# one UE8M0 group of 32 head-dims. +# HEAD_SIZE=128: LANE_CODE_BYTES=16, HALVES=1, SUBCHUNK_HDIMS=32 +# HEAD_SIZE=256: LANE_CODE_BYTES=32, HALVES=2, SUBCHUNK_HDIMS=64 +LANE_CODE_BYTES = KEY_CODE_BYTES // 4 # 32 code bytes per lane per K-tile +HALVES = LANE_CODE_BYTES // 16 # 2 for HEAD_SIZE=256 +HALF_HDIMS = 32 # head-dims per half (= one group) +SUBCHUNK_HDIMS = HEAD_SIZE // 4 # head-dims per chunk_in_tok sub-lane (64) +# Q row-major LDS load geometry (STEP B): each lane loads 8 bf16 (=4 i32), +# so a row spans ``HEAD_SIZE // 8`` column-groups. +Q_GROUPS_PER_ROW = HEAD_SIZE // 8 # 32 for 256 +Q_ROW_SHIFT = Q_GROUPS_PER_ROW.bit_length() - 1 # 5 for 256 +Q_COL_MASK = Q_GROUPS_PER_ROW - 1 # 31 for 256 + +# LDS regions +CENTROID_LDS_BYTES = N_CENTROIDS * 4 # 64 + +# KV LDS rows are padded by 16 B to reduce bank conflicts while preserving +# ds_read/ds_write_b128 alignment. +KV_ROW_PAD_ELEMS = 8 +assert KV_ROW_PAD_ELEMS % 8 == 0, "KV row pad must be a multiple of 8 bf16 (16 B)" +KV_ROW_ELEMS = HEAD_SIZE + KV_ROW_PAD_ELEMS # 264 @ pad 8 +KV_ROW_BYTES = KV_ROW_ELEMS * 2 # 528 +KFP4_ROW_PAD_I32 = 4 +assert KFP4_ROW_PAD_I32 % 4 == 0, "FP4-K row pad must be a multiple of 4 i32 (16 B)" +KFP4_ROW_I32 = HEAD_SIZE // 8 + KFP4_ROW_PAD_I32 +KV_TILE_LDS_BYTES_PADDED = TILE_SIZE * KV_ROW_BYTES # 8448 (V dominates) +SCALE_LDS_BYTES = TILE_SIZE * N_GROUPS * 4 # 16*8*4 = 512 for 256 + +LOG2E = 1.4426950408889634 +NEG_INF_VAL = float("-inf") + + +def _vsplat_mul(vec, scalar): + s = scalar.ir_value() if hasattr(scalar, "ir_value") else scalar + return vec * vector.broadcast(T.f32x4, s) + + +allocator = None + + +def build_ultraquant_decode_hd256_module( + num_seqs: int, + num_kv_heads: int, + num_partitions: int, + padded_slot: int, + max_blocks_per_seq: int = 512, + softmax_scale: float | None = None, + query_group_size: int = QUERY_GROUP_SIZE, + kv_block_size: int = KV_BLOCK_SIZE, + tile_groups_per_partition: int = 1, + fuse_qrot: bool = True, + stride_q_seq: int | None = None, + stride_q_head: int | None = None, +): + """Build the production UltraQuant D=256 decode kernel. + + ``padded_slot`` is the per-(token, head) slot stride in bytes + (``kv_cache.shape[3]``, >= 272). Cache layout is + ``[num_blocks, block_size, num_kv_heads, padded_slot]`` uint8 AoS. + + QK is always scaled FP4×E4M3 MFMA; V always uses native CVT + HW + transpose; tile-groups are always strided. ``fuse_qrot`` selects the + in-kernel WHT when the query layout is eligible. + """ + # QG=6 (Qwen3.6 full-attention layers) reuses the 8-row WHT lane map with + # two inactive rows masked. STEP C may read rows QG..15 from adjacent LDS, + # but those MFMA rows are discarded by the `mfma_row < QG` output gate. + assert query_group_size in (6, 8, 16), ( + f"query_group_size must be 6, 8 or 16; got {query_group_size}" + ) + # head_dim=256 hybrid models (Qwen3.6/3.8) force a larger mamba-aligned KV + # block; the tiled inner loop (_TILES_PER_BLOCK = kv_block_size // TILE_SIZE) + # handles any multiple of TILE_SIZE, so accept 16/32/64/128/256. + assert kv_block_size in (16, 32, 64, 128, 256), ( + f"kv_block_size must be 16/32/64/128/256; got {kv_block_size}" + ) + assert kv_block_size % TILE_SIZE == 0 + assert int(tile_groups_per_partition) >= 1 + + TGPP = int(tile_groups_per_partition) + # Partition p owns tile-groups p, p+P, p+2P, ... This keeps the graph-fixed + # grid wide for short sequences; the reducer is agnostic to the mapping. + NUM_PARTS_C = int(num_partitions) + MFMA_SCALED_K = 128 + MFMA_ISSUES = HEAD_SIZE // MFMA_SCALED_K # 2 for HEAD_SIZE=256 + GRPS_PER_ISSUE = MFMA_SCALED_K // FP8_GROUP_SIZE # 4 + + QG = int(query_group_size) + QG_LOAD_ITERS = (QG * Q_GROUPS_PER_ROW) // WARP_SIZE # 8@QG16, 4@QG8, 3@QG6 + + _BLOCK_THREADS = WARP_SIZE + + # Fuse Q rotation (q_rot = Q @ PiT) into this kernel as an in-register WHT. + # The 8-row lane map masks its two inactive rows for GQA-6. + _FUSE_QROT = bool(fuse_qrot) + + # --- Cache policy on the KV code stream ---------------------------------- + # KV codes are read EXACTLY ONCE per decode (81.6 MB at B=4/seq=75k) and + # never reused, so caching them only evicts Q / block-table / scale lines + # that ARE reused. gfx9 buffer aux bits: bit0=sc0, bit1=nt, bit4=sc1. + # 0 = default (what we ship today: no hint at all) + # 2 = nt (non-temporal: stream, do not pollute) + # 18 = nt|sc1 (also bypass the far cache level) + _KV_AUX = 0 + + # --- Occupancy: size the Q LDS to the ACTUAL query group ------------------- + # STEP B only writes rows [0, QG) and STEP C/F only read those rows. For + # QG=8 the upper 4 KB is allocated-but-never-touched, so shrinking it to + # QG*HEAD_SIZE*2 is bitwise identical and frees LDS/workgroup (17 KB -> 13 KB + # at QG=8), lifting waves/CU (~3 -> ~5) since the kernel is LDS-occupancy + # bound. The launcher pads segm pools to QG=16 regardless, so this is purely + # an in-kernel LDS footprint reduction. + _Q_LDS_BYTES = QG * HEAD_SIZE * 2 + + _BS = int(kv_block_size) + _TILES_PER_BLOCK = _BS // TILE_SIZE + _PADDED_SLOT = int(padded_slot) + assert _PADDED_SLOT >= SLOT_CONTENT_BYTES, ( + f"padded_slot={_PADDED_SLOT} < {SLOT_CONTENT_BYTES} (ultraquant D=256 content)" + ) + assert _PADDED_SLOT % 4 == 0, f"padded_slot={_PADDED_SLOT} must be 4-byte aligned" + + arch = get_hip_arch() + + if softmax_scale is None: + softmax_scale = 1.0 / (HEAD_SIZE**0.5) + _qk_scale = float(softmax_scale) + + # --- Strides --- + _Hq = num_kv_heads * QG + # Defaults describe the pooled [B, Hq, D] q_rot buffer. When the rotation is + # fused in (STEP B'), Q is read straight from the model's query tensor, + # which is a SPLIT OF THE FUSED QKV PROJECTION -- so its row stride is + # (Hq + 2*Hk)*D, not Hq*D. Those strides are constant for a given model and + # TP rank, so the caller passes them in and they bake in like any other + # compile-time constant (they are part of the kernel cache key). + _stride_q_seq = int(stride_q_seq) if stride_q_seq is not None else _Hq * HEAD_SIZE + _stride_q_head = int(stride_q_head) if stride_q_head is not None else HEAD_SIZE + _stride_bt_seq = max_blocks_per_seq + + # ultraquant AoS cache: [num_blocks, block_size, num_kv_heads, padded_slot]. + _stride_cache_head = _PADDED_SLOT + _stride_cache_pos = num_kv_heads * _PADDED_SLOT + _stride_cache_block = _BS * _stride_cache_pos + + _stride_out_part = QG * HEAD_SIZE + _stride_out_head = num_partitions * QG * HEAD_SIZE + _stride_out_seq = num_kv_heads * num_partitions * QG * HEAD_SIZE + _stride_es_seq = num_kv_heads * num_partitions * QG + _stride_ml_seq = _stride_es_seq + + # --- LDS layout --- + # `allocator` is intentionally a BUILD-LOCAL captured by the kernel closure + # (a freevar), NOT a module global. The @flyc.jit drift guard + # (_check_globals_drift) snapshots every module global referenced by the JIT + # launcher's dependency tree (which includes this kernel) and ABORTS if one + # changes between compiles. When more than one partition variant is + # JIT-resident at once (e.g. a ragged/adaptive batch mixes num_partitions + # 256 and 512, or fixed cap=512 yields 294@75k and 512@131k), each build + # created a fresh SmemAllocator; as a reassigned MODULE GLOBAL that tripped + # "global 'allocator' changed since first compile" and crashed the worker. + # Keeping it local + a per-variant symbol name lets the variants coexist. + # The build-local is handed to the launcher via a module attribute set AFTER + # the kernel is built (see end of function); that assignment is safe because + # no @flyc.jit function references the module global anymore. + allocator = SmemAllocator( + None, + arch=arch, + global_sym_name=(f"ultraquant_hd256_smem_p{int(num_partitions)}"), + ) + centroid_off = 0 + allocator.ptr = CENTROID_LDS_BYTES + q_off = allocator.ptr + allocator.ptr += _Q_LDS_BYTES + kv_off = allocator.ptr + allocator.ptr += KV_TILE_LDS_BYTES_PADDED + scale_off = allocator.ptr + allocator.ptr += SCALE_LDS_BYTES + + @flyc.kernel + def ultraquant_decode_hd256_kernel( + out_ptr: fx.Tensor, + exp_sums_ptr: fx.Tensor, + max_logits_ptr: fx.Tensor, + query_ptr: fx.Tensor, + kv_cache_ptr: fx.Tensor, + centroids_ptr: fx.Tensor, + block_tables_ptr: fx.Tensor, + seq_lens_ptr: fx.Tensor, + ): + # ---- IDs --------------------------------------------------------- + tid = gpu.thread_idx.x + seq = gpu.block_idx.x + kv_h = gpu.block_idx.y + part = gpu.block_idx.z + lane = tid # 0..63 + mfma_row = lane & fx.Int32(15) + mfma_col_grp = lane >> fx.Int32(4) # 0..3, K-group dim + + # ---- Buffer resources ------------------------------------------- + q_rsrc = buffer_ops.create_buffer_resource(query_ptr, max_size=True) + bt_rsrc = buffer_ops.create_buffer_resource(block_tables_ptr, max_size=True) + sl_rsrc = buffer_ops.create_buffer_resource(seq_lens_ptr, max_size=True) + cent_rsrc = buffer_ops.create_buffer_resource(centroids_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True) + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True) + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True) + + # ---- LDS pointers ----------------------------------------------- + base = allocator.get_base() + cent_lds = SmemPtr(base, centroid_off, T.f32, shape=(N_CENTROIDS,)) + q_lds_i32 = SmemPtr(base, q_off, T.i32, shape=(_Q_LDS_BYTES // 4,)).get() + q_lds_i64 = SmemPtr(base, q_off, T.i64, shape=(_Q_LDS_BYTES // 8,)).get() + # KV / scale views span ALL warps' private regions (W*stride); the + # per-warp base offset is folded into the element indices below via + # _kvi*/_sci (a literal 0 for W=1, so the single-warp path is unchanged). + kv_lds_i32 = SmemPtr( + base, kv_off, T.i32, shape=(KV_TILE_LDS_BYTES_PADDED // 4,) + ).get() + kv_lds_i64 = SmemPtr( + base, kv_off, T.i64, shape=(KV_TILE_LDS_BYTES_PADDED // 8,) + ).get() + scale_lds_i32 = SmemPtr(base, scale_off, T.i32, shape=(TILE_SIZE * N_GROUPS,)) + + def _kvi32(e): + return e + + def _kvi64(e): + return e + + def _sci(e): + return e + + # ---- Constants --------------------------------------------------- + c_sq = fx.Int32(_stride_q_seq) + c_qh = fx.Int32(_stride_q_head) + c_qg = fx.Int32(QG) + c_bt = fx.Int32(_stride_bt_seq) + c_stride_pos = fx.Int32(_stride_cache_pos) + c_stride_head = fx.Int32(_stride_cache_head) + c_kscale_off = fx.Int32(K_SCALES_OFFSET) # 128 + c_vcode_off = fx.Int32(V_CODES_OFFSET) # 136 + c_vscale_off = fx.Int32(V_SCALES_OFFSET) # 264 + c_w = fx.Int32(WARP_SIZE) + + NEG_INF = arith.constant(NEG_INF_VAL, type=T.f32) + ZERO_F = fx.Float32(0.0) + ONE_F = fx.Float32(1.0) + LOG2E_C = arith.constant(LOG2E, type=T.f32) + QK_SCALE = arith.constant(_qk_scale, type=T.f32) + + def _ival(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + c_zero_i32 = arith.constant(0, type=T.i32) + + def _f32x8_to_fp8_i64(f): + # 8 f32 -> 8 E4M3 bytes packed as i64 (MFMA fp8 operand). + w0 = rocdl.cvt_pk_fp8_f32(T.i32, f[0], f[1], c_zero_i32, 0) + w0 = rocdl.cvt_pk_fp8_f32(T.i32, f[2], f[3], w0, 1) + w1 = rocdl.cvt_pk_fp8_f32(T.i32, f[4], f[5], c_zero_i32, 0) + w1 = rocdl.cvt_pk_fp8_f32(T.i32, f[6], f[7], w1, 1) + pv = vector.from_elements(T.vec(2, T.i32), [w0, w1]) + return vector.extract( + vector.bitcast(T.vec(1, T.i64), pv), static_position=[0] + ) + + def _bf16x8_to_fp8_i64(v_bf16): + f = [ + arith.extf(T.f32, vector.extract(v_bf16, static_position=[i])) + for i in range(8) + ] + return _f32x8_to_fp8_i64(f) + + # ===== STEP A: Load centroids → LDS (cooperative, race-safe) ===== + c_idx_safe = lane & fx.Int32(N_CENTROIDS - 1) + c_val = buffer_ops.buffer_load(cent_rsrc, c_idx_safe, vec_width=1, dtype=T.f32) + cent_lds.store(c_val, [arith.index_cast(T.index, c_idx_safe)]) + gpu.barrier() + + # ===== STEP B': in-kernel Q rotation (fused WHT) ================== + # The shipped path needs a SEPARATE prologue kernel to compute + # q_rot = Q @ PiT. That kernel is ~7 us at B=4 -- essentially ALL fixed + # per-dispatch cost (measured: 2.1M MACs, and sweeping its grid 4x moves + # it 0%). PiT is the orthonormal Sylvester Hadamard with a column + # permutation folded in on the host, so Q @ PiT is a WALSH-HADAMARD + # TRANSFORM: 8 butterfly stages (2048 add/sub per row) instead of a + # 256x256 matmul (65536 MACs) -- 32x less work, cheap enough to redo + # redundantly in every partition and delete the dispatch entirely. + # + # Verified offline: butterfly vs matmul differs by <=1.8e-6 in fp32, and + # after the E4M3 haircut the two agree EXACTLY (0/2048 elements differ), + # because e4m3 keeps only 3 mantissa bits. So this is bit-exact, not an + # approximation. + # + # Lane map: lane L owns row r = L>>3 and the 32 CONSECUTIVE head-dims + # d = (L&7)*32 + t. That puts butterfly stages 0..4 (strides 1..16) + # entirely inside one lane's registers, and leaves only stages 5..7 + # (strides 32/64/128 == chunk XOR 1/2/4) needing a cross-lane exchange. + if const_expr(_FUSE_QROT): + _wrow = lane >> fx.Int32(3) # 0..7 row within the group + _wchk = lane & fx.Int32(7) # 0..7 32-dim chunk + # inv(qperm) is a pure bit permutation of the head-dim index: + # dest = (d & 0x8F) | ((d>>2)&0x10) | ((d<<1)&0x60) + # which splits additively into a compile-time part in t and a + # per-lane part in the chunk index (verified exhaustively): + # ct(t) = (t & 0x0F) | ((t & 0x10) << 1) + # rt(c) = ((c&2)<<3) | ((c&1)<<6) | ((c&4)<<5) + _perm_base = ( + ((_wchk & fx.Int32(2)) << fx.Int32(3)) + | ((_wchk & fx.Int32(1)) << fx.Int32(6)) + | ((_wchk & fx.Int32(4)) << fx.Int32(5)) + ) + _inv_sqrtD = arith.constant(1.0 / math.sqrt(HEAD_SIZE), type=T.f32) + for _qit in range_constexpr((QG + 7) // 8): + _qrow = _wrow + fx.Int32(_qit * 8) + if _qrow < fx.Int32(QG): + _qb = ( + seq * c_sq + (kv_h * c_qg + _qrow) * c_qh + _wchk * fx.Int32(32) + ) + # ---- load this lane's 32 raw-Q head-dims (4 x dwordx4) ---- + _v = [] + for _j in range_constexpr(4): + _raw = buffer_ops.buffer_load( + q_rsrc, + (_qb + fx.Int32(_j * 8)) // fx.Int32(2), + vec_width=4, + dtype=T.i32, + ) + _rb = vector.bitcast(T.vec(8, T.bf16), _raw) + for _e in range_constexpr(8): + _v.append( + arith.extf( + T.f32, vector.extract(_rb, static_position=[_e]) + ) + ) + # ---- stages 0..4: in-register, no LDS, no cross-lane ---- + for _s in range_constexpr(5): + _st = 1 << _s + _nx = list(_v) + for _g in range_constexpr(32 // (2 * _st)): + for _k in range_constexpr(_st): + _i0 = _g * 2 * _st + _k + _i1 = _i0 + _st + _a, _b = _v[_i0], _v[_i1] + _nx[_i0] = _a + _b + _nx[_i1] = _a - _b + _v = _nx + # ---- stages 5..7: chunk XOR 1/2/4 via ds_swizzle ---- + # ds_swizzle_b32 offset (bit15=0, 32-lane groups): + # [4:0]=and_mask [9:5]=or_mask [14:10]=xor_mask + # pure XOR k -> and_mask=0x1F, or_mask=0, xor_mask=k. The + # partner is always within the same 32-lane group (k<=4). + for _s in range_constexpr(3): + _k = 1 << _s + _swz = arith.constant((_k << 10) | 0x1F, type=T.i32) + # High half of the pair computes (partner - mine), the low + # half (mine + partner). Fold that into a single sign + # computed ONCE per stage, so each element costs one + # multiply-add rather than a branch or a per-element select. + _chk_bit = _wchk & fx.Int32(_k) + _is_hi = arith.cmpi( + arith.CmpIPredicate.ne, + _chk_bit.ir_value() + if hasattr(_chk_bit, "ir_value") + else _chk_bit, + fx.Int32(0).ir_value(), + ) + _sf = arith.select( + _is_hi, + arith.constant(-1.0, type=T.f32), + arith.constant(1.0, type=T.f32), + ) + _nx = [] + for _t in range_constexpr(32): + _mi = arith.bitcast(T.i32, _v[_t]) + _pt = arith.bitcast( + T.f32, rocdl.ds_swizzle(T.i32, _mi, _swz) + ) + _nx.append(_v[_t] * _sf + _pt) + _v = _nx + # ---- 1/sqrt(D) normalization + E4M3 haircut -> bf16 ---- + # cvt_pk_fp8_f32 is the SAME instruction the QK operand build + # uses, so this reproduces the host haircut exactly. E4M3 is a + # subset of bf16, so the bf16 store below is lossless and every + # downstream consumer stays byte-identical. + _hc = [] + for _p in range_constexpr(16): + _w = rocdl.cvt_pk_fp8_f32( + T.i32, + _v[2 * _p] * _inv_sqrtD, + _v[2 * _p + 1] * _inv_sqrtD, + fx.Int32(0), + 0, + ) + _u = rocdl.cvt_pk_f32_fp8(T.vec(2, T.f32), _w, 0) + _hc.append( + arith.truncf( + T.bf16, vector.extract(_u, static_position=[0]) + ) + ) + _hc.append( + arith.truncf( + T.bf16, vector.extract(_u, static_position=[1]) + ) + ) + # ---- permuted store ---- + # ct(t) maps t=0..15 -> 0..15 and t=16..31 -> 32..47, i.e. TWO + # contiguous 16-element runs, so the scatter is still 4 aligned + # 16-byte vector stores rather than 32 scalar writes. + _qrow_elem = _qrow * fx.Int32(HEAD_SIZE) + for _half in range_constexpr(2): + for _q4 in range_constexpr(2): + _elems = [_hc[_half * 16 + _q4 * 8 + _e] for _e in range(8)] + _vec = vector.from_elements(T.vec(8, T.bf16), _elems) + _dst = ( + _qrow_elem + _perm_base + fx.Int32(_half * 32 + _q4 * 8) + ) + vector.store( + vector.bitcast(T.vec(4, T.i32), _vec), + q_lds_i32, + [arith.index_cast(T.index, _dst // fx.Int32(2))], + ) + gpu.barrier() + + # ===== STEP B: Load pre-rotated q_rot → row-major Q LDS ========== + for c in range_constexpr(0 if _FUSE_QROT else QG_LOAD_ITERS): + row_chunk = lane + fx.Int32(c * WARP_SIZE) + row = row_chunk >> fx.Int32(Q_ROW_SHIFT) # 0..QG-1 + col_b = row_chunk & fx.Int32(Q_COL_MASK) # 0..Q_GROUPS_PER_ROW-1 + col_elem = col_b * fx.Int32(8) # bf16 elem + q_off_elem = seq * c_sq + (kv_h * c_qg + row) * c_qh + col_elem + q_v = buffer_ops.buffer_load( + q_rsrc, + q_off_elem // fx.Int32(2), + vec_width=4, + dtype=T.i32, + ) + q_lds_byte = row * fx.Int32(HEAD_SIZE * 2) + col_elem * fx.Int32(2) + vector.store( + q_v, + q_lds_i32, + [arith.index_cast(T.index, q_lds_byte // fx.Int32(4))], + ) + gpu.barrier() + + # Scaled-MFMA Q operands (loop-invariant): one per MFMA issue. + q_op_hoisted = [] + for h in range_constexpr(MFMA_ISSUES): + _q_hoist_words = [] + for j in range_constexpr(4): + _q_hoist_idx = ( + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8) + + (fx.Int32(h * GRPS_PER_ISSUE) + mfma_col_grp) * fx.Int32(8) + + fx.Int32(j * 2) + ) + _q_hoist_v = vector.load_op( + T.vec(2, T.i64), + q_lds_i64, + [arith.index_cast(T.index, _q_hoist_idx)], + ) + _q_hoist_words.append( + _bf16x8_to_fp8_i64(vector.bitcast(T.vec(8, T.bf16), _q_hoist_v)) + ) + q_op_hoisted.append( + vector.bitcast( + T.vec(8, T.i32), + vector.from_elements(T.vec(4, T.i64), _q_hoist_words), + ) + ) + + # ===== STEP D: Online softmax + PV state ========================= + running_max = NEG_INF + running_sum = ZERO_F + zero_v4 = arith.constant_vector(0.0, T.f32x4) + acc_pv = [zero_v4 for _ in range(PV_N_CHUNKS)] + + # ===== STEP E: Sequence-len + partition base ==================== + seq_len = buffer_ops.buffer_load(sl_rsrc, seq, vec_width=1, dtype=T.i32) + + # Per-K-tile dequant lane assignment: lane t -> token = t/4, + # chunk_in_tok = t%4 (each chunk = SUBCHUNK_HDIMS=64 head-dims = 2 groups). + tok_in_tile = lane >> fx.Int32(2) + chunk_in_tok = lane & fx.Int32(3) + # Both UE8M0 group bytes for this lane's chunk (groups 2c, 2c+1) live in + # the SAME 4-byte scale word at offset (chunk_in_tok>>1)*4 within the + # 8-byte scale region. Load once; extract per-half below. + c_scale_word_off = (chunk_in_tok >> fx.Int32(1)) * fx.Int32(4) + + # ===== STEP F: K-tile loop ======================================= + c_kcb = fx.Int32(KV_COMPUTE_BLOCK) + c_tgpp = fx.Int32(TGPP) + c_zero_i32 = fx.Int32(0) + c_one_i32 = fx.Int32(1) + c_nparts = fx.Int32(NUM_PARTS_C) + total_tgs = (seq_len + c_kcb - c_one_i32) // c_kcb + avail = total_tgs - part + in_range = avail > c_zero_i32 + trip_raw = (avail + c_nparts - c_one_i32) // c_nparts + trip_clamped = (trip_raw > c_tgpp).select(c_tgpp, trip_raw) + trip_or_zero = in_range.select(trip_clamped, c_zero_i32) + c_zero_idx = arith.constant(0, index=True) + c_one_idx = arith.constant(1, index=True) + trip_idx = arith.index_cast( + T.index, + trip_or_zero.ir_value() + if hasattr(trip_or_zero, "ir_value") + else trip_or_zero, + ) + _init_iter = [ + _ival(running_max), + _ival(running_sum), + *[_ival(p) for p in acc_pv], + ] + _for_op = _scf.ForOp( + c_zero_idx, + trip_idx, + c_one_idx, + _init_iter, + ) + _for_ip = ir.InsertionPoint(_for_op.body) + _for_ip.__enter__() + try: + tg_idx = _for_op.induction_variable + tg_i32 = fx.Int32(arith.index_cast(T.i32, tg_idx)) + partition_start = (part + tg_i32 * c_nparts) * c_kcb + bt_seq_base = seq * c_bt + (partition_start // fx.Int32(_BS)) + running_max = _for_op.inner_iter_args[0] + running_sum = _for_op.inner_iter_args[1] + acc_pv = list(_for_op.inner_iter_args[2:]) + + # ==== SOFTWARE PREFETCH ===================================== + # Emit ONLY the HBM buffer_loads for one tile (K codes, V codes, + # K/V UE8M0 scale words) and return the loaded register values. + # Emitting tile N+1's loads BEFORE tile N's dequant+MFMA lets the + # vmem fetches overlap compute: the compiler's s_waitcnt for these + # values lands at their first use one iteration later, hiding HBM + # latency (the decode kernel is ~85% HBM-stalled at low batch). + # This is bit-identical to the pre-prefetch path -- only the + # *emission order* of independent loads changes, never any + # arithmetic. All addressing locals (slot_base_byte, blk_rsrc, + # ...) are consumed only inside this closure; the dequant below + # needs just the four returned register values. + def _emit_tile_loads(n_tile): + # n_tile is a compile-time int on the single-warp path and a + # runtime fx.Int32 on the multi-warp path (each warp strides the + # 16 tiles). Compute the block/tile/slot geometry accordingly. + if const_expr(isinstance(n_tile, int)): + block_in_part = fx.Int32(n_tile // _TILES_PER_BLOCK) + tile_in_block = n_tile % _TILES_PER_BLOCK + tile_start_tok = partition_start + fx.Int32(n_tile * TILE_SIZE) + slot = fx.Int32(tile_in_block * TILE_SIZE) + tok_in_tile + else: + block_in_part = n_tile // fx.Int32(_TILES_PER_BLOCK) + tile_in_block = n_tile - block_in_part * fx.Int32(_TILES_PER_BLOCK) + tile_start_tok = partition_start + n_tile * fx.Int32(TILE_SIZE) + slot = tile_in_block * fx.Int32(TILE_SIZE) + tok_in_tile + tile_in_seq = tile_start_tok < seq_len + bt_off = bt_seq_base + block_in_part + bt_off_safe = tile_in_seq.select(bt_off, seq * c_bt) + phys_block = buffer_ops.buffer_load( + bt_rsrc, + bt_off_safe, + vec_width=1, + dtype=T.i32, + ) + # A buffer descriptor addresses with a 32-bit voffset, so + # ``phys_block * c_block`` wraps once the cache view exceeds + # 4 GiB. Fold the block base into the descriptor's 64-bit + # base pointer and keep only in-block offsets below. + blk_off_i64 = arith.extui(T.i64, phys_block) * fx.Int64( + _stride_cache_block + ) + blk_rsrc = buffer_ops.create_buffer_resource( + kv_cache_ptr, + max_size=True, + base_byte_offset=blk_off_i64, + ) + block_base = fx.Int32(0) + + slot_base_byte = block_base + slot * c_stride_pos + kv_h * c_stride_head + # K codes: LANE_CODE_BYTES (=32) at slot_base + chunk*32, in + # HALVES x 16-byte buffer_loads. + k_byte0 = slot_base_byte + chunk_in_tok * fx.Int32(LANE_CODE_BYTES) + _kpl = [] + for hf in range_constexpr(HALVES): + _kpl.append( + buffer_ops.buffer_load( + blk_rsrc, + (k_byte0 + fx.Int32(hf * 16)) // fx.Int32(4), + vec_width=4, + dtype=T.i32, + cache_modifier=_KV_AUX, + ) + ) + + # ---- V codes + K/V scale HBM loads ------------------------ + v_byte0 = ( + slot_base_byte + + c_vcode_off + + chunk_in_tok * fx.Int32(LANE_CODE_BYTES) + ) + _vpl = [] + for hf in range_constexpr(HALVES): + _vpl.append( + buffer_ops.buffer_load( + blk_rsrc, + (v_byte0 + fx.Int32(hf * 16)) // fx.Int32(4), + vec_width=4, + dtype=T.i32, + cache_modifier=_KV_AUX, + ) + ) + # UE8M0 group scale words (one 4-byte word covers this lane's + # two groups). scale = 2^(byte-127) = bitcast_f32(byte<<23); + # byte==0 -> +0.0 (zero sentinel). + _ksw = buffer_ops.buffer_load( + blk_rsrc, + (slot_base_byte + c_kscale_off + c_scale_word_off) // fx.Int32(4), + vec_width=1, + dtype=T.i32, + ) + _vsw = buffer_ops.buffer_load( + blk_rsrc, + (slot_base_byte + c_vscale_off + c_scale_word_off) // fx.Int32(4), + vec_width=1, + dtype=T.i32, + ) + return (_kpl, _vpl, _ksw, _vsw) + + # Depth-1 software prefetch: emit the next tile's HBM loads before + # consuming the current tile without changing arithmetic order. + _PF_DEPTH = 1 + # Prime the pipeline with the first _PF_DEPTH tiles' loads. + _pf_queue = [] + for _pi in range_constexpr(min(_PF_DEPTH, 16)): + _pf_queue.append(_emit_tile_loads(_pi)) + _tile_range = range_constexpr(16) + for _jt in _tile_range: + n_tile = _jt + _ntile_tok = fx.Int32(n_tile * TILE_SIZE) + (k_packed_list, v_packed_list, kscale_word, vscale_word) = ( + _pf_queue.pop(0) + ) + _pf_nxt = n_tile + _PF_DEPTH + if const_expr(_pf_nxt < 16): + _pf_queue.append(_emit_tile_loads(_pf_nxt)) + # Store raw FP4 codes (natural contiguous order, + # 16 B = 32 nibbles per UE8M0 group) straight to KV LDS as + # the scaled-MFMA A operand — no dequant here. + for hf in range_constexpr(HALVES): + grp = chunk_in_tok * fx.Int32(2) + fx.Int32(hf) + grp_shift = (grp & fx.Int32(3)) * fx.Int32(8) + kscale_byte = (kscale_word >> grp_shift) & fx.Int32(0xFF) + k_packed = k_packed_list[hf] + vector.store( + k_packed, + kv_lds_i32, + [ + arith.index_cast( + T.index, + _kvi32( + tok_in_tile * fx.Int32(KFP4_ROW_I32) + + grp * fx.Int32(4) + ), + ) + ], + ) + scale_lds_i32.store( + kscale_byte, + [ + arith.index_cast( + T.index, + _sci(tok_in_tile * fx.Int32(N_GROUPS) + grp), + ) + ], + ) + gpu.barrier() + # QK: MFMA_ISSUES native scaled MFMAs over K=128, + # chained through the accumulator. + # A = K raw FP4 codes (vec4 i32 = 16 B) per (token, group) + # B = Q E4M3 (vec8 i32 = 32 B) for the same group + # scaleA = per-(token,group) UE8M0 byte at op_sel byte 0 + # scaleB = 0x7F identity (Q carries no scale) + BLGP_E2M1 = 4 # A operand type code = fp4 e2m1 + CBSZ_E4M3 = 0 # B operand type code = fp8 e4m3 + IDENT = fx.Int32(0x7F) + qk_acc = zero_v4 + for h in range_constexpr(MFMA_ISSUES): + grp = fx.Int32(h * GRPS_PER_ISSUE) + mfma_col_grp + k_op = vector.load_op( + T.vec(4, T.i32), + kv_lds_i32, + [ + arith.index_cast( + T.index, + _kvi32( + mfma_row * fx.Int32(KFP4_ROW_I32) + + grp * fx.Int32(4) + ), + ) + ], + ) + q_op = q_op_hoisted[h] + scbyte = fx.Int32( + scale_lds_i32.load( + [ + arith.index_cast( + T.index, + _sci(mfma_row * fx.Int32(N_GROUPS) + grp), + ) + ] + ) + ) + kscale = fx.Int32(0x7F7F7F00) | scbyte + qk_acc = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), + [ + k_op, + q_op, + qk_acc, + BLGP_E2M1, + CBSZ_E4M3, + 0, + kscale, + 0, + IDENT, + ], + ) + + # Scale + mask out-of-context tokens. + qk_acc = _vsplat_mul(qk_acc, QK_SCALE) + for elem in range_constexpr(4): + kv_tok = ( + partition_start + + _ntile_tok + + mfma_col_grp * fx.Int32(4) + + fx.Int32(elem) + ) + in_b = kv_tok < seq_len + v = vector.extract(qk_acc, static_position=[elem]) + qk_acc = _vector_insert( + in_b.select(v, NEG_INF), + qk_acc, + static_position=[elem], + ) + + # FA2 online softmax: per-query-row reduce. + local_max = vector.reduction(T.f32, "maxnumf", qk_acc) + r1 = local_max.shuffle_xor(fx.Int32(16), c_w) + local_max = local_max.maximumf(r1) + r2 = local_max.shuffle_xor(fx.Int32(32), c_w) + tile_max = local_max.maximumf(r2) + + new_max = running_max.maximumf(tile_max) + max_diff = running_max - new_max + safe_diff = (running_max > NEG_INF).select(max_diff, ZERO_F) + scale = (safe_diff * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast) + running_sum = running_sum * scale + for h in range_constexpr(PV_N_CHUNKS): + acc_pv[h] = _vsplat_mul(acc_pv[h], scale) + running_max = new_max + + tile_sum = ZERO_F + for elem in range_constexpr(4): + s = vector.extract(qk_acc, static_position=[elem]) + d = s - new_max + d = (new_max > NEG_INF).select(d, NEG_INF) + p = (d * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast) + tile_sum = tile_sum + p + qk_acc = _vector_insert(p, qk_acc, static_position=[elem]) + + ts1 = tile_sum.shuffle_xor(fx.Int32(16), c_w) + tile_sum = tile_sum + ts1 + ts2 = tile_sum.shuffle_xor(fx.Int32(32), c_w) + tile_sum = tile_sum + ts2 + running_sum = running_sum + tile_sum + # V dequant → LDS [token][head_dim] (ROW-MAJOR, no transpose) + # written as ds_write_b128 (vec(2,i64)); read back via the + # ds_read_tr16_b64 HW transpose in the PV block below. + v_lds_elem_base = tok_in_tile * fx.Int32( + KV_ROW_ELEMS + ) + chunk_in_tok * fx.Int32(SUBCHUNK_HDIMS) + for hf in range_constexpr(HALVES): + v_packed = v_packed_list[hf] + half_hd = fx.Int32(hf * HALF_HDIMS) + grp_shift = ( + (chunk_in_tok * fx.Int32(2) + fx.Int32(hf)) & fx.Int32(3) + ) * fx.Int32(8) + vscale_byte = (vscale_word >> grp_shift) & fx.Int32(0xFF) + vscale_f32 = arith.bitcast( + T.f32, _ival(vscale_byte << fx.Int32(23)) + ) + for w in range_constexpr(4): + word_i32 = vector.extract(v_packed, static_position=[w]) + # Native CDNA4 scaled convert: 2-wide + # cvt_scalef32_pk_bf16_fp4. srcSel s picks byte s + # of the word (nibbles 2s, 2s+1) → 2 bf16 with the + # UE8M0 group scale fused. 4 cvt calls/word. + bf16_elems = [] + for s in range_constexpr(4): + v2 = rocdl.cvt_scalef32_pk_bf16_fp4( + T.vec(2, T.bf16), + _ival(word_i32), + _ival(vscale_f32), + int(s), + ) + bf16_elems.append(vector.extract(v2, static_position=[0])) + bf16_elems.append(vector.extract(v2, static_position=[1])) + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems) + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16) + v_lds_i64_idx = ( + v_lds_elem_base + half_hd + fx.Int32(w * 8) + ) // fx.Int32(4) + vector.store( + v_i64, + kv_lds_i64, + [arith.index_cast(T.index, _kvi64(v_lds_i64_idx))], + ) + # HW V transpose cross-lane LDS fence (see 128 kernel). + rocdl.sched_barrier(0) + rocdl.s_waitcnt(0xC07F) + gpu.barrier() + + # ---- PV MFMA: A=V[head_dim, token], B=P (=qk_acc bf16) ----- + p_bf16 = arith.trunc_f(T.vec(4, T.bf16), qk_acc) + p_op = vector.bitcast(T.vec(4, T.i16), p_bf16) + # HW-transpose PV: V_lds row-major V[token][head_dim]. + token_idx = lane >> fx.Int32(2) + hd_sub = (lane & fx.Int32(3)) * fx.Int32(4) + v_lane_byte = ( + fx.Int32(kv_off) + + token_idx * fx.Int32(KV_ROW_BYTES) + + hd_sub * fx.Int32(2) + ) + for h in range_constexpr(PV_N_CHUNKS): + v_byte_off = v_lane_byte + fx.Int32(h * 32) + v_byte_i64 = fx.Int64(v_byte_off) + v_ptr = buffer_ops.create_llvm_ptr( + v_byte_i64, + address_space=3, + ) + v_op_raw = rocdl.ds_read_tr16_b64( + T.vec(4, T.i16), + v_ptr, + ).result + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k( + T.f32x4, [v_op_raw, p_op, acc_pv[h], 0, 0, 0] + ) + + _scf.YieldOp( + [ + _ival(running_max), + _ival(running_sum), + *[_ival(p) for p in acc_pv], + ] + ) + finally: + _for_ip.__exit__(None, None, None) + running_max = _for_op.results[0] + running_sum = _for_op.results[1] + acc_pv = list(_for_op.results[2:]) + + # ===== STEP G: Output =========================================== + safe_sum = (running_sum > ZERO_F).select(running_sum, ONE_F) + rcp = ONE_F / safe_sum + + c_os = fx.Int32(_stride_out_seq) + c_oh = fx.Int32(_stride_out_head) + c_op_ = fx.Int32(_stride_out_part) + out_base = seq * c_os + kv_h * c_oh + part * c_op_ + + valid_row_pred = arith.cmpi( + arith.CmpIPredicate.ult, + mfma_row.ir_value() if hasattr(mfma_row, "ir_value") else mfma_row, + arith.constant(QG, type=T.i32), + ) + _if = _scf.IfOp(valid_row_pred) + with ir.InsertionPoint(_if.then_block): + for h in range_constexpr(PV_N_CHUNKS): + pv_norm = _vsplat_mul(acc_pv[h], rcp) + pv_bf16 = arith.trunc_f(T.vec(4, T.bf16), pv_norm) + pv_i32x2 = vector.bitcast(T.vec(2, T.i32), pv_bf16) + head_dim_start = fx.Int32(h * 16) + mfma_col_grp * fx.Int32(4) + out_off_elem = ( + out_base + mfma_row * fx.Int32(HEAD_SIZE) + head_dim_start + ) + buffer_ops.buffer_store( + pv_i32x2, + out_rsrc, + out_off_elem * fx.Int32(2), + offset_is_bytes=True, + ) + + c_npq = fx.Int32(num_partitions * QG) + ml_off = ( + seq * fx.Int32(_stride_ml_seq) + kv_h * c_npq + part * c_qg + mfma_row + ) + es_off = ( + seq * fx.Int32(_stride_es_seq) + kv_h * c_npq + part * c_qg + mfma_row + ) + buffer_ops.buffer_store(running_max, ml_rsrc, ml_off) + buffer_ops.buffer_store(running_sum, es_rsrc, es_off) + _scf.YieldOp([]) + + # Hand the build-local allocator to the launcher. `_get_kernel` reads + # `kmod.allocator` immediately after this returns (before the next build can + # run), so each launcher captures its own variant's allocator. This does NOT + # reintroduce the drift bug: the kernel above closes over the LOCAL + # `allocator`, so no @flyc.jit function references module-global `allocator` + # and it is never snapshotted by _check_globals_drift. + globals()["allocator"] = allocator + return ultraquant_decode_hd256_kernel diff --git a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py new file mode 100644 index 000000000000..7f8c720678ff --- /dev/null +++ b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py @@ -0,0 +1,834 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlyDSL launcher for the optimized UltraQuant 4-bit D=256 decode kernel. + +The production kernel uses scaled FP4×E4M3 QK MFMA, native V conversion, +query hoisting, in-kernel Walsh-Hadamard rotation for GQA-6/8/16, and strided +tile-group scheduling. These choices are fixed and are not environment-tunable. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.triton_utils import triton + +logger = init_logger(__name__) + +_FLYDSL_AVAILABLE: bool | None = None +_ULTRAQUANT_MOD: Any = None +_FLYC: Any = None +_FX: Any = None +_TYPING_T: Any = None +_CC: Any = None +_IR: Any = None + + +def is_ultraquant_flydsl_available() -> bool: + """Return whether the gfx950 UltraQuant FlyDSL kernel can be loaded.""" + global _FLYDSL_AVAILABLE, _ULTRAQUANT_MOD + global _FLYC, _FX, _TYPING_T, _CC, _IR + if _FLYDSL_AVAILABLE is not None: + return _FLYDSL_AVAILABLE + try: + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + _FLYDSL_AVAILABLE = False + return False + + import flydsl.compiler as flyc + import flydsl.expr as fx + from flydsl._mlir import ir + from flydsl.compiler.kernel_function import CompilationContext + from flydsl.expr.typing import T + + from vllm.v1.attention.ops.flydsl_kernels import ( + ultraquant_decode_hd256 as ultraquant_mod, + ) + + _FLYC = flyc + _FX = fx + _TYPING_T = T + _CC = CompilationContext + _IR = ir + _ULTRAQUANT_MOD = ultraquant_mod + _FLYDSL_AVAILABLE = True + logger.info_once("FlyDSL UltraQuant D=256 decode is available") + except Exception as ex: # noqa: BLE001 + _FLYDSL_AVAILABLE = False + logger.warning_once( + "FlyDSL UltraQuant decode is unavailable (%s); using Triton fallback.", + ex, + ) + return _FLYDSL_AVAILABLE + + +_ULTRAQUANT_FLYDSL_GQA = (6, 8, 16) + + +def ultraquant_flydsl_decode_eligible( + *, + head_size: int, + num_kv_groups: int, + has_sinks: bool, + sliding_window: int | None, + flydsl_loaded: bool, +) -> bool: + """Return whether decode should use the gfx950 FlyDSL kernel. + + Production FlyDSL covers D=256 and GQA in {6, 8, 16}. Sinks and sliding + window fall back to unified Triton. ``flydsl_loaded`` is the process-level + compiler/hardware probe so this policy can be unit-tested without a GPU. + """ + return ( + flydsl_loaded + and head_size == 256 + and num_kv_groups in _ULTRAQUANT_FLYDSL_GQA + and not has_sinks + and not (sliding_window and sliding_window > 0) + ) + + +# -- Kernel module cache ------------------------------------------------------- +# Compiling a FlyDSL kernel is expensive; cache by the constexprs that +# parameterize stride math: num_kv_heads, num_partitions, max_blocks_per_seq, +# scale, query_group_size, kv_block_size, padded_slot, hw_v_transpose, +# tile_groups_per_partition. +_KERN_CACHE: dict[tuple, Any] = {} + +_LOG_INVOKED_ONCE: bool = False + +# Per-device cached FP4 E2M1 value table (used as the kernel "centroids" LUT). +_FP4_LUT_CACHE: dict[tuple[str, torch.dtype], torch.Tensor] = {} + + +def _fp4_value_lut(device: torch.device) -> torch.Tensor: + """Return the fixed 16-entry FP4 E2M1 value table as fp32 [16] on device. + + Indexed by the stored 4-bit E2M1 code, so ``lut[nibble]`` is the dequant + value directly (matches ``format.FP4_BITS_TO_VALUE``). + """ + from vllm.v1.attention.ops.ultraquant.format import FP4_BITS_TO_VALUE + + key = (str(device), torch.float32) + t = _FP4_LUT_CACHE.get(key) + if t is None: + t = torch.tensor( + FP4_BITS_TO_VALUE, dtype=torch.float32, device=device + ).contiguous() + _FP4_LUT_CACHE[key] = t + return t + + # Per-device cached qperm head-dim permutation for scaled QK MFMA. + + +# qperm[phys] = the natural head-dim placed at physical operand position phys. +# It realigns the fixed FP4xFP8 scaled-MFMA contraction so that, with K fed +# native-contiguous, Q's columns pair correctly. Folded into PiT on the host +# (PiT[:, qperm]) so the Q rotation matmul emits already-permuted q_rot at zero +# extra runtime cost. Derived + validated in tmp_qk_scaled_mfma.py (contigK). +_QPERM_CACHE: dict[str, torch.Tensor] = {} + + +def _kmap_phys(group: int, s: int) -> int: + # Physical operand position of group's element s for the K=128 fp4/fp8 + # scaled MFMA (matches tmp_qk_scaled_mfma.py::_kmap). + return 32 * (s >> 4) + 64 * (group % 2) + 16 * (group // 2) + (s & 15) + + +def _qperm_index( + device: torch.device, D: int = 128, group_size: int = 32 +) -> torch.Tensor: + # _kmap_phys describes ONE K=128 contraction (4 groups of 32). A head dim + # wider than 128 is covered by ceil(D/128) back-to-back MFMA issues, each + # contracting its own independent 128-wide slice, so the same map repeats + # per slice: qperm = [qperm128, qperm128 + 128, ...]. Applying _kmap_phys + # with group >= 4 instead yields positions outside the slice and is not a + # permutation at all. + key = f"{device}:{D}:{group_size}" + t = _QPERM_CACHE.get(key) + if t is None: + k_per_mfma = 128 + assert D % k_per_mfma == 0, ( + f"scaled QK MFMA needs D to be a multiple of {k_per_mfma}, got {D}" + ) + groups_per_mfma = k_per_mfma // group_size + qperm = [0] * D + for half in range(D // k_per_mfma): + base = half * k_per_mfma + for g in range(groups_per_mfma): + for s in range(group_size): + qperm[base + _kmap_phys(g, s)] = base + g * group_size + s + assert sorted(qperm) == list(range(D)), "qperm is not a permutation" + t = torch.tensor(qperm, dtype=torch.long, device=device).contiguous() + _QPERM_CACHE[key] = t + return t + + +def _detect_max_capture_B() -> int: + try: + from vllm.config import get_current_vllm_config + + cfg = get_current_vllm_config() + candidates: list[int] = [] + sizes = cfg.compilation_config.cudagraph_capture_sizes + if sizes: + candidates.append(int(max(sizes))) + sched = getattr(cfg, "scheduler_config", None) + if sched is not None and getattr(sched, "max_num_seqs", None): + candidates.append(int(sched.max_num_seqs)) + if candidates: + return max(candidates) + except Exception: # noqa: BLE001 + pass + return 512 + + +class _SegmBufPool: + """Single-bucket buffer pool for segm_out/segm_max/segm_sum/output + the + pooled Q-rotation intermediates. Identical to the TQ pool except that + ultraquant needs an extra ``q_fp8`` (float8_e4m3fn) slot for the Q precision + haircut (``q_rot_fp32 -> fp8_e4m3 -> bf16``) so that no fresh allocation + lands in the HIP graph memory pool post-capture. + """ + + __slots__ = ("_bufs", "_max_B") + + def __init__(self) -> None: + self._bufs: dict[tuple, dict[str, torch.Tensor]] = {} + self._max_B: int | None = None + + def get( + self, + B: int, + Hk: int, + Hq: int, + num_partitions: int, + QG: int, + D: int, + device: torch.device, + q_dtype: torch.dtype, + ) -> dict[str, torch.Tensor]: + if self._max_B is None: + self._max_B = _detect_max_capture_B() + B_bucket = max(self._max_B, int(B)) + key = ( + int(Hk), + int(Hq), + int(num_partitions), + int(QG), + int(D), + str(device), + q_dtype, + ) + bufs = self._bufs.get(key) + if bufs is None or bufs["segm_out"].shape[0] < B_bucket: + if bufs is not None and bufs["segm_out"].shape[0] < B_bucket: + # First-time grow: warn so user knows cudagraphs may be + # invalidated. This should only happen during warmup or in + # always-eager configs (a batch larger than any captured + # cudagraph size). + logger.warning_once( + "FlyDSL UltraQuant buffer pool growing from %d " + "to %d (B=%d). If this happens AFTER cudagraph warmup the " + "captured graphs hold stale pointers and will GPU-fault.", + bufs["segm_out"].shape[0], + B_bucket, + B, + ) + bufs = { + "segm_out": torch.empty( + (B_bucket, Hk, num_partitions, QG, D), + dtype=torch.bfloat16, + device=device, + ), + "segm_max": torch.empty( + (B_bucket, Hk, num_partitions, QG), + dtype=torch.float32, + device=device, + ), + "segm_sum": torch.empty( + (B_bucket, Hk, num_partitions, QG), + dtype=torch.float32, + device=device, + ), + "output": torch.empty( + (B_bucket, Hq, D), + dtype=q_dtype, + device=device, + ), + "q_rot": torch.empty( + (B_bucket, Hq, D), + dtype=q_dtype, + device=device, + ), + "q_float": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float32, + device=device, + ), + "q_rot_fp32": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float32, + device=device, + ), + # ultraquant-specific: E4M3 haircut intermediate (pooled so the + # cast does not allocate a fresh tensor post-capture). + "q_fp8": torch.empty( + (B_bucket, Hq, D), + dtype=torch.float8_e4m3fn, + device=device, + ), + } + self._bufs[key] = bufs + self._max_B = B_bucket + logger.info_once( + "FlyDSL UltraQuant buffer pool allocated " + "shape=(Hk=%d, Hq=%d, P=%d, QG=%d, D=%d, dtype=%s) " + "B_bucket=%d. VRAM = %.1f MiB / shape.", + Hk, + Hq, + num_partitions, + QG, + D, + q_dtype, + B_bucket, + sum(t.numel() * t.element_size() for t in bufs.values()) / (1 << 20), + ) + return { + "segm_out": bufs["segm_out"][:B], + "segm_max": bufs["segm_max"][:B], + "segm_sum": bufs["segm_sum"][:B], + "output": bufs["output"][:B], + "q_rot": bufs["q_rot"][:B], + "q_float": bufs["q_float"][:B], + "q_rot_fp32": bufs["q_rot_fp32"][:B], + "q_fp8": bufs["q_fp8"][:B], + } + + +_SEGM_POOL = _SegmBufPool() + +_GET_KERNEL_STATS = {"hits": 0, "misses": 0, "build_total_s": 0.0} + + +def _get_kernel( + num_kv_heads: int, + num_partitions: int, + max_blocks_per_seq: int, + scale: float, + query_group_size: int, + kv_block_size: int, + padded_slot: int, + fuse_qrot: bool, + num_seqs_hint: int, + tile_groups_per_partition: int, + stride_q_seq: int, + stride_q_head: int, +): + key = ( + num_kv_heads, + num_partitions, + max_blocks_per_seq, + round(scale, 8), + query_group_size, + kv_block_size, + padded_slot, + fuse_qrot, + tile_groups_per_partition, + stride_q_seq, + stride_q_head, + ) + cached = _KERN_CACHE.get(key) + if cached is not None: + _GET_KERNEL_STATS["hits"] += 1 + return cached + + assert is_ultraquant_flydsl_available() + import time + + build_start = time.perf_counter() + kmod = _ULTRAQUANT_MOD + assert kmod is not None + kfn = kmod.build_ultraquant_decode_hd256_module( + num_seqs=num_seqs_hint, + num_kv_heads=num_kv_heads, + num_partitions=num_partitions, + padded_slot=padded_slot, + max_blocks_per_seq=max_blocks_per_seq, + softmax_scale=scale, + query_group_size=query_group_size, + kv_block_size=kv_block_size, + tile_groups_per_partition=tile_groups_per_partition, + fuse_qrot=fuse_qrot, + stride_q_seq=stride_q_seq, + stride_q_head=stride_q_head, + ) + allocator = kmod.allocator + flyc = _FLYC + fx = _FX + T = _TYPING_T + CompilationContext = _CC + ir_mod = _IR + + @flyc.jit + def _launch( + out, + es, + ml, + q, + kvc, + cents, + bt, + sl, + gx: fx.Int32, # type: ignore[name-defined] + gy: fx.Int32, # type: ignore[name-defined] + gz: fx.Int32, # type: ignore[name-defined] + stream: fx.Stream, # type: ignore[name-defined] + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir_mod.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + from flydsl.expr import arith + + grid_x = arith.index_cast(T.index, gx.ir_value()) + grid_y = arith.index_cast(T.index, gy.ir_value()) + grid_z = arith.index_cast(T.index, gz.ir_value()) + kfn(out, es, ml, q, kvc, cents, bt, sl).launch( + grid=(grid_x, grid_y, grid_z), + block=(kmod.BLOCK_THREADS, 1, 1), + stream=stream, + ) + + _KERN_CACHE[key] = _launch + build_time = time.perf_counter() - build_start + _GET_KERNEL_STATS["misses"] += 1 + _GET_KERNEL_STATS["build_total_s"] += build_time + logger.info( + "Built FlyDSL UltraQuant kernel #%d in %.2fs (total %.1fs), key=%s", + _GET_KERNEL_STATS["misses"], + build_time, + _GET_KERNEL_STATS["build_total_s"], + key, + ) + return _launch + + +# -- Partition reducer --------------------------------------------------------- +from vllm.triton_utils import tl # noqa: E402 + + +@triton.jit +def _reduce_partitions( + output_ptr, # [N, Hq, D] in OUT_DTYPE + segm_out_ptr, # [N, Hk, P, QG, D] bf16 + segm_max_ptr, # [N, Hk, P, QG] fp32 + segm_sum_ptr, # [N, Hk, P, QG] fp32 + out_stride_n: tl.int64, + out_stride_h: tl.int64, + NUM_KV_HEADS: tl.constexpr, + QG: tl.constexpr, + NUM_PARTS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, +): + n = tl.program_id(0) + hq = tl.program_id(1) + kv_h = hq // QG + qg = hq % QG + # Splitting the head-dim across programs is what gives this kernel enough + # parallelism to use the GPU: the reduction runs over PARTITIONS, so every + # head-dim column is independent. A (B, Hq) grid is only 32 workgroups at + # B=4 -- ~12% of a 256-CU part -- which is why it moved 8.4 MB at just + # ~995 GB/s. The partition reduction itself is untouched (same values, same + # order), so the result is bit-identical; only the columns each program + # owns change. The [P] max/sum vectors get re-read per D-block, which is + # 4 KB of duplicated traffic against a 262 KB tile. + + msum_base = n * (NUM_KV_HEADS * NUM_PARTS * QG) + kv_h * (NUM_PARTS * QG) + qg + so_base = ( + n * (NUM_KV_HEADS * NUM_PARTS * QG * HEAD_SIZE) + + kv_h * (NUM_PARTS * QG * HEAD_SIZE) + + qg * HEAD_SIZE + ) + + p_off = tl.arange(0, NUM_PARTS) + d_off = tl.program_id(2) * BLOCK_D + tl.arange(0, BLOCK_D) + + m_idx = msum_base + p_off * QG + seg_max = tl.load(segm_max_ptr + m_idx) + seg_sum = tl.load(segm_sum_ptr + m_idx) + + valid = seg_max > float("-inf") + overall_max = tl.max(tl.where(valid, seg_max, float("-inf"))) + + rescale = tl.where(valid, tl.exp(seg_max - overall_max), 0.0) + seg_sum_rescaled = seg_sum * rescale + overall_sum = tl.sum(seg_sum_rescaled) + + so_idx = so_base + p_off[:, None] * (QG * HEAD_SIZE) + d_off[None, :] + seg_out = tl.load(segm_out_ptr + so_idx).to(tl.float32) + weighted = seg_out * seg_sum_rescaled[:, None] + acc_sum = tl.sum(weighted, axis=0) + acc = tl.where(overall_sum > 0.0, acc_sum / overall_sum, 0.0) + + out_off = n * out_stride_n + hq * out_stride_h + d_off + tl.store(output_ptr + out_off, acc.to(output_ptr.dtype.element_ty)) + + +# -- Fused Q rotation + E4M3 haircut (Triton) --------------------------------- +# Replaces the 4-op torch prologue (copy bf16->fp32, mm by PiT, copy fp32->e4m3, +# copy e4m3->bf16) with ONE kernel. The win is kernel COUNT, not bytes: each of +# those small ops costs ~2-4us of fixed GPU-side cost here largely independent of +# element count (verified: sweeping the pool's B_bucket 512->1, i.e. 512x fewer +# elements in the copies, moved the floor 0%). Under CUDA graphs the host launch +# cost is already gone, so only per-kernel GPU cost remains and collapsing 4 +# kernels into 1 is the lever that is left. +@triton.jit +def _fused_q_rot_haircut( + q_ptr, # [B, Hq, D] query dtype (bf16/fp16), strided + pit_ptr, # [D, D] fp32 contiguous (already qperm-permuted) + out_ptr, # [B_bucket, Hq, D] out dtype, contiguous + M, # B * Hq (runtime) + q_stride_n: tl.int64, + q_stride_h: tl.int64, + HQ: tl.constexpr, + D: tl.constexpr, + HAIRCUT: tl.constexpr, # round through E4M3 (fp8-consuming QK MFMA paths) + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + m_off = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + m_mask = m_off < M + n = m_off // HQ + h = m_off % HQ + + # Splitting the OUTPUT columns across programs (not just the rows) is what + # gives this kernel parallelism at small batch: M = B*Hq is only 8 at B=1, + # so a row-only grid launches ONE workgroup and serializes the whole + # 256x256 rotation on a single CU. The K reduction is untouched by the + # split, so the result stays bit-identical. + d_off = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) + for k0 in range(0, D, BLOCK_K): + k_off = k0 + tl.arange(0, BLOCK_K) + q_idx = n[:, None] * q_stride_n + h[:, None] * q_stride_h + k_off[None, :] + q_blk = tl.load(q_ptr + q_idx, mask=m_mask[:, None], other=0.0) + p_blk = tl.load(pit_ptr + k_off[:, None] * D + d_off[None, :]) + acc = tl.dot(q_blk.to(tl.float32), p_blk, acc) + + if HAIRCUT: + # fp32 -> E4M3 -> fp32. Mirrors the torch path's + # q_rot_fp32 -> float8_e4m3fn -> bf16 round trip (e4m3 is a subset of + # bf16, so the second hop is exact and folds away here). + acc = acc.to(tl.float8e4nv).to(tl.float32) + + out_idx = m_off[:, None] * D + d_off[None, :] + tl.store(out_ptr + out_idx, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None]) + + +# Prologue Q-rotation fusion (fallback for non-hd256 / in-kernel-WHT-off paths). +# Cleared the accuracy ladder (offline parity -> autoregressive drift -> paired +# GSM8K); pinned ON as part of the validated production config. +_FUSE_Q_ROT = True + +# Stronger form: do the rotation INSIDE the decode kernel as a Walsh-Hadamard +# butterfly and drop the prologue dispatch entirely. The prologue is ~7us at +# B=4 and that is almost all fixed per-dispatch cost (sweeping its grid 4x moves +# it 0%), so deleting the launch -- not optimizing it -- is the win. +# Default ON: verified bit-exact (torch.equal) against the prologue path across +# B={1,2,3,4,8,16} x seq={4k,8k,16k,32k,50021,75k}, including non-power-of-2 +# batch and seq. +_FUSE_QROT_INKERNEL = True + + +def _run_fused_q_rot(query, PiT_used, out, B, Hq, D, haircut): + """Launch the fused prologue. ``out`` is the pooled [B_bucket, Hq, D] buf.""" + M = B * Hq + BLOCK_M = 16 + BLOCK_K = 32 + BLOCK_D = 64 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(D, BLOCK_D)) + _fused_q_rot_haircut[grid]( + query, + PiT_used, + out, + M, + query.stride(0), + query.stride(1), + HQ=int(Hq), + D=int(D), + HAIRCUT=bool(haircut), + BLOCK_M=BLOCK_M, + BLOCK_K=BLOCK_K, + BLOCK_D=BLOCK_D, + ) + + +# -- Public launcher ---------------------------------------------------------- +def flydsl_ultraquant_decode_attention( + query: torch.Tensor, + kv_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + scale: float, + PiT: torch.Tensor | None = None, + max_seq_len: int = 0, + output_buf: torch.Tensor | None = None, + buf_holder: Any = None, + max_num_kv_splits: int = 32, + sinks: torch.Tensor | None = None, +) -> torch.Tensor: + """Launch the optimized FlyDSL UltraQuant D=256 decode.""" + if not is_ultraquant_flydsl_available(): + raise RuntimeError( + "FlyDSL UltraQuant decode requires gfx950 and importable FlyDSL." + ) + if sinks is not None: + raise NotImplementedError("FlyDSL UltraQuant decode does not support sinks") + + B, Hq, D = query.shape + Hk = kv_cache.shape[2] + block_size = kv_cache.shape[1] + padded_slot = int(kv_cache.shape[3]) + QG = Hq // Hk + assert D == 256, f"UltraQuant FlyDSL requires head_dim=256, got {D}" + assert block_size in (16, 32, 64, 128, 256), ( + f"UltraQuant FlyDSL supports kv_block_size 16/32/64/128/256, got {block_size}" + ) + assert QG in _ULTRAQUANT_FLYDSL_GQA, ( + f"UltraQuant FlyDSL supports GQA factor " + f"{', '.join(map(str, _ULTRAQUANT_FLYDSL_GQA))}, got {QG}" + ) + + device = query.device + + # ---- PiT (Hadamard) — cache the contiguous fp32 form on the layer ----- + PiT_f32: torch.Tensor + if buf_holder is not None: + PiT_f32 = getattr(buf_holder, "_ultraquant_PiT_f32", None) + if PiT_f32 is None: + assert PiT is not None, ( + "UltraQuant launcher requires PiT (the Hadamard rotation)" + ) + PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32) + if not PiT_f32.is_contiguous(): + PiT_f32 = PiT_f32.contiguous() + buf_holder._ultraquant_PiT_f32 = PiT_f32 + else: + assert PiT is not None, ( + "UltraQuant launcher requires PiT (the Hadamard rotation)" + ) + PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32) + if not PiT_f32.is_contiguous(): + PiT_f32 = PiT_f32.contiguous() + + centroids_c = _fp4_value_lut(device) + + # ---- Partition count (FA2 split-KV) ----------------------------------- + kv_compute_block = _ULTRAQUANT_MOD.KV_COMPUTE_BLOCK + worst_case_max_seq_len = int(block_table.shape[1]) * int(block_size) + if max_seq_len <= 0: + max_seq_len = worst_case_max_seq_len + sizing_max_seq_len = worst_case_max_seq_len + + # Split-KV parallelism cap. The old default (32) silently serialized long + # contexts: once required_num_partitions exceeds the cap, TGPP grows and each + # workgroup walks MORE tile-groups instead of the grid getting wider. The + # decode kernel is latency-bound at ~1-2 waves/CU (grid = B*Hk*P workgroups + # of a single wavefront), so that serialization dominated long-context decode. + # Raising the cap to 256 keeps TGPP ~1 and measured (B=16, rocprof, per step + # decode+reduce+store): 16k 82.7->40.0us, 32k 152.0->55.1us, 64k 339.4->95.7us, + # 131k 671.3->179.0us (2.1x - 3.8x, growing with context). 512 regresses + # because _reduce_partitions cost grows with partition count. + # Cost: the segm pool scales with P (~8x vs P=32). + # BATCH-ADAPTIVE split-KV cap (low-concurrency occupancy fix). + # rocprof (B=4, 75k): the decode kernel is occupancy/latency-bound, not + # compute-bound -- OccupancyPercent 5.8%, MfmaUtil 4.5%, VALUBusy 14.5% + # (stalled ~85% on HBM), VALUUtilization 90% (no divergence). The grid is + # B*Hk*num_partitions single-wavefront workgroups; at B=4/Hk=1/P=256 that's + # only 1024 wg = 4/CU, too few to hide memory latency. Splitting MORE at low + # batch adds resident wavefronts and cuts ITL; at high batch the grid is + # already full so extra partitions only inflate _reduce_partitions + empty- + # partition overhead. Measured end-to-end (graphs on, TP=8, uniform 75k, + # median ITL vs kv8): cap512 vs cap256 -> C4 +9.0%->+4.9%, C8 +6.6%->+2.5% + # (better), C16 +2.8%->+4.3% (worse). So use 512 at low batch, 256 at high. + # This is only SAFE because the build-local-allocator fix in + # ultraquant_decode_hd256.py lets the two num_partitions variants (256 & 512) + # coexist across CUDA-graph batch buckets without the @flyc.jit + # "global 'allocator' changed since first compile" drift crash. + MAX_PARTITIONS = 512 if B <= 8 else 256 + required_num_partitions = ( + sizing_max_seq_len + kv_compute_block - 1 + ) // kv_compute_block + parallelism_floor = min(MAX_PARTITIONS, max(1, max_num_kv_splits)) + num_partitions_actual = max( + parallelism_floor, + min(MAX_PARTITIONS, required_num_partitions), + ) + # Cap partitions by how much work there actually IS to split. Nothing above + # forces a partition to hold a useful number of tokens -- max_num_kv_splits + # sets a parallelism FLOOR -- so a short context gets split into partitions + # of a handful of tokens each, where every partition is almost entirely + # fixed cost (and each one still writes a full QG x D partial for the + # reduce to read back). Measured sweet spot is ~130-147 tokens/partition, + # roughly flat in between; next_power_of_2 below rounds this target up, so + # the effective range lands at ~96-192 tokens. At B=4 this is worth 27.9 -> + # 24.4 us at seq=2048 and 27.0 -> 24.8 us at seq=16384, and is inert once + # the context is long enough to want every partition (seq >= ~64k). + _tok_per_part = 192 + if _tok_per_part > 0: + _useful_parts = max( + 1, (int(sizing_max_seq_len) + _tok_per_part - 1) // _tok_per_part + ) + num_partitions_actual = min(num_partitions_actual, _useful_parts) + num_partitions = max(2, triton.next_power_of_2(num_partitions_actual)) + _tgpp_required = max( + 1, + (required_num_partitions + num_partitions - 1) // num_partitions, + ) + tile_groups_per_partition = int(triton.next_power_of_2(_tgpp_required)) + + # ---- Pooled buffers --------------------------------------------------- + pool_bufs = _SEGM_POOL.get( + B, + Hk, + Hq, + num_partitions, + QG, + D, + device, + query.dtype, + ) + segm_out = pool_bufs["segm_out"] + segm_max = pool_bufs["segm_max"] + segm_sum = pool_bufs["segm_sum"] + if output_buf is None: + output = pool_bufs["output"] + else: + output = output_buf[:B] if output_buf.shape[0] != B else output_buf + + # Fold the scaled-MFMA head-dim permutation into PiT's columns so the + # WHT / prologue produces native operand order. Always apply the E4M3 + # haircut: scaled QK consumes Q as E4M3. + PiT_used = PiT_f32 + qperm_idx = _qperm_index(device, D=D) + if buf_holder is not None: + PiT_used = getattr(buf_holder, "_ultraquant_PiT_perm_f32", None) + if PiT_used is None: + PiT_used = PiT_f32.index_select(1, qperm_idx).contiguous() + buf_holder._ultraquant_PiT_perm_f32 = PiT_used + else: + PiT_used = PiT_f32.index_select(1, qperm_idx).contiguous() + _q_float = pool_bufs["q_float"] + _q_rot_f32 = pool_bufs["q_rot_fp32"] + _q_fp8 = pool_bufs["q_fp8"] + _q_rot_out = pool_bufs["q_rot"] + + _inkernel_qrot = ( + _FUSE_QROT_INKERNEL + and int(D) == 256 + and QG in _ULTRAQUANT_FLYDSL_GQA + and query.dim() == 3 + and query.stride(2) == 1 + and query.stride(1) == int(D) + ) + if _inkernel_qrot: + q_for_kernel = query + # The fused kernel indexes the head-dim contiguously; fall back if not. + elif _FUSE_Q_ROT and query.stride(-1) == 1: + # One kernel for the whole prologue instead of four ops. + _run_fused_q_rot(query, PiT_used, _q_rot_out, B, Hq, D, haircut=True) + else: + _q_float.copy_(query) + torch.mm( + _q_float.view(B * Hq, D), + PiT_used, + out=_q_rot_f32.view(B * Hq, D), + ) + _q_fp8.copy_(_q_rot_f32) # fp32 -> e4m3 + _q_rot_out.copy_(_q_fp8) # e4m3 -> bf16 + if not _inkernel_qrot: + q_for_kernel = _q_rot_out + + # ---- FlyDSL kernel launch -------------------------------------------- + max_bps = int(block_table.shape[1]) + launch = _get_kernel( + Hk, + num_partitions, + max_bps, + scale, + QG, + block_size, + padded_slot, + fuse_qrot=_inkernel_qrot, + num_seqs_hint=int(B), + tile_groups_per_partition=int(tile_groups_per_partition), + stride_q_seq=int(q_for_kernel.stride(0)), + stride_q_head=int(q_for_kernel.stride(1)), + ) + global _LOG_INVOKED_ONCE + if not _LOG_INVOKED_ONCE: + _LOG_INVOKED_ONCE = True + logger.info( + "FlyDSL UltraQuant launcher invoked: B=%d Hk=%d Hq=%d D=%d QG=%d " + "num_partitions=%d (actual=%d, cap=%d) TGPP=%d max_bps=%d " + "block_size=%d padded_slot=%d max_seq_len=%d " + "(coverage=%d tokens, worst_case=%d tokens)", + B, + Hk, + Hq, + D, + QG, + num_partitions, + num_partitions_actual, + MAX_PARTITIONS, + tile_groups_per_partition, + max_bps, + int(block_size), + padded_slot, + int(max_seq_len), + num_partitions * tile_groups_per_partition * kv_compute_block, + worst_case_max_seq_len, + ) + launch( + segm_out, + segm_sum, + segm_max, + q_for_kernel, + kv_cache, + centroids_c, + block_table, + seq_lens, + B, + Hk, + num_partitions, + torch.cuda.current_stream(), + ) + + # ---- Reduce partitions -> [B, Hq, D] --------------------------------- + # 64 splits the head-dim 4 ways, taking the grid from 32 to 128 workgroups. + # Bit-identical by construction (the partition reduction is per-column). + _red_block_d = min(int(D), 64) + _reduce_partitions[(B, Hq, triton.cdiv(int(D), _red_block_d))]( + output_ptr=output, + segm_out_ptr=segm_out, + segm_max_ptr=segm_max, + segm_sum_ptr=segm_sum, + out_stride_n=output.stride(0), + out_stride_h=output.stride(1), + NUM_KV_HEADS=Hk, + QG=QG, + NUM_PARTS=num_partitions, + HEAD_SIZE=D, + BLOCK_D=_red_block_d, + ) + return output diff --git a/vllm/v1/attention/ops/ultraquant/__init__.py b/vllm/v1/attention/ops/ultraquant/__init__.py new file mode 100644 index 000000000000..dbe17bfdc8bc --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UltraQuant 4-bit KV cache: FP4 E2M1 codes + UE8M0 group-of-32 scales. + +Production decode uses the FlyDSL D=256 kernel on gfx950, with Triton +unified attention as the fallback. Slot size is ``slot_size(head_dim)`` +(272 B at D=256). Format helpers live in ``format``; import them from +there, not this package root. +""" + +__all__ = ["ultraquant_store"] + + +def __getattr__(name: str): + if name == "ultraquant_store": + from vllm.v1.attention.ops.ultraquant.triton_store import ultraquant_store + + return ultraquant_store + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/v1/attention/ops/ultraquant/format.py b/vllm/v1/attention/ops/ultraquant/format.py new file mode 100644 index 000000000000..20977d528fdc --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/format.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FP4 codepoint + UE8M0 scale constants for the UltraQuant KV cache. + +K/V codes are FP4 E2M1. Per-group scales are UE8M0 (one byte, a power of +two): ``s = 2^round(log2(c · absmax))`` with ``c = 0.156``. Decode +consumes Q as FP8 E4M3 so scaled F8F6F4 MFMA can run natively on CDNA4. + +AoS slot layout (group_size=32), per (token, head): + bytes [0 .. D/2) : K codes (FP4 nibbles, 2/byte) + bytes [D/2 .. D/2 + Gk) : K scales (UE8M0, 1 byte × Gk groups) + bytes [D/2 + Gk .. D + Gk) : V codes + bytes [D + Gk .. D + 2*Gk) : V scales + D=256 → Gk=8 → 272 B/slot. D=128 → Gk=4 → 136 B/slot. + +No per-token norm-fold. No V rotation. K is Hadamard-rotated at store. +""" + +from __future__ import annotations + +import math + +# ── 15 unique FP4 E2M1 levels, sorted ascending ──── +FP4_LEVELS_SORTED: tuple[float, ...] = ( + -6.0, + -4.0, + -3.0, + -2.0, + -1.5, + -1.0, + -0.5, + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, +) +assert len(FP4_LEVELS_SORTED) == 15 + +# Midpoints between consecutive sorted levels (14 boundaries). +MIDPOINTS_SORTED: tuple[float, ...] = tuple( + (FP4_LEVELS_SORTED[i] + FP4_LEVELS_SORTED[i + 1]) / 2.0 + for i in range(len(FP4_LEVELS_SORTED) - 1) +) +assert len(MIDPOINTS_SORTED) == 14 + +# FP4 E2M1 bit-pattern decode (16 entries indexed by 4-bit code). +# Sign bit (MSB) | 2 exp bits | 1 mantissa bit (LSB). +FP4_BITS_TO_VALUE: tuple[float, ...] = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) +assert len(FP4_BITS_TO_VALUE) == 16 + +# Sorted index → FP4 E2M1 bit pattern (uint4). +SORTED_TO_BITS: tuple[int, ...] = ( + 0b1111, + 0b1110, + 0b1101, + 0b1100, + 0b1011, + 0b1010, + 0b1001, + 0b0000, + 0b0001, + 0b0010, + 0b0011, + 0b0100, + 0b0101, + 0b0110, + 0b0111, +) +assert len(SORTED_TO_BITS) == 15 +for _i, _bits in enumerate(SORTED_TO_BITS): + _v = FP4_BITS_TO_VALUE[_bits] + _v = 0.0 if _v == 0.0 else _v + assert _v == FP4_LEVELS_SORTED[_i] + +# ── Format constants ────────────────────────────────────────────────────── +GROUP_SIZE: int = 32 +"""Elements per group; one E8M0 scale per group. Fixed at 32 because the +AMD scaled F8F6F4 MFMA instruction also consumes one E8M0 scale per 32 +elements along K. Other group sizes would force a fallback to plain MFMA ++ accumulator-side scale multiply (no hardware fast-path).""" + +FP4_MAX: float = 6.0 +"""Largest representable FP4 magnitude.""" + +DEFAULT_CONSTANT_C: float = 0.156 +"""MSE-optimal scale constant; ``s_raw = c * absmax``, then UE8M0-snapped.""" + + +def get_constant_c() -> float: + """Return the fixed UltraQuant scale constant.""" + return DEFAULT_CONSTANT_C + + +def get_group_size() -> int: + """Return the fixed group size required by scaled MFMA.""" + return GROUP_SIZE + + +# ── Scale encoding ──────────────────────────────────────────────────────── +# Per-group scale `s = pow2_round(c · absmax)` stored as the E8M0 byte; +# codebook is the raw FP4 grid {0, ±0.5, ..., ±6}; dequant = code · s. + + +# ── UE8M0 helpers ────────────────────────────────────────────────────────── +# E8M0 = 8 exponent bits, 0 mantissa bits, no sign. The byte value `e` +# represents `2^(e - 127)` for e in [1, 254]. e=0 and e=255 are reserved +# (zero and NaN respectively in the OCP spec). We use e=0 as our zero +# sentinel for sink-zero groups. + +UE8M0_BIAS: int = 127 + +# fp32 minimum positive normal = 2^-126 ≈ 1.175e-38. Anything smaller +# clamps to the zero sentinel. +_UE8M0_MIN_EXP: int = -126 +_UE8M0_MAX_EXP: int = 127 + + +def ue8m0_encode(s: float) -> int: + """Snap a positive fp32 scale `s` to the nearest power of 2 and + encode as a UE8M0 byte. `s <= 0` encodes as 0 (zero sentinel).""" + if s <= 0.0 or not math.isfinite(s): + return 0 + exp = int(round(math.log2(s))) + if exp < _UE8M0_MIN_EXP: + return 0 + if exp > _UE8M0_MAX_EXP: + exp = _UE8M0_MAX_EXP + return exp + UE8M0_BIAS + + +def ue8m0_decode(byte: int) -> float: + """Decode a UE8M0 byte back to fp32. 0 → 0.0 (zero sentinel).""" + if byte == 0: + return 0.0 + return 2.0 ** (byte - UE8M0_BIAS) + + +# ── Per-slot AoS byte layout ────────────────────────────────────────────── +# Per slot, per head (group_size=GS=32, Gk = D/GS): +# bytes [0 .. D/2) : K codes (2 nibbles/byte) +# bytes [D/2 .. D/2 + Gk) : K scales (E8M0, 1 byte each) +# bytes [D/2 + Gk .. D + Gk) : V codes +# bytes [D + Gk .. D + 2*Gk) : V scales (E8M0, 1 byte each) +# D=256 → 272 B/slot. D=128 → 136 B/slot. + + +def k_codes_bytes(head_dim: int) -> int: + """Bytes for one head's packed K codes.""" + if head_dim % 2 != 0: + raise ValueError(f"head_dim must be even, got {head_dim}") + return head_dim // 2 + + +def n_groups(head_dim: int, group_size: int | None = None) -> int: + gs = group_size if group_size is not None else get_group_size() + if head_dim % gs != 0: + raise ValueError(f"head_dim={head_dim} must be a multiple of group_size={gs}") + return head_dim // gs + + +def k_scales_bytes(head_dim: int, group_size: int | None = None) -> int: + """Bytes for one head's K scales (one E8M0 byte per group).""" + return n_groups(head_dim, group_size) + + +def v_codes_bytes(head_dim: int) -> int: + return k_codes_bytes(head_dim) + + +def v_scales_bytes(head_dim: int, group_size: int | None = None) -> int: + return k_scales_bytes(head_dim, group_size) + + +def slot_size(head_dim: int, group_size: int | None = None) -> int: + """Bytes per (token, head) slot.""" + return ( + k_codes_bytes(head_dim) + + k_scales_bytes(head_dim, group_size) + + v_codes_bytes(head_dim) + + v_scales_bytes(head_dim, group_size) + ) + + +def k_codes_offset(head_dim: int, group_size: int | None = None) -> int: + return 0 + + +def k_scales_offset(head_dim: int, group_size: int | None = None) -> int: + return k_codes_bytes(head_dim) + + +def v_codes_offset(head_dim: int, group_size: int | None = None) -> int: + return k_codes_bytes(head_dim) + k_scales_bytes(head_dim, group_size) + + +def v_scales_offset(head_dim: int, group_size: int | None = None) -> int: + return v_codes_offset(head_dim, group_size) + v_codes_bytes(head_dim) diff --git a/vllm/v1/attention/ops/ultraquant/reference.py b/vllm/v1/attention/ops/ultraquant/reference.py new file mode 100644 index 000000000000..0d7dbe745bab --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/reference.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""PyTorch reference for the UltraQuant KV cache format. + +Scale is UE8M0 (power of two): after ``s = c * absmax`` snap +``s = 2^round(log2(s))`` and store one byte. Q is Hadamard-rotated then +cast to FP8 E4M3 before the QK matmul, matching the kernel launcher. + +Same FP4 E2M1 grid. ``c = 0.156``. No per-token L2 norm. No V rotation. +K is rotated at encode; Q arrives pre-rotated at decode. + +This module is the bit-comparison ground truth for the kernels. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from vllm.v1.attention.ops.ultraquant.format import ( + FP4_LEVELS_SORTED, + GROUP_SIZE, + MIDPOINTS_SORTED, + SORTED_TO_BITS, + UE8M0_BIAS, + get_constant_c, +) + +_LEVELS_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} +_MIDPOINTS_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} +_SORTED_TO_BITS_CACHE: dict[torch.device, torch.Tensor] = {} + + +def _get_levels(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + key = (device, dtype) + t = _LEVELS_CACHE.get(key) + if t is None: + t = torch.tensor(FP4_LEVELS_SORTED, device=device, dtype=dtype) + _LEVELS_CACHE[key] = t + return t + + +def _get_midpoints(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + key = (device, dtype) + t = _MIDPOINTS_CACHE.get(key) + if t is None: + t = torch.tensor(MIDPOINTS_SORTED, device=device, dtype=dtype) + _MIDPOINTS_CACHE[key] = t + return t + + +def _get_sorted_to_bits(device: torch.device) -> torch.Tensor: + t = _SORTED_TO_BITS_CACHE.get(device) + if t is None: + t = torch.tensor(SORTED_TO_BITS, device=device, dtype=torch.uint8) + _SORTED_TO_BITS_CACHE[device] = t + return t + + +_HADAMARD_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def hadamard_matrix( + dim: int, device: torch.device, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Sylvester Hadamard, normalised so H @ H.T = I.""" + if dim <= 0 or (dim & (dim - 1)) != 0: + raise ValueError(f"hadamard_matrix requires power-of-two dim, got {dim}") + key = (dim, device, dtype) + cached = _HADAMARD_CACHE.get(key) + if cached is not None: + return cached + H = torch.tensor([[1.0]], dtype=torch.float64) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _HADAMARD_CACHE[key] = H + return H + + +# ── Core quantization primitives ─────────────────────────────────────────── +def _snap_to_sorted_idx( + x_norm: torch.Tensor, *, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Snap normalised values to nearest FP4 level, returning sorted index + in [0, 14]. Same semantic as `torch.bucketize` against midpoints + (reproduced by 14 `tl.where` in the kernel).""" + boundaries = _get_midpoints(x_norm.device, dtype) + x = x_norm.to(dtype) + # ROCm's bucketize uses int32 element indexing and raises + # hipErrorInvalidConfiguration once numel >= 2**31 (e.g. a [B, Hk, N, D] + # K tensor at B=128, N=16384, Hk=8, D=128 is exactly 2**31). Chunk the + # flattened input so each call stays under the limit (bit-identical output). + _INT32_LIMIT = 2**31 + if x.numel() < _INT32_LIMIT: + return torch.bucketize(x, boundaries) + flat = x.reshape(-1) + out = torch.empty_like(flat, dtype=torch.long) + chunk = _INT32_LIMIT // 2 + for i in range(0, flat.numel(), chunk): + out[i : i + chunk] = torch.bucketize(flat[i : i + chunk], boundaries) + return out.reshape(x.shape) + + +def _pack_nibbles_last_dim(codes: torch.Tensor) -> torch.Tensor: + if codes.shape[-1] % 2 != 0: + raise ValueError(f"pack_nibbles last dim must be even, got {codes.shape[-1]}") + lo = codes[..., 0::2] & 0xF + hi = codes[..., 1::2] & 0xF + return (lo | (hi << 4)).to(torch.uint8) + + +def _unpack_nibbles_last_dim(packed: torch.Tensor, out_dim: int) -> torch.Tensor: + if out_dim != 2 * packed.shape[-1]: + packed_dim = packed.shape[-1] + raise ValueError( + f"unpack_nibbles: out_dim={out_dim} must be 2*packed.last={2 * packed_dim}" + ) + p = packed.to(torch.int32) + lo = (p & 0xF).to(torch.uint8) + hi = ((p >> 4) & 0xF).to(torch.uint8) + out_shape = packed.shape[:-1] + (out_dim,) + out = torch.empty(out_shape, dtype=torch.uint8, device=packed.device) + out[..., 0::2] = lo + out[..., 1::2] = hi + return out + + +def _ue8m0_snap_tensor( + s_raw: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Snap a positive-scale tensor `s_raw` to a power of two and return + ``(s_snapped_fp32, byte_uint8)``. + + Matches the kernel: zero / non-finite inputs map to byte=0 / value=0.0. + Round-half-up via ``floor(log2(s) + 0.5)``, matching Triton's + ``tl.floor(... + 0.5)`` idiom. + """ + is_zero = (s_raw <= 0) | ~torch.isfinite(s_raw) + s_safe = torch.where(is_zero, torch.ones_like(s_raw), s_raw) + log2_s = torch.log2(s_safe) + exp = torch.floor(log2_s + 0.5).to(torch.int32) + # Clamp to UE8M0's representable exponent range (signed pre-bias). + exp = torch.clamp(exp, min=-126, max=127) + byte = (exp + UE8M0_BIAS).to(torch.int32) + byte = torch.where(is_zero, torch.zeros_like(byte), byte).to(torch.uint8) + s_snapped = torch.where( + is_zero, torch.zeros_like(s_raw), torch.exp2(exp.to(torch.float32)) + ) + return s_snapped, byte + + +@dataclass +class UltraQuantEncoded: + """Encoded ultraquant representation of one tensor along its last axis. + + Shapes assume input shape `[..., D]`: + - codes_packed: `[..., D // 2]` uint8 — 2 FP4 nibbles per byte + - scale_bytes: `[..., D // GROUP_SIZE]` uint8 — E8M0 byte per group + Dequant is ``code · 2^(byte - 127)`` (`c` is folded into the byte). + """ + + codes_packed: torch.Tensor # uint8, [..., D//2] + scale_bytes: torch.Tensor # uint8, [..., D//GROUP_SIZE] + + @property + def scales_fp32(self) -> torch.Tensor: + """Decode the E8M0 bytes back to fp32 (for dequant / inspection).""" + b = self.scale_bytes.to(torch.int32) + exp = b - UE8M0_BIAS + is_zero = b == 0 + s = torch.where( + is_zero, + torch.zeros_like(b, dtype=torch.float32), + torch.exp2(exp.to(torch.float32)), + ) + return s + + +def ultraquant_encode( + x: torch.Tensor, + *, + rotate: bool = True, + constant_c: float | None = None, +) -> UltraQuantEncoded: + """Encode `x` (shape [..., D]) to FP4 codes + E8M0 per-group-of-32 scales. + + Steps: + 1. Optional Hadamard rotation: `x_rot = x @ H.T` + 2. Reshape to `[..., G, GROUP_SIZE]`. + 3. `absmax = max(|x_rot|)` per group. + 4. `s_raw = c · absmax`; `s_snapped = 2^round(log2(s_raw))`. + Zero-amax groups encode byte=0. + 5. `sorted_idx = bucketize(x_g / s_snapped, midpoints)` ∈ [0, 14]. + 6. Remap sorted_idx → FP4 E2M1 bit pattern via SORTED_TO_BITS. + 7. Pack pairs of 4-bit codes into bytes. + """ + if x.shape[-1] % GROUP_SIZE != 0: + raise ValueError( + f"ultraquant_encode: last dim {x.shape[-1]} must be a multiple of " + f"GROUP_SIZE={GROUP_SIZE}" + ) + D = x.shape[-1] + G = D // GROUP_SIZE + c = constant_c if constant_c is not None else get_constant_c() + + x_f32 = x.to(torch.float32) + if rotate: + H = hadamard_matrix(D, x.device, torch.float32) + x_rot = x_f32 @ H.T + else: + x_rot = x_f32 + + x_g = x_rot.reshape(*x.shape[:-1], G, GROUP_SIZE) + absmax = x_g.abs().amax(dim=-1, keepdim=True) # [..., G, 1] + s_raw = absmax * c + s_snapped, scale_bytes = _ue8m0_snap_tensor(s_raw.squeeze(-1)) + s_snapped = s_snapped.unsqueeze(-1) # [..., G, 1] + s_for_div = torch.where(s_snapped == 0, torch.ones_like(s_snapped), s_snapped) + + sorted_idx = _snap_to_sorted_idx(x_g / s_for_div) + sorted_idx = torch.where( + (s_snapped == 0).expand_as(sorted_idx), + torch.full_like(sorted_idx, 7), # sorted idx 7 → +0.0 + sorted_idx, + ) + + sorted_to_bits = _get_sorted_to_bits(x.device) + fp4_bits = sorted_to_bits[sorted_idx] # [..., G, GS] + fp4_bits_flat = fp4_bits.reshape(*x.shape[:-1], D) + codes_packed = _pack_nibbles_last_dim(fp4_bits_flat) # [..., D//2] + + return UltraQuantEncoded( + codes_packed=codes_packed, + scale_bytes=scale_bytes, + ) + + +def ultraquant_dequant( + encoded: UltraQuantEncoded, + head_dim: int, +) -> torch.Tensor: + """Decode `encoded` to fp32 (no inverse rotation — result in the + rotated basis if rotation was applied at encode). + + Dequant is ``code · 2^(byte - 127)`` (`c` is already folded into the byte). + """ + G = head_dim // GROUP_SIZE + codes = _unpack_nibbles_last_dim(encoded.codes_packed, head_dim) + bits_to_val = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + device=encoded.codes_packed.device, + dtype=torch.float32, + ) + decoded = bits_to_val[codes.to(torch.long)] # [..., D] fp32 + decoded_g = decoded.reshape(*decoded.shape[:-1], G, GROUP_SIZE) + scales = encoded.scales_fp32.unsqueeze(-1) # [..., G, 1] + return (decoded_g * scales).reshape(*decoded.shape[:-1], head_dim) + + +def ultraquant_encode_decode( + x: torch.Tensor, + *, + rotate: bool = True, + constant_c: float | None = None, +) -> torch.Tensor: + """Full encode→decode round-trip in the rotated basis. Returns fp32.""" + enc = ultraquant_encode( + x, + rotate=rotate, + constant_c=constant_c, + ) + return ultraquant_dequant(enc, x.shape[-1]) + + +# ── Reference attention against the ultraquant encoded cache ────────────────── +def reference_ultraquant_attention( + query: torch.Tensor, # [B, Hq, D] bf16 — RAW (will be rotated + cast to FP8 here) + key: torch.Tensor, # [B, N, Hk, D] bf16 + value: torch.Tensor, # [B, N, Hk, D] bf16 + *, + scale: float, + sinks: torch.Tensor | None = None, + constant_c: float | None = None, +) -> torch.Tensor: + """End-to-end reference attention with ultraquant K/V encode→decode in + the rotated basis. Q is Hadamard-rotated AND cast through FP8 E4M3 + (scale = 1) before the QK matmul — mirrors what the launcher does. + + Output shape: [B, Hq, D] in query.dtype. + """ + B, Hq, D = query.shape + B2, N, Hk, D2 = key.shape + assert B == B2 and D == D2, (B, B2, D, D2) + assert value.shape == key.shape + assert Hq % Hk == 0, f"Hq={Hq} must be a multiple of Hk={Hk}" + g = Hq // Hk + + H = hadamard_matrix(D, query.device, torch.float32) + q_rot = query.to(torch.float32) @ H.T + # Round-trip Q through FP8 E4M3 with scale=1 — same precision haircut + # as the launcher (`q_rot.to(torch.float8_e4m3fn)`). + q_fp8 = q_rot.to(torch.float8_e4m3fn) + q_for_matmul = q_fp8.to(torch.float32) + + k_recon_rot = ultraquant_encode_decode( + key.transpose(1, 2).contiguous(), + rotate=True, + constant_c=constant_c, + ) + v_recon = ultraquant_encode_decode( + value.transpose(1, 2).contiguous(), + rotate=False, + constant_c=constant_c, + ) + + k_recon_rot = k_recon_rot.unsqueeze(2).expand(B, Hk, g, N, D).reshape(B, Hq, N, D) + v_recon = v_recon.unsqueeze(2).expand(B, Hk, g, N, D).reshape(B, Hq, N, D) + + qk = torch.einsum("bhd,bhnd->bhn", q_for_matmul, k_recon_rot) * scale # [B, Hq, N] + + if sinks is not None: + sink_log = sinks.to(torch.float32).reshape(1, Hq, 1).expand(B, Hq, 1) + qk_full = torch.cat([qk, sink_log], dim=-1) + p_full = torch.softmax(qk_full, dim=-1) + p = p_full[..., :N] + else: + p = torch.softmax(qk, dim=-1) + + out = torch.einsum("bhn,bhnd->bhd", p, v_recon) + return out.to(query.dtype) diff --git a/vllm/v1/attention/ops/ultraquant/triton_dequant.py b/vllm/v1/attention/ops/ultraquant/triton_dequant.py new file mode 100644 index 000000000000..9ff8d01bf40f --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_dequant.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Full KV dequant for the UltraQuant cache format. + +Used by continuation prefill: when a chunk brings many new query tokens on +top of a long cached prefix, dequanting the prefix once and running a dense +prefill kernel beats replaying the decode kernel per query token. + +Reads FP4 codes plus UE8M0 group scales and writes K (Hadamard-rotated, as +stored) and V (raw) into pre-allocated fp16/bf16 buffers. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.ultraquant.format import ( + FP4_BITS_TO_VALUE, + UE8M0_BIAS, + get_group_size, + k_scales_offset, + n_groups, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.triton_store import _kv_cache_flat + +_FP4_DECODE_CACHE: dict[tuple[torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_fp4_decode_table(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """16-entry FP4 E2M1 bit-pattern -> value table (cached per device).""" + key = (device, dtype) + t = _FP4_DECODE_CACHE.get(key) + if t is None: + t = torch.tensor(FP4_BITS_TO_VALUE, device=device, dtype=dtype).contiguous() + _FP4_DECODE_CACHE[key] = t + return t + + +@triton.jit +def _ultraquant_full_dequant_kv( + KV_cache_ptr, # uint8 view, flat + Block_table_ptr, # [B, max_num_blocks] int32 + Fp4_decode_ptr, # [16] FP4 bit pattern -> value + K_out_ptr, + V_out_ptr, + stride_ko_b: tl.int64, + stride_ko_h: tl.int64, + stride_ko_s: tl.int64, + stride_vo_b: tl.int64, + stride_vo_h: tl.int64, + stride_vo_s: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + stride_bt_b: tl.int64, + HEAD_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + BLOCK_D: tl.constexpr, + OUT_BF16: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + pos = tl.program_id(0) + bh = tl.program_id(1) + bid = bh // NUM_KV_HEADS + hid = bh % NUM_KV_HEADS + + page_idx = pos // BLOCK_SIZE + page_off = pos % BLOCK_SIZE + block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64) + slot_base = ( + block_num * stride_cache_block + + page_off * stride_cache_pos + + tl.cast(hid, tl.int64) * stride_cache_head + ) + + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + byte_idx = (d_offs // 2).to(tl.int32) + nibble_shift = ((d_offs % 2) * 4).to(tl.int32) + g_offs = tl.arange(0, N_GROUPS_C) + + out_dtype = tl.bfloat16 if OUT_BF16 else tl.float16 + + # K side: codes are packed nibbles, one UE8M0 byte per group of 32. + k_byte_raw = tl.load(KV_cache_ptr + slot_base + byte_idx, mask=d_mask, other=0).to( + tl.int32 + ) + k_codes = (k_byte_raw >> nibble_shift) & 0xF + k_dec = tl.load(Fp4_decode_ptr + k_codes).to(tl.float32) + k_dec = tl.where(d_mask, k_dec, 0.0) + + # UE8M0 scale: value = 2^(byte - bias); byte == 0 encodes exact zero. + k_scale_bytes = tl.load(KV_cache_ptr + slot_base + K_SCALES_OFFSET + g_offs).to( + tl.int32 + ) + k_scales = tl.where( + k_scale_bytes == 0, + 0.0, + tl.exp2(tl.cast(k_scale_bytes - UE8M0_BIAS_C, tl.float32)), + ) + + k_g = tl.reshape(k_dec, [N_GROUPS_C, GROUP_SIZE_C]) + k_recon = tl.reshape(k_g * k_scales[:, None], [BLOCK_D]) + k_recon = tl.where(d_mask, k_recon, 0.0) + + ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(out_dtype), mask=d_mask) + + # V side mirrors K but is stored unrotated. + v_byte_raw = tl.load( + KV_cache_ptr + slot_base + V_CODES_OFFSET + byte_idx, + mask=d_mask, + other=0, + ).to(tl.int32) + v_codes = (v_byte_raw >> nibble_shift) & 0xF + v_dec = tl.load(Fp4_decode_ptr + v_codes).to(tl.float32) + v_dec = tl.where(d_mask, v_dec, 0.0) + + v_scale_bytes = tl.load(KV_cache_ptr + slot_base + V_SCALES_OFFSET + g_offs).to( + tl.int32 + ) + v_scales = tl.where( + v_scale_bytes == 0, + 0.0, + tl.exp2(tl.cast(v_scale_bytes - UE8M0_BIAS_C, tl.float32)), + ) + + v_g = tl.reshape(v_dec, [N_GROUPS_C, GROUP_SIZE_C]) + v_recon = tl.reshape(v_g * v_scales[:, None], [BLOCK_D]) + v_recon = tl.where(d_mask, v_recon, 0.0) + + vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s + tl.store(V_out_ptr + vo_base + d_offs, v_recon.to(out_dtype), mask=d_mask) + + +def ultraquant_full_dequant_kv( + kv_cache: torch.Tensor, + block_table: torch.Tensor, + k_out: torch.Tensor, + v_out: torch.Tensor, + alloc_len: int, +) -> None: + """Dequant ``alloc_len`` cached positions into ``k_out`` / ``v_out``. + + ``k_out`` / ``v_out`` are ``[B, Hk, alloc_len, D]`` in fp16 or bf16. + """ + B = block_table.shape[0] + Hk = kv_cache.shape[2] + D = k_out.shape[3] + block_size = kv_cache.shape[1] + + gs = get_group_size() + _ultraquant_full_dequant_kv[(alloc_len, B * Hk)]( + _kv_cache_flat(kv_cache), + block_table, + _get_fp4_decode_table(kv_cache.device, k_out.dtype), + k_out, + v_out, + k_out.stride(0), + k_out.stride(1), + k_out.stride(2), + v_out.stride(0), + v_out.stride(1), + v_out.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + block_table.stride(0), + HEAD_DIM=D, + BLOCK_SIZE=block_size, + NUM_KV_HEADS=Hk, + K_SCALES_OFFSET=k_scales_offset(D, gs), + V_CODES_OFFSET=v_codes_offset(D, gs), + V_SCALES_OFFSET=v_scales_offset(D, gs), + GROUP_SIZE_C=gs, + N_GROUPS_C=n_groups(D, gs), + BLOCK_D=triton.next_power_of_2(D), + OUT_BF16=1 if k_out.dtype == torch.bfloat16 else 0, + UE8M0_BIAS_C=UE8M0_BIAS, + num_warps=4, + ) diff --git a/vllm/v1/attention/ops/ultraquant/triton_store.py b/vllm/v1/attention/ops/ultraquant/triton_store.py new file mode 100644 index 000000000000..b43f4264e7d2 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_store.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton store kernel for the ultraquant KV cache format. + +Per-token, per-head, per-group-of-32 the kernel computes: + + s_raw = c * absmax(group) # c = 0.156 by default + s = 2^round(log2(s_raw)) # UE8M0 snap (power of 2) + code = quantize_to_fp4(group / s) # OCP E2M1, 4 bits/elt + +and stores ``code`` as packed nibbles + ``s`` as a single UE8M0 byte per +group. The output layout matches ``format.py`` (272 B/slot for D=256). + +K is Hadamard-rotated in-register; V is not. There is no per-token L2 +norm. The UE8M0 snap makes ``s`` a power of two so scaled F8F6F4 MFMA +can consume it on the read side. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.ultraquant.format import ( + MIDPOINTS_SORTED, + SORTED_TO_BITS, + UE8M0_BIAS, + get_constant_c, + get_group_size, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) + + +def _kv_cache_flat(kv_cache: torch.Tensor) -> torch.Tensor: + """1-D uint8 view of the raw KV allocation (handles padded layouts).""" + if kv_cache.is_contiguous(): + return kv_cache.reshape(-1) + n = kv_cache.untyped_storage().nbytes() // kv_cache.element_size() + return torch.empty(0, dtype=kv_cache.dtype, device=kv_cache.device).set_( + kv_cache.untyped_storage(), 0, (n,) + ) + + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16} + +# Built lazily on first launch so importing this module does not require triton. +_MIDPOINTS_C: Any = None + + +def _ensure_const_tables() -> None: + global _MIDPOINTS_C + if _MIDPOINTS_C is None: + _MIDPOINTS_C = tl.constexpr(tuple(float(m) for m in MIDPOINTS_SORTED)) + + +# ── Cached per-device constant tensors ───────────────────────────────────── +_MIDPOINTS_CACHE: dict[torch.device, torch.Tensor] = {} +_SORTED_TO_BITS_CACHE: dict[torch.device, torch.Tensor] = {} +_PIT_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_midpoints_tensor(device: torch.device) -> torch.Tensor: + t = _MIDPOINTS_CACHE.get(device) + if t is None: + t = torch.tensor(MIDPOINTS_SORTED, device=device, dtype=torch.float32) + _MIDPOINTS_CACHE[device] = t + return t + + +def _get_sorted_to_bits_tensor(device: torch.device) -> torch.Tensor: + t = _SORTED_TO_BITS_CACHE.get(device) + if t is None: + t = torch.tensor(SORTED_TO_BITS, device=device, dtype=torch.int32) + _SORTED_TO_BITS_CACHE[device] = t + return t + + +def _get_hadamard( + dim: int, device: torch.device, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """Sylvester Hadamard in fp32 (cached, normalised).""" + key = (dim, device, dtype) + cached = _PIT_CACHE.get(key) + if cached is not None: + return cached + if dim <= 0 or (dim & (dim - 1)) != 0: + raise ValueError(f"ultraquant store requires power-of-two dim, got {dim}") + H = torch.tensor([[1.0]], dtype=torch.float64) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _PIT_CACHE[key] = H + return H + + +# ═══════════════════════════════════════════════════════════════════════════ +# Triton kernel +# ═══════════════════════════════════════════════════════════════════════════ + + +@triton.jit +def _ultraquant_store_kernel( + Key_ptr, # [N*H, D] in raw dtype (bf16 or fp16) + Value_ptr, # [N*H, D] in raw dtype + KV_cache_ptr, # raw uint8 view, flat + Slot_mapping_ptr, # [N] int64 + Midpoints_ptr, # [14] fp32 + SortedToBits_ptr, # [15] int32 + # Cache strides (in bytes) + stride_cache_block: tl.constexpr, + stride_cache_pos: tl.constexpr, + stride_cache_head: tl.constexpr, + # Dimensions + HEAD_DIM: tl.constexpr, + H: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + LOG2_D: tl.constexpr, + # Layout + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + FP4_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + # When 1, the FP4 midpoints and the sorted->bits remap are materialised as + # immediates instead of being read from `Midpoints_ptr`/`SortedToBits_ptr`. + # That removes 28 scalar global loads and two 256-lane divergent gathers + # from the critical path of a kernel whose runtime is pure latency (flat at + # ~22 us from batch 1 to batch 64). Numerically identical: the immediates + # are the same fp32 values, and the remap formula is exact (asserted in + # `format`, verified for all 15 entries). + CONST_TABLES: tl.constexpr = 1, +): + pid = tl.program_id(0) + token_idx = pid // H + head_idx = pid % H + + slot = tl.load(Slot_mapping_ptr + token_idx) + if slot < 0: + return + blk = slot // BLOCK_SIZE + off = slot % BLOCK_SIZE + slot_base = ( + blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head + ).to(tl.int64) + k_code_base = slot_base + k_scale_base = slot_base + K_SCALES_OFFSET + v_code_base = slot_base + V_CODES_OFFSET + v_scale_base = slot_base + V_SCALES_OFFSET + + base = pid * HEAD_DIM + d_offs = tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + shifts_4 = tl.arange(0, 2) * 4 # 0, 4 for nibble packing + g_offs = tl.arange(0, N_GROUPS_C) + + # ── K side ───────────────────────────────────────────────────────── + # 1. Load + Hadamard-rotate K via WHT butterfly. + k_raw = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) + k_rot = tl.where(d_mask, k_raw, 0.0) + + INV_SQRT2: tl.constexpr = 0.7071067811865476 + for _stage in tl.static_range(LOG2_D): + _bmask = 1 << _stage + _pidx = (d_offs ^ _bmask).to(tl.int32) + _pval = tl.gather(k_rot, _pidx, 0) + _sign = tl.where((d_offs & _bmask) == 0, 1.0, -1.0).to(tl.float32) + k_rot = (_sign * k_rot + _pval) * INV_SQRT2 + k_rot = tl.where(d_mask, k_rot, 0.0) + + # 2. Per-group absmax. + k_g = tl.reshape(k_rot, [N_GROUPS_C, GROUP_SIZE_C]) + k_absmax = tl.max(tl.abs(k_g), axis=1) # [N_GROUPS] + k_is_zero = k_absmax == 0.0 + + # 3. UE8M0 snap: s_raw = c * absmax; exp = round(log2(s_raw)); + # divisor = 2^exp (c folded in). Zero groups write byte=0. + k_s_raw = k_absmax * tl.cast(FP4_C, tl.float32) + k_s_safe = tl.where(k_is_zero, 1.0, k_s_raw) + k_log2 = tl.log2(k_s_safe) + k_exp = tl.cast(tl.floor(k_log2 + 0.5), tl.int32) # round-half-up + k_s_snapped = tl.exp2(tl.cast(k_exp, tl.float32)) + k_s_div = tl.where(k_is_zero, 1.0, k_s_snapped) + + # 4. Normalize + snap to sorted FP4 idx via 14 sequential tl.where (`>` for + # bucketize(right=False) parity with the reference). + k_norm = k_g / k_s_div[:, None] + k_sorted = tl.zeros([N_GROUPS_C, GROUP_SIZE_C], dtype=tl.int32) + for i in tl.static_range(14): + mid = _MIDPOINTS_C.value[i] if CONST_TABLES else tl.load(Midpoints_ptr + i) + k_sorted += tl.where(k_norm > mid, 1, 0) + k_sorted = tl.where(k_is_zero[:, None], 7, k_sorted) + + # 5. Remap sorted idx → FP4 E2M1 bit pattern. + if CONST_TABLES: + k_bits = tl.where(k_sorted < 7, 15 - k_sorted, k_sorted - 7) + else: + k_bits = tl.load(SortedToBits_ptr + k_sorted) + + # 6. Pack pairs of 4-bit codes into uint8 bytes. + k_bits_flat = tl.reshape(k_bits, [HEAD_DIM]) + k_pairs = tl.reshape(k_bits_flat, [HEAD_DIM // 2, 2]) + k_packed = tl.sum((k_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + k_code_addrs = k_code_base + tl.arange(0, HEAD_DIM // 2) + tl.store(KV_cache_ptr + k_code_addrs, k_packed) + + # 7. Store K scales as E8M0 bytes (one per group). Zero-amax groups + # write byte=0 (zero sentinel). + k_byte = tl.where(k_is_zero, 0, k_exp + UE8M0_BIAS_C) + k_byte = tl.where((k_byte < 0) | (k_byte > 255), 0, k_byte).to(tl.uint8) + k_scale_addrs = k_scale_base + g_offs + tl.store(KV_cache_ptr + k_scale_addrs, k_byte) + + # ── V side (mirror K but skip rotation) ──────────────────────────── + v_raw = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32) + v_raw = tl.where(d_mask, v_raw, 0.0) + + v_g = tl.reshape(v_raw, [N_GROUPS_C, GROUP_SIZE_C]) + v_absmax = tl.max(tl.abs(v_g), axis=1) + v_is_zero = v_absmax == 0.0 + + v_s_raw = v_absmax * tl.cast(FP4_C, tl.float32) + v_s_safe = tl.where(v_is_zero, 1.0, v_s_raw) + v_log2 = tl.log2(v_s_safe) + v_exp = tl.cast(tl.floor(v_log2 + 0.5), tl.int32) + v_s_snapped = tl.exp2(tl.cast(v_exp, tl.float32)) + v_s_div = tl.where(v_is_zero, 1.0, v_s_snapped) + + v_norm = v_g / v_s_div[:, None] + v_sorted = tl.zeros([N_GROUPS_C, GROUP_SIZE_C], dtype=tl.int32) + for i in tl.static_range(14): + mid = _MIDPOINTS_C.value[i] if CONST_TABLES else tl.load(Midpoints_ptr + i) + v_sorted += tl.where(v_norm > mid, 1, 0) + v_sorted = tl.where(v_is_zero[:, None], 7, v_sorted) + + if CONST_TABLES: + v_bits = tl.where(v_sorted < 7, 15 - v_sorted, v_sorted - 7) + else: + v_bits = tl.load(SortedToBits_ptr + v_sorted) + v_bits_flat = tl.reshape(v_bits, [HEAD_DIM]) + v_pairs = tl.reshape(v_bits_flat, [HEAD_DIM // 2, 2]) + v_packed = tl.sum((v_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8) + v_code_addrs = v_code_base + tl.arange(0, HEAD_DIM // 2) + tl.store(KV_cache_ptr + v_code_addrs, v_packed) + + v_byte = tl.where(v_is_zero, 0, v_exp + UE8M0_BIAS_C) + v_byte = tl.where((v_byte < 0) | (v_byte > 255), 0, v_byte).to(tl.uint8) + v_scale_addrs = v_scale_base + g_offs + tl.store(KV_cache_ptr + v_scale_addrs, v_byte) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Launcher +# ═══════════════════════════════════════════════════════════════════════════ + + +def ultraquant_store( + key: torch.Tensor, # [N, H, D] bf16 or fp16 + value: torch.Tensor, # [N, H, D] bf16 or fp16 + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, slot_size_aligned] uint8 + slot_mapping: torch.Tensor, # [N] int64 + *, + PiT: torch.Tensor | None = None, # unused; kept for call-site compatibility + constant_c: float | None = None, +) -> None: + """Launch the ultraquant store kernel. Writes K/V into `kv_cache` + in-place at the slots specified by `slot_mapping`.""" + if key.dtype not in _SUPPORTED_DTYPES: + raise ValueError( + f"ultraquant_store: key.dtype must be one of " + f"{_SUPPORTED_DTYPES}, got {key.dtype}" + ) + if value.dtype != key.dtype: + raise ValueError( + f"ultraquant_store: key/value dtype mismatch ({key.dtype} vs {value.dtype})" + ) + if slot_mapping.dtype != torch.int64: + slot_mapping = slot_mapping.to(torch.int64) + + N, H, D = key.shape + NH = N * H + block_size = kv_cache.shape[1] + num_kv_heads = kv_cache.shape[2] + padded_slot = kv_cache.shape[3] + + expected_slot = slot_size(D) + if padded_slot < expected_slot: + raise ValueError( + f"ultraquant_store: kv_cache slot {padded_slot} < expected " + f"{expected_slot} for head_dim={D}" + ) + if num_kv_heads != H: + raise ValueError( + f"ultraquant_store: kv_cache num_kv_heads {num_kv_heads} != key heads {H}" + ) + + midpoints = _get_midpoints_tensor(key.device) + sorted_to_bits = _get_sorted_to_bits_tensor(key.device) + + k_flat = key.reshape(NH, D).contiguous() + v_flat = value.reshape(NH, D).contiguous() + + gs = get_group_size() + BLOCK_D = triton.next_power_of_2(D) + N_GROUPS_C = n_groups(D, gs) + K_SCALES_OFFSET = k_scales_offset(D, gs) + V_CODES_OFFSET = v_codes_offset(D, gs) + V_SCALES_OFFSET = v_scales_offset(D, gs) + + stride_block = kv_cache.stride(0) + stride_pos = kv_cache.stride(1) + stride_head = kv_cache.stride(2) + + c = constant_c if constant_c is not None else get_constant_c() + + _ensure_const_tables() + grid = (NH,) + _ultraquant_store_kernel[grid]( + k_flat, + v_flat, + _kv_cache_flat(kv_cache), + slot_mapping, + midpoints, + sorted_to_bits, + stride_cache_block=stride_block, + stride_cache_pos=stride_pos, + stride_cache_head=stride_head, + HEAD_DIM=D, + H=H, + BLOCK_SIZE=block_size, + BLOCK_D=BLOCK_D, + LOG2_D=int(D).bit_length() - 1, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + FP4_C=c, + UE8M0_BIAS_C=UE8M0_BIAS, + CONST_TABLES=1, + num_warps=4, + num_stages=1, + ) diff --git a/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py b/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py new file mode 100644 index 000000000000..69ec744a07f9 --- /dev/null +++ b/vllm/v1/attention/ops/ultraquant/triton_unified_attention.py @@ -0,0 +1,1087 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unified Triton fallback for the UltraQuant 4-bit KV-cache format. + +K and V use packed FP4 E2M1 values with UE8M0 group-of-32 scales. QK uses +native scaled FP4×E4M3 MFMA on CDNA4. Production choices are fixed in code; +there are no environment-variable tuning switches. +""" + +from __future__ import annotations + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_unified_attention import ( + reduce_segments, +) + +# Reuse the sequence-index search and vectorized partition reducer. +from vllm.v1.attention.ops.turboquant_soa.triton_turboquant_unified_attention import ( + _find_seq_idx, +) +from vllm.v1.attention.ops.ultraquant.format import ( + UE8M0_BIAS, + get_group_size, + k_scales_offset, + n_groups, + slot_size, + v_codes_offset, + v_scales_offset, +) +from vllm.v1.attention.ops.ultraquant.triton_store import _kv_cache_flat + +_is_hip = current_platform.is_rocm() + + +@triton.jit +def _ultraquant_fp4_decode_arith(codes): + """Arithmetic FP4 E2M1 decode, bit-exact to format.FP4_BITS_TO_VALUE.""" + mag = codes & 7 + e = mag >> 1 + m = mag & 1 + exp_field = 126 + e + mant_field = tl.where(e != 0, m, 0) << 22 + bits = (exp_field << 23) | mant_field + magval = bits.to(tl.float32, bitcast=True) + magval = tl.where(mag == 0, 0.0, magval) + signf = tl.where((codes & 8) != 0, -1.0, 1.0) + return magval * signf + + +@triton.jit +def _ultraquant_load_k_packed( + KV_cache_ptr, + data_bases, + k_scales_addrs, + d_half_offs, + half_mask, + tile_mask, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, +): + """Load packed FP4 K codes and UE8M0 scales for scaled QK MFMA.""" + addrs = data_bases[:, None] + d_half_offs[None, :] + if UNMASKED: + K_codes = tl.load(KV_cache_ptr + addrs, mask=half_mask[None, :], other=0) + else: + K_codes = tl.load( + KV_cache_ptr + addrs, + mask=tile_mask[:, None] & half_mask[None, :], + other=0, + ) + K_T_packed = tl.trans(K_codes) + grp = tl.arange(0, N_GROUPS_C) + scale_addrs = k_scales_addrs[:, None] + grp[None, :] + if UNMASKED: + K_scales = tl.load(KV_cache_ptr + scale_addrs) + else: + K_scales = tl.load(KV_cache_ptr + scale_addrs, mask=tile_mask[:, None], other=0) + _ = HEAD_DIM + _ = GROUP_SIZE_C + _ = BLOCK_D + return K_T_packed, K_scales + + +@triton.jit +def _ultraquant_load_v_tile( + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + OUT_DTYPE: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + """Load and dequant a V tile with arithmetic FP4 decode.""" + half_idx = d_offs // 2 + nibble_shift = (d_offs % 2) * 4 + addrs = val_bases[:, None] + half_idx[None, :] + if UNMASKED: + byte_raw = tl.load(KV_cache_ptr + addrs, mask=d_mask[None, :], other=0).to( + tl.int32 + ) + else: + byte_raw = tl.load( + KV_cache_ptr + addrs, + mask=tile_mask[:, None] & d_mask[None, :], + other=0, + ).to(tl.int32) + codes = (byte_raw >> nibble_shift[None, :]) & 0xF + grp = tl.arange(0, N_GROUPS_C) + scale_addrs = v_scales_addrs[:, None] + grp[None, :] + if UNMASKED: + scale_bytes = tl.load(KV_cache_ptr + scale_addrs).to(tl.int32) + else: + scale_bytes = tl.load( + KV_cache_ptr + scale_addrs, mask=tile_mask[:, None], other=0 + ).to(tl.int32) + scale_exp = scale_bytes - UE8M0_BIAS_C + fp4_vals = _ultraquant_fp4_decode_arith(codes) + scales = tl.where(scale_bytes == 0, 0.0, tl.exp2(tl.cast(scale_exp, tl.float32))) + V_g = tl.reshape(fp4_vals, [TILE_SIZE, N_GROUPS_C, GROUP_SIZE_C]) + return tl.reshape(V_g * scales[:, :, None], [TILE_SIZE, BLOCK_D]).to(OUT_DTYPE) + + +@triton.jit +def _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_M: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, +): + V = _ultraquant_load_v_tile( + KV_cache_ptr, + val_bases, + v_scales_addrs, + d_offs, + d_mask, + tile_mask, + OUT_DTYPE=tl.bfloat16, + HEAD_DIM=HEAD_DIM, + BLOCK_D=BLOCK_D, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=UNMASKED, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + _ = BLOCK_M + return tl.dot(P.to(tl.bfloat16), V, out_dtype=tl.float32) + + +@triton.jit +def _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D: tl.constexpr, + HEAD_DIM: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + TILE_SIZE: tl.constexpr, + UNMASKED: tl.constexpr, +): + K_T_packed, K_scales = _ultraquant_load_k_packed( + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=BLOCK_D, + HEAD_DIM=HEAD_DIM, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=UNMASKED, + ) + return tl.dot_scaled( + Q, + None, + "e4m3", + K_T_packed, + K_scales, + "e2m1", + out_dtype=tl.float32, + ) + + +@triton.jit +def _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, +): + data_bases = ( + block_base + + slot_within_block * stride_cache_pos + + tl.cast(kv_head_idx, tl.int64) * stride_cache_head + ) + return ( + data_bases, + data_bases + K_SCALES_OFFSET, + data_bases + V_CODES_OFFSET, + data_bases + V_SCALES_OFFSET, + ) + + +# =========================================================================== +# Unified 2D attention kernel (prefill / short-context decode) +# =========================================================================== + + +@triton.jit +def kernel_ultraquant_unified_attention_2d( + output_ptr, + query_ptr, # FP8 E4M3, already Hadamard-rotated + KV_cache_ptr, # uint8 view + block_tables_ptr, + seq_lens_ptr, + query_start_len_ptr, + sinks_ptr, + scale, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, + output_stride_0: tl.int64, + output_stride_1: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_M: tl.constexpr, + num_seqs: tl.int32, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + USE_SINKS: tl.constexpr = 0, + SLIDING_WINDOW: tl.constexpr = 0, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + + seq_idx = _find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_d_half = tl.arange(0, HEAD_SIZE_PADDED // 2) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + half_mask = tl.where(offs_d_half * 2 < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=tl.zeros([], tl.float8e4nv), + ) + + block_table_offset = seq_idx * block_table_stride + + if USE_SINKS: + M = tl.load( + sinks_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + seq_len = tl.load(seq_lens_ptr + seq_idx) + context_len = seq_len - cur_batch_query_len + + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE) + + tile_start = 0 + tile_end = num_tiles + if SLIDING_WINDOW > 0: + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + + query_abs_pos = context_len + query_pos[:, None] + dummy_tile_mask = tl.full([TILE_SIZE], 1, tl.int1) + + for j in range(tile_start, tile_end - 1): + seq_offset = j * TILE_SIZE + offs_t + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + dummy_tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + dummy_tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Tail tile (masked) + if tile_end > tile_start: + j = tile_end - 1 + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + acc = acc / L[:, None] + + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d[None, :] + ) + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +# =========================================================================== +# Unified 3D (split-KV) attention kernel +# =========================================================================== + + +@triton.jit +def kernel_ultraquant_unified_attention_3d( + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + query_ptr, # FP8 E4M3, already Hadamard-rotated + KV_cache_ptr, + block_tables_ptr, + seq_lens_ptr, + query_start_len_ptr, + sinks_ptr, + scale, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, + stride_cache_block: tl.int64, + stride_cache_pos: tl.int64, + stride_cache_head: tl.int64, + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_M: tl.constexpr, + num_seqs: tl.int32, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + K_SCALES_OFFSET: tl.constexpr, + V_CODES_OFFSET: tl.constexpr, + V_SCALES_OFFSET: tl.constexpr, + GROUP_SIZE_C: tl.constexpr, + N_GROUPS_C: tl.constexpr, + UE8M0_BIAS_C: tl.constexpr, + USE_SINKS: tl.constexpr = 0, + SLIDING_WINDOW: tl.constexpr = 0, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) + + seq_idx = _find_seq_idx( + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True + ) + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx + q_block_local_idx = q_block_global_idx - q_block_start_idx + + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx) + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1) + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + seq_len = tl.load(seq_lens_ptr + seq_idx) + tiles_per_segment = tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_SIZE_PADDED) + offs_d_half = tl.arange(0, HEAD_SIZE_PADDED // 2) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d[None, :] + ) + + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1) + half_mask = tl.where(offs_d_half * 2 < HEAD_SIZE, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=tl.zeros([], tl.float8e4nv), + ) + + block_table_offset = seq_idx * block_table_stride + + if USE_SINKS and segm_idx == 0: + M = tl.load( + sinks_ptr + query_offset_1, + mask=query_mask_1, + other=float("-inf"), + ).to(tl.float32) + else: + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32) + + context_len = seq_len - cur_batch_query_len + max_seq_prefix_len = ( + context_len + + q_block_local_idx * BLOCK_Q + + (BLOCK_M - 1) // num_queries_per_kv + + 1 + ) + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE) + + tile_lo = segm_idx * tiles_per_segment + tile_hi = tl.minimum((segm_idx + 1) * tiles_per_segment, num_tiles) + + if SLIDING_WINDOW > 0: + qpos_lo = q_block_local_idx * BLOCK_Q + qpos_hi = tl.minimum( + qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, + cur_batch_query_len - 1, + ) + first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1 + last_allowed_key = context_len + qpos_hi + swa_tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) + swa_tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) + tile_lo = tl.maximum(tile_lo, swa_tile_start) + tile_hi = tl.minimum(tile_hi, swa_tile_end) + + query_abs_pos = context_len + query_pos[:, None] + dummy_tile_mask = tl.full([TILE_SIZE], 1, tl.int1) + tail = tile_hi - 1 + + for j in range(tile_lo, tile_hi - 1): + seq_offset = j * TILE_SIZE + offs_t + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + dummy_tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + dummy_tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=True, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Tail + if tile_lo < tile_hi: + j = tail + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + if TILE_SIZE == BLOCK_SIZE: + physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j).to( + tl.int64 + ) + slot_within_block = offs_t.to(tl.int64) + else: + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64) + block_base = physical_block_idx * stride_cache_block + data_bases, k_scales_addrs, val_bases, v_scales_addrs = _ultraquant_slot_addrs( + block_base, + slot_within_block, + kv_head_idx, + stride_cache_pos, + stride_cache_head, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + ) + + S = scale * _ultraquant_qk( + Q, + KV_cache_ptr, + data_bases, + k_scales_addrs, + offs_d_half, + half_mask, + tile_mask, + BLOCK_D=HEAD_SIZE_PADDED, + HEAD_DIM=HEAD_SIZE, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + ) + seq_mask = seq_offset[None, :] <= query_abs_pos + if SLIDING_WINDOW > 0: + seq_mask = seq_mask & ( + (query_abs_pos - seq_offset[None, :]) < SLIDING_WINDOW + ) + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, + S, + float("-inf"), + ) + m_j = tl.maximum(M, tl.max(S, axis=1)) + m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + P = tl.exp(S - m_j[:, None]) + l_j = tl.sum(P, axis=1) + alpha = tl.exp(M - m_j) + acc = acc * alpha[:, None] + L = L * alpha + l_j + M = m_j + acc += _ultraquant_pv( + P, + KV_cache_ptr, + val_bases, + v_scales_addrs, + offs_d, + dim_mask, + tile_mask, + HEAD_DIM=HEAD_SIZE, + BLOCK_D=HEAD_SIZE_PADDED, + BLOCK_M=BLOCK_M, + GROUP_SIZE_C=GROUP_SIZE_C, + N_GROUPS_C=N_GROUPS_C, + TILE_SIZE=TILE_SIZE, + UNMASKED=False, + UE8M0_BIAS_C=UE8M0_BIAS_C, + ) + + # Write segment partials for stage-2 reduce. + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) + + segm_idx * HEAD_SIZE_PADDED + + tl.arange(0, HEAD_SIZE_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset, + acc, + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + segm_offset = ( + query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_offset_1 * NUM_SEGMENTS_PER_SEQ + + segm_idx + ) + tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1) + tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1) + + +# =========================================================================== +# Launcher +# =========================================================================== + + +_HADAMARD_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.Tensor] = {} + + +def _get_pit(dim: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Sylvester Hadamard projection / sqrt(dim). PiT = PiT.T (symmetric).""" + key = (dim, device, dtype) + H = _HADAMARD_CACHE.get(key) + if H is None: + assert (dim & (dim - 1)) == 0, f"dim={dim} must be power of 2" + H = torch.tensor([[1.0]], dtype=torch.float32) + while H.shape[0] < dim: + H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0) + H = (H / (dim**0.5)).to(device=device, dtype=dtype).contiguous() + _HADAMARD_CACHE[key] = H + return H + + +def ultraquant_unified_attention( + query: torch.Tensor, # [num_tokens, Hq, D] fp16/bf16 (raw) + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 + block_table: torch.Tensor, # [num_seqs, max_num_blocks] int32 + seq_lens: torch.Tensor, # [num_seqs] int32 + query_start_loc: torch.Tensor, # [num_seqs+1] int32 + scale: float, + PiT: torch.Tensor | None = None, + output: torch.Tensor | None = None, + tile_size: int | None = None, + max_query_len: int | None = None, + max_seq_len: int | None = None, + num_kv_splits: int | None = None, + force_2d: bool = False, + sinks: torch.Tensor | None = None, + sliding_window: int | None = None, +) -> torch.Tensor: + """Launch unified UltraQuant attention with scaled F8F6F4 MFMA.""" + assert query.dim() == 3, f"query must be [N, Hq, D], got {query.shape}" + num_tokens, Hq, D = query.shape + Hk = kv_cache.shape[2] + block_size = kv_cache.shape[1] + padded_slot = kv_cache.shape[3] + kv_group_size = Hq // Hk + num_seqs = int(query_start_loc.shape[0] - 1) + device = query.device + + gs = get_group_size() + if padded_slot < slot_size(D, gs): + raise ValueError( + f"ultraquant_unified: cache slot {padded_slot} < expected " + f"{slot_size(D, gs)} for D={D} group_size={gs}" + ) + if Hq % Hk != 0: + raise ValueError(f"Hq={Hq} must be a multiple of Hk={Hk}") + + if PiT is None: + PiT = _get_pit(D, device, torch.float32) + elif PiT.dtype != torch.float32 or not PiT.is_contiguous(): + PiT = PiT.to(torch.float32).contiguous() + + q_rot = (query.float() @ PiT).contiguous() + q_for_kernel = q_rot.to(torch.float8_e4m3fn).contiguous() + + if sinks is not None: + sinks_f32 = sinks if sinks.dtype == torch.float32 else sinks.to(torch.float32) + if not sinks_f32.is_contiguous(): + sinks_f32 = sinks_f32.contiguous() + assert sinks_f32.numel() == Hq, ( + f"sinks must have shape [Hq={Hq}], got numel={sinks_f32.numel()}" + ) + use_sinks = True + else: + sinks_f32 = q_for_kernel # harmless dummy; never dereferenced when USE_SINKS=0 + use_sinks = False + + if output is None: + output = torch.empty_like(query) + + # BLOCK_M heuristic shared with the compressed-KV fallback kernels. + if max_query_len is not None: + is_prefill_like = max_query_len > 1 + else: + is_prefill_like = num_tokens > num_seqs + + if is_prefill_like: + BLOCK_M = max(128, triton.next_power_of_2(kv_group_size)) + else: + BLOCK_M = 16 if kv_group_size <= 16 else triton.next_power_of_2(kv_group_size) + BLOCK_Q = BLOCK_M // kv_group_size + + total_num_q_blocks = num_tokens // BLOCK_Q + num_seqs + + if tile_size is None: + tile_size = 32 if is_prefill_like else 16 + + num_stages_2d = 1 if _is_hip else 2 + num_stages_3d = 3 if _is_hip else 2 + + BLOCK_D = triton.next_power_of_2(D) + N_GROUPS_C = n_groups(D, gs) + K_SCALES_OFFSET = k_scales_offset(D, gs) + V_CODES_OFFSET = v_codes_offset(D, gs) + V_SCALES_OFFSET = v_scales_offset(D, gs) + + scale_for_kernel = float(scale) + + kv_flat = _kv_cache_flat(kv_cache) + + # Dispatch: 2D for prefill / chunked; 3D for pure decode with long KV. + if max_seq_len is None: + max_seq_len_hint = int(block_table.shape[1]) * int(block_size) + else: + max_seq_len_hint = int(max_seq_len) + use_3d = (not force_2d) and (not is_prefill_like) and max_seq_len_hint >= 1024 + + if not use_3d: + kernel_ultraquant_unified_attention_2d[(total_num_q_blocks, Hk)]( + output_ptr=output, + query_ptr=q_for_kernel, + KV_cache_ptr=kv_flat, + block_tables_ptr=block_table, + seq_lens_ptr=seq_lens, + query_start_len_ptr=query_start_loc, + sinks_ptr=sinks_f32, + scale=scale_for_kernel, + num_query_heads=Hq, + num_queries_per_kv=kv_group_size, + block_table_stride=block_table.stride(0), + query_stride_0=q_for_kernel.stride(0), + query_stride_1=q_for_kernel.stride(1), + output_stride_0=output.stride(0), + output_stride_1=output.stride(1), + stride_cache_block=kv_cache.stride(0), + stride_cache_pos=kv_cache.stride(1), + stride_cache_head=kv_cache.stride(2), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + BLOCK_Q=BLOCK_Q, + BLOCK_M=BLOCK_M, + num_seqs=num_seqs, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + UE8M0_BIAS_C=UE8M0_BIAS, + USE_SINKS=1 if use_sinks else 0, + SLIDING_WINDOW=int(sliding_window) + if sliding_window and sliding_window > 0 + else 0, + num_warps=4, + num_stages=num_stages_2d, + ) + return output + + # 3D split-KV path + if num_kv_splits is None: + num_kv_splits = 16 + if num_kv_splits < 1: + num_kv_splits = 1 + if num_kv_splits & (num_kv_splits - 1) != 0: + num_kv_splits = 1 << (num_kv_splits.bit_length() - 1) + max_possible_splits = max(1, (max_seq_len_hint + tile_size - 1) // tile_size) + num_segments = max(1, min(num_kv_splits, max_possible_splits)) + + segm_output = torch.empty( + (num_tokens, Hq, num_segments, BLOCK_D), + dtype=torch.float32, + device=device, + ) + segm_max = torch.empty( + (num_tokens, Hq, num_segments), + dtype=torch.float32, + device=device, + ) + segm_expsum = torch.empty( + (num_tokens, Hq, num_segments), + dtype=torch.float32, + device=device, + ) + + kernel_ultraquant_unified_attention_3d[(total_num_q_blocks, Hk, num_segments)]( + segm_output_ptr=segm_output, + segm_max_ptr=segm_max, + segm_expsum_ptr=segm_expsum, + query_ptr=q_for_kernel, + KV_cache_ptr=kv_flat, + block_tables_ptr=block_table, + seq_lens_ptr=seq_lens, + query_start_len_ptr=query_start_loc, + sinks_ptr=sinks_f32, + scale=scale_for_kernel, + num_query_heads=Hq, + num_queries_per_kv=kv_group_size, + block_table_stride=block_table.stride(0), + query_stride_0=q_for_kernel.stride(0), + query_stride_1=q_for_kernel.stride(1), + stride_cache_block=kv_cache.stride(0), + stride_cache_pos=kv_cache.stride(1), + stride_cache_head=kv_cache.stride(2), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + BLOCK_Q=BLOCK_Q, + BLOCK_M=BLOCK_M, + num_seqs=num_seqs, + NUM_SEGMENTS_PER_SEQ=num_segments, + K_SCALES_OFFSET=K_SCALES_OFFSET, + V_CODES_OFFSET=V_CODES_OFFSET, + V_SCALES_OFFSET=V_SCALES_OFFSET, + GROUP_SIZE_C=gs, + N_GROUPS_C=N_GROUPS_C, + UE8M0_BIAS_C=UE8M0_BIAS, + USE_SINKS=1 if use_sinks else 0, + SLIDING_WINDOW=int(sliding_window) + if sliding_window and sliding_window > 0 + else 0, + num_warps=2, + num_stages=num_stages_3d, + ) + + # Reduce split-KV partials with the shared vectorized reducer. + reduce_segments[(num_tokens, Hq)]( + output_ptr=output, + segm_output_ptr=segm_output, + segm_max_ptr=segm_max, + segm_expsum_ptr=segm_expsum, + seq_lens_ptr=seq_lens, + num_seqs=num_seqs, + num_query_heads=Hq, + out_scale_inv=1.0, + output_stride_0=output.stride(0), + output_stride_1=output.stride(1), + block_table_stride=block_table.stride(0), + TILE_SIZE=tile_size, + HEAD_SIZE=D, + HEAD_SIZE_PADDED=BLOCK_D, + query_start_len_ptr=query_start_loc, + BLOCK_Q=BLOCK_Q, + NUM_SEGMENTS_PER_SEQ=num_segments, + USE_FP8=False, + ) + + return output diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 846e6bcad0de..b6913003be59 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -55,6 +55,7 @@ class KVQuantMode(IntEnum): TURBOQUANT_K3V4_NC = 8 TURBOQUANT_3BIT_NC = 9 NVFP4_DS_MLA = 10 # opaque-bytes NVFP4 DS-MLA layouts (FlashMLA sparse) + ULTRAQUANT_4BIT = 11 # packed FP4 values + UE8M0 group scales @property def is_per_token_head(self) -> bool: @@ -80,6 +81,11 @@ def is_turboquant(self) -> bool: KVQuantMode.TURBOQUANT_3BIT_NC, ) + @property + def is_ultraquant(self) -> bool: + """True for the UltraQuant 4-bit KV-cache format.""" + return self == KVQuantMode.ULTRAQUANT_4BIT + def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`.""" @@ -98,6 +104,8 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: return KVQuantMode.NVFP4 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"): return KVQuantMode[kv_cache_dtype.upper()] + if kv_cache_dtype == "ultraquant_4bit": + return KVQuantMode.ULTRAQUANT_4BIT if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): return KVQuantMode.FP8_PER_TENSOR return KVQuantMode.NONE From ae19af883a3ab7baa8a42ba0a6ef0752e4120b33 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 6 Oct 2026 14:50:10 -0700 Subject: [PATCH 095/952] [CI] Require transformers 5.19 for the EmbeddingGemma2 registry test (#60295) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- tests/models/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 95b424ede580..52ab94e6163d 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -583,7 +583,7 @@ def check_available_online( ), "EmbeddingGemma2Model": _HfExamplesInfo( "google/embeddinggemma-2", - min_transformers_version="5.18.0", + min_transformers_version="5.19.0", is_available_online=False, ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), From b6d8e8afd985f5711eee68e343d2ce908d166488 Mon Sep 17 00:00:00 2001 From: Cheese Cake Date: Wed, 7 Oct 2026 03:32:04 +0530 Subject: [PATCH 096/952] [Bugfix][MRV2][Spec Decode] Accept num_speculative_tokens in NgramGPUSpeculator.propose (#60300) Signed-off-by: cheesecake Co-authored-by: Claude Opus 5.5 --- vllm/v1/worker/gpu/spec_decode/ngram/speculator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 404f576274e8..7a61e3581a11 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -231,6 +231,7 @@ def propose( skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, + num_speculative_tokens: int | None = None, ) -> torch.Tensor: num_reqs = input_batch.num_reqs if dummy_run: From 5ab4e4445bcbe2bf07351b82d1af160234c32273 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Tue, 6 Oct 2026 18:49:57 -0500 Subject: [PATCH 097/952] [Perf][MoE] Skip padded work in block-FP8 DeepGEMM experts (#59128) Signed-off-by: Tyler Michael Smith Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/moe/test_deepgemm.py | 147 +++++++++++++++ ..._mul_per_token_group_quant_fp8_colmajor.py | 89 +++++++++ .../layers/fused_moe/deep_gemm_utils.py | 64 ++++--- .../layers/fused_moe/experts/deep_gemm_moe.py | 42 ++++- .../layers/quantization/utils/fp8_utils.py | 171 +++++++++++++----- 5 files changed, 434 insertions(+), 79 deletions(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index cee0faa59229..385cc596ed9d 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -158,6 +158,153 @@ def make_block_quant_fp8_weights( return w1, w2, w1_s, w2_s +@pytest.mark.parametrize("num_tokens", [16, 129, 512]) +@pytest.mark.parametrize("with_metadata", [False, True]) +@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") +def test_block_fp8_padding_on_graph_replay(num_tokens, with_metadata, monkeypatch): + """Both FP8 GEMMs consume only live rows as routing changes on replay.""" + from unittest.mock import Mock + + import vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe as impl + from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT + + if not ( + impl.current_platform.is_device_capability_family(90) + or ( + impl.current_platform.is_device_capability_family(100) + and DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0 + ) + ): + pytest.skip("Prefix-sum FP8 dispatch requires SM90 or SM10.x") + + torch.manual_seed(0) + e, n, k, topk = 8, 1024, 512, 2 + x = torch.randn(num_tokens, k, device="cuda", dtype=torch.bfloat16) / 10 + aq, scales = per_token_group_quant_fp8(x, 128, use_ue8m0=True) + w1, w2, s1, s2 = make_block_quant_fp8_weights(e, n, k, BLOCK_SIZE) + quant = fp8_w8a8_moe_quant_config(w1_scale=s1, w2_scale=s2, block_shape=BLOCK_SIZE) + experts = impl.DeepGemmExperts(make_dummy_moe_config(), quant) + ids = torch.zeros(num_tokens, topk, device="cuda", dtype=torch.int64) + weights = torch.full((num_tokens, topk), 0.5, device="cuda") + counts = torch.zeros(e, device="cuda", dtype=torch.int32) + metadata = mk.ExpertTokensMetadata(counts, None) if with_metadata else None + shapes = experts.workspace_shapes( + num_tokens, 2 * n, k, topk, e, e, metadata, MoEActivation.SILU + ) + + def route(mode): + ids[:, 0] = torch.arange(num_tokens, device="cuda") % e + ids[:, 1] = (ids[:, 0] + 1) % e + if mode == 1: + ids.fill_(-1) + ids[:1, 0] = e - 1 + elif mode == 2: + ids.fill_(-1) + counts.copy_(torch.bincount(ids[ids >= 0], minlength=e).int()) + + route(0) + variants = [] + for dense in (True, False): + ws1, ws2, output = [ + torch.empty(shape, device="cuda", dtype=torch.bfloat16) for shape in shapes + ] + + def invoke(output=output, ws1=ws1, ws2=ws2): + experts.apply( + output, + aq, + w1, + w2, + weights, + ids, + MoEActivation.SILU, + e, + None, + scales, + None, + ws1, + ws2, + metadata, + False, + ) + + with monkeypatch.context() as mp: + if dense: + platform = Mock(wraps=impl.current_platform) + platform.is_device_capability_family.return_value = False + mp.setattr(impl, "current_platform", platform) + gemm = impl.m_grouped_fp8_gemm_nt_contiguous + calls: list[int] = [] + + def check_layout(*args, dense=dense, calls=calls, gemm=gemm, **kwargs): + assert kwargs.get("use_psum_layout", False) == (not dense) + calls.append(args[3].numel()) + return gemm(*args, **kwargs) + + mp.setattr(impl, "m_grouped_fp8_gemm_nt_contiguous", check_layout) + invoke() + assert len(calls) == 2 + if not dense: + assert calls == [e, e] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + invoke() + variants.append((graph, output, ws1, ws2)) + + for mode in (0, 1, 2, 0): + route(mode) + for graph, _, ws1, ws2 in variants: + ws1.fill_(float("nan")) + ws2.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(variants[0][1], variants[1][1], rtol=0, atol=0) + assert torch.isfinite(variants[1][1]).all() + + +@pytest.mark.parametrize(("n", "k"), [(2048, 512), (512, 1024)]) +@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") +def test_deepgemm_prefix_sum_leaves_unused_capacity_untouched(n, k): + """Both projection shapes skip unused tiles as counts change on replay.""" + from vllm.platforms import current_platform + from vllm.utils.deep_gemm import ( + m_grouped_fp8_gemm_nt_contiguous, + mk_alignment_scope, + ) + + if not ( + current_platform.is_device_capability_family(90) + or current_platform.is_device_capability_family(100) + ): + pytest.skip("Requires SM90 or SM10.x") + + capacity, experts, alignment = 512, 4, 128 + x = torch.randn(capacity, k, device="cuda", dtype=torch.bfloat16) / 10 + aq, scales = per_token_group_quant_fp8(x, 128, use_ue8m0=True) + w, _, ws, _ = make_block_quant_fp8_weights(experts, n // 2, k, BLOCK_SIZE) + ends = torch.tensor([0, 1, 128, 129], device="cuda", dtype=torch.int32) + output = torch.empty(capacity, n, device="cuda", dtype=torch.bfloat16) + + def invoke(): + with mk_alignment_scope(alignment): + m_grouped_fp8_gemm_nt_contiguous( + (aq, scales), (w, ws), output, ends, use_psum_layout=True + ) + + invoke() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + invoke() + for bounds, allocated_end in (([0, 1, 128, 129], 256), ([0] * 4, 0)): + ends.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + output.fill_(91) + graph.replay() + torch.accelerator.synchronize() + assert torch.all(output[allocated_end:] == 91) + if allocated_end: + assert torch.all(output[[0, 128]] != 91) + + def run_single_case(m, n, k, topk, num_experts, block_size): """Run one (M,N,K) configuration on a single GPU and assert DeepGEMM == Triton baseline within tolerance. diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index 204954d1fe46..b7dcb956c8ed 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -7,6 +7,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( _per_token_group_quant_fp8_colmajor, silu_mul_per_token_group_quant_fp8_colmajor, + silu_mul_quant_fp8_packed_triton, ) from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -108,6 +109,94 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int): torch.testing.assert_close(output_scales, ref_output_scales) +@pytest.mark.parametrize( + "counts", [[0, 1, 127, 128, 129, 0, 257, 13], [128] * 8, [0] * 8] +) +@pytest.mark.parametrize("group_size", [32, 128]) +@pytest.mark.parametrize("scale_format", ["packed", "float32", "ceil_ue8m0"]) +@pytest.mark.parametrize("hidden_size", [256, 2048]) +@pytest.mark.parametrize("clamp_limit", [None, 10.0]) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA graph test") +def test_silu_quant_skips_expert_padding_on_graph_replay( + counts: list[int], + group_size: int, + scale_format: str, + hidden_size: int, + clamp_limit: float | None, +): + """Live rows match dense quantization; gaps stay untouched as counts change.""" + alignment = 128 + capacity = sum(triton.cdiv(n, alignment) * alignment for n in counts) + # Exercise unused capacity as well as the gaps between experts. + capacity = max(capacity, alignment) + if counts != [128] * 8: + capacity += alignment + x = torch.randn((capacity, 2 * hidden_size), device=DEVICE, dtype=torch.bfloat16) + x *= 16 + clean_x = x.clone() + ends = torch.empty(len(counts), device=DEVICE, dtype=torch.int32) + output = torch.empty((capacity, hidden_size), device=DEVICE, dtype=FLOAT8_DTYPE) + + def set_counts(values): + positions = [] + mask = torch.zeros(capacity, device=DEVICE, dtype=torch.bool) + start = 0 + for count in values: + mask[start : start + count] = True + positions.append(start + count) + start += triton.cdiv(count, alignment) * alignment + ends.copy_(torch.tensor(positions, device=DEVICE, dtype=torch.int32)) + return mask + + def quantize(x, output=None, **kwargs): + if scale_format == "packed": + return silu_mul_quant_fp8_packed_triton( + x, + group_size=group_size, + output_q=output, + clamp_limit=clamp_limit, + **kwargs, + ) + return silu_mul_per_token_group_quant_fp8_colmajor( + x, + group_size=group_size, + output=output, + clamp_limit=clamp_limit, + use_ue8m0=scale_format == "ceil_ue8m0", + **kwargs, + ) + + def invoke(): + return quantize( + x, + output=output, + expert_ends=ends, + expert_alignment=alignment, + ) + + set_counts(counts) + invoke() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_q, actual_s = invoke() + + for values in (counts, counts[::-1]): + live = set_counts(values) + x.copy_(clean_x) + x[~live] = float("nan") + output.view(torch.uint8).fill_(91) + graph.replay() + reference_q, reference_s = quantize(x) + torch.testing.assert_close( + actual_q.view(torch.uint8)[live], + reference_q.view(torch.uint8)[live], + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual_s[live], reference_s[live], rtol=0, atol=0) + assert torch.all(actual_q.view(torch.uint8)[~live] == 91) + + @pytest.mark.parametrize("T", [128, 256, 512]) @pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2]) @pytest.mark.parametrize("clamp_limit", [7.0, 10.0]) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 33a229695561..624aa9ce86fd 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -119,6 +119,7 @@ def _fwd_kernel_ep_scatter_1( BLOCK_E: tl.constexpr, BLOCK_EXPERT_NUM: tl.constexpr, ALIGN_M: tl.constexpr, + USE_PSUM_LAYOUT: tl.constexpr, ): cur_expert = tl.program_id(0) @@ -140,20 +141,22 @@ def _fwd_kernel_ep_scatter_1( tl.store(expert_start_loc + cur_expert, cur_expert_start) cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) - m_indices_start_ptr = m_indices + cur_expert_start - off_expert = tl.arange(0, BLOCK_E) - - # any rows in the per-expert aligned region that do not correspond to - # real tokens are left untouched here and should remain initialized to - # -1 so DeepGEMM can skip them - for start_m in tl.range(0, cur_expert_token_num, BLOCK_E): - offs = start_m + off_expert - mask = offs < cur_expert_token_num - tl.store( - m_indices_start_ptr + offs, - cur_expert, - mask=mask, - ) + if USE_PSUM_LAYOUT: + # Logical ends include preceding experts' alignment gaps. + tl.store(m_indices + cur_expert, cur_expert_start + cur_expert_token_num) + else: + m_indices_start_ptr = m_indices + cur_expert_start + off_expert = tl.arange(0, BLOCK_E) + + # Unwritten padding retains the -1 initialization. + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E): + offs = start_m + off_expert + mask = offs < cur_expert_token_num + tl.store( + m_indices_start_ptr + offs, + cur_expert, + mask=mask, + ) @triton.jit @@ -285,6 +288,7 @@ def ep_scatter( align_m: int = 128, block_size: int = 128, pack_ue8m0: bool = False, + use_psum_layout: bool = False, ): # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m. BLOCK_E = 128 @@ -295,7 +299,10 @@ def ep_scatter( # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts) grid = num_experts - assert m_indices.shape[0] % align_m == 0 + if use_psum_layout: + assert m_indices.shape == (num_experts,) + else: + assert m_indices.shape[0] % align_m == 0 assert expert_start_loc.shape[0] == num_experts # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is. @@ -311,6 +318,7 @@ def ep_scatter( BLOCK_E=BLOCK_E, BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts), ALIGN_M=align_m, + USE_PSUM_LAYOUT=use_psum_layout, ) grid = min(recv_topk.shape[0], 1024 * 8) @@ -465,6 +473,7 @@ def deepgemm_moe_permute( expert_tokens_meta: mk.ExpertTokensMetadata | None, aq_out: torch.Tensor | None = None, block_size: int | None = None, + use_psum_layout: bool = False, ): assert aq.ndim == 2 assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids" @@ -509,17 +518,15 @@ def deepgemm_moe_permute( else: aq_scale_out = torch.zeros((M_sum, sf_k), device=device, dtype=torch.float32) - # DeepGEMM uses negative values in m_indices (here expert_ids) to mark - # completely invalid / padded blocks that should be skipped. We always - # initialize expert_ids to -1 so any row that is not explicitly written - # by the scatter kernel will be treated as invalid and skipped by - # DeepGEMM's scheduler. - expert_ids = torch.full( - (M_sum,), - fill_value=-1, - device=device, - dtype=torch.int32, - ) + if use_psum_layout: + grouped_layout = torch.empty( + (local_num_experts,), device=device, dtype=torch.int32 + ) + else: + # Negative expert IDs mark padding for the contiguous scheduler. + grouped_layout = torch.full( + (M_sum,), fill_value=-1, device=device, dtype=torch.int32 + ) inv_perm = torch.empty(topk_ids.shape, device=device, dtype=torch.int32) expert_num_tokens = None @@ -541,14 +548,15 @@ def deepgemm_moe_permute( expert_map=expert_map, output_tensor=aq_out, output_tensor_scale=aq_scale_out, - m_indices=expert_ids, + m_indices=grouped_layout, output_index=inv_perm, align_m=align_used, block_size=block_k, pack_ue8m0=pack_ue8m0, + use_psum_layout=use_psum_layout, ) - return aq_out, aq_scale_out, expert_ids, inv_perm, align_used + return aq_out, aq_scale_out, grouped_layout, inv_perm, align_used def deepgemm_unpermute_and_reduce( diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 699b029bad86..0eef7a83356d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -244,7 +244,12 @@ def workspace_shapes( return (workspace1, workspace2, output) def _act_mul_quant( - self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation + self, + input: torch.Tensor, + output: torch.Tensor, + activation: MoEActivation, + expert_ends: torch.Tensor | None = None, + expert_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: assert self.block_shape is not None block_k = self.block_shape[1] @@ -271,6 +276,8 @@ def _act_mul_quant( clamp_limit=self.gemm1_clamp_limit, alpha=self.gemm1_alpha, beta=self.gemm1_beta, + expert_ends=expert_ends, + expert_alignment=expert_alignment, ) act_out = torch.empty( (M_sum, activation_out_dim), dtype=input.dtype, device=input.device @@ -294,6 +301,8 @@ def _act_mul_quant( group_size=block_k, alpha=self.gemm1_alpha, beta=self.gemm1_beta, + expert_ends=expert_ends, + expert_alignment=expert_alignment, ) # 3. fallback path for non-SiLU activations in non‑UE8M0 cases. @@ -349,7 +358,20 @@ def apply( a1q_perm = _resize_cache( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K) ) - a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( + # Both GEMMs and quantization must agree on the live expert ranges. + use_psum_layout = ( + not self.mxfp8 + and ( + current_platform.is_device_capability_family(90) + or ( + current_platform.is_device_capability_family(100) + and DeepGemmQuantScaleFMT.from_oracle() + == DeepGemmQuantScaleFMT.UE8M0 + ) + ) + and activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI_UNINTERLEAVE) + ) + a1q, a1q_scale, grouped_layout, inv_perm, align_used = deepgemm_moe_permute( aq=a1q, aq_scale=a1q_scale, topk_ids=topk_ids, @@ -360,17 +382,21 @@ def apply( # MXFP8 uses a 32-element activation-scale group (block_shape[1]); # FP8-block keeps the default (128) alignment. block_size=self.block_shape[1] if self.mxfp8 else None, + use_psum_layout=use_psum_layout, ) assert a1q.size(0) == M_sum # MXFP8 (1x32) drives the fp8_fp4-aliased grouped GEMM with recipe # (1, 32); the FP8 block path keeps the default (128) recipe. - gemm_kwargs = ( + gemm_kwargs: dict = ( {"recipe_a": (1, self.block_shape[1]), "recipe_b": (1, self.block_shape[1])} if self.mxfp8 else {} ) + if use_psum_layout: + gemm_kwargs["use_psum_layout"] = True + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment; # otherwise the scheduler can pick the wrong expert id from m_indices # under cudagraph replay. @@ -380,7 +406,7 @@ def apply( (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, - expert_ids, + grouped_layout, **gemm_kwargs, ) @@ -389,7 +415,11 @@ def apply( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) ) a2q, a2q_scale = self._act_mul_quant( - input=mm1_out.view(-1, N), output=quant_out, activation=activation + input=mm1_out.view(-1, N), + output=quant_out, + activation=activation, + expert_ends=grouped_layout if use_psum_layout else None, + expert_alignment=align_used if use_psum_layout else 0, ) mm2_out = _resize_cache(workspace2, (M_sum, K)) @@ -397,7 +427,7 @@ def apply( (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, - expert_ids, + grouped_layout, **gemm_kwargs, ) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index c9888dc3cb6f..fb87c3a61450 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -142,6 +142,7 @@ def _silu_mul_quant_fp8_packed_kernel( input_ptr, output_q_ptr, output_scale_ptr, + expert_ends_ptr, M, input_stride_m, output_q_stride_m, @@ -158,6 +159,7 @@ def _silu_mul_quant_fp8_packed_kernel( PACKS_PER_CTA: tl.constexpr, BLOCK_M: tl.constexpr, HAS_CLAMP: tl.constexpr, + EXPERT_ALIGNMENT: tl.constexpr, ): GROUPS_PER_PACK: tl.constexpr = 4 hidden_size: tl.constexpr = N // 2 @@ -166,6 +168,13 @@ def _silu_mul_quant_fp8_packed_kernel( row_start = tl.program_id(1).to(tl.int64) * BLOCK_M row_step = tl.num_programs(1).to(tl.int64) * BLOCK_M + if EXPERT_ALIGNMENT: + expert = tl.program_id(2) + previous_end = tl.load(expert_ends_ptr + expert - 1, expert > 0, other=0) + expert_start = tl.cdiv(previous_end, EXPERT_ALIGNMENT) * EXPERT_ALIGNMENT + row_start += expert_start.to(tl.int64) + M = tl.load(expert_ends_ptr + expert) + groups_per_cta: tl.constexpr = PACKS_PER_CTA * GROUPS_PER_PACK elems_per_cta: tl.constexpr = groups_per_cta * GROUP_SIZE col_start = pack_tile * elems_per_cta @@ -254,7 +263,19 @@ def silu_mul_quant_fp8_packed_triton( clamp_limit: float | None = None, alpha: float = 1.0, beta: float = 0.0, + *, + expert_ends: torch.Tensor | None = None, + expert_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse activation and FP8 quantization, optionally skipping expert padding. + + ``expert_ends`` uses DeepGEMM's prefix-sum layout: each entry is the + exclusive end of an expert's live rows, including earlier alignment gaps. + An expert starts at the preceding end rounded up to ``expert_alignment``. + Padding outputs are left unwritten and must not be consumed as live rows. + Without ``expert_ends``, the original dense launch configuration is used + and the metadata loads are compiled out. + """ assert input.dim() == 2 assert input.is_contiguous() @@ -293,13 +314,29 @@ def silu_mul_quant_fp8_packed_triton( grid_n = triton.cdiv(packs_per_row, packs_per_cta) grid_m = min(triton.cdiv(M, BM), 4096) - grid = (grid_n, grid_m) + grid: tuple[int, ...] = (grid_n, grid_m) + if expert_ends is not None: + assert expert_ends.ndim == 1 and expert_ends.numel() > 0 + assert expert_ends.dtype == torch.int32 + assert expert_ends.device == input.device and expert_ends.is_contiguous() + assert expert_alignment > 0 + num_experts = expert_ends.numel() + # Bound empty CTAs while retaining enough parallelism for full experts. + max_ctas = 4096 if group_size < 128 else 8192 + grid_m = min( + triton.cdiv(M, BM * num_experts), + max(1, triton.cdiv(max_ctas, grid_n * num_experts)), + ) + grid = (grid_n, grid_m, num_experts) + else: + assert expert_alignment == 0 has_clamp = clamp_limit is not None _silu_mul_quant_fp8_packed_kernel[grid]( input, output_q, output_scale_packed, + expert_ends, M, input.stride(0), output_q.stride(0), @@ -316,6 +353,7 @@ def silu_mul_quant_fp8_packed_triton( PACKS_PER_CTA=packs_per_cta, BLOCK_M=BM, HAS_CLAMP=has_clamp, + EXPERT_ALIGNMENT=expert_alignment, num_warps=num_warps, num_stages=num_stages, ) @@ -328,6 +366,7 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( y_ptr, # [M, N] y_q_ptr, # [M, N // 2] y_s_ptr, # [M, (N // 2) // GROUP_SIZE] + expert_ends_ptr, M, # num tokens N, # intermediate size # Stride @@ -345,8 +384,8 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( GROUP_SIZE: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + EXPERT_ALIGNMENT: tl.constexpr, ): - # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks. """Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills the outputs tensors at the right positions. @@ -357,56 +396,77 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( m_offset = pid_m.to(tl.int64) * BLOCK_M n_offset = pid_n.to(tl.int64) * BLOCK_N - if m_offset >= M: - return + row_step = tl.num_programs(0).to(tl.int64) * BLOCK_M + num_iterations = 1 + if EXPERT_ALIGNMENT: + expert = tl.program_id(2) + previous_end = tl.load(expert_ends_ptr + expert - 1, expert > 0, other=0) + expert_start = tl.cdiv(previous_end, EXPERT_ALIGNMENT) * EXPERT_ALIGNMENT + m_offset += expert_start.to(tl.int64) + M = tl.load(expert_ends_ptr + expert) + num_iterations = tl.cdiv(tl.maximum(M - m_offset, 0), row_step) offs_n = tl.arange(0, BLOCK_N).to(tl.int64) offs_m = tl.arange(0, BLOCK_M).to(tl.int64) - base_y_ptr = y_ptr + m_offset * N + n_offset - - act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :] + for _ in range(num_iterations): + base_y_ptr = y_ptr + m_offset * N + n_offset - act_in = tl.load(act_in_ptrs) - mul_in = tl.load(act_in_ptrs + N_2) + act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :] - # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the - # input dtype, run silu in fp32 then narrow, and do the mul at input - # precision so HAS_CLAMP True/False share the same multiplication path. - if HAS_CLAMP: - act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to( - y_ptr.dtype.element_ty - ) - mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to( - y_ptr.dtype.element_ty - ) - # Unified gated activation: silu == swigluoai with alpha=1, beta=0. - # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu - # Keep glu/up at input precision (narrow before the mul) so the alpha=1, - # beta=0 defaults match the C++ silu_and_mul path bit-for-bit. - act_in = act_in.to(tl.float32) - glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty) - up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty) - y = (glu * up).to(tl.float32) - - # quant - _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) - scale_raw = _absmax * (1.0 / fp8_max) - y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw - y_s = tl.reshape(y_s, (BLOCK_M, 1)) - y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + if EXPERT_ALIGNMENT: + live = m_offset + offs_m < M + act_in = tl.load(act_in_ptrs, live[:, None], other=0) + mul_in = tl.load(act_in_ptrs + N_2, live[:, None], other=0) + else: + act_in = tl.load(act_in_ptrs) + mul_in = tl.load(act_in_ptrs + N_2) - # store y_q - base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset - y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :] - tl.store(y_q_ptrs, y_q) + # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the + # input dtype, run silu in fp32 then narrow, and do the mul at input + # precision so HAS_CLAMP True/False share the same multiplication path. + if HAS_CLAMP: + act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to( + y_ptr.dtype.element_ty + ) + mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to( + y_ptr.dtype.element_ty + ) + # Unified gated activation: silu == swigluoai with alpha=1, beta=0. + # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu + # Keep glu/up at input precision (narrow before the mul) so the alpha=1, + # beta=0 defaults match the C++ silu_and_mul path bit-for-bit. + act_in = act_in.to(tl.float32) + glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty) + up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty) + y = (glu * up).to(tl.float32) + + # quant + _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) + scale_raw = _absmax * (1.0 / fp8_max) + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw + y_s = tl.reshape(y_s, (BLOCK_M, 1)) + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + # store y_q + base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset + y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :] + if EXPERT_ALIGNMENT: + tl.store(y_q_ptrs, y_q, live[:, None]) + else: + tl.store(y_q_ptrs, y_q) + + # store y_s + group_id = n_offset // GROUP_SIZE + base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset + y_s_ptrs = base_y_s_ptr + offs_m + y_s = tl.reshape(y_s, (BLOCK_M,)) + if EXPERT_ALIGNMENT: + tl.store(y_s_ptrs, y_s, live) + else: + tl.store(y_s_ptrs, y_s) - # store y_s - group_id = n_offset // GROUP_SIZE - base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset - y_s_ptrs = base_y_s_ptr + offs_m - y_s = tl.reshape(y_s, (BLOCK_M,)) - tl.store(y_s_ptrs, y_s) + m_offset += row_step def silu_mul_per_token_group_quant_fp8_colmajor( @@ -418,9 +478,15 @@ def silu_mul_per_token_group_quant_fp8_colmajor( group_size: int = 128, alpha: float = 1.0, beta: float = 0.0, + *, + expert_ends: torch.Tensor | None = None, + expert_alignment: int = 0, ): """Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate (silu: alpha=1, beta=0; swigluoai: alpha, beta from config). + + Optional expert endpoints follow DeepGEMM's prefix-sum layout. Padding + values and scales are left unwritten and must not be consumed as live rows. """ GROUP_SIZE = group_size assert input.ndim == 2 @@ -458,13 +524,27 @@ def silu_mul_per_token_group_quant_fp8_colmajor( # Force even division so we can avoid edgecases within the kernel. assert M % BLOCK_M == 0 assert N_2 % BLOCK_N == 0 - grid = (M // BLOCK_M, N_2 // BLOCK_N) + grid: tuple[int, ...] = (M // BLOCK_M, N_2 // BLOCK_N) + if expert_ends is not None: + assert expert_ends.ndim == 1 and expert_ends.numel() > 0 + assert expert_ends.dtype == torch.int32 + assert expert_ends.device == input.device and expert_ends.is_contiguous() + assert expert_alignment > 0 and expert_alignment % BLOCK_M == 0 + num_experts = expert_ends.numel() + grid_m = min( + triton.cdiv(M, BLOCK_M * num_experts), + max(1, triton.cdiv(8192, grid[1] * num_experts)), + ) + grid = (grid_m, grid[1], num_experts) + else: + assert expert_alignment == 0 has_clamp = clamp_limit is not None _silu_mul_per_token_group_quant_fp8_colmajor[grid]( input, output, output_scales, + expert_ends, M, N, output_scales.stride(-1), @@ -479,6 +559,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor( GROUP_SIZE, BLOCK_M, BLOCK_N, + expert_alignment, ) return output, output_scales From d77cf8115601f2856aacbfdd82a436ee78a558e6 Mon Sep 17 00:00:00 2001 From: Sriraj <168966222+sricursion@users.noreply.github.com> Date: Wed, 7 Oct 2026 05:51:00 +0530 Subject: [PATCH 098/952] [Bugfix][Frontend] Honor default penalties in batch chat completions (#60307) Signed-off-by: sricursion Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../openai/test_penalty_default_resolution.py | 28 ++++++++++++++++++- .../openai/chat_completion/protocol.py | 4 +-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/test_penalty_default_resolution.py b/tests/entrypoints/openai/test_penalty_default_resolution.py index 2c39f5bef68d..23c57241b740 100644 --- a/tests/entrypoints/openai/test_penalty_default_resolution.py +++ b/tests/entrypoints/openai/test_penalty_default_resolution.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for presence_penalty / frequency_penalty resolution from -default_sampling_params in ChatCompletionRequest and CompletionRequest. +default_sampling_params in ChatCompletionRequest, BatchChatCompletionRequest +and CompletionRequest. Regression test for https://github.com/vllm-project/vllm/issues/50767: these two penalties defaulted to 0.0 (not None) and to_sampling_params() @@ -12,6 +13,7 @@ """ from vllm.entrypoints.openai.chat_completion.protocol import ( + BatchChatCompletionRequest, ChatCompletionRequest, ) from vllm.entrypoints.openai.completion.protocol import ( @@ -29,6 +31,14 @@ def _chat(**kwargs): ) +def _batch(**kwargs): + conversation = [{"role": "user", "content": "hi"}] + request = BatchChatCompletionRequest( + model="test-model", messages=[conversation], **kwargs + ) + return request.to_chat_completion_request(conversation) + + class TestChatCompletionPenaltyDefaults: def test_defaults_applied_when_client_omits(self): """Server-default penalties are applied when the client sends none.""" @@ -49,6 +59,22 @@ def test_falls_back_to_zero_without_default(self): assert sp.frequency_penalty == 0.0 +class TestBatchChatCompletionPenaltyDefaults: + def test_defaults_applied_when_client_omits(self): + """Server-default penalties reach every conversation of a batch.""" + sp = _batch().to_sampling_params(100, _DEFAULTS) + assert sp.presence_penalty == 1.5 + assert sp.frequency_penalty == 0.5 + + def test_client_value_overrides_default(self): + """An explicit client penalty, even 0.0, wins over the server default.""" + sp = _batch(presence_penalty=0.0, frequency_penalty=0.2).to_sampling_params( + 100, _DEFAULTS + ) + assert sp.presence_penalty == 0.0 + assert sp.frequency_penalty == 0.2 + + class TestCompletionPenaltyDefaults: def test_defaults_applied_when_client_omits(self): sp = CompletionRequest(model="test-model", prompt="hi").to_sampling_params( diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index aee1c8a7a63a..15905b8a3da6 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -1093,7 +1093,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel): model: str | None = None # Shared sampling / generation fields — mirror ChatCompletionRequest. - frequency_penalty: float | None = 0.0 + frequency_penalty: float | None = None logit_bias: dict[str, float] | None = None logprobs: bool | None = False top_logprobs: TopLogprobsParam = 0 @@ -1108,7 +1108,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel): max_tokens: int | None = None max_completion_tokens: int | None = None n: int | None = 1 - presence_penalty: float | None = 0.0 + presence_penalty: float | None = None response_format: Any | None = None seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX) stop: StopParam = Field(default_factory=list) From 6c2925fe5fe8156475affbfa995ed272c96fe4fa Mon Sep 17 00:00:00 2001 From: jxj Date: Wed, 7 Oct 2026 08:35:14 +0800 Subject: [PATCH 099/952] [Bugfix][Core] Preserve speculative metrics when coalescing outputs (#55708) Signed-off-by: git-jxj <65210887+git-jxj@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/test_outputs.py | 49 +++++++++++++++++++++++++++++++++++++++++++ vllm/outputs.py | 4 ++++ 2 files changed, 53 insertions(+) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 9acc1f7b74ac..e7b732b3e2ac 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -6,6 +6,7 @@ from vllm.outputs import ( ClassificationRequestOutput, + CompletionOutput, EmbeddingRequestOutput, PoolingOutput, PoolingRequestOutput, @@ -13,6 +14,7 @@ RequestOutput, ScoringRequestOutput, ) +from vllm.v1.metrics.stats import RequestSpecDecodeMetrics pytestmark = pytest.mark.cpu_test @@ -61,3 +63,50 @@ def test_specialized_pooling_outputs_preserve_request_error(output_type, empty_v value = getattr(output.outputs, "probs", value) value = getattr(output.outputs, "score", value) assert value == empty_value + + +def test_request_output_add_preserves_terminal_spec_decode_metrics(): + accumulated = RequestOutput( + request_id="request", + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="first", + token_ids=[1], + cumulative_logprob=None, + logprobs=None, + ) + ], + finished=False, + ) + spec_decode_metrics = RequestSpecDecodeMetrics.new(num_spec_tokens=2) + spec_decode_metrics.observe(num_draft_tokens=2, num_accepted=1) + terminal = RequestOutput( + request_id="request", + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="second", + token_ids=[2], + cumulative_logprob=None, + logprobs=None, + finish_reason="length", + spec_decode_metrics=spec_decode_metrics, + ) + ], + finished=True, + ) + + accumulated.add(terminal, aggregate=True) + + output = accumulated.outputs[0] + assert output.token_ids == [1, 2] + assert output.finish_reason == "length" + assert accumulated.finished + assert output.spec_decode_metrics is spec_decode_metrics diff --git a/vllm/outputs.py b/vllm/outputs.py index 35f71ec30439..8e69ffa91d91 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -212,6 +212,10 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None: completion.sampling_mask.token_ids.extend( next_completion.sampling_mask.token_ids ) + if next_completion.spec_decode_metrics is not None: + completion.spec_decode_metrics = ( + next_completion.spec_decode_metrics + ) completion.cumulative_logprob = ( next_completion.cumulative_logprob ) From db6e3cd8c4f9b84c39fe3c44fab0b1d3117f758b Mon Sep 17 00:00:00 2001 From: Sriraj <168966222+sricursion@users.noreply.github.com> Date: Wed, 7 Oct 2026 06:12:17 +0530 Subject: [PATCH 100/952] [Bugfix][Frontend] Log the full request body when max_log_len is unset (#60310) Signed-off-by: sricursion Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../serve/utils/test_request_logger.py | 18 ++++++++++++++++++ vllm/entrypoints/serve/utils/request_logger.py | 3 +-- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/serve/utils/test_request_logger.py b/tests/entrypoints/serve/utils/test_request_logger.py index 717ff190b3a7..461b954cae0f 100644 --- a/tests/entrypoints/serve/utils/test_request_logger.py +++ b/tests/entrypoints/serve/utils/test_request_logger.py @@ -4,6 +4,9 @@ import logging from unittest.mock import MagicMock, patch +import pytest + +from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.serve.utils.request_logger import RequestLogger @@ -269,3 +272,18 @@ def test_streaming_complete_logs_full_text_content(): assert call_args[1] == "test-streaming-full-text" assert call_args[2] == " (streaming complete)" assert call_args[4] == "streaming_complete" + + +@pytest.mark.parametrize("max_log_len", [None, 10]) +def test_request_logger_log_request_body_truncation(max_log_len): + """The JSON body is logged in full unless max_log_len truncates it.""" + request = CompletionRequest(model="test-model", prompt="Hello, world!") + full_body = request.model_dump_json(exclude_unset=True) + mock_logger = MagicMock() + + with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger): + RequestLogger(max_log_len=max_log_len).log_request_body(request) + + mock_logger.debug.assert_called_once() + logged_body = mock_logger.debug.call_args.args[2] + assert logged_body == full_body[:max_log_len] diff --git a/vllm/entrypoints/serve/utils/request_logger.py b/vllm/entrypoints/serve/utils/request_logger.py index 09b27e067320..6272054a1e99 100644 --- a/vllm/entrypoints/serve/utils/request_logger.py +++ b/vllm/entrypoints/serve/utils/request_logger.py @@ -71,11 +71,10 @@ def log_inputs( def log_request_body(self, request: AnyRequest | AnyPoolingRequest) -> None: if logger.isEnabledFor(logging.DEBUG): - max_log_len = self.max_log_len if self.max_log_len is not None else -1 logger.debug( "Request %s JSON body: %s", getattr(request, "request_id", "N/A"), - request.model_dump_json(exclude_unset=True)[:max_log_len], + request.model_dump_json(exclude_unset=True)[: self.max_log_len], ) def log_outputs( From 703c29179ceea46505f24c8faa3dab42ed3a211e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Wed, 7 Oct 2026 02:59:20 +0200 Subject: [PATCH 101/952] [TEST][XPU][CI] Generalize model runner prefill tail classification tests to all platforms (#60201) Signed-off-by: microslaw --- ...odel_runner_prefill_tail_classification.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index 9375c477a224..92207aa38f6b 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -23,9 +23,10 @@ landed in the same patch. """ +from types import SimpleNamespace + import numpy as np import pytest -import torch from vllm.config import ( CacheConfig, @@ -37,7 +38,9 @@ set_current_vllm_config, ) from vllm.model_executor.layers.attention import Attention +from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -50,7 +53,6 @@ BLOCK_SIZE = 16 NUM_BLOCKS = 10 -DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu" def _get_vllm_config() -> VllmConfig: @@ -76,6 +78,17 @@ def _get_vllm_config() -> VllmConfig: ) +def _make_model_runner(vllm_config: VllmConfig) -> GPUModelRunner: + """Build the V1 runner the platform's worker would, without starting it.""" + worker_cls = resolve_obj_by_qualname(vllm_config.parallel_config.worker_cls) + worker = SimpleNamespace( + vllm_config=vllm_config, + device=current_platform.device_type, + use_v2_model_runner=False, + ) + return worker_cls._make_model_runner(worker) + + def _initialize_kv_cache(runner: GPUModelRunner) -> None: attn_spec = FullAttentionSpec( block_size=BLOCK_SIZE, @@ -124,7 +137,7 @@ def spec_decode_model_runner(dist_init): vllm_config.compilation_config.static_forward_context["layer.0"] = Attention( num_heads, head_size, 0.1 ) - runner = GPUModelRunner(vllm_config, DEVICE_TYPE) + runner = _make_model_runner(vllm_config) _initialize_kv_cache(runner) yield runner From 13f8c453eeb7eaffd417c9011c112e56d1f765e1 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Tue, 6 Oct 2026 20:27:27 -0500 Subject: [PATCH 102/952] [CI][Releases] Trigger `perf-eval` Through Release Pipeline (#52213) Signed-off-by: Micah Williamson Co-authored-by: Kevin H. Luu --- .buildkite/release-pipeline.yaml | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index d6e47715fade..5eba9fd2e114 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -1077,6 +1077,38 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" + - trigger: "perf-eval" + label: "Trigger perf-eval (CUDA)" + key: trigger-perf-eval-cuda + if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" + depends_on: + - build-release-image-x86 + async: true + soft_fail: true + build: + branch: "main" + message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (CUDA)" + env: + NIGHTLY: 1 # run all perf-eval workloads + VLLM_COMMIT: "$BUILDKITE_COMMIT" + VLLM_IMAGE_CUDA: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64" + + - trigger: "perf-eval" + label: "Trigger perf-eval (ROCm)" + key: trigger-perf-eval-rocm + if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" + depends_on: + - build-rocm-release-image + async: true + soft_fail: true + build: + branch: "main" + message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (ROCm)" + env: + NIGHTLY: 1 # run all perf-eval workloads + VLLM_COMMIT: "$BUILDKITE_COMMIT" + VLLM_IMAGE_ROCM: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm" + # ============================================================================= # ROCk Release Pipeline (x86_64 only) # ============================================================================= From 4b4eb6277a43abbc25901de063389c405f0e758f Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 6 Oct 2026 18:30:52 -0700 Subject: [PATCH 103/952] [CI] Remove AMD mirror of CPU-only Rust frontend cargo steps (#60323) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 32 ------------------- .../test_areas/rust_frontend_cargo.yaml | 20 ------------ 2 files changed, 52 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0a050b6979f0..17fe9c0ae58e 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -111,38 +111,6 @@ steps: # # ######################################################################################################################################### -- label: ":computer: (CPU) Rust Frontend Cargo Style + Clippy" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - rust/ - - rust-toolchain.toml - - .buildkite/test_areas/rust_frontend_cargo.yaml - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh - commands: - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy - -- label: ":computer: (CPU) Rust Frontend Cargo" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace" - source_file_dependencies: - - rust/ - - rust-toolchain.toml - - .buildkite/test_areas/rust_frontend_cargo.yaml - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh - commands: - - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test - - label: ":computer: (CPU) Docker Build Metadata" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 53d1da7edf41..11bca6252667 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -14,16 +14,6 @@ steps: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy - mirror: - amd: - label: ":amd: (MI250) Rust Frontend Cargo Style + Clippy" - dind: false - device: mi250_1 - no_gpu: true - working_dir: "/vllm-workspace" - timeout_in_minutes: 35 - depends_on: - - image-build-amd - label: ":computer: (CPU) Rust Frontend Cargo" key: rust-frontend-cargo-tests @@ -38,13 +28,3 @@ steps: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh commands: - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test - mirror: - amd: - label: ":amd: (MI250) Rust Frontend Cargo" - dind: false - device: mi250_1 - no_gpu: true - working_dir: "/vllm-workspace" - timeout_in_minutes: 35 - depends_on: - - image-build-amd From fe51194e2d7965149a655ca20af4f31dd668777d Mon Sep 17 00:00:00 2001 From: vovani Date: Wed, 7 Oct 2026 03:35:12 +0200 Subject: [PATCH 104/952] [Bugfix][MiniMax-M3] Fence decode indexer shared-memory stage reuse (#60194) Signed-off-by: Vladimir Anisimov Co-authored-by: Claude Opus 5.5 Co-authored-by: OpenAI Codex Co-authored-by: Thien Tran --- tests/kernels/attention/test_minimax_m3.py | 40 +++++++++++++++---- .../nvidia/ops/index_decode_score.py | 3 ++ 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 83eef72e31e4..7f0e07309331 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Correctness tests for MiniMax M3 sparse prefill attention kernels.""" +from functools import partial + import pytest import torch @@ -1444,19 +1446,21 @@ def test_decode_index_topk_fp8(num_idx_heads: int): reason="CuteDSL index decode score requires Blackwell.", ) @pytest.mark.parametrize( - ("dtype", "decode_query_len", "max_decode_query_len"), + ("dtype", "decode_query_len", "max_decode_query_len", "concurrent_writes"), [ - (torch.bfloat16, 1, 1), - (torch.bfloat16, 3, 8), - (torch.float8_e4m3fn, 1, 1), - (torch.float8_e4m3fn, 3, 8), - (torch.float8_e4m3fn, 8, 8), + (torch.bfloat16, 1, 1, False), + (torch.bfloat16, 3, 8, False), + (torch.float8_e4m3fn, 1, 1, False), + (torch.float8_e4m3fn, 3, 8, False), + (torch.float8_e4m3fn, 8, 8, False), + pytest.param(torch.float8_e4m3fn, 1, 1, True, id="concurrent-writes"), ], ) def test_decode_index_score_cutedsl_correctness( dtype: torch.dtype, decode_query_len: int, max_decode_query_len: int, + concurrent_writes: bool, ): pytest.importorskip("cutlass") from vllm.models.minimax_m3.nvidia.ops import ( @@ -1466,7 +1470,8 @@ def test_decode_index_score_cutedsl_correctness( torch.manual_seed(0) init_blocks, local_blocks = 0, 0 num_idx_heads, head_dim = 4, 128 - active_seq_lens = torch.tensor((1025, 4097), device="cuda", dtype=torch.int32) + lengths = (131326, 131289, 131252, 131215) if concurrent_writes else (1025, 4097) + active_seq_lens = torch.tensor(lengths, device="cuda", dtype=torch.int32) batch = active_seq_lens.numel() total_q = batch * decode_query_len max_seq_len = int(active_seq_lens.max()) @@ -1488,7 +1493,8 @@ def test_decode_index_score_cutedsl_correctness( ) score = unified_score.transpose(0, 1) - minimax_m3_index_decode_score_cutedsl( + launch = partial( + minimax_m3_index_decode_score_cutedsl, idx_q, index_kv_cache, block_table, @@ -1501,6 +1507,7 @@ def test_decode_index_score_cutedsl_correctness( max_decode_query_len=max_decode_query_len, score_out=score, ) + launch() expected = _reference_decode_index_score( idx_q, index_kv_cache, @@ -1511,6 +1518,23 @@ def test_decode_index_score_cutedsl_correctness( ) torch.testing.assert_close(score, expected) + if concurrent_writes: + # Concurrent stores expose shared-memory reuse before ldmatrix reads finish. + quiet = score.clone() + noise = torch.zeros(1 << 26, device="cuda") + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + try: + for _ in range(20): + with torch.cuda.stream(side): + for _ in range(40): + noise.mul_(1.0) + launch() + torch.testing.assert_close(score, quiet, atol=0, rtol=0) + side.synchronize() + finally: + side.synchronize() + # Sparse attention kernels. def _reference_sparse_attn( diff --git a/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py b/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py index 3d2919e39227..8aa5e5a11e57 100644 --- a/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +++ b/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py @@ -287,6 +287,8 @@ def kernel( cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) for q in cutlass.range_constexpr(Q_TILES): cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + # Order ldmatrix reads before the producer refills this stage with TMA. + cute.arch.fence_proxy("async.shared", space="cta") cute.arch.mbarrier_arrive(tma_empty_mbar) tma_stage = 1 % self.num_stages @@ -342,6 +344,7 @@ def kernel( rC[None, m, n], ) + cute.arch.fence_proxy("async.shared", space="cta") cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) k_start = block_id * BLOCK_K + warp_id * 32 From 283d76d56a45d9bbead5ba650402fb0957abc30c Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Tue, 6 Oct 2026 21:45:04 -0400 Subject: [PATCH 105/952] [Bugfix] Widen ep_gather output offset to int64 (#58373) Signed-off-by: Robert Shaw Co-authored-by: Claude Code --- tests/kernels/moe/test_deepgemm.py | 22 +++++++++++++++++++ .../layers/fused_moe/deep_gemm_utils.py | 2 +- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 385cc596ed9d..e341f12a7107 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -70,6 +70,28 @@ def test_ep_gather_uses_64_bit_row_offsets(): torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) +def test_ep_gather_uses_64_bit_output_row_offsets(): + hidden_size = 6144 + output_row = (1 << 31) // hidden_size + 1 + num_rows = output_row + 1 + output = torch.empty( + (num_rows, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + + ep_gather( + input_tensor=torch.ones((1, hidden_size), device="cuda", dtype=torch.bfloat16), + recv_topk_ids=torch.zeros((num_rows, 1), device="cuda", dtype=torch.int64), + recv_topk_weight=torch.ones((num_rows, 1), device="cuda", dtype=torch.float32), + input_index=torch.zeros((num_rows, 1), device="cuda", dtype=torch.int32), + expert_map=None, + output_tensor=output, + ) + + torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) + + @pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("num_columns", [2048, 6144, 6145]) def test_fp8_workspace_shape(workspace_dtype, num_columns): diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 624aa9ce86fd..1dee9bad54b7 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -415,7 +415,7 @@ def _fwd_kernel_ep_gather( tl.store( output_tensor - + cur_token * output_tensor_stride0 + + tl.cast(cur_token, tl.int64) * output_tensor_stride0 + cur_block * BLOCK_D + off_d, accumulator.to(output_tensor.dtype.element_ty), From 956960f85fd9b5ae0de9faa004738b7a1c81c437 Mon Sep 17 00:00:00 2001 From: Cheese Cake Date: Wed, 7 Oct 2026 07:21:12 +0530 Subject: [PATCH 106/952] [Perf][MRV2] Avoid stream syncs in NonUvaBuffer host-to-device copies (#60215) Signed-off-by: cheesecake Co-authored-by: William Joyner <133391209+joynerwk03@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- tests/kernels/core/test_uva.py | 19 +++++++++++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 4 ++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index 0a183207c859..9d88eeced3ba 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -234,6 +234,25 @@ def test_non_pinned_cpu_tensor(device): assert gpu_view[9, 9] == 198 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only sync check.") +def test_non_uva_buffer_copies_without_stream_sync(): + buffer = buffer_utils.NonUvaBuffer(8, torch.int32) + buffer.np[:] = np.arange(8) + + torch.cuda.set_sync_debug_mode("error") + try: + prefix = buffer.uva(4) + full = buffer.uva() + finally: + torch.cuda.set_sync_debug_mode(0) + # The next step overwrites the host buffer right after the copy is issued. + buffer.np[:] = -1 + torch.accelerator.synchronize() + + assert prefix.tolist() == list(range(4)) + assert full.tolist() == list(range(8)) + + @pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.") @pytest.mark.skipif( not current_platform.is_xpu(), reason="XPU non-contiguous UVA test." diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 0731072e5968..32d997358c13 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -54,8 +54,8 @@ def __init__(self, size: int | Sequence[int], dtype: torch.dtype): def uva(self, n: int | None = None) -> torch.Tensor: if n is None: - return self._uva.copy_(self.cpu) - return self._uva[:n].copy_(self.cpu[:n]) + return self._uva.copy_(self.cpu, non_blocking=True) + return self._uva[:n].copy_(self.cpu[:n], non_blocking=True) class UvaBufferPool: From 8a26869bcad1ae4ab3069a7d2aece86c17c1c695 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 7 Oct 2026 10:32:46 +0800 Subject: [PATCH 107/952] [Bugfix] Index Mamba align state copies by request slot in MRV2 (#60210) Signed-off-by: zjy0516 Co-authored-by: Karl0007 <44745209+Karl0007@users.noreply.github.com> Co-authored-by: Yongqin Wang --- .../kernels/mamba/test_precopy_mamba_align.py | 80 ++++++++++++++++++- tests/v1/worker/test_mamba_utils.py | 46 +++++++++++ vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/model_states/interface.py | 6 ++ .../worker/gpu/model_states/mamba_hybrid.py | 77 +++++++++--------- vllm/v1/worker/mamba_utils.py | 42 ++++++---- 6 files changed, 195 insertions(+), 57 deletions(-) diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py index 22b73c936ece..8a91f8bdace5 100644 --- a/tests/kernels/mamba/test_precopy_mamba_align.py +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -31,7 +31,11 @@ from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils from vllm.platforms import current_platform from vllm.v1.worker import mamba_utils as worker_mamba_utils -from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel +from vllm.v1.worker.mamba_utils import ( + _TEMPORAL_TILES, + postprocess_mamba_fused_kernel, + precopy_mamba_align_fused_kernel, +) _parametrize: Callable[..., Callable[[Any], Any]] @@ -195,7 +199,8 @@ def test_precopy_matches_v1_copy_specs( convs, ssms, device, conv_state_dim_first ) bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + # Non-identity mapping: block-table rows are request slots, not batch rows. + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device).flip(0) grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) precopy_mamba_align_fused_kernel[grid]( dst_col, @@ -225,6 +230,77 @@ def test_precopy_matches_v1_copy_specs( torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) +@_parametrize("conv_state_dim_first", [False, True]) +@_cuda_required +def test_v2_postprocess_indexes_block_tables_by_req_slot(conv_state_dim_first): + """Under PP the V2 postprocess runs after later steps have re-gathered the + batch-ordered tables, so it must index the request-slot tables by slot.""" + device = torch.device("cuda") + torch.manual_seed(0) + num_reqs = 4 + block_size = 8 + num_blocks = num_reqs * MAX_COLS + 1 + bt = torch.arange(1, num_blocks, dtype=torch.int32, device=device).view( + num_reqs, MAX_COLS + ) + # slot 0: copy col 1 -> 0; slot 1: src == dst; slot 2: not aligned; + # slot 3: copy col 2 -> 0 with one accepted draft token. + num_accepted = torch.tensor([1, 1, 1, 2], dtype=torch.int32, device=device) + state_idx = torch.tensor([1, 0, 1, 2], dtype=torch.int32, device=device) + new_num_computed = torch.tensor([8, 8, 5, 8], dtype=torch.int32, device=device) + ref_src = torch.tensor([1, 0, -1, 2], dtype=torch.int32) + ref_dst = torch.zeros(num_reqs, dtype=torch.int32) + ref_bias = torch.tensor([0, 0, 0, 1], dtype=torch.int32) + + convs, ssms = _build_state(num_blocks, device, conv_state_dim_first) + conv_ref, ssm_ref = _reference( + convs, + ssms, + bt.cpu(), + ref_src, + ref_dst, + ref_bias, + num_reqs, + conv_state_dim_first, + ) + base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( + convs, ssms, device, conv_state_dim_first + ) + bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) + idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int32, device=device) + postprocess_mamba_fused_kernel[(num_reqs, NUM_LAYERS * 2, 1)]( + num_accepted.clone(), + state_idx, + None, + new_num_computed, + None, + bt_ptrs, + bt.stride(0), + base, + blk_stride, + elem, + inner, + width, + group, + drc, + drs, + num_accepted, + idx_mapping, + num_reqs, + block_size=block_size, + COPY_BLOCK_SIZE=1024, + CONV_STATE_DIM_FIRST=conv_state_dim_first, + HAS_IDX_MAPPING=True, + PRECOMPUTED_NEW_COMPUTED=True, + TEMPORAL_TILES=1, + ) + torch.accelerator.synchronize() + + for layer in range(NUM_LAYERS): + torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0) + torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0) + + def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch): monkeypatch.setattr( layer_mamba_utils, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 50faf0fb8b2e..4407d6fc5430 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -22,6 +22,7 @@ MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, @@ -638,6 +639,51 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_v2_align_ctx_binds_request_slot_tables_for_state_copies(): + """V2 state copies may run steps after their batch (PP), so they must read + the per-request-slot tables; aligned state indices are computed for the + current batch and keep reading the gathered tables.""" + cfg = _TestConfig(num_layers=1) + device = torch.device("cuda") + kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) + forward_context = { + "layer_0": _make_mock_attention( + torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), + torch.empty(cfg.num_blocks, cfg.temporal_state_dim), + ) + } + block_tables = BlockTables( + block_sizes=[cfg.block_size], + max_num_reqs=cfg.max_num_reqs, + max_num_batched_tokens=64, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[cfg.block_size], + ) + + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._mamba_group_ids = [] + state._mamba_spec = None + state._mamba_state_copy_funcs = _COPY_FUNCS + state.max_num_reqs = cfg.max_num_reqs + state.device = device + state.vllm_config = MagicMock() + state.vllm_config.compilation_config.static_forward_context = forward_context + + state.initialize_kv_cache(kv_cache_config, block_tables) + + ctx = state._mamba_ctx + assert ctx is not None and ctx.is_initialized + assert ctx.block_table_ptrs.tolist() == [ + block_tables.block_tables[0].gpu.data_ptr() + ] + assert ctx.aligned_index_block_table_ptrs.tolist() == [ + block_tables.input_block_tables[0].data_ptr() + ] + + # ----------------------------------------------------------------------------- # stage_postprocess_inputs_to_gpu: single-pass staging into pinned views # ----------------------------------------------------------------------------- diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c3ac543ff290..d0c08707f483 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -778,6 +778,7 @@ def initialize_kv_cache( self.kv_caches = [ cache for cache in kv_caches_dict.values() if cache.device == self.device ] + self.model_state.initialize_kv_cache(self.kv_cache_config, self.block_tables) if is_profiling: self.kv_connector = NO_OP_KV_CONNECTOR else: diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index f03caf4dba83..3db3fc9e7ae6 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -127,6 +127,12 @@ def remove_request(self, req_id: str) -> None: def apply_staged_writes(self) -> None: return None + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, block_tables: BlockTables + ) -> None: + """Hook run after the KV cache tensors are allocated and bound.""" + return None + def capture_inner_cudagraphs( self, input_buffers: InputBuffers, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index f3f3c666e0b9..1f8af07898a1 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState @@ -134,47 +135,42 @@ def _get_mamba_group_info( self._mamba_spec = mamba_spec return self._mamba_group_ids, self._mamba_spec - def _ensure_align_ctx( - self, - kv_cache_config: KVCacheConfig, - mamba_group_ids: list[int], - block_tables: tuple[torch.Tensor, ...], - ) -> MambaSpecDecodeGPUContext: + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, block_tables: BlockTables + ) -> None: + if not self._align_mode: + return + mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) if self._mamba_state_copy_funcs is None: mamba_groups = get_mamba_groups(kv_cache_config) mamba_types = {spec.mamba_type for spec in mamba_groups} copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types) validate_mamba_state_copy_funcs(mamba_groups, copy_funcs) self._mamba_state_copy_funcs = copy_funcs - copy_funcs = self._mamba_state_copy_funcs - if self._mamba_ctx is None: - # Both SD and DS conv layouts support a >0 spec-decode shift: the - # fused pre-copy kernel (``_copy_mamba_state_block``) applies the - # ``token_bias = num_accepted - 1`` window shift per conv layout - # (SD: contiguous slice; DS: per-dim-row strided slice), matching - # the V1 ``get_conv_copy_spec`` semantics. - self._mamba_ctx = MambaSpecDecodeGPUContext.create( - max_num_reqs=self.max_num_reqs, - kv_cache_config=kv_cache_config, - copy_funcs=copy_funcs, - device=self.device, - make_buffer=lambda n, dtype: CpuGpuBuffer( - n, dtype=dtype, device=self.device - ), - ) - ctx = self._mamba_ctx - if not ctx.is_initialized: - forward_context = self.vllm_config.compilation_config.static_forward_context - # block_tables are batch-order slices of the persistent - # input_block_tables (stable data_ptr), so the metadata is captured - # once here and reused across steps. - ctx.initialize_from_forward_context( - kv_cache_config, - forward_context, - copy_funcs, - [block_tables[gid] for gid in mamba_group_ids], - ) - return ctx + # Both SD and DS conv layouts support a >0 spec-decode shift: the + # fused pre-copy kernel (``_copy_mamba_state_block``) applies the + # ``token_bias = num_accepted - 1`` window shift per conv layout + # (SD: contiguous slice; DS: per-dim-row strided slice), matching + # the V1 ``get_conv_copy_spec`` semantics. + self._mamba_ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=self.max_num_reqs, + kv_cache_config=kv_cache_config, + copy_funcs=self._mamba_state_copy_funcs, + device=self.device, + make_buffer=lambda n, dtype: CpuGpuBuffer( + n, dtype=dtype, device=self.device + ), + ) + # Under PP a step's postprocess runs after later steps re-gathered the + # batch-ordered input tables, so state copies read the per-request-slot + # tables. Aligned state indices are for the current batch only. + self._mamba_ctx.initialize_from_forward_context( + kv_cache_config, + self.vllm_config.compilation_config.static_forward_context, + self._mamba_state_copy_funcs, + [block_tables.block_tables[gid].gpu for gid in mamba_group_ids], + [block_tables.input_block_tables[gid] for gid in mamba_group_ids], + ) def preprocess_state( self, @@ -194,8 +190,9 @@ def preprocess_state( num_reqs = input_batch.num_reqs if num_reqs == 0: return - mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) - ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + _, mamba_spec = self._get_mamba_group_info(kv_cache_config) + ctx = self._mamba_ctx + assert ctx is not None # The state-advance + pre-copy kernels run every step; they fast-exit per # request when src_col < 0 or src_col == dst_col, so no copy happens on @@ -301,10 +298,8 @@ def prepare_attn( if hasattr(builder, "mamba_aligned_state_indices"): aligned_index_builders.append((group_idx, builder)) if aligned_index_builders: - ctx = self._ensure_align_ctx( - kv_cache_config, mamba_group_ids, block_tables - ) - all_group_indices = ctx.compute_aligned_state_indices( + assert self._mamba_ctx is not None + all_group_indices = self._mamba_ctx.compute_aligned_state_indices( input_batch.seq_lens, num_reqs ) for group_idx, builder in aligned_index_builders: diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 33a2b8b48017..a3b4eab3405a 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -392,8 +392,8 @@ def postprocess_mamba_fused_kernel( # Output: num_accepted_tokens update (for src==dst case) num_accepted_tokens_out_ptr, # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The - # per-request decision arrays are in req-state-slot order; the block table - # is in batch order, so HAS_IDX_MAPPING splits the two indexings. + # per-request decision arrays and the block tables are both in + # req-state-slot order. idx_mapping_ptr, # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation) num_reqs, @@ -476,10 +476,9 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return - bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx _copy_mamba_state_block( state_idx, - bt_row_idx, + req_idx, src_block_idx, dest_block_idx, accept_token_bias, @@ -610,7 +609,7 @@ def precopy_mamba_align_fused_kernel( token_bias = tl.load(token_bias_ptr + req_idx) _copy_mamba_state_block( state_idx, - batch_idx, + req_idx, src_col, dst_col, token_bias, @@ -806,8 +805,10 @@ class MambaSpecDecodeGPUContext: # Per-group block-table base addresses: int64[num_groups]. Populated in # initialize_from_forward_context from the persistent per-group block - # table tensors (whose data_ptr is stable across steps). + # table tensors (whose data_ptr is stable across steps). Rows are request + # state slots for the state copies and batch rows for the aligned indices. block_table_ptrs: torch.Tensor + aligned_index_block_table_ptrs: torch.Tensor block_table_stride_req: int = 0 # persistent output for the once-per-step, all-group aligned-index launch. @@ -900,6 +901,9 @@ def create( block_table_ptrs=torch.zeros( len(mamba_group_ids), dtype=torch.int64, device=device ), + aligned_index_block_table_ptrs=torch.zeros( + len(mamba_group_ids), dtype=torch.int64, device=device + ), aligned_state_indices=torch.empty( ( len(mamba_group_ids), @@ -924,6 +928,7 @@ def initialize_from_forward_context( forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], + aligned_index_block_tables: list[torch.Tensor] | None = None, ) -> None: """Extract and cache memory layout metadata from Mamba state tensors. @@ -955,8 +960,12 @@ def initialize_from_forward_context( mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to copy functions. block_tables: per-mamba-group persistent block-table tensors, in - the same order as `mamba_group_ids`. Their `data_ptr()` / - `stride(0)` are captured once for the kernel to index into. + the same order as `mamba_group_ids`, with one row per request + state slot. Their `data_ptr()` / `stride(0)` are captured once + for the state-copy kernels to index into. + aligned_index_block_tables: per-mamba-group persistent tables + read by `compute_aligned_state_indices`, with one row per + batch row. Defaults to `block_tables`. """ if self.is_initialized: @@ -968,6 +977,7 @@ def initialize_from_forward_context( forward_context, mamba_state_copy_funcs, block_tables, + aligned_index_block_tables or block_tables, ) def _populate_metadata( @@ -976,6 +986,7 @@ def _populate_metadata( forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], + aligned_index_block_tables: list[torch.Tensor], ) -> None: idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): @@ -1076,16 +1087,19 @@ def _populate_metadata( # `block_tables[i]` is the persistent 2D int32 block-table tensor for # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the # engine's lifetime, so we capture them once here. - assert len(block_tables) == self.num_groups, ( - f"expected {self.num_groups} block tables, got {len(block_tables)}" - ) - strides = {bt.stride(0) for bt in block_tables} + assert len(block_tables) == len(aligned_index_block_tables) == self.num_groups + strides = {bt.stride(0) for bt in (*block_tables, *aligned_index_block_tables)} assert len(strides) == 1, ( f"all mamba block tables must share stride(0), got {strides}" ) self.block_table_stride_req = int(next(iter(strides))) - for i, bt in enumerate(block_tables): + for i, (bt, aligned_bt) in enumerate( + zip(block_tables, aligned_index_block_tables) + ): self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) + self.aligned_index_block_table_ptrs[i] = _reinterpret_u64_as_i64( + aligned_bt.data_ptr() + ) self.is_initialized = True @@ -1107,7 +1121,7 @@ def compute_aligned_state_indices( block_rows = 32 grid = (triton.cdiv(num_reqs, block_rows),) get_aligned_state_indices_multi_group_kernel[grid]( - self.block_table_ptrs, + self.aligned_index_block_table_ptrs, seq_lens, self.aligned_state_indices, self.block_table_stride_req, From 2342343746093e5cd33799cd60186d10c4addea9 Mon Sep 17 00:00:00 2001 From: Hyukjoon Lee Date: Wed, 7 Oct 2026 12:10:54 +0900 Subject: [PATCH 108/952] [ROCm][DCP] Default MLA DCP verify to the round-robin asm decode where it applies (#59965) Signed-off-by: hyukjlee Co-authored-by: Claude Opus 5.5 --- .../attention/test_rocm_aiter_mla_dcp_cprr.py | 134 +++++++++++++++++- vllm/envs.py | 11 +- .../attention/backends/mla/rocm_aiter_mla.py | 21 ++- 3 files changed, 152 insertions(+), 14 deletions(-) diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py index 13b002cda05a..cd88ade8eebf 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py @@ -11,7 +11,11 @@ from __future__ import annotations +import sys +from types import SimpleNamespace + import pytest +import torch rocm_aiter_mla = pytest.importorskip( "vllm.v1.attention.backends.mla.rocm_aiter_mla", @@ -26,9 +30,15 @@ MIN_QLEN = rocm_aiter_mla._MIN_CPRR_QLEN -def test_dcp_verify_env_defaults_to_segmented(monkeypatch): +def test_dcp_verify_env_defaults_to_auto(monkeypatch): monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False) - assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "segmented" + assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "auto" + + +@pytest.mark.parametrize("route", ["auto", "asm", "segmented"]) +def test_dcp_verify_env_accepts_routes(monkeypatch, route): + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route) + assert route == rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY def test_dcp_verify_env_rejects_unknown_route(monkeypatch): @@ -116,6 +126,126 @@ def test_route_needs_gfx950(monkeypatch): ) +# -------------------------------------------------------------------------- +# AiterMLAMetadataBuilder -- the route the env value resolves to +# -------------------------------------------------------------------------- + +DCP = 8 + + +def _build(monkeypatch, route, gathered_heads, segmented_supported=True): + """Run the real builder __init__ on CPU at TP8/DCP8 with speculative decoding. + + Stubs only what needs a GPU, AITER, or a DCP group, so the route decision + under test is the one the builder actually makes. + """ + + def init_common_builder(self, *args, **kwargs): + self.num_heads = gathered_heads // DCP + self.dcp_world_size = DCP + self.reorder_batch_threshold = 5 + + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, "__init__", init_common_builder + ) + monkeypatch.setattr( + rocm_aiter_mla, + "_segmented_dcp_verify_supported", + lambda *args: segmented_supported, + ) + monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) + monkeypatch.setattr( + rocm_aiter_mla, "get_dcp_group", lambda: SimpleNamespace(rank_in_group=0) + ) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda device: SimpleNamespace(multi_processor_count=256), + ) + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace( + dtypes=SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16"), + get_mla_metadata_info_v1=lambda *args, **kwargs: tuple( + (1, torch.int32) for _ in range(6) + ), + ), + ) + if route is None: + monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False) + else: + monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route) + + config = SimpleNamespace( + speculative_config=SimpleNamespace(num_speculative_tokens=4), + parallel_config=SimpleNamespace( + decode_context_parallel_size=DCP, cp_kv_cache_interleave_size=1 + ), + model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16), + scheduler_config=SimpleNamespace(max_num_seqs=2), + cache_config=SimpleNamespace(cache_dtype="fp8_e4m3", num_gpu_blocks=None), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False) + ), + ) + return rocm_aiter_mla.AiterMLAMetadataBuilder( + kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), + layer_names=["layer.0"], + vllm_config=config, + device=torch.device("cpu"), + ) + + +@pytest.mark.parametrize("route", [None, "auto", "AUTO"]) +def test_auto_selects_cprr_for_k3_gathered_heads(on_gfx950, monkeypatch, route): + """K3 at TP8/DCP8 gathers 96 heads; auto must take cprr padded to 128.""" + builder = _build(monkeypatch, route, gathered_heads=96) + assert builder._asm_dcp_verify + assert builder._asm_dcp_verify_heads == 128 + assert builder._num_attention_heads == 128 + + +def test_auto_falls_back_to_segmented_above_largest_native(on_gfx950, monkeypatch): + builder = _build(monkeypatch, "auto", gathered_heads=256) + assert not builder._asm_dcp_verify + assert builder._supports_segmented_dcp_verify + + +def test_auto_without_any_route_names_the_missing_segmented_build( + on_gfx950, monkeypatch +): + with pytest.raises(ValueError) as exc: + _build(monkeypatch, "auto", gathered_heads=256, segmented_supported=False) + assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=auto" in str(exc.value) + assert "lacks segmented MLA decode" in str(exc.value) + + +@pytest.mark.parametrize("route", ["asm", "ASM"]) +def test_explicit_asm_is_case_insensitive(on_gfx950, monkeypatch, route): + """env_with_choices returns the raw string; ASM used to mean segmented.""" + assert _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify + + +def test_explicit_asm_refuses_unservable_heads(on_gfx950, monkeypatch): + with pytest.raises(ValueError) as exc: + _build(monkeypatch, "ASM", gathered_heads=256) + assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm" in str(exc.value) + assert "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented" in str(exc.value) + + +@pytest.mark.parametrize("route", ["segmented", "SEGMENTED"]) +def test_explicit_segmented_never_selects_cprr(on_gfx950, monkeypatch, route): + assert not _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify + + +def test_auto_off_gfx950_stays_segmented(monkeypatch): + import vllm.platforms.rocm as rocm + + monkeypatch.setattr(rocm, "on_gfx950", lambda: False) + assert not _build(monkeypatch, "auto", gathered_heads=96)._asm_dcp_verify + + # -------------------------------------------------------------------------- # qlen floor # -------------------------------------------------------------------------- diff --git a/vllm/envs.py b/vllm/envs.py index f51100fe8b02..851741d5253b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -146,7 +146,7 @@ VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" - VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["asm", "segmented"] = "segmented" + VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["auto", "asm", "segmented"] = "auto" VLLM_ROCM_USE_AITER_MHA: bool = True VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False @@ -1320,12 +1320,13 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1") ), # Kernel for causal multi-token (spec-decode) verify steps under decode - # context parallelism on gfx950: "asm" uses AITER's round-robin ASM - # decode, "segmented" the Triton segmented MLA path. + # context parallelism: "asm" uses AITER's round-robin ASM decode (gfx950), + # "segmented" the Triton segmented MLA path. "auto" (default) takes "asm" + # wherever it can serve the shape and "segmented" otherwise. "VLLM_ROCM_AITER_MLA_DCP_VERIFY": env_with_choices( "VLLM_ROCM_AITER_MLA_DCP_VERIFY", - "segmented", - ["asm", "segmented"], + "auto", + ["auto", "asm", "segmented"], case_sensitive=False, ), # Small-head (<16) AITER MLA decode kernel selection. Small head counts diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index c833a56a8383..32154105a96d 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -601,13 +601,14 @@ def __init__( parallel_config.decode_context_parallel_size, parallel_config.cp_kv_cache_interleave_size, ) + verify_route = envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY.lower() asm_dcp_verify_config = ( _asm_dcp_verify_configured( parallel_config.decode_context_parallel_size, parallel_config.cp_kv_cache_interleave_size, multi_token_decode=vllm_config.speculative_config is not None, ) - and envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "asm" + and verify_route != "segmented" ) super().__init__( kv_cache_spec, @@ -627,14 +628,20 @@ def __init__( self._asm_dcp_verify_heads = _asm_dcp_verify_heads( self.num_heads * self.dcp_world_size ) - if not self._asm_dcp_verify_heads: + if self._asm_dcp_verify_heads: + self._asm_dcp_verify = True + elif verify_route == "asm" or not supports_segmented_dcp_verify: raise ValueError( - "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm, but the round-robin asm " - f"decode has no kernel for {self.num_heads * self.dcp_world_size} " - f"DCP-gathered heads (native counts: {_NATIVE_CPRR_HEADS}). " - "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented." + f"VLLM_ROCM_AITER_MLA_DCP_VERIFY={verify_route}, but the " + "round-robin asm decode has no kernel for " + f"{self.num_heads * self.dcp_world_size} DCP-gathered heads " + f"(native counts: {_NATIVE_CPRR_HEADS}). " + + ( + "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented." + if supports_segmented_dcp_verify + else "This AITER build also lacks segmented MLA decode." + ) ) - self._asm_dcp_verify = True self._supports_segmented_dcp_verify = supports_segmented_dcp_verify self._mla_max_split_per_batch = 0 if self._asm_dcp_verify: From ef014bdf06d4fd468e2bcef63aed98bddaeda123 Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Wed, 7 Oct 2026 05:13:52 +0100 Subject: [PATCH 109/952] [Bugfix][Frontend] Return logprobs in Cohere Chat v2 responses (#59072) Signed-off-by: Mohammad Hijjawi Co-authored-by: Cyrus Leung --- .../cohere/test_serving_conversion.py | 25 +++++++++++++++++++ vllm/entrypoints/cohere/serving.py | 24 ++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index 58fa7fe7932b..e6696c26b42c 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -25,6 +25,7 @@ from vllm.entrypoints.cohere.protocol import ( CohereChatV2Request, CohereChatV2Response, + CohereLogprobItem, ) from vllm.entrypoints.cohere.serving import ( _FINISH_REASON_MAP, @@ -130,6 +131,8 @@ def _build_chat_completion_response( citations: list[Any] | None = None, usage: dict[str, Any] | None = None, kv_transfer_params: dict[str, Any] | None = None, + logprobs: list[tuple[str, float]] | None = None, + token_ids: list[int] | None = None, ) -> ChatCompletionResponse: message: dict[str, Any] = {"role": "assistant"} if content is not None: @@ -151,6 +154,12 @@ def _build_chat_completion_response( "message": message, "finish_reason": finish_reason, "stop_reason": stop_reason, + "logprobs": ( + {"content": [{"token": t, "logprob": lp} for t, lp in logprobs]} + if logprobs is not None + else None + ), + "token_ids": token_ids, } ], # ``usage`` is a required field on ChatCompletionResponse, but the @@ -689,6 +698,8 @@ def test_sampling_and_limits_mapped(self): assert result.top_k == 50 assert result.top_p == 0.95 assert result.logprobs is True + # Needed to fill ``token_ids`` in the Cohere logprobs response. + assert result.return_token_ids is True assert result.priority == 2 assert result.kv_transfer_params == {"x": 1} # ``chat_template_kwargs`` may be expanded by _apply_cohere_*; the @@ -1907,6 +1918,7 @@ def test_text_only(self): assert v2.message.tool_calls is None assert v2.message.tool_plan is None assert v2.usage is None + assert v2.logprobs is None def test_stop_sequence_finish_reason(self): serving = _serving() @@ -1990,6 +2002,19 @@ def test_kv_transfer_params_propagated(self): v2 = serving._chat_completion_to_v2(resp, _make_request()) assert v2.kv_transfer_params == {"k": 1} + def test_logprobs_propagated(self): + serving = _serving() + resp = _build_chat_completion_response( + content="Hello", + logprobs=[("Hel", -0.25), ("lo", -0.5)], + token_ids=[101, 102], + ) + v2 = serving._chat_completion_to_v2(resp, _make_request(logprobs=True)) + assert v2.logprobs == [ + CohereLogprobItem(text="Hel", token_ids=[101], logprobs=[-0.25]), + CohereLogprobItem(text="lo", token_ids=[102], logprobs=[-0.5]), + ] + # ====================================================================== # _build_usage diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index bdaa098659a7..98242cc55d09 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -58,6 +58,7 @@ CohereChatV2Request, CohereChatV2Response, CohereFinishReason, + CohereLogprobItem, CohereUsage, CohereUsageBilledUnits, CohereUsageTokens, @@ -79,6 +80,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, + ChatCompletionResponseChoice, ChatCompletionStreamResponse, ChatCompletionToolsParam, ChatMessage, @@ -562,6 +564,8 @@ def _build_base_chat_completion( frequency_penalty=request.frequency_penalty, presence_penalty=request.presence_penalty, logprobs=request.logprobs, + # Cohere logprob items carry token ids alongside the logprobs. + return_token_ids=request.logprobs, priority=request.priority or 0, kv_transfer_params=request.kv_transfer_params, chat_template_kwargs=request.chat_template_kwargs, @@ -1091,9 +1095,29 @@ def _chat_completion_to_v2( ), message=assistant_msg, usage=usage, + logprobs=self._build_logprobs(choice), kv_transfer_params=response.kv_transfer_params, ) + @staticmethod + def _build_logprobs( + choice: ChatCompletionResponseChoice, + ) -> list[CohereLogprobItem] | None: + if ( + choice.logprobs is None + or not choice.logprobs.content + or choice.token_ids is None + ): + return None + return [ + CohereLogprobItem( + text=entry.token, token_ids=[token_id], logprobs=[entry.logprob] + ) + for entry, token_id in zip( + choice.logprobs.content, choice.token_ids, strict=True + ) + ] + def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage: """Route response construction through the citation-carrying subclass. From aecb717f10bccbaae441d43910ae660a8f2e4e3c Mon Sep 17 00:00:00 2001 From: Turner Jabbour Date: Tue, 6 Oct 2026 22:16:53 -0600 Subject: [PATCH 110/952] [CI] Allowlist-shrink batch 2: wire 9 tests/models/ files into Buildkite, fix 2 stale bugs found (#59050) Signed-off-by: Turner Jabbour Co-authored-by: Claude Sonnet 5 --- .buildkite/test_areas/models_basic.yaml | 22 +++- tests/models/test_dspark_mla.py | 1 + .../transformers/test_layer_registry.py | 115 ++---------------- .../transformers/test_layer_registry_e2e.py | 113 +++++++++++++++++ tests/test_zen_cpu_platform_detection.py | 17 --- tools/pre_commit/test_tethering_allowlist.txt | 11 -- vllm/triton_utils/importing.py | 5 +- 7 files changed, 143 insertions(+), 141 deletions(-) create mode 100644 tests/models/transformers/test_layer_registry_e2e.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 865156f59449..04efae00fb17 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -68,8 +68,10 @@ steps: - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/transformers/test_layer_registry_e2e.py + - tests/models/test_deepseek_v41_replay_start.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -81,8 +83,15 @@ steps: source_file_dependencies: - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py + - tests/models/transformers/test_layer_registry_e2e.py + - tests/models/test_deepseek_v41_replay_start.py + - tests/models/test_deepseek_v41_decoder_replay_layers.py + - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/test_deepseek_v4_dspark_rocm.py + - tests/models/test_deepseek_v4_rocm_wo_a.py + - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" @@ -229,7 +238,7 @@ steps: key: basic-models-test-other-cpu depends_on: - image-build-cpu - timeout_in_minutes: 20 + timeout_in_minutes: 30 source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" @@ -239,6 +248,11 @@ steps: - tests/models/test_qwen3_5_mtp_config.py - tests/models/transformers/fusers/ - tests/model_executor/layers/test_activation.py + - tests/models/test_language_model_cache_is_weak.py + - tests/models/deepseek_v32/test_sequence_parallel.py + - tests/models/test_deepseek_v4_fi_moe_ep.py + - tests/models/test_dspark_mla.py + - tests/models/transformers/test_layer_registry.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index f77832d60854..76af0eb60c10 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -85,6 +85,7 @@ def test_dspark_markov_head_is_replicated( assert head.markov_w2.tp_size == 1 assert head.markov_w1.weight.shape == (128, 8) assert head.markov_w2.weight.shape == (128, 8) + head.markov_w2.quant_method.process_weights_after_loading(head.markov_w2) def fail_collective(*args, **kwargs): raise AssertionError("replicated Markov head must not invoke TP collectives") diff --git a/tests/models/transformers/test_layer_registry.py b/tests/models/transformers/test_layer_registry.py index 9e20e2ba54ad..555d87e379e1 100644 --- a/tests/models/transformers/test_layer_registry.py +++ b/tests/models/transformers/test_layer_registry.py @@ -7,6 +7,9 @@ is set and the symbol exists, and otherwise falls back to `vllm.model_executor.layers.`. These tests pin that contract and the logging that reports which source was used. + +CPU-only: no engine is built. The end-to-end counterpart that serves a model +through `vllm_runner` lives in `test_layer_registry_e2e.py` (GPU job). """ import logging @@ -14,7 +17,6 @@ import types import pytest -import torch from vllm.model_executor.models.transformers import layers @@ -42,16 +44,16 @@ def test_falls_back_to_vllm_when_disabled(monkeypatch, fake_hw_layernorm): assert layers._resolve("layernorm", "RMSNorm") is VllmRMSNorm -def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, caplog): +def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, caplog_vllm): """Enabled and available: the hw-agnostic class is used and logged.""" monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") - with caplog.at_level(logging.INFO): + with caplog_vllm.at_level(logging.INFO): resolved = layers._resolve("layernorm", "RMSNorm") assert resolved is fake_hw_layernorm.RMSNorm - assert "Using hw-agnostic layer: RMSNorm" in caplog.text + assert "Using hw-agnostic layer: RMSNorm" in caplog_vllm.text -def test_falls_back_when_symbol_missing(monkeypatch, caplog): +def test_falls_back_when_symbol_missing(monkeypatch, caplog_vllm): """Enabled but the symbol is not ported: fall back to vLLM and warn.""" monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") # A hw-agnostic module without the requested attribute triggers fallback. @@ -59,10 +61,10 @@ def test_falls_back_when_symbol_missing(monkeypatch, caplog): monkeypatch.setitem(sys.modules, HW_MODULE, empty) from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm - with caplog.at_level(logging.WARNING): + with caplog_vllm.at_level(logging.WARNING): resolved = layers._resolve("layernorm", "RMSNorm") assert resolved is VllmRMSNorm - assert "falling back to default" in caplog.text + assert "falling back to default" in caplog_vllm.text def test_act_and_mul_falls_back_for_unknown_activation( @@ -76,102 +78,3 @@ def test_act_and_mul_falls_back_for_unknown_activation( from vllm.model_executor.layers.activation import GeluAndMul assert isinstance(layers.get_act_and_mul_fn("gelu"), GeluAndMul) - - -@pytest.fixture(scope="module") -def tiny_llama_path(tmp_path_factory): - """A randomly-initialized microscopic Llama saved to disk (with an ungated - tokenizer) so vLLM can load it like any local checkpoint.""" - from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM - - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") - config = LlamaConfig( - vocab_size=tokenizer.vocab_size, - hidden_size=128, - intermediate_size=256, - num_hidden_layers=2, - num_attention_heads=4, - rms_norm_eps=1e-6, - hidden_act="silu", - ) - torch.manual_seed(0) - model = LlamaForCausalLM(config) - - path = tmp_path_factory.mktemp("tiny_llama") - model.save_pretrained(path) - tokenizer.save_pretrained(path) - return str(path) - - -# Registered names of the layers the backend can -# currently route to hw-agnostic implementations. -_COVERED_LAYERS = ("rms_norm", "silu_and_mul") - - -def _layer_providers(model) -> dict[str, str]: - """Map each covered layer type present in the model to the provider its - implementation came from (``hw_agnostic`` or ``vllm``). - """ - - def provider_of(module) -> str | None: - for cls in type(module).__mro__: - if "hw_agnostic.layers" in cls.__module__: - return "hw_agnostic" - if ".model_executor.layers." in cls.__module__: - return "vllm" - return None - - providers: dict[str, str] = {} - for module in model.modules(): - name = getattr(module, "name", None) - if ( - name in _COVERED_LAYERS - and name not in providers - and (prov := provider_of(module)) is not None - ): - providers[name] = prov - return providers - - -def _serve(vllm_runner, model_path, prompts): - """Serve the model through the backend; return (layer_providers, logprobs).""" - with vllm_runner( - model_path, - model_impl="transformers", - max_model_len=64, - enforce_eager=True, - gpu_memory_utilization=0.3, - ) as runner: - assert runner.llm.llm_engine.model_config.using_transformers_backend() - providers = runner.apply_model(_layer_providers)[0] - outputs = runner.generate_greedy_logprobs( - prompts, max_tokens=32, num_logprobs=5 - ) - return providers, outputs - - -def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path): - """Serving the tiny model with hw-agnostic layers matches the vLLM baseline.""" - # spawn: worker re-imports layers with the env set (see docstring). - monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") - # apply_model pickles the introspection function. - monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - from ..utils import check_logprobs_close - - prompts = ["The capital of France is", "vLLM is"] - - monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0") - vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts) - # Both replaceable layers present in a Llama block must be vLLM's here. - assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"} - - monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") - hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts) - assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"} - - check_logprobs_close( - outputs_0_lst=vllm_outputs, - outputs_1_lst=hw_outputs, - name_0="vllm", - name_1="hw_agnostic", - ) diff --git a/tests/models/transformers/test_layer_registry_e2e.py b/tests/models/transformers/test_layer_registry_e2e.py new file mode 100644 index 000000000000..13d4b80a16e7 --- /dev/null +++ b/tests/models/transformers/test_layer_registry_e2e.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end check of the Transformers backend's hw-agnostic layers. + +Serves a tiny Llama through `vllm_runner` with and without +`VLLM_USE_HW_AGNOSTIC` and checks the layer providers and logprobs. Needs a +real engine, so it runs in the GPU job. The CPU-only unit tests for the +resolution logic (`layers._resolve`) are in `test_layer_registry.py`. +""" + +import pytest +import torch + + +@pytest.fixture(scope="module") +def tiny_llama_path(tmp_path_factory): + """A randomly-initialized microscopic Llama saved to disk (with an ungated + tokenizer) so vLLM can load it like any local checkpoint.""" + from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM + + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") + config = LlamaConfig( + vocab_size=tokenizer.vocab_size, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + rms_norm_eps=1e-6, + hidden_act="silu", + ) + torch.manual_seed(0) + model = LlamaForCausalLM(config) + + path = tmp_path_factory.mktemp("tiny_llama") + model.save_pretrained(path) + tokenizer.save_pretrained(path) + return str(path) + + +# Registered names of the layers the backend can +# currently route to hw-agnostic implementations. +_COVERED_LAYERS = ("rms_norm", "silu_and_mul") + + +def _layer_providers(model) -> dict[str, str]: + """Map each covered layer type present in the model to the provider its + implementation came from (``hw_agnostic`` or ``vllm``). + """ + + def provider_of(module) -> str | None: + for cls in type(module).__mro__: + if "hw_agnostic.layers" in cls.__module__: + return "hw_agnostic" + if ".model_executor.layers." in cls.__module__: + return "vllm" + return None + + providers: dict[str, str] = {} + for module in model.modules(): + name = getattr(module, "name", None) + if ( + name in _COVERED_LAYERS + and name not in providers + and (prov := provider_of(module)) is not None + ): + providers[name] = prov + return providers + + +def _serve(vllm_runner, model_path, prompts): + """Serve the model through the backend; return (layer_providers, logprobs).""" + with vllm_runner( + model_path, + model_impl="transformers", + max_model_len=64, + enforce_eager=True, + gpu_memory_utilization=0.3, + ) as runner: + assert runner.llm.llm_engine.model_config.using_transformers_backend() + providers = runner.apply_model(_layer_providers)[0] + outputs = runner.generate_greedy_logprobs( + prompts, max_tokens=32, num_logprobs=5 + ) + return providers, outputs + + +def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path): + """Serving the tiny model with hw-agnostic layers matches the vLLM baseline.""" + # spawn: `layers` resolves its symbols at import time, so the worker must + # re-import it with VLLM_USE_HW_AGNOSTIC already set (fork would inherit + # the parent's already-resolved module). + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + # apply_model pickles the introspection function. + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + from ..utils import check_logprobs_close + + prompts = ["The capital of France is", "vLLM is"] + + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0") + vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts) + # Both replaceable layers present in a Llama block must be vLLM's here. + assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"} + + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1") + hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts) + assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"} + + check_logprobs_close( + outputs_0_lst=vllm_outputs, + outputs_1_lst=hw_outputs, + name_0="vllm", + name_1="hw_agnostic", + ) diff --git a/tests/test_zen_cpu_platform_detection.py b/tests/test_zen_cpu_platform_detection.py index 505d2d3bfda2..b8a39499238e 100644 --- a/tests/test_zen_cpu_platform_detection.py +++ b/tests/test_zen_cpu_platform_detection.py @@ -13,23 +13,6 @@ ) -@pytest.fixture -def caplog_vllm(caplog): - """`caplog`, but also captures vLLM's loggers. - - vllm/logger.py configures the "vllm" logger with propagate=False, so - records from child loggers like "vllm.platforms" reach vLLM's own - handler but never bubble up to the root logger where plain `caplog` - listens. Attach caplog's handler directly to "vllm" to see them too. - """ - logger = logging.getLogger("vllm") - logger.addHandler(caplog.handler) - try: - yield caplog - finally: - logger.removeHandler(caplog.handler) - - def test_is_amd_zen_cpu_detects_amd_with_avx512(): cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw" with ( diff --git a/tools/pre_commit/test_tethering_allowlist.txt b/tools/pre_commit/test_tethering_allowlist.txt index d7f19ec79a25..e738a79154db 100644 --- a/tools/pre_commit/test_tethering_allowlist.txt +++ b/tools/pre_commit/test_tethering_allowlist.txt @@ -48,17 +48,6 @@ tests/kernels/test_shuffle_rows.py # explicitly --ignore'd in kernels.yaml with tests/lora/test_deepseekv2_tp.py # --ignore'd from sharded LoRA run; no dedicated *_tp job tests/lora/test_qwen3moe_tp.py # --ignore'd from sharded LoRA run; no dedicated *_tp job -# tests/models/ - the models_*.yaml jobs run specific paths, not all of models/. -tests/models/deepseek_v32/test_sequence_parallel.py # no models_*.yaml job collects this path -tests/models/test_deepseek_v4_dspark_rocm.py # added by #52362 (2026-09-21), no models_*.yaml job collects this path -tests/models/test_deepseek_v4_fi_moe_ep.py # no models_*.yaml job collects this path -tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py # no models_*.yaml job collects this path -tests/models/test_deepseek_v4_rocm_wo_a.py # added by #54894 (2026-09-20), no models_*.yaml job collects this path -tests/models/test_deepseek_v41_replay_start.py # added by #56227 (2026-09-18), no models_*.yaml job collects this path -tests/models/test_dspark_mla.py # no models_*.yaml job collects this path -tests/models/test_language_model_cache_is_weak.py # no models_*.yaml job collects this path -tests/models/transformers/test_layer_registry.py # no models_*.yaml job collects this path - # tests/v1/ - the v1 subtree is otherwise covered by many `pytest v1/...` jobs. tests/v1/ec_connector/integration/test_epd_correctness.py # run_epd_correctness_test.sh exists but no Buildkite job invokes it tests/v1/ec_connector/integration/test_nixl_failure.py # 2-GPU + NIXL E/PD integration test (#55290); no Buildkite job collects ec_connector/integration/ diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index 12f19cc52b8f..7e501c94f48d 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -3,11 +3,10 @@ import os import types -from importlib.metadata import version from importlib.util import find_spec from vllm.logger import init_logger -from vllm.platforms import current_platform +from vllm.platforms import current_platform, vllm_version_matches_substr from vllm.utils.math_utils import cdiv logger = init_logger(__name__) @@ -35,7 +34,7 @@ ] # Check Triton CPU - if "cpu" in version("vllm"): + if vllm_version_matches_substr("cpu"): if "cpu" in backends: HAS_TRITON = True # Suppress following warnings on CPU-only platforms From 43b4aaea3e40e20ef53ce2bfecfa1b72060bf2b5 Mon Sep 17 00:00:00 2001 From: Sriraj <168966222+sricursion@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:25:34 +0530 Subject: [PATCH 111/952] [Bugfix][Frontend] Return 4xx for client errors in /cohere/v2/chat (#60309) Signed-off-by: sricursion Co-authored-by: Claude --- tests/entrypoints/cohere/test_api_router.py | 28 ++++++++++++++++----- vllm/entrypoints/cohere/api_router.py | 23 ++++++----------- 2 files changed, 29 insertions(+), 22 deletions(-) diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index 1bb2151621c0..209ceb4e3ef4 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -50,7 +50,7 @@ from vllm.entrypoints.serve.exception_handling.handlers.vllm_error import ( vllm_error_handler, ) -from vllm.exceptions import VLLMError +from vllm.exceptions import VLLMError, VLLMValidationError from vllm.sampling_params import SamplingParams @@ -78,7 +78,7 @@ class _Handler: * a :class:`CohereChatV2Response` (non-streaming JSON path); * an async generator yielding SSE frames (streaming path); * an :class:`ErrorResponse` (error envelope path); or - * an exception (router-level 500 path). + * an exception (router-level error path). """ def __init__(self, result): @@ -371,6 +371,22 @@ def test_handler_exception_returns_500_envelope(self): body = r.json() assert body == {"message": "kaboom"} + @pytest.mark.parametrize( + "exc", + [ + VLLMValidationError("bad temperature", parameter="temperature"), + ValueError("bad temperature"), + ], + ids=["vllm_validation_error", "value_error"], + ) + def test_handler_client_error_returns_400_envelope(self, exc): + """Client-caused handler errors are 4xx, matching the OpenAI routes.""" + app = _build_app(handler=_Handler(exc)) + with TestClient(app) as client: + r = client.post("/cohere/v2/chat", json=_minimal_request_body()) + assert r.status_code == HTTPStatus.BAD_REQUEST + assert "bad temperature" in r.json()["message"] + def test_non_json_content_type_rejected(self): """The ``validate_json_request`` dependency raises ``RequestValidationError`` (HTTP 422) for non-JSON content @@ -525,10 +541,10 @@ def test_render_exception_returns_500_envelope(self): assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR assert r.json() == {"message": "tokenizer exploded"} - def test_conversion_exception_returns_500_envelope(self): + def test_conversion_client_error_returns_400_envelope(self): """Conversion runs inside the same ``try`` as the render call, so - a bad v2 body that slips past Pydantic still yields the Cohere - error envelope rather than an unhandled 500 with vLLM's shape. + a bad v2 body that slips past Pydantic still yields a 4xx in the + Cohere error envelope rather than an unhandled 500 with vLLM's shape. """ class _Boom: @@ -538,7 +554,7 @@ def to_chat_completion_request(self, request): app = _build_render_app(_Boom(), _RenderHandler(_generate_request())) with TestClient(app) as client: r = client.post("/cohere/v2/chat/render", json=_minimal_request_body()) - assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert r.status_code == HTTPStatus.BAD_REQUEST assert r.json() == {"message": "unconvertible message"} def test_invalid_body_returns_422(self): diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index db6366df7121..5362b1d9e95d 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -40,6 +40,9 @@ import vllm.envs as envs from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -138,15 +141,9 @@ async def chat_v2(request: CohereChatV2Request, raw_request: Request): try: result = await handler.create_chat_v2(request, raw_request) - except Exception as e: # noqa: BLE001 - report as 500 for parity + except Exception as e: logger.exception("Error in /cohere/v2/chat: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError( - message=sanitize_message(str(e)), - id=_request_id(raw_request), - ).model_dump(exclude_none=True), - ) + return _error_response(create_error_response(e), raw_request) match result: case ErrorResponse(): @@ -194,15 +191,9 @@ async def chat_v2_render(request: CohereChatV2Request, raw_request: Request): try: chat_request = handler.to_chat_completion_request(request) result = await render_handler.render_chat_request(chat_request) - except Exception as e: # noqa: BLE001 - report as 500 for parity + except Exception as e: logger.exception("Error in /cohere/v2/chat/render: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=CohereError( - message=sanitize_message(str(e)), - id=_request_id(raw_request), - ).model_dump(exclude_none=True), - ) + return _error_response(create_error_response(e), raw_request) if isinstance(result, ErrorResponse): return _error_response(result, raw_request) From cbee82324c4ac2e780514122604f8f5fdb7c1d21 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 7 Oct 2026 14:22:10 +0800 Subject: [PATCH 112/952] [Docs][RL] Move pause/resume and batch invariance examples to examples/rl (#59949) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: wang.yuqi Co-authored-by: Claude Opus 5.5 --- .github/CODEOWNERS | 1 - docs/usage/reproducibility.md | 2 +- .../batch_invariance/reproducibility_offline.py | 0 .../{features => rl}/pause_resume/data_parallel_pause_resume.py | 0 examples/{features => rl}/pause_resume/pause_resume_offline.py | 0 mkdocs.yaml | 2 ++ 6 files changed, 3 insertions(+), 2 deletions(-) rename examples/{features => rl}/batch_invariance/reproducibility_offline.py (100%) rename examples/{features => rl}/pause_resume/data_parallel_pause_resume.py (100%) rename examples/{features => rl}/pause_resume/pause_resume_offline.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ac64cf1225d4..71d501eabe7c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -235,7 +235,6 @@ mkdocs.yaml @hmellor /tests/entrypoints/serve/dev/rlhf @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 -/examples/features/pause_resume @aoshen02 /docs/training/rlhf.md @aoshen02 /docs/training/async_rl.md @aoshen02 /docs/training/weight_transfer @aoshen02 diff --git a/docs/usage/reproducibility.md b/docs/usage/reproducibility.md index 680791bbe24a..ef34855ef4ff 100644 --- a/docs/usage/reproducibility.md +++ b/docs/usage/reproducibility.md @@ -7,7 +7,7 @@ reproducible results: or enable [batch invariance](../features/batch_invariance.md) to make the outputs insensitive to scheduling. - In online mode, you can only enable [batch invariance](../features/batch_invariance.md). -Example: [examples/features/batch_invariance/reproducibility_offline.py](../../examples/features/batch_invariance/reproducibility_offline.py) +Example: [examples/rl/batch_invariance/reproducibility_offline.py](../../examples/rl/batch_invariance/reproducibility_offline.py) !!! warning diff --git a/examples/features/batch_invariance/reproducibility_offline.py b/examples/rl/batch_invariance/reproducibility_offline.py similarity index 100% rename from examples/features/batch_invariance/reproducibility_offline.py rename to examples/rl/batch_invariance/reproducibility_offline.py diff --git a/examples/features/pause_resume/data_parallel_pause_resume.py b/examples/rl/pause_resume/data_parallel_pause_resume.py similarity index 100% rename from examples/features/pause_resume/data_parallel_pause_resume.py rename to examples/rl/pause_resume/data_parallel_pause_resume.py diff --git a/examples/features/pause_resume/pause_resume_offline.py b/examples/rl/pause_resume/pause_resume_offline.py similarity index 100% rename from examples/features/pause_resume/pause_resume_offline.py rename to examples/rl/pause_resume/pause_resume_offline.py diff --git a/mkdocs.yaml b/mkdocs.yaml index 8c712257eecd..c272e6a68754 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -102,6 +102,8 @@ plugins: features/quantization/int8.md: features/quantization/llm_compressor/int8_w8a8.md serving/openai_compatible_server.md: serving/online_serving/README.md examples/others/lmcache.md: examples/disaggregated/lmcache.md + examples/features/batch_invariance.md: examples/rl/batch_invariance.md + examples/features/pause_resume.md: examples/rl/pause_resume.md markdown_extensions: - attr_list From e6fc81bc7892f2f58c0e347a701fc060ceef44bb Mon Sep 17 00:00:00 2001 From: Tommy Date: Wed, 7 Oct 2026 03:32:00 -0400 Subject: [PATCH 113/952] [Bugfix][Scheduler] Let pooling chunked prefill use full context (#48039) Signed-off-by: Tommy Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/v1/core/test_scheduler.py | 19 +++++++++++++++++++ tests/v1/core/utils.py | 3 +++ vllm/v1/core/sched/scheduler.py | 10 ++++++++-- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index f78db1917db0..7a8a1a3e4ea5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -35,6 +35,7 @@ MultiModalKwargsItem, PlaceholderRange, ) +from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.utils.hashing import sha256 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager @@ -279,6 +280,24 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None): assert scheduler.running[i] == request +def test_pooling_chunked_prefill_can_finish_at_max_model_len(): + scheduler = create_scheduler( + max_num_seqs=1, max_model_len=8, max_num_batched_tokens=4, runner="pooling" + ) + assert scheduler.max_model_len == 8 + request = Request( + request_id="pool", + prompt_token_ids=[1] * 8, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + scheduler.add_request(request) + + scheduler.schedule() + # The final chunk must reach max_model_len, not stop one token short. + assert scheduler.schedule().num_scheduled_tokens == {"pool": 4} + + def test_scheduler_stats_route_to_existing_output_client(): scheduler = create_scheduler() request = create_requests(num_requests=1)[0] diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 49aaca3d7c6a..b5e1bf89329e 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -18,6 +18,7 @@ SpeculativeConfig, VllmConfig, ) +from vllm.config.model import RunnerOption from vllm.config.scheduler import SchedulerPolicy from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -78,6 +79,7 @@ def create_scheduler( ec_role: str | None = None, use_v2_model_runner: bool | None = None, kv_cache_spec: KVCacheSpec | None = None, + runner: RunnerOption = "auto", per_request_spec_decode_metrics: str = "none", scheduling_policy: SchedulerPolicy = "fcfs", diffusion_canvas_length: int | None = None, @@ -109,6 +111,7 @@ def create_scheduler( # The scheduler reads model_config.max_model_len, not the # SchedulerConfig one, so both must agree. max_model_len=max_model_len, + runner=runner, ) if mm_encoder_only: model_config.multimodal_config.mm_encoder_only = True diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8bdf91da86da..f908a31a1b20 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -147,9 +147,15 @@ def __init__( self.kv_events_config is not None and self.kv_events_config.enable_kv_cache_events ) - # Diffusion models may not sample any tokens for a denoising step. + # Pooling and diffusion models do not append a sampled token after + # prefill, so they should not reserve one context slot while scheduling. self.num_sampled_tokens_per_step = ( - 1 if not vllm_config.model_config.is_diffusion else 0 + 0 + if ( + vllm_config.model_config.runner_type == "pooling" + or vllm_config.model_config.is_diffusion + ) + else 1 ) # Create KVConnector for the Scheduler. Note that each Worker From 7436a7f119dedd0c4d5f4d745aa6c2305ee82252 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 7 Oct 2026 15:45:15 +0800 Subject: [PATCH 114/952] [Rust Frontend] Backtrack from safe text at a complete marker (#59563) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/parser/src/tool/json/mod.rs | 11 ++++++++-- rust/src/parser/src/tool/json/qwen.rs | 17 ++++++++++++++ rust/src/parser/src/unified/hf/parser.rs | 2 +- rust/src/parser/src/utils.rs | 28 ++++++++++++++++++++---- 4 files changed, 51 insertions(+), 7 deletions(-) diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index 1bc50588736c..487be423aef2 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -378,9 +378,16 @@ fn safe_text_event( input: &mut JsonToolInput<'_>, config: JsonToolCallConfig, ) -> ModalResult { + // Stop only where a tool call can start: with exact whitespace, a bare + // start marker without it is text. + let with_whitespace = |marker: &str| match config.marker_whitespace { + JsonToolCallWhitespace::Optional => marker.to_string(), + JsonToolCallWhitespace::Exact(whitespace) => format!("{marker}{whitespace}"), + }; + let start_marker = with_whitespace(config.start_marker); match config.framed_start_marker { - Some(marker) => safe_text_len_mul(input, &[marker, config.start_marker]), - None => safe_text_len(input, config.start_marker), + Some(marker) => safe_text_len_mul(input, &[with_whitespace(marker), start_marker]), + None => safe_text_len(input, start_marker), } .map(|len| JsonToolCallEvent::Text { len }) } diff --git a/rust/src/parser/src/tool/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs index ca7c5919cfcf..c4db8c2fbc8b 100644 --- a/rust/src/parser/src/tool/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -205,6 +205,23 @@ mod tests { assert!(output.calls().is_empty()); } + #[test] + fn qwen_xml_streams_start_marker_without_newline_as_text() { + for chunk in [ + "{\"name\":\"get_weather\"", + "hi\n{\"name\":\"get_weather\"", + ] { + let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut output = ToolParserOutput::default(); + + parser.parse_into(chunk, &mut output).unwrap(); + + // Not a tool-call start: released without waiting for the end of the stream. + assert_eq!(output.normal_text(), chunk); + assert!(output.calls().is_empty()); + } + } + #[test] fn qwen_xml_requires_newline_before_tool_call_end() { let mut parser = Qwen3XmlToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/unified/hf/parser.rs b/rust/src/parser/src/unified/hf/parser.rs index f523c0f31a76..5d5065f109b6 100644 --- a/rust/src/parser/src/unified/hf/parser.rs +++ b/rust/src/parser/src/unified/hf/parser.rs @@ -599,7 +599,7 @@ fn parse_next_event( match safe_text_len_mul(input, &markers) { Ok(_) => return Ok(HfEvent::Text), // A marker, complete or partial, starts here (or the input is empty). - Err(ErrMode::Incomplete(_)) => {} + Err(ErrMode::Backtrack(_) | ErrMode::Incomplete(_)) => {} Err(error) => return Err(error), } diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs index 1324e297fbe5..cb7bdce284c5 100644 --- a/rust/src/parser/src/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -118,7 +118,10 @@ pub fn safe_text_len<'i, I: MarkerStream<'i>, M: MarkerLike>( /// the end of the input for plain markers, and only a genuine special token's /// incomplete marker for guarded ones. /// -/// Returns the text length in bytes, and advances the input. +/// Returns the text length in bytes, and advances the input. Fails with +/// `Backtrack` when a complete marker starts the input, so another branch can +/// consume it, and with `Incomplete` when the input is empty or starts with a +/// marker that needs more input to decide. pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>( input: &mut I, markers: &[M], @@ -128,12 +131,18 @@ pub fn safe_text_len_mul<'i, I: MarkerStream<'i>, M: MarkerLike>( } let start = input.offset(); - let stop = match scan_markers(input, markers, start)? { - Scan::Found(at) | Scan::Pending(at) | Scan::Clear(at) => at, + let (stop, found) = match scan_markers(input, markers, start)? { + Scan::Found(at) => (at, true), + Scan::Pending(at) | Scan::Clear(at) => (at, false), }; let emit_len = stop - start; if emit_len == 0 { - return incomplete(); + // A complete marker starts here: there is no text, whatever input follows. + return if found { + Err(ErrMode::Backtrack(ContextError::new())) + } else { + incomplete() + }; } input.next_slice(emit_len); @@ -529,6 +538,17 @@ mod tests { assert!(matches!(error, ErrMode::Incomplete(_))); } + #[test] + fn safe_text_len_backtracks_at_complete_marker() { + let mut input = Partial::new("body"); + let checkpoint = input.checkpoint(); + + let error = safe_text_len(&mut input, "").unwrap_err(); + + assert!(matches!(error, ErrMode::Backtrack(_))); + assert_eq!(input.offset_from(&checkpoint), 0); + } + #[test] fn safe_text_len_mul_stops_before_earliest_marker() { let mut input = Partial::new("hello<|tool_call>"); From 90f4936d9ee05f6108a401a3acc5c5631f5835f7 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 10:21:47 +0200 Subject: [PATCH 115/952] [CI/Build] Widen the ragged prefill scoring tolerance, drop batch invariance (#60192) Signed-off-by: Stefan Koncarevic --- tests/v1/sample/test_logprobs.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 813bea0bf3aa..a9b73f0ca2a7 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1408,10 +1408,6 @@ def test_prompt_logprobs_with_chunking_and_preemption(): def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): """Per-row scores stay row-aligned across chunked prefill and preemption.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - # The ragged check below compares a second generate call against the first - # one. Default kernel selection varies between processes, which moves bf16 - # logprob tails by up to 0.2; under this the two runs are bit-identical. - monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") prompts = [ "The following numbers of the sequence " @@ -1474,7 +1470,10 @@ def make_params(row_ids): assert float(scores.max()) <= 1e-3, "logprobs must be <= 0" # Ragged rows: each row keeps a prefix of the candidates and must - # reproduce that prefix of the full scores, with -inf padding. + # reproduce that prefix of the full scores, with -inf padding. On ROCm + # the second generate call picks its own kernels, which moves bf16 + # tails by up to 0.19; a misaligned row is off by nats either way. + ragged_atol = 0.2 if current_platform.is_rocm() else 1e-2 ragged_params = make_params(lambda j: candidate_ids(j)[: 1 + j % 4]) ragged_outputs = vllm_model.llm.generate(token_prompts, ragged_params) for output, ragged_output in zip(outputs, ragged_outputs): @@ -1483,7 +1482,7 @@ def make_params(row_ids): n = 1 + j % 4 expected[j, :n] = output.prompt_token_id_logprobs[j, :n] np.testing.assert_allclose( - ragged_output.prompt_token_id_logprobs, expected, atol=1e-2 + ragged_output.prompt_token_id_logprobs, expected, atol=ragged_atol ) # The row count is checked against the prompt at admission. From bc8dd4ad275c838d6315389432b2010d9348c900 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 7 Oct 2026 16:28:27 +0800 Subject: [PATCH 116/952] [Rust Frontend] Share `SchemaRoot` between argument coercion and grammars (#59408) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/parser/src/lib.rs | 1 + .../parser/src/output_grammar/arguments.rs | 236 +------- rust/src/parser/src/schema.rs | 499 +++++++++++++++++ rust/src/parser/src/tool/parameters.rs | 505 ++++++++++-------- rust/src/parser/src/tool/qwen_coder.rs | 35 +- .../src/unified/kimi_k3/structural_tag.rs | 22 +- 6 files changed, 857 insertions(+), 441 deletions(-) create mode 100644 rust/src/parser/src/schema.rs diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs index fd4aa0aa7f71..26256d332103 100644 --- a/rust/src/parser/src/lib.rs +++ b/rust/src/parser/src/lib.rs @@ -5,6 +5,7 @@ pub mod output_grammar; pub mod reasoning; +pub(crate) mod schema; pub mod tool; pub mod unified; pub(crate) mod utils; diff --git a/rust/src/parser/src/output_grammar/arguments.rs b/rust/src/parser/src/output_grammar/arguments.rs index a7d9fb307875..0f621789a3bc 100644 --- a/rust/src/parser/src/output_grammar/arguments.rs +++ b/rust/src/parser/src/output_grammar/arguments.rs @@ -25,14 +25,11 @@ //! - `minProperties` and `maxProperties` are not enforced. //! - `allOf` with more than one schema accepts any value. -use std::collections::HashSet; - -use serde_json::{Map, Value, json}; +use serde_json::{Value, json}; use xgrammar_structural_tag::format::{Format, JsonSchemaFormat}; -/// Bound on `$ref` and combinator nesting within one value, which keeps a long -/// chain of references from overflowing the stack. -const MAX_SCHEMA_DEPTH: usize = 32; +pub use crate::schema::JsonType; +use crate::schema::{self, OptionSource, SchemaOption, SchemaRoot}; /// The schema accepting any value. static ANY_SCHEMA: Value = Value::Bool(true); @@ -77,95 +74,22 @@ impl ParameterKey<'_> { } } -/// JSON type of a value option. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum JsonType { - String, - Number, - Boolean, - Null, - Object, - Array, -} - -impl JsonType { - const ALL: [Self; 6] = [ - Self::String, - Self::Number, - Self::Boolean, - Self::Null, - Self::Object, - Self::Array, - ]; - - /// The type's JSON name. Integers are numbers. - pub fn name(self) -> &'static str { - match self { - Self::String => "string", - Self::Number => "number", - Self::Boolean => "boolean", - Self::Null => "null", - Self::Object => "object", - Self::Array => "array", - } - } - - fn of_schema_type(name: &str) -> Option { - Some(match name { - "string" => Self::String, - "integer" | "number" => Self::Number, - "boolean" => Self::Boolean, - "null" => Self::Null, - "object" => Self::Object, - "array" => Self::Array, - _ => return None, - }) - } - - fn of_value(value: &Value) -> Self { - match value { - Value::String(_) => Self::String, - Value::Number(_) => Self::Number, - Value::Bool(_) => Self::Boolean, - Value::Null => Self::Null, - Value::Object(_) => Self::Object, - Value::Array(_) => Self::Array, - } - } - - /// Schema accepting every JSON value of this type. - fn any_schema(self) -> Value { - match self { - Self::Object => json!({ "type": "object", "additionalProperties": true }), - Self::Array => json!({ "type": "array", "items": true }), - _ => json!({ "type": self.name() }), - } - } -} - /// One option of a parameter value: a single JSON type under a narrowed /// schema, a constant, or any value of the type. pub struct ValueOption<'a> { - /// The option's JSON type. + /// The option's JSON Schema type. pub ty: JsonType, - source: OptionSource, + source: OptionSource<'a>, cx: &'a ArgumentContext<'a>, } -enum OptionSource { - /// A schema narrowed to [`ValueOption::ty`]. - Schema(Map), - /// A `const` or `enum` value. - Literal(Value), - /// Any value of [`ValueOption::ty`]. - Any, -} - impl ValueOption<'_> { /// The option as JSON text. pub fn json(&self) -> Format { match &self.source { - OptionSource::Schema(schema) => self.cx.json(Value::Object(schema.clone())), + OptionSource::Schema(schema) => { + self.cx.json(Value::Object(schema::narrow(schema, self.ty).into_owned())) + } OptionSource::Literal(value) => self.cx.json(json!({ "const": value })), OptionSource::Any => self.cx.json(self.ty.any_schema()), } @@ -216,26 +140,18 @@ pub struct ArgumentOptions { pub excludes: Vec, } -/// Reference targets expanded while collecting one value's options. -#[derive(Default)] -struct Expansion { - expanded: HashSet<*const Value>, - /// Whether a reference was skipped because its target already expanded. - skipped: bool, -} - /// Request state shared by the value options of one call. struct ArgumentContext<'a> { options: &'a ArgumentOptions, - /// The parameters schema, which local `$ref`s resolve against. - root: &'a Value, + root: SchemaRoot<'a>, } impl ArgumentContext<'_> { /// A JSON value under `schema`, which keeps the root's definitions so its /// local `$ref`s still resolve. fn json(&self, mut schema: Value) -> Format { - if let (Some(schema), Some(root)) = (schema.as_object_mut(), self.root.as_object()) { + if let (Some(schema), Some(root)) = (schema.as_object_mut(), self.root.schema().as_object()) + { for key in ["$defs", "definitions"] { if let Some(definitions) = root.get(key) { schema.entry(key).or_insert_with(|| definitions.clone()); @@ -250,122 +166,14 @@ impl ArgumentContext<'_> { }) } - fn resolve_ref(&self, reference: &str) -> Option<&Value> { - self.root.pointer(reference.strip_prefix('#')?) - } - - /// Follow `$ref`s and single-schema `allOf`s at the root. - fn resolve<'s>(&'s self, schema: &'s Value, depth: usize) -> &'s Value { - if depth > MAX_SCHEMA_DEPTH { - return &ANY_SCHEMA; - } - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { - return match self.resolve_ref(reference) { - Some(target) => self.resolve(target, depth + 1), - None => &ANY_SCHEMA, - }; - } - match schema.get("allOf").and_then(Value::as_array).map(Vec::as_slice) { - Some([schema]) => self.resolve(schema, depth + 1), - _ => schema, - } - } - /// Every option of a value under `schema`. - /// - /// Each reference target expands at most once. A union gains no values - /// from a repeated alternative, and a definition shared by several - /// alternatives would otherwise expand once per path to it, exponentially - /// in the nesting. A schema whose references only lead back into it - /// accepts any value. - fn value_options(&self, schema: &Value) -> Vec> { - let mut expansion = Expansion::default(); - let options = self.value_options_at(schema, 0, &mut expansion); - if options.is_empty() && expansion.skipped { - return self.any_value(); - } - options - } - - fn value_options_at( - &self, - schema: &Value, - depth: usize, - expansion: &mut Expansion, - ) -> Vec> { - let schema = match schema { - _ if depth > MAX_SCHEMA_DEPTH => return self.any_value(), - Value::Bool(false) => return vec![], - Value::Object(schema) => schema, - _ => return self.any_value(), - }; - if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { - let Some(target) = self.resolve_ref(reference) else { - return self.any_value(); - }; - if !expansion.expanded.insert(std::ptr::from_ref(target)) { - expansion.skipped = true; - return vec![]; - } - return self.value_options_at(target, depth + 1, expansion); - } - if let Some(value) = schema.get("const") { - return vec![self.literal(value)]; - } - if let Some(values) = schema.get("enum").and_then(Value::as_array) { - return values.iter().map(|value| self.literal(value)).collect(); - } - if let Some(options) = - schema.get("anyOf").or_else(|| schema.get("oneOf")).and_then(Value::as_array) - { - return options - .iter() - .flat_map(|option| self.value_options_at(option, depth + 1, expansion)) - .collect(); - } - if let Some(schemas) = schema.get("allOf").and_then(Value::as_array) { - return match schemas.as_slice() { - [schema] => self.value_options_at(schema, depth + 1, expansion), - _ => self.any_value(), - }; - } - match schema.get("type") { - Some(Value::String(name)) => self.typed(schema.clone(), name).into_iter().collect(), - Some(Value::Array(names)) => names - .iter() - .filter_map(Value::as_str) - .filter_map(|name| { - let mut schema = schema.clone(); - schema.insert("type".to_string(), Value::String(name.to_string())); - self.typed(schema, name) - }) - .collect(), - _ => self.any_value(), - } - } - - fn typed(&self, schema: Map, type_name: &str) -> Option> { - Some(ValueOption { - ty: JsonType::of_schema_type(type_name)?, - source: OptionSource::Schema(schema), - cx: self, - }) - } - - fn literal(&self, value: &Value) -> ValueOption<'_> { - ValueOption { - ty: JsonType::of_value(value), - source: OptionSource::Literal(value.clone()), - cx: self, - } - } - - fn any_value(&self) -> Vec> { - JsonType::ALL + fn value_options<'s>(&'s self, schema: &'s Value) -> Vec> { + let root: SchemaRoot<'s> = self.root; + root.options(schema) .into_iter() - .map(|ty| ValueOption { + .map(|SchemaOption { ty, source }| ValueOption { ty, - source: OptionSource::Any, + source, cx: self, }) .collect() @@ -377,12 +185,12 @@ impl ArgumentContext<'_> { pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format { let cx = ArgumentContext { options, - root: schema, + root: SchemaRoot::new(schema), }; let parameter = |key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema)); - let schema = cx.resolve(schema, 0); + let schema = cx.root.resolved(); let (properties, additional) = match schema { Value::Object(schema) => { let additional = match schema.get("additionalProperties") { @@ -492,6 +300,8 @@ pub fn group_by<'o, 'a, K: PartialEq>( #[cfg(test)] mod tests { + use serde_json::Map; + use super::*; use crate::output_grammar::test_utils::outline; use expect_test::{Expect, expect}; @@ -617,7 +427,7 @@ mod tests { or `celsius` `fahrenheit` - tag `` json(integer(minimum=1)) any_order `` + tag `` json(integer(minimum=1)) any_order `` "#]], ); } @@ -639,13 +449,13 @@ mod tests { sequence optional or - tag `` json(integer) `` + tag `` json(integer) `` tag `` text excluding [``] `` tag `` json(null) `` optional or tag `` `fast` `` - tag `` json(1) `` + tag `` json(1) `` or tag `` /[a-z]+/ `` tag `` json(string[]) `` diff --git a/rust/src/parser/src/schema.rs b/rust/src/parser/src/schema.rs new file mode 100644 index 000000000000..c27427de2e2c --- /dev/null +++ b/rust/src/parser/src/schema.rs @@ -0,0 +1,499 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Views of tool parameter JSON schemas. +//! +//! Tool parsers coerce generated parameter values by their schema, and output +//! grammars constrain the same values by it. Both read a parameter's schema +//! through [`SchemaRoot`], so they agree on which JSON types the value may take. + +use std::borrow::Cow; +use std::collections::HashSet; + +use serde_json::{Map, Value, json}; + +/// Bound on `$ref` and combinator nesting within one value, which keeps a long +/// chain of references from overflowing the stack. +const MAX_SCHEMA_DEPTH: usize = 32; + +/// The schema accepting any value. +static ANY_SCHEMA: Value = Value::Bool(true); + +/// A tool's parameters schema, which local `$ref`s resolve against. +#[derive(Debug, Clone, Copy)] +pub(crate) struct SchemaRoot<'a> { + schema: &'a Value, +} + +/// Reference targets expanded while collecting one value's options. +#[derive(Default)] +struct Expansion { + expanded: HashSet<*const Value>, + /// Whether a reference was skipped because its target already expanded. + skipped: bool, +} + +/// Roots are equal when they are the same schema, which a pointer compare +/// decides without walking it. +impl PartialEq for SchemaRoot<'_> { + fn eq(&self, other: &Self) -> bool { + std::ptr::eq(self.schema, other.schema) + } +} + +impl Eq for SchemaRoot<'_> {} + +/// One option of a value: a single JSON type under a narrowed schema, a +/// constant, or any value of the type. +#[derive(Debug, Clone)] +pub(crate) struct SchemaOption<'a> { + pub(crate) ty: JsonType, + pub(crate) source: OptionSource<'a>, +} + +/// Where a [`SchemaOption`] comes from. +#[derive(Debug, Clone)] +pub(crate) enum OptionSource<'a> { + /// A schema, of which the option takes the values of its type; [`narrow`] + /// sets that type on it. + Schema(&'a Map), + /// A `const` or `enum` value. + Literal(&'a Value), + /// Any value of the option's type. + Any, +} + +/// JSON Schema type of a value option. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JsonType { + String, + Integer, + Number, + Boolean, + Null, + Object, + Array, +} + +impl<'a> SchemaRoot<'a> { + pub(crate) fn new(schema: &'a Value) -> Self { + Self { schema } + } + + /// The parameters schema itself. + pub(crate) fn schema(&self) -> &'a Value { + self.schema + } + + /// The parameters schema after following `$ref`s and single-schema + /// `allOf`s. + pub(crate) fn resolved(&self) -> &'a Value { + self.resolve(self.schema) + } + + /// Follow `$ref`s and single-schema `allOf`s. A reference that does not + /// resolve, or a chain deeper than [`MAX_SCHEMA_DEPTH`], accepts any value. + fn resolve(&self, schema: &'a Value) -> &'a Value { + let mut schema = schema; + for _ in 0..=MAX_SCHEMA_DEPTH { + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { + match self.resolve_ref(reference) { + Some(target) => schema = target, + None => return &ANY_SCHEMA, + } + continue; + } + match schema.get("allOf").and_then(Value::as_array).map(Vec::as_slice) { + Some([inner]) => schema = inner, + _ => return schema, + } + } + &ANY_SCHEMA + } + + /// Every option of a value under `schema`. + /// + /// `const` and `enum` values become literal options, and `anyOf`, `oneOf`, + /// and type arrays one option per alternative, in schema order. An object + /// or array schema without `type` is typed by its `properties`, + /// `additionalProperties`, or `items`. `"nullable": true`, the OpenAPI + /// spelling of an optional value, adds a `null` option. Type names accept + /// the aliases that real tool schemas use, such as `int` or `dict`. + /// + /// Each reference target expands at most once. A union gains no values + /// from a repeated alternative, and a definition shared by several + /// alternatives would otherwise expand once per path to it, exponentially + /// in the nesting. A schema that constrains nothing, or whose references + /// only lead back into it, yields [`Self::any`]. + pub(crate) fn options(&self, schema: &'a Value) -> Vec> { + let mut expansion = Expansion::default(); + let mut options = self.options_at(schema, 0, &mut expansion); + if options.is_empty() && expansion.skipped { + options = Self::any(); + } + let nullable = schema.get("nullable").and_then(Value::as_bool) == Some(true); + if nullable && !options.iter().any(|option| option.ty == JsonType::Null) { + options.push(SchemaOption { + ty: JsonType::Null, + source: OptionSource::Any, + }); + } + options + } + + /// One option per JSON value type, each accepting any value of the type. + pub(crate) fn any() -> Vec> { + JsonType::VALUE_TYPES + .into_iter() + .map(|ty| SchemaOption { + ty, + source: OptionSource::Any, + }) + .collect() + } + + fn resolve_ref(&self, reference: &str) -> Option<&'a Value> { + self.schema.pointer(reference.strip_prefix('#')?) + } + + fn options_at( + &self, + schema: &'a Value, + depth: usize, + expansion: &mut Expansion, + ) -> Vec> { + let schema = match schema { + _ if depth > MAX_SCHEMA_DEPTH => return Self::any(), + Value::Bool(false) => return vec![], + Value::Object(schema) => schema, + _ => return Self::any(), + }; + if let Some(reference) = schema.get("$ref").and_then(Value::as_str) { + let Some(target) = self.resolve_ref(reference) else { + return Self::any(); + }; + if !expansion.expanded.insert(std::ptr::from_ref(target)) { + expansion.skipped = true; + return vec![]; + } + return self.options_at(target, depth + 1, expansion); + } + if let Some(value) = schema.get("const") { + return vec![SchemaOption::literal(value)]; + } + if let Some(values) = schema.get("enum").and_then(Value::as_array) { + return values.iter().map(SchemaOption::literal).collect(); + } + if let Some(options) = + schema.get("anyOf").or_else(|| schema.get("oneOf")).and_then(Value::as_array) + { + return options + .iter() + .flat_map(|option| self.options_at(option, depth + 1, expansion)) + .collect(); + } + if let Some(schemas) = schema.get("allOf").and_then(Value::as_array) { + return match schemas.as_slice() { + [schema] => self.options_at(schema, depth + 1, expansion), + _ => Self::any(), + }; + } + match schema.get("type") { + Some(Value::String(name)) => JsonType::parse(name) + .map(|ty| SchemaOption::typed(schema, ty)) + .into_iter() + .collect(), + Some(Value::Array(names)) => names + .iter() + .filter_map(Value::as_str) + .filter_map(JsonType::parse) + .map(|ty| SchemaOption::typed(schema, ty)) + .collect(), + Some(_) => Self::any(), + None if schema.contains_key("properties") + || schema.contains_key("additionalProperties") => + { + vec![SchemaOption::typed(schema, JsonType::Object)] + } + None if schema.contains_key("items") => { + vec![SchemaOption::typed(schema, JsonType::Array)] + } + None => Self::any(), + } + } +} + +impl<'a> SchemaOption<'a> { + /// The schema of a [`OptionSource::Schema`] option. + pub(crate) fn schema(&self) -> Option<&'a Map> { + match self.source { + OptionSource::Schema(schema) => Some(schema), + OptionSource::Literal(_) | OptionSource::Any => None, + } + } + + /// The option of `schema` that takes its values of type `ty`. + fn typed(schema: &'a Map, ty: JsonType) -> Self { + Self { + ty, + source: OptionSource::Schema(schema), + } + } + + fn literal(value: &'a Value) -> Self { + Self { + ty: JsonType::of_value(value), + source: OptionSource::Literal(value), + } + } +} + +/// `schema` with its `type` set to the canonical name of `ty`, copied only +/// when that changes it. +pub(crate) fn narrow(schema: &Map, ty: JsonType) -> Cow<'_, Map> { + if schema.get("type").and_then(Value::as_str) == Some(ty.name()) { + Cow::Borrowed(schema) + } else { + let mut schema = schema.clone(); + schema.insert("type".to_string(), Value::String(ty.name().to_string())); + Cow::Owned(schema) + } +} + +impl JsonType { + /// The types of JSON values, which do not distinguish integers. + const VALUE_TYPES: [Self; 6] = [ + Self::String, + Self::Number, + Self::Boolean, + Self::Null, + Self::Object, + Self::Array, + ]; + + /// The JSON Schema type name. + pub fn name(self) -> &'static str { + match self { + Self::String => "string", + Self::Integer => "integer", + Self::Number => "number", + Self::Boolean => "boolean", + Self::Null => "null", + Self::Object => "object", + Self::Array => "array", + } + } + + /// Schema accepting every JSON value of this type. + pub(crate) fn any_schema(self) -> Value { + match self { + Self::Object => json!({ "type": "object", "additionalProperties": true }), + Self::Array => json!({ "type": "array", "items": true }), + ty => json!({ "type": ty.name() }), + } + } + + /// Parse a schema type name, including aliases used by real tool schemas. + fn parse(name: &str) -> Option { + let name = name.trim().to_ascii_lowercase(); + Some(match name.as_str() { + "string" | "str" | "text" | "varchar" | "char" | "enum" => Self::String, + "integer" | "int" => Self::Integer, + "number" | "float" | "double" => Self::Number, + "boolean" | "bool" | "binary" => Self::Boolean, + "object" | "dict" | "map" => Self::Object, + "array" | "arr" | "list" | "sequence" => Self::Array, + "null" => Self::Null, + _ if name.starts_with("int") + || name.starts_with("uint") + || name.starts_with("long") + || name.starts_with("short") + || name.starts_with("unsigned") => + { + Self::Integer + } + _ if name.starts_with("num") || name.starts_with("float") => Self::Number, + _ if name.starts_with("dict") => Self::Object, + _ if name.starts_with("list") => Self::Array, + _ => return None, + }) + } + + fn of_value(value: &Value) -> Self { + match value { + Value::String(_) => Self::String, + Value::Number(number) if number.is_f64() => Self::Number, + Value::Number(_) => Self::Integer, + Value::Bool(_) => Self::Boolean, + Value::Null => Self::Null, + Value::Object(_) => Self::Object, + Value::Array(_) => Self::Array, + } + } +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + + use super::*; + + /// A Pydantic v2 `model_json_schema()` for a nested model, an enum, and an + /// optional nested model. + fn pydantic_schema() -> Value { + json!({ + "$defs": { + "Place": { + "properties": { "city": { "title": "City", "type": "string" } }, + "required": ["city"], + "title": "Place", + "type": "object" + }, + "Unit": { "enum": ["celsius", "fahrenheit"], "title": "Unit", "type": "string" } + }, + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "unit": { "$ref": "#/$defs/Unit", "default": "celsius" }, + "backup": { "anyOf": [{ "$ref": "#/$defs/Place" }, { "type": "null" }], "default": null } + }, + "required": ["place"], + "title": "Args", + "type": "object" + }) + } + + fn check_options(root: &Value, schema: &Value, expected: Expect) { + let options = SchemaRoot::new(root) + .options(schema) + .into_iter() + .map(|option| match option.source { + OptionSource::Schema(schema) => format!( + "{} {}", + option.ty.name(), + Value::Object(narrow(schema, option.ty).into_owned()) + ), + OptionSource::Literal(value) => format!("{} = {value}", option.ty.name()), + OptionSource::Any => format!("{} any", option.ty.name()), + }) + .collect::>(); + expected.assert_debug_eq(&options); + } + + #[test] + fn references_resolve_to_their_definitions() { + let root = pydantic_schema(); + let properties = &root["properties"]; + check_options( + &root, + &properties["place"], + expect![[r#" + [ + "object {\"properties\":{\"city\":{\"title\":\"City\",\"type\":\"string\"}},\"required\":[\"city\"],\"title\":\"Place\",\"type\":\"object\"}", + ] + "#]], + ); + check_options( + &root, + &properties["unit"], + expect![[r#" + [ + "string = \"celsius\"", + "string = \"fahrenheit\"", + ] + "#]], + ); + check_options( + &root, + &properties["backup"], + expect![[r#" + [ + "object {\"properties\":{\"city\":{\"title\":\"City\",\"type\":\"string\"}},\"required\":[\"city\"],\"title\":\"Place\",\"type\":\"object\"}", + "null {\"type\":\"null\"}", + ] + "#]], + ); + } + + #[test] + fn unresolvable_and_cyclic_references_accept_any_value() { + let root = json!({ + "$defs": { "a": { "$ref": "#/$defs/b" }, "b": { "$ref": "#/$defs/a" } }, + "properties": { + "missing": { "$ref": "#/$defs/missing" }, + "cycle": { "$ref": "#/$defs/a" }, + "remote": { "$ref": "https://example.com/schema.json" } + } + }); + let schema_root = SchemaRoot::new(&root); + for key in ["missing", "cycle", "remote"] { + assert_eq!( + schema_root.resolve(&root["properties"][key]), + &ANY_SCHEMA, + "{key}" + ); + assert!( + schema_root + .options(&root["properties"][key]) + .iter() + .all(|option| matches!(option.source, OptionSource::Any)), + "{key}" + ); + } + } + + #[test] + fn options_normalize_common_schema_spellings() { + let root = json!({}); + check_options( + &root, + &json!({ "type": ["int", "null"], "minimum": 0 }), + expect![[r#" + [ + "integer {\"type\":\"integer\",\"minimum\":0}", + "null {\"type\":\"null\",\"minimum\":0}", + ] + "#]], + ); + check_options( + &root, + &json!({ "type": "string", "nullable": true }), + expect![[r#" + [ + "string {\"type\":\"string\",\"nullable\":true}", + "null any", + ] + "#]], + ); + check_options( + &root, + &json!({ "enum": ["a", 1, 1.5, null] }), + expect![[r#" + [ + "string = \"a\"", + "integer = 1", + "number = 1.5", + "null = null", + ] + "#]], + ); + check_options( + &root, + &json!({ "allOf": [{ "type": "boolean" }] }), + expect![[r#" + [ + "boolean {\"type\":\"boolean\"}", + ] + "#]], + ); + check_options( + &root, + &json!({ "items": { "type": "string" } }), + expect![[r#" + [ + "array {\"items\":{\"type\":\"string\"},\"type\":\"array\"}", + ] + "#]], + ); + } +} diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index ef59ce08759b..391ca90f105f 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -5,6 +5,7 @@ use std::collections::BTreeMap; use serde_json::{Map, Number, Value}; +use crate::schema::{JsonType, OptionSource, SchemaRoot}; use crate::tool::Tool; /// Normalized parameter schemas for all tools in one request. @@ -21,7 +22,9 @@ pub(crate) struct ToolSchemas { /// downstream tool call execution. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub(super) struct ToolSchema { - params: BTreeMap, + /// The tool's parameters schema, read one level at a time as values are + /// converted, so recursive schemas cost only as deep as the value. + parameters: Value, } /// Parameter input for schema-aware conversion. @@ -47,22 +50,38 @@ pub(crate) struct ParamElement { pub value: ParamInput, } -/// Normalized JSON parameter type used for raw string coercion. +/// One level of a parameter's normalized type, used for raw string coercion. +/// +/// Object and array types normalize their children only when a structured +/// value reaches them. #[derive(Debug, Clone, PartialEq, Eq)] -pub(super) enum JsonParamType { +enum JsonParamType<'a> { String, Integer, Number, Boolean, - Object { - properties: BTreeMap, - additional_properties: Option>, - }, - Array { - items: Option>, - }, + Object(ObjectType<'a>), + Array(ArrayType<'a>), Null, - OneOf(Vec), + /// A non-string `const` or `enum` value: text spelling it as JSON decodes + /// to it. + Literal(&'a Value), + OneOf(Vec>), +} + +/// An object type whose field types are normalized on access. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ObjectType<'a> { + root: SchemaRoot<'a>, + properties: Option<&'a Map>, + additional_properties: Option<&'a Value>, +} + +/// An array type whose item type is normalized on access. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ArrayType<'a> { + root: SchemaRoot<'a>, + items: Option<&'a Value>, } impl ToolSchemas { @@ -70,7 +89,12 @@ impl ToolSchemas { pub(crate) fn from_tools(tools: &[Tool]) -> Self { let tools = tools .iter() - .map(|tool| (tool.name.clone(), ToolSchema::from_schema(&tool.parameters))) + .map(|tool| { + ( + tool.name.clone(), + ToolSchema::from_schema(tool.parameters.clone()), + ) + }) .collect(); Self { tools } @@ -117,25 +141,14 @@ impl ToolSchema { /// parameters to be treated as strings. const fn empty() -> &'static Self { static EMPTY: ToolSchema = ToolSchema { - params: BTreeMap::new(), + parameters: Value::Null, }; &EMPTY } - /// Normalize an OpenAI-style tool parameters JSON schema. - fn from_schema(parameters: &Value) -> Self { - let Some(properties) = parameters.get("properties").and_then(Value::as_object) else { - return Self::default(); - }; - - let params = properties - .iter() - .filter_map(|(name, schema)| { - JsonParamType::from_schema(schema).map(|param_type| (name.clone(), param_type)) - }) - .collect(); - - Self { params } + /// Keep an OpenAI-style tool parameters JSON schema. + fn from_schema(parameters: Value) -> Self { + Self { parameters } } /// Convert one parameter value using its normalized schema type. @@ -144,169 +157,86 @@ impl ToolSchema { /// the value fails to convert, this falls back to returning the raw /// string as a JSON string value, or object-like JSON for structured input. fn convert(&self, name: &str, input: ParamInput) -> Value { - convert_with_optional_schema(self.params.get(name), &input) + let root = SchemaRoot::new(&self.parameters); + let param_type = root + .resolved() + .get("properties") + .and_then(|properties| properties.get(name)) + .and_then(|schema| JsonParamType::from_schema(&root, schema)); + convert_with_optional_schema(param_type.as_ref(), &input) } } -impl JsonParamType { - /// Normalize one parameter property schema. - /// - /// `"nullable": true` (the OpenAPI spelling of an optional parameter) adds - /// `null` to the accepted types. - fn from_schema(schema: &Value) -> Option { - let param_type = Self::from_schema_ignoring_nullable(schema)?; - let nullable = schema.get("nullable").and_then(Value::as_bool) == Some(true); - Some(match param_type { - _ if !nullable || param_type.accepts_null() => param_type, - Self::OneOf(mut types) => { - types.push(Self::Null); - Self::OneOf(types) - } - param_type => Self::OneOf(vec![param_type, Self::Null]), - }) - } - - /// Return whether `null` is one of the accepted types. - fn accepts_null(&self) -> bool { - match self { - Self::Null => true, - Self::OneOf(types) => types.iter().any(Self::accepts_null), - _ => false, +impl<'a> JsonParamType<'a> { + /// Normalize one parameter or nested value schema, or `None` when it does + /// not constrain the value's type. + fn from_schema(root: &SchemaRoot<'a>, schema: &'a Value) -> Option { + let options = root.options(schema); + if options.iter().all(|option| matches!(option.source, OptionSource::Any)) { + return None; } - } - - /// Normalize one parameter property schema, without `nullable`. - fn from_schema_ignoring_nullable(schema: &Value) -> Option { - let schema = schema.as_object()?; - - if let Some(type_value) = schema.get("type") { - return Self::from_type_value(type_value, schema); - } - - if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) { - let param_type = composite - .as_array() - .map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::>()) - .filter(|types| !types.is_empty()) - .map(Self::one_of) - .unwrap_or_else(|| Self::object_from_schema(Some(schema))); - return Some(param_type); - } - - // Typically, these types are already handled by checking the "type" field, but - // we can also infer them from their characteristic fields if "type" is missing. - if let Some(values) = schema.get("enum").and_then(Value::as_array) { - // Enum values are treated as strings, except that a `null` member - // makes the parameter nullable (mirrors Python's enum type - // inference), so a literal "null" coerces to JSON null. - if values.iter().any(Value::is_null) { - return Some(Self::one_of(vec![Self::String, Self::Null])); - } - return Some(Self::String); - } - if schema.contains_key("items") { - return Some(Self::array_from_schema(Some(schema))); - } - if schema.contains_key("properties") || schema.contains_key("additionalProperties") { - return Some(Self::object_from_schema(Some(schema))); - } - - None - } - - /// Normalize a JSON schema `type` value. - fn from_type_value(type_value: &Value, schema: &Map) -> Option { - match type_value { - Value::String(kind) => Self::from_type_name(kind, Some(schema)), - Value::Array(kinds) => { - let types = kinds - .iter() - .filter_map(Value::as_str) - .filter_map(|kind| Self::from_type_name(kind, Some(schema))) - .collect::>(); - if types.is_empty() { - None - } else { - Some(Self::one_of(types)) - } - } - _ => None, - } - } - - /// Normalize one JSON schema type name. - fn from_type_name(kind: &str, schema: Option<&Map>) -> Option { - let kind = kind.trim().to_ascii_lowercase(); - match kind.as_str() { - "string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String), - "integer" | "int" => Some(Self::Integer), - "number" | "float" | "double" => Some(Self::Number), - "boolean" | "bool" | "binary" => Some(Self::Boolean), - "object" | "dict" | "map" => Some(Self::object_from_schema(schema)), - "array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)), - "null" => Some(Self::Null), - _ if kind.starts_with("int") - || kind.starts_with("uint") - || kind.starts_with("long") - || kind.starts_with("short") - || kind.starts_with("unsigned") => + let mut types = Vec::new(); + for option in &options { + // String literals stay with the string type, which keeps text + // verbatim instead of unquoting text that spells one as JSON. + if let OptionSource::Literal(value) = option.source + && !value.is_string() + && !types.contains(&Self::Literal(value)) { - Some(Self::Integer) + types.push(Self::Literal(value)); + } + let schema = option.schema(); + let param_type = match option.ty { + JsonType::String => Self::String, + JsonType::Integer => Self::Integer, + JsonType::Number => Self::Number, + JsonType::Boolean => Self::Boolean, + JsonType::Null => Self::Null, + JsonType::Object => Self::Object(ObjectType { + root: *root, + properties: schema + .and_then(|schema| schema.get("properties")) + .and_then(Value::as_object), + additional_properties: schema + .and_then(|schema| schema.get("additionalProperties")) + .filter(|schema| schema.is_object()), + }), + JsonType::Array => Self::Array(ArrayType { + root: *root, + items: schema.and_then(|schema| schema.get("items")), + }), + }; + if !types.contains(¶m_type) { + types.push(param_type); } - _ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number), - _ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)), - _ if kind.starts_with("list") => Some(Self::array_from_schema(schema)), - _ => None, } - } - - /// Normalize object schema fields. - fn object_from_schema(schema: Option<&Map>) -> Self { - let properties = schema - .and_then(|schema| schema.get("properties")) - .and_then(Value::as_object) - .map(|properties| { - properties - .iter() - .filter_map(|(name, schema)| { - Self::from_schema(schema).map(|param_type| (name.clone(), param_type)) - }) - .collect() - }) - .unwrap_or_default(); - - let additional_properties = - schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| { - if schema.is_object() { - Self::from_schema(schema).map(Box::new) - } else { - None - } - }); - - Self::Object { - properties, - additional_properties, + // Text spelling a literal decodes to it before any type, a string in + // particular, accepts the same text. + types.sort_by_key(|param_type| !matches!(param_type, Self::Literal(_))); + match types.len() { + 0 => None, + 1 => types.pop(), + _ => Some(Self::OneOf(types)), } } +} - /// Normalize array schema fields. - fn array_from_schema(schema: Option<&Map>) -> Self { - let items = schema - .and_then(|schema| schema.get("items")) - .and_then(Self::from_schema) - .map(Box::new); - - Self::Array { items } +impl<'a> ObjectType<'a> { + /// The type of the field named `name`, from its declared schema or the + /// additional properties' schema. + fn field(&self, name: &str) -> Option> { + let schema = self + .properties + .and_then(|properties| properties.get(name)) + .or(self.additional_properties)?; + JsonParamType::from_schema(&self.root, schema) } +} - /// Collapse a candidate type list into one normalized type. - fn one_of(mut types: Vec) -> Self { - if types.len() == 1 { - types.remove(0) - } else { - Self::OneOf(types) - } +impl<'a> ArrayType<'a> { + /// The type of the array's items. + fn items(&self) -> Option> { + JsonParamType::from_schema(&self.root, self.items?) } } @@ -318,7 +248,10 @@ fn is_null_literal(value: &str) -> bool { } /// Convert one parameter input to a normalized JSON value. -fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value { +fn convert_with_optional_schema( + param_type: Option<&JsonParamType<'_>>, + input: &ParamInput, +) -> Value { // Coerce the literal text `null` to JSON null, except for `string`-typed // params, where it must stay the string "null": a model emitting the literal // text "null" for a string field means the string, not a missing value. @@ -341,13 +274,13 @@ fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &Para ParamInput::Text(value) => Value::String(value.clone()), ParamInput::Elements(elements) => { // Convert structured input to object without a schema. - Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None)) + Value::Object(convert_elements_to_object(elements, None)) } } } /// Convert one parameter input to a normalized JSON type. -fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option { +fn try_convert_value(param_type: &JsonParamType<'_>, input: &ParamInput) -> Option { match input { ParamInput::Text(value) => try_convert_text_value(param_type, value), ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements), @@ -358,18 +291,21 @@ fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option Option { +fn try_convert_text_value(param_type: &JsonParamType<'_>, value: &str) -> Option { match param_type { JsonParamType::String => Some(Value::String(value.to_string())), JsonParamType::Integer => convert_integer_text(value), JsonParamType::Number => convert_number_text(value), JsonParamType::Boolean => try_convert_boolean(value), - JsonParamType::Object { .. } if value.trim().is_empty() => Some(Value::Object(Map::new())), - JsonParamType::Array { .. } if value.trim().is_empty() => Some(Value::Array(Vec::new())), + JsonParamType::Object(_) if value.trim().is_empty() => Some(Value::Object(Map::new())), + JsonParamType::Array(_) if value.trim().is_empty() => Some(Value::Array(Vec::new())), // For composite types with string input, interpret the string as JSON of that type. - JsonParamType::Object { .. } => serde_json::from_str(value).ok().filter(Value::is_object), - JsonParamType::Array { .. } => serde_json::from_str(value).ok().filter(Value::is_array), + JsonParamType::Object(_) => serde_json::from_str(value).ok().filter(Value::is_object), + JsonParamType::Array(_) => serde_json::from_str(value).ok().filter(Value::is_array), JsonParamType::Null => is_null_literal(value).then_some(Value::Null), + JsonParamType::Literal(literal) => serde_json::from_str::(value.trim()) + .ok() + .filter(|value| value == *literal), JsonParamType::OneOf(types) => { types.iter().find_map(|param_type| try_convert_text_value(param_type, value)) } @@ -378,25 +314,24 @@ fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option, elements: &[ParamElement], ) -> Option { match param_type { - JsonParamType::Object { - properties, - additional_properties, - } => Some(Value::Object(convert_elements_to_object( + JsonParamType::Object(object) => Some(Value::Object(convert_elements_to_object( elements, - properties, - additional_properties.as_deref(), + Some(object), ))), - JsonParamType::Array { items } => Some(Value::Array( - // Collect all child elements into an array, regardless of their names. - elements - .iter() - .map(|element| convert_with_optional_schema(items.as_deref(), &element.value)) - .collect(), - )), + JsonParamType::Array(array) => { + let items = array.items(); + Some(Value::Array( + // Collect all child elements into an array, regardless of their names. + elements + .iter() + .map(|element| convert_with_optional_schema(items.as_ref(), &element.value)) + .collect(), + )) + } JsonParamType::OneOf(types) => types .iter() .find_map(|param_type| try_convert_elements_value(param_type, elements)), @@ -406,20 +341,20 @@ fn try_convert_elements_value( | JsonParamType::Integer | JsonParamType::Number | JsonParamType::Boolean - | JsonParamType::Null => None, + | JsonParamType::Null + | JsonParamType::Literal(_) => None, } } /// Convert structured elements to an object, using field schemas when present. fn convert_elements_to_object( elements: &[ParamElement], - properties: &BTreeMap, - additional_properties: Option<&JsonParamType>, + object_type: Option<&ObjectType<'_>>, ) -> Map { let mut object = Map::with_capacity(elements.len()); for element in elements { - let param_type = properties.get(&element.name).or(additional_properties); - let value = convert_with_optional_schema(param_type, &element.value); + let param_type = object_type.and_then(|object_type| object_type.field(&element.name)); + let value = convert_with_optional_schema(param_type.as_ref(), &element.value); insert_object_value(&mut object, element.name.clone(), value); } object @@ -508,7 +443,7 @@ mod tests { #[test] fn invalid_schema_converts_everything_as_string() { - let params = ToolSchema::from_schema(&json!({ "type": "object" })); + let params = ToolSchema::from_schema(json!({ "type": "object" })); assert_eq!(params.convert("count", text("42")), json!("42")); assert_eq!(params.convert("count", text("null")), json!(null)); @@ -516,7 +451,7 @@ mod tests { #[test] fn skips_unknown_property_schema_and_unknown_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "unknown_schema": true, @@ -532,7 +467,7 @@ mod tests { #[test] fn converts_supported_types() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string" }, @@ -571,7 +506,7 @@ mod tests { #[test] fn number_conversion_preserves_json_number_spelling_with_legacy_fallback() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": "number" } @@ -618,7 +553,7 @@ mod tests { #[test] fn non_string_values_ignore_surrounding_whitespace() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string" }, @@ -641,7 +576,7 @@ mod tests { #[test] fn numbers_accept_digit_separators() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "count": { "type": "integer" }, @@ -658,7 +593,7 @@ mod tests { #[test] fn nullable_schemas_accept_null() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "text": { "type": "string", "nullable": true }, @@ -675,7 +610,7 @@ mod tests { #[test] fn composite_values_must_decode_to_their_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "payload": { "type": "object" }, @@ -695,7 +630,7 @@ mod tests { #[test] fn converts_upstream_aliases() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "s": { "type": "varchar" }, @@ -717,13 +652,13 @@ mod tests { #[test] fn preserves_union_type_order() { - let integer_first = ToolSchema::from_schema(&json!({ + let integer_first = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": ["integer", "string"] } } })); - let string_first = ToolSchema::from_schema(&json!({ + let string_first = ToolSchema::from_schema(json!({ "type": "object", "properties": { "value": { "type": ["string", "integer"] } @@ -736,7 +671,7 @@ mod tests { #[test] fn converts_composite_schemas() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "choice": { @@ -745,7 +680,7 @@ mod tests { { "type": "string" } ] }, - "fallback_object": { + "unknown_alternatives": { "oneOf": [ { "type": "mystery" } ] @@ -754,15 +689,139 @@ mod tests { })); assert_eq!(params.convert("choice", text("42")), json!(42)); + // Alternatives of unknown types leave the value unconstrained, like an + // unknown `type`. assert_eq!( - params.convert("fallback_object", text(r#"{"x":1}"#)), - json!({ "x": 1 }) + params.convert("unknown_alternatives", text(r#"{"x":1}"#)), + json!(r#"{"x":1}"#) + ); + } + + #[test] + fn resolves_references_in_pydantic_schemas() { + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Place": { + "properties": { "city": { "type": "string" }, "days": { "type": "integer" } }, + "type": "object" + }, + "Count": { "type": "integer" } + }, + "type": "object", + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "backup": { "anyOf": [{ "$ref": "#/$defs/Place" }, { "type": "null" }] }, + "count": { "$ref": "#/$defs/Count" } + } + })); + + assert_eq!( + params.convert("place", text(r#"{"city": "Paris", "days": 3}"#)), + json!({ "city": "Paris", "days": 3 }) ); + assert_eq!( + params.convert("backup", text(r#"{"city": "Paris"}"#)), + json!({ "city": "Paris" }) + ); + assert_eq!(params.convert("backup", text("null")), Value::Null); + assert_eq!(params.convert("count", text("3")), json!(3)); + } + + #[test] + fn recursive_references_convert_without_unbounded_expansion() { + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Node": { + "type": "object", + "properties": { + "id": { "type": "integer" }, + "children": { "type": "array", "items": { "$ref": "#/$defs/Node" } } + } + } + }, + "type": "object", + "properties": { "tree": { "$ref": "#/$defs/Node" } } + })); + + assert_eq!( + params.convert( + "tree", + text(r#"{"id": 1, "children": [{"id": 2, "children": []}]}"#) + ), + json!({ "id": 1, "children": [{ "id": 2, "children": [] }] }) + ); + } + + #[test] + fn branching_recursive_references_convert_structured_values_by_level() { + // Two recursive children per node expand to 2^depth types if the schema + // is walked ahead of the value. + let params = ToolSchema::from_schema(json!({ + "$defs": { + "Tree": { + "type": "object", + "properties": { + "value": { "type": "integer" }, + "left": { "$ref": "#/$defs/Tree" }, + "right": { "$ref": "#/$defs/Tree" } + } + } + }, + "type": "object", + "properties": { "tree": { "$ref": "#/$defs/Tree" } } + })); + + let leaf = |value: &str| ParamInput::Elements(vec![elem("value", text(value))]); + let tree = ParamInput::Elements(vec![ + elem("value", text("1")), + elem( + "left", + ParamInput::Elements(vec![elem("value", text("2")), elem("right", leaf("3"))]), + ), + elem("right", leaf("4")), + ]); + assert_eq!( + params.convert("tree", tree), + json!({ + "value": 1, + "left": { "value": 2, "right": { "value": 3 } }, + "right": { "value": 4 } + }) + ); + } + + #[test] + fn mixed_enums_decode_text_spelling_a_non_string_literal_to_it() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { + "mode": { "enum": ["fast", 1] }, + "flag": { "enum": ["auto", true] } + } + })); + + assert_eq!(params.convert("mode", text("1")), json!(1)); + assert_eq!(params.convert("mode", text("fast")), json!("fast")); + // Text matching no literal keeps the schema order, which tries the + // string first. + assert_eq!(params.convert("mode", text("2")), json!("2")); + assert_eq!(params.convert("flag", text("true")), json!(true)); + assert_eq!(params.convert("flag", text("auto")), json!("auto")); + } + + #[test] + fn integer_enums_convert_to_numbers() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { "level": { "enum": [1, 2, 3] } } + })); + + assert_eq!(params.convert("level", text("2")), json!(2)); } #[test] fn infers_type_from_schema_shape_without_type() { - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "choice": { "enum": ["a", "b"] }, @@ -845,7 +904,7 @@ mod tests { // A `string`-typed param whose value is the literal text "null"/"NULL" // must stay a string (the original case is preserved), rather than being // coerced to JSON null. Non-string types keep coercing "null" to null. - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "name": { "type": "string" }, @@ -868,7 +927,7 @@ mod tests { // must coerce to JSON null (matching Python's `extract_types_from_schema`, // which infers `null` from the enum values), while a non-null enum keeps // "null" as a string. - let params = ToolSchema::from_schema(&json!({ + let params = ToolSchema::from_schema(json!({ "type": "object", "properties": { "mode": { "enum": [null, "auto"] }, diff --git a/rust/src/parser/src/tool/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs index e233f4720c05..e804ce0bb255 100644 --- a/rust/src/parser/src/tool/qwen_coder.rs +++ b/rust/src/parser/src/tool/qwen_coder.rs @@ -306,7 +306,7 @@ mod tests { use serde_json::{Value, json}; use thiserror_ext::AsReport; - use super::{Qwen3CoderToolParser, ToolParser}; + use super::{Qwen3CoderToolParser, Tool, ToolParser}; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::tests::assert_tool_framing_preserves_body_whitespace; use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; @@ -545,6 +545,39 @@ mod tests { ); } + #[test] + fn qwen_coder_parse_complete_resolves_schema_references() { + let tools = [Tool { + name: "plan_trip".to_string(), + description: None, + parameters: json!({ + "$defs": { + "Place": { "type": "object", "properties": { "city": { "type": "string" } } }, + "Days": { "type": "integer" } + }, + "type": "object", + "properties": { + "place": { "$ref": "#/$defs/Place" }, + "days": { "$ref": "#/$defs/Days" } + } + }), + strict: None, + defer_loading: None, + }]; + let mut parser = Qwen3CoderToolParser::new(&tools); + let output = parser + .parse_complete(&build_tool_call( + "plan_trip", + &[("place", "\n{\"city\": \"Paris\"}\n"), ("days", "\n3\n")], + )) + .unwrap(); + + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "place": { "city": "Paris" }, "days": 3 }) + ); + } + #[test] fn qwen_coder_streaming_extracts_single_tool_call() { let mut parser = Qwen3CoderToolParser::new(&test_tools()); diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs index 8f0f098f06d1..4754726c1b5d 100644 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -46,12 +46,12 @@ impl ArgumentSyntax for XtmlArguments { } ParameterKey::Free => ParameterKey::Free, }; - let arguments = group_by(options, |option| option.ty) + let arguments = group_by(options, |option| xtml_type(option.ty)) .into_iter() - .filter_map(|(ty, options)| { + .filter_map(|(type_name, options)| { let values = options .into_iter() - .filter_map(|option| match ty { + .filter_map(|option| match option.ty { // The parser ends the call at its close marker, so a // string value may contain neither close marker. JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), @@ -59,7 +59,7 @@ impl ArgumentSyntax for XtmlArguments { }) .collect::>(); (!values.is_empty()).then(|| { - let suffix = format!("\" type=\"{}\"{SEP}", ty.name()); + let suffix = format!("\" type=\"{type_name}\"{SEP}"); key.tag( &format!("{ARG_OPEN} key=\""), &suffix, @@ -73,6 +73,20 @@ impl ArgumentSyntax for XtmlArguments { } } +/// The `type` attribute K3 renders for a value of type `ty`, as +/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are +/// both `number`. +fn xtml_type(ty: JsonType) -> &'static str { + match ty { + JsonType::String => "string", + JsonType::Integer | JsonType::Number => "number", + JsonType::Boolean => "boolean", + JsonType::Null => "null", + JsonType::Object => "object", + JsonType::Array => "array", + } +} + /// Kimi K3 XTML structural-tag builder. #[derive(Debug, Clone, Copy, Default)] pub struct KimiK3StructuralTagBuilder; From 42ecc8eebe701faf07e7d6f33dbf5a8218825b5b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 7 Oct 2026 16:29:09 +0800 Subject: [PATCH 117/952] [CI][RL] Consolidate RL entrypoint tests under tests/entrypoints/rl (#59948) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: wang.yuqi Co-authored-by: Claude Opus 5.5 --- .buildkite/intel_jobs/entrypoints_intel.yaml | 13 ++++++---- .buildkite/test-amd.yaml | 6 +++-- .buildkite/test_areas/entrypoints.yaml | 2 -- .buildkite/test_areas/rl.yaml | 25 +++++++++++++++++++ .buildkite/test_areas/rust_frontend.yaml | 2 -- .github/CODEOWNERS | 3 +-- .../{serve/dev/rlhf => rl}/__init__.py | 0 .../{serve/dev/rlhf => rl}/conftest.py | 0 .../return_routed_experts}/__init__.py | 0 .../return_routed_experts/test_openai_api.py} | 2 +- .../test_scale_out_api.py} | 0 .../rl/state_transitions/__init__.py | 0 .../state_transitions/test_pause_resume.py | 2 +- .../state_transitions}/test_sleep.py | 0 .../state_transitions/test_weight_checker.py | 2 +- .../test_weight_operation_metrics.py | 2 +- .../{ => rl}/weight_transfer/__init__.py | 0 .../weight_transfer/test_offline.py} | 3 +-- .../test_serving_multimodal_tokens.py | 4 +-- tools/pre_commit/mypy.py | 2 +- 20 files changed, 46 insertions(+), 22 deletions(-) create mode 100644 .buildkite/test_areas/rl.yaml rename tests/entrypoints/{serve/dev/rlhf => rl}/__init__.py (100%) rename tests/entrypoints/{serve/dev/rlhf => rl}/conftest.py (100%) rename tests/entrypoints/{serve/dev/rlhf/state_transitions => rl/return_routed_experts}/__init__.py (100%) rename tests/entrypoints/{openai/test_return_routed_experts.py => rl/return_routed_experts/test_openai_api.py} (98%) rename tests/entrypoints/{scale_out/token_in_token_out/test_return_routed_experts.py => rl/return_routed_experts/test_scale_out_api.py} (100%) create mode 100644 tests/entrypoints/rl/state_transitions/__init__.py rename tests/entrypoints/{serve/dev/rlhf => rl}/state_transitions/test_pause_resume.py (98%) rename tests/entrypoints/{serve/dev => rl/state_transitions}/test_sleep.py (100%) rename tests/entrypoints/{serve/dev/rlhf => rl}/state_transitions/test_weight_checker.py (98%) rename tests/entrypoints/{serve/dev/rlhf => rl}/test_weight_operation_metrics.py (97%) rename tests/entrypoints/{ => rl}/weight_transfer/__init__.py (100%) rename tests/entrypoints/{weight_transfer/test_weight_transfer_llm.py => rl/weight_transfer/test_offline.py} (99%) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 7e823e024563..3cc9c1422d9e 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -19,14 +19,14 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer + - tests/entrypoints/rl/weight_transfer - tests/entrypoints/launchers commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s entrypoints/unit_tests && - pytest -v -s entrypoints/weight_transfer' + pytest -v -s entrypoints/rl/weight_transfer' - label: ":intel: (BMG) Entrypoints Integration (LLM)" key: entrypoints-integration-llm-intel @@ -71,14 +71,17 @@ steps: - vllm/ - "!vllm/distributed/kv_transfer/" - tests/entrypoints/serve + - tests/entrypoints/rl commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc - --ignore=entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py - --deselect="tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && + pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc && + pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + --ignore=entrypoints/rl/return_routed_experts/test_openai_api.py + --ignore=entrypoints/rl/test_weight_operation_metrics.py + --deselect="tests/entrypoints/rl/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && PYTHONPATH=/workspace/vllm pytest -v -s entrypoints/serve/dev/rpc' - label: ":intel: (BMG) Entrypoints Integration (API Server Scale Out)" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 17fe9c0ae58e..8317919ce4b8 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1044,12 +1044,12 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer + - tests/entrypoints/rl/weight_transfer - tests/entrypoints/launchers - vllm/platforms/rocm.py commands: - pytest -v -s entrypoints/unit_tests - - pytest -v -s entrypoints/weight_transfer + - pytest -v -s entrypoints/rl/weight_transfer - pytest -v -s entrypoints/launchers - label: ":amd: (MI300) Entrypoints Integration (Responses API)" # TBD @@ -3746,11 +3746,13 @@ steps: - vllm/ - tests/entrypoints/serve - tests/entrypoints/scale_out + - tests/entrypoints/rl commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 65 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 5f8e1d1c352c..68c6243d479b 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -31,11 +31,9 @@ steps: source_file_dependencies: - vllm/entrypoints - tests/entrypoints/unit_tests - - tests/entrypoints/weight_transfer - tests/entrypoints/launchers commands: - pytest -v -s entrypoints/unit_tests - - pytest -v -s entrypoints/weight_transfer - pytest -v -s entrypoints/launchers mirror: amd: diff --git a/.buildkite/test_areas/rl.yaml b/.buildkite/test_areas/rl.yaml new file mode 100644 index 000000000000..7141a81bde1a --- /dev/null +++ b/.buildkite/test_areas/rl.yaml @@ -0,0 +1,25 @@ +group: RL +depends_on: + - image-build +steps: +- label: ":nvidia: (H200 MIG 35GB) RL Entrypoints" + device: h200_35gb + key: rl-entrypoints-tests + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - "!vllm/distributed/kv_transfer/" + - tests/entrypoints/rl + commands: + - pytest -v -s entrypoints/rl/weight_transfer + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + mirror: + amd: + label: ":amd: (MI355 DPX) RL Entrypoints" + dind: false + device: mi355_dpx + timeout_in_minutes: 50 + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index c5899a9c59bf..59c78525c5e7 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -82,7 +82,6 @@ steps: - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py - # - tests/entrypoints/serve/dev/test_sleep.py - tests/entrypoints/serve/tokenize/test_tokenization.py commands: - export VLLM_USE_RUST_FRONTEND=1 @@ -93,7 +92,6 @@ steps: # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - # - pytest -v -s entrypoints/serve/dev/test_sleep.py # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" mirror: diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 71d501eabe7c..083842c932ff 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -231,8 +231,7 @@ mkdocs.yaml @hmellor /vllm/distributed/weight_transfer @aoshen02 /vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 -/tests/entrypoints/weight_transfer @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 -/tests/entrypoints/serve/dev/rlhf @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas @aoshen02 +/tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 /docs/training/rlhf.md @aoshen02 diff --git a/tests/entrypoints/serve/dev/rlhf/__init__.py b/tests/entrypoints/rl/__init__.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/__init__.py rename to tests/entrypoints/rl/__init__.py diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/rl/conftest.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/conftest.py rename to tests/entrypoints/rl/conftest.py diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py b/tests/entrypoints/rl/return_routed_experts/__init__.py similarity index 100% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py rename to tests/entrypoints/rl/return_routed_experts/__init__.py diff --git a/tests/entrypoints/openai/test_return_routed_experts.py b/tests/entrypoints/rl/return_routed_experts/test_openai_api.py similarity index 98% rename from tests/entrypoints/openai/test_return_routed_experts.py rename to tests/entrypoints/rl/return_routed_experts/test_openai_api.py index 46eeed179531..f4fc4013997a 100644 --- a/tests/entrypoints/openai/test_return_routed_experts.py +++ b/tests/entrypoints/rl/return_routed_experts/test_openai_api.py @@ -7,7 +7,7 @@ import pybase64 as base64 import pytest -from ...utils import RemoteOpenAIServer +from tests.utils import RemoteOpenAIServer MODEL_NAME = "TitanML/tiny-mixtral" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py similarity index 100% rename from tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py rename to tests/entrypoints/rl/return_routed_experts/test_scale_out_api.py diff --git a/tests/entrypoints/rl/state_transitions/__init__.py b/tests/entrypoints/rl/state_transitions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py b/tests/entrypoints/rl/state_transitions/test_pause_resume.py similarity index 98% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py rename to tests/entrypoints/rl/state_transitions/test_pause_resume.py index 3de9e0054963..587b4b73194a 100644 --- a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py +++ b/tests/entrypoints/rl/state_transitions/test_pause_resume.py @@ -10,7 +10,7 @@ import pytest import requests -from tests.entrypoints.serve.dev.rlhf.conftest import ( +from tests.entrypoints.rl.conftest import ( cached_tokens, completion_with_cache_details, gen, diff --git a/tests/entrypoints/serve/dev/test_sleep.py b/tests/entrypoints/rl/state_transitions/test_sleep.py similarity index 100% rename from tests/entrypoints/serve/dev/test_sleep.py rename to tests/entrypoints/rl/state_transitions/test_sleep.py diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py b/tests/entrypoints/rl/state_transitions/test_weight_checker.py similarity index 98% rename from tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py rename to tests/entrypoints/rl/state_transitions/test_weight_checker.py index 7cb58df9d902..ec3c8016a870 100644 --- a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_weight_checker.py +++ b/tests/entrypoints/rl/state_transitions/test_weight_checker.py @@ -9,7 +9,7 @@ import torch from transformers import AutoModelForCausalLM -from tests.entrypoints.serve.dev.rlhf.conftest import ( +from tests.entrypoints.rl.conftest import ( MODEL_NAME, gen, health, diff --git a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py b/tests/entrypoints/rl/test_weight_operation_metrics.py similarity index 97% rename from tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py rename to tests/entrypoints/rl/test_weight_operation_metrics.py index d7758011f925..108dd77b2f1b 100644 --- a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py +++ b/tests/entrypoints/rl/test_weight_operation_metrics.py @@ -7,7 +7,7 @@ from prometheus_client.parser import text_string_to_metric_families from transformers import AutoModelForCausalLM -from tests.entrypoints.serve.dev.rlhf.conftest import MODEL_NAME +from tests.entrypoints.rl.conftest import MODEL_NAME from tests.utils import RemoteOpenAIServer from vllm.distributed.weight_transfer import ( HTTPVLLMWeightSyncClient, diff --git a/tests/entrypoints/weight_transfer/__init__.py b/tests/entrypoints/rl/weight_transfer/__init__.py similarity index 100% rename from tests/entrypoints/weight_transfer/__init__.py rename to tests/entrypoints/rl/weight_transfer/__init__.py diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/rl/weight_transfer/test_offline.py similarity index 99% rename from tests/entrypoints/weight_transfer/test_weight_transfer_llm.py rename to tests/entrypoints/rl/weight_transfer/test_offline.py index e5c3616d6e65..e4a30f17e61d 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/rl/weight_transfer/test_offline.py @@ -15,6 +15,7 @@ import pytest import torch +from tests.utils import create_new_process_for_each_test from vllm.config import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( WeightTransferEngine, @@ -24,8 +25,6 @@ WeightTransferUpdateRequest, ) -from ...utils import create_new_process_for_each_test - # Use a tiny model for fast testing MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 1d017bd3cf1b..2e58d70972f9 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -5,8 +5,8 @@ Mirrors test_serving_tokens.py but exercises the multimodal piping using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render -> /generate -> /detokenize path. Intentionally avoids running the HF -processor in the pytest parent process to keep os.fork() in sibling -tests (e.g. test_weight_transfer_llm.py) deadlock-free. +processor in the pytest parent process to keep os.fork() in other +tests (e.g. entrypoints/rl/weight_transfer/test_offline.py) deadlock-free. """ import httpx diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index a6c8e4641963..245d0daf5e18 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -38,10 +38,10 @@ "tests/entrypoints/generate", "tests/entrypoints/multimodal", "tests/entrypoints/pooling", + "tests/entrypoints/rl", "tests/entrypoints/serve", "tests/entrypoints/tool_parsers", "tests/entrypoints/unit_tests", - "tests/entrypoints/weight_transfer", "tests/kernels/attention", "tests/kernels/core", "tests/kernels/helion", From df417f780a19b68cebc4379942de6ea0e6f6aceb Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:35:56 +0800 Subject: [PATCH 118/952] [Bugfix][MoRIIO] Exclude synchronous READ destinations from KV zeroing (#59164) Signed-off-by: Yichao Zhu Signed-off-by: whx-sjtu Signed-off-by: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Co-authored-by: Yichao Zhu Co-authored-by: Codex Co-authored-by: Yifan Qiao --- .../unit/test_moriio_hma_scheduler.py | 70 ++++++++++++------- .../kv_transfer/kv_connector/v1/base.py | 6 +- .../v1/moriio/moriio_connector.py | 18 +++++ 3 files changed, 67 insertions(+), 27 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 28bc8bebe261..1de729416e90 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -7,32 +7,20 @@ GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE ``N-1`` token accounting, P-side prompt truncation, and offset-template cache wiring. - -Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is -running on ROCm with ``mori`` installed (importing the connector pulls in -``mori``). The authoritative run happens on the MIA recipe image. """ import importlib -import importlib.util +from dataclasses import asdict from types import SimpleNamespace import pytest import torch -from vllm.platforms import current_platform from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs -mori_available = importlib.util.find_spec("mori") is not None - -if not (current_platform.is_rocm() and mori_available): - pytest.skip( - "MoRIIOs are only available on ROCm with mori package installed", - allow_module_level=True, - ) - moriio_connector = importlib.import_module( "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector" ) @@ -525,6 +513,8 @@ def _make_read_scheduler(): request_id_to_transfer_id={}, transfer_id_to_request_id={}, _reqs_need_recv={}, + _reqs_need_save={}, + _reqs_need_send={}, _req_kv_params={}, _max_decode_tail_blocks=1, ) @@ -538,24 +528,52 @@ def _make_read_request(remote_block_ids): "transfer_id": "tx", "remote_engine_id": "prefill", "remote_block_ids": remote_block_ids, + "remote_host": "127.0.0.1", + "remote_handshake_port": 6000, + "remote_notify_port": 6001, }, ) -def test_update_state_drops_decode_recompute_tail_block(): +@pytest.mark.parametrize("external_tokens", [0, 256]) +@pytest.mark.parametrize("has_mamba", [False, True]) +@pytest.mark.parametrize("zero_ids", [None, [], [101], [101, 102, 202, 999]]) +def test_update_state_drops_decode_recompute_tail_block( + external_tokens, has_mamba, zero_ids +): + """Only full hybrid READ destinations skip zeroing; local tails still zero.""" sched = _make_read_scheduler() - request = _make_read_request([[10, 11], [91]]) + sched._has_mamba = has_mamba + remote_blocks = [[10, 11], [91]] if has_mamba else [[10, 11]] + request = _make_read_request(remote_blocks) blocks = _FakeBlocks( - all_groups=([100, 101, 102], [200, 201, 202]), - ) - - sched.update_state_after_alloc(request, blocks, num_external_tokens=256) - - assert sched._reqs_need_recv["req"][1] == [[100, 101], [202]] - assert sched._req_kv_params["req"]["remote_block_ids"] == [ - [10, 11], - [91], - ] + all_groups=([100, 101, 102], [200, 201, 202]) + if has_mamba + else ([100, 101, 102],), + ) + + sched.update_state_after_alloc(request, blocks, num_external_tokens=external_tokens) + + expected_blocks = [[100, 101] if external_tokens else []] + if has_mamba: + expected_blocks.append([202]) + assert sched._reqs_need_recv["req"][1] == expected_blocks + assert sched._req_kv_params["req"]["remote_block_ids"] == remote_blocks + + output = SchedulerOutput.make_empty() + output.new_block_ids_to_zero = None if zero_ids is None else list(zero_ids) + before = asdict(output) + meta = sched.build_connector_meta(output) + assert meta.reqs_to_recv["req"].local_block_ids == expected_blocks + if zero_ids and has_mamba and external_tokens: + before["new_block_ids_to_zero"] = zero_ids[1:] + assert asdict(output) == before + + # A later local allocation of the same pages must still be zeroed. + output.new_block_ids_to_zero = [100, 101, 102] + meta = sched.build_connector_meta(output) + assert not meta.reqs_to_recv + assert output.new_block_ids_to_zero == [100, 101, 102] def test_update_state_pairs_shorter_local_blocks_with_remote_suffix(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index bb9d4597ce3c..71593f59e55c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -583,7 +583,11 @@ def build_connector_meta( """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. - Also, calling this function will reset the state of the connector. + FIXME: one exception: + synchronous READ connectors may remove attention blocks they fully + overwrite this step from new_block_ids_to_zero. They must complete + those loads before the blocks are used and fail the step if a load fails. + Calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 167c7ef3c299..5a8a9176640e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1232,6 +1232,24 @@ def build_connector_meta( local_block_ids=block_ids, kv_transfer_params=kv_params, ) + if ( + self.mode == MoRIIOMode.READ + and self._has_mamba + and scheduler_output.new_block_ids_to_zero + ): + # Hybrid models zero recycled attention pages that held Mamba state. + # Host-submitted READs overwrite these pages and can race zeroing. + # Hybrid READ metadata puts the aligned attention pages first. + read_dst_block_ids = { + b + for _, block_ids in self._reqs_need_recv.values() + for b in block_ids[0] + } + scheduler_output.new_block_ids_to_zero = [ + b + for b in scheduler_output.new_block_ids_to_zero + if b not in read_dst_block_ids + ] for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) From 02d8a94e610f515449d3bc3f44c592e3022e19d8 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 7 Oct 2026 02:38:31 -0700 Subject: [PATCH 119/952] [CI][ROCm] Skip DSv4.1 decoder replay graph test on ROCm (#60378) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- tests/models/test_deepseek_v41_decoder_replay_layers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 9ba2df5c8f4c..1302e0acda30 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -21,6 +21,7 @@ from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( DecoderReplayCudaGraphManager, ) +from vllm.platforms import current_platform from vllm.v1.worker.gpu import cudagraph_utils pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -125,6 +126,9 @@ def _graph_context(num_tokens): ) +@pytest.mark.skipif( + current_platform.is_rocm(), reason="decoder replay CUDA graphs are NVIDIA-only" +) def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) From 6b75dfb8a101dcf23c9e2b9e627433c61030cbcf Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 18:53:17 +0900 Subject: [PATCH 120/952] [Misc] Report a missing DeepSelect extension only when it is requested (#60242) Signed-off-by: lesj0610 --- vllm/model_executor/layers/indexer_topk.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index f0da0312e127..02c4280ea0a1 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -29,12 +29,10 @@ try: import vllm._deepselect_C # noqa: F401 (registers torch.ops.deep_select) -except ImportError as e: - from vllm.logger import init_logger - init_logger(__name__).warning( - "Failed to import the DeepSelect extension (vllm._deepselect_C): %s", e - ) + _HAS_DEEP_SELECT_C = True +except ImportError: + _HAS_DEEP_SELECT_C = False @functools.lru_cache(maxsize=1) @@ -196,6 +194,8 @@ def resolve_backend( failures.append("requires a CUDA platform") elif not self._has_deep_select: failures.append("requires SM100a/SM103a (10.x device family)") + elif not _HAS_DEEP_SELECT_C: + failures.append("vllm._deepselect_C could not be imported") elif not is_deep_select_supported(logits, topk_tokens): failures.append( f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," From 3ca00a8261c7790ba7487e2001ab9004ad8065f2 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 12:22:31 +0200 Subject: [PATCH 121/952] [CI] Skip EmbeddingGemma2 pooling tests needing transformers >= 5.19.0 (#60371) Signed-off-by: Stefan Koncarevic --- .../multimodal/pooling/test_embedding_gemma2.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 701d0914ddbc..7bb69fa69ad8 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -9,11 +9,20 @@ import pytest import torch +from tests.models.registry import HF_EXAMPLE_MODELS + MODEL = os.environ.get("EG2_MODEL_PATH") needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") TEXT_COS, MM_COS = 0.999, 0.999 +@pytest.fixture +def needs_hf_embedding_gemma2(): + """Skip when the installed transformers has no `embedding_gemma2` module.""" + model_info = HF_EXAMPLE_MODELS.get_hf_info("EmbeddingGemma2Model") + model_info.check_transformers_version(on_fail="skip") + + def _cos(a, b): return torch.nn.functional.cosine_similarity( torch.tensor(a, dtype=torch.float32), @@ -239,7 +248,7 @@ def nparams(m): ) -def test_encoder_contract(): +def test_encoder_contract(needs_hf_embedding_gemma2): from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration as G, @@ -423,7 +432,7 @@ def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): assert len(out[0]) == 768 -def test_video_loader_indices(): +def test_video_loader_indices(needs_hf_embedding_gemma2): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) @@ -508,7 +517,7 @@ def test_video_loader_indices(): assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" -def test_video_loader_duplicate_expansion(): +def test_video_loader_duplicate_expansion(needs_hf_embedding_gemma2): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) From d9503dc503e75193d2128908b59c0d141d4b5fb1 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Wed, 7 Oct 2026 21:24:40 +0900 Subject: [PATCH 122/952] [Bugfix][Qwen4Exp] Derive the QSA head_dim from the unsharded head count (#59945) Signed-off-by: lesj0610 Co-authored-by: lesj0610 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/amd/qsa.py | 2 +- vllm/models/qwen4_exp/nvidia/qsa.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 325e0e3d917c..3f03b293b4f0 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -214,7 +214,7 @@ def __init__( elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index a20fcdeec98d..0246db80943d 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -393,7 +393,7 @@ def __init__( elif tp_size % self.total_num_kv_heads: raise ValueError("TP size must be divisible by replicated QSA KV heads") self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.head_dim = int(config.head_dim or self.hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 From 2a54f6b625f28109180b072b704c0b0d372a277d Mon Sep 17 00:00:00 2001 From: Jin Tao Date: Wed, 7 Oct 2026 15:53:40 +0300 Subject: [PATCH 123/952] [ROCm][Perf][GLM-5.3-Flash] Enable AITER fused shared experts on gfx950 (#59221) Signed-off-by: Jin Tao Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .../layers/test_fused_shared_expert.py | 119 +++++++++++++++++- vllm/models/glm5next/common/model.py | 85 ++++++++++++- vllm/models/glm5next/common/mtp.py | 15 ++- 3 files changed, 213 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 6721a4c63b56..d8d5c4fb032a 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -14,7 +14,7 @@ from torch import nn import vllm.config as vllm_config_module -from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -145,6 +145,12 @@ def _stub_quant_config() -> QuantizationConfig: return cast(QuantizationConfig, object()) +class _StubGlm5NextAttention(nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + def get_fse_test_model_config( model_type: str, quantization_config: dict[str, Any], @@ -279,6 +285,33 @@ def get_fse_test_model_config( quantization_config=quantization_config, ) return config, Glm4MoeModel + if model_type == "glm5_next": + from transformers import Glm5NextTextConfig + + from vllm.models.glm5next.common.model import Glm5NextModel + + config = Glm5NextTextConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + layer_types=["linear_attention"], + mlp_layer_types=["sparse"], + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + quantization_config=quantization_config, + ) + return config, Glm5NextModel raise ValueError(f"Unsupported FSE test model: {model_type}") @@ -796,7 +829,7 @@ def __init__(self, enabled: bool) -> None: @pytest.mark.parametrize( "model_type", - ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2"], + ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"], ) @pytest.mark.parametrize( ("use_fse", "exclude"), @@ -850,6 +883,16 @@ def test_models_fse_init( ) else: vllm_config.quant_config = QuarkConfig(quantization_config) + if model_type == "glm5_next": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubGlm5NextAttention) + # The gate allows fusion only at TP4/TP8 on gfx950, and this test + # runs at TP1; the gate has its own test. + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops @@ -902,6 +945,10 @@ def test_models_fse_init( ) mtp = DeepSeekV4MTP(vllm_config=vllm_config) assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) + if model_type == "glm5_next": + assert (model.layers[0].mlp.shared_experts is None) is ( + fse_enabled and not exclude + ) # The dummy quant config here uses mixed mxfp4/fp8 for experts/shared_expert # so should just raise a warning. @@ -926,6 +973,74 @@ def test_models_fse_init( rocm_aiter_ops.refresh_env_variables() +@pytest.mark.parametrize( + ("setting", "value", "reason"), + [ + ("tensor_parallel_size", 4, None), + ("tensor_parallel_size", 8, None), + ("tensor_parallel_size", 2, "tensor_parallel_size is 2"), + ("data_parallel_size", 2, "data_parallel_size is 2"), + ("prefill_context_parallel_size", 2, "prefill_context_parallel_size is 2"), + ("enable_expert_parallel", True, "expert parallelism is enabled"), + ("on_gfx950", False, "the GPU is not gfx950"), + ], +) +def test_glm5_next_fuses_shared_experts_only_in_tuned_setups( + monkeypatch: pytest.MonkeyPatch, + setting: str, + value: object, + reason: str | None, +) -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + parallel_config = SimpleNamespace( + tensor_parallel_size=4, + data_parallel_size=1, + prefill_context_parallel_size=1, + enable_expert_parallel=False, + ) + on_gfx950 = value if setting == "on_gfx950" else True + if setting != "on_gfx950": + setattr(parallel_config, setting, value) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: on_gfx950) + + with patch.object(glm5_next_model.logger, "warning_once") as warning: + tuned = glm5_next_model._fused_shared_experts_tuned( + cast(ParallelConfig, parallel_config) + ) + + assert tuned is (reason is None) + if reason is None: + warning.assert_not_called() + else: + warning.assert_called_once() + assert warning.call_args.args[1] == reason + + +def test_glm5_next_fused_shared_expert_loads_into_the_shared_slot() -> None: + from vllm.models.glm5next.common import model as glm5_next_model + + assert glm5_next_model._num_fused_shared_experts(1, True) == 1 + assert glm5_next_model._num_fused_shared_experts(1, False) == 0 + assert glm5_next_model._num_fused_shared_experts(None, True) == 0 + with pytest.raises(NotImplementedError, match="only 1 shared expert"): + glm5_next_model._num_fused_shared_experts(2, True) + + for shared, fused in [ + ( + "model.layers.3.mlp.shared_experts.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.288.gate_proj.weight_scale_inv", + ), + ( + "model.layers.45.mtp_block.mlp.shared_experts.down_proj.weight", + "model.layers.45.mtp_block.mlp.experts.288.down_proj.weight", + ), + ]: + assert glm5_next_model._fused_shared_expert_name(shared, 288) == fused + routed = "model.layers.3.mlp.experts.7.down_proj.weight" + assert glm5_next_model._fused_shared_expert_name(routed, 288) == routed + + @pytest.mark.parametrize( ("exclude", "expected"), [([], True), ([r"re:.*\.shared_expert\..*"], False)], diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e15c00ecb88b..778738c13712 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -23,6 +23,10 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, +) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -258,7 +262,12 @@ def __init__( ) swiglu_limit = config.swiglu_limit - if config.n_shared_experts is None: + self.is_fused_shared_expert_enabled = False + if config.n_shared_experts is not None: + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix + ) and _fused_shared_experts_tuned(parallel_config) + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -294,7 +303,10 @@ def __init__( enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, - n_shared_experts=None, + n_shared_experts=config.n_shared_experts + if self.is_fused_shared_expert_enabled + else None, + fuse_shared_experts=self.is_fused_shared_expert_enabled, router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) @@ -733,6 +745,9 @@ def get_layer(prefix: str): # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, Glm5NextMoE, "mlp" + ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -834,12 +849,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ), 0, ) + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) expert_params_mapping = fused_moe_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.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=num_redundant_experts, ) else: @@ -868,6 +886,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) # Handle FP8 indexer WK: dequantize to BF16 for fusion with # weights_proj into wk_weights_proj. @@ -1228,6 +1248,65 @@ def get_spec_layer_idx_from_weight_name( return None +def _fused_shared_experts_tuned(parallel_config: ParallelConfig) -> bool: + """AITER has fused-MoE configs tuned for the fused shared-expert shape + (one more expert and one more top-k slot than the routed MoE) only on + gfx950, with every expert on each rank and its weights split by TP4 or + TP8. Data, prefill context and expert parallelism change that split, so + any other GPU or parallel layout would run untuned fallback kernels.""" + from vllm.platforms.rocm import on_gfx950 + + reasons: list[str] = [] + if not on_gfx950(): + reasons.append("the GPU is not gfx950") + if parallel_config.tensor_parallel_size not in (4, 8): + reasons.append( + f"tensor_parallel_size is {parallel_config.tensor_parallel_size}" + ) + if parallel_config.data_parallel_size != 1: + reasons.append(f"data_parallel_size is {parallel_config.data_parallel_size}") + if parallel_config.prefill_context_parallel_size != 1: + reasons.append( + "prefill_context_parallel_size is " + f"{parallel_config.prefill_context_parallel_size}" + ) + if parallel_config.enable_expert_parallel: + reasons.append("expert parallelism is enabled") + + if not reasons: + return True + logger.warning_once( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is ignored for GLM-5.3-Flash: " + "%s. AITER has tuned configs for its fused shared-expert MoE only on " + "gfx950 at TP4 and TP8, without data, prefill context or expert " + "parallelism. Running the shared experts as a separate MLP.", + "; ".join(reasons), + ) + return False + + +def _num_fused_shared_experts(n_shared_experts: int | None, enabled: bool) -> int: + """Expert slots the fused MoE appends for the shared expert; must match the + ``num_fused_shared_experts`` that ``FusedMoE`` allocates.""" + if not enabled or n_shared_experts is None: + return 0 + if n_shared_experts > 1: + raise NotImplementedError( + "Fused shared-expert loading supports only 1 shared expert per " + f"layer, but config.n_shared_experts is {n_shared_experts}. Set " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared " + "experts as a separate MLP." + ) + return n_shared_experts + + +def _fused_shared_expert_name(name: str, n_routed_experts: int) -> str: + """Point a checkpoint ``mlp.shared_experts.*`` tensor at the fused MoE's + shared-expert slot, which follows the routed experts; other names are + returned unchanged.""" + return name.replace("mlp.shared_experts.", f"mlp.experts.{n_routed_experts}.", 1) + + def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): if "indexer.wk." not in name or "wk_weights" in name: return False diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py index 816cc46fdc08..2bb8b10d0746 100644 --- a/vllm/models/glm5next/common/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -10,6 +10,9 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -30,6 +33,8 @@ Glm5NextDecoderLayer, Glm5NextMLAAttention, Glm5NextMoE, + _fused_shared_expert_name, + _num_fused_shared_experts, _try_load_fp8_attn_proj, _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, @@ -217,6 +222,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp" + ) def set_moe_parameters(self): self.num_moe_layers = self.config.num_nextn_predict_layers @@ -298,12 +306,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + num_fused_shared = _num_fused_shared_experts( + self.config.n_shared_experts, self.is_fused_shared_expert_enabled + ) expert_params_mapping = fused_moe_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.n_routed_experts, + num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=self.num_redundant_experts, ) @@ -328,6 +339,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) + if self.is_fused_shared_expert_enabled: + name = _fused_shared_expert_name(name, self.config.n_routed_experts) if _try_load_fp8_indexer_wk( name, From efb8d85f24bd270bf20c143b32b5e7db6df61716 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:19:22 -0400 Subject: [PATCH 124/952] [ROCm][CI] Add FP32 rounding slack to accuracy test near-tie upper bounds (#60292) Signed-off-by: Divakar Verma Co-authored-by: Andreas Karatzas --- tests/v1/e2e/general/accuracy_utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/v1/e2e/general/accuracy_utils.py b/tests/v1/e2e/general/accuracy_utils.py index c8fa770fc5bf..e15ef03ebcfe 100644 --- a/tests/v1/e2e/general/accuracy_utils.py +++ b/tests/v1/e2e/general/accuracy_utils.py @@ -110,8 +110,10 @@ def check_first_divergence( # Log-softmax and reconstructing penalties may round in FP32, but # a near tie cannot excuse selecting the engine's own worse token. assert gap >= -1e-5, f"{detail}: selected token is below its alternative" - assert gap <= tolerance.greedy_atol, ( - f"{detail}: divergence exceeds near-tie bound={tolerance.greedy_atol}" + # Same FP32 rounding slack as above, applied to the upper bound. + assert gap <= tolerance.greedy_atol + 1e-5, ( + f"{detail}: divergence exceeds near-tie bound=" + f"{tolerance.greedy_atol} (+1e-5 slack)" ) @@ -201,9 +203,10 @@ def check_greedy_token( ) best = int(logits.argmax()) gap = float(logits[best] - logits[token]) - assert gap <= tolerance.greedy_atol, ( + # FP32 rounding can push a true near-tie fractionally past the bound. + assert gap <= tolerance.greedy_atol + 1e-5, ( f"{context}: selected token={token}, reference best={best}, " - f"logit gap={gap}, bound={tolerance.greedy_atol}" + f"logit gap={gap}, bound={tolerance.greedy_atol} (+1e-5 slack)" ) return gap From 5281e49908960f490c8faa789eee221d4ab5001d Mon Sep 17 00:00:00 2001 From: 00CC <83797617+tangzzycc@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:21:43 +0800 Subject: [PATCH 125/952] [ROCm][Perf] Enable medium-skinny dispatch in RDNAHybridW4A16LinearKernel (#52619) Signed-off-by: tangzzycc <3081129260@qq.com> --- .../quantization/test_rdna_hybrid_w4a16.py | 81 ++++++++++++++++++- .../mixed_precision/rdna_hybrid_w4a16.py | 16 ++-- 2 files changed, 88 insertions(+), 9 deletions(-) diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index d6045a791c1f..da7ef2d437d5 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -30,6 +30,8 @@ pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE +LDS_CAPACITY_ELEMENTS = hybrid_module.LDS_CAPACITY_ELEMENTS +MEDIUM_SKINNY_LIMIT_ELEMENTS = hybrid_module.MEDIUM_SKINNY_LIMIT_ELEMENTS # --------------------------------------------------------------------------- @@ -100,8 +102,8 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) """Smoke test the registered custom op for both decode and prefill batches. Verifies the dispatch logic in `_rdna_hybrid_w4a16_apply_impl`: - - M <= MAX_SKINNY_BATCH_SIZE: HIP wvSplitK_int4_g - - M > MAX_SKINNY_BATCH_SIZE: Triton prefill kernel + - supported M and K*M: HIP wvSplitK_int4_g + - otherwise: Triton prefill kernel """ if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -742,3 +744,78 @@ def test_process_weights_pads_cliff_rows(K, expected_weight_stride, dist_init): assert tuple(layer.weight_packed.view(torch.int32).shape) == (N, K // 8) # Metadata is untouched by this change. assert layer.weight_scale.is_contiguous() + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize( + "M,K,expected_path", + [ + (4, 8192, "hip"), + (4, 9728, "hip"), + (4, 9856, "triton"), + (MAX_SKINNY_BATCH_SIZE + 1, 1024, "triton"), + ], + ids=["regular_limit", "medium_range", "above_medium", "batch_too_large"], +) +def test_rdna_hybrid_w4a16_dispatch_boundaries(M, K, expected_path, monkeypatch): + """Route only supported regular and medium skinny shapes to HIP.""" + import vllm._custom_ops as ops + + N, G = 16, 128 + x_mk = torch.empty((M, K), dtype=torch.float16) + w_q = torch.empty((N, K // 2), dtype=torch.int8) + scales = torch.empty((N, K // G), dtype=torch.float16) + called = [] + + def fake_hip(*args, **kwargs): + called.append("hip") + return torch.empty((M, N), dtype=x_mk.dtype) + + def fake_triton(*args, **kwargs): + called.append("triton") + return torch.empty((M, N), dtype=x_mk.dtype) + + monkeypatch.setattr(ops, "wvSplitK_int4_g", fake_hip) + monkeypatch.setattr(hybrid_module, "triton_w4a16_skinny_fmt_gemm", fake_triton) + + output = hybrid_module._rdna_hybrid_w4a16_apply_impl( + x_mk, w_q, scales, None, None, 40, G + ) + + assert output.shape == (M, N) + assert called == [expected_path] + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True]) +def test_rdna_hybrid_w4a16_medium_skinny_matches_reference(dtype, has_zp): + """Validate the HIP medium path against an FP32 dequantization reference.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + + from vllm.utils.platform_utils import num_compute_units + + set_random_seed(0) + M, K, N, G = 4, 9728, 256, 128 + assert LDS_CAPACITY_ELEMENTS < K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS + + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4_nk = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) + w_q = pack_int4_exllama_shuffle(w_int4_nk).contiguous().view(torch.int8) + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( + dtype + ) + zp = ( + torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32) + if has_zp + else None + ) + w_zp = _pack_zp_rows_for_kernel(zp) if zp is not None else None + + output = torch.ops.vllm.rdna_hybrid_w4a16_apply( + x_mk, w_q, scales, w_zp, None, num_compute_units(), G + ) + reference = _rdna_hybrid_w4a16_reference(x_mk, w_int4_nk, scales, zp, G, bias=None) + + torch.testing.assert_close(output, reference, rtol=2e-2, atol=2e-2) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 9d5303469945..bed62e89afce 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Hybrid W4A16 kernel: Triton for prefill, HIP skinny for decode. -Routes based on batch size M: - M <= MAX_SKINNY_BATCH_SIZE: HIP skinny GEMM (wvSplitK_int4_g) - M > MAX_SKINNY_BATCH_SIZE: Triton W4A16 fused dequant GEMM +Routes based on the activation shape: + M <= MAX_SKINNY_BATCH_SIZE and K*M <= MEDIUM_SKINNY_LIMIT_ELEMENTS: + HIP skinny GEMM (wvSplitK_int4_g) + otherwise: Triton W4A16 fused dequant GEMM Stores the weights ONCE as int8 [N, K//2] (ExLlama shuffle packed). Both paths read this single buffer: the HIP skinny kernel uses it directly, and @@ -56,13 +57,14 @@ def _on_gfx1151() -> bool: return on_gfx1151() -# Maximum batch size M for the HIP skinny kernel path (C++ supports N_in -# up to 5). When M is below this AND K*M fits in LDS, the skinny kernel is -# used; otherwise the Triton prefill path handles the GEMM. +# Maximum activation rows supported by HIP skinny (Python M maps to C++ N_in). +# When K*M exceeds regular LDS capacity, the C++ medium kernel caches the +# prefix in LDS and reads the remaining activation elements from global memory. MAX_SKINNY_BATCH_SIZE = 5 # 64 KiB per-workgroup LDS limit expressed in fp16 elements. # (AMD RDNA has 128 KiB total LDS per CU, but 64 KiB per workgroup.) LDS_CAPACITY_ELEMENTS = 64 * 1024 // 2 # 32768 fp16 elements +MEDIUM_SKINNY_LIMIT_ELEMENTS = int(LDS_CAPACITY_ELEMENTS * 1.2) # --------------------------------------------------------------------------- @@ -464,7 +466,7 @@ def _rdna_hybrid_w4a16_apply_impl( K = x_2d.shape[1] N = w_q.shape[0] - if M <= MAX_SKINNY_BATCH_SIZE and K * M <= LDS_CAPACITY_ELEMENTS: + if M <= MAX_SKINNY_BATCH_SIZE and K * M <= MEDIUM_SKINNY_LIMIT_ELEMENTS: # record_function is not torch.compile-safe; use nullcontext when # compiling to keep the op traceable. ctx = ( From c83935b3008553c72b1bad09f842b4c99b7fa47f Mon Sep 17 00:00:00 2001 From: viktorkou <105151911+viktorkou@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:04:35 +0300 Subject: [PATCH 126/952] [Bugfix][Model Loader] Support reload_weights with runai_streamer load format (#58124) Signed-off-by: Viktor Koukouliev Co-authored-by: Claude Fable 5.1 Co-authored-by: Isotr0py --- .../test_runai_model_streamer_loader.py | 93 +++++++++++++++++++ .../test_runai_model_streamer_s3.py | 71 ++++++++++++++ tests/v1/worker/test_gpu_model_runner.py | 30 ++++++ .../model_loader/runai_streamer_loader.py | 15 ++- vllm/v1/worker/gpu_model_runner.py | 6 +- 5 files changed, 208 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index 9164b8e4bea9..e0a083b3a8bc 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -6,11 +6,13 @@ from unittest.mock import patch import pytest +import torch from vllm import SamplingParams from vllm.config.load import LoadConfig from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader import runai_streamer_loader as rsl +from vllm.utils.mem_constants import GiB_bytes load_format = "runai_streamer" test_model = "openai-community/gpt2" @@ -122,3 +124,94 @@ def test_runai_invalid_extra_config_leaves_environ_untouched(): with pytest.raises(ValueError, match="memory_limit must be an integer >= -1"): _runai_loader({"concurrency": 16, "memory_limit": -5}) assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ + + +@pytest.mark.parametrize( + "model_config,expected_source", + [ + # Object storage: model_weights holds the URI, model the pulled config dir. + ( + types.SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/weights", + revision="myrev", + ), + "s3://bucket/weights", + ), + # HF repo or local path: model_weights is empty, model is the source. + ( + types.SimpleNamespace( + model="org/model", model_weights="", revision="myrev" + ), + "org/model", + ), + ], + ids=["object_storage", "hf_repo"], +) +def test_runai_get_all_weights_resolves_source(model_config, expected_source): + fake_self = types.SimpleNamespace( + _get_weights_iterator=lambda path, revision: iter( + [(f"{path}@{revision}", None)] + ) + ) + + weights = list( + rsl.RunaiModelStreamerLoader.get_all_weights( + fake_self, model_config, model=None + ) + ) + + assert weights == [(f"{expected_source}@myrev", None)] + + +@pytest.mark.parametrize( + "base_model,mul_model,add_model", + [ + ( + "Qwen/Qwen3-0.6B", + "inference-optimization/Qwen3-0.6B-debug-multiply", + "inference-optimization/Qwen3-0.6B-debug-add", + ), + ], +) +def test_runai_deep_sleep_reload_weights(base_model, mul_model, add_model, vllm_runner): + free, total = torch.accelerator.get_memory_info() + used_bytes_baseline = total - free + + def sleep_and_reload(llm, path): + # Level 2 discards the parameter memory, so after wake_up the weights + # can only come from the reload: a tensor the streamer fails to + # deliver changes the output instead of keeping a stale valid value. + llm.get_llm().sleep(level=2) + free, total = torch.accelerator.get_memory_info() + assert total - free - used_bytes_baseline < 3 * GiB_bytes + llm.get_llm().wake_up(tags=["weights"]) + llm.collective_rpc("reload_weights", kwargs={"weights_path": path}) + free, total = torch.accelerator.get_memory_info() + assert total - free - used_bytes_baseline < 4 * GiB_bytes + llm.get_llm().wake_up(tags=["kv_cache"]) + + with vllm_runner( + model_name=base_model, + load_format=load_format, + enable_sleep_mode=True, + enable_prefix_caching=False, + max_model_len=16, + max_num_seqs=1, + ) as llm: + base_output = llm.generate_greedy(["3 4 ="], max_tokens=4) + + sleep_and_reload(llm, mul_model) + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] + assert mul_perp < add_perp + + sleep_and_reload(llm, add_model) + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] + assert add_perp < mul_perp + + # Round trip to the original weights over discarded memory must + # reproduce the pre-sleep output. + sleep_and_reload(llm, base_model) + assert llm.generate_greedy(["3 4 ="], max_tokens=4) == base_output diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py index d7abc80edbfa..a1440ae5d282 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py @@ -3,7 +3,9 @@ from pathlib import Path +import torch from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher +from safetensors.torch import load_file, save_file from vllm.engine.arg_utils import EngineArgs from vllm.transformers_utils.repo_utils import hf_api @@ -50,3 +52,72 @@ def test_runai_model_loader_download_files_s3_mocked_with_patch( executor = RunaiDummyExecutor(vllm_config) executor.driver_worker.load_model() + + +def test_runai_reload_weights_path_from_mocked_s3(tmp_path: Path, monkeypatch): + # An engine started from object storage keeps the bucket URI in + # model_config.model_weights. A reload with weights_path must stream the + # new checkpoint, not re-stream the one the engine started from. + patcher = StreamerPatcher(str(tmp_path)) + monkeypatch.setattr( + "vllm.transformers_utils.runai_utils.runai_list_safetensors", + patcher.shim_list_safetensors, + ) + monkeypatch.setattr( + "vllm.transformers_utils.runai_utils.runai_pull_files", + patcher.shim_pull_files, + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.weight_utils.SafetensorsStreamer", + patcher.create_mock_streamer, + ) + + base_dir = tmp_path / "gpt2" + hf_api().snapshot_download( + repo_id=test_model, + local_dir=str(base_dir), + allow_patterns=["*.json", "*.txt", "model.safetensors"], + ) + # Second checkpoint under another prefix: the final layer norm scaled, so a + # reload that silently re-streams the original bucket is distinguishable. + tensors = load_file(base_dir / "model.safetensors") + original = tensors["ln_f.weight"].clone() + modified = original * 2 + tensors["ln_f.weight"] = modified + mod_dir = tmp_path / "gpt2-mod" + mod_dir.mkdir() + save_file(tensors, mod_dir / "model.safetensors", metadata={"format": "pt"}) + + vllm_config = EngineArgs( + model="s3://my-mock-bucket/gpt2/", + load_format=load_format, + tensor_parallel_size=1, + ).create_engine_config() + assert vllm_config.model_config.model_weights == "s3://my-mock-bucket/gpt2/" + + executor = RunaiDummyExecutor(vllm_config) + executor.driver_worker.load_model() + model = executor.driver_worker.model_runner.get_model() + (param,) = [p for n, p in model.named_parameters() if n.endswith("ln_f.weight")] + + def current() -> torch.Tensor: + return param.detach().float().cpu() + + torch.testing.assert_close(current(), original, rtol=1e-2, atol=1e-2) + + # No weights_path: re-stream the checkpoint the engine started from, i.e. + # the bucket kept in model_weights. Zero the live tensor first so a reload + # that did nothing is distinguishable from one that streamed the bucket. + param.data.zero_() + executor.driver_worker.reload_weights() + torch.testing.assert_close(current(), original, rtol=1e-2, atol=1e-2) + + # weights_path: the new prefix must win over the stale model_weights. + executor.driver_worker.reload_weights(weights_path="s3://my-mock-bucket/gpt2-mod/") + torch.testing.assert_close(current(), modified, rtol=1e-2, atol=1e-2) + + # After the override, a bare reload means the new checkpoint, not the + # one the engine was started with. + param.data.zero_() + executor.driver_worker.reload_weights() + torch.testing.assert_close(current(), modified, rtol=1e-2, atol=1e-2) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 5cf53899b836..deb4f7b42147 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1015,6 +1015,36 @@ def test_reload_weights_before_load_model(model_runner): model_runner.reload_weights() +def test_reload_weights_path_replaces_object_storage_source(monkeypatch): + # An engine started from an object-storage URI keeps that URI in + # model_config.model_weights, which the runai_streamer loader prefers over + # model_config.model. A new weights_path must replace it, or the loader + # silently re-streams the original checkpoint. + loader = Mock() + loader.get_all_weights.return_value = iter(()) + monkeypatch.setattr(gpu_model_runner_module, "get_model_loader", lambda _: loader) + monkeypatch.setattr(gpu_model_runner_module, "initialize_layerwise_reload", Mock()) + monkeypatch.setattr(gpu_model_runner_module, "finalize_layerwise_reload", Mock()) + + runner = Mock(lora_config=None) + runner.model_config = SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/original", + revision="abc123", + quantization=None, + ) + runner.get_model.return_value.named_parameters.return_value = [] + runner.get_model.return_value.load_weights.return_value = None + + GPUModelRunner.reload_weights(runner, weights_path="org/new-model") + + loader.get_all_weights.assert_called_once_with( + runner.model_config, runner.get_model.return_value + ) + cfg = runner.model_config + assert (cfg.model, cfg.model_weights, cfg.revision) == ("org/new-model", "", None) + + def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): runner = object.__new__(GPUModelRunner) runner.use_async_scheduling = False diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 2f18a2e3aca8..e00116006212 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -134,11 +134,16 @@ def download_model(self, model_config: ModelConfig) -> None: """Download model if necessary.""" self._prepare_weights(model_config.model, model_config.revision) - def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: - """Load weights into a model.""" + def get_all_weights( + self, + model_config: ModelConfig, + model: nn.Module, + ) -> Generator[tuple[str, torch.Tensor], None, None]: model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override - model.load_weights( - self._get_weights_iterator(model_weights, model_config.revision) - ) + yield from self._get_weights_iterator(model_weights, model_config.revision) + + def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: + """Load weights into a model.""" + model.load_weights(self.get_all_weights(model_config, model)) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 76299308f866..0cc3fcc29a8f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5534,9 +5534,11 @@ def reload_weights( ) if weights_path is not None: - # The revision belongs to the model we are reloading away from, - # so it must not be carried over to the new path. + # The revision and any object-storage `model_weights` source + # belong to the model we are reloading away from, so they must + # not be carried over to the new path. self.model_config.model = weights_path + self.model_config.model_weights = "" self.model_config.revision = None weights_iterator = model_loader.get_all_weights(self.model_config, model) weights_iterator = cast( From f6e3944e3100fcb526e5cc24248080c31812be0f Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Wed, 7 Oct 2026 09:35:53 -0500 Subject: [PATCH 127/952] [ROCm][DSv4] AITER MegaMoEV2 Integration For DeepSeek V4 (#59685) Signed-off-by: Micah Williamson Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 + .buildkite/test_areas/kernels.yaml | 2 + .../models/test_deepseek_v4_aiter_mega_moe.py | 186 ++++++++++++++++++ tests/models/test_deepseek_v4_vl_rocm.py | 4 +- vllm/config/kernel.py | 4 + .../device_communicators/all2all.py | 5 + vllm/models/deepseek_v4/amd/dspark.py | 14 +- vllm/models/deepseek_v4/amd/mega_moe.py | 154 +++++++++++++++ vllm/models/deepseek_v4/amd/model.py | 81 ++++++++ vllm/models/deepseek_v4/amd/mtp.py | 8 + 10 files changed, 454 insertions(+), 6 deletions(-) create mode 100644 tests/models/test_deepseek_v4_aiter_mega_moe.py create mode 100644 vllm/models/deepseek_v4/amd/mega_moe.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 8317919ce4b8..9e634ea02b4d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -5858,6 +5858,7 @@ steps: timeout_in_minutes: 20 source_file_dependencies: - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_aiter_mega_moe.py - vllm/models/deepseek_v4/amd/ - vllm/models/deepseek_v4/common/ - vllm/model_executor/layers/fused_moe/ @@ -5871,6 +5872,7 @@ steps: - vllm/envs.py commands: - VLLM_ROCM_USE_AITER_MOE=0 pytest -v -s models/test_deepseek_v4_vl_rocm.py + - pytest -v -s models/test_deepseek_v4_aiter_mega_moe.py dind: false working_dir: /vllm-workspace/tests agent_pool: mi355_1 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 54e7215967cd..cfc628d382c6 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -155,6 +155,7 @@ steps: - image-build-amd source_file_dependencies: - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_aiter_mega_moe.py - vllm/models/deepseek_v4/amd/ - vllm/models/deepseek_v4/common/ - vllm/model_executor/layers/fused_moe/ @@ -170,6 +171,7 @@ steps: # Common qnorm/RoPE/cache kernels already run in the H100 AMD mirror. # Pin the tested BF16 Triton routing and shared-expert path before imports. - VLLM_ROCM_USE_AITER_MOE=0 pytest -v -s models/test_deepseek_v4_vl_rocm.py + - pytest -v -s models/test_deepseek_v4_aiter_mega_moe.py # Own job so mHC never shares a process with FlashInfer's import-time # ptxas override. diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py new file mode 100644 index 000000000000..a8f692e07bfb --- /dev/null +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.models.deepseek_v4.amd import mega_moe as mega_moe_module +from vllm.models.deepseek_v4.amd.mega_moe import ( + DeepseekV4AiterMegaMoEExperts, +) +from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE +from vllm.platforms import current_platform + +HIDDEN = 64 +INTER = 64 + + +def _make_experts( + monkeypatch, + *, + ep_rank=0, + ep_size=1, + num_experts=2, + intermediate_size=INTER, + max_num_batched_tokens=3000, + swiglu_limit=None, +): + monkeypatch.setattr( + mega_moe_module, + "get_ep_group", + lambda: SimpleNamespace( + rank_in_group=ep_rank, world_size=ep_size, cpu_group="ep" + ), + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens) + ) + return DeepseekV4AiterMegaMoEExperts( + vllm_config, + num_experts=num_experts, + top_k=2, + hidden_size=HIDDEN, + intermediate_size=intermediate_size, + swiglu_limit=swiglu_limit, + ) + + +@pytest.mark.cpu_test +def test_experts_load_into_local_slots(monkeypatch): + model = torch.nn.Module() + model.ffn = torch.nn.Module() + model.ffn.experts = _make_experts(monkeypatch, ep_rank=1, ep_size=2, num_experts=4) + params = dict(model.named_parameters()) + mapping = fused_moe_make_expert_params_mapping( + model, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=4, + routed_experts_prefix="", + ) + + shapes = { + ("w1", "weight"): (INTER, HIDDEN // 2), + ("w3", "weight"): (INTER, HIDDEN // 2), + ("w2", "weight"): (HIDDEN, INTER // 2), + ("w1", "weight_scale"): (INTER, HIDDEN // 32), + ("w3", "weight_scale"): (INTER, HIDDEN // 32), + ("w2", "weight_scale"): (HIDDEN, INTER // 32), + } + shard_value = {"w1": 1, "w3": 2, "w2": 3} + loaded = {} + for expert in range(4): + for (shard, suffix), shape in shapes.items(): + name = f"ffn.experts.{expert}.{shard}.{suffix}" + weight = torch.full( + shape, 10 * expert + shard_value[shard], dtype=torch.uint8 + ) + for param_name, weight_name, expert_id, shard_id in mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + param = params[mapped] + if param.weight_loader( + param, + weight, + mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ): + loaded[name] = mapped + break + + assert sorted({n.split(".")[2] for n in loaded}) == ["2", "3"] + assert len(loaded) == 2 * len(shapes) + for slot, expert in enumerate((2, 3)): + for suffix in ("weight", "weight_scale"): + w13 = params[f"ffn.experts.w13_{suffix}"][slot] + w2 = params[f"ffn.experts.w2_{suffix}"][slot] + assert torch.all(w13[:INTER] == 10 * expert + 1) + assert torch.all(w13[INTER:] == 10 * expert + 2) + assert torch.all(w2 == 10 * expert + 3) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") +@pytest.mark.parametrize("hash_routing", [False, True]) +@pytest.mark.parametrize("num_tokens", [1, 37]) +def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_tokens): + num_experts, top_k, vocab = 16, 4, 64 + scaling = 2.5 + device = "cuda" + torch.manual_seed(0) + + class Gate(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.randn(num_experts, HIDDEN, device=device) * 0.1 + self.e_score_correction_bias = ( + None if hash_routing else torch.randn(num_experts, device=device) + ) + self.tid2eid = ( + torch.stack( + [ + torch.randperm(num_experts, device=device)[:top_k] + for _ in range(vocab) + ] + ).to(torch.int32) + if hash_routing + else None + ) + self.bias_vl = None + + def forward(self, x): + return x.float() @ self.weight.t(), None + + captured = {} + + def experts(x, topk_weights, topk_ids): + captured["weights"], captured["ids"] = topk_weights, topk_ids + return 2 * x + + moe = DeepseekV4MoE.__new__(DeepseekV4MoE) + torch.nn.Module.__init__(moe) + moe.gate = Gate() + moe.experts = experts + moe.shared_experts = lambda x: 3 * x + moe.scoring_func = "sqrtsoftplus" + moe.n_activated_experts = top_k + moe.renormalize = True + moe.routed_scaling_factor = scaling + moe.image_sentinel_lo = 0 + + x = torch.randn(num_tokens, HIDDEN, device=device, dtype=torch.bfloat16) + input_ids = torch.randint(0, vocab, (num_tokens,), device=device) + out = moe._forward_mega_moe(x, input_ids) + + ids, weights = captured["ids"], captured["weights"] + assert ids.dtype == torch.int32 and ids.is_contiguous() + assert weights.dtype == torch.float32 and weights.is_contiguous() + + scores = torch.nn.functional.softplus(x.float() @ moe.gate.weight.t()).sqrt() + if hash_routing: + expected_ids = moe.gate.tid2eid[input_ids].long() + else: + expected_ids = (scores + moe.gate.e_score_correction_bias).topk(top_k).indices + expected_weights = scores.gather(1, expected_ids) + expected_weights *= scaling / expected_weights.sum(dim=-1, keepdim=True) + + order = ids.long().sort(dim=-1).indices + expected_order = expected_ids.sort(dim=-1).indices + torch.testing.assert_close( + ids.long().gather(1, order), expected_ids.gather(1, expected_order) + ) + torch.testing.assert_close( + weights.gather(1, order), + expected_weights.gather(1, expected_order), + rtol=1e-3, + atol=1e-4, + ) + torch.testing.assert_close(out, 5 * x) diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 036efbaa4248..762f1cd6c0e0 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -7,7 +7,7 @@ import torch from torch import nn -from vllm.config import ParallelConfig +from vllm.config import KernelConfig, ParallelConfig from vllm.model_executor.models.utils import WeightsMapper from vllm.platforms import current_platform @@ -58,6 +58,7 @@ def test_rocm_moe_routing_and_shared_experts_match_reference( model_config=SimpleNamespace(hf_config=config), quant_config=None, parallel_config=ParallelConfig(), + kernel_config=default_vllm_config.kernel_config, ), prefix="model.layers.0.ffn", ) @@ -230,6 +231,7 @@ def fake_factory(**kwargs): model_config=SimpleNamespace(hf_config=config), quant_config=None, parallel_config=ParallelConfig(), + kernel_config=KernelConfig(), ) hash_moe = rocm_model.DeepseekV4MoE(vllm_config, prefix="model.layers.0.ffn") diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index b8bfca082a4a..fa3ecf5df95d 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -137,6 +137,7 @@ def with_default( "humming", "triton_unfused", "aiter", + "aiter_mega_moe", "aiter_triton_mxfp4_bf16", "flydsl", "hpc", @@ -317,6 +318,9 @@ class KernelConfig: - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels - "aiter": Use AMD AITer kernels (ROCm only) + - "aiter_mega_moe": Use the AITER MegaMoE fused dispatch/GEMM/combine + kernel (ROCm gfx950, requires DP+EP with a MoRI all2all backend; + DeepSeek-V4 only) - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16 (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250) - "flydsl": Use AMD FlyDSL kernels (ROCm only) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index d8790dbf42cb..ae31c0b1d3cd 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1008,6 +1008,11 @@ def __init__(self, cpu_group, all2all_backend: str): self.handle_cache = Cache() torch._C._distributed_c10d._register_process_group("mori", cpu_group) + if get_current_vllm_config().kernel_config.moe_backend == "aiter_mega_moe": + # MegaMoEV2 places its dispatch/combine workspaces on the MoRI + # symmetric heap, which defaults to 2 GB. MegaMoEV2 requires > 4GB + heap_size = os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + logger.info_once("AITER MegaMoE: MORI_SHMEM_HEAP_SIZE=%s", heap_size) mori.shmem.shmem_torch_process_group_init("mori") def _make_all2all_kwargs( diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index e0c69c03287c..e62431d7b62b 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -11,9 +11,8 @@ tilelang / triton / torch) instead of calling the tilelang kernels directly, and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER or TileLang fused MHC is available; False only on the torch fallback); - * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping``, - ``use_native_mega_moe`` and ``finalize_mega_moe_weights`` do not exist in - amd/model.py). + * use the AITER MegaMoE path (``aiter_mega_moe``) instead of the DeepGEMM + mega-MoE path. Everything else — the semi-autoregressive drafting hooks, the Markov head, the sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is @@ -50,6 +49,7 @@ ) from vllm.model_executor.models.utils import maybe_prefix +from .mega_moe import finalize_mega_moe_weights from .model import ( DeepseekV4DecoderLayer, ) @@ -385,14 +385,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: Non-mtp weights (embed/head/main layers) belong to the target model and are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ - # AMD DeepseekV4MoE has no mega-MoE path; always use the standard - # per-expert fused-MoE mapping (mirrors amd/mtp.py). expert_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + routed_experts_prefix=( + "" if self.model.layers[0].ffn.use_mega_moe else "routed_experts" + ), ) expert_scale_suffix = ( ".weight_scale" @@ -496,6 +497,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params + def process_weights_after_loading(self) -> None: + finalize_mega_moe_weights(self) + def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v4/amd/mega_moe.py b/vllm/models/deepseek_v4/amd/mega_moe.py new file mode 100644 index 000000000000..2fdff9335145 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/mega_moe.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AITER MegaMoE experts (fused dispatch, MXFP4 GEMMs, combine) for DeepSeek V4.""" + +from typing import Any + +import torch +import torch.nn as nn + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig +from vllm.distributed import get_ep_group +from vllm.model_executor.utils import set_weight_attrs +from vllm.utils.math_utils import next_power_of_2 + + +class DeepseekV4AiterMegaMoEExperts(nn.Module): + # MegaMoEV2 owns large symmetric-heap buffers, so every layer shares one + # instance and passes its own weights per call. + _runtime_cache: dict[tuple, Any] = {} + + def __init__( + self, + vllm_config: VllmConfig, + *, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + swiglu_limit: float | None, + prefix: str = "", + ): + super().__init__() + ep_group = get_ep_group() + self.ep_rank = ep_group.rank_in_group + self.ep_size = ep_group.world_size + self.num_experts = num_experts + self.num_local_experts = num_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.swiglu_limit = swiglu_limit or 0.0 + self.max_num_tokens = next_power_of_2( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.prefix = prefix + + def make_param(rows: int, cols: int) -> nn.Parameter: + param = nn.Parameter( + torch.zeros(self.num_local_experts, rows, cols, dtype=torch.uint8), + requires_grad=False, + ) + set_weight_attrs(param, {"weight_loader": self.weight_loader}) + return param + + self.w13_weight = make_param(2 * intermediate_size, hidden_size // 2) + self.w13_weight_scale = make_param(2 * intermediate_size, hidden_size // 32) + self.w2_weight = make_param(hidden_size, intermediate_size // 2) + self.w2_weight_scale = make_param(hidden_size, intermediate_size // 32) + self._runtime: Any = None + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool = False, + ) -> bool | None: + local_expert_id = expert_id - self.experts_start_idx + if not 0 <= local_expert_id < self.num_local_experts: + return False if return_success else None + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) + elif shard_id != "w2": + raise ValueError(f"Unsupported expert shard id: {shard_id}") + expert_data.copy_(loaded_weight.view(torch.uint8)) + return True if return_success else None + + def finalize_weights(self) -> None: + if self._runtime is not None: + return + n = self.num_local_experts + self.w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4( + self.w13_weight.data, 16, True + ) + self.w13_weight_scale.data = rocm_aiter_ops.shuffle_scale_a16w4( + self.w13_weight_scale.data.view(-1, self.hidden_size // 32), n, True + ) + self.w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4( + self.w2_weight.data, 16, False + ) + self.w2_weight_scale.data = rocm_aiter_ops.shuffle_scale_a16w4( + self.w2_weight_scale.data.view(-1, self.intermediate_size // 32), n, False + ) + key = ( + id(get_ep_group().cpu_group), + torch.accelerator.current_device_index(), + self.num_experts, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.max_num_tokens, + self.swiglu_limit, + ) + runtime = self._runtime_cache.get(key) + if runtime is None: + from aiter.ops.flydsl.kernels.mega_moe import MegaMoEV2 + + runtime = MegaMoEV2( + rank=self.ep_rank, + world_size=self.ep_size, + model_dim=self.hidden_size, + inter_dim=self.intermediate_size, + experts=self.num_experts, + topk=self.top_k, + quant="a8w4", + w1=self.w13_weight, + w1_scale=self.w13_weight_scale, + w2=self.w2_weight, + w2_scale=self.w2_weight_scale, + max_tok_per_rank=self.max_num_tokens, + swiglu_limit=self.swiglu_limit, + ) + self._runtime_cache[key] = runtime + self._runtime = runtime + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + runtime = self._runtime + assert runtime is not None + # MegaMoEV2 reads these per launch; rebind to this layer's experts. + runtime._s1_w1 = self.w13_weight + runtime._s1_w1_scale = self.w13_weight_scale + runtime.w2 = self.w2_weight + runtime.w2_scale = self.w2_weight_scale + return runtime.forward(hidden_states.contiguous(), topk_weights, topk_ids) + + +DeepseekV4AiterMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +def finalize_mega_moe_weights(model: nn.Module) -> None: + for module in model.modules(): + if isinstance(module, DeepseekV4AiterMegaMoEExperts): + module.finalize_weights() diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index f8306dd12c23..8b26fc0060e7 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -30,6 +30,9 @@ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( rocm_aiter_fused_experts, ) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -71,6 +74,7 @@ make_layers, maybe_prefix, ) +from vllm.models.deepseek_v4.amd.mega_moe import DeepseekV4AiterMegaMoEExperts from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 @@ -529,6 +533,39 @@ def _fuse_shared_experts_enabled(config, parallel_config: ParallelConfig) -> boo ) +def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None: + parallel_config = vllm_config.parallel_config + if not rocm_aiter_ops.is_fused_moe_enabled(): + raise ValueError( + "--moe-backend aiter_mega_moe requires VLLM_ROCM_USE_AITER=1 " + " and VLLM_ROCM_USE_AITER_MOE=1." + ) + if not on_gfx950(): + raise NotImplementedError("AITER MegaMoE requires gfx950.") + if ( + not parallel_config.enable_expert_parallel + or parallel_config.tensor_parallel_size != 1 + or parallel_config.pipeline_parallel_size != 1 + ): + raise NotImplementedError( + "AITER MegaMoE requires data parallel attention with expert " + "parallel MoE: --data-parallel-size N --enable-expert-parallel " + "with TP=1 and PP=1." + ) + if parallel_config.all2all_backend not in ( + "mori_high_throughput", + "mori_low_latency", + ): + raise ValueError( + "AITER MegaMoE runs on MoRI shmem; add " + "--all2all-backend mori_high_throughput." + ) + if parallel_config.enable_eplb: + raise NotImplementedError("AITER MegaMoE does not support EPLB.") + if getattr(vllm_config.model_config.hf_config, "expert_dtype", "fp4") != "fp4": + raise NotImplementedError("AITER MegaMoE only supports fp4 experts.") + + class DeepseekV4MoE(nn.Module): def __init__( self, @@ -643,6 +680,20 @@ def __init__( self.experts_start_idx = self.tp_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" + if self.use_mega_moe: + _validate_aiter_mega_moe_config(vllm_config) + self.experts = DeepseekV4AiterMegaMoEExperts( + vllm_config, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + swiglu_limit=self.swiglu_limit, + prefix=f"{prefix}.experts", + ) + return + fuse_shared_into_routed = ( self.is_fused_shared_expert_enabled or self.fuse_heterogeneous_shared_expert ) @@ -692,6 +743,8 @@ def forward( raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") org_shape = hidden_states.shape + if self.use_mega_moe: + return self._forward_mega_moe(hidden_states, input_ids).view(org_shape) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=hidden_states, @@ -700,6 +753,30 @@ def forward( return final_hidden_states.view(org_shape) + def _forward_mega_moe( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None + ) -> torch.Tensor: + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + topk=self.n_activated_experts, + renormalize=self.renormalize, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + bias_vl=self.gate.bias_vl, + image_sentinel_lo=self.image_sentinel_lo, + ) + final_hidden_states = self.experts(hidden_states, topk_weights, topk_ids) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts( + hidden_states + ) + return final_hidden_states + # Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. _AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168}) @@ -993,6 +1070,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.fuse_heterogeneous_shared_expert = _heterogeneous_shared_expert_enabled( vllm_config ) + self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: DeepseekV4DecoderLayer( @@ -1315,6 +1393,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=num_experts, + routed_experts_prefix="" if self.use_mega_moe else "routed_experts", ) @@ -1453,6 +1532,8 @@ def process_weights_after_loading(self) -> None: module.prepare_attn_preshuffle() elif isinstance(module, DeepseekV4MLP): module.prepare_gateup_preshuffle() + elif isinstance(module, DeepseekV4AiterMegaMoEExperts): + module.finalize_weights() if fused_compressor_layers: logger.info( "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers", diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 48ab036563f1..3746d6e38772 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -46,6 +46,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from .mega_moe import finalize_mega_moe_weights from .model import DeepseekV4DecoderLayer logger = init_logger(__name__) @@ -392,12 +393,16 @@ def _resolve_scale_name(name: str) -> str: head_rank_end = n_local_head * (tp_rank + 1) # Pre-compute expert mapping ONCE. + first_layer = next(iter(self.model.layers.values())) expert_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + routed_experts_prefix=( + "" if first_layer.mtp_block.ffn.use_mega_moe else "routed_experts" + ), ) # FP8 experts register ``..._weight_scale_inv`` (block_quant) while @@ -526,6 +531,9 @@ def _resolve_scale_name(name: str) -> str: logger.info_once("MTP draft model loaded: %d params", len(loaded_params)) return loaded_params + def process_weights_after_loading(self) -> None: + finalize_mega_moe_weights(self) + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: """Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer From 548597367e66dcc90e24576c08e0e25f88c9ee6c Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Wed, 7 Oct 2026 17:44:22 +0300 Subject: [PATCH 128/952] [ROCm][Perf] Add AITER FlyDSL GDN prefill backend (#57560) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- tests/v1/attention/test_gdn_aiter_flydsl.py | 199 ++++++++++++++++++ vllm/_aiter_ops.py | 113 ++++++++++ vllm/engine/arg_utils.py | 6 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 99 ++++++++- vllm/v1/attention/backends/gdn_attn.py | 57 ++++- 5 files changed, 462 insertions(+), 12 deletions(-) create mode 100644 tests/v1/attention/test_gdn_aiter_flydsl.py diff --git a/tests/v1/attention/test_gdn_aiter_flydsl.py b/tests/v1/attention/test_gdn_aiter_flydsl.py new file mode 100644 index 000000000000..3de625e9361d --- /dev/null +++ b/tests/v1/attention/test_gdn_aiter_flydsl.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + ChunkGatedDeltaRule, + _resolve_gdn_prefill_backend, +) +from vllm.platforms import current_platform + + +def _make_config( + *, + backend: str = "aiter_flydsl", + head_k_dim: int | None = 128, + head_v_dim: int | None = 128, + dtype: torch.dtype = torch.bfloat16, +): + return SimpleNamespace( + additional_config={"gdn_prefill_backend": backend}, + model_config=SimpleNamespace( + dtype=dtype, + hf_text_config=SimpleNamespace( + linear_key_head_dim=head_k_dim, + linear_value_head_dim=head_v_dim, + ), + ), + ) + + +def _resolve_on_rocm(config, *, kernels_available: bool = True): + with ( + patch.object( + qwen_gdn_linear_attn.current_platform, "is_rocm", return_value=True + ), + patch.object( + qwen_gdn_linear_attn.rocm_aiter_ops, + "is_gdn_flydsl_prefill_available", + return_value=kernels_available, + ), + ): + return _resolve_gdn_prefill_backend(config) + + +@pytest.mark.parametrize( + "head_k_dim,head_v_dim,dtype,expected", + [ + (128, 128, torch.bfloat16, "aiter_flydsl"), + (64, 128, torch.bfloat16, "triton"), + (128, 64, torch.bfloat16, "triton"), + (128, 128, torch.float16, "triton"), + ], +) +def test_resolve_aiter_flydsl_gdn_prefill_backend( + head_k_dim: int, + head_v_dim: int, + dtype: torch.dtype, + expected: str, +): + config = _make_config(head_k_dim=head_k_dim, head_v_dim=head_v_dim, dtype=dtype) + + requested, active = _resolve_on_rocm(config) + + assert requested == "aiter_flydsl" + assert active == expected + + +def test_kda_style_model_cannot_select_flydsl(): + """Kimi K3 KDA reports head dims through linear_attn_config, not these. + + Its builder overrides _build_chunk_metadata, which the AITER path skips, + so this pins the reason the two never meet: the resolver turns the model + down before the builder is ever constructed. + """ + config = _make_config(head_k_dim=None, head_v_dim=None) + + assert _resolve_on_rocm(config) == ("aiter_flydsl", "triton") + + +@pytest.mark.parametrize("on_rocm", [True, False], ids=["kernels_missing", "not_rocm"]) +def test_explicit_aiter_flydsl_fails_closed(on_rocm: bool): + """An explicit request that cannot be honoured is a configuration error. + + Falling back silently would report the run as healthy while the requested + kernels were never used. + """ + config = _make_config() + if on_rocm: + with pytest.raises(RuntimeError, match="not importable"): + _resolve_on_rocm(config, kernels_available=False) + return + with ( + patch.object( + qwen_gdn_linear_attn.current_platform, "is_rocm", return_value=False + ), + patch.object( + qwen_gdn_linear_attn.current_platform, "is_cuda", return_value=True + ), + pytest.raises(RuntimeError, match="aiter_flydsl"), + ): + _resolve_gdn_prefill_backend(config) + + +def _flydsl_prefill_unavailable() -> bool: + return ( + not current_platform.is_rocm() + or not rocm_aiter_ops.is_gdn_flydsl_prefill_available() + ) + + +@pytest.mark.skipif( + _flydsl_prefill_unavailable(), + reason="needs ROCm with the AITER FlyDSL GDN prefill kernels installed", +) +@pytest.mark.parametrize( + "state_dtype", [torch.float32, torch.bfloat16], ids=["fp32_state", "bf16_state"] +) +def test_aiter_flydsl_prefill_matches_triton_reference(state_dtype: torch.dtype): + """Run both prefill paths on device and compare. + + The reference is the same AITER entry point with the FlyDSL prepare and K5 + kernels switched off, so the only thing that differs between the two calls + is the kernels this backend exists to select. The state dtype follows + --mamba-ssm-cache-dtype, which may be float32 or bfloat16. + """ + from aiter.ops.triton.gated_delta_net import chunk_gated_delta_rule_opt_vk + + torch.manual_seed(0) + device = torch.device("cuda") + seq_lens = [70, 130] + total_tokens = sum(seq_lens) + num_k_heads, num_v_heads, head_dim = 2, 4, 128 + + # use_qk_l2norm_in_kernel=False matches the prefill path, where the model + # has already normalised q and k. Feeding raw normals instead lets the + # delta rule run away to NaN, which says nothing about either kernel. + q = torch.nn.functional.normalize( + torch.randn(1, total_tokens, num_k_heads, head_dim, device=device), dim=-1 + ).to(torch.bfloat16) + k = torch.nn.functional.normalize( + torch.randn(1, total_tokens, num_k_heads, head_dim, device=device), dim=-1 + ).to(torch.bfloat16) + v = torch.randn( + 1, total_tokens, num_v_heads, head_dim, dtype=torch.bfloat16, device=device + ) + # g is a log-domain decay, so keep it negative or the recurrence diverges + # and the comparison stops saying anything about the kernels. + g = -torch.rand(1, total_tokens, num_v_heads, dtype=torch.float32, device=device) + beta = torch.rand(1, total_tokens, num_v_heads, dtype=torch.float32, device=device) + cu_seqlens = torch.tensor([0, 70, 200], dtype=torch.int32, device=device) + initial_state = torch.randn( + len(seq_lens), + num_v_heads, + head_dim, + head_dim, + device=device, + ).to(state_dtype) + + shared = { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "output_final_state": True, + "cu_seqlens": cu_seqlens, + "use_qk_l2norm_in_kernel": False, + } + + reference_out, reference_state = chunk_gated_delta_rule_opt_vk( + initial_state=initial_state.clone(), + state_dtype=state_dtype, + use_chunk_flydsl=False, + use_prepare_flydsl=False, + **shared, + ) + flydsl_out, flydsl_state = ChunkGatedDeltaRule.forward_aiter_flydsl( + None, + initial_state=initial_state.clone(), + aiter_prefill_metadata=rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + seq_lens, cu_seqlens=cu_seqlens + ), + **shared, + ) + + assert flydsl_state.dtype == state_dtype + torch.testing.assert_close( + flydsl_out.float(), reference_out.float(), rtol=2e-2, atol=2e-2 + ) + torch.testing.assert_close( + flydsl_state.float(), reference_state.float(), rtol=2e-2, atol=2e-2 + ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 941f66e5b71f..44b89aeaf12c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1997,6 +1997,8 @@ def get_moe_dispatch_policy(cls) -> int: # Check if the env variable is set _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + # Chunk length the AITER FlyDSL K5 prefill kernels are compiled for. + GDN_FLYDSL_CHUNK_SIZE = 64 _CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE @@ -2467,6 +2469,23 @@ def _gdn_triton_kernels_importable() -> bool: except (ImportError, ModuleNotFoundError): return False + @staticmethod + def _gdn_flydsl_prefill_kernels_importable() -> bool: + try: + from aiter.ops.flydsl.linear_attention_prefill_kernels import ( # noqa: F401 + chunk_gated_delta_rule_fwd_h_flydsl_opt, + gdn_prepare_flydsl_supported, + gdn_prepare_fwd_flydsl, + ) + from aiter.ops.triton.gated_delta_net import ( # noqa: F401 + build_gated_delta_rule_prefill_metadata, + chunk_gated_delta_rule_opt_vk, + ) + + return True + except (ImportError, ModuleNotFoundError): + return False + @classmethod @if_aiter_supported def are_gdn_triton_kernels_available(cls) -> bool: @@ -2478,6 +2497,100 @@ def are_gdn_triton_kernels_available(cls) -> bool: """ return cls._AITER_ENABLED and cls._gdn_triton_kernels_importable() + @classmethod + @functools.cache + def is_gdn_flydsl_prefill_available(cls) -> bool: + """Whether the opt-in AITER FlyDSL GDN prefill path can be used. + + Selecting the backend is an explicit opt-in in itself, but it still + runs AITER kernels, so VLLM_ROCM_USE_AITER remains the one switch that + turns all of them off. + """ + return ( + cls._AITER_ENABLED + and is_aiter_found_and_supported() + and cls._gdn_flydsl_prefill_kernels_importable() + ) + + @classmethod + def build_gdn_flydsl_prefill_metadata( + cls, + seq_lens_cpu: list[int], + cu_seqlens: torch.Tensor, + ) -> object: + """Build the reusable varlen metadata FlyDSL GDN prefill runs against. + + The K5 recurrence is compiled for a fixed 64-token chunk, so the chunk + size is a property of the AITER kernel rather than of FLA, and is not + the caller's to choose. + """ + from aiter.ops.triton.gated_delta_net import ( + build_gated_delta_rule_prefill_metadata, + ) + + return build_gated_delta_rule_prefill_metadata( + seq_lens_cpu, + cu_seqlens=cu_seqlens, + chunk_size=cls.GDN_FLYDSL_CHUNK_SIZE, + ) + + @classmethod + def gdn_flydsl_prefill( + cls, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + aiter_prefill_metadata: object | None, + use_qk_l2norm_in_kernel: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run GDN prefill through AITER's VK path with the FlyDSL kernels. + + Every argument AITER takes is named here rather than forwarded as + ``**kwargs``, so a rename or reordering on the AITER side fails at this + call instead of silently binding to the wrong parameter. + """ + from aiter.ops.flydsl.linear_attention_prefill_kernels import ( + gdn_prepare_flydsl_supported, + ) + from aiter.ops.triton.gated_delta_net import chunk_gated_delta_rule_opt_vk + + if cu_seqlens is not None and aiter_prefill_metadata is None: + raise RuntimeError( + "AITER FlyDSL GDN prefill requires reusable varlen prefill metadata." + ) + if not gdn_prepare_flydsl_supported(k, v): + raise RuntimeError( + "AITER FlyDSL GDN prepare does not support the runtime input shape, " + "dtype, or device." + ) + logger.info_once( + "Dispatching AITER FlyDSL GDN prefill (prepare=flydsl, chunk=flydsl)." + ) + return chunk_gated_delta_rule_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + use_chunk_flydsl=True, + use_prepare_flydsl=True, + state_dtype=initial_state.dtype, + prefill_metadata=aiter_prefill_metadata, + # The caller stages the state densely and writes it back, so the + # kernel gets a plain [N, H, V, K] buffer and never an index into + # the pool; there is correspondingly nothing to update in place. + inplace_final_state=False, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + @classmethod def is_rdna_gdn_triton_kernels_available(cls) -> bool: """RDNA4 (gfx12) analog of are_gdn_triton_kernels_available().""" diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 527b8956d033..800d1b2f2ccf 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -839,7 +839,9 @@ class EngineArgs: ) fail_on_environ_validation: bool = False - gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None + gdn_prefill_backend: ( + Literal["flashinfer", "triton", "cutedsl", "aiter_flydsl"] | None + ) = None kda_prefill_backend: ( Literal["auto", "triton", "flashkda", "flashinfer", "fused"] | None ) = None @@ -1890,7 +1892,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: parser.add_argument( "--gdn-prefill-backend", dest="gdn_prefill_backend", - choices=["flashinfer", "triton", "cutedsl"], + choices=["flashinfer", "triton", "cutedsl", "aiter_flydsl"], default=None, help="Select GDN prefill backend.", ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 7b164f1dd3df..1e7883864027 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -92,7 +92,7 @@ def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, -) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: +) -> tuple[str, Literal["triton", "flashinfer", "cutedsl", "aiter_flydsl"]]: """Resolve GDN prefill backend. FlashInfer's GDN prefill kernel is chosen when: @@ -106,6 +106,11 @@ def _resolve_gdn_prefill_backend( In-tree CuteDSL GDN prefill kernel is chosen when: * "cutedsl" is requested; (opt-in only) * Blackwell (SM10.x) with ``head_k_dim == 128``; + + AITER FlyDSL GDN prefill kernels are chosen when: + * "aiter_flydsl" is requested; (opt-in only) + * ROCm AITER exposes the optimized VK prefill API and FlyDSL kernels; + * the model uses BF16 activations with ``head_k_dim == head_v_dim == 128``. """ additional_config = vllm_config.additional_config backend_cfg = ( @@ -115,12 +120,50 @@ def _resolve_gdn_prefill_backend( ) backend = str(backend_cfg).strip().lower() - if not current_platform.is_cuda(): - return backend, "triton" - head_k_dim = getattr( vllm_config.model_config.hf_text_config, "linear_key_head_dim", None ) + head_v_dim = getattr( + vllm_config.model_config.hf_text_config, "linear_value_head_dim", None + ) + + if current_platform.is_rocm(): + if backend == "aiter_flydsl": + if not rocm_aiter_ops.is_gdn_flydsl_prefill_available(): + raise RuntimeError( + "GDN prefill backend 'aiter_flydsl' was requested but the " + "AITER FlyDSL prefill kernels are not importable. Check " + "that VLLM_ROCM_USE_AITER=1, that this is a CDNA 3 or " + "newer GPU, and that the installed AITER exports them." + ) + if head_k_dim != 128 or head_v_dim != 128: + logger.warning_once( + "GDN prefill backend 'aiter_flydsl' was requested but " + "linear head dims are K=%s V=%s; FlyDSL requires 128/128. " + "Falling back to Triton/FLA.", + head_k_dim, + head_v_dim, + ) + return backend, "triton" + if vllm_config.model_config.dtype != torch.bfloat16: + logger.warning_once( + "GDN prefill backend 'aiter_flydsl' was requested but " + "model dtype is %s; FlyDSL requires bfloat16. " + "Falling back to Triton/FLA.", + vllm_config.model_config.dtype, + ) + return backend, "triton" + return backend, "aiter_flydsl" + return backend, "triton" + + if backend == "aiter_flydsl": + raise RuntimeError( + "GDN prefill backend 'aiter_flydsl' was requested but it is a ROCm " + f"AITER backend and the current platform is {current_platform.device_name}." + ) + + if not current_platform.is_cuda(): + return backend, "triton" supports_flashinfer = False supports_cutedsl = False @@ -172,6 +215,7 @@ def _log_gdn_backend_decision( chosen = { "flashinfer": "FlashInfer", "cutedsl": "CuteDSL", + "aiter_flydsl": "AITER FlyDSL", "triton": "Triton/FLA", }[active_backend] logger.info_once( @@ -259,6 +303,8 @@ def __init__(self) -> None: self._forward_method = self.forward_cuda elif active_backend == "cutedsl": self._forward_method = self.forward_cutedsl + elif active_backend == "aiter_flydsl": + self._forward_method = self.forward_aiter_flydsl else: self._forward_method = self.forward_native @@ -276,6 +322,7 @@ def forward_cuda( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): o, final_state = fi_chunk_gated_delta_rule( q=q, @@ -308,6 +355,7 @@ def forward_native( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): return fla_chunk_gated_delta_rule( q=q, @@ -338,6 +386,7 @@ def forward_cutedsl( chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, ): from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( chunk_gated_delta_rule_cutedsl, @@ -367,6 +416,39 @@ def forward_cutedsl( final_state = None return o, final_state + def forward_aiter_flydsl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + core_attn_out: torch.Tensor | None = None, + aiter_prefill_metadata: object | None = None, + ): + o, final_state = rocm_aiter_ops.gdn_flydsl_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + aiter_prefill_metadata=aiter_prefill_metadata, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + if core_attn_out is not None: + o_flat = o.reshape(-1) + core_attn_out.reshape(-1)[: o_flat.numel()].copy_(o_flat) + return o, final_state + @PluggableLayer.register("qwen_gated_delta_net_attention") class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): @@ -1132,12 +1214,18 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None # CuteDSL kernels require metadata chunk_indices = None chunk_offsets = None + aiter_prefill_metadata = None if self.gdn_prefill_backend == "cutedsl": from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( prepare_metadata_cutedsl, ) chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, T) + elif self.gdn_prefill_backend == "aiter_flydsl": + aiter_prefill_metadata = rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + [T], + cu_seqlens=cu_seqlens, + ) try: self.chunk_gated_delta_rule( @@ -1152,6 +1240,7 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None chunk_indices=chunk_indices, chunk_offsets=chunk_offsets, use_qk_l2norm_in_kernel=False, + aiter_prefill_metadata=aiter_prefill_metadata, ) except Exception: logger.warning( @@ -1182,6 +1271,7 @@ def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None cu_seqlens, chunk_indices, chunk_offsets, + aiter_prefill_metadata, ) torch.accelerator.empty_cache() @@ -1523,6 +1613,7 @@ def _forward_core( chunk_indices=attn_metadata.chunk_indices, chunk_offsets=attn_metadata.chunk_offsets, use_qk_l2norm_in_kernel=False, + aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata, ) # Init cache ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 95f6a808a597..08369adcb1a3 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -7,6 +7,7 @@ import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointBuilder, @@ -79,6 +80,7 @@ class GDNAttentionMetadata: prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + aiter_prefill_metadata: object | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -109,8 +111,11 @@ def __init__( _resolve_gdn_prefill_backend, ) - self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"] + self.gdn_prefill_backend: Literal[ + "triton", "flashinfer", "cutedsl", "aiter_flydsl" + ] _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config) + self._check_chunk_metadata_override(type(self), self.gdn_prefill_backend) if self.speculative_config: assert self.speculative_config.num_speculative_tokens is not None @@ -186,6 +191,32 @@ def __init__( device=device, ) + @staticmethod + def _check_chunk_metadata_override( + builder_cls: type["GDNAttentionMetadataBuilder"], backend: str + ) -> None: + """Reject a subclass whose chunk metadata the AITER path would skip. + + AITER brings its own varlen prefill metadata and never calls + ``_build_chunk_metadata``, so a builder that overrides it would lose + that override without a word. No in-tree subclass can get here -- + ``_resolve_gdn_prefill_backend`` only selects this backend for models + with 128-dim GDN key and value heads, which the KDA builders are not -- + but a future one should be told rather than quietly ignored. + """ + if backend != "aiter_flydsl": + return + if ( + builder_cls._build_chunk_metadata + is GDNAttentionMetadataBuilder._build_chunk_metadata + ): + return + raise RuntimeError( + f"{builder_cls.__name__} builds its own FLA chunk metadata, which " + "the 'aiter_flydsl' GDN prefill backend does not use. Select a " + "different gdn_prefill_backend for this model." + ) + def _build_chunk_metadata( self, prefill_query_start_loc: torch.Tensor, @@ -431,6 +462,7 @@ def build( # type: ignore[override] prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + aiter_prefill_metadata: object | None = None if num_prefills > 0: # In a mixed non-spec batch, decodes are peeled off to the recurrent # kernel (decode-first front slice), so build chunk metadata from the @@ -452,11 +484,23 @@ def build( # type: ignore[override] prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu prefill_state_indices = non_spec_state_indices_tensor - chunk_indices, chunk_offsets = self._build_chunk_metadata( - prefill_query_start_loc, - prefill_query_start_loc_cpu, - query_start_loc.device, - ) + if self.gdn_prefill_backend == "aiter_flydsl": + # AITER carries its own reusable varlen metadata and has no use + # for FLA's chunk indices, so it replaces them rather than + # extending what _build_chunk_metadata returns. + assert prefill_query_start_loc_cpu is not None + aiter_prefill_metadata = ( + rocm_aiter_ops.build_gdn_flydsl_prefill_metadata( + torch.diff(prefill_query_start_loc_cpu).tolist(), + cu_seqlens=prefill_query_start_loc, + ) + ) + else: + chunk_indices, chunk_offsets = self._build_chunk_metadata( + prefill_query_start_loc, + prefill_query_start_loc_cpu, + query_start_loc.device, + ) if num_prefills > 0: context_lens_tensor = m.compute_num_computed_tokens() @@ -568,6 +612,7 @@ def build( # type: ignore[override] prefill_query_start_loc=prefill_query_start_loc, prefill_state_indices=prefill_state_indices, prefill_has_initial_state=prefill_has_initial_state, + aiter_prefill_metadata=aiter_prefill_metadata, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, From 14e3902d5136e2fb00ef86c05fdc8d098f736283 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:46:44 -0700 Subject: [PATCH 129/952] [Frontend][Rust] Add gRPC forbidden token sequences and cache usage (#59837) Signed-off-by: Alec Flowers --- rust/proto/inference.proto | 4 + rust/src/server/src/grpc/convert.rs | 46 +++++++++++- rust/src/server/src/grpc/tests.rs | 2 + .../routes/openai/chat_completions/convert.rs | 1 + .../src/routes/openai/completions/convert.rs | 1 + rust/src/text/src/lower.rs | 75 ++++++++++++++++--- rust/src/text/src/lower/token_ids.rs | 5 ++ rust/src/text/src/request.rs | 3 + 8 files changed, 125 insertions(+), 12 deletions(-) diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 2bc35bcb64a2..8189758212d4 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -102,6 +102,9 @@ message DecodingParameters { bool json_object = 10; string structural_tag = 11; } + + // Token sequences that must not be generated. Each sequence must be nonempty. + repeated TokenIds bad_words_token_ids = 12; } message StoppingCriteria { @@ -210,6 +213,7 @@ message FinishInfo { google.protobuf.Struct kv_transfer_params = 6; //uint64 seed = 7; google.protobuf.Struct ec_transfer_params = 8; + optional uint32 num_cached_tokens = 9; } // Info for candidate tokens other than the input/sampled diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index d5dd1ae64df4..8f2c589f154b 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -203,6 +203,10 @@ fn build_sampling_params( if !d.allowed_token_ids.is_empty() { params.allowed_token_ids = Some(d.allowed_token_ids.clone()); } + if !d.bad_words_token_ids.is_empty() { + params.bad_words_token_ids = + Some(d.bad_words_token_ids.iter().map(|sequence| sequence.ids.clone()).collect()); + } params.structured_outputs = convert_structured_output(d)?; } @@ -403,6 +407,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result>, /// Words to avoid during generation (tokenized to IDs during lowering). pub bad_words: Option>, + /// Pre-tokenized sequences to prohibit, in addition to `bad_words`. + pub bad_words_token_ids: Option>>, /// Specific token IDs for which log probabilities should be returned at /// each position. /// @@ -173,6 +175,7 @@ impl Default for SamplingParams { logit_bias: None, allowed_token_ids: None, bad_words: None, + bad_words_token_ids: None, logprob_token_ids: None, structured_outputs: None, skip_reading_prefix_cache: None, From 71f9a749203c31a263ccad03395623b1d6fc5538 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Wed, 7 Oct 2026 16:57:52 +0200 Subject: [PATCH 130/952] [CI][ROCm] Re-enable the DSv4.1 decoder replay graph test on ROCm (#60392) Signed-off-by: Stefan Koncarevic --- tests/models/test_deepseek_v41_decoder_replay_layers.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 1302e0acda30..f3b124c5ac80 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -21,7 +21,6 @@ from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import ( DecoderReplayCudaGraphManager, ) -from vllm.platforms import current_platform from vllm.v1.worker.gpu import cudagraph_utils pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -126,9 +125,6 @@ def _graph_context(num_tokens): ) -@pytest.mark.skipif( - current_platform.is_rocm(), reason="decoder replay CUDA graphs are NVIDIA-only" -) def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) @@ -136,6 +132,8 @@ def test_replay_graph_matches_eager(monkeypatch): monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4]) cfg = MagicMock(compilation_config=compilation, speculative_config=None) + # Mocking it true would capture into the cuMem pool, which breaks on ROCm. + cfg.use_cumem_cudagraph_pool = False cfg.scheduler_config = MagicMock(max_num_seqs=2, max_num_batched_tokens=64) cfg.cache_config.use_kda_recoverssm = False cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) From 9367d8b9693c3d2fb6c3f99d8ccb4aa7673c0981 Mon Sep 17 00:00:00 2001 From: xaguilar-amd Date: Wed, 7 Oct 2026 17:22:32 +0200 Subject: [PATCH 131/952] [Bugfix][MLA] Scope DSpark non-causal capability to builder layers (#57992) Signed-off-by: Xavier Aguilar Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_noncausal.py | 121 ++++++++++++++++++ .../layers/attention/mla_attention.py | 8 +- 2 files changed, 127 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py index c06ded5c1392..57d27ecd464f 100644 --- a/tests/v1/attention/test_mla_noncausal.py +++ b/tests/v1/attention/test_mla_noncausal.py @@ -19,6 +19,21 @@ class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): supports_non_causal_multi_token_decode = True +class _CausalOnlyMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + pass + + +class _DraftMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + """Triton-like: non-causal decode is supported, DCP draft decode is not.""" + + supports_non_causal_multi_token_decode = True + + +class _DcpCapableMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + supports_non_causal_multi_token_decode = True + supports_non_causal_multi_token_dcp = True + + def _metadata( query_start_loc: list[int], num_tokens: int | None = None, @@ -60,6 +75,54 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: return builder +def _mla_layer(*, non_causal: bool = False) -> SimpleNamespace: + return SimpleNamespace( + non_causal_multi_token_decode=non_causal, + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + prefill_backend=SimpleNamespace(clone=lambda: None), + ) + + +def _merged_mla_spec() -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + non_causal_multi_token_decode=True, + ) + + +def _dspark_dcp_vllm_config( + static_forward_context: dict[str, SimpleNamespace], +) -> SimpleNamespace: + return SimpleNamespace( + speculative_config=SimpleNamespace( + method="dspark", num_speculative_tokens=None + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=2, + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + compilation_config=SimpleNamespace( + static_forward_context=static_forward_context + ), + model_config=SimpleNamespace( + dtype=torch.bfloat16, + max_model_len=128, + get_num_attention_heads=lambda _parallel_config: 16, + ), + cache_config=SimpleNamespace(cache_dtype="auto", block_size=64), + scheduler_config=SimpleNamespace(max_num_seqs=1), + attention_config=SimpleNamespace(use_prefill_query_quantization=False), + ) + + def test_group_capability_keeps_runtime_causality_per_step(): builder = _builder() @@ -137,3 +200,61 @@ def test_mla_cache_marker_is_promoted_to_group_capability(): [unmarked, unmarked] ).non_causal_multi_token_decode assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode + + +def test_builder_scopes_noncausal_capability_to_its_layers(): + merged_spec = _merged_mla_spec() + static_forward_context = { + "target": _mla_layer(non_causal=False), + "draft": _mla_layer(non_causal=True), + # Indexer and compressor caches share an MLA group but predate the flag. + "indexer": SimpleNamespace(), + } + vllm_config = _dspark_dcp_vllm_config(static_forward_context) + device = torch.device("cpu") + + def _flag(layer_names: list[str]) -> bool: + return _DcpCapableMLAMetadataBuilder( + merged_spec, + layer_names, + vllm_config, + device, + supports_dcp_with_varlen=True, + ).non_causal_multi_token_decode + + assert not _flag(["target"]) + assert _flag(["draft"]) + assert _flag(["target", "draft"]) + assert not _flag(["target", "indexer"]) + + +def test_merged_group_spec_does_not_mark_a_target_only_builder(): + # Both builders receive the same merged spec; only the draft layer is + # non-causal. Constructing through __init__ would still fail DCP + # validation if the builder read the group flag. + merged_spec = _merged_mla_spec() + static_forward_context = { + "target": _mla_layer(non_causal=False), + "draft": _mla_layer(non_causal=True), + } + vllm_config = _dspark_dcp_vllm_config(static_forward_context) + device = torch.device("cpu") + assert merged_spec.non_causal_multi_token_decode + + target = _CausalOnlyMLAMetadataBuilder( + merged_spec, + ["target"], + vllm_config, + device, + supports_dcp_with_varlen=True, + ) + assert not target.non_causal_multi_token_decode + + with pytest.raises(ValueError, match="non-causal draft"): + _DraftMLAMetadataBuilder( + merged_spec, + ["draft"], + vllm_config, + device, + supports_dcp_with_varlen=True, + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 00e70d6fd0ae..f550dc3d90d4 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2384,8 +2384,12 @@ def __init__( self.vllm_config = vllm_config self.device = device self.use_pcp = parallel_config.prefill_context_parallel_size > 1 - self.non_causal_multi_token_decode = getattr( - kv_cache_spec, "non_causal_multi_token_decode", False + # merge() unions this flag across the KV group, so a shared target/draft + # group would mark the target builder as non-causal. + ctx = self.compilation_config.static_forward_context + self.non_causal_multi_token_decode = any( + getattr(ctx[name], "non_causal_multi_token_decode", False) + for name in layer_names ) self._validate_dspark_dcp_support(supports_dcp_with_varlen) From f7f62d19160010bdc84b31af407bb8b88e5392d6 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Wed, 7 Oct 2026 18:02:13 +0200 Subject: [PATCH 132/952] [XPU][MoE] Tune Triton fused MoE for Intel XPU (#53065) --- benchmarks/kernels/benchmark_moe.py | 47 ++++- ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 56 ++++++ 4 files changed, 427 insertions(+), 4 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 1fcc7ddde7f8..f746e9aab339 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -303,9 +303,13 @@ def run(): run() torch.accelerator.synchronize() - # Capture 10 invocations with CUDA graph - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + # Capture 10 invocations via graph(); CUDA needs a side stream + graph = ( + torch.cuda.CUDAGraph() + if current_platform.is_cuda_alike() + else torch.xpu.XPUGraph() + ) + with current_platform.graph(graph): for _ in range(10): run() torch.accelerator.synchronize() @@ -361,11 +365,33 @@ def get_rocm_tuning_space(use_fp16): return param_ranges +def get_xpu_tuning_space(use_fp16): + # BLOCK_SIZE_N=32 and num_warps=32 never won; grf_mode=256 beat 128 on B70. + block_m_range = [16, 32, 64, 128, 256] + block_n_range = [64, 128] + block_k_range = [32, 64] + num_warps_range = [4, 8, 16] + group_m_range = [1, 8, 16] + num_stage_range = [3, 4] + + return { + "BLOCK_SIZE_M": block_m_range, + "BLOCK_SIZE_N": block_n_range, + "BLOCK_SIZE_K": block_k_range, + "GROUP_SIZE_M": group_m_range, + "num_warps": num_warps_range, + "num_stages": num_stage_range, + "grf_mode": ["256"], + } + + def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int]]: configs: list[BenchmarkConfig] = [] if current_platform.is_rocm(): param_ranges = get_rocm_tuning_space(use_fp16) + elif current_platform.is_xpu(): + param_ranges = get_xpu_tuning_space(use_fp16) else: # Reduced search space for faster tuning. # TODO(woosuk): Increase the search space and use a performance model to @@ -408,6 +434,14 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int return configs +def prune_xpu_search_space(num_tokens, search_space, topk, num_experts): + # Cap the tile by rows routed to one expert, not total rows: wider tiles + # just pad and are the slowest to compile. + rows_per_expert = (num_tokens * topk + num_experts - 1) // num_experts + max_block_m = max(16, triton.next_power_of_2(rows_per_expert)) + return [c for c in search_space if c["BLOCK_SIZE_M"] <= max_block_m] + + def prune_rocm_search_space( num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk ): @@ -522,7 +556,7 @@ def merge_unique_dicts(list1, list2): @ray.remote(num_gpus=1) class BenchmarkWorker: def __init__(self, seed: int) -> None: - torch.set_default_device("cuda") + torch.set_default_device(current_platform.device_type) set_random_seed(seed) self.seed = seed # Get the device ID to allocate tensors and kernels @@ -619,6 +653,10 @@ def tune( is_fp16, topk, ) + elif current_platform.is_xpu(): + search_space = prune_xpu_search_space( + num_tokens, search_space, topk, num_experts + ) need_device_guard = False if current_platform.is_rocm(): @@ -690,6 +728,7 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: else {} ), **({"kpack": config["kpack"]} if "kpack" in config else {}), + **({"grf_mode": config["grf_mode"]} if "grf_mode" in config else {}), **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}), } diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..6c1788d496d6 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..0cdfc2266482 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 000000000000..46bee68b897d --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,56 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +} From c741bfca70cfb777e2016f827eae31f6e215fe9f Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Wed, 7 Oct 2026 17:05:29 +0100 Subject: [PATCH 133/952] [Docs] Add review areas for andylolu2 (#60416) Signed-off-by: Andy Lo --- docs/community/reviewers.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/community/reviewers.md b/docs/community/reviewers.md index 60d37f0d3b4d..57840ffd8d8e 100644 --- a/docs/community/reviewers.md +++ b/docs/community/reviewers.md @@ -9,7 +9,7 @@ This page lists the active reviewers and describes the role. See the [committers Community members can reach out to the reviewers below for help with PRs in their areas, for example in the `#pr-reviews` Slack channel. Reviewers will add their areas of expertise to this list. Sorted alphabetically by GitHub handle: -- [@andylolu2](https://github.com/andylolu2) +- [@andylolu2](https://github.com/andylolu2): Mistral models, speculative decoding, quantization, scheduler, MLA - [@bnellnm](https://github.com/bnellnm) MoE, Quantization layers - [@BowenBao](https://github.com/BowenBao) - [@cjackal](https://github.com/cjackal): Multimodality; KV Connector and offload From 240ee8de9eb98d4e9269b82054df121adf522c17 Mon Sep 17 00:00:00 2001 From: MohitAMD Date: Wed, 7 Oct 2026 09:33:30 -0700 Subject: [PATCH 134/952] [Bugfix][KV Connector][ROCm] Fix MoRI-IO KV block-offset for speculative decoding (MTP) (#55053) Signed-off-by: Mohit Deopujari Signed-off-by: MohitAMD Co-authored-by: Cursor --- .../unit/test_moriio_connector.py | 348 +++++++++++++++++- .../unit/test_moriio_kv_layout.py | 24 +- .../v1/moriio/moriio_connector.py | 123 +++++-- .../kv_connector/v1/moriio/moriio_layout.py | 14 +- 4 files changed, 452 insertions(+), 57 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index f95e12c601bb..cb33bfbcf03c 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -3,6 +3,8 @@ import importlib.util import socket import uuid +from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch import msgspec @@ -20,11 +22,13 @@ set_current_vllm_config, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ROLE, MoRIIOAgentMetadata, MoRIIOConnectorMetadata, MoRIIOConstants, MoRIIOMode, resolve_host_ip, + set_role, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( @@ -147,6 +151,50 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2): return scheduler +def _write_producer_scheduler(block_size: int = 1) -> Any: + """Bare WRITE-mode PRODUCER scheduler exercising the save path without a + real engine/RDMA stack. Uses a single full-attention KV group (index 0).""" + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.is_producer = True + scheduler.block_size = block_size + scheduler.blocks_per_sw = [0] + scheduler._group_block_sizes = [block_size] + scheduler.transfer_id_to_request_id = {} + scheduler._reqs_need_recv = {} + scheduler._reqs_need_save = {} + scheduler._reqs_need_pending_save = {} + scheduler._req_kv_params = {} + scheduler._reqs_need_send = {} + return scheduler + + +def _spec_kv_params(transfer_id: str = "xfer-spec") -> dict[str, Any]: + # request_id embeds no zmq address, so add_new_req resolves the peer here. + return { + "transfer_id": transfer_id, + "remote_engine_id": "remote-engine", + "remote_block_ids": [], + "remote_host": "127.0.0.1", + "remote_handshake_port": 5000, + "remote_notify_port": 5001, + } + + +def _build_meta( + scheduler: Any, req_ids=None, new_block_ids=None, num_scheduled_tokens=None +) -> Any: + # Final-chunk detection reads num_scheduled_tokens, so expose it. + scheduler_output = SimpleNamespace( + scheduled_cached_reqs=SimpleNamespace( + req_ids=req_ids or [], + new_block_ids=new_block_ids or [], + ), + num_scheduled_tokens=num_scheduled_tokens or {}, + ) + return scheduler.build_connector_meta(scheduler_output) + + class FakeMoRIIOWrapper: # A fake MoRIIOWrapper for testing purposes def __init__(self, *args, **kwargs): @@ -700,6 +748,203 @@ def test_resolve_host_ip_prefers_extra_config(): assert resolve_host_ip({"host_ip": ""}) == fallback +def test_write_mode_excludes_spec_lookahead_blocks(): + # The save path must record only the prompt blocks, clamping off the + # trailing lookahead blocks decode never allocates. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-spec" + prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17] + lookahead_blocks = [18, 19, 20] + local_block_ids = [prompt_blocks + lookahead_blocks] + + # block_size=1 => num_prompt_tokens == prompt block count, all scheduled in + # this single chunk, so final-chunk detection fires. + num_prompt_tokens = len(prompt_blocks) + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._reqs_need_save[req_id] = (req, local_block_ids) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens}) + + assert isinstance(meta, MoRIIOConnectorMetadata) + assert req_id in meta.reqs_to_save + saved = meta.reqs_to_save[req_id] + assert saved.local_block_ids == [prompt_blocks] + + +def test_write_mode_chunked_prefill_clamps_spec_lookahead_blocks(): + # Non-final chunks are buffered untouched; lookahead blocks are only clamped + # on the final-chunk save. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-spec-chunked" + num_prompt_tokens = 8 + first_chunk = [10, 11, 12, 13] + final_chunk = [14, 15, 16, 17, 18, 19, 20] # 4 prompt + 3 lookahead + prompt_blocks = [10, 11, 12, 13, 14, 15, 16, 17] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + + # Step 1: first chunk (4 of 8 tokens) is buffered, not saved. + scheduler._reqs_need_save[req_id] = (req, [first_chunk]) + scheduler._req_kv_params[req_id] = _spec_kv_params() + meta_step1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4}) + assert req_id not in meta_step1.reqs_to_save + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [first_chunk] + + # Step 2: final chunk (computed 4 + scheduled 4 >= 8) assembles the full + # local set, clamps the lookahead, and records only the prompt blocks. + req.num_computed_tokens = 4 + meta_step2 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[[final_chunk]], + num_scheduled_tokens={req_id: 4}, + ) + assert req_id in meta_step2.reqs_to_save + assert meta_step2.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + assert req_id not in scheduler._reqs_need_pending_save + + +@pytest.mark.parametrize( + ("block_size", "num_prompt_tokens", "num_prompt_blocks"), + [ + # block_size=4: exact multiple (2 full blocks) and non-multiple (3 blocks, + # last one partial). + pytest.param(4, 8, 2, id="bs4-exact-multiple"), + pytest.param(4, 9, 3, id="bs4-non-multiple"), + # block_size=16: exact multiple (2 full blocks) and non-multiple (3 + # blocks, last one partial). + pytest.param(16, 32, 2, id="bs16-exact-multiple"), + pytest.param(16, 40, 3, id="bs16-non-multiple"), + ], +) +def test_write_mode_clamps_prompt_blocks_with_block_size_gt_1( + block_size, num_prompt_tokens, num_prompt_blocks +): + # With block_size > 1 the clamp must keep exactly + # ceil(num_prompt_tokens / block_size) prompt blocks and drop the tail. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=block_size) + + req_id = f"req-bs{block_size}-{num_prompt_tokens}" + prompt_blocks = list(range(10, 10 + num_prompt_blocks)) + lookahead_blocks = [10 + num_prompt_blocks] + local_block_ids = [prompt_blocks + lookahead_blocks] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._reqs_need_save[req_id] = (req, local_block_ids) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + meta = _build_meta(scheduler, num_scheduled_tokens={req_id: num_prompt_tokens}) + + assert isinstance(meta, MoRIIOConnectorMetadata) + assert req_id in meta.reqs_to_save + assert len(prompt_blocks) == num_prompt_blocks + assert meta.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + + +def test_write_mode_defers_save_until_prompt_complete_with_lookahead(): + # Under speculative decoding the producer reserves lookahead blocks beyond + # the last-scheduled prompt token, inflating the accumulated block count on + # a non-final chunk. Token-progress detection must keep deferring until the + # prompt is complete, where the lookahead is clamped off. + # + # block_size=1, num_prompt_tokens=10, chunk=4, num_lookahead=3: after 8 + # prompt tokens the producer holds 8 + 3 = 11 blocks, but token progress + # 8 < 10 must keep waiting until the final chunk completes tokens 9-10. + set_role(ROLE.PRODUCER) + scheduler = _write_producer_scheduler(block_size=1) + + req_id = "req-lookahead-final-chunk" + num_prompt_tokens = 10 + + chunk1_blocks = [100, 101, 102, 103] # prompt tokens 1-4 + chunk2_blocks = [104, 105, 106, 107, 108, 109, 110] # 4 prompt + 3 lookahead + prompt_blocks = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] + + req = SimpleNamespace( + request_id=req_id, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=0, + kv_transfer_params=_spec_kv_params(), + ) + scheduler._req_kv_params[req_id] = _spec_kv_params() + + # Chunk 1 (tokens 1-4): computed 0 + scheduled 4 < 10 -> buffered. + scheduler._reqs_need_save[req_id] = (req, [chunk1_blocks]) + meta1 = _build_meta(scheduler, num_scheduled_tokens={req_id: 4}) + assert req_id not in meta1.reqs_to_save, "chunk 1 must not be saved" + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [chunk1_blocks] + + # Chunk 2 (tokens 5-8) inflates the buffered count to 11 via lookahead, but + # token progress 4 + 4 = 8 < 10 keeps the request deferred. + req.num_computed_tokens = 4 + meta2 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[[chunk2_blocks]], + num_scheduled_tokens={req_id: 4}, + ) + assert req_id not in meta2.reqs_to_save, ( + "second-to-last chunk saved prematurely despite incomplete prompt" + ) + assert req_id in scheduler._reqs_need_pending_save + assert scheduler._reqs_need_pending_save[req_id][1] == [ + chunk1_blocks + chunk2_blocks + ] + + # Final chunk (tokens 9-10, no new blocks): computed 8 + scheduled 2 >= 10 + # emits the save with the trailing lookahead block clamped off. + req.num_computed_tokens = 8 + meta3 = _build_meta( + scheduler, + req_ids=[req_id], + new_block_ids=[None], + num_scheduled_tokens={req_id: 2}, + ) + assert req_id in meta3.reqs_to_save, "final chunk must emit the WRITE save" + assert meta3.reqs_to_save[req_id].local_block_ids == [prompt_blocks] + assert req_id not in scheduler._reqs_need_pending_save + + +def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size(): + # Two full-attention groups with different block sizes are each clamped to + # their own prompt-block count; a sliding-window group is left untouched. + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.blocks_per_sw = [0, 0, 3] + scheduler._group_block_sizes = [16, 32, 16] + + req = SimpleNamespace(num_prompt_tokens=40) + block_ids = [ + [0, 1, 2, 3, 4], # ceil(40/16)=3 -> clamp to 3 + [10, 11, 12], # ceil(40/32)=2 -> clamp to 2 + [20, 21, 22, 23, 24], # sliding-window -> untouched + ] + clamped = scheduler._clamp_to_prompt_blocks(req, block_ids) + assert clamped == [[0, 1, 2], [10, 11], [20, 21, 22, 23, 24]] + + def _make_hybrid_kv_cache_config() -> KVCacheConfig: """One full-attention group + one sliding-window group (Gemma-like).""" full_spec = FullAttentionSpec( @@ -731,6 +976,35 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig: ) +def _make_sliding_window_only_kv_cache_config() -> KVCacheConfig: + """A uniform sliding-window model (e.g. Mistral-7B): the single group is a + SlidingWindowSpec, so HMA is not required once hybrid management is + disabled, yet the spec guard still applies.""" + sw_spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=32, + ) + num_blocks = 2 + page = sw_spec.page_size_bytes + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * page, + layers=["sw0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec), + ], + ) + + def _read_scheduler( kv_cache_config: KVCacheConfig, disable_hma: bool = False ) -> MoRIIOConnectorScheduler: @@ -820,15 +1094,6 @@ def test_non_sliding_window_hybrid_is_rejected(): _read_scheduler(config) -def test_token_count_basis_uses_full_attention_group(): - """Chunked-prefill token counting must use an unclipped full-attention - group (blocks_per_sw == 0), not a clipped sliding-window group.""" - scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) - # Group 0 is the full-attention group - assert scheduler._full_attn_group_idx == 0 - assert scheduler._full_attn_block_size == 16 - - def test_get_exchange_clipped_blocks_clips_only_sw_group(): """get_exchange_clipped_blocks keeps the full attn group intact and clips the sliding-window group to its window tail.""" @@ -897,6 +1162,71 @@ def test_single_group_path_unchanged(): assert scheduler.get_exchange_clipped_blocks(blocks) == blocks +class _FakeSpecConfig(SimpleNamespace): + """Minimal speculative_config stand-in. Any drafter-capability probe + (use_dflash/use_eagle/use_dspark/uses_draft_model/...) defaults to False so + VllmConfig.num_lookahead_tokens resolves cleanly as upstream adds methods.""" + + def __getattr__(self, name: str): + if name.startswith(("use_", "uses_")): + return lambda *a, **k: False + raise AttributeError(name) + + +def _fake_speculative_config(num_speculative_tokens: int = 1) -> SimpleNamespace: + return _FakeSpecConfig(num_speculative_tokens=num_speculative_tokens) + + +def _spec_vllm_config() -> VllmConfig: + vllm_config = create_vllm_config(role="kv_producer", read_mode=False) + vllm_config.speculative_config = _fake_speculative_config() + # HMA off => the WRITE+hybrid guard is skipped, so construction reaches the + # spec-decode guard (the only path that hits it). + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + return vllm_config + + +def test_spec_decode_rejected_with_sliding_window_groups(): + """MoRIIO + speculative decoding fails closed on sliding-window models, + since the prompt-block clamp only covers full-attention groups. The only + reachable path is a uniform SWA model with HMA disabled.""" + vllm_config = _spec_vllm_config() + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError, match="speculative decoding"), + ): + MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_sliding_window_only_kv_cache_config(), + ) + + +def test_spec_decode_allowed_with_full_attention_only(): + """Full-attention-only models (e.g. GLM) do not trip the spec guard.""" + vllm_config = _spec_vllm_config() + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_test_kv_cache_config(), + ) + assert connector.connector_scheduler is not None + + +def test_spec_decode_allowed_with_sliding_window_in_read_mode(): + """The spec guard is WRITE-mode only; READ mode is unaffected.""" + vllm_config = create_vllm_config(role="kv_consumer", read_mode=True) + vllm_config.speculative_config = _fake_speculative_config() + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_hybrid_kv_cache_config(), + ) + assert connector.connector_scheduler is not None + + def test_hybrid_write_mode_rejected(): """Hybrid KV cache groups are unsupported in WRITE mode and fail closed.""" vllm_config = create_vllm_config(role="kv_producer", read_mode=False) diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 9860a7b7e1aa..5ec6fb5755d7 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -661,16 +661,34 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match): wrapper._handle_message(payload) -def test_local_block_ids_longer_than_remote_raises_value_error(): - cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) +def test_local_block_ids_longer_than_remote_raises(): + # A longer local list is a genuine bug here and must fail loudly; the WRITE + # producer save path trims lookahead blocks before they reach this layer. + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) # interleaved worker = _worker({"layer": cache}, {"layer": _full_spec()}) with pytest.raises(ValueError, match="longer than remote_block_ids"): moriio_layout.compute_block_transfer_offsets( - "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks + "layer", + cache, + worker.layer_to_spec, + [1, 3, 6, 7], + [4, 5], + _remote_meta().num_blocks, ) +def test_shorter_local_block_ids_still_transfers_only_what_decode_allocated(): + # READ-mode partial transfer keeps its shorter-local behavior. + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + short = moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [1], [4, 5], _remote_meta().num_blocks + ) + assert len(short[0]) == 1 + + @pytest.mark.parametrize( ("local_block_ids", "remote_block_ids", "match"), [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 5a8a9176640e..2d17d7a7ac52 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -581,12 +581,31 @@ def __init__( else (0, self.block_size) for g in kv_cache_config.transfer_groups ] - # add 1 to conservatively account for boundary overlap eg window isn't fully - # aligned with blocks. + # +1 conservatively covers window/block boundary overlap. self.blocks_per_sw = [ cdiv(n_tokens, block_size) + 1 if n_tokens else 0 for n_tokens, block_size in sw_sizes_tokens ] + # Per-group block size over transfer_groups (matching blocks_per_sw and + # the connector's block payload), used to clamp each full-attention group + # to its prompt length on the WRITE save path. + self._group_block_sizes = [ + getattr(g.kv_cache_spec, "block_size", self.block_size) + for g in kv_cache_config.transfer_groups + ] + # The WRITE-mode prompt-block clamp only covers full-attention groups, so + # MoRIIO + speculative decoding is not yet safe on sliding-window / + # hybrid models. TODO: lift once SWA + spec is validated end-to-end. + if ( + self.mode == MoRIIOMode.WRITE + and vllm_config.speculative_config is not None + and any(bps != 0 for bps in self.blocks_per_sw) + ): + raise NotImplementedError( + "MoRIIO with speculative decoding is not yet supported for " + "sliding-window / hybrid models: the prompt-block clamp is only " + "validated for full-attention KV groups." + ) # In WRITE mode and chunked prefill, we perform the write after the last chunk. # Hence we need to track once we are in the last chunk. We do this by computing # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups @@ -1178,6 +1197,45 @@ def update_state_after_alloc( params["do_remote_prefill"] = False + def _clamp_to_prompt_blocks(self, req: "Request", block_ids: BlockIds) -> BlockIds: + """Drop trailing speculative-lookahead blocks the consumer never + allocated, keeping each full-attention group's prompt blocks. + + Args: + req: Request being saved. + block_ids: Per-KV-cache-group local block ids. + + Returns: + ``block_ids`` with every full-attention group truncated to + ``ceil(num_prompt_tokens / block_size)`` leading blocks; + sliding-window groups are left unchanged. + + """ + clamped = list(block_ids) + changed = False + for g, group_blocks in enumerate(block_ids): + if self.blocks_per_sw[g] != 0: + continue + num_prompt_blocks = math.ceil( + req.num_prompt_tokens / self._group_block_sizes[g] + ) + if len(group_blocks) > num_prompt_blocks: + clamped[g] = group_blocks[:num_prompt_blocks] + changed = True + return clamped if changed else block_ids + + def _is_final_prefill_chunk( + self, req: "Request", scheduler_output: SchedulerOutput + ) -> bool: + """Whether this step completes the prompt. + + ``build_connector_meta`` runs before ``_update_after_schedule``, so + ``num_computed_tokens`` excludes the current chunk; add this step's + scheduled tokens. Token progress is lookahead-safe unlike block counts. + """ + num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0) + return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -1191,38 +1249,36 @@ def build_connector_meta( # It places the request metadata into the saving queue. for i, req_id in enumerate(scheduler_output.scheduled_cached_reqs.req_ids): + # Non-disagg requests aren't registered here; skip them + # instead of KeyError-ing. + if req_id not in self._reqs_need_pending_save: + continue + req, existing_blocks = self._reqs_need_pending_save[req_id] + # The final chunk may allocate no new blocks (the previous chunk + # reserved speculative-lookahead slack), so new_block_ids is None; + # still fall through to the final-chunk check below. new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] - if new_block_ids is not None: - # A non-disagg request (no kv_transfer_params, e.g. smoke - # test) is never registered in _reqs_need_pending_save; - # indexing it unconditionally would KeyError and crash the - # EngineCore. Skip it silently. - if req_id not in self._reqs_need_pending_save: - continue - req, existing_blocks = self._reqs_need_pending_save[req_id] - updated_blocks = [ + existing_blocks = [ existing_blocks[g] + new_block_ids[g] for g in range(len(new_block_ids)) ] - self._reqs_need_pending_save[req_id] = (req, updated_blocks) - saved_tokens = ( - len(updated_blocks[self._full_attn_group_idx]) - * self._full_attn_block_size - ) - if saved_tokens >= req.num_prompt_tokens: - # Final chunk: live kv_transfer_params may be cleared, - # so prefer the snapshot from update_state_after_alloc. - kv_params = self._req_kv_params.pop( - req_id, req.kv_transfer_params or {} - ) - meta.add_new_req( - request_id=req_id, - local_block_ids=self._reqs_need_pending_save[req_id][1], - kv_transfer_params=kv_params, - write_mode=True, - ) - del self._reqs_need_pending_save[req_id] + self._reqs_need_pending_save[req_id] = (req, existing_blocks) + if not self._is_final_prefill_chunk(req, scheduler_output): + continue + # Live kv_transfer_params may be cleared; prefer the + # update_state_after_alloc snapshot. + kv_params = self._req_kv_params.pop( + req_id, req.kv_transfer_params or {} + ) + save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks) + meta.add_new_req( + request_id=req_id, + local_block_ids=save_block_ids, + kv_transfer_params=kv_params, + write_mode=True, + ) + del self._reqs_need_pending_save[req_id] # Loop through scheduled reqs and convert to ReqMeta. for req_id, (req, block_ids) in self._reqs_need_recv.items(): @@ -1253,16 +1309,13 @@ def build_connector_meta( for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) - tokens_covered = ( - len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size - ) - if req.num_prompt_tokens > tokens_covered: - # not last chunk prefill + if not self._is_final_prefill_chunk(req, scheduler_output): self._reqs_need_pending_save[req_id] = (req, block_ids) continue + save_block_ids = self._clamp_to_prompt_blocks(req, block_ids) meta.add_new_req( request_id=req_id, - local_block_ids=block_ids, + local_block_ids=save_block_ids, kv_transfer_params=kv_params, write_mode=True, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index b00692dc5188..8c55837dcd6a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -377,10 +377,8 @@ def iter_layer_registration_regions( ) region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len if geometry.regions_per_block == 1: - # With padded or interleaved pages the block stride exceeds the - # meaningful block_len; register the strided span. The span ends with - # the last block's meaningful bytes, not a whole stride, so it never - # runs past the backing allocation of a strided layer view. + # Padded/interleaved pages stride past block_len: register the strided + # span, ending at the last block's bytes so it stays in-allocation. block_stride_bytes = geometry.block_stride * kv_cache.element_size() region_len = max( region_len, @@ -431,12 +429,8 @@ def compute_block_transfer_offsets( [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] ] = merge_contiguous_offsets, ) -> tuple[list[int], list[int], list[int]]: - # A shorter (or empty) local list is the READ-mode "drop the transfer, just - # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled): - # decode pulls fewer blocks than the prefill holds. The zip loop below pairs - # local[i]<->remote[i] and sizes by len(local), so a short local transfers - # only what decode allocated and an empty local is a no-op. A longer local - # list is a genuine bug and still fails loudly. + # A shorter/empty local list transfers only what decode allocated (READ-mode + # partial or full-prefix-hit); a longer local list is a genuine bug. if len(local_block_ids) > len(remote_block_ids): raise ValueError( "local_block_ids longer than remote_block_ids: " From d64a18832bcb97816f3ee88d7594457c2f50d224 Mon Sep 17 00:00:00 2001 From: Daniel Chan Date: Wed, 7 Oct 2026 12:49:59 -0400 Subject: [PATCH 135/952] [ROCm][Triton] Migrating MiniMax-M3 kernels from make_block_ptr to tensor descriptors (#59285) Signed-off-by: jpvillam Signed-off-by: Daniel Co-authored-by: jpvillam Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Rohan Potdar Co-authored-by: Daniel --- vllm/models/minimax_m3/amd/ops/index_topk.py | 15 +-- vllm/models/minimax_m3/amd/ops/sparse_attn.py | 30 ++--- .../minimax_m3/common/ops/sparse_attn.py | 105 ++++++++++-------- 3 files changed, 80 insertions(+), 70 deletions(-) diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index ce27ea122057..d2b4602c56b8 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -229,15 +229,16 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return - q_ptrs = tl.make_block_ptr( + q_desc = tl.make_tensor_descriptor( base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, - shape=(q_len, head_dim), - strides=(stride_q_n, stride_q_d), - offsets=(pid_q * BLOCK_SIZE_Q, 0), - block_shape=(BLOCK_SIZE_Q, head_dim), - order=(1, 0), + shape=[q_len, head_dim], + strides=[stride_q_n, stride_q_d], + block_shape=[BLOCK_SIZE_Q, head_dim], + padding_option="zero", ) - q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + + # 2. Perform the load by providing the multi-dimensional offset directly + q = q_desc.load([pid_q * BLOCK_SIZE_Q, 0]) q_start = prefix_len + pid_q * BLOCK_SIZE_Q off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len diff --git a/vllm/models/minimax_m3/amd/ops/sparse_attn.py b/vllm/models/minimax_m3/amd/ops/sparse_attn.py index 0719374496be..b4355d18f606 100644 --- a/vllm/models/minimax_m3/amd/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/amd/ops/sparse_attn.py @@ -139,15 +139,15 @@ def _gqa_sparse_fwd_kernel( off_t = tl.arange(0, BLOCK_SIZE_T) topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1) real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0) - q_ptrs = tl.make_block_ptr( + q_desc = tl.make_tensor_descriptor( base=q_ptr + q_start * stride_qn + pid_h * stride_qh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_qn, stride_qh, stride_qd), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + shape=[q_len, gqa_group_size, head_dim], + strides=[stride_qn, stride_qh, stride_qd], + block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D], + padding_option="zero", ) - q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") + + q = q_desc.load([pid_q_j * BLOCK_SIZE_Q, 0, 0]) m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32) @@ -228,15 +228,17 @@ def _gqa_sparse_fwd_kernel( lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None] acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( + o_desc = tl.make_tensor_descriptor( base=o_ptr + q_start * stride_on + pid_h * stride_oh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_on, stride_oh, stride_od), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + shape=[q_len, gqa_group_size, head_dim], + strides=[stride_on, stride_oh, stride_od], + block_shape=[BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D], + ) + + o_desc.store( + offsets=[pid_q_j * BLOCK_SIZE_Q, 0, 0], + value=acc_o.to(o_ptr.dtype.element_ty), ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) @torch.no_grad() diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index af703c086226..8dfd6e0c1996 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -110,7 +110,9 @@ def _gqa_sparse_fwd_kernel( bt_row = block_table_ptr + pid_b * stride_bt_b off_n = tl.arange(0, BLOCK_SIZE_K) off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size for j in range(real_q_loop): pid_q_j = pid_q * num_q_loop + j t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th @@ -118,15 +120,20 @@ def _gqa_sparse_fwd_kernel( q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K real_topk = tl.minimum(max_topk, valid_blocks) - q_ptrs = tl.make_block_ptr( - base=q_ptr + q_start * stride_qn + pid_h * stride_qh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_qn, stride_qh, stride_qd), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + off_qi = pid_q_j * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) + qo_mask = ( + (off_qi < q_len)[:, None, None] + & h_mask[None, :, None] + & d_mask[None, None, :] + ) + q = tl.load( + q_ptr + + (q_start + off_qi)[:, None, None] * stride_qn + + (pid_h + off_h)[None, :, None] * stride_qh + + off_d[None, None, :] * stride_qd, + mask=qo_mask, + other=0.0, ) - q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") off_q = ( tl.arange(0, BLOCK_SIZE_Q)[:, None] + pid_q_j * BLOCK_SIZE_Q @@ -203,15 +210,14 @@ def _gqa_sparse_fwd_kernel( lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None] acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( - base=o_ptr + q_start * stride_on + pid_h * stride_oh, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_on, stride_oh, stride_od), - offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), - block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(2, 1, 0), + tl.store( + o_ptr + + (q_start + off_qi)[:, None, None] * stride_on + + (pid_h + off_h)[None, :, None] * stride_oh + + off_d[None, None, :] * stride_od, + acc_o.to(o_ptr.dtype.element_ty), + mask=qo_mask, ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) # --------------------------------------------------------------------------- @@ -306,21 +312,23 @@ def _gqa_sparse_decode_kernel( off_n = tl.arange(0, BLOCK_SIZE_K) off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] bt_row = block_table_ptr + req_id * stride_bt_b m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q_ptrs = tl.make_block_ptr( - base=q_ptr + pid_b * stride_qn + pid_h * stride_qh, - shape=(gqa_group_size, head_dim), - strides=(stride_qh, stride_qd), - offsets=(0, 0), - block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(1, 0), + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h)[:, None] * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, ) - q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero") cur_idx_ptr = idx_base + chunk_start_topk * stride_tk for _ in tl.range(chunk_start_topk, chunk_end_topk): @@ -392,24 +400,23 @@ def _gqa_sparse_decode_kernel( # can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge. scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i)) acc_o = acc_o * scale[:, None] - o_ptrs = tl.make_block_ptr( - base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h, - shape=(gqa_group_size, head_dim), - strides=(stride_o_h, stride_o_d), - offsets=(0, 0), - block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D), - order=(1, 0), + tl.store( + o_ptr + + pid_c * stride_o_c + + pid_b * stride_o_b + + (pid_h + off_h)[:, None] * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, ) - tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1)) - lse_ptrs = tl.make_block_ptr( - base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h, - shape=(gqa_group_size,), - strides=(stride_l_h,), - offsets=(0,), - block_shape=(BLOCK_SIZE_H,), - order=(0,), + tl.store( + lse_ptr + + pid_c * stride_l_c + + pid_b * stride_l_b + + (pid_h + off_h) * stride_l_h, + lse_i.to(lse_ptr.dtype.element_ty), + mask=h_mask, ) - tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,)) @triton.heuristics( @@ -444,16 +451,16 @@ def _merge_topk_attn_out_kernel( off_c = tl.arange(0, NUM_TOPK_CHUNKS) off_d = tl.arange(0, BLOCK_SIZE_D) - o_ptrs = tl.make_block_ptr( - base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h, - shape=(NUM_TOPK_CHUNKS, head_dim), - strides=(stride_o_c, stride_o_d), - offsets=(0, 0), - block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D), - order=(1, 0), + d_mask = off_d < head_dim + o_ptrs = ( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d ) lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c - o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero") + o = tl.load(o_ptrs, mask=d_mask[None, :], other=0.0) lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0 lse_max = tl.max(lse, axis=0) weights = tl.exp2(lse - lse_max) @@ -462,7 +469,7 @@ def _merge_topk_attn_out_kernel( out_ptrs = ( out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d ) - tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=d_mask) # --------------------------------------------------------------------------- From a8260cbc5c020d250b3964141cde8fd839152714 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:01:41 +0100 Subject: [PATCH 136/952] [CI] Bump Transformers version to 5.19.0 (#60381) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- requirements/common.txt | 2 +- requirements/test/cpu.txt | 2 +- requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 2 +- requirements/test/nightly-torch.txt | 2 +- requirements/test/rocm.in | 2 +- requirements/test/rocm.txt | 2 +- requirements/test/xpu.in | 2 +- requirements/test/xpu.txt | 2 +- .../multimodal/pooling/test_embedding_gemma2.py | 15 +++------------ tests/models/registry.py | 1 - vllm/transformers_utils/processor.py | 6 +++++- 12 files changed, 17 insertions(+), 23 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 433494056d0e..ed62ca42f889 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -7,7 +7,7 @@ requests >= 2.33.0 tqdm blake3 py-cpuinfo -transformers >= 5.16.1, < 5.19.0 +transformers >= 5.16.1, < 5.20.0 huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export tokenizers >= 0.23.2 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 45725d903bb3..7d1f13970fbc 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1228,7 +1228,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index b40d80413eee..4a32471652cf 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -40,7 +40,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 1f64120e31b6..29553d78c7b7 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1335,7 +1335,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index a229116859d3..66704d50d3c6 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator>=0.64.0 # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test. # quantization buildkite-test-collector==0.1.9 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 7c0e4848e944..0ad0b255116a 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -36,7 +36,7 @@ datamodel_code_generator>=0.64.0 # required for minicpm3 test nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.18.0 +transformers==5.19.0 schemathesis>=4.0.0 # Required for openai schema test # quantization bitsandbytes==0.50.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index b2919a6e8361..a0ca0fd83310 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1282,7 +1282,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index 4b3d68ddfc9b..c5d08edde9bc 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -20,7 +20,7 @@ nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs modelscope<1.38 runai-model-streamer[s3,gcs,azure]>=0.15.7 tensorizer==2.10.1 -transformers==5.18.0 +transformers==5.19.0 # --- Audio Processing --- av diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 19cd271017ec..96fda8331f33 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1134,7 +1134,7 @@ traitlets==5.16.1 # via # jupyter-client # jupyter-core -transformers==5.18.0 +transformers==5.19.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 7bb69fa69ad8..701d0914ddbc 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -9,20 +9,11 @@ import pytest import torch -from tests.models.registry import HF_EXAMPLE_MODELS - MODEL = os.environ.get("EG2_MODEL_PATH") needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") TEXT_COS, MM_COS = 0.999, 0.999 -@pytest.fixture -def needs_hf_embedding_gemma2(): - """Skip when the installed transformers has no `embedding_gemma2` module.""" - model_info = HF_EXAMPLE_MODELS.get_hf_info("EmbeddingGemma2Model") - model_info.check_transformers_version(on_fail="skip") - - def _cos(a, b): return torch.nn.functional.cosine_similarity( torch.tensor(a, dtype=torch.float32), @@ -248,7 +239,7 @@ def nparams(m): ) -def test_encoder_contract(needs_hf_embedding_gemma2): +def test_encoder_contract(): from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration as G, @@ -432,7 +423,7 @@ def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): assert len(out[0]) == 768 -def test_video_loader_indices(needs_hf_embedding_gemma2): +def test_video_loader_indices(): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) @@ -517,7 +508,7 @@ def test_video_loader_indices(needs_hf_embedding_gemma2): assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" -def test_video_loader_duplicate_expansion(needs_hf_embedding_gemma2): +def test_video_loader_duplicate_expansion(): from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( EmbeddingGemma2VideoProcessor, ) diff --git a/tests/models/registry.py b/tests/models/registry.py index 52ab94e6163d..4255d439cbdc 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -583,7 +583,6 @@ def check_available_online( ), "EmbeddingGemma2Model": _HfExamplesInfo( "google/embeddinggemma-2", - min_transformers_version="5.19.0", is_available_online=False, ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 64af034d9ac4..9aa592084990 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -181,7 +181,11 @@ def get_video_processor_cls_name_from_config( # given type when torchvision is unavailable; callers then use opencv. model_config = get_hf_file_to_dict("config.json", processor_name, revision=revision) if model_config and "model_type" in model_config: - return VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"]) + video_processor = VIDEO_PROCESSOR_MAPPING_NAMES.get(model_config["model_type"]) + # Transformers >= 5.19 maps to {"torchvision": name} + if isinstance(video_processor, dict): + video_processor = video_processor.get("torchvision") + return video_processor return None From 3d6853b32c7c872df00d9e3489f642dc22e974e0 Mon Sep 17 00:00:00 2001 From: pondzikk <142788369+pondzikk@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:09:50 +0200 Subject: [PATCH 137/952] [Bugfix] Seed hybrid mamba state index with mamba_block_size on prefix-cache hits (#55601) Signed-off-by: pondzikk Signed-off-by: zjy0516 Co-authored-by: Claude Fable 5.1 Co-authored-by: zjy0516 Co-authored-by: ptorsten <1652174+ptorsten@users.noreply.github.com> --- tests/v1/worker/test_mamba_hybrid_model_state.py | 16 ++++++++++++++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 7 ++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index f87201e48de2..fa54ccfcedc5 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -67,6 +67,22 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non assert build_attn_metadata.call_args.kwargs["positions"] is positions +def test_add_request_seeds_state_idx_in_mamba_blocks() -> None: + """A drafter group can lower cache_config.block_size below the mamba + block size; the seeded column must still be in mamba blocks.""" + state = object.__new__(MambaHybridModelState) + state.cache_config = SimpleNamespace(block_size=16, mamba_block_size=880) + state._align_mode = True + state.rope_state = None + state.prompt_embeds_state = None + state.num_accepted_tokens_gpu = torch.ones(2, dtype=torch.int32) + state._mamba_state_idx_gpu = torch.zeros(2, dtype=torch.int32) + + state.add_request(1, SimpleNamespace(num_computed_tokens=107_360)) + + assert state._mamba_state_idx_gpu[1] == 121 + + def test_padded_prompt_tail_builds_as_spec_decode( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 1f8af07898a1..c45304efa512 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -115,8 +115,13 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: self.num_accepted_tokens_gpu[req_index].fill_(1) if self._align_mode: # Seed the running state block from the resumed/prefilled position. + # The column is in mamba blocks. cache_config.block_size is the + # min over prefix-cacheable groups, so a drafter group with its own + # smaller block (e.g. DFlash) pulls it below mamba_block_size. + mamba_block_size = self.cache_config.mamba_block_size + assert mamba_block_size is not None self._mamba_state_idx_gpu[req_index].fill_( - (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + (new_req_data.num_computed_tokens - 1) // mamba_block_size ) def _get_mamba_group_info( From 554340f3d3259e321be4c07282be7a02a5aeef83 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Thu, 8 Oct 2026 02:24:21 +0900 Subject: [PATCH 138/952] [Attention] Support NVFP4 KV cache on SM8x and SM12x with FlashInfer (#46963) Signed-off-by: lesj0610 Signed-off-by: mgoin Co-authored-by: lesj0610 Co-authored-by: mgoin Co-authored-by: Claude --- tests/kernels/attention/test_flashinfer.py | 33 +++++ vllm/v1/attention/backends/flashinfer.py | 138 +++++++++++++++++---- 2 files changed, 149 insertions(+), 22 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index c6cba8d530e4..0d9aecd31e7d 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -6,6 +6,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.kv_cache_layout import KVCacheLayout try: import flashinfer @@ -702,3 +703,35 @@ def test_flashinfer_decode_with_paged_fp8_kv( torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2), f"{torch.max(torch.abs(output - ref_output))}", ) + + +@pytest.mark.skipif( + not any(current_platform.is_device_capability_family(f) for f in (80, 120)), + reason="NVFP4 fa2 path", +) +@pytest.mark.parametrize("layout", ["LBNHC", "LBHNC", "BLNHC", "BLHNC"]) +def test_flashinfer_nvfp4_kv_cache_update(default_vllm_config, layout): + from vllm.v1.attention.backends import flashinfer as fi + + torch.set_default_device("cuda") + order = KVCacheLayout[layout].stride_order + # [L, B, 2H, N, 64 + 8] NVFP4 pages for two layers; layer 1 is written. + shape = (2, 6, 4, 16, 64 + 8) + layers = torch.zeros([shape[i] for i in order], dtype=torch.uint8) + layers = layers.permute(*(order.index(i) for i in range(5))) + one = torch.ones(()) + layer = torch.nn.Module() + layer._k_scale = layer._v_scale = one + kv = torch.randn(2, 96, 2, 128, dtype=torch.bfloat16) + impl = fi.FlashInferImpl(4, 128, 1.0, 2, None, None, "nvfp4") + impl.do_kv_cache_update(layer, kv[0], kv[1], layers[1], torch.arange(96)) + + data, scales, kv_layout = fi._nvfp4_kv_views(layers[1], 2) + pages = torch.arange(6, dtype=torch.int32)[None] + seq_lens = torch.tensor([96], dtype=torch.int32) + deq = torch.empty_like(kv)[:, None] + flashinfer.nvfp4_kv_dequantize_paged( + data, scales, pages, seq_lens, one, one, *deq, kv_layout + ) + assert ((deq[:, 0] - kv).norm() / kv.norm()) < 0.2 + assert not layers[0].any() diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 91ff6da3f037..a21720ab26e6 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -16,6 +16,7 @@ BatchPrefillWithRaggedKVCacheWrapper, MultiLevelCascadeAttentionWrapper, get_seq_lens, + nvfp4_quantize_append_paged_kv_cache_with_slot_mapping, ) from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache from flashinfer.prefill import trtllm_batch_context_with_kv_cache @@ -106,6 +107,38 @@ trtllm_workspace_buffer = None +def _nvfp4_kv_on_fa2() -> bool: + return current_platform.is_device_capability_family( + 80 + ) or current_platform.is_device_capability_family(120) + + +_KVPair = tuple[torch.Tensor, torch.Tensor] + + +def _nvfp4_kv_views( + kv_cache: torch.Tensor, num_kv_heads: int +) -> tuple[_KVPair, _KVPair, str]: + """Views of NVFP4 pages laid out as [K data | K scale | V data | V scale].""" + num_blocks, num_slots, block_size, full_dim = kv_cache.shape + hnd = kv_cache.stride()[1:] == (block_size * full_dim, full_dim, 1) + if not hnd and kv_cache.stride()[1:] != (full_dim, num_slots * full_dim, 1): + raise ValueError(f"NVFP4 KV pages must be dense, got {kv_cache.stride()}") + dims = (num_kv_heads, block_size) if hnd else (block_size, num_kv_heads) + side_size = num_kv_heads * block_size * full_dim + (k_data, k_scale), (v_data, v_scale) = ( + nvfp4_split_data_scale( + kv_cache.as_strided( + (num_blocks, *dims, full_dim), + (kv_cache.stride(0), dims[1] * full_dim, full_dim, 1), + kv_cache.storage_offset() + side * side_size, + ) + ) + for side in range(2) + ) + return (k_data, v_data), (k_scale, v_scale), "HND" if hnd else "NHD" + + def _get_trtllm_workspace_buffer(): global trtllm_workspace_buffer if trtllm_workspace_buffer is None: @@ -498,6 +531,9 @@ def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"): + if _nvfp4_kv_on_fa2(): + # The slot-mapping writer has no scale search. + return kv_cache_dtype == "nvfp4" return ( current_platform.is_device_capability_family(100) and supports_trtllm_attention(is_prefill=True) @@ -505,6 +541,30 @@ def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: ) return super().supports_kv_cache_dtype(kv_cache_dtype) + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + if ( + kv_cache_dtype is not None + and kv_cache_dtype.startswith("nvfp4") + and device_capability.major in (8, 12) + ): + if head_size == 64: + return "fa2 prefill misreads an NVFP4 KV cache at head_size 64" + if has_sink: + return "the sink prefill wrapper cannot read an NVFP4 KV cache" + return None + @classmethod def get_supported_head_sizes(cls) -> list[int]: # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157 @@ -777,7 +837,16 @@ def __init__( # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4") - if self.is_kvcache_nvfp4: + # fa2 plans NVFP4 by its uint8 storage dtype. + self.nvfp4_fa2 = self.is_kvcache_nvfp4 and _nvfp4_kv_on_fa2() + self.nvfp4_trtllm = self.is_kvcache_nvfp4 and not self.nvfp4_fa2 + if self.nvfp4_fa2 and self.use_dcp: + raise NotImplementedError( + "The DCP prefill wrapper cannot read NVFP4 block scales." + ) + if self.nvfp4_fa2 and self.head_dim == 64: + raise NotImplementedError("fa2 prefill misreads NVFP4 at head_size 64.") + if self.nvfp4_trtllm: if ( force_use_trtllm_attention() is False or not supports_trtllm_attention(is_prefill=True) @@ -798,6 +867,7 @@ def __init__( else: self.cache_dtype = "auto" self.is_kvcache_nvfp4 = False + self.nvfp4_fa2 = self.nvfp4_trtllm = False assert self.kv_cache_spec.dtype == self.model_config.dtype self.kv_cache_dtype = self.kv_cache_spec.dtype @@ -810,8 +880,11 @@ def __init__( # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel # must be selected statically for FULL cudagraph capture. - can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( - self.num_qo_heads, self.num_kv_heads, is_prefill=False + can_use_xqa_or_trtllm_gen_decode = ( + can_use_trtllm_attention( + self.num_qo_heads, self.num_kv_heads, is_prefill=False + ) + and not self.nvfp4_fa2 ) # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by # get_supported_kernel_block_sizes). @@ -894,7 +967,7 @@ def __init__( self.window_left = self.global_hyperparameters.window_left self.logits_soft_cap = self.global_hyperparameters.logits_soft_cap self.has_sinks = self.global_hyperparameters.has_sinks - if self.has_sinks and not FlashInferBackend.supports_sink(): + if self.has_sinks and (self.nvfp4_fa2 or not FlashInferBackend.supports_sink()): raise NotImplementedError( "FlashInfer backend currently does not support attention " "sinks, please use trtllm on blackwell or flash attention on " @@ -970,6 +1043,8 @@ def get_q_data_type(self, is_prefill: bool) -> torch.dtype: return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype) return self.model_config.dtype if cache_dtype.startswith("nvfp4"): + if self.nvfp4_fa2: + return self.model_config.dtype return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3") return self.kv_cache_spec.dtype @@ -1002,7 +1077,13 @@ def get_cudagraph_support( num_qo_heads=num_qo_heads, num_kv_heads=spec.num_kv_heads, is_prefill=False, - ) or (is_xqa_arch and not _is_xqa_head_dim_supported(spec.head_size)): + ) or ( + is_xqa_arch + and ( + spec.kv_quant_mode.is_nvfp4 + or not _is_xqa_head_dim_supported(spec.head_size) + ) + ): has_uniform_batch_support = False break @@ -1230,8 +1311,8 @@ def _get_prefill_wrapper( ) else: # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + # the wrapper on SM100; SM8x and SM12x read it with fa2. + backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), @@ -1256,8 +1337,8 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): paged_kv_indices = None paged_kv_last_page_len = None # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + # the wrapper on SM100; SM8x and SM12x read it with fa2. + backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), get_flashinfer_layout_string(self.kv_cache_layout), @@ -1710,7 +1791,7 @@ def build( # use FP8 o_data_type so the wrapper matches the # FP8 output buffer allocated in forward(). o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype + FP8_DTYPE if self.nvfp4_trtllm else self.model_config.dtype ) prefill_wrapper.plan( qo_indptr=qo_indptr_prefill_cpu, @@ -1797,9 +1878,7 @@ def build( # NVFP4 trtllm kernel only supports FP8 output; # use FP8 o_data_type so the wrapper matches the # FP8 output buffer allocated in forward(). - o_dtype = ( - FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype - ) + o_dtype = FP8_DTYPE if self.nvfp4_trtllm else self.model_config.dtype paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1] paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[ :num_input_tokens @@ -1891,6 +1970,8 @@ def __init__( ) self.cache_dtype = kv_cache_dtype self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4") + self.nvfp4_fa2 = self.is_kvcache_nvfp4 and _nvfp4_kv_on_fa2() + self.nvfp4_trtllm = self.is_kvcache_nvfp4 and not self.nvfp4_fa2 self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0 self.logits_soft_cap = logits_soft_cap @@ -1937,7 +2018,7 @@ def __init__( self.o_sf_scale: float | None = None # Pre-allocated FP8 output buffer for NVFP4 without fused output quant. - if self.is_kvcache_nvfp4 and vllm_config is not None: + if self.nvfp4_trtllm and vllm_config is not None: max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self._nvfp4_fp8_out = torch.empty( (max_num_tokens, num_heads, head_size), @@ -2195,10 +2276,9 @@ def forward( canonicalize_singleton_dim_strides(k_cache.permute(*stride_order)), canonicalize_singleton_dim_strides(v_cache.permute(*stride_order)), ) - k_data, k_sf = nvfp4_split_data_scale(kv_cache_tuple[0]) - v_data, v_sf = nvfp4_split_data_scale(kv_cache_tuple[1]) - nvfp4_kv_data = (k_data, v_data) - nvfp4_kv_block_scales = (k_sf, v_sf) + nvfp4_kv_data, nvfp4_kv_block_scales, _ = _nvfp4_kv_views( + kv_cache, self.num_kv_heads + ) else: kv_cache_tuple = kv_cache_permute.split(hs, dim=-1) @@ -2276,7 +2356,7 @@ def forward( # Use a pre-allocated FP8 buffer and dequantize # afterwards. needs_fp8_out_prefill = ( - self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE + self.nvfp4_trtllm and output.dtype != FP8_DTYPE ) if needs_fp8_out_prefill: out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens] @@ -2299,7 +2379,8 @@ def forward( prefill_wrapper.run( prefill_query, kv_cache_for_fi, - q_scale=layer._q_scale_float, + # The fa2 NVFP4 query stays in the model dtype. + q_scale=1.0 if self.nvfp4_fa2 else layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, out=out_prefill, @@ -2451,7 +2532,7 @@ def forward( # NVFP4 kernel only supports FP8 output. # Use a pre-allocated FP8 buffer and dequantize afterwards. - needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE + needs_fp8_out = self.nvfp4_trtllm and output.dtype != FP8_DTYPE if needs_fp8_out: out_decode = self._nvfp4_fp8_out[:num_decode_tokens] else: @@ -2488,7 +2569,7 @@ def forward( decode_wrapper.run( decode_query, kv_cache_for_fi, - q_scale=layer._q_scale_float, + q_scale=1.0 if self.nvfp4_fa2 else layer._q_scale_float, k_scale=layer._k_scale_float, v_scale=layer._v_scale_float, out=out_decode, @@ -2668,6 +2749,19 @@ def do_kv_cache_update( # and value[:num_actual_tokens] because the reshape_and_cache_flash # op uses the slot_mapping's shape to determine the number of # actual tokens. + if self.nvfp4_fa2: + data, scales, kv_layout = _nvfp4_kv_views(kv_cache, self.num_kv_heads) + nvfp4_quantize_append_paged_kv_cache_with_slot_mapping( + key, + value, + slot_mapping, + data, + scales, + layer._k_scale, + layer._v_scale, + kv_layout=kv_layout, + ) + return if self.is_kvcache_nvfp4: # (B, 2*H, N, full_dim) -> ((B, N, H, full_dim), # (B, N, H, full_dim)); From fba462e008f0faa16745d1539c339db692d41f13 Mon Sep 17 00:00:00 2001 From: Eduardo Fonseca Date: Wed, 7 Oct 2026 14:01:21 -0400 Subject: [PATCH 139/952] [Bugfix]Fix FlashInfer warmup crash: autotune(tuning_buckets=...) must round M up, not down (Invalid MXFP8 split-K tactic with spec-decode) (#58165) Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- .../test_flashinfer_autotune_warmup.py | 31 +++++++++++++++++-- vllm/model_executor/warmup/kernel_warmup.py | 23 ++++++++++++-- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index d65f2a4ad116..fe89b09d09bc 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -115,8 +115,8 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): assert get_buckets.call_args_list == [call(8192), call(128)] assert autotune.call_args_list == [ - call(tuning_buckets=max_buckets), - call(tuning_buckets=deferred_buckets), + call(tuning_buckets=max_buckets, round_up=True), + call(tuning_buckets=deferred_buckets, round_up=True), ] assert runner._dummy_run.call_args_list == [ call( @@ -136,6 +136,33 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): ] +def test_flashinfer_autotune_buckets_cover_drafter_tokens(): + """The drafter's M can exceed the pass size; buckets must include it.""" + runner = _make_runner([]) + runner.max_num_reqs = 256 + runner.max_num_tokens = 8192 + runner.speculator = SimpleNamespace(num_query_per_req=6) + + with ( + patch( + "vllm.model_executor.warmup.kernel_warmup." + "_flashinfer_autotune_token_counts", + return_value=(8192, 128), + ), + patch( + "vllm.utils.flashinfer.flashinfer_get_hybrid_num_tokens_buckets" + ) as get_buckets, + patch("vllm.utils.flashinfer.autotune"), + ): + _run_flashinfer_autotune_dummy_runs(runner) + + assert get_buckets.call_args_list == [call(8192), call(128 * 6)] + assert [c.kwargs["num_tokens"] for c in runner._dummy_run.call_args_list] == [ + 8192, + 128, + ] + + class _AutotuneGroup: def __init__(self, run, ranks): self.run = run diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 06f04b75d959..d6544333beb5 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -363,13 +363,32 @@ def _run_flashinfer_autotune_dummy_runs( dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False): - tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) + # The drafter may use more tokens than this pass in the same dummy run. + # Include its max M when building the autotune buckets. + max_tuning_tokens = num_tokens + speculator = getattr(runner, "speculator", None) + if speculator is not None: + num_query_per_req = speculator.num_query_per_req + num_draft_reqs = min( + num_tokens, + runner.max_num_reqs, + runner.max_num_tokens // num_query_per_req, + ) + max_tuning_tokens = max( + max_tuning_tokens, num_draft_reqs * num_query_per_req + ) + + tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets( + max_tuning_tokens + ) logger.info( "Running FlashInfer autotune with %d tokens and token buckets %s.", num_tokens, tuning_buckets, ) - with fi_utils.autotune(tuning_buckets=tuning_buckets): + # Round M up to match serving-time bucket selection for non-bucket M. + # See https://github.com/flashinfer-ai/flashinfer/issues/5450 + with fi_utils.autotune(tuning_buckets=tuning_buckets, round_up=True): runner._dummy_run( num_tokens=num_tokens, skip_eplb=True, From 6ee9be6b3c2dd5557ad72cf714b9d8fe3c2009d4 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 7 Oct 2026 11:01:59 -0700 Subject: [PATCH 140/952] [Spec Decode] Rename AR speculators to StandaloneAR / TargetDependentAR (#60335) Signed-off-by: Giancarlo Delfin --- .buildkite/hardware_tests/cpu.yaml | 2 +- tests/v1/cudagraph/test_cudagraph_manager.py | 4 ++-- ... => test_gpu_target_dependent_ar_speculator.py} | 10 +++++----- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 6 +++--- vllm/v1/worker/gpu/spec_decode/eagle/speculator.py | 8 ++++---- .../v1/worker/gpu/spec_decode/gemma4/speculator.py | 6 +++--- vllm/v1/worker/gpu/spec_decode/mtp/speculator.py | 8 ++++---- .../gpu/spec_decode/multi_module_mtp/speculator.py | 6 +++--- .../{autoregressive => standalone_ar}/__init__.py | 0 .../{draft_model => standalone_ar}/speculator.py | 14 +++++++------- .../__init__.py | 0 .../cudagraph_utils.py | 0 .../speculator.py | 12 +++++++++--- 14 files changed, 42 insertions(+), 36 deletions(-) rename tests/v1/worker/{test_gpu_autoregressive_speculator.py => test_gpu_target_dependent_ar_speculator.py} (98%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => standalone_ar}/__init__.py (100%) rename vllm/v1/worker/gpu/spec_decode/{draft_model => standalone_ar}/speculator.py (97%) rename vllm/v1/worker/gpu/spec_decode/{draft_model => target_dependent_ar}/__init__.py (100%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => target_dependent_ar}/cudagraph_utils.py (100%) rename vllm/v1/worker/gpu/spec_decode/{autoregressive => target_dependent_ar}/speculator.py (99%) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index fbd43d80ddb2..3fdce975d3f4 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -87,7 +87,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s \ - tests/v1/worker/test_gpu_autoregressive_speculator.py \ + tests/v1/worker/test_gpu_target_dependent_ar_speculator.py \ tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \ tests/v1/worker/test_gpu_batch_ordering.py \ tests/v1/worker/test_gpu_rejection_sampler_chunking.py \ diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index 668fbdf61b4a..1e4071bcfbfd 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -23,10 +23,10 @@ from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBuffers from vllm.v1.worker.gpu.pcp_manager import PCPManager -from vllm.v1.worker.gpu.spec_decode.autoregressive import ( +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import ( cudagraph_utils as spec_cudagraph_utils, ) -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py similarity index 98% rename from tests/v1/worker/test_gpu_autoregressive_speculator.py rename to tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 0a8cfd1a130e..20e86f4a7827 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -22,17 +22,17 @@ from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module -from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import ( MultiModuleMTPSpeculator, ) from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as spec_module +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class _TestSpeculator(AutoRegressiveSpeculator): +class _TestSpeculator(TargetDependentARSpeculator): def load_draft_model(self, target_model, target_attn_layer_names): return self.test_draft_model diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d0c08707f483..373042a41502 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -820,7 +820,7 @@ def _dummy_run( "skip_attn must only be True for initial memory profiling." ) - # Create a dummy scheduler output. Plain draft-model speculation adds + # Create a dummy scheduler output. Standalone AR speculation adds # one correction slot per request during prefill. The scheduler # accounts for these slots, while dummy runs bypass the scheduler. # Adjust the dummy token count only when the expanded draft batch diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 902a4f085ca9..fc028941cf0d 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -71,11 +71,11 @@ def init_speculator( return EagleSpeculator(vllm_config, device) elif speculative_config.uses_draft_model(): - from vllm.v1.worker.gpu.spec_decode.draft_model.speculator import ( - PlainDraftModelSpeculator, + from vllm.v1.worker.gpu.spec_decode.standalone_ar.speculator import ( + StandaloneARSpeculator, ) - return PlainDraftModelSpeculator(vllm_config, device) + return StandaloneARSpeculator(vllm_config, device) elif speculative_config.use_ngram(): from vllm.v1.worker.gpu.spec_decode.ngram.speculator import ( NgramGPUSpeculator, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index e878872e622b..0611d42b1ce5 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -3,13 +3,13 @@ import torch.nn as nn -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class EagleSpeculator(AutoRegressiveSpeculator): +class EagleSpeculator(TargetDependentARSpeculator): def load_draft_model( self, target_model: nn.Module, diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index 6f76e5055a82..9d553830ae17 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py @@ -17,8 +17,8 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model from vllm.model_executor.model_loader.utils import get_draft_load_config -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, ) logger = init_logger(__name__) @@ -45,7 +45,7 @@ def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None: getattr(attn, scale_name).copy_(getattr(target_attn, scale_name)) -class Gemma4Speculator(AutoRegressiveSpeculator): +class Gemma4Speculator(TargetDependentARSpeculator): @property def advance_draft_positions(self) -> bool: # Gemma4 MTP is Q-only and reads K/V from the target's existing cache. diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index d94702612d9b..a0064197f756 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -3,13 +3,13 @@ import torch.nn as nn -from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( - AutoRegressiveSpeculator, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) -class MTPSpeculator(AutoRegressiveSpeculator): +class MTPSpeculator(TargetDependentARSpeculator): share_mtp_topk_indices: bool = False def load_draft_model( diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 71f7f89416be..919e939c03cc 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -14,11 +14,11 @@ from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( - SpeculatorCudaGraphManager, -) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( + SpeculatorCudaGraphManager, +) from vllm.v1.worker.utils import get_uniform_decode_token_count logger = init_logger(__name__) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/__init__.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/__init__.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/__init__.py rename to vllm/v1/worker/gpu/spec_decode/standalone_ar/__init__.py diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py similarity index 97% rename from vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py rename to vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 5be7b0c86ddd..8bdfb8530be6 100644 --- a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -26,22 +26,22 @@ logger = init_logger(__name__) -class PlainDraftModelSpeculator(DraftModelSpeculator): - """Speculative decoding using a separate smaller draft LM. +class StandaloneARSpeculator(DraftModelSpeculator): + """Standalone autoregressive (AR) drafting with a separate smaller LM. Unlike Eagle, the draft model runs fully independently of the target model. Step 0 builds an expanded buffer (accepted + correction token + rejected slots) via a Triton kernel; steps 1..k-1 are single-token decode steps. """ - # Plain draft model needs one extra slot per request for correction. + # Standalone AR drafting needs one extra slot per request for correction. num_extra_query_per_req = 1 def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) # draft_max_seq_len is read by the parent's attention metadata builder. - # Plain draft model doesn't do per-batch adjustment; cap at max. + # Standalone AR drafting doesn't do per-batch adjustment; cap at max. self.draft_max_seq_len = self.max_model_len self.last_token_indices = torch.zeros( @@ -65,10 +65,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.supports_mm_inputs = False def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - pass # CUDA graph not yet supported for plain draft model speculator. + pass # CUDA graph not yet supported for standalone AR speculator. def capture(self) -> None: - pass # CUDA graph not yet supported for plain draft model speculator. + pass # CUDA graph not yet supported for standalone AR speculator. def set_attn( self, @@ -415,7 +415,7 @@ def _prepare_prefill_inputs_kernel( max_model_len, BLOCK_SIZE: tl.constexpr, ): - """Per-request step-0 input preparation for the plain draft-model speculator. + """Per-request step-0 input preparation for the standalone AR speculator. Output layout for request i (out_start = query_start_loc[i] + i): [out_start, out_start + num_valid) accepted tokens diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/__init__.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/__init__.py diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/cudagraph_utils.py similarity index 100% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/cudagraph_utils.py diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py similarity index 99% rename from vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py rename to vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 46ff82681561..9a8ed7b8b4ca 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -17,16 +17,22 @@ from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) -from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.utils import AttentionGroup, get_uniform_decode_token_count logger = init_logger(__name__) -class AutoRegressiveSpeculator(DraftModelSpeculator): +class TargetDependentARSpeculator(DraftModelSpeculator): + """Autoregressive (AR) drafting that depends on target model state. + + Base for Eagle, MTP, and Gemma4 drafters, which consume the target's + hidden states (and, for Gemma4, its KV cache) at each draft step. + """ + def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) From 32fbfa15e8bacc64182cb1286831bd63d7e4fc12 Mon Sep 17 00:00:00 2001 From: Jinzhen Lin Date: Thu, 8 Oct 2026 02:04:43 +0800 Subject: [PATCH 141/952] [Bugfix][Quantization] Accept folded NVFP4 scales for Humming A16 (#60337) Signed-off-by: jinzhen.ljz Co-authored-by: Kevin H. Luu --- .../quantization/utils/humming/schema.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/humming/schema.py b/vllm/model_executor/layers/quantization/utils/humming/schema.py index c84e8a370d22..aaf6d4a66c87 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/schema.py +++ b/vllm/model_executor/layers/quantization/utils/humming/schema.py @@ -322,6 +322,22 @@ def check_and_fallback_input_schema( input_group_size = input_schema.input_scale_group_size input_scale_dtype = input_schema.input_scale_dtype + def is_compatible(candidate: "HummingInputSchema") -> bool: + check_weight_schema = weight_schema + if ( + candidate.a_dtype in (None, humming_dtypes.float16, humming_dtypes.bfloat16) + and weight_schema.b_dtype == humming_dtypes.float4e2m1 + and weight_group_size == 16 + and weight_schema.bs_dtype + in (None, humming_dtypes.float16, humming_dtypes.bfloat16) + ): + # Humming rejects folded NVFP4 scales despite A16 software dequant + # supporting them. Relax only this check, preserving kernel metadata. + check_weight_schema = dataclasses.replace( + weight_schema, bs_dtype=humming_dtypes.float8e4m3 + ) + return candidate.is_compatible_with(check_weight_schema, param_dtype) + def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: if dtype is None or dtype.num_bits == 16: return False @@ -333,7 +349,7 @@ def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120 return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated - if input_schema.is_compatible_with(weight_schema, param_dtype): + if is_compatible(input_schema): if not allow_fallback: if is_deprecated(a_dtype): msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now." @@ -390,7 +406,7 @@ def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: input_scale_dtype=scale_dtype, input_quant_mode=quant_mode, ) - if candidate.is_compatible_with(weight_schema, param_dtype): + if is_compatible(candidate): return candidate raise ValueError( From b9767299db6d5c0226e6879dadb079a477524d3a Mon Sep 17 00:00:00 2001 From: Cheese Cake Date: Thu, 8 Oct 2026 00:38:11 +0530 Subject: [PATCH 142/952] [Bugfix] Respect Inductor deterministic mode in combo-kernel defaults (#59074) Signed-off-by: cheesecake Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 --- docs/design/debug_vllm_compile.md | 12 ++++++ tests/compile/test_config.py | 67 +++++++++++++++++++++++++++++++ vllm/config/compilation.py | 8 +++- 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md index 7edda6fa6476..af9b3f530d4c 100644 --- a/docs/design/debug_vllm_compile.md +++ b/docs/design/debug_vllm_compile.md @@ -251,6 +251,18 @@ that may call 1+ triton kernels. On rare (but unfortunate) occasions, it may produce an incorrect triton kernel. This may manifest as silent incorrectness, CUDA illegal memory accesses, or loud errors. +### Deterministic compilation + +When vLLM selects combo-kernel defaults, it disables combo-kernel benchmarking +if Inductor's `deterministic` setting is enabled. A `deterministic` value in +`inductor_compile_config` takes precedence over the global Inductor setting. +On torch 2.13, `TORCHINDUCTOR_DETERMINISTIC=1` alone is reset by the first +compiled frame ([pytorch#198563](https://github.com/pytorch/pytorch/issues/198563)), +so set `deterministic` in `inductor_compile_config` instead. +Explicit `combo_kernels` or `benchmark_combo_kernel` settings are preserved. +This avoids conflicting compilation defaults; it does not guarantee reproducible +model outputs. + ### Inductor runtime assertions By default (on torch < 2.12), vLLM disables Inductor's runtime assertions diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 24cb816b5d9d..5a1ba13a57f1 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -1366,3 +1366,70 @@ def test_inductor_asserts_user_override(monkeypatch): assert config.inductor_compile_config.get("size_asserts") is True if not _is_torch_equal_or_newer(torch.__version__, "2.12.0.dev"): assert config.inductor_compile_config.get("alignment_asserts") is False + + +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize("override", [None, False, True]) +def test_combo_kernel_benchmarking_respects_deterministic(deterministic, override): + from torch._inductor import config as inductor_config + + overrides = {} if override is None else {"deterministic": override} + with ( + inductor_config.patch(deterministic=deterministic), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig(inductor_compile_config=overrides) + + assert config.inductor_compile_config["combo_kernels"] is True + effective_deterministic = deterministic if override is None else override + assert config.inductor_compile_config["benchmark_combo_kernel"] is ( + not effective_deterministic + ) + + +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize( + "overrides", + [ + {"combo_kernels": False}, + {"combo_kernels": True}, + {"benchmark_combo_kernel": False}, + {"benchmark_combo_kernel": True}, + ], +) +def test_combo_kernel_explicit_settings_preserved(deterministic, overrides): + from torch._inductor import config as inductor_config + + with ( + inductor_config.patch(deterministic=deterministic), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig(inductor_compile_config=overrides.copy()) + + for key in ("combo_kernels", "benchmark_combo_kernel"): + assert config.inductor_compile_config.get(key) == overrides.get(key) + + +@pytest.mark.parametrize("is_cpu,torch_version", [(True, "2.13.0"), (False, "2.8.0")]) +def test_combo_kernel_defaults_require_supported_platform_and_torch( + is_cpu, torch_version +): + with ( + patch("vllm.config.compilation.current_platform.is_cpu", return_value=is_cpu), + patch("torch.__version__", torch_version), + ): + config = CompilationConfig() + + assert "combo_kernels" not in config.inductor_compile_config + assert "benchmark_combo_kernel" not in config.inductor_compile_config + + +def test_combo_kernel_defaults_without_inductor_deterministic_setting(): + with ( + patch("torch._inductor.config", SimpleNamespace()), + patch("vllm.config.compilation.current_platform.is_cpu", return_value=False), + ): + config = CompilationConfig() + + assert config.inductor_compile_config["combo_kernels"] is True + assert config.inductor_compile_config["benchmark_combo_kernel"] is True diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index afe576daa9d3..cbd132ab5b68 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1002,10 +1002,16 @@ def __post_init__(self) -> None: # (fixme @boyuan) combo kernel does not support cpu yet. and not current_platform.is_cpu() ): + from torch._inductor import config as inductor_config + # use horizontal fusion, which is useful for fusing qk-norm and # qk-rope when query and key have different shapes. self.inductor_compile_config["combo_kernels"] = True - self.inductor_compile_config["benchmark_combo_kernel"] = True + + deterministic = self.inductor_compile_config.get( + "deterministic", getattr(inductor_config, "deterministic", False) + ) + self.inductor_compile_config["benchmark_combo_kernel"] = not deterministic if self.use_inductor_graph_partition and not is_torch_equal_or_newer( "2.9.0.dev" From 17c70c895359f7c941199bdf49b723c602ba75de Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:14:07 -0700 Subject: [PATCH 143/952] [KV Offload] Update KVCR adapter to pool-and-index API (#59899) --- .buildkite/scripts/install-kv-offload.sh | 2 +- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 32 ++++++++------- vllm/v1/kv_offload/tiering/kvcr/manager.py | 40 +++++++++---------- 3 files changed, 36 insertions(+), 38 deletions(-) diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh index 46d756b621b1..eb8bab9e393f 100644 --- a/.buildkite/scripts/install-kv-offload.sh +++ b/.buildkite/scripts/install-kv-offload.sh @@ -10,7 +10,7 @@ fi # TODO: Use the PyPI wheel once it is published. uv pip install --system \ - "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.1" # Keep only the NIXL wheel matching the CI image's CUDA runtime. NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index 5cbc85230ffb..fc4421eb06b2 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -18,7 +18,7 @@ BlockKey, CacheTier, InventoryEvent, - MemDescriptor, + MemoryRef, OpEntryResult, OpEntryStatus, OpHandle, @@ -78,11 +78,9 @@ def __init__(self) -> None: self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.discard_hint_calls: list[str] = [] self.deliver_calls: list[ - tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] - ] = [] - self.deposit_calls: list[ - tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + tuple[OpHandle, dict[BlockKey, list[MemoryRef]], str | None] ] = [] + self.deposit_calls: list[tuple[OpHandle, dict[BlockKey, list[MemoryRef]]]] = [] self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] self._next_op_handle = 1 @@ -107,7 +105,7 @@ def query( def deliver( self, - blocks: Mapping[BlockKey, list[MemDescriptor]], + blocks: Mapping[BlockKey, list[MemoryRef]], request_id: str | None = None, ) -> OpHandle: op_handle = self._next_op_handle @@ -118,7 +116,7 @@ def deliver( def deposit( self, - blocks: Mapping[BlockKey, list[MemDescriptor]], + blocks: Mapping[BlockKey, list[MemoryRef]], ) -> OpHandle: op_handle = self._next_op_handle self._next_op_handle += 1 @@ -248,7 +246,11 @@ def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): compatibility_digest="Opaque-Digest", ) assert kvcr.config is not None - assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.config.pool_layouts == [("", tier.block_size_bytes)] + (region,) = kvcr.backend_configs.framework_regions + assert region.addr == np.asarray(tier._primary_kv_view).ctypes.data + assert (region.size, region.stride or region.size, region.count) == (16, 16, 4) + assert (region.label, region.mem_type, region.device_Id) == ("", "DRAM", 0) assert kvcr.backend_configs.local_dram is None @@ -306,9 +308,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch): assert request_id == "req" assert list(blocks) == [key] (descriptor,) = blocks[key] - assert descriptor.end_point_name == kvcr.nixl_agent_name - assert descriptor.addr == tier._primary_base_addr + 2 * 16 - assert descriptor.size == 16 + assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) assert list(tier.get_finished_jobs()) == [JobResult(7, True)] # Here we verify that request cleanup discards the request-scoped hint in @@ -351,7 +351,7 @@ def test_kvcr_tier_serves_primary_pin_request(monkeypatch): tier = _make_tier(monkeypatch, kvcr) keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) hit_keys = (keys[0], keys[2]) - chunk_ids = {keys[0]: 1, keys[2]: 5} + chunk_ids = {keys[0]: 1, keys[2]: 3} lifecycle: list[str] = [] class Parent: @@ -388,7 +388,7 @@ def on_request_finished(self, req_context): block_descriptors = descriptors[keys[2]] assert block_descriptors is not None (descriptor,) = block_descriptors - assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=3) assert lifecycle == ["new", "finished"] polls = 0 @@ -482,14 +482,16 @@ def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): local_dram = kvcr.backend_configs.local_dram assert local_dram is not None assert [(name, size) for name, _, size in local_dram.pools] == [ - ("", 2 * tier._primary_row_stride) + ("", 2 * tier.block_size_bytes) ] key = OffloadKey(b"k0") tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) _, blocks = kvcr.deposit_calls[0] - assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert blocks[key] == [ + MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) + ] assert list(tier.get_finished_jobs()) == [JobResult(11, True)] assert kvcr.inventory_sink is not None diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 18da2a6c2764..4b6f81b3b71e 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -22,7 +22,6 @@ KVCRBindings, ) from kvcr.config import ( - FrameworkDramInput, G3Options, KeyAdapter, KVCRBackendConfigs, @@ -43,11 +42,12 @@ BlockKey, CacheTier, InventoryEvent, - MemDescriptor, + MemoryRef, OpHandle, PinRequestId, PinResult, QueryStatus, + RegionDescriptor, ) from typing_extensions import override @@ -177,7 +177,7 @@ def cancel_pin_request(self, request: PinRequestId) -> None: def _resolve_pin( self, parent: ParentManager, keys: Collection[BlockKey] - ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + ) -> tuple[str, dict[BlockKey, list[MemoryRef] | None]] | None: offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) request_id = f"kvcr-source:{self._next_request_id}" self._next_request_id += 1 @@ -202,7 +202,7 @@ def _resolve_pin( if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): return None - descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + descriptors: dict[BlockKey, list[MemoryRef] | None] = { BlockKey(bytes(key)): None for key in offload_keys } descriptors.update( @@ -350,12 +350,6 @@ def __init__( with socket.socket() as _s: _s.bind(("", 0)) _nixl_listen_port = _s.getsockname()[1] - self._primary_base_addr = ctypes.addressof( - ctypes.c_char.from_buffer(primary_kv_view) - ) - if primary_kv_view.strides is None: - raise ValueError("primary KV memoryview must expose strides") - self._primary_row_stride = int(primary_kv_view.strides[0]) if secondary_g2_slots < 0: raise ValueError("secondary_g2_slots must be non-negative") if ( @@ -375,7 +369,7 @@ def __init__( ) else: local_mapping = mmap.mmap( - -1, secondary_g2_slots * self._primary_row_stride + -1, secondary_g2_slots * self.block_size_bytes ) local_dram = LocalDramOptions( pools=[ @@ -408,7 +402,7 @@ def __init__( try: self._kvcr = KVCR( KVCRConfig( - pool_layouts=[("", self._primary_row_stride)], + pool_layouts=[("", self.block_size_bytes)], enable_telemetry=enable_telemetry, operation_timeout_ms=operation_timeout_ms, abandon_timeout_ms=abandon_timeout_ms, @@ -429,9 +423,15 @@ def __init__( policy=selected_policy, ), KVCRBackendConfigs( - framework_dram=FrameworkDramInput( - self._primary_base_addr, primary_kv_view.nbytes - ), + framework_regions=[ + RegionDescriptor( + addr=ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ), + size=self.block_size_bytes, + count=len(primary_kv_view), + ) + ], local_dram=local_dram, g3=g3_config, remote_fw_dram=RemoteFWDramOptions( @@ -621,12 +621,8 @@ def _record_inventory(self, event: InventoryEvent) -> None: ) ) - def _make_descriptor(self, chunk_id: int) -> MemDescriptor: - return MemDescriptor( + def _make_descriptor(self, chunk_id: int) -> MemoryRef: + return MemoryRef( end_point_name=self._kvcr.config.nixl_agent_name, - mem_type="DRAM", - addr=self._primary_base_addr + chunk_id * self._primary_row_stride, - size=self._primary_row_stride, - device_Id=0, - info="", + element_index=chunk_id, ) From e572e02d21b2b3559b3ce33adb117dca3f50ccd7 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:16:19 -0400 Subject: [PATCH 144/952] [UX] Skip checkpoint shards an MTP head does not load (#60152) Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw Co-authored-by: Claude Opus 5.5 --- .../test_filter_duplicate_safetensors.py | 60 +++++++++++++++++++ .../model_loader/default_loader.py | 12 +++- .../model_loader/weight_utils.py | 36 +++++++++++ vllm/model_executor/models/nemotron_h_mtp.py | 15 +++-- vllm/model_executor/models/qwen3_5_mtp.py | 14 +++-- vllm/model_executor/models/qwen3_next_mtp.py | 11 ++-- vllm/models/deepseek_v32/amd/mtp.py | 8 ++- vllm/models/deepseek_v32/nvidia/mtp.py | 8 ++- vllm/models/deepseek_v4/nvidia/mtp.py | 40 +++++++------ 9 files changed, 165 insertions(+), 39 deletions(-) diff --git a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py index babc6c061a9c..da04d0108fdf 100644 --- a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py +++ b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py @@ -4,11 +4,17 @@ import json import os import tempfile +from types import SimpleNamespace import pytest +import torch +from safetensors.torch import save_file +from vllm.config.load import LoadConfig +from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import ( filter_duplicate_safetensors_files, + filter_safetensors_files_by_weight_name, ) @@ -77,3 +83,57 @@ def test_filter_duplicate_safetensors_files_all_exist(): if __name__ == "__main__": test_filter_duplicate_safetensors_files_missing_weight() test_filter_duplicate_safetensors_files_all_exist() + + +def _skip_non_mtp(name: str) -> bool: + return not name.startswith("mtp.") + + +@pytest.fixture +def mtp_checkpoint(tmp_path): + """Three shards: target only, target + MTP, MTP only.""" + shards = { + "model-00001-of-00003.safetensors": ["model.layers.0.weight"], + "model-00002-of-00003.safetensors": ["model.layers.1.weight", "mtp.a.weight"], + "model-00003-of-00003.safetensors": ["mtp.b.weight"], + } + for filename, names in shards.items(): + save_file({name: torch.zeros(1) for name in names}, tmp_path / filename) + return tmp_path + + +def test_filter_by_weight_name_drops_shards_without_wanted_weights(mtp_checkpoint): + files = sorted(str(f) for f in mtp_checkpoint.glob("*.safetensors")) + + assert filter_safetensors_files_by_weight_name(files, _skip_non_mtp) == files[1:] + # A rule that rejects everything must not leave the loader without files. + assert filter_safetensors_files_by_weight_name(files, lambda _: True) == files + + +def test_loader_does_not_read_shards_the_model_skips(mtp_checkpoint, monkeypatch): + """Whole-file loaders such as InstantTensor only see the kept shards.""" + opened: list[list[str]] = [] + + def fake_iterator(hf_weights_files, *args, **kwargs): + opened.append(sorted(os.path.basename(f) for f in hf_weights_files)) + return iter(()) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.default_loader." + "instanttensor_weights_iterator", + fake_iterator, + ) + + class Drafter(torch.nn.Module): + def is_unused_checkpoint_weight(self, name: str) -> bool: + return _skip_non_mtp(name) + + loader = DefaultModelLoader(LoadConfig(load_format="instanttensor")) + model_config = SimpleNamespace(model=str(mtp_checkpoint), revision=None) + list(loader.get_all_weights(model_config, Drafter())) + list(loader.get_all_weights(model_config, torch.nn.Module())) + + assert opened == [ + ["model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors"], + sorted(f.name for f in mtp_checkpoint.glob("*.safetensors")), + ] diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index c3e09bccd425..9f1c25dd2e34 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -6,7 +6,7 @@ import glob import os import time -from collections.abc import Generator, Iterable +from collections.abc import Callable, Generator, Iterable from typing import TYPE_CHECKING, cast import torch @@ -28,6 +28,7 @@ filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, filter_mm_encoder_only_safetensors_files, + filter_safetensors_files_by_weight_name, get_quant_config, instanttensor_weights_iterator, maybe_download_from_modelscope, @@ -75,6 +76,10 @@ class Source: allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" + is_unused_weight: Callable[[str], bool] | None = None + """If defined, safetensors shards holding only weights whose checkpoint + name (before *prefix*) it accepts are not read.""" + counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 @@ -283,6 +288,10 @@ def _get_weights_iterator( f"`{source.model_or_path}`; check language_model prefixes " f"{self._encoder_only_lm_prefixes}" ) + if source.is_unused_weight is not None and use_safetensors: + hf_weights_files = filter_safetensors_files_by_weight_name( + hf_weights_files, source.is_unused_weight + ) if self.load_config.load_format == "npcache": # Currently np_cache only support *.bin checkpoints assert use_safetensors is False @@ -359,6 +368,7 @@ def get_all_weights( prefix="", fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), + is_unused_weight=getattr(model, "is_unused_checkpoint_weight", None), ) yield from self._get_weights_iterator(primary_weights) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 12d557fa1ca2..c6a72c4c33c6 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -785,6 +785,42 @@ def _is_lm_only(weight_file: str) -> bool: return kept +def filter_safetensors_files_by_weight_name( + hf_weights_files: list[str], + is_unused_weight: Callable[[str], bool], +) -> list[str]: + """Drop safetensors shards in which `is_unused_weight` accepts every tensor. + + Loaders that read whole files (InstantTensor, fastsafetensors, multi-thread, + eager, prefetch) cannot skip single tensors, so dropping shards is their + only way to avoid reading weights the model does not load. + + Args: + hf_weights_files: Safetensors shard paths. + is_unused_weight: Returns True for checkpoint weight names the model does + not load. + + Returns: + The shards holding at least one wanted tensor, or `hf_weights_files` + unchanged when no shard does. + + """ + kept: list[str] = [] + for st_file in hf_weights_files: + with safe_open(st_file, framework="pt") as f: + if not all(is_unused_weight(name) for name in f.keys()): # noqa: SIM118 + kept.append(st_file) + if not kept: + return hf_weights_files + if len(kept) < len(hf_weights_files): + logger.info_once( + "Skipped %d/%d safetensors shard(s) holding no weights the model loads", + len(hf_weights_files) - len(kept), + len(hf_weights_files), + ) + return kept + + # explicitly use pure text format, with a newline at the end # this makes it impossible to see the animation in the progress bar # but will avoid messing up with ray or multiprocessing, which wraps diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index 519f85a63109..af93cc62ae25 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -426,6 +426,11 @@ def compute_logits( ) return self.logits_processor(self.lm_head, hidden_states) + def is_unused_checkpoint_weight(self, name: str) -> bool: + """Only MTP, embedding and LM head weights are loaded.""" + name = name.removeprefix("language_model.") + return not name.startswith(("mtp.", "lm_head.")) and "embeddings" not in name + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load MTP weights with proper name remapping.""" stacked_params_mapping = [ @@ -456,17 +461,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # MTP weights are nested in "language_model." # in Multimodal Nemotron-H checkpoints. name = name.removeprefix("language_model.") + if self.is_unused_checkpoint_weight(name): + continue is_lm_head_weight = name.startswith("lm_head.") if is_lm_head_weight: self.has_own_lm_head = True - # Only process MTP and LM head weights - - # skip all non-MTP and non-LM head weights - if ( - not name.startswith("mtp.") - and "embeddings" not in name - and not is_lm_head_weight - ): - continue # Skip rotary embeddings (computed, not loaded) if "rotary_emb.inv_freq" in name: continue diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index c17d17014f8d..d7fb1d6e09ea 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -300,16 +300,20 @@ def compute_logits( ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) + def is_unused_checkpoint_weight(self, name: str) -> bool: + return not name.startswith("mtp.") and not any( + key in name for key in ["embed_tokens", "lm_head"] + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(weights): for name, weight in weights: + if self.is_unused_checkpoint_weight(name): + continue if name.startswith("mtp."): name = name.replace("mtp.", "model.") - elif any(key in name for key in ["embed_tokens", "lm_head"]): - if "embed_tokens" in name: - name = name.replace("language_model.", "") - else: - continue + elif "embed_tokens" in name: + name = name.replace("language_model.", "") yield name, weight loader = AutoWeightsLoader(self) diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index e5e306bb861c..ef08878eaa7f 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -227,15 +227,18 @@ def compute_logits( ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - shared_weight_names = ["embed_tokens", "lm_head"] + def is_unused_checkpoint_weight(self, name: str) -> bool: + return not name.startswith("mtp.") and not any( + key in name for key in ["embed_tokens", "lm_head"] + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(weights): for name, weight in weights: + if self.is_unused_checkpoint_weight(name): + continue if name.startswith("mtp."): name = name.replace("mtp.", "model.") - elif not any(key in name for key in shared_weight_names): - continue yield name, weight loader = AutoWeightsLoader(self) diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index 70ddeae374c8..6230348ef44c 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -251,6 +251,12 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: name = name.replace(f"model.layers.{spec_layer}.", "model.") return name + def is_unused_checkpoint_weight(self, name: str) -> bool: + return ( + get_spec_layer_idx_from_weight_name(self.config, name) is None + and "embed_tokens" not in name + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), @@ -278,7 +284,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() _pending_wk_fp8: dict = {} for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: + if "rotary_emb.inv_freq" in name or self.is_unused_checkpoint_weight(name): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index bdc6b95c15e4..d4dee7b9d462 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -356,6 +356,12 @@ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: name = name.replace(f"model.layers.{spec_layer}.", "model.") return name + def is_unused_checkpoint_weight(self, name: str) -> bool: + return ( + get_spec_layer_idx_from_weight_name(self.config, name) is None + and "embed_tokens" not in name + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), @@ -378,7 +384,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() _pending_wk_fp8: dict = {} for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: + if "rotary_emb.inv_freq" in name or self.is_unused_checkpoint_weight(name): continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 64547af038aa..c0f3471dd5b6 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -340,6 +340,26 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def _to_spec_layer_name(self, name: str) -> str: + """Remap V4 `mtp.{i}.*` names to `model.layers.{num_hidden_layers + i}.*` + so that get_spec_layer_idx_from_weight_name can identify them.""" + mtp_layer_idx = 0 + for subname in name.split("."): + try: + # we use the first encountered integer + mtp_layer_idx = int(subname) + break + except ValueError: + continue + return name.replace( + f"mtp.{mtp_layer_idx}.", + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", + ) + + def is_unused_checkpoint_weight(self, name: str) -> bool: + name = self._to_spec_layer_name(name) + return get_spec_layer_idx_from_weight_name(self.config, name) is None + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Weight name remapping for checkpoint compatibility. # Maps checkpoint weight paths to model parameter paths. @@ -356,16 +376,6 @@ def _remap_weight_name(name: str) -> str: name = name.replace(old_pattern, new_pattern) return name - def _find_mtp_layer_idx(name: str) -> int: - subnames = name.split(".") - for subname in subnames: - try: - # we return the first encountered integer - return int(subname) - except ValueError: - continue - return 0 - stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "w1", 0), @@ -409,15 +419,7 @@ def _find_mtp_layer_idx(name: str) -> int: ) for name, loaded_weight in weights: - mtp_layer_idx = _find_mtp_layer_idx(name) - # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to - # `model.layers.{num_hidden_layers + i}.*` so that - # get_spec_layer_idx_from_weight_name can identify them. - name = name.replace( - f"mtp.{mtp_layer_idx}.", - f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", - ) - + name = self._to_spec_layer_name(name) spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: continue From 08567505b09324f891797f4a13c1bd73fc75a827 Mon Sep 17 00:00:00 2001 From: Rui <57304024+ai-yang@users.noreply.github.com> Date: Thu, 8 Oct 2026 03:36:57 +0800 Subject: [PATCH 145/952] [Bugfix][KV Offload] Restore KVCR construction through the secondary-tier factory (#58088) --- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 53 +++++++++++++------ vllm/v1/kv_offload/tiering/kvcr/manager.py | 6 ++- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index fc4421eb06b2..0907293fd883 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -36,7 +36,9 @@ ReqContext, make_offload_key, ) +from vllm.v1.kv_offload.tiering.backpressure import EMABackpressureDetector from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager @@ -180,6 +182,7 @@ def _make_tier( g3: dict[str, object] | None = None, control_ports: list[int] | None = None, data_parallel_rank_local: int | None = None, + backpressure: dict[str, Any] | None = None, ) -> KVCRSecondaryTierManager: def make_control(_bind_host, bind_port, advertise_host): return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") @@ -196,7 +199,25 @@ def make_kvcr(config, bindings, backend_configs, guard_config): monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) - return KVCRSecondaryTierManager( + tier_config = { + "type": "kvcr", + "router_capabilities": ["router_hint"], + "control_host": "127.0.0.1", + "control_ports": control_ports if control_ports is not None else [7777], + "control_advertise_host": "127.0.0.1", + "enable_telemetry": enable_telemetry, + "secondary_g2_slots": secondary_g2_slots, + "kvcr_service_socket_path": kvcr_service_socket_path, + "compatibility_digest": compatibility_digest, + "policy": policy, + "g3": g3, + "local_dram_backend": "UCX", + "remote_fw_dram_backend": "UCX", + } + if backpressure is not None: + tier_config["backpressure"] = backpressure + tier = SecondaryTierFactory.create_secondary_tier( + tier_config=tier_config, offloading_spec=SimpleNamespace( config=SimpleNamespace( parallel=SimpleNamespace( @@ -209,20 +230,9 @@ def make_kvcr(config, bindings, backend_configs, guard_config): ), ), primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), - tier_type="kvcr", - router_capabilities=["router_hint"], - control_host="127.0.0.1", - control_ports=control_ports if control_ports is not None else [7777], - control_advertise_host="127.0.0.1", - enable_telemetry=enable_telemetry, - secondary_g2_slots=secondary_g2_slots, - kvcr_service_socket_path=kvcr_service_socket_path, - compatibility_digest=compatibility_digest, - policy=policy, - g3=g3, - local_dram_backend="UCX", - remote_fw_dram_backend="UCX", ) + assert isinstance(tier, KVCRSecondaryTierManager) + return tier def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): @@ -271,10 +281,19 @@ def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): ) -def test_kvcr_tier_adapts_request_and_load(monkeypatch): - """Check hint forwarding, key conversion, load descriptors, and cleanup.""" +@pytest.mark.parametrize( + "backpressure", + [None, {"backpressure_cls": "EMABackpressureDetector"}], + ids=["default", "backpressure"], +) +def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): + """Factory-created tiers preserve configuration and complete hinted loads.""" kvcr = RecordingKVCR() - tier = _make_tier(monkeypatch, kvcr) + tier = _make_tier(monkeypatch, kvcr, backpressure=backpressure) + if backpressure is None: + assert tier.bp_detector is None + else: + assert isinstance(tier.bp_detector, EMABackpressureDetector) kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[]) ctx = ReqContext( req_id="req", diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 4b6f81b3b71e..4804f193e3e8 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -82,6 +82,7 @@ if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.backpressure import BackpressureDetector _REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES @@ -295,8 +296,11 @@ def __init__( g3: dict[str, Any] | None = None, local_dram_backend: str = "UCX", remote_fw_dram_backend: str = "UCX", + backpressure_detector: "BackpressureDetector | None" = None, ) -> None: - super().__init__(offloading_spec, primary_kv_view, tier_type) + super().__init__( + offloading_spec, primary_kv_view, tier_type, backpressure_detector + ) selected_policy = _resolve_policy(policy) if (kvcr_service_socket_path is None) != (compatibility_digest is None): raise ValueError( From a982c81e3ed337f76b116338192c83f776a3225c Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 7 Oct 2026 16:37:13 -0400 Subject: [PATCH 146/952] [Bugfix] Limit GlmMoeDsaForCausalLM fp8 KV cache default to SM100 (#60281) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/config.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 97b9bfa0d38b..1a1bf9206166 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -46,14 +46,22 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: from vllm.platforms import current_platform cache_config = vllm_config.cache_config - if cache_config.cache_dtype == "auto": - if current_platform.is_xpu(): - cache_config.cache_dtype = "bfloat16" - elif current_platform.is_cuda_alike(): - capability = current_platform.get_device_capability() - cache_config.cache_dtype = ( - "fp8_e4m3" if capability and capability.major >= 10 else "bfloat16" - ) + if ( + cache_config.cache_dtype == "auto" + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + cache_config.cache_dtype = "fp8_e4m3" + logger.info_once("Using fp8 kv-cache for GlmMoeDsaForCausalLM on SM10x") + # MTP shares the target's KV cache layout; a separate draft model + # must not inherit a default it was never validated with. + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.method != "mtp" + and spec_config.kv_cache_dtype is None + ): + spec_config.kv_cache_dtype = "auto" # For Glm-Moe-DSA, qrep + a2a is better than the default all-gather + ag-rs # in most cases. From b47ec440bc0f2342848d6f1d99942cc3e909a66a Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:15:39 -0400 Subject: [PATCH 147/952] [CI] Stamp provenance labels on the torch-nightly image so Initialized Snapshot E2E can run (#59487) Signed-off-by: Thang Nguyen Co-authored-by: Kimi --- .buildkite/image_build/image_build_torch_nightly.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh index 56c7116e14ac..4d978b5299c8 100755 --- a/.buildkite/image_build/image_build_torch_nightly.sh +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -61,6 +61,8 @@ NIGHTLY_UBUNTU_VERSION="24.04" NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu${NIGHTLY_UBUNTU_VERSION}" echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" +# Provenance labels match the bake `_labels` target in docker/docker-bake.hcl; +# the initialized-snapshot E2E refuses images whose commit labels differ. docker buildx build --file docker/Dockerfile \ --build-arg max_jobs=16 \ --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ @@ -70,6 +72,10 @@ docker buildx build --file docker/Dockerfile \ --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ + --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT}" \ + --label "ai.vllm.build.commit=${BUILDKITE_COMMIT}" \ + --label "ai.vllm.build.pipeline=${BUILDKITE_PIPELINE_SLUG:-local}" \ + --label "ai.vllm.build.url=${BUILDKITE_BUILD_URL:-}" \ --tag "$IMAGE_TAG" \ --push \ --target test \ From 947dd62c0700851a79b0a8f32cb823e25a003b31 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Wed, 7 Oct 2026 23:18:19 +0200 Subject: [PATCH 148/952] [CI] Fix failing test_deepseek_v4_rocm_dspark* tests (#60462) Signed-off-by: simondanielsson --- .buildkite/test_areas/models_basic.yaml | 1 + tests/models/test_deepseek_v4_dspark_rocm.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 04efae00fb17..a25c9a78a86c 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -81,6 +81,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/models/deepseek_v4/ - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py index efa97e10d7f9..a5f8426e4ea5 100644 --- a/tests/models/test_deepseek_v4_dspark_rocm.py +++ b/tests/models/test_deepseek_v4_dspark_rocm.py @@ -15,7 +15,10 @@ def _make_uninitialized_model(confidence_head): object.__setattr__( model, "model", - SimpleNamespace(confidence_head=confidence_head), + SimpleNamespace( + confidence_head=confidence_head, + layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + ), ) return model From b803edbc9fa5adbbf31136bf128d88533f7d39c2 Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:56:40 -0300 Subject: [PATCH 149/952] [Model] Streamline EmbeddingGemma 2 config resolution and harmonize Triton prefill attention (#60289) Signed-off-by: Luciano Martins Co-authored-by: Luciano Martins Co-authored-by: Matthew Bonanni --- .../pooling/test_embedding_gemma2.py | 125 +++++++----------- vllm/model_executor/models/config.py | 88 +----------- vllm/transformers_utils/config.py | 2 +- 3 files changed, 54 insertions(+), 161 deletions(-) diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py index 701d0914ddbc..0dcb72406da5 100644 --- a/tests/models/multimodal/pooling/test_embedding_gemma2.py +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -573,9 +573,10 @@ def test_video_media_io_embedding_gemma2(): assert meta_conn.get("do_sample_frames") is False -def test_config_capping_and_explicit_limits(caplog): - """Verify EmbeddingGemma2ModelConfig capping and explicit user overrides.""" - import logging +def test_config_capping_and_explicit_limits(): + """Verify EmbeddingGemma2ModelConfig backend selection, default capping, + and preservation of explicit limits. + """ from types import SimpleNamespace from unittest.mock import patch @@ -585,19 +586,13 @@ def test_config_capping_and_explicit_limits(caplog): ) from vllm.v1.attention.backends.registry import AttentionBackendEnum - # (a) Defaults: uncapped max_model_len=262144 -> capped to 8192, - # and derived scheduler limits (>= 262144) lowered to 8192. + # (a) When attention backend is None, default to TRITON_ATTN, + # and default capping: max_model_len > 8192 without explicit limit capped to 8192 cfg_default = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=262144, - max_num_encoder_input_tokens=262144, - encoder_cache_size=262144, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=262144, + original_max_model_len=None, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): @@ -605,83 +600,61 @@ def test_config_capping_and_explicit_limits(caplog): assert cfg_default.attention_config.backend == AttentionBackendEnum.TRITON_ATTN assert cfg_default.model_config.max_model_len == 8192 - assert cfg_default.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_default.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_default.scheduler_config.encoder_cache_size == 8192 - - # (b) Reachable explicit limits: explicit max_model_len=8192 and explicit - # max_num_batched_tokens=32768 stays 32768. - # Note: omitting max_model_len while passing max_num_batched_tokens < 262144 - # is rejected earlier by SchedulerConfig validation (32768 < 262144) before - # this hook runs. - cfg_explicit_batched = SimpleNamespace( + + # (b) Default capping with original_max_model_len == -1 + cfg_default_minus1 = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=8192, original_max_model_len=8192 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=32768, - max_num_encoder_input_tokens=32768, - encoder_cache_size=32768, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=131072, + original_max_model_len=-1, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_batched) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_default_minus1) - assert cfg_explicit_batched.model_config.max_model_len == 8192 - assert cfg_explicit_batched.scheduler_config.max_num_batched_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.max_num_encoder_input_tokens == 32768 - assert cfg_explicit_batched.scheduler_config.encoder_cache_size == 32768 + assert cfg_default_minus1.model_config.max_model_len == 8192 - # (c) Explicit value > uncapped_len (e.g. 300000) with uncapped model len is lowered - # to 8192 and logs a warning because it exceeds uncapped_len. - cfg_explicit_large = SimpleNamespace( + # (c) Explicit limit 4096 preserved + cfg_explicit_4k = SimpleNamespace( attention_config=SimpleNamespace(backend=None), model_config=SimpleNamespace( - model="dummy", max_model_len=262144, original_max_model_len=None + max_model_len=4096, + original_max_model_len=4096, ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=300000, - max_num_encoder_input_tokens=300000, - encoder_cache_size=300000, - max_num_seqs=256, - verify_max_model_len=lambda x: None, - ), - ) - with ( - caplog.at_level(logging.WARNING), - patch.object(Gemma4Config, "verify_and_update_config"), - ): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_large) - - assert cfg_explicit_large.model_config.max_model_len == 8192 - assert cfg_explicit_large.scheduler_config.max_num_batched_tokens == 8192 - assert cfg_explicit_large.scheduler_config.max_num_encoder_input_tokens == 8192 - assert cfg_explicit_large.scheduler_config.encoder_cache_size == 8192 - assert any( - "lowering explicitly set scheduler_config.max_num_batched_tokens" - " from 300000 to 8192" in rec.message - for rec in caplog.records ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_4k) - # (d) Explicit max_model_len=16384 is respected and not capped to 8192 - cfg_explicit_model_len = SimpleNamespace( - attention_config=SimpleNamespace(backend=None), + assert cfg_explicit_4k.model_config.max_model_len == 4096 + + # (d) Explicit limit 16384 preserved and explicit backend preserved + cfg_explicit_16k = SimpleNamespace( + attention_config=SimpleNamespace(backend=AttentionBackendEnum.FLASH_ATTN), model_config=SimpleNamespace( - model="dummy", max_model_len=16384, original_max_model_len=16384 - ), - scheduler_config=SimpleNamespace( - max_num_batched_tokens=16384, - max_num_encoder_input_tokens=16384, - encoder_cache_size=16384, - max_num_seqs=256, - verify_max_model_len=lambda x: None, + max_model_len=16384, + original_max_model_len=16384, ), ) with patch.object(Gemma4Config, "verify_and_update_config"): - EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_model_len) + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_16k) + + assert cfg_explicit_16k.attention_config.backend == AttentionBackendEnum.FLASH_ATTN + assert cfg_explicit_16k.model_config.max_model_len == 16384 + + +def test_sentence_transformer_tokenizer_config_max_seq_length(tmp_path): + """Verify get_sentence_transformer_tokenizer_config accepts config + with only max_seq_length. + """ + import json + + from vllm.transformers_utils.config import ( + get_sentence_transformer_tokenizer_config, + ) - assert cfg_explicit_model_len.model_config.max_model_len == 16384 - assert cfg_explicit_model_len.scheduler_config.max_num_batched_tokens == 16384 + # 1. Synthetic model directory with only max_seq_length in sentence_bert_config.json + cfg_file = tmp_path / "sentence_bert_config.json" + cfg_file.write_text(json.dumps({"max_seq_length": 8192})) + loaded = get_sentence_transformer_tokenizer_config(str(tmp_path)) + assert loaded is not None + assert loaded["max_seq_length"] == 8192 diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 1a1bf9206166..7f6a3aef40a7 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -309,90 +309,10 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) model_config = vllm_config.model_config - orig_len = getattr(model_config, "original_max_model_len", None) - if orig_len in (None, -1): - default_len = 8192 - # Check sentence_bert_config.json if available - try: - import json - from pathlib import Path - - cfg_path = Path(model_config.model) / "sentence_bert_config.json" - if cfg_path.exists(): - data = json.loads(cfg_path.read_text()) - if "max_seq_length" in data and isinstance( - data["max_seq_length"], int - ): - default_len = data["max_seq_length"] - except Exception: - pass - - if model_config.max_model_len > default_len: - uncapped_len = model_config.max_model_len - logger.info( - "EmbeddingGemma2: max_model_len not explicitly set; capping from " - "%d to %d.", - model_config.max_model_len, - default_len, - ) - model_config.max_model_len = default_len - scheduler_config = vllm_config.scheduler_config - if scheduler_config is not None: - new_batched_tokens = max( - model_config.max_model_len, scheduler_config.max_num_seqs - ) - if scheduler_config.max_num_batched_tokens >= uncapped_len: - if scheduler_config.max_num_batched_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - else: - logger.info( - "EmbeddingGemma2: lowering " - "scheduler_config.max_num_batched_tokens from %d to %d " - "to match capped max_model_len.", - scheduler_config.max_num_batched_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_batched_tokens = new_batched_tokens - if hasattr(vllm_config, "_set_compile_ranges"): - vllm_config._set_compile_ranges() - if scheduler_config.max_num_encoder_input_tokens >= uncapped_len: - if scheduler_config.max_num_encoder_input_tokens > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.max_num_encoder_input_tokens " - "from %d to %d to match capped max_model_len.", - scheduler_config.max_num_encoder_input_tokens, - new_batched_tokens, - ) - scheduler_config.max_num_encoder_input_tokens = ( - new_batched_tokens - ) - if scheduler_config.encoder_cache_size >= uncapped_len: - if scheduler_config.encoder_cache_size > uncapped_len: - logger.warning( - "EmbeddingGemma2: lowering explicitly set " - "scheduler_config.encoder_cache_size from %d to %d " - "to match capped max_model_len.", - scheduler_config.encoder_cache_size, - new_batched_tokens, - ) - scheduler_config.encoder_cache_size = new_batched_tokens - if hasattr(scheduler_config, "verify_max_model_len"): - scheduler_config.verify_max_model_len( - model_config.max_model_len - ) - elif model_config.max_model_len > 32768: - logger.warning( - "EmbeddingGemma2: max_model_len=%d is large; " - "consider --max-model-len 8192 to bound encoder memory.", - model_config.max_model_len, - ) + if model_config.max_model_len > 8192 and getattr( + model_config, "original_max_model_len", None + ) in (None, -1): + model_config.max_model_len = 8192 class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 49312047144f..6a339297eec5 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1139,7 +1139,7 @@ def get_sentence_transformer_tokenizer_config( logger.info("Found sentence-transformers tokenize configuration.") - if all(k in encoder_dict for k in ("max_seq_length", "do_lower_case")): + if "max_seq_length" in encoder_dict: return encoder_dict return None From f38c6778ce4cf62b1ca7a8e091892f5bfdf79f98 Mon Sep 17 00:00:00 2001 From: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com> Date: Thu, 8 Oct 2026 06:01:09 +0800 Subject: [PATCH 150/952] [quantization] Add per-token NVFP4 quantization MoE support for ReLU2 (#56740) Signed-off-by: Xuanteng Huang Co-authored-by: Claude Fable 5.1 Co-authored-by: OpenAI Codex --- ...ron-3-Super-120B-A12B-NVFP4-per-token.yaml | 14 ++ .../gsm8k/configs/models-blackwell-ep.txt | 1 + tests/kernels/moe/test_trtllm_nvfp4_moe.py | 130 +++++++++++++++--- .../model_loader/test_weight_cache.py | 46 +++++++ tests/quantization/test_blackwell_moe.py | 13 ++ tests/quantization/test_modelopt.py | 77 ++++++++++- tests/quantization/test_online.py | 32 ++++- .../test_quantization_config_args.py | 51 +++++-- vllm/config/quantization.py | 2 + .../layers/fused_moe/oracle/nvfp4.py | 2 + .../layers/quantization/modelopt.py | 57 ++++++-- .../layers/quantization/online/nvfp4.py | 17 ++- .../quantization/utils/flashinfer_fp4_moe.py | 7 +- .../quantization/utils/flashinfer_utils.py | 14 ++ .../model_loader/weight_cache/protocol.py | 19 ++- .../model_loader/weight_utils.py | 19 ++- vllm/models/deepseek_v4/quant_config.py | 5 + vllm/models/deepseek_v41/quant_config.py | 5 + 18 files changed, 450 insertions(+), 61 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml new file mode 100644 index 000000000000..62587ae10a7d --- /dev/null +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml @@ -0,0 +1,14 @@ +model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" +accuracy_threshold: 0.93 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1200 +server_args: >- + --enforce-eager + --max-model-len 4096 + --tensor-parallel-size 2 + --enable-expert-parallel + --mamba-backend flashinfer + --moe-backend flashinfer_trtllm + --quantization-config.moe.activation nvfp4_per_token + --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/models-blackwell-ep.txt b/tests/evals/gsm8k/configs/models-blackwell-ep.txt index c2a19051986f..34090d720512 100644 --- a/tests/evals/gsm8k/configs/models-blackwell-ep.txt +++ b/tests/evals/gsm8k/configs/models-blackwell-ep.txt @@ -1,3 +1,4 @@ Qwen3-Next-80B-A3B-NVFP4-EP2.yaml Qwen3-Next-FP8-EP2.yaml Nemotron-3-Super-120B-A12B-NVFP4.yaml +Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 292e4549bb7d..4ee095505196 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -16,8 +16,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config from tests.kernels.quantization.nvfp4_utils import ( - FLOAT4_E2M1_MAX, - FLOAT8_E4M3_MAX, + convert_swizzled_to_linear, dequantize_nvfp4_to_dtype, ) from tests.kernels.utils import torch_moe @@ -39,6 +38,13 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_static_weights_for_trtllm_fp4_moe, + reorder_w1w3_to_w3w1, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe from vllm.utils.math_utils import next_power_of_2 @@ -96,10 +102,14 @@ def __init__(self, *, compile_native: bool = True) -> None: def _make_trtllm_fp4_moe_kernel( - moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, ) -> mk.FusedMoEKernel: experts = TrtLlmNvFp4ExpertsModular( - moe_config=moe_config, quant_config=quant_config + moe_config=moe_config, + quant_config=quant_config, + per_token_activation=per_token_activation, ) # Mimic the production weight-loader path so per-expert tensors that # are normally precomputed in process_weights_after_loading (g1_scale_c @@ -126,16 +136,37 @@ def _make_trtllm_fp4_moe_kernel( ACTIVATION_CASES = [ - pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"), - pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"), - pytest.param(MoEActivation.SITU, SITU, None, id="situ"), + pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, False, id="silu"), + pytest.param( + MoEActivation.SILU, + MoEActivation.SILU, + None, + True, + id="silu_per_token", + ), + pytest.param( + MoEActivation.SILU, + SILU_WITH_CLAMP, + _SWIGLU_LIMIT, + False, + id="silu_clamp", + ), + pytest.param(MoEActivation.SITU, SITU, None, False, id="situ"), pytest.param( MoEActivation.RELU2_NO_MUL, MoEActivation.RELU2_NO_MUL, None, + False, id="relu2_no_mul", ), - pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"), + pytest.param( + MoEActivation.RELU2_NO_MUL, + MoEActivation.RELU2_NO_MUL, + None, + True, + id="relu2_no_mul_per_token", + ), + pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, False, id="gelu"), ] @@ -143,7 +174,10 @@ def _make_trtllm_fp4_moe_kernel( @pytest.mark.parametrize("e", [128]) @pytest.mark.parametrize("topk", [8]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("activation,torch_activation,swiglu_limit", ACTIVATION_CASES) +@pytest.mark.parametrize( + "activation,torch_activation,swiglu_limit,per_token_activation", + ACTIVATION_CASES, +) @torch.inference_mode() def test_trtllm_fp4_moe_no_graph( m: int, @@ -155,6 +189,7 @@ def test_trtllm_fp4_moe_no_graph( activation: MoEActivation, torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul], swiglu_limit: float | None, + per_token_activation: bool, workspace_init, ): # FlashInfer's trtllm_batched_gemm_runner has no precompiled tile @@ -202,6 +237,42 @@ def test_trtllm_fp4_moe_no_graph( # clamp/output-scale coupling in the FlashInfer kernel wrapper. quant_config.g1_alphas.fill_(_LARGE_OUTPUT1_SCALE) + w1_q_linear = w1_q + w2_q_linear = w2_q + w1_scale_reference = quant_config.w1_scale + w2_scale_reference = quant_config.w2_scale + assert w1_scale_reference is not None + assert w2_scale_reference is not None + w1_scale_linear = torch.stack( + [ + convert_swizzled_to_linear(scale, w1_q.size(1), k, quant_blocksize) + for scale in w1_scale_reference + ] + ) + w2_scale_linear = torch.stack( + [ + convert_swizzled_to_linear(scale, w2_q.size(1), n, quant_blocksize) + for scale in w2_scale_reference + ] + ) + w1_q_for_kernel = w1_q + if is_gated_act: + w1_q_for_kernel, w1_scale_linear = reorder_w1w3_to_w3w1( + w1_q.clone(), w1_scale_linear + ) + w1_q, w1_scale, w2_q, w2_scale = prepare_static_weights_for_trtllm_fp4_moe( + w1_q_for_kernel, + w2_q, + w1_scale_linear, + w2_scale_linear, + hidden_size=k, + intermediate_size=n, + num_experts=e, + is_gated_activation=is_gated_act, + ) + quant_config._w1.scale = w1_scale + quant_config._w2.scale = w2_scale + score = torch.randn((m, e), device="cuda", dtype=dtype) topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) @@ -224,7 +295,11 @@ def test_trtllm_fp4_moe_no_graph( ), ) - trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) + trtllm_experts = _make_trtllm_fp4_moe_kernel( + moe_config, + quant_config, + per_token_activation=per_token_activation, + ) if activation == MoEActivation.SITU: torch.testing.assert_close( trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale @@ -245,10 +320,20 @@ def test_trtllm_fp4_moe_no_graph( # Reference: round-trip activations and weights through FP4 # quant/dequant so the comparison isolates kernel/activation behavior # from quantization error. - a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / a.abs().max()).to( - torch.float32 - ) - a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) + if per_token_activation: + from flashinfer import SfLayout, nvfp4_quantize + + a_fp4, a_scale_interleaved, per_token_scale = nvfp4_quantize( + a, + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + a_global_scale = 1.0 / per_token_scale.unsqueeze(1) + else: + assert quant_config.a1_gscale is not None + a_global_scale = quant_config.a1_gscale[0] + a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) a_in_dtype = dequantize_nvfp4_to_dtype( a_fp4, a_scale_interleaved, @@ -256,6 +341,7 @@ def test_trtllm_fp4_moe_no_graph( dtype=a.dtype, device=a.device, block_size=quant_blocksize, + is_sf_linear_layout=per_token_activation, ) w1_d = torch.empty( @@ -264,16 +350,16 @@ def test_trtllm_fp4_moe_no_graph( w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) for idx in range(e): w1_d[idx] = dequantize_nvfp4_to_dtype( - w1_q[idx], - quant_config.w1_scale[idx], + w1_q_linear[idx], + w1_scale_reference[idx], (1 / quant_config.g1_alphas[idx]), dtype=dtype, device=w1_q.device, block_size=quant_blocksize, ) w2_d[idx] = dequantize_nvfp4_to_dtype( - w2_q[idx], - quant_config.w2_scale[idx], + w2_q_linear[idx], + w2_scale_reference[idx], (1 / quant_config.g2_alphas[idx]), dtype=dtype, device=w2_q.device, @@ -284,6 +370,13 @@ def test_trtllm_fp4_moe_no_graph( a_in_dtype, w1_d, w2_d, score, topk, activation=torch_activation ) + if swiglu_limit is None: + cosine = torch.nn.functional.cosine_similarity( + torch_output.float().flatten(), + trtllm_output.float().flatten(), + dim=0, + ) + assert cosine > 0.99 torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1) @@ -416,5 +509,6 @@ def test_trtllm_fp4_moe_deferred_finalize( MoEActivation.GELU, MoEActivation.GELU, None, + False, None, ) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 731c73262586..7960a237914b 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -356,6 +356,52 @@ def test_weight_cache_key_distinguishes_dp_ranks(): assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"] +@pytest.mark.parametrize("hf_quant_config", [None, {"quant_algo": "NVFP4"}]) +def test_weight_cache_key_distinguishes_nvfp4_activation_override( + tmp_path, hf_quant_config +): + """Cache keys isolate activation modes and reject ambiguous legacy hashes.""" + import json + from dataclasses import replace + from types import SimpleNamespace + + import torch + + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey + from vllm.utils.hashing import safe_hash + + model_config = SimpleNamespace( + model=str(tmp_path), + dtype=torch.bfloat16, + quantization="modelopt_fp4", + quantization_config=None, + revision=None, + hf_config=SimpleNamespace( + architectures=["NemotronHForCausalLM"], + quantization_config=hf_quant_config, + ), + ) + static_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0) + model_config.quantization_config = QuantizationConfigArgs( + moe={"activation": "nvfp4_per_token"} + ) + per_token_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0) + assert static_key.mismatched_fields(per_token_key) == ["quant_config_hash"] + + legacy_hash = ( + "" + if hf_quant_config is None + else safe_hash( + json.dumps(hf_quant_config, sort_keys=True).encode(), + usedforsecurity=False, + ).hexdigest() + ) + assert static_key.mismatched_fields( + replace(static_key, quant_config_hash=legacy_hash) + ) == ["quant_config_hash"] + + def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): """Copy mode clones the weights into the engine, so nothing is external (vllm_config is never touched).""" diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index da70491bbc27..3a4e2b704b81 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -297,3 +297,16 @@ def test_nemotron_fp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): hf_overrides=HF_OVERRIDE_TEXT, extra_args=["--moe-backend=flashinfer_trtllm"], ) + + +def test_nemotron_nvfp4_per_token_moe_flashinfer_trtllm( + monkeypatch: pytest.MonkeyPatch, +): + can_initialize( + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=[ + "--moe-backend=flashinfer_trtllm", + "--quantization-config.moe.activation=nvfp4_per_token", + ], + ) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 7b1131a00e39..33c798cba09b 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -6,6 +6,7 @@ """ import os +from contextlib import nullcontext from typing import Any, NoReturn from unittest.mock import MagicMock, Mock, patch @@ -17,7 +18,9 @@ load_model_without_vllm_runner, ) from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.load import LoadConfig from vllm.config.model import ModelConfig +from vllm.config.quantization import QuantizationConfigArgs from vllm.forward_context import set_forward_context from vllm.model_executor.kernels.linear import ( FlashInferCuteDslNvFp4W4A16LinearKernel, @@ -40,12 +43,14 @@ kMxfp8Dynamic, kMxfp8Static, kNvfp4Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) +from vllm.model_executor.model_loader.weight_utils import get_quant_config from vllm.platforms import current_platform @@ -113,11 +118,13 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon ) -def test_modelopt_nvfp4_quantizes_parallel_lm_head(): +@pytest.mark.parametrize("moe_activation", [None, "nvfp4_per_token"]) +def test_modelopt_nvfp4_quantizes_parallel_lm_head(moe_activation): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, kv_cache_quant_algo=None, exclude_modules=[], + quantization_args=QuantizationConfigArgs(moe={"activation": moe_activation}), ) method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") @@ -790,6 +797,74 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16( assert kwargs["activation_key"] is kNvfp4Dynamic +@pytest.mark.parametrize("quantization", ["modelopt_fp4", "modelopt_mixed"]) +@pytest.mark.parametrize("per_token", [False, True]) +def test_modelopt_nvfp4_moe_activation_override(quantization, per_token): + """Runtime overrides reach NVFP4 configs without changing HF metadata.""" + hf_config = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION" + if quantization == "modelopt_mixed" + else "NVFP4", + "quantized_layers": {"experts": {"quant_algo": "NVFP4"}}, + } + args = ( + QuantizationConfigArgs(moe={"activation": "nvfp4_per_token"}) + if per_token + else None + ) + config = get_quant_config( + Mock( + quantization=quantization, + quantization_config=args, + hf_config=Mock(quantization_config=hf_config), + ), + LoadConfig(), + ) + if isinstance(config, ModelOptMixedPrecisionConfig): + config = config.nvfp4_config + assert isinstance(config, ModelOptNvFp4Config) + assert config.moe_activation_override == (kNvfp4DynamicToken if per_token else None) + assert "_online_quantization_args" not in hf_config + + +@pytest.mark.parametrize( + "quant_method,activation,backend,error", + [ + ("NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", None), + ("W4A16_NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", "W4A4 checkpoint"), + ( + "NVFP4", + kMxfp8Dynamic, + "FLASHINFER_TRTLLM", + "Unsupported.*activation override", + ), + ("NVFP4", kNvfp4DynamicToken, "MARLIN", "requires the FlashInfer TRTLLM"), + ], +) +def test_modelopt_nvfp4_moe_validates_activation_override( + quant_method, activation, backend, error +): + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend + from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE + + config = ModelOptNvFp4Config( + quant_method=quant_method, + quantization_args=QuantizationConfigArgs(moe={"activation": activation}), + ) + expected = pytest.raises(ValueError, match=error) if error else nullcontext() + with ( + patch( + "vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend", + return_value=(NvFp4MoeBackend[backend], Mock()), + ) as select_backend, + expected, + ): + method = ModelOptNvFp4FusedMoE(config, Mock()) + assert method.per_token_activation + assert select_backend.call_args.kwargs["activation_key"] == activation + + @pytest.mark.parametrize( "per_layer_algo, expected_weight, expected_activation", [ diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 3f27757b9017..c6b9894009cf 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -103,6 +103,8 @@ amax_for_tp_weight_quant, kMxfp8Dynamic, kMxfp8Static, + kNvfp4DynamicToken, + kNvfp4Static, weight_amax, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -211,15 +213,26 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( assert args.ignore == [] +@pytest.mark.parametrize( + ("is_act_and_mul", "expected_hidden_alignment"), + [(True, 256), (False, 512)], +) def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( - monkeypatch, + monkeypatch, is_act_and_mul, expected_hidden_alignment ) -> None: - method = object.__new__(Nvfp4OnlineMoEMethod) - method.moe = SimpleNamespace(is_act_and_mul=True) - method.nvfp4_backend = object() - method.experts_cls = object - method.moe_quant_config = None - method.moe_kernel = None + monkeypatch.setattr(current_platform, "is_device_capability_family", lambda _: True) + select_backend = Mock(return_value=(object(), object)) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.online.nvfp4.select_nvfp4_moe_backend", + select_backend, + ) + moe_config = SimpleNamespace(is_act_and_mul=is_act_and_mul) + method = Nvfp4OnlineMoEMethod(moe=moe_config) + select_backend.assert_called_once_with( + config=moe_config, + weight_key=kNvfp4Static, + activation_key=kNvfp4DynamicToken, + ) layer = Mock() converted_weights = [ @@ -274,7 +287,12 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( assert method.moe_kernel is kernel assert convert_weights.call_count == 3 + assert all( + call.kwargs["trtllm_hidden_alignment"] == expected_hidden_alignment + for call in convert_weights.call_args_list + ) make_kernel.assert_called_once() + assert make_kernel.call_args.kwargs["per_token_activation"] is True get_quant_config.assert_called_once() assert process_weights.call_count == 3 diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 1d6ced0d48f4..7e2e21a3acbf 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -26,6 +26,8 @@ kInt4Static32, kInt8StaticChannelSym, kMxfp8Dynamic, + kNvfp4DynamicToken, + kNvfp4Static, ) # ---- QuantSpec ------------------------------------------------------------ @@ -68,13 +70,17 @@ def test_args_linear_string_resolves_via_quant_key_names(): assert args.moe is None -def test_args_moe_string_resolves_via_online_shorthand(): +@pytest.mark.parametrize( + "shorthand,weight", + [("fp8_per_block", kFp8Static128BlockSym), ("nvfp4_per_token", kNvfp4Static)], +) +def test_args_moe_string_resolves_via_online_shorthand(shorthand, weight): # An online-shorthand name pulls the matching slot from _ONLINE_SHORTHANDS # (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the # same per-layer-kind spec the `--quantization fp8_per_block` shorthand # would). - args = quant_config_args(moe="fp8_per_block") - assert args.moe == QuantSpec(weight=kFp8Static128BlockSym) + args = quant_config_args(moe=shorthand) + assert args.moe == QuantSpec(weight=weight) def test_args_string_shorthand_missing_slot_raises(): @@ -84,9 +90,14 @@ def test_args_string_shorthand_missing_slot_raises(): quant_config_args(linear="int8_per_channel_weight_only") -def test_args_accepts_dict_form(): - args = quant_config_args(moe={"activation": "mxfp8"}) - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) +@pytest.mark.parametrize( + "activation,expected", + [("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)], +) +def test_args_accepts_dict_form(activation, expected): + args = quant_config_args(moe={"activation": activation}) + assert args.moe == QuantSpec(weight=None, activation=expected) + assert args.linear is None def test_targets_reject_non_string_keys(): @@ -110,22 +121,34 @@ def test_resolve_colliding_shorthand_is_deferred(quantization: str): assert resolve_quantization_config(quantization, None) is None -def test_resolve_int8_shorthand_leaves_linear_unset(): - # int8_per_channel_weight_only is MoE-only; linear stays None so that - # OnlineQuantizationConfig leaves Linear layers in full precision. - args = resolve_quantization_config("int8_per_channel_weight_only", None) +@pytest.mark.parametrize( + "shorthand,weight", + [ + ("int8_per_channel_weight_only", kInt8StaticChannelSym), + ("nvfp4_per_token", kNvfp4Static), + ], +) +def test_resolve_moe_only_shorthand_leaves_linear_unset(shorthand, weight): + args = resolve_quantization_config(shorthand, None) assert args is not None assert args.linear is None - assert args.moe == QuantSpec(weight=kInt8StaticChannelSym) + assert args.moe == QuantSpec(weight=weight) -def test_resolve_quantization_config_only(): +@pytest.mark.parametrize("quantization", [None, "modelopt_fp4"]) +@pytest.mark.parametrize( + "activation,expected", + [("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)], +) +def test_resolve_activation_override(quantization, activation, expected): # When only `quantization_config` is given (e.g. for an already-quantized # checkpoint that needs an activation override), it's returned as-is. - args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}}) + args = resolve_quantization_config( + quantization, {"moe": {"activation": activation}} + ) assert args is not None assert args.linear is None - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) + assert args.moe == QuantSpec(weight=None, activation=expected) def test_resolve_merges_explicit_over_shorthand(): diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index d9500847268b..face2c3a7ba7 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -27,6 +27,7 @@ kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) @@ -40,6 +41,7 @@ "fp8_per_block_dynamic": kFp8Dynamic128Sym, "mxfp8": kMxfp8Dynamic, "mxfp4": kMxfp4Dynamic, + "nvfp4_per_token": kNvfp4DynamicToken, "int8_per_channel_static": kInt8StaticChannelSym, # Load-time MXFP4-to-int4 for Kimi-K3 on gfx942. "int4_per_group_32": kInt4Static32, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 330e7b6aa13a..80aa9122f447 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -334,6 +334,7 @@ def convert_to_nvfp4_moe_kernel_format( a2_scale: torch.Tensor | None, is_act_and_mul: bool, use_a16: bool = False, + trtllm_hidden_alignment: int = 256, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -419,6 +420,7 @@ def convert_to_nvfp4_moe_kernel_format( w2_scale_2=w2_scale_2, a2_scale=a2_scale, is_act_and_mul=is_act_and_mul, + trtllm_hidden_alignment=trtllm_hidden_alignment, ) elif nvfp4_backend == NvFp4MoeBackend.HUMMING: from vllm.model_executor.layers.quantization.utils.humming import ( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e33226cad3c8..98af93c30fad 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,7 +12,7 @@ import vllm.envs as envs from vllm.config import get_current_vllm_config_or_none -from vllm.config.quantization import QuantSpec +from vllm.config.quantization import QuantizationConfigArgs, QuantSpec from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, @@ -63,6 +63,9 @@ ) from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + trtllm_nvfp4_hidden_alignment, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( FP8_SCALE_SENTINEL, process_fp8_input_tensor_strategy_moe, @@ -90,6 +93,7 @@ kMxfp8Dynamic, kMxfp8Static, kNvfp4Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( @@ -724,12 +728,18 @@ def __init__( kv_cache_quant_algo: str | None = None, exclude_modules: list[str] | None = None, group_size: int = 16, + *, + quantization_args: QuantizationConfigArgs | None = None, ) -> None: if exclude_modules is None: exclude_modules = [] super().__init__(exclude_modules) self.quant_method = quant_method self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + moe_spec = quantization_args.moe if quantization_args is not None else None + self.moe_activation_override: QuantKey | None = ( + moe_spec.activation if moe_spec is not None else None + ) if is_checkpoint_nvfp4_serialized: self.group_size = group_size self.kv_cache_quant_algo = kv_cache_quant_algo @@ -815,6 +825,7 @@ def _from_config( kv_cache_quant_method, exclude_modules, group_size, + quantization_args=original_config.get("_online_quantization_args"), ) @@ -835,18 +846,36 @@ def __init__( ) -> None: super().__init__(moe_config) self.quant_config = quant_config - # W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With - # activation_key=None every W4A4 backend's _supports_quant_scheme - # rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic) - # exactly); only Marlin survives. Marlin's MoE path drops - # activation scales in convert_to_nvfp4_moe_kernel_format, so no - # other change is needed. self.use_a16 = quant_config.quant_method == "W4A16_NVFP4" + activation_key = None if self.use_a16 else kNvfp4Dynamic + + if quant_config.moe_activation_override is not None: + if self.use_a16: + raise ValueError( + "NVFP4 activation overrides require a W4A4 checkpoint." + ) + activation_key = quant_config.moe_activation_override + if activation_key not in (kNvfp4Dynamic, kNvfp4DynamicToken): + raise ValueError( + "Unsupported ModelOpt NVFP4 MoE activation override: " + f"{activation_key}. Use quantization_config.moe.activation=" + "'nvfp4_per_token' or omit the override." + ) + self.per_token_activation = activation_key == kNvfp4DynamicToken + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=None if self.use_a16 else kNvfp4Dynamic, + activation_key=activation_key, ) + if ( + self.per_token_activation + and self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM + ): + raise ValueError( + "Per-token NVFP4 activation for pre-quantized weights " + "requires the FlashInfer TRTLLM MoE backend." + ) self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( self.nvfp4_backend @@ -984,6 +1013,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: self._build_moe_kernel(layer) return + if self.per_token_activation: + layer.w13_input_scale.data.fill_(1.0) + layer.w2_input_scale.data.fill_(1.0) + # Use a single gscale for w13. if self.moe.is_act_and_mul and not torch.allclose( layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1] @@ -1016,6 +1049,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, use_a16=self.use_a16, + trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment( + per_token_activation=self.per_token_activation, + is_act_and_mul=self.moe.is_act_and_mul, + ), ) replace_parameter(layer, "w13_weight", w13) @@ -1039,6 +1076,7 @@ def _build_moe_kernel(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + per_token_activation=self.per_token_activation, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -1588,6 +1626,7 @@ def _from_config( group_size: int | None, **kwargs: Any, ) -> "ModelOptMixedPrecisionConfig": + quantization_args = original_config.get("_online_quantization_args") if "quantization" in original_config: quantized_layers = original_config["quantization"].get( "quantized_layers", {} @@ -1626,6 +1665,7 @@ def _from_config( kv_cache_quant_algo=kv_cache_quant_method, exclude_modules=[], group_size=group_size, + quantization_args=quantization_args, ) # Sibling config for layers that declare quant_algo: "W4A16_NVFP4". # get_quant_method resolves this sub-config to the (kNvfp4Static, None) @@ -1638,6 +1678,7 @@ def _from_config( kv_cache_quant_algo=kv_cache_quant_method, exclude_modules=[], group_size=group_size, + quantization_args=quantization_args, ) mxfp8_config = ModelOptMxFp8Config( diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index efc97705c538..6c780c2c46aa 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -19,6 +19,9 @@ from vllm.model_executor.layers.quantization.online.moe_base import ( OnlineMoEMethodBase, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + trtllm_nvfp4_hidden_alignment, +) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( FLOAT4_E2M1_MAX, ) @@ -87,6 +90,8 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): (SM100) only. """ + activation_quant_key = kNvfp4DynamicToken + def __init__( self, *, @@ -100,9 +105,13 @@ def __init__( self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=kNvfp4DynamicToken, + activation_key=self.activation_quant_key, ) + @property + def per_token_activation(self) -> bool: + return self.activation_quant_key == kNvfp4DynamicToken + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return @@ -156,6 +165,10 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: w2_scale_2=layer.w2_weight_scale_2, a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, + trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment( + per_token_activation=self.per_token_activation, + is_act_and_mul=self.moe.is_act_and_mul, + ), ) replace_parameter(layer, "w13_weight", w13) @@ -176,7 +189,7 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), - per_token_activation=True, + per_token_activation=self.per_token_activation, ) else: # Reload creates new scale tensors; derived kernel scales must use diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 24b95ef89aee..c5a299eea6e9 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -344,6 +344,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( w2_scale_2: torch.Tensor, a2_scale: torch.Tensor, is_act_and_mul: bool, + trtllm_hidden_alignment: int = 256, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -398,7 +399,11 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi( - w13, w13_scale, w2, w2_scale + w13, + w13_scale, + w2, + w2_scale, + min_alignment=trtllm_hidden_alignment, ) if layer.moe_config.hidden_dim_unpadded is None: layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index be28de9b66d6..5a6ee764286b 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -280,6 +280,20 @@ def align_fp4_moe_weights_for_fi( return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate +def trtllm_nvfp4_hidden_alignment( + per_token_activation: bool, is_act_and_mul: bool +) -> int: + """Hidden-dim alignment for the TRTLLM-Gen NVFP4 MoE cubins. + + The per-token (dynamic activation scale) non-gated variant has far fewer + tile configs for 256-aligned K and its heuristic default tactic fails with + "No valid config found" (e.g. Nemotron's 2688 -> 2816), so pad to 512. + """ + if per_token_activation and not is_act_and_mul: + return 512 + return 256 + + def align_fp4_moe_hidden_dim_for_fi( w13: torch.Tensor, w13_scale: torch.Tensor, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 72053805cd8b..535c4de96205 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -28,6 +28,7 @@ import vllm.version from vllm.config import ModelConfig +from vllm.config.utils import normalize_value from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_cache.utils import ( format_socket_role_suffix, @@ -227,12 +228,17 @@ def verify_peer_is_owner(conn: socket.socket) -> None: raise PermissionError(f"Rejecting weight cache connection from uid {peer_uid}") -def _hash_quant_config(quant_config: Any) -> str: - if quant_config is None: - return "" +def _hash_quant_config(quant_config: Any, runtime_quant_config: Any = None) -> str: if hasattr(quant_config, "to_dict"): quant_config = quant_config.to_dict() - payload = json.dumps(quant_config, sort_keys=True, default=str) + payload = json.dumps( + { + "checkpoint": quant_config, + "runtime": normalize_value(runtime_quant_config), + }, + sort_keys=True, + default=str, + ) return safe_hash(payload.encode(), usedforsecurity=False).hexdigest() @@ -333,7 +339,10 @@ def from_model_config( pp_rank=pp_rank, dtype=str(model_config.dtype), quantization=model_config.quantization, - quant_config_hash=_hash_quant_config(quant_config), + quant_config_hash=_hash_quant_config( + quant_config=quant_config, + runtime_quant_config=model_config.quantization_config, + ), revision=model_config.revision, vllm_version=vllm.version.__version__, is_draft=is_draft, diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index c6a72c4c33c6..ac1da8d6987b 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -200,6 +200,9 @@ def get_quant_config( QuantizationConfigArgs, resolve_quantization_config, ) + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptQuantConfigBase, + ) from vllm.model_executor.layers.quantization.online.base import ( OnlineQuantizationConfig, ) @@ -220,6 +223,14 @@ def maybe_compose_online_quantization( ) return checkpoint_config + def make_checkpoint_config(config: dict[str, Any]) -> QuantizationConfig: + if issubclass(quant_cls, ModelOptQuantConfigBase): + assert online_args is None or isinstance( + online_args, QuantizationConfigArgs + ) + config = {**config, "_online_quantization_args": online_args} + return maybe_compose_online_quantization(quant_cls.from_config(config)) + # Read the quantization config from the HF model config, if available. hf_quant_config = getattr(model_config.hf_config, "quantization_config", None) # some vision model may keep quantization_config in their text_config @@ -260,9 +271,7 @@ def maybe_compose_online_quantization( ): pass # fall through to file-based loading below else: - return maybe_compose_online_quantization( - quant_cls.from_config(hf_quant_config) - ) + return make_checkpoint_config(hf_quant_config) # if hf_quant_config is None, we will try to get config from # hf_overrides @@ -376,14 +385,14 @@ def maybe_compose_online_quantization( if model_config.quantization in ("modelopt", "modelopt_mixed"): if config.get("producer", {}).get("name") == "modelopt": - return maybe_compose_online_quantization(quant_cls.from_config(config)) + return make_checkpoint_config(config) else: raise ValueError( f"Unsupported quantization config" f" found for {model_config.quantization} in {f}." ) - return maybe_compose_online_quantization(quant_cls.from_config(config)) + return make_checkpoint_config(config) def get_sparse_attention_config( diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index be8846fc79c5..f3aaf72d9b56 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -110,6 +110,11 @@ def _get_nvfp4_config(self) -> ModelOptNvFp4Config: kv_cache_quant_algo=None, exclude_modules=[], group_size=16, + quantization_args=( + self.online_quantization_config.args + if self.online_quantization_config is not None + else None + ), ) return self._nvfp4_config diff --git a/vllm/models/deepseek_v41/quant_config.py b/vllm/models/deepseek_v41/quant_config.py index bed60ae75874..384e9488793c 100644 --- a/vllm/models/deepseek_v41/quant_config.py +++ b/vllm/models/deepseek_v41/quant_config.py @@ -114,6 +114,11 @@ def _get_nvfp4_config(self) -> ModelOptNvFp4Config: kv_cache_quant_algo=None, exclude_modules=[], group_size=16, + quantization_args=( + self.online_quantization_config.args + if self.online_quantization_config is not None + else None + ), ) return self._nvfp4_config From d169c5a90354233d2acec687aec201971fa135fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 00:16:05 +0200 Subject: [PATCH 151/952] [Docs] Add disagg PD flow-example guidance to /pr-checklist skill (#60111) Signed-off-by: NickLucche Co-authored-by: Mistral Vibe --- .agents/skills/pr-checklist/SKILL.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.agents/skills/pr-checklist/SKILL.md b/.agents/skills/pr-checklist/SKILL.md index 4d73be6adf3e..ae3099b6a076 100644 --- a/.agents/skills/pr-checklist/SKILL.md +++ b/.agents/skills/pr-checklist/SKILL.md @@ -181,6 +181,8 @@ Summarize model evaluations and benchmarks in compact tables with the configurat When fixing bugs or issues, provide a clear root-cause analysis. Explain what caused the problem, how it was identified, and why the chosen solution effectively addresses it. A patch which resolves the issue without a proper root-cause analysis may lead to recurring problems and increased maintenance overhead. +For fixes that involve disaggregated prefill/decode (PD) logic, such as the NIXL connector (`vllm/distributed/kv_transfer/kv_connector/v1/nixl/`), make a best effort to include a minimal example of the flow being changed and/or the initial flow that caused the failure. A short numbered sequence (or diagram) of a request's path across the P and D instances, e.g. scheduler admission, prefill execution, KV transfer registration and write, decode-side read, token generation, lets reviewers trace the root cause and the fix without reconstructing the distributed flow from the diff. Keep the example minimal: only the steps and components involved in the change. See the [disaggregated prefill](../../../docs/features/disagg_prefill.md), [NIXL KV push](../../../docs/design/nixl_kv_push_connector.md), and [NIXL KV cache lease](../../../docs/design/nixl_kv_cache_lease.md) docs for the flows involved. + #### 4.4: Implementation Details Provide a detailed explanation of the implementation, including design decisions, trade-offs, and any relevant technical details. Highlight any non-obvious choices, tradeoffs, potential limitations, and areas where maintainer attention is required. From ef63c23d35acccbba8e014fc88da19d541da38a9 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Wed, 7 Oct 2026 15:28:42 -0700 Subject: [PATCH 152/952] [Bugfix][DSv4.1] Skip SWA bounded replay for KV loads that carry the window (#59197) Signed-off-by: Yifan Qiao Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 --- .../aux_output_connector/test_store.py | 16 +++++ tests/v1/core/test_prefix_caching.py | 4 ++ tests/v1/core/test_prefix_replay.py | 27 ++++++++ .../unit/test_nixl_connector_hma.py | 35 +++++++--- tests/v1/kv_connector/unit/utils.py | 3 + .../aux_output_connector/connector.py | 6 +- .../v1/mooncake/mooncake_connector.py | 19 ++++- .../kv_connector/v1/nixl/base_scheduler.py | 13 +++- vllm/v1/core/sched/scheduler.py | 69 ++++++++++++++----- 9 files changed, 161 insertions(+), 31 deletions(-) diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py index 39277f3e7e0e..24d67c3d60e1 100644 --- a/tests/distributed/aux_output_connector/test_store.py +++ b/tests/distributed/aux_output_connector/test_store.py @@ -96,6 +96,10 @@ class _SchedulerRequest: default_factory=lambda: SimpleNamespace(routed_experts_prompt_start=0) ) + @property + def num_prompt_tokens(self) -> int: + return self.num_tokens - self.num_output_tokens + def is_finished(self) -> bool: return self.finished @@ -1439,6 +1443,18 @@ def test_scheduler_starts_worker_output_at_requested_prompt_token(): assert metadata.requests == {request.request_id: 2} + # A P/D prefiller cuts its prompt short of the last token, which can pass + # the requested start; the rows then start at the end of the cut prompt. + connector = _make_connector() + request = _scheduler_request( + "cut", [b"a" * 32], num_tokens=3, num_output_tokens=0, prompt_start=4 + ) + metadata = connector.build_connector_meta( + _step_output([request.request_id], [0], [3]), + {request.request_id: request}, + ) + assert metadata.requests == {request.request_id: 3} + def test_scheduler_connector_preserves_request_finish_order(): connector = _make_connector() diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 3972fe5230da..427201a7ceef 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -594,8 +594,12 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens): failed_recving_kv_req_ids=set(), finished_recving_kv_req_ids={request.request_id}, prefix_replay_tokens=0, + prefix_replay_group_ids=(), ) scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler) + scheduler._load_restores_replay_window = MethodType( + Scheduler._load_restores_replay_window, scheduler + ) Scheduler._update_waiting_for_remote_kv(scheduler, request) assert request.num_tokens - request.num_computed_tokens == 1 assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py index da5718580848..323413a86dce 100644 --- a/tests/v1/core/test_prefix_replay.py +++ b/tests/v1/core/test_prefix_replay.py @@ -247,6 +247,33 @@ def test_remote_kv_hit_is_taken_in_whole_blocks(): assert swa_blocks[new_req.replay_start // BLOCK_SIZE] is not swa_manager._null_block +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("matched", [WINDOW, 4 * BLOCK_SIZE - 1]) +def test_kv_load_carrying_window_does_not_replay(is_async, matched): + """A load of the request's own blocks (P/D) restores every transfer group, + window included, so the hit is taken as is: no replay, no cut to whole + blocks, and a hit no longer than the window is kept.""" + scheduler = _replay_scheduler( + use_kv_connector=MockKVConfig(matched_tokens=matched, is_async=is_async) + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.connector.get_loaded_kv_cache_group_ids = lambda request: (FULL, SWA) + scheduler.add_request(request) + out = scheduler.schedule() + if is_async: + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler.update_from_output( + out, create_model_runner_output([], finished_recving={request.request_id}) + ) + out = scheduler.schedule() + new_req = _new_req_data(out, request) + assert new_req.num_computed_tokens == matched + assert new_req.replay_start == 0 + assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS - matched + + @pytest.mark.parametrize("via_connector", [False, True]) def test_hit_no_longer_than_window_is_ignored(via_connector): """Such a hit would be recomputed in full anyway. It is not adopted (an diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index f96595572162..b4b5e063445f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1559,17 +1559,22 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks(): @pytest.mark.cpu_test @pytest.mark.parametrize( - "has_mamba,is_hma_required,expected_count", + "has_mamba,is_hma_required,bounded_replay,expected_count", [ - (True, True, 9), - (False, False, 10), - (False, True, 10), + (True, True, False, 9), + (False, False, False, 10), + (False, True, False, 10), + (False, True, True, 9), ], - ids=["mamba", "fa_only", "swa_only"], + ids=["mamba", "fa_only", "swa_only", "swa_bounded_replay"], ) -def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count): - """D-side: Mamba gets N-1 matched tokens, non-Mamba gets N.""" - sched = make_nixl_scheduler(has_mamba=has_mamba, is_hma_required=is_hma_required) +def test_mamba_n1_d_side(has_mamba, is_hma_required, bounded_replay, expected_count): + """D-side: Mamba and SWA bounded replay get N-1 matched tokens, others N.""" + sched = make_nixl_scheduler( + has_mamba=has_mamba, + is_hma_required=is_hma_required, + bounded_replay=bounded_replay, + ) req = create_request(num_tokens=10, do_remote_prefill=True) count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) @@ -1638,6 +1643,20 @@ def test_mamba_n1_p_side_truncation(): fa_sched.on_new_request(fa_req) assert len(fa_req.prompt_token_ids) == fa_original + # SWA bounded replay: the window is not replayed after the load, so the + # prefiller stops short of the last token too. + swa_sched = make_nixl_scheduler(is_hma_required=True, bounded_replay=True) + swa_req = create_request(num_tokens=10, do_remote_decode=True) + swa_sched.on_new_request(swa_req) + assert len(swa_req.prompt_token_ids) == 9 + + # A request that skips reading the prefix cache (prompt logprobs) is not + # cut: the decoder loads nothing and recomputes its whole prompt. + logprobs_req = create_request(num_tokens=10, do_remote_decode=True) + logprobs_req.sampling_params.skip_reading_prefix_cache = True + swa_sched.on_new_request(logprobs_req) + assert len(logprobs_req.prompt_token_ids) == 10 + @pytest.mark.cpu_test @pytest.mark.parametrize( diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 86062fc47848..d738e183952f 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -532,6 +532,7 @@ def make_kv_cache_config( def make_nixl_scheduler( has_mamba: bool = False, is_hma_required: bool = False, + bounded_replay: bool = False, heartbeat: bool = False, kv_lease_duration: int = 30, ): @@ -548,6 +549,7 @@ def make_nixl_scheduler( sched = object.__new__(NixlConnectorScheduler) sched._has_mamba = has_mamba + sched._bounded_replay = bounded_replay sched._is_hma_required = is_hma_required sched.kv_cache_config = make_kv_cache_config( block_size=16, @@ -615,6 +617,7 @@ def make_nixl_push_scheduler( sched.side_channel_port = 5600 sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled sched._has_mamba = has_mamba + sched._bounded_replay = False sched.kv_cache_config = make_kv_cache_config( block_size=16, mamba_enabled=has_mamba, diff --git a/vllm/distributed/aux_output_connector/connector.py b/vllm/distributed/aux_output_connector/connector.py index bd8cbf4d183d..c27e23ea4be6 100644 --- a/vllm/distributed/aux_output_connector/connector.py +++ b/vllm/distributed/aux_output_connector/connector.py @@ -67,8 +67,12 @@ def build_connector_meta( block_hashes_by_request[request_id] = packed self._sent_hash_counts[request_id] = len(request.block_hashes) assert request.sampling_params is not None + # A P/D prefiller may cut the prompt short of the requested start. scheduled_requests[request_id] = max( - request.sampling_params.routed_experts_prompt_start, + min( + request.sampling_params.routed_experts_prompt_start, + request.num_prompt_tokens, + ), 0 if request.num_output_tokens == 0 else request.num_tokens - 1, ) # A settled token can complete a hash block after the next async schedule diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 5f064259e8b2..8300cd950b0d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -897,6 +897,10 @@ def __init__( # GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec # path is currently tested with GDN; Mamba2 is not validated yet. self._has_mamba = kv_cache_config.has_mamba_layers + self._bounded_replay = any( + g.kv_cache_spec.prefix_replay_tokens > 0 + for g in kv_cache_config.kv_cache_groups + ) # Requests that need to start recv/send. # New requests are added by update_state_after_alloc in @@ -941,8 +945,10 @@ def get_sw_clipped_blocks( def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: + always recomputes the last token and must start from h(N-1), and for + DSV41 SWA bounded replay, whose window is not replayed after a P/D load + and must be laid out for that token.""" + if (self._has_mamba or self._bounded_replay) and num_prompt_tokens > 1: return num_prompt_tokens - 1 return num_prompt_tokens @@ -958,6 +964,9 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: params is not None and not params.get("_p_side_truncated") and request.num_prompt_tokens > 1 + # A request that skips reading the prefix cache (prompt logprobs) + # loads nothing on the decoder, which recomputes its whole prompt. + and not request.get_skip_reading_prefix_cache() ): if request.prompt_token_ids is not None: request.prompt_token_ids.pop() @@ -973,7 +982,11 @@ def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: def on_new_request(self, request: "Request") -> None: params = request.kv_transfer_params - if params is not None and params.get("do_remote_decode") and self._has_mamba: + if ( + params is not None + and params.get("do_remote_decode") + and (self._has_mamba or self._bounded_replay) + ): self._truncate_mamba_request_for_prefill(request) def get_num_new_matched_tokens( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 92f0cbc58a0b..55ee98d46596 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -101,6 +101,10 @@ def __init__( ) ) self._has_mamba = kv_cache_config.has_mamba_layers + self._bounded_replay = any( + g.kv_cache_spec.prefix_replay_tokens > 0 + for g in kv_cache_config.kv_cache_groups + ) logger.info("Initializing NIXL Scheduler %s", engine_id) if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: @@ -376,14 +380,16 @@ def _prefill_backoff(self) -> int: """Trailing prompt tokens the prefiller must not compute; the decoder recomputes them locally. - Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module + Mamba needs h(N-1) so the decoder can derive h(N) itself. DSV41's sliding + window under bounded replay is not replayed after a P/D load, so it must + be laid out for the token the decoder recomputes. Multi-module MTP needs to keep its whole lookahead window off of the prefiller, which would otherwise embed the unverified drafts in the MTP layer's KV cache. The decoder would never rebuild them, because the update is sized by the rejection count, which is zero for the first decode. """ return max( - 1 if self._has_mamba else 0, + 1 if self._has_mamba or self._bounded_replay else 0, self.vllm_config.num_prefill_lookahead_tokens - 1, ) @@ -412,6 +418,9 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: # Guard against repeated truncation after preemption/reschedule. and not params.get("_p_side_truncated") and request.num_prompt_tokens > backoff + # A request that skips reading the prefix cache (prompt logprobs) + # loads nothing on the decoder, which recomputes its whole prompt. + and not request.get_skip_reading_prefix_cache() ): if request.prompt_token_ids is not None: del request.prompt_token_ids[-backoff:] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f908a31a1b20..91ea8e93e42c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -290,10 +290,14 @@ def __init__( # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt # after a prefix hit by recomputing its trailing tokens. One window # for all such groups, so the rewind matches every group's allocation. - replay_windows = { - group.kv_cache_spec.prefix_replay_tokens - for group in kv_cache_config.kv_cache_groups + self.prefix_replay_group_ids = tuple( + group_id + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) if group.kv_cache_spec.prefix_replay_tokens > 0 + ) + replay_windows = { + kv_cache_config.kv_cache_groups[group_id].kv_cache_spec.prefix_replay_tokens + for group_id in self.prefix_replay_group_ids } assert len(replay_windows) <= 1, ( f"Prefix replay windows should agree: {sorted(replay_windows)}" @@ -973,14 +977,17 @@ def skip_request(from_queue: RequestQueue) -> None: skip_request(request_queue) continue - if self.prefix_replay_tokens: - # SWA bounded replay recomputes the hit's last + if self.prefix_replay_tokens and not ( + ext_tokens and self._load_restores_replay_window(request) + ): + # DSV41 SWA bounded replay recomputes the hit's last # window, from the block holding its first # token; the sliding-window groups retire # whole blocks below the window of the hit's # next token. Hits end on a block boundary, or # a hit ending one token short of one would - # retire the block the replay starts in. + # retire the block the replay starts in. A load + # that carries the window replays nothing. ext_tokens -= ext_tokens % self.block_size load_kv_async = load_kv_async and ext_tokens > 0 @@ -1026,9 +1033,17 @@ def skip_request(from_queue: RequestQueue) -> None: num_new_local_computed_tokens + num_external_computed_tokens ) assert num_computed_tokens <= request.num_tokens - if 0 < num_computed_tokens <= self.prefix_replay_tokens: - # SWA bounded replay: a hit no longer than the replayed - # window would be recomputed in full and save nothing. + # DSV41 SWA bounded replay: a KV load that carries the + # window is taken as is; any other hit no longer than the + # replayed window would be recomputed in full and save + # nothing. + window_loaded = bool( + num_external_computed_tokens + ) and self._load_restores_replay_window(request) + if ( + not window_loaded + and 0 < num_computed_tokens <= self.prefix_replay_tokens + ): new_computed_blocks = ( self.kv_cache_manager.empty_kv_cache_blocks ) @@ -1091,7 +1106,7 @@ def skip_request(from_queue: RequestQueue) -> None: if did_prefix_cache_lookup: # Fresh admission: local and/or sync external hit. num_replay_tokens = self._mark_prefix_replay( - request, num_computed_tokens + request, num_computed_tokens, window_loaded ) num_computed_tokens -= num_replay_tokens # Number of tokens to be scheduled. @@ -3083,16 +3098,31 @@ def _inflight_prefill_reserved_blocks(self) -> int: self._request_remaining_blocks(req) for req in self._inflight_prefills ) - def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int: + def _load_restores_replay_window(self, request: Request) -> bool: + """Whether the KV connector's load for ``request`` restores every DSV41 + SWA bounded-replay group, as a P/D transfer of the request's own blocks + does; a prefix-cache store holds no window.""" + assert self.connector is not None + return bool(self.prefix_replay_group_ids) and set( + self.prefix_replay_group_ids + ).issubset(self.connector.get_loaded_kv_cache_group_ids(request)) + + def _mark_prefix_replay( + self, request: Request, num_hit_tokens: int, window_loaded: bool + ) -> int: """Record where a prefix hit's replay starts and return the number of hit tokens to recompute; the caller rewinds the computed count by it. The replayed tokens are the hit's last window (see ``Request.replay_start``): the worker rebuilds their sliding-window KV - and leaves their cached KV alone. + and leaves their cached KV alone. A hit whose KV load carries the + window replays nothing. """ if not self.prefix_replay_tokens: return 0 + if window_loaded: + request.replay_start = 0 + return 0 num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens) request.replay_start = num_hit_tokens - num_replay_tokens return num_replay_tokens @@ -3106,7 +3136,8 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: """ assert self.connector is not None - if request.request_id in self.failed_recving_kv_req_ids: + load_failed = request.request_id in self.failed_recving_kv_req_ids + if load_failed: # Request had KV load failures; num_computed_tokens was already # updated in _update_requests_with_invalid_blocks if request.num_computed_tokens: @@ -3132,11 +3163,15 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: # This will cache the blocks iff caching is enabled. self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) - # SWA bounded replay recomputes the tail of the hit, which covers the - # last token; otherwise a full prompt hit re-computes that token so the - # next one can be sampled. + # DSV41 SWA bounded replay recomputes the tail of a hit loaded without + # its window, or of a failed load, which may have lost it; the replay + # covers the last token. Otherwise a full prompt hit re-computes that + # token so the next one can be sampled. num_replay_tokens = self._mark_prefix_replay( - request, request.num_computed_tokens + request, + request.num_computed_tokens, + window_loaded=not load_failed + and self._load_restores_replay_window(request), ) if num_replay_tokens > 0: request.num_computed_tokens -= num_replay_tokens From 4ac5a36a8e38b79c8084589b0bc4f2a249bf7c2a Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Thu, 8 Oct 2026 01:44:17 +0300 Subject: [PATCH 153/952] [CI/Build][NVIDIA] Build NIXL with NIXL EP from source in Rubin image (#60090) Signed-off-by: Artem Perevedentsev --- .buildkite/release-pipeline.yaml | 5 + docker/Dockerfile | 42 ++ docker/versions.json | 3 + docs/contributing/dockerfile/dockerfile.md | 5 + .../installation/gpu.cuda.inc.md | 5 + tools/build_nixl_with_torch215_for_rubin.py | 619 ++++++++++++++++++ 6 files changed, 679 insertions(+) create mode 100644 tools/build_nixl_with_torch215_for_rubin.py diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 5eba9fd2e114..6064f93104f5 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -355,6 +355,10 @@ steps: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + # TODO(arpera): + # Remove BUILD_NIXL=true once NIXL has a wheel built against torch 2.15 + # For details: https://github.com/vllm-project/vllm/pull/60090 + # Tracking issue report: https://github.com/ai-dynamo/nixl/issues/2297 - | DOCKER_BUILDKIT=1 docker buildx build --load \ $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \ @@ -366,6 +370,7 @@ steps: --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ --build-arg INSTALL_KV_CONNECTORS=true \ --build-arg INSTALL_RUBIN_PRERELEASE=true \ + --build-arg BUILD_NIXL=true \ --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg RUN_WHEEL_CHECK=false \ diff --git a/docker/Dockerfile b/docker/Dockerfile index d8a1258a74c5..fb8fa6b264ae 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -92,6 +92,11 @@ ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER} # Flag enables built-in KV-connector dependency libs into docker images ARG INSTALL_KV_CONNECTORS=false +# Flag builds NIXL, including its NIXL EP extension for the installed PyTorch, +# from source and installs it in place of the nixl packages from the KV-connector +# requirements +ARG BUILD_NIXL=false + #################### BASE BUILD IMAGE #################### # prepare basic build environment FROM ${BUILD_BASE_IMAGE} AS base @@ -548,6 +553,24 @@ RUN --mount=type=cache,target=/opt/uv/cache \ (cd "${TRITON_WORKSPACE}/dist" \ && sha256sum build-state > wheel.sha256); \ fi + +# Build NIXL and its NIXL EP extension against the PyTorch installed here when +# requested. The wheels are installed by the final image, which has no compiler. +ARG BUILD_NIXL +COPY tools/build_nixl_with_torch215_for_rubin.py /tmp/build_nixl_for_rubin.py +RUN --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \ + set -eux; \ + NIXL_WORKSPACE=/tmp/nixl_workspace; \ + mkdir -p "${NIXL_WORKSPACE}/dist"; \ + if [ "${BUILD_NIXL}" = "true" ]; then \ + dnf install -y --setopt=install_weak_deps=False autoconf automake libtool; \ + dnf clean all; \ + rm -rf /var/cache/dnf; \ + NIXL_REQUIREMENTS_FILE=/tmp/vllm-requirements/kv_connectors.txt \ + WHEELS_CACHE_HOME="${NIXL_WORKSPACE}/dist" \ + MAX_JOBS="${max_jobs:-}" \ + python3 /tmp/build_nixl_for_rubin.py --skip-system-deps --build-only; \ + fi #################### EXTENSIONS BUILD IMAGE #################### #################### WHEEL BUILD IMAGE #################### @@ -1260,6 +1283,25 @@ RUN --mount=type=cache,target=/opt/uv/cache \ /tmp/install_triton_from_source.sh \ /tmp/triton-source-wheel.sha256 +# Install source-built NIXL last, replacing the nixl packages installed above, so +# the image has a NIXL EP build for its PyTorch version. There is no GPU driver +# at build time, so check for the extension instead of importing it. +ARG BUILD_NIXL +RUN --mount=type=cache,target=/opt/uv/cache \ + --mount=type=bind,from=extensions-build,source=/tmp/nixl_workspace/dist,target=/tmp/nixl-dist,ro \ + set -eu; \ + if [ "${BUILD_NIXL}" = "true" ]; then \ + ls /tmp/nixl-dist/*.whl; \ + uv pip uninstall --system nixl nixl-cu12 nixl-cu13; \ + uv pip install --system --no-deps /tmp/nixl-dist/*.whl; \ + python3 -c "import glob, importlib.util as u, os, torch; \ + spec = u.find_spec('nixl_ep_cu' + os.environ['CUDA_VERSION'].split('.')[0]); \ + mm = ''.join(torch.__version__.split('.')[:2]); \ + ext = glob.glob(os.path.join(spec.submodule_search_locations[0], 'nixl_ep_cpp_torch' + mm + '.*')) if spec else []; \ + assert ext, 'NIXL EP has no build for torch ' + torch.__version__; \ + print('NIXL EP extension:', ext[0])"; \ + fi + ENV VLLM_USAGE_SOURCE production-docker-image ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \ VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \ diff --git a/docker/versions.json b/docker/versions.json index bcec27fd5590..68f426d0c7b5 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -37,6 +37,9 @@ "INSTALL_KV_CONNECTORS": { "default": "false" }, + "BUILD_NIXL": { + "default": "false" + }, "SCCACHE_BUCKET_NAME": { "default": "vllm-build-sccache" }, diff --git a/docs/contributing/dockerfile/dockerfile.md b/docs/contributing/dockerfile/dockerfile.md index db46540f9f2c..279c8f0c6a34 100644 --- a/docs/contributing/dockerfile/dockerfile.md +++ b/docs/contributing/dockerfile/dockerfile.md @@ -24,6 +24,11 @@ The `extensions-build` stage can also produce an optional source-built Triton wheel. `vllm-openai-base` installs that wheel after its other Python dependencies so dependency resolution cannot restore an older Triton version. +With `BUILD_NIXL=true`, `extensions-build` also builds NIXL, including its NIXL +EP extension for the installed PyTorch, with `tools/build_nixl_with_torch215_for_rubin.py`. +`vllm-openai-base` then replaces the NIXL packages installed from the KV-connector +requirements with these wheels. + >
> ![](../../assets/contributing/dockerfile-stages-dependency.png){ align="center" alt="query" width="100%" } >
diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 0d0d0c0c85ab..86e38c570788 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -439,6 +439,11 @@ BuildKit does not automatically invalidate cached layers when a mutable Git ref changes. Use `--no-cache-filter extensions-build` to refresh an empty, branch, or tag revision. +Set `BUILD_NIXL=true` to build NIXL from source. NIXL's release wheels do not +include the NIXL EP extension for the PyTorch nightly used by the Rubin build. +The build uses the NIXL version pinned in `requirements/kv_connectors.txt` and +replaces the NIXL packages installed from the KV-connector requirements. + For `FINAL_BASE_IMAGE`, use the public, multi-arch `nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or newer, the first NCCL release with Rubin (SM107) support. diff --git a/tools/build_nixl_with_torch215_for_rubin.py b/tools/build_nixl_with_torch215_for_rubin.py new file mode 100644 index 000000000000..ee9fa72ae959 --- /dev/null +++ b/tools/build_nixl_with_torch215_for_rubin.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# TODO(arpera): +# This script is temporary. It builds nixl from source code against torch 2.15 +# because currently there is no nixl wheel built against torch 2.15 which +# vLLM docker image for Rubin requires. +# Remove this script from vLLM once there is nixl wheel built against torch 2.15. + +import argparse +import glob +import os +import shutil +import subprocess +import sys +import urllib.request + +# --- Configuration --- +WHEELS_CACHE_HOME = os.environ.get("WHEELS_CACHE_HOME", "/tmp/wheels_cache") +ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) +UCX_DIR = os.path.join("/tmp", "ucx_source") +NIXL_DIR = os.path.join("/tmp", "nixl_source") +UCX_INSTALL_DIR = os.path.join("/tmp", "ucx_install") +NIXL_BUILD_DIR = os.path.join("/tmp", "nixl_build") +BUILD_TOOLS_DIR = os.path.join("/tmp", "nixl_build_tools") +UCX_REPO_URL = "https://github.com/openucx/ucx.git" +NIXL_REPO_URL = "https://github.com/ai-dynamo/nixl.git" +UCX_REF = os.environ.get("UCX_REF", "v1.23.x") +TEMP_WHEEL_DIR = os.path.join("/tmp", "nixl_temp_wheelhouse") +NIXL_REQUIREMENTS_FILE = os.environ.get( + "NIXL_REQUIREMENTS_FILE", + os.path.join(ROOT_DIR, os.pardir, "requirements", "kv_connectors.txt"), +) +MAX_JOBS = int(os.environ.get("MAX_JOBS") or os.cpu_count() or 1) +CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda") +# SM targets NIXL EP can run on (it requires sm_90 or newer). +NIXL_EP_CUDA_ARCHS = ("90", "100", "103", "107", "110", "120") + +# Libraries that must come from the host or from torch, not from the wheel. +AUDITWHEEL_EXCLUDES = [ + "libcuda*", + "libcufile*", + "libcuobjclient*", + "libssl*", + "libcrypto*", + "libefa*", + "libhwloc*", + "libfabric*", + "libtorch*", + "libc10*", + "libdoca*", + "libred_client*", + "libred_async*", + "liblz4*", +] + + +# --- Helper Functions --- +def get_latest_nixl_version(): + """Helper function to get latest release version of NIXL. + + Raises: + RuntimeError: If the latest release cannot be determined. + + """ + # /releases/latest redirects to the tag page. Unlike the REST API it is not + # rate limited for anonymous clients, which are shared on compute nodes. + request = urllib.request.Request( + f"{NIXL_REPO_URL.removesuffix('.git')}/releases/latest", method="HEAD" + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return response.geturl().rsplit("/", 1)[-1] + except OSError as e: + raise RuntimeError( + "Cannot determine the latest NIXL release, set NIXL_VERSION." + ) from e + + +def get_pinned_nixl_version(): + """Returns the NIXL version pinned in the vLLM requirements, if there is one.""" + try: + with open(NIXL_REQUIREMENTS_FILE) as requirements: + for line in requirements: + requirement = line.split("#")[0].split(";")[0] + name, separator, version = requirement.partition("==") + if separator and name.strip() == "nixl" and version.split(): + return version.split()[0] + except OSError: + pass + return None + + +def resolve_nixl_version(): + """Returns the NIXL version to build and where it comes from.""" + if os.environ.get("NIXL_VERSION"): + return os.environ["NIXL_VERSION"], "the NIXL_VERSION environment variable" + pinned_version = get_pinned_nixl_version() + if pinned_version: + return pinned_version, os.path.normpath(NIXL_REQUIREMENTS_FILE) + return get_latest_nixl_version(), "the latest NIXL release" + + +NIXL_VERSION, NIXL_VERSION_SOURCE = resolve_nixl_version() +# Version as it appears in wheel file names (tags may carry a "v" prefix). +NIXL_WHEEL_VERSION = NIXL_VERSION.removeprefix("v") + + +def run_command(command, cwd=".", env=None): + """Helper function to run a shell command and check for errors.""" + print(f"--> Running command: {' '.join(command)} in '{cwd}'", flush=True) + subprocess.check_call(command, cwd=cwd, env=env) + + +def is_pip_package_installed(package_name): + """Checks if a package is installed via pip without raising an exception.""" + result = subprocess.run( + [sys.executable, "-m", "pip", "show", package_name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def find_nixl_wheel_in_cache(cache_dir): + """Finds the nixl-cuXX wheel file in the specified cache directory.""" + # The repaired wheel will have a 'manylinux' tag, but this glob still works. + search_pattern = os.path.join(cache_dir, f"nixl_*{NIXL_WHEEL_VERSION}*.whl") + wheels = glob.glob(search_pattern) + if wheels: + # Sort to get the most recent/highest version if multiple exist + wheels.sort() + return wheels[-1] + return None + + +def find_nixl_meta_wheel(wheel_dir): + """Finds the `nixl` meta package wheel (the nixl_cuXX dispatcher).""" + wheels = glob.glob(os.path.join(wheel_dir, f"nixl-{NIXL_WHEEL_VERSION}-*.whl")) + return wheels[-1] if wheels else None + + +def get_cuda_version(): + """Returns the (major, minor) CUDA toolkit version, or None without CUDA.""" + nvcc = os.path.join(CUDA_HOME, "bin", "nvcc") + if not os.path.exists(nvcc): + return None + output = subprocess.check_output([nvcc, "--version"], text=True) + _, separator, release = output.partition("release ") + try: + if not separator: + raise ValueError + major, minor = release.split(",")[0].split(".")[:2] + return int(major), int(minor) + except ValueError: + raise RuntimeError( + f"Cannot parse the CUDA version from `{nvcc} --version`." + ) from None + + +def get_default_cuda_arch_list(): + """Returns the SM targets of NIXL EP that the installed nvcc can compile.""" + nvcc = os.path.join(CUDA_HOME, "bin", "nvcc") + supported = subprocess.check_output([nvcc, "--list-gpu-code"], text=True).split() + return ",".join(sm for sm in NIXL_EP_CUDA_ARCHS if f"sm_{sm}" in supported) + + +def is_torch_available(): + """Checks whether torch can be imported by the current interpreter.""" + result = subprocess.run( + [sys.executable, "-c", "import torch"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def is_nixl_installed(build_nixl_ep): + """Checks that the requested NIXL version, and its EP extension for the + installed torch when requested, is already installed.""" + cuda_version = get_cuda_version() + if cuda_version is None: + return is_pip_package_installed("nixl") + checks = [ + "import importlib, importlib.metadata as m", + f"assert m.version('nixl') == '{NIXL_WHEEL_VERSION}'", + f"assert m.version('nixl-cu{cuda_version[0]}') == '{NIXL_WHEEL_VERSION}'", + ] + if build_nixl_ep: + # The package picks its extension by torch version and fails to import + # when none was built for the installed torch. + checks.append(f"importlib.import_module('nixl_ep_cu{cuda_version[0]}')") + result = subprocess.run( + [sys.executable, "-c", "\n".join(checks)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return result.returncode == 0 + + +def install_system_dependencies(build_nixl_ep=False): + """Installs required system packages using apt-get if run as root.""" + if shutil.which("apt-get") is None: + print("--> apt-get not found, relying on the build tools in PATH.", flush=True) + return + if os.geteuid() != 0: + print("\n---", flush=True) + print( + "WARNING: Not running as root. \ + Skipping system dependency installation.", + flush=True, + ) + print( + "Please ensure the listed packages are installed on your system:", + flush=True, + ) + print( + " patchelf build-essential git cmake ninja-build \ + autotools-dev automake libtool libtool-bin \ + libibverbs-dev", + flush=True, + ) + print("---\n", flush=True) + return + + print("--- Running as root. Installing system dependencies... ---", flush=True) + apt_packages = [ + "patchelf", # <-- Add patchelf here + "build-essential", + "git", + "cmake", + "ninja-build", + "autotools-dev", + "automake", + "libtool", + "libtool-bin", + "pkg-config", + "libibverbs-dev", + ] + # -base CUDA images lack the headers UCX (NVML) and the torch extension + # (cuSPARSE/cuSOLVER/cuBLAS, included by ATen/cuda/CUDAContext.h) need. + cuda_version = get_cuda_version() + if cuda_version is not None: + major, minor = cuda_version + if not os.path.exists(os.path.join(CUDA_HOME, "include", "nvml.h")): + apt_packages.append(f"cuda-nvml-dev-{major}-{minor}") + if build_nixl_ep: + apt_packages += [ + f"{package}-{major}-{minor}" + for package in ("libcusparse-dev", "libcublas-dev", "libcusolver-dev") + ] + run_command(["apt-get", "update"]) + run_command(["apt-get", "install", "-y"] + apt_packages) + print("--- System dependencies installed successfully. ---\n", flush=True) + + +def install_python_build_dependencies(): + """Installs the Python tools needed to build and repair the NIXL wheel.""" + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "auditwheel", + "meson-python", + "ninja", + "patchelf", + "pybind11", + "pyyaml", + "setuptools>=80.9.0", + "tomlkit", + ] + ) + # An older meson already on the system (e.g. from the distribution) makes + # meson-python fail to package NIXL and pip would keep it, so use our own. + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "--upgrade", + f"--target={BUILD_TOOLS_DIR}", + "meson", + ] + ) + + +def build_ucx(ucx_install_path): + """Builds UCX from source, with CUDA support if CUDA is installed.""" + if not os.path.exists(UCX_DIR): + run_command(["git", "clone", UCX_REPO_URL, UCX_DIR]) + ucx_source_path = os.path.abspath(UCX_DIR) + run_command(["git", "checkout", UCX_REF], cwd=ucx_source_path) + run_command(["./autogen.sh"], cwd=ucx_source_path) + configure_command = [ + "./contrib/configure-release-mt", + f"--prefix={ucx_install_path}", + "--enable-shared", + "--disable-static", + "--disable-doxygen-doc", + "--enable-experimental-api", + "--enable-optimizations", + "--without-avx", + "--enable-cma", + "--enable-devel-headers", + *([f"--with-cuda={CUDA_HOME}"] if get_cuda_version() else []), + "--with-verbs", + "--without-gdrcopy", + "--without-dc", + "--without-rdmacm", + "--without-gga", + "--with-ze=no", + ] + run_command(configure_command, cwd=ucx_source_path) + run_command(["make", "-j", str(MAX_JOBS)], cwd=ucx_source_path) + run_command(["make", "install"], cwd=ucx_source_path) + + +def build_nixl_wheel(ucx_path, build_env, wheel_dir, args): + """Builds the (unrepaired) NIXL wheel against the UCX installed in ucx_path.""" + if not os.path.exists(NIXL_DIR): + run_command(["git", "clone", NIXL_REPO_URL, NIXL_DIR]) + else: + run_command(["git", "fetch", "--tags"], cwd=NIXL_DIR) + nixl_source_path = os.path.abspath(NIXL_DIR) + # Tags are "v" for current releases and bare for older ones. + for tag in (f"v{NIXL_WHEEL_VERSION}", NIXL_WHEEL_VERSION): + tag_exists = subprocess.run( + ["git", "rev-parse", "--verify", "--quiet", f"refs/tags/{tag}"], + cwd=nixl_source_path, + stdout=subprocess.DEVNULL, + ) + if tag_exists.returncode == 0: + run_command(["git", "checkout", "--force", tag], cwd=nixl_source_path) + print(f"--> Checked out NIXL version: {tag}", flush=True) + break + else: + raise RuntimeError(f"NIXL has no tag for version {NIXL_WHEEL_VERSION}.") + + cuda_version = get_cuda_version() + if cuda_version is not None: + # The wheel name selects which of nixl-cu12 / nixl-cu13 is built. + run_command( + [ + sys.executable, + "contrib/tomlutil.py", + "--wheel-name", + f"nixl-cu{cuda_version[0]}", + "pyproject.toml", + ], + cwd=nixl_source_path, + ) + + meson_options = { + "build_tests": "false", + "ucx_path": ucx_path, + } + cuda_arch_list = args.cuda_arch_list + if args.build_nixl_ep: + meson_options["build_nixl_ep"] = "true" + meson_options["build_examples"] = "true" + cuda_arch_list = cuda_arch_list or get_default_cuda_arch_list() + if cuda_arch_list: + meson_options["nixl_cuda_arch_list"] = cuda_arch_list + if args.nixl_plugins: + meson_options["enable_plugins"] = args.nixl_plugins + config_settings = [f"-Csetup-args=-D{k}={v}" for k, v in meson_options.items()] + # Keep the build tree: auditwheel needs the subproject libraries it holds. + config_settings.append(f"-Cbuild-dir={NIXL_BUILD_DIR}") + config_settings.append(f"-Ccompile-args=-j{MAX_JOBS}") + + run_command( + [ + sys.executable, + "-m", + "pip", + "wheel", + ".", + "--no-deps", + # Build against the torch already installed, not a fresh isolated one. + "--no-build-isolation", + f"--wheel-dir={wheel_dir}", + *config_settings, + ], + cwd=nixl_source_path, + env=build_env, + ) + + +def repair_nixl_wheel(unrepaired_wheel, repaired_dir, ucx_install_path, build_env): + """Bundles the shared libraries and UCX modules into a self-contained wheel.""" + auditwheel_command = ["auditwheel", "repair"] + for pattern in AUDITWHEEL_EXCLUDES: + auditwheel_command += ["--exclude", pattern] + auditwheel_command += [unrepaired_wheel, f"--wheel-dir={repaired_dir}"] + # The wheel ships prometheus-cpp as libcore.so.1.3.0 but its users ask for the + # soname libcore.so.1.3, which only exists as a symlink in the build tree. + repair_env = build_env.copy() + prometheus_dir = os.path.join(NIXL_BUILD_DIR, "subprojects", "prometheus-cpp") + repair_env["LD_LIBRARY_PATH"] = f"{prometheus_dir}:{build_env['LD_LIBRARY_PATH']}" + run_command(auditwheel_command, env=repair_env) + + repaired_wheel = find_nixl_wheel_in_cache(repaired_dir) + if not repaired_wheel: + raise RuntimeError("Failed to find the repaired NIXL wheel.") + # UCX loads its transports (CUDA, IB, ...) with dlopen, so auditwheel does not + # see them as dependencies; the NIXL plugin itself is already in the wheel. + run_command( + [ + sys.executable, + "contrib/wheel_add_ucx_plugins.py", + "--skip-nixl-plugins", + f"--ucx-plugins-dir={os.path.join(ucx_install_path, 'lib', 'ucx')}", + repaired_wheel, + ], + cwd=os.path.abspath(NIXL_DIR), + env=build_env, + ) + return repaired_wheel + + +def install_nixl_wheels(wheels): + """Installs the NIXL wheels and drops the nixl-cuXX of other CUDA majors.""" + # --force-reinstall: the wheels usually have the version of the release that + # is already installed, which pip would otherwise keep. + # w/o "no-deps", it will install cuda-torch + run_command( + [ + sys.executable, + "-m", + "pip", + "install", + "--force-reinstall", + "--no-deps", + *wheels, + ] + ) + + cuda_version = get_cuda_version() + if cuda_version is None: + return + # The dispatcher only loads nixl-cu, and the meta + # package built here does not depend on the other variant. + for cuda_major in (12, 13): + package = f"nixl-cu{cuda_major}" + if cuda_major != cuda_version[0] and is_pip_package_installed(package): + run_command([sys.executable, "-m", "pip", "uninstall", "-y", package]) + + +def build_and_install_prerequisites(args): + """Builds UCX and NIXL from source, creating a self-contained wheel.""" + print(f"--> NIXL version {NIXL_WHEEL_VERSION} (from {NIXL_VERSION_SOURCE})") + if args.build_nixl_ep is None: + args.build_nixl_ep = get_cuda_version() is not None and is_torch_available() + reuse_existing = not (args.force_reinstall or args.build_only) + if reuse_existing and is_nixl_installed(args.build_nixl_ep): + print("--> NIXL is already installed. Nothing to do.", flush=True) + return + + cached_wheel = find_nixl_wheel_in_cache(WHEELS_CACHE_HOME) + if reuse_existing and cached_wheel: + print( + f"\n--> Found self-contained wheel: \ + {os.path.basename(cached_wheel)}.", + flush=True, + ) + print("--> Installing from cache, skipping all source builds.", flush=True) + cached_wheels = [cached_wheel] + cached_meta_wheel = find_nixl_meta_wheel(WHEELS_CACHE_HOME) + if cached_meta_wheel: + cached_wheels.append(cached_meta_wheel) + install_nixl_wheels(cached_wheels) + print("\n--- Installation from cache complete. ---", flush=True) + return + + print( + "\n--> No installed package or cached wheel found. \ + Starting full build process...", + flush=True, + ) + if not args.skip_system_deps: + install_system_dependencies(args.build_nixl_ep) + print("\n--> Installing Python build dependencies...", flush=True) + install_python_build_dependencies() + print(f"--> Using wheel cache directory: {WHEELS_CACHE_HOME}", flush=True) + os.makedirs(WHEELS_CACHE_HOME, exist_ok=True) + + # -- Step 1: Build UCX from source, unless an existing install is given -- + # With an external UCX (e.g. HPC-X) the wheel keeps depending on it at runtime. + external_ucx = args.ucx_path is not None + if external_ucx: + ucx_install_path = os.path.abspath(args.ucx_path) + print(f"\n[1/3] Using the existing UCX in {ucx_install_path}", flush=True) + else: + ucx_install_path = os.path.abspath(UCX_INSTALL_DIR) + print("\n[1/3] Configuring and building UCX from source...", flush=True) + build_ucx(ucx_install_path) + print("--- UCX build and install complete ---", flush=True) + + # -- Step 2: Build NIXL wheel from source -- + print("\n[2/3] Building NIXL wheel from source...", flush=True) + ucx_lib_path = os.path.join(ucx_install_path, "lib") + ucx_plugin_path = os.path.join(ucx_lib_path, "ucx") + build_env = os.environ.copy() + # pip-installed build tools live next to the running interpreter. + build_env["PATH"] = os.pathsep.join( + [ + os.path.join(BUILD_TOOLS_DIR, "bin"), + os.path.join(CUDA_HOME, "bin"), + os.path.dirname(sys.executable), + build_env["PATH"], + ] + ) + build_env["PYTHONPATH"] = os.pathsep.join( + filter(None, [BUILD_TOOLS_DIR, build_env.get("PYTHONPATH")]) + ) + build_env["PKG_CONFIG_PATH"] = os.path.join(ucx_lib_path, "pkgconfig") + existing_ld_path = os.environ.get("LD_LIBRARY_PATH", "") + build_env["LD_LIBRARY_PATH"] = ( + f"{ucx_lib_path}:{ucx_plugin_path}:{existing_ld_path}".strip(":") + ) + print(f"--> Using LD_LIBRARY_PATH: {build_env['LD_LIBRARY_PATH']}", flush=True) + + build_nixl_wheel(ucx_install_path, build_env, TEMP_WHEEL_DIR, args) + + # -- Step 3: Repair the wheel by copying UCX libraries -- + unrepaired_wheel = find_nixl_wheel_in_cache(TEMP_WHEEL_DIR) + if not unrepaired_wheel: + raise RuntimeError("Failed to find the NIXL wheel after building it.") + if external_ucx: + print("\n[3/3] Skipping the repair step: UCX is provided externally.") + built_wheels = [unrepaired_wheel] + else: + print("\n[3/3] Repairing NIXL wheel to include UCX libraries...", flush=True) + built_wheels = [ + repair_nixl_wheel( + unrepaired_wheel, + os.path.join(TEMP_WHEEL_DIR, "repaired"), + ucx_install_path, + build_env, + ) + ] + + meta_wheel = find_nixl_meta_wheel( + os.path.join(NIXL_BUILD_DIR, "src", "bindings", "python", "nixl-meta") + ) + if meta_wheel: + built_wheels.append(meta_wheel) + else: + print("WARNING: the nixl meta package was not built.", flush=True) + + new_wheels = [] + for wheel in built_wheels: + run_command(["cp", wheel, WHEELS_CACHE_HOME]) + new_wheels.append(os.path.join(WHEELS_CACHE_HOME, os.path.basename(wheel))) + run_command(["rm", "-rf", TEMP_WHEEL_DIR]) + + wheel_names = ", ".join(os.path.basename(w) for w in new_wheels) + if args.build_only: + print(f"--> Built wheels in {WHEELS_CACHE_HOME}: {wheel_names}", flush=True) + return + print( + f"--> Successfully built wheels: {wheel_names}. Now installing...", flush=True + ) + install_nixl_wheels(new_wheels) + print("--- NIXL installation complete ---", flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Build and install UCX and NIXL dependencies." + ) + parser.add_argument( + "--force-reinstall", + action="store_true", + help="Force rebuild and reinstall of UCX and NIXL \ + even if they are already installed.", + ) + parser.add_argument( + "--build-only", + action="store_true", + help="Build the wheels into WHEELS_CACHE_HOME and stop, without \ + installing them or looking at what is already installed.", + ) + parser.add_argument( + "--skip-system-deps", + action="store_true", + help="Do not install system packages with apt-get; the build tools \ + (git, autotools, cmake, patchelf, ...) must already be on PATH.", + ) + parser.add_argument( + "--build-nixl-ep", + action=argparse.BooleanOptionalAction, + default=None, + help="Build the NIXL EP extension for the installed torch version. \ + Defaults to on when CUDA and torch are available.", + ) + parser.add_argument( + "--cuda-arch-list", + default=os.environ.get("NIXL_CUDA_ARCH_LIST"), + help="Comma-separated CUDA SM targets for NIXL, e.g. 100,103,107. \ + Defaults to the NIXL EP targets the installed nvcc supports.", + ) + parser.add_argument( + "--nixl-plugins", + default=os.environ.get("NIXL_PLUGINS", "UCX"), + help="Comma-separated NIXL plugins to build. Pass an empty string to \ + build every plugin whose dependencies are available.", + ) + parser.add_argument( + "--ucx-path", + default=os.environ.get("UCX_PATH"), + help="Prefix of an existing UCX to build against instead of building one. \ + The resulting wheel is not self-contained.", + ) + args = parser.parse_args() + build_and_install_prerequisites(args) From d6fe5dca686b3f6c5f018d399f4f119fedf842f1 Mon Sep 17 00:00:00 2001 From: Dhairyashil R G Date: Thu, 8 Oct 2026 04:51:51 +0530 Subject: [PATCH 154/952] [Bugfix][Parser] Emit buffered post-reasoning text when no tool parser is configured (#58911) Signed-off-by: Dhairyashil R G <18738735+dhairyashilRG@users.noreply.github.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- tests/parser/test_streaming.py | 72 ++++++++++++++++++++++++++++++++++ vllm/parser/abstract_parser.py | 13 +++++- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index cad16746be44..d5d43398039f 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -663,6 +663,78 @@ def test_engine_reasoning_no_tool_batched_content_passthrough(tokenizer, request assert len(tool_calls) == 0 +@pytest.mark.parametrize( + "text_chunks", + [ + # A trailing "<" that the engine holds back (it may start a tag). + ["let me thinka <", "x> b"], + # Whitespace after "" in the same delta as the marker. + ["let me think\n\n", "bold done"], + # The delta that ends reasoning carries only held-back text. + ["let me think", "<", "b>bold done"], + ], +) +def test_engine_reasoning_no_tool_batched_holdback_not_dropped( + tokenizer, request_obj, text_chunks +): + """Text the engine buffers when reasoning ends inside a batched delta is + returned by finish_streaming(); with no tool parser it must still be + emitted as content, not dropped.""" + parser = Qwen3ReasoningNoToolParser(tokenizer) + chunks = [tokenizer.encode(t, add_special_tokens=False) for t in text_chunks] + + results = stream_chunks(parser, tokenizer, chunks, request_obj) + reasoning, content, tool_calls = collect_fields(results) + + assert reasoning == "let me think" + assert content == "".join(text_chunks).split("", 1)[1] + assert len(tool_calls) == 0 + + +@pytest.mark.parametrize( + ("text_chunks", "include_reasoning", "expected_content"), + [ + # The delta that ends reasoning is also the last one: the held-back + # "<" must be emitted exactly once (the end-of-stream flush skips the + # reasoning parser once reasoning has ended). + (["let me thinka <"], True, "a <"), + # Whitespace-only content after the marker at the end of the stream. + (["let me think\n\n"], True, "\n\n"), + # Hiding reasoning must not hide the flushed content. + (["let me think\n\n", "bold"], False, "\n\nbold"), + ], +) +def test_engine_reasoning_no_tool_holdback_finished_and_hidden_reasoning( + tokenizer, text_chunks, include_reasoning, expected_content +): + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + include_reasoning=include_reasoning, + ) + parser = Qwen3ReasoningNoToolParser(tokenizer) + chunks = [tokenizer.encode(t, add_special_tokens=False) for t in text_chunks] + + results: list[DeltaMessage | None] = [] + prompt_token_ids: list[int] | None = [] + for i, chunk in enumerate(chunks): + results.append( + parser.parse_delta( + tokenizer.decode(chunk), + chunk, + request, + prompt_token_ids=prompt_token_ids, + finished=i == len(chunks) - 1, + ) + ) + prompt_token_ids = None + reasoning, content, tool_calls = collect_fields(results) + + assert reasoning == ("let me think" if include_reasoning else "") + assert content == expected_content + assert len(tool_calls) == 0 + + def _decode_stream_deltas(tokenizer, groups): """Decode token-ID groups into ``(delta_text, group)`` pairs via the real incremental ``DecodeStream``. diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index c28563fabf65..f75d7ce9028d 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -904,8 +904,17 @@ def parse_delta( (delta_message.content if delta_message else None) or "" ) + ((flush_delta.content if flush_delta else None) or "") if self._engine_based: - if delta_message and self._tool_parser is not None: - delta_message.content = None + if self._tool_parser is not None: + # The tool-call phase below consumes current_text. + if delta_message: + delta_message.content = None + elif current_text: + # No tool phase follows to consume the flushed text, + # and the engine is not fed again after the + # transition, so emit it as content here. + if delta_message is None: + delta_message = DeltaMessage() + delta_message.content = current_text else: delta_text = current_text From 38dc8ee50055622ab0d5fe6bb32a5aee526bcbf5 Mon Sep 17 00:00:00 2001 From: www6v Date: Thu, 8 Oct 2026 07:46:41 +0800 Subject: [PATCH 155/952] [Bugfix] Validate tool-parser/tokenizer compatibility at startup (#59749) Signed-off-by: www6v Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- vllm/parser/parser_manager.py | 58 +++++++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 719b396130af..cf0be908ed63 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -73,6 +73,29 @@ def get_reasoning_parser( raise TypeError(f"{reasoning_parser_name=} has not been registered") from e return parser + @classmethod + def validate_parser_tokenizer( + cls, + parser_cls: type[Parser], + tokenizer: TokenizerLike, + tool_parser_name: str | None = None, + reasoning_parser_name: str | None = None, + model_name: str | None = None, + ) -> None: + """Fail at startup if the tokenizer lacks what the parsers need.""" + try: + parser_cls(tokenizer) + except Exception as e: + flags = [] + if parser_cls.tool_parser_cls is not None: + flags.append(f"--tool-call-parser {tool_parser_name}") + if parser_cls.reasoning_parser_cls is not None: + flags.append(f"--reasoning-parser {reasoning_parser_name}") + raise TypeError( + f"{' and '.join(flags)} cannot be used with the tokenizer of " + f"{model_name!r}: {e}" + ) from e + @classmethod def get_parser( cls, @@ -97,13 +120,44 @@ def get_parser( If True, HarmonyParser is always returned. tool_strict_level: Server-side floor for tool-call structural tags (``--tool-strict-level``). - tokenizer: Tokenizer whose `response_template` metadata is - validated when the `hf` parser is selected. + tokenizer: Tokenizer the composed parser is validated against + at startup (see :meth:`validate_parser_tokenizer`). ``None`` + skips the check (e.g. ``skip_tokenizer_init``). Returns: A Parser class, or None if neither parser is specified. """ + parser_cls = cls._compose_parser( + tool_parser_name=tool_parser_name, + reasoning_parser_name=reasoning_parser_name, + enable_auto_tools=enable_auto_tools, + model_name=model_name, + is_harmony=is_harmony, + tool_strict_level=tool_strict_level, + tokenizer=tokenizer, + ) + if parser_cls is not None and tokenizer is not None: + cls.validate_parser_tokenizer( + parser_cls, + tokenizer, + tool_parser_name=tool_parser_name, + reasoning_parser_name=reasoning_parser_name, + model_name=model_name, + ) + return parser_cls + + @classmethod + def _compose_parser( + cls, + tool_parser_name: str | None, + reasoning_parser_name: str | None, + enable_auto_tools: bool, + model_name: str | None, + is_harmony: bool, + tool_strict_level: str, + tokenizer: TokenizerLike | None, + ) -> type[Parser] | None: if not tool_parser_name and not reasoning_parser_name: return None From ed825aafc0afb36929e837035dd70760c9bc6953 Mon Sep 17 00:00:00 2001 From: YueshenZ <844936239@qq.com> Date: Thu, 8 Oct 2026 08:07:52 +0800 Subject: [PATCH 156/952] [Bugfix][Frontend] Fix Responses namespace aliases in named tool calling (#59723) Signed-off-by: Shen <312058264+YueshenZ@users.noreply.github.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- docs/features/tool_calling.md | 2 + tests/tool_use/test_tool_choice_required.py | 74 ++++++++++++++++++- vllm/entrypoints/openai/responses/protocol.py | 70 ++++++++++++------ vllm/tool_parsers/utils.py | 2 - 4 files changed, 124 insertions(+), 24 deletions(-) diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 8dccd3a8b857..3081047dcec4 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -94,6 +94,8 @@ For best results, we recommend ensuring that the expected output format / schema To use a named function, you need to define the functions in the `tools` parameter of the chat completion request, and specify the `name` of one of the tools in the `tool_choice` parameter of the chat completion request. +In the Responses API, `tool_choice={"type": "function", "name": "..."}` can select a function inside a namespace by its full `namespace__function` name or by its local name. Names are matched literally first, so a plain function named `other__add` never stands in for `add`. A local name defined in more than one namespace must be selected by its full name. + ## Required Function Calling vLLM supports the `tool_choice='required'` option in the chat completion API. Similar to the named function calling, it also uses structured outputs, so this is enabled by default and will work with any supported model. However, support for alternative decoding backends are on the [roadmap](../usage/v1_guide.md#features) for the V1 engine. diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index 26978fbb7c0e..ebdea5aa0b9f 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -5,7 +5,11 @@ import pytest import regex as re -from openai.types.responses import FunctionTool, ToolChoiceFunction, WebSearchTool +from openai.types.responses import ( + FunctionTool, + ToolChoiceFunction, + WebSearchTool, +) from pydantic import TypeAdapter from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -13,6 +17,7 @@ ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.exceptions import VLLMValidationError from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.tool_parsers.utils import ( @@ -581,3 +586,70 @@ def test_responses_empty_params_constrains_object(self, params): choice = ToolChoiceFunction(type="function", name="ping") schema = get_json_schema_from_tools(choice, [tool]) assert schema == {"type": "object", "properties": {}} + + +class TestForcedFunctionShortNameAlias: + """Forced choices must use the selected function's parameter schema.""" + + PARAMS = {"type": "object", "properties": {"x": {"type": "integer"}}} + + @pytest.mark.parametrize("plain_first", [True, False]) + @pytest.mark.parametrize( + "namespace,function_name,choice_name", + [ + ("math", "add", "add"), + ("math", "add", "math__add"), + ("math", "vector__add", "vector__add"), + ("math__vector", "vector__add", "vector__add"), + ], + ) + def test_namespace_schema_is_not_shadowed( + self, plain_first, namespace, function_name, choice_name + ): + tools = [ + { + "type": "function", + "name": "other__add", + "parameters": { + "type": "object", + "properties": {"y": {"type": "string"}}, + }, + }, + { + "type": "namespace", + "name": namespace, + "description": "math tools", + "tools": [ + { + "type": "function", + "name": function_name, + "parameters": self.PARAMS, + } + ], + }, + ] + if not plain_first: + tools.reverse() + request = ResponsesRequest.model_validate( + { + "model": "test-model", + "input": "Add the numbers.", + "tools": tools, + "tool_choice": {"type": "function", "name": choice_name}, + } + ) + assert ( + get_json_schema_from_tools(request.tool_choice, request.tools) + == self.PARAMS + ) + + def test_plain_function_double_underscore_name_has_no_phantom_alias(self): + tool = FunctionTool(type="function", name="math__add", parameters=self.PARAMS) + choice = ToolChoiceFunction(type="function", name="add") + with pytest.raises(ValueError, match="has not been passed in `tools`"): + get_json_schema_from_tools(choice, [tool]) + + def test_plain_function_full_name_still_resolves(self): + tool = FunctionTool(type="function", name="math__add", parameters=self.PARAMS) + choice = ToolChoiceFunction(type="function", name="math__add") + assert get_json_schema_from_tools(choice, [tool]) == self.PARAMS diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 5d4a7411ac21..35e6d3aa3dcd 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -159,6 +159,51 @@ def _default_input_image_details(value: Any) -> Any: return {**value, "content": new_content} if changed else value +def _tool_field(tool: Any, field: str) -> Any: + return tool.get(field) if isinstance(tool, dict) else getattr(tool, field, None) + + +def _resolve_named_tool_choice(tool_name: str, tools: list[Any]) -> str: + """Resolve a named tool choice to the flat ``__`` form.""" + from vllm.tool_parsers.utils import flat_namespace_tool_name + + exact_names: set[str] = set() + local_names: dict[str, list[str]] = {} + for tool in tools: + if _tool_field(tool, "type") == "namespace": + namespace = _tool_field(tool, "name") + namespaced_tools = _tool_field(tool, "tools") + if not isinstance(namespaced_tools, list) or not isinstance(namespace, str): + return tool_name + for namespaced_tool in namespaced_tools: + local_name = _tool_field(namespaced_tool, "name") + if not isinstance(local_name, str): + continue + flat_name = flat_namespace_tool_name(namespace, local_name) + exact_names.add(flat_name) + local_names.setdefault(local_name, []).append(flat_name) + else: + name = _tool_field(tool, "name") + if isinstance(name, str): + exact_names.add(name) + if tool_name in exact_names: + return tool_name + candidates = local_names.get(tool_name, []) + if len(candidates) == 1: + return candidates[0] + if len(candidates) > 1: + raise VLLMValidationError( + f"Tool choice '{tool_name}' is ambiguous: it is defined in more " + f"than one namespace ({', '.join(candidates)}). Select it by its " + "full '__' name.", + parameter="tool_choice", + ) + raise VLLMValidationError( + "Tool choice 'function' not found in 'tools' parameter.", + parameter="tool_choice", + ) + + class ResponsesRequest(OpenAIBaseModel): # Ordered by official OpenAI API documentation # https://platform.openai.com/docs/api-reference/responses/create @@ -662,31 +707,14 @@ def check_tool_usage(cls, data): ) elif is_named_tool_choice and tools is not None: tool_name = tool_choice.get("name") - tool_names = set() - for tool in tools: - if isinstance(tool, dict): - if tool.get("type") == "namespace": - namespace = tool.get("name") - namespaced_tools = tool.get("tools") - if not isinstance(namespaced_tools, list): - return data - for namespaced_tool in namespaced_tools: - namespaced_name = ( - namespaced_tool.get("name") - if isinstance(namespaced_tool, dict) - else getattr(namespaced_tool, "name", None) - ) - tool_names.add(namespaced_name) - tool_names.add(f"{namespace}__{namespaced_name}") - else: - tool_names.add(tool.get("name")) - else: - tool_names.add(getattr(tool, "name", None)) - if not tool_name or tool_name not in tool_names: + if not tool_name: raise VLLMValidationError( "Tool choice 'function' not found in 'tools' parameter.", parameter="tool_choice", ) + resolved = _resolve_named_tool_choice(tool_name, tools) + if resolved != tool_name: + data["tool_choice"] = {**tool_choice, "name": resolved} return data diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 86d68a9b5130..7a93c843b8e3 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -432,8 +432,6 @@ def get_json_schema_from_tools( continue for name, params in iter_response_function_tool_info(tool): responses_tool_map[name] = params - if "__" in name: - responses_tool_map.setdefault(name.rsplit("__", 1)[1], params) if tool_name not in responses_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") return _params_or_empty_object(responses_tool_map[tool_name]) From 8faac3dd2af39fd631f018277387f01c89ad5782 Mon Sep 17 00:00:00 2001 From: v-shobhit <161510941+v-shobhit@users.noreply.github.com> Date: Wed, 7 Oct 2026 17:10:32 -0700 Subject: [PATCH 157/952] [Bugfix] [Spec Decode] DeepSeek V4/V4.1 NVFP4: keep ignored MTP/DSpark experts on MXFP4 (#59655) Signed-off-by: Shobhit Verma Signed-off-by: v-shobhit <161510941+v-shobhit@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Benjamin Chislett --- .buildkite/test_areas/lm_eval.yaml | 2 + .buildkite/test_areas/quantization.yaml | 1 + .../DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml | 15 ++++ .../gsm8k/configs/models-spec-decode.txt | 1 + tests/quantization/test_fp8.py | 76 +++++++++++++++++++ vllm/models/deepseek_v41/quant_config.py | 32 +++++++- 6 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index bcb111084843..d1a115a1e1b0 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -322,9 +322,11 @@ steps: optional: true source_file_dependencies: - tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml + - tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml - tests/evals/gsm8k/configs/Kimi-K3-pruned75-DSpark-TP4.yaml - tests/evals/gsm8k/configs/models-spec-decode.txt - vllm/models/deepseek_v4/ + - vllm/models/deepseek_v41/ - vllm/models/kimi_k3/ - vllm/model_executor/models/qwen3_dspark.py - vllm/v1/spec_decode/metrics.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index e73054a4bcaf..b49b4b92c736 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -12,6 +12,7 @@ steps: source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization + - vllm/models/deepseek_v41/quant_config.py - tests/quantization commands: # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. diff --git a/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml b/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml new file mode 100644 index 000000000000..78d899569644 --- /dev/null +++ b/tests/evals/gsm8k/configs/DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml @@ -0,0 +1,15 @@ +model_name: "nvidia/DeepSeek-V4.1-Flash-NVFP4" +accuracy_threshold: 0.88 +tolerance: 0.03 +# The bundled DSpark draft (mtp.*) stays MXFP4 in this checkpoint. Loading it +# as NVFP4 keeps accuracy but drops acceptance length to 1.0 (3.57 when correct). +min_acceptance_length: 2.5 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1800 +server_args: >- + --tokenizer-mode deepseek_v41 + --trust-remote-code + --max-model-len 8192 + --tensor-parallel-size 4 + --speculative-config '{"method":"dspark","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/models-spec-decode.txt b/tests/evals/gsm8k/configs/models-spec-decode.txt index 427c35d912ab..96cead25858c 100644 --- a/tests/evals/gsm8k/configs/models-spec-decode.txt +++ b/tests/evals/gsm8k/configs/models-spec-decode.txt @@ -1,2 +1,3 @@ DeepSeek-V4-Flash-DSpark-confidence-TEP4.yaml Kimi-K3-pruned75-DSpark-TP4.yaml +DeepSeek-V4.1-Flash-NVFP4-DSpark-TP4.yaml diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 1cb53f223929..64341843d951 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -318,6 +318,82 @@ def test_deepseek_v41_declines_quark_configs(): ) +_DSV41_NVFP4 = { + "moe_quant_algo": "NVFP4", + "ignore": ["*.attn.*", "*.ffn.shared_experts.*", "head", "mtp.*"], +} +_DSV4_NVFP4_DSPARK = { + "moe_quant_algo": "NVFP4", + "ignore": ["*.attn.*", "*.ffn.shared_experts.*", "head"], +} + + +@pytest.mark.parametrize( + ("ckpt_quant_cfg", "prefix", "expected"), + [ + pytest.param(_DSV41_NVFP4, "model.layers.1.ffn.experts", "nvfp4", id="target"), + pytest.param( + _DSV41_NVFP4, + "language_model.model.layers.0.ffn.experts", + "nvfp4", + id="target-vl-prefix", + ), + pytest.param( + _DSV41_NVFP4, "model.layers.40.ffn.experts", "mxfp4", id="ignored-mtp0" + ), + pytest.param( + _DSV41_NVFP4, "model.layers.42.ffn.experts", "mxfp4", id="ignored-mtp2" + ), + pytest.param( + _DSV4_NVFP4_DSPARK, + "model.layers.41.ffn.experts", + "nvfp4", + id="quantized-mtp1", + ), + pytest.param( + {"moe_quant_algo": "NVFP4"}, + "model.layers.40.ffn.experts", + "nvfp4", + id="no-ignore-list", + ), + pytest.param({}, "model.layers.1.ffn.experts", "mxfp4", id="not-nvfp4"), + ], +) +def test_deepseek_v41_nvfp4_export_keeps_ignored_experts_mxfp4( + monkeypatch, ckpt_quant_cfg, prefix, expected +): + """NVFP4 exports keep the routed experts they ``ignore`` in MXFP4, e.g. the + bundled DSpark draft (``mtp.*``, built as ``layers.{40 + k}``) of + DeepSeek-V4.1-Flash-NVFP4. The ignore list is read from the checkpoint + passed to ``from_config``: the current hf_config below carries none.""" + import vllm.model_executor.layers.quantization.modelopt as modelopt + import vllm.models.deepseek_v41.quant_config as dsv41_qc + from vllm.model_executor.layers.fused_moe import RoutedExperts + + hf_config = SimpleNamespace( + expert_dtype="fp4", + quantization_config={"moe_quant_algo": ckpt_quant_cfg.get("moe_quant_algo")}, + ) + model_config = SimpleNamespace( + hf_config=hf_config, hf_text_config=SimpleNamespace(num_hidden_layers=40) + ) + monkeypatch.setattr( + dsv41_qc, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=model_config), + ) + monkeypatch.setattr(dsv41_qc, "Mxfp4MoEMethod", lambda moe_config: "mxfp4") + monkeypatch.setattr(modelopt, "ModelOptNvFp4FusedMoE", lambda **kwargs: "nvfp4") + + config = dsv41_qc.DeepseekV4FP8Config.from_config( + {"quant_method": "fp8", "activation_scheme": "dynamic", **ckpt_quant_cfg} + ) + monkeypatch.setattr(config, "_get_nvfp4_config", lambda: None) + layer = RoutedExperts.__new__(RoutedExperts) + layer.moe_config = None + assert config.get_quant_method(layer, prefix) == expected + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA") @pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu]) @pytest.mark.parametrize( diff --git a/vllm/models/deepseek_v41/quant_config.py b/vllm/models/deepseek_v41/quant_config.py index 384e9488793c..1d92943db425 100644 --- a/vllm/models/deepseek_v41/quant_config.py +++ b/vllm/models/deepseek_v41/quant_config.py @@ -4,8 +4,11 @@ from __future__ import annotations +import fnmatch from typing import TYPE_CHECKING, cast +import regex as re + from vllm.config import get_current_vllm_config from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.fused_moe import ( @@ -23,6 +26,7 @@ ) _DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") +_EXPERTS_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)\.ffn\.experts$") if TYPE_CHECKING: from vllm.model_executor.layers.quantization.modelopt import ( @@ -54,6 +58,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._resolved_expert_dtype: str | None = None self._resolved_moe_quant_algo: str | None = None + self._modelopt_ignore: tuple[str, ...] = () self._nvfp4_config: ModelOptNvFp4Config | None = None # ``is_scale_e8m0`` is a property that resolves on first read, # by which time the current vllm_config has been set. @@ -103,6 +108,21 @@ def moe_quant_algo(self) -> str: self._resolve_moe_overrides() return self._resolved_moe_quant_algo or "" + def _is_nvfp4_expert_layer(self, prefix: str) -> bool: + if self.moe_quant_algo != "NVFP4": + return False + match = _EXPERTS_LAYER_RE.search(prefix) + if match is None or not self._modelopt_ignore: + return True + # MTP/draft layers are built as layers.{num_hidden_layers + k} but are + # named mtp.{k} in some checkpoints. + model_config = get_current_vllm_config().model_config + num_layers = model_config.hf_text_config.num_hidden_layers + idx = int(match.group(1)) + layer = f"mtp.{idx - num_layers}" if idx >= num_layers else f"layers.{idx}" + name = f"{layer}.ffn.experts" + return not any(fnmatch.fnmatchcase(name, p) for p in self._modelopt_ignore) + def _get_nvfp4_config(self) -> ModelOptNvFp4Config: if self._nvfp4_config is None: from vllm.model_executor.layers.quantization.modelopt import ( @@ -173,6 +193,9 @@ def override_quantization_method( @classmethod def from_config(cls, config: dict) -> DeepseekV4FP8Config: + # ModelOpt NVFP4 exports keep the experts they ``ignore`` in the native + # MXFP4 format, e.g. the bundled MTP/DSpark layers (``mtp.*``). + ignore = config.get("ignore") # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8 # path: the runtime layout matches the DeepSeek-native fp8 checkpoint, # so translate the schema into format Fp8Config.from_config expects. @@ -188,7 +211,12 @@ def from_config(cls, config: dict) -> DeepseekV4FP8Config: name for name in quark_exclude if isinstance(name, str) ], } - return cast("DeepseekV4FP8Config", super().from_config(config)) + quant_config = cast("DeepseekV4FP8Config", super().from_config(config)) + if isinstance(ignore, list): + quant_config._modelopt_ignore = tuple( + p for p in ignore if isinstance(p, str) + ) + return quant_config def get_quant_method(self, layer, prefix): if ( @@ -221,7 +249,7 @@ def get_quant_method(self, layer, prefix): ): return UnquantizedFusedMoEMethod(layer.moe_config) if self.expert_dtype == "fp4": - if self.moe_quant_algo == "NVFP4": + if self._is_nvfp4_expert_layer(prefix): from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4FusedMoE, ) From a58bdd0d4ec62764a8d6c5d4f9ed6a1dd4bf9b80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 02:11:59 +0200 Subject: [PATCH 158/952] [Perf] JIT monitor: hook TileLang JITImpl.compile instead of __call__ (#58685) Signed-off-by: NickLucche --- tests/jit_monitor/test_hooks.py | 5 ++++- vllm/utils/jit_monitor.py | 18 ++++++++---------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py index 452fe9699504..6d4f25314c08 100644 --- a/tests/jit_monitor/test_hooks.py +++ b/tests/jit_monitor/test_hooks.py @@ -65,10 +65,13 @@ def __call__(self, *args, **kwargs): key, _ = self.func.parse_args(*args, **kwargs) kernel = self._kernel_cache.get(key) if kernel is None: - kernel = "compiled" + kernel = self.compile(*args, **kwargs) self._kernel_cache[key] = kernel return kernel + def compile(self, *args, **kwargs): + return "compiled" + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) fake_kernel.JITKernel = FakeJITKernel diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index ffb63bc493d4..387189c793fb 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -499,32 +499,30 @@ def _init_with_monitor(self, *args, **kwargs): jit_kernel_cls.__init__ = _init_with_monitor if jit_impl_cls is not None: - original_call = jit_impl_cls.__call__ + # JITImpl.__call__ only calls compile() on a kernel-cache miss, so + # hooking compile() keeps the per-call hot path untouched. + original_compile = jit_impl_cls.compile - @functools.wraps(original_call) - def _call_with_monitor(self, *args, **kwargs): + @functools.wraps(original_compile) + def _compile_with_monitor(self, *args, **kwargs): global _tilelang_jitimpl_compile_depth - cache_key = _tilelang_cache_miss_key(self, args, kwargs) - if cache_key is None: - return original_call(self, *args, **kwargs) - _tilelang_jitimpl_compile_depth += 1 try: detail = None if _verbose: detail = _format_verbose_tilelang_compile_details( - self, args, kwargs, cache_key + self, args, kwargs, _tilelang_cache_miss_key(self, args, kwargs) ) func = getattr(self, "func", None) orig_func = getattr(func, "orig_func", None) _log_tilelang_jit_compile( _tilelang_kernel_name(orig_func or func), detail ) - return original_call(self, *args, **kwargs) + return original_compile(self, *args, **kwargs) finally: _tilelang_jitimpl_compile_depth -= 1 - jit_impl_cls.__call__ = _call_with_monitor + jit_impl_cls.compile = _compile_with_monitor _tilelang_hook_installed = True From 1283362e7618c79dedb5a7bdd9308d3f0e0d389c Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 7 Oct 2026 17:17:56 -0700 Subject: [PATCH 159/952] [CI] Add headroom and selection coverage for H200 fast lanes (#56775) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- .buildkite/test_areas/benchmarks.yaml | 2 +- .buildkite/test_areas/cuda.yaml | 7 ++++++- .buildkite/test_areas/misc.yaml | 4 ++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 4b681cd589a9..40f8f9796324 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Benchmarks CLI" key: benchmarks-cli-test - timeout_in_minutes: 45 + timeout_in_minutes: 55 device: h200_18gb source_file_dependencies: - vllm/ diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 02ac11879fcb..45b3dc768e6d 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -33,14 +33,19 @@ steps: - label: ":nvidia: (H200 MIG 18GB) CUDAGraph" device: h200_18gb key: cudagraph - timeout_in_minutes: 30 + timeout_in_minutes: 45 source_file_dependencies: - tests/v1/cudagraph - vllm/v1/cudagraph_dispatcher.py - vllm/config/compilation.py + - vllm/config/vllm.py - vllm/compilation + - vllm/v1/worker/gpu/cudagraph_utils.py + - vllm/v1/worker/gpu/dp_utils.py - vllm/v1/worker/encoder_cudagraph.py - vllm/v1/worker/encoder_cudagraph_defs.py + - vllm/v1/worker/gpu/model_runner.py + - vllm/v1/worker/gpu/ubatch_utils.py commands: - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py - pytest -v -s v1/cudagraph/test_cudagraph_mode.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cbcfaeaea7ca..e1490139c052 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) V1 Sample" key: v1-sample - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: &v1-sample-logits-deps - vllm/config/ @@ -38,7 +38,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) V1 Logits + Oracle" key: v1-logits-oracle - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: *v1-sample-logits-deps commands: From 1995c0fcd02f88d5daa3fdbcd47e0f016e7b89b9 Mon Sep 17 00:00:00 2001 From: Kang Nahoon <49896157+ejrtks1020@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:24:20 +0900 Subject: [PATCH 160/952] [Bugfix][Frontend] Apply output_config and thinking in Anthropic count_tokens (#59180) Signed-off-by: Kang Nahoon Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_anthropic_messages_conversion.py | 26 +++++++++++++++++++ vllm/entrypoints/anthropic/protocol.py | 2 ++ vllm/entrypoints/anthropic/serving.py | 8 +++--- 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 1bccee46016e..399c898d0b01 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1987,6 +1987,32 @@ def test_claude_code_payload(self): assert result.include_reasoning is True assert result.thinking_token_budget is None + @pytest.mark.parametrize( + "config", + [ + pytest.param({"output_config": {"effort": "low"}}, id="effort"), + pytest.param({"thinking": {"type": "disabled"}}, id="disabled"), + pytest.param( + {"thinking": {"type": "enabled", "budget_tokens": 2048}}, + id="enabled", + ), + ], + ) + def test_count_tokens_matches_messages(self, config): + """Effort and thinking change the rendered prompt, so count_tokens must + apply them the same way /v1/messages does.""" + messages = [{"role": "user", "content": "Hello"}] + expected = _convert( + AnthropicMessagesRequest( + model="test-model", max_tokens=4096, messages=messages, **config + ) + ) + result = _convert( + AnthropicCountTokensRequest(model="test-model", messages=messages, **config) + ) + assert result.reasoning_effort == expected.reasoning_effort + assert result.thinking_token_budget == expected.thinking_token_budget + class TestProbeDisabledThinkingEffort: """``auto`` falls back to ``low`` when ``none`` cannot turn thinking off.""" diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 2a14c6293327..5997d2e3d2e6 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -385,7 +385,9 @@ class AnthropicCountTokensRequest(BaseModel): model: str messages: list[AnthropicMessage] + output_config: AnthropicOutputConfig | None = None system: str | list[AnthropicContentBlock] | None = None + thinking: AnthropicThinkingConfig | None = None tool_choice: AnthropicToolChoice | None = None tools: list[AnthropicTool] | None = None diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 05c1e24f2a48..f11697fa165e 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -633,8 +633,6 @@ def _handle_thinking( ``display`` is intentionally ignored: suppressing reasoning would mark it ended for structured outputs and drop it from multi-turn history. """ - if isinstance(anthropic_request, AnthropicCountTokensRequest): - return thinking: AnthropicThinkingConfig | None = anthropic_request.thinking if thinking is None: return @@ -655,8 +653,6 @@ def _handle_output_config( anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest, ) -> None: """Handle output configuration such as output format and effort.""" - if isinstance(anthropic_request, AnthropicCountTokensRequest): - return output_config: AnthropicOutputConfig | None = anthropic_request.output_config if output_config and output_config.format and output_config.format.json_schema: req.response_format = ResponseFormat( @@ -1238,9 +1234,13 @@ async def count_tokens( raw_request: Request | None = None, ) -> AnthropicCountTokensResponse | ErrorResponse: """Implements Anthropic's messages.count_tokens endpoint.""" + disabled_thinking_effort: AnthropicDisabledThinkingEffort = "none" + if request.thinking is not None and request.thinking.type == "disabled": + disabled_thinking_effort = await self._get_disabled_thinking_effort() chat_req = self.to_chat_completion_request( request, merge_inline_system=self._merge_inline_system, + disabled_thinking_effort=disabled_thinking_effort, ) result = await self.render_chat_request(chat_req) if isinstance(result, ErrorResponse): From 9e3e37cb3c041fa3cb7dc578bca1fa0a004b715a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Thu, 8 Oct 2026 02:56:28 +0200 Subject: [PATCH 161/952] [TEST][XPU][CI] disable tests of unsupported quantization (#60383) Signed-off-by: microslaw Co-authored-by: microslaw Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- .buildkite/scripts/hardware_ci/run-intel-ci-test.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 08f8dcdeab89..b1030e372c41 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -153,7 +153,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s quantization/test_auto_round.py && - pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights' + pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' - label: "XPU compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: diff --git a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh index c679d0aff242..44f0d9944adc 100644 --- a/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh +++ b/.buildkite/scripts/hardware_ci/run-intel-ci-test.sh @@ -57,7 +57,7 @@ case "${test_suite}" in cd tests pytest -v -s quantization/test_auto_round.py - pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights + pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]" ;; compressed-tensors-fp8) cd tests From 2dd887426e28dd7dfc96465acc45e06e849f3d44 Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Thu, 8 Oct 2026 09:03:08 +0800 Subject: [PATCH 162/952] [Bugfix][Frontend] Do not modify caller params in LLM.score() (#59959) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 叮当猫 <2756823972@qq.com> Signed-off-by: apex-mochen <2756823972@qq.com> Co-authored-by: DeepSeek Harness Co-authored-by: Taneem Ibrahim Co-authored-by: Codex --- .../entrypoints/pooling/test_io_processor.py | 37 +++++++++++++++++++ vllm/entrypoints/pooling/offline.py | 2 + 2 files changed, 39 insertions(+) diff --git a/tests/entrypoints/pooling/test_io_processor.py b/tests/entrypoints/pooling/test_io_processor.py index 9d2207fdf15d..79bddc0ffe8f 100644 --- a/tests/entrypoints/pooling/test_io_processor.py +++ b/tests/entrypoints/pooling/test_io_processor.py @@ -7,6 +7,8 @@ from vllm import PoolingParams from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor from vllm.entrypoints.pooling.typing import OfflineEncodeInputsContext from vllm.exceptions import VLLMValidationError from vllm.renderers import TokenizeParams @@ -221,3 +223,38 @@ def test_qwen3_reranker_no_warning_when_template_provided(monkeypatch): ) mock_logger.warning.assert_not_called() + + +def _make_score_llm(monkeypatch) -> PoolingOfflineMixin: + """A scoring LLM with no model behind it. + + Input validation runs for real; only the rendering, engine and + post-processing boundaries are stubbed out. The pooling task is filled in + before any of them, so the mutation is observable without inference. + """ + proc = CrossEncoderIOProcessor.__new__(CrossEncoderIOProcessor) + proc.is_multimodal_model = False + proc.architecture = "CrossEncoder" + monkeypatch.setattr(proc, "get_request_factory_offline", lambda ctx: (None, 0)) + monkeypatch.setattr(proc, "post_process_offline", lambda ctx: []) + + llm = PoolingOfflineMixin.__new__(PoolingOfflineMixin) + llm.runner_type = "pooling" + llm.pooling_task = "classify" # SCORE_TYPE_MAP -> "cross-encoder" + llm.model_config = SimpleNamespace(hf_config=SimpleNamespace(num_labels=1)) + llm.pooling_io_processors = {"cross-encoder": proc} + monkeypatch.setattr(llm, "_run_tiling_engine", lambda *args, **kwargs: []) + return llm + + +def test_score_leaves_caller_pooling_params_untouched(monkeypatch): + llm = _make_score_llm(monkeypatch) + + params = PoolingParams() + llm.score("query", "doc", pooling_params=params) + + assert params.task is None, ( + "LLM.score() wrote the model's own pooling task into the caller's " + "PoolingParams; reusing that object with another pooling model now " + "fails with 'You cannot overwrite ...' before any inference runs." + ) diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 0a0304e7a4c5..27388f1ff0a9 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -362,6 +362,8 @@ def score( assert isinstance(pooling_params, PoolingParams) pooling_task = io_processor.pooling_task + # Clone to avoid modifying the caller's pooling parameters. + pooling_params = pooling_params.clone() if pooling_params.task is None: pooling_params.task = pooling_task elif pooling_params.task != pooling_task: From fdfc171a7495bed2997609a0758d7e69c69440c0 Mon Sep 17 00:00:00 2001 From: Junyi Yao Date: Wed, 7 Oct 2026 18:07:35 -0700 Subject: [PATCH 163/952] [Bugfix][ToolParser] Stream every tool call when one delta completes several (#60354) Signed-off-by: Junyi Yao Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: sfeng33 <4florafeng@gmail.com> --- tests/tool_use/test_tool_choice_required.py | 152 +++++--- vllm/parser/abstract_parser.py | 15 +- vllm/renderers/online_derenderer.py | 6 - vllm/tool_parsers/streaming.py | 373 ++++++++++++-------- 4 files changed, 341 insertions(+), 205 deletions(-) diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index ebdea5aa0b9f..be58ffa83aa9 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -295,45 +295,61 @@ def test_structured_outputs_json_without_parameters( ) -def _collect_required_tool_streaming_json(output_json: str, delta_len: int) -> str: +def _stream_required_tool_calls( + output_json: str, + deltas: list[str], + tool_call_id_type: str = "random", + tool_call_idx: int | None = None, +) -> list[dict]: + """Feed ``deltas`` (which concatenate to ``output_json``) through the + required-tool streaming helper and rebuild the calls per array index, + checking that every index gets exactly one id/name chunk first.""" + assert "".join(deltas) == output_json previous_text = "" - function_name_returned = False - messages = [] - for i in range(0, len(output_json), delta_len): - delta_text = output_json[i : i + delta_len] + calls: dict[int, dict] = {} + for delta_text in deltas: current_text = previous_text + delta_text - - delta_message, function_name_returned = extract_required_tool_call_streaming( + delta_message, _ = extract_required_tool_call_streaming( previous_text=previous_text, current_text=current_text, - delta_text=delta_text, - function_name_returned=function_name_returned, - tool_call_idx=None, - tool_call_id_type="random", + tool_call_idx=tool_call_idx, + tool_call_id_type=tool_call_id_type, ) - - if delta_message: - messages.append(delta_message) - previous_text = current_text + if delta_message is None: + continue + assert delta_message.tool_calls + for tc in delta_message.tool_calls: + assert tc.function is not None + if tc.id is not None: + assert tc.index not in calls, "id/name sent twice for one index" + assert tc.type == "function" + assert tc.function.name + calls[tc.index] = { + "id": tc.id, + "name": tc.function.name, + "arguments": tc.function.arguments or "", + } + else: + assert tc.index in calls, "arguments before id/name" + assert tc.function.name is None + assert tc.function.arguments + calls[tc.index]["arguments"] += tc.function.arguments + assert list(calls) == list(range(len(calls))) + return [calls[i] for i in range(len(calls))] - assert len(messages) > 0 - combined_messages = "[" - for message in messages: - fn = message.tool_calls[0].function - assert fn is not None - if fn.name: - if len(combined_messages) > 1: - combined_messages += "}," +def _fixed_len_deltas(text: str, delta_len: int) -> list[str]: + return [text[i : i + delta_len] for i in range(0, len(text), delta_len)] - combined_messages += ( - '{"name": "' + fn.name + '", "parameters": ' + (fn.arguments or "") - ) - else: - combined_messages += fn.arguments or "" - combined_messages += "}]" - return combined_messages + +def _assert_streams_to(output: list[dict], deltas: list[str]) -> None: + calls = _stream_required_tool_calls(json.dumps(output), deltas) + assert [ + {"name": c["name"], "parameters": json.loads(c["arguments"])} for c in calls + ] == output + for call, expected in zip(calls, output): + assert call["arguments"] == json.dumps(expected["parameters"]) @pytest.mark.parametrize("output", VALID_TOOLS) @@ -343,11 +359,7 @@ def test_streaming_output_valid(output, empty_params, delta_len): output = deepcopy(output) if empty_params: output = [{"name": o["name"], "parameters": {}} for o in output] - output_json = json.dumps(output) - - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len) - assert json.loads(combined_messages) == output - assert json.dumps(json.loads(combined_messages)) == output_json + _assert_streams_to(output, _fixed_len_deltas(json.dumps(output), delta_len)) @pytest.mark.parametrize( @@ -358,22 +370,78 @@ def test_streaming_output_valid(output, empty_params, delta_len): "a }} b", 'a " } b', r"a \ } b", + # Text that looks like the end of one call and the start of another. + '"}}, {"name": "get_forecast", "parameters": {"city": "x', ], ) @pytest.mark.parametrize("delta_len", [1, 2, 3, 8, 9999]) def test_streaming_output_valid_with_braces_in_string(city, delta_len): - output = [{"name": "get_current_weather", "parameters": {"city": city}}] - output_json = json.dumps(output) - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len) - assert json.loads(combined_messages) == output - assert json.dumps(json.loads(combined_messages)) == output_json + output = [ + {"name": "get_current_weather", "parameters": {"city": city}}, + {"name": "get_forecast", "parameters": {"city": city, "days": 1}}, + ] + _assert_streams_to(output, _fixed_len_deltas(json.dumps(output), delta_len)) def test_streaming_output_valid_with_trailing_extra_data(): output = [{"name": "get_current_weather", "parameters": {"city": "Vienna"}}] output_json = json.dumps(output) + "\nDONE" - combined_messages = _collect_required_tool_streaming_json(output_json, delta_len=3) - assert json.loads(combined_messages) == output + calls = _stream_required_tool_calls(output_json, _fixed_len_deltas(output_json, 3)) + assert [ + {"name": c["name"], "parameters": json.loads(c["arguments"])} for c in calls + ] == output + + +TWO_CALLS = [ + {"name": "get_current_weather", "parameters": {"city": "Dallas"}}, + {"name": "get_forecast", "parameters": {"city": "Dallas", "days": 3}}, +] + + +def _split_at(text: str, cuts: list[int]) -> list[str]: + bounds = [0, *cuts, len(text)] + return [text[a:b] for a, b in zip(bounds, bounds[1:])] + + +def test_streaming_multiple_calls_completed_in_one_delta(): + """Regression test for #60351: a delta that completes more than one + call must stream every call, not only the last one.""" + text = json.dumps(TWO_CALLS) + _assert_streams_to(TWO_CALLS, [text[:1], text[1:]]) + _assert_streams_to(TWO_CALLS, [text]) + + +def test_streaming_delta_crossing_call_boundary(): + """A delta that closes one call and opens the next must attribute the + closing argument text to the first call and still announce the next.""" + text = json.dumps(TWO_CALLS) + boundary = text.index("}, {") + 1 + name_end = text.index('"get_forecast"') + 6 + for cuts in ( + [boundary - 1, name_end], # '}, {"name": "get_f' in one delta + [boundary - 2, boundary + 2], # '"}}, {' then the rest + [boundary + 2], # opening brace of the next call ends a delta + ): + _assert_streams_to(TWO_CALLS, _split_at(text, cuts)) + + +def test_streaming_invalid_name_string_is_never_announced(): + """A closed name that is not a valid JSON string must not be streamed + (the non-streaming path rejects such an array as well).""" + text = '[{"name": "bad\\x", "parameters": {"city": "x"}}]' + assert _stream_required_tool_calls(text, _fixed_len_deltas(text, 4)) == [] + + +def test_streaming_tool_call_idx_increments_per_started_call(): + """kimi_k2 ids number the calls started in one delta consecutively.""" + text = json.dumps(TWO_CALLS) + calls = _stream_required_tool_calls( + text, [text], tool_call_id_type="kimi_k2", tool_call_idx=5 + ) + assert [c["id"] for c in calls] == [ + "functions.get_current_weather:5", + "functions.get_forecast:6", + ] FUNCTION_TOOL = FunctionTool( diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f75d7ce9028d..4af133c9882c 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -741,8 +741,6 @@ def _extract_tool_calls_streaming( extract_required_tool_call_streaming( previous_text=previous_text, current_text=current_text, - delta_text=delta_text, - function_name_returned=function_name_returned, tool_call_idx=tool_call_idx, tool_call_id_type=tool_call_id_type, ) @@ -952,12 +950,13 @@ def parse_delta( elif not delta_message.reasoning: delta_message.reasoning = reasoning_from_this_batch - if ( - delta_message - and delta_message.tool_calls - and delta_message.tool_calls[0].id is not None - ): - state.history_tool_call_cnt += 1 + if delta_message and delta_message.tool_calls: + # A delta may start several tool calls at once (e.g. the + # "required" path when one delta completes more than one + # call), so count every id rather than only the first. + state.history_tool_call_cnt += sum( + tc.id is not None for tc in delta_message.tool_calls + ) # No phase active: pass through as content. # Skip when reasoning just ended in this delta — the engine already diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index bd5bf88fbe8c..66b73fb4912c 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -640,12 +640,6 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: if tc.index < len(last_tool_call_ids): # Pin: reuse the ID already recorded for this index # rather than one a from scratch replay regenerated. - # Real trigger not just defensive with - # tool_choice="required", - # extract_required_tool_call_streaming resets - # function_name_returned to False whenever the - # partial JSON transiently fails to parse which - # re-emits id+name for the same index on replay. tc.id = last_tool_call_ids[tc.index] else: last_tool_call_ids.append(tc.id) diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 70da7f3cc6b8..07c1a2e031eb 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -2,19 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +from dataclasses import dataclass from typing import TYPE_CHECKING -import partial_json_parser -import regex as re -from partial_json_parser.core.options import Allow - from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, ) -from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: @@ -23,63 +19,6 @@ TokenizerLike = object -def _bracket_level_state( - s: str, opening: str = "{", closing: str = "}" -) -> tuple[int, bool, bool]: - level = 0 - in_string = False - escaped = False - for char in s: - if escaped: - escaped = False - continue - if in_string and char == "\\": - escaped = True - continue - if char == '"': - in_string = not in_string - continue - if not in_string: - if char == opening: - level += 1 - elif char == closing: - level -= 1 - return level, in_string, escaped - - -def filter_delta_text( - delta_text: str, - previous_text: str, -) -> tuple[str, bool]: - """Trim trailing tool-list delimiters from required-tool streaming text.""" - bracket_level, in_string, escaped = _bracket_level_state(previous_text) - updated_delta = "" - passed_zero = False - for char in delta_text: - if escaped: - escaped = False - elif in_string: - if char == "\\": - escaped = True - elif char == '"': - in_string = False - elif char == '"': - in_string = True - elif char == "{": - bracket_level += 1 - passed_zero = bracket_level == 0 - elif char == "}": - bracket_level -= 1 - passed_zero = bracket_level == 0 - - if bracket_level != 0: - updated_delta += char - else: - if not in_string and char == ",": - break - return updated_delta, passed_zero - - def extract_named_tool_call_streaming( *, delta_text: str, @@ -124,100 +63,236 @@ def extract_named_tool_call_streaming( ) +@dataclass +class _RequiredToolCallSpan: + """Where one element of a ``tool_choice="required"`` array sits in the text. + + The model output for ``tool_choice="required"`` is a JSON array of + ``{"name": ..., "parameters": ...}`` objects. ``name`` is only set once + its string value is closed and decodes as JSON; ``name_end`` is the + offset just past that closing quote. ``args_start``/``args_end`` are the + character offsets of the ``parameters`` value (``args_end`` is ``None`` + while the value is still being generated). + + All offsets are prefix-stable: scanning any prefix of the text yields the + same spans truncated to that prefix. That is what lets one scan of the + current text tell what an earlier, shorter text had already revealed. + """ + + name: str | None = None + name_end: int | None = None + args_start: int | None = None + args_end: int | None = None + + def ready_at(self, length: int) -> bool: + """Whether the first ``length`` characters of the text already made + this call streamable: its name was complete and its parameters value + had started.""" + return ( + self.name is not None + and self.name_end is not None + and self.name_end <= length + and self.args_start is not None + and self.args_start < length + ) + + @property + def ready(self) -> bool: + return self.name is not None and self.args_start is not None + + +# Parse states for one element object of the required-tool-call array. +_EXPECT_KEY = 0 +_EXPECT_COLON = 1 +_EXPECT_VALUE = 2 +_IN_STRING_VALUE = 3 +_IN_CONTAINER_VALUE = 4 +_IN_SCALAR_VALUE = 5 +_AFTER_VALUE = 6 + + +def _scan_required_tool_calls(text: str) -> list[_RequiredToolCallSpan]: + """Locate every tool call of a (possibly partial) required-tool-call array. + + A single string- and nesting-aware pass over ``text`` that records, for + each top-level array element, its ``name`` (once complete) and the + character span of its ``parameters`` value. Text after the closing ``]`` + of the array is ignored. Only the element level of the array is + interpreted, so braces, brackets and quotes inside nested values and + strings cannot confuse it. + """ + spans: list[_RequiredToolCallSpan] = [] + depth = 0 # nesting depth of ``{}`` and ``[]``; element objects sit at 2 + in_string = False + escaped = False + string_start = 0 + cur: _RequiredToolCallSpan | None = None + expect = _EXPECT_KEY + key: str | None = None + value_start = 0 + + def end_value(pos: int) -> None: + nonlocal expect + assert cur is not None + if key == "parameters": + cur.args_end = pos + elif key == "name" and expect == _IN_STRING_VALUE: + try: + cur.name = json.loads(text[value_start:pos]) + cur.name_end = pos + except json.JSONDecodeError: + # Not a valid JSON string: the call can never become ready, + # which matches the non-streaming path rejecting the array. + cur.name = None + expect = _AFTER_VALUE + + def start_value(pos: int, state: int) -> None: + nonlocal expect, value_start + assert cur is not None + value_start = pos + if key == "parameters": + cur.args_start = pos + expect = state + + for i, ch in enumerate(text): + if in_string: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_string = False + if depth == 2 and cur is not None: + if expect == _EXPECT_KEY: + key = text[string_start + 1 : i] + expect = _EXPECT_COLON + elif expect == _IN_STRING_VALUE: + end_value(i + 1) + continue + + if ch == '"': + in_string = True + string_start = i + if depth == 2 and cur is not None and expect == _EXPECT_VALUE: + start_value(i, _IN_STRING_VALUE) + continue + + if ch in "{[": + if depth == 0: + if ch != "[": + return spans + elif depth == 1: + if ch != "{": + return spans + cur = _RequiredToolCallSpan() + spans.append(cur) + expect = _EXPECT_KEY + key = None + elif depth == 2 and cur is not None and expect == _EXPECT_VALUE: + start_value(i, _IN_CONTAINER_VALUE) + depth += 1 + continue + + if ch in "}]": + if depth == 2 and cur is not None and expect == _IN_SCALAR_VALUE: + end_value(i) + depth -= 1 + if depth == 2 and cur is not None and expect == _IN_CONTAINER_VALUE: + end_value(i + 1) + elif depth == 1: + cur = None + elif depth <= 0: + return spans + continue + + if depth != 2 or cur is None: + continue + + if ch == ":": + if expect == _EXPECT_COLON: + expect = _EXPECT_VALUE + elif ch == ",": + if expect == _IN_SCALAR_VALUE: + end_value(i) + expect = _EXPECT_KEY + key = None + elif ch.isspace(): + if expect == _IN_SCALAR_VALUE: + end_value(i) + elif expect == _EXPECT_VALUE: + start_value(i, _IN_SCALAR_VALUE) + + return spans + + def extract_required_tool_call_streaming( *, previous_text: str, current_text: str | None, - delta_text: str, - function_name_returned: bool, tool_call_idx: int | None, tool_call_id_type: str, ) -> tuple[DeltaMessage | None, bool]: - if current_text is None or current_text == "": - # if the current text is empty, we cannot parse it - return None, function_name_returned - try: - flags = Allow.ALL - obj, _ = partial_json_loads(current_text, flags) - except ( - partial_json_parser.core.exceptions.MalformedJSON, - json.JSONDecodeError, - ): - obj = None - - # check if the current text is a valid array - # containing a partial tool calling object - # if not repeat - if obj is None or not isinstance(obj, list) or not len(obj) > 0: - function_name_returned = False - delta_message = None - else: - _, finishes_previous_tool = filter_delta_text(delta_text, previous_text) - # take the last tool call from the generated list - current_tool_call = obj[-1] - - # once parameters have been generated the name is complete as well - if not finishes_previous_tool and ( - "name" not in current_tool_call or "parameters" not in current_tool_call - ): - function_name_returned = False - delta_message = None - else: - if not function_name_returned: - # get partly generated arguments from the latest tool call - param_match = re.search( - r'.*"parameters":\s*(.*)', current_text, re.DOTALL - ) - if param_match: - arguments = param_match.group(1) - arguments_prefix = current_text[: param_match.start(1)] - arguments, _ = filter_delta_text(arguments, arguments_prefix) - else: - arguments = "" - - # if this iteration finishes a previous tool call but a - # new incomplete tool is already generated, take the - # previous from the list - if finishes_previous_tool and "parameters" not in current_tool_call: - current_tool_call = obj[-2] - - function_name_returned = True - tool_call_id = make_tool_call_id( - id_type=tool_call_id_type, - func_name=current_tool_call["name"], - idx=tool_call_idx, - ) - delta_message = DeltaMessage( - tool_calls=[ - DeltaToolCall( - id=tool_call_id, - function=DeltaFunctionCall( - name=current_tool_call["name"], arguments=arguments - ), - index=len(obj) - 1, - type="function", - ) - ] - ) + """Stream the tool calls of a ``tool_choice="required"`` JSON array. + + ``previous_text`` must be a prefix of ``current_text`` (the accumulated + text before and after this delta). Every array element that has something + not yet sent is emitted, in array order, so a delta may carry several + ``DeltaToolCall`` entries: the first chunk of a call carries its id, name + and the arguments generated so far; later chunks carry only new argument + text. What was already sent is read off the same scan, since the spans + are prefix-stable, so no state is needed beyond the two texts. - else: - delta_text, _ = filter_delta_text(delta_text, previous_text) - - if delta_text != "": - delta_message = DeltaMessage( - tool_calls=[ - DeltaToolCall( - function=DeltaFunctionCall( - # OpenAI API returns None - # instead of name every time - name=None, - arguments=delta_text, - ), - index=len(obj) - 1, - ) - ] + Returns the delta (``None`` when nothing is new) and whether at least one + call has been announced so far. + """ + if not current_text: + return None, False + + spans = _scan_required_tool_calls(current_text) + sent_len = len(previous_text) + + tool_calls: list[DeltaToolCall] = [] + started = 0 + for index, span in enumerate(spans): + if not span.ready: + # Elements are generated in order: nothing later can be ready. + break + assert span.args_start is not None + end = span.args_end if span.args_end is not None else len(current_text) + if span.ready_at(sent_len): + # Already announced: send only the new argument text. + sent_end = min(end, sent_len) + new_arguments = current_text[sent_end:end] + if new_arguments: + tool_calls.append( + DeltaToolCall( + index=index, + function=DeltaFunctionCall( + # OpenAI API returns None instead of name every time + name=None, + arguments=new_arguments, + ), ) - else: - delta_message = None + ) + continue - return delta_message, function_name_returned + idx = None if tool_call_idx is None else tool_call_idx + started + started += 1 + tool_calls.append( + DeltaToolCall( + id=make_tool_call_id( + id_type=tool_call_id_type, func_name=span.name, idx=idx + ), + type="function", + index=index, + function=DeltaFunctionCall( + name=span.name, + arguments=current_text[span.args_start : end], + ), + ) + ) + + function_name_returned = any(span.ready for span in spans) + if not tool_calls: + return None, function_name_returned + return DeltaMessage(tool_calls=tool_calls), function_name_returned From 8733a21d5517d56fb7ef573b379eaeb7c08e1d03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 03:22:04 +0200 Subject: [PATCH 164/952] [Core] Bound execute_dummy_batch RPC waits by the execute-model timeout (#60398) Signed-off-by: NickLucche --- tests/v1/executor/test_multiproc_executor.py | 22 +++++++++++++------- vllm/v1/executor/multiproc_executor.py | 6 +++++- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/v1/executor/test_multiproc_executor.py b/tests/v1/executor/test_multiproc_executor.py index 706bfae38189..6a3b3d516ebf 100644 --- a/tests/v1/executor/test_multiproc_executor.py +++ b/tests/v1/executor/test_multiproc_executor.py @@ -70,16 +70,24 @@ def fail(): @pytest.mark.parametrize("stalled", [False, True]) -def test_take_draft_token_ids_uses_execute_model_timeout(monkeypatch, stalled): - """A wedged draft-token readback must not block EngineCore forever.""" +@pytest.mark.parametrize( + "method,result", + [ + ("take_draft_token_ids", DraftTokenIds(["req"], [[1, 2, 3]])), + ("execute_dummy_batch", None), + ], +) +def test_model_rpc_uses_execute_model_timeout(monkeypatch, method, result, stalled): + """A wedged model RPC must not block EngineCore forever. For + execute_dummy_batch this is an idle DP rank stuck in a collective with a + hung peer.""" monkeypatch.setenv("VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", "7") - draft = DraftTokenIds(["req"], [[1, 2, 3]]) def dequeue(*, timeout=None): assert timeout is not None and 0 < timeout <= 7 if stalled: raise TimeoutError - return WorkerProc.ResponseStatus.SUCCESS, draft + return WorkerProc.ResponseStatus.SUCCESS, result executor: Any = MultiprocExecutor.__new__(MultiprocExecutor) executor.is_failed = False @@ -89,7 +97,7 @@ def dequeue(*, timeout=None): executor.response_mqs = [None, SimpleNamespace(dequeue=dequeue)] if stalled: - with pytest.raises(TimeoutError, match="take_draft_token_ids timed out"): - executor.take_draft_token_ids() + with pytest.raises(TimeoutError, match=f"{method} timed out"): + getattr(executor, method)() else: - assert executor.take_draft_token_ids() is draft + assert getattr(executor, method)() is result diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 78290d25bafc..cad34368546d 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -364,7 +364,11 @@ def sample_tokens( # type: ignore[override] ) def execute_dummy_batch(self) -> None: - self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank) + self.collective_rpc( + "execute_dummy_batch", + unique_reply_rank=self.output_rank, + timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, + ) def take_draft_token_ids(self) -> DraftTokenIds | None: # OPTIMIZATION: Get output only from a single worker (output_rank) From 9dba339d07eaf16c9aa874eb042838f20c3ee0c8 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 18:38:06 -0700 Subject: [PATCH 165/952] [Bugfix][ROCm] Keep MLA derived weights in place on weight reload (#60389) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) --- .../model_loader/test_reload.py | 43 ++++++++++++++++++- .../layers/attention/mla_attention.py | 38 ++++++++++------ 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 85ddea39469a..aaad16778f12 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1051,6 +1051,37 @@ def loaded(value): assert not hasattr(rnorm, "qnorm_weight") +def _device_tensor_ptrs(worker) -> dict[str, int]: + """Addresses of every device tensor held by a module; graphs capture them.""" + ptrs = {} + for name, module in worker.get_model().named_modules(): + tensors = {**vars(module), **module._parameters, **module._buffers} + for attr, tensor in tensors.items(): + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + ptrs[f"{name}.{attr}"] = tensor.data_ptr() + return ptrs + + +def _moved_tensors(llm, ptrs: list[dict[str, int]]) -> list[list[str]]: + after = llm.collective_rpc(_device_tensor_ptrs) + return [ + sorted(n for n, p in old.items() if new.get(n) != p) + for old, new in zip(ptrs, after) + ] + + +@pytest.mark.parametrize( + "use_aiter", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER is ROCm-only" + ), + ), + ], +) @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) @@ -1094,13 +1125,19 @@ def loaded(value): ), ], ) -def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): +def test_reload_weights( + base_model, mul_model, add_model, tp_size, use_aiter, vllm_runner, monkeypatch +): if current_platform.device_count() < tp_size: pytest.skip(reason="Not enough CUDA devices") if "FP8" in base_model and _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + if use_aiter: + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + with vllm_runner( model_name=base_model, tensor_parallel_size=tp_size, @@ -1109,12 +1146,16 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): max_model_len=16, max_num_seqs=1, ) as llm: + # Reloading must update tensors in place: captured graphs keep pointers. + ptrs = llm.collective_rpc(_device_tensor_ptrs) llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert mul_perp < add_perp llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model}) + assert _moved_tensors(llm, ptrs) == [[]] * len(ptrs) mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] assert add_perp < mul_perp diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f550dc3d90d4..4286502ea9fe 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1272,23 +1272,34 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): quark_quantize_weight_to_mxfp4, ) - self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) - # Convert from (L, N, P) to (N, L, P) - self.W_K = self.W_K.transpose(0, 1) - self.W_K_scale = self.W_K_scale.transpose(0, 1) - - self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4( - W_UV.permute(1, 2, 0) - ) + W_K, W_K_scale = quark_quantize_weight_to_mxfp4(W_UK) + W_V, W_V_scale = quark_quantize_weight_to_mxfp4(W_UV.permute(1, 2, 0)) + # Copy in place on reload so captured CUDA graphs stay valid. + # Convert W_K from (L, N, P) to (N, L, P). + for name, tensor in ( + ("W_K", W_K.transpose(0, 1)), + ("W_K_scale", W_K_scale.transpose(0, 1)), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) elif self.is_aiter_triton_fp8_bmm_enabled: W_K = W_UK.transpose(0, 1) # 16 512 128 W_V = W_UV.permute(1, 2, 0) # 16 128 512 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( + W_K, W_K_scale = dynamic_per_batched_tensor_quant( W_K, dtype=current_platform.fp8_dtype() ) - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( + W_V, W_V_scale = dynamic_per_batched_tensor_quant( W_V, dtype=current_platform.fp8_dtype() ) + # Copy in place on reload so captured CUDA graphs stay valid. + for name, tensor in ( + ("W_K", W_K), + ("W_K_scale", W_K_scale), + ("W_V", W_V), + ("W_V_scale", W_V_scale), + ): + replace_parameter(self, name, tensor, prefer_copy=True) # The kernel operates on non-padded inputs. Hence, pre-compiling # triton kernel to avoid runtime compilation for unseen batch sizes @@ -1327,8 +1338,11 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): # Convert from (L, N, P) to (N, P, L) replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) if self.dcp_q_replicate: - self.W_UK_T_dcp_qrep = get_dcp_group().all_gather( - self.W_UK_T.contiguous(), dim=0 + replace_parameter( + self, + "W_UK_T_dcp_qrep", + get_dcp_group().all_gather(self.W_UK_T.contiguous(), dim=0), + prefer_copy=True, ) # If we should not load quant weights, we initialize the scales to 1.0 From a4b14e57b6d5ad25a2215eff93d6f93ba05005df Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Thu, 8 Oct 2026 09:56:19 +0800 Subject: [PATCH 166/952] [Platform] Allow BackendEnum dynamic register (#41091) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_attention_backend_registry.py | 141 ++++++++++++++++++ .../v1/attention/test_mla_prefill_registry.py | 120 +++++++++++---- .../backends/mla/prefill/registry.py | 54 ++++++- vllm/v1/attention/backends/registry.py | 85 ++++++++++- 4 files changed, 364 insertions(+), 36 deletions(-) diff --git a/tests/test_attention_backend_registry.py b/tests/test_attention_backend_registry.py index 5885411ed96a..14fe06699d01 100644 --- a/tests/test_attention_backend_registry.py +++ b/tests/test_attention_backend_registry.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest import torch from vllm.platforms.interface import DeviceCapability @@ -71,6 +72,48 @@ def get_builder_cls(): return None +class StringDecoratorAttentionBackend(AttentionBackend): + """Mock backend defined at module level so get_class() can resolve + the auto-generated class path (local classes are not importable).""" + + @staticmethod + def get_name(): + return "DECORATED" + + @staticmethod + def get_impl_cls(): + return CustomAttentionImpl + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_required_kv_cache_layout(): + return None + + +class StringDecoratorMambaBackend(AttentionBackend): + """Mock mamba backend defined at module level so get_class() can + resolve the auto-generated class path.""" + + @staticmethod + def get_name(): + return "DECORATED" + + @staticmethod + def get_impl_cls(): + return CustomMambaAttentionImpl + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_required_kv_cache_layout(): + return None + + def test_custom_attention_backend_requires_dcp_opt_in(): """A new implementation remains usable without implicitly opting into DCP.""" config = AttentionSelectorConfig( @@ -178,3 +221,101 @@ def test_register_custom_mamba_backend_with_class_path(): backend_cls = MambaAttentionBackendEnum.CUSTOM.get_class() assert backend_cls.get_name() == "CUSTOM_MAMBA" assert backend_cls.get_impl_cls() == CustomMambaAttentionImpl + + +@pytest.fixture(params=[AttentionBackendEnum, MambaAttentionBackendEnum]) +def _enum_cls(request): + return request.param + + +def _remove_dynamic_member(enum_cls, member): + value = member.value + name = member.name + enum_cls._member_map_.pop(name, None) + enum_cls._member_names_.remove(name) + delattr(enum_cls, name) + if enum_cls._value2member_map_.get(value) is member: + del enum_cls._value2member_map_[value] + + +def test_register_dynamic_enum_member(_enum_cls): + backend_path = { + AttentionBackendEnum: ( + "tests.test_attention_backend_registry.CustomAttentionBackend" + ), + MambaAttentionBackendEnum: ( + "tests.test_attention_backend_registry.CustomMambaAttentionBackend" + ), + }[_enum_cls] + + member = _enum_cls.register("DYNAMIC_TEST", backend_path) + assert member.name == "DYNAMIC_TEST" + assert member.value == backend_path + assert member is _enum_cls.DYNAMIC_TEST + assert _enum_cls["DYNAMIC_TEST"] is member + _remove_dynamic_member(_enum_cls, member) + + +def test_register_dynamic_enum_member_duplicate_raises(_enum_cls): + _enum_cls.register("DUP_TEST", "some.module.Class") + with pytest.raises(ValueError, match="already exists"): + _enum_cls.register("DUP_TEST", "other.module.OtherClass") + _remove_dynamic_member(_enum_cls, _enum_cls.DUP_TEST) + + +def test_register_dynamic_enum_member_reserved_name_raises(_enum_cls): + with pytest.raises(ValueError, match="Invalid or reserved backend name"): + _enum_cls.register("get_path", "some.module.Class") + + +@pytest.fixture( + params=[ + pytest.param((AttentionBackendEnum, False), id="attention"), + pytest.param((MambaAttentionBackendEnum, True), id="mamba"), + ] +) +def _enum_is_mamba(request): + return request.param + + +def _backend_path(enum_cls): + if enum_cls is AttentionBackendEnum: + return "tests.test_attention_backend_registry.CustomAttentionBackend" + return "tests.test_attention_backend_registry.CustomMambaAttentionBackend" + + +def test_register_backend_with_string_name_direct(_enum_is_mamba): + enum_cls, is_mamba = _enum_is_mamba + path = _backend_path(enum_cls) + register_backend("STRING_DIRECT", path, is_mamba=is_mamba) + + member = enum_cls.STRING_DIRECT + assert member.is_overridden() + assert member.get_path() == path + expected = "CUSTOM" if not is_mamba else "CUSTOM_MAMBA" + assert member.get_class().get_name() == expected + member.clear_override() + _remove_dynamic_member(enum_cls, member) + + +def test_register_backend_with_string_name_decorator(_enum_is_mamba): + enum_cls, is_mamba = _enum_is_mamba + # The decorated class must be module-level: the decorator derives a + # class path from it, and get_class() re-imports that path. A class + # defined inside this test would have a non-importable path. + decorated_cls = ( + StringDecoratorMambaBackend if is_mamba else StringDecoratorAttentionBackend + ) + + result = register_backend("STRING_DECORATOR", is_mamba=is_mamba)(decorated_cls) + assert result is decorated_cls + + member = enum_cls.STRING_DECORATOR + assert member.is_overridden() + assert member.get_path() == ( + f"{decorated_cls.__module__}.{decorated_cls.__qualname__}" + ) + assert member.get_class() is decorated_cls + assert member.get_class().get_name() == "DECORATED" + member.clear_override() + _remove_dynamic_member(enum_cls, member) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index fc5efe3a5be0..0068d53b88cf 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -54,10 +54,20 @@ def test_prefill_backend_clone_has_isolated_metadata(): def cleanup_overrides(): """Clear any overrides after each test.""" yield - for member in MLAPrefillBackendEnum: + for member in list(MLAPrefillBackendEnum): member.clear_override() +def _remove_dynamic_member(enum_cls, member): + value = member.value + name = member.name + enum_cls._member_map_.pop(name, None) + enum_cls._member_names_.remove(name) + delattr(enum_cls, name) + if enum_cls._value2member_map_.get(value) is member: + del enum_cls._value2member_map_[value] + + def test_custom_is_not_alias_of_any_backend(): all_backends = list(MLAPrefillBackendEnum) @@ -83,21 +93,11 @@ def test_custom_unregistered_raises(): def test_register_custom_backend_with_class_path(): register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.CUSTOM, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.CUSTOM, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - assert MLAPrefillBackendEnum.CUSTOM.is_overridden() - - class_path = MLAPrefillBackendEnum.CUSTOM.get_path() - assert class_path == ( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ) - - backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class() - assert backend_cls.get_name() == "CUSTOM" + assert MLAPrefillBackendEnum.CUSTOM.get_class().get_name() == "CUSTOM" def test_register_custom_backend_as_decorator(): @@ -120,33 +120,20 @@ def run_prefill_context_chunk(self, chunk, q, k, v, out=None): def test_override_existing_backend(): - original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path() - register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.FLASH_ATTN, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - + assert MLAPrefillBackendEnum.FLASH_ATTN.get_class().get_name() == "CUSTOM" assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() - assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path - - backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() - assert backend_cls.get_name() == "CUSTOM" def test_clear_override(): original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path() - register_mla_prefill_backend( - backend=MLAPrefillBackendEnum.FLASH_ATTN, - class_path=( - "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" - ), + MLAPrefillBackendEnum.FLASH_ATTN, + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", ) - assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() - MLAPrefillBackendEnum.FLASH_ATTN.clear_override() assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden() assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path @@ -172,3 +159,72 @@ def test_rocm_aiter_fa_registered(): # The AITER FA path is the fp16/bf16 generic-varlen prefill path. assert backend_cls.supports_dtype(torch.bfloat16) assert backend_cls.supports_dtype(torch.float16) + + +def test_register_dynamic_member(): + p = "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend" + member = MLAPrefillBackendEnum.register("DYNAMIC_TEST", p) + assert member.name == "DYNAMIC_TEST" + assert member is MLAPrefillBackendEnum["DYNAMIC_TEST"] + + _remove_dynamic_member(MLAPrefillBackendEnum, member) + + +def test_register_dynamic_member_duplicate_raises(): + path = "some.module.Class" + MLAPrefillBackendEnum.register("DUP_TEST", path) + with pytest.raises(ValueError, match="already exists"): + MLAPrefillBackendEnum.register("DUP_TEST", "other.module.OtherClass") + _remove_dynamic_member(MLAPrefillBackendEnum, MLAPrefillBackendEnum.DUP_TEST) + + +class StringDecoratorMLAPrefillBackend(MLAPrefillBackend): + """Mock backend defined at module level so get_class() can resolve + the auto-generated class path (local classes are not importable).""" + + supported_dtypes = [torch.bfloat16] + requires_r1_mla_dimensions = False + + @staticmethod + def get_name() -> str: + return "DECORATED" + + def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): + raise NotImplementedError + + def run_prefill_context_chunk(self, chunk_idx, q, k, v): + raise NotImplementedError + + +def test_register_mla_prefill_backend_with_string_name_direct(): + register_mla_prefill_backend( + "STRING_DIRECT", + "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend", + ) + member = MLAPrefillBackendEnum.STRING_DIRECT + assert member.is_overridden() + assert member.get_class().get_name() == "CUSTOM" + + member.clear_override() + _remove_dynamic_member(MLAPrefillBackendEnum, member) + + +def test_register_mla_prefill_backend_with_string_name_decorator(): + # The decorated class must be module-level: the decorator derives a + # class path from it, and get_class() re-imports that path. A class + # defined inside this test would have a non-importable path. + decorated_cls = StringDecoratorMLAPrefillBackend + + result = register_mla_prefill_backend("STRING_DECORATOR")(decorated_cls) + assert result is decorated_cls + + member = MLAPrefillBackendEnum.STRING_DECORATOR + assert member.is_overridden() + assert member.get_path() == ( + f"{decorated_cls.__module__}.{decorated_cls.__qualname__}" + ) + assert member.get_class() is decorated_cls + assert member.get_class().get_name() == "DECORATED" + + member.clear_override() + _remove_dynamic_member(MLAPrefillBackendEnum, member) diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index f7ec31e15361..19d37fbc225b 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -10,11 +10,14 @@ from enum import Enum, EnumMeta from typing import TYPE_CHECKING +from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname if TYPE_CHECKING: from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +logger = init_logger(__name__) + class _MLAPrefillBackendEnumMeta(EnumMeta): """Metaclass for MLAPrefillBackendEnum to provide better error messages.""" @@ -61,6 +64,37 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None + @classmethod + def register(cls, name: str, value: str) -> "MLAPrefillBackendEnum": + """Dynamically register a new MLA prefill backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"MLA prefill backend {name} already exists in {cls.__name__}. " + f"Use register_mla_prefill_backend(" + f"{cls.__name__}.{name}, '{value}') to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new MLA prefill backend: %s -> %s", name, value) + return member + def get_path(self) -> str: """Get the class path for this backend (respects overrides). @@ -107,13 +141,15 @@ def clear_override(self) -> None: def register_mla_prefill_backend( - backend: MLAPrefillBackendEnum, + backend: MLAPrefillBackendEnum | str, class_path: str | None = None, ) -> Callable[[type], type]: """Register or override an MLA prefill backend implementation. Args: - backend: The MLAPrefillBackendEnum member to register. + backend: The MLAPrefillBackendEnum member to register, + or a string name for a new custom backend + (e.g., "CUSTOM_MLA_PREFILL"). class_path: Optional class path. If not provided and used as decorator, will be auto-generated from the class. @@ -131,13 +167,27 @@ class MyCustomFlashAttn(MLAPrefillBackend): class MyCustomPrefillBackend(MLAPrefillBackend): ... + # Register a new custom MLA prefill backend with a dynamic enum name + @register_mla_prefill_backend("CUSTOM_MLA_PREFILL") + class CustomMLAPrefillBackend(MLAPrefillBackend): + ... + # Direct registration register_mla_prefill_backend( MLAPrefillBackendEnum.CUSTOM, "my.module.MyCustomPrefillBackend" ) + # Direct registration with string name + register_mla_prefill_backend( + "CUSTOM_MLA_PREFILL", + "custom.attention.mla_prefill.CustomMLAPrefillBackend" + ) + """ + # Handle dynamic enum creation for string backend names + if isinstance(backend, str): + backend = MLAPrefillBackendEnum.register(backend, class_path or "") def decorator(cls: type) -> type: _MLA_PREFILL_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}" diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index ad184f091329..207c112a384c 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -157,6 +157,38 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None + @classmethod + def register(cls, name: str, value: str) -> "AttentionBackendEnum": + """Dynamically register a new backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"Backend {name} already exists in {cls.__name__}. " + f"Use register_backend({cls.__name__}.{name}, '{value}') " + f"to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + # Create new enum member dynamically + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new attention backend: %s -> %s", name, value) + return member + def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). @@ -214,6 +246,37 @@ class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): backend.get_class() """ + @classmethod + def register(cls, name: str, value: str) -> "MambaAttentionBackendEnum": + """Dynamically register a new mamba backend enum member. + + Args: + name: The name for the new enum member + value: The fully qualified class path string + + Returns: + The newly created enum member + + """ + if name in cls._member_map_: + raise ValueError( + f"Backend {name} already exists in {cls.__name__}. " + f"Use register_backend({cls.__name__}.{name}, '{value}') " + f"to override." + ) + if not name.isidentifier() or hasattr(cls, name): + raise ValueError(f"Invalid or reserved backend name: {name}") + + member = object.__new__(cls) + member._name_ = name + member._value_ = value + setattr(cls, name, member) + cls._member_map_[name] = member + cls._value2member_map_[value] = member + cls._member_names_.append(name) + logger.info("Registered new attention backend: %s -> %s", name, value) + return member + MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend" MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend" SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend" @@ -275,14 +338,15 @@ def clear_override(self) -> None: def register_backend( - backend: AttentionBackendEnum | MambaAttentionBackendEnum, + backend: AttentionBackendEnum | MambaAttentionBackendEnum | str, class_path: str | None = None, is_mamba: bool = False, ) -> Callable[[type], type]: """Register or override a backend implementation. Args: - backend: The AttentionBackendEnum member to register + backend: The AttentionBackendEnum/MambaAttentionBackendEnum member to register, + or a string name for a new custom backend (e.g., "CUSTOM_MLA"). class_path: Optional class path. If not provided and used as decorator, will be auto-generated from the class. is_mamba: Whether the backend is a Mamba attention backend. @@ -306,13 +370,30 @@ class MyCustomMambaAttn: class MyCustomBackend: ... + # Register a new custom attention backend with a dynamic enum name + @register_backend("CUSTOM_MLA") + class CustomMLABackend: + ... + # Direct registration register_backend( AttentionBackendEnum.CUSTOM, "my.module.MyCustomBackend" ) + # Direct registration with string name + register_backend( + "CUSTOM_MLA", + "custom.attention.mla_v1.CustomMLAABackend" + ) + """ + # Handle dynamic enum creation for string backend names + if isinstance(backend, str): + if is_mamba: + backend = MambaAttentionBackendEnum.register(backend, class_path or "") + else: + backend = AttentionBackendEnum.register(backend, class_path or "") def decorator(cls: type) -> type: if is_mamba: From caff2b44cb4e37c81f2cc09234f557fd7cd96505 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 10:30:35 +0800 Subject: [PATCH 167/952] [BugFix] Pick the DP world-group port at bind time via the coordination store (#51018) Signed-off-by: aoshen02 Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Co-authored-by: Codex --- .../engine/test_core_engine_actor_manager.py | 69 ++++++++++- tests/v1/engine/test_dp_world_group_port.py | 109 ++++++++++++++++++ vllm/distributed/parallel_state.py | 48 +++++++- vllm/v1/engine/utils.py | 18 +++ 4 files changed, 240 insertions(+), 4 deletions(-) create mode 100644 tests/v1/engine/test_dp_world_group_port.py diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index f1a67ad9ba28..768f51f4379b 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -13,8 +13,9 @@ import pytest import ray import zmq +from torch.distributed import TCPStore -from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import get_open_port, make_zmq_socket, split_zmq_path from vllm.v1.engine.core import EngineCoreActorMixin from vllm.v1.engine.core_client import BackgroundResources from vllm.v1.engine.utils import ( @@ -134,6 +135,72 @@ def _make_addresses() -> EngineZmqAddresses: ) +@pytest.mark.parametrize( + ("overrides", "expect_store"), + [ + ({"data_parallel_backend": "mp"}, True), + ({}, True), + ( + { + "data_parallel_backend": "mp", + "data_parallel_rank": 1, + "local_engines_only": True, + }, + False, + ), + ({"data_parallel_backend": "mp", "data_parallel_rank_local": 0}, False), + ({"data_parallel_backend": "mp", "enable_elastic_ep": True}, False), + ({"data_parallel_backend": "mp", "data_parallel_size": 1}, False), + ], + ids=["mp", "ray", "non-master-node", "offline", "elastic-ep", "single-engine"], +) +def test_coordination_store_held_only_by_the_online_dp_master( + monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any], expect_store: bool +) -> None: + """The DP master node holds a coordination store for both engine backends, + so engines there pick world-group ports at bind time. Every other + deployment keeps the pre-allocated ports and must not hold one.""" + + class FakeManager: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + vllm_config = _make_vllm_config_ray_dp_multinode() + parallel_config = vllm_config.parallel_config + for name, value in overrides.items(): + setattr(parallel_config, name, value) + # A fresh port per case so a lingering handshake socket cannot fail the next. + parallel_config.data_parallel_rpc_port = get_open_port() + parallel_config._coord_store_port = 0 + monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineProcManager", FakeManager) + monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineActorManager", FakeManager) + monkeypatch.setattr( + "vllm.v1.engine.utils.wait_for_engine_startup", lambda *args, **kwargs: None + ) + + with launch_core_engines( + vllm_config, + executor_class=_DummyExecutor, + log_stats=False, + addresses=_make_addresses(), + ) as engine_launch: + assert engine_launch.engine_manager is not None + if not expect_store: + assert not parallel_config._coord_store_port + return + assert parallel_config._coord_store_port + # Engines look the store up while they start, so it must be reachable + # for as long as this frame is alive. + client = TCPStore( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + is_master=False, + wait_for_workers=False, + ) + client.set("probe", b"1") + assert client.get("probe") == b"1" + + def _make_cpu_placement_group(): pg = ray.util.placement_group( [{"CPU": 0.001}, {"CPU": 1.0}], diff --git a/tests/v1/engine/test_dp_world_group_port.py b/tests/v1/engine/test_dp_world_group_port.py new file mode 100644 index 000000000000..49f5a59ab5c4 --- /dev/null +++ b/tests/v1/engine/test_dp_world_group_port.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DP world group picks its port when it binds, not ahead of time.""" + +import multiprocessing as mp +import socket +import traceback +from multiprocessing.queues import Queue + +import pytest +import torch + +from vllm.distributed.utils import create_tcp_store +from vllm.utils.network_utils import get_open_port + +RESULT_TIMEOUT_S = 240 + + +def _join_dp_world_group( + dp_rank: int, coord_store_port: int, dp_master_port: int, result: Queue +) -> None: + """Bring up one DP rank of the world group and report what it saw.""" + try: + import torch.distributed as dist + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.parallel import ParallelConfig + from vllm.distributed.parallel_state import init_distributed_environment + + parallel_config = ParallelConfig( + data_parallel_size=2, data_parallel_rank=dp_rank + ) + parallel_config.data_parallel_master_ip = "127.0.0.1" + parallel_config._data_parallel_master_port_list = [dp_master_port] + parallel_config._coord_store_port = coord_store_port + + with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)): + init_distributed_environment( + world_size=1, rank=0, local_rank=0, backend="gloo" + ) + + reduced = torch.tensor([dp_rank + 1]) + dist.all_reduce(reduced) + result.put((dp_rank, int(reduced.item()), None)) + except Exception: + result.put((dp_rank, None, traceback.format_exc())) + + +@pytest.mark.skipif( + not torch.distributed.is_gloo_available(), reason="gloo backend required" +) +@pytest.mark.parametrize("with_coord_store", [True, False]) +def test_dp_world_group_survives_a_taken_preallocated_port(with_coord_store: bool): + """With a coordination store, the two ranks agree on a port the group binds + itself, so a pre-allocated port taken in the meantime cannot stop startup. + Without one, the pre-allocated port is still used as before. + """ + squatter = None + coord_store = None + try: + if with_coord_store: + # Something else grabbed the pre-allocated port after it was picked. + squatter = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + squatter.bind(("127.0.0.1", 0)) + squatter.listen(1) + dp_master_port = squatter.getsockname()[1] + coord_store = create_tcp_store( + "127.0.0.1", 0, is_master=True, world_size=-1, wait_for_workers=False + ) + coord_store_port = coord_store.port + else: + dp_master_port = get_open_port() + coord_store_port = 0 + + ctx = mp.get_context("spawn") + result: Queue = ctx.Queue() + procs = [ + ctx.Process( + target=_join_dp_world_group, + args=(dp_rank, coord_store_port, dp_master_port, result), + ) + for dp_rank in range(2) + ] + for proc in procs: + proc.start() + reports: list[tuple[int, int | None, str | None]] = [] + try: + while len(reports) < len(procs): + reports.append(result.get(timeout=RESULT_TIMEOUT_S)) + if reports[-1][2] is not None: + break # the peer would wait on this rank; do not wait for it + finally: + for proc in procs: + proc.join(timeout=30) + if proc.is_alive(): + proc.kill() + + for dp_rank, reduced, error in sorted(reports): + assert error is None, f"DP rank {dp_rank} failed:\n{error}" + # Both ranks joined the same group: 1 + 2 over the two ranks. + assert reduced == 3, f"DP rank {dp_rank} all-reduced to {reduced}" + + if with_coord_store: + assert coord_store is not None + bound_port = int(coord_store.get("world_pg_port").decode()) + assert bound_port != dp_master_port + finally: + if squatter is not None: + squatter.close() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index de6680e02745..feb8f25d1eeb 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -42,6 +42,7 @@ import torch.distributed._functional_collectives as funcol import torch.distributed._symmetric_memory from torch.distributed import Backend, ProcessGroup, Store +from torch.distributed.constants import default_pg_timeout import vllm.envs as envs from vllm.distributed.device_communicators.base_device_communicator import ( @@ -49,6 +50,7 @@ ) from vllm.distributed.utils import ( StatelessProcessGroup, + create_tcp_store, get_cached_tcp_store_client, ) from vllm.logger import init_logger @@ -1829,6 +1831,7 @@ def init_distributed_environment( config = get_current_vllm_config_or_none() enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep + world_pg_store: Store | None = None if ( config is not None and config.parallel_config.distributed_executor_backend != "external_launcher" @@ -1852,7 +1855,46 @@ def init_distributed_environment( distributed_init_method = get_distributed_init_method(ip, port) else: ip = parallel_config.data_parallel_master_ip - port = parallel_config.get_next_dp_init_port() + if ( + parallel_config._coord_store_port + and not torch.distributed.is_initialized() + ): + # Group rank 0 lets the world-group TCPStore bind port 0 + # and publishes the kernel-assigned port via the + # coordination store. A one-shot client of our own: the + # cached one is inherited across fork and shared with + # sibling workers. + coord_store = create_tcp_store( + ip, + parallel_config._coord_store_port, + is_master=False, + wait_for_workers=False, + ) + if rank == 0: + # Workers can only connect once the port is published. + world_pg_store = create_tcp_store( + ip, + 0, + world_size=world_size, + is_master=True, + wait_for_workers=False, + timeout=timeout or default_pg_timeout, + multi_tenant=True, + ) + port = world_pg_store.port + coord_store.set("world_pg_port", str(port).encode()) + else: + port = int(coord_store.get("world_pg_port").decode()) + world_pg_store = create_tcp_store( + ip, + port, + world_size=world_size, + is_master=False, + timeout=timeout or default_pg_timeout, + multi_tenant=True, + ) + else: + port = parallel_config.get_next_dp_init_port() distributed_init_method = get_distributed_init_method(ip, port) logger.debug( "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP", @@ -1882,8 +1924,8 @@ def init_distributed_environment( "Fallback Gloo backend is not available." ) backend = "gloo" - store = None - if distributed_init_method.startswith("file://"): + store = world_pg_store + if store is None and distributed_init_method.startswith("file://"): store = torch.distributed.FileStore( distributed_init_method.removeprefix("file://"), world_size ) diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 203d1b407c11..49e6884ac9f1 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -1195,6 +1195,23 @@ def launch_core_engines( else: coordinator = None + # Hold a coordination TCPStore (alive for this frame) so engines pick DP + # master ports at bind time; pre-allocated ports can be taken before use. + coord_store = None + if ( + dp_size > 1 + and not offline_mode + and dp_rank == 0 + and local_engine_count > 0 + and not parallel_config.enable_elastic_ep + ): + from vllm.distributed.utils import create_tcp_store + + coord_store = create_tcp_store( + host, 0, is_master=True, world_size=-1, wait_for_workers=False + ) + parallel_config._coord_store_port = coord_store.port + if parallel_config.data_parallel_backend == "ray": logger.info("Starting ray-based data parallel backend") @@ -1415,6 +1432,7 @@ def wait_for_engine_startup( "data_parallel_master_ip", "data_parallel_master_port", "_data_parallel_master_port_list", + "_coord_store_port", "data_parallel_size", ) } From 3722645539f4cbf11028e69eea0643dee5cdc773 Mon Sep 17 00:00:00 2001 From: bnkaiteam Date: Thu, 8 Oct 2026 11:32:35 +0900 Subject: [PATCH 168/952] [Bugfix][Frontend] Accept successful 2xx responses for batch output uploads (#60507) Signed-off-by: bnkaiteam Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/entrypoints/launchers/test_run_batch.py | 29 ++++++++++++++----- vllm/entrypoints/launchers/run_batch.py | 4 +-- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py index f7a250580727..1376eb0645f3 100644 --- a/tests/entrypoints/launchers/test_run_batch.py +++ b/tests/entrypoints/launchers/test_run_batch.py @@ -1071,17 +1071,32 @@ def _make_aiohttp_put_session(status: int = 200, body_text: str = ""): @pytest.mark.asyncio -async def test_upload_data_uploads_once_on_success(): - """A successful upload must not be retried (regression guard).""" - session = _make_aiohttp_put_session(status=200) - with patch( - "vllm.entrypoints.launchers.run_batch.aiohttp.ClientSession", - return_value=session, +@pytest.mark.parametrize("status", [200, 201, 204]) +@pytest.mark.parametrize("from_file", [False, True]) +async def test_upload_data_uploads_once_on_success(status, from_file, tmp_path): + """Successful PUT responses must not retry either upload path.""" + session = _make_aiohttp_put_session(status=status) + data_or_file = "payload" + if from_file: + path = tmp_path / "output.jsonl" + path.write_text(data_or_file, encoding="utf-8") + data_or_file = str(path) + + with ( + patch( + "vllm.entrypoints.launchers.run_batch.aiohttp.ClientSession", + return_value=session, + ), + patch( + "vllm.entrypoints.launchers.run_batch.asyncio.sleep", + new_callable=AsyncMock, + ) as sleep, ): await upload_data( - "https://example.com/output.jsonl", "payload", from_file=False + "https://example.com/output.jsonl", data_or_file, from_file=from_file ) assert session.put.call_count == 1 + sleep.assert_not_awaited() @pytest.mark.asyncio diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py index f54e87dafefc..266592a731a5 100644 --- a/vllm/entrypoints/launchers/run_batch.py +++ b/vllm/entrypoints/launchers/run_batch.py @@ -385,7 +385,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No if from_file: with open(data_or_file, "rb") as file: async with session.put(output_url, data=file) as response: - if response.status != 200: + if not 200 <= response.status < 300: error_text = await response.text() raise Exception( f"Failed to upload file.\n" @@ -394,7 +394,7 @@ async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> No ) else: async with session.put(output_url, data=data_or_file) as response: - if response.status != 200: + if not 200 <= response.status < 300: error_text = await response.text() raise Exception( f"Failed to upload data.\n" From 87d9996abe7ae0173fa8c1b6ccc2411041bcac81 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 8 Oct 2026 04:37:51 +0200 Subject: [PATCH 169/952] [ROCm][Perf][GLM-5.3-Flash] BF16 splitk kernel for sparse MLA decode (#58584) Signed-off-by: simondanielsson Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../benchmark_sparse_mla_decode_bf16.py | 130 +++++++ .../attention/test_rocm_aiter_mla_sink.py | 6 + .../attention/test_rocm_triton_attn_dsv4.py | 96 +++++ .../backends/mla/rocm_aiter_mla_sparse.py | 53 ++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 355 ++++++++++++++++++ 5 files changed, 627 insertions(+), 13 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py diff --git a/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py b/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py new file mode 100644 index 000000000000..c32c8519aadb --- /dev/null +++ b/benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the ROCm bf16 sparse-MLA decode split-K kernel in isolation. + +Example usage: +python3 benchmarks/kernels/benchmark_sparse_mla_decode_bf16.py \ + --rows 1 4 16 64 --lens 2048 --splits auto 4 16 32 +""" + +import argparse +import itertools +import json + +import torch + +from vllm.triton_utils import triton +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_bf16_triton, + _rocm_sparse_attn_prefill_ragged_triton, + rocm_sparse_decode_bf16_num_splits, +) + +BLOCK_H = 16 + + +def make_inputs(rows: int, live: int, args): + torch.manual_seed(args.seed) + kv = ( + torch.randn(args.capacity, args.head_dim, device="cuda", dtype=torch.bfloat16) + * 0.125 + ) + q = ( + torch.randn( + rows, args.heads, args.head_dim, device="cuda", dtype=torch.bfloat16 + ) + * 0.125 + ) + slots = torch.stack( + [torch.randperm(args.capacity, device="cuda")[:live] for _ in range(rows)] + ) + indices = slots.reshape(-1).to(torch.int32) + indptr = torch.arange(rows + 1, device="cuda", dtype=torch.int32) * live + sink = ( + torch.randn(args.heads, device="cuda", dtype=torch.float32) + if args.sink + else None + ) + return q, kv, indices, indptr, sink + + +def resolve_splits(spec: str, rows: int, live: int, args) -> int: + if spec != "auto": + return int(spec) + return rocm_sparse_decode_bf16_num_splits(rows, args.heads, live) + + +def benchmark(backend: str, rows: int, live: int, spec: str, args) -> dict: + q, kv, indices, indptr, sink = make_inputs(rows, live, args) + nope, rope = args.head_dim, 0 + splits = resolve_splits(spec, rows, live, args) + + def ragged(): + return _rocm_sparse_attn_prefill_ragged_triton( + q, kv, indices, indptr, args.scale, sink, nope, rope + ) + + def splitk(): + return _rocm_sparse_attn_decode_ragged_bf16_triton( + q, kv, indices, indptr, args.scale, sink, nope, rope, splits + ) + + run = ragged if backend == "ragged" else splitk + torch.testing.assert_close(run(), ragged(), atol=2e-2, rtol=2e-2) + ms = triton.testing.do_bench(run, warmup=args.warmup, rep=args.rep) + return dict( + backend=backend, + rows=rows, + live=live, + heads=args.heads, + head_dim=args.head_dim, + splits=splits if backend == "splitk" else 1, + workgroups=rows * triton.cdiv(args.heads, BLOCK_H) * splits, + us=ms * 1000, + kv_gbps=rows * live * args.head_dim * 2 / (ms * 1e6), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8, 16, 64]) + parser.add_argument( + "--lens", + type=int, + nargs="+", + default=[2048], + help="Selected KV slots per row; GLM-5.3-Flash uses index_topk=2048", + ) + parser.add_argument( + "--splits", + nargs="+", + default=["auto"], + help="Split counts to sweep, or auto for the runtime heuristic", + ) + parser.add_argument( + "--backends", + nargs="+", + choices=["ragged", "splitk"], + default=["ragged", "splitk"], + ) + parser.add_argument( + "--heads", type=int, default=8, help="Heads per rank; 64 at TP8 gives 8" + ) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument("--capacity", type=int, default=65536) + parser.add_argument("--scale", type=float, default=0.044194173824159216) + parser.add_argument("--sink", action="store_true") + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--seed", type=int, default=123) + args = parser.parse_args() + + print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args)))) + for backend, rows, live in itertools.product(args.backends, args.rows, args.lens): + specs = ["1"] if backend == "ragged" else args.splits + for spec in specs: + print(json.dumps(benchmark(backend, rows, live, spec, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index 40b888c3622e..0b007dbc875c 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -156,6 +156,8 @@ def test_sparse_mla_sink_matches_ragged_reference( num_decodes=batch_size, num_decode_tokens=batch_size, max_query_len=1, + max_seq_len=max(seq_lens), + topk_tokens=max(seq_lens), ) sinks = torch.linspace(-2.0, 6.0, real_heads, device=device) @@ -367,12 +369,15 @@ def test_sparse_mla_sink_matches_dense_attention_with_empty_rows_and_paged_cache for page, offset in rows ] lengths = torch.tensor([0, *(len(rows) for rows in selected_rows)]) + longest_row = max(len(rows) for rows in selected_rows) metadata = SimpleNamespace( block_size=block_size, num_prefills=0, num_decodes=len(selected_rows), num_decode_tokens=len(selected_rows), max_query_len=1, + max_seq_len=longest_row, + topk_tokens=longest_row, qo_indptr=torch.arange( len(selected_rows) + 1, dtype=torch.int32, device="cuda" ), @@ -479,6 +484,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): num_decodes=num_tokens, num_decode_tokens=num_tokens, max_query_len=1, + max_seq_len=kv.shape[0], block_size=block_size, topk_tokens=impl.topk_indices_buffer.shape[1], req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device="cuda"), diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 3b200be514d5..ca2cf2eb7017 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -596,6 +596,87 @@ def test_sparse_attn_prefill_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +NOPE_ONLY_HEAD_DIM = 512 + + +def _nope_only_ragged_inputs( + kv_lens: list[int], num_heads: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build a GLM-like 512+0 NoPE decode batch with ragged KV segments.""" + num_kv = 4096 + kv = ( + torch.randn(num_kv, NOPE_ONLY_HEAD_DIM, dtype=torch.bfloat16, device=device) + * 0.125 + ) + q = ( + torch.randn( + len(kv_lens), + num_heads, + NOPE_ONLY_HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + rows = [torch.randperm(num_kv)[:length].tolist() for length in kv_lens] + indices, indptr = _ragged_from_rows(rows, device) + return q, kv, indices, indptr + + +# Lengths below the split count leave splits empty; non-multiples force a tail. +_NOPE_ONLY_KV_LENS = [0, 1, 3, 37, 129, 2048] + + +@requires_split_decode_arch +@pytest.mark.parametrize("num_splits", [2, 32]) +@pytest.mark.parametrize("with_sink", [True, False]) +@torch.inference_mode() +def test_sparse_attn_decode_bf16_split_k_matches_ragged( + num_splits: int, with_sink: bool +) -> None: + """Split-K decode must reproduce the single-workgroup ragged kernel.""" + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_bf16_triton, + _rocm_sparse_attn_prefill_ragged_triton, + ) + + device = torch.device("cuda") + set_random_seed(0) + num_heads = 16 + q, kv, indices, indptr = _nope_only_ragged_inputs( + _NOPE_ONLY_KV_LENS, num_heads, device + ) + attn_sink = ( + torch.randn(num_heads, dtype=torch.float32, device=device) + if with_sink + else None + ) + scale = NOPE_ONLY_HEAD_DIM**-0.5 + + expected = _rocm_sparse_attn_prefill_ragged_triton( + q=q, + kv=kv, + indices=indices, + indptr=indptr, + scale=scale, + attn_sink=attn_sink, + nope_head_dim=NOPE_ONLY_HEAD_DIM, + rope_head_dim=0, + ) + actual = _rocm_sparse_attn_decode_ragged_bf16_triton( + q=q, + kv=kv, + indices=indices, + indptr=indptr, + scale=scale, + attn_sink=attn_sink, + nope_head_dim=NOPE_ONLY_HEAD_DIM, + rope_head_dim=0, + num_splits=num_splits, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict: device = torch.device("cuda") set_random_seed(7) @@ -1105,6 +1186,21 @@ def test_decode_num_splits_gfx950(monkeypatch) -> None: assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32 assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 + assert mod._decode_gfx950_num_splits(4, 1, 2048.0, 0.0, 32) == 32 + + +def test_sparse_decode_bf16_num_splits_floor_and_tile_clamp(monkeypatch) -> None: + """Short rows skip splitting, and splits never exceed the tiles they walk.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args: 32) + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 32) + + floor = mod._SPARSE_DECODE_BF16_MIN_SPLIT_LEN + block_k = mod._SPARSE_DECODE_BF16_BLOCK_K + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, floor - 1) == 1 + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, floor) == floor // block_k + assert mod.rocm_sparse_decode_bf16_num_splits(1, 1, 2048) == 32 @requires_split_decode_arch diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e473f859480b..e3d982e882bb 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -37,7 +37,9 @@ ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_sparse_attn_decode_bf16, rocm_sparse_attn_prefill, + rocm_sparse_decode_bf16_num_splits, ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout from vllm.v1.worker.workspace import current_workspace_manager @@ -904,20 +906,45 @@ def _forward_mla( self.sinks.reshape(1, self.num_heads, 1), q.shape[1], ).reshape(-1) - rocm_sparse_attn_prefill( - q=q, - kv=kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]), - indices=None, - topk_length=None, - scale=self.scale, - head_dim=q.shape[-1], - nope_head_dim=self.kv_lora_rank, - rope_head_dim=q.shape[-1] - self.kv_lora_rank, - attn_sink=triton_sinks, - output=output, - ragged_indices=attn_metadata.paged_kv_indices, - ragged_indptr=attn_metadata.paged_kv_indptr, + kv = kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]) + decode_num_splits = ( + rocm_sparse_decode_bf16_num_splits( + num_tokens, + q.shape[1], + min(attn_metadata.max_seq_len, attn_metadata.topk_tokens), + ) + if attn_metadata.num_decode_tokens == num_tokens + else 1 ) + if decode_num_splits > 1: + rocm_sparse_attn_decode_bf16( + q=q, + kv=kv, + scale=self.scale, + head_dim=q.shape[-1], + nope_head_dim=self.kv_lora_rank, + rope_head_dim=q.shape[-1] - self.kv_lora_rank, + attn_sink=triton_sinks, + output=output, + ragged_indices=attn_metadata.paged_kv_indices, + ragged_indptr=attn_metadata.paged_kv_indptr, + num_splits=decode_num_splits, + ) + else: + rocm_sparse_attn_prefill( + q=q, + kv=kv, + indices=None, + topk_length=None, + scale=self.scale, + head_dim=q.shape[-1], + nope_head_dim=self.kv_lora_rank, + rope_head_dim=q.shape[-1] - self.kv_lora_rank, + attn_sink=triton_sinks, + output=output, + ragged_indices=attn_metadata.paged_kv_indices, + ragged_indptr=attn_metadata.paged_kv_indptr, + ) output = AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output) return output, None diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index a3412aa83508..90a73c101c0d 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2208,6 +2208,119 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _sparse_attn_decode_ragged_bf16_partial_kernel( + q_ptr, + kv_ptr, + kv_indices_ptr, + kv_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride_t, + q_stride_h, + q_stride_d, + kv_stride_n, + kv_stride_d, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + head_dim, + num_kv, + scale, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, BLOCK_D) + head_mask = head_offsets < num_heads + dim_mask = dim_offsets < head_dim + + q = tl.load( + q_ptr + + query_idx * q_stride_t + + head_offsets[:, None] * q_stride_h + + dim_offsets[None, :] * q_stride_d, + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) + + kv_start = tl.load(kv_indptr_ptr + query_idx) + kv_end = tl.load(kv_indptr_ptr + query_idx + 1) + kv_len = kv_end - kv_start + chunk = (kv_len + NUM_SPLITS - 1) // NUM_SPLITS + kv_lo = split_id * chunk + kv_hi = tl.minimum(kv_lo + chunk, kv_len) + + k_offsets = tl.arange(0, BLOCK_K) + slot = tl.load( + kv_indices_ptr + kv_start + kv_lo + k_offsets, + mask=kv_lo + k_offsets < kv_hi, + other=-1, + ) + for k_start in tl.range(kv_lo, kv_hi, BLOCK_K): + k_pos = k_start + k_offsets + in_range = k_pos < kv_hi + valid = in_range & (slot >= 0) & (slot < num_kv) + safe_slot = tl.where(valid, slot, 0) + + kv = tl.load( + kv_ptr + + _sparse_kv_row_offset(safe_slot[:, None], kv_stride_n) + + dim_offsets[None, :] * kv_stride_d, + mask=valid[:, None] & dim_mask[None, :], + other=0.0, + ) + + next_k_pos = k_start + BLOCK_K + k_offsets + slot = tl.load( + kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_hi, other=-1 + ) + + scores = tl.dot(q, tl.trans(kv)) * scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv) + m_i = m_new + l_i = l_new + + # Every launched slot is written; empty splits store the (neg_large, 0.0) + # sentinel the reduce combines to nothing. + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets + tl.store(part_m_ptr + pm_base, m_i, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + tl.store( + part_acc_ptr + + query_idx * pa_stride0 + + split_id * pa_stride_s + + head_offsets[:, None] * pa_stride_h + + dim_offsets[None, :], + acc, + mask=head_mask[:, None] & dim_mask[None, :], + ) + + @triton.jit def _decode_e8m0_scales_triton(encoded_scales): scale_bits = encoded_scales.to(tl.int32) << 23 @@ -3664,6 +3777,182 @@ def _decode_gfx950_num_splits( return num_splits +# Empirically single-pass bf16 decode is faster than split-K below this number of +# selected slots per query, so we do not split. +_SPARSE_DECODE_BF16_MIN_SPLIT_LEN = 512 +_SPARSE_DECODE_BF16_BLOCK_H = 16 +_SPARSE_DECODE_BF16_BLOCK_K = 32 + + +def rocm_sparse_decode_bf16_num_splits( + num_queries: int, num_heads: int, sparse_len: int +) -> int: + """Number or kv splits in splitK for the sparse bf16 decode, or 1 for single-pass. + + Args: + num_queries: Decode rows in the batch. + num_heads: Query heads per row. + sparse_len: Longest selected KV run any decode row can walk. + + Returns: + The split count, or 1 when the caller should use the single-pass kernel. + + """ + if sparse_len < _SPARSE_DECODE_BF16_MIN_SPLIT_LEN: + return 1 + block_k = _SPARSE_DECODE_BF16_BLOCK_K + heads_blocks = triton.cdiv(num_heads, _SPARSE_DECODE_BF16_BLOCK_H) + select = _decode_gfx950_num_splits if _ON_GFX950 else _decode_num_splits + num_splits = select(num_queries, heads_blocks, sparse_len, 0.0, block_k) + # Number of splits cannot exceed the available k tiles + return max(1, min(num_splits, math.ceil(sparse_len / block_k))) + + +def _rocm_sparse_attn_decode_ragged_bf16_triton( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + indptr: torch.Tensor, + scale: float, + attn_sink: torch.Tensor | None, + nope_head_dim: int, + rope_head_dim: int, + num_splits: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Split-K decode over an bf16 ragged KV cache. + + Partitions each query's selected tokens across different workgroups and combines + the partials through reduction. + + Args: + q: Queries laid out as ``[sq, h, d]``. + kv: Unquantized KV rows laid out as ``[skv, d]``. + indices: Flattened per-query KV slots. + indptr: Segment offsets into ``indices``, ``[sq + 1]``. + scale: Softmax scale. + attn_sink: Optional per-head sink logits. + nope_head_dim: NoPE width of ``d``. + rope_head_dim: RoPE width of ``d``. + num_splits: Number of KV splits per query. + out: Optional destination with ``d`` trailing elements. + + Returns: + The attention output, ``out`` when provided. + + """ + assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}" + assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}" + assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}" + assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}" + assert not q.is_cpu and not kv.is_cpu and not indices.is_cpu and not indptr.is_cpu + + indices = _as_int32_contiguous_1d(indices) + indptr = _as_int32_contiguous_1d(indptr) + has_attn_sink = attn_sink is not None + if attn_sink is None: + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) + else: + attn_sink = attn_sink.contiguous() + + num_queries, num_heads, head_dim = q.shape + assert indptr.numel() == num_queries + 1, ( + f"expected indptr shape [{num_queries + 1}], got {indptr.shape}" + ) + _validate_sparse_dims( + head_dim, + nope_head_dim, + rope_head_dim, + "_rocm_sparse_attn_decode_ragged_bf16_triton", + ) + if out is None: + out = torch.empty_like(q) + assert out.shape[-1] == head_dim, ( + f"expected out trailing dim {head_dim}, got {out.shape[-1]}" + ) + + block_h = _SPARSE_DECODE_BF16_BLOCK_H + block_d = triton.next_power_of_2(head_dim) + block_k = _SPARSE_DECODE_BF16_BLOCK_K + heads_blocks = triton.cdiv(num_heads, block_h) + comb_dim = nope_head_dim + rope_head_dim + + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, comb_dim), + dtype=torch.float32, + device=q.device, + ) + + _sparse_attn_decode_ragged_bf16_partial_kernel[ + (num_queries, num_splits, heads_blocks) + ]( + q, + kv, + indices, + indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + q.stride(2), + kv.stride(0), + kv.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + head_dim, + kv.shape[0], + float(scale), + BLOCK_H=block_h, + BLOCK_D=block_d, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + num_warps=4, + ) + + _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + None, + None, + None, + out.stride(0), + out.stride(1), + 0, + head_dim // 32, + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + 0, + num_heads, + HAS_ATTN_SINK=has_attn_sink, + ADAPTIVE_SPLITS=False, + COMB_DIM=comb_dim, + BLOCK_H=1, + NUM_SPLITS=num_splits, + SPLITS_PAD=triton.next_power_of_2(num_splits), + FUSE_INV_ROPE=False, + NOPE=nope_head_dim, + HALF=rope_head_dim // 2, + QUANT_OUT=False, + num_warps=4, + ) + return out + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -4143,6 +4432,72 @@ def rocm_sparse_attn_prefill( ) +def rocm_sparse_attn_decode_bf16( + q: torch.Tensor, + kv: torch.Tensor, + scale: float, + head_dim: int, + nope_head_dim: int, + rope_head_dim: int, + attn_sink: torch.Tensor | None, + output: torch.Tensor, + ragged_indices: torch.Tensor, + ragged_indptr: torch.Tensor, + num_splits: int, +) -> None: + """Run split-K sparse attention over decode rows using an unquantized KV cache. + + Args: + q: Decode queries laid out as ``[sq, h, d]``. + kv: KV cache laid out as ``[skv, 1, d]``. + scale: Softmax scale. + head_dim: Post-absorption head width. + nope_head_dim: NoPE width of ``head_dim``. + rope_head_dim: RoPE width of ``head_dim``. + attn_sink: Optional per-head sink logits. + output: Destination, written in place. + ragged_indices: Flattened per-query KV slots. + ragged_indptr: Segment offsets into ``ragged_indices``, ``[sq + 1]``. + num_splits: KV splits per query, from + :func:`rocm_sparse_decode_bf16_num_splits`. + + """ + assert kv.ndim == 3 and kv.shape[1] == 1, ( + f"ROCm Triton sparse decode expects kv=[skv,1,d], got {kv.shape}" + ) + _validate_sparse_dims( + head_dim, + nope_head_dim, + rope_head_dim, + "rocm_sparse_attn_decode_bf16", + ) + num_queries, num_heads = q.shape[0], q.shape[1] + direct = output.shape[-1] == head_dim + out = ( + output + if direct + else torch.empty( + (num_queries, num_heads, head_dim), + dtype=output.dtype, + device=output.device, + ) + ) + _rocm_sparse_attn_decode_ragged_bf16_triton( + q=q, + kv=kv.squeeze(1), + indices=ragged_indices, + indptr=ragged_indptr, + scale=scale, + attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]], + nope_head_dim=nope_head_dim, + rope_head_dim=rope_head_dim, + num_splits=num_splits, + out=out, + ) + if not direct: + output.copy_(out[..., : output.shape[-1]]) + + def rocm_sparse_attn_decode( q: torch.Tensor, kv_cache: torch.Tensor | None, From 8352b2427f704652da1910f0d53f3fdcc2fa2b0d Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 8 Oct 2026 11:38:09 +0800 Subject: [PATCH 170/952] [Multimodal] Optimize encoder cudagraph padding logic (#58925) Signed-off-by: Isotr0py Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/v1/cudagraph/test_encoder_cudagraph.py | 132 ++++++++++++++++++ vllm/model_executor/models/kimi_k25.py | 34 ++--- .../deepseek_v41/common/vl_cudagraph.py | 29 +--- vllm/utils/math_utils.py | 8 ++ vllm/v1/worker/encoder_cudagraph.py | 51 ++++--- 5 files changed, 188 insertions(+), 66 deletions(-) diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 749f9aa20dee..348944387886 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -245,6 +245,138 @@ def test_num_graphs_to_capture_tracks_budgets(self): assert mgr.get_num_graphs_to_capture() == 3 +# --------------------------------------------------------------------------- +# _execute_local greedy packing with budget-cliff deferral (no GPU required) +# --------------------------------------------------------------------------- + + +def _run_packing( + mgr: EncoderCudaGraphManager, specs: list[EncoderItemSpec] +) -> list[tuple[list[int], str, int]]: + """Run _execute_local with graph replay and postprocess mocked out. + + Returns: + (batch_indices, path, token_budget) for every graph replay. + + """ + mgr.model = MagicMock() + + runs: list[tuple[list[int], str, int]] = [] + + def fake_run(mm_kwargs, token_budget, path="default", axis_keys=()): + runs.append((list(mm_kwargs["indices"]), path, token_budget)) + n_tokens = sum( + specs[i].get_path_output_tokens(path) for i in mm_kwargs["indices"] + ) + return torch.zeros(n_tokens, 32) + + def fake_postprocess( + graph_outputs, + batch_indices, + per_item_out_tokens, + outputs_by_orig_idx, + clone, + batch_mm_kwargs, + ): + for i in batch_indices: + outputs_by_orig_idx[i] = torch.zeros(per_item_out_tokens[i], 32) + + mgr.model.postprocess_encoder_output = fake_postprocess + with ( + patch.object(mgr, "_get_item_specs", lambda mm_kwargs: specs), + patch.object( + mgr, + "_select_items", + lambda mm_kwargs, indices: ({"indices": indices}, ()), + ), + patch.object(mgr, "_run_budget_graph", fake_run), + ): + result = mgr._execute_local({}) + assert len(result) == len(specs) + return runs + + +class TestExecuteLocalPacking: + """Greedy packing defers items that cross a budget cliff.""" + + @staticmethod + def _spec(tokens: int) -> EncoderItemSpec: + return EncoderItemSpec(input_size=tokens, output_tokens=tokens) + + def test_cliff_deferral_splits_batch(self): + # Merging would force 4200 tokens into the 8192 budget; deferring + # the 4000-token item costs 256 + 4096 instead. + mgr = _make_manager_with_budgets([256, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(4000), self._spec(200)]) + assert runs == [([1], "default", 256), ([0], "default", 4096)] + + def test_tie_does_not_defer(self): + # 4096 + 4096 ties the merged 8192 budget; merging wins because it + # saves one replay. + mgr = _make_manager_with_budgets([256, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(4000), self._spec(4000)]) + assert runs == [([0, 1], "default", 8192)] + + def test_deferral_requires_item_outweighing_batch(self): + # The third 1366-token item would save budget if deferred (4096+2048 + # < 8192) but is smaller than the current batch (4096), so it merges: + # deferring mid-size items fragments flood packing into extra replays. + mgr = _make_manager_with_budgets([256, 2048, 4096, 8192]) + runs = _run_packing(mgr, [self._spec(1366)] * 3) + assert runs == [([0, 1, 2], "default", 8192)] + + def test_packing_without_cliff_unchanged(self): + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100), self._spec(100)]) + assert runs == [([0, 1], "default", 256)] + + def test_small_items_split_to_avoid_cliff(self): + # Merging 3+ items of 100 tokens crosses 256 -> 4096; splitting + # into pairs costs 256 + 256 < 4096. + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100)] * 4) + assert runs == [([0, 1], "default", 256), ([2, 3], "default", 256)] + + def test_max_batch_size_split_unchanged(self): + mgr = _make_manager_with_budgets([1024, 4096]) + mgr.max_batch_size = 3 + runs = _run_packing(mgr, [self._spec(100)] * 4) + assert runs == [([0, 1, 2], "default", 1024), ([3], "default", 1024)] + + def test_oversized_item_still_falls_back_to_eager(self): + mgr = _make_manager_with_budgets([256, 4096]) + runs = _run_packing(mgr, [self._spec(100), self._spec(9000)]) + assert runs == [([0], "default", 256)] + assert mgr.graph_misses == 1 + + def test_multi_path_deferral_skips_zero_token_paths(self): + mgr = _make_manager_with_budgets([256, 4096, 8192]) + mgr.path_token_budgets = { + "a": [0, 256, 4096, 8192], + "b": [256, 4096, 8192], + } + specs = [ + EncoderItemSpec( + input_size=300, + output_tokens=300, + path_output_tokens={"a": 200, "b": 100}, + ), + EncoderItemSpec( + input_size=4000, + output_tokens=4000, + path_output_tokens={"a": 0, "b": 4000}, + ), + ] + runs = _run_packing(mgr, specs) + # Path b alone triggers deferral (256 + 4096 < 8192); the deferred + # item has zero tokens on path a, so that path is skipped. + assert runs == [ + ([0], "a", 256), + ([0], "b", 256), + ([1], "b", 4096), + ] + + # --------------------------------------------------------------------------- # get_cumulative_stats # --------------------------------------------------------------------------- diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index d66bed3642bc..8cd354b903af 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -479,30 +479,16 @@ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: # Image-only (t == 1). Video chunks (t > 1) fall back to eager. # # ------------------------------------------------------------------ # - @property - def _encoder_cudagraph_pad_totals(self) -> dict[int, int]: - """Row count of each captured buffer set, keyed by cu_seqlens ptr.""" - totals = self.__dict__.get("_encoder_cg_pad_totals") - if totals is None: - totals = {} - self.__dict__["_encoder_cg_pad_totals"] = totals - return totals - def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig - pad_totals = self._encoder_cudagraph_pad_totals - def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: - # Varlen attention requires cu_seqlens[-1] to equal the number of - # rows actually passed in. The captured buffers are sized for the - # full token budget, so a smaller real batch has to be completed - # with one trailing padding sequence; declaring fewer rows than the - # buffer holds is undefined behaviour and returns NaN on FlashAttn. - total = pad_totals.get(dst.data_ptr()) + # Zero-length padding (same as Qwen): tail slots repeat the last + # real offset, so FlashAttn schedules no work for padding rows. n = min(src.shape[0], dst.shape[0]) dst[:n].copy_(src[:n]) - dst[n:] = total if total is not None else src[-1] + if n < dst.shape[0]: + dst[n:] = src[-1] if n else 0 return EncoderCudaGraphConfig( modalities=["vision_chunk"], @@ -633,8 +619,8 @@ def prepare_encoder_cudagraph_capture_inputs( # max_seqlen must cover the worst case: one item consuming the full # budget, i.e. token_budget * kh * kw patches. - # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can append - # a padding sequence covering rows the real batch does not fill. + # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can + # pad the tail with zero-length sequences. metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( grid_thw_list, max_batch_size=max_batch_size + 1, @@ -645,10 +631,6 @@ def prepare_encoder_cudagraph_capture_inputs( values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values} values.update({k: v for k, v in metadata.items() if v is not None}) - cu_seqlens = values.get("cu_seqlens") - if cu_seqlens is not None: - self._encoder_cudagraph_pad_totals[cu_seqlens.data_ptr()] = total_patches - return EncoderCudaGraphCaptureInputs(values=values) def prepare_encoder_cudagraph_replay_buffers( @@ -663,8 +645,8 @@ def prepare_encoder_cudagraph_replay_buffers( grid_thw_list = self._get_grid_thw_list(mm_kwargs) pixel_values = mm_kwargs["pixel_values"] - # Unpadded: pad_cu_seqlens completes the tail with the padding sequence - # so cu_seqlens[-1] matches the captured buffer's row count. + # Unpadded: pad_cu_seqlens pads the tail with zero-length sequences + # whose entries repeat the real total. metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( grid_thw_list, max_batch_size=None, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py index bd8ed5a3fd9a..19baf339eb17 100644 --- a/vllm/models/deepseek_v41/common/vl_cudagraph.py +++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py @@ -94,8 +94,6 @@ def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: grid = grid.tolist() return [[int(v) for v in row] for row in grid] - _encoder_cg_pad_totals: dict[int, int] = {} - # -- SupportsEncoderCudaGraph protocol -- def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: @@ -129,24 +127,13 @@ def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: "cudagraph_mm_encoder." ) - pad_totals = self._encoder_cg_pad_totals - def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: - # Varlen attention requires cu_seqlens[-1] to equal the number of - # rows actually passed in. The captured buffers are sized for the - # full token budget, so a smaller real batch is completed with one - # trailing padding sequence; declaring fewer rows than the buffer - # holds is undefined behaviour and returns NaN on FlashAttn. - total = pad_totals.get(dst.data_ptr()) - if total is None: - raise RuntimeError( - "cu_seqlens replay buffer was not registered at capture " - "time; the manager must replay into the capture-time " - "buffers." - ) + # Zero-length padding (same as Qwen): tail slots repeat the last + # real offset, so FlashAttn schedules no work for padding rows. n = min(src.shape[0], dst.shape[0]) dst[:n].copy_(src[:n]) - dst[n:] = total + if n < dst.shape[0]: + dst[n:] = src[-1] if n else 0 return EncoderCudaGraphConfig( modalities=["image"], @@ -268,8 +255,8 @@ def prepare_encoder_cudagraph_capture_inputs( max_seqlen_override=total_patches, cached=False, ) - # Spare cu_seqlens slots let replay append a padding sequence that - # covers the rows a smaller real batch does not fill. + # Spare cu_seqlens slots: replay pads smaller real batches with + # zero-length sequences (tail slots repeat the last real offset). real_cu = metadata.pop("cu_seqlens") cu_seqlens = torch.full( (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device @@ -278,8 +265,6 @@ def prepare_encoder_cudagraph_capture_inputs( merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) - self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches - return EncoderCudaGraphCaptureInputs( values={ "patches": patches, @@ -303,7 +288,7 @@ def prepare_encoder_cudagraph_replay_buffers( patches = patches.to(dtype) # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and - # pad_cu_seqlens appends the padding sequence covering the tail rows. + # pad_cu_seqlens pads the tail with zero-length sequences. metadata = build_packed_vit_metadata( vit_grid, rope_dim=self.vision.rope_dim, diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py index 8fb0009d60bf..9854249af28b 100644 --- a/vllm/utils/math_utils.py +++ b/vllm/utils/math_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Math utility functions for vLLM.""" +from bisect import bisect_left + # Approximate value of 1/ln(2), used for log/exp base conversion # Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B) RCP_LN2 = 1.4426950216 @@ -30,3 +32,9 @@ def round_down(x: int, y: int) -> int: 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) + + +def round_up_to(values: list[int], n: int) -> int | None: + """Round `n` up to the smallest member of sorted `values` >= n, or None.""" + i = bisect_left(values, n) + return values[i] if i < len(values) else None diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 3b587ad71f80..51c832001ca8 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -23,6 +23,7 @@ from vllm.model_executor.models.utils import scatter_output_slices from vllm.model_executor.models.vision import get_load_balance_assignment from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import round_up_to from vllm.utils.torch_utils import current_stream from vllm.v1.worker.encoder_cudagraph_defs import ( ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, @@ -347,17 +348,17 @@ def _capture_budget_graph( def _find_smallest_fitting_budget_given_tokens( self, total_tokens: int, budgets: list[int] | None = None ) -> int | None: - """Find smallest budget >= total_tokens. - - Returns: - Token budget if found, None if no fitting budget. - - """ + """Find smallest budget >= total_tokens, or None if none fits.""" budgets = budgets if budgets is not None else self.token_budgets - for budget in budgets: - if budget >= total_tokens: - return budget - return None + return round_up_to(budgets, total_tokens) + + def _path_fitting_budget(self, path: str, tokens: int) -> int | None: + """Smallest captured budget on a path; 0 when idle, None if oversized.""" + if tokens == 0: + return 0 + return self._find_smallest_fitting_budget_given_tokens( + tokens, self.path_token_budgets[path] + ) def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: """Get item specs from the model.""" @@ -479,24 +480,38 @@ def append_current_batch() -> None: if not current_batch: return path_budgets = { - path: ( - 0 - if current_tokens[path] == 0 - else self._find_smallest_fitting_budget_given_tokens( - current_tokens[path], self.path_token_budgets[path] - ) - ) + path: self._path_fitting_budget(path, current_tokens[path]) for path in paths } batches.append((list(current_batch), path_budgets)) + def should_defer_item(item_tokens: dict[str, int]) -> bool: + """Whether closing the batch strictly reduces summed budgets.""" + for path in paths: + item = item_tokens[path] + if item == 0: + continue + cur = current_tokens[path] + b_cur = self._path_fitting_budget(path, cur) + b_alone = self._path_fitting_budget(path, item) + b_merged = self._path_fitting_budget(path, cur + item) + # Only called when the item fits; lookups land on a budget. + assert ( + b_cur is not None and b_alone is not None and b_merged is not None + ) + # Only defer items that alone outweigh the whole current + # batch; splitting smaller ones fragments flood packing. + if b_cur + b_alone < b_merged and b_alone >= b_cur: + return True + return False + for orig_idx in sorted_indices: item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths} fits = len(current_batch) < self.max_batch_size and all( current_tokens[path] + item_tokens[path] <= max_path_budgets[path] for path in paths ) - if current_batch and not fits: + if current_batch and (not fits or should_defer_item(item_tokens)): append_current_batch() current_batch = [] current_tokens = dict.fromkeys(paths, 0) From 5dd4a628dde34d5b5737067f131a2499f672f684 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:26:12 -0700 Subject: [PATCH 171/952] [ROCm][Bugfix] Map sleep-mode cuMem memory for the host (#60367) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) --- csrc/cumem_allocator.cpp | 23 +++++++++++++++---- .../memory/sleep_mode/test_sleep_mode.py | 1 - 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 720c990450b6..6965d2124a65 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -203,12 +203,25 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, } #endif - CUmemAccessDesc accessDesc = {}; - accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - accessDesc.location.id = device; - accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + // Also map for the host, as PyTorch does for its expandable segments: on + // ROCm >= 7.2, Tensor.item() on a large-BAR device dereferences VRAM + // directly from the host, which segfaults on a device-only mapping. + constexpr int num_desc = 2; +#else + constexpr int num_desc = 1; +#endif + CUmemAccessDesc accessDesc[num_desc] = {}; + accessDesc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + accessDesc[0].location.id = device; + accessDesc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + accessDesc[1].location.type = hipMemLocationTypeHost; + accessDesc[1].location.id = 0; // ignored + accessDesc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#endif - CUDA_CHECK(cuMemSetAccess(d_mem, size, &accessDesc, 1)); + CUDA_CHECK(cuMemSetAccess(d_mem, size, accessDesc, num_desc)); if (error_code != 0) { return; } diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 73af6e39103d..9f6180335ab1 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -323,7 +323,6 @@ def _attn_kv_scales(model) -> dict[str, list[float]]: } -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Segfaults on ROCm.") @requires_fp8 @create_new_process_for_each_test() def test_deep_sleep_compressed_tensors_kv_scales( From 1fbbfa165709c490dae48884eebdfbf6a1aef330 Mon Sep 17 00:00:00 2001 From: Hyukjoon Lee Date: Thu, 8 Oct 2026 13:33:32 +0900 Subject: [PATCH 172/952] [Bugfix][Kimi-K3] Declare max_tp_shards on the DSpark MLA KV cache spec (#59733) Signed-off-by: hyukjlee Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Summer Yang --- tests/models/test_dspark_mla.py | 33 +++++++++++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/mla.py | 1 + 2 files changed, 34 insertions(+) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 76af0eb60c10..53a93b940cd5 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -336,3 +336,36 @@ def __call__(self, main_x): assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2) assert calls[0][3] is slot_mappings[0] assert calls[1][3] is slot_mappings[2] + + +@pytest.mark.cpu_test +def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): + """The draft's MLA layers must share a KV cache group with the target's.""" + from vllm.model_executor.layers.attention.mla_attention import MLAAttention + from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention + from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + + vllm_config = SimpleNamespace( + model_config=None, cache_config=SimpleNamespace(block_size=64) + ) + target_attn = SimpleNamespace( + kv_cache_dtype="fp8", + head_size=576, + sliding_window=None, + indexer=None, + non_causal_multi_token_decode=False, + attn_backend=SimpleNamespace(get_name=lambda: "ROCM_AITER_MLA"), + _uses_flat_kv_cache=lambda: False, + ) + draft_attn = SimpleNamespace( + kv_cache_dtype="fp8", head_size=576, non_causal_multi_token_decode=True + ) + target_spec = MLAAttention.get_kv_cache_spec(target_attn, vllm_config) + draft_spec = MultiHeadLatentAttention.get_kv_cache_spec(draft_attn, vllm_config) + + kv_cache_spec = {f"model.layers.{i}.attn": target_spec for i in range(24)} + kv_cache_spec |= {f"draft.layers.{i}.attn": draft_spec for i in range(5)} + + groups = _get_kv_cache_groups_uniform_page_size(kv_cache_spec) + assert len(groups) == 1 + assert len(groups[0].layer_names) == 29 diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index bf7c53e9e515..18fd93579be7 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -422,6 +422,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: return MLAAttentionSpec( # type: ignore[call-arg] block_size=vllm_config.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_size, dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, From 4187ade7b87e4c683ff60e41ad27cc6659f8739d Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 8 Oct 2026 12:37:11 +0800 Subject: [PATCH 173/952] [Frontend] Consolidate rl entrypoints. (#57849) Signed-off-by: wang.yuqi Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .github/CODEOWNERS | 2 +- .../test_route_request_bodies.py} | 4 +- .../rl/weight_transfer_metrics}/__init__.py | 0 .../test_across_api_servers_metrics.py} | 0 .../test_weight_operation_metrics.py} | 20 +++--- vllm/distributed/weight_transfer/clients.py | 2 +- .../launchers/api_server/routers.py | 4 ++ vllm/entrypoints/llm.py | 63 ++---------------- vllm/entrypoints/rl/__init__.py | 0 vllm/entrypoints/rl/factories.py | 11 ++++ vllm/entrypoints/rl/offline.py | 64 +++++++++++++++++++ vllm/entrypoints/rl/online/__init__.py | 0 .../dev/rlhf => rl/online}/api_router.py | 8 +-- .../{serve/dev/rlhf => rl/online}/metrics.py | 0 vllm/entrypoints/serve/__init__.py | 4 -- 15 files changed, 99 insertions(+), 83 deletions(-) rename tests/entrypoints/{unit_tests/test_dev_route_request_bodies.py => rl/return_routed_experts/test_route_request_bodies.py} (96%) rename {vllm/entrypoints/serve/dev/rlhf => tests/entrypoints/rl/weight_transfer_metrics}/__init__.py (100%) rename tests/entrypoints/rl/{test_weight_operation_metrics.py => weight_transfer_metrics/test_across_api_servers_metrics.py} (100%) rename tests/entrypoints/{unit_tests/test_rl_weight_operation_metrics.py => rl/weight_transfer_metrics/test_weight_operation_metrics.py} (90%) create mode 100644 vllm/entrypoints/rl/__init__.py create mode 100644 vllm/entrypoints/rl/factories.py create mode 100644 vllm/entrypoints/rl/offline.py create mode 100644 vllm/entrypoints/rl/online/__init__.py rename vllm/entrypoints/{serve/dev/rlhf => rl/online}/api_router.py (97%) rename vllm/entrypoints/{serve/dev/rlhf => rl/online}/metrics.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 083842c932ff..677e37ff8b60 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -229,7 +229,7 @@ mkdocs.yaml @hmellor # Reinforcement learning /vllm/distributed/weight_transfer @aoshen02 -/vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02 +/vllm/entrypoints/rl @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 /tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 diff --git a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py similarity index 96% rename from tests/entrypoints/unit_tests/test_dev_route_request_bodies.py rename to tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index 78fe5eef5dde..b51d08db0758 100644 --- a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py +++ b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py @@ -9,7 +9,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient -from vllm.entrypoints.serve.dev.rlhf.api_router import router as rlhf_router +from vllm.entrypoints.rl.online.api_router import router as rl_router from vllm.entrypoints.serve.dev.rpc.api_router import router as rpc_router from vllm.entrypoints.serve.exception_handling.register import init_exception_handler @@ -30,7 +30,7 @@ async def call(*args, **kwargs): @pytest.fixture def client(): app = FastAPI() - app.include_router(rlhf_router) + app.include_router(rl_router) app.include_router(rpc_router) init_exception_handler(app) app.state.engine_client = Engine() diff --git a/vllm/entrypoints/serve/dev/rlhf/__init__.py b/tests/entrypoints/rl/weight_transfer_metrics/__init__.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/__init__.py rename to tests/entrypoints/rl/weight_transfer_metrics/__init__.py diff --git a/tests/entrypoints/rl/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py similarity index 100% rename from tests/entrypoints/rl/test_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py diff --git a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py similarity index 90% rename from tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index bdddf3d247ba..22ec5e611252 100644 --- a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py +++ b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py @@ -13,9 +13,9 @@ from fastapi.testclient import TestClient from prometheus_client import CollectorRegistry -from vllm.entrypoints.serve.dev.rlhf import api_router -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics -from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics +from vllm.entrypoints.rl.online import api_router +from vllm.entrypoints.rl.online import metrics as rl_metrics +from vllm.entrypoints.rl.online.metrics import WeightOperationMetrics from vllm.entrypoints.serve.exception_handling.register import init_exception_handler from vllm.v1.metrics.prometheus import get_prometheus_registry @@ -52,13 +52,13 @@ async def call(*args, **kwargs): @pytest.fixture def registry(monkeypatch): registry = CollectorRegistry() - monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry)) + monkeypatch.setattr(rl_metrics, "_metrics", WeightOperationMetrics(registry)) return registry @pytest.mark.parametrize("fail", [False, True]) def test_record_releases_in_flight_and_observes_duration(registry, fail): - metrics = rlhf_metrics.weight_operation_metrics() + metrics = rl_metrics.weight_operation_metrics() with ( pytest.raises(RuntimeError) if fail else nullcontext(), metrics.record("update"), @@ -108,7 +108,7 @@ async def test_cancelled_route_releases_in_flight(registry): HOLD_OPERATION = """ import sys -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics with weight_operation_metrics().record("update"): print("entered", flush=True) sys.stdin.readline() @@ -163,19 +163,19 @@ def test_default_recorder_is_created_lazily_on_the_default_registry(): probe = """ from prometheus_client import REGISTRY -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics +from vllm.entrypoints.rl.online import metrics as rl_metrics NAMES = [ "vllm:rl_weight_update_operation_duration_seconds", "vllm:rl_weight_update_operations_in_flight", ] -assert rlhf_metrics._metrics is None +assert rl_metrics._metrics is None assert not set(NAMES) & set(REGISTRY._names_to_collectors) -metrics = rlhf_metrics.weight_operation_metrics() +metrics = rl_metrics.weight_operation_metrics() assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight -assert rlhf_metrics.weight_operation_metrics() is metrics +assert rl_metrics.weight_operation_metrics() is metrics print("lazy registration ok") """ result = subprocess.run( diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 1bd2f98937a9..745e67779e0a 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -58,7 +58,7 @@ def _json_safe_update_payload( class HTTPVLLMWeightSyncClient: """Talks to a vLLM server over the RLHF HTTP routes. - Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`: + Mirrors `vllm/entrypoints/rl/online/api_router.py`: `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, `/finish_weight_update`. """ diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py index 39d670c019b7..40d04862afb2 100644 --- a/vllm/entrypoints/launchers/api_server/routers.py +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -36,6 +36,10 @@ def register_api_routers( register_vllm_dev_api_routers(app) + from vllm.entrypoints.rl.factories import register_rl_api_routers + + register_rl_api_routers(app) + if "generate" in supported_tasks: from vllm.entrypoints.generate.api_router import ( register_generate_api_routers, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 45b6aede583c..c38789a2e29c 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -28,10 +28,6 @@ TokenizerMode, ) from vllm.config.quantization import QuantizationConfigArgs -from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - WeightTransferUpdateRequest, -) from vllm.engine.arg_utils import EngineArgs from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -40,6 +36,7 @@ ) from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.rl.offline import RLOfflineMixin from vllm.entrypoints.serve.utils.api_utils import log_non_default_args from vllm.inputs import PromptType from vllm.logger import configure_logging_if_needed, init_logger @@ -64,7 +61,9 @@ logger = init_logger(__name__) -class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): +class LLM( + BeamSearchOfflineMixin, PoolingOfflineMixin, RLOfflineMixin, OfflineInferenceMixin +): """An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed @@ -884,60 +883,6 @@ def get_metrics(self) -> list["Metric"]: """ return self.llm_engine.get_metrics() - def init_weight_transfer_engine( - self, request: WeightTransferInitRequest | dict - ) -> None: - """Initialize weight transfer for RL training. - - Args: - request: Weight transfer initialization request with backend-specific info - - """ - init_info_dict = ( - request["init_info"] if isinstance(request, dict) else request.init_info - ) - - self.llm_engine.collective_rpc( - "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} - ) - - def start_weight_update(self) -> None: - """Start a new weight update.""" - self.llm_engine.collective_rpc("start_weight_update") - - def start_draft_weight_update(self) -> None: - """Start a new weight update targeting the speculative draft model.""" - self.llm_engine.collective_rpc("start_draft_weight_update") - - def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: - """Update the weights of the model. - - Args: - request: Weight update request with backend-specific update info - - """ - update_info_dict = ( - request["update_info"] if isinstance(request, dict) else request.update_info - ) - - self.llm_engine.collective_rpc( - "update_weights", kwargs={"update_info": update_info_dict} - ) - - def finish_weight_update(self, weight_version: str | None = None) -> None: - """Finish the weight update and set its version if provided.""" - self.llm_engine.collective_rpc("finish_weight_update") - if weight_version is not None: - self.llm_engine.set_weight_version(weight_version) - - def update_weight_version(self, new_version: str) -> None: - """Set the weight version without updating weights.""" - self.llm_engine.set_weight_version(new_version) - - def get_weight_version(self) -> str: - """Return the latest committed weight version.""" - return self.llm_engine.get_weight_version() - def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" # Cache the result to avoid repeated collective_rpc calls diff --git a/vllm/entrypoints/rl/__init__.py b/vllm/entrypoints/rl/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/rl/factories.py b/vllm/entrypoints/rl/factories.py new file mode 100644 index 000000000000..a5d237ec0bf6 --- /dev/null +++ b/vllm/entrypoints/rl/factories.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from fastapi import FastAPI + + +def register_rl_api_routers(app: FastAPI): + from .online.api_router import router as rl_router + + app.include_router(rl_router) diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py new file mode 100644 index 000000000000..02fc63104d8f --- /dev/null +++ b/vllm/entrypoints/rl/offline.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) +from vllm.entrypoints.offline_utils import OfflineInferenceMixin + + +class RLOfflineMixin(OfflineInferenceMixin): + def init_weight_transfer_engine( + self, request: WeightTransferInitRequest | dict + ) -> None: + """Initialize weight transfer for RL training. + + Args: + request: Weight transfer initialization request with backend-specific info + + """ + init_info_dict = ( + request["init_info"] if isinstance(request, dict) else request.init_info + ) + + self.llm_engine.collective_rpc( + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + ) + + def start_weight_update(self) -> None: + """Start a new weight update.""" + self.llm_engine.collective_rpc("start_weight_update") + + def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + self.llm_engine.collective_rpc("start_draft_weight_update") + + def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: + """Update the weights of the model. + + Args: + request: Weight update request with backend-specific update info + + """ + update_info_dict = ( + request["update_info"] if isinstance(request, dict) else request.update_info + ) + + self.llm_engine.collective_rpc( + "update_weights", kwargs={"update_info": update_info_dict} + ) + + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() diff --git a/vllm/entrypoints/rl/online/__init__.py b/vllm/entrypoints/rl/online/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/rl/online/api_router.py similarity index 97% rename from vllm/entrypoints/serve/dev/rlhf/api_router.py rename to vllm/entrypoints/rl/online/api_router.py index 7e52e512170c..ef9269749f3b 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -5,7 +5,7 @@ from http import HTTPStatus from typing import Annotated, Any, Literal -from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -14,7 +14,7 @@ WeightTransferUpdateRequest, ) from vllm.engine.protocol import EngineClient -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics from vllm.logger import init_logger from vllm.v1.engine import PauseMode @@ -280,7 +280,3 @@ async def get_world_size( else: world_size = parallel_config.world_size return JSONResponse(content={"world_size": world_size}) - - -def attach_router(app: FastAPI): - app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/rlhf/metrics.py b/vllm/entrypoints/rl/online/metrics.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/metrics.py rename to vllm/entrypoints/rl/online/metrics.py diff --git a/vllm/entrypoints/serve/__init__.py b/vllm/entrypoints/serve/__init__.py index c358b9a6db26..d8254cb33548 100644 --- a/vllm/entrypoints/serve/__init__.py +++ b/vllm/entrypoints/serve/__init__.py @@ -50,10 +50,6 @@ def register_vllm_dev_api_routers(app: FastAPI): attach_cache_router(app) - from .dev.rlhf.api_router import attach_router as attach_rlhf_router - - attach_rlhf_router(app) - from .dev.rpc.api_router import attach_router as attach_rpc_router attach_rpc_router(app) From 775792f2d5cdbdb0dcf433943259744b17452a7e Mon Sep 17 00:00:00 2001 From: aniskumar-nv Date: Wed, 7 Oct 2026 21:40:55 -0700 Subject: [PATCH 174/952] [Bugfix] Reland #53020: tie lm_head.weight for Nemotron Parse when checkpoint omits it (#60272) Signed-off-by: Anish Kumar Co-authored-by: Cyrus Leung Co-authored-by: Claude Fable 5.1 --- .../multimodal/generation/test_nemotron_parse.py | 11 ++++++++++- vllm/model_executor/models/nemotron_parse.py | 9 +++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py index 69e41b21e342..57e40430c63c 100644 --- a/tests/models/multimodal/generation/test_nemotron_parse.py +++ b/tests/models/multimodal/generation/test_nemotron_parse.py @@ -104,7 +104,16 @@ def run_test( ) -@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.2"]) +@pytest.mark.parametrize( + "model", + [ + "nvidia/NVIDIA-Nemotron-Parse-v1.2", + # v2.0's checkpoint ties lm_head to the decoder's input embeddings + # instead of shipping a separate lm_head.weight tensor; this + # exercises that path (see tie_weights() in nemotron_parse.py). + "nvidia/NVIDIA-Nemotron-Parse-2.0", + ], +) @pytest.mark.parametrize("dtype", ["bfloat16"]) @pytest.mark.parametrize("num_logprobs", [5]) def test_models( diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 9ac38c0021db..3ba731d2e0f2 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -568,6 +568,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.lm_head = ParallelLMHead( config.decoder.vocab_size, config.decoder.d_model, quant_config=quant_config ) + # Some checkpoints (e.g. compact exports) tie the output head to the + # decoder's input embeddings instead of materializing a separate + # lm_head.weight tensor. + tie_word_embeddings = bool( + getattr(config, "tie_word_embeddings", False) + or getattr(config.decoder, "tie_word_embeddings", False) + ) + if tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.decoder.embed_tokens) self.logits_processor = LogitsProcessor( self.vocab_size, config.decoder.vocab_size ) From 4ac0179f8a4c6f459cd73dc6ac2e0bd22a5ad294 Mon Sep 17 00:00:00 2001 From: Mehdi Ghanimifard Date: Thu, 8 Oct 2026 06:49:49 +0200 Subject: [PATCH 175/952] [ROCm][Perf] Use native merge_attn_states kernel on gfx942/gfx950 (#60159) Signed-off-by: mghanimi Co-authored-by: Cursor Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- vllm/v1/attention/ops/merge_attn_states.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py index d28c78bdb9fa..b16786a13b47 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -139,8 +139,19 @@ def supported_headdim(prefix: torch.Tensor) -> bool: return headdim % 4 == 0 return headdim % 8 == 0 + # CDNA3 and newer (gfx942, gfx950, gfx1250). gfx90a is CDNA2 and + # stays on the Triton merge. + def supported_platform() -> bool: + if current_platform.is_cuda(): + return True + if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() >= 3 + return False + if ( - current_platform.is_cuda() + supported_platform() and supported_dtypes(prefix_output) and supported_headdim(prefix_output) ): From 81198e97ba7eee2a22540caaa756b7fdddcb4d93 Mon Sep 17 00:00:00 2001 From: Hui Shi <70356247+huishi-hs@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:05:39 +0800 Subject: [PATCH 176/952] [Refactor] Abstract DeepGemm MegaMOE backend (#58053) Signed-off-by: Hui Shi Co-authored-by: LuneExplorer Co-authored-by: ArmageddonKnight --- tests/models/test_deepseek_v4_mega_moe.py | 15 +- .../layers/fused_moe/deep_gemm_mega_moe.py | 466 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 306 ++++-------- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 107 +++- vllm/models/kimi_k3/nvidia/model.py | 53 +- 5 files changed, 681 insertions(+), 266 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 104ad8e89514..76915b94a32d 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -6,6 +6,9 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmSm100MegaMoEBackend, +) from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( bind_routed_experts_capturer, ) @@ -451,7 +454,9 @@ def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( ): param.data.random_(1, 128) originals.append(param.data.clone()) - monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) monkeypatch.setattr( "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace( @@ -547,7 +552,9 @@ def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): intermediate_size=intermediate_size, num_shared_experts=1, ) - experts._check_runtime_supported = lambda: None + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) def fp8_parameter(*shape): return torch.nn.Parameter( @@ -664,6 +671,10 @@ def forward(self, hidden_states): class FakeExperts(torch.nn.Module): has_fused_shared_experts = fused + _backend = DeepGemmSm100MegaMoEBackend() + + def _ensure_backend(self): + return self._backend def forward(self, hidden_states, *args, **kwargs): return torch.ones_like(hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py new file mode 100644 index 000000000000..633ced108208 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM MegaMoE backends: fused EP dispatch, expert GEMMs and combine.""" + +from inspect import signature +from typing import Any + +import torch +import torch.nn as nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp8Dynamic, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: + """Reinterpret a uint8 UE8M0 exponent tensor as float32 scale values.""" + return (sf.to(torch.int32) << 23).view(torch.float32) + + +class DeepGemmMegaMoEBackend: + """Abstract MegaMoE backend that hides deep_gemm-specific details. + + A backend advertises its expected hidden-state quantization layout via + ``hidden_quant`` and its expected MMA type via ``mma_type``. It is + responsible for the routed-expert weight transform run at load time, the + optional shared-expert weight transform, and the actual MegaMoE kernel + invocation at inference time. Runtime capability checks live in + :func:`get_deep_gemm_mega_moe_backend`: obtaining a backend already implies + that the current device and shapes are supported. + """ + + # Human-readable MMA type tag used by deep_gemm helpers. SM100 default. + mma_type: str = "fp8xfp4" + hidden_quant: QuantKey = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + raise NotImplementedError + + def supports_shared_experts(self, deep_gemm) -> bool: + return False + + def transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + """Return ``(l1, l2)`` transformed shared-expert weights, or ``None``. + + ``shared_experts`` must expose ``gate_up_proj`` and ``down_proj`` + linear layers. A ``None`` return means the backend cannot fuse shared + experts for this module; the caller should fall back to a separate + shared MLP. + """ + return None + + def run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + raise NotImplementedError + + def supports_bf16_mega_gate(self) -> bool: + """Whether this backend can run the DeepGEMM bf16 MegaMoE gate kernel. + + Defaults to ``False`` (the symbol is not guaranteed to exist in every + deep_gemm build). Backends that target a deep_gemm build exposing the + ``bf16_mega_gate`` symbol override this and :meth:`bf16_mega_gate`. + """ + return False + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + +class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): + """Default MegaMoE backend targeting SM100 via deep_gemm fp8_fp4_mega_moe.""" + + mma_type = "fp8xfp4" + hidden_quant = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + w13_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w13_weight_scale).contiguous(), + 2 * intermediate_size, + hidden_size, + (1, 32), + num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w2_weight_scale).contiguous(), + hidden_size, + intermediate_size, + (1, 32), + num_local_experts, + ) + kwargs = {} if activation is None else {"activation": activation} + return deep_gemm.transform_weights_for_mega_moe( + (w13_weight.view(torch.int8).contiguous(), w13_scale), + (w2_weight.view(torch.int8).contiguous(), w2_scale), + **kwargs, + ) + + def supports_shared_experts(self, deep_gemm) -> bool: + """Check the Python API before touching a symmetric-memory group. + + This also gives users of an older precompiled vLLM wheel a safe serial + fallback instead of failing halfway through multi-rank buffer setup. + """ + try: + buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters + kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters + except (TypeError, ValueError): + return False + return ( + hasattr(deep_gemm, "get_block_m_for_mega_moe") + and hasattr(deep_gemm, "transform_weights_for_mega_moe") + and "num_shared_experts" in buffer_params + and "shared_l1_weights" in kernel_params + and "shared_l2_weights" in kernel_params + ) + + def transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + gate_up = shared_experts.gate_up_proj + down = shared_experts.down_proj + gate_up_weight = gate_up.weight.data + gate_up_scale = ( + gate_up.weight_scale + if hasattr(gate_up, "weight_scale") + else gate_up.weight_scale_inv + ).data + down_weight = down.weight.data + down_scale = ( + down.weight_scale + if hasattr(down, "weight_scale") + else down.weight_scale_inv + ).data + + # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, + # while the checkpoint uses coarser block-FP8 scales (usually + # 128x128). Build a dedicated, numerically equivalent scale view. + ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) + if ( + gate_up_scale.dtype in ue8m0_scale_dtypes + and down_scale.dtype in ue8m0_scale_dtypes + ): + gate_up_scale = self._prepare_shared_expert_scale( + deep_gemm, + gate_up, + gate_up_scale, + gate_up_weight.shape[0], + gate_up_weight.shape[1], + prefix, + ) + down_scale = self._prepare_shared_expert_scale( + deep_gemm, + down, + down_scale, + down_weight.shape[0], + down_weight.shape[1], + prefix, + ) + + if gate_up_scale is None or down_scale is None: + return None + + shared_intermediate_size = intermediate_size * num_shared_experts + expected_gate_up_shape = ( + 2 * shared_intermediate_size, + hidden_size, + ) + expected_down_shape = (hidden_size, shared_intermediate_size) + if ( + gate_up_weight.dtype != torch.float8_e4m3fn + or down_weight.dtype != torch.float8_e4m3fn + or gate_up_scale.dtype != torch.int32 + or down_scale.dtype != torch.int32 + or tuple(gate_up_weight.shape) != expected_gate_up_shape + or tuple(down_weight.shape) != expected_down_shape + ): + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: expected " + "replicated block-FP8 weights with gate_up=%s, down=%s, and " + "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", + prefix, + expected_gate_up_shape, + expected_down_shape, + tuple(gate_up_weight.shape), + gate_up_weight.dtype, + gate_up_scale.dtype, + tuple(down_weight.shape), + down_weight.dtype, + down_scale.dtype, + ) + return None + + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (gate_up_weight, gate_up_scale), + (down_weight, down_scale), + ) + # L1 interleaving allocates a full copy. Re-home the loader Parameter on + # that storage so the original 2*intermediate*hidden FP8 tensor can be + # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. + # The generic linear post-load hook may still repack the serial scales, + # but this shared MLP is never called after native fusion is enabled. + gate_up.weight.data = transformed_l1[0] + transformed_l1 = (gate_up.weight.data, transformed_l1[1]) + return transformed_l1, transformed_l2 + + @staticmethod + def _prepare_shared_expert_scale( + deep_gemm, + linear: nn.Module, + scale: torch.Tensor, + mn: int, + k: int, + prefix: str, + ) -> torch.Tensor | None: + block_size = getattr(linear, "weight_block_size", None) + if block_size is None or len(block_size) != 2: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "shared FP8 weight block size is unavailable.", + prefix, + ) + return None + + block_m, block_k = block_size + num_k_blocks = (k + block_k - 1) // block_k + # The linear hook's DeepGEMM layout packs four k-block bytes per int32. + packed = scale.dtype == torch.int32 + expected_shape = ( + (mn, (num_k_blocks + 3) // 4) + if packed + else ((mn + block_m - 1) // block_m, num_k_blocks) + ) + if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: + logger.warning( + "Disabling native MegaMoE shared-expert fusion for %s: " + "cannot convert shared scale shape %s with block size %s " + "to MegaMoE's 1x32 layout for weight (%d, %d).", + prefix, + tuple(scale.shape), + tuple(block_size), + mn, + k, + ) + return None + + if packed: + row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] + else: + row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] + scale_1x32 = ( + ue8m0_uint8_to_float(row_scale) + .repeat_interleave(block_k // 32, dim=1)[:, : k // 32] + .contiguous() + ) + # The grouped API is used with a singleton dimension to request the + # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the + # 2D layout required for a shared expert. + return deep_gemm.transform_sf_into_required_layout( + scale_1x32.unsqueeze(0), + mn, + k, + (1, 32), + 1, + ).squeeze(0) + + def run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + kwargs: dict[str, Any] = { + name: value + for name, value in ( + ("activation", activation), + ("activation_alpha", activation_alpha), + ("activation_beta", activation_beta), + ) + if value is not None + } + if shared_l1_weights is not None and shared_l2_weights is not None: + kwargs["shared_l1_weights"] = shared_l1_weights + kwargs["shared_l2_weights"] = shared_l2_weights + deep_gemm.fp8_fp4_mega_moe( + y, + l1_weights, + l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + **kwargs, + ) + + def supports_bf16_mega_gate(self) -> bool: + """SM100 deep_gemm build is expected to provide the bf16 gate kernel.""" + return True + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.utils.deep_gemm import bf16_mega_gate + + return bf16_mega_gate( + x, + weight, + num_topk, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + ) + + +def get_deep_gemm_mega_moe_backend( + device: torch.device, + hidden_size: int, + intermediate_size: int, +) -> DeepGemmMegaMoEBackend: + """Return the DeepGEMM MegaMoE backend for ``device``. + + A successful return implies the backend can run on this device with the + given ``hidden_size`` / ``intermediate_size``. All runtime capability + checks (arch + shape divisibility) live here; there is no separate check + on the returned backend. + + Currently only the SM100 deep_gemm backend is registered here. SM90 and + other backends are expected to be provided by new backends or plugins + that override this factory. + + Raises: + NotImplementedError: If the device architecture is unsupported. + ValueError: If the shapes are unsupported. + + """ + if not current_platform.is_device_capability_family( + 100, device_id=device.index or 0 + ): + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + return DeepGemmSm100MegaMoEBackend() diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b9996d007c99..fc10cd07c1cf 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing from collections.abc import Callable, Iterable -from inspect import signature from itertools import islice import regex as re @@ -33,6 +32,11 @@ FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmMegaMoEBackend, + get_deep_gemm_mega_moe_backend, + ue8m0_uint8_to_float, +) from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) @@ -205,7 +209,9 @@ def make_deepseek_v4_expert_params_mapping( class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {} + _symm_buffer_cache: dict[ + tuple[int, int, int, int, int, int, int, int, str], object + ] = {} def __init__( self, @@ -296,6 +302,8 @@ def __init__( tuple[torch.Tensor, torch.Tensor] | None ) = None + self._backend: DeepGemmMegaMoEBackend | None = None + # Register in the static forward context so the custom-op wrapper # can look up this module by name from within a torch.compile graph. compilation_config = vllm_config.compilation_config @@ -357,214 +365,71 @@ def weight_loader( @staticmethod def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - return (sf.to(torch.int32) << 23).view(torch.float32) + return ue8m0_uint8_to_float(sf) - def _check_runtime_supported(self) -> None: - device = self.w13_weight.device - if torch.cuda.get_device_capability(device)[0] != 10: - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) - - @staticmethod - def _deep_gemm_supports_shared_experts(deep_gemm) -> bool: - """Check the Python API before touching a symmetric-memory group. + def _ensure_backend(self) -> DeepGemmMegaMoEBackend: + """Lazily resolve the MegaMoE backend. - This also gives users of an older precompiled vLLM wheel a safe serial - fallback instead of failing halfway through multi-rank buffer setup. + The backend is looked up via :func:`get_deep_gemm_mega_moe_backend`, which + performs all runtime capability checks. Failure propagates directly + to the caller. This must be called before :meth:`finalize_weights` + releases the loader-side parameters. """ - try: - buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters - kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters - except (TypeError, ValueError): - return False - return ( - hasattr(deep_gemm, "get_block_m_for_mega_moe") - and hasattr(deep_gemm, "transform_weights_for_mega_moe") - and "num_shared_experts" in buffer_params - and "shared_l1_weights" in kernel_params - and "shared_l2_weights" in kernel_params + if self._backend is not None: + return self._backend + assert self.w13_weight is not None, ( + "_ensure_backend() must be called before finalize_weights() " + "releases the loader-side parameters." ) - - def _finalize_shared_expert_weights( - self, deep_gemm, shared_experts: DeepseekV4MLP - ) -> None: - gate_up = shared_experts.gate_up_proj - down = shared_experts.down_proj - gate_up_weight = gate_up.weight.data - gate_up_scale = ( - gate_up.weight_scale - if hasattr(gate_up, "weight_scale") - else gate_up.weight_scale_inv - ).data - down_weight = down.weight.data - down_scale = ( - down.weight_scale - if hasattr(down, "weight_scale") - else down.weight_scale_inv - ).data - - # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row, - # while the checkpoint uses coarser block-FP8 scales (usually - # 128x128). Build a dedicated, numerically equivalent scale view. - ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) - if ( - gate_up_scale.dtype in ue8m0_scale_dtypes - and down_scale.dtype in ue8m0_scale_dtypes - ): - gate_up_scale = self._prepare_shared_expert_scale( - deep_gemm, - gate_up, - gate_up_scale, - gate_up_weight.shape[0], - gate_up_weight.shape[1], - ) - down_scale = self._prepare_shared_expert_scale( - deep_gemm, - down, - down_scale, - down_weight.shape[0], - down_weight.shape[1], - ) - - if gate_up_scale is None or down_scale is None: - self.num_shared_experts = 0 - return - - shared_intermediate_size = self.intermediate_size * self.num_shared_experts - expected_gate_up_shape = ( - 2 * shared_intermediate_size, + self._backend = get_deep_gemm_mega_moe_backend( + self.w13_weight.device, self.hidden_size, + self.intermediate_size, ) - expected_down_shape = (self.hidden_size, shared_intermediate_size) - if ( - gate_up_weight.dtype != torch.float8_e4m3fn - or down_weight.dtype != torch.float8_e4m3fn - or gate_up_scale.dtype != torch.int32 - or down_scale.dtype != torch.int32 - or tuple(gate_up_weight.shape) != expected_gate_up_shape - or tuple(down_weight.shape) != expected_down_shape - ): - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: expected " - "replicated block-FP8 weights with gate_up=%s, down=%s, and " - "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.", - self.prefix, - expected_gate_up_shape, - expected_down_shape, - tuple(gate_up_weight.shape), - gate_up_weight.dtype, - gate_up_scale.dtype, - tuple(down_weight.shape), - down_weight.dtype, - down_scale.dtype, - ) - self.num_shared_experts = 0 - return + return self._backend - transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( - (gate_up_weight, gate_up_scale), - (down_weight, down_scale), - ) - # L1 interleaving allocates a full copy. Re-home the loader Parameter on - # that storage so the original 2*intermediate*hidden FP8 tensor can be - # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. - # The generic linear post-load hook may still repack the serial scales, - # but this shared MLP is never called after native fusion is enabled. - gate_up.weight.data = transformed_l1[0] - self._transformed_shared_l1_weights = ( - gate_up.weight.data, - transformed_l1[1], - ) - self._transformed_shared_l2_weights = transformed_l2 + def _check_runtime_supported(self) -> None: + """Preserved for API compatibility; delegates to the backend check.""" + self._ensure_backend() - def _prepare_shared_expert_scale( - self, - deep_gemm, - linear: nn.Module, - scale: torch.Tensor, - mn: int, - k: int, - ) -> torch.Tensor | None: - block_size = getattr(linear, "weight_block_size", None) - if block_size is None or len(block_size) != 2: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "shared FP8 weight block size is unavailable.", - self.prefix, - ) - return None - - block_m, block_k = block_size - num_k_blocks = (k + block_k - 1) // block_k - # The linear hook's DeepGEMM layout packs four k-block bytes per int32. - packed = scale.dtype == torch.int32 - expected_shape = ( - (mn, (num_k_blocks + 3) // 4) - if packed - else ((mn + block_m - 1) // block_m, num_k_blocks) - ) - if block_k % 32 != 0 or tuple(scale.shape) != expected_shape: - logger.warning( - "Disabling native MegaMoE shared-expert fusion for %s: " - "cannot convert shared scale shape %s with block size %s " - "to MegaMoE's 1x32 layout for weight (%d, %d).", - self.prefix, - tuple(scale.shape), - tuple(block_size), - mn, - k, - ) - return None + def _finalize_shared_expert_weights(self, shared_experts: DeepseekV4MLP) -> None: + """Delegate shared-expert weight transform to the backend and update state. - if packed: - row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] - else: - row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] - scale_1x32 = ( - self._ue8m0_uint8_to_float(row_scale) - .repeat_interleave(block_k // 32, dim=1)[:, : k // 32] - .contiguous() + On success, populates ``self._transformed_shared_l{1,2}_weights``. If + the backend cannot fuse shared experts (returns ``None``), disables + native fusion by setting ``self.num_shared_experts = 0``. + """ + backend = self._ensure_backend() + result = backend.transform_shared_expert_weights( + shared_experts=shared_experts, + num_shared_experts=self.num_shared_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + prefix=self.prefix, ) - # The grouped API is used with a singleton dimension to request the - # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the - # 2D layout required for a shared expert. - return deep_gemm.transform_sf_into_required_layout( - scale_1x32.unsqueeze(0), - mn, - k, - (1, 32), - 1, - ).squeeze(0) + if result is None: + self.num_shared_experts = 0 + return + ( + self._transformed_shared_l1_weights, + self._transformed_shared_l2_weights, + ) = result def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: + backend = self._ensure_backend() from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - if self._transformed_l1_weights is None: - self._check_runtime_supported() - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, - self.hidden_size, - (1, 32), - self.num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, - ) self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, ) ) # Drop the original loader-side parameters: the MegaMoE kernels only @@ -583,7 +448,7 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: return if self._transformed_shared_l1_weights is not None: return - if not self._deep_gemm_supports_shared_experts(deep_gemm): + if not backend.supports_shared_experts(deep_gemm): logger.warning_once( "Disabling native MegaMoE shared-expert fusion because the " "installed DeepGEMM Python API is older than the vLLM " @@ -591,7 +456,8 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: ) self.num_shared_experts = 0 return - self._finalize_shared_expert_weights(deep_gemm, shared_experts) + + self._finalize_shared_expert_weights(shared_experts) @property def has_fused_shared_experts(self) -> bool: @@ -602,6 +468,7 @@ def get_symm_buffer(self): deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -613,6 +480,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, self.num_shared_experts if self.has_fused_shared_experts else 0, + backend.mma_type, ) symm_buffer = self._symm_buffer_cache.get(key) if symm_buffer is None: @@ -626,6 +494,7 @@ def get_symm_buffer(self): num_shared_experts=( self.num_shared_experts if self.has_fused_shared_experts else 0 ), + mma_type=backend.mma_type, ) self._symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -700,7 +569,6 @@ def forward( deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -732,6 +600,8 @@ def forward( else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() shared_x_sf = None shared_block_m = None if self.has_fused_shared_experts: @@ -742,7 +612,7 @@ def forward( symm_buffer.num_max_tokens_per_rank, num_tokens, self.top_k, - "fp8xfp4", + backend.mma_type, ) prepare_megamoe_inputs( @@ -756,27 +626,28 @@ def forward( is_padding=is_padding, shared_x_sf=shared_x_sf, shared_block_m=shared_block_m, + hidden_quant=backend.hidden_quant, ) assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None if self.has_fused_shared_experts: - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, shared_l1_weights=self._transformed_shared_l1_weights, shared_l2_weights=self._transformed_shared_l2_weights, activation_clamp=activation_clamp, fast_math=fast_math, ) else: - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, fast_math=fast_math, ) @@ -1039,9 +910,18 @@ def forward( return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape + # The MegaMoE experts module owns the backend; resolve and assert it. + self.experts._ensure_backend() + assert self.experts._backend is not None, ( + "MegaMoE backend must be resolved before routing." + ) + mega_backend = self.experts._backend # Small local padded batches favor GateLinear; 128-expert gates cross earlier. gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 - if hidden_states.shape[0] <= gate_threshold: + if ( + not mega_backend.supports_bf16_mega_gate() + or hidden_states.shape[0] <= gate_threshold + ): router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, @@ -1060,8 +940,6 @@ def forward( image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, ) else: - from vllm.utils.deep_gemm import bf16_mega_gate - if mega_gate_metadata is None: raise RuntimeError( "MegaMoE routing metadata must be prepared once by the model." @@ -1079,10 +957,10 @@ def forward( ] fix_routing_mask = mega_gate_metadata.hash_token_mask - topk_weights, topk_ids = bf16_mega_gate( - hidden_states, - self.gate.weight, - self.n_activated_experts, + topk_weights, topk_ids = mega_backend.bf16_mega_gate( + x=hidden_states, + weight=self.gate.weight, + num_topk=self.n_activated_experts, scoring_func=self.scoring_func, routed_scaling_factor=self.routed_scaling_factor, ep_rank=self.ep_rank, diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index aa9482bed643..2d93b0d43032 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -9,8 +9,15 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + QuantKey, + kMxfp8Dynamic, +) from vllm.triton_utils import tl, triton +_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) + @triton.jit(do_not_specialize=["num_tokens"]) def _prepare_megamoe_inputs_kernel( @@ -48,6 +55,7 @@ def _prepare_megamoe_inputs_kernel( GROUP_K: tl.constexpr, BLOCK_TOPK: tl.constexpr, SHARED_BLOCK_M: tl.constexpr, + USE_UE8M0: tl.constexpr, ) -> None: token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) token_mask = token_id < num_tokens @@ -69,15 +77,18 @@ def _prepare_megamoe_inputs_kernel( amax = tl.maximum(amax, 1.0e-4) scale = amax / 448.0 - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + if USE_UE8M0: + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + else: + quant_scale = scale hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None] + scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) fp8 = scaled.to(tl.float8e4nv) tl.store( @@ -86,21 +97,32 @@ def _prepare_megamoe_inputs_kernel( mask=k_mask, ) - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( - tl.int32 - ) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - mask=token_mask, - ) + if USE_UE8M0: + scale_offsets = tl.arange(0, num_groups) + packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( + tl.int32 + ) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + else: + sf_offset = tl.arange(0, num_groups) + tl.store( + x_sf + + token_id[:, None] * x_sf_stride_m + + k_block_id * x_sf_stride_k + + sf_offset[None, :], + quant_scale, + mask=token_mask[:, None], + ) # DeepGEMM's SM100 shared-expert TMA loads require the activation scales # in an MN-major layout whose row permutation depends on the MegaMoE # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale # is already resident, avoiding another kernel and temporary tensor. - if shared_x_sf is not None: + if USE_UE8M0 and shared_x_sf is not None: m_block_id = token_id // SHARED_BLOCK_M m_in_block = token_id % SHARED_BLOCK_M aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128 @@ -160,6 +182,27 @@ def _prepare_megamoe_inputs_kernel( ) +def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: + """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" + scale = hidden_quant.scale + group_shape = scale.group_shape + if ( + hidden_quant.dtype != FP8_DTYPE + or not hidden_quant.symmetric + or hidden_quant.scale2 is not None + or scale.static + or not group_shape.is_per_group() + or block_k % group_shape.col != 0 + or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) + ): + raise ValueError( + "DeepSeek V4 MegaMoE input staging requires symmetric dynamic fp8 " + f"per-group quantization with a group size dividing {block_k} and " + f"E8M0 or fp32 scales, got {hidden_quant}." + ) + return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES + + def prepare_megamoe_inputs( hidden_states: torch.Tensor, topk_weights: torch.Tensor, @@ -171,7 +214,26 @@ def prepare_megamoe_inputs( is_padding: torch.Tensor | None = None, shared_x_sf: torch.Tensor | None = None, shared_block_m: int | None = None, + hidden_quant: QuantKey = kMxfp8Dynamic, ) -> None: + """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. + + Args: + hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x_fp8: Output buffer for the fp8-quantized hidden states. + x_sf: Output buffer for the hidden-state scale factors. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x_sf: Optional output buffer for shared-expert scale factors. + Must be given together with ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed + four per int32 into ``x_sf``; fp32 scales are stored directly. + + """ num_tokens, hidden_size = hidden_states.shape if num_tokens == 0: return @@ -180,6 +242,8 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE input staging requires hidden_size to be " "a multiple of 128." ) + block_k = 128 + group_k, use_ue8m0 = _resolve_hidden_quant(hidden_quant, block_k) top_k = topk_ids.shape[1] if topk_weights.shape != topk_ids.shape: raise ValueError( @@ -191,6 +255,11 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE shared input staging requires both " "shared_x_sf and shared_block_m." ) + if shared_x_sf is not None and not use_ue8m0: + raise ValueError( + "DeepSeek V4 MegaMoE shared input staging currently requires " + "UE8M0-packed hidden scales." + ) if shared_x_sf is not None: assert shared_block_m is not None if shared_block_m <= 0: @@ -209,7 +278,6 @@ def prepare_megamoe_inputs( f"{required_rows}, got {shared_x_sf.shape[0]}." ) - block_k = 128 # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. block_m = 8 if num_tokens >= 64 else 1 grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) @@ -247,8 +315,9 @@ def prepare_megamoe_inputs( top_k, BLOCK_M=block_m, BLOCK_K=block_k, - GROUP_K=32, + GROUP_K=group_k, BLOCK_TOPK=block_topk, SHARED_BLOCK_M=shared_block_m or 1, + USE_UE8M0=use_ue8m0, num_warps=4, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index cc1ca609a6ba..84e93b82fbb4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -365,27 +365,16 @@ def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: return self._check_runtime_supported() - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - w13_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), - 2 * self.intermediate_size, - self.hidden_size, - (1, 32), - self.num_local_experts, - ) - w2_scale = deep_gemm.transform_sf_into_required_layout( - self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), - self.hidden_size, - self.intermediate_size, - (1, 32), - self.num_local_experts, - ) + backend = self._ensure_backend() self._transformed_l1_weights, self._transformed_l2_weights = ( - deep_gemm.transform_weights_for_mega_moe( - (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), - (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, activation=self.activation, ) ) @@ -407,6 +396,7 @@ def get_symm_buffer(self): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -418,6 +408,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, self.activation, + backend.mma_type, ) symm_buffer = self._kimi_symm_buffer_cache.get(key) if symm_buffer is None: @@ -429,6 +420,7 @@ def get_symm_buffer(self): self.hidden_size, self.intermediate_size, activation=self.activation, + mma_type=backend.mma_type, ) self._kimi_symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -449,10 +441,6 @@ def forward( f"but its symmetric buffer supports {self.max_num_tokens}." ) y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -483,6 +471,8 @@ def forward( else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() prepare_megamoe_inputs( hidden_states, topk_weights, @@ -492,22 +482,23 @@ def forward( symm_buffer.topk_idx[:num_tokens], symm_buffer.topk_weights[:num_tokens], is_padding=is_padding, + hidden_quant=backend.hidden_quant, ) self.finalize_weights() assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, - activation=self.activation, + fast_math=fast_math, # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation=self.activation, activation_alpha=self.activation_beta or 1.0, activation_beta=self.activation_linear_beta or 0.0, - fast_math=fast_math, ) return y From 816173ca9d9256f85c36d09d133a7ed219ccbd8e Mon Sep 17 00:00:00 2001 From: "Jiahan Chang (Cyrus)" <173873397+jiahanc@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:38:07 +0800 Subject: [PATCH 177/952] [Kernel] Remove WAR for LL CuTeDSL BF16 GEMM (#60434) Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> --- tests/kernels/test_flashinfer_bf16_gemm.py | 2 + .../test_flashinfer_autotune_warmup.py | 9 ++-- vllm/model_executor/layers/utils.py | 5 +- vllm/model_executor/warmup/kernel_warmup.py | 50 +------------------ vllm/utils/flashinfer.py | 43 +++------------- 5 files changed, 16 insertions(+), 93 deletions(-) diff --git a/tests/kernels/test_flashinfer_bf16_gemm.py b/tests/kernels/test_flashinfer_bf16_gemm.py index 63a5a39c2e67..28520a6083e4 100644 --- a/tests/kernels/test_flashinfer_bf16_gemm.py +++ b/tests/kernels/test_flashinfer_bf16_gemm.py @@ -25,6 +25,8 @@ def require_flashinfer_bf16_cutedsl() -> None: (8, True, False), (16, False, True), (32, True, True), + (64, True, True), + (512, False, False), ], ) def test_flashinfer_bf16_cutedsl_correctness( diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index fe89b09d09bc..f7ea6847f26f 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -39,12 +39,11 @@ def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): return moe -def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): +def _make_runner(modules, *, max_tokens: int = 8192): """Create a runner carrying only the state used by FlashInfer warmup.""" return SimpleNamespace( scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), vllm_config=SimpleNamespace( - kernel_config=SimpleNamespace(linear_backend=linear_backend), attention_config=SimpleNamespace(hisparse_config=None), parallel_config=SimpleNamespace(data_parallel_rank=0), ), @@ -63,21 +62,19 @@ def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): _make_moe(max_deferred_tokens=128), _make_moe(max_deferred_tokens=64, enabled=False), _make_moe(max_deferred_tokens=-1), - ], - linear_backend="flashinfer_cutedsl", + ] ) with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): token_counts = _flashinfer_autotune_token_counts(runner) - assert token_counts == (8192, 32, 128) + assert token_counts == (8192, 128) def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): runner = _make_runner( [_make_moe(max_deferred_tokens=4096)], max_tokens=32, - linear_backend="flashinfer_cutedsl", ) with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index d4ee153fb251..80eef7c75a12 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -128,16 +128,13 @@ def _can_use_flashinfer_cutedsl_bf16( if ( k <= 0 or n <= 0 + or x.numel() == 0 or weight.shape[1] != k or k % 128 != 0 or x.data_ptr() % 32 != 0 or weight.data_ptr() % 32 != 0 ): return False - - m = x.numel() // k - if not 1 <= m <= 32: - return False return bias is None or ( bias.is_cuda and bias.device == x.device diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d6544333beb5..e5751730e85d 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -316,9 +316,6 @@ def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: return None -_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS = 32 - - def _flashinfer_deferred_moe_token_counts( runner: "GPUModelRunner", ) -> tuple[int, ...]: @@ -339,19 +336,10 @@ def _flashinfer_deferred_moe_token_counts( return tuple(dict.fromkeys(token_counts)) -def _flashinfer_autotune_token_counts( - runner: "GPUModelRunner", *, include_bf16: bool = True -) -> tuple[int, ...]: +def _flashinfer_autotune_token_counts(runner: "GPUModelRunner") -> tuple[int, ...]: max_tokens = runner.scheduler_config.max_num_batched_tokens # Tune the widest bucket set first so bounded passes reuse its configs. token_counts = [max_tokens] - linear_backend = runner.vllm_config.kernel_config.linear_backend - if ( - include_bf16 - and linear_backend == "flashinfer_cutedsl" - and max_tokens > _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - ): - token_counts.append(_FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS) token_counts.extend(_flashinfer_deferred_moe_token_counts(runner)) return tuple(dict.fromkeys(token_counts)) @@ -362,7 +350,7 @@ def _run_flashinfer_autotune_dummy_runs( import vllm.utils.flashinfer as fi_utils dummy_run_kwargs = {"skip_attn": True} if skip_attn else {} - for num_tokens in _flashinfer_autotune_token_counts(runner, include_bf16=False): + for num_tokens in _flashinfer_autotune_token_counts(runner): # The drafter may use more tokens than this pass in the same dummy run. # Include its max M when building the autotune buckets. max_tuning_tokens = num_tokens @@ -398,36 +386,6 @@ def _run_flashinfer_autotune_dummy_runs( ) -def _run_flashinfer_bf16_autotune_dummy_run( - runner: "GPUModelRunner", - *, - skip_ops: set[str] | None = None, - skip_attn: bool = False, -) -> None: - import vllm.utils.flashinfer as fi_utils - - if ( - runner.vllm_config.kernel_config.linear_backend != "flashinfer_cutedsl" - or runner.scheduler_config.max_num_batched_tokens - <= _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - ): - return - - num_tokens = _FLASHINFER_BF16_AUTOTUNE_MAX_TOKENS - tuning_buckets = fi_utils.flashinfer_get_hybrid_num_tokens_buckets(num_tokens) - logger.info("Running FlashInfer BF16-only autotune with %d tokens.", num_tokens) - # Other ops execute from their existing cache or fallback, without the - # BF16 bucket cap. No full-model autotune context may enclose this pass. - with fi_utils.autotune_bf16_only(tuning_buckets, skip_ops=skip_ops): - runner._dummy_run( - num_tokens=num_tokens, - skip_eplb=True, - is_profile=True, - randomize_inputs=True, - **({"skip_attn": True} if skip_attn else {}), - ) - - def _autotune_cache_fingerprint(path: Path) -> tuple[str, int] | None: """Identify a saved autotune file by its FlashInfer metadata and size.""" try: @@ -534,10 +492,6 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) - with torch.inference_mode(): - _run_flashinfer_bf16_autotune_dummy_run( - runner, skip_ops=skip_ops, skip_attn=hisparse_enabled - ) finally: set_autotune_process_group(None) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index d7fa807c33b3..f9d6c4439f56 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -11,8 +11,7 @@ import importlib.util import os import shutil -from collections.abc import Callable, Iterator -from contextvars import ContextVar +from collections.abc import Callable from typing import Any, NoReturn import requests @@ -27,24 +26,6 @@ logger = init_logger(__name__) -_bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar( - "flashinfer_bf16_autotune_buckets", default=None -) - - -@contextlib.contextmanager -def autotune_bf16_only( - tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None -) -> Iterator[None]: - """Tune BF16 calls with bounded buckets, outside full-model autotuning.""" - token = _bf16_autotune_buckets.set(tuning_buckets) - try: - with autotune(tune_mode=False, skip_ops=skip_ops): - yield - finally: - _bf16_autotune_buckets.reset(token) - - def flashinfer_bf16_mm( a: torch.Tensor, b: torch.Tensor, @@ -54,21 +35,14 @@ def flashinfer_bf16_mm( ) -> torch.Tensor: from flashinfer import mm_bf16 - tuning_buckets = _bf16_autotune_buckets.get() - tuning = ( - autotune(tune_mode=True, tuning_buckets=tuning_buckets) - if tuning_buckets is not None - else contextlib.nullcontext() + return mm_bf16( + a, + b, + bias=bias, + pdl=pdl, + out_dtype=torch.bfloat16, + backend=backend, ) - with tuning: - return mm_bf16( - a, - b, - bias=bias, - pdl=pdl, - out_dtype=torch.bfloat16, - backend=backend, - ) # This is the storage path for the cubins, it can be replaced @@ -1295,7 +1269,6 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: __all__ = [ "has_flashinfer", "flashinfer_bf16_mm", - "autotune_bf16_only", "has_flashinfer_bf16_gemm", "is_flashinfer_bf16_gemm_supported", "is_flashinfer_cutedsl_bf16_gemm_supported", From 95307ac6d4cd8be2c1e38f236e9104c569af9025 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Thu, 8 Oct 2026 14:00:58 +0800 Subject: [PATCH 178/952] [Doc] Update ViT CUDA graph design doc (#60519) Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- docs/design/cuda_graphs_multimodal.md | 96 ++++++++++++--------------- 1 file changed, 43 insertions(+), 53 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index dca200e2de03..0a8a831f0164 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -7,27 +7,20 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili !!! note Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md). -## Compatibility Matrix - -!!! note - The symbols used below have the following meanings: - - - ✅ = Full compatibility - - 🟠 = Partial compatibility - - ❌ = No compatibility - - ❔ = Unknown or TBD - -### Model x Feature +## Supported Models | Architecture | Models | CG for Image | CG for Video | Multi-Path Graph | | ------------ | ------ | ------------ | ------------ | --------------- | | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `DeepseekV41ForCausalLM` | `DeepSeek-V4.1-Flash` | ✅︎ | ❌︎ | ❌︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | | `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | +| `KimiK25ForConditionalGeneration` | `Kimi-K2.5, Kimi-K2.6` | ✅︎ | ❌︎ | ❌︎ | +| `KimiK3ForConditionalGeneration` | `Kimi-K3` | ✅︎ | ❌︎ | ❌︎ | | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | @@ -38,26 +31,6 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili | `MiniCPMV` | `MiniCPMV2.5` | ✅︎ | ❌︎ | ❌︎ | | `MiniCPMV` | `MiniCPMV2.6`, `MiniCPMV4.0` | ✅︎ | ✅︎ | ❌︎ | -### Model x Hardware - -| Architecture | NV Blackwell | NV Ampere | AMD MI300X | AMD MI350X / MI355X | -| ------------ | ---------------- | ------------- | -------------- | --------------------- | -| `DeepseekOCRForCausalLM` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Ernie4_5_VLMoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Gemma3ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Glm4vForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Gemma4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `InternVLChatModel` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `KimiVLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Llama4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen2VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen2_5_VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | -| `MiniCPMV` | ✅︎ | ✅︎ | ❔ | ❔ | - !!! note Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. @@ -80,7 +53,7 @@ The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, m * [EncoderItemSpec][vllm.v1.worker.encoder_cudagraph_defs.EncoderItemSpec]: describes a single encoder input item (image or video) with its input size and output token count. * [BudgetGraphMetadata][vllm.v1.worker.encoder_cudagraph.BudgetGraphMetadata]: holds the captured CUDA Graph and its associated I/O buffers for a single token budget level. -### Budget-based graph capture +### Budget-Based Graph Capture Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g., `[2048, 4096, 8192, 13824]`). Each budget defines a fixed token capacity, and all budgets share the same maximum batch size (number of images). The `BudgetGraphMetadata` for each level stores the graph along with pre-allocated input, metadata, and output buffers: @@ -99,20 +72,7 @@ Budgets are auto-generated as power-of-2 levels from a model-provided range via Each entry in `EncoderCudaGraphConfig.paths` defines an independently captured encoder path. A path can provide its own minimum token budget and opt into zero-token batches; the manager generates and stores a separate budget graph set for every configured path. -### Greedy bin-packing at runtime - -When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. - -For multi-path models, the same greedy packing loop constrains every configured path simultaneously (see [Multi-Path graph capture](#multi-path-graph-capture)). - -For each graph replay: - -1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. -2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them. -3. Replay the CUDA Graph. -4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). - -### Multi-Path graph capture +### Multi-Path Graph Capture `EncoderCudaGraphConfig.paths` maps path names to `EncoderCudaGraphPathConfig` capture policies. For example, DeepSeek-OCR configures a **global** image path and a **local** patch path, which are captured independently under `budget_graphs["global"]` and `budget_graphs["local"]`. @@ -127,7 +87,7 @@ Both lists are capped at the same `max_budget`. * Sort images by total output tokens (global + local), smallest first. * Greedily pack images: an image is added to the current sub-batch only if every accumulated path token count is within that path's maximum budget, with the image count ≤ `max_batch_size`. -* Once any path constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path. +* Once any path constraint would overflow, or the budget-cliff deferral condition fires on any path, finalize the sub-batch and find the smallest fitting budget **independently** for each path. * Repeat until all images are packed. **Partial graph fallback.** For each non-empty path, the manager replays the smallest fitting graph or runs only that path eagerly when no graph fits. Paths with zero tokens are skipped; a `0`-budget graph is never captured or replayed. @@ -139,11 +99,41 @@ Both lists are capped at the same `max_budget`. !!! note The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image. -### Data-parallel support +### Greedy Bin-Packing at Runtime + +When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint, or deferring it is cheaper), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution. + +For multi-path models, the same greedy packing loop constrains every configured path simultaneously. + +For each graph replay: + +1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. +2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them. +3. Replay the CUDA Graph. +4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). + +**Budget-cliff deferral:** + +Greedy packing can be wasteful when a large image arrives after small ones: merging it into the current sub-batch may force the whole batch into a much larger budget graph even though splitting would cost less. Because token budgets are discrete (e.g., `[256, 4096, 8192]`), one extra token can double the required graph level — a "budget cliff". Images that overflow the largest budget still finalize the sub-batch unconditionally; for images that *fit*, the packer additionally applies a deferral test: finalize the current sub-batch (replaying it at its own small budget) and start a new sub-batch with the deferred image. + +For each candidate image, the packer compares three budgets per path: `b_cur` (the smallest budget fitting the current sub-batch), `b_alone` (the smallest budget fitting the image alone), and `b_merged` (the smallest budget fitting both). Merging costs `b_merged`; deferring costs `b_cur + b_alone` plus one extra replay. The image is deferred when, on any path, both conditions hold: + +1. `b_cur + b_alone < b_merged` — splitting strictly saves budget. The strict inequality makes ties merge: with budgets `[256, 4096, 8192]`, two 4000-token images tie (`4096 + 4096 == 8192`), and merging wins because a single replay avoids an extra graph launch. +2. `b_alone >= b_cur` — the image alone occupies at least as large a budget bucket as the entire current sub-batch. + +**Condition 1 is the cost test.** For example, a 200-token image followed by a 4000-token image with budgets `[256, 4096, 8192]`: merging takes `b_merged = 8192` (4200 tokens round up to 8192), while deferring takes `256 + 4096 = 4352` — the image is deferred, saving 3840 tokens of padding compute. + +**Condition 2 separates "whales" from "tag-alongs".** Because images are sorted smallest first, a whale arrives last, and deferring it lets the small sub-batch replay at its own small budget instead of being dragged into a large bucket for one image's sake. Without condition 2, deferring mid-size images fragments a flood of similar images into extra replays: with budgets `[256, 2048, 4096, 8192]` and six equal 1366-token images, condition 1 alone defers every third image and yields three 4096 replays (12288 budget tokens in total), while with condition 2 five images merge into one 8192 replay plus a single 2048 replay — fewer replays *and* fewer budget tokens (10240). + +The conditions compose cleanly: when `b_alone == b_cur`, condition 1 requires `b_merged > 2 * b_cur`, i.e. deferral fires only when merging jumps at least two budget levels, so ties always merge. This also lets small images defer across a huge cliff: with budgets `[256, 4096]` and four 100-token images, adding the third image (300 tokens) would jump from 256 to 4096, so the packer defers it and packs the images as two 256-budget pairs. + +For multi-path items, the conditions are evaluated per path and OR-ed: any path on which the image is a whale defers the whole image, while paths where the image contributes zero tokens are skipped. + +### ViT DP Mode Support When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`. -### Video inference support +### Video Inference Support Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. @@ -152,7 +142,7 @@ Following (ViT full CUDA graph Mixed inputs (image+video) per prompt are also supported now. -## Model integration via `SupportsEncoderCudaGraph` +## Model Integration via `SupportsEncoderCudaGraph` Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods: @@ -180,9 +170,9 @@ Four fields in `CompilationConfig` control encoder CUDA Graphs: Multi-path mode is configured at the model level through `EncoderCudaGraphConfig.paths`. Each `EncoderCudaGraphPathConfig` can set a path-specific minimum budget and whether zero-token batches are allowed. The manager automatically generates separate budget lists and uses the same execution loop for single- and multi-path models. -## Usage guide +## Usage Guide -### Image inference +### Image Inference Enable encoder CUDA Graphs via `compilation_config`: @@ -226,7 +216,7 @@ model = vllm.LLM( The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets). -### Video inference +### Video Inference Enable encoder CUDA Graphs via `compilation_config`: From fb2ac82458cf6bff45b21378a85fa659213b51e5 Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Wed, 7 Oct 2026 23:15:18 -0700 Subject: [PATCH 179/952] [ROCm][DSv4.1][Perf] Replace the Torch prefill combine with the shared combine_topk_swa_indices kernel (#56638) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Ren Yuzhou <54501155+yuzhouo7@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 158 ++++++++++++++++- .../attention/test_deepseek_v4_swa_visible.py | 23 ++- vllm/models/deepseek_v41/amd/rocm.py | 164 +----------------- .../deepseek_v41/common/ops/cache_utils.py | 10 +- 4 files changed, 185 insertions(+), 170 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index ca2cf2eb7017..362ac6b7ce0a 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -41,7 +41,7 @@ def _on_gfx950() -> bool: # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), - reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", + reason="split-K decode and the V4.1 combine tests need an AMD gfx942/gfx950", ) requires_gfx950 = pytest.mark.skipif( not _on_gfx950(), @@ -2407,3 +2407,159 @@ def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None: torch.testing.assert_close( gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3 ) + + +def _combine_topk_swa_indices_torch( + topk_indices, + query_start_loc, + seq_lens, + gather_lens, + window_size, + compress_ratio, + topk, + M, + N, +): + """The Torch path this kernel replaces, copied from before the switch. + + Kept verbatim so the reference cannot drift toward the kernel. + """ + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + _SPARSE_PREFILL_TOPK_ALIGNMENT, + ) + + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() + num_tokens = topk_indices.shape[0] + combined_topk = ( + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + // _SPARSE_PREFILL_TOPK_ALIGNMENT + * _SPARSE_PREFILL_TOPK_ALIGNMENT + ) + combined_indices = torch.full( + (num_tokens, combined_topk), + fill_value=-1, + dtype=torch.int32, + device=topk_indices.device, + ) + combined_lens = torch.empty( + num_tokens, dtype=torch.int32, device=topk_indices.device + ) + + # query_start_loc may have a non-zero base for a narrowed mixed batch. + query_lens = query_start_loc[1:] - query_start_loc[:-1] + req_ids = torch.repeat_interleave( + torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens + ) + query_starts = query_start_loc[:-1] - query_start_loc[0] + token_offsets = torch.arange(num_tokens, device=seq_lens.device) - ( + torch.repeat_interleave(query_starts, query_lens) + ) + positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets + + logical_topk_width = min(topk, topk_indices.shape[1]) + topk_lens = torch.minimum( + (positions + 1) // compress_ratio, + torch.full_like(positions, logical_topk_width), + ).clamp_min(0) + topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device) + topk_mask = topk_offsets[None, :] < topk_lens[:, None] + topk_values = topk_indices[:, :logical_topk_width].to(torch.int32) + topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N) + combined_indices[:, :logical_topk_width] = torch.where( + topk_valid, + topk_values + (M * req_ids).to(torch.int32)[:, None], + -1, + ) + + swa_lens = torch.minimum( + positions + 1, torch.full_like(positions, window_size) + ).clamp_min(0) + swa_offsets = torch.arange(window_size, device=seq_lens.device) + swa_mask = swa_offsets[None, :] < swa_lens[:, None] + swa_columns = topk_lens[:, None] + swa_offsets[None, :] + gather_starts = seq_lens - gather_lens + swa_values = ( + M * req_ids[:, None] + + N + + swa_offsets[None, :] + + positions[:, None] + - swa_lens[:, None] + + 1 + - gather_starts[req_ids, None] + ).to(torch.int32) + rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as( + swa_columns + ) + flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask] + combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask]) + combined_lens.copy_((topk_lens + swa_lens).to(torch.int32)) + return combined_indices, combined_lens + + +def _v41_combine_case(case, window): + dev = "cuda" + ratio, topk = 2, 512 + if case == "seq_len_below_query_len": + qsl = torch.tensor([0, 64], dtype=torch.int32, device=dev) + sl = torch.tensor([32], dtype=torch.int32, device=dev) + rows = 64 + elif case == "seq_len_zero": + # The warmup batch: every pos is negative, so both lengths clamp to 0. + qsl = torch.tensor([0, 32, 64], dtype=torch.int32, device=dev) + sl = torch.zeros(2, dtype=torch.int32, device=dev) + rows = 64 + elif case == "compress_ratio_zero": + qsl = torch.tensor([0, 4, 8], dtype=torch.int32, device=dev) + sl = torch.tensor([4, 4], dtype=torch.int32, device=dev) + rows, ratio, topk = 8, 0, 0 + elif case == "nonzero_base": + # Prefill rows of a mixed batch: query_start_loc starts past the decodes. + qsl = torch.tensor([32, 64, 128], dtype=torch.int32, device=dev) + sl = torch.tensor([48, 80], dtype=torch.int32, device=dev) + rows = 96 + else: + qsl = torch.tensor([0, 32, 96], dtype=torch.int32, device=dev) + sl = torch.tensor([48, 80], dtype=torch.int32, device=dev) + rows = 96 + # What the SWA builder gathers: the query plus up to window - 1 rows before it. + gl = torch.minimum(sl, qsl[1:] - qsl[:-1] + window - 1) + set_random_seed(0) + ti = torch.randint(0, 4096, (rows, max(topk, 1)), dtype=torch.int32) + if case == "invalid_topk": + # Failed candidates stay -1 instead of pointing into another request. + ti[:, ::7] = -1 + return ti.to(dev), qsl, sl, gl, window, ratio, topk, 100000, 4096 + + +@requires_split_decode_arch +@pytest.mark.parametrize( + "case,window", + [ + ("seq_len_below_query_len", 128), + ("seq_len_zero", 128), + ("compress_ratio_zero", 128), + ("nonzero_base", 128), + ("invalid_topk", 128), + ("plain", 64), + ("plain", 128), + ("plain", 256), + ], +) +def test_v41_combine_topk_swa_indices_matches_torch(case, window) -> None: + from vllm.models.deepseek_v41.amd.rocm import combine_topk_swa_indices + + args = _v41_combine_case(case, window) + got_indices, got_lens = combine_topk_swa_indices(*args) + ref_indices, ref_lens = _combine_topk_swa_indices_torch(*args) + assert torch.equal(got_lens, ref_lens), ( + f"lens differ: ref={ref_lens.tolist()} got={got_lens.tolist()}" + ) + assert torch.equal(got_indices, ref_indices) + + +def test_v41_rocm_combine_uses_shared_kernel() -> None: + """ROCm prefill takes the shared combine, so it inherits the replay bound.""" + from vllm.models.deepseek_v41.amd import rocm + from vllm.models.deepseek_v41.common.ops import cache_utils + + assert rocm.combine_topk_swa_indices is cache_utils.combine_topk_swa_indices diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 7a74ce8c81be..8f08929549a8 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -302,8 +302,13 @@ def combine_case( with_image: bool, replay_starts: list[int] | None = None, combine_fn=combine_topk_swa_indices, + in_range_topk: bool = False, ): - """Run combine_topk_swa_indices and return (indices, lens, expected).""" + """Run combine_topk_swa_indices and return (indices, lens, expected). + + ``in_range_topk`` draws top-k indices below N, for combiners that drop + out-of-range ones. + """ device = torch.device("cuda") num_reqs = len(seq_lens) replay_starts = replay_starts or [0] * num_reqs @@ -326,7 +331,11 @@ def combine_case( M = N + int(gather_lens.max()) + 8 gen = torch.Generator(device="cpu").manual_seed(0) topk_indices = torch.randint( - 0, 4096, (num_tokens, max(topk, 1)), generator=gen, dtype=torch.int32 + 0, + N if in_range_topk else 4096, + (num_tokens, max(topk, 1)), + generator=gen, + dtype=torch.int32, ).to(device) topk_indices = topk_indices[:, : max(topk, 1)] @@ -415,24 +424,26 @@ def test_combine_topk_swa_with_image_spans(cfg): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("cfg", COMBINE_CASES) -def test_v41_combine_topk_swa_stops_at_replay_start(cfg): +@pytest.mark.parametrize("query_len", [24, 20]) +def test_v41_combine_topk_swa_stops_at_replay_start(cfg, query_len): """SWA bounded replay: the gathered buffer starts at replay_start, so the window never indexes below it.""" from vllm.models.deepseek_v41.common.ops.cache_utils import ( combine_topk_swa_indices as combine_v41, ) - # Request 0 replays [16, 40): the windows of its first rows would - # otherwise reach below 16. + # Request 0 replays from 16: a 24-token chunk starts there, a 20-token + # chunk leaves only 4 context tokens above it (fewer than the window). indices, lens, rows, exp_lens = combine_case( cfg["compress_ratio"], cfg["topk"], seq_lens=[40, 12], - query_lens=[24, 12], + query_lens=[query_len, 12], spans=[[], []], with_image=False, replay_starts=[16, 0], combine_fn=combine_v41, + in_range_topk=True, ) assert lens.cpu().tolist() == exp_lens assert indices.cpu().tolist() == rows diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 64b4a9d3d1f2..3aa54f9b0e97 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -26,7 +26,10 @@ DeepseekV4IndexerCache, _replace_layer_index, ) -from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache +from vllm.models.deepseek_v41.common.ops import ( + combine_topk_swa_indices, + dequantize_and_gather_k_cache, +) from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, @@ -129,165 +132,6 @@ def apply_pre_quantized_block_scaled_mm( return out.to(dtype=kernel.config.out_dtype) -# ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes -# do not touch the shared DeepSeek V4 cache utilities. -_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 - - -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - TOPK_WIDTH: tl.constexpr, - PADDED_TOP_K: tl.constexpr, -): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) - - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - gather_start = seq_len - gather_len - - for token_idx in range(query_start + worker_id, query_end, num_workers): - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - - topk_offset = tl.arange(0, PADDED_TOP_K) - topk_mask = topk_offset < topk_len - safe_topk_offset = tl.where(topk_offset < TOPK_WIDTH, topk_offset, 0) - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + safe_topk_offset, - mask=topk_mask, - other=-1, - ) - valid_topk = (topk_indices >= 0) & (topk_indices < N) - topk_indices = tl.where(valid_topk, topk_indices + M * batch_idx, -1) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + topk_offset, - topk_indices, - mask=topk_mask, - ) - - swa_offset = tl.arange(0, WINDOW_SIZE) - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + swa_offset, - M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start, - mask=swa_offset < swa_len, - ) - - tl.store(combined_lens_ptr + token_idx, topk_len + swa_len) - - -def combine_topk_swa_indices( - topk_indices: torch.Tensor, - query_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor, - window_size: int, - compress_ratio: int, - topk: int, - M: int, - N: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Combine compressed-attention and sliding-window indices with Torch. - - The Triton implementation inherited from DeepSeek V4 launches a - two-dimensional grid with 128 workers per request. On gfx950 it can issue - an out-of-bounds access for V4.1's mixed prefill metadata (including the - synthetic mixed-token warmup). This path is prefill-only and the tensors - are small, so use ordinary Torch indexing until a gfx950-safe fused kernel - is available. - """ - topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous() - num_tokens = topk_indices.shape[0] - combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) - // _SPARSE_PREFILL_TOPK_ALIGNMENT - * _SPARSE_PREFILL_TOPK_ALIGNMENT - ) - combined_indices = torch.full( - (num_tokens, combined_topk), - fill_value=-1, - dtype=torch.int32, - device=topk_indices.device, - ) - combined_lens = torch.empty( - num_tokens, dtype=torch.int32, device=topk_indices.device - ) - - # query_start_loc may have a non-zero base for a narrowed mixed batch. - query_lens = query_start_loc[1:] - query_start_loc[:-1] - req_ids = torch.repeat_interleave( - torch.arange(seq_lens.shape[0], device=seq_lens.device), query_lens - ) - query_starts = query_start_loc[:-1] - query_start_loc[0] - token_offsets = torch.arange(num_tokens, device=seq_lens.device) - ( - torch.repeat_interleave(query_starts, query_lens) - ) - positions = seq_lens[req_ids] - query_lens[req_ids] + token_offsets - - logical_topk_width = min(topk, topk_indices.shape[1]) - topk_lens = torch.minimum( - (positions + 1) // compress_ratio, - torch.full_like(positions, logical_topk_width), - ).clamp_min(0) - topk_offsets = torch.arange(logical_topk_width, device=seq_lens.device) - topk_mask = topk_offsets[None, :] < topk_lens[:, None] - topk_values = topk_indices[:, :logical_topk_width].to(torch.int32) - topk_valid = topk_mask & (topk_values >= 0) & (topk_values < N) - combined_indices[:, :logical_topk_width] = torch.where( - topk_valid, - topk_values + (M * req_ids).to(torch.int32)[:, None], - -1, - ) - - swa_lens = torch.minimum( - positions + 1, torch.full_like(positions, window_size) - ).clamp_min(0) - swa_offsets = torch.arange(window_size, device=seq_lens.device) - swa_mask = swa_offsets[None, :] < swa_lens[:, None] - swa_columns = topk_lens[:, None] + swa_offsets[None, :] - gather_starts = seq_lens - gather_lens - swa_values = ( - M * req_ids[:, None] - + N - + swa_offsets[None, :] - + positions[:, None] - - swa_lens[:, None] - + 1 - - gather_starts[req_ids, None] - ).to(torch.int32) - rows = torch.arange(num_tokens, device=seq_lens.device)[:, None].expand_as( - swa_columns - ) - flat_dst = rows[swa_mask] * combined_topk + swa_columns[swa_mask] - combined_indices.view(-1).index_copy_(0, flat_dst, swa_values[swa_mask]) - combined_lens.copy_((topk_lens + swa_lens).to(torch.int32)) - return combined_indices, combined_lens - - @triton.jit def _compute_topk_lens_kernel( topk_lens_ptr, diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index b7f07d8d84c0..a6139c71f836 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -1089,7 +1089,8 @@ def kernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query if COMPRESS_RATIO > 0: - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + # Warmup batches can put pos below zero; keep the length >= 0. + topk_len = tl.maximum(tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K), 0) else: # SWA-only layers pass TOP_K=0; skip the division entirely # (integer div by constexpr 0 is UB and yields garbage, which @@ -1100,7 +1101,7 @@ def kernel( # start (SWA bounded replay), so the window cannot start before # the gathered buffer does. swa_start = tl.maximum(swa_start, gather_start) - swa_len = pos - swa_start + 1 + swa_len = tl.maximum(pos - swa_start + 1, 0) offset = tl.arange(0, PADDED_TOP_K) mask = offset < topk_len @@ -1108,9 +1109,12 @@ def kernel( topk_indices_ptr + token_idx * topk_indices_stride + offset, mask=mask, ) + # A failed (-1) or out-of-pool candidate stays -1 instead of landing + # in the previous request's rows or this request's SWA rows. + valid = (topk_indices >= 0) & (topk_indices < N) tl.store( combined_indices_ptr + token_idx * combined_indices_stride + offset, - topk_indices + M * batch_idx, + tl.where(valid, topk_indices + M * batch_idx, -1), mask=mask, ) # Index into gathered buffer: N + (position - gather_start) From f226af0e114bb67bf59fe6ee2fc762dcbf956b4c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 8 Oct 2026 14:25:27 +0800 Subject: [PATCH 180/952] [Bugfix] Eliminate frontend ZMQ port TOCTOU with inherited listeners (#54113) Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 --- rust/Cargo.lock | 62 +++- rust/Cargo.toml | 2 +- rust/src/cmd/src/cli.rs | 22 +- rust/src/cmd/src/cli/tests.rs | 192 ++++++------ rust/src/engine-core-client/Cargo.toml | 2 + rust/src/engine-core-client/src/client.rs | 25 +- .../engine-core-client/src/tests/client.rs | 38 ++- rust/src/engine-core-client/src/transport.rs | 68 +++- rust/src/server/src/grpc/tests.rs | 10 +- .../api_server/_api_server_spawn_workers.py | 14 - .../test_api_server_process_manager.py | 294 +++++++++--------- tests/utils_/test_network_utils.py | 47 +++ .../engine/test_core_engine_actor_manager.py | 90 ++---- vllm/entrypoints/cli/serve.py | 35 +-- vllm/utils/network_utils.py | 103 +++++- vllm/v1/engine/core_client.py | 57 ++-- vllm/v1/engine/utils.py | 85 ++++- vllm/v1/utils.py | 235 +++++++------- 18 files changed, 793 insertions(+), 588 deletions(-) delete mode 100644 tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9dcdb0bdbc2c..331ebfa8dfaa 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -593,6 +593,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.44" @@ -1614,6 +1625,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -3451,6 +3463,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -3489,6 +3512,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_distr" version = "0.5.1" @@ -3936,9 +3965,9 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "saa" -version = "5.5.0" +version = "5.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16c7f49c9d5caa3bf4b3106900484b447b9253fe99670ceb81cb6cb5027855e1" +checksum = "bd85992a680de788217d9c70b74aa7685b108cf1aba02f728fb27f1b27199ad1" [[package]] name = "same-file" @@ -3960,12 +3989,12 @@ dependencies = [ [[package]] name = "scc" -version = "3.6.9" +version = "3.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45bb5ce9efd4a6e7b0f86c2697fe4c1d78d1f4e6d988c54b752d577cafe22fe8" +checksum = "5d5f8ca67c4f6208caf824037a6a9df0211b4e0937fccc4f543441f294d04511" dependencies = [ "saa", - "sdd 4.7.3", + "sdd 4.8.11", ] [[package]] @@ -4039,9 +4068,12 @@ checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" [[package]] name = "sdd" -version = "4.7.3" +version = "4.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b21a75f5913ab130e4b369fb8693be25f29b983e2ecad4279df9bfa5dd8aaf3e" +checksum = "a31877eee03686b905845993b5d9809a220c9dd5a10e9d6021d7ca67384c9070" +dependencies = [ + "saa", +] [[package]] name = "secrecy" @@ -4991,7 +5023,6 @@ dependencies = [ "bytes", "libc", "mio", - "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -5799,6 +5830,7 @@ dependencies = [ "half", "hex", "itertools 0.14.0", + "libc", "parking_lot", "rmp-serde", "rmpv", @@ -5808,6 +5840,7 @@ dependencies = [ "serde_repr", "serde_tuple", "serde_with", + "socket2", "task-local", "tempfile", "thiserror 2.0.18", @@ -6643,22 +6676,21 @@ checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" [[package]] name = "zeromq" version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efb2c254fd8f366755335c9e43b865f8484fe3bd717d65ffe7c3f28852863030" +source = "git+https://github.com/zeromq/zmq.rs?rev=67dcba97559ebe8e8206dab753710d8b9cb07591#67dcba97559ebe8e8206dab753710d8b9cb07591" dependencies = [ "async-trait", "asynchronous-codec", "bytes", "crossbeam-queue", "futures", + "futures-channel", "log", - "num-traits", - "once_cell", "parking_lot", - "rand 0.9.5", + "rand 0.10.3", "regex", - "scc 3.6.9", - "thiserror 1.0.69", + "scc 3.8.8", + "socket2", + "thiserror 2.0.18", "tokio", "tokio-util", "uuid", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 38436467f862..57f680cc147e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -155,7 +155,7 @@ vllm-tokenizer = { path = "src/tokenizer" } vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.4.1" -zeromq = { version = "0.6.0", default-features = false, features = [ +zeromq = { git = "https://github.com/zeromq/zmq.rs", rev = "67dcba97559ebe8e8206dab753710d8b9cb07591", default-features = false, features = [ "tokio-runtime", "all-transport", ] } diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 1269e0b8afa2..d3e369bbe388 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -501,8 +501,8 @@ impl SharedRuntimeArgs { self, listen_fd: i32, grpc_listen_fd: Option, - input_address: String, - output_address: String, + input_listener_fd: i32, + output_listener_fd: i32, coordinator_address: Option, engine_start_index: u32, engine_count: usize, @@ -518,8 +518,8 @@ impl SharedRuntimeArgs { Config { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, data_parallel_size, @@ -697,14 +697,12 @@ pub struct FrontendArgs { /// supervisor. When not set, no gRPC server is started. #[arg(long)] pub grpc_listen_fd: Option, - /// Frontend input ROUTER socket address that the Python engines will - /// connect to. + /// Inherited frontend input ROUTER listener file descriptor. #[arg(long)] - pub input_address: String, - /// Frontend output PULL socket address that the Python engines will push - /// responses to. + pub input_listener_fd: i32, + /// Inherited frontend output PULL listener file descriptor. #[arg(long)] - pub output_address: String, + pub output_listener_fd: i32, /// Optional Python-owned frontend-side DP coordinator socket address for /// external coordinator mode in the bootstrapped frontend path, i.e., /// `stats_update_address`. @@ -733,8 +731,8 @@ impl FrontendArgs { self.runtime.into_bootstrapped_config( self.listen_fd, self.grpc_listen_fd, - self.input_address, - self.output_address, + self.input_listener_fd, + self.output_listener_fd, self.coordinator_address, self.engine_start_index, self.engine_count, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 944febc36c10..aa164bc618f8 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -668,10 +668,10 @@ fn frontend_args_json_passes_tls_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_keyfile":"/tmp/key.pem"}"#, ]) @@ -693,10 +693,10 @@ fn frontend_args_json_rejects_out_of_range_cert_reqs() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_cert_reqs":5}"#, ]) @@ -718,10 +718,10 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#, ]) @@ -807,10 +807,10 @@ fn frontend_args_json_passes_lora_modules_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","lora_modules":[{"name":"alice","path":"org/alice","base_model_name":null,"is_3d_lora_weight":true}]}"#, ]) @@ -862,10 +862,10 @@ fn frontend_args_json_accepts_api_key_string() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":"secret"}"#, ]) @@ -885,10 +885,10 @@ fn frontend_args_json_accepts_api_key_list() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":["secret-a","secret-b"]}"#, ]) @@ -991,10 +991,10 @@ fn frontend_args_accept_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--args-json", @@ -1008,8 +1008,8 @@ fn frontend_args_accept_json() { FrontendArgs { listen_fd: 3, grpc_listen_fd: None, - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, coordinator_address: Some( "tcp://127.0.0.1:7000", ), @@ -1083,10 +1083,10 @@ fn frontend_args_json_applies_defaults() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--engine-count", "4", "--args-json", @@ -1115,10 +1115,10 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","max_model_len":-1}"#, ]) @@ -1165,10 +1165,10 @@ fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", &payload.to_string(), ]) @@ -1204,10 +1204,10 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) @@ -1286,10 +1286,10 @@ fn frontend_args_json_ignores_unknown_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","uds":"/tmp/vllm.sock","nested_unknown":{"x":1}}"#, ]) @@ -1308,10 +1308,10 @@ fn frontend_args_json_sets_prompt_tokens_details_flag() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2,"enable_prompt_tokens_details":true}"#, ]) @@ -1371,10 +1371,10 @@ fn frontend_args_json_parses_cors_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","allowed_origins":["http://a.com"],"allow_credentials":true}"#, ]) @@ -1396,10 +1396,10 @@ fn frontend_args_json_rejects_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/prefix"}"#, ]) @@ -1423,10 +1423,10 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","root_path":"/prefix"}"#, ]) @@ -1452,10 +1452,10 @@ fn frontend_args_json_rejects_malformed_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B""#, ]) @@ -1999,10 +1999,10 @@ fn frontend_args_reject_legacy_handshake_flags() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, "--handshake-address", @@ -2043,10 +2043,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--engine-start-index", @@ -2068,8 +2068,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present expect![[r#" Config { transport_mode: Bootstrapped { - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, engine_start_index: 3, engine_count: 1, data_parallel_size: 4, @@ -2200,10 +2200,10 @@ fn frontend_config_adopts_inherited_grpc_listener() { "3", "--grpc-listen-fd", "4", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","grpc_port":50051}"#, ]) @@ -2227,10 +2227,10 @@ fn frontend_args_json_enables_profiling_when_profiler_config_set() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#, ]) @@ -2251,10 +2251,10 @@ fn frontend_args_json_disables_profiling_when_profiler_config_absent() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, ]) @@ -2275,10 +2275,10 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#, ]) diff --git a/rust/src/engine-core-client/Cargo.toml b/rust/src/engine-core-client/Cargo.toml index 14b9fe142349..86ffe4c4cdae 100644 --- a/rust/src/engine-core-client/Cargo.toml +++ b/rust/src/engine-core-client/Cargo.toml @@ -18,6 +18,7 @@ futures.workspace = true half.workspace = true hex.workspace = true itertools.workspace = true +libc.workspace = true parking_lot.workspace = true rmp-serde.workspace = true rmpv.workspace = true @@ -27,6 +28,7 @@ serde_json.workspace = true serde_repr.workspace = true serde_tuple.workspace = true serde_with.workspace = true +socket2.workspace = true task-local.workspace = true tempfile = { workspace = true, optional = true } thiserror.workspace = true diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 36a4fbffa5e1..b3516936fef2 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -53,16 +53,15 @@ pub enum TransportMode { local_output_address: Option, }, - /// The Python supervisor has already chosen the frontend transport - /// addresses, and the Rust process only needs to bind them and wait for - /// engine registration frames. + /// The Python supervisor has already bound the frontend transport + /// listeners. Rust adopts them and waits for engine registration frames. + /// `EngineCoreClient::connect` consumes each descriptor exactly once; + /// copied integer values do not create additional descriptor ownership. Bootstrapped { - /// Input ROUTER socket address that engines will connect to for - /// requests. - input_address: String, - /// Output PULL socket address that engines will connect to for - /// responses. - output_address: String, + /// Raw input ROUTER listener descriptor inherited from the supervisor. + input_listener_fd: i32, + /// Raw output PULL listener descriptor inherited from the supervisor. + output_listener_fd: i32, /// First data-parallel engine rank expected to register on this /// transport. engine_start_index: u32, @@ -336,8 +335,8 @@ impl EngineCoreClient { } TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, ready_timeout, @@ -348,8 +347,8 @@ impl EngineCoreClient { } transport::connect_bootstrapped( - input_address, - output_address, + *input_listener_fd, + *output_listener_fd, *engine_start_index, *engine_count, *ready_timeout, diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 0ac81764a9ed..dbd3fa69ceb2 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::convert::TryFrom; use std::io::Cursor; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::path::PathBuf; use std::process::Command; use std::sync::Once; @@ -311,10 +313,17 @@ fn bootstrapped_test_config( client_index: u32, coordinator_mode: Option, ) -> EngineCoreClientConfig { + fn listener(address: &str) -> i32 { + let path = address.strip_prefix("ipc://").expect("bootstrapped tests use IPC listeners"); + let _ = std::fs::remove_file(path); + let listener = UnixListener::bind(path).expect("bind inherited test listener"); + listener.into_raw_fd() + } + EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd: listener(&input_address), + output_listener_fd: listener(&output_address), engine_start_index: 0, engine_count, data_parallel_size: engine_count, @@ -362,15 +371,22 @@ fn bootstrapped_test_config_with_start_index( #[test] fn client_config_validates_bootstrapped_dp_range() { - let mut config = bootstrapped_test_config_with_start_index( - "ipc://unused-input".to_string(), - "ipc://unused-output".to_string(), - 1, - 2, - Duration::from_secs(1), - 0, - None, - ); + // Validation never consumes the descriptors, so placeholders avoid opening + // listeners that nothing would close. + let mut config = EngineCoreClientConfig { + transport_mode: TransportMode::Bootstrapped { + input_listener_fd: -1, + output_listener_fd: -1, + engine_start_index: 1, + engine_count: 2, + data_parallel_size: 3, + ready_timeout: Duration::from_secs(1), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + engine_stats_enabled: true, + }; config.validate().expect("frontend may own a subset of global DP ranks"); let TransportMode::Bootstrapped { diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index 7aac8774cac7..e453cafdc8fb 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; use std::ops::Deref; +use std::os::fd::{FromRawFd, OwnedFd, RawFd}; +use std::os::unix::net::UnixListener as StdUnixListener; use std::time::Duration; use bytes::Bytes; @@ -327,8 +329,8 @@ pub async fn connect_handshake( /// identities are synthesized from contiguous rank order instead of being /// discovered through a Rust-owned handshake. pub async fn connect_bootstrapped( - input_address: &str, - output_address: &str, + input_listener_fd: RawFd, + output_listener_fd: RawFd, engine_start_index: u32, engine_count: usize, ready_timeout: Duration, @@ -354,10 +356,16 @@ pub async fn connect_bootstrapped( } let mut input_socket = RouterSocket::new(); - let input_address = input_socket.bind(input_address).await?.to_string(); + let input_address = input_socket + .bind_listener(adopt_listener(input_listener_fd)?) + .await? + .to_string(); let mut output_socket = PullSocket::new(); - let output_address = output_socket.bind(output_address).await?.to_string(); + let output_address = output_socket + .bind_listener(adopt_listener(output_listener_fd)?) + .await? + .to_string(); let engines = wait_for_input_registrations( &mut input_socket, @@ -385,6 +393,34 @@ pub async fn connect_bootstrapped( }) } +fn adopt_listener(fd: RawFd) -> Result { + if fd < 0 || unsafe { libc::fcntl(fd, libc::F_GETFD) } < 0 { + return Err(Error::InvalidClientConfig { + message: format!("inherited ZMQ listener fd {fd} is not open"), + }); + } + + // SAFETY: Python's Popen(pass_fds) gives the frontend its own descriptor. + // The bootstrapped startup path calls this function once per descriptor, + // transferring ownership to the listener constructed below. + let fd = unsafe { OwnedFd::from_raw_fd(fd) }; + let socket: socket2::Socket = fd.into(); + let domain = socket.domain()?; + + if domain == socket2::Domain::IPV4 || domain == socket2::Domain::IPV6 { + let listener: std::net::TcpListener = socket.into(); + return Ok(listener.into()); + } + if domain == socket2::Domain::UNIX { + let listener: StdUnixListener = socket.into(); + return Ok(listener.into()); + } + + Err(Error::InvalidClientConfig { + message: format!("unsupported inherited listener domain {domain:?}"), + }) +} + /// Bind new input and output sockets. async fn bind_local_sockets( local_host: &str, @@ -603,7 +639,12 @@ pub async fn run_output_loop( #[cfg(test)] mod tests { - use super::bind_local_sockets; + use std::os::fd::IntoRawFd; + + use zeromq::RouterSocket; + use zeromq::prelude::Socket; + + use super::{adopt_listener, bind_local_sockets}; #[tokio::test] async fn bind_local_sockets_resolves_zero_port_bindings() { @@ -614,4 +655,21 @@ mod tests { assert!(output_address.starts_with("tcp://127.0.0.1:")); assert_ne!(input_address, output_address); } + + #[tokio::test] + async fn inherited_tcp_listener_is_adopted_without_rebinding() { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).expect("bind listener"); + let expected = listener.local_addr().expect("listener address"); + let mut router = RouterSocket::new(); + + let endpoint = router + .bind_listener( + adopt_listener(listener.into_raw_fd()).expect("adopt inherited listener fd"), + ) + .await + .expect("bind inherited listener"); + + assert_eq!(endpoint.to_string(), format!("tcp://{expected}")); + assert!(router.close().await.is_empty()); + } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d8d0fb4e6cdd..631eb76b36ab 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -5,6 +5,8 @@ use std::collections::BTreeMap; use std::fs; use std::future::Future; use std::io; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -2230,10 +2232,14 @@ async fn control_aggregates_multi_engine_capacity() { ready_1.world_size = 12; ready_1.data_parallel_rank = start_rank + 1; + let listener_fd = |address: &str| { + let path = address.strip_prefix("ipc://").expect("IPC test endpoint"); + UnixListener::bind(path).expect("bind inherited test listener").into_raw_fd() + }; let client_config = EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address: input_address.clone(), - output_address: output_address.clone(), + input_listener_fd: listener_fd(&input_address), + output_listener_fd: listener_fd(&output_address), engine_start_index: start_rank, engine_count: 2, data_parallel_size: global_size, diff --git a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py b/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py deleted file mode 100644 index c8ff969a8a36..000000000000 --- a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Spawn worker targets kept free of heavy imports. - -``multiprocessing`` with the ``spawn`` start method re-imports the module that -defines a process target in the child. Housing these stubs in a stdlib-only -module keeps child startup fast and deterministic, instead of paying a multi- -second ``import vllm`` before the child can run. -""" - - -def exit_before_report_worker(listen_address, sock, args, client_config=None): - """Exit immediately without touching ``actual_address_pipe``.""" - return diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 6c579df8383c..318f57b86ff7 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing +import os import socket import threading import time @@ -11,14 +12,10 @@ import pytest import zmq -from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( - exit_before_report_worker, -) -from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import make_zmq_listener, make_zmq_socket from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, - get_engine_client_zmq_addr, wait_for_completion_or_failure, ) @@ -46,36 +43,30 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): args.results.put(stats.get_num_requests()) -# Module-level stub for the gather_actual_addresses test. Must be -# importable by `multiprocessing.spawn` (no closures, no nesting). -def defer_addresses_stub_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit.""" +def _adopt_zmq_listener_worker(listen_address, sock, args, client_config): ctx = zmq.Context() + router = make_zmq_socket( + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], + ) + pull = make_zmq_socket( + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], + ) try: - in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True - ) - out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True - ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + identity, payload = router.recv_multipart() + assert payload == b"input" + assert pull.recv() == b"output" + router.send_multipart([identity, b"ok"]) finally: + router.close(linger=0) + pull.close(linger=0) ctx.term() @@ -89,15 +80,11 @@ def api_server_args(): "sock": sock, "args": "test_args", # Simple string to avoid pickling issues "num_servers": 3, - "input_addresses": [ - "tcp://127.0.0.1:5001", - "tcp://127.0.0.1:5002", - "tcp://127.0.0.1:5003", + "input_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) for _ in range(3) ], - "output_addresses": [ - "tcp://127.0.0.1:6001", - "tcp://127.0.0.1:6002", - "tcp://127.0.0.1:6003", + "output_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) for _ in range(3) ], "stats_update_address": "tcp://127.0.0.1:7000", } @@ -120,6 +107,13 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): try: # Verify the manager was initialized correctly assert len(manager.processes) == 3 + assert all( + listener.socket.fileno() == -1 + for listener in ( + *args["input_listeners"], + *args["output_listeners"], + ) + ) # Verify all processes are running for proc in manager.processes: @@ -145,6 +139,70 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +def _assert_endpoint_reserved(address: str) -> None: + """A competing bind must fail while the endpoint has a listener owner.""" + if address.startswith("ipc://"): + competitor = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + bind_address: str | tuple[str, int] = address.removeprefix("ipc://") + else: + competitor = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + bind_address = ("127.0.0.1", int(address.rsplit(":", 1)[1])) + with competitor, pytest.raises(OSError, match="Address already in use"): + competitor.bind(bind_address) + + +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_api_server_child_adopts_inherited_zmq_listeners(tmp_path, scheme): + if scheme == "tcp": + input_path = output_path = "tcp://127.0.0.1:0" + else: + input_path = f"ipc://{tmp_path}/input.sock" + output_path = f"ipc://{tmp_path}/output.sock" + input_listener = make_zmq_listener(input_path, zmq.ROUTER) + output_listener = make_zmq_listener(output_path, zmq.PULL) + input_address = input_listener.address + output_address = output_listener.address + # The supervisor reserves each endpoint as soon as it allocates it. + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + http_socket = socket.socket() + manager = APIServerProcessManager( + target_server_fn=_adopt_zmq_listener_worker, + listen_address="localhost:8000", + sock=http_socket, + args="test_args", + num_servers=1, + input_listeners=[input_listener], + output_listeners=[output_listener], + ) + + ctx = zmq.Context() + dealer = make_zmq_socket( + ctx, input_address, zmq.DEALER, bind=False, identity=b"test" + ) + push = make_zmq_socket(ctx, output_address, zmq.PUSH, bind=False) + dealer.setsockopt(zmq.RCVTIMEO, 30_000) + try: + assert input_listener.socket.fileno() == -1 + assert output_listener.socket.fileno() == -1 + # After the parent closes its copies, the child's inherited descriptors + # keep both endpoints reserved; closing must not unlink IPC paths. + assert manager.processes[0].is_alive() + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + dealer.send(b"input") + push.send(b"output") + assert dealer.recv() == b"ok" + manager.processes[0].join(timeout=10) + assert manager.processes[0].exitcode == 0 + finally: + manager.shutdown() + dealer.close(linger=0) + push.close(linger=0) + ctx.term() + http_socket.close() + + @pytest.mark.skip_global_cleanup def test_api_server_processes_share_admission_stats(api_server_args): args = api_server_args.copy() @@ -354,96 +412,7 @@ def run_with_exception_capture(): WORKER_RUNTIME_SECONDS = prev_worker_runtime -@pytest.mark.timeout(60) -def test_gather_actual_addresses_end_to_end(): - """Each child binds ROUTER/PULL with a kernel-picked port and reports - the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`.""" - host = "127.0.0.1" - num_servers = 4 - - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - for addr in placeholder_inputs + placeholder_outputs: - assert addr == f"tcp://{host}:0", addr - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - target_server_fn=defer_addresses_stub_worker, - ) - - try: - assert len(manager.processes) == num_servers - actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - assert len(actual_inputs) == num_servers - assert len(actual_outputs) == num_servers - - for addr in actual_inputs + actual_outputs: - scheme, parsed_host, port = split_zmq_path(addr) - assert scheme == "tcp", addr - assert parsed_host == host, addr - assert port and int(port) > 0, addr - - all_addrs = actual_inputs + actual_outputs - assert len(set(all_addrs)) == len(all_addrs), all_addrs - - -@pytest.mark.timeout(30) -def test_gather_actual_addresses_child_crash_before_report(): - """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots.""" - host = "127.0.0.1" - num_servers = 2 - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - # exit_before_report_worker exits without touching - # ``actual_address_pipe`` — simulates a child that dies before - # reporting its bound addresses. - target_server_fn=exit_before_report_worker, - ) - try: - # Sentinel-first vs pipe-EOF-first both produce "reporting". - with pytest.raises(RuntimeError, match="reporting"): - manager.gather_actual_addresses(timeout=10.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - -def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): +def test_rust_frontend_launch_log_redacts_credentials(monkeypatch): """The Rust frontend command carries every non-default arg as JSON. Credentials must not reach the log line.""" import subprocess as subprocess_mod @@ -469,15 +438,17 @@ def wait(self, timeout=None): return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) - with caplog.at_level("INFO", logger="vllm.v1.utils"): + with patch("vllm.v1.utils.logger.info") as log_info: RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -485,11 +456,13 @@ def wait(self, timeout=None): finally: sock.close() - message = caplog.text + log_format, *log_args = log_info.call_args.args + message = log_format % tuple(log_args) assert "Launching Rust frontend:" in message assert hf_token not in message assert api_key not in message assert '"hf_token": "***"' in message + assert '"api_key": "***"' in message def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_path): @@ -519,6 +492,8 @@ def wait(self, timeout=None): return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) with caplog.at_level("INFO", logger="vllm.v1.utils"): @@ -526,8 +501,8 @@ def wait(self, timeout=None): binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -568,14 +543,18 @@ def fake_popen(cmd, **kwargs): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) grpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() try: monkeypatch.setattr(subprocess_mod, "Popen", fake_popen) RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -585,7 +564,44 @@ def fake_popen(cmd, **kwargs): grpc_fd = grpc_sock.fileno() flag_index = cmd.index("--grpc-listen-fd") assert cmd[flag_index + 1] == str(grpc_fd) - assert list(kwargs["pass_fds"]) == [sock.fileno(), grpc_fd] + assert list(kwargs["pass_fds"]) == [ + sock.fileno(), + input_fd, + output_fd, + grpc_fd, + ] finally: sock.close() grpc_sock.close() + + +def test_rust_frontend_spawn_failure_cleans_up_listeners(tmp_path): + """No child owns the listeners when Popen fails, so the parent removes them.""" + from vllm.entrypoints.launchers.cli_args import make_arg_parser + from vllm.utils.argparse_utils import FlexibleArgumentParser + from vllm.v1.utils import RustFrontendProcessManager + + args = make_arg_parser(FlexibleArgumentParser()).parse_args( + ["--model", "org/model"] + ) + input_listener = make_zmq_listener(f"ipc://{tmp_path}/input.sock", zmq.ROUTER) + output_listener = make_zmq_listener(f"ipc://{tmp_path}/output.sock", zmq.PULL) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + with pytest.raises(FileNotFoundError): + RustFrontendProcessManager( + binary_path=str(tmp_path / "missing-vllm-rs"), + sock=sock, + args=args, + input_listener=input_listener, + output_listener=output_listener, + engine_start_index=0, + engine_count=1, + data_parallel_size=1, + ) + finally: + sock.close() + + for listener in (input_listener, output_listener): + assert listener.socket.fileno() == -1 + assert not os.path.exists(listener.address.removeprefix("ipc://")) diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index a6c359e26930..7ba3d6ae3ecc 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -13,6 +13,7 @@ get_open_ports_list, get_tcp_uri, join_host_port, + make_zmq_listener, make_zmq_path, make_zmq_socket, split_host_port, @@ -205,6 +206,52 @@ def test_make_zmq_socket_ipv6(): ctx.term() +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_make_zmq_socket_adopts_listener(tmp_path, scheme): + path = ( + "tcp://127.0.0.1:0" if scheme == "tcp" else f"ipc://{tmp_path}/inherited.sock" + ) + listener = make_zmq_listener(path, zmq.ROUTER) + ctx = zmq.Context() + router = dealer = None + try: + router = make_zmq_socket( + ctx, + listener.address, + zmq.ROUTER, + bind=True, + listener=listener.socket, + ) + assert listener.socket.fileno() == -1 + dealer = make_zmq_socket( + ctx, listener.address, zmq.DEALER, bind=False, identity=b"client" + ) + dealer.send(b"ping") + assert router.recv_multipart() == [b"client", b"ping"] + finally: + if router is not None: + router.close(linger=0) + if dealer is not None: + dealer.close(linger=0) + ctx.term() + listener.cleanup() + + +def test_make_zmq_listener_failed_ipc_bind_keeps_existing_path(tmp_path): + path = tmp_path / "taken.sock" + owner = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + owner.bind(str(path)) + owner.listen() + try: + with pytest.raises(OSError, match="Address already in use"): + make_zmq_listener(f"ipc://{path}", zmq.ROUTER) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(path)) + finally: + owner.close() + + def test_make_zmq_path(): assert make_zmq_path("tcp", "127.0.0.1", 5555) == "tcp://127.0.0.1:5555" assert make_zmq_path("tcp", "::1", 5555) == "tcp://[::1]:5555" diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 768f51f4379b..1b23356134b0 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -23,7 +23,7 @@ EngineZmqAddresses, _dp_nodes_master_first, _node_ip_from_resources, - get_engine_zmq_addresses, + bind_engine_zmq_listeners, launch_core_engines, ) from vllm.v1.utils import APIServerProcessManager @@ -67,35 +67,27 @@ def get_addresses(self) -> tuple[list[str], list[str]]: # Module-level stub worker for the Ray-DP regression test. Must be importable -# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker -# in ``tests/entrypoints/test_api_server_process_manager.py``. -def _bind_and_report_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via ``actual_address_pipe``, then exit.""" +# by ``multiprocessing.spawn`` (no closures, no nesting). +def _adopt_listener_worker(listen_address, sock, args, client_config): + """Adopt the supervisor-bound ROUTER/PULL listeners, then exit.""" ctx = zmq.Context() try: in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], ) out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + in_sock.close(linger=0) + out_sock.close(linger=0) finally: ctx.term() @@ -327,17 +319,9 @@ def test_ray_dp_addresses_resolved_before_actor_creation( ``launch_core_engines`` Ray branch pickles ``addresses`` into each engine actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes`` is a no-op, so the actor uses that pickled snapshot for the rest of its - life. If ``run_multi_api_server`` allocates ``addresses`` as - ``tcp://host:0`` placeholders (its default), the actors hold placeholders - forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does - not raise, so the failure mode is a deterministic hang. - - The Ray-DP carve-out in ``run_multi_api_server`` forces - ``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"`` - so addresses are pre-allocated in the driver and Ray pickles real ports - into each actor. This test mirrors that call-site logic and asserts the - actors hold real (non-placeholder) endpoints. If the carve-out is - removed without an alternative fix, the test fails. + life. The supervisor therefore binds the frontend listeners before actor + construction and Ray pickles their resolved endpoints. This test asserts + that each actor holds real (non-placeholder) endpoints. """ created_placement_groups: list[Any] = [] @@ -356,15 +340,10 @@ def create_dp_placement_groups(vllm_config: Any): vllm_config = _make_vllm_config_ray_dp_multinode() - # Mirror run_multi_api_server's address-allocation logic. The Ray DP - # carve-out forces pre-allocation so the addresses pickled into engine - # actors at .remote() time are real, not ``tcp://host:0``. - is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers=2, - defer_api_server_ports=not is_ray_dp, - ) + # Mirror run_multi_api_server: bind listeners in the supervisor so the + # addresses pickled into Ray actors contain kernel-assigned ports. + listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers=2) + addresses = listeners.addresses sock = socket.socket() engine_manager: CoreEngineActorManager | None = None @@ -381,26 +360,17 @@ def create_dp_placement_groups(vllm_config: Any): engine_manager = engine_launch.engine_manager assert isinstance(engine_manager, CoreEngineActorManager) - # API-server children bind to the pre-allocated ports. + # API-server children adopt the already-bound listener FDs. api_server_manager = APIServerProcessManager( listen_address="tcp://127.0.0.1:0", sock=sock, args="test_args", num_servers=2, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, - target_server_fn=_bind_and_report_worker, + input_listeners=listeners.inputs, + output_listeners=listeners.outputs, + target_server_fn=_adopt_listener_worker, ) - # run_multi_api_server skips ``gather_actual_addresses`` for - # Ray DP (addresses are already real). Mirror that. - if not is_ray_dp: - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses(timeout=15.0) - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Snapshot what each Ray actor actually holds. actors = ( engine_manager.local_engine_actors + engine_manager.remote_engine_actors @@ -427,9 +397,9 @@ def create_dp_placement_groups(vllm_config: Any): assert scheme == "tcp", url assert port and int(port) > 0, ( f"Ray actor was pickled with placeholder address {url!r}; " - "``run_multi_api_server`` must pre-allocate ports for the " - "Ray DP backend so the actors hold real endpoints by the " - "time they DEALER-connect. See PR #42585 / Ray-DP " + "``run_multi_api_server`` must bind frontend listeners before " + "constructing Ray DP actors so they hold real endpoints when " + "they DEALER-connect. See PR #42585 / Ray-DP " "multi-API-server regression." ) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 4c9410488b3f..94e6545c0899 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -328,19 +328,10 @@ def signal_handler(signum, frame): None ) - from vllm.v1.engine.utils import get_engine_zmq_addresses - - # Defer port allocation to the child's bind() to avoid TOCTOU, except - # for Rust front-end and Ray DP, which can't see the post-bind rebind - # (CLI-arg subprocess / pickled-into-actor snapshot respectively) and - # so pre-allocate driver-side -- reintroducing the original race only - # there. - is_ray_dp = parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers, - defer_api_server_ports=not (rust_frontend_path or is_ray_dp), - ) + from vllm.v1.engine.utils import bind_engine_zmq_listeners + + zmq_listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers) + addresses = zmq_listeners.addresses with launch_core_engines( vllm_config, executor_class, log_stats, addresses @@ -365,8 +356,8 @@ def signal_handler(signum, frame): sock=sock, grpc_sock=grpc_sock, args=args, - input_address=addresses.inputs[0], - output_address=addresses.outputs[0], + input_listener=zmq_listeners.inputs[0], + output_listener=zmq_listeners.outputs[0], engine_start_index=expected_engine_start_index, engine_count=expected_engine_count, data_parallel_size=parallel_config.data_parallel_size, @@ -379,22 +370,12 @@ def signal_handler(signum, frame): sock=sock, args=args, num_servers=num_api_servers, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, + input_listeners=zmq_listeners.inputs, + output_listeners=zmq_listeners.outputs, stats_update_address=stats_update_address, tensor_queue=engine_launch.tensor_queue, ) - if not is_ray_dp: - # Forward each child's bound endpoints to the engine handshake - # (runs on ``with`` exit). Skipped for Ray DP, where addresses - # are pre-allocated above and Ray actors already hold them. - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses() - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Set frontend processes to watch during engine startup. # If any of these processes exit before the engines are up, the engine startup # will be aborted with an error. diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py index fa24d4d5f9ab..6bc4f9c9fdf8 100644 --- a/vllm/utils/network_utils.py +++ b/vllm/utils/network_utils.py @@ -11,6 +11,7 @@ Iterator, Sequence, ) +from dataclasses import dataclass from typing import Any from uuid import uuid4 @@ -25,6 +26,23 @@ logger = init_logger(__name__) +@dataclass +class ZmqListener: + """A supervisor-bound listener that a ZMQ socket can adopt.""" + + address: str + socket: socket.socket + + def close(self) -> None: + self.socket.close() + + def cleanup(self) -> None: + self.close() + if self.address.startswith("ipc://"): + with contextlib.suppress(FileNotFoundError): + os.unlink(self.address.removeprefix("ipc://")) + + def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]): for sock in sockets: if sock is not None: @@ -307,6 +325,27 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501 +# Apply the existing make_zmq_socket memory policy to every socket owner. +_ZMQ_LARGE_BUFFER_SIZE = 512 * 1024**2 +_ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY = 32 * 1024**3 +_ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY = 16 * 1024**3 + + +def _get_zmq_socket_buffer_size() -> int: + """Choose the shared libzmq and inherited-listener buffer policy.""" + mem = psutil.virtual_memory() + # For systems with substantial memory (>32GB total, >16GB available): + # - Set a large 0.5GB buffer to improve throughput + # For systems with less memory: + # - Use system default (-1) to avoid excessive memory consumption + if ( + mem.total > _ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY + and mem.available > _ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY + ): + return _ZMQ_LARGE_BUFFER_SIZE + return -1 + + def make_zmq_socket( ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined] path: str, @@ -315,22 +354,19 @@ def make_zmq_socket( identity: bytes | None = None, linger: int | None = None, router_handover: bool = False, + listener: socket.socket | None = None, ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined] - """Make a ZMQ socket with the proper bind/connect semantics.""" - mem = psutil.virtual_memory() - socket = ctx.socket(socket_type) + """Make a ZMQ socket with the proper bind/connect semantics. - # Calculate buffer size based on system memory - total_mem = mem.total / 1024**3 - available_mem = mem.available / 1024**3 - # For systems with substantial memory (>32GB total, >16GB available): - # - Set a large 0.5GB buffer to improve throughput - # For systems with less memory: - # - Use system default (-1) to avoid excessive memory consumption - buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1 + When supplied, ``listener`` is detached and its fd moves to libzmq. + """ + socket = ctx.socket(socket_type) + buf_size = _get_zmq_socket_buffer_size() if bind is None: bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB) + if listener is not None and not bind: + raise ValueError("An inherited ZMQ listener requires bind=True") if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): socket.setsockopt(zmq.RCVHWM, 0) @@ -353,6 +389,9 @@ def make_zmq_socket( if socket_type == zmq.XPUB: socket.setsockopt(zmq.XPUB_VERBOSE, True) + if listener is not None: + socket.setsockopt(zmq.USE_FD, listener.detach()) + # Determine if the path is a TCP socket with an IPv6 address. # Enable IPv6 on the zmq socket if so. scheme, host, _ = split_zmq_path(path) @@ -367,6 +406,48 @@ def make_zmq_socket( return socket +def make_zmq_listener(path: str, socket_type: Any) -> ZmqListener: + """Bind and listen on a raw socket for later ZMQ adoption.""" + scheme, host, port = split_zmq_path(path) + if scheme == "tcp": + family = socket.AF_INET6 if is_valid_ipv6_address(host) else socket.AF_INET + bind_address: str | tuple[str, int] = (host, int(port)) + elif scheme == "ipc": + family = socket.AF_UNIX + bind_address = path.removeprefix("ipc://") + os.makedirs(os.path.dirname(bind_address), exist_ok=True) + else: + raise ValueError(f"Cannot inherit a {scheme} ZMQ listener") + + listener = socket.socket(family, socket.SOCK_STREAM) + bound = False + try: + buf_size = _get_zmq_socket_buffer_size() + if buf_size >= 0: + # Accepted stream sockets inherit these settings from the listener. + # ZMQ_USE_FD adopts this pre-created socket, so the policy belongs + # here as well as on the eventual libzmq socket. + if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, buf_size) + if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, buf_size) + + listener.bind(bind_address) + bound = True + listener.listen() + if scheme == "tcp": + path = get_tcp_uri(host, listener.getsockname()[1]) + return ZmqListener(address=path, socket=listener) + except BaseException: + listener.close() + # A failed bind leaves any existing pathname owned by its listener. + if scheme == "ipc" and bound: + assert isinstance(bind_address, str) + with contextlib.suppress(FileNotFoundError): + os.unlink(bind_address) + raise + + @contextlib.contextmanager def zmq_socket_ctx( path: str, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 88fdc1d1e7f1..74520036a7c8 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -11,7 +11,6 @@ from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass -from multiprocessing.connection import Connection from multiprocessing.queues import Queue from threading import Thread from typing import Any, TypeAlias, TypeVar @@ -56,7 +55,8 @@ from vllm.v1.engine.utils import ( CoreEngineActorManager, CoreEngineProcManager, - get_engine_zmq_addresses, + EngineZmqAddresses, + get_engine_zmq_bind_addresses, launch_core_engines, ) from vllm.v1.executor import Executor @@ -634,63 +634,50 @@ def __init__( self.stats_update_address = client_addresses.get("stats_update_address") # Tensor queues passed via client_addresses for multi-API-server case tensor_queue = client_addresses.get("tensor_queue") + input_listener = client_addresses.get("input_listener") + output_listener = client_addresses.get("output_listener") self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, input_address, zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, + listener=input_listener, ) self.resources.output_socket = make_zmq_socket( - self.ctx, output_address, zmq.PULL + self.ctx, + output_address, + zmq.PULL, + listener=output_listener, ) - # Report bound endpoints back so the parent can forward - # them to engines (mirrors the DPCoordinator pattern). - actual_address_pipe: Connection | None = client_addresses.get( - "actual_address_pipe" - ) - if actual_address_pipe is not None: - try: - actual_input = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_output = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_address_pipe.send( - { - "input_address": actual_input, - "output_address": actual_output, - } - ) - finally: - actual_address_pipe.close() # Engines are managed externally: this process does not fork # them, so there is no fork race; start the MM warmup now. self._start_mm_warmup() else: # Engines are managed by this client. - addresses = get_engine_zmq_addresses(vllm_config) + bind_addresses = get_engine_zmq_bind_addresses(vllm_config) self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, - addresses.inputs[0], + bind_addresses.inputs[0], zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, ) self.resources.output_socket = make_zmq_socket( - self.ctx, addresses.outputs[0], zmq.PULL + self.ctx, bind_addresses.outputs[0], zmq.PULL ) - # Resolve ``tcp://host:0`` placeholders to bound endpoints - # before engines DEALER-connect. No-op for IPC. - addresses.inputs[0] = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - addresses.outputs[0] = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() + # EngineZmqAddresses contains only endpoints resolved by the + # sockets that own the bind operation. + addresses = EngineZmqAddresses( + inputs=[self.input_socket.getsockopt(zmq.LAST_ENDPOINT).decode()], + outputs=[ + self.resources.output_socket.getsockopt( + zmq.LAST_ENDPOINT + ).decode() + ], + ) with launch_core_engines( vllm_config, executor_class, log_stats, addresses diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 49e6884ac9f1..4f1be3bacd1c 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -23,9 +23,10 @@ from vllm.ray.ray_env import get_env_vars_to_copy from vllm.utils import numa_utils from vllm.utils.network_utils import ( - get_open_port, + ZmqListener, get_open_zmq_ipc_path, get_tcp_uri, + make_zmq_listener, zmq_socket_ctx, ) from vllm.utils.system_utils import get_mp_context, set_env_var @@ -82,8 +83,18 @@ def __init__(self, index: int = 0, local: bool = True): self.state = CoreEngineState.NEW +@dataclass +class EngineZmqBindAddresses: + """Frontend endpoints before their owning sockets bind.""" + + inputs: list[str] + outputs: list[str] + + @dataclass class EngineZmqAddresses: + """Resolved frontend endpoints advertised to engine processes.""" + # ZMQ input socket addresses for each front-end client (requests) inputs: list[str] # ZMQ output socket addresses for each front-end client (responses) @@ -98,6 +109,27 @@ class EngineZmqAddresses: frontend_stats_publish_address: str | None = None +@dataclass +class EngineZmqListeners: + inputs: list[ZmqListener] + outputs: list[ZmqListener] + + @property + def addresses(self) -> EngineZmqAddresses: + return EngineZmqAddresses( + inputs=[listener.address for listener in self.inputs], + outputs=[listener.address for listener in self.outputs], + ) + + def cleanup(self) -> None: + for listener in (*self.inputs, *self.outputs): + listener.cleanup() + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.cleanup() + + @dataclass class EngineHandshakeMetadata: """Metadata sent to each engine process during startup handshake, @@ -1079,22 +1111,16 @@ def shutdown(self, timeout: float | None = None) -> None: ray.util.remove_placement_group(pg) -def get_engine_zmq_addresses( +def get_engine_zmq_bind_addresses( vllm_config: VllmConfig, num_api_servers: int = 1, - *, - defer_api_server_ports: bool = True, -) -> EngineZmqAddresses: - """Allocate ZMQ addresses for engine-client communication. - - By default each TCP address is a ``tcp://host:0`` placeholder; the - consumer (API-server child or single-process ``MPClient``) binds, then - recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)`` - and writes it back into ``addresses`` before the engine handshake. - - Set ``defer_api_server_ports=False`` only when the consumer cannot - report a bound port back (e.g. the Rust front-end). IPC paths are - unaffected.""" +) -> EngineZmqBindAddresses: + """Create ZMQ endpoint templates for engine-client communication. + + Each TCP address is a ``tcp://host:0`` placeholder. A single-process + ``MPClient`` binds its own ZMQ sockets. Cross-process frontends use + ``bind_engine_zmq_listeners`` so the supervisor binds raw listeners before + launching engines and transfers listener ownership to each frontend.""" parallel_config = vllm_config.parallel_config local_engine_count = parallel_config.data_parallel_size_local local_start_index = parallel_config.data_parallel_rank_local @@ -1119,14 +1145,39 @@ def get_engine_zmq_addresses( def _addr() -> str: if client_local_only: return get_open_zmq_ipc_path() - return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port()) + return get_tcp_uri(host, 0) - return EngineZmqAddresses( + return EngineZmqBindAddresses( inputs=[_addr() for _ in range(num_api_servers)], outputs=[_addr() for _ in range(num_api_servers)], ) +def bind_engine_zmq_listeners( + vllm_config: VllmConfig, + num_api_servers: int = 1, +) -> EngineZmqListeners: + """Bind frontend transport listeners before launching engines.""" + bind_addresses = get_engine_zmq_bind_addresses( + vllm_config, + num_api_servers, + ) + inputs: list[ZmqListener] = [] + outputs: list[ZmqListener] = [] + try: + inputs.extend( + make_zmq_listener(address, zmq.ROUTER) for address in bind_addresses.inputs + ) + outputs.extend( + make_zmq_listener(address, zmq.PULL) for address in bind_addresses.outputs + ) + return EngineZmqListeners(inputs=inputs, outputs=outputs) + except BaseException: + for listener in (*inputs, *outputs): + listener.cleanup() + raise + + FrontendProcess = BaseProcess | _SubprocessWrapper diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index ac877263a5f6..868864f0ca41 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -29,7 +29,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message -from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri +from vllm.utils.network_utils import ZmqListener, get_open_zmq_ipc_path, get_tcp_uri from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.core.sched.output import SchedulerOutput @@ -177,28 +177,22 @@ def __init__( sock: Any, args: argparse.Namespace, num_servers: int, - input_addresses: list[str], - output_addresses: list[str], + input_listeners: list[ZmqListener], + output_listeners: list[ZmqListener], target_server_fn: Callable | None = None, stats_update_address: str | None = None, tensor_queue: Queue | None = None, ): """Initialize and start API server worker processes. - ``input_addresses``/``output_addresses`` may contain - ``tcp://host:0`` placeholders; each child must report the actual - bound endpoint over its ``actual_address_pipe`` in ``client_config`` - and the parent collects them via - :py:meth:`gather_actual_addresses`. - Args: target_server_fn: Override function to call for each API server process listen_address: Address to listen for client connections sock: Socket for client connections args: Command line arguments num_servers: Number of API server processes to start - input_addresses: Input addresses for each API server - output_addresses: Output addresses for each API server + input_listeners: Input listeners for each API server + output_listeners: Output listeners for each API server stats_update_address: Optional stats update address tensor_queue: Optional tensor IPC queue for sharing MM tensors @@ -209,7 +203,6 @@ def __init__( spawn_context = multiprocessing.get_context("spawn") self.processes: list[BaseProcess] = [] - self._address_pipes: list[connection.Connection] = [] admission_counters = None if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: @@ -220,126 +213,65 @@ def __init__( SharedAdmissionStats.num_counters(num_servers), ) - for i, in_addr, out_addr in zip( - range(num_servers), input_addresses, output_addresses - ): - client_config: dict[str, Any] = { - "input_address": in_addr, - "output_address": out_addr, - "client_count": num_servers, - "client_index": i, - } - if admission_counters is not None: - client_config["mp_admission_counters"] = admission_counters - if stats_update_address is not None: - client_config["stats_update_address"] = stats_update_address - if tensor_queue is not None: - client_config["tensor_queue"] = tensor_queue - - parent_recv, child_send = spawn_context.Pipe(duplex=False) - self._address_pipes.append(parent_recv) - client_config["actual_address_pipe"] = child_send - - proc = spawn_context.Process( - target=target_server_fn or run_api_server_worker_proc, - name=f"ApiServer_{i}", - args=(listen_address, sock, args, client_config), + listeners = [*input_listeners, *output_listeners] + self._listeners = listeners + try: + listener_pairs = zip( + range(num_servers), input_listeners, output_listeners, strict=True ) - self.processes.append(proc) - proc.start() - - # Drop parent's write end so reader sees EOF on child death. - child_send.close() + for i, input_listener, output_listener in listener_pairs: + client_config: dict[str, Any] = { + "input_address": input_listener.address, + "output_address": output_listener.address, + "input_listener": input_listener.socket, + "output_listener": output_listener.socket, + "client_count": num_servers, + "client_index": i, + } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters + if stats_update_address is not None: + client_config["stats_update_address"] = stats_update_address + if tensor_queue is not None: + client_config["tensor_queue"] = tensor_queue + + proc = spawn_context.Process( + target=target_server_fn or run_api_server_worker_proc, + name=f"ApiServer_{i}", + args=(listen_address, sock, args, client_config), + ) + self.processes.append(proc) + proc.start() + input_listener.close() + output_listener.close() + except BaseException: + _shutdown_processes_and_cleanup_listeners(self.processes, listeners) + raise logger.info("Started %d API server processes", len(self.processes)) # Shutdown only the API server processes on garbage collection # The extra processes are managed by their owners - self._finalizer = weakref.finalize(self, shutdown, self.processes) - - def gather_actual_addresses( - self, - timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S, - ) -> tuple[list[str], list[str]]: - """Return (inputs, outputs) reported by each child, indexed by - ``client_index``. Raises ``RuntimeError`` on timeout or premature - child exit.""" - n = len(self._address_pipes) - inputs: list[str | None] = [None] * n - outputs: list[str | None] = [None] * n - pending: dict[connection.Connection, int] = { - pipe: i for i, pipe in enumerate(self._address_pipes) - } - sentinel_to_idx: dict[Any, int] = { - proc.sentinel: i for i, proc in enumerate(self.processes) - } - - deadline = time.monotonic() + timeout - try: - while pending: - remaining = deadline - time.monotonic() - if remaining <= 0: - missing = [self.processes[i].name for i in pending.values()] - raise RuntimeError( - f"Timed out after {timeout:.1f}s waiting for " - f"API server(s) to report bound ZMQ addresses: " - f"{missing}" - ) - waitables: list[Any] = list(pending.keys()) + list( - sentinel_to_idx.keys() - ) - ready = connection.wait(waitables, timeout=remaining) - # Drain pipes before checking sentinels: a child that sent - # its message and then exited can surface both events in - # the same poll, and we must record the success first. - for item in ready: - if isinstance(item, connection.Connection) and item in pending: - idx = pending.pop(item) - try: - msg: dict[str, str] = item.recv() - except EOFError as e: - raise RuntimeError( - f"API server {self.processes[idx].name} " - f"closed its address pipe without " - f"reporting its bound ZMQ addresses" - ) from e - inputs[idx] = msg["input_address"] - outputs[idx] = msg["output_address"] - item.close() - for item in ready: - if item in sentinel_to_idx: - idx = sentinel_to_idx.pop(item) - pipe = self._address_pipes[idx] - if pipe in pending: - proc = self.processes[idx] - raise RuntimeError( - f"API server process {proc.name} exited " - f"(code={proc.exitcode}) before reporting " - f"its bound ZMQ addresses" - ) - finally: - for pipe in pending: - with contextlib.suppress(Exception): - pipe.close() - - return inputs, outputs # type: ignore[return-value] + self._finalizer = weakref.finalize( + self, + _shutdown_processes_and_cleanup_listeners, + self.processes, + listeners, + ) def shutdown(self, timeout: float | None = None) -> None: """Shutdown API server processes with configurable timeout""" - for pipe in self._address_pipes: - with contextlib.suppress(Exception): - pipe.close() - self._address_pipes = [] - if self._finalizer.detach() is not None: - shutdown(self.processes, timeout=timeout) + _shutdown_processes_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class RustFrontendProcessManager: """Manages a single Rust frontend subprocess. Launches the Rust vllm-rs binary in 'frontend' mode, passing the - listening socket fd and ZMQ transport addresses. Provides the same + HTTP and ZMQ listener file descriptors. Provides the same interface as APIServerProcessManager for process monitoring. """ @@ -348,8 +280,8 @@ def __init__( binary_path: str, sock: Any, args: argparse.Namespace, - input_address: str, - output_address: str, + input_listener: ZmqListener, + output_listener: ZmqListener, engine_start_index: int, engine_count: int, data_parallel_size: int, @@ -359,18 +291,22 @@ def __init__( import os import subprocess - fd = sock.fileno() - os.set_inheritable(fd, True) + listen_fd = sock.fileno() + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() + inherited_fds = [listen_fd, input_fd, output_fd] + for fd in inherited_fds: + os.set_inheritable(fd, True) cmd = [ binary_path, "frontend", "--listen-fd", - str(fd), - "--input-address", - input_address, - "--output-address", - output_address, + str(listen_fd), + "--input-listener-fd", + str(input_fd), + "--output-listener-fd", + str(output_fd), "--engine-start-index", str(engine_start_index), "--engine-count", @@ -378,12 +314,11 @@ def __init__( "--data-parallel-size", str(data_parallel_size), ] - pass_fds = [fd] if grpc_sock is not None: grpc_fd = grpc_sock.fileno() os.set_inheritable(grpc_fd, True) cmd.extend(["--grpc-listen-fd", str(grpc_fd)]) - pass_fds.append(grpc_fd) + inherited_fds.append(grpc_fd) if stats_update_address is not None: cmd.extend(["--coordinator-address", stats_update_address]) from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args @@ -425,18 +360,34 @@ def __init__( redacted_json = json.dumps(redact_sensitive_args(args_dict), sort_keys=True) logger.info("Launching Rust frontend: %s", " ".join(cmd[:-1] + [redacted_json])) - self._proc = subprocess.Popen(cmd, pass_fds=pass_fds) + self._listeners = [input_listener, output_listener] + try: + self._proc = subprocess.Popen(cmd, pass_fds=inherited_fds) + except BaseException: + for listener in self._listeners: + listener.cleanup() + raise + else: + input_listener.close() + output_listener.close() # Create a process wrapper with a sentinel fd for monitoring self.processes: list[_SubprocessWrapper] = [ _SubprocessWrapper(self._proc, "RustFrontend") ] - self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes) + self._finalizer = weakref.finalize( + self, + _shutdown_subprocesses_and_cleanup_listeners, + self.processes, + self._listeners, + ) def shutdown(self, timeout: float | None = None) -> None: if self._finalizer.detach() is not None: - _shutdown_subprocesses(self.processes, timeout=timeout) + _shutdown_subprocesses_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class _SubprocessWrapper: @@ -533,6 +484,18 @@ def _shutdown_subprocesses( logger.debug_once("[shutdown] Subprocess manager: complete") +def _shutdown_subprocesses_and_cleanup_listeners( + procs: list[_SubprocessWrapper], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + _shutdown_subprocesses(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def run_api_server_worker_proc( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: @@ -687,6 +650,18 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: logger.debug_once("[shutdown] Process manager: complete") +def _shutdown_processes_and_cleanup_listeners( + procs: list[BaseProcess], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + shutdown(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def copy_slice( from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int ) -> torch.Tensor: From ee406adf0cbcd1e7b073a6b0ca6a3afeec2b9171 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:28:44 +0400 Subject: [PATCH 181/952] [Bugfix] Fix ZeroDivisionError in Qwen3-VL video sampling when the source fps is unknown (#60501) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/multimodal/test_video.py | 33 ++++++++++++++++++++++++++++++++- vllm/multimodal/video.py | 10 ++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index d593edba2912..407bffed20a5 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -13,7 +13,7 @@ import numpy.typing as npt import pytest import torch -from transformers import AutoVideoProcessor +from transformers import AutoVideoProcessor, Qwen3VLVideoProcessor from transformers.video_utils import VideoMetadata from vllm.assets.base import get_vllm_public_assets @@ -35,6 +35,7 @@ get_video_loader_backend_for_processor, ) from vllm.multimodal.video_decoders import decode_video, resolve_video_backend_kwargs +from vllm.multimodal.video_decoders.opencv import OpenCVVideoBackendMixin from vllm.multimodal.video_decoders.pynvvideocodec import ( PYNVVIDEOCODEC_DECODER_CACHE_SIZE, PyNvVideoCodecDecoderSlot, @@ -730,6 +731,36 @@ def test_cosmos3_edge_uses_qwen3_vl_video_backend(): assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend) +def test_qwen3_vl_unknown_source_fps_matches_hf( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """An unknown source fps (reported as 0) samples like HF's fps=None, and + the metadata carries the fps HF assumes, which the frame timestamps use.""" + get_video_metadata = OpenCVVideoBackendMixin.get_video_metadata + + def unknown_fps(cap) -> VideoSourceMetadata: + return get_video_metadata(cap)._replace(original_fps=0.0, duration=0.0) + + monkeypatch.setattr( + OpenCVVideoBackendMixin, "get_video_metadata", staticmethod(unknown_fps) + ) + _, metadata = Qwen3VLVideoBackend.load_bytes( + dummy_video_path.read_bytes(), fps=2, backend="opencv" + ) + + total_frames = metadata["total_num_frames"] + hf_metadata = VideoMetadata(total_num_frames=total_frames, fps=None) + hf_indices = Qwen3VLVideoProcessor().sample_frames(hf_metadata, fps=2).tolist() + assert metadata["frames_indices"] == hf_indices + assert metadata["fps"] == hf_metadata.fps + + source = VideoSourceMetadata(total_frames, original_fps=0.0, duration=0.0) + target = VideoTargetMetadata(num_frames=-1, fps=2, max_duration=300) + assert ( + Qwen3VLVideoBackend.compute_frames_index_to_sample(source, target) == hf_indices + ) + + @pytest.mark.parametrize( "model_repo, expected_loader_cls, hf_sample_kwargs", [ diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 8f22bffe958c..fc0cdc28d478 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -361,6 +361,15 @@ class Qwen3VLVideoBackend(VideoBackend): _MAX_FRAMES: ClassVar[int] = 768 _MAX_FPS: ClassVar[int] = 30 + @classmethod + def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: + # vLLM reports an unknown source fps as 0. HF's Qwen3VLVideoProcessor + # assumes a 24 fps source then, and keeps it in the metadata for the + # frame timestamps. + if source.original_fps > 0: + return source + return VideoSourceMetadata(source.total_frames_num, 24.0, source.duration) + @classmethod def compute_frames_index_to_sample( cls, @@ -368,6 +377,7 @@ def compute_frames_index_to_sample( target: VideoTargetMetadata, **kwargs, ) -> list[int]: + source = cls._prepare_source(source) total_frames_num = source.total_frames_num original_fps = source.original_fps fps = min(target.fps, cls._MAX_FPS) From 814301b26800a9d22a7d967a161aa07da4cbb5f2 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 08:38:41 +0200 Subject: [PATCH 182/952] [Bugfix][Quant] Read safetensors headers of checkpoints with other file names (#60411) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas --- tests/quantization/test_auto_gptq.py | 18 ++++++++- tests/transformers_utils/test_config.py | 29 ++++++++++++++ .../layers/quantization/utils/gptq_utils.py | 5 +++ vllm/transformers_utils/config.py | 40 ++++++++++++++++++- 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 6e7791654c41..7e54599c8098 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -23,7 +23,10 @@ AutoGPTQLinearMethod, AutoGPTQMoEMethod, ) -from vllm.model_executor.layers.quantization.utils.gptq_utils import override_config +from vllm.model_executor.layers.quantization.utils.gptq_utils import ( + is_layer_gptq_quantized, + override_config, +) from vllm.platforms import current_platform PROMPT = "On the surface of Mars, we found" @@ -71,6 +74,19 @@ def test_auto_gptq_config_get_name(): assert AutoGPTQConfig.get_name() == "auto_gptq" +def test_auto_gptq_quantizes_every_layer_when_the_module_list_is_unknown(): + """An empty module list means the safetensors metadata could not be read, + not that the checkpoint holds no quantized layer.""" + assert is_layer_gptq_quantized( + prefix="model.layers.0.mlp.down_proj", + quantized_layers=[], + ) + assert not is_layer_gptq_quantized( + prefix="model.layers.0.mlp.down_proj", + quantized_layers=["self_attn.q_proj"], + ) + + @pytest.mark.parametrize( "dynamic", [ diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 4fd27ba27879..9b13e725604a 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -241,6 +241,7 @@ def test_safetensors_metadata_of_repo_without_safetensors(): ) api = SimpleNamespace( get_safetensors_metadata=get_safetensors_metadata, + list_repo_files=MagicMock(return_value=["pytorch_model.bin"]), snapshot_download=MagicMock(side_effect=LocalEntryNotFoundError("no cache")), ) @@ -250,6 +251,34 @@ def test_safetensors_metadata_of_repo_without_safetensors(): get_safetensors_metadata.assert_called_once() +def test_safetensors_metadata_of_repo_with_a_nonstandard_file_name(): + """`get_safetensors_metadata` only knows `model.safetensors` and its index, + so older checkpoints are read through the file listing instead.""" + from huggingface_hub.errors import NotASafetensorsRepoError + from huggingface_hub.utils import TensorInfo + + weights = "gptq_model-4bit-128g.safetensors" + tensor = TensorInfo(dtype="I32", shape=[5632], data_offsets=(0, 22528)) + parse_safetensors_file_metadata = MagicMock( + return_value=SimpleNamespace(tensors={"layers.0.mlp.down_proj.qweight": tensor}) + ) + api = SimpleNamespace( + get_safetensors_metadata=MagicMock( + side_effect=NotASafetensorsRepoError("not a safetensors repo") + ), + list_repo_files=MagicMock(return_value=["config.json", weights]), + parse_safetensors_file_metadata=parse_safetensors_file_metadata, + ) + + with patch.object(config_module, "hf_api", lambda: api): + metadata = get_safetensors_params_metadata("some/old-gptq-model") + + assert metadata["layers.0.mlp.down_proj.qweight"]["dtype"] == "I32" + parse_safetensors_file_metadata.assert_called_once_with( + "some/old-gptq-model", weights, revision=None + ) + + @pytest.mark.parametrize( ("section_key", "mrope_section", "expected_num_dims"), [ diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index d134abf8dcec..024564e6f34c 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -133,6 +133,11 @@ def is_layer_gptq_quantized( quantized_layers = flatten_list(quantized_layers) + # Empty means no metadata could be read at all, not that the checkpoint + # holds no quantized layer, and loading GPTQ weights unquantized cannot work. + if not quantized_layers: + return True + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 6a339297eec5..c22f6fc3c2ea 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1355,6 +1355,30 @@ def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]: } +def _read_safetensors_metadata_in_repo( + model: str, + *, + revision: str | None = None, +) -> dict[str, Any]: + api = hf_api() + try: + filenames = [ + filename + for filename in api.list_repo_files(model, revision=revision) + if filename.endswith(".safetensors") + ] + return { + param_name: asdict(info) + for filename in filenames + for param_name, info in api.parse_safetensors_file_metadata( + model, filename, revision=revision + ).tensors.items() + } + except Exception as e: + logger.debug("Could not read safetensors headers for %s: %s", model, e) + return {} + + def get_safetensors_params_metadata( model: str, *, @@ -1372,6 +1396,12 @@ def get_safetensors_params_metadata( for param_name, info in file_mt.tensors.items() } + # `get_safetensors_metadata` only looks for `model.safetensors` and its + # index, so older checkpoints naming their weights differently land here + # even though their headers are perfectly readable. + if metadata := _read_safetensors_metadata_in_repo(model, revision=revision): + return metadata + # Hub fetch failed (e.g. 429, network unreachable). Fall back to the # local HF cache: weights may already be cached from a prior run, and # weight loading itself uses the same cache. @@ -1390,7 +1420,15 @@ def get_safetensors_params_metadata( str(e), ) return {} - return _read_safetensors_metadata_in_dir(Path(local_dir)) + + metadata = _read_safetensors_metadata_in_dir(Path(local_dir)) + if not metadata: + logger.warning_once( + "Could not retrieve safetensors metadata for %s " + "(Hub fetch failed and the cached snapshot holds no weights yet).", + model, + ) + return metadata @cache From 032d39e9b451e858ac9cb18eb0f587ca005361b5 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:53:37 +0800 Subject: [PATCH 183/952] [XPU][CI] Split Intel jobs to improve B50 utilization (#58817) Signed-off-by: zengxian Signed-off-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> --- .buildkite/intel_jobs/lora_intel.yaml | 37 +++++++-- .buildkite/intel_jobs/misc_intel.yaml | 82 +++++++++++++++++-- .../intel_jobs/model_executor_intel.yaml | 32 +++++++- .../intel_jobs/models_multimodal_intel.yaml | 13 +-- .buildkite/intel_jobs/samplers_intel.yaml | 29 ++++++- .buildkite/intel_jobs/test-intel.yaml | 35 ++++++-- 6 files changed, 201 insertions(+), 27 deletions(-) diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 2cecbf18f773..885e4b546691 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -33,7 +33,6 @@ steps: pytest -v -s lora/test_peft_helper.py && pytest -v -s lora/test_resolver.py && pytest -v -s lora/test_utils.py && - pytest -v -s lora/test_add_lora.py && pytest -v -s lora/test_worker.py' - label: LoRA Fused/MoE Kernels @@ -120,7 +119,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -138,8 +137,7 @@ steps: pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" && pytest -v -s lora/test_transformers_model.py && pytest -v -s lora/test_chatglm3_tp.py && - pytest -v -s lora/test_llama_tp.py::test_llama_lora && - pytest -s -v lora/test_minicpmv_tp.py' + pytest -v -s lora/test_llama_tp.py::test_llama_lora' - label: LoRA Multimodal key: lora-multimodal @@ -148,7 +146,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -163,5 +161,30 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - pytest -v -s lora/test_default_mm_loras.py && pytest -v -s lora/test_whisper.py' + +- label: LoRA Add + Multimodal + key: lora-add-multimodal + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/lora + - tests/lora + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + pytest -v -s lora/test_add_lora.py && + pytest -v -s lora/test_default_mm_loras.py && + pytest -s -v lora/test_minicpmv_tp.py' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 5fa4d833d7f0..b7ae9f907846 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -10,7 +10,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -37,7 +37,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -61,7 +61,40 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && pytest -v -s v1/executor && - VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit' + VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit + --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' + +- label: EC CPU Connector (XPU) + key: ec-cpu-connector-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/v1/core + - tests/v1/executor + - tests/v1/kv_offload + - tests/v1/worker + - tests/v1/kv_connector/unit + - tests/v1/ec_connector/unit + - tests/v1/metrics + - tests/entrypoints/openai/correctness/test_lmeval.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'pip install -r requirements/kv_connectors.txt && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + cd tests && + pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' - label: V1 Sample + Logits key: v1-sample-logits-intel @@ -70,7 +103,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -265,7 +298,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -293,10 +326,47 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s detokenizer && + pytest -v -s detokenizer --ignore=detokenizer/test_stop_strings.py && pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' +- label: Detokenizer Stop Strings (XPU) + key: detokenizer-stop-strings-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/assets/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tokenizers/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/detokenizer + - tests/multimodal + - tests/utils_ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s detokenizer/test_stop_strings.py' + - label: Fusion Unit Tests key: fusion-unit-tests timeout_in_minutes: 30 diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 98e4b0ff4511..5da075e235b6 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -28,7 +28,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]" @@ -38,3 +38,31 @@ steps: --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes" --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"' + +- label: Tensorizer Model Loader Tests (XPU) + key: tensorizer-model-loader-tests-xpu + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/engine/arg_utils.py + - vllm/config/model.py + - vllm/model_executor + - tests/model_executor + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'apt-get update && apt-get install -y curl libsodium23 && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + export PYTHONFAULTHANDLER=1 && + cd tests && + pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"' diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index eb36e0c02eb8..061be1b136ed 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -33,7 +33,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -72,7 +72,8 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"' + pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && + pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model' - label: "Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper @@ -81,7 +82,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: @@ -95,7 +96,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' + pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_gemma_ragged_audio.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' - label: Multi-Modal Processor # 44min key: multi-modal-processor @@ -131,7 +132,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml index 41ec3c4df77b..027237225d3a 100644 --- a/.buildkite/intel_jobs/samplers_intel.yaml +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -6,6 +6,33 @@ steps: key: samplers-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - tests/samplers + - tests/conftest.py + - vllm/entrypoints/generate/beam_search + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers + --ignore=samplers/test_beam_search.py' + +- label: Samplers Beam Search Test (XPU Kernel) + key: samplers-beam-search-test-xpu-kernel-intel + timeout_in_minutes: 40 + device: intel_gpu agent_tags: label: production gpu: 1+ @@ -26,4 +53,4 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers' + VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers/test_beam_search.py' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index b1030e372c41..a2a0e14be18c 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -23,7 +23,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -46,7 +46,6 @@ steps: python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel && python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 && - VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC -tp 2 --max-model-len 8192 ' - label: "XPU W8A8 FP8 Linear Examples" @@ -58,7 +57,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -86,7 +85,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 24+ + mem: 16+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" @@ -106,8 +105,34 @@ steps: pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 + - label: "XPU Qwen3-30B + V1 KV Connector" + key: xpu-qwen3-30b-v1-kv-connector + depends_on: + - image-build-xpu + timeout_in_minutes: 90 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - vllm/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pip install dpctl && + pytest -v -s v1/kv_connector/unit/test_multi_connector.py && + pytest -v -s v1/kv_connector/unit/test_example_connector.py && + cd .. && + VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' - label: "XPU server test" key: xpu-server-test depends_on: From e58f8dde7f02631829142f11eaa89099bcd94a98 Mon Sep 17 00:00:00 2001 From: ARAVINDHAN T Date: Thu, 8 Oct 2026 00:04:32 -0700 Subject: [PATCH 184/952] [Bugfix] Use FusedMoEWithLoRA for shared MoE LoRAs on 3D models (#60098) Signed-off-by: ARAVINDHAN T Co-authored-by: linitra24 --- vllm/config/lora.py | 5 +++-- vllm/lora/model_manager.py | 10 ++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/vllm/config/lora.py b/vllm/config/lora.py index f32b89ae5996..6c994a82f8de 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -81,8 +81,9 @@ class LoRAConfig: """If True, load MoE expert adapters in the "shared-outer" layout, where the gate/up (`w1`/`w3`) lora_A and the down (`w2`) lora_B are shared across all experts (stored once with expert-dim 1) instead of per-expert. The shared - factors are broadcast to the expert count at kernel time. Only meaningful for - MoE models whose adapters use this layout; ignored otherwise.""" + factors are broadcast to the expert count at kernel time. 3D-weight MoE + models use `FusedMoEWithLoRA` for this layout. Only meaningful for MoE + models whose adapters use this layout; ignored otherwise.""" max_lora_cls_labels: int | None = Field(default=None, ge=1) """Maximum output size for LoRA classification heads. Defaults to the base classification head size.""" diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 802e5b9233d7..b97a0a412a16 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -144,20 +144,18 @@ def __init__( else set() ) - # When the engine is started with enable_mixed_moe_lora_format=True - # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of - # the model's 3D flag, so 2D and 3D adapters can coexist. + # Mixed-format and shared-outer adapters both use FusedMoEWithLoRA. + # FusedMoE3DWithLoRA only implements the fused gate/up PEFT pair. self._enable_mixed_moe_lora_format = ( is_moe and lora_config.enable_mixed_moe_lora_format ) + self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras self._is_3d_moe_model = ( self._is_moe and self.model.is_3d_moe_weight and not self._enable_mixed_moe_lora_format + and not self._enable_moe_shared_loras ) - # Shared MoE adapters: w13 lora_A / w2 lora_B shared across experts, - # stored as pre-stacked experts.w{1,2,3} tensors (startup opt-in). - self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras self.packed_modules_mapping = process_packed_modules_mapping( self.model, force_2d_moe=self._enable_mixed_moe_lora_format, From fd0216d3b598501e28041d84ed0b7dd16453aac6 Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Thu, 8 Oct 2026 12:45:53 +0530 Subject: [PATCH 185/952] [Bugfix][CPU] Fix fused Gumbel-max wrapping-window bias (#60221) Signed-off-by: Rehan Khan Co-authored-by: Bob --- .buildkite/hardware_tests/cpu.yaml | 2 + .../scripts/hardware_ci/run-cpu-test-arm.sh | 1 + .buildkite/test_areas/kernels.yaml | 1 + csrc/cpu/sampling_kernels.cpp | 151 +++++++++-------- tests/kernels/test_cpu_fused_sampling.py | 153 ++++++++++++++++-- vllm/v1/sample/ops/topk_topp_sampler.py | 8 +- 6 files changed, 236 insertions(+), 80 deletions(-) diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 3fdce975d3f4..d54688b57491 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -23,6 +23,7 @@ steps: - tests/test_zentorch_version_pin.py - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py - tests/kernels/test_onednn.py + - tests/kernels/test_cpu_fused_sampling.py - tests/kernels/test_awq_int4_to_int8.py - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - tests/kernels/quantization/test_zen_da8w4.py @@ -49,6 +50,7 @@ steps: tests/test_zentorch_version_pin.py \ tests/kernels/mamba/test_cpu_short_conv.py \ tests/kernels/test_onednn.py \ + tests/kernels/test_cpu_fused_sampling.py \ tests/kernels/test_awq_int4_to_int8.py \ tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \ tests/kernels/quantization/test_zen_da8w4.py \ diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 689ca92d6a05..1f598b277bb2 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -43,6 +43,7 @@ kernel_tests() { docker exec "$CONTAINER_NAME" bash -c " set -e pytest -x -v -s tests/kernels/test_onednn.py + pytest -x -v -s tests/kernels/test_cpu_fused_sampling.py pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index cfc628d382c6..d3c8f657902c 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -256,6 +256,7 @@ steps: --ignore=kernels/test_fused_qk_norm_rope_gate.py --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py --ignore=kernels/test_top_k_per_row.py + --ignore=kernels/test_cpu_fused_sampling.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py --ignore=kernels/test_bf16x3_router_gemm_rocm.py --ignore=kernels/test_ll_bf16_gemm.py diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp index 3e2f00a0ba95..1d30b92a8d9a 100644 --- a/csrc/cpu/sampling_kernels.cpp +++ b/csrc/cpu/sampling_kernels.cpp @@ -1,82 +1,73 @@ #include "cpu_types.hpp" -#include +#include #include #include #include -#include -#include namespace { -constexpr int GUMBEL_TABLE_SIZE = 1 << 20; -constexpr int GUMBEL_TABLE_MASK = GUMBEL_TABLE_SIZE - 1; +// Gumbel-max: argmax_i(logit_i + g_i) with g_i ~ Gumbel(0, 1) matches +// softmax. Noise is SplitMix64(seed, i), so tokens are independent and +// consecutive seeds are not a 1-token table shift (vllm#59786). +constexpr uint64_t SPLITMIX_GAMMA = 0x9E3779B97F4A7C15ULL; +constexpr double GUMBEL_MAX = 37.5; // 53-bit u in (0,1) => -log(-log u) < 37.5 -static float g_gumbel_table[GUMBEL_TABLE_SIZE]; -static std::once_flag g_gumbel_init_flag; - -static void init_gumbel_table() { - std::mt19937 rng(42); - std::uniform_real_distribution uniform(1.0e-7f, 1.0f - 1.0e-7f); - for (int i = 0; i < GUMBEL_TABLE_SIZE; ++i) { - float u = uniform(rng); - g_gumbel_table[i] = -std::log(-std::log(u)); - } +static inline uint64_t splitmix64(uint64_t z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); } -static inline void ensure_gumbel_table() { - std::call_once(g_gumbel_init_flag, init_gumbel_table); +static void fused_gumbel_argmax_row(int64_t* __restrict__ output, + const float* __restrict__ row, + uint64_t seed, int64_t vocab_size) { + const uint64_t key = splitmix64(seed); + + // Compare in (logit, noise) space: x + g > best_x + best_g is + // (x - best_x) + g > best_g. Adding g into a 1e38 logit loses it in + // double ULP, so ties at the float max would always keep the lower index. + double best_logit = 0.0; + double best_g = 0.0; + int64_t best_idx = 0; + bool have_best = false; + for (int64_t i = 0; i < vocab_size; ++i) { + const double x = row[i]; + if (x == -std::numeric_limits::infinity()) continue; + if (have_best && (x - best_logit) + GUMBEL_MAX <= best_g) continue; + const uint64_t bits = + splitmix64(key + uint64_t(i + 1) * SPLITMIX_GAMMA) >> 11; + const double one_minus_u = + (double((1ULL << 53) - 1 - bits) + 0.5) * 0x1.0p-53; + if (have_best && one_minus_u >= std::exp((x - best_logit) - best_g)) + continue; + const double u = (double(bits) + 0.5) * 0x1.0p-53; + const double g = -std::log(-std::log(u)); + if (!have_best || (x - best_logit) + g > best_g) { + have_best = true; + best_logit = x; + best_g = g; + best_idx = i; + } + } + *output = best_idx; } -// Fused Gumbel-max kernel - static void fused_gumbel_argmax_kernel(int64_t* __restrict__ output, const float* __restrict__ logits, const int64_t* __restrict__ seeds, const int64_t batch_size, const int64_t vocab_size) { - ensure_gumbel_table(); - constexpr int VEC_ELEM_NUM = vec_op::FP32Vec8::VEC_ELEM_NUM; - const int64_t vec_end = vocab_size - (vocab_size % VEC_ELEM_NUM); - + if (batch_size <= 0) return; + if (batch_size == 1) { + fused_gumbel_argmax_row(output, logits, uint64_t(seeds[0]), vocab_size); + return; + } #pragma omp parallel for schedule(static) for (int64_t b = 0; b < batch_size; ++b) { - const float* row = logits + b * vocab_size; - const uint32_t seed = static_cast(seeds[b]); - - float best_score = -std::numeric_limits::infinity(); - int64_t best_idx = 0; - - for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) { - vec_op::FP32Vec8 logit_vec(row + i); - - float gumbel_buf[VEC_ELEM_NUM]; - for (int k = 0; k < VEC_ELEM_NUM; ++k) - gumbel_buf[k] = g_gumbel_table[(seed + i + k) & GUMBEL_TABLE_MASK]; - vec_op::FP32Vec8 gumbel_vec(gumbel_buf); - - vec_op::FP32Vec8 score = logit_vec + gumbel_vec; - - float score_buf[VEC_ELEM_NUM]; - score.save(score_buf); - for (int k = 0; k < VEC_ELEM_NUM; ++k) { - if (score_buf[k] > best_score) { - best_score = score_buf[k]; - best_idx = i + k; - } - } - } - - for (int64_t i = vec_end; i < vocab_size; ++i) { - float s = row[i] + g_gumbel_table[(seed + i) & GUMBEL_TABLE_MASK]; - if (s > best_score) { - best_score = s; - best_idx = i; - } - } - - output[b] = best_idx; + fused_gumbel_argmax_row(output + b, logits + b * vocab_size, + uint64_t(seeds[b]), vocab_size); } } @@ -91,22 +82,41 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output, for (int64_t b = 0; b < batch_size; ++b) { const float* row = logits + b * vocab_size; + // Vector max has no consistent NaN rule, so NaNs are caught via the sum: + // any NaN (or inf + -inf, a harmless false positive) makes it NaN. vec_op::FP32Vec16 vmax(-std::numeric_limits::infinity()); + vec_op::FP32Vec16 vsum(0.0f); for (int64_t i = 0; i < vec_end; i += VEC_ELEM_NUM) { - vmax = vmax.max(vec_op::FP32Vec16(row + i)); + vec_op::FP32Vec16 v(row + i); + vmax = vmax.max(v); + vsum = vsum + v; } float best_val = vmax.reduce_max(); + bool maybe_nan = std::isnan(vsum.reduce_sum()); for (int64_t i = vec_end; i < vocab_size; ++i) { + maybe_nan |= std::isnan(row[i]); if (row[i] > best_val) { best_val = row[i]; } } - int64_t best_idx = 0; - for (int64_t i = 0; i < vocab_size; ++i) { - if (row[i] == best_val) { - best_idx = i; - break; + int64_t best_idx = -1; + if (maybe_nan) { + // torch.argmax treats NaN as the maximum: return the first NaN. + for (int64_t i = 0; i < vocab_size; ++i) { + if (std::isnan(row[i])) { + best_idx = i; + break; + } + } + } + if (best_idx < 0) { + best_idx = 0; + for (int64_t i = 0; i < vocab_size; ++i) { + if (row[i] == best_val) { + best_idx = i; + break; + } } } output[b] = best_idx; @@ -117,21 +127,28 @@ static void greedy_argmax_kernel(int64_t* __restrict__ output, torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits, const torch::Tensor& seeds) { + TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]"); TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "logits must be float32"); + TORCH_CHECK(logits.size(1) > 0, "vocab_size must be positive"); + TORCH_CHECK(seeds.device().is_cpu(), "seeds must be a CPU tensor"); + TORCH_CHECK(seeds.scalar_type() == torch::kInt64, "seeds must be int64"); TORCH_CHECK(seeds.dim() == 1 && seeds.size(0) == logits.size(0), "seeds must be 1-D with batch_size elements"); auto logits_contig = logits.contiguous(); + auto seeds_contig = seeds.contiguous(); auto output = torch::empty({logits_contig.size(0)}, torch::kInt64); - fused_gumbel_argmax_kernel( - output.data_ptr(), logits_contig.data_ptr(), - seeds.data_ptr(), logits_contig.size(0), logits_contig.size(1)); + fused_gumbel_argmax_kernel(output.data_ptr(), + logits_contig.data_ptr(), + seeds_contig.data_ptr(), + logits_contig.size(0), logits_contig.size(1)); return output; } torch::Tensor greedy_argmax(const torch::Tensor& logits) { + TORCH_CHECK(logits.device().is_cpu(), "logits must be a CPU tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [batch, vocab]"); TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "logits must be float32"); diff --git a/tests/kernels/test_cpu_fused_sampling.py b/tests/kernels/test_cpu_fused_sampling.py index 8a3857fe9ee7..cdb10ab7e32e 100644 --- a/tests/kernels/test_cpu_fused_sampling.py +++ b/tests/kernels/test_cpu_fused_sampling.py @@ -43,20 +43,20 @@ def test_negative_logits(self): class TestFusedGumbelArgmax: - @pytest.mark.parametrize("vocab_size", VOCAB_SIZES) - def test_distribution_chi_squared(self, vocab_size: int): + def test_distribution_chi_squared(self): """Verify sampling distribution via chi-squared goodness of fit.""" - small_vocab = min(vocab_size, 100) + # Small support is enough for chi-squared power; large-vocab bias + # is covered by test_noise_not_sliding_window (#59786). + small_vocab = 100 logits = torch.randn(1, small_vocab, dtype=torch.float32) probs = logits.softmax(dim=-1).squeeze(0) n_samples = 100_000 - counts = torch.zeros(small_vocab) - for trial in range(n_samples): - seed = torch.tensor([trial * 7 + 13], dtype=torch.long) - tile = logits.expand(1, -1).contiguous() - idx = torch.ops._C.fused_gumbel_argmax(tile, seed) - counts[idx.item()] += 1 + seeds = torch.arange(n_samples, dtype=torch.long) * 7 + 13 + idx = torch.ops._C.fused_gumbel_argmax( + logits.expand(n_samples, -1).contiguous(), seeds + ) + counts = torch.bincount(idx, minlength=small_vocab).float() expected = probs * n_samples mask = expected > 5 @@ -70,6 +70,39 @@ def test_distribution_chi_squared(self, vocab_size: int): f"dof={dof}, p={p_value:.6f}" ) + def test_noise_not_sliding_window(self): + """Table[(seed+i)&MASK] made consecutive seeds a 1-token shift. + + On uniform logits that yields P(w(s+1) == w(s)-1) ~ 1-2/V. Independent + per-token noise yields ~1/V. See #59786. + """ + vocab, n = 128, 2048 + logits = torch.zeros(n, vocab, dtype=torch.float32) + seeds = torch.arange(n, dtype=torch.long) + winners = torch.ops._C.fused_gumbel_argmax(logits, seeds) + shifted = (winners[:-1] - 1) % vocab + frac = (winners[1:] == shifted).float().mean().item() + assert frac < 0.05, ( + f"consecutive-seed winners follow a table shift (frac={frac:.3f})" + ) + + def test_zipf_head_matches_softmax(self): + """Zipf head probability should match softmax (#59786).""" + vocab, n_samples, batch = 8192, 16384, 256 + logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log() + p0 = torch.softmax(logits.double(), dim=-1)[0].item() + row = logits.view(1, -1).expand(batch, vocab).contiguous() + n_top = 0 + for start in range(0, n_samples, batch): + seeds = torch.arange(start, start + batch, dtype=torch.long) + tokens = torch.ops._C.fused_gumbel_argmax(row, seeds) + n_top += int((tokens == 0).sum()) + q0 = n_top / n_samples + se = (p0 * (1.0 - p0) / n_samples) ** 0.5 + assert abs(q0 - p0) < 8.0 * se, ( + f"P(top-1) kernel={q0:.4f} exact={p0:.4f} se={se:.4f}" + ) + def test_deterministic_same_seed(self): """Same seed produces same result.""" logits = torch.randn(4, 32000, dtype=torch.float32) @@ -95,6 +128,108 @@ def test_output_in_range(self, vocab_size: int, batch_size: int): assert result.min() >= 0 assert result.max() < vocab_size + def test_samples_only_from_finite_logits(self): + """top-k/top-p mask with -inf; those IDs must never win.""" + allowed = torch.tensor([3, 17, 99, 1024, 8191], dtype=torch.long) + logits = torch.full((64, 32000), float("-inf"), dtype=torch.float32) + logits[:, allowed] = torch.randn(64, allowed.numel()) + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(64, dtype=torch.long) + ) + assert torch.isin(tokens, allowed).all() + + def test_single_unmasked_token(self): + logits = torch.full((8, 4096), float("-inf"), dtype=torch.float32) + logits[:, 123] = 0.0 + seeds = torch.tensor([0, -1, -(2**63), 2**63 - 1, 42, 99, 1000, 7]) + tokens = torch.ops._C.fused_gumbel_argmax(logits, seeds) + assert (tokens == 123).all() + + def test_tokens_beyond_gumbel_max_cannot_win(self): + """A gap larger than GUMBEL_MAX is unreachable.""" + logits = torch.zeros(256, 4, dtype=torch.float32) + logits[:, 1:] = -80.0 + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(256, dtype=torch.long) + ) + assert (tokens == 0).all() + + def test_tied_huge_logits_not_index_biased(self): + """Noise must still break ties when logits sit at float32 max.""" + logits = torch.full( + (4096, 2), torch.finfo(torch.float32).max, dtype=torch.float32 + ) + tokens = torch.ops._C.fused_gumbel_argmax( + logits, torch.arange(4096, dtype=torch.long) + ) + frac0 = (tokens == 0).float().mean().item() + assert set(tokens.tolist()) == {0, 1}, tokens.unique().tolist() + assert 0.35 < frac0 < 0.65, f"P(index 0)={frac0:.3f} expected ~0.5" + + def test_serial_batch1_matches_batched(self): + """OpenMP rows must match the batch=1 serial path.""" + logits = torch.randn(8, 2048, dtype=torch.float32) + seeds = torch.arange(8, dtype=torch.long) + batched = torch.ops._C.fused_gumbel_argmax(logits, seeds) + serial = torch.stack( + [ + torch.ops._C.fused_gumbel_argmax(logits[i : i + 1], seeds[i : i + 1])[0] + for i in range(8) + ] + ) + torch.testing.assert_close(batched, serial) + + def test_noncontiguous_matches_contiguous(self): + logits = torch.randn(8, 1024, dtype=torch.float32)[:, ::2] + seeds = torch.arange(16, dtype=torch.long)[::2] + assert not logits.is_contiguous() and not seeds.is_contiguous() + expected = torch.ops._C.fused_gumbel_argmax( + logits.contiguous(), seeds.contiguous() + ) + actual = torch.ops._C.fused_gumbel_argmax(logits, seeds) + torch.testing.assert_close(actual, expected) + + def test_empty_batch(self): + tokens = torch.ops._C.fused_gumbel_argmax( + torch.empty(0, 7), torch.empty(0, dtype=torch.long) + ) + assert tokens.shape == (0,) and tokens.dtype == torch.int64 + + def test_masked_support_matches_softmax(self): + """Finite subset after -inf mask should match softmax.""" + from scipy.stats import chisquare + + logits = torch.tensor( + [0.0, -0.2, -0.7, -1.5, -3.0, float("-inf")], dtype=torch.float32 + ) + n_samples = 20_000 + tokens = torch.ops._C.fused_gumbel_argmax( + logits.expand(n_samples, -1).contiguous(), + torch.arange(n_samples, dtype=torch.long), + ) + counts = torch.bincount(tokens, minlength=logits.numel()).double() + assert counts[-1] == 0 + probs = logits.double().softmax(-1) + _, p_value = chisquare(counts[:-1].numpy(), (probs[:-1] * n_samples).numpy()) + assert p_value > 1e-7, f"masked-support p-value: {p_value}" + + def test_zipf_tail_matches_softmax(self): + """Independent 53-bit Gumbel must not starve the Zipf tail (#59786).""" + vocab, n_samples, batch = 8192, 16384, 256 + logits = -1.1 * torch.arange(1, vocab + 1, dtype=torch.float32).log() + p_tail = torch.softmax(logits.double(), dim=-1)[vocab // 2 :].sum().item() + row = logits.view(1, -1).expand(batch, vocab).contiguous() + n_tail = 0 + for start in range(0, n_samples, batch): + seeds = torch.arange(start, start + batch, dtype=torch.long) + tokens = torch.ops._C.fused_gumbel_argmax(row, seeds) + n_tail += int((tokens >= vocab // 2).sum()) + q_tail = n_tail / n_samples + se = (p_tail * (1.0 - p_tail) / n_samples) ** 0.5 + assert abs(q_tail - p_tail) < 8.0 * se, ( + f"P(tail) kernel={q_tail:.4f} exact={p_tail:.4f} se={se:.4f}" + ) + class TestMemory: def test_fused_no_intermediate_allocs(self): diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index a3084a09c28e..daa94ebc3af2 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -319,11 +319,11 @@ def forward_cpu( k: torch.Tensor | None, p: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Fused Gumbel-max sampling for CPU. + """CPU Gumbel-max sampling with per-token SplitMix64 noise. - Uses a precomputed Gumbel table + single-pass argmax over logits, - skipping softmax and intermediate allocations entirely. - Falls back to the native path when fp64 Gumbel noise is requested. + Independent (seed, token) streams avoid the wrapping-window bias + (#59786). Falls back to native sampling when fp64 Gumbel noise + is requested. """ logits = apply_top_k_top_p(logits, k, p) logits_to_return = None From 29ae81fe83712bb74c00a4ce0515d031f1225c9b Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Thu, 8 Oct 2026 00:27:47 -0700 Subject: [PATCH 186/952] [Rust Frontend] Preserve explicit zero sampling values in gRPC requests (#60115) Signed-off-by: Alec Flowers --- rust/proto/inference.proto | 6 +-- rust/src/server/src/grpc/convert.rs | 63 +++++++++++++++++++++++------ rust/src/server/src/grpc/tests.rs | 34 ++++++++-------- 3 files changed, 71 insertions(+), 32 deletions(-) diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 8189758212d4..79c1d0a147e2 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -75,9 +75,9 @@ message KVHintAction { message RandomSampling { uint32 num_sequences = 1; // "n", default (0) means 1 - uint32 top_k = 2; // 0 means default - float top_p = 3; // 0 means default - float min_p = 4; // 0 means default + optional uint32 top_k = 2; // Omitted means default; 0 disables top-k. + optional float top_p = 3; // Omitted means default; must be in (0, 1]. + optional float min_p = 4; // Omitted means default; 0 disables min-p. optional int64 seed = 5; } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 8f2c589f154b..a409cc7f79ee 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -161,10 +161,7 @@ fn build_sampling_params( ..SamplingParams::default() }; - // RandomSampling: for every remaining sampling field the protobuf default (`0`) - // is treated as "unset" and leaves the resolved value to the lowering - // stage, which falls back to the model-provided default or a - // neutral/disabled value otherwise. + // Preserve explicit values; omitted fields inherit model defaults during lowering. if let Some(s) = sampling { // num_sequences (n > 1) is not supported yet by the TextLlm layer; the response // path also hardcodes SequenceOutput.index = 0, so accepting >1 would silently @@ -174,15 +171,9 @@ fn build_sampling_params( "num_sequences > 1 is not supported", )); } - if s.top_k != 0 { - params.top_k = Some(s.top_k); - } - if s.top_p != 0.0 { - params.top_p = Some(s.top_p); - } - if s.min_p != 0.0 { - params.min_p = Some(s.min_p); - } + params.top_k = s.top_k; + params.top_p = s.top_p; + params.min_p = s.min_p; params.seed = s.seed; } @@ -790,6 +781,52 @@ mod tests { } } + #[test] + fn grpc_sampling_preserves_explicit_values_and_model_defaults_through_lowering() { + for ((top_k, top_p, min_p), expected) in [ + ((None, None, None), (8, 0.9, 0.2)), + ((Some(0), Some(1.0), Some(0.0)), (0, 1.0, 0.0)), + ((Some(50), Some(0.8), Some(0.1)), (50, 0.8, 0.1)), + ] { + let request = pb::GenerateRequest { + temperature: Some(0.7), + sampling: Some(pb::RandomSampling { + top_k, + top_p, + min_p, + ..Default::default() + }), + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(decoded, stream, &["test-model".to_string()]).unwrap(); + let engine = vllm_text::lower_text_request( + text, + vec![1], + SamplingHints { + default_top_k: Some(8), + default_top_p: Some(0.9), + default_min_p: Some(0.2), + ..Default::default() + }, + SamplingLimits { + max_model_len: 256, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + &TestTokenizer::new(), + ) + .unwrap() + .generate_request; + let params = engine.sampling_params; + assert_eq!((params.top_k, params.top_p, params.min_p), expected); + } + } + } + #[test] fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() { let req = pb::GenerateRequest { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 631eb76b36ab..0ddbfb096e43 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1366,22 +1366,24 @@ async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { ) .await; - let status = client - .generate(pb::GenerateRequest { - request_id: "test-invalid-sampling".to_string(), - model: "test-model".to_string(), - prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), - sampling: Some(pb::RandomSampling { - top_p: 2.0, + for top_p in [0.0, 2.0] { + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-sampling".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + sampling: Some(pb::RandomSampling { + top_p: Some(top_p), + ..Default::default() + }), ..Default::default() - }), - ..Default::default() - }) - .await - .expect_err("should fail when top_p is out of range"); + }) + .await + .expect_err("should fail when top_p is out of range"); - assert_eq!(status.code(), tonic::Code::InvalidArgument); - assert!(status.message().contains("top_p")); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("top_p")); + } server_task.abort(); } @@ -1519,8 +1521,8 @@ async fn unary_generate_with_sampling_params() { prompt: Some(pb::generate_request::Prompt::Text("test".to_string())), temperature: Some(0.7), sampling: Some(pb::RandomSampling { - top_k: 50, - top_p: 0.9, + top_k: Some(50), + top_p: Some(0.9), seed: Some(42), ..Default::default() }), From 0f112d180033f3757cf283723f5593a40b02ed82 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Thu, 8 Oct 2026 09:33:28 +0200 Subject: [PATCH 187/952] [Bugfix][Pooling] Preserve cache_salt for chunked embeddings (#47696) Signed-off-by: Ting Sun --- vllm/entrypoints/pooling/embed/io_processor.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 3c5b75cee62c..1b23a77ec60d 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -142,13 +142,16 @@ def maybe_pre_process_chunked(self, ctx: PoolingServeContext) -> None: ) prompt_token_ids = cast(list[int], token_ids) + cache_salt = cast(str | None, engine_input["prompts"].get("cache_salt")) for chunk_idx, chunk_tokens in enumerate( chunk_list(prompt_token_ids, max_model_len) ): chunked_engine_inputs.append( PoolingEngineInput( - prompts=tokens_input(prompt_token_ids=chunk_tokens), + prompts=tokens_input( + prompt_token_ids=chunk_tokens, cache_salt=cache_salt + ), params=engine_input["params"], lora_requests=engine_input["lora_requests"], priorities=engine_input["priorities"], From 4cf1307e72cb1648e4e3c3ab4e0bf85d5616fcb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Ma=C5=84czak?= Date: Thu, 8 Oct 2026 10:36:05 +0200 Subject: [PATCH 188/952] [XPU] Add tuned Mamba SSU configs for 2 more B70 shapes (#57565) Signed-off-by: pmanczak --- ..._Pro_B70_Graphics,cache_dtype=float16.json | 47 +++++++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float32.json | 39 +++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float16.json | 47 +++++++++++++++++++ ..._Pro_B70_Graphics,cache_dtype=float32.json | 47 +++++++++++++++++++ 4 files changed, 180 insertions(+) create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..d51b085769b1 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "48": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "96": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "1536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "3072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "6144": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "12288": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "24576": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "49152": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..666506768c01 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,39 @@ +{ + "triton_version": "3.8.0", + "48": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "768": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "1536": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "3072": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "6144": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "12288": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "24576": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "49152": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..aed0a3c63515 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "80": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "640": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2560": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "5120": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "10240": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "20480": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "40960": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "81920": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "122880": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "163840": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..7c642c06bc2e --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,47 @@ +{ + "triton_version": "3.8.0", + "80": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "640": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "1280": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2560": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "5120": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "10240": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "20480": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "40960": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "81920": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "122880": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "163840": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + } +} \ No newline at end of file From 7d47ac2ddba430510f1503dae424bb298fe5a300 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Thu, 8 Oct 2026 04:52:59 -0400 Subject: [PATCH 189/952] [CI] Enroll five steps in automatic sharding (#60492) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 Co-authored-by: khluu --- .buildkite/test_areas/entrypoints.yaml | 1 + .buildkite/test_areas/kernels.yaml | 1 + .buildkite/test_areas/quantization.yaml | 1 + .buildkite/test_areas/spec_decode.yaml | 1 + 4 files changed, 4 insertions(+) diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 68c6243d479b..34a190e810b9 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -231,6 +231,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Pooling)" device: h200_35gb key: entrypoints-integration-pooling + automatic_shard: true timeout_in_minutes: 75 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d3c8f657902c..e0d6c4ce40a6 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -504,6 +504,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Mamba Kernels" device: h200_35gb key: kernels-mamba-test + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - csrc/mamba/ diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index b49b4b92c736..429b01f27320 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -127,6 +127,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Quantized Models" device: h200_35gb key: quantized-models-test + automatic_shard: true timeout_in_minutes: 65 env: VLLM_USE_V2_MODEL_RUNNER: "0" diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 6882c24b3745..d031d4fdea9b 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) V1 Spec Decode" device: h200_35gb key: v1-spec-decode + automatic_shard: true timeout_in_minutes: 40 source_file_dependencies: - vllm/config/ From 70bdce6133d0094f420a3714bb9ee186238c82f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 11:05:52 +0200 Subject: [PATCH 190/952] [Model][GLM-5.3-Flash] Warm up the DCP top-k merge kernels for the kpool indexer (#60032) Signed-off-by: NickLucche --- vllm/models/glm5next/nvidia/sparse_indexer.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 806b633c5836..77fa4c409152 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -26,6 +26,7 @@ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform from vllm.utils.deep_gemm import has_deep_gemm +from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( LayerNameType, _resolve_layer_name, @@ -717,6 +718,14 @@ def __init__( ) self.dcp_world_size = _parallel.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group + if current_platform.is_cuda() and has_cutedsl(): + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( # noqa: E501 + _PACK_DCP_TOPK_CANDIDATES_KERNEL, + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL, + ) + + _PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup() + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup() @property def cp_kv_cache_interleave_size(self) -> int: From 0004536dd00f8237fd6dd68171fbda9b680c07b2 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Thu, 8 Oct 2026 17:13:48 +0800 Subject: [PATCH 191/952] [XPU][Graph] Replace boolean-mask assignment with `torch.where` (#60529) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- vllm/model_executor/models/mimo_mtp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/mimo_mtp.py b/vllm/model_executor/models/mimo_mtp.py index 65566230a562..5e331fa8f002 100644 --- a/vllm/model_executor/models/mimo_mtp.py +++ b/vllm/model_executor/models/mimo_mtp.py @@ -73,7 +73,7 @@ def forward( ) -> torch.Tensor: assert inputs_embeds is not None # masking inputs at position 0, as not needed by MTP - inputs_embeds[positions == 0] = 0 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) inputs_embeds = self.token_layernorm(inputs_embeds) previous_hidden_states = self.hidden_layernorm(previous_hidden_states) From 21d97584593de08f7a1a44dad8aca32aa4a01c51 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:30:52 +0100 Subject: [PATCH 192/952] Remove Python 3.10 support now that it is EOL (#60402) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .pre-commit-config.yaml | 28 +++------- AGENTS.md | 3 -- CMakeLists.txt | 48 ++++++++--------- docker/Dockerfile.cpu | 2 +- docs/contributing/README.md | 10 +--- docs/deployment/frameworks/skypilot.md | 8 +-- docs/getting_started/installation/cpu.md | 2 +- docs/getting_started/installation/gpu.md | 2 +- docs/getting_started/quickstart.md | 2 +- .../structured_outputs/pyproject.toml | 2 +- pyproject.toml | 5 +- requirements/cpu.txt | 1 - setup.py | 14 ++--- .../launchers/test_dp_supervisor.py | 2 +- tests/entrypoints/launchers/test_shutdown.py | 6 +-- .../tensorizer_loader/test_tensorizer.py | 2 +- tests/multimodal/media/test_connector.py | 7 ++- tests/v1/distributed/test_async_llm_dp.py | 2 +- tools/ci/kernel_symbol_map.py | 2 +- tools/pre_commit/check_forbidden_imports.py | 6 ++- tools/pre_commit/mypy.py | 4 +- vllm/benchmarks/serve.py | 4 +- vllm/benchmarks/sweep/plot.py | 4 +- vllm/benchmarks/sweep/serve_workload.py | 3 +- vllm/benchmarks/sweep/server.py | 2 +- vllm/config/engram.py | 3 +- vllm/config/lora.py | 3 +- vllm/config/parallel.py | 3 +- vllm/config/profiler.py | 3 +- vllm/config/scheduler.py | 3 +- vllm/config/speculative.py | 3 +- vllm/config/structured_outputs.py | 3 +- vllm/config/watermarking.py | 3 +- vllm/distributed/utils.py | 12 ++--- vllm/distributed/weight_transfer/base.py | 2 +- .../distributed/weight_transfer/ipc_engine.py | 3 +- .../weight_transfer/m2n_trainer.py | 3 +- .../weight_transfer/nccl_engine.py | 3 +- .../weight_transfer/sharded_rdt_trainer.py | 3 +- .../weight_transfer/sparse_nccl_engine.py | 3 +- vllm/entrypoints/chat_utils.py | 3 +- vllm/entrypoints/launchers/dp_supervisor.py | 2 +- vllm/entrypoints/llm.py | 3 +- vllm/entrypoints/pooling/embed/serving.py | 3 +- vllm/entrypoints/pooling/pooling/serving.py | 3 +- vllm/entrypoints/pooling/scoring/typing.py | 4 +- vllm/entrypoints/serve/tokenize/api_router.py | 2 +- vllm/inputs/engine.py | 4 +- vllm/inputs/llm.py | 4 +- vllm/model_executor/kernels/linear/base.py | 3 +- .../scaled_mm/BlockScaledMMLinearKernel.py | 3 +- vllm/model_executor/model_loader/utils.py | 3 +- vllm/model_executor/models/interfaces.py | 3 +- .../deepseek_v4/common/mm_preprocess.py | 3 +- vllm/multimodal/parse.py | 2 +- vllm/multimodal/processing/processor.py | 3 +- vllm/renderers/hf.py | 2 +- vllm/snapshot/manifest.py | 3 +- vllm/snapshot/runtime.py | 4 +- vllm/tokenizers/registry.py | 4 +- vllm/usage/usage_lib.py | 2 +- vllm/utils/collection_utils.py | 4 +- vllm/utils/import_utils.py | 3 +- vllm/v1/kv_cache_interface.py | 3 +- vllm/v1/metrics/perf.py | 3 +- vllm/v1/structured_output/backend_outlines.py | 53 +++++++------------ vllm/v1/worker/workspace.py | 7 +-- 67 files changed, 141 insertions(+), 219 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3e418e005ec8..92963cf69a22 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -199,29 +199,13 @@ repos: additional_dependencies: [PyYAML, regex] verbose: true # full-scan `note:` lines are advisory, so the hook passes require_serial: true # one full scan, not one per pre-commit file batch - - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.10 - entry: python tools/pre_commit/mypy.py "3.10" - <<: &mypy_common - language: python - types_or: [python, pyi] - require_serial: true - additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] - - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.11 + - id: mypy # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward + name: Run mypy for the oldest supported Python version entry: python tools/pre_commit/mypy.py "3.11" - <<: *mypy_common - stages: [manual] # Only run in CI - - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.12 - entry: python tools/pre_commit/mypy.py "3.12" - <<: *mypy_common - stages: [manual] # Only run in CI - - id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward - name: Run mypy for Python 3.13 - entry: python tools/pre_commit/mypy.py "3.13" - <<: *mypy_common - stages: [manual] # Only run in CI + language: python + types_or: [python, pyi] + require_serial: true + additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh diff --git a/AGENTS.md b/AGENTS.md index 933fa05cfcab..a31d156b4b63 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,9 +114,6 @@ pre-commit run --all-files # Run a specific hook: pre-commit run ruff-check --all-files - -# Run mypy as it is in CI: -pre-commit run mypy-3.12 --all-files --hook-stage manual ``` The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check. diff --git a/CMakeLists.txt b/CMakeLists.txt index 0a5414b4bf41..4ed5ff2573c1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -46,7 +46,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) # Supported python versions. These versions will be searched in order, the # first match will be selected. These should be kept in sync with setup.py. # -set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") +set(PYTHON_SUPPORTED_VERSIONS "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") @@ -134,40 +134,34 @@ endif() # # spinloop extension (pure CXX; must stay above the non-CUDA device branch so # CPU builds define the target before the early return) -# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading -# failure is handled gracefully on vLLM side for lower Python versions. -# Skip the target entirely on Python < 3.11 so the build doesn't break. +# This extension requires SABI 3.11 since it relies on Py_buffer support. # -if(Python_VERSION VERSION_GREATER_EQUAL "3.11") - set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") - set(SPINLOOP_COMPILE_FLAGS "") - if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") - list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") - endif() - define_extension_target( - spinloop - DESTINATION vllm - LANGUAGE CXX - SOURCES ${VLLM_SPINLOOP_EXT_SRC} - COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} - USE_SABI 3.11 - WITH_SOABI) +set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") +set(SPINLOOP_COMPILE_FLAGS "") +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") + list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx") endif() +define_extension_target( + spinloop + DESTINATION vllm + LANGUAGE CXX + SOURCES ${VLLM_SPINLOOP_EXT_SRC} + COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS} + USE_SABI 3.11 + WITH_SOABI) # # fs_io extension (pure CXX; must stay above the non-CUDA device branch # so CPU builds define the target before the early return). # GIL-releasing filesystem helpers for FileSystemTierManager. # -if(Python_VERSION VERSION_GREATER_EQUAL "3.11") - define_extension_target( - fs_io_C - DESTINATION vllm - LANGUAGE CXX - SOURCES csrc/fs_io.cpp - USE_SABI 3.11 - WITH_SOABI) -endif() +define_extension_target( + fs_io_C + DESTINATION vllm + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) # # Forward the non-CUDA device extensions to external CMake scripts. diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 655a956ba7cd..b3943a122564 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -13,7 +13,7 @@ # vllm-dev: used for development # # Build arguments: -# PYTHON_VERSION=3.13|3.12 (default)|3.11|3.10 +# PYTHON_VERSION=3.13|3.12 (default)|3.11 # VLLM_CPU_X86=false (default)|true (for cross-compilation) # VLLM_CPU_ARM_BF16=false (default)|true (for cross-compilation) # diff --git a/docs/contributing/README.md b/docs/contributing/README.md index f8e9cfcbb1a5..e7a6448ca63e 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -74,7 +74,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment For JIT kernel warmup conventions, see [JIT Kernel Warmup](./jit_kernel_warmup.md). !!! tip - vLLM is compatible with Python versions 3.10 to 3.13. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. + vLLM is compatible with Python versions 3.11 to 3.14. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. Therefore, we recommend developing with Python 3.12 to minimise the chance of your local environment clashing with our CI environment. @@ -102,14 +102,6 @@ integration, use the configuration files in `tools/pre_commit/`: pre-commit run -a # runs on all files (short for --all-files) ``` - --- - - Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with: - - ```bash - pre-commit run --hook-stage manual mypy-3.11 - ``` - ### Documentation MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, [mkdocs.yaml](../../mkdocs.yaml). diff --git a/docs/deployment/frameworks/skypilot.md b/docs/deployment/frameworks/skypilot.md index e032769f15e8..309367068d6b 100644 --- a/docs/deployment/frameworks/skypilot.md +++ b/docs/deployment/frameworks/skypilot.md @@ -37,7 +37,7 @@ See the vLLM SkyPilot YAML for serving, [serving.yaml](https://github.com/skypil HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -136,7 +136,7 @@ SkyPilot can scale up the service to multiple service replicas with built-in aut HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -248,7 +248,7 @@ This will scale the service up to when the QPS exceeds 2 for each replica. HF_TOKEN: # Change to your own huggingface token, or use --env to pass. setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm pip install vllm==0.4.0.post1 @@ -293,7 +293,7 @@ It is also possible to access the Llama-3 service with a separate GUI frontend, cpus: 2 setup: | - conda create -n vllm python=3.10 -y + conda create -n vllm python=3.12 -y conda activate vllm # Install Gradio for web UI. diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index fbd20f334a6c..d27753413823 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -30,7 +30,7 @@ When open a Github issue about the CPU backend, please add `[CPU Backend]` in th ## Requirements -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 === "Intel/AMD x86" diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 91d933dd4e86..fc8d9148c56b 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -25,7 +25,7 @@ vLLM is a Python library that supports the following GPU variants. Select your G ## Requirements - OS: Linux -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 !!! note vLLM does not support Windows natively. To run vLLM on Windows, you can use the Windows Subsystem for Linux (WSL) with a compatible Linux distribution, or use some community-maintained forks, e.g. [https://github.com/SystemPanic/vllm-windows](https://github.com/SystemPanic/vllm-windows). diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 6f17094a9f47..afa9713341b2 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -8,7 +8,7 @@ This guide will help you quickly get started with vLLM to perform: ## Prerequisites - OS: Linux -- Python: 3.10 -- 3.13 +- Python: 3.11 -- 3.14 !!! note vLLM also works on macOS with [vLLM-Metal](https://github.com/vllm-project/vllm-metal) for Apple Silicon GPU acceleration. See the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab. diff --git a/examples/features/structured_outputs/pyproject.toml b/examples/features/structured_outputs/pyproject.toml index 5e366ab0a03d..d5ffe3c8d462 100644 --- a/examples/features/structured_outputs/pyproject.toml +++ b/examples/features/structured_outputs/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "examples-online-structured-outputs" -requires-python = ">=3.10, <3.14" +requires-python = ">=3.11, <3.15" dependencies = ["openai==1.78.1", "pydantic==2.11.4"] version = "0.0.0" diff --git a/pyproject.toml b/pyproject.toml index 53ce6dfbe06d..bc924440fa59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ license-files = ["LICENSE"] readme = "README.md" description = "A high-throughput and memory-efficient inference and serving engine for LLMs" classifiers = [ - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -32,7 +31,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Information Analysis", ] -requires-python = ">=3.10,<3.15" +requires-python = ">=3.11,<3.15" dynamic = [ "version", "dependencies", "optional-dependencies"] [project.urls] @@ -103,6 +102,8 @@ ignore = [ "B007", # f-string format "UP032", + # `StrEnum` changes `str()` and `format()` output of members + "UP042", ] [tool.ruff.lint.pydocstyle] diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 191ed26ae047..0570517e74cb 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -22,7 +22,6 @@ torchcodec >= 0.14; platform_machine != "riscv64" and platform_machine != "ppc64 intel-openmp==2024.2.1; platform_machine == "x86_64" # triton-cpu, only for x86_64 CPUs -triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.10" triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.11" triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version >= "3.12" diff --git a/setup.py b/setup.py index 0ac53a45b5b1..5e49f74e6fcd 100644 --- a/setup.py +++ b/setup.py @@ -268,8 +268,8 @@ def configure(self, ext: CMakeExtension) -> None: cfg = envs.CMAKE_BUILD_TYPE or default_cfg cmake_args = [ - "-DCMAKE_BUILD_TYPE={}".format(cfg), - "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE), + f"-DCMAKE_BUILD_TYPE={cfg}", + f"-DVLLM_TARGET_DEVICE={VLLM_TARGET_DEVICE}", ] verbose = envs.VERBOSE @@ -293,7 +293,7 @@ def configure(self, ext: CMakeExtension) -> None: # Pass the python executable to cmake so it can find an exact # match. - cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)] + cmake_args += [f"-DVLLM_PYTHON_EXECUTABLE={sys.executable}"] # Pass the python path to cmake so it can reuse the build dependencies # on subsequent calls to python. @@ -305,7 +305,7 @@ def configure(self, ext: CMakeExtension) -> None: # To override this, set the FETCHCONTENT_BASE_DIR environment variable. fc_base_dir = os.path.join(ROOT_DIR, ".deps") fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir) - cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)] + cmake_args += [f"-DFETCHCONTENT_BASE_DIR={fc_base_dir}"] # # Setup parallelism and build tool @@ -313,13 +313,13 @@ def configure(self, ext: CMakeExtension) -> None: num_jobs, nvcc_threads = self.compute_num_jobs() if nvcc_threads: - cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] + cmake_args += [f"-DNVCC_THREADS={nvcc_threads}"] if is_ninja_available(): build_tool = ["-G", "Ninja"] cmake_args += [ "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", - "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs), + f"-DCMAKE_JOB_POOLS:STRING=compile={num_jobs}", ] else: # Default build tool to whatever cmake picks. @@ -1374,7 +1374,7 @@ def _read_requirements(filename: str) -> list[str]: # copying the relevant .py files from the source repository. ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True)) -if not _is_xpu() and sys.version_info >= (3, 11): +if not _is_xpu(): ext_modules.append(CMakeExtension(name="vllm.spinloop")) ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) diff --git a/tests/entrypoints/launchers/test_dp_supervisor.py b/tests/entrypoints/launchers/test_dp_supervisor.py index c10f7ee69899..174157f12a5d 100644 --- a/tests/entrypoints/launchers/test_dp_supervisor.py +++ b/tests/entrypoints/launchers/test_dp_supervisor.py @@ -630,7 +630,7 @@ async def _poll_until_api_server_running( timeout=request_timeout, ): return - except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + except (TimeoutError, aiohttp.ClientError) as exc: last_exc = exc print("Test detected not started yet, sleeping for 1s") remaining_s = deadline - time.monotonic() diff --git a/tests/entrypoints/launchers/test_shutdown.py b/tests/entrypoints/launchers/test_shutdown.py index f03afbff09d9..2a78b4e80ed6 100644 --- a/tests/entrypoints/launchers/test_shutdown.py +++ b/tests/entrypoints/launchers/test_shutdown.py @@ -285,7 +285,7 @@ async def test_wait_timeout_completes_requests(): sigterm_sent.set() await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True @@ -463,7 +463,7 @@ async def test_abort_timeout_fails_inflight_requests(): try: await asyncio.wait_for(request_task, timeout=5) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True @@ -592,7 +592,7 @@ async def test_multi_api_server_shutdown(): try: await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT) - except asyncio.TimeoutError: + except TimeoutError: pass finally: state.stop_requesting = True diff --git a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py index 070c8329674b..c19765c462c4 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py +++ b/tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py @@ -579,7 +579,7 @@ async def test_serialize_and_serve_entrypoints(tmp_path): try: await asyncio.wait_for(fut, 180) - except asyncio.TimeoutError: + except TimeoutError: pytest.fail("Server did not start successfully") finally: proc.terminate() diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index 5949c66b562a..f688b51534c8 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio import mimetypes import os import shutil @@ -292,7 +291,7 @@ async def test_fetch_video_http(video_url: str, num_frames: int): try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) @@ -324,7 +323,7 @@ async def test_fetch_video_http_with_dynamic_loader( try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) @@ -402,7 +401,7 @@ async def test_allowed_media_domains(video_url: str, num_frames: int): try: video_sync, metadata_sync = connector.fetch_video(video_url) video_async, metadata_async = await connector.fetch_video_async(video_url) - except (TimeoutError, asyncio.TimeoutError) as e: + except TimeoutError as e: pytest.skip(f"Timeout fetching video (CI network flakiness): {e}") assert np.array_equal(video_sync, video_async) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index ecdef1f38b0b..fb17ccb1d0fc 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -719,7 +719,7 @@ async def consume_gen(req_id: str) -> None: # Drive the staggered barrier. Old code deadlocks here. try: await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30) - except asyncio.TimeoutError: + except TimeoutError: for t in mid_barrier_tasks: t.cancel() pytest.fail( diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index 2a1e8964d055..c88f2d9ad551 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -255,7 +255,7 @@ def main() -> int: or os.environ.get("VLLM_BUILD_COMMIT") or "", "cuda": cuda_release(nvcc), - "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), + "generated_at": dt.datetime.now(dt.UTC).isoformat(timespec="seconds"), "build_dirs": [], "incomplete": False, "errors": [], diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 126d20094a2a..531f6fb8a2b4 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -88,7 +88,11 @@ class ForbiddenImport: allowed_pattern=re.compile(r"^\s*import\s+pybase64(\s*|\s+as\s+base64\s*)$"), ), "re": ForbiddenImport( - pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)", + # regex has no equivalent of the private sre modules + pattern=( + r"^\s*(?:import\s+re(?:$|\s|,)" + r"|from\s+re\s+import(?!(?:\s*,?\s*_(?:constants|parser))+\s*(?:#.*)?$))" + ), tip="Replace 'import re' with 'import regex as re' or 'import regex'.", allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"), allowed_files={"setup.py"}, diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 245d0daf5e18..876e19e1fdb9 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -10,7 +10,7 @@ python tools/pre_commit/mypy.py Args: - python_version: Python version to use (e.g., "3.10") or "local" to use + python_version: Python version to use (e.g., "3.11") or "local" to use the local Python version. changed_files: List of changed files to check. @@ -152,7 +152,7 @@ def mypy( Args: targets: List of files or directories to check. - python_version: Python version to use (e.g., "3.10") or None to use + python_version: Python version to use (e.g., "3.11") or None to use the default mypy version. follow_imports: Value for the --follow-imports option or None to use the default mypy behavior. diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 9e3c0b57e45f..895e5c6c1890 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -253,7 +253,7 @@ async def fetch_spec_decode_metrics( num_accepted_tokens=num_accepted_tokens, accepted_per_pos=accepted_per_pos, ) - except (aiohttp.ClientError, asyncio.TimeoutError): + except (TimeoutError, aiohttp.ClientError): return None @@ -315,7 +315,7 @@ async def fetch_diffusion_metrics( num_canvas_positions=num_canvas_positions, num_committed_tokens=num_committed_tokens, ) - except (aiohttp.ClientError, asyncio.TimeoutError): + except (TimeoutError, aiohttp.ClientError): return None diff --git a/vllm/benchmarks/sweep/plot.py b/vllm/benchmarks/sweep/plot.py index 59a90636f0e0..6b36a7acc601 100644 --- a/vllm/benchmarks/sweep/plot.py +++ b/vllm/benchmarks/sweep/plot.py @@ -8,9 +8,9 @@ from functools import partial from pathlib import Path from types import TracebackType -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self -from typing_extensions import Self, override +from typing_extensions import override from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.collection_utils import full_groupby diff --git a/vllm/benchmarks/sweep/serve_workload.py b/vllm/benchmarks/sweep/serve_workload.py index a488837e2e6f..acf0ce5f7f3d 100644 --- a/vllm/benchmarks/sweep/serve_workload.py +++ b/vllm/benchmarks/sweep/serve_workload.py @@ -4,10 +4,9 @@ import math from dataclasses import dataclass, fields from pathlib import Path -from typing import ClassVar, Literal, get_args +from typing import ClassVar, Literal, assert_never, get_args import numpy as np -from typing_extensions import assert_never from vllm.benchmarks.datasets import DEFAULT_NUM_PROMPTS from vllm.utils.argparse_utils import FlexibleArgumentParser diff --git a/vllm/benchmarks/sweep/server.py b/vllm/benchmarks/sweep/server.py index 87d841ac827f..dff7e069242b 100644 --- a/vllm/benchmarks/sweep/server.py +++ b/vllm/benchmarks/sweep/server.py @@ -6,9 +6,9 @@ import subprocess import time from types import TracebackType +from typing import Self import requests -from typing_extensions import Self class ServerProcess: diff --git a/vllm/config/engram.py b/vllm/config/engram.py index fd95bc2acd2b..7a7ed8f277ef 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from pydantic import model_validator -from typing_extensions import Self from vllm.config.utils import config, get_hash_factors, hash_factors diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 6c994a82f8de..a77660f3b130 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self import torch from pydantic import ConfigDict, Field, model_validator -from typing_extensions import Self from vllm import envs from vllm.config.utils import config diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 7df8f4e0a4ed..37affbaf9ebf 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -5,13 +5,12 @@ import socket from collections.abc import Callable from datetime import timedelta -from typing import TYPE_CHECKING, Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, Self, overload import regex as re import torch from pydantic import Field, field_validator, model_validator from torch.distributed import ProcessGroup, ReduceOp, Store -from typing_extensions import Self import vllm.envs as envs from vllm.config.fault_tolerance import FaultToleranceConfig diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 7547546d140d..0ad920e21eda 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import Field, model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 1fb94d3efd07..0d10103ed4ec 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -3,10 +3,9 @@ from collections.abc import Callable from dataclasses import InitVar -from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self, cast from pydantic import Field, field_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cd42206f983a..4c61ead82dab 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -4,10 +4,9 @@ import copy import functools from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Any, Literal, get_args +from typing import TYPE_CHECKING, Any, Literal, Self, get_args from pydantic import Field, SkipValidation, field_validator, model_validator -from typing_extensions import Self from vllm.config import LoadConfig from vllm.config.cache import CacheDType diff --git a/vllm/config/structured_outputs.py b/vllm/config/structured_outputs.py index 059581e37e65..4a13900b5622 100644 --- a/vllm/config/structured_outputs.py +++ b/vllm/config/structured_outputs.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.utils.hashing import safe_hash diff --git a/vllm/config/watermarking.py b/vllm/config/watermarking.py index 8701c7b72e4d..4d559c29ce64 100644 --- a/vllm/config/watermarking.py +++ b/vllm/config/watermarking.py @@ -2,10 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import hashlib -from typing import Literal +from typing import Literal, Self from pydantic import Field, model_validator -from typing_extensions import Self from vllm.config.utils import config from vllm.logger import init_logger diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index c780316bb4e1..3618361f2a66 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -38,9 +38,7 @@ # We prefer to use os.sched_yield as it results in tighter polling loops, # measured to be around 3e-7 seconds. However on earlier versions of Python # os.sched_yield() does not release the GIL, so we fall back to time.sleep(0) -USE_SCHED_YIELD = (sys.version_info[:3] >= (3, 11, 1)) or ( - sys.version_info[:2] == (3, 10) and sys.version_info[2] >= 8 -) +USE_SCHED_YIELD = sys.version_info[:3] >= (3, 11, 1) def sched_yield(): @@ -52,8 +50,8 @@ def sched_yield(): def ensure_divisibility(numerator, denominator): """Ensure that numerator is divisible by the denominator.""" - assert numerator % denominator == 0, "{} is not divisible by {}".format( - numerator, denominator + assert numerator % denominator == 0, ( + f"{numerator} is not divisible by {denominator}" ) @@ -146,9 +144,7 @@ def get_pp_indices( try: partitions = [int(layer) for layer in partition_list_str.split(",")] except ValueError as err: - raise ValueError( - "Invalid partition string: {}".format(partition_list_str) - ) from err + raise ValueError(f"Invalid partition string: {partition_list_str}") from err if len(partitions) != pp_size: raise ValueError(f"{len(partitions)=} does not match {pp_size=}.") if sum(partitions) != num_hidden_layers: diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2fb60fb71b41..3140b70e9656 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -11,12 +11,12 @@ ClassVar, Generic, Protocol, + Self, TypeVar, runtime_checkable, ) import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 15864811e96c..13da117b99e2 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -4,12 +4,11 @@ import pickle from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Self import pybase64 as base64 import torch from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor -from typing_extensions import Self from vllm import envs from vllm.config.weight_transfer import WeightTransferConfig diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 68c62f72dbac..2d1d9f47ec02 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -15,10 +15,9 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Self, cast import torch -from typing_extensions import Self from vllm.distributed.weight_transfer.base import ( TrainerInitInfo, diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 5b98845870fd..77d8b608727e 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -5,10 +5,9 @@ from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 6b755e772e18..51ea401ef92d 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -24,7 +24,7 @@ from collections.abc import Callable, Collection from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass, field -from typing import Any, ClassVar +from typing import Any, ClassVar, Self import ray import torch @@ -33,7 +33,6 @@ rebuild_cuda_tensor, reduce_tensor, ) -from typing_extensions import Self from vllm.distributed.weight_transfer.base import ( ParamMeta, diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 9c704ab344ea..01b29529deb7 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -11,10 +11,9 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, ClassVar, Self import torch -from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 2d875a974bef..f18407e535e9 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -17,6 +17,7 @@ Final, Generic, Literal, + Required, TypeAlias, TypeVar, Union, @@ -48,7 +49,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter # pydantic needs the TypedDict from typing_extensions -from typing_extensions import Required, TypedDict, override +from typing_extensions import TypedDict, override from vllm import envs from vllm.config import ModelConfig diff --git a/vllm/entrypoints/launchers/dp_supervisor.py b/vllm/entrypoints/launchers/dp_supervisor.py index 192bbb169814..f2a84c4d9dfd 100644 --- a/vllm/entrypoints/launchers/dp_supervisor.py +++ b/vllm/entrypoints/launchers/dp_supervisor.py @@ -200,7 +200,7 @@ async def _probe_endpoint( # vLLM returns 503 on EngineDeadError, so we should return # immediately if vLLM responds with a non-200 status code. return response.status == HTTPStatus.OK - except (aiohttp.ClientError, asyncio.TimeoutError) as e: + except (TimeoutError, aiohttp.ClientError) as e: # Allow retry of connection errors. logger.debug( "Probe attempt %d/%d failed on port %d: %r", diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index c38789a2e29c..9dfff78ed9ff 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -3,13 +3,12 @@ from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, overload import cloudpickle import torch.nn as nn from pydantic import ValidationError from tqdm.auto import tqdm -from typing_extensions import overload from vllm.config import ( AttentionConfig, diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index a9f6dd310ec7..4fdb73d363df 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TypeAlias, cast +from typing import TypeAlias, assert_never, cast from fastapi.responses import JSONResponse, Response, StreamingResponse -from typing_extensions import assert_never from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 56bc7c7865b1..4e4ea1fe1bd5 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import cast +from typing import assert_never, cast from fastapi.responses import JSONResponse, Response, StreamingResponse -from typing_extensions import assert_never from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py index a67843e57770..9afaada6d666 100644 --- a/vllm/entrypoints/pooling/scoring/typing.py +++ b/vllm/entrypoints/pooling/scoring/typing.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass -from typing import TypeAlias +from typing import Required, TypeAlias -from typing_extensions import Required, TypedDict +from typing_extensions import TypedDict from vllm.entrypoints.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, diff --git a/vllm/entrypoints/serve/tokenize/api_router.py b/vllm/entrypoints/serve/tokenize/api_router.py index 9cdd82fee3f3..cd515d31d2e8 100644 --- a/vllm/entrypoints/serve/tokenize/api_router.py +++ b/vllm/entrypoints/serve/tokenize/api_router.py @@ -3,10 +3,10 @@ from http import HTTPStatus +from typing import assert_never from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse -from typing_extensions import assert_never from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.tokenize.protocol import ( diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index 26b840f76f3f..aa8c9054335e 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -3,9 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Literal, TypeAlias +from typing import TYPE_CHECKING, Literal, NotRequired, TypeAlias, assert_never -from typing_extensions import NotRequired, TypedDict, assert_never +from typing_extensions import TypedDict from vllm.exceptions import VLLMValidationError diff --git a/vllm/inputs/llm.py b/vllm/inputs/llm.py index a95b274c1853..1ef43fd23b17 100644 --- a/vllm/inputs/llm.py +++ b/vllm/inputs/llm.py @@ -3,9 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, final +from typing import TYPE_CHECKING, Any, NotRequired, TypeAlias, TypeVar, final -from typing_extensions import NotRequired, TypedDict +from typing_extensions import TypedDict if TYPE_CHECKING: import torch diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py index 0db595a3ca42..af1f597ec60c 100644 --- a/vllm/model_executor/kernels/linear/base.py +++ b/vllm/model_executor/kernels/linear/base.py @@ -3,10 +3,9 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any, ClassVar, Generic, TypeVar +from typing import Any, ClassVar, Generic, Self, TypeVar import torch -from typing_extensions import Self from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey diff --git a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py index 3a053412f3cf..865bebeada54 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py @@ -3,10 +3,9 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import ClassVar +from typing import ClassVar, Self import torch -from typing_extensions import Self from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.fp8_utils import ( diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index c12b8e76c5f5..65d1ff2808d7 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -6,11 +6,10 @@ import time import warnings from contextlib import contextmanager, nullcontext -from typing import Any +from typing import Any, assert_never import torch from torch import nn -from typing_extensions import assert_never import vllm.envs as envs from vllm.config import ( diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 433c373c8e35..272b78da50a4 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -21,6 +21,7 @@ ClassVar, Literal, Protocol, + Self, TypeAlias, overload, runtime_checkable, @@ -31,7 +32,7 @@ import torch.nn as nn from torch import Tensor from transformers.models.whisper.tokenization_whisper import LANGUAGES -from typing_extensions import Self, TypeIs +from typing_extensions import TypeIs from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index 62a5c2070f5b..9c1a4664c83d 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -22,13 +22,12 @@ import math from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, assert_never import numpy as np import torch from PIL import Image, ImageOps from transformers import BatchFeature -from typing_extensions import assert_never from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 98f8a8b2f3e9..546eb6abfda2 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -13,11 +13,11 @@ TypeAlias, TypeGuard, TypeVar, + assert_never, ) import numpy as np import torch -from typing_extensions import assert_never from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalUUIDDict from vllm.utils.collection_utils import is_list_of diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index 726702035355..abdb58540255 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -13,10 +13,11 @@ NamedTuple, Protocol, TypeAlias, + assert_never, ) import torch -from typing_extensions import TypeVar, assert_never +from typing_extensions import TypeVar from vllm.config import SchedulerConfig from vllm.inputs import ( diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 4301c6917613..4a8bcd0c03d7 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1000,7 +1000,7 @@ async def _render_with_timeout(self, coro): return await coro try: return await asyncio.wait_for(coro, timeout=timeout) - except asyncio.TimeoutError: + except TimeoutError: self._replace_executor() raise TimeoutError( f"Chat template rendering timed out after " diff --git a/vllm/snapshot/manifest.py b/vllm/snapshot/manifest.py index a029c544c52e..c550ee80c5cf 100644 --- a/vllm/snapshot/manifest.py +++ b/vllm/snapshot/manifest.py @@ -6,7 +6,7 @@ import os import stat from pathlib import Path -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, Self from pydantic import ( BaseModel, @@ -19,7 +19,6 @@ field_validator, model_validator, ) -from typing_extensions import Self class SnapshotCompatibilityError(RuntimeError): diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py index fdced5116cdf..83c70e473833 100644 --- a/vllm/snapshot/runtime.py +++ b/vllm/snapshot/runtime.py @@ -24,7 +24,7 @@ from collections.abc import Iterable from contextlib import suppress from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import UTC, datetime from pathlib import Path from typing import NoReturn @@ -872,7 +872,7 @@ def make_manifest( return SnapshotManifest( schema_version=1, boundary="post-engine-init-reloadable-state-released", - created_at=datetime.now(timezone.utc).isoformat(), + created_at=datetime.now(UTC).isoformat(), artifact_bytes=self._artifact_bytes(workdir), **identity.model_dump(), model=source_model, diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index e7064fc252f2..ebf85a11fd25 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -4,10 +4,10 @@ from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, assert_never import huggingface_hub -from typing_extensions import TypeVar, assert_never +from typing_extensions import TypeVar from vllm import envs from vllm.logger import init_logger diff --git a/vllm/usage/usage_lib.py b/vllm/usage/usage_lib.py index 5197aada59ac..b75a4fa16184 100644 --- a/vllm/usage/usage_lib.py +++ b/vllm/usage/usage_lib.py @@ -69,7 +69,7 @@ def is_usage_stats_enabled(): def _get_current_timestamp_ns() -> int: - return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e9) + return int(datetime.datetime.now(datetime.UTC).timestamp() * 1e9) def _detect_cloud_provider() -> str: diff --git a/vllm/utils/collection_utils.py b/vllm/utils/collection_utils.py index de32ed0c15c2..8a77da86cd56 100644 --- a/vllm/utils/collection_utils.py +++ b/vllm/utils/collection_utils.py @@ -8,9 +8,9 @@ import math from collections import defaultdict from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence -from typing import Generic, Literal, TypeVar +from typing import Generic, Literal, TypeVar, assert_never, overload -from typing_extensions import TypeIs, assert_never, overload +from typing_extensions import TypeIs T = TypeVar("T") diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index aaacccc7cb5a..04a7ebec2eab 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -12,10 +12,9 @@ import sys from functools import cache from types import ModuleType -from typing import Any +from typing import Any, Never import regex as re -from typing_extensions import Never from vllm.logger import init_logger diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b6913003be59..cf707611b47d 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -11,10 +11,9 @@ from fractions import Fraction from functools import cached_property from math import prod -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, Self, TypeVar import torch -from typing_extensions import Self from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 1cd8fa79f1d3..a2b3f534d88c 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -10,12 +10,11 @@ from abc import ABC, abstractmethod from collections.abc import Iterable from dataclasses import asdict, dataclass, field -from typing import Any, Protocol +from typing import Any, Protocol, Self import prometheus_client import torch from pydantic import BaseModel, Field, ValidationError, model_validator -from typing_extensions import Self import vllm.envs as envs from vllm.config import VllmConfig diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index dab4f2aa0462..655b8b14631f 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -4,10 +4,9 @@ from __future__ import annotations import ast -import importlib import json -import sys from dataclasses import dataclass, field +from re import _constants, _parser # type: ignore[attr-defined] from typing import TYPE_CHECKING import torch @@ -36,19 +35,6 @@ oc = LazyLoader("oc", globals(), "outlines_core") json_schema = LazyLoader("json_schema", globals(), "outlines_core.json_schema") -# Python 3.11+ sre_parse and sre_constants -# are deprecated, so we must import them from re -if sys.version_info >= (3, 11): - # Hack to get around pre-commit regex module rule - # because going through re is the only way to get sre_parse - # and sre_constants in Python 3.11+ - _re = importlib.import_module("re") - sre_parse = _re._parser - sre_constants = _re._constants -else: - import sre_constants - import sre_parse - @dataclass class OutlinesBackend(StructuredOutputBackend): @@ -261,20 +247,20 @@ def subpattern_consumes(parsed) -> bool: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # literal, character class, or dot always consumes - if ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): + if ttype in (_parser.LITERAL, _parser.IN, _parser.ANY): return True # quantified subpattern: check inner pattern - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: _, mx, sub = tval if mx != 0 and subpattern_consumes(sub): return True # alternation: if any branch consumes, the whole does - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: _, branches = tval if any(subpattern_consumes(br) for br in branches): return True # grouped subpattern: recurse into its contents - elif ttype == sre_parse.SUBPATTERN and subpattern_consumes(tval[3]): + elif ttype == _parser.SUBPATTERN and subpattern_consumes(tval[3]): return True # No consumers, return False return False @@ -282,14 +268,11 @@ def subpattern_consumes(parsed) -> bool: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # Direct anchors or look-around - if ttype == sre_parse.AT or ttype in ( - sre_constants.ASSERT, - sre_constants.ASSERT_NOT, - ): + if ttype == _parser.AT or ttype in (_constants.ASSERT, _constants.ASSERT_NOT): return True # Nested subpattern: check - if ttype == sre_parse.SUBPATTERN: + if ttype == _parser.SUBPATTERN: # tval: (group, add_flags, del_flags, subpattern) if _prefix_needs_context(tval[3]): return True @@ -298,7 +281,7 @@ def subpattern_consumes(parsed) -> bool: # if any branch has a prefix anchor => True, # else if at least one branch consumes => prefix ends => False - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: saw_consumer = False for br in tval[1]: if _prefix_needs_context(br): @@ -309,11 +292,11 @@ def subpattern_consumes(parsed) -> bool: return False # Immediate consumer tokens - elif ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): + elif ttype in (_parser.LITERAL, _parser.IN, _parser.ANY): return False # if subpattern has anchor => True, if it can consume => stop - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: if _prefix_needs_context(tval[2]): return True if subpattern_consumes(tval[2]): @@ -327,25 +310,25 @@ def _check_unsupported(parsed) -> None: tokens = parsed.data if hasattr(parsed, "data") else parsed for ttype, tval in tokens: # backreference - if ttype in (sre_parse.GROUPREF, sre_parse.GROUPREF_EXISTS): + if ttype in (_parser.GROUPREF, _parser.GROUPREF_EXISTS): raise ValueError("Backreferences are unsupported.") # look-around assertion - elif ttype in (sre_constants.ASSERT, sre_constants.ASSERT_NOT): + elif ttype in (_constants.ASSERT, _constants.ASSERT_NOT): raise ValueError("Look-Around assertion are unsupported.") # unicode word boundaries - elif ttype == sre_parse.AT: - if tval in (sre_constants.AT_BOUNDARY, sre_constants.AT_NON_BOUNDARY): + elif ttype == _parser.AT: + if tval in (_constants.AT_BOUNDARY, _constants.AT_NON_BOUNDARY): raise ValueError("Unicode word boundaries are unsupported.") - elif ttype == sre_parse.BRANCH: + elif ttype == _parser.BRANCH: # tval is (None, branches) for branch in tval[1]: _check_unsupported(branch) # tval is (min, max, subpattern) - elif ttype == sre_parse.MAX_REPEAT: + elif ttype == _parser.MAX_REPEAT: _check_unsupported(tval[2]) @@ -357,9 +340,9 @@ def validate_regex_is_buildable(pattern: str) -> None: https://docs.rs/regex-automata/latest/regex_automata/dfa/trait.Automaton.html#method.universal_start_state """ try: - parsed = sre_parse.parse(pattern) + parsed = _parser.parse(pattern) - except sre_constants.error as e: + except _constants.error as e: raise VLLMValidationError(f"Error parsing regex: {e}") from e try: diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 4bae1f52bec6..e69e55bd8216 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -229,12 +229,7 @@ def get_caller_info() -> str: # Walk up the stack skipping WorkspaceManager frames curr_frame = curr_frame.f_back while curr_frame is not None: - # TODO: This only catches instance methods (self), missing - # classmethods and staticmethods. Once Python 3.11+ is the - # minimum supported version, use co_qualname instead: - # qualname = curr_frame.f_code.co_qualname - # if qualname.startswith("WorkspaceManager."): - if isinstance(curr_frame.f_locals.get("self"), WorkspaceManager): + if curr_frame.f_code.co_qualname.startswith("WorkspaceManager."): curr_frame = curr_frame.f_back continue filename = os.path.basename(curr_frame.f_code.co_filename) From c4cd88d9fee92ff080837c6c1cdf1de6fe4013fe Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 8 Oct 2026 02:31:13 -0700 Subject: [PATCH 193/952] [CI] Make (H100) Distributed DP + EP optional (#60540) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 3710f80bfdb8..78a6af2e0ca4 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -397,6 +397,7 @@ steps: - label: ":nvidia: (H100) Distributed DP + EP" key: distributed-tests-8xh100 timeout_in_minutes: 20 + optional: true device: h100 num_devices: 8 working_dir: "/vllm-workspace/tests" @@ -419,6 +420,7 @@ steps: amd: label: ":amd: (MI355) Distributed DP + EP" dind: false + optional: false device: mi355_8 timeout_in_minutes: 50 depends_on: From 7816afdb259dbb18f9e98d11b05194dd3446d0fa Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 11:33:41 +0200 Subject: [PATCH 194/952] [Bugfix][Core] Drain async KV loads in pause(mode="wait"), and pause the drain test only once a request is in flight (#60572) Signed-off-by: aoshen02 Signed-off-by: Stefan Koncarevic Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/v1/core/test_scheduler.py | 90 +++++++++++++++++++++++ tests/v1/distributed/test_async_llm_dp.py | 7 +- vllm/v1/core/sched/interface.py | 2 +- vllm/v1/core/sched/scheduler.py | 12 ++- 4 files changed, 103 insertions(+), 8 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7a8a1a3e4ea5..615dd82a1c7b 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2573,6 +2573,96 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): return initial_ecos +def _runner_output( + req_ids: list[str], finished_recving: set[str] | None = None +) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[1] for _ in req_ids], + kv_connector_output=KVConnectorOutput(finished_recving=finished_recving), + ) + + +@pytest.mark.parametrize( + "load_done_before_pause", [False, True], ids=["load-in-flight", "load-done"] +) +def test_pause_new_drains_async_kv_loads(load_done_before_pause: bool): + """A PAUSED_NEW (pause mode="wait") drain must finish requests that hold KV + blocks for an async load, keep blockless queued requests queued, and leave + no blocks behind, so the post-drain cache reset succeeds.""" + block_size = 16 + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv(matched_tokens=2 * block_size, is_async=True), + block_size=block_size, + ) + loading, queued = create_requests( + num_requests=2, + num_tokens=3 * block_size, + max_tokens=2, + block_size=block_size, + req_ids=["loading", "queued"], + ) + scheduler.add_request(loading) + output = scheduler.schedule() + assert loading.status == RequestStatus.WAITING_FOR_REMOTE_KVS + pending_recv = {"loading"} + if load_done_before_pause: + scheduler.update_from_output(output, _runner_output([], pending_recv)) + pending_recv = None + else: + scheduler.update_from_output(output, _runner_output([])) + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + scheduler.add_request(queued) + + num_steps = 0 + while scheduler.has_requests(): + num_steps += 1 + assert num_steps <= 4, "pause drain did not finish" + output = scheduler.schedule() + assert set(output.num_scheduled_tokens) <= {"loading"} + scheduler.update_from_output( + output, _runner_output(list(output.num_scheduled_tokens), pending_recv) + ) + pending_recv = None + + # What pause runs once drained; it raises if the load still holds blocks. + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert loading.status == RequestStatus.FINISHED_LENGTH_CAPPED + assert loading.request_id not in scheduler.requests + assert list(scheduler.waiting) == [queued] + assert queued.status == RequestStatus.WAITING + + scheduler.set_pause_state(PauseState.UNPAUSED) + scheduler.schedule() + assert queued.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_pause_new_ignores_streaming_session_waiting_for_input(): + """A session idle in WAITING_FOR_STREAMING_REQ cannot progress, so it must + not keep a PAUSED_NEW drain from completing.""" + stop_token = 7 + scheduler = create_scheduler() + (session,) = create_requests( + num_requests=1, num_tokens=4, stop_token_ids=[stop_token], req_ids=["s"] + ) + session.resumable = True + scheduler.add_request(session) + output = scheduler.schedule() + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=["s"], req_id_to_index={"s": 0}, sampled_token_ids=[[stop_token]] + ), + ) + assert session in scheduler.kv_holding_waiting + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + assert scheduler.get_num_unfinished_requests() == 0 + + @pytest.mark.parametrize( ("load_modes", "expected_has_sync_loads"), [ diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index fb17ccb1d0fc..f4b4ca2a7902 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -747,12 +747,13 @@ async def test_dp_pause_wait_mode_drains_in_flight(): prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=64), ) + # Wait for the first output, so drain has something to observe: a + # request still queued holds no blocks and is not waited on. + out = await asyncio.wait_for(collector.get(), timeout=30) await engine.pause_generation(mode="wait") assert await engine.is_paused() - while True: + while not out.finished: out = await asyncio.wait_for(collector.get(), timeout=30) - if out.finished: - break await engine.resume_generation() async for out in engine.generate( diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index c868d249c8a5..dd88d99c1f79 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -26,7 +26,7 @@ class PauseState(enum.IntEnum): - UNPAUSED: Normal operation - PAUSE_NEW: No new requests are scheduled, requests already in - running state are scheduled. + running state or holding KV blocks are scheduled. - PAUSE_ALL: No requests are scheduled """ diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 91ea8e93e42c..9739dbb3d516 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -887,7 +887,7 @@ def schedule(self, throttle_prefills: bool = False) -> 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: + if not preempted_reqs and self._pause_state != PauseState.PAUSED_ALL: step_skipped_waiting: deque[Request] = deque() step_skipped_kv_holding: deque[Request] = deque() @@ -900,8 +900,11 @@ def skip_request(from_queue: RequestQueue) -> None: step_skipped_waiting.appendleft(request_to_skip) while token_budget > 0: - # Requests holding KV blocks are always drained first. - request_queue = self.kv_holding_waiting or self.waiting + # Requests holding KV blocks are always drained first, and are + # the only ones admitted under PAUSED_NEW (e.g. async KV loads). + request_queue = self.kv_holding_waiting + if not request_queue and self._pause_state == PauseState.UNPAUSED: + request_queue = self.waiting if not request_queue or input_budget <= draft_slots: break # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not @@ -2802,7 +2805,8 @@ def get_num_unfinished_requests(self) -> int: return 0 num_running, num_waiting = self.get_request_counts() if self._pause_state == PauseState.PAUSED_NEW: - return num_running + # Requests holding KV blocks still drain; queued ones stay queued. + num_waiting = len(self.kv_holding_waiting) num_waiting -= self.num_waiting_for_streaming_input return num_waiting + num_running From e265b77982e5e2724ce0b2841a2b05fe339d62ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Thu, 8 Oct 2026 17:39:56 +0800 Subject: [PATCH 195/952] [CI][Mamba] Use well-conditioned prompts for FlashInfer ReplaySSM parity tests (#60515) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 --- tests/v1/e2e/test_replayssm_decode.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 1a369c781f03..ed3c6b0217c8 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -25,6 +25,12 @@ "The capital of France is", "Once upon a time, in a small village,", ] +# The NVFP4 model is not reproducible across engine instances; these prompts +# avoid near-ties and decode confidently past the replay window. +FLASHINFER_PROMPTS = [ + "The capital of France is", + "It was the best of times, it was the worst of times,", +] @pytest.fixture(autouse=True) @@ -43,6 +49,7 @@ def _check_replayssm_parity( *, tensor_parallel_size=1, mamba_backend: str = "triton", + prompts: list[str] = PROMPTS, name_1: str = "replayssm", expected_v2: bool | None = None, ): @@ -60,7 +67,7 @@ def _check_replayssm_parity( with vllm_runner(model_name, **common) as llm: if expected_v2 is not None: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 - baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + baseline = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5) with vllm_runner( model_name, use_replayssm=True, @@ -69,7 +76,7 @@ def _check_replayssm_parity( ) as llm: if expected_v2 is not None: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 - replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + replay = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5) check_logprobs_close( outputs_0_lst=baseline, @@ -105,6 +112,7 @@ def test_replayssm_flashinfer_decode_matches_baseline( vllm_runner, model_name, mamba_backend="flashinfer", + prompts=FLASHINFER_PROMPTS, name_1="replayssm_flashinfer", expected_v2=use_v2_model_runner, ) @@ -127,7 +135,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na }, ) with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: - baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + baseline = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) with vllm_runner( model_name, use_replayssm=True, @@ -135,7 +145,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na mamba_backend="flashinfer", **common, ) as llm: - replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + replay = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) check_logprobs_close( outputs_0_lst=baseline, From 24be8c0bca2fd806548ad84c0999885a65f0805d Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Thu, 8 Oct 2026 02:42:07 -0700 Subject: [PATCH 196/952] [CI] Remove dead in-tree TPU (torch_xla) CI scripts (#60591) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .../scripts/hardware_ci/run-cpu-test.sh | 7 +- .../hardware_ci/run-tpu-v1-test-part2.sh | 166 ----------------- .../scripts/hardware_ci/run-tpu-v1-test.sh | 172 ------------------ .buildkite/scripts/tpu/cleanup_docker.sh | 24 --- .buildkite/scripts/tpu/config_v6e_1.env | 14 -- .buildkite/scripts/tpu/docker_run_bm.sh | 91 --------- .buildkite/scripts/tpu/quantized_v6e_1.env | 14 -- .buildkite/scripts/tpu/run_bm.sh | 93 ---------- .github/mergify.yml | 4 - .../run_tpu_disagg_accuracy_test.sh | 157 ---------------- .../run_tpu_edge_case_test.sh | 125 ------------- 11 files changed, 3 insertions(+), 864 deletions(-) delete mode 100755 .buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh delete mode 100755 .buildkite/scripts/hardware_ci/run-tpu-v1-test.sh delete mode 100755 .buildkite/scripts/tpu/cleanup_docker.sh delete mode 100644 .buildkite/scripts/tpu/config_v6e_1.env delete mode 100755 .buildkite/scripts/tpu/docker_run_bm.sh delete mode 100644 .buildkite/scripts/tpu/quantized_v6e_1.env delete mode 100755 .buildkite/scripts/tpu/run_bm.sh delete mode 100644 tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh delete mode 100644 tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index eecf70ea2bde..5f9355f9e943 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -19,10 +19,9 @@ DISK_USAGE_THRESHOLD=${DISK_USAGE_THRESHOLD:-80} CACHE_MAX_AGE=${CACHE_MAX_AGE:-24h} # Reclaim disk only when the host is under pressure, aging out anything unused -# for less than CACHE_MAX_AGE so hot layers survive -- same `--filter -# until=h` pattern the TPU CI scripts already rely on. `docker buildx -# prune --max-used-space` is a no-op on this host's `docker` driver (BuildKit -# embedded in dockerd never enforces the size cap), so we don't use it. +# for less than CACHE_MAX_AGE so hot layers survive (`--filter until=h`). +# `docker buildx prune --max-used-space` is a no-op on this host's `docker` +# driver (BuildKit embedded in dockerd never enforces the size cap), so we don't use it. prune_if_disk_pressure() { local docker_root disk_usage docker_root=$(docker info -f '{{.DockerRootDir}}' 2>/dev/null || true) diff --git a/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh b/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh deleted file mode 100755 index f300b7f80991..000000000000 --- a/.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh +++ /dev/null @@ -1,166 +0,0 @@ -#!/bin/bash - -set -xu - - -remove_docker_container() { - docker rm -f tpu-test || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -# Build the docker image. -docker build -f docker/Dockerfile.tpu -t vllm-tpu . - -# Set up cleanup. -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} -cleanup_docker - -# For HF_TOKEN. -source /etc/environment - -docker run --privileged --net host --shm-size=16G -it \ - -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \ - vllm-tpu /bin/bash -c ' -set -e # Exit immediately if a command exits with a non-zero status. -set -u # Treat unset variables as an error. - -echo "--- Starting script inside Docker container ---" - -# Create results directory -RESULTS_DIR=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $RESULTS_DIR" - -# Install dependencies -echo "--- Installing Python dependencies ---" -python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \ - && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \ - && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \ - && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0 -echo "--- Python dependencies installed ---" - -export VLLM_XLA_CHECK_RECOMPILATION=1 -export VLLM_XLA_CACHE_PATH= - -echo "--- Hardware Information ---" -# tpu-info -echo "--- Starting Tests ---" -set +e -overall_script_exit_code=0 - -# --- Test Definitions --- -# If a test fails, this function will print logs and will not cause the main script to exit. -run_test() { - local test_num=$1 - local test_name=$2 - local test_command=$3 - local log_file="$RESULTS_DIR/test_${test_num}.log" - local actual_exit_code - - echo "--- TEST_$test_num: Running $test_name ---" - - # Execute the test command. - eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2) - actual_exit_code=$? - - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log - - if [ "$actual_exit_code" -ne 0 ]; then - echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2 - echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2 - if [ -f "$log_file" ]; then - cat "$log_file" >&2 - else - echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2 - fi - echo "--- End of log for TEST_$test_num ($test_name) ---" >&2 - return "$actual_exit_code" # Return the failure code - else - echo "TEST_$test_num ($test_name) PASSED." - return 0 # Return success - fi -} - -# Helper function to call run_test and update the overall script exit code -run_and_track_test() { - local test_num_arg="$1" - local test_name_arg="$2" - local test_command_arg="$3" - - # Run the test - run_test "$test_num_arg" "$test_name_arg" "$test_command_arg" - local test_specific_exit_code=$? - - # If the test failed, set the overall script exit code to 1 - if [ "$test_specific_exit_code" -ne 0 ]; then - # No need for extra echo here, run_test already logged the failure. - overall_script_exit_code=1 - fi -} - -# --- Actual Test Execution --- -run_and_track_test 1 "test_struct_output_generate.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\"" -run_and_track_test 2 "test_moe_pallas.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py" -run_and_track_test 3 "test_lora.py" \ - "VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py" -run_and_track_test 4 "test_tpu_qkv_linear.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py" -run_and_track_test 5 "test_spmd_model_weight_loading.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py" -run_and_track_test 6 "test_kv_cache_update_kernel.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_kv_cache_update_kernel.py" -run_and_track_test 7 "test_tpu_int8.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_int8.py" - -# After all tests have been attempted, exit with the overall status. -if [ "$overall_script_exit_code" -ne 0 ]; then - echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---" -else - echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---" -fi -exit "$overall_script_exit_code" -' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct. - -# Capture the exit code of the docker run command -DOCKER_RUN_EXIT_CODE=$? - -# The trap will run for cleanup. -# Exit the main script with the Docker run command's exit code. -if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then - echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE." - exit "$DOCKER_RUN_EXIT_CODE" -else - echo "Docker run command completed successfully." - exit 0 -fi -# TODO: This test fails because it uses RANDOM_SEED sampling -# pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \ diff --git a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh b/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh deleted file mode 100755 index 6f4e0d37db7d..000000000000 --- a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh +++ /dev/null @@ -1,172 +0,0 @@ -#!/bin/bash - -set -xu - - -remove_docker_container() { - docker rm -f tpu-test || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -# Build the docker image. -docker build -f docker/Dockerfile.tpu -t vllm-tpu . - -# Set up cleanup. -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} -cleanup_docker - -# For HF_TOKEN. -source /etc/environment - -docker run --privileged --net host --shm-size=16G -it \ - -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \ - vllm-tpu /bin/bash -c ' -set -e # Exit immediately if a command exits with a non-zero status. -set -u # Treat unset variables as an error. - -echo "--- Starting script inside Docker container ---" - -# Create results directory -RESULTS_DIR=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $RESULTS_DIR" - -# Install dependencies -echo "--- Installing Python dependencies ---" -python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \ - && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \ - && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \ - && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0 -echo "--- Python dependencies installed ---" - -export VLLM_XLA_CHECK_RECOMPILATION=1 -export VLLM_XLA_CACHE_PATH= - -echo "--- Hardware Information ---" -# tpu-info -echo "--- Starting Tests ---" -set +e -overall_script_exit_code=0 - -# --- Test Definitions --- -# If a test fails, this function will print logs and will not cause the main script to exit. -run_test() { - local test_num=$1 - local test_name=$2 - local test_command=$3 - local log_file="$RESULTS_DIR/test_${test_num}.log" - local actual_exit_code - - echo "--- TEST_$test_num: Running $test_name ---" - - # Execute the test command. - eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2) - actual_exit_code=$? - - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log - echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log - - if [ "$actual_exit_code" -ne 0 ]; then - echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2 - echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2 - if [ -f "$log_file" ]; then - cat "$log_file" >&2 - else - echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2 - fi - echo "--- End of log for TEST_$test_num ($test_name) ---" >&2 - return "$actual_exit_code" # Return the failure code - else - echo "TEST_$test_num ($test_name) PASSED." - return 0 # Return success - fi -} - -# Helper function to call run_test and update the overall script exit code -run_and_track_test() { - local test_num_arg="$1" - local test_name_arg="$2" - local test_command_arg="$3" - - # Run the test - run_test "$test_num_arg" "$test_name_arg" "$test_command_arg" - local test_specific_exit_code=$? - - # If the test failed, set the overall script exit code to 1 - if [ "$test_specific_exit_code" -ne 0 ]; then - # No need for extra echo here, run_test already logged the failure. - overall_script_exit_code=1 - fi -} - -# --- Actual Test Execution --- -run_and_track_test 0 "test_perf.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_perf.py" -run_and_track_test 1 "test_compilation.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_compilation.py" -run_and_track_test 2 "test_basic.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_basic.py" -run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \ - "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine" -run_and_track_test 4 "test_quantization_accuracy.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py" -run_and_track_test 6 "test_tpu_model_runner.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py" -run_and_track_test 7 "test_sampler.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py" -run_and_track_test 8 "test_topk_topp_sampler.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py" -run_and_track_test 9 "test_multimodal.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_multimodal.py" -run_and_track_test 10 "test_pallas.py" \ - "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py" - -# After all tests have been attempted, exit with the overall status. -if [ "$overall_script_exit_code" -ne 0 ]; then - echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---" -else - echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---" -fi -exit "$overall_script_exit_code" -' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct. - -# Capture the exit code of the docker run command -DOCKER_RUN_EXIT_CODE=$? - -# The trap will run for cleanup. -# Exit the main script with the Docker run command's exit code. -if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then - echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE." - exit "$DOCKER_RUN_EXIT_CODE" -else - echo "Docker run command completed successfully." - exit 0 -fi -# TODO: This test fails because it uses RANDOM_SEED sampling -# pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \ diff --git a/.buildkite/scripts/tpu/cleanup_docker.sh b/.buildkite/scripts/tpu/cleanup_docker.sh deleted file mode 100755 index 740d81fb39bb..000000000000 --- a/.buildkite/scripts/tpu/cleanup_docker.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/bash - -set -euo pipefail - -docker_root=$(docker info -f '{{.DockerRootDir}}') -if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 -fi -echo "Docker root directory: $docker_root" -# Check disk usage of the filesystem where Docker's root directory is located -disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') -# Define the threshold -threshold=70 -if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=24h" --all - echo "Docker images and volumes cleanup completed." -else - echo "Disk usage is below $threshold%. No cleanup needed." -fi diff --git a/.buildkite/scripts/tpu/config_v6e_1.env b/.buildkite/scripts/tpu/config_v6e_1.env deleted file mode 100644 index c9e3c26571e7..000000000000 --- a/.buildkite/scripts/tpu/config_v6e_1.env +++ /dev/null @@ -1,14 +0,0 @@ -# Environment config -TEST_NAME=llama8b -CONTAINER_NAME=tpu-test - -# vllm config -MODEL=meta-llama/Llama-3.1-8B-Instruct -MAX_NUM_SEQS=256 -MAX_NUM_BATCHED_TOKENS=1024 -TENSOR_PARALLEL_SIZE=1 -MAX_MODEL_LEN=2048 -DOWNLOAD_DIR=/mnt/disks/persist -EXPECTED_THROUGHPUT=8.0 -INPUT_LEN=1800 -OUTPUT_LEN=128 diff --git a/.buildkite/scripts/tpu/docker_run_bm.sh b/.buildkite/scripts/tpu/docker_run_bm.sh deleted file mode 100755 index efb632e0a854..000000000000 --- a/.buildkite/scripts/tpu/docker_run_bm.sh +++ /dev/null @@ -1,91 +0,0 @@ -#!/bin/bash - -if [ ! -f "$1" ]; then - echo "Error: The env file '$1' does not exist." - exit 1 # Exit the script with a non-zero status to indicate an error -fi - -ENV_FILE=$1 - -# For testing on local vm, use `set -a` to export all variables -source /etc/environment -# shellcheck source=/dev/null -source "$ENV_FILE" - -remove_docker_container() { - docker rm -f "$CONTAINER_NAME" || true; -} - -trap remove_docker_container EXIT - -# Remove the container that might not be cleaned up in the previous run. -remove_docker_container - -LOG_ROOT=$(mktemp -d) -# If mktemp fails, set -e will cause the script to exit. -echo "Results will be stored in: $LOG_ROOT" - -if [ -z "$HF_TOKEN" ]; then - echo "Error: HF_TOKEN is not set or is empty." - exit 1 -fi - -# Make sure mounted disk or dir exists -if [ ! -d "$DOWNLOAD_DIR" ]; then - echo "Error: Folder $DOWNLOAD_DIR does not exist. This is useually a mounted drive. If no mounted drive, just create a folder." - exit 1 -fi - -echo "Run model $MODEL" -echo - -echo "starting docker...$CONTAINER_NAME" -echo -docker run \ - -v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \ - --env-file "$ENV_FILE" \ - -e HF_TOKEN="$HF_TOKEN" \ - -e TARGET_COMMIT="$BUILDKITE_COMMIT" \ - -e MODEL="$MODEL" \ - -e WORKSPACE=/workspace \ - --name "$CONTAINER_NAME" \ - -d \ - --privileged \ - --network host \ - -v /dev/shm:/dev/shm \ - vllm/vllm-tpu-bm tail -f /dev/null - -echo "run script..." -echo -docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/tpu/run_bm.sh" - -echo "copy result back..." -VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt -BM_LOG="$LOG_ROOT/$TEST_NAME"_bm_log.txt -docker cp "$CONTAINER_NAME:/workspace/vllm_log.txt" "$VLLM_LOG" -docker cp "$CONTAINER_NAME:/workspace/bm_log.txt" "$BM_LOG" - -throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g') -echo "throughput for $TEST_NAME at $BUILDKITE_COMMIT: $throughput" - -if [ "$BUILDKITE" = "true" ]; then - echo "Running inside Buildkite" - buildkite-agent artifact upload "$VLLM_LOG" - buildkite-agent artifact upload "$BM_LOG" -else - echo "Not running inside Buildkite" -fi - -# -# compare the throughput with EXPECTED_THROUGHPUT -# and assert meeting the expectation -# -if [[ -z "$throughput" || ! "$throughput" =~ ^[0-9]+([.][0-9]+)?$ ]]; then - echo "Failed to get the throughput" - exit 1 -fi - -if (( $(echo "$throughput < $EXPECTED_THROUGHPUT" | bc -l) )); then - echo "Error: throughput($throughput) is less than expected($EXPECTED_THROUGHPUT)" - exit 1 -fi diff --git a/.buildkite/scripts/tpu/quantized_v6e_1.env b/.buildkite/scripts/tpu/quantized_v6e_1.env deleted file mode 100644 index ecb98d4516bd..000000000000 --- a/.buildkite/scripts/tpu/quantized_v6e_1.env +++ /dev/null @@ -1,14 +0,0 @@ -# Environment config -TEST_NAME=llama8bw8a8 -CONTAINER_NAME=tpu-test - -# vllm config -MODEL=RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8 -MAX_NUM_SEQS=128 -MAX_NUM_BATCHED_TOKENS=1024 -TENSOR_PARALLEL_SIZE=1 -MAX_MODEL_LEN=2048 -DOWNLOAD_DIR=/mnt/disks/persist -EXPECTED_THROUGHPUT=8.7 -INPUT_LEN=1800 -OUTPUT_LEN=128 diff --git a/.buildkite/scripts/tpu/run_bm.sh b/.buildkite/scripts/tpu/run_bm.sh deleted file mode 100755 index b5d001bea0fe..000000000000 --- a/.buildkite/scripts/tpu/run_bm.sh +++ /dev/null @@ -1,93 +0,0 @@ -#!/bin/bash - -set -euo pipefail - -VLLM_LOG="$WORKSPACE/vllm_log.txt" -BM_LOG="$WORKSPACE/bm_log.txt" - -if [ -n "$TARGET_COMMIT" ]; then - head_hash=$(git rev-parse HEAD) - if [ "$TARGET_COMMIT" != "$head_hash" ]; then - echo "Error: target commit $TARGET_COMMIT does not match HEAD: $head_hash" - exit 1 - fi -fi - -echo "model: $MODEL" -echo - -# -# create a log folder -# -mkdir "$WORKSPACE/log" - -# TODO: Move to image building. -pip install pandas -pip install datasets - -# -# create sonnet_4x -# -echo "Create sonnet_4x.txt" -echo "" > benchmarks/sonnet_4x.txt -for _ in {1..4} - do - cat benchmarks/sonnet.txt >> benchmarks/sonnet_4x.txt -done - -# -# start vllm service in backend -# -echo "lanching vllm..." -echo "logging to $VLLM_LOG" -echo - -vllm serve "$MODEL" \ - --seed 42 \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ - --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \ - --no-enable-prefix-caching \ - --download_dir "$DOWNLOAD_DIR" \ - --max-model-len "$MAX_MODEL_LEN" > "$VLLM_LOG" 2>&1 & - - -echo "wait for 20 minutes.." -echo -# sleep 1200 -# wait for 10 minutes... -for _ in {1..120}; do - # TODO: detect other type of errors. - if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then - echo "Detected RuntimeError, exiting." - exit 1 - elif grep -Fq "Application startup complete" "$VLLM_LOG"; then - echo "Application started" - break - else - echo "wait for 10 seconds..." - sleep 10 - fi -done - -# -# run test -# -echo "run benchmark test..." -echo "logging to $BM_LOG" -echo -vllm bench serve \ - --backend vllm \ - --model "$MODEL" \ - --dataset-name sonnet \ - --dataset-path benchmarks/sonnet_4x.txt \ - --sonnet-input-len "$INPUT_LEN" \ - --sonnet-output-len "$OUTPUT_LEN" \ - --ignore-eos > "$BM_LOG" - -echo "completed..." -echo - -throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g') -echo "throughput: $throughput" -echo diff --git a/.github/mergify.yml b/.github/mergify.yml index 0798654d34a0..98c7d0327ab1 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -549,8 +549,6 @@ pull_request_rules: - label != stale - or: - files~=tpu.py - - files~=_tpu - - files~=tpu_ - files~=/tpu/ actions: label: @@ -564,8 +562,6 @@ pull_request_rules: - label != stale - and: - -files~=tpu.py - - -files~=_tpu - - -files~=tpu_ - -files~=/tpu/ actions: label: diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh deleted file mode 100644 index 9274e3c573cb..000000000000 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh +++ /dev/null @@ -1,157 +0,0 @@ -#!/bin/bash -set -xe - -# Hosts / ports -PREFILL_HOST=${PREFILL_HOST:-"localhost"} -PREFILL_PORT=${PREFILL_PORT:-8100} -PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577} -DECODE_HOST=${DECODE_HOST:-"localhost"} -DECODE_PORT=${DECODE_PORT:-8200} -PROXY_HOST=${PROXY_HOST:-"localhost"} -PROXY_PORT=${PROXY_PORT:-8192} -BASELINE_HOST=${BASELINE_HOST:-"localhost"} -BASELINE_PORT=${BASELINE_PORT:-9290} - - -# Model to run. -MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024} -BLOCK_SIZE=${BLOCK_SIZE:-32} - - -# execution env -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" -GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" -EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" -CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} -CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} - -OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"} - -# Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT - - -# Waits for vLLM server to start. -wait_for_server() { - local host=$1 - local port=$2 - timeout 1200 bash -c " - until curl -s ${host}:${port}/v1/completions > /dev/null; do - sleep 1 - done" && return 0 || return 1 -} - -# Cleanup function -cleanup() { - echo "Caught Ctrl+C, cleaning up..." - # Cleanup commands - pgrep python | xargs kill -9 || true - # pkill -f python || true - echo "Cleanup complete. Exiting." -} - -launch_baseline() { - BASELINE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${BASELINE_HOST} \ - --port ${BASELINE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --gpu-memory-utilization 0.5 \ - --enforce-eager" - echo "${BASELINE_BASE_CMD}" - ssh -tt "${BASELINE_HOST}" "${BASELINE_BASE_CMD}" & -} - -launch_pd() { - PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \ - VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${PREFILL_HOST} \ - --port ${PREFILL_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - - DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${DECODE_HOST} \ - --port ${DECODE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - echo "${PREFILL_BASE_CMD}" - echo "${DECODE_BASE_CMD}" - sleep 2 - - # execute on hosts - ssh -tt "${PREFILL_HOST}" "${PREFILL_BASE_CMD}" & - ssh -tt "${DECODE_HOST}" "${DECODE_BASE_CMD}" & - sleep 1 - wait_for_server "${PREFILL_HOST}" "${PREFILL_PORT}" - sleep 1 - wait_for_server "${DECODE_HOST}" "${DECODE_PORT}" - sleep 1 -} - -launch_pd_proxy(){ - PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - python3 ${EXP_ROOT}/toy_proxy_server.py \ - --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \ - --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \ - --host=${PROXY_HOST} --port ${PROXY_PORT}" - echo "${PROXY_BASE_CMD}" - ssh -tt "${PROXY_HOST}" "${PROXY_BASE_CMD}" & -} - -run_tests(){ - local service_url=$1 - local mode=$2 - python3 "${EXP_ROOT}"/test_disagg_accuracy.py --service_url="${service_url}" --model_name="${MODEL_NAME}" --mode="${mode}" --file_name="${OUTPUT_FILE}" -} - - -# run non-disagg. baseline & save outputs -launch_baseline -sleep 2 -wait_for_server "${BASELINE_HOST}" "${BASELINE_PORT}" -run_tests "http://${BASELINE_HOST}:${BASELINE_PORT}" "baseline" -cleanup -sleep 10 - - -# run disagg. & do exact-match with the outputs from baseline -launch_pd -launch_pd_proxy -sleep 10 -run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg" -echo "-----P/D success----" - -rm "${OUTPUT_FILE}" -cleanup - -exit 0 diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh deleted file mode 100644 index 5969455025e0..000000000000 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh +++ /dev/null @@ -1,125 +0,0 @@ -#!/bin/bash -set -xe - -# Hosts / ports -PREFILL_HOST=${PREFILL_HOST:-"localhost"} -PREFILL_PORT=${PREFILL_PORT:-8100} -PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577} -DECODE_HOST=${DECODE_HOST:-"localhost"} -DECODE_PORT=${DECODE_PORT:-8200} -PROXY_HOST=${PROXY_HOST:-"localhost"} -PROXY_PORT=${PROXY_PORT:-8192} -BASELINE_HOST=${BASELINE_HOST:-"localhost"} -BASELINE_PORT=${BASELINE_PORT:-9290} - - -# Model to run. -MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"} -MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024} -BLOCK_SIZE=${BLOCK_SIZE:-32} - - -# execution env -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" -GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" -EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" -CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} -CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} - -OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"} - -# Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT - -# Waits for vLLM server to start. -wait_for_server() { - local host=$1 - local port=$2 - timeout 1200 bash -c " - until curl -s ${host}:${port}/v1/completions > /dev/null; do - sleep 1 - done" && return 0 || return 1 -} - -# Cleanup function -cleanup() { - echo "Caught Ctrl+C, cleaning up..." - # Cleanup commands - pgrep python | xargs kill -9 || true - # pkill -f python || true - echo "Cleanup complete. Exiting." -} - - -launch_pd() { - PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \ - VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${PREFILL_HOST} \ - --port ${PREFILL_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - - DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - UCX_TLS=tcp \ - VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \ - VLLM_LOGGING_LEVEL=DEBUG \ - PJRT_DEVICE=TPU \ - VLLM_WORKER_MULTIPROC_METHOD=spawn \ - VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \ - --host ${DECODE_HOST} \ - --port ${DECODE_PORT} \ - --max-model-len ${MAX_MODEL_LEN}\ - --seed 42 \ - --block-size ${BLOCK_SIZE} \ - --enforce-eager \ - --gpu-memory-utilization 0.5 \ - --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'" - - echo "${PREFILL_BASE_CMD}" - echo "${DECODE_BASE_CMD}" - sleep 2 - - # execute on hosts - ssh -tt "${PREFILL_HOST}" "${PREFILL_BASE_CMD}" & - ssh -tt "${DECODE_HOST}" "${DECODE_BASE_CMD}" & - sleep 1 - wait_for_server "${PREFILL_HOST}" "${PREFILL_PORT}" - sleep 1 - wait_for_server "${DECODE_HOST}" "${DECODE_PORT}" - sleep 1 -} - -launch_pd_proxy(){ - PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME}; - python3 ${EXP_ROOT}/toy_proxy_server.py \ - --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \ - --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \ - --host=${PROXY_HOST} --port ${PROXY_PORT}" - echo "${PROXY_BASE_CMD}" - ssh -tt "${PROXY_HOST}" "${PROXY_BASE_CMD}" & -} - - -# run disagg. & do exact-match with the outputs from baseline -launch_pd -launch_pd_proxy -sleep 10 - -PREFILL_HOST=${PREFILL_HOST} \ -PREFILL_PORT=${PREFILL_PORT} \ -DECODE_HOST=${DECODE_HOST} \ -DECODE_PORT=${DECODE_PORT} \ -PROXY_HOST=${PROXY_HOST} \ -PROXY_PORT=${PROXY_PORT} python -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py From dfaabf91ac3dfa7e8e6bd0f56c7478026a533b2f Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Thu, 8 Oct 2026 17:48:47 +0800 Subject: [PATCH 197/952] [XPU] [CI] Fix stale ignore path for weight_transfer_metrics tests on Intel CI (#60558) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/entrypoints_intel.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 3cc9c1422d9e..96cbd2002a2f 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -80,7 +80,7 @@ steps: pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc && pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer --ignore=entrypoints/rl/return_routed_experts/test_openai_api.py - --ignore=entrypoints/rl/test_weight_operation_metrics.py + --ignore=entrypoints/rl/weight_transfer_metrics --deselect="tests/entrypoints/rl/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && PYTHONPATH=/workspace/vllm pytest -v -s entrypoints/serve/dev/rpc' From e73895a5dfc242ed8243bacb394ab75438b591d8 Mon Sep 17 00:00:00 2001 From: Muhammad Usman Mateen <74776275+usmanmateen@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:52:50 +0100 Subject: [PATCH 198/952] [Bugfix] Make `suppress_stdout` redirect fd 1 instead of `sys.stdout.fileno()` (#59336) Signed-off-by: usmanmateen Co-authored-by: Misha Goin --- tests/utils_/test_system_utils.py | 25 ++++++++++++++++++++++++- vllm/utils/system_utils.py | 3 ++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/utils_/test_system_utils.py b/tests/utils_/test_system_utils.py index 5ef55877a7c7..55ed35eeb56d 100644 --- a/tests/utils_/test_system_utils.py +++ b/tests/utils_/test_system_utils.py @@ -1,11 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import io import os +import sys import tempfile from pathlib import Path -from vllm.utils.system_utils import _maybe_force_spawn, unique_filepath +from vllm.utils.system_utils import ( + _maybe_force_spawn, + suppress_stdout, + unique_filepath, +) def test_unique_filepath(): @@ -25,3 +32,19 @@ def test_numa_bind_forces_spawn(monkeypatch): monkeypatch.setattr("sys.argv", ["vllm", "serve", "--numa-bind"]) _maybe_force_spawn() assert os.environ["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn" + + +def test_suppress_stdout_when_sys_stdout_has_no_fd(capfd): + with contextlib.redirect_stdout(io.StringIO()), suppress_stdout(): + os.write(1, b"c library output\n") + assert "c library output" not in capfd.readouterr().out + + +def test_suppress_stdout_keeps_stderr_when_sys_stdout_is_stderr(capfd, monkeypatch): + monkeypatch.setattr(sys, "stdout", sys.stderr) + with suppress_stdout(): + os.write(1, b"c library output\n") + os.write(sys.stderr.fileno(), b"error message\n") + out, err = capfd.readouterr() + assert "c library output" not in out + assert "error message" in err diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 5b90438dca49..a15dbf3c6cdf 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -72,7 +72,8 @@ def suppress_stdout(): yield return - stdout_fd = sys.stdout.fileno() + # C libraries write to fd 1, whatever sys.stdout is bound to. + stdout_fd = 1 stdout_dup = os.dup(stdout_fd) devnull_fd = os.open(os.devnull, os.O_WRONLY) From 459764104d2f8ac1b8a947f9a5a2603bdb07e360 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 8 Oct 2026 17:56:05 +0800 Subject: [PATCH 199/952] [Frontend] Move offline_utils.py to entrypoints/common. (#58052) Signed-off-by: wang.yuqi --- .github/CODEOWNERS | 2 +- tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py | 4 ++-- tests/entrypoints/unit_tests/test_offline_utils.py | 2 +- vllm/entrypoints/common/__init__.py | 0 vllm/entrypoints/{offline_utils.py => common/offline.py} | 0 vllm/entrypoints/generate/beam_search/offline.py | 2 +- vllm/entrypoints/llm.py | 2 +- vllm/entrypoints/pooling/offline.py | 2 +- vllm/entrypoints/rl/offline.py | 2 +- 9 files changed, 8 insertions(+), 8 deletions(-) create mode 100644 vllm/entrypoints/common/__init__.py rename vllm/entrypoints/{offline_utils.py => common/offline.py} (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 677e37ff8b60..65d9454b283b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -45,7 +45,7 @@ /vllm/entrypoints/serve @njhill /vllm/entrypoints/*.py @njhill /vllm/entrypoints/chat_utils.py @DarkLight1337 -/vllm/entrypoints/offline_utils.py @DarkLight1337 +/vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 # Rust Frontend diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index c03efa8f10da..55990ed1dda0 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -212,7 +212,7 @@ def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer( llm.renderer = renderer monkeypatch.setattr( - "vllm.entrypoints.offline_utils.parse_model_prompt", + "vllm.entrypoints.common.offline.parse_model_prompt", lambda _model_config, parsed_prompt: parsed_prompt, ) @@ -246,7 +246,7 @@ def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override( llm.renderer = renderer monkeypatch.setattr( - "vllm.entrypoints.offline_utils.parse_model_prompt", + "vllm.entrypoints.common.offline.parse_model_prompt", lambda _model_config, parsed_prompt: parsed_prompt, ) diff --git a/tests/entrypoints/unit_tests/test_offline_utils.py b/tests/entrypoints/unit_tests/test_offline_utils.py index 39f77672a41f..a37969c9b366 100644 --- a/tests/entrypoints/unit_tests/test_offline_utils.py +++ b/tests/entrypoints/unit_tests/test_offline_utils.py @@ -4,7 +4,7 @@ import pytest from vllm import SamplingParams -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.exceptions import VLLMValidationError from vllm.lora.request import LoRARequest diff --git a/vllm/entrypoints/common/__init__.py b/vllm/entrypoints/common/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/offline_utils.py b/vllm/entrypoints/common/offline.py similarity index 100% rename from vllm/entrypoints/offline_utils.py rename to vllm/entrypoints/common/offline.py diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 1ad451ac13ff..0acd0b43a52f 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -8,7 +8,7 @@ from tqdm import tqdm from vllm import RequestOutput, TextPrompt, TokensPrompt -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.pooling_params import PoolingParams diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 9dfff78ed9ff..a879b614844e 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -33,6 +33,7 @@ ChatTemplateContentFormatOption, load_chat_template, ) +from vllm.entrypoints.common.offline import _O, _R, OfflineInferenceMixin from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin from vllm.entrypoints.rl.offline import RLOfflineMixin @@ -52,7 +53,6 @@ from vllm.v1.sample.logits_processor import LogitsProcessor from ..renderers import ChatParams -from .offline_utils import _O, _R, OfflineInferenceMixin if TYPE_CHECKING: from vllm.v1.metrics.reader import Metric diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 27388f1ff0a9..81d78101c429 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -7,7 +7,7 @@ from tqdm.auto import tqdm from vllm.entrypoints.chat_utils import ChatTemplateConfig -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.inputs import DataPrompt, PromptType from vllm.logger import init_logger from vllm.lora.request import LoRARequest diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py index 02fc63104d8f..1db6a82533dd 100644 --- a/vllm/entrypoints/rl/offline.py +++ b/vllm/entrypoints/rl/offline.py @@ -5,7 +5,7 @@ WeightTransferInitRequest, WeightTransferUpdateRequest, ) -from vllm.entrypoints.offline_utils import OfflineInferenceMixin +from vllm.entrypoints.common.offline import OfflineInferenceMixin class RLOfflineMixin(OfflineInferenceMixin): From e161a0798c2305252e2b9bae483c537e9d0fc027 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 12:17:11 +0200 Subject: [PATCH 200/952] [Perf][Core] Reduce free KV cache queue overhead (#60269) Signed-off-by: NickLucche --- tests/v1/core/test_kv_cache_utils.py | 18 ++++++++++ vllm/v1/core/kv_cache_utils.py | 53 +++++++++++++++++----------- 2 files changed, 50 insertions(+), 21 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 7a34d07c07c6..53833ba83711 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -966,6 +966,24 @@ def test_free_kv_cache_block_queue_popleft_n(): assert block.next_free_block is None +@pytest.mark.parametrize("num_blocks", [1, 2, 16]) +def test_free_kv_cache_block_queue_popleft_n_drain_and_refill(num_blocks): + """Draining a batch must leave the sentinels ready for reuse.""" + blocks = [KVCacheBlock(block_id=i) for i in range(num_blocks)] + queue = FreeKVCacheBlockQueue(blocks) + + assert queue.popleft_n(num_blocks) == blocks + assert queue.num_free_blocks == 0 + assert queue.fake_free_list_head.next_free_block is queue.fake_free_list_tail + assert queue.fake_free_list_tail.prev_free_block is queue.fake_free_list_head + assert all(b.prev_free_block is None and b.next_free_block is None for b in blocks) + assert queue.popleft_n(0) == [] + + queue.append_n(blocks) + assert queue.popleft_n(num_blocks) == blocks + assert queue.num_free_blocks == 0 + + def test_free_kv_cache_block_queue_get_all_free_blocks(): # Create a list of KVCacheBlock objects blocks = [KVCacheBlock(block_id=i) for i in range(5)] diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 46934d530b54..79ec2fc5e4ed 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -303,19 +303,17 @@ def popleft(self) -> KVCacheBlock: The first free block. """ - if ( - self.fake_free_list_head.next_free_block is self.fake_free_list_tail - or self.fake_free_list_head.next_free_block is None - ): + head = self.fake_free_list_head + first_block = head.next_free_block + if first_block is self.fake_free_list_tail or first_block is None: assert self.num_free_blocks == 0, ( f"num_free_blocks ({self.num_free_blocks}) is out of sync " "with the free list." ) raise ValueError("No free blocks available") - first_block: KVCacheBlock = self.fake_free_list_head.next_free_block - - if first_block.next_free_block is None: + next_block = first_block.next_free_block + if next_block is None: # This should not happen if the block is from the free list. # It indicates a bug in the caller's logic. raise RuntimeError( @@ -325,8 +323,8 @@ def popleft(self) -> KVCacheBlock: # Connect fake_head and the next block of first_block (i.e. second block # or fake tail). - self.fake_free_list_head.next_free_block = first_block.next_free_block - first_block.next_free_block.prev_free_block = self.fake_free_list_head + head.next_free_block = next_block + next_block.prev_free_block = head # Remove the block from the linked list. first_block.prev_free_block = first_block.next_free_block = None @@ -350,6 +348,16 @@ def popleft_n(self, n: int) -> list[KVCacheBlock]: self.num_free_blocks -= n curr_block = self.fake_free_list_head.next_free_block + if n == 1: + assert curr_block is not None + next_block = curr_block.next_free_block + curr_block.prev_free_block = None + curr_block.next_free_block = None + if next_block is not None: + self.fake_free_list_head.next_free_block = next_block + next_block.prev_free_block = self.fake_free_list_head + return [curr_block] + # Pop n blocks from the head of the list ret = [] for _ in range(n): @@ -375,15 +383,17 @@ def remove(self, block: KVCacheBlock) -> None: block: The block to remove. """ - if block.prev_free_block is None or block.next_free_block is None: + prev_block = block.prev_free_block + next_block = block.next_free_block + if prev_block is None or next_block is None: # This should not happen if the block is from the free list. # It indicates a bug in the caller's logic. raise RuntimeError(f"remove() called on an invalid block: {block}") # Link the previous block to the next block. - block.prev_free_block.next_free_block = block.next_free_block + prev_block.next_free_block = next_block # Link the next block to the previous block. - block.next_free_block.prev_free_block = block.prev_free_block + next_block.prev_free_block = prev_block # Remove the block from the linked list. block.prev_free_block = block.next_free_block = None @@ -397,19 +407,19 @@ def append(self, block: KVCacheBlock) -> None: block: The block to append. """ - if self.fake_free_list_tail.prev_free_block is None: + tail = self.fake_free_list_tail + last_block = tail.prev_free_block + if last_block is None: raise RuntimeError( "prev_free_block of fake_free_list_tail should always exist" ) - last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block - # Connect the new block after the last block. last_block.next_free_block = block block.prev_free_block = last_block # Connect the fake tail after the new block. - block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = block + block.next_free_block = tail + tail.prev_free_block = block self.num_free_blocks += 1 @@ -441,10 +451,11 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None: blocks: The blocks to append. """ - if len(blocks) == 0: + if not blocks: return - last_block = self.fake_free_list_tail.prev_free_block + tail = self.fake_free_list_tail + last_block = tail.prev_free_block assert last_block is not None, ( "prev_free_block of fake_free_list_tail should always exist" ) @@ -455,8 +466,8 @@ def append_n(self, blocks: list[KVCacheBlock]) -> None: last_block = block # Connect the last block of to the fake tail - last_block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = last_block + last_block.next_free_block = tail + tail.prev_free_block = last_block self.num_free_blocks += len(blocks) From ba77c4c13018aa19244545cdc7ae8e016d66955e Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Thu, 8 Oct 2026 11:35:22 +0100 Subject: [PATCH 201/952] [Bugfix][Attention] Preserve small FP8 softmax weights in Triton attention (#60156) Signed-off-by: Andy Lo --- .../test_triton_unified_attention.py | 47 +++++++++++++++++++ .../attention/ops/triton_unified_attention.py | 10 +++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index e30d9225eac8..d936e377040c 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -181,6 +181,53 @@ def ref_paged_attn( return torch.cat(outputs, dim=0) +@torch.inference_mode() +def test_fp8_softmax_preserves_small_probabilities() -> None: + """Keep exp(-8) contributions that underflow when cast directly to E4M3.""" + device = torch.device(DEVICE_TYPE) + num_tokens = block_size = 32 + head_size = 128 + + query = torch.zeros(1, 1, head_size, dtype=FP8_DTYPE, device=device) + query[..., 0] = 1 + key_cache = torch.zeros(1, block_size, 1, head_size, dtype=FP8_DTYPE, device=device) + key_cache[:, 1:, :, 0] = -8 + value_cache = torch.ones_like(key_cache) + value_cache[:, 0] = 0 + output = torch.empty(1, 1, head_size, dtype=torch.bfloat16, device=device) + + cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device=device) + seqused_k = torch.tensor([num_tokens], dtype=torch.int32, device=device) + block_table = torch.tensor([[0]], dtype=torch.int32, device=device) + scale = torch.ones(1, dtype=torch.float32, device=device) + + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + seqused_k=seqused_k, + max_seqlen_k=num_tokens, + softmax_scale=1.0, + causal=True, + window_size=(-1, -1), + block_table=block_table, + softcap=0, + q_descale=scale, + k_descale=scale, + v_descale=scale, + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, + ) + + scores = torch.einsum("qhd,thd->qht", query.float(), key_cache[0].float()) + probabilities = torch.softmax(scores, dim=-1) + expected = torch.einsum("qht,thd->qhd", probabilities, value_cache[0].float()) + + torch.testing.assert_close(output.float(), expected, atol=5e-5, rtol=1e-2) + + def ref_paged_clamped_mm_attn( query: torch.Tensor, key_cache: torch.Tensor, diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 4d73e49e91b3..63a9cdc7f956 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -586,7 +586,15 @@ def kernel_unified_attention( P_v = (P * v_token_head_scales[None, :]).to(V.dtype) acc += tl.dot(P_v, V) else: - acc += tl.dot(P.to(V.dtype), V) + # Keep softmax weights representable when the P @ V dot uses FP8 + # tensor cores. Undo the power-of-two scale after the dot. + if V.dtype == tl.float8e4nv: + p_scale = 256.0 + elif V.dtype.is_fp8(): + p_scale = 128.0 + else: + p_scale = 1.0 + acc += tl.dot((P * p_scale).to(V.dtype), V) / p_scale # ---- Epilogue --------------------------------------------------------- if IS_3D: From b49c921a27083c25c16f308be1326055dbc87bbf Mon Sep 17 00:00:00 2001 From: devshah-cohere Date: Thu, 8 Oct 2026 07:11:55 -0400 Subject: [PATCH 202/952] [Bugfix][Rust Frontend] Always include token_ids in completion stream choices (#60286) Co-authored-by: Bugen Zhao Signed-off-by: devshah-cohere --- .../server/src/routes/openai/completions.rs | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 7074b8c6293c..646a67d1fb8f 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -271,6 +271,13 @@ async fn completion_chunk_stream( macro_rules! yield_chunk { ($chunk:expr) => {{ let mut chunk = $chunk; + if return_token_ids { + for choice in &mut chunk.choices { + if choice.token_ids.is_none() { + choice.token_ids = Some(Vec::new()); + } + } + } if include_continuous_usage { chunk.usage = Some(continuous_usage.to_usage()); } @@ -787,6 +794,70 @@ mod tests { } } + #[tokio::test] + async fn completion_stream_serializes_requested_token_ids_on_every_choice() { + let stream = stream::iter(vec![ + Ok(DecodedTextEvent::Start { + prompt_token_ids: vec![1, 2].into(), + prompt_logprobs: None, + }), + Ok(decoded_delta( + "ok", + vec![7, 8], + None, + Some(Finished { + usage: vllm_llm::TokenUsage { + prompt_token_count: 2, + output_token_count: 2, + cached_token_count: 0, + }, + finish_reason: FinishReason::Length, + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: None, + }), + )), + ]); + + let chunks = completion_chunk_stream( + stream, + "cmpl-1".to_string(), + "model".to_string(), + 1, + ApiServerOptions::default(), + ResponseOptions { + return_token_ids: true, + ..Default::default() + }, + ) + .collect::>() + .await; + + let chunks: Vec<_> = chunks.into_iter().try_collect().expect("stream should succeed"); + let payloads: Vec<_> = chunks + .iter() + .map(|chunk| serde_json::to_value(chunk).expect("chunk should serialize")) + .collect(); + + assert_eq!(payloads.len(), 3); + assert_eq!( + payloads[0]["choices"][0]["prompt_token_ids"], + serde_json::json!([1, 2]) + ); + assert_eq!( + payloads[0]["choices"][0]["token_ids"], + serde_json::json!([]) + ); + assert_eq!( + payloads[1]["choices"][0]["token_ids"], + serde_json::json!([7, 8]) + ); + assert_eq!( + payloads[2]["choices"][0]["token_ids"], + serde_json::json!([]) + ); + } + #[tokio::test] async fn collect_completion_hides_internal_prompt_only_token() { let stream = stream::iter(vec![ From 73c742b09c14f0260ccc9e9e7449444bdaf5475a Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Thu, 8 Oct 2026 04:41:14 -0700 Subject: [PATCH 203/952] [Bugfix] Make disable_any_whitespace actually disable whitespace on xgrammar (#58067) Signed-off-by: Stu Cao Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar.py | 93 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 27 +++++- 2 files changed, 118 insertions(+), 2 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar.py b/tests/v1/structured_output/test_backend_xgrammar.py index 339f0de17f03..bc29112077cc 100644 --- a/tests/v1/structured_output/test_backend_xgrammar.py +++ b/tests/v1/structured_output/test_backend_xgrammar.py @@ -2,12 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import ClassVar + import pytest +from transformers import AutoTokenizer, PreTrainedTokenizerBase from xgrammar import Grammar from xgrammar.testing import _is_grammar_accept_string +from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.structured_output.backend_types import StructuredOutputOptions from vllm.v1.structured_output.backend_xgrammar import ( + XgrammarBackend, has_xgrammar_unsupported_json_features, validate_xgrammar_grammar, ) @@ -462,3 +468,90 @@ def test_validate_xgrammar_preserves_multiline_choice(self): assert _is_grammar_accept_string(grammar, "yes") assert _is_grammar_accept_string(grammar, "no\nplease") assert not _is_grammar_accept_string(grammar, r"no\nplease") + + +# ================================================ +# Test XgrammarBackend functionality +# ================================================ + + +class TestXgrammarBackend: + class TestPR58067Regressions: + """Here we check behavior logic of our flag + disable_any_whitespace when used with xgrammar backend. + """ + + _TOKENIZER = "openai-community/gpt2" + _VOCAB_SIZE = 50257 + tokenizer: ClassVar[PreTrainedTokenizerBase] + + @classmethod + @pytest.fixture(scope="class", autouse=True) + def _shared_tokenizer(cls): + cls.tokenizer = AutoTokenizer.from_pretrained(cls._TOKENIZER) + + @classmethod + def _backend(cls, *, disable_any_whitespace: bool) -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig( + backend="xgrammar", + disable_any_whitespace=disable_any_whitespace, + ) + ) + return XgrammarBackend( + vllm_config, tokenizer=cls.tokenizer, vocab_size=cls._VOCAB_SIZE + ) + + @classmethod + def _accepts_json(cls, backend: XgrammarBackend, json_text: str) -> bool: + grammar = backend.compile_grammar(StructuredOutputOptions.JSON_OBJECT, "") + return grammar.accept_tokens("req", cls.tokenizer.encode(json_text)) + + def test_disable_any_whitespace_is_false(self): + """Verify expected behavior of disable_any_whitespace=False""" + backend = self._backend(disable_any_whitespace=False) + + # Check that grammar accepts json with NO space after colon + assert self._accepts_json(backend, '{"no_space_after_me":"yes_ofcourse"}') + + # Check that grammar accepts json with space after colon + assert self._accepts_json(backend, '{"yes_space_after_me": true}') + + # Check that grammar accepts json with NO space after comma + assert self._accepts_json( + backend, + '{"never_space_after_comma":"accepted","good_boy":":3"}', + ) + + # Check that grammar accepts json with space after comma + assert self._accepts_json( + backend, + '{"where_space_after_comma":42, "wow_thats_cool":4242}', + ) + + def test_disable_any_whitespace_is_true(self): + """Verify expected behavior of disable_any_whitespace=True""" + backend = self._backend(disable_any_whitespace=True) + + # Check that grammar accepts json with NO space after colon + assert self._accepts_json(backend, '{"no_space_after_me":"yes_ofcourse"}') + + # Check that grammar rejects json with space after colon + assert not self._accepts_json(backend, '{"yes_space_after_me": true}') + + # FIXME(arpera): + # Comma + separators=(",", ":") is wrong on xgrammar <=0.2.8 + # Seehttps://github.com/mlc-ai/xgrammar/issues/945 + # Re-enable both tests once vLLM pins xgrammar with the fix (>=0.2.9). + # + # # Check that grammar accepts json with NO space after comma + # assert self._accepts_json( + # backend, + # '{"never_space_after_comma":"accepted","good_boy":":3"}', + # ) + # + # # Check that grammar rejects json with space after comma + # assert not self._accepts_json( + # backend, + # '{"where_space_after_comma":42, "wow_thats_cool":4242}', + # ) diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 733ed7db092b..b675a9a3f193 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -89,13 +89,36 @@ def compile_grammar( grammar_spec: str, stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: + # Note(arpera): + # Our flag disable_any_whitespace does NOT map directly to + # xgrammar's flag any_whitespace + # To achieve desired behavior of disable_any_whitespace + # we have to set not only any_whitespace + # but also specify a list of separators after which + # xgrammar must not insert spaces. + # This is a requirement of xgrammar's API, so we must comply with it. + # + # FIXME(arpera): + # Currently xgrammar v0.2.8 DOES emit spaces after comma + # even if we specify it in separators list. + # The bug has been reported to xgrammar team: + # https://github.com/mlc-ai/xgrammar/issues/945 + # Please, track that issue, and once it is resolved remove this comment. + # Upd. this bug was fixed in xgrammar main branch on Oct 8, 2026 + # and will be available in next release. + # So, remove this comment once xgrammar updates to v0.2.9 + separators = (",", ":") if self.disable_any_whitespace else None if request_type == StructuredOutputOptions.JSON: ctx = self.compiler.compile_json_schema( - grammar_spec, any_whitespace=not self.disable_any_whitespace + grammar_spec, + any_whitespace=not self.disable_any_whitespace, + separators=separators, ) elif request_type == StructuredOutputOptions.JSON_OBJECT: ctx = self.compiler.compile_json_schema( - '{"type": "object"}', any_whitespace=not self.disable_any_whitespace + '{"type": "object"}', + any_whitespace=not self.disable_any_whitespace, + separators=separators, ) elif request_type == StructuredOutputOptions.GRAMMAR: if grammar_is_likely_lark(grammar_spec): From 458ba2edf85bc9b7ebc0d5141f34ea658520faf5 Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:01:47 +0800 Subject: [PATCH 204/952] [ROCm][Kimi-K3] Fuse AttnRes output with per-token FP8 input quantization (#59069) Signed-off-by: whx-sjtu Co-authored-by: OpenAI Codex Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- .../benchmark_kimi_k3_amd_fp8_fusions.py | 165 ++++++++++++++++++ tests/models/kimi_k3/test_amd_attn_res.py | 99 +++++++++++ .../kimi_k3/test_amd_mla_direct_return.py | 81 +++++++++ vllm/models/kimi_k3/amd/kda.py | 16 +- vllm/models/kimi_k3/amd/linear.py | 49 +++++- vllm/models/kimi_k3/amd/mla.py | 29 ++- vllm/models/kimi_k3/amd/ops/attn_res.py | 31 +++- 7 files changed, 455 insertions(+), 15 deletions(-) create mode 100644 benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py diff --git a/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py b/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py new file mode 100644 index 000000000000..dd210ae55ddc --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_amd_fp8_fusions.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Eight-GPU AttnRes+FP8 projection comparison with rotated graph timing.""" + +import argparse +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.platforms import current_platform + + +def capture(fn, repeats): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(5): + fn() + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(repeats): + fn() + return graph + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 256, 1024, 4096, 16384] + ) + parser.add_argument("--blocks", type=int, nargs="+", default=[1, 4, 8]) + args = parser.parse_args() + assert int(os.environ["WORLD_SIZE"]) == 8 + assert torch.accelerator.device_count() == 8 + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + torch.manual_seed(1729 + rank) + dist.init_process_group("gloo") + args.output.mkdir(parents=True, exist_ok=True) + dtype = current_platform.fp8_dtype() + rows = [] + with set_current_vllm_config(VllmConfig()): + weight = torch.randn(2112, 7168, device="cuda", dtype=torch.bfloat16) * 0.02 + weight, ws = rocm_aiter_ops.per_token_quant(weight, dtype) + weight = rocm_aiter_ops.shuffle_weight(weight) + for m in args.tokens: + prefix = torch.randn(m, 7168, device="cuda", dtype=torch.bfloat16) + delta = torch.zeros_like(prefix) + blocks = torch.randn(m, 9, 7168, device="cuda", dtype=torch.bfloat16) + norm = torch.ones(7168, device="cuda", dtype=torch.bfloat16) + score = torch.randn_like(norm) / 7168**0.5 + for b in args.blocks: + + def produce( + fused, + prefix=prefix, + delta=delta, + blocks=blocks, + norm=norm, + score=score, + b=b, + ): + result = attn_res( + prefix, + delta, + blocks, + norm, + score, + norm, + b, + -1, + 1e-5, + 1e-5, + quant_dtype=dtype if fused else None, + ) + return ( + result + if fused + else rocm_aiter_ops.per_token_quant(result, dtype) + ) + + def segment(fused, produce=produce): + x, scale = produce(fused) + return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm( + x, weight, scale, ws, output_dtype=torch.bfloat16 + ) + + expected, actual = segment(False), segment(True) + relative_l2 = ( + (actual.float() - expected.float()).norm() / expected.float().norm() + ).item() + assert relative_l2 < 0.002, relative_l2 + repeats = 20 if m < 1024 else 4 + functions = { + "baseline_producer": lambda produce=produce: produce(False), + "fused_producer": lambda produce=produce: produce(True), + "baseline_segment": lambda segment=segment: segment(False), + "fused_segment": lambda segment=segment: segment(True), + } + graphs = {k: capture(fn, repeats) for k, fn in functions.items()} + samples = {k: [] for k in functions} + names = list(functions) + for sample in range(9): + for name in names[sample % 4 :] + names[: sample % 4]: + dist.barrier() + graphs[name].replay() + torch.accelerator.synchronize() + start, end = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + start.record() + graphs[name].replay() + end.record() + end.synchronize() + samples[name].append(start.elapsed_time(end) * 1000 / repeats) + row = { + "tokens": m, + "blocks": b, + "relative_l2": relative_l2, + "samples_us": samples, + "median_us": {k: statistics.median(v) for k, v in samples.items()}, + } + rows.append(row) + print( + json.dumps( + { + "rank": rank, + "tokens": m, + "blocks": b, + "median_us": row["median_us"], + } + ), + flush=True, + ) + del graphs + (args.output / f"rank-{rank}.json").write_text( + json.dumps( + { + "rank": rank, + "world_size": 8, + "torch": torch.__version__, + "rows": rows, + "scope": "AttnRes + FP8 quant + TP8 MLA input projection", + }, + indent=2, + ) + + "\n" + ) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py index 80d5e1936f4f..f4efdbd5b21d 100644 --- a/tests/models/kimi_k3/test_amd_attn_res.py +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -191,3 +191,102 @@ def test_amd_attn_res_fused_contract( original_blocks[:, block_write_idx].copy_(expected_prefix) torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) assert actual.is_contiguous() + + +@pytest.mark.parametrize( + "num_tokens,num_blocks,has_delta,write_block", + [ + (0, 0, False, False), + (1, 0, True, True), + (17, 4, True, True), + (320, 8, False, False), + ], +) +def test_amd_attn_res_fp8_preserves_prefix_and_quantized_output( + num_tokens, + num_blocks, + has_delta, + write_block, + default_vllm_config, +): + from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + torch.manual_seed(41) + hidden_size = 7168 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=7) + delta = _randn_with_row_padding(num_tokens, hidden_size) if has_delta else None + blocks = _randn_with_row_padding(num_tokens, 9, hidden_size, padding=13) + norm = 1 + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) * 0.1 + score = torch.randn_like(norm) / hidden_size**0.5 + out_norm = 1 + torch.randn_like(norm) * 0.1 + ref_prefix, ref_blocks = prefix.clone(), blocks.clone() + kwargs = dict( + num_blocks=num_blocks, + block_write_idx=num_blocks if write_block else -1, + eps=1e-5, + output_norm_eps=1e-5, + ) + reference = attn_res(ref_prefix, delta, ref_blocks, norm, score, out_norm, **kwargs) + output, scale = attn_res( + prefix, + delta, + blocks, + norm, + score, + out_norm, + quant_dtype=current_platform.fp8_dtype(), + **kwargs, + ) + assert output.dtype == current_platform.fp8_dtype() + assert scale.shape == (num_tokens, 1) + if num_tokens: + # Keep reference rounding independent of Inductor's fused quantization. + quant = QuantFP8( + static=False, group_shape=GroupShape.PER_TOKEN, compile_native=False + ) + ref_output, ref_scale = quant(reference) + torch.testing.assert_close(scale, ref_scale, atol=1e-7, rtol=1e-6) + # Floating-point fusion can move values at FP8 rounding midpoints. + # Bound each change to the adjacent code and bound how often it occurs. + code_delta = ( + output.view(torch.uint8).to(torch.int16) + - ref_output.view(torch.uint8).to(torch.int16) + ).abs() + assert code_delta.max().item() <= 1 + assert (code_delta != 0).float().mean().item() < 1e-4 + torch.testing.assert_close(prefix, ref_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks, ref_blocks, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "matching,requires_unquantized", [(True, False), (False, False), (True, True)] +) +def test_mla_input_quant_requires_all_consumers(matching, requires_unquantized): + from types import SimpleNamespace + + from torch import nn + + from vllm.model_executor.layers.fusion.quant_activation import ( + expose_input_quant_key, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + ) + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + + attention = object.__new__(KimiK3MultiHeadLatentAttentionWrapper) + nn.Module.__init__(attention) + attention.indexer = None + attention.q_lora_rank = 1536 + attention.fused_qkv_a_proj = nn.Identity() + attention.g_proj = nn.Identity() + kernel = SimpleNamespace(input_quant_key=lambda: kFp8DynamicTokenSym) + expose_input_quant_key(attention.fused_qkv_a_proj, kernel) + if matching: + expose_input_quant_key(attention.g_proj, kernel) + attention.g_proj.requires_unquantized_input = requires_unquantized + actual = attention.get_input_quant_key() + assert actual == ( + kFp8DynamicTokenSym if matching and not requires_unquantized else None + ) diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py index a1c120737df8..50a52f4fe148 100644 --- a/tests/models/kimi_k3/test_amd_mla_direct_return.py +++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest import torch from torch import nn @@ -36,3 +39,81 @@ def test_mla_self_attention_returns_projection_storage_directly() -> None: assert output is projected assert output.data_ptr() == projected.data_ptr() + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("fused_decode", [False, True]) +def test_mla_forward_preserves_output_shape_with_fp8_input( + monkeypatch, quantized: bool, fused_decode: bool +) -> None: + """Both MLA paths accept producer-quantized input after decode fusion.""" + from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTokenSym, + ) + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + + hidden_states = torch.randn(4, 8, dtype=torch.bfloat16) + if quantized: + hidden_states = QuantizedActivation( + hidden_states.to(kFp8DynamicTokenSym.dtype), + torch.ones(4, 1), + hidden_states.dtype, + hidden_states.shape, + kFp8DynamicTokenSym, + ) + projected = torch.randn(4, 8, dtype=torch.bfloat16) + paths = [] + + def input_projection(value): + assert value is hidden_states + return torch.zeros(4, 8, dtype=torch.bfloat16), None + + def attention_output(path, output_shape): + assert output_shape == projected.shape + paths.append(path) + return projected + + attention = SimpleNamespace( + q_lora_rank=2, + kv_lora_rank=4, + qk_rope_head_dim=2, + qk_head_dim=4, + num_heads=2, + v_head_dim=4, + fused_qkv_a_proj=input_projection, + q_a_layernorm=nn.Identity(), + _normalize_q_kv=lambda q, kv: (q, kv), + q_b_proj=lambda q: (torch.zeros(4, 8, dtype=q.dtype), None), + rotary_emb=None, + indexer=None, + dcp_q_replicate=False, + _fused_qk_prep=fused_decode, + _fused_decode=lambda q, kv, pe, pos, shape, *args: attention_output( + "fused", shape + ), + g_proj=None, + o_proj=lambda value: (value, None), + ) + if fused_decode: + attention.mla_attn = SimpleNamespace(layer_name="test") + monkeypatch.setattr( + "vllm.models.kimi_k3.amd.mla.get_attention_context", + lambda _: ( + SimpleNamespace(num_actual_tokens=4, num_decode_tokens=4), + None, + torch.zeros(1), + torch.arange(4), + ), + ) + else: + attention.mla_attn = lambda q, kv, pe, **kwargs: attention_output( + "regular", kwargs["output_shape"] + ) + + output = KimiK3MultiHeadLatentAttentionWrapper.forward( + attention, torch.arange(4), hidden_states + ) + + assert output is projected + assert paths == ["fused" if fused_decode else "regular"] diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index e2dea74b6b9e..409cf6b1bf90 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -13,6 +13,10 @@ from vllm.distributed import divide from vllm.forward_context import get_forward_context from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -37,6 +41,7 @@ causal_conv1d_fn, causal_conv1d_update, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm.models.kimi_k3.amd.kda_metadata import KimiK3ROCmKDABackend @@ -300,13 +305,16 @@ def rearrange_mixed_qkv( qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1) return qkv.unbind(0) + def get_input_quant_key(self) -> QuantKey | None: + return get_input_quant_key(self.in_proj_qkvgfab) + def forward( self, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, positions: torch.Tensor, ) -> torch.Tensor: - num_tokens = hidden_states.size(0) projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + num_tokens = projected_qkvgfab.shape[0] split_sizes = [ 3 * self.local_projection_size, @@ -326,8 +334,8 @@ def forward( core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, + dtype=projected_qkvgfab.dtype, + device=projected_qkvgfab.device, ) self._forward( diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 65dad661ab12..3a1d827abe47 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -2,11 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -from typing import Any +from typing import Any, cast import torch from torch import nn +from vllm import envs from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_pp_group, @@ -19,6 +20,10 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -38,6 +43,10 @@ ) from vllm.model_executor.layers.mla import MLAModules from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8DynamicTokenSym, +) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -114,6 +123,9 @@ def __init__( "Only silu and situ are supported." ) + def get_input_quant_key(self) -> QuantKey | None: + return get_input_quant_key(self.gate_up_proj) + def forward(self, x): gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) @@ -148,8 +160,10 @@ def _apply_attn_res( delta: torch.Tensor | None = None, output_norm: RMSNorm | None = None, block_write_idx: int = -1, -) -> torch.Tensor: - return attn_res( + quant_key: QuantKey | None = None, +) -> torch.Tensor | QuantizedActivation: + fuse_quant = quant_key == kFp8DynamicTokenSym and not envs.VLLM_BATCH_INVARIANT + result = attn_res( prefix_sum, delta, block_residual, @@ -160,7 +174,18 @@ def _apply_attn_res( block_write_idx, norm.variance_epsilon, 0.0 if output_norm is None else output_norm.variance_epsilon, + quant_dtype=cast(torch.dtype, kFp8DynamicTokenSym.dtype) + if fuse_quant + else None, ) + if fuse_quant: + assert isinstance(result, tuple) + data, scale = result + return QuantizedActivation( + data, scale, prefix_sum.dtype, prefix_sum.shape, kFp8DynamicTokenSym + ) + assert isinstance(result, torch.Tensor) + return result class KimiMoE(nn.Module): @@ -452,10 +477,13 @@ def __init__( prefix, ) + def get_input_quant_key(self) -> QuantKey | None: + return self.mla_attn.get_input_quant_key() + def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: return self.mla_attn(positions, hidden_states) @@ -579,7 +607,7 @@ def __init__( def _run_self_attn( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: return self.self_attn( hidden_states=hidden_states, @@ -626,6 +654,11 @@ def forward_attn_residual( prefix_delta: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: prefix_sum = hidden_states + attention_quant_key = ( + self.self_attn.get_input_quant_key() + if isinstance(self.self_attn, (KimiK3DeltaAttention, KimiMLAAttention)) + else None + ) hidden_states = _apply_attn_res( prefix_sum, block_residual, @@ -635,6 +668,7 @@ def forward_attn_residual( delta=prefix_delta, output_norm=self.input_layernorm, block_write_idx=(self.block_write_idx if self.is_block_write_layer else -1), + quant_key=attention_quant_key, ) if self.is_block_write_layer: @@ -659,6 +693,11 @@ def forward_attn_residual( mlp_valid_blocks, delta=prefix_delta, output_norm=self.post_attention_layernorm, + quant_key=( + self.mlp.get_input_quant_key() + if isinstance(self.mlp, KimiMLP) + else None + ), ) hidden_states = self.mlp(hidden_states) diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index 11df7d035f31..1e4a51e8bf50 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -8,8 +8,13 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.attention.attention import get_attention_context +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms import current_platform _OPT_KV_LORA_RANK = 512 @@ -49,6 +54,21 @@ def _fused_qk_prep_supported(self) -> bool: not attn.is_aiter_triton_fp8_bmm_enabled and attn.W_UK_T is not None ) + def get_input_quant_key(self) -> QuantKey | None: + if self.indexer is not None: + return None + consumers = ( + [self.fused_qkv_a_proj] + if self.q_lora_rank is not None + else [self.kv_a_proj_with_mqa, self.q_proj] + ) + if self.g_proj is not None: + consumers.append(self.g_proj) + if any(layer is None for layer in consumers): + return None + keys = [get_input_quant_key(layer) for layer in consumers if layer is not None] + return keys[0] if all(key == keys[0] for key in keys) else None + def _normalize_q_kv( self, q_c: torch.Tensor, @@ -173,10 +193,15 @@ def _fused_decode( def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: q_c = None + num_tokens = ( + hidden_states.orig_shape[0] + if isinstance(hidden_states, QuantizedActivation) + else hidden_states.shape[0] + ) if self.q_lora_rank is not None: assert self.fused_qkv_a_proj is not None, ( @@ -238,7 +263,7 @@ def forward( if self.dcp_q_replicate: q_dcp_replicated, q = q, q_proj_layer._local_view(q) - output_shape = (hidden_states.shape[0], self.num_heads * self.v_head_dim) + output_shape = (num_tokens, self.num_heads * self.v_head_dim) # The fused decode path covers only the MQA slice, so it is taken for # decode-only batches; anything else falls through to MLAAttention. diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py index 5802e603670a..5c4fe7a0b0b2 100644 --- a/vllm/models/kimi_k3/amd/ops/attn_res.py +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -22,6 +22,7 @@ def _attn_res_kernel( qk_weight_ptr, output_norm_weight_ptr, output_ptr, + output_scale_ptr, stride_prefix_m: tl.constexpr, stride_delta_m: tl.constexpr, stride_block_m: tl.constexpr, @@ -35,6 +36,7 @@ def _attn_res_kernel( HAS_DELTA: tl.constexpr, WRITE_BLOCK: tl.constexpr, APPLY_OUTPUT_NORM: tl.constexpr, + QUANT_MAX: tl.constexpr, BLOCK_L: tl.constexpr, BLOCK_D: tl.constexpr, ): @@ -129,6 +131,14 @@ def _attn_res_kernel( output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 ).to(tl.float32) output = mixed * output_reciprocal_std * output_norm_weight + if QUANT_MAX > 0: + # Preserve the rounding of the original BF16 output before quantizing. + output = output.to(prefix_ptr.dtype.element_ty).to(tl.float32) + amax = tl.max(tl.where(d_mask, tl.abs(output), 0.0), axis=0) + scale = tl.maximum(amax, 1e-10) * (1.0 / QUANT_MAX) + inv_scale = 1.0 / scale + output = tl.minimum(tl.maximum(output * inv_scale, -QUANT_MAX), QUANT_MAX) + tl.store(output_scale_ptr + row_idx, scale) tl.store( output_ptr + row_idx * stride_output_m + d_offsets, output, @@ -147,7 +157,9 @@ def attn_res( block_write_idx: int, eps: float, output_norm_eps: float, -) -> torch.Tensor: + *, + quant_dtype: torch.dtype | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: num_tokens, hidden_size = prefix.shape assert prefix.stride(-1) == 1 assert delta is None or delta.stride(-1) == 1 @@ -155,9 +167,18 @@ def attn_res( assert norm_weight.stride(-1) == 1 assert qk_weight.stride(-1) == 1 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 - output = prefix.new_empty(prefix.shape) + if quant_dtype is not None: + assert quant_dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + output = torch.empty_like( + prefix, dtype=quant_dtype or prefix.dtype, memory_format=torch.contiguous_format + ) + scale = ( + torch.empty((num_tokens, 1), device=prefix.device, dtype=torch.float32) + if quant_dtype is not None + else None + ) if num_tokens == 0: - return output + return output if scale is None else (output, scale) # Source tiling helps decode, while one-source tiles scale better for prefill. if num_tokens >= 256 or num_blocks <= 1: @@ -172,6 +193,7 @@ def attn_res( qk_weight, output_norm_weight, output, + scale, prefix.stride(0), 0 if delta is None else delta.stride(0), blocks.stride(0), @@ -185,9 +207,10 @@ def attn_res( HAS_DELTA=delta is not None, WRITE_BLOCK=block_write_idx >= 0, APPLY_OUTPUT_NORM=output_norm_weight is not None, + QUANT_MAX=0.0 if quant_dtype is None else torch.finfo(quant_dtype).max, BLOCK_L=block_l, BLOCK_D=triton.next_power_of_2(hidden_size), num_warps=num_warps, num_stages=2, ) - return output + return output if scale is None else (output, scale) From 194da61de1769c28bfad2b6642c60cc7eca544dd Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Thu, 8 Oct 2026 06:01:25 -0700 Subject: [PATCH 205/952] [Bugfix] Reject negative --device-ids indices instead of selecting the wrong device (#49053) Signed-off-by: Vineeth Sai Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 5 +++-- vllm/engine/arg_utils.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 924d2bafec77..fc9f0229c46c 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -857,13 +857,14 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): class TestDeviceIds: - def test_device_ids_with_cvd_out_of_range(self, monkeypatch): + @pytest.mark.parametrize("device_ids", [[0, 2], [-1]]) + def test_device_ids_with_cvd_out_of_range(self, monkeypatch, device_ids): """--device-ids index beyond the CVD set raises ValueError.""" from vllm.platforms import current_platform key = current_platform.device_control_env_var monkeypatch.setenv(key, "4,5") - args = EngineArgs(model="m", device_ids=[0, 2]) + args = EngineArgs(model="m", device_ids=device_ids) with pytest.raises(ValueError, match="out of range"): args._resolve_device_ids() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 800d1b2f2ccf..94dd3e2ef1c8 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2113,7 +2113,7 @@ def _resolve_device_ids(self) -> list[int] | None: for x in cvd.split(",") ] for i in int_ids: - if i >= len(cvd_ids): + if not 0 <= i < len(cvd_ids): raise ValueError( f"--device-ids index {i} is out of range for " f"{current_platform.device_control_env_var}" From 4ed4c4433b87e45d72f401f27f1b56c932c90071 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 8 Oct 2026 10:00:49 -0400 Subject: [PATCH 206/952] [Bugfix][HiSparse] Drop the unaccounted, write-only prefill mirror staging buffer (#60429) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../v1/attention/test_sparse_mla_backends.py | 12 +--- tests/v1/worker/test_attn_utils.py | 2 +- .../kv_connector/v1/hisparse/worker.py | 1 - .../layers/attention/mla_attention.py | 12 ---- vllm/v1/hisparse/binding.py | 4 -- vllm/v1/hisparse/runtime.py | 55 +------------------ 6 files changed, 5 insertions(+), 81 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 3d32aed5d6e7..4d33075c9303 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -2578,7 +2578,7 @@ def test_hisparse_multi_step_writes_request_major_output(): @requires_hisparse_ops -def test_hisparse_kv_update_writes_resident_and_staging_caches(): +def test_hisparse_kv_update_writes_resident_cache(): device = torch.device(DEVICE_TYPE) block_size = 4 row_width = 8 @@ -2597,18 +2597,11 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches(): block_table=torch.tensor([[1]], dtype=torch.int32, device=device), slot_mapping=resident_slots, ) - cache_handle.mirror_staging_cache = torch.empty( - (1, block_size, row_width), dtype=torch.float32, device=device - ) - cache_handle.mirror_staging_slots = torch.arange( - block_size, dtype=torch.int64, device=device - ) slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device) kv_c = torch.randn(8, row_width - 2, device=device) k_pe = torch.randn(8, 1, 2, device=device) cache_handle.num_actual_tokens = slots.numel() cache_handle.decode_batch = False - cache_handle.host_mirror_required = True source_cache = torch.zeros_like(cache_handle.view.cache) impl = object.__new__(FlashMLASparseImpl) layer = SimpleNamespace( @@ -2635,9 +2628,6 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches(): expected.to(device), ) torch.testing.assert_close(source_cache, torch.zeros_like(source_cache)) - staged = cache_handle.mirror_staging_cache.view(-1, row_width) - staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1) - torch.testing.assert_close(staged[:3], staged_expected) @requires_hisparse_ops diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 41148070528a..9c9c26c3c2c0 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -434,7 +434,7 @@ def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase): cache_config=SimpleNamespace( get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC ), - scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1), + scheduler_config=SimpleNamespace(max_num_seqs=1), ) def allocate(*args): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index bf6329b3bfdf..b9dd01f05420 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -417,7 +417,6 @@ def start_step( self._clear_forward_mirror_state() for handle in self.cache_handles: handle.all_context_pages_resident = metadata.all_context_pages_resident - handle.mirror_from_resident = True self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event) transfers = ( metadata.command.page_transfers if metadata.command is not None else [] diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4286502ea9fe..2a3048096c98 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -802,18 +802,6 @@ def update_kv_cache( kv_cache_dtype, k_scale, ) - if cache is not None: - mirror_target = cache.mirror_write_target(kv_c_normed.shape[0]) - if mirror_target is not None: - mirror_cache, mirror_slots = mirror_target - self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_c_normed, - k_pe, - mirror_cache, - mirror_slots, - kv_cache_dtype, - k_scale, - ) def prepare_kv_cache_update( self, attn_metadata: "MLACommonMetadata | None" diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index d021d8b7be09..b8c0d5dff269 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -147,7 +147,6 @@ def init_hisparse_kv_cache( initialize_hisparse_runtime_buffers( cache_handles, max_num_reqs=vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens, ) return kv_caches except Exception: @@ -197,9 +196,6 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None: ) ) release_pinned_state(list(runtimes.values()), registered_pools, shared_region) - for cache in cache_handles: - cache.mirror_staging_cache = None - cache.mirror_staging_slots = None def bind_hisparse_kv_caches( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index fc415fa2eccf..a6378ca9141a 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -1120,10 +1120,7 @@ def __init__(self, runtime: HiSparseRuntime) -> None: self.num_decode_tokens = 0 self.req_id_per_token: torch.Tensor | None = None self.host_mirror_required = False - self.mirror_from_resident = False self.mirror_slot_mapping: torch.Tensor | None = None - self.mirror_staging_cache: torch.Tensor | None = None - self.mirror_staging_slots: torch.Tensor | None = None self.submit_layer_mirror: Callable[[], None] | None = None # Speculator layers write their rows after the target forward. self.draft_layer = False @@ -1160,26 +1157,6 @@ def write_target( num_rows = min(num_rows, self.runtime.max_num_reqs) return self.view.cache, self.slot_mapping[:num_rows], num_rows - def mirror_write_target( - self, num_rows: int - ) -> tuple[torch.Tensor, torch.Tensor] | None: - if ( - not self.host_mirror_required - or self.decode_batch - or self.mirror_from_resident - ): - return None - cache = self.mirror_staging_cache - slots = self.mirror_staging_slots - if cache is None or slots is None: - raise RuntimeError("HiSparse prefill mirror staging is not bound.") - if num_rows > slots.numel(): - raise RuntimeError( - "HiSparse prefill mirror exceeds staging capacity: " - f"{num_rows} > {slots.numel()}." - ) - return cache, slots[:num_rows] - def finish_kv_update(self) -> None: if self.dummy_batch: return @@ -1242,9 +1219,8 @@ def initialize_hisparse_runtime_buffers( cache_handles: list[HiSparseCacheHandle], *, max_num_reqs: int, - max_num_batched_tokens: int, ) -> None: - """Allocate shared request state and per-layer prefill staging buffers.""" + """Allocate the request state shared by every HiSparse runtime.""" assert cache_handles resident = cache_handles[0].view assert resident is not None @@ -1252,26 +1228,8 @@ def initialize_hisparse_runtime_buffers( request_state_indices = torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ) - staging_blocks = ( - max_num_batched_tokens + resident.block_size - 1 - ) // resident.block_size - mirror_staging_caches = torch.empty( - ( - len(cache_handles), - staging_blocks, - resident.block_size, - resident.cache.shape[-1], - ), - dtype=resident.cache.dtype, - device=device, - ) - mirror_staging_slots = torch.arange( - max_num_batched_tokens, dtype=torch.int64, device=device - ) - for layer_index, cache_handle in enumerate(cache_handles): + for cache_handle in cache_handles: cache_handle.runtime.request_state_indices = request_state_indices - cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index] - cache_handle.mirror_staging_slots = mirror_staging_slots def create_hisparse_cache_handle( @@ -1325,11 +1283,4 @@ def create_hisparse_cache_handle( config.device_buffer_size, max_num_reqs, ) - handle = HiSparseCacheHandle(runtime) - speculative_config = vllm_config.speculative_config - handle.mirror_from_resident = bool( - vllm_config.scheduler_config.async_scheduling - and speculative_config is not None - and speculative_config.uses_draft_kv_cache() - ) - return handle + return HiSparseCacheHandle(runtime) From daa9085143d65896ba3fa77635b0b110aab3e585 Mon Sep 17 00:00:00 2001 From: Pravein Govindan Kannan Date: Thu, 8 Oct 2026 19:55:02 +0530 Subject: [PATCH 207/952] [KV Connector][NIXL] Register HiSparse host pool under single TP Rank (#60161) Signed-off-by: Pravein Govindan Kannan Signed-off-by: Matthew Bonanni Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../kv_connector/unit/test_nixl_connector.py | 62 ++++++++++++++++ .../unit/test_nixl_connector_hma.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 72 ++++++++++++++----- .../kv_connector/v1/nixl/pull_worker.py | 11 ++- 4 files changed, 128 insertions(+), 18 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 18097d220e79..52ac746391ff 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2083,6 +2083,7 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): worker.kv_caches_base_addr = {"local": {0: [100, 200]}} worker._has_mamba = False worker._mixed_mem_types = True + worker._skip_dram_xfer = False worker.region_mem_types = ["DRAM", "VRAM"] worker.region_num_blocks = [2, 2] worker._desc_is_dram_by_block_size = {} @@ -2150,6 +2151,7 @@ def recv_worker(): worker.enable_permute_local_kv = False worker.enable_heterogeneous_attn_post_process = False worker.tp_rank = 0 + worker._skip_dram_xfer = False worker._log_failure = MagicMock() # type: ignore[method-assign] worker.xfer_stats = NixlKVConnectorStats() worker.nixl_wrapper = MagicMock() @@ -2207,6 +2209,66 @@ def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker): worker.nixl_wrapper.send_notif.assert_not_called() +def test_hisparse_dram_rank0_only_read_keeps_device_part(recv_worker): + worker = recv_worker + worker.tp_rank = 1 + worker._skip_dram_xfer = True + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + worker._dram_src_handles_by_block_size = {16: None} + worker.nixl_wrapper.make_prepped_xfer.side_effect = [102] + worker.nixl_wrapper.check_xfer_state.return_value = "DONE" + + worker._read_blocks_mixed( + request_id="request", + local_block_size_key=16, + local_device_handle=20, + local_dram_handle=None, + remote_xfer_side_handle=30, + local_block_descs_ids=np.array([0, 2]), + remote_block_descs_ids=np.array([5, 7]), + notif_agent="prefill", + notif_id=b"request:1", + ) + + (device_read,) = worker.nixl_wrapper.make_prepped_xfer.call_args_list + assert device_read.args[:2] == ("READ", 20) + np.testing.assert_array_equal(device_read.args[2], [0]) + np.testing.assert_array_equal(device_read.args[4], [7]) + + assert worker.get_finished() == (set(), {"request"}) + worker.nixl_wrapper.send_notif.assert_called_once_with( + "prefill", notif_msg=b"request:1" + ) + + +def test_hisparse_dram_rank0_only_dram_only_read_completes(recv_worker): + """With nothing left to read, the request still completes and notifies.""" + worker = recv_worker + worker.tp_rank = 1 + worker._skip_dram_xfer = True + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + + worker._read_blocks_mixed( + request_id="request", + local_block_size_key=16, + local_device_handle=20, + local_dram_handle=None, + remote_xfer_side_handle=30, + local_block_descs_ids=np.array([0, 1]), + remote_block_descs_ids=np.array([5, 6]), + notif_agent="prefill", + notif_id=b"request:1", + ) + + worker.nixl_wrapper.make_prepped_xfer.assert_not_called() + assert worker.get_finished() == (set(), {"request"}) + worker.nixl_wrapper.send_notif.assert_called_once_with( + "prefill", notif_msg=b"request:1" + ) + + def element_byte_addrs(view: torch.Tensor) -> list[int]: """Absolute byte addresses of every element of a (possibly strided) view.""" offsets = torch.zeros(view.shape, dtype=torch.int64) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index b4b5e063445f..f016b26475f0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -267,6 +267,7 @@ def test_region_pull_completion_zeros_only_own_padding(region_pull_worker): worker.region_num_blocks = [4, 4] worker.block_len_per_layer = [8, 8] worker.block_stride_per_layer = [24, 24] + worker._skip_dram_xfer = False worker._recving_metadata = { "request": ReqMeta( local_block_ids=[[1, 2], [0, 1]], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index fa29d41a825e..2a2f660cf15e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -789,9 +789,12 @@ def __init__( self.region_names: list[str] = [] self.region_num_blocks: list[int] = [] self._mixed_mem_types = False + # True on TP ranks that share a HiSparse host pool with TP rank 0 and + # leave its NIXL registration and DRAM reads to rank 0. + self._skip_dram_xfer = False self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {} self._desc_pos_by_block_size: dict[int, np.ndarray] = {} - self._dram_src_handles_by_block_size: dict[int, int] = {} + self._dram_src_handles_by_block_size: dict[int, int | None] = {} # PP>1 (push mode): this worker holds a contiguous layer slice and # transfers into the matching sub-range of a PP=1 remote's regions. @@ -828,7 +831,7 @@ def __init__( # Populated dynamically during handshake based on remote configuration. # Per-source split handles, keyed by (tp_ratio, remote_block_size). self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} - self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} + self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int | None]] = {} # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) @@ -1818,12 +1821,15 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.num_descs = sum(xfer_region_num_blocks) self._mixed_mem_types = len(set(region_mem_types)) > 1 + self._skip_dram_xfer = self._leaves_dram_xfer_to_rank0() if self._mixed_mem_types: assert self.use_mla and not self._has_mamba, ( "Mixed-device KV registration is only supported for MLA " "models without Mamba layers." ) for mem_type in sorted(set(region_mem_types)): + if mem_type == "DRAM" and self._skip_dram_xfer: + continue ranges_for_mem_type = [ (start, end - start, device_id, "") for (_, cache_mem_type), ( @@ -2191,10 +2197,13 @@ def register_local_xfer_handler( blocks_data[desc_is_dram, 2] = 0 dram_blocks = blocks_data[dram_idx] vram_blocks = blocks_data[vram_idx] - dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") - self._dram_src_handles_by_block_size[block_size] = ( - self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) - ) + if self._skip_dram_xfer: + self._dram_src_handles_by_block_size[block_size] = None + else: + dram_descs = self.nixl_wrapper.get_xfer_descs(dram_blocks, "DRAM") + self._dram_src_handles_by_block_size[block_size] = ( + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs) + ) descs = self.nixl_wrapper.get_xfer_descs(vram_blocks, self.nixl_memory_type) return ( self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), @@ -2423,12 +2432,15 @@ def add_remote_agent( if self._mixed_mem_types: handle_data = np.asarray(handle_data, dtype=np.uint64) desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size] - dram_descs = self.nixl_wrapper.get_xfer_descs( - handle_data[desc_is_dram], "DRAM" - ) - dram_handle = self.nixl_wrapper.prep_xfer_dlist( - "NIXL_INIT_AGENT", dram_descs - ) + if self._skip_dram_xfer: + dram_handle = None + else: + dram_descs = self.nixl_wrapper.get_xfer_descs( + handle_data[desc_is_dram], "DRAM" + ) + dram_handle = self.nixl_wrapper.prep_xfer_dlist( + "NIXL_INIT_AGENT", dram_descs + ) self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle) handle_data = handle_data[~desc_is_dram] descs = self.nixl_wrapper.get_xfer_descs( @@ -2887,6 +2899,27 @@ def post_process_device_kv_on_receive_heterogeneous_attn( indices=indices, ) + def _leaves_dram_xfer_to_rank0(self) -> bool: + """Whether this rank leaves the DRAM regions' NIXL work to TP rank 0. + + A HiSparse host pool shared by the local TP ranks (one mmap per DP + rank) is registered with NIXL and filled by TP rank 0 alone: every + rank would otherwise register the same pages on every NIC and READ the + same host KV into them. Private pools are unaffected. + """ + if not ( + self._mixed_mem_types and self.kv_cache_config.hisparse_shared_host_pool + ): + return False + skip = self.tp_rank != 0 + logger.info( + "HiSparse host pool: NIXL DRAM registration and reads on TP rank 0 " + "only (tp_rank=%d, skip=%s)", + self.tp_rank, + skip, + ) + return skip + def _zero_region_blocks(self, block_ids: BlockIds) -> None: """Clear clipped physical pages in their owning region only.""" if not any(block_ids): @@ -2895,6 +2928,9 @@ def _zero_region_blocks(self, block_ids: BlockIds) -> None: for region, blocks in enumerate(block_ids): if not blocks: continue + if self._skip_dram_xfer and self.region_mem_types[region] == "DRAM": + # The shared pool is rank 0's to write. + continue cache = self.device_kv_caches[self.region_names[region]] storage = cache.untyped_storage() pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_( @@ -3795,11 +3831,13 @@ def _finish_shutdown(self) -> None: for handles in self.src_xfer_handles_by_tp_ratio.values(): for handle in handles: self.nixl_wrapper.release_dlist_handle(handle) - for handles in self._dram_src_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - for handle in self._dram_src_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) + for dram_handles in self._dram_src_handles_by_tp_ratio.values(): + for dram_handle in dram_handles: + if dram_handle is not None: + self.nixl_wrapper.release_dlist_handle(dram_handle) + for dram_handle in self._dram_src_handles_by_block_size.values(): + if dram_handle is not None: + self.nixl_wrapper.release_dlist_handle(dram_handle) except Exception: logger.exception("NIXL dlist-handle release failed at shutdown.") self.src_xfer_handles_by_block_size.clear() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 5bcfdb9560a0..16a25149ea22 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -593,7 +593,16 @@ def _read_blocks_mixed( remote_ids = np.asarray(remote_block_descs_ids) is_dram = desc_is_dram[local_ids] - assert local_dram_handle is not None + if self._skip_dram_xfer: + # TP rank 0 reads the DRAM (shared host pool) part for every + # local rank; keep only the device part. Register the request + # even if nothing is left, so it still completes and notifies. + keep = ~is_dram + local_ids = local_ids[keep] + remote_ids = remote_ids[keep] + is_dram = is_dram[keep] + self._recving_transfers.setdefault(request_id, []) + assert local_dram_handle is not None or not is_dram.any() reads = ( (is_dram, local_dram_handle), (~is_dram, local_device_handle), From 2f7f05c256d671a7bdf04bcf4b4fc6a1566bbfcb Mon Sep 17 00:00:00 2001 From: zzaebok <44357534+zzaebok@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:28:16 +0800 Subject: [PATCH 208/952] [Bugfix][MoRIIO] Prevent remote KV reload after decoder preemption (#53250) Signed-off-by: Jaebok Lee --- tests/v1/kv_connector/unit/test_moriio_connector.py | 11 +++++++++++ .../v1/kv_connector/unit/test_moriio_hma_scheduler.py | 4 ++-- .../kv_connector/v1/moriio/moriio_connector.py | 3 +++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index cb33bfbcf03c..73cc5f0583d0 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -195,6 +195,17 @@ def _build_meta( return scheduler.build_connector_meta(scheduler_output) +def test_write_mode_does_not_reload_remote_kv_after_preemption(): + scheduler = _write_consumer_scheduler_for_finished_request() + request = create_request(request_id=1, num_tokens=32, do_remote_prefill=True) + + assert scheduler.get_num_new_matched_tokens(request, 0) == (32, True) + + assert request.kv_transfer_params is not None + request.kv_transfer_params["do_remote_prefill"] = False + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + + class FakeMoRIIOWrapper: # A fake MoRIIOWrapper for testing purposes def __init__(self, *args, **kwargs): diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 1de729416e90..4bd32f786a28 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -824,7 +824,7 @@ def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens(): req = SimpleNamespace( num_prompt_tokens=10, prompt_token_ids=list(range(10)), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=2) # Pure-attention WRITE: no N-1 drop; full length minus already-computed. @@ -848,7 +848,7 @@ def test_get_num_new_matched_tokens_supports_embeds_only_prompts( num_prompt_tokens=10, prompt_token_ids=None, prompt_embeds=object(), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) assert sched.get_num_new_matched_tokens(req, num_computed_tokens) == ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 2d17d7a7ac52..082643778e7d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -769,6 +769,9 @@ def get_num_new_matched_tokens( num_prompt_tokens = request.num_prompt_tokens num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0) if self.mode == MoRIIOMode.WRITE: + params = request.kv_transfer_params or {} + if not params.get("do_remote_prefill"): + return 0, False # MoriiO in write mode, no remote prefill. Hybrid models never get # here: register_kv_caches refuses WRITE mode for them, so there is # no recurrent-state accounting to do. From 273e099693a6c92e1bc516cd10ffc38b18197eec Mon Sep 17 00:00:00 2001 From: "yhcheong (vllmellm)" Date: Thu, 8 Oct 2026 22:44:58 +0800 Subject: [PATCH 209/952] [Bugfix][ROCm] Select page-aligned kernel blocks for pooled indexers so the block table addresses storage pages (GLM-5.3-Flash) (#59412) Signed-off-by: vllmellm --- .../test_kpool_indexer_block_sizes.py | 135 ++++++++++++++++++ vllm/v1/worker/utils.py | 67 +++++---- 2 files changed, 178 insertions(+), 24 deletions(-) create mode 100644 tests/v1/attention/test_kpool_indexer_block_sizes.py diff --git a/tests/v1/attention/test_kpool_indexer_block_sizes.py b/tests/v1/attention/test_kpool_indexer_block_sizes.py new file mode 100644 index 000000000000..a737988dde66 --- /dev/null +++ b/tests/v1/attention/test_kpool_indexer_block_sizes.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for storage-block kernel-block selection in prepare_kernel_block_sizes.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.models.glm5next.common.attention import Glm5NextIndexerCache +from vllm.v1.attention.backend import AttentionBackend, MultipleOf +from vllm.v1.attention.backends.mla import indexer as indexer_mod +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerBackend, + KpoolTailBackend, +) +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, +) +from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.worker.utils import prepare_kernel_block_sizes + +pytestmark = pytest.mark.cpu_test + +# ``index_kpool`` of zai-org/GLM-5.3-Flash: one indexer state pools 4 +# compressed states (hence storage_block_size = page * 4). +INDEX_KPOOL = 4 + +# The hybrid KDA/mamba manager floors at TP8..TP1; 640 is the geometry of the +# original report ("the 640-token block table"). +TP_FLOOR_MANAGER_BLOCKS = [640, 1152, 2176, 4352] + + +def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None: + # Platform predicates are mutually exclusive in production. Override both + # so these ROCm policy tests do not inherit CUDA capabilities from the CI + # host running them (same approach as test_deepseek_v4_rocm_adaptive). + monkeypatch.setattr(indexer_mod.current_platform, "is_cuda", lambda: False) + monkeypatch.setattr(indexer_mod.current_platform, "is_rocm", lambda: True) + + +def _kpool_storage_spec(manager_block: int) -> MLAAttentionSpec: + """Real Glm5NextIndexerCache spec for the given cache block size.""" + with set_current_vllm_config(VllmConfig()): + cache = Glm5NextIndexerCache( + head_dim=128, + dtype=torch.bfloat16, + prefix="model.layers.0.indexer.k_cache_probe", + cache_config=SimpleNamespace(block_size=manager_block), + index_kpool=INDEX_KPOOL, + ) + spec = cache.get_kv_cache_spec(VllmConfig()) + assert isinstance(spec, MLAAttentionSpec) + return spec + + +def _prepare( + spec: MLAAttentionSpec, backends: list[type[AttentionBackend]] +) -> list[int]: + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)] + ) + attn_groups = [[SimpleNamespace(backend=backend) for backend in backends]] + return prepare_kernel_block_sizes(kv_cache_config, attn_groups) + + +@pytest.mark.parametrize("manager_block", TP_FLOOR_MANAGER_BLOCKS) +def test_prepare_uses_storage_block_for_the_kpool_group( + monkeypatch: pytest.MonkeyPatch, manager_block: int +): + """The bug: on ROCm, kpool groups must get 128/256 - not the 640-4352 + manager block - as their kernel block size. + + Manager-granular selection is what made the block table address manager + blocks while the kpool writer and the index-cache gather read pool pages, + aliasing every cache column past the request's row. + """ + _mock_rocm_platform(monkeypatch) + spec = _kpool_storage_spec(manager_block) + expected_storage = 256 if manager_block % 256 == 0 else 128 + assert spec.storage_block_size == expected_storage + + selected = _prepare( + spec, + [ROCMAiterMLASparseBackend, DeepseekV32IndexerBackend, KpoolTailBackend], + ) + assert selected == [expected_storage] + assert selected != [manager_block] + # The hybrid block-table split stays integral. + assert manager_block % selected[0] == 0 + + +def test_prepare_falls_back_when_storage_block_is_not_supported(): + """No storage block, or a storage block the backends reject (e.g. the + CUDA indexer's exact ``[64]`` vs a 128-token storage block): selection is + `select_common_block_size`'s answer, unchanged from no-fix behavior. + """ + + class Fixed64Backend: + @staticmethod + def get_supported_kernel_block_sizes(): + return [64] + + @staticmethod + def get_name() -> str: + return "FIXED64" + + # Storage block the group's backend does not accept -> backend vote (64). + spec = _kpool_storage_spec(640) + assert spec.storage_block_size == 128 + assert _prepare(spec, [Fixed64Backend]) == [64] + + # No storage block at all (any non-kpool MLA model) -> backend vote. + plain_spec = MLAAttentionSpec( + block_size=640, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + ) + assert plain_spec.storage_block_size is None + assert _prepare(plain_spec, [Fixed64Backend]) == [64] + + # ...and with backends that accept the manager block, the manager wins, + # exactly as on main. + class AcceptAllBackend: + @staticmethod + def get_supported_kernel_block_sizes(): + return [MultipleOf(1)] + + @staticmethod + def get_name() -> str: + return "ACCEPTALL" + + assert _prepare(plain_spec, [AcceptAllBackend]) == [640] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 422ba8885fd7..b61c4cebed18 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -327,6 +327,30 @@ def update_draft_decode_metadata( self.get_metadata_builder().update_draft_decode_metadata(metadata) +def _block_size_is_supported( + backends: list[type[AttentionBackend]], block_size: int +) -> bool: + """Check if the block size is supported by all backends. + + An exact ``int`` declaration must match exactly; a ``MultipleOf`` + declaration accepts any multiple of its base. + """ + for backend in backends: + is_supported = False + for supported_size in backend.get_supported_kernel_block_sizes(): + if isinstance(supported_size, int): + if block_size == supported_size: + is_supported = True + elif isinstance(supported_size, MultipleOf): + if block_size % supported_size.base == 0: + is_supported = True + else: + raise ValueError(f"Unknown supported size: {supported_size}") + if not is_supported: + return False + return True + + def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], @@ -348,27 +372,7 @@ def select_common_block_size( ValueError: If no valid block size found. """ - - def block_size_is_supported( - backends: list[type[AttentionBackend]], block_size: int - ) -> bool: - """Check if the block size is supported by all backends.""" - for backend in backends: - is_supported = False - for supported_size in backend.get_supported_kernel_block_sizes(): - if isinstance(supported_size, int): - if block_size == supported_size: - is_supported = True - elif isinstance(supported_size, MultipleOf): - if block_size % supported_size.base == 0: - is_supported = True - else: - raise ValueError(f"Unknown supported size: {supported_size}") - if not is_supported: - return False - return True - - if block_size_is_supported(backends, kv_manager_block_size): + if _block_size_is_supported(backends, kv_manager_block_size): return kv_manager_block_size # MultipleOf constraints also accept the manager size if they accept a divisor. @@ -381,7 +385,7 @@ def block_size_is_supported( } for size in sorted(candidates, reverse=True): - if block_size_is_supported(backends, size): + if _block_size_is_supported(backends, size): return size raise ValueError( f"No common block size for {kv_manager_block_size} (" @@ -493,9 +497,24 @@ def prepare_kernel_block_sizes( # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] - selected_kernel_size = select_common_block_size( - kv_manager_block_size, group_backends + storage_block_size = ( + kv_cache_spec.storage_block_size + if isinstance(kv_cache_spec, MLAAttentionSpec) + else None ) + if storage_block_size is not None and _block_size_is_supported( + group_backends, storage_block_size + ): + # Storage-block specs (e.g. the GLM-5.3-Flash kpool indexer + # cache) address the cache in pool pages, and every other + # consumer (cache views, metadata builders, hisparse) already + # uses storage_block_size as the kernel block. Fall back to the + # backend vote when the group's backends do not accept it. + selected_kernel_size = storage_block_size + else: + selected_kernel_size = select_common_block_size( + kv_manager_block_size, group_backends + ) kernel_block_sizes.append(selected_kernel_size) elif isinstance(kv_cache_spec, MambaSpec): # This is likely Mamba or other non-attention cache, no splitting. From d33872dfeb79db89b734ce13e07af4630e8bdce2 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 8 Oct 2026 08:48:12 -0600 Subject: [PATCH 210/952] [Feat] Structured decisions endpoint (`/v1/systemone`) (#59299) Signed-off-by: Matt Mastracci Signed-off-by: Lucas Wilkinson Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Chauncey --- docs/serving/online_serving/README.md | 3 + .../online_serving/structured_decisions.md | 110 ++++++++ docs/usage/security.md | 1 + .../generate/structured_decisions/__init__.py | 0 .../test_question_types.py | 105 ++++++++ .../structured_decisions/test_strategies.py | 65 +++++ .../test_structured_decisions_e2e.py | 217 ++++++++++++++++ vllm/entrypoints/generate/api_router.py | 33 +++ .../generate/generative_scoring/serving.py | 75 ++---- vllm/entrypoints/generate/label_reads.py | 76 ++++++ .../generate/structured_decisions/__init__.py | 2 + .../structured_decisions/api_router.py | 53 ++++ .../generate/structured_decisions/protocol.py | 88 +++++++ .../structured_decisions/question_types.py | 160 ++++++++++++ .../generate/structured_decisions/serving.py | 121 +++++++++ .../structured_decisions/strategies.py | 235 ++++++++++++++++++ vllm/entrypoints/launchers/cli_args.py | 2 + 17 files changed, 1285 insertions(+), 61 deletions(-) create mode 100644 docs/serving/online_serving/structured_decisions.md create mode 100644 tests/entrypoints/generate/structured_decisions/__init__.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_question_types.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_strategies.py create mode 100644 tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py create mode 100644 vllm/entrypoints/generate/label_reads.py create mode 100644 vllm/entrypoints/generate/structured_decisions/__init__.py create mode 100644 vllm/entrypoints/generate/structured_decisions/api_router.py create mode 100644 vllm/entrypoints/generate/structured_decisions/protocol.py create mode 100644 vllm/entrypoints/generate/structured_decisions/question_types.py create mode 100644 vllm/entrypoints/generate/structured_decisions/serving.py create mode 100644 vllm/entrypoints/generate/structured_decisions/strategies.py diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 8e1a4f3398ec..fc46a482c157 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -76,6 +76,9 @@ For further details on speech to text, please refer to [this page](speech_to_tex - [Generative Scoring API](generative_scoring.md) (`/generative_scoring`) - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`). - Computes next-token probabilities for specified `label_token_ids`. +- [Structured Decisions API](structured_decisions.md) (`/v1/systemone`) + - Applicable to [text generation models](../../models/generative_models.md) with a [chat template](#chat-template). + - Answers typed questions about a state with a probability for every allowed answer. ## Instrumentator APIs diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md new file mode 100644 index 000000000000..dd26f2130164 --- /dev/null +++ b/docs/serving/online_serving/structured_decisions.md @@ -0,0 +1,110 @@ +# Structured Decisions + +The `/v1/systemone` endpoint answers a set of typed questions about a state and +returns a probability for every allowed answer. + +The endpoint is off unless the server starts with +`--enable-structured-decisions`. With the flag, startup fails for a model that +does not support structured decisions, and with `--logprobs-mode raw_logits` or +`processed_logits`, since `label_mass` needs log probabilities. + +## How it works + +1. Each question is one read. The user message is the state, followed by the + question and its options, labeled `A`, `B`, `C`, ... in order, and + "Answer with the letter of one option only." +2. The read is the reply's first token. Each label must be one token there, + or the request returns 400. +3. The read returns the logprobs of that question's label tokens. Softmax over + the labels gives the answer's probabilities. + +Thinking is off unless `chat_template_kwargs` turns it on. + +This endpoint does not serve diffusion models yet. + +Every read of a request starts with the state, so with +`--enable-prefix-caching` the state is prefilled once and each further +question prefills only its own text, then reads one token. + +## Question types + +| type | criteria | answer | +| --- | --- | --- | +| `choice` | map of option name to a description or `null` | `choice`, `probabilities` by option name, `confidence` | + +A choice has at least one option. A question id is a non-empty string. + +## Example + +```bash +curl -s http://localhost:8000/v1/systemone \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen3-0.6B", + "state": {"ticket": "My card was charged twice for one order."}, + "questions": { + "team": { + "type": "choice", + "instructions": "Which team should handle this ticket?", + "criteria": { + "billing": "payments and refunds", + "shipping": "deliveries", + "security": "account access" + } + } + }, + "chat_template_kwargs": {"enable_thinking": false} + }' +``` + +Response, with probabilities that depend on the model: + +```json +{ + "id": "decision-...", + "object": "structured_decision", + "created": 1790000000, + "model": "Qwen/Qwen3-0.6B", + "answers": { + "team": { + "type": "choice", + "choice": "billing", + "probabilities": {"billing": 0.97, "shipping": 0.01, "security": 0.02}, + "confidence": 0.90 + } + }, + "usage": {"input_tokens": 142, "output_tokens": 1}, + "diagnostics": { + "team": {"label_mass": 0.93, "argmax_is_label": true} + } +} +``` + +`diagnostics.label_mass` is the probability the model put on the labels across +its whole vocabulary. A low value means most of the probability went to tokens +that are not labels. `confidence` is the chosen option's probability over the +whole vocabulary: its share of the labels times `label_mass`. + +## Labels + +Options are labeled `A` to `Z` in the order the request lists them, so a +`choice` has at most 26 options. + +## Limits + +| limit | value | +| --- | --- | +| questions per request | 64 | +| options per `choice` | 26 | + +## Request fields + +| field | meaning | +| --- | --- | +| `state` | what the questions are about: a string, or JSON that is sent as its JSON text | +| `questions` | question id to `{type, instructions, criteria}`, asked in this order | +| `instructions` | optional context placed ahead of the questions | +| `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | + +A question with any field other than `type`, `instructions` and `criteria` is +rejected with a 400. diff --git a/docs/usage/security.md b/docs/usage/security.md index 3dad19d6cb03..c0d2c79700e6 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -175,6 +175,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/responses/{response_id}/cancel` - Cancel a response - `/v1/score` - Scoring API - `/v1/rerank` - Reranking API +- `/v1/systemone` - Structured decisions API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) diff --git a/tests/entrypoints/generate/structured_decisions/__init__.py b/tests/entrypoints/generate/structured_decisions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py new file mode 100644 index 000000000000..e5906accf025 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import pytest + +from vllm.entrypoints.generate.structured_decisions.question_types import ( + QUESTION_TYPES, + Option, + Question, + QuestionType, + StructuredDecisionError, + build_question, + register_question_type, +) + + +def choice(qid="bucket", criteria=None, instructions="Which team?", max_options=128): + return build_question( + qid, + "choice", + instructions, + criteria or {"billing": "money", "outage": None, "other": None}, + max_options, + ) + + +def test_choice_labels_follow_option_order(): + q = choice() + assert q.labels == ("A", "B", "C") + assert [a.name for a in q.options] == ["billing", "outage", "other"] + assert q.options[0].description == "money" + wide = choice(criteria={str(i): None for i in range(26)}) + assert wide.labels[-1] == "Z" + + +def test_choice_prompt(): + q = choice() + assert q.type.prompt(q) == ( + "Question: Which team?\n" + "A: billing - money\n" + "B: outage\n" + "C: other\n" + "Answer with the letter of one option only." + ) + + +def test_choice_answer_shape(): + q = choice() + answer = q.type.answer(q, [0.2, 0.7, 0.1], 0.5) + assert answer == { + "type": "choice", + "choice": "outage", + "probabilities": {"billing": 0.2, "outage": 0.7, "other": 0.1}, + "confidence": 0.35, + } + + +@pytest.mark.parametrize( + "qid,type_name,criteria,match", + [ + ("", "choice", {"x": None, "y": None}, "non-empty"), + ("q", "nope", {"x": None, "y": None}, "unknown question type"), + ("q", "choice", ["x", "y"], "must map option names"), + ], +) +def test_build_question_rejects(qid, type_name, criteria, match): + with pytest.raises(StructuredDecisionError, match=match): + build_question(qid, type_name, "", criteria, 128) + + +def test_one_option_choice(): + q = choice(criteria={"only": None}) + assert q.labels == ("A",) + assert q.type.answer(q, [1.0], 0.4)["confidence"] == 0.4 + + +def test_option_limit(): + with pytest.raises(StructuredDecisionError, match="at most 2 options"): + choice(criteria={"x": None, "y": None, "z": None}, max_options=2) + with pytest.raises(StructuredDecisionError, match="at most 26 options"): + choice(criteria={str(i): None for i in range(27)}) + + +def test_registered_type_plugs_in(): + class BinaryQuestion(QuestionType): + name = "test_binary" + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + return [Option("yes"), Option("no")] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + return {"type": self.name, "yes": probs[0]} + + register_question_type(BinaryQuestion) + try: + q = build_question("ok", "test_binary", "Is it fine?", None, 128) + assert q.labels == ("A", "B") + assert q.type.answer(q, [0.9, 0.1], 1.0) == {"type": "test_binary", "yes": 0.9} + with pytest.raises(ValueError, match="already registered"): + register_question_type(BinaryQuestion) + finally: + del QUESTION_TYPES["test_binary"] diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py new file mode 100644 index 000000000000..c4f4bce61be1 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from typing import Any + +import pytest +from fastapi import FastAPI + +from vllm.entrypoints.generate.structured_decisions.api_router import ( + register_structured_decisions_api_router, +) +from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.strategies import ( + NextTokenStrategy, + reply_label_ids, + select_read_strategy, +) + + +def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: + return SimpleNamespace(architecture=architecture, logprobs_mode=logprobs_mode) + + +def test_strategy_selection(): + qwen = "Qwen3ForCausalLM" + assert select_read_strategy(model(qwen)) is NextTokenStrategy + assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): + select_read_strategy(model("LlamaForCausalLM")) + with pytest.raises(ValueError, match="not raw_logits"): + select_read_strategy(model(qwen, "raw_logits")) + + +def test_route_needs_the_flag(): + for enabled in (False, True): + app = FastAPI() + app.state.args = SimpleNamespace(enable_structured_decisions=enabled) + register_structured_decisions_api_router(app) + paths = {getattr(route, "path", None) for route in app.routes} + assert ("/v1/systemone" in paths) == enabled + + +@pytest.fixture(scope="module") +def qwen(): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + prompt_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "state"}], + add_generation_prompt=True, + enable_thinking=False, + tokenize=True, + return_dict=False, + ) + return tokenizer, prompt_ids + + +def test_labels_start_the_reply(qwen): + tokenizer, prompt_ids = qwen + tail, ids = reply_label_ids(tokenizer, prompt_ids) + assert tokenizer.decode(tail) == "\n\n" + assert [tokenizer.decode([i]) for i in ids] == list(LABELS) + # After a colon, Qwen writes ":A" as one token, so "A" is not one token. + with pytest.raises(ValueError, match="not one distinct token"): + reply_label_ids(tokenizer, tokenizer.encode("team:")) diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py new file mode 100644 index 000000000000..e060b096ec51 --- /dev/null +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import requests + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen3-0.6B" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "1024", + "--enforce-eager", + "--max-num-seqs", + "32", + "--enable-prefix-caching", + "--enable-structured-decisions", + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +def post(server, body): + return requests.post(server.url_for("v1/systemone"), json=body) + + +def choice(instructions, *names): + return { + "type": "choice", + "instructions": instructions, + "criteria": dict.fromkeys(names), + } + + +COPY = { + "team": choice( + "Which team does the message name?", "billing", "shipping", "security" + ), + "lang": choice("Which language does the message name?", "English", "French"), +} + + +def test_choice_decision(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": {"team": "billing", "language": "French"}, + "questions": COPY, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["object"] == "structured_decision" + assert list(body["answers"]) == ["team", "lang"] + team = body["answers"]["team"] + assert team["type"] == "choice" + assert set(team["probabilities"]) == {"billing", "shipping", "security"} + assert sum(team["probabilities"].values()) == pytest.approx(1.0, abs=1e-4) + assert team["confidence"] == pytest.approx( + max(team["probabilities"].values()) * body["diagnostics"]["team"]["label_mass"] + ) + for diag in body["diagnostics"].values(): + assert 0.0 < diag["label_mass"] <= 1.0 + 1e-6 + assert body["usage"]["output_tokens"] == 2 + assert body["usage"]["input_tokens"] > 0 + + +def test_one_option_choice(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": "Team: billing.", + "questions": { + "team": choice("Which team does the message name?", "billing") + }, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + team = body["answers"]["team"] + assert team["choice"] == "billing" + assert team["probabilities"] == {"billing": 1.0} + assert team["confidence"] == pytest.approx( + body["diagnostics"]["team"]["label_mass"] + ) + + +def test_one_option_coin_flip(server): + # A fair coin with only "heads" allowed. The answer is forced, so its + # confidence must come from the model's probability of the label. + response = post( + server, + { + "model": MODEL_NAME, + "state": "I flip a fair coin.", + "questions": {"coin": choice("How did the coin land?", "heads")}, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + coin = body["answers"]["coin"] + assert coin["probabilities"] == {"heads": 1.0} + assert coin["confidence"] == pytest.approx( + body["diagnostics"]["coin"]["label_mass"] + ) + + +@pytest.mark.parametrize( + "state,questions,expected", + [ + ( + {"team": "billing", "language": "French"}, + COPY, + {"team": "billing", "lang": "French"}, + ), + ( + "My card was charged twice for one order.", + { + "card": choice("Does the message mention a card?", "yes", "no"), + "mood": choice("How does the customer feel?", "happy", "angry"), + }, + {"card": "yes", "mood": "angry"}, + ), + ( + "Hello.", + { + "sky": choice( + "What color is the sky on a clear day?", "blue", "green", "red" + ) + }, + {"sky": "blue"}, + ), + ], +) +def test_single_read_answers(server, state, questions, expected): + response = post( + server, + { + "model": MODEL_NAME, + "state": state, + "questions": questions, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + answers = response.json()["answers"] + assert {qid: a["choice"] for qid, a in answers.items()} == expected + + +def test_option_limit(server): + body = { + "model": MODEL_NAME, + "state": "Pick option 7.", + "chat_template_kwargs": {"enable_thinking": False}, + } + ok = post( + server, {**body, "questions": {"n": choice("Which?", *map(str, range(26)))}} + ) + assert ok.status_code == 200, ok.text + assert len(ok.json()["answers"]["n"]["probabilities"]) == 26 + over = post( + server, {**body, "questions": {"n": choice("Which?", *map(str, range(27)))}} + ) + assert over.status_code == 400 + assert "at most 26 options" in over.json()["error"]["message"] + + +def test_thinking_is_refused(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": "x", + "questions": {"q": choice("Which?", "a", "b")}, + "chat_template_kwargs": {"enable_thinking": True}, + }, + ) + assert response.status_code == 400 + assert "thinking must be off" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + "questions,match", + [ + ( + { + "q": { + "type": "choice", + "criteria": {"a": None, "b": None}, + "depends_on": ["x"], + } + }, + "unknown field(s)", + ), + ( + { + f"q{i}": {"type": "choice", "criteria": {"a": None, "b": None}} + for i in range(65) + }, + "at most 64", + ), + ], +) +def test_validation_errors(server, questions, match): + response = post(server, {"model": MODEL_NAME, "state": "x", "questions": questions}) + assert response.status_code == 400 + assert match in response.json()["error"]["message"] diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index c5703515d82b..70869b61b5d2 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -53,6 +53,12 @@ def register_generate_api_routers(app: FastAPI): register_generative_scoring_api_router(app) + from .structured_decisions.api_router import ( + register_structured_decisions_api_router, + ) + + register_structured_decisions_api_router(app) + async def init_generate_state( engine_client: "EngineClient", @@ -220,3 +226,30 @@ async def init_generate_state( state.openai_serving_models, request_logger=request_logger, ) + + from .structured_decisions.serving import ServingStructuredDecisions + from .structured_decisions.strategies import ReadContext, select_read_strategy + + strategy_cls = ( + select_read_strategy(engine_client.model_config) + if "generate" in supported_tasks + and getattr(args, "enable_structured_decisions", False) + else None + ) + state.serving_structured_decisions = ( + ServingStructuredDecisions( + state.openai_serving_models, + strategy_cls( + ReadContext( + engine_client=engine_client, + online_renderer=state.online_renderer, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + ), + request_logger=request_logger, + ) + if strategy_cls is not None + else None + ) diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 4d5a15c2f54c..bd46158edcab 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -11,13 +11,14 @@ import asyncio import math import time -from collections.abc import AsyncGenerator, Mapping +from collections.abc import Mapping from typing import Literal from fastapi import Request from pydantic import Field from vllm.engine.protocol import EngineClient +from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.protocol import ( ErrorResponse, @@ -28,7 +29,6 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import EngineInput, tokens_input from vllm.logger import init_logger -from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tracing import ( @@ -37,7 +37,6 @@ log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.async_utils import merge_async_iterators logger = init_logger(__name__) @@ -271,37 +270,28 @@ async def create_generative_scoring( else await self._get_trace_headers(raw_request.headers) ) - # Schedule requests for all inputs - generators: list[AsyncGenerator[RequestOutput, None]] = [] for i, engine_input in enumerate(engine_inputs): - request_id_item = f"{request_id}-{i}" - self._log_inputs( - request_id_item, + f"{request_id}-{i}", engine_input, params=sampling_params, lora_request=lora_request, ) - generator = self.engine_client.generate( - engine_input, - sampling_params, - request_id_item, + try: + reads = await next_token_label_reads( + self.engine_client, + engine_inputs, + [sampling_params] * len(engine_inputs), + request_id, lora_request=lora_request, trace_headers=trace_headers, priority=request.priority, ) - generators.append(generator) - - # Collect results - result_generator = merge_async_iterators(*generators) - results: list[RequestOutput | None] = [None] * len(engine_inputs) - - try: - async for i, res in result_generator: - results[i] = res except asyncio.CancelledError: return self.create_error_response("Client disconnected") + except ValueError as e: + return self.create_error_response(e) except Exception as e: logger.exception("Error during generation") return self.create_error_response(e) @@ -311,45 +301,8 @@ async def create_generative_scoring( total_prompt_tokens = 0 total_completion_tokens = 0 - for i, result in enumerate(results): - if result is None: - return self.create_error_response( - f"Failed to generate result for item {i}" - ) - - # Check for errors - if result.outputs and result.outputs[0].finish_reason == "error": - return self.create_error_response(f"Generation error for item {i}") - - # Get logprobs from the generated token - if not result.outputs or len(result.outputs) == 0: - return self.create_error_response(f"No output generated for item {i}") - - output = result.outputs[0] - if output.logprobs is None or len(output.logprobs) == 0: - return self.create_error_response( - f"No logprobs available for item {i}. " - "This might indicate an issue with logprobs configuration." - ) - - # The logprobs dict maps token_id -> Logprob object - # For logprobs=-1, this contains all vocab tokens - logprobs_dict = output.logprobs[0] - - # Extract logprobs for label tokens - label_logprobs: dict[int, float] = {} - missing_tokens = [] - for token_id in request.label_token_ids: - if token_id in logprobs_dict: - label_logprobs[token_id] = logprobs_dict[token_id].logprob - else: - missing_tokens.append(token_id) - - if missing_tokens: - return self.create_error_response( - f"Token IDs {missing_tokens} not found in logprobs for item {i}. " - "This might indicate the tokens are outside the model's vocabulary." - ) + for i, read in enumerate(reads): + label_logprobs = dict(zip(request.label_token_ids, read.logprobs)) # Compute probabilities based on apply_softmax setting token_probs = self._compute_probabilities( @@ -370,7 +323,7 @@ async def create_generative_scoring( # Update token counts total_prompt_tokens += prompt_token_counts[i] - total_completion_tokens += len(output.token_ids) + total_completion_tokens += len(read.result.outputs[0].token_ids) # Build response model_name = self.models.model_name(lora_request) diff --git a/vllm/entrypoints/generate/label_reads.py b/vllm/entrypoints/generate/label_reads.py new file mode 100644 index 000000000000..eda61b22bc5a --- /dev/null +++ b/vllm/entrypoints/generate/label_reads.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-token reads that return the logprobs of chosen label tokens.""" + +from collections.abc import AsyncGenerator, Mapping, Sequence +from dataclasses import dataclass + +from vllm.engine.protocol import EngineClient +from vllm.inputs import EngineInput +from vllm.lora.request import LoRARequest +from vllm.outputs import RequestOutput +from vllm.sampling_params import SamplingParams +from vllm.utils.async_utils import merge_async_iterators + + +@dataclass(frozen=True) +class LabelRead: + result: RequestOutput + logprobs: list[float] + """Full-vocabulary logprob of each label token, in ``logprob_token_ids`` + order.""" + + +async def next_token_label_reads( + engine_client: EngineClient, + engine_inputs: Sequence[EngineInput], + sampling_params: Sequence[SamplingParams], + request_id: str, + *, + lora_request: LoRARequest | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, +) -> list[LabelRead]: + """Runs one read per input at once, as request ``{request_id}-{i}``. Each + read's params set ``max_tokens=1`` and the label tokens as + ``logprob_token_ids``. Raises ValueError naming the item when a read has + no output or lacks a label's logprob.""" + generators: list[AsyncGenerator[RequestOutput, None]] = [ + engine_client.generate( + engine_input, + params, + f"{request_id}-{i}", + lora_request=lora_request, + trace_headers=trace_headers, + priority=priority, + ) + for i, (engine_input, params) in enumerate(zip(engine_inputs, sampling_params)) + ] + results: list[RequestOutput | None] = [None] * len(generators) + async for i, res in merge_async_iterators(*generators): + results[i] = res + + reads = [] + for i, (result, params) in enumerate(zip(results, sampling_params)): + if result is None: + raise ValueError(f"Failed to generate result for item {i}") + if not result.outputs: + raise ValueError(f"No output generated for item {i}") + output = result.outputs[0] + if output.finish_reason == "error": + raise ValueError(f"Generation error for item {i}") + if not output.logprobs: + raise ValueError( + f"No logprobs available for item {i}. " + "This might indicate an issue with logprobs configuration." + ) + logprobs = output.logprobs[0] + label_ids = params.logprob_token_ids or [] + missing = [t for t in label_ids if t not in logprobs] + if missing: + raise ValueError( + f"Token IDs {missing} not found in logprobs for item {i}. " + "This might indicate the tokens are outside the model's vocabulary." + ) + reads.append(LabelRead(result, [logprobs[t].logprob for t in label_ids])) + return reads diff --git a/vllm/entrypoints/generate/structured_decisions/__init__.py b/vllm/entrypoints/generate/structured_decisions/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py new file mode 100644 index 000000000000..a0fbaa518e07 --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from http import HTTPStatus + +from fastapi import APIRouter, Depends, FastAPI, Request +from fastapi.responses import JSONResponse + +from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + validate_json_request, + with_cancellation, +) + +from .protocol import StructuredDecisionRequest +from .serving import ServingStructuredDecisions + +router = APIRouter() + + +def structured_decisions(request: Request) -> ServingStructuredDecisions | None: + return getattr(request.app.state, "serving_structured_decisions", None) + + +@router.post( + "/v1/systemone", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse}, + }, +) +@with_cancellation +@load_aware_call +async def create_structured_decision( + request: StructuredDecisionRequest, raw_request: Request +): + handler = structured_decisions(raw_request) + if handler is None: + raise NotImplementedError( + "The model does not support the structured decision API" + ) + result = await handler.create_decision(request, raw_request) + if isinstance(result, ErrorResponse): + return JSONResponse(content=result.model_dump(), status_code=result.error.code) + return JSONResponse(content=result.model_dump()) + + +def register_structured_decisions_api_router(app: FastAPI): + if getattr(app.state.args, "enable_structured_decisions", False): + app.include_router(router) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py new file mode 100644 index 000000000000..3a193e203bad --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request and response models of the structured decisions API (/v1/systemone).""" + +import time +from typing import Any, Literal + +from pydantic import Field + +from vllm.config import ModelConfig +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel +from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.utils import random_uuid + + +class QuestionSpec(OpenAIBaseModel): + type: str + instructions: Any = "" + criteria: Any = None + + +class StructuredDecisionRequest(OpenAIBaseModel): + model: str | None = None + state: Any = Field(..., description="What the questions are about: text or JSON.") + questions: dict[str, QuestionSpec] = Field( + ..., description="Question id to question, asked in this order." + ) + instructions: str | None = Field( + default=None, description="Context placed ahead of the questions." + ) + chat_template_kwargs: dict[str, Any] | None = None + priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1) + request_id: str = Field(default_factory=random_uuid) + + +class DecisionUsage(OpenAIBaseModel): + input_tokens: int + output_tokens: int + + +class QuestionDiagnostics(OpenAIBaseModel): + label_mass: float = Field( + description="Probability the model put on the labels, over the full " + "vocabulary. A low value means most of the probability went to tokens " + "that are not labels." + ) + argmax_is_label: bool + + +class StructuredDecisionResponse(OpenAIBaseModel): + id: str + object: Literal["structured_decision"] = "structured_decision" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + answers: dict[str, dict[str, Any]] + usage: DecisionUsage + diagnostics: dict[str, QuestionDiagnostics] + + +class ReadPromptRequest(OpenAIBaseModel): + """Chat options for one read's prompt: the state and the question in the + user turn, ending at the generation prompt so the label is the reply's + first token. Thinking is off unless the request turns it on, or the label + would follow a thought rather than start the reply.""" + + chat_template_kwargs: dict[str, Any] | None = None + + def build_chat_params( + self, + default_template: str | None, + default_template_content_format: ChatTemplateContentFormatOption, + ) -> ChatParams: + return ChatParams( + chat_template=default_template, + chat_template_content_format=default_template_content_format, + chat_template_kwargs=merge_kwargs( + merge_kwargs({"enable_thinking": False}, self.chat_template_kwargs), + dict(add_generation_prompt=True, continue_final_message=False), + ), + ) + + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: + return TokenizeParams( + max_total_tokens=model_config.max_model_len, + max_output_tokens=1, + add_special_tokens=False, + ) diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py new file mode 100644 index 000000000000..a9a6514a0373 --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Question types for structured decisions. + +A question type turns a request's criteria into options, writes the question +for the model with a label per option, and builds the answer from the label +probabilities. +""" + +import math +import string +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, ClassVar + + +class StructuredDecisionError(ValueError): + """An invalid request. Returned as a 400.""" + + +#: Options are labeled A to Z in order. A single capital letter is one token +#: at the start of a reply for the tokenizers tested; every read checks it. +LABELS = tuple(string.ascii_uppercase) + + +@dataclass(frozen=True) +class Option: + name: str + description: str | None = None + + +@dataclass(frozen=True) +class Question: + id: str + type: "QuestionType" + instructions: str + options: tuple[Option, ...] + labels: tuple[str, ...] + + +class QuestionType(ABC): + name: ClassVar[str] + + @abstractmethod + def parse_options(self, qid: str, criteria: Any) -> list[Option]: ... + + @abstractmethod + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + """The answer for ``question``. ``probs[i]`` is the probability of + ``question.labels[i]`` among the labels, and the list sums to 1. + ``label_mass`` is the labels' total probability over the vocabulary.""" + + def prompt(self, question: Question) -> str: + """The question as the model reads it, after the state.""" + lines = [f"Question: {question.instructions}"] if question.instructions else [] + for label, o in zip(question.labels, question.options): + lines.append( + f"{label}: {o.name} - {o.description}" + if o.description + else f"{label}: {o.name}" + ) + lines.append("Answer with the letter of one option only.") + return "\n".join(lines) + + +QUESTION_TYPES: dict[str, QuestionType] = {} + + +def register_question_type(cls: type[QuestionType]) -> type[QuestionType]: + if cls.name in QUESTION_TYPES: + raise ValueError(f"question type {cls.name!r} is already registered") + QUESTION_TYPES[cls.name] = cls() + return cls + + +def get_question_type(name: str) -> QuestionType: + try: + return QUESTION_TYPES[name] + except KeyError: + raise StructuredDecisionError( + f"unknown question type {name!r}; supported: {sorted(QUESTION_TYPES)}" + ) from None + + +def build_question( + qid: str, + type_name: str, + instructions: Any, + criteria: Any, + max_options: int, +) -> Question: + if not qid: + raise StructuredDecisionError("question ids must be non-empty") + qtype = get_question_type(type_name) + options = qtype.parse_options(qid, criteria) + if not options: + raise StructuredDecisionError(f"question {qid!r}: needs at least one option") + names = [o.name for o in options] + if len(set(names)) != len(names): + raise StructuredDecisionError(f"question {qid!r}: duplicate option names") + limit = min(max_options, len(LABELS)) + if len(options) > limit: + raise StructuredDecisionError( + f"question {qid!r}: at most {limit} options for this model" + ) + if not isinstance(instructions, str): + instructions = "" if instructions is None else str(instructions) + return Question( + id=qid, + type=qtype, + instructions=instructions, + options=tuple(options), + labels=LABELS[: len(options)], + ) + + +def argmax(values: list[float]) -> int: + return max(range(len(values)), key=values.__getitem__) + + +def label_softmax(logprobs: list[float]) -> list[float]: + """Softmax of the label logprobs.""" + top = max(logprobs) + weights = [math.exp(lp - top) for lp in logprobs] + total = sum(weights) + return [w / total for w in weights] + + +@register_question_type +class ChoiceQuestion(QuestionType): + """Pick one option. ``criteria`` maps each option name to a description, + or to null.""" + + name = "choice" + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if not isinstance(criteria, dict) or not criteria: + raise StructuredDecisionError( + f"question {qid!r}: choice criteria must map option names to " + "a description or null" + ) + return [ + Option(str(name), None if desc is None else str(desc)) + for name, desc in criteria.items() + ] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + top = argmax(probs) + return { + "type": self.name, + "choice": question.options[top].name, + "probabilities": {a.name: p for a, p in zip(question.options, probs)}, + # The chosen label's probability over the whole vocabulary, so a + # read that mostly wanted a non-label token reports low confidence. + "confidence": probs[top] * label_mass, + } diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py new file mode 100644 index 000000000000..7e2136b8a25e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import json +from typing import Any + +from fastapi import Request + +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.serving import BaseServing +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.logger import init_logger + +from .protocol import ( + DecisionUsage, + QuestionDiagnostics, + StructuredDecisionRequest, + StructuredDecisionResponse, +) +from .question_types import Question, StructuredDecisionError, build_question +from .strategies import DecisionLimits, ReadStrategy + +logger = init_logger(__name__) + + +def state_text(state: Any) -> str: + return state if isinstance(state, str) else json.dumps(state) + + +def parse_questions( + request: StructuredDecisionRequest, + limits: DecisionLimits, +) -> list[Question]: + if not request.questions: + raise StructuredDecisionError("questions: needs at least one question") + if len(request.questions) > limits.max_questions: + raise StructuredDecisionError( + f"questions: at most {limits.max_questions} for this model" + ) + questions = [] + for qid, spec in request.questions.items(): + if spec.model_extra: + raise StructuredDecisionError( + f"question {qid!r}: unknown field(s) {sorted(spec.model_extra)}" + ) + questions.append( + build_question( + qid, + spec.type, + spec.instructions, + spec.criteria, + limits.max_options, + ) + ) + return questions + + +class ServingStructuredDecisions(BaseServing): + def __init__( + self, + models: OpenAIServingModels, + strategy: ReadStrategy, + *, + request_logger: RequestLogger | None = None, + ) -> None: + model_config = strategy.context.engine_client.model_config + super().__init__( + models=models, model_config=model_config, request_logger=request_logger + ) + self.strategy = strategy + self.limits = strategy.limits() + + async def create_decision( + self, + request: StructuredDecisionRequest, + raw_request: Request | None = None, + ) -> StructuredDecisionResponse | ErrorResponse: + if (error := await self._check_model(request)) is not None: # type: ignore[arg-type] + return error + engine_client = self.strategy.context.engine_client + if engine_client.errored: + raise engine_client.dead_error + + base_id = self._base_request_id(raw_request, default=request.request_id) + request_id = f"decision-{base_id}" + try: + questions = parse_questions(request, self.limits) + lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] + engine_client.check_admission(len(questions)) + reads = await self.strategy.read( + questions, + request.instructions, + state_text(request.state), + request_id=request_id, + chat_template_kwargs=request.chat_template_kwargs, + lora_request=lora_request, + priority=request.priority, + ) + except StructuredDecisionError as e: + return self.create_error_response(e) + except asyncio.CancelledError: + return self.create_error_response("Client disconnected") + + answers: dict[str, dict[str, Any]] = {} + diagnostics: dict[str, QuestionDiagnostics] = {} + for q, read in zip(questions, reads): + answers[q.id] = q.type.answer(q, read.probs, read.label_mass) + diagnostics[q.id] = QuestionDiagnostics( + label_mass=read.label_mass, argmax_is_label=read.argmax_is_label + ) + return StructuredDecisionResponse( + id=request_id, + model=self.models.model_name(lora_request), + answers=answers, + usage=DecisionUsage( + input_tokens=sum(r.input_tokens for r in reads), + output_tokens=sum(r.output_tokens for r in reads), + ), + diagnostics=diagnostics, + ) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py new file mode 100644 index 000000000000..22233ff91b0e --- /dev/null +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read strategies: how a model is asked for label probabilities. + +Only models in NEXT_TOKEN_ARCHITECTURES are currently supported. +""" + +import math +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.generate.label_reads import next_token_label_reads +from vllm.inputs import EngineInput +from vllm.lora.request import LoRARequest +from vllm.renderers.inputs.preprocess import extract_prompt_components +from vllm.renderers.online_renderer import OnlineRenderer +from vllm.sampling_params import SamplingParams +from vllm.tokenizers import TokenizerLike + +from .protocol import ReadPromptRequest +from .question_types import LABELS, Question, StructuredDecisionError, label_softmax + + +@dataclass +class ReadContext: + engine_client: EngineClient + online_renderer: OnlineRenderer + chat_template: str | None + chat_template_content_format: ChatTemplateContentFormatOption + default_chat_template_kwargs: dict[str, Any] + + +@dataclass(frozen=True) +class DecisionLimits: + max_questions: int + max_options: int + + +@dataclass +class QuestionRead: + probs: list[float] + label_mass: float + argmax_is_label: bool + input_tokens: int + output_tokens: int + + +class ReadStrategy(ABC): + def __init__(self, context: ReadContext): + self.context = context + + @abstractmethod + def limits(self) -> DecisionLimits: ... + + @abstractmethod + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + ) -> list[QuestionRead]: + """One read per question, in order.""" + + +def reply_label_ids( + tokenizer: TokenizerLike, prompt_ids: Sequence[int] +) -> tuple[list[int], list[int]]: + """The prompt's ids after its last added token, and the token each of LABELS + adds after them as the first token of the reply. Raises ValueError if a + label is not one distinct token there.""" + added = set(tokenizer.get_added_vocab().values()) + start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) + tail = list(prompt_ids[start:]) + text = tokenizer.decode(tail) + ids: list[int] = [] + for label in LABELS: + extended = tokenizer.encode(text + label, add_special_tokens=False) + if extended[:-1] != tail or extended[-1] in ids: + raise ValueError( + f"label {label!r} is not one distinct token after this model's " + "chat prompt" + ) + ids.append(extended[-1]) + return tail, ids + + +class NextTokenStrategy(ReadStrategy): + """Autoregressive models. Each question is one request: the state, then the + question with its labeled options, and the logprobs of the label tokens as + the reply's first token. The requests share the state. With prefix caching, + it is prefilled once.""" + + def __init__(self, context: ReadContext): + super().__init__(context) + tokenizer = context.online_renderer.renderer.get_tokenizer() + probe = tokenizer.apply_chat_template( + [{"role": "user", "content": "x"}], + **{ + **context.default_chat_template_kwargs, + "chat_template": context.chat_template, + "add_generation_prompt": True, + "tokenize": True, + "return_dict": False, + "enable_thinking": False, + }, + ) + if isinstance(probe, str): + probe = tokenizer.encode(probe, add_special_tokens=False) + # Every prompt ends with the same generation prompt, so the labels' + # tokens are the same for every question. + self.tail, self.label_ids = reply_label_ids(tokenizer, probe) + + def limits(self) -> DecisionLimits: + return DecisionLimits(max_questions=64, max_options=len(LABELS)) + + def _read_request( + self, chat_template_kwargs: dict[str, Any] | None + ) -> ReadPromptRequest: + if (chat_template_kwargs or {}).get("enable_thinking"): + raise StructuredDecisionError( + "a read is the reply's first token, so thinking must be off" + ) + return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs) + + async def _render( + self, read_request: ReadPromptRequest, messages: list[dict[str, Any]] + ) -> tuple[EngineInput, list[int]]: + ctx = self.context + _, (engine_input,) = await ctx.online_renderer.preprocess_chat( + read_request, + messages, + default_template=ctx.chat_template, + default_template_content_format=ctx.chat_template_content_format, + default_template_kwargs=ctx.default_chat_template_kwargs, + ) + prompt_ids = extract_prompt_components( + ctx.engine_client.model_config, engine_input + ).token_ids + return engine_input, list(prompt_ids or []) + + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + ) -> list[QuestionRead]: + ctx = self.context + read_request = self._read_request(chat_template_kwargs) + + slots, engine_inputs = [], [] + for q in questions: + messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] + if instructions: + messages.insert(0, {"role": "system", "content": instructions}) + engine_input, prompt_ids = await self._render(read_request, messages) + if prompt_ids[len(prompt_ids) - len(self.tail) :] != self.tail: + raise StructuredDecisionError( + "these chat options end the prompt differently, so the " + "labels' tokens are unknown" + ) + slots.append(self.label_ids[: len(q.labels)]) + engine_inputs.append(engine_input) + + label_reads = await next_token_label_reads( + ctx.engine_client, + engine_inputs, + [ + SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids) + for ids in slots + ], + request_id, + lora_request=lora_request, + priority=priority, + ) + + reads = [] + for ids, label_read in zip(slots, label_reads): + output = label_read.result.outputs[0] + reads.append( + QuestionRead( + probs=label_softmax(label_read.logprobs), + label_mass=sum(math.exp(lp) for lp in label_read.logprobs), + argmax_is_label=bool(output.token_ids) + and output.token_ids[0] in ids, + input_tokens=len(label_read.result.prompt_token_ids or ()), + output_tokens=len(output.token_ids), + ) + ) + return reads + + +NEXT_TOKEN_ARCHITECTURES = frozenset( + { + "Qwen3ForCausalLM", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + } +) + + +# label_mass sums the labels' full-vocabulary probabilities. The logprobs modes +# return those probabilities. The logits modes return raw logits instead. +LOGPROBS_MODES = frozenset({"raw_logprobs", "processed_logprobs"}) + + +def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: + """Raises ValueError when the model cannot serve structured decisions. The + server opted in with --enable-structured-decisions, so startup fails.""" + if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: + raise ValueError( + "--enable-structured-decisions does not support " + f"{model_config.architecture}. Supported architectures: " + f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" + ) + if model_config.logprobs_mode not in LOGPROBS_MODES: + raise ValueError( + "--enable-structured-decisions needs --logprobs-mode raw_logprobs or " + f"processed_logprobs, not {model_config.logprobs_mode}" + ) + return NextTokenStrategy diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 28fe79fd691e..18f8cb00b503 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -163,6 +163,8 @@ class BaseFrontendArgs: enable_tokenizer_info_endpoint: bool = False """Enable the `/tokenizer_info` endpoint. May expose chat templates and other tokenizer configuration.""" + enable_structured_decisions: bool = False + """Enable the `/v1/systemone` structured decisions endpoint.""" enable_log_outputs: bool = False """If set to True, log model outputs (generations). Requires `--enable-log-requests`. Output text and finish reasons are logged at INFO, From d80ddaac73217770af1548f43e0cd7e68d4ec33d Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:55:20 +0800 Subject: [PATCH 211/952] [ROCm][Perf] Kimi-K3 Store only the current rank's shards in latent MoE up-proj (#59591) Signed-off-by: kliuae --- .../kimi_k3/test_amd_latent_moe_runner.py | 96 +++++++++++++++---- vllm/models/kimi_k3/amd/latent_moe_runner.py | 16 +++- vllm/models/kimi_k3/amd/linear.py | 41 ++++++-- 3 files changed, 125 insertions(+), 28 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index fd5bcfd22b78..e47b5fc08e87 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -22,7 +22,7 @@ from vllm.distributed import get_tp_group from vllm.model_executor.layers.fused_moe.runner import moe_runner from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear from vllm.models.kimi_k3.amd import latent_moe_runner from vllm.models.kimi_k3.amd.latent_moe_runner import ROCmLatentMoERunner from vllm.models.kimi_k3.amd.linear import KimiRoutedOutputTransform @@ -40,20 +40,42 @@ DTYPE = torch.bfloat16 -def _build_transform(device: torch.device) -> KimiRoutedOutputTransform: +def _build_transform( + device: torch.device, *, row_sharded: bool = False +) -> tuple[KimiRoutedOutputTransform, torch.Tensor]: norm = RMSNorm(LATENT_SIZE, eps=EPS).to(device=device, dtype=DTYPE) - up_proj = ReplicatedLinear( - LATENT_SIZE, - HIDDEN_SIZE, - bias=False, - params_dtype=DTYPE, - prefix="routed_expert_up_proj", - ).to(device=device) torch.manual_seed(0) norm.weight.data.copy_(1 + 0.1 * torch.randn_like(norm.weight)) - up_proj.weight.data.copy_(torch.randn_like(up_proj.weight) / LATENT_SIZE**0.5) - return KimiRoutedOutputTransform(norm, up_proj) + full_weight = ( + torch.randn(HIDDEN_SIZE, LATENT_SIZE, device=device, dtype=DTYPE) + / LATENT_SIZE**0.5 + ) + + up_proj: ReplicatedLinear | ColumnParallelLinear + if row_sharded: + up_proj = ColumnParallelLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + gather_output=True, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + up_proj.weight_loader(up_proj.weight, full_weight) + else: + up_proj = ReplicatedLinear( + LATENT_SIZE, + HIDDEN_SIZE, + bias=False, + params_dtype=DTYPE, + prefix="routed_expert_up_proj", + ).to(device=device) + up_proj.weight.data.copy_(full_weight) + return ( + KimiRoutedOutputTransform(norm, up_proj, row_sharded=row_sharded), + full_weight, + ) def _tail_runner( @@ -72,6 +94,7 @@ def _tail_runner( runner = object.__new__(ROCmLatentMoERunner) attrs = { "routed_output_transform": transform, + "_up_proj_preshard": bool(getattr(transform, "row_sharded", False)), "_up_proj_shard_size": HIDDEN_SIZE // tp_world, "_logged_sharded_tail": False, "moe_config": SimpleNamespace( @@ -101,9 +124,13 @@ def _rank_partials( def _check_matches_replicated( - device: torch.device, tp_world: int, rank: int, use_ep: bool = False + device: torch.device, + tp_world: int, + rank: int, + use_ep: bool = False, + row_sharded: bool = False, ) -> None: - transform = _build_transform(device) + transform, up_weight = _build_transform(device, row_sharded=row_sharded) runner = _tail_runner(transform, tp_world, use_ep) group = get_tp_group().device_group @@ -118,7 +145,7 @@ def _check_matches_replicated( transform.norm.weight, EPS, ), - transform.up_proj.weight, + up_weight, ) expected.add_(_all_reduced(shared_output, group)) @@ -136,12 +163,37 @@ def _check_matches_replicated_under_ep( _check_matches_replicated(device, tp_world, rank, use_ep=True) +def _check_presharded_matches_replicated( + device: torch.device, tp_world: int, rank: int +) -> None: + _check_matches_replicated(device, tp_world, rank, row_sharded=True) + + +def _check_unfused_tail_gathers_full_hidden( + device: torch.device, tp_world: int, rank: int +) -> None: + transform, up_weight = _build_transform(device, row_sharded=True) + assert transform.up_proj.weight.shape[0] == HIDDEN_SIZE // tp_world + + # all-gather reconstructs the full projection when all ranks feed the same + # input, so use one seed for all ranks. + torch.manual_seed(1234) + latent = torch.randn(4, LATENT_SIZE, device=device, dtype=DTYPE).mul_(0.01) + + torch.testing.assert_close( + transform(latent), + F.linear(transform.norm(latent), up_weight), + atol=8e-2, + rtol=3e-2, + ) + + def _check_writes_only_its_own_shard( device: torch.device, tp_world: int, rank: int ) -> None: """Two ranks that swapped both weight rows and write offsets still sum to the right total, so inspect each slice before the final collective.""" - transform = _build_transform(device) + transform, up_weight = _build_transform(device) runner = _tail_runner(transform, tp_world) group = get_tp_group().device_group @@ -169,7 +221,7 @@ def _capture(states, trunc_size, output_is_reduced=None): shard = HIDDEN_SIZE // tp_world start, end = rank * shard, (rank + 1) * shard local = captured["states"] - projected = F.linear(latent, transform.up_proj.weight[start:end]) + projected = F.linear(latent, up_weight[start:end]) torch.testing.assert_close( local[:, start:end], before[:, start:end] + projected, atol=8e-2, rtol=3e-2 @@ -181,6 +233,8 @@ def _capture(states, trunc_size, output_is_reduced=None): _CHECKS = { "matches_replicated": _check_matches_replicated, "matches_replicated_ep": _check_matches_replicated_under_ep, + "presharded_matches_replicated": _check_presharded_matches_replicated, + "unfused_gathers": _check_unfused_tail_gathers_full_hidden, "own_shard_only": _check_writes_only_its_own_shard, } @@ -219,6 +273,16 @@ def test_sharded_tail_tp8_ep8_matches_replicated_projection() -> None: _run_ranks("matches_replicated_ep", 8) +@multi_gpu_test(num_gpus=4) +def test_presharded_tail_tp4_matches_replicated_projection() -> None: + _run_ranks("presharded_matches_replicated", 4) + + +@multi_gpu_test(num_gpus=4) +def test_presharded_unfused_tail_gathers_the_full_hidden_dim() -> None: + _run_ranks("unfused_gathers", 4) + + @multi_gpu_test(num_gpus=4) def test_sharded_tail_tp4_writes_only_its_own_shard() -> None: _run_ranks("own_shard_only", 4) diff --git a/vllm/models/kimi_k3/amd/latent_moe_runner.py b/vllm/models/kimi_k3/amd/latent_moe_runner.py index dbebfd418203..f6508d6b559f 100644 --- a/vllm/models/kimi_k3/amd/latent_moe_runner.py +++ b/vllm/models/kimi_k3/amd/latent_moe_runner.py @@ -36,18 +36,23 @@ def __init__( up_proj = getattr(transform, "up_proj", None) tp_size = get_tensor_model_parallel_world_size() + self._up_proj_preshard = bool(getattr(transform, "row_sharded", False)) self._up_proj_shard_size = 0 self._tail_shardable = ( up_proj is not None and tp_size > 1 - and up_proj.weight.shape[0] % tp_size == 0 + and (self._up_proj_preshard or up_proj.weight.shape[0] % tp_size == 0) and self._shared_experts is not None and not self.moe_config.is_sequence_parallel and self.routed_scaling_factor == 1.0 ) if self._tail_shardable: assert up_proj is not None - self._up_proj_shard_size = up_proj.weight.shape[0] // tp_size + self._up_proj_shard_size = ( + up_proj.weight.shape[0] + if self._up_proj_preshard + else up_proj.weight.shape[0] // tp_size + ) else: logger.warning_once( "K3 latent-MoE tail is not shardable under this config, " @@ -80,7 +85,12 @@ def _shard_up_proj_tail( shard_size = self._up_proj_shard_size shard_start = get_tensor_model_parallel_rank() * shard_size - up_proj_shard = transform.up_proj.weight.narrow(0, shard_start, shard_size) + weight = transform.up_proj.weight + up_proj_shard = ( + weight + if self._up_proj_preshard + else weight.narrow(0, shard_start, shard_size) + ) hidden_shard = shared_output.narrow(-1, shard_start, shard_size) # hidden_shard += latent @ up_proj_shard.T, accumulated in the GEMM's diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py index 3a1d827abe47..682773189352 100644 --- a/vllm/models/kimi_k3/amd/linear.py +++ b/vllm/models/kimi_k3/amd/linear.py @@ -137,11 +137,13 @@ class KimiRoutedOutputTransform(nn.Module): def __init__( self, norm: RMSNorm | None, - up_proj: ReplicatedLinear, + up_proj: ReplicatedLinear | ColumnParallelLinear, + row_sharded: bool = False, ) -> None: super().__init__() self.norm = norm self.up_proj = up_proj + self.row_sharded = row_sharded def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.norm is not None: @@ -265,7 +267,7 @@ def __init__( self.routed_expert_down_proj: ReplicatedLinear | None self.routed_expert_norm: RMSNorm | None - self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_expert_up_proj: ReplicatedLinear | ColumnParallelLinear | None self.routed_output_transform: KimiRoutedOutputTransform | None if self.use_latent_moe: self.routed_expert_down_proj = ReplicatedLinear( @@ -280,15 +282,36 @@ def __init__( if self.latent_moe_use_norm else None ) - self.routed_expert_up_proj = ReplicatedLinear( - self.moe_hidden_size, - hidden_size, - bias=False, - quant_config=None, - prefix=f"{prefix}.routed_expert_up_proj", + shard_up_proj = ( + self.tp_size > 1 + and hidden_size % self.tp_size == 0 + and self.num_shared_experts is not None + and self.routed_scaling_factor == 1.0 ) + if shard_up_proj: + # gather_output is a fallback restoring the full hidden dim for + # the unfused tail. The fused tail reads from .weight directly + # so it is not triggered + self.routed_expert_up_proj = ColumnParallelLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + gather_output=True, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + else: + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) self.routed_output_transform = KimiRoutedOutputTransform( - self.routed_expert_norm, self.routed_expert_up_proj + self.routed_expert_norm, + self.routed_expert_up_proj, + row_sharded=shard_up_proj, ) else: self.routed_expert_down_proj = None From 994b8b77be74e65c869c1ca8979fefb17b162b11 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 8 Oct 2026 17:06:39 +0200 Subject: [PATCH 212/952] [Bugfix][ROCm] RDNA3 W4A16 MoE: handle the N-first CT weight layout (#60388) Signed-off-by: JartX --- .../quantization/test_rdna3_moe_w4a16.py | 81 +++++++++++++++++++ .../layers/fused_moe/oracle/int_wna16.py | 17 ++-- .../compressed_tensors_moe_wna16.py | 6 +- 3 files changed, 95 insertions(+), 9 deletions(-) diff --git a/tests/kernels/quantization/test_rdna3_moe_w4a16.py b/tests/kernels/quantization/test_rdna3_moe_w4a16.py index 6212e6f08c68..58aa90aa07a9 100644 --- a/tests/kernels/quantization/test_rdna3_moe_w4a16.py +++ b/tests/kernels/quantization/test_rdna3_moe_w4a16.py @@ -18,6 +18,7 @@ import pytest import torch +import torch.nn.functional as F from vllm.platforms import current_platform @@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402 moe_align_block_size, ) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( # noqa: E402 + _process_weights_rdna3, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 pack_quantized_values_into_int32, ) @@ -95,6 +99,16 @@ def _make_qzeros(E, groups, N): return qz.unsqueeze(0).expand(E, -1, -1).contiguous() +def _make_ct_weights(E, N, K, group_size, dtype): + """Compressed-tensors canonical N-first layout: packed ``[E, N, K // 8]`` + and scales ``[E, N, K // group_size]``, plus the dequantized weights.""" + q = torch.randint(0, 16, (E, N, K), dtype=torch.int32, device=device) + s = torch.rand(E, N, K // group_size, dtype=dtype, device=device) * 0.1 + 0.01 + packed = pack_quantized_values_into_int32(q, scalar_types.uint4b8, packed_dim=2) + w = (q - 8).float() * s.float().repeat_interleave(group_size, dim=2) + return packed.contiguous(), s, w + + @gfx1100_only @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) @pytest.mark.parametrize("M", NUM_TOKENS) @@ -359,3 +373,70 @@ def test_expert_id_minus_one(): # Output should remain zero (expert skipped) assert torch.equal(out, torch.zeros_like(out)) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_process_weights_from_compressed_tensors_layout(dtype): + """Weights in the layout compressed-tensors hands to the oracle go through + the RDNA3 post-processing and a full MoE forward, against a dequantized + reference. The tests above build the kernel layout directly, so they do + not notice when the layout upstream of the kernel changes (#52798).""" + E, K, N_inter, top_k, group_size, M = 8, 512, 256, 2, 32, 16 + torch.manual_seed(5) + w13, w13_s, w13_ref = _make_ct_weights(E, 2 * N_inter, K, group_size, dtype) + w2, w2_s, w2_ref = _make_ct_weights(E, K, N_inter, group_size, dtype) + w13, w2, w13_s, w2_s, w13_z, w2_z, *_ = _process_weights_rdna3( + w13, w2, w13_s, w2_s, group_size + ) + assert w13_z is not None and w2_z is not None + + x = torch.randn(M, K, dtype=dtype, device=device) + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + topk_w = torch.softmax(torch.randn(M, top_k, device=device), dim=-1).float() + si, ei, ntp = moe_align_block_size(topk_ids, 1, E) + + w1_out = torch.zeros(M * top_k, 2 * N_inter, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + w1_out, + w13, + w13_s, + w13_z, + torch.empty(0, device=device), + si, + ei, + ntp, + top_k, + 1, + False, + 0, + ) + act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device) + apply_moe_activation(MoEActivation.SILU, act_out, w1_out) + out = torch.zeros(M, K, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + act_out, + out, + w2, + w2_s, + w2_z, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + top_k, + ) + + ref = torch.zeros(M, K, device=device) + for m in range(M): + for k in range(top_k): + e = int(topk_ids[m, k]) + h = w13_ref[e] @ x[m].float() + a = F.silu(h[:N_inter]) * h[N_inter:] + ref[m] += (w2_ref[e] @ a) * topk_w[m, k] + rel_l2 = (torch.norm(out.float() - ref) / torch.norm(ref)).item() + assert rel_l2 < (0.02 if dtype == torch.float16 else 0.05), rel_l2 diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 3e91fd9e041c..9b4b1633790f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -926,15 +926,18 @@ def _process_weights_rdna3( ]: """RDNA3 (gfx1100) W4A16 weight post-processing. - Interleaves the packed nibbles per expert (the exllama shuffle the dense - RDNA3 kernel also uses) and synthesizes the symmetric zero points that - ``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout - ``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the - kernel wants, so neither is repacked. + Transposes the canonical N-first inputs to the K-first layout of + ``moe_gptq_gemm_rdna3`` (weights ``[E, K // 8, N]``, scales + ``[E, groups, N]``), interleaves the packed nibbles per expert (the exllama + shuffle the dense RDNA3 kernel also uses) and synthesizes the symmetric + zero points the kernel dequantizes with. """ device = w13.device num_experts = w13.size(0) + w13 = w13.transpose(1, 2).contiguous() + w2 = w2.transpose(1, 2).contiguous() + for e in range(num_experts): w13_e = w13[e].contiguous() ops.gptq_shuffle(w13_e, 4) @@ -950,8 +953,8 @@ def _qzeros(w: torch.Tensor) -> torch.Tensor: return ( w13, w2, - w13_scale.contiguous(), - w2_scale.contiguous(), + w13_scale.transpose(1, 2).contiguous(), + w2_scale.transpose(1, 2).contiguous(), _qzeros(w13), _qzeros(w2), None, # w13_input_global_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index fdd37d120354..1e06f474133a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -450,10 +450,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if w2_bias is not None: replace_parameter(layer, "w2_bias", w2_bias) - # CPU fused_experts_cpu requires zero points even for symmetric quant. + # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points + # even for symmetric quant (the oracle synthesizes them). # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None. if ( - not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU + not self.symmetric + or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3) ) and self.wna16_backend != WNA16MoEBackend.EMULATION: assert w13_qzeros is not None and w2_qzeros is not None replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) From fd928e5f3cf17bce1ba090040beec48f86c84ff6 Mon Sep 17 00:00:00 2001 From: yifanFengg Date: Thu, 8 Oct 2026 08:19:57 -0700 Subject: [PATCH 213/952] [Bugfix] Check batch-invariance support for every attention backend (#60271) Signed-off-by: yifanFengg Co-authored-by: Claude Fable 5.1 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/v1/worker/gpu/attn_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 4fa308bcf471..8aa9ba3cb8b2 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -8,6 +8,7 @@ import numpy as np import torch +import vllm.envs as envs from vllm.config import ( VllmConfig, get_layers_from_vllm_config, @@ -224,6 +225,14 @@ def init_attn_backend( for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() + if ( + envs.VLLM_BATCH_INVARIANT + and not attn_backend.supports_batch_invariance() + ): + raise RuntimeError( + "VLLM batch_invariant mode is not supported for " + f"{attn_backend.get_name()}." + ) if layer_name in fast_prefill_eligible_layers: attn_backend = create_fast_prefill_custom_backend( "FastPrefill", attn_backend From 072d49bb6eb017a1e33bf3d016a979972438a90f Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Thu, 8 Oct 2026 23:25:21 +0800 Subject: [PATCH 214/952] [Platform] Support platform-specific MLA prefill backend selection (#56959) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../gen_files/generate_attention_backends.py | 96 +++--- .../v1/attention/test_mla_prefill_selector.py | 297 ++++++++++++------ vllm/platforms/cpu.py | 20 ++ vllm/platforms/interface.py | 84 +++++ vllm/platforms/rocm.py | 55 ++++ .../backends/mla/prefill/aiter_flash_attn.py | 5 +- .../v1/attention/backends/mla/prefill/base.py | 18 +- .../backends/mla/prefill/flashinfer.py | 7 +- .../backends/mla/prefill/selector.py | 142 ++------- .../backends/mla/prefill/tokenspeed_mla.py | 7 +- .../backends/mla/prefill/trtllm_ragged.py | 7 +- 11 files changed, 462 insertions(+), 276 deletions(-) diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py index 725d7072ca92..135dcf3a35a2 100644 --- a/docs/mkdocs/gen_files/generate_attention_backends.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -51,7 +51,7 @@ MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" -MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" +PLATFORM_INTERFACE_FILE = REPO_ROOT / "vllm" / "platforms" / "interface.py" # --------------------------------------------------------------------------- @@ -320,67 +320,83 @@ def parse_mla_prefill_registry() -> dict[str, str]: def parse_mla_prefill_priorities() -> dict[str, list[str]]: - """Parse MLA prefill backend priorities from selector.py. + """Parse MLA prefill backend priorities from platforms/interface.py. Returns: A dict with keys like 'blackwell' and 'default' containing lists of backend enum names in priority order. """ - if not MLA_PREFILL_SELECTOR_FILE.exists(): + if not PLATFORM_INTERFACE_FILE.exists(): return {} try: - tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text()) + tree = ast.parse(PLATFORM_INTERFACE_FILE.read_text()) except Exception: return {} + def extract_backends(stmt: ast.AST) -> list[str] | None: + """Read a backend enum list from a return/assignment statement.""" + if isinstance(stmt, ast.Return) or ( + isinstance(stmt, ast.Assign) + and len(stmt.targets) == 1 + and isinstance(stmt.targets[0], ast.Name) + ): + value = stmt.value + else: + return None + if not isinstance(value, ast.List): + return None + return [elt.attr for elt in value.elts if isinstance(elt, ast.Attribute)] + + def extract_last_list(nodes: list[ast.stmt], recurse: bool) -> list[str] | None: + """Return the last backend list found among the given statements.""" + result = None + for stmt in nodes: + backends = extract_backends(stmt) + if backends is not None: + result = backends + elif recurse and isinstance(stmt, ast.If): + nested = extract_last_list([*stmt.body, *stmt.orelse], True) + if nested is not None: + result = nested + return result + + def is_blackwell_check(stmt: ast.AST) -> bool: + return ( + isinstance(stmt, ast.If) + and isinstance(stmt.test, ast.Compare) + and isinstance(stmt.test.left, ast.Attribute) + and stmt.test.left.attr == "major" + and stmt.test.comparators + and isinstance(stmt.test.comparators[0], ast.Constant) + and stmt.test.comparators[0].value == 10 + ) + priorities: dict[str, list[str]] = {} for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef): + if not isinstance(node, ast.ClassDef) or node.name != "Platform": continue - if node.name != "_get_mla_prefill_backend_priorities": + + method = find_method(node, "get_mla_prefill_backend_cls") + if method is None: continue # Look for if statements checking device_capability.major - for stmt in ast.walk(node): + for stmt in ast.walk(method): if not isinstance(stmt, ast.If): continue - # Check if it's a capability.major == 10 check (Blackwell) - is_blackwell = ( - isinstance(stmt.test, ast.Compare) - and isinstance(stmt.test.left, ast.Attribute) - and stmt.test.left.attr == "major" - and stmt.test.comparators - and isinstance(stmt.test.comparators[0], ast.Constant) - and stmt.test.comparators[0].value == 10 - ) - - # Extract backends from return statements - for body_stmt in stmt.body: - if isinstance(body_stmt, ast.Return) and isinstance( - body_stmt.value, ast.List - ): - backends = [] - for elt in body_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) - if is_blackwell: - priorities["blackwell"] = backends - else: - priorities["default"] = backends - - # Extract from else branch - for else_stmt in stmt.orelse: - if isinstance(else_stmt, ast.Return) and isinstance( - else_stmt.value, ast.List - ): - backends = [] - for elt in else_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) + if is_blackwell_check(stmt): + # The Blackwell lists live in the nested if/else inside the + # major == 10 branch. + backends = extract_last_list(stmt.body, True) + if backends is not None: + priorities["blackwell"] = backends + else: + backends = extract_last_list([*stmt.body, *stmt.orelse], False) + if backends is not None: priorities["default"] = backends return priorities diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d5b81695aa41..3c6e1a929571 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend selector.""" +from contextlib import suppress from unittest.mock import MagicMock, patch import pytest @@ -9,13 +10,12 @@ from vllm.config import AttentionConfig, ModelConfig, VllmConfig from vllm.platforms import current_platform -from vllm.platforms.interface import DeviceCapability +from vllm.platforms.interface import DeviceCapability, Platform from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, _auto_select_mla_prefill_backend, - _get_mla_prefill_backend_priorities, get_mla_prefill_backend, ) @@ -55,51 +55,170 @@ def _make_vllm_config( return mock_vllm_config -class TestGetMLAPrefillBackend: - """Tests for get_mla_prefill_backend (public API).""" +class TestPlatformMLAPrefillHooks: + """Tests for the Platform.get_mla_prefill_backend_cls hook itself.""" - def test_cpu_uses_sdpa_prefill(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + def _get_tried_backends(self, capability, mla_dimensions): + """Run the default hook, returning the backend order tried.""" + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=mla_dimensions, + ) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls - # A Zen CPU with zentorch selects the Zen backend, so expect whichever - # one this host offers. - expected = ( - zen_cls if zen_cls.is_available() else MLAPrefillBackendEnum.CPU.get_class() + with ( + patch.object(Platform, "get_device_capability", return_value=capability), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + Platform.get_mla_prefill_backend_cls(selector_config) + return tried + + def test_priorities_no_capability(self): + tried = self._get_tried_backends( + None, + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] - with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_cpu.return_value = True + def test_priorities_blackwell_glm_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256), + ) + assert tried == [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] - backend = get_mla_prefill_backend(vllm_config) - assert backend is expected + def test_priorities_blackwell_deepseek_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] - def test_zen_cpu_prefers_zentorch_prefill(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + def test_priorities_hopper(self): + tried = self._get_tried_backends( + DeviceCapability(major=9, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] + + def test_skips_import_error_backends(self): + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ), + ) + valid = MagicMock() + valid.validate_configuration.return_value = [] + + def mock_get_class(backend_enum): + if backend_enum is MLAPrefillBackendEnum.FLASH_ATTN: + return valid + raise ImportError("not installed") with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=True), + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), ): - mock_platform.is_cpu.return_value = True + selected = Platform.get_mla_prefill_backend_cls(selector_config) + assert selected is valid - backend = get_mla_prefill_backend(vllm_config) - assert backend is zen_cls + def test_all_invalid_raises_error(self): + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - def test_zen_cpu_falls_back_when_zentorch_unavailable(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + invalid = MagicMock() + invalid.validate_configuration.return_value = ["not available"] with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=False), + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", return_value=invalid), + pytest.raises(ValueError, match="No valid MLA"), ): - mock_platform.is_cpu.return_value = True + Platform.get_mla_prefill_backend_cls(selector_config) + + def test_cpu_platform_returns_fixed_backend(self): + """CPU backends are returned via the fixed-backend hook.""" + try: + cpu_cls = MLAPrefillBackendEnum.CPU.get_class() + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + except ImportError: + pytest.skip("CPU MLA prefill backend not available") + return + + from vllm.platforms.cpu import CpuPlatform + + expected = zen_cls if zen_cls.is_available() else cpu_cls + assert ( + CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + is expected + ) + + +class TestGetMLAPrefillBackend: + """Tests for get_mla_prefill_backend (public API).""" + + def test_cpu_uses_sdpa_prefill(self): + vllm_config = _make_vllm_config() + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = ( + MLAPrefillBackendEnum.CPU.get_class() + ) backend = get_mla_prefill_backend(vllm_config) assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_zen_cpu_prefers_zentorch_prefill(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=True): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is zen_cls + + def test_zen_cpu_falls_back_when_zentorch_unavailable(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=False): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_no_device_capability_returns_flash_attn(self): vllm_config = _make_vllm_config() @@ -108,16 +227,8 @@ class FlashAttnBackend: def get_name(): return "FLASH_ATTN" - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - MLAPrefillBackendEnum.FLASH_ATTN, - "get_class", - return_value=FlashAttnBackend, - ), - ): - mock_platform.get_device_capability.return_value = None - mock_platform.is_cpu.return_value = False + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = FlashAttnBackend backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -214,7 +325,7 @@ def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int): mock_platform.get_device_capability.return_value = DeviceCapability( major=9, minor=0 ) - mock_platform.is_rocm.return_value = False + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -224,7 +335,6 @@ class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" def test_blackwell_glm_dimensions_use_trtllm(self): - capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -235,7 +345,6 @@ def test_blackwell_glm_dimensions_use_trtllm(self): ) try: - flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class() except ImportError: pytest.skip("MLA prefill backend not available") @@ -243,19 +352,13 @@ def test_blackwell_glm_dimensions_use_trtllm(self): with ( patch("vllm.platforms.current_platform") as mock_platform, - patch.object(flash_attn_cls, "is_available", return_value=True), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): - # Force the non-ROCm priority on the Blackwell. - mock_platform.is_rocm.return_value = False - backend = _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + mock_platform.get_mla_prefill_backend_cls.return_value = trtllm_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "TRTLLM_RAGGED" - def test_all_fail_raises_error(self): - capability = DeviceCapability(major=10, minor=0) + def test_platform_hook_raises_propagates(self): selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -265,18 +368,14 @@ def test_all_fail_raises_error(self): ), ) - def mock_get_class(backend_enum): # noqa: ARG001 - cls = MagicMock() - cls.validate_configuration.return_value = ["not available"] - return cls - - with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class): - _auto_select_mla_prefill_backend.cache_clear() - with pytest.raises(ValueError, match="No valid MLA"): - _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + with ( + patch("vllm.platforms.current_platform") as mock_platform, + pytest.raises(ValueError, match="No valid MLA"), + ): + mock_platform.get_mla_prefill_backend_cls.side_effect = ValueError( + "No valid MLA prefill backend found" + ) + _auto_select_mla_prefill_backend(selector_config) class TestBackendValidation: @@ -391,18 +490,29 @@ class TestROCmAiterFAPrefillSelection: def test_rocm_priorities_prefer_aiter_fa(self): """On ROCm, ROCM_AITER_FA is tried first, FLASH_ATTN as fallback.""" - with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_rocm.return_value = True - priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5), - MLADimensions( - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - ), - ) + from vllm.platforms.rocm import RocmPlatform - assert priorities == [ + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls + + with ( + patch.object( + RocmPlatform, + "get_device_capability", + return_value=MagicMock(), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + RocmPlatform.get_mla_prefill_backend_cls(selector_config) + + assert tried == [ MLAPrefillBackendEnum.ROCM_AITER_FA, MLAPrefillBackendEnum.FLASH_ATTN, ] @@ -459,50 +569,29 @@ def test_auto_select_prefers_aiter_fa_on_rocm(self): AiterFlashAttnPrefillBackend, ) - # gfx gating is simulated via the mocked validate_configuration, - # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=[], - ), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = ( + AiterFlashAttnPrefillBackend + ) + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "ROCM_AITER_FA" def test_auto_select_falls_back_to_flash_attn_when_aiter_invalid(self): - from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( - AiterFlashAttnPrefillBackend, - ) - try: flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() except ImportError: pytest.skip("FLASH_ATTN backend not available") return - # the fallback is forced by the mocked validate_configuration, + # the fallback is forced by the mocked get_mla_prefill_backend_cls, # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=["compute capability not supported"], - ), - patch.object(flash_attn_cls, "validate_configuration", return_value=[]), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "FLASH_ATTN" diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 25d28daf937f..e0257acb8007 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -18,6 +18,7 @@ get_visible_memory_node, ) from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import CpuArchEnum, Platform, PlatformEnum @@ -26,6 +27,10 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backends.mla.prefill.base import ( + MLAPrefillBackend, + ) + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: VllmConfig = None @@ -173,6 +178,21 @@ def get_attn_backend_cls( logger.info("Cannot use %s backend on CPU.", selected_backend) return AttentionBackendEnum.CPU_ATTN.get_path() + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device.""" + for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): + try: + cpu_backend_cls = backend_enum.get_class() + except ImportError: + continue + if cpu_backend_cls.is_available(): + return cpu_backend_cls + raise ValueError("No valid CPU MLA prefill backend found.") + @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: meminfo = get_memory_node_info(device_id) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index f4583daba26f..dc3200759c82 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -26,6 +26,8 @@ from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: FlexibleArgumentParser = object @@ -429,6 +431,88 @@ def get_vit_attn_backend( ) return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device. + + Platforms that have a single, fixed MLA prefill backend can override + this to return it directly. The default implementation selects the + best backend from a device-capability-based priority list. + + Returns: + The selected backend class. + + Raises: + ValueError: If no MLA prefill backend is valid for the + configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.base import MLADimensions + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + if device_capability is None: + return MLAPrefillBackendEnum.FLASH_ATTN.get_class() + + if device_capability.major == 10: # Blackwell + if mla_selector_config.mla_dimensions == MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ): + priorities = [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: + priorities = [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: # Hopper (SM90) and older + priorities = [MLAPrefillBackendEnum.FLASH_ATTN] + + all_invalid_reasons: dict[str, list[str]] = {} + + for backend_enum in priorities: + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def get_device_capability( cls, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index f394e56bf640..5840ffe3b792 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -24,6 +24,10 @@ from vllm.config.kernel import IrOpPriorityConfig from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, + ) from vllm.v1.attention.selector import AttentionSelectorConfig logger = init_logger(__name__) @@ -904,6 +908,57 @@ def get_vit_attn_backend( logger.info_once("Using Torch SDPA backend for ViT model.") return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """On ROCm, prefer the AITER FlashAttention backend, falling back to + FlashAttention. + + Raises: + ValueError: If neither backend is valid for the configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + all_invalid_reasons: dict[str, list[str]] = {} + for backend_enum in ( + MLAPrefillBackendEnum.ROCM_AITER_FA, + MLAPrefillBackendEnum.FLASH_ATTN, + ): + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def set_device(cls, device: torch.device) -> None: """Set the device for the current platform.""" diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 1d03b9e77392..8b0e337bb221 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -31,7 +31,10 @@ def get_name() -> str: return "ROCM_AITER_FA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: if not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_mi3xx diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 1b589cb295f4..4b4361214f84 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -49,7 +49,10 @@ def get_name() -> str: raise NotImplementedError @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: return True @classmethod @@ -90,16 +93,19 @@ def supports_out(self) -> bool: @classmethod def validate_configuration( cls, - device_capability: "DeviceCapability", + device_capability: "DeviceCapability | None", selector_config: "MLAPrefillSelectorConfig", ) -> list[str]: invalid_reasons: list[str] = [] if not cls.supports_compute_capability(device_capability): - invalid_reasons.append( - f"compute capability {device_capability.major}." - f"{device_capability.minor} not supported" - ) + if device_capability is None: + invalid_reasons.append("device capability not available") + else: + invalid_reasons.append( + f"compute capability {device_capability.major}." + f"{device_capability.minor} not supported" + ) if not cls.supports_dtype(selector_config.dtype): invalid_reasons.append(f"dtype {selector_config.dtype} not supported") diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 777dd33756c7..5cf47ea2caeb 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -51,8 +51,11 @@ def get_name() -> str: return "FLASHINFER" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index 0cb29f51b49a..f0f001580964 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -12,9 +12,7 @@ import torch from vllm.logger import init_logger -from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions -from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum if TYPE_CHECKING: from vllm.config import VllmConfig @@ -45,54 +43,8 @@ def __repr__(self): ) -def _get_mla_prefill_backend_priorities( - device_capability: DeviceCapability, - mla_dimensions: MLADimensions, -) -> list[MLAPrefillBackendEnum]: - """Get MLA prefill backend priorities based on device capability. - - Args: - device_capability: The device's compute capability. - mla_dimensions: The model's MLA head dimensions. - - Returns: - List of backends in priority order (highest priority first). - - """ - from vllm.platforms import current_platform - - if current_platform.is_rocm(): - return [ - MLAPrefillBackendEnum.ROCM_AITER_FA, - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - if device_capability.major == 10: # Blackwell - if mla_dimensions == MLADimensions( - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - ): - return [ - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - else: # Hopper (SM90) and older - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - def get_mla_prefill_backend( - vllm_config: "VllmConfig", + vllm_config: "VllmConfig | None", ) -> "type[MLAPrefillBackend]": """Select the MLA prefill backend based on configuration and device. @@ -101,7 +53,8 @@ def get_mla_prefill_backend( priority-based selection. Args: - vllm_config: The vLLM configuration. + vllm_config: The vLLM configuration. May be None before model + resolution, in which case defaults are used. Returns: The selected prefill backend class. @@ -109,30 +62,7 @@ def get_mla_prefill_backend( """ from vllm.platforms import current_platform - if current_platform.is_cpu(): - # CPUs have no compute capability, so the capability-driven priority - # path below does not apply. Prefer an accelerator-specific CPU backend - # when its kernels are present, else the generic SDPA one. - for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): - try: - cpu_backend_cls = backend_enum.get_class() - except ImportError: - continue - if cpu_backend_cls.is_available(): - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return cpu_backend_cls - raise ValueError("No valid CPU MLA prefill backend found.") - - device_capability = current_platform.get_device_capability() - if device_capability is None: - logger.info_once( - "Device capability not available, using FlashAttention MLA prefill backend." - ) - return MLAPrefillBackendEnum.FLASH_ATTN.get_class() - - attention_config = vllm_config.attention_config - - model_config = vllm_config.model_config + model_config = vllm_config.model_config if vllm_config is not None else None if model_config is None: selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype()) else: @@ -146,11 +76,21 @@ def get_mla_prefill_backend( ), ) - if attention_config.mla_prefill_backend is not None: + if current_platform.is_cpu(): + # CPUs have a single fixed MLA prefill backend, so explicit backend + # configuration does not apply; always auto-select via the platform. + return _auto_select_mla_prefill_backend(selector_config) + + attention_config = vllm_config.attention_config if vllm_config is not None else None + if ( + attention_config is not None + and attention_config.mla_prefill_backend is not None + ): selected_backend = attention_config.mla_prefill_backend backend_cls: type[MLAPrefillBackend] | None = None try: backend_cls = selected_backend.get_class() + device_capability = current_platform.get_device_capability() invalid_reasons = backend_cls.validate_configuration( device_capability, selector_config ) @@ -166,63 +106,27 @@ def get_mla_prefill_backend( logger.info_once("Using %s MLA prefill backend.", selected_backend.name) return backend_cls - return _auto_select_mla_prefill_backend( - device_capability, - selector_config, - ) + return _auto_select_mla_prefill_backend(selector_config) @cache def _auto_select_mla_prefill_backend( - device_capability: DeviceCapability, selector_config: MLAPrefillSelectorConfig, ) -> "type[MLAPrefillBackend]": """Auto-select the best available MLA prefill backend. Args: - device_capability: The device's compute capability. selector_config: Hashable configuration for backend selection. Returns: The selected prefill backend class. - """ - priorities = _get_mla_prefill_backend_priorities( - device_capability, - selector_config.mla_dimensions, - ) - all_invalid_reasons: dict[str, list[str]] = {} + Raises: + ValueError: If the platform has no valid MLA prefill backend. - for backend_enum in priorities: - backend_cls: type[MLAPrefillBackend] | None = None - try: - backend_cls = backend_enum.get_class() - invalid_reasons = backend_cls.validate_configuration( - device_capability, selector_config - ) - except ImportError: - invalid_reasons = ["ImportError"] - if not invalid_reasons: - assert backend_cls is not None - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return backend_cls - all_invalid_reasons[backend_enum.name] = invalid_reasons - - reasons_str = ( - "{" - + ", ".join( - f"{name}: [{', '.join(reasons)}]" - for name, reasons in all_invalid_reasons.items() - ) - + "}" - ) - config_str = repr(selector_config) - logger.debug_once( - "Some MLA prefill backends are not valid with %s. Reasons: %s.", - config_str, - reasons_str, - ) + """ + from vllm.platforms import current_platform - raise ValueError( - f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}." - ) + selected_backend = current_platform.get_mla_prefill_backend_cls(selector_config) + logger.info_once("Using %s MLA prefill backend.", selected_backend.get_name()) + return selected_backend diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 6985b0a70c24..d65721b81b5f 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -35,8 +35,11 @@ def get_name() -> str: return "TOKENSPEED_MLA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 _INSTALL_HINT = ( "tokenspeed_mla package is not installed. " diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 167738b07af4..3cb498be2b16 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -43,8 +43,11 @@ def get_name() -> str: return "TRTLLM_RAGGED" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: From 194bd72be7b9eb1c406e090328ddaf09a4e2d6eb Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 8 Oct 2026 08:28:23 -0700 Subject: [PATCH 215/952] [Compile] Remove torch.compile-based sequence parallelism and AsyncTP (#60517) --- .buildkite/test-amd.yaml | 108 -- .buildkite/test_areas/compile.yaml | 128 --- docs/design/cuda_graphs.md | 4 +- docs/design/fusions.md | 76 +- docs/design/vllm_ir.md | 12 +- docs/features/batch_invariance.md | 5 +- .../compile/correctness_e2e/test_async_tp.py | 270 ----- .../correctness_e2e/test_sequence_parallel.py | 414 ------- tests/compile/fusions_e2e/common.py | 6 - tests/compile/fusions_e2e/conftest.py | 73 +- tests/compile/fusions_e2e/models.py | 19 - .../compile/fusions_e2e/test_tp2_async_tp.py | 339 ------ .../fusions_e2e/test_tp2_rocm_async_tp.py | 58 - .../compile/fusions_e2e/test_tp2_rocm_moe.py | 60 - .../passes/distributed/test_async_tp.py | 394 ------- .../distributed/test_sequence_parallelism.py | 343 ------ tests/compile/rocm/__init__.py | 2 - tests/compile/rocm/test_async_tp.py | 75 -- tests/compile/test_config.py | 184 +-- .../test_sequence_parallelism_threshold.py | 192 ---- .../test_routed_experts_capture.py | 5 +- tests/test_config.py | 12 +- tests/v1/worker/test_gpu_worker.py | 3 - .../passes/fusion/collective_fusion.py | 1013 ----------------- .../passes/fusion/sequence_parallelism.py | 620 ---------- vllm/compilation/passes/pass_manager.py | 7 - vllm/config/compilation.py | 60 +- vllm/config/vllm.py | 159 --- vllm/platforms/xpu.py | 1 - vllm/utils/flashinfer.py | 75 -- vllm/v1/worker/gpu/model_runner.py | 1 - vllm/v1/worker/gpu_model_runner.py | 42 +- vllm/v1/worker/gpu_worker.py | 38 - vllm/v1/worker/utils.py | 31 +- 34 files changed, 32 insertions(+), 4797 deletions(-) delete mode 100644 tests/compile/correctness_e2e/test_async_tp.py delete mode 100644 tests/compile/correctness_e2e/test_sequence_parallel.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_async_tp.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py delete mode 100644 tests/compile/fusions_e2e/test_tp2_rocm_moe.py delete mode 100644 tests/compile/passes/distributed/test_async_tp.py delete mode 100644 tests/compile/passes/distributed/test_sequence_parallelism.py delete mode 100644 tests/compile/rocm/__init__.py delete mode 100644 tests/compile/rocm/test_async_tp.py delete mode 100644 tests/compile/test_sequence_parallelism_threshold.py delete mode 100644 vllm/compilation/passes/fusion/collective_fusion.py delete mode 100644 vllm/compilation/passes/fusion/sequence_parallelism.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 9e634ea02b4d..4ae92315b4c9 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -693,28 +693,8 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition" -# - label: ":amd: (MI300) Sequence Parallel Correctness" # TBD -# timeout_in_minutes: 180 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/" -# source_file_dependencies: -# - vllm/model_executor/layers/ -# - vllm/compilation/ -# - vllm/v1/worker/ -# - vllm/v1/cudagraph_dispatcher.py -# - tests/compile/correctness_e2e/test_sequence_parallel.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_TEST_CLEAN_GPU_MEMORY=1 -# - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - - label: ":amd: (MI300) Distributed Compile + RPC" # TBD timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3180,7 +3160,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ":amd: (MI355) Fusion and Compile FP8" timeout_in_minutes: 45 @@ -3205,64 +3184,6 @@ steps: - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py -- label: ":amd: (MI355) Sequence Parallel Correctness" - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - -- label: ":amd: (MI355) AsyncTP Correctness" - timeout_in_minutes: 60 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":amd: (MI355 DPX) Fusion E2E FP8 Config Sweep" timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -5916,35 +5837,6 @@ steps: - amdmi355 optional: true -- label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - timeout_in_minutes: 30 - working_dir: /vllm-workspace/ - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - dind: false - agent_pool: mi355_2 - num_gpus: 2 - mirror_hardwares: - - amdexperimental - - amdproduction - - amdgfx950nightly - - amdmi355 - optional: true - label: ":amd: (MI355 DPX) MRV2 Sampler JIT Warmup" timeout_in_minutes: 10 diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index acc1bafa7f2a..4a5a4f1a3cea 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,84 +2,6 @@ group: Compile depends_on: - image-build steps: -- label: ":nvidia: (B200) Sequence Parallel Correctness" - key: sequence-parallel-correctness-tests-2xb200 - timeout_in_minutes: 45 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - mirror: - amd: - label: ":amd: (MI355) Sequence Parallel Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 180 - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama.py - - vllm/v1/attention/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - vllm/platforms/rocm.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - - tests/models/registry.py - - tests/utils.py - -- label: ":nvidia: (B200) AsyncTP Correctness" - key: asynctp-correctness-tests-b200 - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/" - device: b200-k8s - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py - mirror: - amd: - label: ":amd: (MI355) AsyncTP Correctness" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - depends_on: - - image-build-amd - source_file_dependencies: - - .buildkite/test_areas/compile.yaml - - vllm/compilation/ - - vllm/config/ - - vllm/distributed/ - - vllm/model_executor/layers/ - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/llama.py - - vllm/model_executor/models/qwen3.py - - vllm/v1/attention/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/transformers_utils/repo_utils.py - - tests/compile/correctness_e2e/test_async_tp.py - - tests/compile/fusions_e2e/common.py - - tests/models/registry.py - - tests/utils.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_correctness -k RedHatAI - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_async_tp_pass_nvfp4_correctness - - label: ":nvidia: (H100) Distributed Compile" key: distributed-compile-unit-tests-2xh100 timeout_in_minutes: 45 @@ -300,7 +222,6 @@ steps: - nvidia-smi # Run all models and attn backends but only Inductor partition and native custom ops - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 Quick" @@ -388,53 +309,6 @@ steps: - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and Llama-3 and ROCM_AITER_UNIFIED_ATTN and (dynamo_partition or +quant_fp8 or +rms_norm)' -- label: ":nvidia: (H100) Fusion E2E TP2 AsyncTP Config Sweep" - key: fusion-e2e-tp2-asynctp-config-sweep-h100 - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/" - device: h100 - num_devices: 2 - source_file_dependencies: - - csrc/quantization/ - - vllm/compilation/ - # can affect pattern matching - - vllm/model_executor/layers/layernorm.py - - vllm/model_executor/layers/activation.py - - vllm/model_executor/layers/attention/attention.py - - vllm/model_executor/layers/quantization/input_quant_fp8.py - - tests/compile/fusions_e2e/ - commands: - - nvidia-smi - # Run just llama3 (fp8 & bf16) for all config combinations - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3" - mirror: - amd: - label: ':amd: (MI355) Fusion E2E TP2 AsyncTP BF16 Config Sweep' - dind: false - device: mi355_2 - num_devices: 2 - working_dir: /vllm-workspace/ - timeout_in_minutes: 30 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/compilation/ - - vllm/model_executor/layers/linear.py - - vllm/model_executor/layers/utils.py - - vllm/distributed/ - - vllm/v1/attention/backends/triton_attn.py - - tests/compile/fusions_e2e/ - - tests/compile/passes/distributed/test_async_tp.py - - tests/compile/rocm/test_async_tp.py - - tests/compile/correctness_e2e/test_async_tp.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/models/qwen3.py - commands: - - rocm-smi - - pytest -v -s tests/compile/rocm/test_async_tp.py - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py::test_rocm_async_tp_bf16_output_correctness - - label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 timeout_in_minutes: 45 @@ -453,7 +327,6 @@ steps: # include qwen/deepseek with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported # for ar-rms-quant-fp4, also sweep llama3 - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4" - - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites" @@ -480,7 +353,6 @@ steps: # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported. - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)' key: fusion-e2e-tp2-ar-rms-inductor-partition-amd diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index abebdd54ae46..9d53a98f6585 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -225,9 +225,9 @@ outputs = model.generate( ) ``` -### Piecewise compilation and full graph custom passes (attention fusion, sequence parallelism) +### Piecewise compilation and full graph custom passes (attention fusion) -Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. +Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs. Long term, we've added the ability to partition the graph in Inductor instead of right after Dynamo. It can be enabled with `CompilationConfig.use_inductor_graph_partition=True` but is currently experimental and only available with `torch>=2.9`. This also increases compilation time as it has to compile the whole graph and cannot reuse piecewise compilation artifacts. Once vLLM supports 2.9, we plan to make this the default approach as it will also speed up piecewise cudagraph capture. diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 002984c113a1..4095c41a7e63 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -25,8 +25,6 @@ or just on the low or high end. | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low | -| [Sequence Parallelism](#sequence-parallelism-enable_sp) | `enable_sp` | AllReduce → ReduceScatter + AllGather | Off by default | Prereq for AsyncTP | Yes | High | -| [AsyncTP GEMM + collective](#asynctp-gemm--collective-overlap-fuse_gemm_comms) | `fuse_gemm_comms` | GEMM → reduce-scatter / all-gather → GEMM | Off by default | 7-10% | Yes | High | | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always | | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always | @@ -44,9 +42,7 @@ The table below lists the quantization schemes supported by each fusion on each | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | | `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | -| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — | -| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | -| `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | +| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — | | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | @@ -56,10 +52,6 @@ The table below lists the quantization schemes supported by each fusion on each fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant) for per-backend details. -† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today; -other architectures support requires setting `PassConfig.sp_min_token_num` explicitly. -SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`. - ## Enabling / Disabling Fusions Fusions are exposed through `PassConfig`, which is nested inside `CompilationConfig`: @@ -186,72 +178,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py) -### Sequence Parallelism (`enable_sp`) - -**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather, -splitting the sequence dimension across TP ranks. This restructures the graph so the subsequent AsyncTP -pass can fuse the reduce-scatter / all-gather with the surrounding GEMMs. - -Sequence Parallelism itself does not directly improve performance; it is a prerequisite for the -AsyncTP pass (`fuse_gemm_comms`). SP is only applied above a minimum token threshold that is -autoconfigured based on device capability and model `hidden_size`. Currently only active on -H100/SM90 for models with `hidden_size >= 8192`. The threshold is configurable via -`PassConfig.sp_min_token_num`. - -The general transformation: - -```text -Input → AllReduce → RMSNorm → Output -becomes: -Input → ReduceScatter → local RMSNorm → AllGather → Output -``` - -Patterns covered: - -- First block: `AllReduce → RMSNorm` → `ReduceScatter → RMSNorm → AllGather` -- Middle blocks: `AllReduce → fused_add_RMSNorm` → `ReduceScatter → fused_add_RMSNorm → AllGather` -- Both with optional `→ FP8 static quant` suffix - -Requires: `use_inductor_graph_partition=True` **or** piecewise compilation with static sizes -divisible by `tensor_parallel_size`. - -Supported hardware: Only tested on NVIDIA CUDA, possibly works on ROCm. FP8 all-gather requires sm90+. - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - -### AsyncTP GEMM + Collective Overlap (`fuse_gemm_comms`) - -!!! info - Requires `enable_sp=True` (enabled automatically). This pass is a no-op if Sequence Parallelism has not been applied. - -**What it fuses.** After Sequence Parallelism transforms the graph, fuses GEMM kernels with the -surrounding reduce-scatter (output projection) and all-gather (input projection) using -`torch.ops.symm_mem` symmetric-memory primitives, overlapping communication and computation. -This overlap is only profitable for large `num_tokens`, so the fusion (and preceding SP) -is only performed in the higher compiled range above `PassConfig.sp_min_token_num`. - -Patterns covered: - -- `GEMM → reduce-scatter` → `fused_matmul_reduce_scatter` -- `all-gather → GEMM` → `all_gather_matmul` -- FP8 scaled variants of both patterns - -Supported hardware: NVIDIA CUDA with symmetric-memory (`torch.distributed._symmetric_memory`) support. - -On B200, pattern-matching fp8 FlashInfer scaled MM is not supported, so it must be disabled -([#27893](https://github.com/vllm-project/vllm/issues/27893)) - -```shell -VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ... -``` - -**Code locations.** - -- Pass: [`vllm/compilation/passes/fusion/collective_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/collective_fusion.py) -- Sequence parallelism pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py) - ### QK Norm + RoPE (`enable_qk_norm_rope_fusion`) !!! info diff --git a/docs/design/vllm_ir.md b/docs/design/vllm_ir.md index 82628f3762fe..e324ad25d03a 100644 --- a/docs/design/vllm_ir.md +++ b/docs/design/vllm_ir.md @@ -233,19 +233,17 @@ storing this information in vLLM's `PassContext` for later use in clone eliminat ### 3. IR Fusion and Transformation Passes After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations. -These passes can perform fusion, distribute operations for sequence parallelism, and other transformations: +These passes can perform fusion and other transformations: ```python -# Example: Sequence Parallelism (see SequenceParallelismPass) -# Before SP pass +# Example: AllReduce + RMSNorm fusion (see AllReduceFusionPass) +# Before all_reduce = torch.ops.vllm.all_reduce(x, "tp:0") rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -# after SP pass -reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0") -rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) -all_gather = torch.ops.vllm.all_gather(x, "tp:0") +# After +rms_norm = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm(x, weight, ...) ``` Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations, diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index e82336d24117..b40ff4da8cc6 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -148,11 +148,10 @@ When batch invariance is enabled, vLLM: 1. Uses deterministic kernel implementations for attention and other operations 2. Ensures consistent numerical behavior across different batch sizes -3. Disables certain optimizations that may introduce non-determinism (such as sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant) +3. Disables certain optimizations that may introduce non-determinism 4. Under tensor parallelism, keeps custom all-reduce on with a fixed reduction order (the 1-stage kernel is pinned, and large inputs are reduced in fixed-size chunks), and disables FlashInfer, AITER and QuickReduce all-reduce 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell), - runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out - (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes). + runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out. 6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics !!! warning diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py deleted file mode 100644 index 307f39039b39..000000000000 --- a/tests/compile/correctness_e2e/test_async_tp.py +++ /dev/null @@ -1,270 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json -import logging -from contextlib import nullcontext - -import pytest - -from tests.compile.fusions_e2e.common import FUSION_LOG_PATTERNS -from tests.models.registry import HF_EXAMPLE_MODELS -from tests.utils import ( - compare_two_settings, - create_new_process_for_each_test, - multi_gpu_test, -) -from vllm.config import ( - CompilationMode, -) -from vllm.platforms import current_platform -from vllm.transformers_utils.repo_utils import hf_api -from vllm.utils.flashinfer import has_flashinfer - -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_HF_OVERRIDES = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, -} - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -def test_rocm_async_tp_bf16_output_correctness(caplog_mp_spawn): - """Trained-model outputs must agree with unfused TP, including padded decode.""" - # A snapshot path pins both servers and the comparison helper's tokenizer. - model = hf_api().snapshot_download( - "Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca" - ) - common_args = [ - "--dtype", - "bfloat16", - "--load-format", - "auto", - "--tensor-parallel-size", - "2", - "--distributed-executor-backend", - "mp", - "--max-model-len", - "512", - "--max-num-batched-tokens", - "512", - "--max-num-seqs", - "8", - "--kv-cache-memory-bytes", - str(64 * 1024**2), - "--gpu-memory-utilization", - "0.01", - "--attention-backend", - "TRITON_ATTN", - "--disable-custom-all-reduce", - ] - - def settings(enabled): - return [ - *common_args, - "--compilation-config", - json.dumps( - { - "mode": CompilationMode.VLLM_COMPILE, - "cudagraph_mode": "NONE", - "splitting_ops": [], - "inductor_compile_config": {"force_disable_caches": True}, - "pass_config": { - "enable_sp": enabled, - "fuse_gemm_comms": enabled, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - ), - ] - - env = {"VLLM_ROCM_USE_AITER": "0", "VLLM_DISABLE_COMPILE_CACHE": "1"} - with caplog_mp_spawn(logging.DEBUG) as logs: - compare_two_settings( - model, - settings(False), - settings(True), - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - matches = [ - int(count) for count in FUSION_LOG_PATTERNS["async_tp"].findall(logs.text) - ] - print(f"AsyncTP rewrite counts: {matches}") - assert len(matches) >= 2 and all(count > 0 for count in matches), ( - f"Expected positive AsyncTP rewrites on both TP workers, found {matches}" - ) - - -@create_new_process_for_each_test() -@pytest.mark.parametrize( - "model_id", - ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"], -) -@pytest.mark.parametrize("tp_size", [2]) -@pytest.mark.parametrize("async_tp_enabled", [True]) -@pytest.mark.parametrize("distributed_backend", ["mp"]) -def test_async_tp_pass_correctness( - model_id: str, - tp_size: int, - async_tp_enabled: bool, - distributed_backend: str, - num_gpus_available: int, - monkeypatch, -): - # Disable FlashInfer FP8 scaled_mm kernel as it is incompatible with - # async TP patterns. No-op on H100 (kernel requires CC >= 100). - monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") - - model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - model_info.check_available_online(on_fail="skip") - - pp_size = 1 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": {"fuse_gemm_comms": async_tp_enabled}, - } - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - compare_two_settings( - model_id, - async_tp_args, - tp_args, - method="generate", - force_v1_runner=True, - ) - - -@create_new_process_for_each_test() -def test_async_tp_pass_nvfp4_correctness(num_gpus_available: int, caplog_mp_spawn): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 AsyncTP requires Blackwell or ROCm CDNA") - if is_blackwell and not has_flashinfer(): - pytest.skip("FlashInfer is required for the NVFP4 AsyncTP path") - - tp_size = 2 - if num_gpus_available < tp_size: - pytest.skip(f"Need at least {tp_size} GPUs") - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - "--load-format", - "dummy", - *(["--linear-backend", "flashinfer_cutlass"] if is_blackwell else []), - "--hf-overrides", - json.dumps(NVFP4_HF_OVERRIDES), - ] - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [2, 4, 8], - "splitting_ops": [], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": True, - "fuse_allreduce_rms": False, - "sp_min_token_num": 1, - }, - } - if is_rocm_cdna: - compilation_config["inductor_compile_config"] = {"force_disable_caches": True} - - async_tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - env = {"VLLM_DISABLE_COMPILE_CACHE": "1"} if is_rocm_cdna else None - log_context = caplog_mp_spawn(logging.DEBUG) if is_rocm_cdna else nullcontext() - with log_context as logs: - compare_two_settings( - NVFP4_MODEL_ID, - async_tp_args, - tp_args, - env1=env, - env2=env, - method="generate", - force_v1_runner=True, - ) - if is_rocm_cdna: - assert logs is not None - matches_by_rank = { - rank: [ - int(count) - for line in logs.text.splitlines() - if line.startswith(f"(Worker_TP{rank} pid=") - for count in FUSION_LOG_PATTERNS["async_tp"].findall(line) - ] - for rank in range(tp_size) - } - assert all( - any(count > 0 for count in matches) for matches in matches_by_rank.values() - ), ( - "Expected positive NVFP4 AsyncTP rewrites on both TP workers, " - f"found {matches_by_rank}" - ) diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py deleted file mode 100644 index 10d25320c012..000000000000 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ /dev/null @@ -1,414 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""WARNING: This test runs in both single-node (4 GPUs) and multi-node -(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is -important to set the distributed backend to "mp" to avoid Ray scheduling -all workers in a node other than the head node, which can cause the test -to fail. -""" - -import json -import os -from dataclasses import dataclass -from typing import Literal, NamedTuple - -import pytest - -from vllm.config.compilation import CompilationMode -from vllm.config.model import RunnerOption -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_torch_equal_or_newer - -from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo -from ...utils import compare_all_settings, create_new_process_for_each_test - -logger = init_logger("test_sequence_parallel") - -VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1" -NVFP4_MODEL_ID = "nvidia/Llama-3.1-8B-Instruct-NVFP4" -NVFP4_MODEL_INFO = _HfExamplesInfo(NVFP4_MODEL_ID) - - -class ParallelSetup(NamedTuple): - tp_size: int - pp_size: int - fuse_norm_quant: bool - fuse_act_quant: bool - eager_mode: bool - - -class SPTestOptions(NamedTuple): - multi_node_only: bool - load_format: str | None = None - model_info: _HfExamplesInfo | None = None - - -@dataclass -class SPTestSettings: - parallel_setups: list[ParallelSetup] - distributed_backends: list[str] - runner: RunnerOption - test_options: SPTestOptions - - @staticmethod - def fast( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [ - ParallelSetup( - tp_size=tp_base, - pp_size=pp_multiplier * pp_base, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_multiplier in [1, 2] - ] - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - @staticmethod - def fp8_quant( - *, - tp_base: int = 2, - pp_base: int = 1, - runner: RunnerOption = "auto", - multi_node_only: bool = False, - load_format: str | None = None, - ): - parallel_setups = [] - for fusion_val in [False, True]: - parallel_setups.append( - ParallelSetup( - tp_size=tp_base, - pp_size=pp_base, - fuse_norm_quant=fusion_val, - fuse_act_quant=fusion_val, - eager_mode=True, - ) - ) - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - - def iter_params(self, model_id: str): - opts = self.test_options - - for parallel_setup in self.parallel_setups: - for backend in self.distributed_backends: - yield ( - model_id, - parallel_setup, - backend, - self.runner, - opts, - ) - - -def _build_sp_args( - model_id: str, - parallel_setup: ParallelSetup, - distributed_backend: str, - runner: RunnerOption, - test_options: SPTestOptions, - num_gpus_available: int, - use_inductor_graph_partition: bool, - fuse_gemm_comms: bool, - enable_prompt_embeds: bool, - *, - is_multimodal: bool, -) -> tuple[list[str], list[str]] | None: - ( - tp_size, - pp_size, - fuse_norm_quant, - fuse_act_quant, - eager_mode, - ) = parallel_setup - - multi_node_only = test_options.multi_node_only - load_format = test_options.load_format - - model_info = test_options.model_info or HF_EXAMPLE_MODELS.find_hf_info(model_id) - model_info.check_transformers_version(on_fail="skip") - - trust_remote_code = model_info.trust_remote_code - tokenizer_mode = model_info.tokenizer_mode - hf_overrides = dict(model_info.hf_overrides) - require_embed_inputs = model_info.require_embed_inputs - - if load_format == "dummy": - # Avoid OOM - text_overrides = { - "num_hidden_layers": 4, - "hidden_size": 512, - "intermediate_size": 800, - "num_attention_heads": 4, - "num_key_value_heads": 1, - } - - if is_multimodal: - hf_overrides.update({"text_config": text_overrides}) - else: - hf_overrides.update(text_overrides) - else: - model_info.check_available_online(on_fail="skip") - - if num_gpus_available < tp_size * pp_size: - return None - if VLLM_MULTI_NODE and distributed_backend == "mp": - return None - if multi_node_only and not VLLM_MULTI_NODE: - return None - - common_args = [ - "--dtype", - "bfloat16", - "--max-model-len", - "2048", - "--max-num-seqs", - "8", - ] - if eager_mode: - common_args.append("-cc.cudagraph_mode=none") - if runner != "auto": - common_args.extend(["--runner", runner]) - if trust_remote_code: - common_args.append("--trust-remote-code") - if tokenizer_mode: - common_args.extend(["--tokenizer-mode", tokenizer_mode]) - if load_format: - common_args.extend(["--load-format", load_format]) - if hf_overrides: - common_args.extend(["--hf-overrides", json.dumps(hf_overrides)]) - if require_embed_inputs: - common_args.extend( - [ - "--skip-tokenizer-init", - "--enable-prompt-embeds", - "--enable-mm-embeds", - ] - ) - elif enable_prompt_embeds: - common_args.append("--enable-prompt-embeds") - - compilation_config = { - "mode": CompilationMode.VLLM_COMPILE, - "compile_sizes": [4, 8], - "pass_config": { - "enable_sp": True, - "fuse_gemm_comms": fuse_gemm_comms, - "fuse_norm_quant": fuse_norm_quant, - "fuse_act_quant": fuse_act_quant, - "fuse_allreduce_rms": False, - "eliminate_noops": True, - "sp_min_token_num": 0, - }, - "use_inductor_graph_partition": use_inductor_graph_partition, - } - if not use_inductor_graph_partition: - compilation_config["splitting_ops"] = [] - - tp_sp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--pipeline-parallel-size", - str(pp_size), - "--distributed-executor-backend", - distributed_backend, - "--compilation_config", - json.dumps(compilation_config), - ] - - tp_args = [ - *common_args, - "--tensor-parallel-size", - str(tp_size), - "--distributed-executor-backend", - "mp", - ] - - return tp_args, tp_sp_args - - -def _compare_sp_settings( - model_id: str, - settings: list[tuple[list[str], list[str]]], - *, - method: Literal["generate", "encode"], -) -> None: - if not settings: - pytest.skip("No supported sequence-parallel configurations") - - settings_by_tp_args: dict[tuple[str, ...], list[list[str]]] = {} - for tp_args, tp_sp_args in settings: - settings_by_tp_args.setdefault(tuple(tp_args), []).append(tp_sp_args) - - for tp_args_key, grouped_tp_sp_args in settings_by_tp_args.items(): - all_args = [list(tp_args_key), *grouped_tp_sp_args] - compare_all_settings( - model_id, - all_args, - [None] * len(all_args), - method=method, - force_v1_runner=True, - ) - - -SP_TEXT_GENERATION_MODELS = { - # [Decoder-only] - "hmellor/tiny-random-LlamaForCausalLM": SPTestSettings.fast(), - "RedHatAI/Llama-3.2-1B-Instruct-FP8": SPTestSettings.fp8_quant(), -} - - -@pytest.mark.parametrize( - ("model_id", "settings"), - [ - pytest.param(model_id, settings, id=model_id) - for model_id, settings in SP_TEXT_GENERATION_MODELS.items() - ], -) -@create_new_process_for_each_test() -def test_tp_sp_generation( - model_id: str, - settings: SPTestSettings, - num_gpus_available, -): - # Skip FP8 SP-only test on sm89 (compute capability 8.9). - # An unknown capability (None) is left to fail in the test body rather than - # being silently skipped here. - capability = current_platform.get_device_capability() - if "fp8" in model_id.lower() and capability is not None and capability < (9, 0): - pytest.skip("FP8 reduction support begins with sm90 capable devices.") - - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for _, parallel_setup, backend, runner, test_options in settings.iter_params( - model_id - ): - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - backend, - runner, - test_options, - num_gpus_available, - use_inductor_graph_partition, - fuse_gemm_comms=False, # TODO: enable async TP - enable_prompt_embeds=False, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -# Focused regression test for the SP + prompt_embeds graph-rewrite path. -# Covers pp_size=1 (SP only) and pp_size=2 (SP + PP); kept small on purpose so -# we don't double the matrix of `test_tp_sp_generation` above. -SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [ - ParallelSetup( - tp_size=2, - pp_size=pp_size, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=False, - ) - for pp_size in [1, 2] -] - - -@create_new_process_for_each_test() -def test_tp_sp_generation_prompt_embeds( - num_gpus_available, -): - model_id = "hmellor/tiny-random-LlamaForCausalLM" - graph_partition_options = [False] - if is_torch_equal_or_newer("2.9.0.dev"): - graph_partition_options.append(True) - - comparisons = [] - for parallel_setup in SP_PROMPT_EMBEDS_PARALLEL_SETUPS: - for use_inductor_graph_partition in graph_partition_options: - comparison = _build_sp_args( - model_id, - parallel_setup, - distributed_backend="mp", - runner="auto", - test_options=SPTestOptions(multi_node_only=False, load_format=None), - num_gpus_available=num_gpus_available, - use_inductor_graph_partition=use_inductor_graph_partition, - fuse_gemm_comms=False, - enable_prompt_embeds=True, - is_multimodal=False, - ) - if comparison is not None: - comparisons.append(comparison) - - _compare_sp_settings(model_id, comparisons, method="generate") - - -@create_new_process_for_each_test() -def test_tp_sp_nvfp4_generation(num_gpus_available: int): - is_blackwell = current_platform.is_cuda() and ( - current_platform.is_device_capability_family(100) - ) - is_rocm_cdna = False - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_cdna - - is_rocm_cdna = on_cdna() - if not (is_blackwell or is_rocm_cdna): - pytest.skip("NVFP4 sequence parallelism requires Blackwell or ROCm CDNA") - - comparison = _build_sp_args( - NVFP4_MODEL_ID, - ParallelSetup( - tp_size=2, - pp_size=1, - fuse_norm_quant=True, - fuse_act_quant=True, - eager_mode=True, - ), - "mp", - "auto", - SPTestOptions( - multi_node_only=False, - load_format="dummy", - model_info=NVFP4_MODEL_INFO, - ), - num_gpus_available, - use_inductor_graph_partition=False, - fuse_gemm_comms=False, - enable_prompt_embeds=False, - is_multimodal=False, - ) - _compare_sp_settings( - NVFP4_MODEL_ID, - [] if comparison is None else [comparison], - method="generate", - ) diff --git a/tests/compile/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 1f92ff5a1988..a4d996b48ddd 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/tests/compile/fusions_e2e/common.py @@ -21,8 +21,6 @@ class Matches(NamedTuple): # distributed ar_rms_fusion: int = 0 aiter_ar_rms_fusion: int = 0 - sequence_parallel: int = 0 - async_tp: int = 0 class ModelFusionInfo(NamedTuple): @@ -143,8 +141,4 @@ def has_cuda_graph_wrapper_metadata() -> bool: "aiter_ar_rms_fusion": re.compile( r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns" ), - "sequence_parallel": re.compile( - r"sequence_parallelism.py:\d+] Replaced (\d+) patterns" - ), - "async_tp": re.compile(r"collective_fusion.py:\d+] Replaced (\d+) patterns"), } diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index db7ba05a5734..8da7aa2a5a94 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -170,26 +170,8 @@ def run( for match_name in matches_check: log_matches = list(int(ms) for ms in log_matches_dict[match_name]) - # AR+RMS skips the largest range; SP skips the smallest. - # When both are enabled, AR+RMS activation count is - # model-dependent (hidden_size affects threshold), so derive - # from log data. - if ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - assert ( - len(log_matches) >= tp_size and len(log_matches) % tp_size == 0 - ), ( - f"Expected multiple of {tp_size} ar_rms log entries, " - f"found {len(log_matches)}" - ) - num_ranges_activated = len(log_matches) // tp_size - elif ( - match_name in ("ar_rms_fusion", "sequence_parallel") - and num_compile_ranges >= 2 - ): + # AR+RMS skips the largest range. + if match_name == "ar_rms_fusion" and num_compile_ranges >= 2: num_ranges_activated = num_compile_ranges - 1 else: num_ranges_activated = num_compile_ranges @@ -233,44 +215,6 @@ def run( f"Expecting at least {expected_matches - matches.ar_rms_fusion} " f"where ar+rms+quant was activated" ) - elif ( - match_name == "async_tp" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # AsyncTP only finds patterns on ranges where SP ran. - n_sp_ranges = num_compile_ranges - 1 - assert ( - sum(m == expected_matches for m in log_matches) - == tp_size * n_sp_ranges - ), ( - f"Expecting {expected_matches} async_tp on " - f"{tp_size * n_sp_ranges} SP-range entries, " - f"found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size, ( - f"Expecting 0 async_tp on {tp_size} small-range entries " - f"(no SP), found: {log_matches}" - ) - elif ( - match_name == "ar_rms_fusion" - and "sequence_parallel" in matches_check - and num_compile_ranges >= 2 - ): - # SP consumes allreduce patterns first, so AR+RMS finds - # full matches only on the smallest range (no SP). - assert sum(m == expected_matches for m in log_matches) == tp_size, ( - f"Expecting {expected_matches} ar_rms on " - f"{tp_size} small-range entries, found: {log_matches}" - ) - assert sum(m == 0 for m in log_matches) == tp_size * ( - num_ranges_activated - 1 - ), ( - f"Expecting 0 ar_rms on " - f"{tp_size * (num_ranges_activated - 1)} large-range " - f"entries (SP took precedence), found: {log_matches}" - ) - elif match_name == "act_quant_fusion": actual_match = match_table.get("activation_quant_fusion_pass", 0) assert actual_match == expected_matches * n_expected, ( @@ -305,17 +249,4 @@ def run( f"(found {len(log_matches)}) in:\n {log_holder.text}" ) - if match_name == "sequence_parallel" and num_compile_ranges >= 2: - log_matches = re.findall( - r"pass_manager.py:\d+] Skipping " - r".*SequenceParallelismPass.* with compile range", - log_holder.text, - ) - - n_expected = tp_size * (num_compile_ranges - num_ranges_activated) - assert len(log_matches) == n_expected, ( - f'Could not find {n_expected} "Skipping SequenceParallelismPass" ' - f"(found {len(log_matches)}) in:\n {log_holder.text}" - ) - return run diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 082ab0bf5296..d80d171984c3 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -73,8 +73,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -85,8 +83,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -96,8 +92,6 @@ act_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 4, ), ) @@ -127,8 +121,6 @@ rms_quant_fusion=n_layers, attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -138,8 +130,6 @@ matches=lambda n_layers: Matches( attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2, - sequence_parallel=n_layers * 2, - async_tp=n_layers * 2 - 1, ), ) @@ -149,8 +139,6 @@ norm_rope_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -161,8 +149,6 @@ norm_rope_fusion=n_layers, attn_quant_fusion=0, # attn + group quant not supported ar_rms_fusion=n_layers * 2 + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), ) @@ -192,9 +178,6 @@ # MLA attn + per-group FP8 quant attn_quant_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, - # TODO - # sequence_parallel= n_layers * 2 + 1, - # async_tp=n_layers * 2, ), ) @@ -213,8 +196,6 @@ matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, aiter_ar_rms_fusion=n_layers + 1, - sequence_parallel=n_layers * 2 + 1, - async_tp=n_layers * 2, ), model_kwargs=( {"quantization_config": {"moe": {"activation": "mxfp8"}}} diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py deleted file mode 100644 index b0d32dad1c17..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_async_tp.py +++ /dev/null @@ -1,339 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable - -import pytest - -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from ...utils import multi_gpu_test -from .common import ( - INDUCTOR_GRAPH_PARTITION, - AttentionBackendCase, - Matches, - custom_ops_combos, - is_blackwell, - nvfp4_kernel_exposes_input_quant_key, -) -from .models import ( - FLASHINFER_ATTN, - TRITON_ATTN, - llama3_8b, - llama3_8b_fp4, - llama3_8b_fp8, - llama4_scout_fp8, - qwen3_a3b, -) - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4], -) -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") -def test_tp2_async_tp_nvfp4_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - if nvfp4_kernel_exposes_input_quant_key(): - pytest.skip( - "NVFP4 kernel exposes input_quant_key; manual fusion fires " - "instead of compiler pass-based fusion" - ) - - # NVFP4 currently wires the all-gather + GEMM path only. - matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_act_quant=False, - fuse_attn_quant=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "attn_quant_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_async_tp_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp8, llama4_scout_fp8], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fp8_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - fuse_norm_quant=True, - fuse_act_quant=False, - fuse_attn_quant=True, - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "rms_quant_fusion", - "norm_rope_fusion", - "attn_quant_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b], -) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -@pytest.mark.parametrize("n_layers", [4]) -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -def test_tp2_sp_ar_rms_fusions( - model_name: str, - matches_fn: Callable[[int], Matches], - model_kwargs: dict, - hf_overrides: Callable[[int], dict], - attn_backend: AttentionBackendCase, - n_layers: int, - custom_ops: str, - inductor_graph_partition: bool, - run_e2e_fusion_test, -): - matches = matches_fn(n_layers) - - # Reduce size of model and skip weight loading time - model_kwargs["hf_overrides"] = hf_overrides(n_layers) - model_kwargs["load_format"] = "dummy" - model_kwargs["max_model_len"] = 1024 - - compilation_config = dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=True, - # Override threshold for testing (models have small hidden_size) - sp_min_token_num=512, - ), - ) - - matches_check = [ - "norm_rope_fusion", - "ar_rms_fusion", - "sequence_parallel", - "async_tp", - ] - - run_e2e_fusion_test( - model_name, - matches, - model_kwargs, - attn_backend, - compilation_config, - matches_check, - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py b/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py deleted file mode 100644 index 3143143bec86..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_async_tp.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Bounded ROCm model coverage for actual sequence-parallel AsyncTP rewrites.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION, custom_ops_combos -from .models import TRITON_ATTN, llama3_8b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) -@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm"))) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_async_tp_bf16_fusions( - inductor_graph_partition, custom_ops, attn_backend, run_e2e_fusion_test -): - n_layers = 4 - model_kwargs = dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ) - run_e2e_fusion_test( - llama3_8b.model_name, - llama3_8b.matches(n_layers), - model_kwargs, - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=custom_ops.split(","), - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py b/tests/compile/fusions_e2e/test_tp2_rocm_moe.py deleted file mode 100644 index 9a228ca76c50..000000000000 --- a/tests/compile/fusions_e2e/test_tp2_rocm_moe.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm MoE coverage for the TP2 GEMM/collective fusion contract.""" - -import pytest - -from tests.utils import multi_gpu_test -from vllm.config import PassConfig -from vllm.platforms import current_platform - -from .common import INDUCTOR_GRAPH_PARTITION -from .models import TRITON_ATTN, qwen3_a3b - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract") -@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION[:1]) -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) -def test_tp2_rocm_moe_async_tp_bf16_fusions( - inductor_graph_partition, attn_backend, run_e2e_fusion_test -): - """MoE dense projections retain the Qwen SP/AsyncTP rewrite counts.""" - n_layers = 4 - run_e2e_fusion_test( - qwen3_a3b.model_name, - qwen3_a3b.matches(n_layers), - dict( - hf_overrides=dict( - num_hidden_layers=n_layers, - hidden_size=512, - intermediate_size=1024, - moe_intermediate_size=128, - num_attention_heads=8, - num_key_value_heads=2, - head_dim=64, - num_experts=8, - num_experts_per_tok=2, - ), - load_format="dummy", - max_model_len=1024, - kv_cache_memory_bytes=64 * 1024**2, - gpu_memory_utilization=0.01, - max_num_batched_tokens=1024, - max_num_seqs=32, - ), - attn_backend, - dict( - use_inductor_graph_partition=inductor_graph_partition, - custom_ops=["-rms_norm"], - pass_config=PassConfig( - enable_qk_norm_rope_fusion=True, - enable_sp=True, - fuse_gemm_comms=True, - fuse_allreduce_rms=False, - sp_min_token_num=512, - ), - ), - ["norm_rope_fusion", "sequence_parallel", "async_tp"], - tp_size=2, - ) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py deleted file mode 100644 index c0fd4c00a8f7..000000000000 --- a/tests/compile/passes/distributed/test_async_tp.py +++ /dev/null @@ -1,394 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import ( - multi_gpu_test, -) -from vllm.compilation.passes.fusion.collective_fusion import AsyncTPPass -from vllm.config import ( - CompilationConfig, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type -FP8_DTYPE = current_platform.fp8_dtype() - -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestMMRSModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.gate_proj = torch.nn.Parameter( - torch.empty((self.hidden_size * 2, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.gate_proj, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + reduce scatter in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - - # matrix multiplication - permute = self.gate_proj.permute(1, 0) - mm = torch.mm(view, permute) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_matmul_reduce_scatter.default] - - -class TestAGMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = torch.nn.Parameter( - torch.empty((hidden_size, hidden_size)), requires_grad=False - ) - # Initialize weights - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - """Forward pass implementing the mm + all gather in the FX graph.""" - # Reshape input - view = hidden_states.reshape(-1, self.hidden_size) - all_gather = tensor_model_parallel_all_gather(view, dim=0) - permute = self.weight.permute(1, 0) - mm = torch.mm(all_gather, permute) - return mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_matmul.default] - - -class _BaseScaledMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.weight = ( - torch.empty([hidden_size, hidden_size], dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - # Initialize scale_b for _scaled_mm. - self.scale_b = torch.ones(1, self.hidden_size, dtype=torch.float32) - - -class TestScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the scaled_mm + reduce scatter in the FX graph.""" - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - fp8_input, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(scaled_mm, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + scaled_mm in the FX graph.""" - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - scaled_mm = torch._scaled_mm( - all_gather, - self.weight, - scale_a=scale_a, - scale_b=self.scale_b, - out_dtype=self.dtype, - ) - return scaled_mm - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -class TestCutlassScaledMMRSModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the cutlass_scaled_mm + reduce scatter - in the FX graph - - """ - fp8_input = input.to(FP8_DTYPE) - scale_a = torch.ones(input.shape[0], 1, dtype=torch.float32) - mm_out = torch.empty( - (fp8_input.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=input.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, fp8_input, self.weight, scale_a, self.scale_b, None - ) - reduce_scatter = tensor_model_parallel_reduce_scatter(mm_out, dim=0) - return reduce_scatter - - def ops_in_model_before(self): - return [torch.ops.vllm.reduce_scatter.default] - - def ops_in_model_after(self): - return [torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter.default] - - -class TestAGCutlassScaledMMModel(_BaseScaledMMModel): - def forward(self, input: torch.Tensor): - """Forward pass implementing the all gather + cutlass_scaled_mm - in the FX graph - """ - # Reshape input - fp8_input = input.to(FP8_DTYPE) - all_gather = tensor_model_parallel_all_gather(fp8_input, dim=0) - - scale_a = torch.ones(all_gather.shape[0], 1, dtype=torch.float32) - - mm_out = torch.empty( - (all_gather.shape[0], self.weight.shape[1]), - dtype=self.dtype, - device=all_gather.device, - ) - torch.ops._C.cutlass_scaled_mm( - mm_out, all_gather, self.weight, scale_a, self.scale_b, None - ) - return mm_out - - def ops_in_model_before(self): - return [torch.ops.vllm.all_gather.default] - - def ops_in_model_after(self): - return [torch.ops.symm_mem.fused_all_gather_scaled_matmul.default] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model", - [ - TestMMRSModel, - TestAGMMModel, - TestScaledMMRSModel, - TestAGScaledMMModel, - pytest.param( - TestCutlassScaledMMRSModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - pytest.param( - TestAGCutlassScaledMMModel, - marks=pytest.mark.skipif( - not hasattr(torch.ops._C, "cutlass_scaled_mm"), - reason="Requires cutlass_scaled_mm", - ), - ), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("dynamic", [True, False]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_async_tp_pass_replace( - test_model: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, -): - num_processes = 2 - distributed_init_method = get_file_store_init_method() - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model, - batch_size, - seq_len, - hidden_size, - dtype, - dynamic, - distributed_init_method, - ), - nprocs=nprocs, - ) - - run_torch_spawn(async_tp_pass_on_test_model, num_processes) - - -def test_async_tp_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - async_tp_pass = object.__new__(AsyncTPPass) - async_tp_pass.compilation_config = vllm_config.compilation_config - - with pytest.raises( - AssertionError, match="AsyncTPPass requires full-graph compilation" - ): - async_tp_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def async_tp_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: torch.nn.Module, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - dynamic: bool, - distributed_init_method: str, - check_numerics: bool = False, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - } - ) - - # initialize distributed - init_distributed_environment( - world_size=world_size, - rank=local_rank, - distributed_init_method=distributed_init_method, - local_rank=local_rank, - ) - - # configure vllm config for SequenceParallelismPass - vllm_config = VllmConfig() - vllm_config.compilation_config = CompilationConfig( - pass_config=PassConfig( - fuse_gemm_comms=True, - ), - ) - vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - vllm_config.model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - - async_tp_pass = AsyncTPPass(vllm_config) - backend = TestBackend(async_tp_pass) - - assert ( - async_tp_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - async_tp_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - - model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor - - hidden_states = torch.randn( - (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False - ) - if check_numerics: - hidden_states += local_rank * 0.125 - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - actual = compiled_model(hidden_states) - - if check_numerics: - # These are new checks, so no existing tolerance is relaxed. - # Allow BF16 rounding vs. eager, but keep the lifetime check exact. - expected = model(hidden_states) - torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) - snapshot = actual.clone() - changed = compiled_model(-hidden_states) - torch.testing.assert_close(changed, -expected, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(actual, snapshot, rtol=0, atol=0) - - assert async_tp_pass.matched_count == 1 - - # In pre-nodes, all gather or reduce scatter should exist, - # fused_matmul_reduce_scatter or fused_all_gather_matmul should not - backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) - - # In post-nodes, fused_matmul_reduce_scatter or \ - # fused_all_gather_matmul should exist - backend.check_after_ops(model.ops_in_model_after()) diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py deleted file mode 100644 index c40d75f6754a..000000000000 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ /dev/null @@ -1,343 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -import torch - -import vllm.envs as envs -from tests.compile.backend import TestBackend -from tests.utils import TestFP8Layer, multi_gpu_test -from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass -from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass -from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass -from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass -from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass -from vllm.config import ( - CompilationConfig, - CUDAGraphMode, - DeviceConfig, - ModelConfig, - PassConfig, - VllmConfig, - get_current_vllm_config, - set_current_vllm_config, -) -from vllm.config.utils import Range -from vllm.distributed import tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - init_distributed_environment, - initialize_model_parallel, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) -from vllm.platforms import current_platform -from vllm.utils.system_utils import update_environment_variables -from vllm.utils.torch_utils import set_random_seed - -DEVICE_TYPE = current_platform.device_type - -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") - -FP8_DTYPE = current_platform.fp8_dtype() -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] - - -class TestAllReduceRMSNormModel(torch.nn.Module): - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)] - - def forward(self, x): - z = torch.relu(x) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = torch.mm(y, self.w[0]) - x2 = tensor_model_parallel_all_reduce(z2) - - y2, resid = self.norm[1](x2, resid) - - z3 = torch.mm(y2, self.w[1]) - x3 = tensor_model_parallel_all_reduce(z3) - - y3, resid = self.norm[2](x3, resid) - - z4 = torch.mm(y3, self.w[2]) - x4 = tensor_model_parallel_all_reduce(z4) - - y4, resid = self.norm[3](x4, resid) - return y4 - - def ops_in_model_before(self): - return [torch.ops.vllm.all_reduce.default] - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model(self): - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - ] - - -class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module): - quant_key = kFp8StaticTensorSym - - def __init__(self, hidden_size=16, eps=1e-6): - super().__init__() - self.vllm_config = get_current_vllm_config() - self.hidden_size = hidden_size - self.eps = eps - self.norm = [RMSNorm(hidden_size, eps) for i in range(4)] - self.fp8_linear_layers = [ - TestFP8Layer( - weight_shape=(hidden_size, hidden_size), - activation_quant_key=self.quant_key, - weight_quant_key=self.quant_key, - input_dtype=self.vllm_config.model_config.dtype, - ) - for i in range(3) - ] - - def forward(self, hidden_states): - # avoid having graph input be an arg to a pattern directly - z = torch.relu(hidden_states) - x = resid = tensor_model_parallel_all_reduce(z) - y = self.norm[0](x) - - z2 = self.fp8_linear_layers[0](y) - - x2 = tensor_model_parallel_all_reduce(z2) - y2, resid = self.norm[1](x2, resid) - - z3 = self.fp8_linear_layers[1](y2) - - x3 = tensor_model_parallel_all_reduce(z3) - y3, resid = self.norm[2](x3, resid) # use resid here - - z4 = self.fp8_linear_layers[2](y3) - x4 = tensor_model_parallel_all_reduce(z4) - y4, resid = self.norm[3](x4, resid) # use resid here - return y4 - - def ops_in_model_after(self): - return [ - torch.ops.vllm.all_gather.default, - torch.ops.vllm.reduce_scatter.default, - ] - - def ops_in_model_before(self): - return [ - torch.ops.vllm.all_reduce.default, - ] - - def ops_in_model(self): - if self.vllm_config.compilation_config.pass_config.fuse_norm_quant: - return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default] - else: - quant_ops = ( - [torch.ops._C.static_scaled_fp8_quant.default] - if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers) - else [torch.ops.aten.reciprocal] - ) - return [ - torch.ops.vllm_ir.rms_norm, - torch.ops.vllm_ir.fused_add_rms_norm, - *quant_ops, - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize( - "test_model_cls, custom_ops", - [ - (TestAllReduceRMSNormModel, "+rms_norm"), - (TestAllReduceRMSNormModel, "-rms_norm"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "+rms_norm,-quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,+quant_fp8"), - (TestAllReduceRMSNormStaticQuantFP8Model, "-rms_norm,-quant_fp8"), - ], -) -@pytest.mark.parametrize("batch_size", [8]) -@pytest.mark.parametrize("seq_len", [16]) -@pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("fuse_norm_quant", [True, False]) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") -def test_sequence_parallelism_pass( - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - num_processes = 2 - - def run_torch_spawn(fn, nprocs): - # need to use torch.mp.spawn otherwise will have problems with - # torch.distributed and cuda - torch.multiprocessing.spawn( - fn, - args=( - num_processes, - test_model_cls, - custom_ops, - batch_size, - seq_len, - hidden_size, - dtype, - fuse_norm_quant, - dynamic, - ), - nprocs=nprocs, - ) - - run_torch_spawn(sequence_parallelism_pass_on_test_model, num_processes) - - -def test_sequence_parallelism_pass_requires_full_graph_compilation(): - vllm_config = VllmConfig() - vllm_config.compilation_config.use_inductor_graph_partition = False - vllm_config.compilation_config.splitting_ops = [ - "vllm::unified_attention_with_output" - ] - - sequence_parallelism_pass = object.__new__(SequenceParallelismPass) - sequence_parallelism_pass.compilation_config = vllm_config.compilation_config - sequence_parallelism_pass.min_token_num = 1 - - with pytest.raises( - AssertionError, - match="SequenceParallelismPass requires full-graph compilation", - ): - sequence_parallelism_pass.is_applicable_for_range(Range(start=8, end=8)) - - -def sequence_parallelism_pass_on_test_model( - local_rank: int, - world_size: int, - test_model_cls: type[torch.nn.Module], - custom_ops: str, - batch_size: int, - seq_len: int, - hidden_size: int, - dtype: torch.dtype, - fuse_norm_quant: bool, - dynamic: bool, -): - set_random_seed(0) - - device = torch.device(f"{DEVICE_TYPE}:{local_rank}") - torch.accelerator.set_device_index(device) - torch.set_default_device(device) - torch.set_default_dtype(dtype) - - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": "12345", - } - ) - - # initialize distributed - init_distributed_environment() - - # configure vllm config for SequenceParallelismPass - custom_ops_list = custom_ops.split(",") if custom_ops else [] - compilation_config = CompilationConfig( - splitting_ops=[], # avoid automatic rms_norm enablement - cudagraph_mode=CUDAGraphMode.NONE, # avoid piecewise warnings - custom_ops=custom_ops_list, - pass_config=PassConfig( - enable_sp=True, - fuse_norm_quant=fuse_norm_quant, - eliminate_noops=True, - ), - ) # NoOp needed for fusion - device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) - - # this is a fake model name to construct the model config - # in the vllm_config, it's not really used. - model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8" - model_config = ModelConfig( - model=model_name, trust_remote_code=True, dtype=dtype, seed=42 - ) - - vllm_config = VllmConfig( - model_config=model_config, - device_config=device_config, - compilation_config=compilation_config, - ) - - with set_current_vllm_config(vllm_config): - initialize_model_parallel(tensor_model_parallel_size=world_size) - noop_pass = NoOpEliminationPass(vllm_config) - sequence_parallelism_pass = SequenceParallelismPass(vllm_config) - cleanup_pass = PostCleanupPass(vllm_config) - assert ( - sequence_parallelism_pass.compilation_config.splitting_ops - == vllm_config.compilation_config.splitting_ops - ) - assert ( - sequence_parallelism_pass.compilation_config.use_inductor_graph_partition - == vllm_config.compilation_config.use_inductor_graph_partition - ) - passes_for_backend: list[VllmInductorPass] = [ - noop_pass, - sequence_parallelism_pass, - ] - - if fuse_norm_quant: - fusion_pass = RMSNormQuantFusionPass(vllm_config) - passes_for_backend.append(fusion_pass) - - passes_for_backend.append(cleanup_pass) - - backend = TestBackend(*passes_for_backend) - - model = test_model_cls(hidden_size) - - hidden_states = torch.randn((batch_size * seq_len, hidden_size), dtype=dtype) - - if dynamic: - torch._dynamo.mark_dynamic(hidden_states, 0) - - compiled_model = torch.compile(model, backend=backend) - compiled_model(hidden_states) - - assert sequence_parallelism_pass.matched_count == 4 - - # In pre-nodes, all reduce should be there, - # reduce scatter and all gather should not - for op in model.ops_in_model_before(): - assert backend.op_count(op, before=True) == 4 - - # In post-nodes, reduce scatter and all gather should be there, - # all reduce should not - for op in model.ops_in_model_after(): - assert backend.op_count(op, before=False) == 4 - - for op in model.ops_in_model(): - assert backend.op_count(op, before=False) > 0 diff --git a/tests/compile/rocm/__init__.py b/tests/compile/rocm/__init__.py deleted file mode 100644 index 208f01a7cb5e..000000000000 --- a/tests/compile/rocm/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/compile/rocm/test_async_tp.py b/tests/compile/rocm/test_async_tp.py deleted file mode 100644 index a4462265f8e7..000000000000 --- a/tests/compile/rocm/test_async_tp.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""ROCm AsyncTP numerical coverage kept outside the shared distributed suite.""" - -import pytest -import torch - -from tests.compile.passes.distributed.test_async_tp import ( - TestAGMMModel, - TestMMRSModel, - async_tp_pass_on_test_model, -) -from tests.utils import multi_gpu_test -from vllm.distributed import ( - tensor_model_parallel_all_gather, - tensor_model_parallel_reduce_scatter, -) -from vllm.platforms import current_platform -from vllm.utils.network_utils import get_file_store_init_method - - -class TestRocmMMRSModel(TestMMRSModel): - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None) - return tensor_model_parallel_reduce_scatter(mm, dim=0) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -class TestRocmAGMMModel(TestAGMMModel): - def __init__(self, hidden_size=16, dtype=torch.bfloat16): - super().__init__(hidden_size, dtype) - self.weight = torch.nn.Parameter( - torch.empty((hidden_size * 2, hidden_size)), requires_grad=False - ) - torch.nn.init.normal_(self.weight, std=0.02) - - def forward(self, hidden_states): - view = hidden_states.reshape(-1, self.hidden_size) - gathered = tensor_model_parallel_all_gather(view, dim=0) - return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None) - - def ops_in_model_before(self): - return super().ops_in_model_before() + [ - torch.ops.vllm.rocm_unquantized_gemm.default - ] - - -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract") -@pytest.mark.parametrize( - "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel] -) -@pytest.mark.parametrize("dynamic", [False, True]) -def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic): - """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives.""" - torch.multiprocessing.spawn( - async_tp_pass_on_test_model, - args=( - 2, - test_model, - 8, - 16, - 128, - torch.bfloat16, - dynamic, - get_file_store_init_method(), - True, - ), - nprocs=2, - ) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 5a1ba13a57f1..4e7a756f8cbe 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -464,49 +464,41 @@ def test_should_split(): ( "cudagraph_capture_sizes", "max_cudagraph_capture_size", - "tp_size", - "enable_sp", "max_num_batched_tokens", "cudagraph_mode", "expected_max_size", ), [ - (None, None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([1, 2, 4], 4, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + (None, None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([1, 2, 4], 4, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), ( [1, 2, 4], 8, - 1, - False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError, ), - ([1, 256], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), - ([], None, 1, False, 2048, CUDAGraphMode.NONE, 0), - (None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0), + ([1, 256], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + ([], None, 2048, CUDAGraphMode.NONE, 0), + (None, 0, 2048, CUDAGraphMode.NONE, 0), # truncated to nearest multiple of 8 or 16 - (None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), + (None, 257, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256), # max_num_batched_tokens <= max_cudagraph_capture_size should always be # captured even if not landing on a 16-stride step - (None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), + (None, 2048, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257), # max from list - ([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), - # SP forces full-graph compilation, sizes are filtered by TP - ([1, 2, 4, 15], None, 2, True, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15), # limited by the max_tokens - ([1, 2, 4, 15], None, 1, False, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), + ([1, 2, 4, 15], None, 8, CUDAGraphMode.FULL_AND_PIECEWISE, 4), # the list should contain at least 1 element when use cudagraph - ([], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + ([], None, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), # the max capturing size should be >= 1 when use cudagraph - (None, 0, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), + (None, 0, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, ValidationError), ], ) def test_cudagraph_sizes_post_init( cudagraph_capture_sizes, max_cudagraph_capture_size, - tp_size, - enable_sp, max_num_batched_tokens, cudagraph_mode, expected_max_size, @@ -515,10 +507,7 @@ def test_cudagraph_sizes_post_init( if expected_max_size == ValidationError: ctx = pytest.raises(expected_max_size) - with ( - ctx, - patch.object(current_platform, "device_count", return_value=tp_size), - ): + with ctx: kwargs = {} if cudagraph_capture_sizes is not None: kwargs["cudagraph_capture_sizes"] = cudagraph_capture_sizes @@ -526,18 +515,15 @@ def test_cudagraph_sizes_post_init( kwargs["max_cudagraph_capture_size"] = max_cudagraph_capture_size compilation_config = CompilationConfig( pass_config=PassConfig( - enable_sp=enable_sp, fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True, - sp_min_token_num=512 if enable_sp else None, ), cudagraph_mode=cudagraph_mode, **kwargs, ) engine_args = EngineArgs( model="facebook/opt-125m", - tensor_parallel_size=tp_size, max_num_seqs=min(max_num_batched_tokens, 128), max_num_batched_tokens=max_num_batched_tokens, compilation_config=compilation_config, @@ -728,35 +714,6 @@ def test_default_cudagraph_capture_sizes_keep_all_sizes_bounded(): ) -def test_cudagraph_capture_sizes_respect_sequence_parallelism(): - """Sequence-parallel capture sizes stay divisible by tensor parallel size.""" - compilation_config = CompilationConfig( - cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE - ) - compilation_config.pass_config.enable_sp = True - config = _mock_config_for_cudagraph_sizes( - max_num_seqs=32, - num_speculative_tokens=16, - max_num_batched_tokens=32768, - compilation_config=compilation_config, - ) - config.parallel_config = SimpleNamespace(tensor_parallel_size=2) - config.update_sizes_for_sequence_parallelism = lambda sizes: ( - VllmConfig.update_sizes_for_sequence_parallelism(config, sizes) - ) - - with patch.object( - current_platform, - "is_device_capability_family", - return_value=False, - ): - VllmConfig._set_cudagraph_sizes(config) - - assert all(size % 2 == 0 for size in compilation_config.cudagraph_capture_sizes) - assert 544 not in compilation_config.cudagraph_capture_sizes - assert compilation_config.max_cudagraph_capture_size == 512 - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") def test_real_vllm_config_caps_widest_ngram_decode_batch(): """Real config post-init preserves the default capture ceiling.""" @@ -1046,123 +1003,6 @@ def test_blackwell_cudagraph_default(is_blackwell, expected_max_size): ) -@pytest.mark.skipif( - not current_platform.support_static_graph_mode(), - reason="Skip if not cudagraph mode supported", -) -@pytest.mark.parametrize( - ( - "cudagraph_mode", - "use_inductor_graph_partition", - "expected_enable_sp", - "expected_cudagraph_mode", - "expected_piecewise_compile", - "expected_capture_sizes", - "expected_max_size", - ), - [ - (CUDAGraphMode.PIECEWISE, False, True, CUDAGraphMode.FULL, False, [2, 4], 4), - ( - CUDAGraphMode.FULL_DECODE_ONLY, - False, - True, - CUDAGraphMode.FULL_DECODE_ONLY, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - False, - True, - CUDAGraphMode.FULL, - False, - [2, 4], - 4, - ), - ( - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - True, - CUDAGraphMode.FULL_AND_PIECEWISE, - True, - [2, 4], - 4, - ), - ], -) -def test_sequence_parallelism_requires_full_graph_compilation( - cudagraph_mode: CUDAGraphMode, - use_inductor_graph_partition: bool, - expected_enable_sp: bool, - expected_cudagraph_mode: CUDAGraphMode, - expected_piecewise_compile: bool, - expected_capture_sizes: list[int], - expected_max_size: int, -): - with patch.object(current_platform, "device_count", return_value=2): - vllm_config = VllmConfig( - parallel_config=ParallelConfig(tensor_parallel_size=2), - scheduler_config=SchedulerConfig( - max_num_seqs=128, - max_num_batched_tokens=2048, - max_model_len=2048, - is_encoder_decoder=False, - ), - ) - vllm_config.model_config = MagicMock( - dtype=torch.float16, - enforce_eager=False, - is_moe=False, - disable_cascade_attn=False, - get_hidden_size=MagicMock(return_value=4096), - ) - vllm_config.compilation_config = CompilationConfig( - mode=CompilationMode.VLLM_COMPILE, - cudagraph_capture_sizes=[1, 2, 4, 15], - max_cudagraph_capture_size=None, - compile_sizes=["cudagraph_capture_sizes"], - use_inductor_graph_partition=use_inductor_graph_partition, - pass_config=PassConfig( - enable_sp=True, - fuse_gemm_comms=True, - fuse_norm_quant=True, - fuse_act_quant=True, - eliminate_noops=True, - sp_min_token_num=512, - ), - cudagraph_mode=cudagraph_mode, - ) - vllm_config.compilation_config.set_splitting_ops_for_v1( - all2all_backend=vllm_config.parallel_config.all2all_backend, - data_parallel_size=1, - ) - vllm_config._set_compile_ranges() - vllm_config._set_cudagraph_sizes() - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - == use_inductor_graph_partition - ) - assert ( - bool(vllm_config.compilation_config.splitting_ops) == expected_piecewise_compile - ) - assert vllm_config.compilation_config.pass_config.enable_sp == expected_enable_sp - assert ( - vllm_config.compilation_config.pass_config.fuse_gemm_comms == expected_enable_sp - ) - assert vllm_config.compilation_config.cudagraph_mode == expected_cudagraph_mode - assert ( - vllm_config.compilation_config.cudagraph_capture_sizes == expected_capture_sizes - ) - assert ( - vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size - ) - assert ( - 511 in vllm_config.compilation_config.compile_ranges_endpoints - ) == expected_enable_sp - - def test_cached_compilation_config(default_vllm_config): import torch from torch._inductor.utils import run_and_get_code diff --git a/tests/compile/test_sequence_parallelism_threshold.py b/tests/compile/test_sequence_parallelism_threshold.py deleted file mode 100644 index 090b77b330ab..000000000000 --- a/tests/compile/test_sequence_parallelism_threshold.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm.compilation.passes.fusion.sequence_parallelism import ( - SP_MIN_HIDDEN_SIZE, - SP_MIN_PER_GPU_SIZE_MB, - get_sequence_parallelism_threshold, -) - - -class TestGetSequenceParallelismThreshold: - """Tests for get_sequence_parallelism_threshold function.""" - - def test_non_cuda_returns_none(self, mock_cuda_platform): - """Non-CUDA platforms should return None.""" - with mock_cuda_platform(is_cuda=False): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_unsupported_device_capability_returns_none(self, mock_cuda_platform): - """Unsupported device capabilities (e.g., sm80) should return None.""" - with mock_cuda_platform(capability=(8, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=8192, tp_size=2, element_size=2 - ) - assert result is None - - def test_small_hidden_size_returns_none(self, mock_cuda_platform): - """H100 with hidden_size below threshold should return None.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=4096, - tp_size=2, - element_size=2, # 4096 < 8192 - ) - assert result is None - - def test_h100_large_model_returns_threshold(self, mock_cuda_platform): - """H100 with large enough hidden_size should return calculated threshold.""" - with mock_cuda_platform(capability=(9, 0)): - hidden_size = 8192 - tp_size = 2 - element_size = 2 # float16/bfloat16 - - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - - # Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - MiB = 1024 * 1024 - expected = int( - (SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB) - // (hidden_size * element_size) - ) - assert result == expected - assert result == 1024 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # Boundary: exactly at min hidden size threshold, tp_size=1 - # (8 * 1 * 1024 * 1024) // (8192 * 2) = 512 - (8192, 1, 2, 512), - # Larger hidden size reduces token threshold - # (8 * 1 * 1024 * 1024) // (16384 * 2) = 256 - (16384, 1, 2, 256), - # Larger tp_size increases token threshold - # (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048 - (8192, 4, 2, 2048), - # Larger element_size (fp32) reduces token threshold - # (8 * 2 * 1024 * 1024) // (8192 * 4) = 512 - (8192, 2, 4, 512), - ], - ) - def test_threshold_calculation_variations( - self, mock_cuda_platform, hidden_size, tp_size, element_size, expected - ): - """Test threshold calculation with various parameter combinations.""" - with mock_cuda_platform(capability=(9, 0)): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_hidden_size_boundary(self, mock_cuda_platform): - """Test behavior at the exact hidden_size boundary.""" - with mock_cuda_platform(capability=(9, 0)): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=SP_MIN_HIDDEN_SIZE[90], - tp_size=2, - element_size=2, - ) - assert result is not None - - -# XPU-specific constants (must match sequence_parallelism.py values) -_XPU_MIN_HIDDEN_SIZE = 4096 -_XPU_MIN_PER_GPU_SIZE_MB = 8.0 - - -class TestGetSequenceParallelismThresholdXPU: - """Tests for get_sequence_parallelism_threshold on XPU platform.""" - - def test_xpu_small_hidden_size_returns_none(self, mock_xpu_platform): - """XPU with hidden_size below threshold should return None.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - def test_xpu_large_model_returns_threshold(self, mock_xpu_platform): - """XPU with hidden_size >= threshold should return calculated value.""" - with mock_xpu_platform(): - hidden_size = _XPU_MIN_HIDDEN_SIZE - tp_size = 2 - element_size = 2 - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - # (8 * 2 * 1024 * 1024) // (4096 * 2) = 2048 - MiB = 1024 * 1024 - expected = int( - (_XPU_MIN_PER_GPU_SIZE_MB * tp_size * MiB) // (hidden_size * element_size) - ) - assert result == expected - assert result == 2048 - - @pytest.mark.parametrize( - "hidden_size,tp_size,element_size,expected", - [ - # (8 * 1 * 1024 * 1024) // (4096 * 2) = 1024 - (4096, 1, 2, 1024), - # (8 * 4 * 1024 * 1024) // (4096 * 2) = 4096 - (4096, 4, 2, 4096), - # (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 - (8192, 2, 2, 1024), - # (8 * 2 * 1024 * 1024) // (4096 * 4) = 1024 - (4096, 2, 4, 1024), - ], - ) - def test_xpu_threshold_calculation_variations( - self, mock_xpu_platform, hidden_size, tp_size, element_size, expected - ): - """Test XPU threshold calculation with various parameter combinations.""" - with mock_xpu_platform(): - result = get_sequence_parallelism_threshold( - hidden_size=hidden_size, - tp_size=tp_size, - element_size=element_size, - ) - assert result == expected - - def test_xpu_hidden_size_boundary(self, mock_xpu_platform): - """Test behavior at the exact XPU hidden_size boundary.""" - with mock_xpu_platform(): - # Just below threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE - 1, - tp_size=2, - element_size=2, - ) - assert result is None - - # Exactly at threshold - result = get_sequence_parallelism_threshold( - hidden_size=_XPU_MIN_HIDDEN_SIZE, - tp_size=2, - element_size=2, - ) - assert result is not None diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 2d57236f8efa..5416fdc8ea6d 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -582,10 +582,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch): use_ubatching=False, enable_elastic_ep=False, ), - compilation_config=SimpleNamespace( - mode=CompilationMode.NONE, - pass_config=SimpleNamespace(enable_sp=False), - ), + compilation_config=SimpleNamespace(mode=CompilationMode.NONE), cache_config=SimpleNamespace( kv_sharing_fast_prefill=False, mamba_cache_mode="none", diff --git a/tests/test_config.py b/tests/test_config.py index 0098ecdc3dac..ebf76c76a0c6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -843,7 +843,6 @@ def test_breakable_cudagraph_platform_default( ("bi-off", 2048, (16, 8), 6144, 1, False), ("opt-out", 2048, (16, 8), 6144, 1, False), ("eager", 2048, (16, 8), 6144, 1, False), - ("sp", 2048, (16, 8), 6144, 1, False), ("fp16", 2048, (16, 8), 6144, 1, False), ("quantized", 2048, (16, 8), 6144, 1, False), ("untuned", 4096, (32, 32), 11008, 1, False), @@ -886,11 +885,7 @@ def test_batch_invariant_breakable_cudagraph( get_num_kv_heads=lambda pc: heads[1], ) config.parallel_config = SimpleNamespace(tensor_parallel_size=tp) - config.compilation_config = ( - CompilationConfig(pass_config=PassConfig(enable_sp=True)) - if case == "sp" - else CompilationConfig() - ) + config.compilation_config = CompilationConfig() try: assert config._maybe_enable_breakable_cudagraph() is expected if expected: @@ -1021,7 +1016,6 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( "FakeAttentionBackend", uniform_decode_query_len=4, use_v2_model_runner=use_v2_model_runner, - tensor_parallel_size=1, ) assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE @@ -1104,7 +1098,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, ) @@ -1117,7 +1110,6 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling(): "FakeAttentionBackend", uniform_decode_query_len=1, use_v2_model_runner=True, - tensor_parallel_size=1, kv_cache_config=kv_cache_config, max_num_reqs=256, is_profiling=True, @@ -3359,8 +3351,6 @@ def test_vllm_config_explicit_overrides(): take precedence over callable defaults, across different models and optimization levels. """ - from vllm.config.compilation import PassConfig - quantized_model = ModelConfig("RedHatAI/Llama-3.2-1B-FP8") moe_model = ModelConfig("deepseek-ai/DeepSeek-V2-Lite") regular_model = ModelConfig("Qwen/Qwen1.5-7B") diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 90adce187448..6dc20a817794 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -378,9 +378,6 @@ def run_model(scheduler_output, intermediate_tensors): worker = SimpleNamespace( vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace( - pass_config=SimpleNamespace(enable_sp=False) - ), parallel_config=SimpleNamespace( pipeline_parallel_size=2, distributed_executor_backend="mp" ), diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py deleted file mode 100644 index f2aa6d7d2f1d..000000000000 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ /dev/null @@ -1,1013 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Callable -from contextlib import suppress - -import torch -import torch._inductor.pattern_matcher as pm -import torch.distributed.distributed_c10d as c10d -import torch.fx as fx -from torch._inductor.pattern_matcher import PatternMatcherPass -from torch.distributed._symmetric_memory import enable_symm_mem_for_group - -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op - -from ..inductor_pass import enable_fake_mode -from ..vllm_inductor_pass import ( - VllmFusionPatternMatcherPass, - VllmInductorPass, - VllmPatternMatcherPass, - VllmPatternReplacement, -) - -FP8_DTYPE = current_platform.fp8_dtype() - -logger = init_logger(__name__) - - -def _flashinfer_scaled_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - bias: torch.Tensor | None = None, - scale_result: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, - use_fast_accum: bool = False, -) -> None: - # Import lazily to avoid a circular import during module initialization - # when docs or other tooling import the pass without FlashInfer. - from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm_out - - assert bias is None, "FlashInfer symm_mem adapter does not support bias" - assert scale_result is None, ( - "FlashInfer symm_mem adapter does not support result scaling" - ) - assert not use_fast_accum, ( - "FlashInfer symm_mem adapter does not support use_fast_accum" - ) - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer symm_mem adapter expects 2D inputs and output" - ) - assert scale_a.numel() == 1 and scale_b.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - - flashinfer_scaled_fp8_mm_out( - A, - B, - scale_a, - scale_b, - out=out, - out_dtype=out_dtype or out.dtype, - ) - - -def _flashinfer_fp4_mm_out( - A: torch.Tensor, - B: torch.Tensor, - *, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - alpha: torch.Tensor, - out_dtype: torch.dtype | None = None, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> None: - from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm_out - - assert A.ndim == 2 and B.ndim == 2 and out.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs and output" - ) - flashinfer_scaled_fp4_mm_out( - A, - B, - scale_a, - scale_b, - alpha, - out=out, - out_dtype=out_dtype or out.dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter_fake( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - result_shape = list(output_shape) - result_shape[orig_scatter_dim] //= world_size - return torch.empty( - result_shape, - dtype=out_dtype or torch.bfloat16, - device=A.device, - ) - - -def fused_flashinfer_scaled_matmul_reduce_scatter( - A: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - reduce_op: str, - orig_scatter_dim: int, - scatter_dim_after_maybe_reshape: int, - group_name: str, - output_shape: list[int], - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert orig_scatter_dim == 0 and scatter_dim_after_maybe_reshape == 0, ( - "FlashInfer symm_mem adapter currently only supports scatter_dim=0" - ) - world_size = c10d._resolve_process_group(group_name).size() - assert A.ndim == 2 and B.ndim == 2, "FlashInfer symm_mem adapter expects 2D inputs" - assert A.is_contiguous(), "FlashInfer symm_mem adapter expects contiguous A" - assert A_scale.numel() == 1 and B_scale.numel() == 1, ( - "FlashInfer symm_mem adapter only supports tensor-wise FP8 scales" - ) - assert A.shape[0] % world_size == 0, ( - "FlashInfer symm_mem adapter expects M divisible by world size" - ) - - kwargs = { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - return torch.distributed._symmetric_memory._fused_scaled_matmul_reduce_scatter_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A=A, - B=B, - A_scale=A_scale, - kwargs=kwargs, - out_dtype=out_dtype, - reduce_op=reduce_op, - orig_scatter_dim=orig_scatter_dim, - scatter_dim_after_maybe_reshape=scatter_dim_after_maybe_reshape, - group_name=group_name, - output_shape=output_shape, - ) - - -def fused_all_gather_flashinfer_scaled_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_scaled_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale: torch.Tensor, - B_scale: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer symm_mem adapter currently only supports gather_dim=0" - ) - _, outputs = torch.distributed._symmetric_memory._fused_all_gather_matmul_impl( - mm_out_op=_flashinfer_scaled_mm_out, - A_shard=A_shard, - Bs=[B], - A_scale=A_scale, - kwargs_list=[ - { - "scale_b": B_scale, - "bias": None, - "scale_result": None, - "out_dtype": out_dtype, - "use_fast_accum": False, - } - ], - out_dtypes=[out_dtype], - gather_dim=gather_dim, - group_name=group_name, - return_A=False, - ) - return outputs[0] - - -def fused_all_gather_flashinfer_fp4_matmul_fake( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - world_size = c10d._resolve_process_group(group_name).size() - output_shape = list(A_shard.shape) - output_shape[gather_dim] *= world_size - output_shape[-1] = B.shape[1] - return torch.empty( - output_shape, - dtype=out_dtype or torch.bfloat16, - device=A_shard.device, - ) - - -def fused_all_gather_flashinfer_fp4_matmul( - A_shard: torch.Tensor, - B: torch.Tensor, - A_scale_shard: torch.Tensor, - B_scale: torch.Tensor, - alpha: torch.Tensor, - gather_dim: int, - group_name: str, - out_dtype: torch.dtype | None = None, - view_a_scale_as_fp8: bool = False, - use_8x4_sf_layout: bool = False, - backend: str = "cutlass", -) -> torch.Tensor: - assert gather_dim == 0, ( - "FlashInfer FP4 symm_mem adapter currently only supports gather_dim=0" - ) - assert A_shard.ndim == 2 and A_scale_shard.ndim == 2 and B.ndim == 2, ( - "FlashInfer FP4 symm_mem adapter expects 2D inputs" - ) - if view_a_scale_as_fp8: - A_scale_shard = A_scale_shard.view(torch.float8_e4m3fn) - - group = c10d._resolve_process_group(group_name) - world_size = group.size() - output = A_shard.new_empty( - A_shard.shape[0] * world_size, - B.shape[1], - dtype=out_dtype or torch.bfloat16, - ) - output_shards = output.chunk(world_size) - - A = A_shard.new_empty(A_shard.shape[0] * world_size, A_shard.shape[1]) - A_scale = A_scale_shard.new_empty( - A_scale_shard.shape[0] * world_size, - A_scale_shard.shape[1], - ) - - def fp4_shard_consumer(shards: list[torch.Tensor], rank: int) -> None: - _flashinfer_fp4_mm_out( - shards[0], - B, - scale_a=shards[1], - scale_b=B_scale, - alpha=alpha, - out=output_shards[rank], - out_dtype=out_dtype, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - - torch.distributed._symmetric_memory._pipelined_multi_all_gather_and_consume( - [A_shard, A_scale_shard], - fp4_shard_consumer, - [A, A_scale], - group_name, - False, - ) - return output - - -direct_register_custom_op( - op_name="fused_flashinfer_scaled_matmul_reduce_scatter", - op_func=fused_flashinfer_scaled_matmul_reduce_scatter, - fake_impl=fused_flashinfer_scaled_matmul_reduce_scatter_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_scaled_matmul", - op_func=fused_all_gather_flashinfer_scaled_matmul, - fake_impl=fused_all_gather_flashinfer_scaled_matmul_fake, -) - -direct_register_custom_op( - op_name="fused_all_gather_flashinfer_fp4_matmul", - op_func=fused_all_gather_flashinfer_fp4_matmul, - fake_impl=fused_all_gather_flashinfer_fp4_matmul_fake, -) - - -class BasePattern: - def __init__(self, dtype: torch.dtype, device: str | None) -> None: - self.dtype = dtype - self.device = device - self.tp = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - - -class GEMMReduceScatterPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - mul = torch.empty([16, 4], device=self.device, dtype=self.dtype) - mm_weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - return [mul, mm_weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - if self.rocm_gemm: - mm = torch.ops.vllm.rocm_unquantized_gemm(mul, mm_weight, None) - else: - mm = torch.ops.aten.mm.default(mul, mm_weight) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement(mul: torch.Tensor, mm_weight: torch.Tensor) -> torch.Tensor: - gemm_rs = torch.ops.symm_mem.fused_matmul_reduce_scatter( - mul, - mm_weight.t() if self.rocm_gemm else mm_weight, - "sum", - scatter_dim=0, - group_name=self.tp.device_group.group_name, - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmGEMMReduceScatterPattern(GEMMReduceScatterPattern): - rocm_gemm = True - - -class AllGatherGEMMPattern(BasePattern): - rocm_gemm = False - - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([4, 4], device=self.device, dtype=self.dtype) - weight = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [x, weight] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - if self.rocm_gemm: - return torch.ops.vllm.rocm_unquantized_gemm(all_gather, weight, None) - return torch.ops.aten.mm.default(all_gather, weight) - - def replacement(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_matmul( - x, - [weight.t() if self.rocm_gemm else weight], - gather_dim=0, - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class RocmAllGatherGEMMPattern(AllGatherGEMMPattern): - rocm_gemm = True - - -class ScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - return [input, mm_weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - scaled_mm = torch.ops.aten._scaled_mm.default( - input, - mat2=mat2, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - scaled_mm, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - return [x, weight, scale_a, scale_b] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - return torch.ops.aten._scaled_mm.default( - all_gather, - mat2=weight, - scale_a=scale_a, - scale_b=scale_b, - bias=None, - scale_result=None, - out_dtype=self.dtype, - ) - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class CutlassScaledMMReduceScatterPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - input = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - mm_weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - scale_a = torch.empty([16, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - cutlass_mm_output = torch.empty([16, 16], device=self.device, dtype=self.dtype) - return [input, mm_weight, scale_a, scale_b, cutlass_mm_output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=cutlass_mm_output, - a=input, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - - reduce_scatter = torch.ops.vllm.reduce_scatter.default( - cutlass_scaled_mm[1], - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return reduce_scatter - - def replacement( - input: torch.Tensor, - mat2: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - cutlass_mm_output: torch.Tensor, - ) -> torch.Tensor: - # Calculate output shape: input @ mat2 with scatter_dim reduced - output_shape = [*input.shape[:-1], mat2.shape[1]] - scatter_dim = 0 - gemm_rs = torch.ops.vllm.patched_fused_scaled_matmul_reduce_scatter( - input, - mat2, - scale_a, - scale_b, - "sum", - scatter_dim, # orig_scatter_dim - scatter_dim, # scatter_dim_after_maybe_reshape - self.tp.device_group.group_name, - output_shape, - None, # bias - None, # result_scale - self.dtype, # out_dtype - False, # use_fast_accum - ) - - return gemm_rs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class AllGatherCutlassScaledMMPattern(BasePattern): - def get_inputs(self) -> list[torch.Tensor]: - x = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - weight = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - - s1 = x.shape[0] * self.tp_size - - scale_a = torch.empty([s1, 1], device=self.device, dtype=torch.float32) - scale_b = torch.empty([1, 16], device=self.device, dtype=torch.float32) - - s2 = weight.shape[1] - output = torch.empty([s1, s2], device=self.device, dtype=self.dtype) - - return [x, weight, scale_a, scale_b, output] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp.unique_name - ) - - cutlass_scaled_mm = torch.ops.higher_order.auto_functionalized( - torch.ops._C.cutlass_scaled_mm.default, - out=output, - a=all_gather, - b=weight, - a_scales=scale_a, - b_scales=scale_b, - bias=None, - ) - return cutlass_scaled_mm[1] - - def replacement( - x: torch.Tensor, - weight: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - output: torch.Tensor, - ) -> torch.Tensor: - ag_output, mm_outputs = torch.ops.symm_mem.fused_all_gather_scaled_matmul( # noqa - x, - [weight], - scale_a, - [scale_b], - gather_dim=0, - biases=[None], - result_scales=[None], - out_dtypes=[self.dtype], - use_fast_accum=[False], - group_name=self.tp.device_group.group_name, - ) - return mm_outputs - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class FlashInferBMMFP8ReduceScatterPattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_2d = torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - bmm = torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(a_2d, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - output = torch.ops.aten.reshape.default(bmm, list(bmm.shape[1:])) - return torch.ops.vllm.reduce_scatter.default( - output, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_flashinfer_scaled_matmul_reduce_scatter.default( - a_2d, - b_2d, - a_scale, - b_scale, - "sum", - 0, - 0, - self.tp.device_group.group_name, - [a_2d.shape[0], b_2d.shape[1]], - self.dtype, - ) - - return _replacement - - -class FlashInferAllGatherBMMFP8Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 16], device=self.device, dtype=FP8_DTYPE) - b_2d = ( - torch.empty([16, 16], device=self.device, dtype=FP8_DTYPE) - .contiguous() - .transpose(0, 1) - ) - a_scale = torch.empty([1], device=self.device, dtype=torch.float32) - b_scale = torch.empty([1], device=self.device, dtype=torch.float32) - return [a_shard_2d, b_2d, a_scale, b_scale] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - all_gather = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - return torch.ops.vllm.bmm_fp8.default( - torch.ops.aten.unsqueeze.default(all_gather, 0), - torch.ops.aten.unsqueeze.default(b_2d, 0), - a_scale, - b_scale, - self.dtype, - "auto", - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale: torch.Tensor, - b_scale: torch.Tensor, - ) -> torch.Tensor: - fused = torch.ops.vllm.fused_all_gather_flashinfer_scaled_matmul.default( - a_shard_2d, - b_2d, - a_scale, - b_scale, - 0, - self.tp.device_group.group_name, - self.dtype, - ) - return torch.ops.aten.unsqueeze.default(fused, 0) - - return _replacement - - -class FlashInferAllGatherFP4Pattern( - BasePattern, VllmPatternReplacement[..., torch.Tensor] -): - def __init__( - self, - dtype: torch.dtype, - device: str | None, - backend: str, - use_8x4_sf_layout: bool, - a_scale_view: str, - ) -> None: - super().__init__(dtype, device) - self.backend = backend - self.use_8x4_sf_layout = use_8x4_sf_layout - self.a_scale_view = a_scale_view - - def get_inputs(self) -> list[torch.Tensor]: - a_shard_2d = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - b_2d = torch.empty([8, 16], device=self.device, dtype=torch.uint8) - a_scale_shard = torch.empty([128, 4], device=self.device, dtype=torch.int32) - b_scale = torch.empty([4, 128], device=self.device, dtype=torch.uint8) - alpha = torch.empty([], device=self.device, dtype=torch.float32) - return [ - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - ] - - @property - def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - all_gather_a = torch.ops.vllm.all_gather.default( - a_shard_2d, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - all_gather_a_scale = torch.ops.vllm.all_gather.default( - a_scale_shard, - dim=0, - world_size=self.tp_size, - group_name=self.tp.unique_name, - ) - a_scale = all_gather_a_scale - if self.a_scale_view in ("float8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.float8_e4m3fn) - if self.a_scale_view in ("uint8", "float8_uint8"): - a_scale = torch.ops.aten.view.dtype(a_scale, torch.uint8) - return torch.ops.vllm.flashinfer_mm_fp4.default( - all_gather_a, - b_2d, - a_scale, - b_scale, - alpha, - self.dtype, - self.use_8x4_sf_layout, - self.backend, - ) - - return _pattern - - @property - def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - a_shard_2d: torch.Tensor, - b_2d: torch.Tensor, - a_scale_shard: torch.Tensor, - b_scale: torch.Tensor, - alpha: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.vllm.fused_all_gather_flashinfer_fp4_matmul.default( - a_shard_2d, - b_2d, - a_scale_shard, - b_scale, - alpha, - 0, - self.tp.device_group.group_name, - self.dtype, - self.a_scale_view in ("float8", "float8_uint8"), - self.use_8x4_sf_layout, - self.backend, - ) - - return _replacement - - -class AsyncTPPass(VllmFusionPatternMatcherPass): - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config, pass_name="async_tp_pass") - - enable_symm_mem_for_group(get_tp_group().device_group.group_name) - GEMMReduceScatterPattern(self.model_dtype, self.device).register(self.pm_pass) - - AllGatherGEMMPattern(self.model_dtype, self.device).register(self.pm_pass) - if current_platform.is_rocm() and self.model_dtype == torch.bfloat16: - RocmGEMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - RocmAllGatherGEMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - - # These fusions are enabled only for bfloat16 models because - # `scaled_mm` or `cutlass_scaled_mm` with per-token (row-wise) scaling - # only supports bfloat16 as the output dtype. - if self.model_dtype == torch.bfloat16: - ScaledMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - AllGatherScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - if hasattr(torch.ops._C, "cutlass_scaled_mm"): - CutlassScaledMMReduceScatterPattern( - self.model_dtype, self.device - ).register(self.pm_pass) - AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - with suppress(ImportError): - import vllm.utils.flashinfer # noqa: F401 - if hasattr(torch.ops.vllm, "bmm_fp8"): - self.register( - FlashInferAllGatherBMMFP8Pattern(self.model_dtype, self.device) - ) - self.register( - FlashInferBMMFP8ReduceScatterPattern(self.model_dtype, self.device) - ) - if hasattr(torch.ops.vllm, "flashinfer_mm_fp4"): - for backend in ("cutlass", "cudnn"): - for a_scale_view in ("float8_uint8", "uint8"): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - backend, - use_8x4_sf_layout=False, - a_scale_view=a_scale_view, - ) - ) - for use_8x4_sf_layout in (False, True): - for a_scale_view in ("float8",): - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "trtllm", - use_8x4_sf_layout=use_8x4_sf_layout, - a_scale_view=a_scale_view, - ) - ) - self.register( - FlashInferAllGatherFP4Pattern( - self.model_dtype, - self.device, - "cute-dsl", - use_8x4_sf_layout=False, - a_scale_view="float8", - ) - ) - # NVFP4 reduce-scatter does not need scale communication: FP4 - # scales are consumed by the local GEMM and only BF16 partial - # outputs are reduced. Keep this PR scoped to the all-gather - # path; reduce-scatter needs a dedicated FP4 producer rather - # than the existing FP8-style helper. - - self.dump_patterns(config, self.pm_pass) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - # This pass is applied on top of the sequence parallelism pass, - # which is only supported in fullgraph compilation mode. - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "AsyncTPPass requires full-graph compilation" - return True - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.pm_pass.apply(graph) - VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count - logger.debug("Replaced %s patterns", self.matched_count) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py deleted file mode 100644 index b1a39b1e0146..000000000000 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ /dev/null @@ -1,620 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import functools -from collections.abc import Callable, Sequence -from typing import Any - -import torch -import torch._inductor.pattern_matcher as pm -import torch.fx as fx -from torch._higher_order_ops.auto_functionalize import auto_functionalized -from torch._inductor.pattern_matcher import PatternMatcherPass - -import vllm.ir.ops -from vllm.config import VllmConfig -from vllm.config.utils import Range -from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8StaticTensorSym, -) - -from ..inductor_pass import enable_fake_mode -from ..utility.noop_elimination import NoOpEliminationPass -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass -from .matcher_utils import MatcherQuantFP8 - -logger = init_logger(__name__) - -if hasattr(torch.ops._C, "scaled_fp4_quant"): - SCALED_FP4_QUANT_OUT_OVERLOAD = torch.ops._C.scaled_fp4_quant.out - SCALED_FP4_QUANT_DEFAULT_OVERLOAD = torch.ops._C.scaled_fp4_quant.default - -# Min hidden size per device capability for sequence parallelism -# Only apply sequence parallelism for models with hidden_size >= threshold -SP_MIN_HIDDEN_SIZE: dict[int, int] = { - 90: 8192, # H100: only for models with hidden_size >= 8192 - 100: 8192, # Blackwell family: only for models with hidden_size >= 8192 -} - -# Min size per GPU per device capability for sequence parallelism -# Total min size = min_per_gpu_size * tp_size -# This ensures the threshold scales appropriately with tensor parallelism -SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = { - 90: 8, # 8MB per GPU for H100 - # Use a more conservative threshold on Blackwell so TP8 starts later. - 100: 32, -} - - -def get_sequence_parallelism_threshold( - hidden_size: int, - tp_size: int, - element_size: int, -) -> int | None: - """Calculate the minimum token threshold for applying sequence parallelism. - - Returns None if sequence parallelism should not be applied based on model size. - - Branching logic based on device capability: - - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability] - - If not, returns None (SP disabled for small models on this device) - - If yes, calculates threshold based on per-GPU size - - Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) // - (hidden_size * element_size) - """ - from vllm.platforms import current_platform - - if current_platform.is_xpu(): - min_hidden_size = 4096 - min_per_gpu_size_mb = 8.0 - elif current_platform.is_cuda(): - capability = current_platform.get_device_capability() - if capability is None: - return None - - # Collapse Blackwell variants (sm100/sm103/...) into one policy bucket. - if current_platform.is_device_capability_family(100): - device_capability = 100 - else: - device_capability = capability.to_int() - - # Check if device has configured thresholds - _hidden = SP_MIN_HIDDEN_SIZE.get(device_capability) - _gpu_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) - if _hidden is None or _gpu_mb is None: - return None - min_hidden_size, min_per_gpu_size_mb = _hidden, _gpu_mb - else: - return None - - # Only apply sequence parallelism for models meeting the size threshold - if hidden_size < min_hidden_size: - return None - - MiB = 1024 * 1024 - min_size = min_per_gpu_size_mb * MiB * tp_size - return int(min_size // (hidden_size * element_size)) - - -def get_first_out_wrapper( - fn: Callable[..., Sequence[torch.Tensor]], -) -> Callable[..., torch.Tensor]: - @functools.wraps(fn) - def wrapper(*args: Any) -> torch.Tensor: - return fn(*args)[0] - - return wrapper - - -class _SequenceParallelPatternHelper: - """Helper for sequence parallelism patterns.""" - - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - self.epsilon = epsilon - self.dtype = dtype - self.device = device - self.tp_group = get_tp_group() - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() - - def _all_reduce(self, x: torch.Tensor) -> torch.Tensor: - return tensor_model_parallel_all_reduce(x) - - def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.reduce_scatter.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def _all_gather(self, x: torch.Tensor) -> torch.Tensor: - return torch.ops.vllm.all_gather.default( - x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name - ) - - def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs) - - def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs) - - -class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight - return [self.empty([1, 8, 4]), self.empty([4])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rmsnorm = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - - return rmsnorm, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - - rmsnorm = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - all_gather = self._all_gather(rmsnorm) - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - - return [ - residual, - mm_1, - rms_norm_weights, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - return rmsnorm[0], rmsnorm[1] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # The pattern matcher replaces from the end of the graph - # (last layer first). At the time each match is replaced, - # the preceding layer has NOT been replaced yet, so - # `residual` is still full-size and the slice below is - # correct. Once the preceding layer IS replaced, its - # residual output shrinks to [local_len, H], and this - # slice becomes semantically incorrect (e.g. for rank > 0, - # the indices would be out of bounds). However, since the - # symbolic output shape equals the input shape, - # NoOpEliminationPass (called at the end of - # SequenceParallelismPass.__call__) removes these slices - # before the graph is ever executed or compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rmsnorm = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - all_gather = self._all_gather(rmsnorm[0]) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, rmsnorm[1] - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__( - self, - epsilon: float, - dtype: torch.dtype, - device: str | None, - ) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - # input, weight, scale - return [self.empty([1, 8, 4]), self.empty([4]), self.empty_f32([1, 1])] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - return quant, all_reduce - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - - return all_gather, reduce_scatter - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): - def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: - super().__init__(epsilon, dtype, device) - self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) - - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) - residual = torch.empty([4, 4], device=self.device, dtype=self.dtype) - rms_norm_weights = torch.empty([4, 4], device=self.device, dtype=self.dtype) - scale = torch.empty([1, 1], device=self.device, dtype=torch.float32) - - return [residual, mm_1, rms_norm_weights, scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - return quant, residual_out - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # See MiddleAllReduceRMSNormPattern.replacement for a - # detailed explanation of the temporary slice below: - # it is correct when first inserted, becomes semantically - # incorrect after the preceding layer is replaced, and is - # removed by NoOpEliminationPass before the graph is compiled. - reduce_scatter = self._reduce_scatter(mm_1) - local_len = reduce_scatter.size(0) - # when the preceding VocabParallelEmbedding is excluded - # from the FX graph (e.g., passing `inputs_embeds` directly in VLMs), - # the FirstAllReduceRMSNorm pattern is never matched. we must - # perform a proper TP-aware slice here. simply using `[0:local_len]` - # would incorrectly cause all ranks to process rank 0's chunk. - residual = residual[ - self.tp_rank * local_len : self.tp_rank * local_len + local_len, ... - ] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - quant, _ = self.quant_matcher(rms, scale) - all_gather = self._all_gather(quant) - # residual output is now [local_len, H]; the next layer's - # slice on it is semantically incorrect until - # NoOpEliminationPass removes it. - return all_gather, residual_out - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - pm.register_replacement( - get_first_out_wrapper(pattern), - get_first_out_wrapper(replacement), - self.get_inputs(), - pm.fwd_only, - pm_pass, - ) - - -class FirstAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - input = self.empty([8, 16]) - weight = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [input, weight, input_global_scale, quant_output, output_scale] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(input) - rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], all_reduce, quant[2] - - def replacement( - input: torch.Tensor, - weight: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - reduce_scatter = self._reduce_scatter(input) - rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return ( - self._all_gather(quant[0]), - reduce_scatter, - self._all_gather(quant[1]), - ) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class MiddleAllReduceRMSNormStaticNVFP4Pattern(_SequenceParallelPatternHelper): - def get_inputs(self) -> list[torch.Tensor]: - mm_1 = self.empty([8, 16]) - residual = self.empty([8, 16]) - rms_norm_weights = self.empty([16]) - input_global_scale = self.empty_f32([1, 1]) - quant_output = torch.empty([8, 8], device=self.device, dtype=torch.uint8) - output_scale = torch.empty([128, 4], device=self.device, dtype=torch.int32) - return [ - residual, - mm_1, - rms_norm_weights, - input_global_scale, - quant_output, - output_scale, - ] - - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - all_reduce = self._all_reduce(mm_1) - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - all_reduce, residual, rms_norm_weights, self.epsilon - ) - quant = auto_functionalized( - SCALED_FP4_QUANT_OUT_OVERLOAD, - input=rms, - input_scale=input_global_scale, - is_sf_swizzled_layout=True, - output=quant_output, - output_scale=output_scale, - ) - return quant[1], residual_out, quant[2] - - def replacement( - residual: torch.Tensor, - mm_1: torch.Tensor, - rms_norm_weights: torch.Tensor, - input_global_scale: torch.Tensor, - quant_output: torch.Tensor, - output_scale: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # Keep this slice in sync with the non-quantized SP replacement: - # once the previous SP pattern fires, it becomes a no-op. - reduce_scatter = self._reduce_scatter(mm_1) - residual = residual[0 : reduce_scatter.size(0), ...] - rms, residual_out = vllm.ir.ops.fused_add_rms_norm( - reduce_scatter, residual, rms_norm_weights, self.epsilon - ) - rms = torch.ops.aten.view.default(rms, [-1, rms.shape[-1]]) - quant = SCALED_FP4_QUANT_DEFAULT_OVERLOAD( - rms, - input_global_scale, - True, - ) - return self._all_gather(quant[0]), residual_out, self._all_gather(quant[1]) - - pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass - ) - - -class SequenceParallelismPass(VllmPatternMatcherPass): - """This pass enables sequence parallelism for models. - It identifies patterns where an AllReduce operation is followed by - an RMSNorm (or RMSNorm and then Quantization) operation. - These patterns are replaced with a ReduceScatter operation, followed by - a local RMSNorm/Quantization, and then an AllGather operation. - - The general transformation is: - Input -> AllReduce -> RMSNorm -> Output - becomes - Input -> ReduceScatter -> RMSNorm -> AllGather -> Output - - While this pass itself does not directly yield performance improvements, - it lays the groundwork for subsequent fusion passes, such as - GEMM + ReduceScatter and AllGather + GEMM fusions. These fusions can - significantly reduce communication overhead and improve overall model - performance. - - This pass is only supported when compiling the whole graph (fullgraph - mode, i.e. using Inductor graph partition or empty splitting_ops). - Piecewise compilation is not supported because the residual tensor - gets split across TP ranks, causing size mismatches at subgraph - boundaries. - - This pass splits up the residual tensor across TP ranks and hence divides - its size. The pattern matcher starts at the end of the graph (last layer - first), so when each replacement inserts a residual slice, the preceding - layer has not been replaced yet and the slice is correct. Once the - preceding layer IS replaced, its residual output shrinks and the slice - becomes semantically incorrect (out-of-bounds indices for rank > 0). - The graph is never executed in this intermediate state — - NoOpEliminationPass removes these slices based on symbolic shape equality - (input shape == output shape) before the graph is compiled. - """ - - @enable_fake_mode - def __init__(self, config: VllmConfig) -> None: - super().__init__(config) - - # Get min_token_num threshold - # Read min_token_num from config (calculated during config init) - self.min_token_num = None - if config.model_config is not None: - pass_config = config.compilation_config.pass_config - self.min_token_num = pass_config.sp_min_token_num - - if self.min_token_num is not None: - # Take the min to avoid exceeding max_num_batched_tokens - max_batched = config.scheduler_config.max_num_batched_tokens - if max_batched is not None: - self.min_token_num = min(self.min_token_num, max_batched) - logger.debug_once( - f"Sequence parallelism min token threshold: {self.min_token_num}", - scope="global", - ) - - # Used to clean up redundant views created temporarily - # to circumvent residual shape change issues - self.noop_cleanup = NoOpEliminationPass(config) - self.noop_cleanup.pass_name = f"{self.pass_name}.{self.noop_cleanup.pass_name}" - - self.patterns: PatternMatcherPass = PatternMatcherPass( - pass_name="sequence_parallelism_pass" - ) - - for epsilon in [1e-5, 1e-6]: - # RMSNorm + Static FP8 quantization patterns - FirstAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticFP8Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - if "SCALED_FP4_QUANT_OUT_OVERLOAD" in globals(): - FirstAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - MiddleAllReduceRMSNormStaticNVFP4Pattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - # Normal RMSNorm patterns - FirstAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - MiddleAllReduceRMSNormPattern( - epsilon, self.model_dtype, self.device - ).register(self.patterns) - - self.dump_patterns(config, self.patterns) - - def is_applicable_for_range(self, compile_range: Range) -> bool: - """Determines if sequence parallelism should be applied for the given - compile range. - - SP is only beneficial for larger batch sizes where the communication - overhead is amortized. For small batches, the overhead of splitting - and gathering tensors across TP ranks outweighs the benefits. - - Returns False (SP disabled) when: - - min_token_num is None (SP disabled for this device/config) - - The compile range starts below the minimum token threshold - """ - assert ( - self.compilation_config.use_inductor_graph_partition - or not self.compilation_config.splitting_ops - ), "SequenceParallelismPass requires full-graph compilation" - - # min_token_num is None when SP is disabled for this device/config - # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size) - if self.min_token_num is None: - return False - - # Only apply SP when batch size meets the minimum threshold - return compile_range.start >= self.min_token_num - - @VllmInductorPass.time_and_log - def __call__(self, graph: fx.Graph) -> None: - self.matched_count = self.patterns.apply(graph) - logger.debug("Replaced %s patterns", self.matched_count) - # Clean up reshape nodes - self.noop_cleanup(graph) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index df0e58b21081..197f2208a15c 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -35,13 +35,11 @@ RMSNormReshapeFusionPass, ) from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass - from .fusion.sequence_parallelism import SequenceParallelismPass from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass - from .fusion.collective_fusion import AsyncTPPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass @@ -156,11 +154,6 @@ def configure(self, config: VllmConfig) -> None: if self.pass_config.eliminate_noops: self.passes += [NoOpEliminationPass(config)] - if self.pass_config.enable_sp: - self.passes += [SequenceParallelismPass(config)] - if self.pass_config.fuse_gemm_comms: - self.passes += [AsyncTPPass(config)] - if enable_transformers_norm_canonicalization: self.passes += [AddRMSNormFusionPass(config)] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index cbd132ab5b68..fc371d94b493 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -133,12 +133,6 @@ class PassConfig: """Fuse the custom Attention and MLAAttention + quant ops.""" eliminate_noops: bool = Field(default=True) """Eliminate no-op ops.""" - enable_sp: bool = None # type: ignore[assignment] - """Enable sequence parallelism. Requires TP>1. Automatically disabled - if the model's hidden_size is too small for SP to be beneficial - (threshold is device-capability dependent).""" - fuse_gemm_comms: bool = None # type: ignore[assignment] - """Enable async TP.""" fuse_allreduce_rms: bool = None # type: ignore[assignment] """Enable flashinfer allreduce fusion.""" enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment] @@ -183,12 +177,6 @@ class PassConfig: 8: 1, # 1MB }, }, where key is the device capability""" - sp_min_token_num: int | None = None - """The minimum number of tokens above which vllm should use - sequence parallelism. Specified as an integer token count. - Unspecified will fallback to default values which are compute - capability and world size dependent.""" - # TODO(luka) better pass enabling system. def flashinfer_max_size(self, world_size: int) -> int | None: @@ -231,8 +219,6 @@ def compute_hash(self) -> str: "fuse_norm_quant", "fuse_act_quant", "fuse_attn_quant", - "enable_sp", - "fuse_gemm_comms", "fuse_allreduce_rms", "fuse_act_padding", "fuse_mla_dual_rms_norm", @@ -1244,25 +1230,6 @@ def set_splitting_ops_for_v1( self.cudagraph_mode = CUDAGraphMode.FULL self.splitting_ops = [] - if ( - not self.use_inductor_graph_partition - and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms) - and self.splitting_ops - ): - logger.warning_once( - "Sequence parallelism requires full-graph compilation when " - "use_inductor_graph_partition is off. Setting splitting_ops " - "to an empty list to preserve SP and async TP." - ) - self.splitting_ops = [] - if self.cudagraph_mode.has_piecewise_cudagraphs(): - logger.warning_once( - "Sequence parallelism is incompatible with piecewise " - "cudagraph when use_inductor_graph_partition is off. " - "Setting cudagraph_mode to FULL." - ) - self.cudagraph_mode = CUDAGraphMode.FULL - # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible if ( all2all_backend == "deepep_high_throughput" @@ -1403,7 +1370,6 @@ def resolve_cudagraph_mode_and_sizes( min_cg_attn_backend: str | None, uniform_decode_query_len: int = 1, use_v2_model_runner: bool = False, - tensor_parallel_size: int = 1, kv_cache_config: "KVCacheConfig | None" = None, max_num_reqs: int | None = None, is_profiling: bool = False, @@ -1530,10 +1496,7 @@ def resolve_cudagraph_mode_and_sizes( and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and uniform_decode_query_len > 1 ): - self.adjust_cudagraph_sizes_for_spec_decode( - uniform_decode_query_len, - tensor_parallel_size, - ) + self.adjust_cudagraph_sizes_for_spec_decode(uniform_decode_query_len) # For Mamba models with FULL decode cudagraphs, each decode # sequence needs one Mamba cache block. The decode cudagraph @@ -1562,25 +1525,8 @@ def resolve_cudagraph_mode_and_sizes( self.cudagraph_mode = cudagraph_mode return cudagraph_mode - def adjust_cudagraph_sizes_for_spec_decode( - self, uniform_decode_query_len: int, tensor_parallel_size: int - ): + def adjust_cudagraph_sizes_for_spec_decode(self, uniform_decode_query_len: int): multiple_of = uniform_decode_query_len - if tensor_parallel_size > 1 and self.pass_config.enable_sp: - multiple_of = max(uniform_decode_query_len, tensor_parallel_size) - if ( - multiple_of % uniform_decode_query_len != 0 - or multiple_of % tensor_parallel_size != 0 - ): - raise ValueError( - f"Can't determine cudagraph shapes that are both a " - f"multiple of {uniform_decode_query_len} " - f"(num_speculative_tokens + 1) required by spec-decode " - f"and {tensor_parallel_size} (tensor_parallel_size) " - f"required by sequence parallelism please adjust " - f"num_speculative_tokens or disable sequence parallelism" - ) - if not self.cudagraph_capture_sizes or multiple_of <= 1: return @@ -1600,7 +1546,7 @@ def adjust_cudagraph_sizes_for_spec_decode( if len(rounded_sizes) == 0: raise ValueError( f"No valid cudagraph sizes after rounding to multiple of {multiple_of} " - f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)" + f"(num_speculative_tokens + 1)" f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}" f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})" f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8fc89f1591cd..3fee812700b0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -145,12 +145,10 @@ class OptimizationLevel(IntEnum): PerformanceMode = Literal["balanced", "interactivity", "throughput"] IS_QUANTIZED = False -IS_DENSE = False # The optimizations that depend on these properties currently set to False # in all cases. # if model_config is not None: # IS_QUANTIZED = lambda c: c.model_config.is_quantized() -# IS_DENSE = lambda c: not c.model_config.is_model_moe() # See https://github.com/vllm-project/vllm/issues/25689. @@ -259,8 +257,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": False, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": False, "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, @@ -282,8 +278,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": False, "fuse_attn_quant": False, - "enable_sp": False, - "fuse_gemm_comms": False, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, @@ -305,8 +299,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -328,8 +320,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: "fuse_act_quant": enable_act_fusion, "fuse_allreduce_rms": enable_allreduce_rms_fusion, "fuse_attn_quant": IS_QUANTIZED, - "enable_sp": IS_DENSE, - "fuse_gemm_comms": IS_DENSE, "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, @@ -828,9 +818,6 @@ def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool: or model.dtype != torch.bfloat16 or model.quantization is not None or not current_platform.is_cuda() - # Sequence parallelism / async TP are torch.compile passes. - or self.compilation_config.pass_config.enable_sp - or self.compilation_config.pass_config.fuse_gemm_comms ): return False family = bi._get_tuned_matmul_arch_family( @@ -1770,26 +1757,6 @@ def __post_init__(self): ) self.compilation_config.mode = CompilationMode.NONE - # TODO: This is a stopgap that turns sequence parallelism / async TP - # off under batch invariance. Make the sequence-parallel reduce-scatter - # path batch-invariant instead so they can stay enabled (#56370). - # Sequence parallelism / async TP rewrite all_reduce + rms_norm into a - # reduce-scatter whose reduction order depends on the batch, so they - # are not batch-invariant. Decide this before the breakable-CUDA-graph - # default below, which declines to auto-enable when they are on. - pass_config = self.compilation_config.pass_config - if envs.VLLM_BATCH_INVARIANT and ( - pass_config.enable_sp or pass_config.fuse_gemm_comms - ): - logger.warning_once( - "Disabling sequence parallelism and async TP " - "(pass_config.enable_sp / fuse_gemm_comms) when " - "VLLM_BATCH_INVARIANT is enabled: the reduce-scatter path " - "is not batch-invariant (see vllm-project/vllm#56370)." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph() if not breakable_cudagraph_enabled and ( @@ -1920,38 +1887,6 @@ def has_blocked_weights(): breakable_cudagraph_enabled=breakable_cudagraph_enabled ) - # async tp is built on top of sequence parallelism and requires it. - pass_config = self.compilation_config.pass_config - if pass_config.fuse_gemm_comms: - pass_config.enable_sp = True - if pass_config.enable_sp: - if self.parallel_config.tensor_parallel_size == 1: - logger.warning_once("Sequence Parallelism requires TP>1, disabling") - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - else: - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - if pass_config.sp_min_token_num is None: - logger.warning_once( - "Model hidden_size too small for the SP " - "threshold heuristic, disabling. To force SP, " - "set pass_config.sp_min_token_num manually." - ) - pass_config.enable_sp = False - pass_config.fuse_gemm_comms = False - from vllm.utils.torch_utils import HAS_OPAQUE_TYPE if HAS_OPAQUE_TYPE: @@ -2184,30 +2119,6 @@ def has_blocked_weights(): data_parallel_size=effective_dp_size, ) - if self.compilation_config.pass_config.enable_sp: - # With pipeline parallelism, native rms norm tracing errors due to - # incorrect residual shape. - # Use custom rms norm to unblock. In the future, - # the pass will operate on higher-level IR to avoid the issue. - # TODO: https://github.com/vllm-project/vllm/issues/27894 - if self.compilation_config.mode != CompilationMode.VLLM_COMPILE: - logger.warning_once( - "Sequence parallelism is enabled, but running in wrong " - "vllm compile mode: %s.", - self.compilation_config.mode, - ) - - if self.parallel_config.pipeline_parallel_size > 1: - if "-rms_norm" not in self.compilation_config.custom_ops: - self.compilation_config.custom_ops.append("+rms_norm") - else: - logger.warning_once( - "Sequence parallelism not supported with " - "native rms_norm when using %s, " - "this will likely lead to an error.", - "pipeline parallelism", - ) - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( @@ -2395,29 +2306,6 @@ def has_blocked_weights(): # Log the custom passes that are enabled self.compilation_config.pass_config.log_enabled_passes() - def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list: - # remove the sizes that not multiple of tp_size when - # enable sequence parallelism - removed_sizes = [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size != 0 - ] - if removed_sizes: - logger.warning( - "Batch sizes %s are removed because they are not " - "multiple of tp_size %d when " - "sequence parallelism is enabled", - removed_sizes, - self.parallel_config.tensor_parallel_size, - ) - - return [ - size - for size in possible_sizes - if size % self.parallel_config.tensor_parallel_size == 0 - ] - def _set_max_num_scheduled_tokens(self): """In most cases, the scheduler may schedule a batch with as many tokens as the worker is configured to handle. @@ -2665,16 +2553,6 @@ def request_counts(max_reqs: int) -> list[int]: # de-duplicate and sort the sizes cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes)) - if ( - self.parallel_config.tensor_parallel_size > 1 - and self.compilation_config.pass_config.enable_sp - ): - # Sequence parallelism only captures TP-divisible sizes, so a - # wider non-divisible decode batch cannot be captured under SP. - cudagraph_capture_sizes = self.update_sizes_for_sequence_parallelism( - cudagraph_capture_sizes - ) - # user-specific compilation_config.max_cudagraph_capture_size get # truncated to valid_max_size when they are inconsistent. valid_max_size = ( @@ -2764,37 +2642,6 @@ def _set_compile_ranges(self): "allreduce-rms fusion will be enabled for all num_tokens." ) - # Add the compile ranges for sequence parallelism - if compilation_config.pass_config.enable_sp: - pass_config = compilation_config.pass_config - - # Calculate min_token_num if not explicitly provided - # User override works regardless of hidden_size - if pass_config.sp_min_token_num is None: - from vllm.compilation.passes.fusion.sequence_parallelism import ( - get_sequence_parallelism_threshold, - ) - - tp_size = self.parallel_config.tensor_parallel_size - hidden_size = self.model_config.get_hidden_size() - assert isinstance(self.model_config.dtype, torch.dtype) - element_size = self.model_config.dtype.itemsize - pass_config.sp_min_token_num = get_sequence_parallelism_threshold( - hidden_size, tp_size, element_size - ) - - min_token_num = pass_config.sp_min_token_num - max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens - if min_token_num is not None and ( - max_num_batched_tokens is not None - and min_token_num < max_num_batched_tokens - and min_token_num > 1 - ): - # Add endpoint at min_token_num - 1 to ensure SP applies - # starting from min_token_num - # This creates ranges: [1, min-1] (no SP), [min, max] (SP applies) - computed_compile_ranges_endpoints.append(min_token_num - 1) - if compilation_config.pass_config.fuse_rope_kvcache: max_token_num = ( compilation_config.pass_config.rope_kvcache_fusion_max_token_num @@ -3124,12 +2971,6 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported: list[str] = [] speculative_config = self.speculative_config - if ( - self.compilation_config.pass_config.enable_sp - and self.parallel_config.tensor_parallel_size > 1 - ): - unsupported.append("sequence parallelism") - # V2 does not implement the external_launcher (torchrun) PP-output # broadcast that V1 uses to keep all ranks in sync (broadcast_pp_output). if ( diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 6b7db6d0dc55..355a15094db0 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -352,7 +352,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: pass_config = compilation_config.pass_config fusion_passes_to_disable = { - "fuse_gemm_comms": "Async TP", "fuse_allreduce_rms": "AllReduce + RMSNorm fusion", "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index f9d6c4439f56..9b615cd2fa4b 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -1065,47 +1065,6 @@ def flashinfer_scaled_fp4_mm( ) -def flashinfer_scaled_fp4_mm_out( - a: torch.Tensor, - b: torch.Tensor, - block_scale_a: torch.Tensor, - block_scale_b: torch.Tensor, - alpha: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None, - use_8x4_sf_layout: bool, - backend: str, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2 - assert a.stride(-1) == 1 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert out.device.type == "cuda" - - if backend in ("cutlass", "cudnn"): - if block_scale_a.dtype != torch.uint8: - block_scale_a = block_scale_a.view(torch.uint8) - if block_scale_b.dtype != torch.uint8: - block_scale_b = block_scale_b.view(torch.uint8) - - from flashinfer import mm_fp4 as flashinfer_mm_fp4_ - - flashinfer_mm_fp4_( - a, - b, - block_scale_a, - block_scale_b, - alpha, - out_dtype or out.dtype, - out=out, - block_size=16, - use_8x4_sf_layout=use_8x4_sf_layout, - backend=backend, - ) - return out - - def flashinfer_scaled_fp8_mm( a: torch.Tensor, b: torch.Tensor, @@ -1136,38 +1095,6 @@ def flashinfer_scaled_fp8_mm( return output -def flashinfer_scaled_fp8_mm_out( - a: torch.Tensor, - b: torch.Tensor, - scale_a: torch.Tensor, - scale_b: torch.Tensor, - out: torch.Tensor, - out_dtype: torch.dtype | None = None, -) -> torch.Tensor: - assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 - assert a.shape[1] == b.shape[0] - assert out.shape == (a.shape[0], b.shape[1]) - assert scale_a.numel() == 1 and scale_b.numel() == 1 - assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn - assert out.device.type == "cuda" - assert a.is_contiguous() - - from flashinfer import bmm_fp8 as bmm_fp8_ - - bmm_fp8_( - a.unsqueeze(0), - # FlashInfer expects the weight in the same column-major view layout - # consumed by flashinfer_scaled_fp8_mm, so keep the transposed view. - b.unsqueeze(0), - scale_a, - scale_b, - out_dtype or out.dtype, - out.unsqueeze(0), - "auto", - ) - return out - - def flashinfer_quant_nvfp4_8x4_sf_layout( a: torch.Tensor, a_global_sf: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1310,9 +1237,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "use_trtllm_attention", "flashinfer_mxfp4_quantize", "flashinfer_scaled_fp4_mm", - "flashinfer_scaled_fp4_mm_out", "flashinfer_scaled_fp8_mm", - "flashinfer_scaled_fp8_mm_out", "flashinfer_quant_nvfp4_8x4_sf_layout", "flashinfer_fp8_blockscale_gemm", "should_use_flashinfer_for_blockscale_fp8_gemm", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 373042a41502..61d6197ceb9d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -730,7 +730,6 @@ def initialize_kv_cache( attn_cg_support.min_cg_attn_backend, self.decode_query_len, use_v2_model_runner=True, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0cc3fcc29a8f..21b0f67e2fe8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -116,7 +116,7 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks from vllm.utils.platform_utils import num_compute_units @@ -222,7 +222,6 @@ ) from vllm.v1.worker.utils import ( EncoderTimingStats, - is_residual_scattered_for_sp, raise_if_nan_logits, ) from vllm.v1.worker.workspace import lock_workspace @@ -3340,20 +3339,9 @@ def sync_and_gather_intermediate_tensors( ) -> IntermediateTensors: assert self.intermediate_tensors is not None - tp = self.vllm_config.parallel_config.tensor_parallel_size - is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens) - - # When sequence parallelism is enabled, the "residual" tensor is - # sharded across TP ranks. All-gather it here because downstream - # QKV + Attention needs the full residual before the SP split point. if sync_self: assert intermediate_tensors is not None for k, v in intermediate_tensors.items(): - is_scattered = k == "residual" and is_rs - if is_scattered: - local_len = num_tokens // tp - v = get_tp_group().all_gather(v[:local_len], dim=0) - self.intermediate_tensors[k][:num_tokens].copy_( v[:num_tokens], non_blocking=True ) @@ -3448,14 +3436,6 @@ def _pool( async_output_copy_stream=self._get_or_create_async_output_copy_stream(), ) - def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: - # Pad tokens to multiple of tensor_parallel_size when - # enabled collective fusion for SP - tp_size = self.vllm_config.parallel_config.tensor_parallel_size - if self.compilation_config.pass_config.enable_sp and tp_size > 1: - return round_up(num_scheduled_tokens, tp_size) - return num_scheduled_tokens - def _prepare_padding_mask( self, num_tokens_unpadded: int, num_tokens_padded: int ) -> torch.Tensor: @@ -3944,8 +3924,6 @@ def _determine_batch_execution_and_padding( ) has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora - num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens) - def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): return self.cudagraph_dispatcher.dispatch( num_tokens=num_tokens, @@ -3957,18 +3935,9 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) cudagraph_mode, batch_descriptor = dispatch_cudagraph( - num_tokens_padded, disable_full=use_cascade_attn or has_encoder_output + num_tokens, disable_full=use_cascade_attn or has_encoder_output ) num_tokens_padded = batch_descriptor.num_tokens - if self.compilation_config.pass_config.enable_sp: - assert ( - batch_descriptor.num_tokens - % self.vllm_config.parallel_config.tensor_parallel_size - == 0 - ), ( - "Sequence parallelism requires num_tokens to be " - "a multiple of tensor parallel size" - ) # Extra coordination when running data-parallel since we need to coordinate # across ranks @@ -4453,15 +4422,9 @@ def execute_model( sample_hidden_states = hidden_states[logits_indices] if not get_pp_group().is_last_rank: - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, num_tokens_padded - ) - } get_pp_group().send_tensor_dict( hidden_states.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) logits = None else: @@ -7135,7 +7098,6 @@ def _check_and_update_cudagraph_mode( min_cg_attn_backend, self.uniform_decode_query_len, use_v2_model_runner=False, - tensor_parallel_size=self.parallel_config.tensor_parallel_size, kv_cache_config=self.kv_cache_config, max_num_reqs=self.max_num_reqs, is_profiling=is_profiling, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d4c64c554273..70427d85a103 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -12,7 +12,6 @@ from types import NoneType from typing import TYPE_CHECKING, Any, cast -import numpy as np import torch import torch.nn as nn @@ -98,7 +97,6 @@ maybe_apply_startup_plan, maybe_save_startup_plan, ) -from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -1273,45 +1271,10 @@ def execute_model( intermediate_tensors = None forward_pass = scheduler_output.total_num_scheduled_tokens > 0 - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - all_gather_tensors = {} - compilation_config = self.vllm_config.compilation_config - parallel_config = self.vllm_config.parallel_config - - if ( - parallel_config.pipeline_parallel_size > 1 - and compilation_config.pass_config.enable_sp - and forward_pass - ): - # currently only supported by V1 GPUModelRunner - assert not self.use_v2_model_runner - num_scheduled_tokens_np = np.array( - list(scheduler_output.num_scheduled_tokens.values()), - dtype=np.int32, - ) - # TODO(lucas): This is pretty gross; ideally we should only ever call - # `_determine_batch_execution_and_padding` once (will get called again - # in `execute_model`) but this requires a larger refactor of PP. - _, batch_desc, _, _, _ = ( - self.model_runner._determine_batch_execution_and_padding( - num_tokens=num_scheduled_tokens, - num_reqs=len(num_scheduled_tokens_np), - num_scheduled_tokens_np=num_scheduled_tokens_np, - max_num_scheduled_tokens=num_scheduled_tokens_np.max(), - use_cascade_attn=False, # TODO(lucas): Handle cascade attention - ) - ) - all_gather_tensors = { - "residual": not is_residual_scattered_for_sp( - self.vllm_config, batch_desc.num_tokens - ) - } - if forward_pass and not get_pp_group().is_first_rank: tensor_dict, comm_handles, comm_postprocess = ( get_pp_group().irecv_tensor_dict( all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) ) assert tensor_dict is not None @@ -1349,7 +1312,6 @@ def execute_model( handles = get_pp_group().isend_tensor_dict( output.tensors, all_gather_group=get_tp_group(), - all_gather_tensors=all_gather_tensors, ) self._pp_send_work = handles[1:] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b61c4cebed18..db0f81b81c36 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -10,7 +10,7 @@ import numpy as np import torch -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -813,35 +813,6 @@ def get_uniform_decode_token_count( return None -def is_residual_scattered_for_sp( - vllm_config: VllmConfig, num_input_tokens: int -) -> bool: - """Check if the residual tensor is scattered for sequence parallelism. - - The residual tensor is scattered across tensor parallel ranks when sequence - parallelism and tensor parallelism is enabled. SP is only supported in - full-graph compilation mode. - """ - if not vllm_config.compilation_config.pass_config.enable_sp: - return False - - tp = vllm_config.parallel_config.tensor_parallel_size - - if tp == 1: - return False - - assert ( - vllm_config.compilation_config.use_inductor_graph_partition - or not vllm_config.compilation_config.splitting_ops - ), "Sequence parallelism requires full-graph compilation" - - # When sequence parallelism is enabled, we always pad num_input_tokens - # to be a multiple of tensor_parallel_size (tp) earlier. - assert num_input_tokens % tp == 0 - - return True - - @dataclass class EncoderTimingStats: """Per-request timing statistics for encoder forward pass.""" From 0751f5272fb323e0a5c62a6157f0a3c1d93e6bc5 Mon Sep 17 00:00:00 2001 From: Haokai Ma Date: Thu, 8 Oct 2026 08:29:04 -0700 Subject: [PATCH 216/952] [Bugfix] Scope deprecated CLI argument warnings to the parser that defines them (#59917) Signed-off-by: Tyler Ma <158780265+TylerMa-debugg@users.noreply.github.com> Co-authored-by: Claude --- tests/utils_/test_argparse_utils.py | 23 +++++++++++++++++++++++ vllm/utils/argparse_utils.py | 13 +++++++++---- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/utils_/test_argparse_utils.py b/tests/utils_/test_argparse_utils.py index 89e0f9977d0c..98f2fe790c12 100644 --- a/tests/utils_/test_argparse_utils.py +++ b/tests/utils_/test_argparse_utils.py @@ -4,6 +4,7 @@ import json import os +import sys from argparse import BooleanOptionalAction import pytest @@ -252,6 +253,28 @@ def test_duplicate_dict_args(caplog_vllm, parser): assert "--optimization-level" in caplog_vllm.text +@pytest.mark.skipif( + sys.version_info >= (3, 13), reason="argparse handles deprecated natively" +) +def test_deprecated_arg_scoped_to_subcommand(caplog_vllm, disable_log_dedup): + """A deprecated arg of one subcommand must not warn for a sibling + subcommand that defines a non-deprecated arg with the same dest.""" + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + run_batch = subparsers.add_parser("run-batch") + run_batch.add_argument_group("Frontend").add_argument( + "--url", default="0.0.0.0", deprecated=True + ) + chat = subparsers.add_parser("chat") + chat.add_argument("--url", default="http://localhost:8000/v1") + + parser.parse_args(["chat", "--url", "http://localhost:9000/v1"]) + assert "is deprecated" not in caplog_vllm.text + + parser.parse_args(["run-batch", "--url", "1.2.3.4"]) + assert caplog_vllm.text.count("argument 'url' is deprecated") == 1 + + def test_model_specification( parser_with_config, cli_config_file, cli_config_file_with_model ): diff --git a/vllm/utils/argparse_utils.py b/vllm/utils/argparse_utils.py index b91b1f26fd13..9ad03453fd20 100644 --- a/vllm/utils/argparse_utils.py +++ b/vllm/utils/argparse_utils.py @@ -119,7 +119,6 @@ def add_arguments(self, actions): class FlexibleArgumentParser(ArgumentParser): """ArgumentParser that allows both underscore and dash in names.""" - _deprecated: set[Action] = set() _json_tip: str = ( "When passing JSON CLI arguments, the following sets of arguments " "are equivalent:\n" @@ -132,6 +131,8 @@ class FlexibleArgumentParser(ArgumentParser): _search_keyword: str | None = None def __init__(self, *args, **kwargs): + # Must exist before super().__init__(), which already adds arguments. + self._deprecated: set[Action] = set() # Set the default "formatter_class" to SortedHelpFormatter if "formatter_class" not in kwargs: kwargs["formatter_class"] = SortedHelpFormatter @@ -160,7 +161,7 @@ def error(self, message: str) -> NoReturn: def parse_known_args(self, args=None, namespace=None): namespace, args = super().parse_known_args(args, namespace) - for action in FlexibleArgumentParser._deprecated: + for action in self._deprecated: if ( hasattr(namespace, dest := action.dest) and getattr(namespace, dest) != action.default @@ -172,15 +173,19 @@ def add_argument(self, *args, **kwargs): deprecated = kwargs.pop("deprecated", False) action = super().add_argument(*args, **kwargs) if deprecated: - FlexibleArgumentParser._deprecated.add(action) + self._deprecated.add(action) return action class _FlexibleArgumentGroup(_ArgumentGroup): + def __init__(self, container, *args, **kwargs): + super().__init__(container, *args, **kwargs) + self._deprecated: set[Action] = container._deprecated + def add_argument(self, *args, **kwargs): deprecated = kwargs.pop("deprecated", False) action = super().add_argument(*args, **kwargs) if deprecated: - FlexibleArgumentParser._deprecated.add(action) + self._deprecated.add(action) return action def add_argument_group(self, *args, **kwargs): From e41d5d55439808e79326009599d5d83e157bf2d1 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Thu, 8 Oct 2026 17:29:20 +0200 Subject: [PATCH 217/952] [Bugfix] Don't warn about missing watermark config for default requests (#60412) Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson --- tests/watermarking/test_compatibility.py | 5 +++-- vllm/config/vllm.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/watermarking/test_compatibility.py b/tests/watermarking/test_compatibility.py index 748223d94b56..48c27db09de2 100644 --- a/tests/watermarking/test_compatibility.py +++ b/tests/watermarking/test_compatibility.py @@ -93,7 +93,7 @@ def test_default_watermarking_resolves_when_configured(): assert enabled.watermarking is None -@pytest.mark.parametrize("watermarking", [None, True]) +@pytest.mark.parametrize("watermarking", [True]) def test_requested_watermarking_without_engine_config_warns_and_disables( watermarking, caplog_vllm, reset_warning_once ): @@ -148,6 +148,7 @@ def test_incompatible_modes_are_allowed_when_watermarking_is_disabled(params): @pytest.mark.parametrize( "params", [ + SamplingParams(), SamplingParams(temperature=0), SamplingParams(trace_decode_token_ids=[1]), ], @@ -158,7 +159,7 @@ def test_requests_without_engine_config_are_unwatermarked( with caplog_vllm.at_level("WARNING"): resolved = _validate(params, server_uses_watermarking=False) - assert "engine has no watermark configuration" in caplog_vllm.text + assert "engine has no watermark configuration" not in caplog_vllm.text assert resolved is False assert params.watermarking is None diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3fee812700b0..5d38b6765308 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1310,7 +1310,7 @@ def _check_supports_watermarking( ) -> bool: watermark_config = getattr(self, "watermark_config", None) if watermark_config is None: - if config is not None and config.watermarking is not False: + if config is not None and config.watermarking is True: logger.warning_once( "Watermarking is enabled for this request, but the engine has no " "watermark configuration. This and subsequent requests will run " From 0fc0cefb4eecf685c4bbd20912db9c1294d84c1e Mon Sep 17 00:00:00 2001 From: Justin Barlow <7@netlobo.com> Date: Thu, 8 Oct 2026 09:34:55 -0600 Subject: [PATCH 218/952] [Bugfix] Don't int-sort attention-type names in --kv-cache-dtype-skip-layers for packed KV dtypes (#60364) Signed-off-by: lobo235 <7@netlobo.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/engine/arg_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 94dd3e2ef1c8..5f32ec8deb9f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2285,9 +2285,11 @@ def create_engine_config( ) boundary = TurboQuantConfig.get_boundary_skip_layers(model_config) - existing = set(cache_config.kv_cache_dtype_skip_layers) + # Sorted so the order does not depend on set iteration: the list is + # a CacheConfig hash factor. Entries can also be attention type + # names (e.g. "sliding_window"), so sort them as strings. cache_config.kv_cache_dtype_skip_layers = sorted( - existing | set(boundary), key=int + set(cache_config.kv_cache_dtype_skip_layers) | set(boundary) ) ray_runtime_env = None From e2414ec1dbb177b54ac1bbce60e94769ebacbb28 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Thu, 8 Oct 2026 18:47:56 +0300 Subject: [PATCH 219/952] [ROCm][Qwen4Exp] Use fused PLE Triton kernels on the AMD backend (#60021) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- tests/models/qwen4_exp/test_ple.py | 57 +- vllm/models/qwen4_exp/amd/model.py | 18 +- vllm/models/qwen4_exp/amd/mtp.py | 4 +- vllm/models/qwen4_exp/amd/ple_layer.py | 784 +++++------------- vllm/models/qwen4_exp/common/ops/__init__.py | 4 + .../qwen4_exp/{nvidia => common}/ops/ple.py | 0 .../qwen4_exp/nvidia/ngram_embedding.py | 2 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 2 +- 8 files changed, 279 insertions(+), 592 deletions(-) create mode 100644 vllm/models/qwen4_exp/common/ops/__init__.py rename vllm/models/qwen4_exp/{nvidia => common}/ops/ple.py (100%) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 1ad44e830457..d7ee461d8f0d 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -45,6 +45,7 @@ Qwen4ExpPLEUnquantizedEmbeddingMethod, ) from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer +from vllm.platforms import current_platform from vllm.utils.torch_utils import weak_ref_tensor from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionMetadata, @@ -1102,7 +1103,7 @@ def test_fused_ngram_ids_correctness( contexts: list[list[int]], first_token_id: int, ) -> None: - from vllm.models.qwen4_exp.nvidia.ops.ple import ple_ngram_ids + from vllm.models.qwen4_exp.common.ops.ple import ple_ngram_ids device = torch.device("cuda") query_start_loc = torch.tensor( @@ -1857,15 +1858,20 @@ def _make_conv_metadata( ), ], ) +@pytest.mark.parametrize( + "layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"] +) def test_fused_conv_correctness( case: _ConvBatchCase, state_layout: str, + layer_cls: type[nn.Module], ) -> None: device = torch.device("cuda") metadata, num_real_tokens = _make_conv_metadata(case, device) if case.state_index_stride > 1: assert metadata.state_indices_tensor.stride(0) == case.state_index_stride - module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + # layer_cls is type[nn.Module], so layer_cls.__new__ is type.__new__. + module = object.__new__(layer_cls) nn.Module.__init__(module) module.conv_state_len = (case.kernel_size - 1) * case.dilation module.short_conv_dilation = case.dilation @@ -1938,7 +1944,7 @@ def test_fused_conv_correctness( def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None: import math - from vllm.models.qwen4_exp.nvidia.ops.ple import ple_gate + from vllm.models.qwen4_exp.common.ops.ple import ple_gate device = torch.device("cuda") hc, h = 4, 2560 @@ -2038,7 +2044,6 @@ def _build_amd_ngram_embedding( embedding_dim=config.ple_embed_dim, ple_dense_layer_id=0, max_total_tokens=4, - max_num_reqs=2, prefix="test.ple_embedding", layer_name="test.ple", quant_config=( @@ -2224,3 +2229,47 @@ def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor: expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="only the ROCm backend runs the n-gram embedding under torch.compile", +) +@pytest.mark.parametrize( + ("query_lens", "contexts"), + [([3, 1], [[5, 6], [0, 7]]), ([1, 0, 2], [[1, 2], [3, 4], [0, 0]])], + ids=["two-requests", "empty-request"], +) +def test_amd_ngram_embedding_matches_reference_under_compile( + monkeypatch: pytest.MonkeyPatch, + query_lens: list[int], + contexts: list[list[int]], +) -> None: + """The AMD n-gram ID op must run under Dynamo with a symbolic request count.""" + module, loaded_weight = _build_amd_ngram_embedding( + monkeypatch, device="cuda:0", cpu_offload=False, fp8_checkpoint=False + ) + device = torch.device("cuda:0") + query_start_loc = torch.tensor( + [0, *accumulate(query_lens)], dtype=torch.int32, device=device + ) + input_ids = torch.tensor( + [9, 0, 11, 12][: sum(query_lens)], dtype=torch.int32, device=device + ) + ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device) + + compiled = torch.compile(module, fullgraph=True, dynamic=True) + output = compiled(input_ids, query_start_loc, ngram_context) + + ngram_ids = _reference_ngram_ids( + input_ids, + query_start_loc, + ngram_context, + module.layer_multipliers, + module.ngram_heads_vocab_sizes, + module.ngram_heads_offsets, + module.eos_token_id, + module.heads_per_ngram, + ) + expected = loaded_weight[ngram_ids.cpu()].flatten(-2).to(device) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index b3a2b7505335..f36f70bee2ac 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -140,9 +140,9 @@ def _remap_qsa_cache_scale_name( "_input_scale", ] -# The checkpoint keeps down and injection projections separate; runtime packs -# them into adjacent logical shards of one MergedColumnParallelLinear. -_HC_WEIGHTS_MAPPER = WeightsMapper( +# The checkpoint stores these projections separately; runtime packs each group +# into adjacent logical shards of a MergedColumnParallelLinear. +_EXTRA_WEIGHTS_MAPPER = WeightsMapper( orig_to_new_stacked={ "hyper_connection.input_mix_weight_down.weight": ( "hyper_connection.input_mix_weight_down_block_inject.weight", @@ -152,6 +152,8 @@ def _remap_qsa_cache_scale_name( "hyper_connection.input_mix_weight_down_block_inject.weight", 1, ), + "ple.key_proj": ("ple.kv_proj", 0), + "ple.value_proj": ("ple.kv_proj", 1), } ) @@ -276,7 +278,7 @@ def forward( if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - hidden_states = hidden_states + self.ple( + hidden_states = self.ple( hidden_states, input_ids, query_start_loc, @@ -371,7 +373,7 @@ def update_physical_experts_metadata( } ) class Qwen4ExpModel(nn.Module): - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -594,6 +596,7 @@ class Qwen4ExpForCausalLM( packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], + "kv_proj": ["key_proj", "value_proj"], "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], "in_proj_ba": ["in_proj_b", "in_proj_a"], "input_mix_weight_down_block_inject": [ @@ -644,8 +647,6 @@ def forward( inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs - # path (query_start_loc / ngram_context) reaches Qwen4ExpModel. return self.model( input_ids, positions, @@ -825,11 +826,12 @@ class Qwen4ExpForConditionalGeneration( requires_raw_input_tokens = True packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | { + "kv_proj": ["key_proj", "value_proj"], "input_mix_weight_down_block_inject": [ "input_mix_weight_down", "block_inject_weight", "_input_mix_padding", - ] + ], } @staticmethod diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index d06f6436e77f..b6513e802d73 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -50,7 +50,7 @@ from .hyperconnection import GatedResidual, HyperConnectionConfig from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm from .model import ( - _HC_WEIGHTS_MAPPER, + _EXTRA_WEIGHTS_MAPPER, _QWEN4_EXP_IGNORED_MISSING_SUFFIXES, Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, @@ -172,7 +172,7 @@ def _make_draft_vllm_config( } ) class Qwen4ExpMultiTokenPredictor(nn.Module): - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index bd0607b61338..fbc07b73ddf8 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -2,18 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GPU-resident Qwen4Exp position-learning enhancement layers.""" -import math from collections.abc import Iterable, Sequence import torch -import torch.nn.functional as F from torch import nn from transformers import Qwen4ExpTextConfig from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, @@ -30,13 +28,13 @@ PleShortConvAttentionBackend, PleShortConvAttentionMetadata, ) -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, ) +from ..common.ops.ple import ple_conv, ple_gate, ple_ngram_ids logger = init_logger(__name__) @@ -177,7 +175,6 @@ def __init__( embedding_dim: int, ple_dense_layer_id: int, max_total_tokens: int, - max_num_reqs: int, prefix: str, layer_name: str, *, @@ -270,55 +267,6 @@ def __init__( self.ngram_embedding.weight.device, self.ngram_embedding.weight.is_pinned(), ) - self.register_buffer( - "positions_buffer", - torch.arange(max_total_tokens, dtype=torch.int64), - persistent=False, - ) - self.register_buffer( - "padded_buffer", - torch.full( - (max_num_reqs, max_total_tokens), - self.eos_token_id, - dtype=torch.int64, - ), - persistent=False, - ) - - @staticmethod - def _shift_precompute( - tokens: torch.Tensor, eos_token_id: int - ) -> tuple[torch.Tensor, torch.Tensor]: - if tokens.dim() != 2: - raise ValueError("tokens must be a 2D tensor") - batch_size, seq_len = tokens.shape - positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) - eos_positions = torch.where(tokens == eos_token_id, positions, -1) - previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values - previous_eos = torch.cat( - [ - eos_positions.new_full((batch_size, 1), -1), - previous_eos_inclusive[:, :-1], - ], - dim=1, - ) - return positions, positions.unsqueeze(0) - previous_eos - 1 - - @staticmethod - def _shift_apply( - tokens: torch.Tensor, - positions: torch.Tensor, - position_in_segment: torch.Tensor, - shift: int, - eos_token_id: int, - ) -> torch.Tensor: - if shift == 0: - return tokens - source = positions - shift - gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) - shifted = tokens.gather(1, gather_indices) - valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) - return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) def forward( self, @@ -326,64 +274,17 @@ def forward( query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: - input_ids = input_ids.reshape(-1).long() - query_start_loc = query_start_loc.long() - num_reqs = query_start_loc.numel() - 1 - num_tokens = input_ids.shape[0] - if num_tokens > self.positions_buffer.numel(): - raise ValueError( - f"PLE received {num_tokens} tokens, but its workspace supports " - f"at most {self.positions_buffer.numel()}" - ) - if num_reqs > self.padded_buffer.shape[0]: - raise ValueError( - f"PLE received {num_reqs} requests, but its workspace supports " - f"at most {self.padded_buffer.shape[0]}" - ) - - positions = self.positions_buffer[:num_tokens] - packed = self.padded_buffer[:num_reqs] - packed.fill_(self.eos_token_id) - request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 - request_indices.clamp_(max=num_reqs - 1) - columns = (positions - query_start_loc[request_indices]).clamp( - 0, packed.shape[1] - 1 - ) - packed[request_indices, columns] = input_ids - ngram_context = ngram_context[:num_reqs].to( - device=input_ids.device, dtype=torch.long + input_ids = input_ids.reshape(-1) + ngram_ids = input_ids.new_empty( + (input_ids.shape[0], self.ngram_heads), dtype=torch.int64 ) - - context = torch.cat([ngram_context, packed], dim=-1) - positions_2d, position_in_segment = self._shift_precompute( - context, self.eos_token_id + torch.ops.vllm.qwen4_exp_amd_ple_ngram_ids( + input_ids, + query_start_loc, + ngram_context, + ngram_ids, + self.layer_name, ) - shifted = [context] - for shift in range(1, self.ngram_size): - shifted.append( - self._shift_apply( - context, - positions_2d, - position_in_segment, - shift, - self.eos_token_id, - ) - ) - adjusted_columns = columns + self.ngram_size - 1 - id_blocks = [] - for ngram in range(2, self.ngram_size + 1): - start = (ngram - 2) * self.heads_per_ngram - end = start + self.heads_per_ngram - mixed = shifted[0] * self.layer_multipliers[0] - for index in range(1, ngram): - mixed = torch.bitwise_xor( - mixed, shifted[index] * self.layer_multipliers[index] - ) - sizes = self.ngram_heads_vocab_sizes[start:end] - offsets = self.ngram_heads_offsets[start:end] - ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets - id_blocks.append(ids[request_indices, adjusted_columns]) - ngram_ids = torch.cat(id_blocks, dim=-1) embedding = self.ngram_embedding if embedding.supports_prefetch: output = ngram_ids.new_empty( @@ -508,26 +409,21 @@ def __init__( int(config.ple_embed_dim), self.ple_dense_layer_id, vllm_config.scheduler_config.max_num_batched_tokens, - vllm_config.scheduler_config.max_num_seqs, f"{prefix}.ple_embedding", prefix, data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, quant_config=quant_config, params_dtype=model_config.dtype, ) - self.key_proj = ReplicatedLinear( + # The PLE cache is TP-replicated, so this merged projection is too. + self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), - self.hc_hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.key_proj", - ) - self.value_proj = ReplicatedLinear( - int(config.ple_embed_dim), - self.hidden_size, + [self.hc_hidden_size, self.hidden_size], bias=False, + params_dtype=model_config.dtype, quant_config=quant_config, - prefix=f"{prefix}.value_proj", + prefix=f"{prefix}.kv_proj", + disable_tp=True, ) norm_args = ( self.hc_hidden_size, @@ -580,355 +476,15 @@ def get_state_shape(self) -> Sequence[tuple[int, ...]]: num_spec=self.num_spec_tokens, ) - def _apply_norm( - self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor - ) -> torch.Tensor: - shape = hidden_states.shape - return norm(hidden_states.flatten(-2)).reshape(shape) - - def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor: - # Profiling / CUDA graph capture only; conv state is not updated. - inputs_t = inputs.transpose(0, 1).unsqueeze(0) - output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)] - return F.silu(output).squeeze(0).transpose(0, 1) - - def _short_conv_dilated_decode_batched( - self, - x_d: torch.Tensor, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - state_indices_tensor_d: torch.Tensor, - has_initial_states_d: torch.Tensor | None, - ) -> torch.Tensor: - state_indices = state_indices_tensor_d.to( - device=conv_state.device, dtype=torch.int64 - ) - # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to - # slot 0 for a safe gather, then zero output and skip write-back. - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - if has_initial_states_d is None: - has_initial_state = valid_state - else: - if has_initial_states_d.numel() < state_indices_tensor_d.numel(): - raise ValueError( - "has_initial_states_d size mismatch: " - f"got {has_initial_states_d.numel()}, " - f"need >= {state_indices_tensor_d.numel()}." - ) - has_initial_state = has_initial_states_d[ - : state_indices_tensor_d.numel() - ].to(device=conv_state.device, dtype=torch.bool) - has_initial_state = has_initial_state & valid_state - - cached_state = conv_state.index_select(0, state_indices) - state = cached_state[..., : self.conv_state_len].to(x_d.dtype) - if self.conv_state_len > 0: - initial_state = torch.where( - has_initial_state.view(-1, 1, 1), - state, - torch.zeros_like(state), - ) - history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1) - else: - history = x_d.unsqueeze(-1) - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ).squeeze(-1) - output = F.silu(conv_output) - output = output * valid_state.view(-1, 1).to(output.dtype) - - if self.conv_state_len > 0: - next_state = history[..., -self.conv_state_len :] - # Padded rows are remapped to the reserved null slot. Preserve its - # existing value while writing the new states for valid rows. - existing_base_state = cached_state[..., : self.conv_state_len] - safe_next_state = torch.where( - valid_state.view(-1, 1, 1), - next_state.to(conv_state.dtype), - existing_base_state, - ) - cached_state[..., : self.conv_state_len] = safe_next_state - conv_state.index_copy_(0, state_indices, cached_state) - - return output - - def _short_conv_dilated_prefill_batched( - self, - x_p: torch.Tensor, - metadata: PleShortConvAttentionMetadata, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - state_indices_tensor_p: torch.Tensor, - num_prefills: int, - num_decode_tokens: int, - num_prefill_tokens: int, - ) -> torch.Tensor: - query_start_loc_p = metadata.query_start_loc_p - if query_start_loc_p is None: - raise ValueError("query_start_loc is required for prefill short-conv") - # The metadata builder guarantees that the prefill query offsets start - # at 0 and end at num_prefill_tokens. Avoid reading those values here, - # since doing so would force a device-to-host synchronization. - has_initial_states_p = metadata.has_initial_states_p - if has_initial_states_p is None: - raise ValueError("has_initial_states_p is required for prefill short-conv") - - output = torch.empty_like(x_p) - q_starts = query_start_loc_p.to(torch.int64) - if state_indices_tensor_p.numel() < num_prefills: - raise ValueError( - "state_indices_tensor_p size mismatch: " - f"got {state_indices_tensor_p.numel()}, " - f"need >= {num_prefills}." - ) - if has_initial_states_p.numel() < num_prefills: - raise ValueError( - "has_initial_states_p size mismatch: " - f"got {has_initial_states_p.numel()}, " - f"need >= {num_prefills}." - ) - if num_prefills == 0 or x_p.numel() == 0: - return output - lengths = q_starts[1:] - q_starts[:-1] - # Use the CPU-computed packing width from the metadata builder instead - # of synchronizing on lengths.max(). - max_len = metadata.max_prefill_query_len - if max_len <= 0: - return output - - hidden_size = x_p.shape[1] - positions = torch.arange( - num_prefill_tokens, device=x_p.device, dtype=torch.int64 - ) - req_indices = torch.searchsorted(q_starts[1:], positions, right=True) - col_indices = positions - q_starts[req_indices] - - packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size)) - packed_tokens[req_indices, col_indices] = x_p - packed_tokens = packed_tokens.transpose(1, 2).contiguous() - - state_indices = state_indices_tensor_p[:num_prefills].to( - device=conv_state.device, dtype=torch.int64 - ) - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - has_initial = has_initial_states_p[:num_prefills].to( - device=conv_state.device, dtype=torch.bool - ) - if self.conv_state_len > 0: - if conv_state.shape[0] == 0: - state = conv_state.new_zeros( - (num_prefills, hidden_size, self.conv_state_len), - dtype=x_p.dtype, - ) - else: - state = conv_state.index_select(0, state_indices)[ - ..., : self.conv_state_len - ].to(x_p.dtype) - use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1) - initial_state = torch.where( - use_initial_mask, - state, - torch.zeros_like(state), - ) - history = torch.cat((initial_state, packed_tokens), dim=-1) - else: - history = packed_tokens - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ) - conv_output = F.silu(conv_output).transpose(1, 2).contiguous() - - token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64) - valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1) - valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view( - num_prefills, 1 - ) - conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0) - output.copy_(conv_output[req_indices, col_indices]) - - if self.conv_state_len > 0 and conv_state.shape[0] > 0: - state_starts = lengths.to(device=history.device, dtype=torch.int64).view( - num_prefills, 1, 1 - ) - state_offsets = torch.arange( - self.conv_state_len, device=history.device, dtype=torch.int64 - ).view(1, 1, self.conv_state_len) - next_state = history.gather( - dim=2, - index=(state_starts + state_offsets).expand(-1, history.size(1), -1), - ) - # Write back without a host synchronization. Valid, non-empty rows - # receive their new state; padding and zero-length rows keep the - # current cache value. - existing_state = conv_state.index_select(0, state_indices) - existing_base_state = existing_state[..., : self.conv_state_len] - update_mask = valid_state & (lengths.to(device=conv_state.device) > 0) - safe_next_state = torch.where( - update_mask.view(num_prefills, 1, 1), - next_state.to(conv_state.dtype), - existing_base_state, - ) - existing_state[..., : self.conv_state_len] = safe_next_state - conv_state.index_copy_(0, state_indices, existing_state) - return output - - def _short_conv_dilated_spec_batched( - self, - x_spec: torch.Tensor, - conv_state: torch.Tensor, - conv_weights: torch.Tensor, - spec_state_indices_tensor: torch.Tensor, - spec_query_start_loc: torch.Tensor, - num_accepted_tokens: torch.Tensor, - spec_query_len: int, - ) -> torch.Tensor: - """Dilated short-conv for speculative-decode (MTP) requests. - - Each spec request feeds multiple (draft + 1) query tokens. The conv - outputs are computed causally after rolling back the previous draft - state by ``num_accepted_tokens - 1``. The current candidate inputs stay - in the extended cache for the next forward, matching - ``causal_conv1d_update``. - - ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query - length and is a Python int, so no host synchronization is needed; this - keeps the path safe for full CUDA-graph capture/replay where the buffers - are padded at the request level. - """ - num_reqs = spec_state_indices_tensor.numel() - hidden_size = x_spec.size(-1) - # Use a fixed packing width instead of synchronizing on lengths.max(). - max_len = spec_query_len - # Full CUDA graphs can pad these buffers. Only the first num_reqs - # accepted-token counts belong to actual speculative requests. - num_accepted_tokens = num_accepted_tokens[:num_reqs] - q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64) - # Keep the number of real speculative tokens on the device. - total_real_tokens = q_starts[num_reqs] - - state_indices = spec_state_indices_tensor.to( - device=conv_state.device, dtype=torch.int64 - ) - valid_state = state_indices != NULL_BLOCK_ID - state_indices = torch.where( - valid_state, state_indices, torch.zeros_like(state_indices) - ) - positions = torch.arange( - x_spec.size(0), device=x_spec.device, dtype=torch.int64 - ) - # Route graph-padded token rows to the discarded dummy request so that - # they cannot overwrite real packed data. - req_indices = torch.searchsorted(q_starts[1:], positions, right=True) - valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs) - clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0)) - col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1) - pack_req_indices = torch.where( - valid_tokens, - clamped_req_indices, - torch.full_like(req_indices, num_reqs), - ) - pack_col_indices = torch.where( - valid_tokens, col_indices, torch.zeros_like(col_indices) - ) - - # The last request row is the dummy sink for graph padding. - packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size)) - packed[pack_req_indices, pack_col_indices] = x_spec - packed = packed.transpose(1, 2).contiguous() - - if self.conv_state_len > 0: - cached_state = conv_state.index_select(0, state_indices) - rollback_offsets = num_accepted_tokens.to( - device=conv_state.device, dtype=torch.int64 - ).sub(1) - rollback_offsets = torch.where( - valid_state, - rollback_offsets.clamp_(0, max_len - 1), - torch.zeros_like(rollback_offsets), - ) - state_offsets = torch.arange( - self.conv_state_len, device=conv_state.device, dtype=torch.int64 - ).view(1, 1, self.conv_state_len) - rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets - state = cached_state.gather( - 2, rollback_indices.expand(-1, hidden_size, -1) - ).to(x_spec.dtype) - state = torch.where( - valid_state.view(num_reqs, 1, 1), - state, - torch.zeros_like(state), - ) - # Append a zeroed dummy-row state to match the [num_reqs + 1] pack. - dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len)) - state_full = torch.cat((state, dummy_state), dim=0) - history = torch.cat((state_full, packed), dim=-1) - else: - history = packed - - conv_output = F.conv1d( - history, - conv_weights.unsqueeze(1).contiguous(), - groups=history.size(1), - dilation=self.short_conv_dilation, - ) - conv_output = F.silu(conv_output).transpose(1, 2).contiguous() - - output = conv_output[pack_req_indices, pack_col_indices] - output = output * valid_tokens.view(-1, 1).to(output.dtype) - - # Keep all current candidate inputs in the extended state. On the next - # target forward, ``num_accepted_tokens - 1`` selects the rollback - # window before processing the newly scheduled tokens. - if self.conv_state_len > 0: - state_capacity = self.conv_state_len + max_len - 1 - if conv_state.size(-1) < state_capacity: - raise RuntimeError( - "PLE short-conv cache cannot retain speculative tokens: " - f"got {conv_state.size(-1)}, need {state_capacity}." - ) - candidate_state = history[:num_reqs, :, 1 : state_capacity + 1] - query_lengths = q_starts[1:] - q_starts[:-1] - state_positions = torch.arange( - state_capacity, device=history.device, dtype=torch.int64 - ).view(1, 1, state_capacity) - update_lengths = (self.conv_state_len + query_lengths - 1).view( - num_reqs, 1, 1 - ) - update_mask = valid_state.view(num_reqs, 1, 1) & ( - state_positions < update_lengths - ) - existing_state = cached_state[..., :state_capacity] - next_state = torch.where( - update_mask, - candidate_state.to(conv_state.dtype), - existing_state, - ) - cached_state[..., :state_capacity] = next_state - conv_state.index_copy_(0, state_indices, cached_state) - - return output - def _short_conv_dilated_dispatch( self, inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, - ) -> torch.Tensor: + ) -> None: num_prefills = metadata.num_prefills num_decodes = metadata.num_decodes num_decode_tokens = metadata.num_decode_tokens @@ -936,110 +492,145 @@ def _short_conv_dilated_dispatch( has_prefill = num_prefills > 0 has_decode = num_decodes > 0 has_spec = metadata.spec_sequence_masks is not None - x = inputs[: metadata.num_actual_tokens] - - # Split spec / non-spec tokens. - if has_spec: - if has_prefill or has_decode: - assert metadata.spec_token_indx is not None - assert metadata.non_spec_token_indx is not None - x_spec = x.index_select(0, metadata.spec_token_indx.long()) - x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long()) - else: - x_spec = x - x_non_spec = None - else: - x_spec = None - x_non_spec = x + has_non_spec = has_prefill or has_decode + inputs = inputs[: metadata.num_actual_tokens] + residual = residual[: metadata.num_actual_tokens] + outer_residual = outer_residual[: metadata.num_actual_tokens] + + spec_token_indices = None + non_spec_token_indices = None + if has_spec and has_non_spec: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + spec_token_indices = metadata.spec_token_indx + non_spec_token_indices = metadata.non_spec_token_indx - spec_output = None - # 1. Run the multi-query speculative-decode part. if has_spec: assert metadata.spec_state_indices_tensor is not None - assert metadata.spec_query_start_loc is not None - assert metadata.num_accepted_tokens is not None - spec_output = self._short_conv_dilated_spec_batched( - x_spec=x_spec, + query_start_loc = metadata.spec_query_start_loc + num_accepted_tokens = metadata.num_accepted_tokens + assert query_start_loc is not None + assert num_accepted_tokens is not None + spec_state_indices = metadata.spec_state_indices_tensor[ + : metadata.num_spec_decodes + ] + # Mixed batches stay in their original row order; the kernels map + # logical spec/non-spec rows instead of materializing both groups. + ple_conv( + inputs=inputs, + residual=residual, conv_state=conv_state, conv_weights=conv_weights, - spec_state_indices_tensor=metadata.spec_state_indices_tensor[ - : metadata.num_spec_decodes - ], - spec_query_start_loc=metadata.spec_query_start_loc, - num_accepted_tokens=metadata.num_accepted_tokens, + state_indices=spec_state_indices, + outer_residual=outer_residual, + mode="spec", + dilation=self.short_conv_dilation, + query_start_loc=query_start_loc, + num_accepted_tokens=num_accepted_tokens, spec_query_len=metadata.spec_query_len, + token_indices=spec_token_indices, ) - # 2. Run regular decode and prefill requests. - conv_out_non_spec = None - state_indices_tensor = metadata.state_indices_tensor - if x_non_spec is not None: - assert state_indices_tensor is not None - if has_prefill: - state_indices_tensor_d, state_indices_tensor_p = torch.split( - state_indices_tensor, - [num_decodes, num_prefills], - dim=0, + if not has_non_spec: + return + + state_indices = metadata.state_indices_tensor + assert state_indices is not None + if has_prefill: + state_indices_d, state_indices_p = torch.split( + state_indices, [num_decodes, num_prefills], dim=0 + ) + if non_spec_token_indices is None: + inputs_d, inputs_p = torch.split( + inputs, [num_decode_tokens, num_prefill_tokens], dim=0 ) - x_d, x_p = torch.split( - x_non_spec, + residual_d, residual_p = torch.split( + residual, [num_decode_tokens, num_prefill_tokens], dim=0 + ) + outer_residual_d, outer_residual_p = torch.split( + outer_residual, [num_decode_tokens, num_prefill_tokens], dim=0, ) - non_spec_parts: list[torch.Tensor] = [] - if has_decode: - non_spec_parts.append( - self._short_conv_dilated_decode_batched( - x_d=x_d, - conv_state=conv_state, - conv_weights=conv_weights, - state_indices_tensor_d=state_indices_tensor_d, - has_initial_states_d=metadata.has_initial_states_d, - ) - ) - non_spec_parts.append( - self._short_conv_dilated_prefill_batched( - x_p=x_p, - metadata=metadata, - conv_state=conv_state, - conv_weights=conv_weights, - state_indices_tensor_p=state_indices_tensor_p, - num_prefills=num_prefills, - num_decode_tokens=num_decode_tokens, - num_prefill_tokens=num_prefill_tokens, - ) - ) - conv_out_non_spec = torch.vstack(non_spec_parts) + token_indices_d = None + token_indices_p = None else: - conv_out_non_spec = self._short_conv_dilated_decode_batched( - x_d=x_non_spec, + inputs_d = inputs_p = inputs + residual_d = residual_p = residual + outer_residual_d = outer_residual_p = outer_residual + token_indices_d, token_indices_p = torch.split( + non_spec_token_indices, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) + + if has_decode: + ple_conv( + inputs=inputs_d, + residual=residual_d, conv_state=conv_state, conv_weights=conv_weights, - state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)], - has_initial_states_d=metadata.has_initial_states_d, + state_indices=state_indices_d, + outer_residual=outer_residual_d, + mode="decode", + dilation=self.short_conv_dilation, + has_initial_states=metadata.has_initial_states_d, + token_indices=token_indices_d, ) - # 3. Merge both parts back into the original token order. - if has_spec and conv_out_non_spec is not None: - assert metadata.spec_token_indx is not None - assert metadata.non_spec_token_indx is not None - assert spec_output is not None - output = x.new_empty((metadata.num_actual_tokens, x.size(-1))) - output.index_copy_(0, metadata.spec_token_indx, spec_output) - output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec) - return output - elif has_spec: - assert spec_output is not None - return spec_output - if conv_out_non_spec is None: - return x - return conv_out_non_spec + query_start_loc = metadata.query_start_loc_p + if query_start_loc is None: + raise ValueError("query_start_loc is required for prefill short-conv") + has_initial_states = metadata.has_initial_states_p + if has_initial_states is None: + raise ValueError( + "has_initial_states_p is required for prefill short-conv" + ) + ple_conv( + inputs=inputs_p, + residual=residual_p, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices=state_indices_p, + outer_residual=outer_residual_p, + mode="prefill", + dilation=self.short_conv_dilation, + query_start_loc=query_start_loc, + has_initial_states=has_initial_states, + token_indices=token_indices_p, + ) + else: + num_decode_rows = ( + non_spec_token_indices.numel() + if non_spec_token_indices is not None + else inputs.size(0) + ) + ple_conv( + inputs=inputs, + residual=residual, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices=state_indices[:num_decode_rows], + outer_residual=outer_residual, + mode="decode", + dilation=self.short_conv_dilation, + has_initial_states=metadata.has_initial_states_d, + token_indices=non_spec_token_indices, + ) - def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: + def _short_conv( + self, + inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, + ) -> None: forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata + # Profiling omits all metadata or this Mamba entry. Short convolution + # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - return self._short_conv_fallback(inputs) + residual.add_(outer_residual) + return if not isinstance(attn_metadata, dict): raise RuntimeError( @@ -1049,8 +640,8 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - # MRV2 omits Mamba-family metadata during profile warmup. - return self._short_conv_fallback(inputs) + residual.add_(outer_residual) + return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( "Expected PleShortConvAttentionMetadata for layer " @@ -1059,6 +650,7 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: ) conv_state = self.kv_cache[0] + # Canonicalize both backend cache layouts to [slot, channel, window]. if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) conv_weights = self.conv1d.weight.squeeze(1) @@ -1072,11 +664,13 @@ def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: f"expect at least {state_capacity}." ) conv_state = conv_state[..., -state_capacity:] - return self._short_conv_dilated_dispatch( - inputs, - layer_attn_metadata, - conv_state, - conv_weights.to(dtype=inputs.dtype), + self._short_conv_dilated_dispatch( + inputs=inputs, + residual=residual, + outer_residual=outer_residual, + metadata=layer_attn_metadata, + conv_state=conv_state, + conv_weights=conv_weights.to(dtype=inputs.dtype), ) def forward( @@ -1095,27 +689,58 @@ def forward( ) embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) embeddings = self.ple_embedding.ngram_embedding.dequantize( - embeddings, - self.ple_embedding.ngram_embedding.params_dtype, - ) - key, _ = self.key_proj(embeddings) - value, _ = self.value_proj(embeddings) - token_count = hidden_states.shape[0] - key = key.reshape(token_count, self.hc_count, self.hidden_size) - query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size) - key = self._apply_norm(self.norm_key, key) - query = self._apply_norm(self.norm_query, query) - gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) - gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt()) - gated_value = gate * value.unsqueeze(-2) - normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) - conv_output = torch.zeros_like(normalized) + embeddings, hidden_states.dtype + ) + kv, _ = self.kv_proj(embeddings) + key, value = kv.split(self.kv_proj.output_sizes, dim=-1) + gated_output, conv_input = ple_gate( + key, + value, + hidden_states, + self.norm_key.weight, + self.norm_query.weight, + self.norm_conv.weight, + self.norm_key.eps, + ) + # The short-conv op is a piecewise-graph splitting op: state routing + # reads runtime request metadata. It accumulates the convolution and + # the outer residual into gated_output. torch.ops.vllm.qwen4_exp_ple_short_conv( - normalized, - conv_output, + conv_input, + gated_output, + hidden_states, self.prefix, ) - return gated_value.flatten(-2) + conv_output + return gated_output + + +def qwen4_exp_amd_ple_ngram_ids( + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + """Hash the current request layout into n-gram embedding indices. + + The launch sizes its request binary search from the request count, which + is symbolic under Dynamo, so it runs outside Inductor's FX graph. + """ + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpPLELayer): + raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner") + embedding = layer.ple_embedding + ple_ngram_ids( + input_ids=input_ids, + query_start_loc=query_start_loc, + ngram_context=ngram_context, + layer_multipliers=embedding.layer_multipliers, + ngram_heads_vocab_sizes=embedding.ngram_heads_vocab_sizes, + ngram_heads_offsets=embedding.ngram_heads_offsets, + eos_token_id=embedding.eos_token_id, + heads_per_ngram=embedding.heads_per_ngram, + output=output, + ) def qwen4_exp_amd_ple_ngram_embedding( @@ -1155,12 +780,19 @@ def qwen4_exp_amd_ple_ngram_embedding_pinned( def qwen4_exp_ple_short_conv( inputs: torch.Tensor, - output: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, layer_name: str, ) -> None: layer = get_forward_context().no_compile_layers[layer_name] - result = layer._short_conv(inputs) - output[: result.shape[0]].copy_(result) + layer._short_conv(inputs, residual, outer_residual) + + +direct_register_custom_op( + op_name="qwen4_exp_amd_ple_ngram_ids", + op_func=qwen4_exp_amd_ple_ngram_ids, + mutates_args=["output"], +) direct_register_custom_op( @@ -1180,7 +812,7 @@ def qwen4_exp_ple_short_conv( direct_register_custom_op( op_name="qwen4_exp_ple_short_conv", op_func=qwen4_exp_ple_short_conv, - mutates_args=["output"], + mutates_args=["residual"], ) diff --git a/vllm/models/qwen4_exp/common/ops/__init__.py b/vllm/models/qwen4_exp/common/ops/__init__.py new file mode 100644 index 000000000000..a87d2ca4bed3 --- /dev/null +++ b/vllm/models/qwen4_exp/common/ops/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp kernels shared by the NVIDIA and AMD backends; import leaf modules +directly.""" diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py similarity index 100% rename from vllm/models/qwen4_exp/nvidia/ops/ple.py rename to vllm/models/qwen4_exp/common/ops/ple.py diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 3be6ea6013dd..cf03cbb40d47 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -23,7 +23,7 @@ Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, ) -from .ops.ple import ple_ngram_ids +from ..common.ops.ple import ple_ngram_ids logger = init_logger(__name__) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index ee8a2b47dd76..bdb16b3adef0 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -24,8 +24,8 @@ PleShortConvAttentionMetadata, ) +from ..common.ops.ple import ple_conv, ple_gate from .ngram_embedding import Qwen4ExpNGramEmbedding -from .ops.ple import ple_conv, ple_gate class Qwen4ExpPLEGroupedNorm(nn.Module): From 9f3d70e08157a95911f15809f2f1446424825ea7 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Thu, 8 Oct 2026 18:04:46 +0200 Subject: [PATCH 220/952] [Quantization] Enable shared expert fusion compatibility with online `shared_expert` quantization (showcase: along Quark MXFP4 routed experts) (#55686) Signed-off-by: Felix Marty Signed-off-by: fxmarty-amd Co-authored-by: OpenAI Codex Co-authored-by: Misha Goin --- .../layers/test_fused_shared_expert.py | 89 ++- .../model_loader/test_reload.py | 1 + tests/models/kimi_k3/test_gfx942_int4.py | 10 +- tests/quantization/test_auto_gptq.py | 1 + tests/quantization/test_config_utils.py | 5 +- tests/quantization/test_online.py | 9 +- .../test_online_shared_expert_fusion.py | 637 ++++++++++++++++++ .../test_quantization_config_args.py | 6 +- tests/quantization/utils.py | 29 +- .../model_executor/layers/fused_moe/config.py | 1 + vllm/model_executor/layers/fused_moe/layer.py | 3 + .../layers/fused_moe/routed_experts.py | 56 ++ vllm/model_executor/layers/fused_moe/utils.py | 34 + .../layers/quantization/base_config.py | 2 +- .../layers/quantization/online/base.py | 44 +- .../layers/quantization/online/fp8.py | 23 +- .../layers/quantization/online/mxfp8.py | 3 + .../layers/quantization/utils/config_utils.py | 98 ++- vllm/model_executor/models/AXK1.py | 1 + vllm/model_executor/models/deepseek_v2.py | 1 + vllm/model_executor/models/glm4_moe.py | 1 + vllm/model_executor/models/qwen3_next.py | 1 + vllm/models/glm5next/common/model.py | 1 + vllm/models/minimax_m3/amd/model.py | 1 + 24 files changed, 984 insertions(+), 73 deletions(-) create mode 100644 tests/quantization/test_online_shared_expert_fusion.py diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index d8d5c4fb032a..78aa563925ff 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -15,10 +15,12 @@ import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.config_utils import ( get_quark_ocp_mx_group_size, @@ -141,10 +143,6 @@ def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]: return quantization_config -def _stub_quant_config() -> QuantizationConfig: - return cast(QuantizationConfig, object()) - - class _StubGlm5NextAttention(nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__() @@ -340,7 +338,7 @@ def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) @@ -355,7 +353,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( assert ( fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) is False ) @@ -364,7 +362,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( def test_resolve_layer_fused_shared_expert_passes_module_prefixes( monkeypatch: pytest.MonkeyPatch, ) -> None: - quant_config = _stub_quant_config() + quant_config = ExpertsInt8Config() monkeypatch.setattr( fused_moe_utils.rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", @@ -413,7 +411,7 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) assert "shared experts are excluded" in caplog.text @@ -423,6 +421,7 @@ def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( ) -> None: class DeepseekV4Config: expert_dtype = "fp4" + online_quantization_config = None def _is_quark_mxfp4_ocp(self, hf_config: object) -> bool: return True @@ -1070,6 +1069,73 @@ def test_quark_shared_expert_fse_compatibility( ) +@pytest.mark.parametrize("exclude_shared_expert", [False, True]) +def test_online_targets_off_shared_expert_defer_to_checkpoint( + exclude_shared_expert: bool, +) -> None: + """Unrelated online targets preserve the checkpoint's FSE decision.""" + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig( + { + **_QUARK_FSE_CONFIG, + "exclude": ( + [f"{shared_prefix}.down_proj"] if exclude_shared_expert else [] + ), + } + ) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*linear_attn.out_proj": "mxfp4"}) + ) + + compatible, reason = is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) + + assert compatible is not exclude_shared_expert + assert reason == ( + f"Quark excludes shared experts at {shared_prefix}" + if exclude_shared_expert + else None + ) + + +def test_online_targeting_only_one_shared_projection_rejects_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}.down_proj": "mxfp4"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_prefix}", + ) + + +def test_online_fp8_shared_expert_rejects_quark_mxfp4_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}*": "fp8_per_tensor"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online shared-expert quantization keys do not match the routed expert keys", + ) + + def test_quark_shared_expert_fse_exclude_is_scoped_to_the_layer() -> None: """Excluding one layer's shared experts must not disable FSE elsewhere. @@ -1335,14 +1401,15 @@ def test_quark_packed_layer_config_must_match_global_config() -> None: def test_non_quark_shared_expert_fse_is_incompatible() -> None: compatible, reason = is_shared_expert_quant_fse_compatible( - _stub_quant_config(), + ExpertsInt8Config(), "model.layers.0.mlp.experts", "model.layers.0.mlp.shared_experts", ) assert not compatible assert reason == ( - "shared-expert FSE quantization compatibility is not implemented for object" + "shared-expert FSE quantization compatibility is not implemented for " + "ExpertsInt8Config" ) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index aaad16778f12..66c5b490508a 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -528,6 +528,7 @@ def test_padded_moe_reload_releases_each_layer( layer.quant_method = method layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) layer._loaded_expert_biases = set() + layer._fused_shared_expert_quantizer = None method.create_weights( layer, experts, diff --git a/tests/models/kimi_k3/test_gfx942_int4.py b/tests/models/kimi_k3/test_gfx942_int4.py index 20bef41eb89d..e38a8ecff944 100644 --- a/tests/models/kimi_k3/test_gfx942_int4.py +++ b/tests/models/kimi_k3/test_gfx942_int4.py @@ -4,7 +4,7 @@ import sys from collections.abc import Callable from types import ModuleType, SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest @@ -91,9 +91,7 @@ def test_int4_per_group_32_is_not_an_online_moe_method(): QuantizationConfigArgs(moe={"weight": "int4_per_group_32"}) ) assert ( - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") is None ) @@ -103,6 +101,4 @@ def test_unknown_online_moe_weight_still_raises(): QuantizationConfigArgs(moe={"weight": "fp8_per_token"}) ) with pytest.raises(ValueError, match="online quantization"): - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 7e54599c8098..4f60a974caed 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -162,6 +162,7 @@ def test_routed_experts_loads_per_expert_biases(): class Loader: quant_config = None quant_method = object() + _fused_shared_expert_quantizer = None moe_config = SimpleNamespace( is_act_and_mul=True, tp_rank=0, diff --git a/tests/quantization/test_config_utils.py b/tests/quantization/test_config_utils.py index 6a0c27748f90..a0cb12a9f8ac 100644 --- a/tests/quantization/test_config_utils.py +++ b/tests/quantization/test_config_utils.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests quantization configuration matching utilities.""" -from unittest.mock import Mock - import pytest from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( should_ignore_layer, ) @@ -58,7 +57,7 @@ def test_online_ignore_supports_fnmatch_patterns(): ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls(Mock(), layer_name) + config.resolve_quant_method_cls(LinearBase, layer_name) @pytest.mark.parametrize( diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index c6b9894009cf..ce13aa9a65d1 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -195,7 +195,6 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( config, ) - layer = Mock(spec=LinearBase) for prefix in ( "vision_model.encoder.layers.0.attn.qkv", "vision_model.encoder.layers.1.mlp.fc2", @@ -203,10 +202,10 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( "mlp1.1", "mlp1.3", ): - assert config.resolve_quant_method_cls(layer, prefix) is None + assert config.resolve_quant_method_cls(LinearBase, prefix) is None resolved = config.resolve_quant_method_cls( - layer, "language_model.model.layers.0.self_attn.qkv_proj" + LinearBase, "language_model.model.layers.0.self_attn.qkv_proj" ) assert resolved is not None and resolved[-1] is method_cls assert other_config.ignored_layers == [] @@ -1320,9 +1319,7 @@ def test_online_quantization_targets_ignore_collision() -> None: ) ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") def test_online_quantization_targets_reject_unsupported_layer() -> None: diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py new file mode 100644 index 000000000000..2dc3d5da76bf --- /dev/null +++ b/tests/quantization/test_online_shared_expert_fusion.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests ROCm AITER fused shared experts with online quantization.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from tests.quantization.utils import load_model_without_vllm_runner +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.model import ModelConfig +from vllm.config.quantization import ( + QuantizationConfigArgs, + resolve_quantization_config, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig +from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkOCP_MX_MoEMethod, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, +) +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) +from vllm.model_executor.model_loader.utils import get_model_architecture +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm AITER shared-expert tests require ROCm.", +) + +_QUARK_MXFP4_CONFIG: dict[str, Any] = { + "global_quant_config": { + "input_tensors": { + "dtype": "fp4", + "is_dynamic": True, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "output_tensors": None, + "weight": { + "dtype": "fp4", + "is_dynamic": False, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "bias": None, + "target_device": None, + }, + "algo_config": None, + "softmax_quant_spec": None, + "quant_method": "quark", + "layer_type_quant_config": {}, + "layer_quant_config": {}, + "exclude": ["re:^(?!.*\\.mlp\\.experts(?:\\.|$)).*$"], + "kv_cache_quant_config": {}, + "kv_cache_post_rope": False, + "quant_mode": "eager_mode", + "version": "0.12+9d3d471cdf1", + "export": { + "kv_cache_group": [], + "min_kv_scale": 0.0, + "pack_method": "reorder", + "weight_format": "real_quantized", + "weight_merge_groups": None, + }, +} + +_QUARK_MXFP4_LAYER_CONFIG: dict[str, Any] = _QUARK_MXFP4_CONFIG["global_quant_config"] +_QUARK_MXFP4_CONFIG = { + **_QUARK_MXFP4_CONFIG, + "global_quant_config": { + "input_tensors": {}, + "output_tensors": None, + "weight": {}, + "bias": None, + "target_device": None, + }, + "layer_quant_config": {"*mlp.experts*": _QUARK_MXFP4_LAYER_CONFIG}, +} + + +class _StubAttention(torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + +@pytest.fixture(autouse=True) +def reset_aiter_shared_expert_topk_metadata(monkeypatch: pytest.MonkeyPatch): + """Isolate AITER's process-global shared-expert routing buffer.""" + monkeypatch.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None) + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + yield + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + + +def _write_minimal_moe_config(model_path: Path, architecture: str) -> None: + """Write a compact MoE config with one shared expert.""" + config = { + "architectures": [architecture], + "hidden_size": 256, + "intermediate_size": 512, + "moe_intermediate_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "num_hidden_layers": 1, + "vocab_size": 256, + "max_position_embeddings": 64, + "q_lora_rank": 64, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "n_routed_experts": 4, + "num_experts": 4, + "num_experts_per_tok": 1, + "n_group": 1, + "topk_group": 1, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "moe_layer_freq": 1, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "tie_word_embeddings": False, + "rope_theta": 10000, + "quantization_config": _QUARK_MXFP4_CONFIG, + } + if architecture == "AXK1ForCausalLM": + config["model_type"] = "axk1" + elif architecture == "MiniMaxM3SparseForCausalLM": + text_config = { + **config, + "model_type": "minimax_m3_text", + "dense_intermediate_size": 512, + "hidden_act": "swigluoai", + "num_local_experts": 4, + "n_shared_experts": 1, + "moe_layer_freq": [1], + "sparse_attention_config": None, + "quantization_config": { + **_QUARK_MXFP4_CONFIG, + "layer_quant_config": { + "*block_sparse_moe.experts*": _QUARK_MXFP4_LAYER_CONFIG + }, + "exclude": [r"re:^(?!.*\.block_sparse_moe\.experts(?:\.|$)).*$"], + }, + } + config = { + "architectures": [architecture], + "model_type": "minimax_m3_vl", + "text_config": text_config, + } + elif architecture == "Glm4MoeForCausalLM": + config["model_type"] = "glm4_moe" + elif architecture == "Glm4MoeLiteForCausalLM": + config["model_type"] = "glm4_moe_lite" + elif architecture in ("Qwen3NextForCausalLM", "Qwen3_5MoeForCausalLM"): + config.update( + model_type=( + "qwen3_next" + if architecture == "Qwen3NextForCausalLM" + else "qwen3_5_moe_text" + ), + layer_types=[ + "linear_attention" + if architecture == "Qwen3NextForCausalLM" + else "full_attention" + ], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + ) + elif architecture == "Qwen4ExpForCausalLM": + config.update( + model_type="qwen4_exp_text", + layer_types=["full_attention"], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + eos_token_id=1, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + ) + elif architecture == "Glm5NextForCausalLM": + config.update( + model_type="glm5_next", + layer_types=["linear_attention"], + mlp_layer_types=["sparse"], + qk_rope_head_dim=0, + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + ) + else: + config["model_type"] = ( + "deepseek_v3" + if "V3" in architecture + else "deepseek_v32" + if architecture == "DeepseekV32ForCausalLM" + else "deepseek_v2" + ) + if architecture in ("DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"): + config.update( + index_topk=1, + index_kpool=1, + index_n_heads=1, + index_head_dim=32, + index_kv_lora_rank=32, + ) + (model_path / "config.json").write_text(json.dumps(config)) + + +@pytest.mark.parametrize( + "load_path", ["load_weights", "fused_gate_up", "weight_loader"] +) +def test_online_shared_expert_loads_bf16_weights_into_mxfp4_slot( + default_vllm_config, + dist_init, + load_path: str, +) -> None: + """A BF16 shared expert is quantized while routed MXFP4 weights are loaded. + + `weight_loader` covers DeepSeek-style model loaders, which bypass + `RoutedExperts.load_weights`. + """ + default_vllm_config.model_config = ModelConfig() + hidden_size = intermediate_size = 64 + num_routed_experts = 2 + device = current_platform.device_type + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(device): + runner = FusedMoEFactory( + num_experts=num_routed_experts, + top_k=1, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ) + layer = runner.routed_experts + assert isinstance(layer.quant_method, QuarkOCP_MX_MoEMethod) + + routed_gate, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_up, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_down, _ = mxfp4_quantize( + torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + ) + shared_gate = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_up = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_down = torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + + weights = [ + ("0.gate_proj.weight", routed_gate), + ("0.up_proj.weight", routed_up), + ("0.down_proj.weight", routed_down), + ] + if load_path == "load_weights": + weights += [ + ("2.gate_proj.weight", shared_gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + elif load_path == "fused_gate_up": + weights += [ + ("2.gate_up_proj.weight", torch.cat([shared_gate, shared_up])), + ("2.down_proj.weight", shared_down), + ] + list(layer.load_weights(weights)) + + if load_path == "weight_loader": + for stem, shard_id, weight in ( + ("w13", "w1", shared_gate), + ("w13", "w3", shared_up), + ("w2", "w2", shared_down), + ): + param = getattr(layer, f"{stem}_weight") + param.weight_loader( + param, weight, f"{layer.layer_name}.{stem}_weight", shard_id, 2 + ) + + expected_gate, expected_gate_scale = mxfp4_quantize(shared_gate) + expected_up, expected_up_scale = mxfp4_quantize(shared_up) + expected_down, expected_down_scale = mxfp4_quantize(shared_down) + assert torch.equal(layer.w13_weight[2, :intermediate_size], expected_gate) + assert torch.equal(layer.w13_weight[2, intermediate_size:], expected_up) + assert torch.equal(layer.w2_weight[2], expected_down) + assert torch.equal( + layer.w13_weight_scale[2, :intermediate_size], expected_gate_scale + ) + assert torch.equal(layer.w13_weight_scale[2, intermediate_size:], expected_up_scale) + assert torch.equal(layer.w2_weight_scale[2], expected_down_scale) + + +def test_online_shared_expert_reload_compatibility( + default_vllm_config, + dist_init, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reload a BF16 shared weight into its MXFP4 slot via meta staging.""" + default_vllm_config.model_config = ModelConfig() + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(current_platform.device_type): + layer = FusedMoEFactory( + num_experts=2, + top_k=1, + hidden_size=64, + intermediate_size=64, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ).routed_experts + monkeypatch.setattr( + layer.quant_method, "process_weights_after_loading", lambda _: None + ) + + routed_weights = [] + for expert_id in range(2): + for projection in ("gate_proj", "up_proj", "down_proj"): + weight, scale = mxfp4_quantize( + torch.randn(64, 64, dtype=torch.bfloat16) + ) + prefix = f"{expert_id}.{projection}" + routed_weights.extend( + [(f"{prefix}.weight", weight), (f"{prefix}.weight_scale", scale)] + ) + + shared_gate = torch.randn(64, 64, dtype=torch.bfloat16) + shared_up = torch.randn(64, 64, dtype=torch.bfloat16) + shared_down = torch.randn(64, 64, dtype=torch.bfloat16) + + def checkpoint_weights(gate: torch.Tensor): + return [ + *routed_weights, + ("2.gate_proj.weight", gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + + list(layer.load_weights(checkpoint_weights(shared_gate))) + old_weight = layer.w13_weight[2, :64].clone() + reloaded_gate = shared_gate + 1 + expected_weight, expected_scale = mxfp4_quantize(reloaded_gate) + assert not torch.equal(old_weight, expected_weight) + + record_metadata_for_reloading(layer) + initialize_layerwise_reload(layer) + + assert layer.w13_weight.is_meta + list(layer.load_weights(checkpoint_weights(reloaded_gate))) + finalize_layerwise_reload(layer, default_vllm_config.model_config) + assert not layer.w13_weight.is_meta + assert torch.equal(layer.w13_weight[2, :64], expected_weight) + assert torch.equal(layer.w13_weight_scale[2, :64], expected_scale) + + +# TODO: Add DeepseekV4ForConditionalGeneration once it uses +# resolve_layer_fused_shared_expert. +@pytest.mark.parametrize( + "architecture", + [ + "AXK1ForCausalLM", + "DeepseekForCausalLM", + "DeepseekV2ForCausalLM", + "DeepseekV32ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", + "Glm5NextForCausalLM", + "GlmMoeDsaForCausalLM", + "MiniMaxM3SparseForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3_5MoeForCausalLM", + "Qwen4ExpForCausalLM", + ], +) +def test_online_quantization( + architecture: str, + tmp_path: Path, + monkeypatch, + dist_init, + workspace_init, +) -> None: + """Online MXFP4 fuses each architecture's shared expert into its MoE.""" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + if architecture == "Glm5NextForCausalLM": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubAttention) + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) + elif architecture == "Glm4MoeLiteForCausalLM": + from vllm.model_executor.models import glm4_moe_lite + + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteAttention", _StubAttention) + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteMLAAttention", _StubAttention) + elif architecture == "DeepseekV32ForCausalLM": + from vllm.models.deepseek_v32.amd import model as deepseek_v32_model + + monkeypatch.setattr( + deepseek_v32_model, "DeepseekV32MLAAttention", _StubAttention + ) + elif architecture == "Qwen3_5MoeForCausalLM": + from vllm.model_executor.models import qwen3_5 + + monkeypatch.setattr(qwen3_5, "Qwen3NextAttention", _StubAttention) + elif architecture == "Qwen4ExpForCausalLM": + from vllm.models.qwen4_exp.amd import model as qwen4_exp_model + + monkeypatch.setattr(qwen4_exp_model, "Qwen3NextAttention", _StubAttention) + _write_minimal_moe_config(tmp_path, architecture) + expected_shared_expert_name = ( + "shared_expert" if architecture.startswith("Qwen") else "shared_experts" + ) + target_pattern = f"*{expected_shared_expert_name}*" + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, vllm_config = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + expected_model_cls, resolved_architecture = get_model_architecture( + vllm_config.model_config + ) + assert resolved_architecture == architecture + assert isinstance(model, expected_model_cls) + + fused_moes = [ + module + for module in model.modules() + if getattr(module, "is_fused_shared_expert_enabled", False) + and hasattr(module, "experts") + ] + assert len(fused_moes) == 1, f"{architecture} did not enable FSE" + moe = fused_moes[0] + shared_expert_name = ( + "shared_expert" if hasattr(moe, "shared_expert") else "shared_experts" + ) + assert shared_expert_name == expected_shared_expert_name + assert getattr(moe, shared_expert_name) is None + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts.quant_method, QuarkOCP_MX_MoEMethod) + shared_expert_prefix = routed_experts.moe_config.shared_expert_prefix + assert shared_expert_prefix is not None + assert shared_expert_prefix.endswith(shared_expert_name) + assert routed_experts._fused_shared_expert_quantizer is not None + expert_map_manager = routed_experts.expert_map_manager + assert expert_map_manager.num_fused_shared_experts == 1 + assert expert_map_manager.map_global_to_local( + routed_experts.moe_config.num_experts + ) == (routed_experts.moe_config.num_experts) + expected_weight_dtype = { + Mxfp4MoeBackend.AITER_MXFP4_BF16: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.AITER_MXFP4_MXFP4: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.EMULATION: torch.uint8, + }[routed_experts.quant_method.mxfp4_backend] + assert routed_experts.w13_weight.dtype == expected_weight_dtype + assert routed_experts.w13_weight_scale.dtype == torch.uint8 + assert routed_experts.w2_weight.dtype == expected_weight_dtype + assert routed_experts.w2_weight_scale.dtype == torch.uint8 + assert vllm_config.quant_config is not None + assert vllm_config.quant_config.online_quantization_config is not None + assert not any( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled" in warning + for warning in logged_warnings + ) + if architecture == "MiniMaxM3SparseForCausalLM": + # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove + # this branch once its virtual fused projections are registered. + assert logged_messages == ["Quantizing 0 layers of types: "] + else: + assert logged_messages == [ + "Quantizing 2 layers of types: " + f"mlp.{shared_expert_name}.down_proj: 1 " + f"(from targets: {target_pattern}, mxfp4); " + f"mlp.{shared_expert_name}.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + + +def test_online_gate_up_proj_target_disables_fse( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + dist_init, + workspace_init, +) -> None: + """A partial shared-expert target quantizes while disabling FSE.""" + architecture = "DeepseekForCausalLM" + target_pattern = "*shared_experts.gate_up_proj*" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + _write_minimal_moe_config(tmp_path, architecture) + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, _ = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + + moe = next( + module + for module in model.modules() + if hasattr(module, "is_fused_shared_expert_enabled") + and hasattr(module, "shared_experts") + ) + assert not moe.is_fused_shared_expert_enabled + assert moe.shared_experts is not None + assert logged_messages == [ + "Quantizing 1 layers of types: " + "mlp.shared_experts.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + assert logged_warnings == [ + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be " + "enabled - skipping for this layer: online quantization targets only " + "part of the shared expert at model.layers.0.mlp.shared_experts." + ] diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 7e2e21a3acbf..f24566f5fa8a 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for QuantizationConfigArgs parsing.""" -from unittest.mock import Mock - import pytest from tests.quantization.utils import quant_config_args, quant_spec @@ -244,6 +242,4 @@ def test_targets_reject_moe_only_shorthand_for_linear_layer(): ) with pytest.raises(ValueError, match="does not define a QuantSpec"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index ed9776ad1c25..d8446952f0d7 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -15,7 +15,11 @@ VllmConfig, set_current_vllm_config, ) -from vllm.config.quantization import QuantizationConfigArgs, QuantSpec +from vllm.config.quantization import ( + QuantizationConfigArgs, + QuantSpec, + resolve_quantization_config, +) from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.platforms import current_platform @@ -40,10 +44,17 @@ def _limit_num_hidden_layers( original_load_weights = model.load_weights def should_load_weight(name: str) -> bool: - for prefix in ("model.layers.", "layers."): - if name.startswith(prefix): - layer_idx = int(name.removeprefix(prefix).split(".", 1)[0]) - return layer_idx < num_hidden_layers + layers_prefix = "layers." + if name.startswith(layers_prefix): + layer_name = name.removeprefix(layers_prefix) + elif ".layers." in name: + layer_name = name.split(".layers.", maxsplit=1)[1] + else: + return True + + layer_idx_str = layer_name.split(".", maxsplit=1)[0] + if layer_idx_str.isdigit(): + return int(layer_idx_str) < num_hidden_layers return True def load_weights(weights): @@ -71,13 +82,19 @@ def load_model_without_vllm_runner( quantization=quantization, **(model_config_kwargs or {}), ) + model_config.quantization_config = resolve_quantization_config( + model_config.quantization, model_config.quantization_config + ) vllm_config_args = dict(vllm_config_kwargs or {}) vllm_config_args.setdefault( "compilation_config", CompilationConfig(mode=CompilationMode.NONE) ) vllm_config = VllmConfig(model_config=model_config, **vllm_config_args) hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {} - num_hidden_layers = hf_overrides.get("num_hidden_layers") + text_config_overrides = hf_overrides.get("text_config") or {} + num_hidden_layers = text_config_overrides.get( + "num_hidden_layers", hf_overrides.get("num_hidden_layers") + ) with set_current_vllm_config(vllm_config): model_loader = model_loader_cls(vllm_config.load_config) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 158bc90769f3..f20af8a6f826 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,7 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate # shared-expert reduce), returning the partial per-rank sum. The caller is diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 294c538c7b02..6b7c64bdc429 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -122,6 +122,7 @@ def FusedMoEFactory( is_fused_checkpoint_transposed: bool = False, n_shared_experts: int | None = None, fuse_shared_experts: bool = False, + shared_expert_prefix: str | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -197,6 +198,7 @@ def FusedMoEFactory( n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) fuse_shared_experts: Whether to enable shared-expert fusion. + shared_expert_prefix: Checkpoint prefix for the fused shared expert. router_logits_dtype: Data type for router logits buffers gate: Pre-configured gate module shared_experts: Pre-configured shared experts module @@ -360,6 +362,7 @@ def FusedMoEFactory( elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, has_bias=has_bias, is_lora_enabled=vllm_config.lora_config is not None, + shared_expert_prefix=shared_expert_prefix, activation=moe_activation, device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index adc8a2701007..ae5829a311d9 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, resolve_quant_method, @@ -153,6 +154,7 @@ def __init__( self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters self._loaded_expert_biases: set[str] = set() + self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer() self.quant_method = self._get_quant_method( self.layer_name, @@ -630,6 +632,31 @@ def _load_single_value( # _to_scalar's reshape(()) would reject the size-2 weight_shape. param_data[expert_id] = loaded_weight + def _get_fused_shared_expert_quantizer( + self, + ) -> Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] | None: + """Return the quantizer for an online-quantized fused shared expert.""" + online_config = getattr(self.quant_config, "online_quantization_config", None) + prefix = self.moe_config.shared_expert_prefix + if ( + online_config is None + or prefix is None + or self.expert_map_manager.num_fused_shared_experts == 0 + ): + return None + # FSE compatibility checks require every shared-expert projection to + # use the routed experts' weight key. + resolved = online_config.resolve_quant_method_cls( + LinearBase, f"{prefix}.down_proj" + ) + if resolved is None: + return None + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + return ONLINE_SHARED_EXPERT_QUANTIZERS[resolved[3].weight] + @overload def weight_loader( self, @@ -687,6 +714,35 @@ def weight_loader( return False if return_success else None # Hereafter, `expert_id` is local physical id + # A full-precision fused shared expert is quantized into the checkpoint + # layout, then loaded like a pre-quantized expert weight and scale. + if ( + self._fused_shared_expert_quantizer is not None + and self.moe_config.num_logical_experts + <= global_expert_id + < self.moe_config.num_logical_experts + + self.expert_map_manager.num_fused_shared_experts + and weight_name.endswith("_weight") + and loaded_weight.is_floating_point() + and loaded_weight.dtype != param.dtype + ): + weight, weight_scale = self._fused_shared_expert_quantizer( + loaded_weight.to(self.moe_config.device) + ) + stem = "w2" if shard_id == "w2" else "w13" + self.weight_loader( + getattr(self, f"{stem}_weight_scale"), + weight_scale, + f"{weight_name}_scale", + shard_id, + global_expert_id, + return_success=True, + ) + loaded = self.weight_loader( + param, weight, weight_name, shard_id, global_expert_id, True + ) + return loaded if return_success else None + # is_transposed: if the dim to shard the weight # should be flipped. Required by GPTQ/AWQ (K-first format). # should be whatever dimension intermediate_size_per_partition is diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index a98e181caeb2..c724a7c7cb21 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -84,6 +84,40 @@ def resolve_layer_fused_shared_expert( else (True, None) ) is_fused_shared_expert_enabled = fse_requested and fse_compatible + + # In online shared_expert quantization is used, register it to + # `online_quantization_config.quantized_layers` ahead of weight loading. + if is_fused_shared_expert_enabled and quant_config is not None: + online_quantization_config = quant_config.online_quantization_config + if online_quantization_config is not None: + from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, + ) + + for projection_name in ("gate_up_proj", "down_proj"): + projection_prefix = f"{prefix}.{shared_expert_name}.{projection_name}" + quant_method_metadata = ( + online_quantization_config.resolve_quant_method_cls( + LinearBase, projection_prefix + ) + ) + + if quant_method_metadata is None: + continue + + source, quant_key_str, target_pattern, _, quant_method_cls = ( + quant_method_metadata + ) + + if quant_method_cls in (None, UnquantizedLinearMethod): + continue + + online_quantization_config.quantized_layers[projection_prefix] = ( + source.value, + quant_key_str, + target_pattern, + ) if fse_requested and not is_fused_shared_expert_enabled: logger.warning( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 6c1f93473906..3ec2e728e16c 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -320,7 +320,7 @@ def resolve_quant_method( base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) ) online_target = quant_config.online_quantization_config.resolve_quant_method_cls( - layer, prefix + type(layer), prefix ) if checkpoint_is_quantized: diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 3f35b56a31b8..511714837faa 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Any @@ -60,6 +60,7 @@ find_matching_patterns, get_layer_name_after_index, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Static128BlockSym, @@ -104,6 +105,15 @@ class OnlineQuantizationSource(str, Enum): kNvfp4Static: Nvfp4OnlineMoEMethod, } +# Quantizers for full-precision fused shared-expert weights, keyed by the routed +# experts' weight QuantKey. Each shard is quantized to the checkpoint layout on +# arrival, so only block-local schemes are supported. +ONLINE_SHARED_EXPERT_QUANTIZERS: dict[ + QuantKey, Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] +] = { + kMxfp4Static: mxfp4_quantize, +} + def _find_matching_targets( prefix: str, @@ -217,14 +227,14 @@ def _get_method_cls( self, spec: QuantSpec | None, table: dict[QuantKey, type], - layer: torch.nn.Module, + layer_cls: type[torch.nn.Module], ) -> type | None: """Resolve the online method class for a layer's quantization spec. Args: spec: Quantization specification to resolve. table: Mapping from weight quantization keys to method classes. - layer: Layer that will use the resolved method. + layer_cls: Type of the layer that will use the resolved method. Returns: The matching method class, or None when ``spec`` has no weight @@ -239,7 +249,7 @@ def _get_method_cls( cls = table.get(spec.weight) if cls is None: raise ValueError( - f"online quantization for {type(layer).__name__} with " + f"online quantization for {layer_cls.__name__} with " f"weight={spec.weight} is not supported; supported weight " f"keys: {sorted(str(k) for k in table)}" ) @@ -254,12 +264,12 @@ def _get_method_cls( return cls def resolve_quant_method_cls( - self, layer: torch.nn.Module, prefix: str + self, layer_cls: type[torch.nn.Module], prefix: str ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None: """Resolve quantization metadata and method class without instantiating it. Args: - layer: Layer for which to resolve online quantization. + layer_cls: Type of layer for which to resolve online quantization. prefix: Fully qualified layer name. Returns: @@ -271,17 +281,17 @@ def resolve_quant_method_cls( quant_spec: QuantSpec | None if self.args.targets is not None: resolved_pattern = self._resolve_targets_quant_method_metadata( - prefix, layer + prefix, layer_cls ) if resolved_pattern is None: return None source, quant_key_str, target_pattern, quant_spec, table = resolved_pattern else: - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): source = OnlineQuantizationSource.linear quant_spec = self.args.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): source = OnlineQuantizationSource.moe quant_spec = self.args.moe table = _ONLINE_MOE_METHODS @@ -298,14 +308,14 @@ def resolve_quant_method_cls( quant_key_str = str(quant_spec) target_pattern = None - quant_method_cls = self._get_method_cls(quant_spec, table, layer) + quant_method_cls = self._get_method_cls(quant_spec, table, layer_cls) if quant_method_cls is None: return None assert quant_spec is not None return source, quant_key_str, target_pattern, quant_spec, quant_method_cls def _resolve_targets_quant_method_metadata( - self, prefix: str, layer: torch.nn.Module + self, prefix: str, layer_cls: type[torch.nn.Module] ) -> ( tuple[OnlineQuantizationSource, str, str, QuantSpec, dict[QuantKey, type]] | None @@ -314,7 +324,7 @@ def _resolve_targets_quant_method_metadata( Args: prefix: Fully qualified layer name. - layer: Layer matched against configured target patterns. + layer_cls: Type of layer matched against configured target patterns. Returns: A tuple of source, quantization key string, target pattern, spec, @@ -349,22 +359,22 @@ def _resolve_targets_quant_method_metadata( target_pattern = matches[0] quant_key_str = self.args.targets[target_pattern] shorthand = _ONLINE_SHORTHANDS[quant_key_str] - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): quant_spec = shorthand.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): quant_spec = shorthand.moe table = _ONLINE_MOE_METHODS else: raise ValueError( f"Layer {prefix} was matched by quantization_config.targets " f"({target_pattern}), but online quantization is not supported for " - f"{type(layer).__name__}." + f"{layer_cls.__name__}." ) if quant_spec is None: raise ValueError( f"targets pattern {target_pattern} = {quant_key_str} does " - f"not define a QuantSpec for {type(layer).__name__} layers " + f"not define a QuantSpec for {layer_cls.__name__} layers " f"(matched at {prefix})." ) return ( @@ -379,7 +389,7 @@ def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": # `targets` takes precedence over `moe` and `linear` and is exclusive. - resolved = self.resolve_quant_method_cls(layer, prefix) + resolved = self.resolve_quant_method_cls(type(layer), prefix) if resolved is not None: source, quant_key_str, target_pattern, _, quant_method_cls = resolved self.quantized_layers[prefix] = ( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 23a36c8fbc3b..f4662ad5a66e 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -12,8 +12,6 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend - from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey - import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config @@ -40,6 +38,7 @@ ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, amax_for_moe_weight_quant, amax_for_tp_weight_quant, create_fp8_quant_key, @@ -122,6 +121,7 @@ class OnlineLinearBase(LinearMethodBase): weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True + activation_quant_key: QuantKey | None def __init__(self): self.out_dtype = torch.get_default_dtype() @@ -203,6 +203,10 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + activation_quant_key = ( + kFp8DynamicTokenSym if cutlass_fp8_supported() else kFp8DynamicTensorSym + ) + def __init__(self): super().__init__() @@ -211,11 +215,6 @@ def __init__(self): self.use_marlin = False self.marlin_input_dtype = None self.weight_quant_key = kFp8StaticTensorSym - # Use per-token quantization for better perf if dynamic and cutlass - if cutlass_fp8_supported(): - self.activation_quant_key = kFp8DynamicTokenSym - else: - self.activation_quant_key = kFp8DynamicTensorSym def create_weights( self, @@ -237,6 +236,8 @@ def create_weights( **extra_weight_attrs, ) + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, @@ -305,13 +306,11 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-block during loading.""" + activation_quant_key = kFp8Dynamic128Sym + def __init__(self): super().__init__() self.weight_block_size = [128, 128] - self.activation_quant_key = create_fp8_quant_key( - static=False, - group_shape=GroupShape(1, self.weight_block_size[0]), - ) self.weight_quant_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.weight_block_size) ) @@ -337,6 +336,8 @@ def create_weights( ) layer.weight_block_size = self.weight_block_size + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 9cddab85c12a..ecc196e9f185 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -9,6 +9,7 @@ from torch.nn import Module from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey if TYPE_CHECKING: import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -42,6 +43,8 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): FP8 with block-32 scales) during weight loading. """ + activation_quant_key: QuantKey | None = None + def create_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py index eb50ed1a23db..6e294adb3a65 100644 --- a/vllm/model_executor/layers/quantization/utils/config_utils.py +++ b/vllm/model_executor/layers/quantization/utils/config_utils.py @@ -8,9 +8,13 @@ import regex as re +from vllm.logger import init_logger + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +logger = init_logger(__name__) + def get_quark_ocp_mx_group_size( quant_config: "QuantizationConfig | None", @@ -70,19 +74,101 @@ def is_shared_expert_quant_fse_compatible( A compatibility flag and, when incompatible, the reason. """ - if projection_names is None: - projection_names = ["gate_up_proj", "down_proj"] - - if quant_config is None: - return True, None - from vllm.model_executor.layers.quantization.fp8 import Fp8Config + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + if projection_names is None: + projection_names = ["gate_up_proj", "down_proj"] + + if quant_config is None: + return True, None + + online_quant_config = quant_config.online_quantization_config + shared_expert_targets = None + if online_quant_config is not None: + from vllm.model_executor.layers.linear import LinearBase + + online_quant_config.packed_modules_mapping = quant_config.packed_modules_mapping + shared_expert_targets = [ + online_quant_config.resolve_quant_method_cls( + LinearBase, f"{shared_expert_prefix}.{projection_name}" + ) + for projection_name in projection_names + ] + if all(target is None for target in shared_expert_targets): + shared_expert_targets = None + + if shared_expert_targets is not None: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + # NOTE: online shared experts quantization check is only implemented for quark + # quant method at the moment. This can be extended here for other quant methods. + # Requirement is: + # - Retrieve routed weight quant key + # - Retrieve routed activation quant key + # and ensure the online quantization set on shared expert match them. + if isinstance(quant_config, QuarkConfig): + routed_weight_key, routed_activation_key, _ = ( + quant_config.get_quant_method_target(expert_prefix, RoutedExperts) + ) + else: + reason = ( + "online quantization FSE compatibility check is not implemented for " + f"{type(quant_config).__name__}" + ) + logger.warning_once(reason) + return False, reason + + if routed_weight_key not in ONLINE_SHARED_EXPERT_QUANTIZERS: + return ( + False, + "online shared-expert quantization is not supported for routed " + f"expert weight key {routed_weight_key}", + ) + + for shared_expert_target in shared_expert_targets: + if shared_expert_target is None: + return ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_expert_prefix}", + ) + + _, _, _, shared_quant_spec, shared_method_cls = shared_expert_target + if shared_method_cls in (None, UnquantizedLinearMethod): + return ( + False, + "online quantization excludes FSE weights at " + f"{shared_expert_prefix}", + ) + assert issubclass(shared_method_cls, OnlineLinearBase) + + # NOTE: We can not rely on shared_quant_spec.activation_quant_key + # as _ONLINE_SHORTHANDS: dict[str, QuantSpec] in quantization.py + # does not define the default activation quant key. + # `online_quant_config.resolve_quant_method_cls` should return the + # activation quant key as well. Only some online methods define it + # on the class, so compare the weight key first. + if ( + shared_quant_spec.weight != routed_weight_key + or shared_method_cls.activation_quant_key != routed_activation_key + ): + return ( + False, + "online shared-expert quantization keys do not match the " + "routed expert keys", + ) + return True, None + if isinstance(quant_config, DeepseekV4FP8Config): from vllm.config import get_current_vllm_config from vllm.model_executor.models.utils import extract_layer_index diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index da38035db9bd..d8ceb5083c2f 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -199,6 +199,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 70d87529f663..a1a919fae644 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -396,6 +396,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, ) diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 5e5974e209a3..4e541a49413e 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -204,6 +204,7 @@ def __init__( config.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 19fc8c156286..9580894e22ee 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -232,6 +232,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate if self.shared_expert is None else None, diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 778738c13712..26becc9619e7 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -307,6 +307,7 @@ def __init__( if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 9105f97af848..3bf89286034f 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -465,6 +465,7 @@ def __init__( self.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", quant_config=quant_config, prefix=f"{prefix}.experts", ) From f4dde3132c94cccce485a1a59acac1d06739a59d Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 8 Oct 2026 18:18:05 +0200 Subject: [PATCH 221/952] [CI][ROCm] Reuse AITER's prebuilt MoE module in the modular-kernel harness (#60633) Signed-off-by: Stefan Koncarevic --- .../moe/modular_kernel_tools/common.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 9ae6092c944c..5f6349a32880 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from typing import Any @@ -702,6 +702,21 @@ def make_modular_kernel( return modular_kernel +def _shuffle_weights_for_aiter(rank_weights: WeightTensors) -> WeightTensors: + """Pre-shuffle weights so AITER selects its prebuilt `preshuffle_on` module. + + Production shuffles in `process_weights_after_loading`; this harness builds its + tensors directly, so without this every call asks for a `preshuffle_off` module + that is not prebuilt, and JIT-compiles a kernel the image already ships. + """ + from vllm._aiter_ops import rocm_aiter_ops + + w1, w2 = rocm_aiter_ops.shuffle_weights(rank_weights.w1, rank_weights.w2) + w1.is_shuffled = True + w2.is_shuffled = True + return replace(rank_weights, w1=w1, w2=w2) + + def _maybe_convert_weights_for_experts( config: Config, rank_weights: WeightTensors, @@ -715,6 +730,15 @@ def _maybe_convert_weights_for_experts( fe_type = config.fused_experts_type fe_name = getattr(fe_type, "__name__", "") + # AITER's prebuilt modules carry a gfx950 instance table even on gfx942 and reject + # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor + # and per-token schemes. Those keep JIT-compiling until + # https://github.com/ROCm/aiter/issues/5766 is fixed. + if fe_name == "AiterExperts" and ( + config.quant_dtype is None or config.quant_block_shape is not None + ): + return _shuffle_weights_for_aiter(rank_weights) + backend: Fp8MoeBackend | None = None if fe_name == "TrtLlmFp8ExpertsModular": backend = Fp8MoeBackend.FLASHINFER_TRTLLM From 8674a794c788d4c7dac10c6f7ddfdae677588d18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Thu, 8 Oct 2026 18:42:15 +0200 Subject: [PATCH 222/952] [Perf] Enable Kimi K3 shared expert sharding with DeepEP v2 (#56875) Co-authored-by: Misha Goin --- .../models/kimi_k3/test_sequence_parallel.py | 21 ++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 24 +++++++++++++++---- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index 9e8faea099cb..3057f85c7e44 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -301,6 +301,27 @@ def test_shard_sequence_parallel_mlp_gating( ) +@pytest.mark.parametrize( + ("moe_backend", "all2all_backend", "expected"), + [ + ("deep_gemm_mega_moe", "naive", True), + ("auto", "deepep_v2", True), + ("auto", "deepep_low_latency", False), + ], +) +def test_shared_expert_sharding_backend_gate( + moe_backend: str, + all2all_backend: str, + expected: bool, +): + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend=moe_backend), + parallel_config=SimpleNamespace(all2all_backend=all2all_backend), + ) + + assert kimi_model.can_shard_sequence_parallel_shared_expert(vllm_config) is expected + + def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config): """Sharded SP MLP must reproduce the replicated result for every token. diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 84e93b82fbb4..22d2c4a1a9a1 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -156,6 +156,20 @@ def shard_sequence_parallel_mlp( ) +def can_shard_sequence_parallel_shared_expert(vllm_config: VllmConfig) -> bool: + """Whether the active MoE backend preserves the SP shared-expert contract. + + MegaMoE invokes the shared expert directly. DeepEP v2 invokes it through + ``MoERunner``, but its EP combine returns the routed output in the local + sequence-sharded layout, so the shared expert can use the same + all-gather/partial-GEMM/reduce-scatter implementation. + """ + return ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + or vllm_config.parallel_config.all2all_backend == "deepep_v2" + ) + + def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: # Both feature flags may be enabled; the worker's static SP topology binds # its singleton to exactly one mode. @@ -617,10 +631,12 @@ def __init__( quant_config=quant_config, reduce_results=False, use_sequence_parallel=use_sequence_parallel, - # Only the MegaMoE path calls the shared experts directly; the - # FusedMoE path below hands them to the runner, which fuses - # their reduction and assumes the replicated layout. - can_shard_sequence_parallel=self.use_mega_moe, + # MegaMoE calls the shared expert directly. DeepEP v2 hands it + # to MoERunner, but preserves the local SP output layout needed + # by the sharded shared-expert path. + can_shard_sequence_parallel=can_shard_sequence_parallel_shared_expert( + vllm_config + ), run_gemm_rs_ar=run_gemm_rs_ar, prefix=f"{prefix}.shared_experts", activation_situ_beta=activation_situ_beta, From 242e4213fc9845ff6fe607af1aee626fd8acc990 Mon Sep 17 00:00:00 2001 From: Doug Smith Date: Thu, 8 Oct 2026 12:43:39 -0400 Subject: [PATCH 223/952] [Doc] Require GCC 13 for CUDA source builds (#60259) Signed-off-by: Doug Smith --- CMakeLists.txt | 12 ++++++++++++ docs/getting_started/installation/gpu.cuda.inc.md | 9 ++------- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4ed5ff2573c1..b4dd9ada5eb6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1568,6 +1568,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") include(cmake/external_projects/triton_kernels.cmake) endif() +# DeepGEMM and DeepSelect host code uses C++20 , which libstdc++ +# provides only from GCC 13 onward. +if (VLLM_GPU_LANG STREQUAL "CUDA" + AND CMAKE_CXX_COMPILER_ID STREQUAL "GNU" + AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS "13") + message(FATAL_ERROR + "GCC >= 13 is required to build vLLM with CUDA (found ${CMAKE_CXX_COMPILER_VERSION}). " + "The DeepGEMM and DeepSelect dependencies use C++20 in host code, " + "which libstdc++ does not provide before GCC 13. " + "See: https://github.com/vllm-project/vllm/issues/58158") +endif() + # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 86e38c570788..6ad581ac97d1 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -140,13 +140,8 @@ You can find more information about vLLM's wheels in [Install the latest code](# #### Full build (with compilation) {#full-build} !!! note "Compiler requirement" - Building from source requires GCC/G++ ≥ 11.3. PyTorch's C++20 headers are - not compatible with GCC 10 or GCC < 11.3. On Ubuntu 22.04: - ```bash - sudo apt-get install -y gcc-11 g++-11 - sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 \ - --slave /usr/bin/g++ g++ /usr/bin/g++-11 - ``` + Building from source requires GCC/G++ ≥ 13 ([#58158](https://github.com/vllm-project/vllm/issues/58158)). + Ubuntu 24.04, the base of the default vLLM image, ships GCC 13 by default. If you want to modify C++ or CUDA code, you'll need to build vLLM from source. This can take several minutes: From 705e124c05cffe7c00dd19434922c94e0610d674 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 8 Oct 2026 14:00:03 -0400 Subject: [PATCH 224/952] [Bugfix] Followup to #59299: preserve Unicode in structured decision state prompts (#60651) Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Misha Goin --- .../generate/structured_decisions/test_strategies.py | 6 ++++++ vllm/entrypoints/generate/structured_decisions/serving.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index c4f4bce61be1..017788f7b1d8 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -10,6 +10,7 @@ register_structured_decisions_api_router, ) from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.serving import state_text from vllm.entrypoints.generate.structured_decisions.strategies import ( NextTokenStrategy, reply_label_ids, @@ -40,6 +41,11 @@ def test_route_needs_the_flag(): assert ("/v1/systemone" in paths) == enabled +def test_structured_state_keeps_unicode_in_prompt(): + state = {"message": "Français 日本語"} + assert state_text(state) == '{"message": "Français 日本語"}' + + @pytest.fixture(scope="module") def qwen(): from transformers import AutoTokenizer diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 7e2136b8a25e..4921bfcf3574 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -25,7 +25,7 @@ def state_text(state: Any) -> str: - return state if isinstance(state, str) else json.dumps(state) + return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) def parse_questions( From 972195f4f589688fd4cfbbcfb55a5b5502f5af35 Mon Sep 17 00:00:00 2001 From: Tialo <65392801+Tialo@users.noreply.github.com> Date: Thu, 8 Oct 2026 21:19:20 +0300 Subject: [PATCH 225/952] [CI][Build] build macos wheels using abi3 (#60312) Co-authored-by: Misha Goin --- .buildkite/scripts/build-macos-wheel.sh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/build-macos-wheel.sh b/.buildkite/scripts/build-macos-wheel.sh index cb95d92ff14f..d956f71038d3 100755 --- a/.buildkite/scripts/build-macos-wheel.sh +++ b/.buildkite/scripts/build-macos-wheel.sh @@ -25,6 +25,7 @@ export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}" uv venv --python 3.12 uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match -uv build --wheel --no-build-isolation -o artifacts/dist +uv build --wheel --no-build-isolation -o artifacts/dist \ + --config-setting=--build-option=--py-limited-api=cp38 ls -l artifacts/dist/*.whl From 1664c4f44416804dc47e3677e50933e6e9b4f45e Mon Sep 17 00:00:00 2001 From: kyleliang-nv Date: Thu, 8 Oct 2026 11:27:57 -0700 Subject: [PATCH 226/952] [Bugfix][MiniMax-M3] Guard sparse decode sentinel indices (#54597) Signed-off-by: Kyle Liang Signed-off-by: Kyle Liang Co-authored-by: Kyle Liang Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 17 +++- tests/kernels/attention/test_minimax_m3.py | 97 +++++++++++++++++++ .../minimax_m3/common/ops/sparse_attn.py | 24 +++-- 3 files changed, 128 insertions(+), 10 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e0d6c4ce40a6..aa0971b4c55e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -552,6 +552,21 @@ steps: - pytest -v -s kernels/attention/test_deepgemm_attention.py - pytest -v -s quantization/test_cutlass_w4a16.py +- label: ":nvidia: (B200) MiniMax-M3 Attention Kernels" + key: kernels-minimax-m3-attention-b200 + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/" + device: b200-k8s + num_devices: 1 + source_file_dependencies: + - vllm/models/minimax_m3/ + - vllm/config/cache.py + - vllm/v1/kv_cache_interface.py + - tests/v1/attention/utils.py + - tests/kernels/attention/test_minimax_m3.py + commands: + - pytest -v -s tests/kernels/attention/test_minimax_m3.py + - label: ":nvidia: (B200) Kernels Shard %N" key: kernels-b200 timeout_in_minutes: 30 @@ -601,7 +616,6 @@ steps: - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - - tests/kernels/attention/test_minimax_m3.py - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py @@ -647,7 +661,6 @@ steps: pytest -v -s tests/kernels/moe/test_flashinfer.py pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py - pytest -v -s tests/kernels/attention/test_minimax_m3.py ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" >&2 diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 7f0e07309331..7ee900fd3dd1 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -2612,6 +2612,103 @@ def test_decode_sparse_attention_correctness( assert error.max().item() < 1.7e-2 +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True) +@pytest.mark.parametrize("capture_graph", [False, True]) +@pytest.mark.parametrize("num_reqs", [2, 17], ids=["split-k", "single-chunk"]) +@pytest.mark.parametrize( + "invalid_block", [-1, 1_000_000], ids=["sentinel", "out-of-range"] +) +def test_decode_sparse_attention_ignores_invalid_topk( + kv_layout: KVCacheLayout, capture_graph: bool, num_reqs: int, invalid_block: int +): + """Unused top-k entries must not be dereferenced as block-table indices.""" + torch.manual_seed(0) + decode_query_len = 4 + seq_lens = torch.full( + (num_reqs,), decode_query_len, device="cuda", dtype=torch.int32 + ) + seq_lens[1] = 4096 + q = torch.randn( + seq_lens.numel() * decode_query_len, + NUM_Q_HEADS, + HEAD_DIM, + device="cuda", + dtype=DTYPE, + ) + kv_cache = _allocate_main_kv_via_contract(1, kv_layout) + + # All logical blocks alias the single valid page. With a contiguous multi-row + # table, row 1's raw pointer offset -1 resolves to the final element of row + # 0. Poison that exact location so an unsafe load deterministically forms an + # out-of-range KV-cache pointer. + clean_block_table = torch.zeros(num_reqs, 32, device="cuda", dtype=torch.int32) + block_table = clean_block_table.clone() + block_table[0, -1] = 1_000_000 + topk_idx = torch.full( + (NUM_KV_HEADS, q.shape[0], TOPK), + -1, + device="cuda", + dtype=torch.int32, + ) + topk_idx[..., 0] = 0 + # Request 1 has invalid entries both before and after its valid block. + # Two requests produce empty split-K chunks; 17 put the full top-k in one + # chunk, exercising neutral updates before and after a nonempty softmax. + request_topk = topk_idx[:, decode_query_len : 2 * decode_query_len] + request_topk.fill_(invalid_block) + request_topk[..., TOPK // 2] = 0 + + # A valid-only run isolates sentinel handling from BF16-vs-FP32 rounding. + # Include block 0 exactly once; replacing sentinels with 0 would duplicate + # its contribution. The general correctness test covers reference accuracy. + valid_topk = torch.zeros( + (NUM_KV_HEADS, q.shape[0], 1), device="cuda", dtype=torch.int32 + ) + expected = torch.empty_like(q) + minimax_m3_sparse_attn_decode( + q, + kv_cache, + valid_topk, + clean_block_table, + seq_lens, + NUM_KV_HEADS, + SM_SCALE, + expected, + decode_query_len, + ) + actual = torch.empty_like(q) + + def run_decode() -> None: + minimax_m3_sparse_attn_decode( + q, + kv_cache, + topk_idx, + block_table, + seq_lens, + NUM_KV_HEADS, + SM_SCALE, + actual, + decode_query_len, + ) + + if capture_graph: + # An unfixed kernel faults during capture and leaves the CUDA context + # unusable. Run negative-control reproductions in a dedicated process. + run_decode() + current_platform.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_decode() + for _ in range(3): + graph.replay() + else: + run_decode() + current_platform.synchronize() + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_decode_wrong_layout_breaks_parity(): """Negative (AC-3/AC-5): consuming the physical HND buffer as if it were already contiguous-NHD (i.e. skipping the allocator's inverse permute) diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 8dfd6e0c1996..6426115ae356 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -304,7 +304,7 @@ def _gqa_sparse_decode_kernel( # attention range instead of letting padded rows produce negative lengths. kv_len = tl.maximum(query_pos + 1, 0) - # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)). + # Upper bound on top-k entries from causal length; entries may be sentinels. idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K real_topk = tl.minimum(max_topk, num_blocks) @@ -334,10 +334,15 @@ def _gqa_sparse_decode_kernel( for _ in tl.range(chunk_start_topk, chunk_end_topk): blk = tl.load(cur_idx_ptr).to(tl.int32) cur_idx_ptr = cur_idx_ptr + stride_tk - c = blk * BLOCK_SIZE_K - page = tl.load(bt_row + blk).to(tl.int64) + # topk_idx uses -1 for unused entries. Select a safe logical block + # before dereferencing the block table; invalid and future blocks then + # contribute a neutral online-softmax update without control flow. + valid_blk = (blk >= 0) & (blk < num_blocks) + safe_blk = tl.where(valid_blk, blk, 0) + c = safe_blk * BLOCK_SIZE_K + page = tl.load(bt_row + safe_blk).to(tl.int64) pos = c + off_n - pos_mask = pos < kv_len + pos_mask = (pos < kv_len) & valid_blk k = tl.load( kv_cache_ptr + page * stride_kv_blk @@ -364,9 +369,11 @@ def _gqa_sparse_decode_kernel( qk += tl.where(pos_mask[None, :], 0, float("-inf")) qk += tl.dot(q, k) * sm_scale_log2e m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp2(qk - m_ij[:, None]) + m_safe = tl.where(valid_blk, m_ij, 0.0) + p = tl.where(pos_mask[None, :], tl.exp2(qk - m_safe[:, None]), 0.0) l_ij = tl.sum(p, axis=1) - acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + alpha = tl.where(valid_blk, tl.exp2(m_i - m_safe), 1.0) + acc_o = acc_o * alpha[:, None] v = tl.load( kv_cache_ptr + page * stride_kv_blk @@ -390,8 +397,9 @@ def _gqa_sparse_decode_kernel( ) v = (v * v_scale[:, None]).to(q.dtype) acc_o += tl.dot(p.to(v.dtype), v) - m_i = m_ij - lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_ij = m_safe + tl.log2(tl.exp2(lse_i - m_safe) + l_ij) + m_i = tl.where(valid_blk, m_safe, m_i) + lse_i = tl.where(valid_blk, lse_ij, lse_i) if USE_PDL: tl.extra.cuda.gdc_launch_dependents() From d8fa463e2ba751759a3534dd3ce11af92f1eed83 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:33:48 -0400 Subject: [PATCH 227/952] [Bugfix][Prefix caching] Fix inconsistent hybrid Mamba prefix cache boundaries (#60533) Signed-off-by: Wei Zhao Signed-off-by: zjy0516 Co-authored-by: Wei Zhao Co-authored-by: Codex Co-authored-by: zjy0516 Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../models/kimi_k3/test_amd_kda_checkpoint.py | 7 +- tests/models/kimi_k3/test_kda_metadata.py | 60 ++++- tests/models/kimi_k3/test_prefix_cache.py | 40 +++ .../v1/attention/test_gdn_metadata_builder.py | 8 +- .../test_mamba2_prefill_checkpoints.py | 9 +- .../test_partial_prefix_cache_hits.py | 253 ++++++++++++++---- tests/v1/core/test_prefix_caching.py | 38 +-- vllm/config/cache.py | 11 + .../model_executor/layers/mamba/checkpoint.py | 11 +- vllm/v1/core/block_pool.py | 7 +- vllm/v1/core/kv_cache_coordinator.py | 69 ++--- vllm/v1/core/kv_cache_utils.py | 108 ++++++-- vllm/v1/core/sched/scheduler.py | 35 ++- vllm/v1/core/single_type_kv_cache_manager.py | 60 +++-- vllm/v1/engine/core.py | 14 + vllm/v1/kv_cache_interface.py | 9 +- vllm/v1/worker/gpu_worker.py | 9 +- 17 files changed, 547 insertions(+), 201 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 0a622bfa0450..6f3129a224e5 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -75,7 +75,7 @@ def _builder( vllm_config.cache_config.prefix_match_unit = prefix_match_unit if num_spec: vllm_config.speculative_config = _StubSpeculativeConfig(num_spec) - return KimiK3ROCmKDAMetadataBuilder( + builder = KimiK3ROCmKDAMetadataBuilder( kv_cache_spec=MambaSpec( block_size=block_size, shapes=((16, 64),), @@ -91,6 +91,11 @@ def _builder( vllm_config=vllm_config, device=torch.device("cuda"), ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + prefix_match_unit or block_size + ) + return builder def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw): diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 74414279552c..63024634d305 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -100,6 +100,8 @@ def _make_builder( prefix_match_unit: int | None = None, use_eagle: bool = False, disable_eagle_block_drop: bool = False, + cache_hit_alignment_tokens: int | None = None, + hash_block_size: int | None = None, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -122,6 +124,11 @@ def _make_builder( vllm_config.cache_config.use_replayssm = use_recoverssm vllm_config.cache_config.use_kda_recoverssm = use_recoverssm vllm_config.cache_config.prefix_match_unit = prefix_match_unit + hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size + vllm_config.cache_config.hash_block_size = hash_block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + cache_hit_alignment_tokens or hash_block_size + ) builder = builder_cls( kv_cache_spec=MambaSpec( block_size=mamba_block_size, @@ -220,6 +227,8 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): vllm_config=builder_config, device=DEVICE, ) + builder_config.cache_config.hash_block_size = BLOCK_SIZE + builder_config.cache_config.cache_hit_alignment_tokens = BLOCK_SIZE # An all-prefill speculative batch used to leave an all-false spec mask # alive, then access active_non_spec_mask_cpu before it was initialized. @@ -287,18 +296,45 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): +@pytest.mark.parametrize( + "batch,prefix_match_unit,hit_alignment,expected_offsets,expected_states", + [ + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + None, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + 16, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + (BatchSpec(seq_lens=[530], query_lens=[530]), None, 128, [512], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), None, 128, [32], [32]), + (BatchSpec(seq_lens=[530], query_lens=[530]), 16, 16, [528], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), 16, 16, [48], [32]), + ], +) +def test_internal_checkpoint_metadata_targets_last_aligned_boundary( + batch, prefix_match_unit, hit_alignment, expected_offsets, expected_states +): device = torch.device("cuda") - batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True - ).replace(is_prefilling=torch.tensor([True, True])) + ).replace(is_prefilling=torch.tensor([True] * len(batch.seq_lens))) builder = _make_builder( KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=prefix_match_unit, + cache_hit_alignment_tokens=hit_alignment, + hash_block_size=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) @@ -313,26 +349,33 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): assert actual.checkpoint is not None torch.testing.assert_close( actual.checkpoint.state_indices, - torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device), + torch.tensor(expected_states, dtype=torch.int32, device=device), ) torch.testing.assert_close( actual.checkpoint.checkpoint_offsets, - torch.tensor([48, 0], dtype=torch.int32, device=device), + torch.tensor(expected_offsets, dtype=torch.int32, device=device), ) @pytest.mark.parametrize( - ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), - [(False, 16, 80), (True, 16, 96), (False, 8, None)], + ("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), + [ + (100, False, 16, 80), + (100, True, 16, 96), + (100, False, 8, None), + (96, False, 16, 80), + (96, True, 16, 80), + ], ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_spec_internal_checkpoint_metadata_targets_replay_boundary( + seq_len: int, disable_eagle_block_drop: bool, prefix_match_unit: int, expected_offset: int | None, ) -> None: device = torch.device("cuda") - batch = BatchSpec(seq_lens=[100], query_lens=[100]) + batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True ) @@ -390,6 +433,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset(): full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 3525550fe970..20e5d6670e5c 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -8,6 +8,7 @@ from collections.abc import Iterator from dataclasses import dataclass from math import lcm +from uuid import uuid4 import pytest import requests @@ -353,6 +354,45 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) +@pytest.mark.parametrize( + "prompt_len", + [ + # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: + # ensure lookup reads the cached attention proof at 7040 despite + # the final hit limit of 7039, then applies EAGLE's 128-token drop. + pytest.param(7040, id="eagle-proof-scan"), + # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: + # ensure the Mamba checkpoint is 7168, not 7040 from applying both + # the P - 1 cap and EAGLE's 128-token drop. + pytest.param(7296, id="eagle-double-cap"), + ], +) +def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: + """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" + mode = Mode(128, True) + salt = str(uuid4()) + rng = random.Random(0) + prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)] + deployment = DEPLOYMENTS["plain"] + with _serve(deployment, mode) as servers: + url = servers[0].url_root + cold = _complete(url, prompt, salt, max_tokens=1) + resend = _complete(url, prompt, salt, max_tokens=1) + recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) + + # EAGLE's dropped hash unit already leaves tokens to recompute on resend. + expected = (prompt_len // 128 - 1) * 128 + actual = _cached(resend) + assert _cached(cold) == _cached(recompute) == 0 + assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}" + check_logprobs_close( + outputs_0_lst=[_tokens_text_logprobs(recompute)], + outputs_1_lst=[_tokens_text_logprobs(resend)], + name_0="recompute", + name_1="resend", + ) + + @pytest.mark.parametrize( "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] ) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index cd01fba49e71..4b52eb9c0950 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -385,7 +385,7 @@ def test_cudagraph_capture_batch_stays_decode_only(): def _create_checkpoint_builder_and_batch( - num_spec: int, prefix_match_unit: int = 16 + num_spec: int, prefix_match_unit: int | None = 16 ) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -413,6 +413,8 @@ def _create_checkpoint_builder_and_batch( vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or 64 + vllm_config.cache_config.cache_hit_alignment_tokens = prefix_match_unit or 64 batch = BatchSpec( seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99] ) @@ -431,7 +433,9 @@ def _create_checkpoint_builder_and_batch( @pytest.mark.parametrize("num_spec", [0, 3]) -@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)]) +@pytest.mark.parametrize( + "prefix_match_unit, expected_offset", [(None, 64), (16, 96), (8, 96)] +) def test_checkpoint_metadata_preserves_non_spec_order( num_spec, prefix_match_unit, expected_offset ): diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py index e5276786e23c..a3968b5a0bc8 100644 --- a/tests/v1/attention/test_mamba2_prefill_checkpoints.py +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -77,7 +77,9 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks: int = 1, prefill_checkpoint_alignment: int | None = 1, ) -> Mamba2AttentionMetadataBuilder: - vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE) + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=MAMBA_BLOCK_SIZE + ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode # get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`. vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE @@ -89,12 +91,15 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, prefill_checkpoint_alignment=prefill_checkpoint_alignment, ) - return Mamba2AttentionMetadataBuilder( + builder = Mamba2AttentionMetadataBuilder( kv_cache_spec=spec, layer_names=["layer.0"], vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = MAMBA_BLOCK_SIZE + vllm_config.cache_config.cache_hit_alignment_tokens = MAMBA_BLOCK_SIZE + return builder def _build(builder, seq_lens, query_lens): diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 76beffce452a..51317deb8d01 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -22,6 +22,7 @@ get_block_hash, get_group_id, init_none_hash, + resolve_cache_hit_alignment_tokens, ) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( @@ -148,29 +149,84 @@ def make_full_mamba_manager( ) -def test_dcp_fine_hit_retention_uses_hash_alignment_without_eagle(): - """DCP must not discard a reusable Mamba state at hash alignment.""" - hash_block_size = 2 +@pytest.mark.parametrize( + "dcp_world_size,hash_block_size,block_size,prompt_tokens,checkpoint_blocks", + [ + (4, 2, 2, 7, 0), + (8, 16, 16, 512, 1), + (8, 16, 16, 513, 1), + *[(8, 128, 896, p, 1) for p in (127, 128, 129, 7040, 7168, 7296, 7449)], + ], +) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("extension_tokens", [0, 32]) +def test_dcp_fine_hit_retention_uses_hash_alignment( + dcp_world_size, + hash_block_size, + block_size, + prompt_tokens, + checkpoint_blocks, + use_eagle, + extension_tokens, +): + """Resends and extensions reuse the retained hash-aligned checkpoint.""" manager = make_full_mamba_manager( - dcp_world_size=4, + dcp_world_size=dcp_world_size, hash_block_size=hash_block_size, - full_block_size=hash_block_size, - mamba_block_size=hash_block_size, - num_blocks=64, + full_block_size=block_size, + mamba_block_size=block_size, + num_blocks=128, + use_eagle=use_eagle, + num_speculative_blocks=4 if use_eagle else 0, + num_prefill_checkpoint_blocks=checkpoint_blocks, ) manager.coordinator.retention_interval = 0 - - token_ids = list(range(7)) - producer = make_request("producer", token_ids, hash_block_size, sha256) - computed_blocks, num_computed, _ = manager.get_computed_blocks(producer) - assert num_computed == 0 - assert manager.allocate_slots(producer, 6, 0, computed_blocks) is not None + assert manager.coordinator.enable_partial_hash_hits + scheduler = SimpleNamespace( + block_size=block_size * dcp_world_size, + cache_config=SimpleNamespace(block_size=block_size), + hash_block_size=hash_block_size, + max_num_scheduled_tokens=8192, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + use_eagle_block_drop=use_eagle, + mamba_partial_cache_hit=True, + mamba_shared_prefix_checkpoint=False, + mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks), + mamba_prefill_checkpoint_alignment=16, + ) + producer = make_request( + "producer", list(range(prompt_tokens)), hash_block_size, sha256 + ) + while producer.num_computed_tokens < prompt_tokens: + scheduled = Scheduler._mamba_block_aligned_split( + scheduler, producer, prompt_tokens - producer.num_computed_tokens + ) + assert scheduled > 0 + manager.new_step_starts() + assert manager.allocate_slots(producer, scheduled) is not None + producer.num_computed_tokens += scheduled manager.free(producer) manager.new_step_starts() - consumer = make_request("consumer", token_ids, hash_block_size, sha256) - _, num_computed, _ = manager.get_computed_blocks(consumer) - assert num_computed == 6 + replay = make_request( + "replay", list(range(prompt_tokens + extension_tokens)), hash_block_size, sha256 + ) + _, hit, _ = manager.get_computed_blocks(replay) + expected = ( + max((prompt_tokens // hash_block_size - 1) * hash_block_size, 0) + if use_eagle + else ((prompt_tokens - 1) // hash_block_size * hash_block_size) + ) + assert hit == expected + cached_states = { + end + for end in range(hash_block_size, prompt_tokens + 1, hash_block_size) + if manager.block_pool.get_cached_block( + producer.block_hashes[end // hash_block_size - 1], [1] + ) + is not None + } + assert cached_states == ({expected} if expected else set()) @pytest.mark.parametrize("dcp_world_size", [1, 4]) @@ -339,7 +395,7 @@ def test_hybrid_mamba_align_partial_hash_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks) @@ -488,7 +544,7 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -538,7 +594,7 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks(): num_prefill_checkpoint_blocks=1, ) - owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + owner = make_request("owner", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(owner) assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None manager.free(owner) @@ -619,44 +675,70 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 -@pytest.mark.parametrize( - "retention_interval,transient_published", - [(None, True), (0, False)], -) +@pytest.mark.parametrize("retention_interval", [None, 0, 32, 64, 256]) +@pytest.mark.parametrize("hash_block_size", [16, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) def test_internal_checkpoint_publication_respects_retention( - retention_interval, transient_published + retention_interval, hash_block_size, use_eagle ): - """With retention_interval=0, a mid-prompt internal checkpoint is - request-local and must not enter the prefix cache; the prompt-end - checkpoint keeps the existing publication behavior. Dense retention - (None) is unchanged. - """ - hash_block_size = 16 + """Publish only retained block boundaries and the final prompt checkpoint.""" manager = make_full_mamba_manager( dcp_world_size=1, hash_block_size=hash_block_size, full_block_size=hash_block_size, mamba_block_size=32, num_prefill_checkpoint_blocks=1, + use_eagle=use_eagle, ) manager.coordinator.retention_interval = retention_interval request = make_request("producer", list(range(240)), hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(request) - # Mid-prompt chunk: exports a transient checkpoint at 112. + # The EAGLE adjustment can make a chunk checkpoint block-aligned or partial. assert manager.allocate_slots(request, 128, num_computed, computed_blocks) - transient_hash = request.block_hashes[112 // hash_block_size - 1] + checkpoint = (128 if use_eagle else 127) // hash_block_size * hash_block_size - ( + hash_block_size if use_eagle else 0 + ) + retained_checkpoint = checkpoint % 32 == 0 and ( + retention_interval is None + or (retention_interval > 0 and checkpoint % retention_interval == 0) + ) + transient_hash = request.block_hashes[checkpoint // hash_block_size - 1] transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) - assert (transient_hit is not None) == transient_published + assert (transient_hit is not None) == retained_checkpoint + offloads = drain_boundary_state_offloads(manager).get(request.request_id, []) + assert any(position == checkpoint for _, _, position in offloads) == ( + retained_checkpoint + ) + assert all( + position % 32 == 0 + and ( + retention_interval is None + or (retention_interval > 0 and position % retention_interval == 0) + ) + for _, _, position in offloads + ) + if retention_interval in (None, 32, 64): + retained_hash = request.block_hashes[128 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(retained_hash, [1]) is not None request.num_computed_tokens = 128 manager.new_step_starts() - # Final chunk: the prompt-end checkpoint at 224 stays published. + # The final checkpoint is reusable regardless of the retention interval. assert manager.allocate_slots(request, 112) is not None - end_hash = request.block_hashes[224 // hash_block_size - 1] + prompt_checkpoint = ( + 240 if use_eagle else 239 + ) // hash_block_size * hash_block_size - (hash_block_size if use_eagle else 0) + end_hash = request.block_hashes[prompt_checkpoint // hash_block_size - 1] assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + assert any( + position == prompt_checkpoint + for _, _, position in drain_boundary_state_offloads(manager).get( + request.request_id, [] + ) + ) def test_transient_checkpoint_evicts_retained_boundary_hash(): @@ -802,9 +884,9 @@ def test_hash_aligned_query_end_uses_regular_partial_tail(): mamba_block_size=mamba_block_size, num_prefill_checkpoint_blocks=1, ) - request = make_request("producer", list(range(14)), hash_block_size, sha256) + request = make_request("producer", list(range(15)), hash_block_size, sha256) - new_blocks = manager.allocate_slots(request, request.num_tokens) + new_blocks = manager.allocate_slots(request, 14) assert new_blocks is not None mamba_manager = manager.coordinator.single_type_managers[1] @@ -928,7 +1010,7 @@ def test_boundary_state_offloads_returns_cow_target(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1004,7 +1086,7 @@ def test_finished_partial_tail_uses_table_source_once(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None source_block_id = manager.get_blocks("0").get_block_ids()[1][1] @@ -1024,7 +1106,7 @@ def test_finished_partial_tail_uses_table_source_once(): enable_caching=True, hash_block_size=hash_block_size, ) - req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1) assert ( advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks) @@ -1212,7 +1294,7 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1267,7 +1349,7 @@ def test_truncate_computed_blocks_allows_short_mamba_group_only(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1331,7 +1413,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1411,7 +1493,7 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1470,7 +1552,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1544,7 +1626,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1700,7 +1782,7 @@ def test_cow_retained_blocks_returned_for_release(): enable_caching=True, hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -2133,6 +2215,72 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed +@pytest.mark.parametrize("with_sliding_window", [False, True]) +def test_resolved_cache_hit_alignment_matches_coordinator(with_sliding_window): + """Workers place Mamba checkpoints on the engine-resolved alignment, so it + must be where the coordinator lands hits, including when a group without + fine-grained lookups disables partial hits.""" + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + if with_sliding_window: + kv_cache_groups.append( + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=2 * hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=2 * hash_block_size, + ), + ) + ) + kv_cache_config = KVCacheConfig( + num_blocks=64, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + ) + coordinator = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ).coordinator + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(enable_prefix_caching=True), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + assert coordinator.enable_partial_hash_hits != with_sliding_window + assert ( + resolve_cache_hit_alignment_tokens( + kv_cache_config, + vllm_config, + coordinator.scheduler_block_size, + hash_block_size, + ) + == coordinator._cache_hit_alignment_tokens + ) + + def test_opted_out_scratch_group_keeps_partial_hash_hits(): hash_block_size = 2 mamba_block_size = 2 * hash_block_size @@ -2179,7 +2327,7 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2240,7 +2388,7 @@ def test_kpool_tail_supports_128_token_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits shared_prefix = [10] * (12 * hash_block_size) - req0 = make_request("0", shared_prefix, hash_block_size, sha256) + req0 = make_request("0", shared_prefix + [11], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert ( manager.allocate_slots(req0, len(shared_prefix), num_computed, computed_blocks) @@ -2266,7 +2414,7 @@ def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): ) assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2, 3], 2, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2318,7 +2466,7 @@ def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( assert manager.coordinator.single_type_managers[1].block_size == block_size prefix = list(range(12)) - req0 = make_request("snapshot-owner", prefix, block_size, sha256) + req0 = make_request("snapshot-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None manager.free(req0) @@ -2355,7 +2503,7 @@ def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): mamba_block_size=block_size, ) prefix = list(range(12)) - req0 = make_request("joint-owner", prefix, block_size, sha256) + req0 = make_request("joint-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None manager.new_step_starts() @@ -2702,6 +2850,7 @@ def test_checkpoint_reservation_matches_worker( mamba_block_size=mamba_block_size, checkpoint_alignment=16, drop_eagle_block=use_eagle, + cache_hit_alignment_tokens=mamba_manager.cache_hit_alignment_tokens, ) reserved = mamba_manager._checkpoints.get(request.request_id) if reserved is not None: diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 427201a7ceef..a89fd085f1fd 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -5047,19 +5047,10 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): - """An identical resend and a longer sibling resume at DIFFERENT positions. - - How far a lookup matches depends on who is asking. A resend of the same - prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed - for logits), while a sibling whose prompt merely starts with this one caps - above the prompt. The two coincide unless the prompt length is an exact - multiple of the alignment -- there they differ by one alignment unit, and - under the EAGLE drop BOTH are reachable. - - Retaining only the higher one leaves the resend with every retained state - above every candidate its lookup can produce, and the reconciled hit - collapses to 0 -- the same zero-hit failure sparse retention already fixes - at unaligned prompt lengths. + """EAGLE resend and extension reuse the same retained checkpoint. + + Attention may inspect the prompt's final block before dropping it, so an + aligned resend must not lose another block to the P - 1 lookup limit. """ block_size = 32 num_spec = 3 @@ -5097,16 +5088,12 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): use_eagle=True, ) - # 128 tokens, an exact multiple of the 32-token alignment. A longer sibling - # matches 128 and drops to 96; this prompt's own resend caps at 127, matches - # 96 and drops to 64. Both states must survive retention. + # Both requests prove 128 attention tokens, drop one block, and reuse 96. token_ids = [i for i in range(4) for _ in range(block_size)] req0 = make_request("0", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - # Prefill in block-aligned chunks the way the align-mode scheduler does: a - # state only materializes as a chunk's running-state block, so a - # single-shot prefill could not retain the lower one. + # Materialize the checkpoint through normal block-aligned prefill chunks. for chunk_end in (32, 64, 96, 128): blocks = manager.allocate_slots( req0, @@ -5118,8 +5105,7 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert blocks is not None req0.num_computed_tokens = chunk_end - # Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are - # mamba blocks 1 and 2. + # Both materialized replay checkpoints survive; resend and extension use 96. pool = manager.block_pool expected_mamba_cached = {1, 2} for i in range(4): @@ -5130,15 +5116,13 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert cached is None, f"mamba hash {i} should not be cached" manager.free(req0) - # The identical resend: full attention matches blocks 0-2 (96 tokens, capped - # by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without - # the lower state retained the reconciled hit would be 0. + # The identical resend must reuse the retained checkpoint, not miss to zero. req1 = make_request("1", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) - assert num_computed_tokens == 2 * block_size - assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2] + assert num_computed_tokens == 3 * block_size + assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3] - # The longer sibling resumes one alignment unit higher, off the same prompt. + # The longer sibling reuses that same checkpoint. longer = make_request("2", token_ids + [9] * block_size, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer) assert num_computed_tokens == 3 * block_size diff --git a/vllm/config/cache.py b/vllm/config/cache.py index e6436f98f9cb..5eb26726c718 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -230,6 +230,15 @@ def device_memory_utilization(self, value: float) -> None: for hybrid models where requests occupy multiple KV cache groups.""" kv_cache_max_concurrency: float | None = field(default=None, init=False) """Per-DP-engine maximum concurrency at max_model_len tokens.""" + # The two differ when hashes are finer than the scheduler block but partial + # hash hits are off (e.g. a sliding-window group cannot look up sub-block + # hits), so hits only land on the scheduler block. + # TODO(zjy0516): Merge into one value + hash_block_size: int | None = field(default=None, init=False) + """Tokens per prefix-cache block hash.""" + cache_hit_alignment_tokens: int | None = field(default=None, init=False) + """Token granularity at which prefix-cache hits land: `hash_block_size` + with partial hash hits, otherwise the scheduler block size.""" kv_sharing_fast_prefill: bool = False """In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254), @@ -298,6 +307,8 @@ def compute_hash(self) -> str: "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", + "hash_block_size", + "cache_hit_alignment_tokens", # Feature toggles not impacting compiled graph shape "kv_sharing_fast_prefill", "swa_bounded_replay", diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index 45888113e152..2f3122f87058 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -25,6 +25,7 @@ def compute_mamba_prefill_checkpoints( mamba_block_size: int, checkpoint_alignment: int | None, drop_eagle_block: bool, + cache_hit_alignment_tokens: int, ) -> tuple[list[int], list[int]]: """Per-row internal prefill checkpoint offsets and cache block columns. @@ -43,7 +44,7 @@ def compute_mamba_prefill_checkpoints( for seq_len, query_len in zip(seq_lens, query_lens): query_start = seq_len - query_len position = get_mamba_prefill_checkpoint_position( - seq_len, hash_block_size, drop_eagle_block=drop_eagle_block + seq_len, cache_hit_alignment_tokens, drop_eagle_block=drop_eagle_block ) valid = is_mamba_prefill_checkpoint_valid( query_start=query_start, @@ -106,12 +107,17 @@ def build( return None if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0: return None + cache_config = self.vllm_config.cache_config + hash_block_size = cache_config.hash_block_size + cache_hit_alignment_tokens = cache_config.cache_hit_alignment_tokens + if hash_block_size is None or cache_hit_alignment_tokens is None: + # Unresolved only for the minimal KV cache built for profiling. + return None assert m.seq_lens_cpu_upper_bound is not None all_query_lens = m.query_start_loc_cpu.diff().tolist() query_lens = [all_query_lens[row] for row in request_rows] seq_lens = m.seq_lens_cpu_upper_bound.tolist() block_size = self.kv_cache_spec.block_size - hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size speculative_config = self.vllm_config.speculative_config drop_eagle_block = ( speculative_config is not None and speculative_config.use_eagle_block_drop() @@ -123,6 +129,7 @@ def build( mamba_block_size=block_size, checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment, drop_eagle_block=drop_eagle_block, + cache_hit_alignment_tokens=cache_hit_alignment_tokens, ) if not any(checkpoint_offsets): return None diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 75bdb40c24be..3b0102fe84aa 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -479,8 +479,8 @@ def cache_partial_block( kv_cache_group_id: KV cache group that owns the partial entry. block_size: Cache block size for the owning group. The partial entry hash itself is always the prefix-chain hash at - ``num_tokens``; ``block_size`` is used to assert that the - entry is partial within the owning cache block. + ``num_tokens``; ``block_size`` must be a multiple of the hash + block size. replace_existing_hashes: Whether the block contents were replaced and all existing cache entries must be removed before the new entry is registered. @@ -494,9 +494,6 @@ def cache_partial_block( return None assert block_size % self.hash_block_size == 0 - assert replace_existing_hashes or ( - block_size > self.hash_block_size and num_tokens % block_size != 0 - ) block_hash = self._get_partial_block_hash(request, num_tokens) num_hash_blocks = num_tokens // self.hash_block_size block_hash_with_group_id = make_block_hash_with_group_id( diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index deb7f89967b2..c32173a63554 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -4,13 +4,14 @@ from collections.abc import Sequence from typing import NamedTuple -from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + partial_hash_hits_enabled, ) from vllm.v1.core.single_type_kv_cache_manager import ( CrossAttentionManager, @@ -28,8 +29,6 @@ ) from vllm.v1.request import Request -logger = init_logger(__name__) - def _validate_prefix_cache_retention_interval( retention_interval: int | None, @@ -695,37 +694,15 @@ def __init__( # managers in every group. TP needs hashing finer than the Mamba block; # DCP accepts equality because it scales the effective full-attention # block instead. - has_partial_mamba_group = any( - isinstance(g.kv_cache_spec, MambaSpec) - and g.kv_cache_spec.mamba_cache_mode == "align" - and ( - (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) - or ( - dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size - ) - ) - for g in kv_cache_config.kv_cache_groups - ) self.enable_partial_hash_hits = ( - allow_partial_hash_hits and has_partial_mamba_group + allow_partial_hash_hits + and partial_hash_hits_enabled( + kv_cache_config.kv_cache_groups, + hash_block_size, + dcp_world_size, + [type(manager) for manager in self.single_type_managers], + ) ) - if self.enable_partial_hash_hits: - unsupported_partial_hit_managers = { - type(manager).__name__ - for manager, group in zip( - self.single_type_managers, kv_cache_config.kv_cache_groups - ) - if group.kv_cache_spec.prefix_cacheable - and not manager.supports_fine_grained_hash_lookup - and manager.block_size != hash_block_size - } - if unsupported_partial_hit_managers: - self.enable_partial_hash_hits = False - logger.warning_once( - "Disabling fine-grained prefix-cache hits because these KV " - "cache managers require block-aligned lookups: %s.", - ", ".join(sorted(unsupported_partial_hit_managers)), - ) cache_hit_alignment_tokens = self._cache_hit_alignment_tokens for manager in self.single_type_managers: manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens @@ -899,15 +876,15 @@ def find_longest_cache_hit( # mamba: its finder never drops (draft models have no mamba # layers), so the hit would grow past the candidate. if drop_eagle_block and not isinstance(spec, MambaSpec): - eagle_margin = ( - self.hash_block_size - if self.enable_partial_hash_hits - and manager_cls.supports_fine_grained_hash_lookup - and group_block_size > self.hash_block_size - else group_block_size + eagle_margin = eagle_proof_margin( + group_block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, ) _max_length = min( - curr_hit_length + eagle_margin, max_cache_hit_length + curr_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, ) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, @@ -983,9 +960,21 @@ def find_longest_cache_hit_per_group( for spec, group_ids, manager_cls, use_eagle in self.attention_groups: manager = self.single_type_managers[group_ids[0]] + lookup_length = max_cache_hit_length + if use_eagle and not isinstance(spec, MambaSpec): + eagle_margin = eagle_proof_margin( + manager.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, + ) + lookup_length = min( + max_cache_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, + ) blocks, group_hit = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, - max_length=max_cache_hit_length, + max_length=lookup_length, kv_cache_group_ids=group_ids, block_pool=manager.block_pool, kv_cache_spec=spec, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 79ec2fc5e4ed..53a01a5cabe2 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -53,6 +53,7 @@ if TYPE_CHECKING: from vllm.v1.core.block_pool import BlockPool + from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager # BlockHash represents the hash of a single KV-cache block used for @@ -747,6 +748,60 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache return replace(spec, block_size=block_size) +def partial_hash_hits_enabled( + kv_cache_groups: Sequence[KVCacheGroupSpec], + hash_block_size: int, + dcp_world_size: int = 1, + manager_classes: Sequence[type["SingleTypeKVCacheManager"]] | None = None, +) -> bool: + """Whether aligned Mamba states support sub-block prefix-cache hits. + + If ``manager_classes`` (one per group) is given, also require every other + prefix-cacheable group to support fine-grained lookups. + """ + if not any( + isinstance(spec, MambaSpec) + and spec.mamba_cache_mode == "align" + and ( + (dcp_world_size == 1 and spec.block_size > hash_block_size) + or (dcp_world_size > 1 and spec.block_size >= hash_block_size) + ) + for group in kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ): + return False + if manager_classes is None: + return True + unsupported_managers = { + manager_cls.__name__ + for group, manager_cls in zip(kv_cache_groups, manager_classes) + if group.kv_cache_spec.prefix_cacheable + and not manager_cls.supports_fine_grained_hash_lookup + and resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size) + != hash_block_size + } + if unsupported_managers: + logger.warning_once( + "Disabling fine-grained prefix-cache hits because these KV " + "cache managers require block-aligned lookups: %s.", + ", ".join(sorted(unsupported_managers)), + ) + return False + return True + + +def eagle_proof_margin( + block_size: int, hash_block_size: int, fine_grained_lookup: bool +) -> int: + """Tokens an EAGLE group matches past a cache hit before dropping them. + + Fine-grained lookups drop one hash unit; others drop one cache block. + """ + if fine_grained_lookup and block_size > hash_block_size: + return hash_block_size + return block_size + + def resolve_kv_cache_block_sizes( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, @@ -775,8 +830,7 @@ def resolve_kv_cache_block_sizes( if groups and not groups[0].kv_cache_spec.dcp_sharded: dcp = 1 bs = cache_config.block_size * dcp - # The Mamba prefill checkpoint builder reads prefix_match_unit directly, - # so a value dropped here puts its checkpoints off the scheduler's grid. + # Reject a prefix_match_unit that would otherwise be silently ignored. if ( cache_config.prefix_match_unit not in (None, bs) and len(groups) == 1 @@ -797,8 +851,8 @@ def resolve_kv_cache_block_sizes( ] scheduler_block_size = math.lcm(*group_block_sizes) - logger.info("kv cache group sizes %s", group_block_sizes) - logger.info("kv lcm block sizes %s", scheduler_block_size) + logger.info_once("kv cache group sizes %s", tuple(group_block_sizes)) + logger.info_once("kv lcm block sizes %s", scheduler_block_size) # Block hashes are only consumed by prefix caching and KV connectors # (P/D, offloading); when neither is active, keep hash_block_size equal @@ -839,18 +893,10 @@ def resolve_kv_cache_block_sizes( and isinstance(spec.tokens_per_state, int) and spec.tokens_per_state > 1 } - has_partial_mamba_group = any( - isinstance(spec, MambaSpec) - and spec.mamba_cache_mode == "align" - and ( - (dcp == 1 and block_size > hash_block_size) - or (dcp > 1 and block_size >= hash_block_size) - ) - for group, block_size in zip(groups, group_block_sizes) - for spec in iter_layer_specs(group.kv_cache_spec) - ) cache_hit_alignment = ( - hash_block_size if has_partial_mamba_group else scheduler_block_size + hash_block_size + if partial_hash_hits_enabled(groups, hash_block_size, dcp) + else scheduler_block_size ) if any(cache_hit_alignment % alignment for alignment in prefix_alignments): raise ValueError( @@ -861,6 +907,38 @@ def resolve_kv_cache_block_sizes( return scheduler_block_size, hash_block_size +def resolve_cache_hit_alignment_tokens( + kv_cache_config: KVCacheConfig, + vllm_config: VllmConfig, + scheduler_block_size: int, + hash_block_size: int, +) -> int: + """Token granularity at which prefix-cache hits land. + + Mirrors the alignment the hybrid KV cache coordinator pushes to its + managers: fine-grained hits land on ``hash_block_size``, all others on + ``scheduler_block_size``. + """ + groups = kv_cache_config.kv_cache_groups + if not vllm_config.cache_config.enable_prefix_caching or len(groups) <= 1: + return scheduler_block_size + manager_classes = [] + for group in groups: + manager_cls = KVCacheSpecRegistry.get_manager_class( + group.kv_cache_spec, group.role + ) + assert manager_cls is not None + manager_classes.append(manager_cls) + if partial_hash_hits_enabled( + groups, + hash_block_size, + vllm_config.parallel_config.decode_context_parallel_size, + manager_classes, + ): + return hash_block_size + return scheduler_block_size + + def get_request_block_hasher( hash_block_size: int, caching_hash_fn: Callable[[Any], bytes], diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9739dbb3d516..0758c6122565 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -381,9 +381,8 @@ def __init__( or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 for group in kv_cache_config.kv_cache_groups ) - # A finer prefix_match_unit is configured: a mamba partial tail entry - # can only be registered by a step ending exactly at the prompt's last - # hash boundary, so the split adds that stop. + # Fine-grained hits reuse the prompt's hash-aligned Mamba checkpoint. + # Split there only if the backend cannot save it within the final chunk. self.mamba_partial_cache_hit = ( self.need_mamba_block_aligned_split and self.hash_block_size < self.block_size @@ -460,14 +459,19 @@ def _mamba_block_aligned_split( # The last block-aligned position whose state can be cached. With # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. - last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle_block_drop: - last_cache_position = max(last_cache_position - block_size, 0) + ckpt_block_size = ( + self.hash_block_size if self.mamba_partial_cache_hit else self.block_size + ) + last_cache_position = get_mamba_prefill_checkpoint_position( + request.num_tokens, + block_size, + self.use_eagle_block_drop, + ) end = start + num_new_tokens checkpoint_position = get_mamba_prefill_checkpoint_position( prefill_end, - self.hash_block_size, + ckpt_block_size, drop_eagle_block=self.use_eagle_block_drop, ) use_internal_checkpoint = ( @@ -500,18 +504,14 @@ def _mamba_block_aligned_split( next_block_boundary = (start // block_size + 1) * block_size tail_boundary = ( - request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + self.hash_block_size, + self.use_eagle_block_drop, + ) if self.mamba_partial_cache_hit and not use_internal_checkpoint else 0 ) - if tail_boundary and self.use_eagle_block_drop: - # Eagle matches one hash unit past the candidate and drops it, so - # nothing proves the prompt's own last hash boundary. Materialize - # the state one unit lower, where the hit can actually land. Keyed on - # the block-drop bit, not plain use_eagle: this shift exists only to - # compensate for the drop, and the Mamba manager's matching gate - # reads the same bit (the coordinator is handed use_eagle_block_drop). - tail_boundary = max(tail_boundary - self.hash_block_size, 0) junction = request.shared_prefix_boundary # Block-floored: a sub-block junction's state is not separately cacheable. block_floored = start + (junction - start) // block_size * block_size @@ -532,8 +532,7 @@ def _mamba_block_aligned_split( else 0, # Never run past the last cacheable block boundary mid-chunk. last_cache_position, - # Fine-grained hits: the prompt's partial-tail entry can only be - # registered by a chunk ending exactly at its last hash boundary. + # Stop at the prompt checkpoint unless the backend saves it internally. tail_boundary if last_cache_position < tail_boundary < request.num_prompt_tokens else 0, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 5940a1cd0a3f..3e7ae1b7063a 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -876,7 +876,14 @@ def _cache_partial_tail_block( block are intentionally skipped. """ hash_block_size = self.block_pool.hash_block_size - boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + # A resend matches at most `prompt_len - 1` tokens, but EAGLE lookups + # read the whole prompt before dropping a hash unit. + token_limit = ( + request.num_prompt_tokens + if self.use_eagle + else request.num_prompt_tokens - 1 + ) + boundary_tokens = token_limit // hash_block_size * hash_block_size if boundary_tokens == 0 or boundary_tokens > num_tokens: return if boundary_tokens % self.block_size == 0: @@ -1788,7 +1795,7 @@ def get_num_blocks_to_allocate( # is kept, so `prefill_end` skips reserving blocks for the others. checkpoint_position = get_mamba_prefill_checkpoint_position( num_tokens, - self.block_pool.hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) if not self._needs_internal_checkpoint( @@ -2073,20 +2080,29 @@ def _cache_partial_tail_block( checkpoint_position == get_mamba_prefill_checkpoint_position( request.num_prompt_tokens, - hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) ) + is_retained_checkpoint = False + if checkpoint_position % self.block_size == 0: + block_idx = checkpoint_position // self.block_size - 1 + retained = self.reachable_block_mask( + block_idx, + block_idx + 1, + self.cache_hit_alignment_tokens, + self.kv_cache_spec, + self.drop_eagle_checkpoint_block, + retention_interval=retention_interval, + ) + is_retained_checkpoint = retained is None or retained[0] if ( - retention_interval == 0 - and num_tokens < request.num_prompt_tokens - and not is_prompt_checkpoint + not is_prompt_checkpoint + and not is_retained_checkpoint and checkpoint_position != request.shared_prefix_boundary ): - # retention_interval == 0 keeps this transient checkpoint - # request-local. The slot may carry a hash from this step's - # full-block pass; that must go too, since the checkpoint - # state is about to overwrite the block. + # The transient state overwrites this slot, so neither it nor + # the full-block hash registered earlier this step is reusable. self.block_pool._maybe_evict_cached_block(checkpoint_block) return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: @@ -2099,21 +2115,13 @@ def _cache_partial_tail_block( block_size=self.block_size, replace_existing_hashes=True, ) - if self.block_size == hash_block_size: + if num_tokens <= 0 or num_tokens % hash_block_size != 0: return None - if num_tokens % self.block_size == 0: - return None - if num_tokens % hash_block_size != 0: - return None - latest_prompt_hash_boundary = ( - request.num_prompt_tokens // hash_block_size - ) * hash_block_size - if self.drop_eagle_checkpoint_block: - # Eagle groups match one hash unit past the candidate and drop it, - # so register the tail one unit lower. - latest_prompt_hash_boundary = max( - latest_prompt_hash_boundary - hash_block_size, 0 - ) + latest_prompt_hash_boundary = get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + hash_block_size, + self.drop_eagle_checkpoint_block, + ) # The junction is the other position a sibling resumes at: where one was # observed to stop, and where the scheduler already ends a chunk. Bounded # to the prompt chunk being computed -- during decode the target is the @@ -2126,7 +2134,7 @@ def _cache_partial_tail_block( ): return None - block_idx = num_tokens // self.block_size + block_idx = cdiv(num_tokens, self.block_size) - 1 blocks = self.req_to_blocks[request.request_id] if block_idx >= len(blocks): return None @@ -2141,7 +2149,7 @@ def _cache_partial_tail_block( kv_cache_group_id=self.kv_cache_group_id, block_size=self.block_size, ) - if partial_hash is not None: + if partial_hash is not None and num_tokens % self.block_size: self._partial_hit_reqs[request.request_id] = (block_idx, source_block) self.num_cached_block[request.request_id] = block_idx # Producer of this partial tail: the boundary state currently lives diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index e0a003af6254..679bbbfd2415 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -53,6 +53,7 @@ get_kv_cache_configs, get_request_block_hasher, init_none_hash, + resolve_cache_hit_alignment_tokens, resolve_kv_cache_block_sizes, update_kv_cache_capacity, ) @@ -352,6 +353,19 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.validate_block_size() + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + scheduler_kv_cache_config, vllm_config + ) + cache_hit_alignment_tokens = resolve_cache_hit_alignment_tokens( + scheduler_kv_cache_config, + vllm_config, + scheduler_block_size, + hash_block_size, + ) + for kv_cache_config in kv_cache_configs: + kv_cache_config.hash_block_size = hash_block_size + kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens + self.model_executor.initialize_from_config(kv_cache_configs) if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self.model_executor.compile_or_warm_up_model() diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cf707611b47d..ca3f485917eb 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1127,7 +1127,10 @@ def get_mamba_prefill_checkpoint_position( drop_eagle_block: bool, ) -> int: """Return the reusable Mamba checkpoint boundary for a prefill.""" - checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size + # Without EAGLE, leave one token to recompute on resend. + # EAGLE's block drop already leaves tokens to recompute. + proof_limit = num_tokens if drop_eagle_block else num_tokens - 1 + checkpoint_position = proof_limit // hash_block_size * hash_block_size if drop_eagle_block: checkpoint_position -= hash_block_size return max(checkpoint_position, 0) @@ -1487,6 +1490,10 @@ class KVCacheConfig: """Resolved retention policy for local prefix-cache checkpoints.""" kv_cache_layout: str | None = None """The KV cache layout resolved by the engine core, adopted by all workers.""" + hash_block_size: int | None = None + """Tokens per prefix-cache block hash, resolved by the engine core.""" + cache_hit_alignment_tokens: int | None = None + """Token granularity of prefix-cache hits, resolved by the engine core.""" hisparse_host_num_blocks: int | None = None """Capacity of the dedicated HiSparse host-block manager, when enabled.""" diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 70427d85a103..f2dfc8a689f0 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -821,10 +821,15 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: # so that it's available to the warmup stage. self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks - # Adopt the engine core's layout; workers spawned after resolution - # (e.g. elastic EP scale-up) only see it through the config. + # Adopt the engine core's layout and prefix-cache granularity; workers + # spawned after resolution (e.g. elastic EP scale-up) only see them + # through the config. if kv_cache_config.kv_cache_layout is not None: record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout) + self.cache_config.hash_block_size = kv_cache_config.hash_block_size + self.cache_config.cache_hit_alignment_tokens = ( + kv_cache_config.cache_hit_alignment_tokens + ) # Init kv cache connector here, because it requires # `kv_cache_config`. From 30c740589055c86f29859f61b13ae6620449f3fc Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 14:39:09 -0400 Subject: [PATCH 228/952] [MoE] MXFP4 oracle: add CUTLASS W4A4 backend and BF16 activation fallback (#60079) Signed-off-by: mgoin Co-authored-by: Claude --- tests/kernels/moe/test_cutlass_moe.py | 84 +++++++++++++++++ tests/kernels/moe/test_ocp_mx_moe.py | 52 ++++++++++- .../layers/fused_moe/oracle/mxfp4.py | 92 ++++++++++++++++--- 3 files changed, 214 insertions(+), 14 deletions(-) diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 191079662326..050e1bc10487 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -37,9 +37,12 @@ run_cutlass_moe_fp4, run_cutlass_moe_fp8, run_cutlass_moe_mxfp4, + swizzle_mxfp4_scales, ) +from vllm.model_executor.layers.fused_moe.oracle import mxfp4 as mxfp4_oracle from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -294,6 +297,87 @@ def run_moe( torch.testing.assert_close(actual[1:], torch.zeros_like(actual[1:])) +@torch.inference_mode() +def test_cutlass_mxfp4_oracle_matches_direct_kernel(monkeypatch, workspace_init): + """The MXFP4 oracle selects CUTLASS for W4A4 and swizzles checkpoint + scales into the layout the kernel expects.""" + if not CutlassExpertsMxfp4._supports_current_device(): + pytest.skip("CUTLASS mxfp4 MoE is not supported on this GPU") + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + + set_random_seed(7) + e, m, n, k, topk = 8, 33, 256, 512, 2 + device = torch.device("cuda") + w1 = torch.randint(0, 256, (e, 2 * n, k // 2), device=device, dtype=torch.uint8) + w2 = torch.randint(0, 256, (e, k, n // 2), device=device, dtype=torch.uint8) + w1_scale = torch.randint(118, 124, (e, 2 * n, k // 32), device=device).byte() + w2_scale = torch.randint(118, 124, (e, k, n // 32), device=device).byte() + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) + score = torch.randn((m, e), device=device, dtype=torch.bfloat16) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + def swizzle(scale: torch.Tensor) -> torch.Tensor: + rows, cols = scale.shape[1:] + return torch.stack( + [swizzle_mxfp4_scales(s, rows, cols * 32).view(rows, cols) for s in scale] + ) + + expected = torch.empty_like(a) + run_cutlass_moe_mxfp4( + output=expected, + a=a, + w1_fp4=w1, + w1_blockscale=swizzle(w1_scale), + w2_fp4=w2, + w2_blockscale=swizzle(w2_scale), + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + workspace13=torch.empty((m * topk, max(2 * n, k)), device=device).to(a), + workspace2=torch.empty((m * topk, n), device=device).to(a), + m=m, + n=n, + k=k, + e=e, + device=device, + apply_router_weight_on_input=False, + ) + + moe_config = make_dummy_moe_config( + num_experts=e, experts_per_token=topk, hidden_dim=k, intermediate_size=n + ) + with set_current_vllm_config(vllm_config): + backend, experts_cls = mxfp4_oracle.select_mxfp4_moe_backend( + moe_config, activation_key=kMxfp4Dynamic + ) + assert backend == mxfp4_oracle.Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4 + + w13, w2, w13_scale, w2_scale, _, _ = ( + mxfp4_oracle.convert_weight_to_mxfp4_moe_kernel_format( + backend, torch.nn.Module(), w1, w2, w1_scale, w2_scale + ) + ) + kernel = mxfp4_oracle.make_mxfp4_moe_kernel( + mxfp4_oracle.make_mxfp4_moe_quant_config(backend, w13_scale, w2_scale), + moe_config, + experts_cls, + backend, + ) + actual = kernel.apply( + a, + w13, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @dataclasses.dataclass class MOETensors: a: torch.Tensor diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 5a7053a8b1d4..f20927565ebd 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -7,7 +7,11 @@ import pytest import torch -from tests.kernels.moe.utils import check_accuracy, check_deferred_moe_finalize +from tests.kernels.moe.utils import ( + check_accuracy, + check_deferred_moe_finalize, + make_dummy_moe_config, +) from vllm._aiter_ops import ( is_aiter_found, is_aiter_found_and_supported, @@ -2009,6 +2013,52 @@ def is_supported_config( mxfp4_oracle.select_mxfp4_moe_backend(moe_config) +@pytest.mark.parametrize( + "moe_backend,w4a4,supported,expected", + [ + # The checkpoint's activation wins over BF16 backends listed first. + ("auto", True, {"CUTLASS_MXFP4_MXFP4", "MARLIN"}, "CUTLASS_MXFP4_MXFP4"), + # No W4A4 kernel: run at BF16 rather than another quantized activation. + ("auto", True, {"FLASHINFER_TRTLLM_MXFP4_MXFP8", "MARLIN"}, "MARLIN"), + ("marlin", True, {"MARLIN"}, "MARLIN"), + # A BF16 checkpoint never gets its activations quantized. + ("auto", False, {"CUTLASS_MXFP4_MXFP4", "MARLIN"}, "MARLIN"), + ], +) +def test_select_mxfp4_moe_backend_respects_checkpoint_activation( + monkeypatch: pytest.MonkeyPatch, + moe_backend: str, + w4a4: bool, + supported: set[str], + expected: str, +): + """Activations run at the checkpoint's precision when a backend supports + it and fall back to BF16 otherwise, never to a lower precision.""" + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + ) + + def backend_to_kernel_cls(backend): + class Experts: + @staticmethod + def is_supported_config(*args): + return backend.name in supported, None + + return [Experts] + + monkeypatch.setattr(mxfp4_oracle, "backend_to_kernel_cls", backend_to_kernel_cls) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + moe_config = make_dummy_moe_config() + moe_config.moe_backend = moe_backend + + backend, _ = mxfp4_oracle.select_mxfp4_moe_backend( + moe_config, activation_key=kMxfp4Dynamic if w4a4 else None + ) + + assert backend.name == expected + + @pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"]) @pytest.mark.parametrize("use_ep", [False, True]) def test_aiter_mxfp4_monolithic_rejects_expert_parallel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 5463853caba1..f5fccc03436f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -114,6 +114,18 @@ def _pack_deepgemm_mxfp4_scales( ) +def _swizzle_cutlass_mxfp4_scales(scale: torch.Tensor) -> torch.Tensor: + """Swizzle [E, N, K // 32] block scales into the CUTLASS tiled layout.""" + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + swizzle_mxfp4_scales, + ) + + n, scale_k = scale.shape[1:] + return torch.stack( + [swizzle_mxfp4_scales(s, n, scale_k * 32).reshape(n, scale_k) for s in scale] + ) + + class Mxfp4MoeBackend(Enum): NONE = "None" # b12x backends @@ -127,6 +139,8 @@ class Mxfp4MoeBackend(Enum): # FlashInfer CUTLASS backends FLASHINFER_CUTLASS_MXFP4_MXFP8 = "FLASHINFER_CUTLASS_MXFP4_MXFP8" FLASHINFER_CUTLASS_MXFP4_BF16 = "FLASHINFER_CUTLASS_MXFP4_BF16" + # vLLM CUTLASS backend + CUTLASS_MXFP4_MXFP4 = "CUTLASS_MXFP4_MXFP4" # Marlin BATCHED_MARLIN = "BATCHED_MARLIN" MARLIN = "MARLIN" @@ -218,6 +232,13 @@ def backend_to_kernel_cls( return [FlashInferExperts] + elif backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + + return [CutlassExpertsMxfp4] + elif backend == Mxfp4MoeBackend.TRITON: from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 OAITritonExperts, @@ -378,6 +399,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.TRITON, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, # TRITON_UNFUSED has bug with MTP support # TODO re-enable after kernel is fixed # TRITON_UNFUSED @@ -440,7 +462,10 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: return kMxfp8Dynamic if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: return kFp8StaticTensorSym - if backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: + if backend in ( + Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, + ): return kMxfp4Dynamic return None # BF16 activation @@ -474,6 +499,22 @@ def _make_log_backend(backend: Mxfp4MoeBackend) -> str: return f"Using '{backend.value}' Mxfp4 MoE backend." +def _warn_on_bf16_fallback( + backend: Mxfp4MoeBackend, requested_activation_key: QuantKey | None +) -> None: + if ( + requested_activation_key is not None + and backend != Mxfp4MoeBackend.EMULATION + and _backend_activation_key(backend) is None + ): + logger.warning_once( + "No Mxfp4 MoE backend supports activation=%s for the deployment " + "configuration. Falling back to '%s' with BF16 activations.", + requested_activation_key, + backend.value, + ) + + def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str: base = ( f"Mxfp4 MoE backend '{backend.value}' does not support the " @@ -505,17 +546,19 @@ def _filter_by_activation( backends: list[Mxfp4MoeBackend], requested_activation_key: QuantKey | None, ) -> list[Mxfp4MoeBackend]: - """Pick variants matching ``requested_activation_key``; without one, - prefer BF16 if the list has any, else keep the list as-is so explicit - non-BF16 picks (e.g. the ``_afp8`` aliases) still land.""" + """Pick variants matching ``requested_activation_key``, followed by BF16 + variants as a higher-precision fallback; without one, prefer BF16 if the + list has any, else keep the list as-is so explicit non-BF16 picks (e.g. + the ``_afp8`` aliases) still land.""" + bf16 = [b for b in backends if _backend_activation_key(b) is None] if requested_activation_key is not None: - return [ + matching = [ b for b in backends if _backend_activation_key(b) == requested_activation_key or b == Mxfp4MoeBackend.EMULATION ] - bf16 = [b for b in backends if _backend_activation_key(b) is None] + return matching + [b for b in bf16 if b not in matching] return bf16 if bf16 else backends @@ -576,8 +619,9 @@ def select_mxfp4_moe_backend( Args: config: MoE configuration - activation_key: Optional activation quantization key. If provided, - overrides the default activation key for backend selection. + activation_key: Optional activation quantization key declared by + the checkpoint. Backends that match it are preferred, with BF16 + activation backends as a fallback. Use kFp8StaticTensorSym for W4A8 scheme. Note: Shape-specific fallbacks may still occur at runtime. @@ -614,7 +658,7 @@ def select_mxfp4_moe_backend( else _backend_activation_key(requested_backend) ) try: - return _return_or_raise( + selected = _return_or_raise( requested_backend, config, kMxfp4Static, @@ -623,6 +667,9 @@ def select_mxfp4_moe_backend( ) except ValueError as e: last_error = e + else: + _warn_on_bf16_fallback(requested_backend, requested_activation_key) + return selected assert last_error is not None raise last_error @@ -648,10 +695,10 @@ def select_mxfp4_moe_backend( unsupported_reasons = [] for backend in AVAILABLE_BACKENDS: - # Use requested_activation_key if provided, otherwise use backend default + # EMULATION takes the requested key, other backends use their own act_key = ( requested_activation_key - if requested_activation_key is not None + if backend == Mxfp4MoeBackend.EMULATION else _backend_activation_key(backend) ) for k_cls in backend_to_kernel_cls(backend): @@ -660,6 +707,7 @@ def select_mxfp4_moe_backend( ) if supported: logger.info_once(_make_log_backend(backend)) + _warn_on_bf16_fallback(backend, requested_activation_key) return backend, k_cls else: logger.debug_once(_make_log_unsupported(backend, reason)) @@ -789,6 +837,10 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( # DeepGEMM requires M/N/K alignment intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 128) + elif backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + # CUTLASS scale tiles cover 128 rows by 128 elements + intermediate_size = round_up(intermediate_size, 128) + hidden_size = round_up(hidden_size, 128) elif backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): intermediate_size = round_up(intermediate_size, 128) if current_platform.is_xpu(): @@ -1433,7 +1485,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( ]: """Convert loaded weights into backend-specific kernel format. - Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends. + Supports CUTLASS, DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin + backends. """ is_gfx1250 = False if current_platform.is_rocm(): @@ -1451,6 +1504,16 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) + if mxfp4_backend == Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4: + return ( + w13_weight.data, + w2_weight.data, + _swizzle_cutlass_mxfp4_scales(w13_weight_scale.data), + _swizzle_cutlass_mxfp4_scales(w2_weight_scale.data), + w13_bias, + w2_bias, + ) + if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, @@ -2005,7 +2068,10 @@ def make_mxfp4_moe_quant_config( block_shape=None, gemm1_clamp_limit=swiglu_limit, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: + elif mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4, + ): return ocp_mx_moe_quant_config( quant_dtype="mxfp4", w1_bias=w1_bias, From 2274e86842d6e59a5f9f9d1fb6cc1dd2d2bf8350 Mon Sep 17 00:00:00 2001 From: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:41:10 +0800 Subject: [PATCH 229/952] [Spec Decode][Model] Support DFlash2 draft models with GLM-5.3-Flash (#56983) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> --- vllm/models/glm5next/common/model.py | 51 +++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 26becc9619e7..fd7c0263bdda 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -66,9 +66,11 @@ Glm4vForConditionalGeneration, ) from vllm.model_executor.models.interfaces import ( + EagleModelMixin, HasInnerState, IsHybrid, MixtureOfExperts, + SupportsEagle3, SupportsPP, ) from vllm.model_executor.models.utils import ( @@ -685,7 +687,7 @@ def hc_fused_post_pre( ) -class Glm5NextModel(nn.Module): +class Glm5NextModel(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -743,6 +745,7 @@ def get_layer(prefix: str): get_layer, prefix=f"{prefix}.layers", ) + self._aux_post_op = MHCPostOp() # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] @@ -767,6 +770,26 @@ def get_layer(prefix: str): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + def _aux_hidden_state( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + post: torch.Tensor | None, + comb: torch.Tensor | None, + ) -> torch.Tensor: + """Completed residual stream entering a layer, as one hidden vector. + + mHC layers defer their final ``hc_post`` into the next layer's fused + pre-op, so ``hidden_states`` holds the raw layer output while the + widened stream is completed here and contracted back to + ``hidden_size``. Non-mHC layers already return the summed stream. + """ + if post is None: + return hidden_states + assert residual is not None and comb is not None + completed = self._aux_post_op(hidden_states, residual, post, comb) + return hc_contract(completed, self.config.mhc_num_residual_streams) + def forward( self, input_ids: torch.Tensor | None, @@ -796,7 +819,15 @@ def forward( if self.is_sequence_parallel: hidden_states = sp_shard(hidden_states) - for layer in self._active_layers: + aux_hidden_states: list[torch.Tensor] = [] + for idx, layer in enumerate(self._active_layers, start=self.start_layer): + if idx in self.aux_hidden_state_layers: + aux_hidden_state = self._aux_hidden_state( + hidden_states, residual, post, comb + ) + if self.is_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) hidden_states, residual, post, comb = layer( positions, hidden_states, residual, post, comb ) @@ -811,10 +842,18 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) + if self.end_layer in self.aux_hidden_state_layers: + final_aux = self._aux_hidden_state(hidden_states, residual, post, comb) + if self.is_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_states.append(final_aux) + if self.is_sequence_parallel: hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] hidden_states = self.norm(hidden_states) + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1006,7 +1045,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Glm5NextForCausalLM( - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid, SupportsEagle3 ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1102,7 +1141,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: dummy_inputs=Glm4vDummyInputsBuilder, ) class Glm5NextForConditionalGeneration( - Glm4vForConditionalGeneration, HasInnerState, IsHybrid, MixtureOfExperts + Glm4vForConditionalGeneration, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsEagle3, ): # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The # multimodal wrapper must declare the same interfaces so vLLM treats it as From c4a5ff5c18ceba5f06512ef6b17a48338f275763 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Fri, 9 Oct 2026 02:46:01 +0800 Subject: [PATCH 230/952] [Frontend] shutdown engine core and its sub processes asap when initializing (#52299) Signed-off-by: Andy Xie Co-authored-by: Simon Mo --- .../test_engine_core_shutdown_signal.py | 224 ++++++++++++++++++ vllm/v1/engine/core.py | 68 ++++-- 2 files changed, 275 insertions(+), 17 deletions(-) create mode 100644 tests/v1/engine/test_engine_core_shutdown_signal.py diff --git a/tests/v1/engine/test_engine_core_shutdown_signal.py b/tests/v1/engine/test_engine_core_shutdown_signal.py new file mode 100644 index 000000000000..73bc794088d6 --- /dev/null +++ b/tests/v1/engine/test_engine_core_shutdown_signal.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Regression tests for termination signals that arrive while EngineCoreProc is +# still being constructed. +# +# When this fix matters +# --------------------- +# Engine shutdown is driven from the frontend. If a frontend process dies while +# EngineCore is still starting up, the process manager sends SIGTERM to its +# children and force kills whatever is still alive after a short timeout (see +# shutdown in vllm/v1/utils.py). Construction is slow - it spawns the executor +# processes, loads weights, profiles and captures CUDA graphs - so that +# shutdown often lands inside the construction window. Bigger models and more +# ranks widen the window. This is the case PR 52299 (follow up to 43417) has to +# survive: 43417 stops EngineCore from exiting silently, but the request itself +# still has to outlive the construction window. +# +# run_engine_core installs the signal handler before engine_core exists, so in +# that window the handler runs with engine_core == None. Losing the request +# there is permanent, not a delay: shutdown_state stays RUNNING, and +# wakeup_engine returns early while consuming the one-shot SignalCallback. The +# busy loop then starts RUNNING and blocks on input_queue forever, so no later +# signal can wake it. Graceful shutdown is gone for good and the process only +# stops when something force kills it, which skips executor shutdown and GPU +# release. In a log this shows up as the signal being received, followed by a +# force kill of EngineCore after the process manager timeout. +# +# What these tests pin +# -------------------- +# A request arriving in the window must be remembered somewhere other than +# engine_core, applied before run_busy_loop starts, and must not consume the +# one-shot wakeup callback, so a later signal can still wake an idle loop. + +import queue +import signal as signal_module +from typing import Any + +from vllm.v1.engine import EngineCoreRequestType +from vllm.v1.engine import core as core_module +from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState + +WAKEUP_ITEM = (EngineCoreRequestType.WAKEUP, None) +CONSTRUCTION_SIGNAL = signal_module.SIGTERM + + +class _FakeSignal: + # Minimal stand-in for the signal module. It records the installed handler + # so the test can fire it by hand, and avoids touching the real process + # signal disposition. Signals is reused from the real module so the handler + # can still resolve signal.Signals(signum).name. + SIGTERM = signal_module.SIGTERM + SIGINT = signal_module.SIGINT + SIG_DFL = signal_module.SIG_DFL + Signals = signal_module.Signals + + def __init__(self): + self.handlers = {} + + def signal(self, signum, handler): + self.handlers[signum] = handler + + def fire(self): + self.handlers[CONSTRUCTION_SIGNAL](CONSTRUCTION_SIGNAL, None) + + +class _FakeSignalCallback: + # One-shot stand-in for SignalCallback. The real class runs the callback on + # a dedicated thread and exits after a single trigger, so every later + # trigger is a no-op. Running the callback inline keeps the test + # deterministic instead of depending on thread scheduling. + instances: list["_FakeSignalCallback"] = [] + + def __init__(self, callback): + self.callback = callback + self.trigger_count = 0 + self.consumed = False + self.stopped = False + _FakeSignalCallback.instances.append(self) + + def trigger(self): + self.trigger_count += 1 + if self.stopped or self.consumed: + return + self.consumed = True + callback, self.callback = self.callback, None + callback() + + def stop(self): + self.stopped = True + + +class _StubParallelConfig: + def __init__(self): + self.data_parallel_size = 1 + self.data_parallel_rank_local = 0 + self.data_parallel_index = 0 + self.data_parallel_rank = 0 + self.numa_bind = False + + def reconfigure_for_independent_dp_rank(self): + pass + + +class _StubModelConfig: + is_moe = False + + +class _StubVllmConfig: + def __init__(self): + self.parallel_config = _StubParallelConfig() + self.model_config = _StubModelConfig() + self.logging_config = None + self.kv_transfer_config = None + self.shutdown_timeout = 0 + + +def _run_engine_core( + monkeypatch, + *, + signal_during_construction=False, + signal_while_idle=False, +): + _FakeSignalCallback.instances.clear() + fake_signal = _FakeSignal() + trace = {} + + class StubEngineCoreProc: + def __init__(self, *args, **kwargs): + self.shutdown_state = EngineShutdownState.RUNNING + # Same queue type as EngineCore, so the wakeup path under test is + # the real one. + self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]() + self.vllm_config = kwargs.get("vllm_config") + if signal_during_construction: + # The termination signal lands mid construction, while the + # caller still sees engine_core == None. + fake_signal.fire() + trace["constructed"] = True + + def run_busy_loop(self): + # What the real loop sees on its first _handle_shutdown() check. + trace["loop_entered"] = True + trace["state_at_loop_entry"] = self.shutdown_state + if signal_while_idle: + fake_signal.fire() + try: + trace["wakeup"] = self.input_queue.get_nowait() + except queue.Empty: + trace["wakeup"] = None + + def shutdown(self): + pass + + def _send_engine_dead(self): + pass + + def has_work(self): + return False + + monkeypatch.setattr(core_module, "signal", fake_signal) + monkeypatch.setattr(core_module, "SignalCallback", _FakeSignalCallback) + monkeypatch.setattr(core_module, "EngineCoreProc", StubEngineCoreProc) + monkeypatch.setattr(core_module, "DPEngineCoreProc", StubEngineCoreProc) + monkeypatch.setattr(core_module, "configure_logging", lambda *a, **k: None) + monkeypatch.setattr( + core_module, "maybe_register_config_serialize_by_value", lambda *a, **k: None + ) + monkeypatch.setattr(core_module, "set_process_title", lambda *a, **k: None) + monkeypatch.setattr(core_module, "maybe_init_worker_tracer", lambda *a, **k: None) + monkeypatch.setattr(core_module, "decorate_logs", lambda *a, **k: None) + + EngineCoreProc.run_engine_core(vllm_config=_StubVllmConfig()) + return trace + + +def test_signal_during_construction_is_applied_before_loop(monkeypatch): + trace = _run_engine_core(monkeypatch, signal_during_construction=True) + + assert trace["constructed"] is True + assert trace["loop_entered"] is True + # The request arrived while engine_core was None, so it could not be stored + # there. It must still be applied before run_busy_loop starts, otherwise the + # loop comes up as RUNNING, the shutdown is silently dropped, and force kill + # becomes the only way out. + assert trace["state_at_loop_entry"] == EngineShutdownState.REQUESTED, ( + "a signal received during construction must be applied before the loop " + f"starts, got {trace['state_at_loop_entry']!r}" + ) + + +def test_signal_during_construction_keeps_later_wakeup_working(monkeypatch): + trace = _run_engine_core( + monkeypatch, + signal_during_construction=True, + signal_while_idle=True, + ) + + assert trace["state_at_loop_entry"] == EngineShutdownState.REQUESTED, ( + "the construction time signal must be applied before the loop starts, " + f"got {trace['state_at_loop_entry']!r}" + ) + # SignalCallback is one-shot. If the construction time signal consumed it, + # no wakeup could be delivered afterwards and the idle busy loop would block + # on input_queue forever instead of shutting down, so the process would sit + # there holding GPU memory until something force kills it. + assert trace["wakeup"] == WAKEUP_ITEM, ( + "a signal after construction must still wake the idle loop, " + f"got {trace['wakeup']!r}" + ) + + +def test_signal_while_idle_still_wakes_the_loop(monkeypatch): + trace = _run_engine_core(monkeypatch, signal_while_idle=True) + + assert trace["state_at_loop_entry"] == EngineShutdownState.RUNNING + assert trace["wakeup"] == WAKEUP_ITEM + + +def test_no_signal_leaves_the_loop_running(monkeypatch): + trace = _run_engine_core(monkeypatch) + + assert trace["state_at_loop_entry"] == EngineShutdownState.RUNNING + assert trace["wakeup"] is None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 679bbbfd2415..bfbaa38d9478 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1410,40 +1410,74 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): vllm_config.kv_transfer_config.engine_id, ) - parallel_config.data_parallel_index = dp_rank - if data_parallel and vllm_config.model_config.is_moe: - # Set data parallel rank for this engine process. - parallel_config.data_parallel_rank = dp_rank - engine_core = DPEngineCoreProc(*args, **kwargs) - else: - # Non-MoE DP ranks are completely independent, so treat like DP=1. - # Note that parallel_config.data_parallel_index will still reflect - # the original DP rank. - parallel_config.reconfigure_for_independent_dp_rank() - engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) - - assert engine_core is not None + # declare engine core as None and register signal handler before + # engine_core initialization in case engine core process is terminated + # in the middle of initialization. For example, frontend process exits + # unexpectedly, if signal handler is not registered, engine core process + # will use default signal handler which will terminate the process + # silently. If engine core model executor has been initialized, it will + # not be terminated asap. Register signal handler early and only log a + # message will make engine core process alive. engine core process and + # its subprocesses(model executor processes) will be force killed in + # CoreEngineProcManager.shutdown. + engine_core = None + # A termination signal can arrive while engine_core is still being + # constructed. At that point the request cannot be recorded on + # engine_core, so keep it in its own flag and apply it later. + shutdown_requested = False def wakeup_engine(): # Wakes up idle engine via input_queue when shutdown is requested # Not safe in a signal handler - we may interrupt the main thread # while it is holding the non-reentrant input_queue.mutex + assert engine_core is not None engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) - signal_callback = SignalCallback(wakeup_engine) - def signal_handler(signum, frame): + nonlocal shutdown_requested signal_name = signal.Signals(signum).name logger.info( "[shutdown] EngineCore: trigger received signal=%s", signal_name, ) - engine_core.shutdown_state = EngineShutdownState.REQUESTED - signal_callback.trigger() + # Record the request first so it survives the construction window. + shutdown_requested = True + if engine_core is not None: + engine_core.shutdown_state = EngineShutdownState.REQUESTED + # signal_callback is armed only after engine_core exists. Don't + # trigger it earlier: it is one-shot, and a wakeup cannot be + # delivered while there is no input_queue to wake. + # For more info: https://github.com/vllm-project/vllm/pull/52299#pullrequestreview-5370255271 + if signal_callback is not None: + signal_callback.trigger() signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + parallel_config.data_parallel_index = dp_rank + if data_parallel and vllm_config.model_config.is_moe: + # Set data parallel rank for this engine process. + parallel_config.data_parallel_rank = dp_rank + engine_core = DPEngineCoreProc(*args, **kwargs) + else: + # Non-MoE DP ranks are completely independent, so treat like DP=1. + # Note that parallel_config.data_parallel_index will still reflect + # the original DP rank. + parallel_config.reconfigure_for_independent_dp_rank() + engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) + + assert engine_core is not None + + # Apply a signal received during construction before entering the + # loop. run_busy_loop calls _handle_shutdown first, so the loop exits + # right away instead of blocking on an idle input_queue. + if shutdown_requested: + engine_core.shutdown_state = EngineShutdownState.REQUESTED + + # Arm the wakeup callback only now that input_queue exists, so a + # later signal can still wake an idle busy loop. + signal_callback = SignalCallback(wakeup_engine) + engine_core.run_busy_loop() except SystemExit as e: From 56e6da2c4717f6c6c2faa0159f8c0fa2cf508d6f Mon Sep 17 00:00:00 2001 From: Yeonwoo Sung Date: Fri, 9 Oct 2026 03:47:30 +0900 Subject: [PATCH 231/952] [Bugfix][Frontend] Honor stop_token_ids on the Responses API (#56592) Signed-off-by: YeonwooSung Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/entrypoints/openai/responses/test_sampling_params.py | 2 ++ vllm/entrypoints/openai/responses/protocol.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 877e8724e8b2..01fdfa4594ea 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -47,6 +47,7 @@ def test_extra_sampling_params(self): repetition_penalty=1.2, seed=42, stop=["END", "STOP"], + stop_token_ids=[42, 43], ignore_eos=True, vllm_xargs={"custom": "value"}, ) @@ -56,6 +57,7 @@ def test_extra_sampling_params(self): assert sampling_params.repetition_penalty == 1.2 assert sampling_params.seed == 42 assert sampling_params.stop == ["END", "STOP"] + assert sampling_params.stop_token_ids == [42, 43] assert sampling_params.ignore_eos is True assert sampling_params.extra_args == {"custom": "value"} diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 35e6d3aa3dcd..12e40e98a882 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -363,6 +363,7 @@ class ResponsesRequest(OpenAIBaseModel): repetition_penalty: float | None = None seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX) stop: StopParam = [] + stop_token_ids: list[int] | None = [] ignore_eos: bool = False vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, @@ -521,6 +522,7 @@ def to_sampling_params( max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, stop=stop, + stop_token_ids=self.stop_token_ids, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, repetition_penalty=repetition_penalty, From 9bbc58156095f45bdb91ddbdbd244119d3afb12c Mon Sep 17 00:00:00 2001 From: Aarushi Jain Date: Thu, 8 Oct 2026 13:53:00 -0500 Subject: [PATCH 232/952] [Bugfix] Close MessageQueue zmq resources via finalizer to avoid exit (#60472) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- tests/distributed/test_shm_broadcast.py | 22 ++++++++++++++++ .../device_communicators/shm_broadcast.py | 25 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 7debed4cbbff..d28732261a13 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc import io import pickle import random import threading import time +import weakref from types import SimpleNamespace from unittest import mock @@ -634,6 +636,26 @@ def test_acquire_read_releases_slot_when_reader_raises(): reader.shutdown() +def test_reader_in_reference_cycle_is_collected_without_blocking(): + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=1024 * 1024, + max_chunks=1, + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + reader.cycle = reader + reader_ref = weakref.ref(reader) + del reader + + collector = threading.Thread(target=gc.collect, daemon=True) + collector.start() + collector.join(timeout=10) + + assert not collector.is_alive() + assert reader_ref() is None + + def test_warning_logs(caplog_vllm): """Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals when indefinite=False, and are not emitted when indefinite=True. diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index bafc5728b85e..da2288f77b54 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -9,6 +9,7 @@ import sys import threading import time +import weakref from contextlib import contextmanager from dataclasses import dataclass, field from multiprocessing import shared_memory @@ -472,6 +473,12 @@ class Handle: remote_addr_ipv6: bool = False +def _close_zmq(context: zmq.Context, sockets: list[zmq.Socket]) -> None: + for socket in sockets: + socket.close(linger=0) + context.term() + + class MessageQueue: def __init__( self, @@ -561,9 +568,26 @@ def __init__( remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, ) + self._close_zmq_on_collect(context) logger.debug("vLLM message queue communication handle: %s", self.handle) + def _close_zmq_on_collect(self, context: zmq.Context) -> None: + # pyzmq's Context.__del__ can block forever in term() when the context + # and its sockets are collected in the same GC cycle, so the finalizer + # keeps them alive and closes the sockets before terminating. + sockets = [self.local_socket, self.remote_socket] + if self._spin_condition is not None: + sockets += [ + self._spin_condition.local_notify_socket, + self._spin_condition.read_cancel_socket, + self._spin_condition.write_cancel_socket, + ] + finalizer = weakref.finalize( + self, _close_zmq, context, [s for s in sockets if s is not None] + ) + finalizer.atexit = False # type: ignore[misc] + def export_handle(self) -> Handle: return self.handle @@ -613,6 +637,7 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._spin_condition = None # type: ignore self.shutting_down = False + self._close_zmq_on_collect(context) return self def wait_until_ready(self): From 18751bc6a7b4b24fa5e6bbd5299100dd0b19a22b Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:11:45 -0400 Subject: [PATCH 233/952] [ROCm][CI] Add coverage for AiterExperts token padding: inf/nan garbage rows (#59334) Signed-off-by: Divakar Verma Co-authored-by: Andreas Karatzas --- .../moe/test_modular_kernel_combinations.py | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 4b769931ea04..f29c22085240 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -17,6 +17,10 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, + kFp8StaticTensorSym, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm @@ -988,6 +992,173 @@ def test_aiter_moe_padding_matrix_mxfp4(): padding matrix once MI350/gfx950 hardware is available in CI.""" +# --- 4c: HIP-graph token padding -- inf/nan garbage rows (topk_ids=-1) ------ + +# AiterExperts isolates garbage rows from real rows under both AITER sorting +# backends, and reference_moe_impl's dispatch (`mask = topk_ids == i`) is +# likewise row-local, so comparing real rows against it is valid. Exception: +# AiterExperts's a1 quant scale is normally fixed before garbage injection +# (computed in RankTensors.make()), but "fp8_tensor_token" computes it live +# from the (already-injected) hidden_states -- see the skip below. + +# {mode name -> {row index: fill value}}. Covers multiple garbage-row counts +# and positions, and both inf and nan. +_TOKEN_PADDING_GARBAGE_MODES: dict[str, dict[int, float]] = { + "single_inf": {3: float("inf")}, + "single_nan": {5: float("nan")}, + "multiple_mixed": { + 1: float("inf"), + 4: float("-inf"), + 7: float("nan"), + 12: float("inf"), + }, +} + + +def _apply_token_padding_garbage_rows( + rank_tensors: RankTensors, garbage_rows: dict[int, float] +) -> None: + """Simulate the HIP-graph padding-token contract in-place: garbage rows + get inf/nan hidden_states and their routing masked to topk_ids=-1 / + topk_weights=0, mirroring grouped_topk_router.py's `is_padding` masking. + """ + for row, fill in garbage_rows.items(): + rank_tensors.hidden_states[row].fill_(fill) + rank_tensors.topk_ids[row].fill_(-1) + rank_tensors.topk_weights[row].fill_(0.0) + + +def _aiter_token_padding_worker( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + config: Config, + weights: WeightTensors, + verbose: bool, + garbage_rows: dict[int, float], +): + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + set_random_seed(pgi.rank) + + weights = copy.deepcopy(weights) + weights.to_current_device() + + rank_tensors = RankTensors.make(config, pgi) + _apply_token_padding_garbage_rows(rank_tensors, garbage_rows) + + with set_current_vllm_config(vllm_config): + mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors) + ref_out = reference_moe_impl(config, weights, rank_tensors) + + real_rows = [i for i in range(config.Ms) if i not in garbage_rows] + + assert torch.isfinite(mk_out[real_rows]).all(), ( + "Garbage (padding) token rows leaked NaN/Inf into real rows' output " + "through AiterExperts." + ) + + # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out + # mk_out would still pass the checks below -- guard against that. + ref_scale = ref_out[real_rows].abs().mean() + mk_scale = mk_out[real_rows].abs().mean() + assert mk_scale > 0.5 * ref_scale, ( + f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks " + f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})." + ) + + if config.quant_config is not None: + check_accuracy( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2, percent=0.9 + ) + else: + torch.testing.assert_close( + ref_out[real_rows], mk_out[real_rows], atol=3e-2, rtol=3e-2 + ) + + +@require_aiter_moe +@pytest.mark.parametrize("quant_config", _PADDING_QUANT_CONFIGS, ids=_PADDING_QUANT_IDS) +@pytest.mark.parametrize("garbage_mode", list(_TOKEN_PADDING_GARBAGE_MODES)) +def test_aiter_moe_token_padding_garbage_rows( + garbage_mode: str, quant_config: TestMoEQuantConfig | None +): + """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). + + Simulates HIP-graph token padding: a subset of rows get inf/nan + hidden_states with topk_ids=-1 / topk_weights=0, mirroring the router's + padding-token masking contract (grouped_topk_router.py). Verifies the + real (non-garbage) rows' output through AiterExperts remains finite and + matches reference_moe_impl, across multiple garbage-row counts/positions + and both inf and nan fill values, for unquantized + the 4 fp8 quant + schemes AiterExperts supports. + """ + from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoDPEPModular, + ) + + garbage_rows = _TOKEN_PADDING_GARBAGE_MODES[garbage_mode] + + # Reuses 4b's already-confirmed 128-aligned ("padded") K/N -- these are + # known to work across all 5 quant schemes here, including AITER's + # block-quant kernel (which 4b's _PADDING_K_UNPADDED/_PADDING_N_UNPADDED + # sizes do not support). Token-padding safety doesn't depend on + # hidden_pad/intermediate_pad, so there's no need for a separate shape. + config = Config( + Ms=_PADDING_M, + K=_PADDING_K_PADDED, + N=_PADDING_N_PADDED, + E=_PADDING_E, + topks=_PADDING_TOPK, + dtype=torch.bfloat16, + quant_config=quant_config, + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, + fused_experts_type=AiterExperts, + world_size=1, + ) + assert config.is_valid()[0] + assert config.fe_supports_quant_scheme(), ( + f"AiterExperts does not support quant scheme {quant_config}." + ) + + # This (weight, activation) pair is AiterExperts's real per-tensor + # dynamic-activation-quant scheme -- its batch-wide max-abs scale is + # poisoned by a real inf garbage row. Upstream bug: + # https://github.com/ROCm/aiter/issues/6275. + if ( + quant_config is not None + and config.fp8_quant_key_pair() == (kFp8StaticTensorSym, kFp8DynamicTensorSym) + and any(v in (float("inf"), float("-inf")) for v in garbage_rows.values()) + ): + pytest.skip( + "Batch-wide dynamic per-tensor quant scale is poisoned by a " + "real inf garbage row. See " + "https://github.com/ROCm/aiter/issues/6275." + ) + + weights = WeightTensors.make(config) + vllm_config, env_dict = config.make_env_data() + + parallel_launch_with_config( + config.world_size, + _aiter_token_padding_worker, + vllm_config, + env_dict, + config, + weights, + False, + garbage_rows, + ) + + +@pytest.mark.skip( + reason="MXFP4 AiterExperts token padding requires MI350/gfx950 hardware. " + 'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").' +) +def test_aiter_moe_token_padding_garbage_rows_mxfp4(): + pass + + if __name__ == "__main__": # Ability to test individual PrepareAndFinalize and FusedExperts combination from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser From fc494822eaca100b6339a4911ce36784926573fc Mon Sep 17 00:00:00 2001 From: Kevin Li <110002169+Kevinli24@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:14:35 -0500 Subject: [PATCH 234/952] [Test][Quantization] Gate CUTLASS W4A8 kernel tests to SM90 (#58878) Signed-off-by: Kevin Li Co-authored-by: Misha Goin --- tests/kernels/quantization/test_cutlass_w4a8.py | 11 +++++------ tests/kernels/quantization/test_cutlass_w4a8_moe.py | 3 +-- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/kernels/quantization/test_cutlass_w4a8.py b/tests/kernels/quantization/test_cutlass_w4a8.py index 968d5f2c74b9..fc34a15bde56 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8.py +++ b/tests/kernels/quantization/test_cutlass_w4a8.py @@ -29,8 +29,11 @@ # `is_quant_method_supported` conflates kernels with quantization methods # an assumption which is breaking down as quantizations methods can have # have kernels and some kernels support multiple quantization methods. -capability = current_platform.get_device_capability() -IS_SUPPORTED_BY_GPU = capability is not None and capability[0] >= 9 +# CUTLASS W4A8 kernels are only built for SM90 (9.0a); the kernel selector +# requires exactly SM90 too, so match it here rather than accepting >= 90. +IS_SUPPORTED_BY_GPU = ( + current_platform.is_cuda() and current_platform.is_device_capability(90) +) MNK_SHAPES = [ (1, 128, 128), @@ -106,10 +109,6 @@ class Tensors: # TODO: in future PR refactor this and `is_quant_method_supported` in the kernel # unit tests to a common utility function. Currently the use of -# `is_quant_method_supported` conflates kernels with quantization methods -# an assumption which is breaking down as quantizations methods can have -# have kernels and some kernels support multiple quantization methods. -IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90) # For testing quantized linear kernels diff --git a/tests/kernels/quantization/test_cutlass_w4a8_moe.py b/tests/kernels/quantization/test_cutlass_w4a8_moe.py index b8c10d051056..4050314a1100 100644 --- a/tests/kernels/quantization/test_cutlass_w4a8_moe.py +++ b/tests/kernels/quantization/test_cutlass_w4a8_moe.py @@ -17,9 +17,8 @@ from vllm.scalar_type import ScalarType, scalar_types from vllm.utils.torch_utils import set_random_seed -capability = current_platform.get_device_capability() IS_SUPPORTED_BY_GPU = ( - current_platform.is_cuda() and capability is not None and capability[0] >= 9 + current_platform.is_cuda() and current_platform.is_device_capability(90) ) From fff10d20728780c58a9427986174eb7485ab5c12 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Fri, 9 Oct 2026 04:16:22 +0900 Subject: [PATCH 235/952] [Attention] Allow an NVFP4 KV cache on SM90 with fa2 (#60620) Signed-off-by: lesj0610 Co-authored-by: Misha Goin --- tests/kernels/attention/test_flashinfer.py | 2 +- vllm/v1/attention/backends/flashinfer.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index 0d9aecd31e7d..c56af3a6e852 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -706,7 +706,7 @@ def test_flashinfer_decode_with_paged_fp8_kv( @pytest.mark.skipif( - not any(current_platform.is_device_capability_family(f) for f in (80, 120)), + not any(current_platform.is_device_capability_family(f) for f in (80, 90, 120)), reason="NVFP4 fa2 path", ) @pytest.mark.parametrize("layout", ["LBNHC", "LBHNC", "BLNHC", "BLHNC"]) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index a21720ab26e6..85b7456dd889 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -108,9 +108,7 @@ def _nvfp4_kv_on_fa2() -> bool: - return current_platform.is_device_capability_family( - 80 - ) or current_platform.is_device_capability_family(120) + return any(current_platform.is_device_capability_family(f) for f in (80, 90, 120)) _KVPair = tuple[torch.Tensor, torch.Tensor] @@ -557,7 +555,7 @@ def supports_combination( if ( kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4") - and device_capability.major in (8, 12) + and device_capability.major in (8, 9, 12) ): if head_size == 64: return "fa2 prefill misreads an NVFP4 KV cache at head_size 64" @@ -1311,7 +1309,7 @@ def _get_prefill_wrapper( ) else: # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper on SM100; SM8x and SM12x read it with fa2. + # the wrapper on SM100; SM8x, SM90 and SM12x read it with fa2. backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( self._get_workspace_buffer(), @@ -1337,7 +1335,7 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): paged_kv_indices = None paged_kv_last_page_len = None # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper on SM100; SM8x and SM12x read it with fa2. + # the wrapper on SM100; SM8x, SM90 and SM12x read it with fa2. backend = "trtllm-gen" if self.nvfp4_trtllm else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( self._get_workspace_buffer(), From e3a29aa9b33e894ee576d3d092ae7605fb4e7721 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 15:32:07 -0400 Subject: [PATCH 236/952] [Bugfix][MoE] Size the permute scratch from its inputs (#60447) Signed-off-by: mgoin Co-authored-by: Claude --- .../benchmark_moe_permute_unpermute.py | 8 - tests/kernels/moe/test_moe.py | 30 +--- .../kernels/moe/test_moe_permute_unpermute.py | 68 ++++---- .../layers/fused_moe/experts/cutlass_moe.py | 35 ++-- .../fused_moe/experts/fused_humming_moe.py | 22 +-- .../layers/fused_moe/moe_permute_unpermute.py | 153 +++++++----------- vllm/v1/worker/workspace.py | 3 +- 7 files changed, 128 insertions(+), 191 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe_permute_unpermute.py b/benchmarks/kernels/benchmark_moe_permute_unpermute.py index 6867274ebe28..f5aa21fb4bd5 100644 --- a/benchmarks/kernels/benchmark_moe_permute_unpermute.py +++ b/benchmarks/kernels/benchmark_moe_permute_unpermute.py @@ -56,13 +56,9 @@ def benchmark_permute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(i: int): @@ -135,13 +131,9 @@ def benchmark_unpermute( qhidden_states, input_gating, topk, False ) scratch = MoEPermuteScratch( - max_num_tokens=num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_experts, device=qhidden_states.device, - hidden_size=hidden_size, - hidden_dtype=qhidden_states.dtype, ) def prepare(): diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8d6deac4f025..0f786deb5498 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1437,6 +1437,7 @@ def test_humming_global_valid_shape_m( def test_humming_permute_scratch_is_shared_by_config( monkeypatch: pytest.MonkeyPatch, ): + """Layers share one scratch whose key is independent of batch shape.""" from types import SimpleNamespace from unittest.mock import Mock @@ -1447,34 +1448,19 @@ def test_humming_permute_scratch_is_shared_by_config( manager = workspace.WorkspaceManager(torch.device("cpu")) monkeypatch.setattr(workspace, "_manager", manager) - scratch_topk6 = Mock() - scratch_topk1 = Mock() - scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) + scratch = Mock() + scratch_type = Mock(side_effect=[scratch]) monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) - moe_config.moe_parallel_config.dp_size = 2 experts = SimpleNamespace(moe_config=moe_config) other_layer = SimpleNamespace(moe_config=moe_config) - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - assert ( - humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 - ) - assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - - assert scratch_type.call_count == 2 - first_call, second_call = scratch_type.call_args_list - assert first_call.kwargs["max_num_tokens"] == 1024 - assert first_call.kwargs["topk"] == 6 - assert second_call.kwargs["max_num_tokens"] == 6144 - assert second_call.kwargs["topk"] == 1 - + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert humming.HummingExpertsBase._get_permute_scratch(other_layer) is scratch manager.lock() - assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 - with pytest.raises(AssertionError, match="was not allocated during warmup"): - humming.HummingExpertsBase._get_permute_scratch(experts, 2) + assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch + assert scratch_type.call_count == 1 def test_humming_delegates_to_instance_activation(): @@ -1531,7 +1517,7 @@ def test_humming_grouped_apply_forwards_valid_prefix( num_experts=2, estimate_local_valid_shape_m=lambda _: 6, prepare_buffers=lambda *_: buffers, - _get_permute_scratch=lambda _, *, indices_only=False: object(), + _get_permute_scratch=lambda: object(), process_input=Mock( side_effect=lambda _, **kwargs: ( kwargs["inputs"], diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 76f4b57cf07f..0ab5760c3590 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -169,13 +169,9 @@ def test_moe_permute_unpermute( scratch = None if use_scratch: scratch = get_moe_permute_scratch( - max_num_tokens=n_token, - topk=topk, num_experts=n_expert, num_local_experts=n_local_expert, device=hidden_states.device, - hidden_size=n_hidden, - hidden_dtype=dtype, ) ( @@ -247,13 +243,9 @@ def test_moe_permute_reuses_scratch_buffers( _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False) scratch_config = dict( - max_num_tokens=n_token, - topk=topk, num_experts=n_expert, num_local_experts=n_expert, device=hidden_states.device, - hidden_size=n_hidden, - hidden_dtype=hidden_states.dtype, ) scratch = get_moe_permute_scratch(**scratch_config) @@ -300,15 +292,10 @@ def test_moe_permute_reuses_scratch_buffers( permuted_idx_2, ) = second - torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2) torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2) torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2) torch.testing.assert_close(permuted_idx_1, permuted_idx_2) - assert ( - permuted_hidden_states_1.untyped_storage().data_ptr() - == permuted_hidden_states_2.untyped_storage().data_ptr() - ) assert ( expert_first_token_offset_1.untyped_storage().data_ptr() == expert_first_token_offset_2.untyped_storage().data_ptr() @@ -330,14 +317,11 @@ def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None: device = torch.device("cuda:0") scratch = get_moe_permute_scratch( - max_num_tokens=128, - topk=6, num_experts=16, num_local_experts=16, device=device, - hidden_size=2048, - hidden_dtype=torch.bfloat16, ) + scratch.reserve(128 * 6) current_workspace_manager().lock() runs = [] for n_token in (128, 1, 33): @@ -372,13 +356,9 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non ubatch = 0 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch) scratch_config = dict( - max_num_tokens=64, - topk=6, num_experts=16, num_local_experts=16, device=device, - hidden_size=2048, - hidden_dtype=torch.bfloat16, ) runs = [] for i, n_token in enumerate((1, 7, 33, 64)): @@ -396,7 +376,7 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non runs.append((stream, graph, hidden, topk_ids, output, scratch)) torch.cuda.current_stream().wait_stream(stream) - assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4 + assert len({run[-1].permuted_idx.data_ptr() for run in runs}) == 4 manager.lock() for _ in range(3): for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs): @@ -423,8 +403,6 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: monkeypatch.setattr(workspace, "_manager", None) config = dict( - max_num_tokens=4, - topk=2, num_experts=4, num_local_experts=4, device=torch.device("cuda"), @@ -432,6 +410,8 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None: first = get_moe_permute_scratch(**config) second = get_moe_permute_scratch(**config) + first.reserve(8) + second.reserve(8) assert ( first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr() ) @@ -450,13 +430,9 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda") expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda") scratch = MoEPermuteScratch( - max_num_tokens=5, - topk=1, num_experts=4, num_local_experts=2, device=hidden_states.device, - hidden_size=16, - hidden_dtype=hidden_states.dtype, ) permuted, _, expert_offsets, inverse, _ = moe_permute( @@ -498,3 +474,39 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") torch.testing.assert_close(indices.flatten(), expected_inverse) torch.testing.assert_close(expert_offsets, expected_offsets) + + +def test_moe_permute_scratch_grows_until_workspace_lock(workspace_init) -> None: + """Capacity follows the inputs seen in warmup, whatever their topk or size.""" + if not moe_permute_unpermute_supported(): + pytest.skip("moe_permute_unpermute is not supported on this platform.") + + device = torch.device("cuda:0") + scratch = get_moe_permute_scratch( + num_experts=16, num_local_experts=16, device=device + ) + + def run(n_token: int, topk: int, dtype: torch.dtype) -> None: + hidden = torch.randn((n_token, 64), device=device).to(dtype) + topk_ids = torch.randint(16, (n_token, topk), device=device) + actual = moe_permute(hidden, None, topk_ids, 16, scratch=scratch) + expected = moe_permute(hidden, None, topk_ids, 16) + for a, e in zip(actual, expected): + if a is not None: + torch.testing.assert_close(a.float(), e.float()) + + run(8, 6, torch.bfloat16) + run(64, 6, torch.bfloat16) + run(384, 1, torch.float8_e4m3fn) + assert scratch.max_expanded_rows == 384 + + current_workspace_manager().lock() + run(64, 6, torch.bfloat16) + with pytest.raises(AssertionError, match="Scratch growth is not allowed"): + run(65, 6, torch.bfloat16) + + # Experts rebuilt after the lock (e.g. a weight reload) keep the capacity. + rebuilt = get_moe_permute_scratch( + num_experts=16, num_local_experts=16, device=device + ) + assert rebuilt.max_expanded_rows == 384 diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index a03651c0a93a..6f9a62f282ce 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -20,6 +20,7 @@ ) from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( MoEPermuteScratch, + get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, moe_unpermute, @@ -304,7 +305,6 @@ def __init__( self.ab_strides2 = ab_strides2 self.c_strides1 = c_strides1 self.c_strides2 = ab_strides1_c_strides2 - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def _supports_current_device() -> bool: @@ -335,15 +335,13 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceDelegate() def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def apply( self, @@ -1296,7 +1294,6 @@ def __init__( self.s_strides2[:, 0] = k self.group_size = group_size - self._permute_scratch: MoEPermuteScratch | None = None @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -1357,15 +1354,13 @@ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: return self.out_dtype if self.out_dtype is not None else act_dtype def _get_permute_scratch(self) -> MoEPermuteScratch | None: - if self._permute_scratch is None and moe_permute_unpermute_supported(): - self._permute_scratch = MoEPermuteScratch( - max_num_tokens=self.moe_config.max_num_tokens, - topk=self.moe_config.experts_per_token, - num_experts=self.moe_config.num_experts, - num_local_experts=self.moe_config.num_local_experts, - device=torch.device(self.moe_config.device), - ) - return self._permute_scratch + if not moe_permute_unpermute_supported(): + return None + return get_moe_permute_scratch( + num_experts=self.moe_config.num_experts, + num_local_experts=self.moe_config.num_local_experts, + device=torch.device(self.moe_config.device), + ) def workspace_shapes( self, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index a6d44c3141f7..66e03229a9e1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -258,25 +258,14 @@ def humming_forward( **kwargs, ) - def _get_permute_scratch( - self, topk: int, indices_only: bool = False - ) -> MoEPermuteScratch | None: + def _get_permute_scratch(self) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None - max_expanded_rows = ( - self.moe_config.max_num_tokens - * self.moe_config.dp_size - * self.moe_config.experts_per_token - ) return get_moe_permute_scratch( - max_num_tokens=math.ceil(max_expanded_rows / topk), - topk=topk, num_experts=self.moe_config.num_experts, num_local_experts=self.moe_config.num_local_experts, device=torch.device(self.moe_config.device), - hidden_size=None if indices_only else self.moe_config.hidden_dim, - hidden_dtype=None if indices_only else self.moe_config.in_dtype, ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): @@ -911,14 +900,18 @@ def apply( ) buffers["output"] = output + scratch = self._get_permute_scratch() if a1q_scale is None: - scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True) assert scratch is not None scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch) expert_offsets, scatter_idx = scatter_metadata inv_perm = scatter_idx.flatten() num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None else: + # Prequantized inputs skip w13 quantization, so its buffer is free + # to hold the permuted hidden states. + permuted_hidden_states = buffers["quanted_gate_up_input"] + assert permuted_hidden_states.dtype == hidden_states.dtype hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute( hidden_states=hidden_states, a1q_scale=a1q_scale, @@ -926,7 +919,8 @@ def apply( n_expert=global_num_experts, n_local_expert=self.num_experts, expert_map=expert_map, - scratch=self._get_permute_scratch(topk_ids.size(1)), + permuted_hidden_states=permuted_hidden_states, + scratch=scratch, ) scatter_idx = None num_valid_tokens = None diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index d83152fc696a..b4bf3c89163d 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -13,90 +13,70 @@ @dataclass class MoEPermuteScratch: - # Reused metadata buffers for repeated grouped-MoE permutes. - max_num_tokens: int - topk: int + """Reused metadata buffers for repeated grouped-MoE permutes. + + Capacity follows the largest input seen, like the shared workspace: it grows + during warmup and is fixed once the workspace manager is locked. + """ + num_experts: int num_local_experts: int device: torch.device - hidden_size: int | None = None - hidden_dtype: torch.dtype | None = None + max_expanded_rows: int = field(init=False, default=0) token_expert_indices: torch.Tensor = field(init=False) expert_first_token_offset: torch.Tensor = field(init=False) permuted_idx: torch.Tensor = field(init=False) inv_permuted_idx: torch.Tensor = field(init=False) - permuted_hidden_states: torch.Tensor | None = field(init=False, default=None) sort_workspace: torch.Tensor = field(init=False) permuted_experts_id: torch.Tensor = field(init=False) sorted_row_idx: torch.Tensor = field(init=False) topk_ids_int32: torch.Tensor = field(init=False) topk_ids_for_sort: torch.Tensor = field(init=False) - max_expanded_rows: int = field(init=False) def __post_init__(self) -> None: - assert self.max_num_tokens > 0 - assert self.topk > 0 assert self.num_experts > 0 assert self.num_local_experts > 0 - if self.hidden_size is None: - assert self.hidden_dtype is None - else: - assert self.hidden_dtype is not None - - self.max_expanded_rows = self.max_num_tokens * self.topk - self.token_expert_indices = torch.arange( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) self.expert_first_token_offset = torch.empty( self.num_local_experts + 1, dtype=torch.int64, device=self.device ) - self.permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.inv_permuted_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - if self.hidden_size is not None: - hidden_numel = self.max_expanded_rows * self.hidden_size - self.permuted_hidden_states = torch.empty( - hidden_numel, dtype=self.hidden_dtype, device=self.device + # torch.device("cuda") in config, after initialized, + # will be changed to cuda:{index}, so we need to refresh here. + self.device = self.expert_first_token_offset.device + + def reserve(self, expanded_rows: int) -> None: + """Ensure capacity for ``expanded_rows`` (num_tokens * topk) rows.""" + if expanded_rows <= self.max_expanded_rows: + return + if ( + is_workspace_manager_initialized() + and current_workspace_manager().is_locked() + ): + raise AssertionError( + f"Workspace is locked but MoE permute scratch requires " + f"{expanded_rows} rows, current capacity is " + f"{self.max_expanded_rows} rows. Scratch growth is not allowed " + "after locking." ) - self.permuted_experts_id = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.sorted_row_idx = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_int32 = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device - ) - self.topk_ids_for_sort = torch.empty( - self.max_expanded_rows, dtype=torch.int32, device=self.device + + def empty_rows() -> torch.Tensor: + return torch.empty(expanded_rows, dtype=torch.int32, device=self.device) + + self.token_expert_indices = torch.arange( + expanded_rows, dtype=torch.int32, device=self.device ) + self.permuted_idx = empty_rows() + self.inv_permuted_idx = empty_rows() + self.permuted_experts_id = empty_rows() + self.sorted_row_idx = empty_rows() + self.topk_ids_int32 = empty_rows() + self.topk_ids_for_sort = empty_rows() sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size( - self.max_expanded_rows, self.num_experts + expanded_rows, self.num_experts ) self.sort_workspace = torch.empty( sorter_size, dtype=torch.int8, device=self.device ) - # torch.device("cuda") in config, after initialized, - # will be changed to cuda:{index}, so we need to refresh here. - self.device = self.token_expert_indices.device - - def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: - n_token, n_hidden = hidden_states.shape - assert hidden_states.device == self.device - assert topk_ids.device == self.device - assert n_token <= self.max_num_tokens - assert topk_ids.size(1) == self.topk - assert topk_ids.size(0) == n_token - if self.hidden_size is not None: - assert n_hidden == self.hidden_size - assert hidden_states.dtype == self.hidden_dtype - assert self.permuted_hidden_states is not None - - def token_expert_indices_view(self, n_token: int) -> torch.Tensor: - return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk) + self.max_expanded_rows = expanded_rows def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: if topk_ids.dtype == torch.int32: @@ -114,17 +94,16 @@ def moe_prepare_scatter( ) -> tuple[torch.Tensor, torch.Tensor]: """Generate expert offsets and shared scatter/unpermute destination indices.""" assert topk_ids.device == scratch.device - assert topk_ids.size(0) <= scratch.max_num_tokens - assert topk_ids.size(1) == scratch.topk n_token, topk = topk_ids.shape expanded_rows = topk_ids.numel() + scratch.reserve(expanded_rows) inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) if expanded_rows == 0: scratch.expert_first_token_offset.zero_() return scratch.expert_first_token_offset, inverse torch.ops._moe_C.moe_prepare_scatter( scratch.prepare_topk_ids(topk_ids), - scratch.token_expert_indices_view(n_token), + scratch.token_expert_indices[:expanded_rows].view(n_token, topk), expert_map, scratch.num_experts, scratch.num_local_experts, @@ -140,13 +119,9 @@ def moe_prepare_scatter( def get_moe_permute_scratch( *, - max_num_tokens: int, - topk: int, num_experts: int, num_local_experts: int, device: torch.device, - hidden_size: int | None = None, - hidden_dtype: torch.dtype | None = None, ) -> MoEPermuteScratch: """Share scratch across sequential layers in the current ubatch and lane. @@ -155,28 +130,15 @@ def get_moe_permute_scratch( def create_scratch() -> MoEPermuteScratch: return MoEPermuteScratch( - max_num_tokens=max_num_tokens, - topk=topk, num_experts=num_experts, num_local_experts=num_local_experts, device=device, - hidden_size=hidden_size, - hidden_dtype=hidden_dtype, ) if not is_workspace_manager_initialized(): return create_scratch() - key = ( - MoEPermuteScratch, - max_num_tokens, - topk, - num_experts, - num_local_experts, - device, - hidden_size, - hidden_dtype, - ) + key = (MoEPermuteScratch, num_experts, num_local_experts, device) return current_workspace_manager().get_persistent_resource(key, create_scratch) @@ -204,9 +166,8 @@ def moe_permute( parallel shard. permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor. If None, the output tensor will be created in this function. - scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch - buffers. Validated against hidden_states and topk_ids when given, - otherwise the buffers are allocated in this function. + scratch (Optional[MoEPermuteScratch]): Optional reusable metadata + buffers. Otherwise they are allocated in this function. Returns: - permuted_hidden_states (torch.Tensor): permuted activation. @@ -227,20 +188,11 @@ def moe_permute( if n_local_expert == -1: n_local_expert = n_expert if permuted_hidden_states is None: - if scratch is None: - permuted_hidden_states = torch.empty( - (permuted_row_size, n_hidden), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - else: - scratch.validate(hidden_states, topk_ids) - hidden_numel = permuted_row_size * n_hidden - scratch_hidden_states = scratch.permuted_hidden_states - assert scratch_hidden_states is not None - permuted_hidden_states = scratch_hidden_states[:hidden_numel].view( - permuted_row_size, n_hidden - ) + permuted_hidden_states = torch.empty( + (permuted_row_size, n_hidden), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), ( f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}" f" but got {permuted_hidden_states.size()}" @@ -278,10 +230,15 @@ def moe_permute( permuted_idx, ) else: - scratch.validate(hidden_states, topk_ids) + assert hidden_states.device == scratch.device + assert topk_ids.device == scratch.device + assert topk_ids.size(0) == n_token assert n_expert == scratch.num_experts assert n_local_expert == scratch.num_local_experts - token_expert_indices = scratch.token_expert_indices_view(n_token) + scratch.reserve(permuted_row_size) + token_expert_indices = scratch.token_expert_indices[:permuted_row_size].view( + n_token, topk + ) expert_first_token_offset = scratch.expert_first_token_offset permuted_idx = scratch.permuted_idx[:permuted_row_size] permuted_idx.fill_(permuted_row_size) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index e69e55bd8216..a05534281c11 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -139,7 +139,8 @@ def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T: AssertionError. Keys must identify the resource type and configuration. Callers must not - resize or replace tensor storage held by a cached resource. Uses of the + resize or replace tensor storage held by a cached resource once the + manager is locked. Uses of the same resource must not overlap; include the CUDA stream in the key when streams can execute concurrently. From bab0ac61f8dae8020df6725c90474bdd647f1c40 Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Fri, 9 Oct 2026 03:32:46 +0800 Subject: [PATCH 237/952] [Bugfix][Frontend] Forward the reasoning wiring to the engine for batch chat completions (#59263) Signed-off-by: apex-mochen <2756823972@qq.com> Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_batched_chat_completions.py | 102 ++++++++++++++++++ .../openai/chat_completion/batch_serving.py | 13 ++- .../openai/chat_completion/serving.py | 37 ++++--- 3 files changed, 139 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py index 946615f97105..41c74bb5cef8 100644 --- a/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py +++ b/tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py @@ -10,6 +10,7 @@ from tests.entrypoints.openai.chat_completion.test_serving_chat import ( BASE_MODEL_PATHS, + CHAT_TEMPLATE, MockHFConfig, MockModelConfig, _build_renderer, @@ -24,10 +25,13 @@ ) from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, + ChatCompletionRequest, ) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput +from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_renderer import OnlineRenderer from vllm.v1.engine.async_llm import AsyncLLM @@ -454,3 +458,101 @@ def parse(self, model_output, **kwargs): ) assert prefixes == [[1], [2]] + + +REASONING_PARSER = "kimi_k3" + + +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_batched_reasoning_wiring_matches_single_chat() -> None: + """`/v1/chat/completions/batch` must hand the engine the same reasoning + wiring as `/v1/chat/completions`. + + Both `reasoning_ended` and `reasoning_parser_kwargs` are carried on the + request into `EngineCoreRequest`, where + `StructuredOutputManager._get_reasoner()` rebuilds the request-local + reasoning parser from the forwarded `chat_template_kwargs` -- "so the + structured-output gate observes the same template kwargs used by the + frontend". A divergence makes a template-kwarg-driven parser fall back to + its default (`thinking` on) for a request that disabled thinking, so + `is_reasoning_end()` keeps reporting False, the structured-output gate + never opens, and the grammar is silently never applied. + + Both frontends run over one and the same model and renderer, so the only + thing that can differ between the two calls is the serving code itself. + """ + mock_model_config = MockModelConfig() + shared_renderer = _build_renderer(mock_model_config) + + def _build_serving(cls): + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.model_config = mock_model_config + mock_engine.errored = False + mock_engine.input_processor = MagicMock() + mock_engine.renderer = shared_renderer + mock_engine.generate.side_effect = lambda *args, **kwargs: _generator(0, "ok") + + models = OpenAIServingModels(mock_engine, BASE_MODEL_PATHS) + online_renderer = OnlineRenderer( + model_config=mock_model_config, + renderer=mock_engine.renderer, + request_logger=None, + chat_template=CHAT_TEMPLATE, + chat_template_content_format="auto", + reasoning_parser=REASONING_PARSER, + ) + return mock_engine, cls( + mock_engine, + models, + response_role="assistant", + online_renderer=online_renderer, + chat_template=CHAT_TEMPLATE, + chat_template_content_format="auto", + request_logger=None, + reasoning_parser=REASONING_PARSER, + ) + + conversation = [{"role": "user", "content": "What is the capital of France?"}] + chat_template_kwargs = {"enable_thinking": False} + + single_engine, single_serving = _build_serving(OpenAIServingChat) + await single_serving.create_chat_completion( + ChatCompletionRequest( + model=MOCK_MODEL_NAME, + messages=conversation, + chat_template_kwargs=chat_template_kwargs, + ) + ) + single_call = single_engine.generate.call_args_list[-1] + + batch_engine, batch_serving = _build_serving(OpenAIServingChatBatch) + await batch_serving.create_batch_chat_completion( + BatchChatCompletionRequest( + model=MOCK_MODEL_NAME, + messages=[conversation], + chat_template_kwargs=chat_template_kwargs, + ) + ) + batch_call = batch_engine.generate.call_args_list[-1] + + # Guard against a vacuous comparison: the reference path must be wired up. + assert single_call.kwargs["reasoning_ended"] is not None + assert single_call.kwargs["reasoning_parser_kwargs"] is not None + + assert batch_call.kwargs["reasoning_ended"] == single_call.kwargs["reasoning_ended"] + assert ( + batch_call.kwargs["reasoning_parser_kwargs"] + == single_call.kwargs["reasoning_parser_kwargs"] + ) + + # ...and the forwarded kwargs must actually turn thinking off for the + # engine-side parser, i.e. the structured-output gate opens right away. + engine_chat_template_kwargs = batch_call.kwargs["reasoning_parser_kwargs"][ + "chat_template_kwargs" + ] + reasoner = ReasoningParserManager.get_reasoning_parser(REASONING_PARSER)( + tokenizer=batch_engine.renderer.tokenizer, + chat_template_kwargs=engine_chat_template_kwargs, + ) + assert reasoner.is_reasoning_end([]) is True diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index e630b97f182a..d3ea1a1ac3de 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -169,6 +169,12 @@ async def create_batch_chat_completion( else await self._get_trace_headers(raw_request.headers) ) session_id = self._get_session_id(single_request, raw_request) + chat_template_kwargs = self._effective_chat_template_kwargs(single_request) + parser = ( + self._make_parser(single_request, tokenizer, chat_template_kwargs) + if self.parser_cls is not None + else None + ) generators.append( self.engine_client.generate( engine_prompt, @@ -179,7 +185,12 @@ async def create_batch_chat_completion( priority=request.priority, data_parallel_rank=data_parallel_rank, session_id=session_id, - reasoning_ended=None, + **self._engine_reasoning_kwargs( + single_request, + parser, + chat_template_kwargs, + self._extract_prompt_components(engine_prompt).token_ids, + ), ) ) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9b0737c25d47..69b46ba67f61 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -233,6 +233,28 @@ def _engine_chat_template_kwargs( """ return chat_template_kwargs + def _engine_reasoning_kwargs( + self, + request: ChatCompletionRequest, + parser: Parser | None, + chat_template_kwargs: dict[str, Any], + prompt_token_ids: list[int] | None, + ) -> dict[str, Any]: + """`reasoning_ended` and `reasoning_parser_kwargs` for `generate()`.""" + reasoning_parser_kwargs = None + if parser is not None and parser.reasoning_parser is not None: + reasoning_parser_kwargs = { + "chat_template_kwargs": self._engine_chat_template_kwargs( + chat_template_kwargs + ), + } + return { + "reasoning_ended": request.resolve_reasoning_ended( + parser, prompt_token_ids or [] + ), + "reasoning_parser_kwargs": reasoning_parser_kwargs, + } + async def render_chat_request( self, request: ChatCompletionRequest, @@ -364,10 +386,6 @@ async def _create_chat_completion( session_id=session_id, ) else: - reasoning_ended = request.resolve_reasoning_ended( - parser, prompt_token_ids or [] - ) - generator = self.engine_client.generate( engine_input, sampling_params, @@ -377,14 +395,9 @@ async def _create_chat_completion( priority=self._get_priority(request, raw_request), data_parallel_rank=data_parallel_rank, session_id=session_id, - reasoning_ended=reasoning_ended, - reasoning_parser_kwargs={ - "chat_template_kwargs": self._engine_chat_template_kwargs( - chat_template_kwargs - ), - } - if parser is not None and parser.reasoning_parser is not None - else None, + **self._engine_reasoning_kwargs( + request, parser, chat_template_kwargs, prompt_token_ids + ), ) generators.append(generator) From c23ca06b6d944579ad9d220297441236755b4744 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Thu, 8 Oct 2026 21:38:33 +0200 Subject: [PATCH 238/952] [Bugfix][Structured Output] Feed the token that implicitly ends reasoning to the grammar (#59627) Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Misha Goin Co-authored-by: Yifan Zong <263738792+yzong-rh@users.noreply.github.com> --- tests/parser/engine/test_deepseek_v4.py | 4 +- tests/parser/engine/test_parser_engine.py | 45 ++++++++--- .../spec_decode/test_mtp_structured_output.py | 7 +- .../test_structured_output_manager.py | 48 +++++++++++- vllm/parser/engine/adapters.py | 6 +- vllm/parser/engine/parser_engine.py | 16 +++- vllm/v1/structured_output/__init__.py | 78 +++++++++++-------- 7 files changed, 151 insertions(+), 53 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 2cf3a8bb85f3..f5b66d9dc1cb 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -643,8 +643,8 @@ def test_reasoning_end_token_ids_drop_multi_token_tool_opener( # token-ID fast path; the implicit tool-call exit still goes through # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} - assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 + assert thinking_parser.find_reasoning_end([7, _THINK_END_ID, 8]) == (1, False) + assert thinking_parser.find_reasoning_end([7, 8]) == (2, False) def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 8fd7c4d2623f..dc91cef74028 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -28,12 +28,13 @@ from vllm.parser.abstract_parser import DelegatingParser from vllm.parser.engine.adapters import make_adapters from vllm.parser.engine.events import EventType, SemanticEvent -from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine import ParserEngine, ReasoningEnd from vllm.parser.engine.parser_engine_config import ( ParserEngineConfig, ParserState, Transition, ) +from vllm.parser.glm47_moe import glm47_moe_config from vllm.parser.parser_manager import ParserManager # ── Shared test configs ────────────────────────────────────────────── @@ -158,7 +159,10 @@ def test_every_reasoning_end_exit_contributes(self): (EventType.REASONING_END, EventType.TOOL_CALL_START), ), ) - assert _make_engine(cfg).reasoning_end_token_ids == {201, 202} + engine = _make_engine(cfg) + assert engine.reasoning_end_token_ids == {201, 202} + assert engine.find_reasoning_end([5, 201]) == ReasoningEnd(1, False) + assert engine.find_reasoning_end([5, 202]) == ReasoningEnd(1, True) def test_transition_staying_in_reasoning_is_ignored(self): cfg = _with_reasoning_exits(("THINK_START", ParserState.REASONING, ())) @@ -182,17 +186,40 @@ def test_unresolved_think_end_fails_closed(self): def test_config_without_reasoning_has_empty_set(self): assert _make_engine(_hermes_config()).reasoning_end_token_ids == frozenset() - def test_find_reasoning_end_offset_returns_first_match(self): + def test_find_reasoning_end_returns_first_match(self): engine = _make_engine() - assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) == 2 - assert engine.find_reasoning_end_offset([]) == 0 + assert engine.find_reasoning_end([5, 201, 6, 201]) == ReasoningEnd(1, False) + assert engine.find_reasoning_end([5, 6]) == ReasoningEnd(2, False) + assert engine.find_reasoning_end([]) == ReasoningEnd(0, False) # Rejected-draft placeholders never match. - assert engine.find_reasoning_end_offset([-1, 201]) == 1 + assert engine.find_reasoning_end([-1, 201]) == ReasoningEnd(1, False) - def test_find_reasoning_end_offset_with_empty_set_returns_none(self): + def test_find_reasoning_end_with_empty_set_returns_none(self): engine = _make_engine(_hermes_config()) - assert engine.find_reasoning_end_offset([201]) is None + assert engine.find_reasoning_end([201]) is None + + +@pytest.mark.parametrize( + ("config", "implicit_token"), + [ + pytest.param(glm47_moe_config(), "", id="glm47_moe"), + ], +) +def test_implicit_reasoning_end_of_parsers( + config: ParserEngineConfig, implicit_token: str +): + literals = sorted(set(config.token_id_terminals.values())) + vocab = {literal: 1000 + i for i, literal in enumerate(literals)} + engine = _make_engine(config, vocab=vocab) + implicit_ids = set() + for token_id in engine.reasoning_end_token_ids: + end = engine.find_reasoning_end([5, token_id, 6]) + assert end is not None and end.offset == 1 + kept = engine.extract_content_ids([5, token_id, 6]) == [token_id, 6] + assert end.implicit == kept + if end.implicit: + implicit_ids.add(token_id) + assert implicit_ids == {vocab[implicit_token]} # ── TestEventsToDelta ──────────────────────────────────────────────── diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 5e4312227d7e..8a72765db150 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -11,6 +11,7 @@ from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.engine.parser_engine import ReasoningEnd from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -113,12 +114,12 @@ def __init__(self, end_token_ids): def reasoning_end_token_ids(self): return self._end_token_ids - def find_reasoning_end_offset(self, token_ids): + def find_reasoning_end(self, token_ids): self.windows.append(list(token_ids)) for offset, token in enumerate(token_ids): if token in self._end_token_ids: - return offset - return len(token_ids) + return ReasoningEnd(offset, False) + return ReasoningEnd(len(token_ids), False) def is_reasoning_end(self, input_ids): return any(token in self._end_token_ids for token in input_ids) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 6082bcd6a7b1..fbceb3bd07e4 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -11,6 +11,8 @@ from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.config.structured_outputs import StructuredOutputsBackend +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.engine.parser_engine import ReasoningEnd from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -18,6 +20,7 @@ TOKENIZER = "gpt2" THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +IMPLICIT_THINK_END = "z" # accepted in strings, rejected at the object start EOS = "<|eos|>" # resolved to tokenizer.eos_token_id JSON_SCHEMA = '{"type": "object"}' BACKENDS = ("xgrammar", "guidance") @@ -43,6 +46,7 @@ class FlowCase: prefix: str = "" expected_validated: tuple[str, ...] | None = None reasoning_ended: bool | None = False + engine_reasoner: bool = False xfail_guidance: str | None = None @@ -59,6 +63,23 @@ def is_reasoning_end_streaming(self, input_ids, delta_ids): return self.is_reasoning_end(delta_ids) +class MockEngineReasoner(ParserEngineReasoningAdapter): + """Ends reasoning implicitly on ``marker``, which is also content.""" + + def __init__(self, tokenizer, marker: int): + self.marker = marker + + @property + def reasoning_end_token_ids(self): + return frozenset({self.marker}) + + def find_reasoning_end(self, token_ids): + ids = list(token_ids) + if self.marker in ids: + return ReasoningEnd(ids.index(self.marker), True) + return ReasoningEnd(len(ids), False) + + def _single_token(tokenizer, text: str) -> int: token_ids = tokenizer.encode(text) assert len(token_ids) == 1, (text, token_ids) @@ -331,6 +352,28 @@ def _run_real_flow( ), id="becomes_active_terminates", ), + pytest.param( + FlowCase( + prefix='{"a": "', + raw_drafts=(IMPLICIT_THINK_END,), + expected_row_pattern="UC", + expected_reasoning=True, + expect_terminated=False, + engine_reasoner=True, + ), + id="implicit_end_accepted", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", IMPLICIT_THINK_END, "{", "}"), + expected_validated=(" ",), + expected_row_pattern="UUUUU", + expected_reasoning=False, + expect_terminated=False, + engine_reasoner=True, + ), + id="implicit_end_rejected", + ), ] @@ -344,7 +387,8 @@ def test_real_flow( if backend == "guidance" and case.xfail_guidance: pytest.xfail(case.xfail_guidance) - reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + marker = IMPLICIT_THINK_END if case.engine_reasoner else THINK_END + reasoner_kwargs = {"marker": _single_token(tokenizer, marker)} manager, request = _build_harness( tokenizer, backend, @@ -352,6 +396,8 @@ def test_real_flow( reasoning_ended=case.reasoning_ended, reasoning_parser_kwargs=reasoner_kwargs, ) + if case.engine_reasoner: + manager.reasoner_cls = MockEngineReasoner raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) expected_texts = ( diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index 4f4aa9f5a1e0..cd7aaa729785 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -27,7 +27,7 @@ ChatCompletionRequest, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest - from vllm.parser.engine.parser_engine import ParserEngine + from vllm.parser.engine.parser_engine import ParserEngine, ReasoningEnd from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.utils import Tool @@ -65,8 +65,8 @@ def _skip_tool_parsing(self) -> Iterator[None]: def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: return self._parser_engine.is_reasoning_end(list(input_ids)) - def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: - return self._parser_engine.find_reasoning_end_offset(token_ids) + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: + return self._parser_engine.find_reasoning_end(token_ids) def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids) diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 527c310c1485..b04fdaa7944d 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -9,7 +9,7 @@ import json from collections.abc import Sequence from functools import cached_property -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import regex as re @@ -45,6 +45,13 @@ logger = init_logger(__name__) +class ReasoningEnd(NamedTuple): + """Where reasoning ends in a token sequence.""" + + offset: int # index of the first end token, or the sequence length + implicit: bool # the end token is also content, e.g. a tool-call start + + class ToolCallSlot: __slots__ = ( "id", @@ -638,14 +645,15 @@ def _derive_reasoning_end_token_ids( def reasoning_end_token_ids(self) -> frozenset[int]: return self._reasoning_end_token_ids - def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: + def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: + """First reasoning end in `token_ids`, and if its token is content.""" end_ids = self._reasoning_end_token_ids if not end_ids: return None for offset, token_id in enumerate(token_ids): if token_id in end_ids: - return offset - return len(token_ids) + return ReasoningEnd(offset, token_id != self._reasoning_end_token_id) + return ReasoningEnd(len(token_ids), False) def is_reasoning_end(self, input_ids: list[int]) -> bool: config = self.parser_engine_config diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 3b35d39d1447..2163a800a563 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -3,7 +3,7 @@ import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple from vllm.config import VllmConfig from vllm.logger import init_logger @@ -33,6 +33,17 @@ logger = init_logger(__name__) +class ConstraintBounds(NamedTuple): + """Where tokens start being fed to the grammar and being masked.""" + + grammar_start: int + constraint_start: int + + +_CONSTRAINED_FROM_START = ConstraintBounds(0, 0) +_NO_SPEC_TOKENS = ConstraintBounds(1, 1) + + class StructuredOutputManager: """Engine-level manager for structured output requests.""" @@ -217,32 +228,34 @@ def _async_submit_fill_bitmask( ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) - def _get_constraint_start( + def _get_constraint_bounds( self, request: "Request", spec_tokens: Sequence[int], spec_tokens_committed: bool = False, - ) -> int: + ) -> ConstraintBounds: """Return the index into `spec_tokens` where tokens should start being constrained by the grammar. Assumes `spec_tokens` are stripped of -1 padding tokens. Returns `len(spec_tokens) + 1` if no token should be constrained after accepting all `spec_tokens`. + This is `constraint_start`; `grammar_start` is one less if reasoning ends + implicitly on a content token (e.g. a tool-call start). `spec_tokens_committed` is True if spec_tokens is already in `request.all_token_ids`. """ if self.enable_in_reasoning: - return 0 + return _CONSTRAINED_FROM_START structured_req = request.structured_output_request assert structured_req is not None if structured_req.reasoning_ended: - return 0 + return _CONSTRAINED_FROM_START reasoner = self._get_reasoner(request) if reasoner is None: - return 0 + return _CONSTRAINED_FROM_START if structured_req.reasoning_ended is None: # This should be removed here, but since `openai_gptoss` @@ -250,30 +263,33 @@ def _get_constraint_start( # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. if reasoner.is_reasoning_end(request.prompt_token_ids or []): - return 0 + return _CONSTRAINED_FROM_START structured_req.reasoning_ended = False num_spec_tokens = len(spec_tokens) if num_spec_tokens <= 0: - return num_spec_tokens + 1 + return _NO_SPEC_TOKENS - # Use `find_reasoning_end_offset` to find constraint start if supported + # Engine-based parsers also report whether the end token is content. if ( isinstance(reasoner, ParserEngineReasoningAdapter) and reasoner.reasoning_end_token_ids ): - offset = reasoner.find_reasoning_end_offset(spec_tokens) - if offset is not None: - return offset + 1 - - # Fallback to `find_reasoning_end_offset` + end = reasoner.find_reasoning_end(spec_tokens) + if end is not None: + constraint_start = end.offset + 1 + if end.implicit: + return ConstraintBounds(end.offset, constraint_start) + return ConstraintBounds(constraint_start, constraint_start) + + # Fallback: the grammar starts after the end token. # TODO: Build a read-only Sequence view over all_token_ids # instead of copying the entire all_token_ids into input_ids. if spec_tokens_committed: if not spec_tokens or not reasoner.is_reasoning_end_streaming( request.all_token_ids, spec_tokens ): - return num_spec_tokens + 1 + return ConstraintBounds(num_spec_tokens + 1, num_spec_tokens + 1) input_ids = request.all_token_ids.copy() delta_ids = list(spec_tokens) @@ -281,15 +297,15 @@ def _get_constraint_start( input_ids = request.all_token_ids.copy() input_ids.extend(spec_tokens) if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): - return num_spec_tokens + 1 + return ConstraintBounds(num_spec_tokens + 1, num_spec_tokens + 1) delta_ids = list(spec_tokens) for i in range(num_spec_tokens - 1, 0, -1): input_ids.pop() delta_ids.pop() if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): - return i + 1 - return 1 + return ConstraintBounds(i + 1, i + 1) + return ConstraintBounds(1, 1) def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" @@ -297,8 +313,8 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in return spec_tokens spec_tokens = strip_speculative_padding(spec_tokens) - constraint_start = self._get_constraint_start(request, spec_tokens) - if constraint_start >= len(spec_tokens): + grammar_start, _ = self._get_constraint_bounds(request, spec_tokens) + if grammar_start >= len(spec_tokens): return spec_tokens structured_req = request.structured_output_request @@ -307,8 +323,8 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in grammar = structured_req.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - prefix = spec_tokens[:constraint_start] - validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + prefix = spec_tokens[:grammar_start] + validated = grammar.validate_tokens(spec_tokens[grammar_start:]) return prefix + validated def grammar_bitmask( @@ -358,7 +374,9 @@ def grammar_bitmask( if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self._get_constraint_start(request, ()) == 0 + apply_bitmask = ( + self._get_constraint_bounds(request, ()).constraint_start == 0 + ) batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -384,7 +402,7 @@ def grammar_bitmask( assert isinstance(grammar, StructuredOutputGrammar) req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) - constraint_start = self._get_constraint_start( + grammar_start, constraint_start = self._get_constraint_bounds( request, strip_speculative_padding(req_tokens) ) state_advancements = 0 @@ -402,7 +420,7 @@ def grammar_bitmask( self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) if token == -1: seen_padding = True - elif apply_bitmask: + elif not seen_padding and i >= grammar_start: if not grammar.is_terminated() and grammar.accept_tokens( req_id, [token] ): @@ -453,17 +471,15 @@ def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - constraint_start = self._get_constraint_start( + grammar_start, _ = self._get_constraint_bounds( request, new_token_ids, spec_tokens_committed=True ) - # Early return only when the constraint hasn't started. + # Early return only when no reasoning-end token is present. # Otherwise, latch structured_req.reasoning_ended. - if constraint_start > len(new_token_ids): + if grammar_start > len(new_token_ids): return True structured_req.reasoning_ended = True - return grammar.accept_tokens( - request.request_id, new_token_ids[constraint_start:] - ) + return grammar.accept_tokens(request.request_id, new_token_ids[grammar_start:]) def clear_backend(self) -> None: if self.backend is not None: From 1afb029839a7c84f14dae30bb08f2fea79184431 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 8 Oct 2026 15:45:44 -0400 Subject: [PATCH 239/952] [Bugfix][HiSparse] Route by swap-row capacity, not the index-group workspace (#60682) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- tests/v1/attention/test_sparse_mla_backends.py | 12 ++++++++++-- vllm/v1/attention/backends/mla/index_group.py | 2 +- vllm/v1/hisparse/runtime.py | 9 ++++++--- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 4d33075c9303..973a9e03e88b 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3296,7 +3296,8 @@ def run_kernel(self, *, q, **kwargs): # noqa: ARG001 cache = SimpleNamespace( runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) + hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)), + max_swap_rows=num_tokens, ), source_block_table=torch.empty( num_decodes, 1, dtype=torch.int32, device=device @@ -3604,15 +3605,22 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): def test_hisparse_resident_prefill_uses_attention_block_stride(): + """A batch beyond the swap-row capacity must bypass swap_in. + + The index-group workspace has one more row than the HiSparse swap rows, so + routing on its size sent a 2-token batch with max_num_seqs=1 (a piecewise + CUDA-graph capture) into swap_in, which overflowed its capacity of 1. + """ expected = torch.tensor([[19]], dtype=torch.int32) cache_handle = SimpleNamespace( all_context_pages_resident=True, view=SimpleNamespace(block_size=64, attention_block_stride=832), block_table=torch.tensor([[3]], dtype=torch.int32), + runtime=SimpleNamespace(max_swap_rows=1), ) index_group = object.__new__(HiSparseMLAIndexGroup) index_group.caches = [cache_handle] - index_group.physical_topk_indices = torch.empty((1, 1), dtype=torch.int32) + index_group.physical_topk_indices = torch.empty((2, 1), dtype=torch.int32) index_group._convert_once = MagicMock(return_value=expected) topk = torch.zeros((2, 1), dtype=torch.int32) metadata = SimpleNamespace( diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 5a18332e41e1..39f205b91d97 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -233,7 +233,7 @@ def convert_logical_to_physical_topk( cache = self.cache(layer_index) num_tokens = logical_topk_indices.shape[0] req_id_per_token = attn_metadata.req_id_per_token[:num_tokens] - if num_tokens > self.physical_topk_indices.shape[0]: + if num_tokens > cache.runtime.max_swap_rows: # Prefill-sized batches do not fit the decode residency workspace. # Non-resident prefills are staged before reaching this path. assert cache.all_context_pages_resident diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index a6378ca9141a..1f7a5ad5ca3c 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -894,14 +894,17 @@ def invalidate_written_slots( def begin_forward(self) -> None: self._swap_step = 0 + @property + def max_swap_rows(self) -> int: + return self.index_group.shared_topk.physical_topk_indices.shape[0] + def _step_rows(self, num_tokens: int) -> slice: start = self._swap_step * num_tokens stop = start + num_tokens - if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]: + if stop > self.max_swap_rows: raise ValueError( "HiSparse swap rows exceed the configured speculative decode " - f"capacity: stop={stop}, capacity=" - f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}." + f"capacity: stop={stop}, capacity={self.max_swap_rows}." ) self._swap_step += 1 return slice(start, stop) From 72d59adc2c76de8f77aad640c7a04d663a970735 Mon Sep 17 00:00:00 2001 From: Alberto Massidda Date: Thu, 8 Oct 2026 21:58:13 +0200 Subject: [PATCH 240/952] [Bugfix][KV Connector][NIXL] Transfer the whole PLE short-conv page in disaggregated P/D (#59997) Signed-off-by: Alberto Massidda Co-authored-by: Alberto Massidda Co-authored-by: Claude Fable 5.1 Co-authored-by: Artem Perevedentsev --- .../unit/test_nixl_desc_geometry.py | 60 +++++++++++++++---- .../kv_connector/v1/nixl/base_worker.py | 28 +++++++-- .../kv_connector/v1/nixl/metadata.py | 5 +- 3 files changed, 77 insertions(+), 16 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 029ffafd3ceb..9f6bc56deeb5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -1245,8 +1245,12 @@ def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2) -def _make_csa_linear_ple_worker(scratch_aliases: str = "compressed"): - """``scratch_aliases`` selects which pages the compressor ring overlays.""" +def _make_csa_linear_ple_worker( + scratch_aliases: str = "compressed", ple_page_size: int = 256 +): + """``scratch_aliases`` selects which pages the compressor ring overlays; + ``ple_page_size`` widens the PLE page beyond the page it shares a region + with (block-outer layout: both start at byte 0 of the block).""" from unittest.mock import MagicMock from vllm.config import set_current_vllm_config @@ -1305,7 +1309,7 @@ def _make_csa_linear_ple_worker(scratch_aliases: str = "compressed"): block_size=1, shapes=((6, 3),), dtypes=(torch.float16,), - page_size_padded=256, + page_size_padded=ple_page_size, mamba_type=MambaAttentionBackendEnum.SHORT_CONV, tp_replicated=True, ) @@ -1339,17 +1343,22 @@ def group(specs): ("compressed.0",), ("compressed.1",), ) - region_size = 512 page_size = 256 + # Region 0 hosts the PLE page: its blocks are as wide as the widest page. + block_strides = [max(page_size, ple_page_size)] + [page_size] * ( + len(tensor_regions) - 1 + ) + region_sizes = [2 * stride for stride in block_strides] + region_offsets = [sum(region_sizes[:index]) for index in range(len(region_sizes))] kv_cache_config = KVCacheConfig( num_blocks=2, kv_cache_tensors=[ KVCacheTensor( - size=len(tensor_regions) * region_size, + size=sum(region_sizes), layers=[layer_name], - layer_stride=region_size, - block_stride=page_size, - offset=region_index * region_size, + layer_stride=region_sizes[region_index], + block_stride=block_strides[region_index], + offset=region_offsets[region_index], ) for region_index, layer_names in enumerate(tensor_regions) for layer_name in layer_names @@ -1386,10 +1395,14 @@ def group(specs): set_current_vllm_config(vllm_config), ): worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) - tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(4)] + tensors = [ + torch.zeros((2, stride), dtype=torch.uint8) for stride in block_strides + ] worker.register_kv_caches( { - layer_name: tensors[region_index] + layer_name: tensors[region_index][ + :, : (ple_page_size if layer_name == "mamba.ple" else page_size) + ] for region_index, layer_names in enumerate(tensor_regions) for layer_name in layer_names } @@ -1523,6 +1536,7 @@ def test_csa_linear_remote_ple_is_copied_whole(): kv_cache_layout="HND", block_size=4, ssm_sizes=(24, 32), + ple_block_len=256, attn_backend_name="test", physical_blocks_per_logical_kv_block=1, ) @@ -1535,7 +1549,7 @@ def test_csa_linear_remote_ple_is_copied_whole(): assert descriptors[-2:, 0].tolist() == [0x10000, 0x10100] assert descriptors[-2:, 1].tolist() == [256, 256] - metadata.block_lens[0] = 128 + metadata.ple_block_len = 128 with pytest.raises(ValueError, match="PLE pages require identical"): worker._build_mamba_remote( metadata, @@ -1544,6 +1558,30 @@ def test_csa_linear_remote_ple_is_copied_whole(): ) +@pytest.mark.cpu_test +def test_csa_linear_ple_page_wider_than_the_shared_region_page(): + """The PLE page is discovered in the region of the main KV page registered + first (both start at byte 0 of the block); its descriptors must still cover + the whole PLE page, not the region's block_len.""" + worker = _make_csa_linear_ple_worker(ple_page_size=384) + + assert worker._ple_region_index == 0 + assert worker.block_len_per_layer[0] == 256 + assert worker.block_stride_per_layer[0] == 384 + assert worker._ple_block_len == 384 + + bases = [0x10000, 0x20000, 0x30000, 0x40000] + mamba = worker._build_mamba_local(bases) + assert mamba[-2:, 0].tolist() == [0x10000, 0x10000 + 384] + assert mamba[-2:, 1].tolist() == [384, 384] + + metadata = msgspec.msgpack.decode( + worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata + ) + assert metadata.ple_block_len == 384 + assert metadata.block_lens[0] == 256 + + def _make_ring_worker(): """Paged MLA group plus a per-layer ring group, no Mamba: the shape of a model whose sliding-window KV lives in per-request rings.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 2a2f660cf15e..7ce6e93ebb46 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -935,6 +935,9 @@ def __init__( # page it overlays, so its regions cannot be recovered from spec type. self._scratch_region_indices = list[int]() self._ple_region_index: int | None = None + # Page length of the PLE short-conv state; its region may be shared with + # pages of a different (smaller) size, so it cannot use the region's. + self._ple_block_len: int | None = None # Enable different block lengths for different layers *only* when MLA is used. # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. @@ -1464,6 +1467,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): seen_base_addresses: list[int] = [] self._ssm_region_indices = [] self._scratch_region_indices = [] + self._ple_block_len = None self._ple_region_index = None packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values())) @@ -1751,6 +1755,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): assert self._is_csa_linear assert self._ple_region_index in (None, region_index) self._ple_region_index = region_index + self._ple_block_len = block_len elif region_index not in self._ssm_region_indices: self._ssm_region_indices.append(region_index) elif ( @@ -1885,6 +1890,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): else self.host_buffer_kv_cache_layout, block_size=self.block_size, ssm_sizes=self._mamba_ssm_size, + ple_block_len=self._ple_block_len, attn_backend_name=self.backend_name, physical_blocks_per_logical_kv_block=( self._physical_blocks_per_logical_kv_block @@ -1906,6 +1912,19 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): agent_metadata_bytes=encoder.encode(agent_metadata), ) + def _ple_page_len(self) -> int: + """Transfer length of one PLE state page. + + The PLE page shares its region with the first page of other cache + groups (block-outer layout), whose block_len may be smaller; the + descriptors must cover the whole PLE page or part of the conv state + never reaches the decode side.""" + assert self._ple_region_index is not None + assert self._ple_block_len is not None, ( + "PLE region registered without its page length" + ) + return self._ple_block_len + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for local mamba blocks with DS conv layout, as an Nx3 uint64 array. @@ -1960,7 +1979,7 @@ def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id)) if (region_index := self._ple_region_index) is not None: - block_len = self.block_len_per_layer[region_index] * physical_per_logical + block_len = self._ple_page_len() * physical_per_logical block_stride = ( self.block_stride_per_layer[region_index] * physical_per_logical ) @@ -2019,11 +2038,12 @@ def _build_mamba_remote( if (region_index := self._ple_region_index) is not None: local_block_len = ( - self.block_len_per_layer[region_index] - * self._physical_blocks_per_logical_kv_block + self._ple_page_len() * self._physical_blocks_per_logical_kv_block ) + # Same NIXL_CONNECTOR_VERSION on both sides, so the field is set. + assert nixl_agent_meta.ple_block_len is not None remote_block_len = ( - nixl_agent_meta.block_lens[region_index] * remote_physical_per_logical + nixl_agent_meta.ple_block_len * remote_physical_per_logical ) if local_block_len != remote_block_len: raise ValueError( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 8d0aad763358..a33f620a0a0d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -49,8 +49,9 @@ # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata # 12: Add per-region member names for PP push # 13: Add packed-member layouts and order-independent packed-push backend hashes +# 14: Add Qwen3.8-Flash-Next PLE page length - #59997 # -NIXL_CONNECTOR_VERSION: int = 13 +NIXL_CONNECTOR_VERSION: int = 14 @dataclass @@ -73,6 +74,8 @@ class NixlAgentMetadata: region_mem_types: list[str] | None = None dcp_size: int = 1 pcp_size: int = 1 + # Page length of the PLE short-conv state (None when the model has none). + ple_block_len: int | None = None # Layer names sharing each advertised region, in region order. region_members: list[list[str]] = field(default_factory=list) # Packed member -> (byte offset in its region's block, bytes per page). From ae9adec9059b6c6d831cfa46ce966a3d629a4526 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 8 Oct 2026 16:21:05 -0400 Subject: [PATCH 241/952] [Frontend] Expose /v1/systemone without an opt-in flag (#60677) Signed-off-by: Francisco Javier Arceo <4163062+franciscojavierarceo@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Cyrus Leung --- .../online_serving/structured_decisions.md | 8 ++--- .../structured_decisions/test_strategies.py | 35 +++++++++++++++---- .../test_structured_decisions_e2e.py | 1 - .../entrypoints/openai/test_openai_schema.py | 8 ++++- vllm/entrypoints/generate/api_router.py | 25 ++++++------- .../structured_decisions/api_router.py | 3 +- .../structured_decisions/strategies.py | 7 ++-- vllm/entrypoints/launchers/cli_args.py | 2 -- 8 files changed, 56 insertions(+), 33 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index dd26f2130164..7de935a62f2b 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -3,10 +3,10 @@ The `/v1/systemone` endpoint answers a set of typed questions about a state and returns a probability for every allowed answer. -The endpoint is off unless the server starts with -`--enable-structured-decisions`. With the flag, startup fails for a model that -does not support structured decisions, and with `--logprobs-mode raw_logits` or -`processed_logits`, since `label_mass` needs log probabilities. +The endpoint is available by default on generation models. Models without a +supported read strategy return 501. Supported architectures require +`--logprobs-mode raw_logprobs` or `processed_logprobs`, since `label_mass` needs +log probabilities. ## How it works diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 017788f7b1d8..51bff19c8ca2 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -5,6 +5,7 @@ import pytest from fastapi import FastAPI +from fastapi.testclient import TestClient from vllm.entrypoints.generate.structured_decisions.api_router import ( register_structured_decisions_api_router, @@ -16,6 +17,9 @@ reply_label_ids, select_read_strategy, ) +from vllm.entrypoints.serve.exception_handling.register import ( + init_exception_handler, +) def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: @@ -32,13 +36,30 @@ def test_strategy_selection(): select_read_strategy(model(qwen, "raw_logits")) -def test_route_needs_the_flag(): - for enabled in (False, True): - app = FastAPI() - app.state.args = SimpleNamespace(enable_structured_decisions=enabled) - register_structured_decisions_api_router(app) - paths = {getattr(route, "path", None) for route in app.routes} - assert ("/v1/systemone" in paths) == enabled +def test_route_is_registered_by_default(): + app = FastAPI() + register_structured_decisions_api_router(app) + paths = {getattr(route, "path", None) for route in app.routes} + assert "/v1/systemone" in paths + + +def test_unsupported_model_returns_501(): + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.serving_structured_decisions = None + init_exception_handler(app) + register_structured_decisions_api_router(app) + + with TestClient(app) as client: + response = client.post( + "/v1/systemone", + json={ + "model": "unsupported", + "state": "x", + "questions": {"answer": {"type": "choice", "criteria": {"yes": None}}}, + }, + ) + assert response.status_code == 501 def test_structured_state_keeps_unicode_in_prompt(): diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py index e060b096ec51..3e0f9bc0fb66 100644 --- a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -19,7 +19,6 @@ def server(): "--max-num-seqs", "32", "--enable-prefix-caching", - "--enable-structured-decisions", ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 18b01b85a8a3..a405550c4cf4 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -7,6 +7,7 @@ import schemathesis from hypothesis import HealthCheck, settings from schemathesis import GenerationMode +from schemathesis.checks import not_a_server_error from schemathesis.config import ( ChecksConfig, CoveragePhaseConfig, @@ -161,8 +162,13 @@ def test_openapi_stateless(case: schemathesis.Case): }.get(key, DEFAULT_TIMEOUT_SECONDS) # No need to verify SSL certificate for localhost - case.call_and_validate( + systemone = case.operation.path == "/v1/systemone" + response = case.call_and_validate( verify=False, timeout=timeout, headers={"Content-Type": "application/json"}, + excluded_checks=[not_a_server_error] if systemone else None, ) + if systemone: + # SmolVLM does not support structured decisions. + assert response.status_code < 500 or response.status_code == 501 diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 70869b61b5d2..67a256192443 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -230,16 +230,11 @@ async def init_generate_state( from .structured_decisions.serving import ServingStructuredDecisions from .structured_decisions.strategies import ReadContext, select_read_strategy - strategy_cls = ( - select_read_strategy(engine_client.model_config) - if "generate" in supported_tasks - and getattr(args, "enable_structured_decisions", False) - else None - ) - state.serving_structured_decisions = ( - ServingStructuredDecisions( - state.openai_serving_models, - strategy_cls( + strategy = None + if "generate" in supported_tasks: + try: + strategy_cls = select_read_strategy(engine_client.model_config) + strategy = strategy_cls( ReadContext( engine_client=engine_client, online_renderer=state.online_renderer, @@ -247,9 +242,15 @@ async def init_generate_state( chat_template_content_format=args.chat_template_content_format, default_chat_template_kwargs=default_chat_template_kwargs, ) - ), + ) + except ValueError: + strategy = None + state.serving_structured_decisions = ( + ServingStructuredDecisions( + state.openai_serving_models, + strategy, request_logger=request_logger, ) - if strategy_cls is not None + if strategy is not None else None ) diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py index a0fbaa518e07..5729b58b2e50 100644 --- a/vllm/entrypoints/generate/structured_decisions/api_router.py +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -49,5 +49,4 @@ async def create_structured_decision( def register_structured_decisions_api_router(app: FastAPI): - if getattr(app.state.args, "enable_structured_decisions", False): - app.include_router(router) + app.include_router(router) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 22233ff91b0e..864443cbb0b9 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -219,17 +219,16 @@ async def read( def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: - """Raises ValueError when the model cannot serve structured decisions. The - server opted in with --enable-structured-decisions, so startup fails.""" + """Raise ValueError when the model cannot serve structured decisions.""" if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: raise ValueError( - "--enable-structured-decisions does not support " + "The structured decisions API does not support " f"{model_config.architecture}. Supported architectures: " f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" ) if model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( - "--enable-structured-decisions needs --logprobs-mode raw_logprobs or " + "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" ) return NextTokenStrategy diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 18f8cb00b503..28fe79fd691e 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -163,8 +163,6 @@ class BaseFrontendArgs: enable_tokenizer_info_endpoint: bool = False """Enable the `/tokenizer_info` endpoint. May expose chat templates and other tokenizer configuration.""" - enable_structured_decisions: bool = False - """Enable the `/v1/systemone` structured decisions endpoint.""" enable_log_outputs: bool = False """If set to True, log model outputs (generations). Requires `--enable-log-requests`. Output text and finish reasons are logged at INFO, From c5c116138267ac738bc262495326d28ffff834a2 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:28:38 -0400 Subject: [PATCH 242/952] [ROCm][Build] Fail fast when `setup.py develop` deps aren't pip-preinstalled (#59447) Signed-off-by: Divakar Verma --- setup.py | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/setup.py b/setup.py index 5e49f74e6fcd..16a3d35ecad6 100644 --- a/setup.py +++ b/setup.py @@ -1311,6 +1311,74 @@ def get_vllm_version() -> str: return version +def _check_requirements_preinstalled(requirements: list[str]) -> None: + """`python setup.py develop` satisfies missing `install_requires` via + setuptools' legacy easy_install path instead of pip's wheel-aware + resolver. That path has repeatedly tried (and failed) to compile old + sdists of fast-moving C-extension packages such as aiohttp from source + against the running Python's headers, e.g. + https://github.com/vllm-project/vllm/issues/34073. Fail fast with an + actionable message instead of a cryptic compiler error. + + Skipped under `--no-deps`, which makes setuptools skip `install_requires` + processing entirely, so the easy_install path this guards against is + unreachable anyway. + """ + from importlib.metadata import PackageNotFoundError + from importlib.metadata import requires as installed_requires + from importlib.metadata import version as installed_version + + from packaging.requirements import Requirement + + missing = [] + checked = set() + + def _check(req_str: str, extra: str | None = None) -> None: + # Strip inline "# comment" suffixes (as pip does when reading + # requirements files); packaging.Requirement cannot parse them. + req_str = re.split(r"\s+#", req_str, maxsplit=1)[0].strip() + try: + req = Requirement(req_str) + except Exception: + return + # `extra` is only set when recursing into an extra's own deps + # below; it selects which `; extra == "..."` markers apply. + if req.marker is not None and not req.marker.evaluate( + {"extra": extra} if extra else None + ): + return + if (req.name, extra) in checked: + return + checked.add((req.name, extra)) + try: + installed = installed_version(req.name) + except PackageNotFoundError: + missing.append(req_str) + return + if req.specifier and not req.specifier.contains(installed, prereleases=True): + missing.append(f"{req_str} (found {req.name}=={installed})") + return + # A requested extra (e.g. "fastapi[standard]") pulls in packages + # that easy_install resolves too, so check those as well. + for extra_name in req.extras: + for dep in installed_requires(req.name) or []: + _check(dep, extra=extra_name) + + for req_str in requirements: + _check(req_str) + + if missing: + raise RuntimeError( + "The following dependencies are missing or outdated:\n " + + "\n ".join(missing) + + "\n\n`python setup.py develop` cannot reliably install these " + "itself (see https://github.com/vllm-project/vllm/issues/34073). " + "Install them with pip first, e.g.:\n" + " pip install -r requirements/rocm.txt\n" + "then re-run `python setup.py develop`." + ) + + def get_requirements() -> list[str]: """Get Python package dependencies from requirements.txt.""" requirements_dir = ROOT_DIR / "requirements" @@ -1355,6 +1423,8 @@ def _read_requirements(filename: str) -> list[str]: requirements = modified_requirements elif _is_hip(): requirements = _read_requirements("rocm.txt") + if "develop" in sys.argv[1:] and "--no-deps" not in sys.argv[1:]: + _check_requirements_preinstalled(requirements) elif _is_tpu(): requirements = _read_requirements("tpu.txt") elif _is_cpu(): From 2ddde19e8e0ba1b9d06a3f9b6201f7fc915d27ab Mon Sep 17 00:00:00 2001 From: sungbin1015 Date: Fri, 9 Oct 2026 05:38:17 +0900 Subject: [PATCH 243/952] [Frontend] Support min_p in the Responses API (#48084) Signed-off-by: sungbin1015 Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../online_serving/openai_compatible_server.md | 6 ++++++ tests/benchmarks/test_responses_request_func.py | 4 +--- .../openai/responses/test_sampling_params.py | 13 +++++++++++++ vllm/benchmarks/serve.py | 5 ----- vllm/entrypoints/openai/responses/protocol.py | 7 +++++++ 5 files changed, 27 insertions(+), 8 deletions(-) diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index c50910040d32..7b4e84a17fee 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -204,6 +204,12 @@ To get prompt token IDs before generation, use a self-contained Responses request without inference and returns the prompt token IDs, so a separate `/tokenize` call is not needed. +#### Sampling parameters + +In addition to the OpenAI request fields, the Responses API forwards vLLM +[sampling parameters](../../api/README.md#inference-parameters) such as `top_k`, +`min_p`, `repetition_penalty`, `seed`, `stop`, and `ignore_eos` to the engine. + #### Extra parameters The following extra parameters in the request object are supported: diff --git a/tests/benchmarks/test_responses_request_func.py b/tests/benchmarks/test_responses_request_func.py index 38f1bdc50781..1ecf6115fff0 100644 --- a/tests/benchmarks/test_responses_request_func.py +++ b/tests/benchmarks/test_responses_request_func.py @@ -183,6 +183,7 @@ def test_request_payload_is_accepted_by_the_server_schema(): "temperature": 0.7, "top_p": 0.9, "top_k": 20, + "min_p": 0.05, "frequency_penalty": 0.2, "presence_penalty": 0.1, "repetition_penalty": 1.05, @@ -217,9 +218,6 @@ def test_request_payload_is_accepted_by_the_server_schema(): assert set(seen["payload"]) <= set(ResponsesRequest.model_fields) ResponsesRequest.model_validate(seen["payload"]) - # `vllm bench serve` rejects --min-p for this backend for the same reason. - assert "min_p" not in ResponsesRequest.model_fields - def test_latency_accounting_matches_the_other_endpoints(): """`latency - ttft` must equal `sum(itl)`. diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 01fdfa4594ea..24d54b8960de 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -29,6 +29,7 @@ def test_basic_sampling_params(self): temperature=0.8, top_p=0.95, top_k=50, + min_p=0.1, max_output_tokens=100, ) @@ -37,8 +38,20 @@ def test_basic_sampling_params(self): assert sampling_params.temperature == 0.8 assert sampling_params.top_p == 0.95 assert sampling_params.top_k == 50 + assert sampling_params.min_p == 0.1 assert sampling_params.max_tokens == 100 + def test_min_p_default(self): + """Test that min_p defaults to 0.0 when unset (parity with top_k).""" + request = ResponsesRequest( + model="test-model", + input="test input", + ) + + sampling_params = request.to_sampling_params(default_max_tokens=1000) + + assert sampling_params.min_p == 0.0 + def test_extra_sampling_params(self): """Test extra sampling parameters are correctly mapped.""" request = ResponsesRequest( diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 895e5c6c1890..01a0f62aeee2 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -2196,11 +2196,6 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: "Sampling parameters are only supported by openai-compatible backends." ) - # The Responses API accepts every sampling parameter above except - # min_p, which it would silently drop as an unknown field. - if args.backend == "openai-responses" and "min_p" in sampling_params: - raise ValueError("--min-p is not supported by the Responses API.") - if "temperature" not in sampling_params: print( "WARNING: vllm bench serve no longer sets temperature==0 (greedy) " diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 12e40e98a882..d1a4ea4f6c21 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -259,6 +259,7 @@ class ResponsesRequest(OpenAIBaseModel): top_logprobs: TopLogprobsParam = 0 top_p: float | None = None top_k: int | None = None + min_p: float | None = None truncation: Literal["auto", "disabled"] | None = "disabled" user: str | None = None skip_special_tokens: bool = True @@ -444,6 +445,7 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: "temperature": 1.0, "top_p": 1.0, "top_k": 0, + "min_p": 0.0, } def extract_structured_outputs(self) -> StructuredOutputsParams | None: @@ -494,6 +496,10 @@ def to_sampling_params( top_k = default_sampling_params.get( "top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"] ) + if (min_p := self.min_p) is None: + min_p = default_sampling_params.get( + "min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"] + ) if (repetition_penalty := self.repetition_penalty) is None: repetition_penalty = default_sampling_params.get("repetition_penalty", 1.0) @@ -519,6 +525,7 @@ def to_sampling_params( watermarking=self.watermarking, top_p=top_p, top_k=top_k, + min_p=min_p, max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, stop=stop, From 240785b82c299c6c08bc5239e997cf62d67cf7bb Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 8 Oct 2026 13:44:04 -0700 Subject: [PATCH 244/952] [Bugfix] Keep the GLM-5.3 kpool tail and Qwen4 QSA ring out of the null block (#59528) Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- .../qwen4_exp/test_qsa_pre_indexer_amd.py | 5 ++-- tests/models/qwen4_exp/test_qsa_reference.py | 18 ++++++----- .../attention/test_kpool_tail_slot_mapping.py | 30 ++++++++++++------- vllm/models/qwen4_exp/common/qsa_cache.py | 7 +++-- vllm/v1/attention/backends/mla/indexer.py | 9 ++++-- 5 files changed, 45 insertions(+), 24 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py index 8a39504e0067..7698198923ce 100644 --- a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py +++ b/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py @@ -55,9 +55,10 @@ def _reference_gemma_rmsnorm( def _make_block_table(block_counts): - num_blocks = sum(block_counts) + # Block 0 is the null block, which no request owns. + num_blocks = sum(block_counts) + 1 table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32) - physical_blocks = torch.randperm(num_blocks) + physical_blocks = torch.randperm(sum(block_counts)) + 1 offset = 0 for request, count in enumerate(block_counts): table[request, :count] = physical_blocks[offset : offset + count] diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index bdcc4e342370..a55cc028f9db 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -315,7 +315,7 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device) token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device) - block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device) + block_table = torch.tensor([[1], [3], [2]], dtype=torch.int32, device=device) common = SimpleNamespace( num_actual_tokens=16, num_reqs=3, @@ -340,10 +340,10 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: 4, -1, -1, - 3, - 0, - 1, - 2, + 15, + 12, + 13, + 14, -1, -1, -1, @@ -351,6 +351,10 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: assert metadata.slot_mapping.tolist() == expected + # A dummy batch puts every request on the null block, which owns no ring. + block_table.zero_() + assert builder.build(0, common).slot_mapping.tolist() == [-1] * 16 + @pytest.mark.parametrize("chunk_start", list(range(8))) def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None: @@ -367,7 +371,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> query_len = num_spec + 1 slots = qsa_cache.circular_qsa_slot_mapping( - torch.tensor([[0]], dtype=torch.int32), + torch.tensor([[1]], dtype=torch.int32), torch.zeros(query_len, dtype=torch.int32), torch.arange(chunk_start, chunk_start + query_len), capacity, @@ -375,7 +379,7 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> ) committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start) - assert set(slots.tolist()).isdisjoint((committed % capacity).tolist()) + assert set((slots % capacity).tolist()).isdisjoint((committed % capacity).tolist()) def _qsa_key_cache( diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index da25eaefb23b..183c63ec2ffd 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -227,8 +227,10 @@ def test_circular_mapping_matches_generic_for_short_requests(prompt_len): def test_circular_mapping_preserves_padding_and_empty_batch(): - own_blocks = [5, 9] - per_req = [list(range(10)), list(range(12))] + # The last request is a padding request on the null block, which owns no + # tail block and must never be written. + own_blocks = [5, 9, 0] + per_req = [list(range(10)), list(range(12)), [0, 1]] padded_len = sum(len(p) for p in per_req) + 8 positions, qsl, slot_mapping, num_actual, num_reqs = make_batch( per_req, padded_len=padded_len @@ -237,7 +239,8 @@ def test_circular_mapping_preserves_padding_and_empty_batch(): out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) assert out.shape == slot_mapping.shape - assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1)) + pad_start = int(qsl[-2]) + assert torch.equal(out[pad_start:], torch.full_like(out[pad_start:], -1)) empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs) assert torch.equal(empty, slot_mapping) @@ -430,25 +433,27 @@ def test_interleaved_decode_pollution_legacy_vs_circular(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @pytest.mark.parametrize( - "per_req,num_actual,padded_len", + "per_req,own_blocks,num_actual,padded_len", [ - ([list(range(10)), list(range(12))], 22, 22), - ([list(range(10)), list(range(12))], 22, 30), - ([[3, 4], [0], [7, 8, 9]], 6, 8), - ([[5]], 1, 1), + ([list(range(10)), list(range(12))], [5, 6], 22, 22), + ([list(range(10)), list(range(12))], [5, 6], 22, 30), + ([[3, 4], [0], [7, 8, 9]], [5, 0, 7], 6, 8), + ([[5]], [5], 1, 1), + ([[0], [0]], [0, 0], 2, 4), # dummy run: every row on the null block ], ) -def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): +def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len): """The CUDA (Triton) path must match the CPU torch reference, including tokens between the last request boundary and num_actual_tokens (mapped to - the last request) and untouched padding beyond num_actual.""" + the last request), untouched padding beyond num_actual, and PAD for + requests on the null block.""" positions, qsl, slot_mapping, _, num_reqs = make_batch( per_req, padded_len=padded_len ) # Replace the all--1 placeholder slots with sentinel values to check the # padding range is copied through untouched. slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 - bt = make_tail_block_table(list(range(5, 5 + num_reqs))) + bt = make_tail_block_table(own_blocks) ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) got = circular_tail_slots( @@ -460,6 +465,9 @@ def test_triton_mapping_matches_cpu(per_req, num_actual, padded_len): num_reqs, ) torch.testing.assert_close(got.cpu(), ref) + for req, blk in enumerate(own_blocks): + if blk == 0: + assert (got[int(qsl[req]) : int(qsl[req + 1])] == -1).all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index 13228185ec3c..f960dea07194 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -127,7 +127,9 @@ def circular_qsa_slot_mapping( valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) safe_requests = requests.clamp(0, block_table.shape[0] - 1) physical_blocks = block_table[safe_requests, 0].long() - valid &= physical_blocks >= 0 + # Dummy and padding requests own no ring and sit on the null block, + # which must never be written. + valid &= physical_blocks > 0 slots = physical_blocks * compressor_state_size + positions.remainder( compressor_state_size ) @@ -288,7 +290,8 @@ def _build_qsa_metadata_kernel( mask=valid, other=-1, ) - valid &= physical_block >= 0 + # Dummy and padding requests own no ring and sit on the null block. + valid &= physical_block > 0 slot = physical_block * circular_buffer_size + ( logical_position % circular_buffer_size ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 8b186502f2a9..1005b8e4246a 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -702,11 +702,15 @@ def _kpool_tail_slot_mapping_kernel( end = tl.load(query_start_loc_ptr + pid + 1) end = tl.where(pid == num_reqs - 1, num_actual_tokens, end) own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64) + # Every scheduled request owns a tail block; the runners give dummy and + # padding requests the null block, which must never be written. + has_block = own_block != 0 for i in range(start, end, BLOCK): offs = i + tl.arange(0, BLOCK) mask = offs < end pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64) - tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask) + slot = tl.where(has_block, own_block * kpool + pos % kpool, -1) + tl.store(out_ptr + offs, slot, mask=mask) else: offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK) mask = offs < num_tokens @@ -759,7 +763,8 @@ def compute_kpool_tail_slot_mapping( req = req.clamp_(min=0, max=num_reqs - 1) own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64) pos = positions[:num_actual_tokens].to(torch.int64) - out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool) + slots = own_block * kpool + torch.remainder(pos, kpool) + out[:num_actual_tokens] = torch.where(own_block != 0, slots, -1) return out From 8d10c38b7304aa42432f4117f0d197ff46258ecb Mon Sep 17 00:00:00 2001 From: Kamran <79936503+spa5k@users.noreply.github.com> Date: Fri, 9 Oct 2026 03:12:56 +0530 Subject: [PATCH 245/952] [Perf] Reduce fallback decode CPU by 69% at 32K (#60497) Signed-off-by: spa5k <79936503+spa5k@users.noreply.github.com> Co-authored-by: Codex --- tests/tokenizers_/test_detokenize.py | 48 +++++++++++++++++++++++++++- vllm/tokenizers/detokenizer_utils.py | 20 ++++++++---- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/tests/tokenizers_/test_detokenize.py b/tests/tokenizers_/test_detokenize.py index d0117e4711c0..1dd2a000a6fa 100644 --- a/tests/tokenizers_/test_detokenize.py +++ b/tests/tokenizers_/test_detokenize.py @@ -8,7 +8,11 @@ from transformers import AutoTokenizer, PythonBackend, TokenizersBackend from vllm.sampling_params import SamplingParams -from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens +from vllm.tokenizers.detokenizer_utils import ( + convert_ids_list_to_tokens, + convert_prompt_ids_to_tokens, + detokenize_incrementally, +) from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.protocol import TokenizerLike from vllm.v1.engine import EngineCoreRequest @@ -98,6 +102,48 @@ def tokenizer(tokenizer_name): ) +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_slow_decode_long_history_preserves_text_and_token_ids(tokenizer): + truth = "A map. 我很感谢你的热情 🙂🌶️\n" * 32 + ids = tokenizer.encode(truth, add_special_tokens=False) + text, output_ids = _run_incremental_decode( + tokenizer, ids, skip_special_tokens=True, starting_index=0, fast=False + ) + assert text == tokenizer.decode(ids, skip_special_tokens=True) + assert output_ids == ids + + +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_first_incremental_decode_returns_the_prompt_window(tokenizer): + ids = tokenizer.encode("A map in a library. " * 8, add_special_tokens=False) + prompt_tokens, _, _ = convert_prompt_ids_to_tokens( + tokenizer, ids[:-1], skip_special_tokens=True + ) + new_token = tokenizer.convert_ids_to_tokens([ids[-1]], skip_special_tokens=True) + tokens, _, _, _ = detokenize_incrementally( + tokenizer, ids, None, 0, 0, skip_special_tokens=True + ) + assert tokens == prompt_tokens + new_token + + +@pytest.mark.parametrize( + "tokenizer_name", ["openai-community/gpt2", "mistralai/Pixtral-12B-2409"] +) +def test_incremental_decode_past_the_window_emits_no_text(tokenizer): + ids = tokenizer.encode("A map in a library.", add_special_tokens=False) + tokens = tokenizer.convert_ids_to_tokens(ids[:-1], skip_special_tokens=True) + offset = len(tokens) + 1 + _, text, prefix_offset, read_offset = detokenize_incrementally( + tokenizer, ids, tokens, offset, offset, skip_special_tokens=True + ) + assert text == "" + assert (prefix_offset, read_offset) == (offset, offset) + + @pytest.mark.parametrize("tokenizer_name", ["mistralai/Pixtral-12B-2409"]) @pytest.mark.parametrize( "truth", diff --git a/vllm/tokenizers/detokenizer_utils.py b/vllm/tokenizers/detokenizer_utils.py index ddaf512442c2..fc2ce4d44c82 100644 --- a/vllm/tokenizers/detokenizer_utils.py +++ b/vllm/tokenizers/detokenizer_utils.py @@ -230,7 +230,13 @@ def detokenize_incrementally( _replace_none_with_empty(new_tokens) # type: ignore[arg-type] else: new_tokens = [""] - output_tokens = prev_tokens + new_tokens + output_token_count = len(prev_tokens) + len(new_tokens) + window_start = ( + 0 if is_first_iter else min(prefix_offset, read_offset, len(prev_tokens)) + ) + output_tokens = ( + prev_tokens[window_start:] if window_start else prev_tokens + ) + new_tokens # If this is the first iteration, return all tokens. if is_first_iter: @@ -241,19 +247,21 @@ def detokenize_incrementally( # surrounding ids. if tokenizer.is_fast or not tokenizer.get_added_vocab(): prefix_text = tokenizer.convert_tokens_to_string( - output_tokens[prefix_offset:read_offset] + output_tokens[prefix_offset - window_start : read_offset - window_start] + ) + new_text = tokenizer.convert_tokens_to_string( + output_tokens[prefix_offset - window_start :] ) - new_text = tokenizer.convert_tokens_to_string(output_tokens[prefix_offset:]) else: prefix_text = _convert_tokens_to_string_with_added_encoders( tokenizer, - output_tokens[prefix_offset:read_offset], + output_tokens[prefix_offset - window_start : read_offset - window_start], skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, ) new_text = _convert_tokens_to_string_with_added_encoders( tokenizer, - output_tokens[prefix_offset:], + output_tokens[prefix_offset - window_start :], skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, ) @@ -266,4 +274,4 @@ def detokenize_incrementally( return new_tokens, "", prefix_offset, read_offset new_text = new_text[len(prefix_text) :] - return new_tokens, new_text, read_offset, len(output_tokens) + return new_tokens, new_text, read_offset, output_token_count From ab905a885dfbfc60a2c02286cc9c608c93884de3 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Thu, 8 Oct 2026 18:22:20 -0400 Subject: [PATCH 246/952] [CI] Run GH200 test on the shared arm64 CI image (#60253) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .buildkite/hardware_tests/gh200.yaml | 13 ++++++--- .../scripts/hardware_ci/run-gh200-test.sh | 29 ------------------- 2 files changed, 9 insertions(+), 33 deletions(-) delete mode 100644 .buildkite/scripts/hardware_ci/run-gh200-test.sh diff --git a/.buildkite/hardware_tests/gh200.yaml b/.buildkite/hardware_tests/gh200.yaml index 2d42aa45416c..54f2e276327e 100644 --- a/.buildkite/hardware_tests/gh200.yaml +++ b/.buildkite/hardware_tests/gh200.yaml @@ -2,10 +2,15 @@ group: Hardware steps: - label: "GH200 Test" key: gh200-test + timeout_in_minutes: 30 soft_fail: true device: gh200 - no_plugin: true optional: true - commands: - - nvidia-smi - - bash .buildkite/scripts/hardware_ci/run-gh200-test.sh + working_dir: "/vllm-workspace" + depends_on: + - arm64-image-build + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - nvidia-smi + - python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B diff --git a/.buildkite/scripts/hardware_ci/run-gh200-test.sh b/.buildkite/scripts/hardware_ci/run-gh200-test.sh deleted file mode 100644 index 20295fc25e90..000000000000 --- a/.buildkite/scripts/hardware_ci/run-gh200-test.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash - -# This script build the GH200 docker image and run the offline inference inside the container. -# It serves a sanity check for compilation and basic model usage. -set -ex - -# Skip the new torch installation during build since we are using the specified version for arm64 in the Dockerfile -python3 tools/use_existing_torch.py - -# Try building the docker image -DOCKER_BUILDKIT=1 docker build . \ - --file docker/Dockerfile \ - --target vllm-openai \ - --platform "linux/arm64" \ - -t gh200-test \ - --build-arg max_jobs=66 \ - --build-arg nvcc_threads=2 \ - --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0-b8b5f17a7d9ccfc25bbc5cf17b3fcea12964a042 \ - --build-arg torch_cuda_arch_list="9.0+PTX" - -# Setup cleanup -remove_docker_container() { docker rm -f gh200-test || true; } -trap remove_docker_container EXIT -remove_docker_container - -# Run the image and test offline inference -docker run -e HF_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c ' - python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B -' From 4274ae956c7df16b07f6d68ba400e6321de1f47e Mon Sep 17 00:00:00 2001 From: DGlan <162426117+DGlan@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:38:16 -0400 Subject: [PATCH 247/952] [Bugfix][EC Connector] Release shared memory on NIXL setup failure (#59875) Signed-off-by: DGlan <162426117+DGlan@users.noreply.github.com> --- .../unit/test_scheduler_nixl_ctor.py | 68 +++++++++++++++++++ .../ec_connector/cpu/scheduler/__init__.py | 53 +++++++-------- 2 files changed, 92 insertions(+), 29 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py index 7d2622c44ebc..8cad9ec65c4d 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import uuid +from pathlib import Path import pytest import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod +import vllm.distributed.nixl_utils as nixl_utils from tests.v1.ec_connector.unit.utils import create_ec_vllm_config from vllm.config.ec_transfer import ECTransferConfig from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( @@ -57,6 +59,72 @@ def test_string_false_in_extra_config_leaves_nixl_off(monkeypatch): s.shutdown() +@pytest.mark.cpu_test +def test_missing_nixl_releases_shared_region(monkeypatch): + """A failed constructor leaves no scheduler for the caller to shut down.""" + # Import before patching to avoid caching the missing-NIXL values. + from vllm.distributed.ec_transfer.ec_connector.cpu.data import nixl # noqa: F401 + + region = _region() + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) + monkeypatch.setattr(nixl_utils, "NixlWrapper", None) + monkeypatch.setattr(nixl_utils, "nixl_agent_config", None) + cfg = create_ec_vllm_config(ec_role="ec_both") + cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True + + try: + with pytest.raises(RuntimeError, match="ec_enable_nixl requires NIXL"): + ECCPUScheduler(cfg) + + assert not Path(region._mmap_path).exists() + assert region._fd is None + assert region._mmap_obj is None + finally: + region.cleanup() + + +@pytest.mark.cpu_test +def test_control_transport_failure_releases_nixl_and_shared_region(monkeypatch): + import vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq as zmq_mod + import vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl as data_mod + + class RegisteredData: + registered = True + + def deregister(self): + self.registered = False + + data = RegisteredData() + error = RuntimeError("control transport failed") + + def fail_control_transport(): + raise error + + region = _region() + monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) + monkeypatch.setattr(nixl_utils, "NixlWrapper", object()) + monkeypatch.setattr(nixl_utils, "nixl_agent_config", object()) + monkeypatch.setattr(data_mod, "NixlDataTransport", lambda **kwargs: data) + monkeypatch.setattr(zmq_mod, "ZmqClientTransport", fail_control_transport) + cfg = create_ec_vllm_config(ec_role="ec_consumer") + cfg.model_config.hf_config = None + cfg.model_config.get_inputs_embeds_size.return_value = 8 + cfg.model_config.model = "test-model" + cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True + + try: + with pytest.raises(RuntimeError, match="control transport failed") as exc: + ECCPUScheduler(cfg) + + assert exc.value is error + assert not data.registered + assert not Path(region._mmap_path).exists() + assert region._fd is None + assert region._mmap_obj is None + finally: + region.cleanup() + + @pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package") def test_gate_on_wires_data_transport_and_producer_session(monkeypatch): # Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index f1d22a5714a6..65403980b9d0 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -115,7 +115,11 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._element_size: int = 0 self._ack_timeout_s: float = 0.0 if self._nixl_enabled: - self._setup_nixl(vllm_config) + try: + self._setup_nixl(vllm_config) + except Exception: + self.shutdown() + raise def _setup_nixl(self, vllm_config: "VllmConfig") -> None: # Lazy imports keep nixl/zmq off the gate-off path. @@ -165,35 +169,26 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT - # Registering the region with NIXL and binding the control sockets are - # the first side effects here. __init__ propagates a failure, so the - # caller never receives a scheduler it could shut down: unwind through - # the same teardown shutdown() uses. - try: - self._data = NixlDataTransport( - agent_name=engine_id, - base_ptr=self._region.blocks.data_ptr(), - num_blocks=self._region.num_blocks, - block_size_bytes=self._region.block_size_bytes, - total_size_bytes=self._region.num_blocks - * self._region.block_size_bytes, + self._data = NixlDataTransport( + agent_name=engine_id, + base_ptr=self._region.blocks.data_ptr(), + num_blocks=self._region.num_blocks, + block_size_bytes=self._region.block_size_bytes, + total_size_bytes=self._region.num_blocks * self._region.block_size_bytes, + ) + if self._is_producer: + assert self._peer_host is not None + assert self._peer_port is not None + self._producer_session = ProducerSession( + transport=ZmqServerTransport( + host=self._peer_host, port=self._peer_port + ), + data=self._data, + cache=self._cache, + compat_hash=self._compat_hash, ) - if self._is_producer: - assert self._peer_host is not None - assert self._peer_port is not None - self._producer_session = ProducerSession( - transport=ZmqServerTransport( - host=self._peer_host, port=self._peer_port - ), - data=self._data, - cache=self._cache, - compat_hash=self._compat_hash, - ) - if self._is_consumer: - self._transport = ZmqClientTransport() - except Exception: - self._teardown_nixl() - raise + if self._is_consumer: + self._transport = ZmqClientTransport() def has_cache_item(self, identifier: str) -> bool: if not self._is_consumer: From 581692cc0fd564546ca1f58475e795e2a154cb3f Mon Sep 17 00:00:00 2001 From: Nishant Revur Date: Thu, 8 Oct 2026 18:44:50 -0700 Subject: [PATCH 248/952] [Bugfix][Spec Decode] Fix multimodal embedding merge in Llama 4 and Mistral Large 3 EAGLE drafters (#57633) Signed-off-by: Nishant Revur Co-authored-by: Claude --- ...test_gpu_target_dependent_ar_speculator.py | 24 +++++++++++++++ vllm/model_executor/models/llama4_eagle.py | 23 +++++++++++++-- .../models/mistral_large_3_eagle.py | 29 +++++++++++++++++-- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 20e86f4a7827..b63bd7a2ff05 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -337,6 +337,30 @@ def test_draft_model_multimodal_embedding_capability(model_cls, expected): assert supports_multimodal_embeddings(model_cls) is expected +@pytest.mark.parametrize( + "model_cls", [EagleLlama4ForCausalLM, EagleMistralLarge3ForCausalLM] +) +def test_mm_capable_drafter_merges_multimodal_embeddings(model_cls): + """Drafters advertising multimodal embeddings must embed text ids and + scatter the multimodal embeddings themselves; they do not inherit + SupportsMultiModal's helpers.""" + table = torch.arange(8, dtype=torch.float32).repeat(3, 1).T # [8, 3] + drafter = object.__new__(model_cls) + drafter.model = SimpleNamespace(embed_input_ids=lambda ids: table[ids]) + input_ids = torch.tensor([1, 5, 6, 2]) + + text_only = model_cls.embed_input_ids(drafter, input_ids) + assert torch.equal(text_only, table[input_ids]) + + mm_embeds = [torch.full((2, 3), -1.0)] + is_multimodal = torch.tensor([False, True, True, False]) + merged = model_cls.embed_input_ids( + drafter, input_ids, mm_embeds, is_multimodal=is_multimodal + ) + assert torch.equal(merged[[0, 3]], table[[1, 2]]) + assert torch.equal(merged[[1, 2]], mm_embeds[0]) + + def test_run_model_unpacks_tuple_return_for_mtp(monkeypatch): logits_hidden = torch.full((4, 3), 1.0) feedback_hidden = torch.full((4, 3), 2.0) diff --git a/vllm/model_executor/models/llama4_eagle.py b/vllm/model_executor/models/llama4_eagle.py index 975ee28dc22e..8634f113d063 100644 --- a/vllm/model_executor/models/llama4_eagle.py +++ b/vllm/model_executor/models/llama4_eagle.py @@ -36,10 +36,15 @@ from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM from vllm.model_executor.models.utils import extract_layer_index -from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModalEmbeddings, + _require_is_multimodal, +) from .utils import ( AutoWeightsLoader, WeightsMapper, + _merge_multimodal_embeddings, maybe_prefix, process_eagle_weight, ) @@ -204,7 +209,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def get_language_model(self) -> torch.nn.Module: return self.model - embed_input_ids = SupportsMultiModal.embed_input_ids # type: ignore + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) def forward( self, diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index d2e263f3e7d1..1398fd19fc91 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -23,8 +23,17 @@ ) from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM -from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings -from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModalEmbeddings, + _require_is_multimodal, +) +from .utils import ( + WeightsMapper, + _merge_multimodal_embeddings, + make_empty_intermediate_tensors_factory, + maybe_prefix, +) logger = init_logger(__name__) @@ -148,7 +157,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def get_language_model(self) -> torch.nn.Module: return self.model - embed_input_ids = SupportsMultiModal.embed_input_ids # type: ignore + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) def forward( self, From 201e537d98aa2ec75f12134e78f728fb0b80d542 Mon Sep 17 00:00:00 2001 From: "Brian Weston, PhD" Date: Thu, 8 Oct 2026 18:49:13 -0700 Subject: [PATCH 249/952] [Bugfix][LoRA] Validate incompatible MoE adapter formats (#59924) Signed-off-by: Brian Weston Co-authored-by: Codex Co-authored-by: Jee Jee Li Co-authored-by: linitra24 --- docs/features/lora.md | 13 ++++--- tests/lora/test_lora_manager.py | 64 +++++++++++++++++++++++++++++++++ vllm/lora/model_manager.py | 22 ++++++++++++ vllm/lora/worker_manager.py | 1 + 4 files changed, 96 insertions(+), 4 deletions(-) diff --git a/docs/features/lora.md b/docs/features/lora.md index 4c95bde8cd9f..c05bc8c0fa7a 100644 --- a/docs/features/lora.md +++ b/docs/features/lora.md @@ -280,10 +280,10 @@ curl -X POST http://localhost:8000/v1/load_lora_adapter \ !!! warning "You must know your adapter's layout" Under `--enable-mixed-moe-lora-format`, vLLM trusts whatever - `is_3d_lora_weight` the caller declares — it does **not** inspect the - checkpoint to verify. A wrong declaration will load weights into the - wrong stacked buffers and silently produce garbage outputs, with no - error at load time. Confirm the layout before serving: + `is_3d_lora_weight` the caller declares. Loading checks reject fused 3D + expert weights routed to a 2D wrapper without the required flags, but + do not validate every layout mismatch. A wrong declaration can still + cause loading failures or incorrect outputs. Confirm the layout before serving: - **2D (per-expert, megatron-style)** → set `is_3d_lora_weight: false`. Adapter keys look like `...experts.{idx}.gate_proj.lora_A.weight`, @@ -299,6 +299,11 @@ is ignored: vLLM picks the wrapper from the base model's `is_3d_moe_weight` and the adapter is required to match. The field is also ignored for non-MoE models. +For a model using a 2D MoE wrapper (such as Qwen3-MoE), loading fused 3D +expert adapter weights requires both `enable_mixed_moe_lora_format=True` +on the engine and `is_3d_lora_weight=True` on the request. Otherwise, adapter +loading raises a `ValueError` naming these settings before activation. + ## LoRA model lineage in model card The new format of `--lora-modules` is mainly to support the display of parent model information in the model card. Here's an explanation of how your current response supports this: diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 43fa1975bd79..ffb064ebfbdc 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -16,10 +16,12 @@ ReplicatedLinearWithLoRA, RowParallelLinearWithLoRA, ) +from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA from vllm.lora.lora_model import LoRAModel from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights from vllm.lora.model_manager import ( DEFAULT_LANGUAGE_WRAPPER_KEY, + AdapterLRUCache, LoRAMapping, LoRAModelManager, LRUCacheLoRAModelManager, @@ -47,6 +49,68 @@ DEFAULT_DTYPE = torch.get_default_dtype() +@pytest.mark.parametrize( + "mixed,declared_3d", [(False, False), (False, True), (True, False)] +) +def test_reject_fused_moe_lora_before_registration(mixed, declared_3d): + """Missing format flags should produce guidance without caching the adapter.""" + module_name = "model.layers.0.mlp.experts" + manager = LoRAModelManager.__new__(LoRAModelManager) + manager.lora_config = LoRAConfig(max_cpu_loras=1) + manager._registered_adapters = AdapterLRUCache(1, manager.deactivate_adapter) + manager._enable_mixed_moe_lora_format = mixed + manager.is_pooling_model = False + manager.modules = {module_name: FusedMoEWithLoRA.__new__(FusedMoEWithLoRA)} + adapter = LoRAModel( + 1, + 8, + { + module_name: LoRALayerWeights( + module_name, 8, 16, torch.zeros(16, 4), torch.zeros(4, 16) + ) + }, + ) + adapter.is_3d_lora_weight = declared_3d + + with pytest.raises(ValueError, match="enable_mixed_moe_lora_format=True") as exc: + manager.add_adapter(adapter) + assert "is_3d_lora_weight=True" in str(exc.value) + assert module_name in str(exc.value) + assert manager.list_adapters() == {} + + +@pytest.mark.parametrize( + "wrapper,mixed,declared_3d,weight_name", + [ + (FusedMoE3DWithLoRA, False, False, "experts"), + (FusedMoEWithLoRA, True, True, "experts"), + (FusedMoEWithLoRA, False, False, "experts.0.down_proj"), + (FusedMoEWithLoRA, True, False, "experts.0.down_proj"), + (FusedMoEWithLoRA, False, False, "q_proj"), + ], +) +def test_moe_format_validation_allows_supported_layouts( + wrapper, mixed, declared_3d, weight_name +): + """Preserve native 3D, mixed-format, per-expert and attention-only adapters.""" + manager = LoRAModelManager.__new__(LoRAModelManager) + manager._enable_mixed_moe_lora_format = mixed + manager.is_pooling_model = False + manager.modules = {"experts": wrapper.__new__(wrapper)} + adapter = LoRAModel( + 1, + 8, + { + weight_name: LoRALayerWeights( + weight_name, 8, 16, torch.zeros(8, 4), torch.zeros(4, 8) + ) + }, + ) + adapter.is_3d_lora_weight = declared_3d + + manager._validate_moe_lora_format(adapter) + + @pytest.mark.parametrize("device", DEVICES) def test_from_lora_tensors(qwen3_lora_files, device): tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors")) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index b97a0a412a16..8dc7b5605600 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -414,6 +414,7 @@ def _deactivate_adapter(self, lora_id: int): pass def _add_adapter(self, lora: LoRAModel): + self._validate_moe_lora_format(lora) self._create_merged_loras_inplace(lora) self._registered_adapters[lora.id] = lora @@ -1252,6 +1253,27 @@ def _get_module_to_save_weights( ) return weights + def _validate_moe_lora_format(self, lora_model: LoRAModel) -> None: + if self._enable_mixed_moe_lora_format and getattr( + lora_model, "is_3d_lora_weight", False + ): + return + for module_name, module in self.modules.items(): + if not isinstance(module, FusedMoEWithLoRA) or isinstance( + module, FusedMoE3DWithLoRA + ): + continue + module_lora = self._get_lora_layer_weights(lora_model, module_name) + if module_lora is not None and torch.is_tensor(module_lora.lora_a): + raise ValueError( + f"LoRA adapter {lora_model.id} contains fused 3D MoE weights " + f"for {module_name!r}, but the model uses a 2D MoE LoRA " + "wrapper. Start the engine with " + "enable_mixed_moe_lora_format=True " + "(--enable-mixed-moe-lora-format) and set " + "is_3d_lora_weight=True on the LoRA request." + ) + def _validate_modules_to_save(self, lora_model: LoRAModel) -> None: if not lora_model.modules_to_save: return diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index fe266879a859..9427b3e1341f 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -153,6 +153,7 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: # adapter manager can route 3D-format checkpoints through the # 3D->2D conversion when running under the universal 2D wrapper. lora.is_3d_lora_weight = lora_request.is_3d_lora_weight + self._adapter_manager._validate_moe_lora_format(lora) # Validate classification-head weights. self._adapter_manager._validate_modules_to_save(lora) From 4611c2c3d20e6f55f65c96d87f921b1be3005be9 Mon Sep 17 00:00:00 2001 From: Mia Ojeda Date: Thu, 8 Oct 2026 20:16:51 -0600 Subject: [PATCH 250/952] [BugFix][DiffusionGemma][CPU] Copy async output snapshots instead of aliasing reused buffers (#59829) Signed-off-by: ojeda-e Co-authored-by: Lucas Wilkinson --- tests/distributed/aux_output_connector/test_store.py | 3 --- tests/v1/worker/test_gpu_model_runner_v2.py | 10 ++++++++++ vllm/v1/worker/gpu/async_utils.py | 2 +- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py index 24d67c3d60e1..62d4472811fc 100644 --- a/tests/distributed/aux_output_connector/test_store.py +++ b/tests/distributed/aux_output_connector/test_store.py @@ -251,9 +251,6 @@ def test_next_step_does_not_consume_pending_output(monkeypatch): event = Mock() monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event) monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext()) - monkeypatch.setattr( - async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy() - ) worker = _make_worker(1) worker.begin_step( _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index a10f54a97c75..dec382725bc1 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -18,6 +18,7 @@ MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner @@ -357,3 +358,12 @@ def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer assert torch.equal(out, buffer[:4]) else: assert out is hidden_states + + +def test_async_copy_to_np_does_not_alias_reused_buffer(): + buffer = torch.zeros(4, dtype=torch.int64) + + snapshot = async_copy_to_np(buffer) + buffer.fill_(1) + + assert snapshot.tolist() == [0, 0, 0, 0] diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 647b0c297ccc..0d6481585e21 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -262,7 +262,7 @@ def get_output(self) -> ModelRunnerOutput: def async_copy_to_np(x: torch.Tensor) -> np.ndarray: - return x.to("cpu", non_blocking=True).numpy() + return x.to("cpu", non_blocking=True, copy=True).numpy() @contextlib.contextmanager From 295c062db8a48ea73bbd05d1174cb37d24bf86f9 Mon Sep 17 00:00:00 2001 From: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:24:56 -0700 Subject: [PATCH 251/952] Remove sudo from dockerfile command (#60488) Signed-off-by: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- docker/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index fb8fa6b264ae..d524d0fddd47 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -153,7 +153,7 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ && export SCCACHE_DOWNLOAD_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \ && tar -xzf sccache.tar.gz \ - && sudo mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && mv sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ && rm -rf sccache.tar.gz sccache-v0.8.1-${SCCACHE_ARCH}-unknown-linux-musl; \ fi From 702ce313f4f85f5f2e76b7796aa5cdc80f3595a1 Mon Sep 17 00:00:00 2001 From: zzaebok <44357534+zzaebok@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:55:36 +0800 Subject: [PATCH 252/952] [Perf][Multimodal] Sort multimodal encoder inputs by modality (#60196) Signed-off-by: Jaebok Lee Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_encoder_runner.py | 17 ++++++- tests/v1/worker/test_gpu_model_runner.py | 64 ++++++++++++++++++++++++ vllm/v1/worker/gpu/mm/encoder_runner.py | 7 ++- vllm/v1/worker/gpu/mm/lora.py | 35 +++++++------ vllm/v1/worker/gpu_model_runner.py | 16 +++--- 5 files changed, 112 insertions(+), 27 deletions(-) diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 98aaa568f0b6..69afc6ea93e7 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -78,7 +78,7 @@ def test_encoder_graph_unsupported_modality_uses_eager_output(): "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs", return_value=[("audio", 1, {"audio_values": torch.zeros(2)})], ): - result = runner.execute_mm_encoder([]) + result = runner.execute_mm_encoder([("audio", MagicMock())]) assert len(result) == 1 assert result[0] is output manager.execute.assert_not_called() @@ -290,6 +290,21 @@ def test_execute_mm_encoder_caches_outputs_without_gathering(): state.encoder_runner.gather_mm_embeddings.assert_not_called() +def test_execute_mm_encoder_sorts_by_modality(): + """Items are encoded sorted by modality; outputs keep the input order.""" + mm_kwargs = [("video", 0), ("image", 1), ("video", 2), ("image", 3)] + runner = _make_runner([], []) + runner.model = MagicMock() + runner.model.embed_multimodal.side_effect = lambda x: [torch.full((1, 1), x)] + with patch( + "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs", + side_effect=lambda items, **_: [(m, 1, {"x": x}) for m, x in items], + ) as group: + outputs = runner.execute_mm_encoder(mm_kwargs) + assert group.call_args.args[0] == [mm_kwargs[i] for i in (1, 3, 0, 2)] + assert [int(o) for o in outputs] == [0, 1, 2, 3] + + def test_execute_mm_encoder_is_a_noop_without_scheduled_items(): """A step that schedules no encoder input must not touch the encoder.""" cache = EncoderCache() diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index deb4f7b42147..3851d5dad6bd 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -624,6 +624,70 @@ def test_set_active_mm_loras_builds_tower_and_connector_mappings(): assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12) +def test_set_active_mm_loras_follows_modality_sorted_encoder_order(): + model = Mock() + model.get_mm_lora_token_counts.side_effect = ( + lambda *, modality, mm_kwargs, num_mm_embeds: (num_mm_embeds, None) + ) + lora_manager = Mock() + lora_manager.supports_tower_connector_lora.return_value = True + + encoder_cache = EncoderCache() + for req_id, modality, length in [("req-a", "video", 1), ("req-b", "image", 2)]: + encoder_cache.mm_features[req_id] = [ + MultiModalFeatureSpec( + data=None, + modality=modality, + identifier=f"{req_id}-0", + mm_position=PlaceholderRange(offset=0, length=length), + ) + ] + lora_state = LoraState(max_num_reqs=2) + lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora") + lora_state.add_request("req-a", 0, lora_request) + lora_state.add_request("req-b", 1, None) + + set_active_mm_loras( + model=model, + lora_manager=lora_manager, + encoder_cache=encoder_cache, + req_id_to_index={"req-a": 0, "req-b": 1}, + lora_state=lora_state, + scheduled_encoder_inputs={"req-a": [0], "req-b": [0]}, + ) + + # The image is encoded before the video, so its tokens come first. + _, tower_mapping = lora_manager.set_active_adapters.call_args_list[0].args + assert tower_mapping.index_mapping == (0, 0, 7) + + +def test_batch_mm_inputs_from_scheduler_sorts_by_modality(): + features = { + f"req{i}": MultiModalFeatureSpec( + data=Mock(), + modality=m, + identifier=f"hash{i}", + mm_position=PlaceholderRange(offset=i, length=1), + ) + for i, m in enumerate(["video", "image", "video"]) + } + runner = SimpleNamespace( + requests={r: SimpleNamespace(mm_features=[f]) for r, f in features.items()} + ) + scheduler_output = SimpleNamespace( + scheduled_encoder_inputs={r: [0] for r in features} + ) + + mm_hashes, mm_kwargs, mm_lora_refs = GPUModelRunner._batch_mm_inputs_from_scheduler( + runner, scheduler_output + ) + + order = ["req1", "req0", "req2"] + assert mm_hashes == [features[r].identifier for r in order] + assert mm_kwargs == [(features[r].modality, features[r].data) for r in order] + assert mm_lora_refs == [(r, features[r].mm_position) for r in order] + + def test_update_states_new_request(model_runner, dist_init): req_id = "req_0" diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 16e42d140af2..4d7ab9201394 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -150,9 +150,12 @@ def profile_encoder_cache( def execute_mm_encoder( self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]] ) -> list[torch.Tensor]: + # Stable-sort items by modality so each modality is encoded in as few + # batches as possible. The returned outputs are restored to input order. + order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0]) encoder_outputs: list[torch.Tensor] = [] for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( - mm_kwargs, device=self.device, pin_memory=PIN_MEMORY + [mm_kwargs[i] for i in order], device=self.device, pin_memory=PIN_MEMORY ): cg_manager = self.cudagraph_manager cudagraph_output = ( @@ -169,7 +172,7 @@ def execute_mm_encoder( ) sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) encoder_outputs.extend(batch_outputs) - return encoder_outputs + return [encoder_outputs[i] for i in np.argsort(order)] @contextmanager def timed_encoder_operation(self, request_ids: Collection[str]): diff --git a/vllm/v1/worker/gpu/mm/lora.py b/vllm/v1/worker/gpu/mm/lora.py index 819f3b1cb36a..e73e028f30f0 100644 --- a/vllm/v1/worker/gpu/mm/lora.py +++ b/vllm/v1/worker/gpu/mm/lora.py @@ -6,6 +6,7 @@ from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.lora.worker_manager import WorkerLoRAManager +from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -31,7 +32,7 @@ def set_active_mm_loras( encoder_token_counts: list[int] = [] connector_token_counts: list[int] = [] - # iterate through images + mm_items: list[tuple[int, MultiModalFeatureSpec]] = [] for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): req_idx = req_id_to_index.get(req_id) if req_idx is None: @@ -39,27 +40,29 @@ def set_active_mm_loras( lora_id = int(lora_state.lora_ids[req_idx]) mm_features = encoder_cache.mm_features[req_id] - - # iterate through visual tokens - for mm_input_id in encoder_input_ids: - pos_info = mm_features[mm_input_id].mm_position - - tower_tokens, connector_tokens = model.get_mm_lora_token_counts( - modality=mm_features[mm_input_id].modality, - mm_kwargs=mm_features[mm_input_id].data, - num_mm_embeds=pos_info.get_num_embeds(), - ) - - prompt_lora_mapping.append(lora_id) - token_lora_mapping.extend([lora_id] * tower_tokens) - encoder_token_counts.append(tower_tokens) - connector_token_counts.append(connector_tokens) + mm_items.extend((lora_id, mm_features[i]) for i in encoder_input_ids) if lora_id > 0: lora_request = lora_state.lora_requests.get(req_id) if lora_request is not None: lora_requests.add(lora_request) + # Runs items sorted by modality. + mm_items.sort(key=lambda x: x[1].modality) + for lora_id, mm_feature in mm_items: + pos_info = mm_feature.mm_position + + tower_tokens, connector_tokens = model.get_mm_lora_token_counts( + modality=mm_feature.modality, + mm_kwargs=mm_feature.data, + num_mm_embeds=pos_info.get_num_embeds(), + ) + + prompt_lora_mapping.append(lora_id) + token_lora_mapping.extend([lora_id] * tower_tokens) + encoder_token_counts.append(tower_tokens) + connector_token_counts.append(connector_tokens) + if not prompt_lora_mapping: return diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 21b0f67e2fe8..c02bdc74e22b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2929,7 +2929,14 @@ def _batch_mm_inputs_from_scheduler( mm_kwargs.append((mm_feature.modality, mm_feature.data)) mm_lora_refs.append((req_id, mm_feature.mm_position)) - return mm_hashes, mm_kwargs, mm_lora_refs + # Stable-sort by modality so each modality is encoded in as few batches + # as possible. Encoder outputs are cached by mm_hash, so order is free. + order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0]) + return ( + [mm_hashes[i] for i in order], + [mm_kwargs[i] for i in order], + [mm_lora_refs[i] for i in order], + ) def _cache_encoder_output( self, @@ -2990,13 +2997,6 @@ def _execute_mm_encoder( and scheduler_output.scheduled_encoder_inputs ) - # Batch mm inputs as much as we can: if a request in the batch has - # multiple modalities or a different modality than the previous one, - # we process it separately to preserve item order. - # FIXME(ywang96): This is a hacky way to deal with multiple modalities - # in the same batch while still being able to benefit from batching - # multimodal inputs. The proper solution should be reordering the - # encoder outputs. model = cast(SupportsMultiModal, self.model) if self.lora_config and self.lora_manager.supports_tower_connector_lora(): From a1cd0ec650852bf2ce124b3dea930a38f0d243da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Delhaye?= Date: Fri, 9 Oct 2026 05:01:51 +0200 Subject: [PATCH 253/952] [Model] Add zerank-2 reranker (#60415) Signed-off-by: rdlhio --- docs/models/pooling_models/scoring.md | 9 +- examples/pooling/score/template/zerank2.jinja | 13 +++ .../pooling/score/using_template_offline.py | 2 + .../pooling/score/using_template_online.py | 1 + .../language/pooling_mteb_test/test_zerank.py | 95 +++++++++++++++++++ 5 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 examples/pooling/score/template/zerank2.jinja create mode 100644 tests/models/language/pooling_mteb_test/test_zerank.py diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 1c33c6852d7e..0198810cbf68 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -49,7 +49,7 @@ The three supported scoring functions are as illustrated in the figure below. | `LlamaBidirectionalForSequenceClassification`C | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2`, etc. | [nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) | ✅︎ | ✅︎ | | `ModernBertForSequenceClassification` | ModernBERT-based | `Alibaba-NLP/gte-reranker-modernbert-base`, etc. | N/A | | | | `Qwen2ForSequenceClassification`C | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2`(see note), etc. | [mxbai_rerank_v2.jinja](../../../examples/pooling/score/template/mxbai_rerank_v2.jinja) | ✅︎ | ✅︎ | -| `Qwen3ForSequenceClassification`C | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja) | ✅︎ | ✅︎ | +| `Qwen3ForSequenceClassification`C | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), `polaria-tech/zerank-2-reranker-vllm`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja), [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja) | ✅︎ | ✅︎ | | `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | N/A | ✅︎ | | | `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | N/A | ✅︎ | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | N/A | \* | \* | @@ -88,6 +88,13 @@ The three supported scoring functions are as illustrated in the figure below. vllm serve Qwen/Qwen3-Reranker-0.6B --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' ``` +!!! note + `polaria-tech/zerank-2-reranker-vllm` is [zeroentropy/zerank-2-reranker](https://huggingface.co/zeroentropy/zerank-2-reranker) converted to `Qwen3ForSequenceClassification`; it returns zerank-2's calibrated score `sigmoid(logit / 5)`. Serve it with [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja), which also supports the `instruction` request field. + + ```bash + vllm serve polaria-tech/zerank-2-reranker-vllm --chat-template examples/pooling/score/template/zerank2.jinja + ``` + #### Multimodal Models !!! note diff --git a/examples/pooling/score/template/zerank2.jinja b/examples/pooling/score/template/zerank2.jinja new file mode 100644 index 000000000000..8bd16e430121 --- /dev/null +++ b/examples/pooling/score/template/zerank2.jinja @@ -0,0 +1,13 @@ +{#- zerank-2 (zeroentropy/zerank-2-reranker): query in the system turn, document in the user turn. + The assistant suffix is emitted unconditionally, since score templates are rendered + without add_generation_prompt. An optional `instruction` is wrapped as recommended by + ZeroEntropy; without it, the prompt is the official zerank-2 prompt. -#} +{%- set query = messages | selectattr("role", "eq", "query") | map(attribute="content") | list -%} +{%- set document = messages | selectattr("role", "eq", "document") | map(attribute="content") | list -%} +{%- set query = query[0] if query else '' -%} +{%- if instruction is defined and instruction is string and instruction | trim -%} +{%- set query = '' + query + '\n' + instruction + '' -%} +{%- endif -%} +{{- '<|im_start|>system\n' + query + '<|im_end|>\n' -}} +{{- '<|im_start|>user\n' + (document[0] if document else '') + '<|im_end|>\n' -}} +{{- '<|im_start|>assistant\n' -}} diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py index abe8157d4c2f..1b7fe9f7c0ad 100644 --- a/examples/pooling/score/using_template_offline.py +++ b/examples/pooling/score/using_template_offline.py @@ -53,6 +53,7 @@ def get_chat_template(model: str) -> str: "mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja", "mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja", "nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja", + "polaria-tech/zerank-2-reranker-vllm": "zerank2.jinja", } # Get the template filename for the specified model @@ -110,6 +111,7 @@ def get_hf_overrides(model: str) -> dict[str, Any]: "method": "from_2_way_softmax", }, "nvidia/llama-nemotron-rerank-1b-v2": {}, + "polaria-tech/zerank-2-reranker-vllm": {}, } hf_overrides = model_name_to_hf_overrides_map.get(model) diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py index f7dc109dd81f..9cf0703ca475 100644 --- a/examples/pooling/score/using_template_online.py +++ b/examples/pooling/score/using_template_online.py @@ -19,6 +19,7 @@ vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja + vllm serve polaria-tech/zerank-2-reranker-vllm --runner pooling --chat-template examples/pooling/score/template/zerank2.jinja vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja """ diff --git a/tests/models/language/pooling_mteb_test/test_zerank.py b/tests/models/language/pooling_mteb_test/test_zerank.py new file mode 100644 index 000000000000..c90c80bc4d40 --- /dev/null +++ b/tests/models/language/pooling_mteb_test/test_zerank.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import mteb +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +from tests.conftest import HfRunner +from tests.models.utils import RerankModelInfo + +from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models + +RERANK_MODELS = [ + RerankModelInfo( + "polaria-tech/zerank-2-reranker-vllm", + architecture="Qwen3ForSequenceClassification", + chat_template_name="zerank2.jinja", + seq_pooling_type="LAST", + attn_type="decoder", + is_prefix_caching_supported=True, + is_chunked_prefill_supported=True, + mteb_score=0.34305, + mteb_tol=1e-2, + enable_test=True, + ), +] + + +class ZerankHfRunner(MtebCrossEncoderMixin, HfRunner): + # The converted checkpoint has no lm_head, so the reference is the original + # causal LM: the logit of the "Yes" token at the last position. The converted + # checkpoint returns sigmoid(logit / 5), which ranks documents identically. + original_model = "zeroentropy/zerank-2-reranker" + + def __init__( + self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any + ) -> None: + from transformers import AutoModelForCausalLM, AutoTokenizer + + HfRunner.__init__( + self, + model_name=self.original_model, + auto_cls=AutoModelForCausalLM, + dtype=dtype, + **kwargs, + ) + + self.tokenizer = AutoTokenizer.from_pretrained( + self.original_model, padding_side="left" + ) + self.token_true_id = self.tokenizer.convert_tokens_to_ids("Yes") + + @torch.no_grad + def predict( + self, + inputs1: DataLoader[mteb.types.BatchedInput], + inputs2: DataLoader[mteb.types.BatchedInput], + *args, + **kwargs, + ) -> np.ndarray: + queries = [text for batch in inputs1 for text in batch["text"]] + corpus = [text for batch in inputs2 for text in batch["text"]] + + tokenizer = self.tokenizer + prompts = [] + for query, document in zip(queries, corpus): + conversation = [ + {"role": "query", "content": query}, + {"role": "document", "content": document}, + ] + + prompt = tokenizer.apply_chat_template( + conversation=conversation, + tools=None, + chat_template=self.chat_template, + tokenize=False, + ) + prompts.append(prompt) + + scores = [] + for prompt in prompts: + inputs = tokenizer([prompt], return_tensors="pt") + inputs = self.wrap_device(inputs) + logits = self.model(**inputs).logits[:, -1, :] + scores.append(logits[0, self.token_true_id].item()) + return torch.Tensor(scores) + + +@pytest.mark.flaky(reruns=2) +@pytest.mark.parametrize("model_info", RERANK_MODELS) +def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: + mteb_test_rerank_models(vllm_runner, model_info, hf_runner=ZerankHfRunner) From 51ef2397685b49ebd90b39bb4b4806a509df19b5 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:11:19 -0700 Subject: [PATCH 254/952] [Rust Frontend] Carry reasoning controls over gRPC (#59836) Signed-off-by: Alec Flowers Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/Cargo.lock | 2 +- rust/proto/Cargo.toml | 2 +- rust/proto/README.md | 2 +- rust/proto/inference.proto | 22 ++++++++ rust/src/server/src/grpc/convert.rs | 85 ++++++++++++++++++++++++++++- 5 files changed, 108 insertions(+), 5 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 331ebfa8dfaa..c7b3d66e4c1b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5947,7 +5947,7 @@ dependencies = [ [[package]] name = "vllm-proto" -version = "0.4.0" +version = "0.5.0" dependencies = [ "prost", "prost-types", diff --git a/rust/proto/Cargo.toml b/rust/proto/Cargo.toml index 492cb5a06d67..c2e32d43a716 100644 --- a/rust/proto/Cargo.toml +++ b/rust/proto/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vllm-proto" -version = "0.4.0" +version = "0.5.0" edition.workspace = true license.workspace = true description = "Protobuf messages and gRPC bindings for the vLLM Rust frontend" diff --git a/rust/proto/README.md b/rust/proto/README.md index d312e27acac6..3a60c05dfdf0 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -16,7 +16,7 @@ the schemas on pull requests, including on forks. ```toml [dependencies] -vllm-proto = "0.4" +vllm-proto = "0.5" ``` For example, import `vllm_proto::inference_client::InferenceClient` or diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 79c1d0a147e2..2daacbfc35ba 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -58,6 +58,25 @@ message GenerateRequest { optional bool watermarking = 16; KVHintsEnvelope kv_hints = 17; + + // Transitional, to be deprecated: only consulted when the engine runs its own + // reasoning parser and the structured-output constraint covers only the + // output after reasoning. Prefer a constraint that covers the full output + // from the first generated token. + EngineReasoningGate engine_reasoning_gate = 18; +} + +// Inputs to the engine's reasoning gate, which keeps a structured-output +// constraint inactive until the engine's reasoning parser detects the end of +// reasoning. +message EngineReasoningGate { + // Omitted: the engine infers the initial state from the prompt. + // true: apply the constraint from the first generated token. + // false: generation starts inside reasoning. + optional bool reasoning_ended = 1; + // Chat template kwargs used to render the prompt, so the engine's reasoning + // parser observes the same thinking mode. + google.protobuf.Struct chat_template_kwargs = 2; } message KVHintsEnvelope { @@ -118,6 +137,9 @@ message StoppingCriteria { bool include_stop_strings = 5; bool ignore_eos = 6; + + // Omitted or -1 disables the budget; zero immediately closes reasoning. + optional int64 thinking_token_budget = 7; } message ResponseOptions { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index a409cc7f79ee..d495fea485c7 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -9,6 +9,7 @@ use tonic::Status; use uuid::Uuid; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_text::{ DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, PromptTruncation, @@ -84,6 +85,7 @@ pub fn to_text_request( }; let session_id = req.session_id.filter(|s| !s.is_empty()); let kv_hints = req.kv_hints.map(kv_hints_from_proto); + let reasoning_gate = req.engine_reasoning_gate.unwrap_or_default(); let sampling = req.sampling.as_ref(); let decoding = req.decoding.as_ref(); @@ -137,8 +139,15 @@ pub fn to_text_request( data_parallel_rank: None, session_id, kv_hints, - reasoning_parser_kwargs: Default::default(), - reasoning_ended: None, + reasoning_parser_kwargs: ReasoningParserKwargs { + chat_template_kwargs: reasoning_gate + .chat_template_kwargs + .map(|kwargs| { + kwargs.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect() + }) + .unwrap_or_default(), + }, + reasoning_ended: reasoning_gate.reasoning_ended, lora_request: None, arrival_time: None, }) @@ -213,6 +222,7 @@ fn build_sampling_params( params.stop_token_ids = Some(s.stop_token_ids.clone()); } params.ignore_eos = s.ignore_eos; + params.thinking_token_budget = s.thinking_token_budget; } // ResponseOptions → logprobs @@ -611,6 +621,77 @@ mod tests { } } + #[test] + fn grpc_reasoning_controls_reach_engine_request_with_presence_intact() { + let chat_template_kwargs = + serde_json::json!({"enable_thinking": false, "reasoning_effort": "high"}); + for (gate, budget) in [ + (None, None), + ( + Some(pb::EngineReasoningGate { + reasoning_ended: Some(false), + chat_template_kwargs: json_to_proto_struct(&chat_template_kwargs), + }), + Some(0), + ), + ( + Some(pb::EngineReasoningGate { + reasoning_ended: Some(true), + chat_template_kwargs: None, + }), + Some(128), + ), + (Some(pb::EngineReasoningGate::default()), Some(-1)), + ] { + let expected = ( + serde_json::json!({ + "chat_template_kwargs": gate + .as_ref() + .and_then(|gate| gate.chat_template_kwargs.as_ref()) + .map_or_else(|| serde_json::json!({}), |_| chat_template_kwargs.clone()), + }), + gate.as_ref().and_then(|gate| gate.reasoning_ended), + budget.filter(|v| *v >= 0).map(|v| v as u64), + ); + let request = pb::GenerateRequest { + engine_reasoning_gate: gate, + stopping: Some(pb::StoppingCriteria { + thinking_token_budget: budget, + ..Default::default() + }), + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let request = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(request, stream, &["test-model".to_string()]) + .expect("reasoning controls accepted"); + let engine = vllm_text::lower_text_request( + text, + vec![1], + SamplingHints::default(), + SamplingLimits { + max_model_len: 256, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + &TestTokenizer::new(), + ) + .expect("reasoning controls lower to engine") + .generate_request; + assert_eq!( + ( + serde_json::to_value(engine.reasoning_parser_kwargs).unwrap(), + engine.reasoning_ended, + engine.sampling_params.thinking_token_budget, + ), + expected, + ); + } + } + } + #[test] fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() { for watermarking in [None, Some(true), Some(false)] { From 112a2829157cc4267afae8c5577df20a36136c22 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Varikuntla Date: Thu, 8 Oct 2026 21:20:43 -0700 Subject: [PATCH 255/952] [Bugfix][Frontend] Keep streamed completion text_offset in step with tokens (#60668) Signed-off-by: Vineeth Sai Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../completion/test_completion_text_offset.py | 97 +++++++++++++++++++ vllm/entrypoints/openai/completion/serving.py | 11 ++- 2 files changed, 105 insertions(+), 3 deletions(-) create mode 100644 tests/entrypoints/openai/completion/test_completion_text_offset.py diff --git a/tests/entrypoints/openai/completion/test_completion_text_offset.py b/tests/entrypoints/openai/completion/test_completion_text_offset.py new file mode 100644 index 000000000000..c8ba7df62f8a --- /dev/null +++ b/tests/entrypoints/openai/completion/test_completion_text_offset.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput + +TOKENS = [(1, "abc"), (2, " de"), (3, "!")] +# Text deltas as the detokenizer streams them with stop=["\n\n"]: one char is +# held back until the request finishes, so the text lags the tokens. +TEXT_DELTAS = ["ab", "c d", "e!"] + + +def _logprobs(i: int) -> dict[int, Logprob]: + token_id, token = TOKENS[i] + return {token_id: Logprob(logprob=-1.0, rank=1, decoded_token=token)} + + +def _output(text: str, idxs: list[int], finished: bool) -> RequestOutput: + return RequestOutput( + request_id="r", + prompt="Hi", + prompt_token_ids=[9], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=[TOKENS[i][0] for i in idxs], + cumulative_logprob=None, + logprobs=[_logprobs(i) for i in idxs], + finish_reason="stop" if finished else None, + ) + ], + finished=finished, + ) + + +def _serving() -> OpenAIServingCompletion: + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.enable_prompt_tokens_details = False + serving.enable_force_include_usage = False + serving.enable_per_request_metrics = False + serving.system_fingerprint = None + serving.return_tokens_as_token_ids = False + return serving + + +def _request(stream: bool) -> CompletionRequest: + return CompletionRequest( + model="m", prompt="Hi", max_tokens=8, logprobs=0, stop=["\n\n"], stream=stream + ) + + +@pytest.mark.asyncio +async def test_stream_text_offset_matches_non_stream_with_stop_holdback(): + async def results(): + for i, text in enumerate(TEXT_DELTAS): + yield 0, _output(text, [i], finished=i == len(TEXT_DELTAS) - 1) + + serving = _serving() + stream_offsets = [] + async for line in serving.completion_stream_generator( + _request(stream=True), + [{"prompt": "Hi"}], + results(), + "r", + 0, + "m", + num_prompts=1, + tokenizer=None, + request_metadata=RequestResponseMetadata(request_id="r"), + ): + data = line.removeprefix("data: ").strip() + if data == "[DONE]": + continue + for choice in json.loads(data)["choices"]: + stream_offsets += choice["logprobs"]["text_offset"] + + response = serving.request_output_to_completion_response( + [_output("".join(TEXT_DELTAS), [0, 1, 2], finished=True)], + _request(stream=False), + "r", + 0, + "m", + None, + RequestResponseMetadata(request_id="r"), + ) + assert response.choices[0].logprobs is not None + assert response.choices[0].logprobs.text_offset == [0, 3, 6] + assert stream_offsets == [0, 3, 6] diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 81dd00f0b3bd..7f7bdac0757b 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -289,7 +289,7 @@ async def completion_stream_generator( request_metadata: RequestResponseMetadata, ) -> AsyncGenerator[str, None]: num_choices = 1 if request.n is None else request.n - previous_text_lens = [0] * num_choices * num_prompts + next_text_offsets = [0] * num_choices * num_prompts previous_num_tokens = [0] * num_choices * num_prompts has_echoed = [False] * num_choices * num_prompts num_prompt_tokens = [0] * num_prompts @@ -386,13 +386,18 @@ async def completion_stream_generator( num_output_top_logprobs=request.logprobs, tokenizer=tokenizer, logprob_token_ids=request.logprob_token_ids, - initial_text_offset=previous_text_lens[i], + initial_text_offset=next_text_offsets[i], return_as_token_id=request.return_tokens_as_token_ids, ) + # Advance by the tokens just reported, not by delta_text: + # text held back for stop strings lags the tokens. + if logprobs.tokens: + next_text_offsets[i] = logprobs.text_offset[-1] + len( + logprobs.tokens[-1] + ) else: logprobs = None - previous_text_lens[i] += len(delta_text) previous_num_tokens[i] += len(output.token_ids) stop_reason = output.stop_reason From 9824159d94082deda7e112bfdbadbd441eb8084a Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Fri, 9 Oct 2026 08:20:49 +0400 Subject: [PATCH 256/952] [Bugfix] Fix InternLM2 tool parser dropping characters from streamed arguments (#60709) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_internlm2_tool_parser.py | 28 +++++++++++++++++++ vllm/tool_parsers/internlm2_tool_parser.py | 17 +++++++---- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 881933890dfb..28465654e4fa 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -10,6 +10,7 @@ ToolParserTestConfig, ToolParserTests, ) +from tests.tool_parsers.utils import run_tool_extraction_streaming from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser @@ -162,3 +163,30 @@ def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) - streamed += arguments assert json.loads(streamed) == {"city": "Dallas", "state": "TX"} + + +@pytest.mark.parametrize( + "arguments_deltas", + [ + ['{"location": "Beijing', '", "days": 3}}<|action_end|>'], + [*'{"location": "Beijing", "days": 3}}', "<|action_end|>"], + ], + ids=["split_inside_string", "char_by_char"], +) +def test_streaming_arguments_match_final_arguments( + default_tokenizer: TokenizerLike, arguments_deltas: list[str] +) -> None: + """Streamed argument deltas must add up to the final arguments, wherever + the chunk boundaries fall.""" + parser = Internlm2ToolParser(default_tokenizer) + deltas = [ + '<|action_start|><|plugin|>{"name": "get_weather", "parameters": ', + *arguments_deltas, + ] + + reconstructor = run_tool_extraction_streaming(parser, deltas) + + assert len(reconstructor.tool_calls) == 1 + streamed = reconstructor.tool_calls[0].function.arguments + streamed += parser.get_remaining_unstreamed_args() + assert streamed == '{"location": "Beijing", "days": 3}' diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index 5a99be6afd0e..709fd6c8f2a8 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ Tool, ToolParser, ) -from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json +from vllm.tool_parsers.utils import find_common_prefix, is_complete_json logger = init_logger(__name__) @@ -172,10 +172,17 @@ def extract_tool_calls_streaming( elif cur_arguments and prev_arguments: cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) - - argument_diff = extract_intermediate_diff( - cur_args_json, prev_args_json - ) + sent = len(self.streamed_args_for_tool[self.current_tool_id]) + + # Partial JSON is serialized with synthetic closers, so + # only the part shared with the previous parse is final. + if is_complete_json(parsable_arr): + argument_diff = cur_args_json[sent:] + elif cur_args_json != prev_args_json: + prefix = find_common_prefix(prev_args_json, cur_args_json) + argument_diff = prefix[sent:] + else: + argument_diff = "" delta = DeltaMessage( tool_calls=[ From e31b823bf59bdd96d09fdcf310c5fb02e934d3bf Mon Sep 17 00:00:00 2001 From: Kai Hoang <56827800+hpdkhoa@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:20:54 +0700 Subject: [PATCH 257/952] [Bugfix][Tool Parser] Use LongCat tags in the LongCat structural tag (#60747) Signed-off-by: hpdkhoa Co-authored-by: Claude Opus 5.5 Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structural_tag_registry.py | 39 ++++++++++++ vllm/tool_parsers/longcat_tool_parser.py | 2 + vllm/tool_parsers/structural_tag_registry.py | 60 +++++++++++++------ 3 files changed, 84 insertions(+), 17 deletions(-) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 70f1921e459a..817c48a4401d 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -48,6 +48,7 @@ from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser +from vllm.tool_parsers.longcat_tool_parser import LongcatFlashToolParser from vllm.tool_parsers.mimo_tool_parser import MiMoToolParser from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser from vllm.tool_parsers.plamo3_engine_tool_parser import Plamo3EngineToolParser @@ -850,6 +851,44 @@ class StepParser(DelegatingParser): ] +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "function", "function": {"name": "get_weather"}}], +) +def test_longcat_forced_tool_choice_round_trips_native_tags( + sample_tools: list[ChatCompletionToolsParam], tool_choice +): + """Forced LongCat calls must use , not Hermes tags.""" + + class LongcatParser(DelegatingParser): + tool_parser_cls = LongcatFlashToolParser + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=[tool.model_dump(exclude_none=True) for tool in sample_tools], + tool_choice=tool_choice, + ) + parser = LongcatParser(MagicMock(), tools=sample_tools) + request = parser.adjust_request(request) + + grammar = Grammar.from_structural_tag(request.structured_outputs.structural_tag) + # Rendered exactly as the LongCat-Flash chat template renders tool calls. + output = ( + '\n{"name": "get_weather", ' + '"arguments": {"city": "Dallas"}}\n' + ) + assert _is_grammar_accept_string(grammar, output) + assert not _is_grammar_accept_string(grammar, output.replace("longcat_", "")) + + _, content, tool_calls = parser.parse(output, request, enable_auto_tools=True) + + assert not content + assert [(call.name, json.loads(call.arguments)) for call in tool_calls] == [ + ("get_weather", {"city": "Dallas"}) + ] + + def test_tool_parsers_without_structural_tag_support_required_and_named(): class NonStructuralTagToolParser(ToolParser): pass diff --git a/vllm/tool_parsers/longcat_tool_parser.py b/vllm/tool_parsers/longcat_tool_parser.py index ccfa94167256..919f80143737 100644 --- a/vllm/tool_parsers/longcat_tool_parser.py +++ b/vllm/tool_parsers/longcat_tool_parser.py @@ -9,6 +9,8 @@ class LongcatFlashToolParser(Hermes2ProToolParser): + structural_tag_model = "longcat" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 1d3a7c4a4c7f..f3148a827eb0 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -79,7 +79,7 @@ } ) VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( - {"glm_4_7", "hermes", "hy_v4", "kimi_k3", "plamo3"} + {"glm_4_7", "hermes", "hy_v4", "kimi_k3", "longcat", "plamo3"} ) SUPPORTED_STRUCTURAL_TAG_MODELS = ( XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS @@ -346,15 +346,19 @@ def get_function_parameters(function) -> dict[str, Any] | bool: return function.parameters if function.parameters is not None else True -def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: +def _hermes_tool_tags( + tools: list[FunctionToolParam], + start_tag: str, + end_tag: str, +) -> list[TagFormat]: arguments_field_prefix = '", "arguments": ' formats = [ # # {"name": "t1", "arguments": {"q": "v"}} # - ('\n{"name": "', "}\n"), + (start_tag + '\n{"name": "', "}\n" + end_tag), # {"name": "t1", "arguments": {"q": "v"}} - ('{"name": "', "}"), + (start_tag + '{"name": "', "}" + end_tag), ] return [ @@ -370,35 +374,29 @@ def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: ] -@register_vllm_structural_tag("hermes") -def get_hermes_structural_tag( +def _hermes_style_structural_tag( tools: list[FunctionToolParam], - builtin_tools: list[BuiltinToolParam], tool_choice: SimplifiedToolChoice, - reasoning: bool, - token_suffix: str = "", + start_tag: str, + end_tag: str, ) -> StructuralTag: - del builtin_tools, reasoning, token_suffix - - tool_call_trigger = "" - + tags = _hermes_tool_tags(tools, start_tag, end_tag) if tool_choice == "auto": - tags = _hermes_tool_tags(tools) suffix_tag = ( - TriggeredTagsFormat(triggers=[tool_call_trigger], tags=tags) + TriggeredTagsFormat(triggers=[start_tag], tags=tags) if tags else AnyTextFormat() ) elif tool_choice == "forced": suffix_tag = TagsWithSeparatorFormat( - tags=_hermes_tool_tags(tools), + tags=tags, separator="", at_least_one=True, stop_after_first=True, ) else: suffix_tag = TagsWithSeparatorFormat( - tags=_hermes_tool_tags(tools), + tags=tags, separator="", at_least_one=True, ) @@ -406,6 +404,34 @@ def get_hermes_structural_tag( return StructuralTag(format=suffix_tag) +@register_vllm_structural_tag("hermes") +def get_hermes_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, + token_suffix: str = "", +) -> StructuralTag: + del builtin_tools, reasoning, token_suffix + return _hermes_style_structural_tag( + tools, tool_choice, "", "" + ) + + +@register_vllm_structural_tag("longcat") +def get_longcat_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, + token_suffix: str = "", +) -> StructuralTag: + del builtin_tools, reasoning, token_suffix + return _hermes_style_structural_tag( + tools, tool_choice, "", "" + ) + + def _minimax_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: return [ TagFormat( From 81d9c45e4b87b8abc82fd69f81be306f5ceae49a Mon Sep 17 00:00:00 2001 From: Rajeev Warrier Date: Fri, 9 Oct 2026 09:53:07 +0530 Subject: [PATCH 258/952] =?UTF-8?q?[Docs]=20[Multimodal]=20Note=20that=20p?= =?UTF-8?q?er-request=20media=5Fio=5Fkwargs=20needs=20--trust-request-mm?= =?UTF-8?q?=E2=80=A6=20(#60598)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Rajeev Warrier --- docs/features/multimodal_inputs.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 40ceb8158b52..cb091c1c79d6 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -1020,6 +1020,13 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction. +!!! warning + Per-request `media_io_kwargs` (and `mm_processor_kwargs`) are rejected by + default because they let a client change media loading and preprocessing + resource usage. Start the server with `--trust-request-mm-kwargs` only when + your API clients are trusted; see + [Per-request multimodal arguments](../usage/security.md#per-request-multimodal-arguments). + **Supported Parameters:** | Parameter | Type | Description | @@ -1042,7 +1049,7 @@ When you extract video frames on the client side and send them as `video/jpeg` ( frames_b64 = ",".join([encode_image(f) for f in frames]) video_url = f"data:video/jpeg;base64,{frames_b64}" - # Pass video metadata via media_io_kwargs + # Pass video metadata via media_io_kwargs (server started with --trust-request-mm-kwargs) response = client.chat.completions.create( model="your-multimodal-model", messages=[{ From fa312090cb1d972aa96f6b0ea174b7a32b4b0548 Mon Sep 17 00:00:00 2001 From: Aaron Hao Date: Thu, 8 Oct 2026 21:40:13 -0700 Subject: [PATCH 259/952] [Frontend] Accept each connection on the least loaded API worker (#60723) Signed-off-by: hao-aaron Co-authored-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../launchers/test_balanced_accepts.py | 218 ++++++++++++++++++ .../entrypoints/launchers/api_server/entry.py | 11 +- vllm/entrypoints/launchers/launcher.py | 130 ++++++++++- vllm/v1/utils.py | 10 + 4 files changed, 365 insertions(+), 4 deletions(-) create mode 100644 tests/entrypoints/launchers/test_balanced_accepts.py diff --git a/tests/entrypoints/launchers/test_balanced_accepts.py b/tests/entrypoints/launchers/test_balanced_accepts.py new file mode 100644 index 000000000000..3ac1025ca853 --- /dev/null +++ b/tests/entrypoints/launchers/test_balanced_accepts.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NoSignalServer(peer_loads=...): API workers sharing one listening socket +accept each connection on the worker with the fewest open connections.""" + +import asyncio +import datetime +import multiprocessing +import os +import socket +import ssl +from collections import Counter + +import pytest +import uvicorn +import uvloop + +from vllm.entrypoints.launchers.launcher import NoSignalServer + +TIMEOUT_S = 30 + + +def _bound_socket() -> socket.socket: + # Like setup_server: bound, not yet listening. + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + return sock + + +def _app(hold_s: float = 0.0): + pid = str(os.getpid()).encode() + + async def app(scope, receive, send): + if scope["type"] == "lifespan": + while (await receive())["type"] != "lifespan.shutdown": + await send({"type": "lifespan.startup.complete"}) + await send({"type": "lifespan.shutdown.complete"}) + return + await asyncio.sleep(hold_s) + headers = [(b"content-length", str(len(pid)).encode())] + await send({"type": "http.response.start", "status": 200, "headers": headers}) + await send({"type": "http.response.body", "body": pid}) + + return app + + +async def _get(port: int, requests: int = 1, ssl_ctx=None) -> list[bytes]: + """``requests`` requests on one keep-alive connection; returns the bodies.""" + reader, writer = await asyncio.open_connection("127.0.0.1", port, ssl=ssl_ctx) + bodies = [] + for _ in range(requests): + writer.write(b"GET / HTTP/1.1\r\nHost: t\r\n\r\n") + head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), TIMEOUT_S) + assert head.startswith(b"HTTP/1.1 200") + length = int(head.split(b"content-length: ")[1].split(b"\r\n")[0]) + bodies.append(await asyncio.wait_for(reader.readexactly(length), TIMEOUT_S)) + writer.close() + return bodies + + +async def _serve(sock, requests, loads=None, **config_kwargs): + config = uvicorn.Config(_app(), lifespan="on", log_level="warning", **config_kwargs) + loads = loads if loads is not None else [0] + server = NoSignalServer(config, peer_loads=(loads, 0)) + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + assert not task.done(), task.result() + await asyncio.sleep(0.01) + try: + return await requests(sock.getsockname()[1], server) + finally: + server.should_exit = True + await asyncio.wait_for(task, TIMEOUT_S) + # The listener is closed, nothing is left accepting, and peers no + # longer see this worker as a target. + assert sock.fileno() == -1 + assert not server._accept_tasks + assert loads[0] == 2**31 - 1 + + +def _run(coro, loop_impl): + return uvloop.run(coro) if loop_impl == "uvloop" else asyncio.run(coro) + + +@pytest.mark.parametrize("loop_impl", ["asyncio", "uvloop"]) +def test_requests_and_keep_alive_served(loop_impl): + async def requests(port, server): + return await _get(port), await _get(port, requests=3) + + one, three = _run(_serve(_bound_socket(), requests), loop_impl) + assert one == [str(os.getpid()).encode()] + assert len(three) == 3 + + +def _self_signed(tmp_path): + x509 = pytest.importorskip("cryptography.x509") + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.x509.oid import NameOID + + key = ec.generate_private_key(ec.SECP256R1()) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "127.0.0.1")]) + now = datetime.datetime.now(datetime.UTC) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=1)) + .not_valid_after(now + datetime.timedelta(hours=1)) + .sign(key, hashes.SHA256()) + ) + keyfile, certfile = tmp_path / "key.pem", tmp_path / "cert.pem" + keyfile.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + certfile.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + return str(keyfile), str(certfile) + + +def test_tls_handshake_failure_keeps_serving(tmp_path): + keyfile, certfile = _self_signed(tmp_path) + client_ctx = ssl.create_default_context() + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + + async def requests(port, server): + # A plain-text client fails the handshake; the next TLS client is served. + reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.write(b"GET / HTTP/1.1\r\nHost: t\r\n\r\n") + await asyncio.wait_for(reader.read(), TIMEOUT_S) + writer.close() + return await _get(port, ssl_ctx=client_ctx) + + bodies = uvloop.run( + _serve(_bound_socket(), requests, ssl_keyfile=keyfile, ssl_certfile=certfile) + ) + assert bodies == [str(os.getpid()).encode()] + + +def test_busier_worker_still_serves(): + # This worker is index 0; peer 1 claims no connections but never accepts. + loads = [0, 0] + + async def requests(port, server): + reader, writer = await asyncio.open_connection("127.0.0.1", port) + bodies = await _get(port) + assert loads[0] >= 1 # the held connection was published + writer.close() + return bodies + + assert uvloop.run(_serve(_bound_socket(), requests, loads)) == [ + str(os.getpid()).encode() + ] + + +def _worker(sock, loads, index, ready, stop): + async def main(): + config = uvicorn.Config(_app(hold_s=1), lifespan="on", log_level="warning") + server = NoSignalServer(config, peer_loads=(loads, index)) + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + await asyncio.sleep(0.01) + ready.release() + await asyncio.get_running_loop().run_in_executor(None, stop.wait) + server.should_exit = True + await task + + # API workers run on uvloop, whose own server accepts a whole burst at once. + uvloop.run(main()) + + +@pytest.mark.skipif((os.cpu_count() or 1) < 8, reason="needs idle cores per worker") +@pytest.mark.parametrize("arrival_s", [0, 1]) +def test_connections_balanced_over_workers(arrival_s): + """Both a burst and a steady stream of connections end up balanced; with + the event loop's own server one worker takes most of either.""" + workers, connections = 4, 400 + sock = _bound_socket() + ctx = multiprocessing.get_context("spawn") + ready, stop = ctx.Semaphore(0), ctx.Event() + loads = ctx.RawArray("i", workers) + procs = [ + ctx.Process(target=_worker, args=(sock, loads, i, ready, stop)) + for i in range(workers) + ] + for p in procs: + p.start() + try: + for _ in range(workers): + assert ready.acquire(timeout=120) + + async def get_after(delay): + await asyncio.sleep(delay) + return await _get(sock.getsockname()[1]) + + async def arrive(): + delays = [arrival_s * i / connections for i in range(connections)] + results = await asyncio.gather(*(get_after(d) for d in delays)) + return Counter(body for (body,) in results) + + counts = uvloop.run(arrive()) + assert sum(counts.values()) == connections + # Fair share is 100. + assert max(counts.values()) <= 5 * connections // (4 * workers), counts + finally: + stop.set() + for p in procs: + p.join(30) + if p.is_alive(): + p.kill() + sock.close() diff --git a/vllm/entrypoints/launchers/api_server/entry.py b/vllm/entrypoints/launchers/api_server/entry.py index 28bb696f2c4e..e5ba536e1128 100644 --- a/vllm/entrypoints/launchers/api_server/entry.py +++ b/vllm/entrypoints/launchers/api_server/entry.py @@ -8,7 +8,7 @@ import socket import tempfile from argparse import Namespace -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, MutableSequence from contextlib import asynccontextmanager from typing import Any @@ -118,6 +118,7 @@ async def build_and_serve( listen_address: str, sock: socket.socket, args: Namespace, + peer_loads: tuple[MutableSequence[int], int] | None = None, **uvicorn_kwargs, ) -> asyncio.Task: """Build FastAPI app, initialize state, and start serving. @@ -142,6 +143,7 @@ async def build_and_serve( app, sock=sock, enable_ssl_refresh=args.enable_ssl_refresh, + peer_loads=peer_loads, host=args.host, port=args.port, log_level=args.uvicorn_log_level, @@ -185,12 +187,17 @@ async def run_server_worker( if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + peer_loads = None + if client_config and "api_server_loads" in client_config: + loads = client_config.pop("api_server_loads") + peer_loads = (loads, client_config["client_index"]) + async with build_async_engine_client( args, client_config=client_config, ) as engine_client: shutdown_task = await build_and_serve( - engine_client, listen_address, sock, args, **uvicorn_kwargs + engine_client, listen_address, sock, args, peer_loads, **uvicorn_kwargs ) # NB: Await server shutdown only after the backend context is exited try: diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 9fabfa3277af..44161439b690 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -5,7 +5,7 @@ import contextlib import signal import socket -from collections.abc import Generator +from collections.abc import Generator, MutableSequence from functools import partial from typing import Any @@ -34,23 +34,149 @@ logger = init_logger(__name__) +# A worker with more connections than another leaves a pending connection to +# the less loaded one, rechecking every 1 ms, for up to 5 ms (it may be busy). +_ACCEPT_DEFER_CHECKS = 5 +# The load of a worker that is not accepting (also the initial value of the +# shared array), so it is never the least loaded. +_NOT_ACCEPTING = 2**31 - 1 + class NoSignalServer(uvicorn.Server): """Uvicorn server that never installs its own SIGINT/SIGTERM handlers. Callers register their own handlers on the event loop for graceful shutdown; uvicorn's would race with and override them (see #49668). + + With ``peer_loads`` (several API workers on one shared socket: a shared + array with one slot per worker, and this worker's index), each worker + publishes its number of open connections, and a pending connection is + accepted by the least loaded worker. The event loop's own server would + accept every queued connection in whichever worker wakes first. """ + def __init__( + self, + config: uvicorn.Config, + peer_loads: tuple[MutableSequence[int], int] | None = None, + ): + super().__init__(config) + self.peer_loads = peer_loads + self._accept_tasks: set[asyncio.Task] = set() + self._connecting = 0 + @contextlib.contextmanager def capture_signals(self) -> Generator[None, None, None]: yield + def _publish_load(self) -> None: + assert self.peer_loads is not None + loads, index = self.peer_loads + loads[index] = len(self.server_state.connections) + self._connecting + + def _least_loaded(self) -> bool: + assert self.peer_loads is not None + loads, index = self.peer_loads + return loads[index] <= min(loads) + + async def startup(self, sockets: list[socket.socket] | None = None) -> None: + if not (self.peer_loads and sockets): + return await super().startup(sockets) + # With no sockets uvicorn starts no asyncio server; _accept serves + # them with uvicorn's protocol instead. + await super().startup(sockets=[]) + if not self.started: # lifespan startup failed (uvicorn < 0.50) + return + self._publish_load() + for sock in sockets: + sock.listen(self.config.backlog) + sock.setblocking(False) + self._track(self._accept(sock)) + + def _track(self, coro) -> None: + task = asyncio.create_task(coro) + self._accept_tasks.add(task) + task.add_done_callback(self._accept_tasks.discard) + + async def _accept(self, sock: socket.socket) -> None: + loop = asyncio.get_running_loop() + while True: + await _readable(loop, sock) + # Connections closed since the last accept; peers refresh likewise. + self._publish_load() + for _ in range(_ACCEPT_DEFER_CHECKS): + if self._least_loaded(): + break + await asyncio.sleep(0.001) + try: + conn, _ = sock.accept() + except (BlockingIOError, InterruptedError, ConnectionAbortedError): + continue # another worker took it + except OSError: + # e.g. EMFILE; asyncio's own server also retries after 1 s. + logger.exception("Error accepting a connection") + await asyncio.sleep(1) + continue + self._connecting += 1 + self._publish_load() + # Its own task, so a TLS handshake does not hold up the next accept. + self._track(self._connect(loop, conn)) + + async def _connect(self, loop: asyncio.AbstractEventLoop, conn: socket.socket): + config = self.config + try: + await loop.connect_accepted_socket( + lambda: config.http_protocol_class( # type: ignore[call-arg] + config=config, + server_state=self.server_state, + app_state=self.lifespan.state, + _loop=loop, + ), + conn, + ssl=config.ssl, + ) + except asyncio.CancelledError: + conn.close() + raise + except Exception: + # e.g. a failed TLS handshake + logger.debug("Accepted connection failed to start", exc_info=True) + conn.close() + finally: + self._connecting -= 1 + self._publish_load() + + async def shutdown(self, sockets: list[socket.socket] | None = None) -> None: + # Before super() closes the listening sockets. + tasks = list(self._accept_tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + if self.peer_loads is not None: + loads, index = self.peer_loads + loads[index] = _NOT_ACCEPTING + await super().shutdown(sockets) + + +async def _readable(loop: asyncio.AbstractEventLoop, sock: socket.socket) -> None: + ready: asyncio.Future[None] = loop.create_future() + + def wake() -> None: + if not ready.done(): + ready.set_result(None) + + loop.add_reader(sock, wake) + try: + await ready + finally: + loop.remove_reader(sock) + async def serve_http( app: FastAPI, sock: socket.socket | None, enable_ssl_refresh: bool = False, + peer_loads: tuple[MutableSequence[int], int] | None = None, **uvicorn_kwargs: Any, ): """Start a FastAPI app using Uvicorn, with support for custom Uvicorn config @@ -96,7 +222,7 @@ async def serve_http( config.h11_max_incomplete_event_size = h11_max_incomplete_event_size config.h11_max_header_count = h11_max_header_count config.load() - server = NoSignalServer(config) + server = NoSignalServer(config, peer_loads=peer_loads) app.state.server = server loop = asyncio.get_running_loop() diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index 868864f0ca41..b8d98ff1cf74 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -213,6 +213,14 @@ def __init__( SharedAdmissionStats.num_counters(num_servers), ) + # Open connections per API server, to accept on the least loaded one; + # a server that has not started yet is never the least loaded. + api_server_loads = ( + spawn_context.RawArray("i", [2**31 - 1] * num_servers) + if num_servers > 1 + else None + ) + listeners = [*input_listeners, *output_listeners] self._listeners = listeners try: @@ -228,6 +236,8 @@ def __init__( "client_count": num_servers, "client_index": i, } + if api_server_loads is not None: + client_config["api_server_loads"] = api_server_loads if admission_counters is not None: client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: From 88a419649e36f3412d2650c7162135a3b12bc5a1 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Thu, 8 Oct 2026 23:45:13 -0500 Subject: [PATCH 260/952] [ROCm][Kimi-K3] Use vLLM's FP8 max in the fused AttnRes quantization (#60717) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/models/kimi_k3/test_amd_attn_res.py | 12 ++++++++---- vllm/models/kimi_k3/amd/ops/attn_res.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py index f4efdbd5b21d..871d53faae12 100644 --- a/tests/models/kimi_k3/test_amd_attn_res.py +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -247,12 +247,16 @@ def test_amd_attn_res_fp8_preserves_prefix_and_quantized_output( ) ref_output, ref_scale = quant(reference) torch.testing.assert_close(scale, ref_scale, atol=1e-7, rtol=1e-6) + # Floating-point fusion can move values at FP8 rounding midpoints. # Bound each change to the adjacent code and bound how often it occurs. - code_delta = ( - output.view(torch.uint8).to(torch.int16) - - ref_output.view(torch.uint8).to(torch.int16) - ).abs() + # Codes are sign-magnitude, so order them by value to keep a step + # across zero adjacent. + def ordinal(x: torch.Tensor) -> torch.Tensor: + code = x.view(torch.uint8).to(torch.int16) + return torch.where(code >= 0x80, 0x80 - code, code) + + code_delta = (ordinal(output) - ordinal(ref_output)).abs() assert code_delta.max().item() <= 1 assert (code_delta != 0).float().mean().item() < 1e-4 torch.testing.assert_close(prefix, ref_prefix, atol=0, rtol=0) diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py index 5c4fe7a0b0b2..95f748e72d1f 100644 --- a/vllm/models/kimi_k3/amd/ops/attn_res.py +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -10,6 +10,7 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max from vllm.triton_utils import tl, triton @@ -207,7 +208,7 @@ def attn_res( HAS_DELTA=delta is not None, WRITE_BLOCK=block_write_idx >= 0, APPLY_OUTPUT_NORM=output_norm_weight is not None, - QUANT_MAX=0.0 if quant_dtype is None else torch.finfo(quant_dtype).max, + QUANT_MAX=0.0 if quant_dtype is None else get_fp8_min_max()[1], BLOCK_L=block_l, BLOCK_D=triton.next_power_of_2(hidden_size), num_warps=num_warps, From 8cbd5d03006c33185f402249ff2b448efd594986 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 8 Oct 2026 22:48:16 -0600 Subject: [PATCH 261/952] [Frontend] structured decisions: noul and score question types (#60696) Signed-off-by: Matt Mastracci Co-authored-by: DeepSeek Co-authored-by: Claude Opus 5.5 Co-authored-by: Misha Goin --- .../online_serving/structured_decisions.md | 44 ++++++-- .../test_question_types.py | 74 +++++++++++++ .../structured_decisions/test_strategies.py | 3 + .../test_structured_decisions_e2e.py | 44 ++++++++ .../structured_decisions/question_types.py | 100 ++++++++++++++++-- .../structured_decisions/strategies.py | 32 ++++-- 6 files changed, 270 insertions(+), 27 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 7de935a62f2b..2194b54e47b4 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -11,8 +11,8 @@ log probabilities. ## How it works 1. Each question is one read. The user message is the state, followed by the - question and its options, labeled `A`, `B`, `C`, ... in order, and - "Answer with the letter of one option only." + question and its options with the type's labels, and an instruction to + answer with one label. 2. The read is the reply's first token. Each label must be one token there, or the request returns 400. 3. The read returns the logprobs of that question's label tokens. Softmax over @@ -31,8 +31,11 @@ question prefills only its own text, then reads one token. | type | criteria | answer | | --- | --- | --- | | `choice` | map of option name to a description or `null` | `choice`, `probabilities` by option name, `confidence` | +| `noul` | `null`, or an object with optional `true` and `false` descriptions | `noul` (probability of yes), `probabilities`, `confidence` | +| `score` | ordered list of level names | `score` (expected 0-indexed level), `legend`, `probabilities` by level index, `confidence` | -A choice has at least one option. A question id is a non-empty string. +A choice has at least one option, a score at least two levels. A question id +is a non-empty string. ## Example @@ -51,6 +54,16 @@ curl -s http://localhost:8000/v1/systemone \ "shipping": "deliveries", "security": "account access" } + }, + "urgent": { + "type": "noul", + "instructions": "Does the customer need a reply within the hour?", + "criteria": {"true": "the account is locked or money is moving"} + }, + "tone": { + "type": "score", + "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "furious"] } }, "chat_template_kwargs": {"enable_thinking": false} @@ -71,11 +84,26 @@ Response, with probabilities that depend on the model: "choice": "billing", "probabilities": {"billing": 0.97, "shipping": 0.01, "security": 0.02}, "confidence": 0.90 + }, + "urgent": { + "type": "noul", + "noul": 0.02, + "probabilities": {"yes": 0.02, "no": 0.98}, + "confidence": 0.95 + }, + "tone": { + "type": "score", + "score": 0.6, + "legend": {"0": "calm", "1": "annoyed", "2": "furious"}, + "probabilities": {"0": 0.7, "1": 0.2, "2": 0.1}, + "confidence": 0.65 } }, - "usage": {"input_tokens": 142, "output_tokens": 1}, + "usage": {"input_tokens": 142, "output_tokens": 3}, "diagnostics": { - "team": {"label_mass": 0.93, "argmax_is_label": true} + "team": {"label_mass": 0.93, "argmax_is_label": true}, + "urgent": {"label_mass": 0.98, "argmax_is_label": true}, + "tone": {"label_mass": 0.91, "argmax_is_label": true} } } ``` @@ -87,8 +115,9 @@ whole vocabulary: its share of the labels times `label_mass`. ## Labels -Options are labeled `A` to `Z` in the order the request lists them, so a -`choice` has at most 26 options. +Labels follow the type: a `choice` labels its options `A` to `Z` in the order +the request lists them, a `noul` labels its options `yes` and `no`, and a +`score` labels its levels `0` to `9`, so a level's label is its score. ## Limits @@ -96,6 +125,7 @@ Options are labeled `A` to `Z` in the order the request lists them, so a | --- | --- | | questions per request | 64 | | options per `choice` | 26 | +| levels per `score` | 10 | ## Request fields diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py index e5906accf025..3b9ad7ba60a9 100644 --- a/tests/entrypoints/generate/structured_decisions/test_question_types.py +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -25,6 +25,20 @@ def choice(qid="bucket", criteria=None, instructions="Which team?", max_options= ) +def noul(qid="urgent", criteria=None, instructions="Reply within the hour?"): + return build_question(qid, "noul", instructions, criteria, 128) + + +def score(qid="tone", criteria=None, instructions="How angry?", max_options=128): + return build_question( + qid, + "score", + instructions, + criteria or ["calm", "annoyed", "furious"], + max_options, + ) + + def test_choice_labels_follow_option_order(): q = choice() assert q.labels == ("A", "B", "C") @@ -62,6 +76,10 @@ def test_choice_answer_shape(): ("", "choice", {"x": None, "y": None}, "non-empty"), ("q", "nope", {"x": None, "y": None}, "unknown question type"), ("q", "choice", ["x", "y"], "must map option names"), + ("q", "noul", ["yes", "no"], "must be an object"), + ("q", "noul", {"maybe": "same day"}, "must be an object"), + ("q", "score", {"0": "calm", "1": "furious"}, "ordered list of levels"), + ("q", "score", ["calm"], "ordered list of levels"), ], ) def test_build_question_rejects(qid, type_name, criteria, match): @@ -82,6 +100,62 @@ def test_option_limit(): choice(criteria={str(i): None for i in range(27)}) +def test_noul_labels_and_prompt(): + q = noul() + assert q.labels == ("yes", "no") + assert [a.name for a in q.options] == ["yes", "no"] + assert q.type.prompt(q) == ( + "Question: Reply within the hour?\nAnswer with yes or no only." + ) + described = noul(criteria={"true": "same day", "false": "later"}) + assert described.type.prompt(described) == ( + "Question: Reply within the hour?\n" + "yes: same day\n" + "no: later\n" + "Answer with yes or no only." + ) + + +def test_noul_answer_shape(): + q = noul() + answer = q.type.answer(q, [0.8, 0.2], 0.5) + assert answer == { + "type": "noul", + "noul": 0.8, + "probabilities": {"yes": 0.8, "no": 0.2}, + "confidence": 0.4, + } + + +def test_score_labels_are_levels(): + assert score().labels == ("0", "1", "2") + with pytest.raises(StructuredDecisionError, match="at most 10"): + score(criteria=[str(i) for i in range(11)]) + + +def test_score_prompt(): + q = score() + assert q.type.prompt(q) == ( + "Question: How angry?\n" + "0: calm\n" + "1: annoyed\n" + "2: furious\n" + "Answer with the number of one level only." + ) + + +def test_score_answer_shape(): + q = score() + answer = q.type.answer(q, [0.5, 0.25, 0.25], 0.4) + assert answer == { + "type": "score", + "score": 0.75, + "legend": {"0": "calm", "1": "annoyed", "2": "furious"}, + "probabilities": {"0": 0.5, "1": 0.25, "2": 0.25}, + "confidence": 0.2, + } + + def test_registered_type_plugs_in(): class BinaryQuestion(QuestionType): name = "test_binary" diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 51bff19c8ca2..3d9538e4dc8b 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -87,6 +87,9 @@ def test_labels_start_the_reply(qwen): tail, ids = reply_label_ids(tokenizer, prompt_ids) assert tokenizer.decode(tail) == "\n\n" assert [tokenizer.decode([i]) for i in ids] == list(LABELS) + # A noul's and a score's labels are one token here too. + _, ids = reply_label_ids(tokenizer, prompt_ids, ("yes", "no", *"0123456789")) + assert len(set(ids)) == 12 # After a colon, Qwen writes ":A" as one token, so "A" is not one token. with pytest.raises(ValueError, match="not one distinct token"): reply_label_ids(tokenizer, tokenizer.encode("team:")) diff --git a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py index 3e0f9bc0fb66..a206c56ccc06 100644 --- a/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py +++ b/tests/entrypoints/generate/structured_decisions/test_structured_decisions_e2e.py @@ -188,6 +188,48 @@ def test_thinking_is_refused(server): assert "thinking must be off" in response.json()["error"]["message"] +def test_noul_and_score(server): + response = post( + server, + { + "model": MODEL_NAME, + "state": {"ticket": "My card was charged twice for one order."}, + "questions": { + "urgent": {"type": "noul", "instructions": "Is money moving?"}, + "tone": { + "type": "score", + "instructions": "How upset is the customer?", + "criteria": ["calm", "annoyed", "furious"], + }, + }, + "chat_template_kwargs": {"enable_thinking": False}, + }, + ) + assert response.status_code == 200, response.text + body = response.json() + urgent = body["answers"]["urgent"] + assert urgent["type"] == "noul" + assert set(urgent["probabilities"]) == {"yes", "no"} + assert urgent["probabilities"]["yes"] + urgent["probabilities"]["no"] == ( + pytest.approx(1.0, abs=1e-4) + ) + assert urgent["noul"] == urgent["probabilities"]["yes"] + assert urgent["confidence"] == pytest.approx( + max(urgent["probabilities"].values()) + * body["diagnostics"]["urgent"]["label_mass"] + ) + tone = body["answers"]["tone"] + assert tone["type"] == "score" + assert tone["legend"] == {"0": "calm", "1": "annoyed", "2": "furious"} + assert set(tone["probabilities"]) == {"0", "1", "2"} + assert sum(tone["probabilities"].values()) == pytest.approx(1.0, abs=1e-4) + assert tone["score"] == pytest.approx( + sum(i * p for i, p in enumerate(tone["probabilities"].values())) + ) + assert 0 <= tone["score"] <= 2 + assert body["usage"]["output_tokens"] == 2 + + @pytest.mark.parametrize( "questions,match", [ @@ -208,6 +250,8 @@ def test_thinking_is_refused(server): }, "at most 64", ), + ({"q": {"type": "noul", "criteria": ["yes", "no"]}}, "must be an object"), + ({"q": {"type": "score", "criteria": ["calm"]}}, "ordered list of levels"), ], ) def test_validation_errors(server, questions, match): diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py index a9a6514a0373..0e80694b14eb 100644 --- a/vllm/entrypoints/generate/structured_decisions/question_types.py +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -18,8 +18,8 @@ class StructuredDecisionError(ValueError): """An invalid request. Returned as a 400.""" -#: Options are labeled A to Z in order. A single capital letter is one token -#: at the start of a reply for the tokenizers tested; every read checks it. +#: A choice labels its options A to Z in order. A single capital letter is one +#: token at the start of a reply for the tokenizers tested. Every read checks it. LABELS = tuple(string.ascii_uppercase) @@ -40,6 +40,9 @@ class Question: class QuestionType(ABC): name: ClassVar[str] + #: The labels of a question's options, in option order. + label_set: ClassVar[tuple[str, ...]] = LABELS + reply_instruction: ClassVar[str] = "Answer with the letter of one option only." @abstractmethod def parse_options(self, qid: str, criteria: Any) -> list[Option]: ... @@ -52,16 +55,22 @@ def answer( ``question.labels[i]`` among the labels, and the list sums to 1. ``label_mass`` is the labels' total probability over the vocabulary.""" + def reply_line(self, label: str, option: Option) -> str | None: + """The option's line in the prompt, or None to leave it out.""" + return ( + f"{label}: {option.name} - {option.description}" + if option.description + else f"{label}: {option.name}" + ) + def prompt(self, question: Question) -> str: """The question as the model reads it, after the state.""" lines = [f"Question: {question.instructions}"] if question.instructions else [] for label, o in zip(question.labels, question.options): - lines.append( - f"{label}: {o.name} - {o.description}" - if o.description - else f"{label}: {o.name}" - ) - lines.append("Answer with the letter of one option only.") + line = self.reply_line(label, o) + if line is not None: + lines.append(line) + lines.append(self.reply_instruction) return "\n".join(lines) @@ -100,7 +109,7 @@ def build_question( names = [o.name for o in options] if len(set(names)) != len(names): raise StructuredDecisionError(f"question {qid!r}: duplicate option names") - limit = min(max_options, len(LABELS)) + limit = min(max_options, len(qtype.label_set)) if len(options) > limit: raise StructuredDecisionError( f"question {qid!r}: at most {limit} options for this model" @@ -112,7 +121,7 @@ def build_question( type=qtype, instructions=instructions, options=tuple(options), - labels=LABELS[: len(options)], + labels=qtype.label_set[: len(options)], ) @@ -158,3 +167,74 @@ def answer( # read that mostly wanted a non-label token reports low confidence. "confidence": probs[top] * label_mass, } + + +@register_question_type +class NoulQuestion(QuestionType): + """Yes or no. ``criteria`` may describe what true and false mean.""" + + name = "noul" + label_set = ("yes", "no") + reply_instruction = "Answer with yes or no only." + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if criteria is not None and ( + not isinstance(criteria, dict) or not criteria.keys() <= {"true", "false"} + ): + raise StructuredDecisionError( + f"question {qid!r}: noul criteria must be an object with true and false" + ) + criteria = criteria or {} + true = criteria.get("true") + false = criteria.get("false") + return [ + Option("yes", None if true is None else str(true)), + Option("no", None if false is None else str(false)), + ] + + def reply_line(self, label: str, option: Option) -> str | None: + return f"{label}: {option.description}" if option.description else None + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + return { + "type": self.name, + "noul": probs[0], + "probabilities": {"yes": probs[0], "no": probs[1]}, + # The chosen label's probability over the whole vocabulary, like a + # choice's confidence. + "confidence": max(probs) * label_mass, + } + + +@register_question_type +class ScoreQuestion(QuestionType): + """Rate the state on an ordered scale. ``criteria`` is the ordered list of + level names, from the first to the last level of the scale.""" + + name = "score" + # 0-indexed like the answer, so a level's label is its score. + label_set = tuple(string.digits) + reply_instruction = "Answer with the number of one level only." + + def parse_options(self, qid: str, criteria: Any) -> list[Option]: + if not isinstance(criteria, list) or len(criteria) < 2: + raise StructuredDecisionError( + f"question {qid!r}: score criteria must be an ordered list of levels" + ) + return [Option(str(level)) for level in criteria] + + def answer( + self, question: Question, probs: list[float], label_mass: float + ) -> dict[str, Any]: + top = argmax(probs) + return { + "type": self.name, + # The expected level, 0-indexed, so a level's index in the request's + # criteria is its value. + "score": sum(i * p for i, p in enumerate(probs)), + "legend": {str(i): o.name for i, o in enumerate(question.options)}, + "probabilities": {str(i): p for i, p in enumerate(probs)}, + "confidence": probs[top] * label_mass, + } diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 864443cbb0b9..af588e13a289 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -23,7 +23,13 @@ from vllm.tokenizers import TokenizerLike from .protocol import ReadPromptRequest -from .question_types import LABELS, Question, StructuredDecisionError, label_softmax +from .question_types import ( + LABELS, + QUESTION_TYPES, + Question, + StructuredDecisionError, + label_softmax, +) @dataclass @@ -73,17 +79,19 @@ async def read( def reply_label_ids( - tokenizer: TokenizerLike, prompt_ids: Sequence[int] + tokenizer: TokenizerLike, + prompt_ids: Sequence[int], + labels: Sequence[str] = LABELS, ) -> tuple[list[int], list[int]]: - """The prompt's ids after its last added token, and the token each of LABELS - adds after them as the first token of the reply. Raises ValueError if a - label is not one distinct token there.""" + """The prompt's ids after its last added token, and the token each of + ``labels`` adds after them as the first token of the reply. Raises + ValueError if a label is not one distinct token there.""" added = set(tokenizer.get_added_vocab().values()) start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) tail = list(prompt_ids[start:]) text = tokenizer.decode(tail) ids: list[int] = [] - for label in LABELS: + for label in labels: extended = tokenizer.encode(text + label, add_special_tokens=False) if extended[:-1] != tail or extended[-1] in ids: raise ValueError( @@ -116,9 +124,13 @@ def __init__(self, context: ReadContext): ) if isinstance(probe, str): probe = tokenizer.encode(probe, add_special_tokens=False) - # Every prompt ends with the same generation prompt, so the labels' - # tokens are the same for every question. - self.tail, self.label_ids = reply_label_ids(tokenizer, probe) + # Every prompt ends with the same generation prompt, so a label's token + # is the same for every question. + self.tail, _ = reply_label_ids(tokenizer, probe, ()) + self.label_ids: dict[str, int] = {} + for qtype in QUESTION_TYPES.values(): + _, ids = reply_label_ids(tokenizer, probe, qtype.label_set) + self.label_ids.update(zip(qtype.label_set, ids)) def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -173,7 +185,7 @@ async def read( "these chat options end the prompt differently, so the " "labels' tokens are unknown" ) - slots.append(self.label_ids[: len(q.labels)]) + slots.append([self.label_ids[label] for label in q.labels]) engine_inputs.append(engine_input) label_reads = await next_token_label_reads( From ad73a4740dd780c5620099261738a30b979b262c Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 9 Oct 2026 14:07:49 +0800 Subject: [PATCH 262/952] [Perf][DSv4.1] Trim the last KV-source layer in decoder replay (#59894) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 Co-authored-by: Juntian Liu --- ...test_deepseek_v41_decoder_replay_layers.py | 37 +++++++--- .../models/test_deepseek_v41_replay_batch.py | 45 +++++++++-- vllm/models/deepseek_v41/attention.py | 28 +++++++ .../deepseek_v41/decoder_replay_layers.py | 48 ++++++++---- vllm/models/deepseek_v41/nvidia/model.py | 62 ++++++++++++---- .../models/deepseek_v41/nvidia/model_state.py | 74 +++++++++++-------- 6 files changed, 219 insertions(+), 75 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index f3b124c5ac80..614d903ca10d 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -67,10 +67,8 @@ def _states(num_tokens: int): ) -def test_run_gathers_states_and_realigns_shared_indexer_buffers(): - topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int() - candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int() - topk_before, candidates_before = topk.clone(), candidates.clone() +def test_run_gathers_states(): + """Every row's KV is written first; the layers then run on the replay rows.""" seen = {} def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual): @@ -80,7 +78,8 @@ def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual seen["is_padding"] = replay_context.is_padding return residual, pre_mix - layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates], set()) + write_batch_kv = MagicMock() + layers = DecoderReplayLayers(WINDOW, run_layers, write_batch_kv, set(), "swa") states = _states(NUM_TOKENS) hidden, residual = states[0], states[-1] full, sub = object(), object() @@ -90,11 +89,10 @@ def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual outputs = layers(*states) assert get_forward_context() is context + write_batch_kv.assert_called_once_with(*states) rows = torch.tensor(REPLAY_ROWS, device=DEVICE) assert torch.equal(seen["hidden_states"], hidden[rows]) assert seen["attn_metadata"] is sub and seen["is_padding"] is None - assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows]) - assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows]) assert len(outputs) == 2 and outputs[0].shape == residual.shape assert torch.equal(outputs[0][rows], residual[rows]) assert outputs[0][1:173].abs().sum() == 0 @@ -107,12 +105,14 @@ def run_layers(hidden_states, *rest): seen["attn_metadata"] = get_forward_context().attn_metadata return (hidden_states,) - layers = DecoderReplayLayers(WINDOW, run_layers, [], set()) + write_batch_kv = MagicMock() + layers = DecoderReplayLayers(WINDOW, run_layers, write_batch_kv, set(), "swa") states = _states(NUM_TOKENS) full = object() with override_forward_context(_context(full)): (output,) = layers(*states) assert output is states[0] and seen["attn_metadata"] is full + write_batch_kv.assert_not_called() def _graph_context(num_tokens): @@ -125,6 +125,25 @@ def _graph_context(num_tokens): ) +def test_piecewise_graphs_break_out_by_size(): + """A PIECEWISE graph's KV write and break are fixed at capture, when no replay + batch is set, so both follow the graph's size rather than the replay batch.""" + write_batch_kv = MagicMock() + layers = DecoderReplayLayers( + WINDOW, lambda *states: ("whole",), write_batch_kv, set(), "swa" + ) + layers._run_in_graph_break = lambda *states: ("break",) # type: ignore[method-assign] + states = _states(NUM_TOKENS) + with override_forward_context(_graph_context(1024)): + assert layers(*states) == ("whole",) # no threshold: no break + layers.trim_threshold = 768 + assert layers(*states) == ("break",) + _set_replay_batch(layers, REPLAY_ROWS, object()) + with override_forward_context(_graph_context(576)): + assert not layers.uses_replay_batch() # a replay batch alone does not break out + write_batch_kv.assert_called_once_with(*states) + + def test_replay_graph_matches_eager(monkeypatch): """The replay graph of the next captured size pads and runs the rows.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) @@ -142,7 +161,7 @@ def test_replay_graph_matches_eager(monkeypatch): def run_layers(hidden, positions, ids, pre, post, mix, residual): return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] - layers = DecoderReplayLayers(4, run_layers, [], set()) + layers = DecoderReplayLayers(4, run_layers, MagicMock(), set(), "swa") manager = DecoderReplayCudaGraphManager(cfg, DEVICE, layers) def prepare(desc): diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index 810a1dbd9a02..bcd4db64b25c 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -11,6 +11,12 @@ import torch from vllm.config import CUDAGraphMode +from vllm.forward_context import ( + BatchDescriptor, + ForwardContext, + override_forward_context, +) +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers @@ -59,7 +65,7 @@ def state(monkeypatch): cfg.parallel_config.data_parallel_size = 1 cfg.compilation_config.fast_moe_cold_start = False cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) + layers = DecoderReplayLayers(WINDOW, MagicMock(), MagicMock(), set(), "swa_first") model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) builds: list = [] @@ -72,7 +78,7 @@ def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, ** attn_metadata=object(), ) builds.append(build) - return {"swa": build.attn_metadata} + return {"swa": build.attn_metadata, "swa_first": build.attn_metadata} monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) state = DeepseekV41ModelState(cfg, model, None, DEVICE) @@ -120,7 +126,8 @@ def _slot_mappings(num_tokens: int) -> torch.Tensor: def _prepare(state, batch, cg_mode): - """Runs prepare_attn; returns the replay batch and the replay build's inputs.""" + """Runs prepare_attn; returns the replay batch and the replay build's inputs. + The first replay layer's sliding-window build follows the replay build.""" state.prepare_attn( batch, cg_mode, @@ -130,7 +137,7 @@ def _prepare(state, batch, cg_mode): GROUPS, ) replay = state.decoder_replay_layers.replay_batch - return (None, None) if replay is None else (replay, state.builds[-1]) + return (None, None) if replay is None else (replay, state.builds[-2]) def test_replay_batch_keeps_each_request_window(state): @@ -151,8 +158,16 @@ def test_replay_batch_keeps_each_request_window(state): assert build.replay_start.tolist() == [0, 300 - WINDOW, 50] assert torch.equal(sub.positions, batch.positions[rows]) assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows]) + # The first replay layer keeps the encoder-side window starts: it wrote + # every row's window KV. + first = state.builds[-1] + assert first.batch is sub and first.cg_mode == CUDAGraphMode.NONE + assert first.replay_start.tolist() == [0, 0, 50] context = replay.forward_context - assert context.attn_metadata == {"swa": build.attn_metadata} + assert context.attn_metadata == { + "swa": build.attn_metadata, + "swa_first": first.attn_metadata, + } assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows]) assert context.dp_metadata is None and context.is_padding is None @@ -190,7 +205,8 @@ def test_replay_batch_keeps_adaptive_verification_query_bound(state): def test_graph_steps_trim_only_piecewise_at_threshold(state): """A CUDA graph keeps the layers on its whole batch, unless it is a PIECEWISE - one at the trim threshold or above; an eager step trims a batch to its rows.""" + one at the trim threshold or above, as its capture decided; an eager step + trims a batch to its rows.""" batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512) for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL): assert _prepare(state, batch, cg_mode) == (None, None) @@ -199,8 +215,21 @@ def test_graph_steps_trim_only_piecewise_at_threshold(state): assert replay.rows.tolist() == REPLAY_ROWS assert build.cg_mode == CUDAGraphMode.NONE assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS) - state._trim_threshold = 256 - replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE) + layers = state.decoder_replay_layers + graph = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE, + batch_descriptor=BatchDescriptor(num_tokens=512), + ) + for threshold in (513, 512): + layers.trim_threshold = threshold + replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE) + with override_forward_context(graph): # the capture's decision + assert ( + layers.uses_replay_batch() == (replay is not None) == (threshold == 512) + ) assert replay is not None and replay.rows.tolist() == REPLAY_ROWS assert _prepare(state, batch, CUDAGraphMode.FULL) == (None, None) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 73d9c0c7233d..559d874cc483 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -686,6 +686,34 @@ def forward( ) return self._o_proj(attn_out, positions) + def forward_kv(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> None: + """Write only the KV ``forward`` writes (SWA KV, compressed KV, index K), + skipping wq_b, the indexer's top-k and attention.""" + attn_metadata = get_forward_context().attn_metadata + qr_kv, kv_score, _ = self._run_parallel_input_projections(hidden_states) + _, _, kv = self._split_qkv_and_norm(qr_kv) + if isinstance(attn_metadata, dict): + # The KV half of _fused_qnorm_rope_kv_insert, bit for bit, without + # the unused query; same as dspark's _insert_context_kv. + cache = self.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + cast(Any, attn_metadata[self.swa_cache_layer.prefix]).slot_mapping, + positions, + self.rotary_emb.cos_sin_cache, + self.swa_cache_layer.block_size, + self._flashinfer_fp8_kv_scale + if cache.dtype == torch.float8_e4m3fn + else None, + self.kv_mxfp8, + ) + if self.compressor is not None: + latent = self.compressor(kv_score, positions) + self.compressor.insert_cache(latent, positions, self.rotary_emb) + if self.indexer is not None and self.indexer.owns_k: + self.indexer._produce_k(latent, positions, self.indexer_rotary_emb) + def bind_gemm_rs(self) -> None: """Fuse ``wo_b`` with the sequence-parallel TP reduce-scatter. diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index d2d7ffd15364..28aa5e399f63 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Decoder-side SWA bounded replay: running the replay layers on their batch. -The layers past the last KV-source layer own nothing but sliding-window KV, so -in eager prefill steps they run on each request's last ``window`` rows only. +Once the last KV-source layer has written every row's KV, it and the layers past +it (which own only SWA KV) run on each request's last ``window`` rows only. ``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention metadata and a forward context of its own, like a microbatch; ``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers @@ -21,6 +21,7 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import ( ForwardContext, + get_forward_context, in_piecewise_cudagraph, override_forward_context, ) @@ -39,35 +40,57 @@ class ReplayBatch: class DecoderReplayLayers: """Runs the replay layers on the step's replay batch. - ``run_layers`` takes a batch's layer inputs and returns its per-row - outputs. ``row_buffers`` hold per-row results the source layer's indexer - publishes for the layers after it; they are compacted to the replay rows - in place. ``metadata_prefixes`` are the attention metadata keys the layers - read, so the replay batch builds only those. + ``run_layers`` takes a batch's layer inputs and returns its per-row outputs; + ``write_batch_kv`` takes them too and writes the first replay layer's KV for + every row. ``metadata_prefixes`` are the attention metadata keys the layers + read, so the replay batch builds only those. ``first_swa_prefix`` keys the + first replay layer's SWA metadata, built with the batch's window starts. """ def __init__( self, window: int, run_layers: Callable[..., tuple[torch.Tensor, ...]], - row_buffers: list[torch.Tensor], + write_batch_kv: Callable[..., None], metadata_prefixes: set[str], + first_swa_prefix: str, ) -> None: self.window = window self.run_layers = run_layers - self.row_buffers = row_buffers + self.write_batch_kv = write_batch_kv self.metadata_prefixes = metadata_prefixes + self.first_swa_prefix = first_swa_prefix # Set by the model state every step; None runs the layers on the batch. self.replay_batch: ReplayBatch | None = None + # Set by the model state with replay graphs: PIECEWISE model graphs of + # this many tokens or more break out to the replay batch. + self.trim_threshold: int | None = None self._break_outputs: list[torch.Tensor] | None = None def __call__( self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None ) -> tuple[torch.Tensor, ...]: - if self.replay_batch is not None and in_piecewise_cudagraph(): - return self._run_in_graph_break(hidden_states, *states) + if self.uses_replay_batch(): + # Outside the break, so a PIECEWISE graph captures it. + self.write_batch_kv(hidden_states, *states) + if in_piecewise_cudagraph(): + return self._run_in_graph_break(hidden_states, *states) return self._run(hidden_states, *states) + def uses_replay_batch(self) -> bool: + """Whether this forward uses the replay batch: in a PIECEWISE graph by its + size, fixed at capture (no replay batch is set then); else if one is set.""" + if in_piecewise_cudagraph(): + batch_descriptor = get_forward_context().batch_descriptor + assert batch_descriptor is not None + return self.graph_uses_replay_batch(batch_descriptor.num_tokens) + return self.replay_batch is not None + + def graph_uses_replay_batch(self, num_tokens: int) -> bool: + """Whether a PIECEWISE graph of ``num_tokens`` breaks out to the replay + batch; capture and the model state's runtime replay batch share it.""" + return self.trim_threshold is not None and num_tokens >= self.trim_threshold + @eager_break_during_capture def _run_in_graph_break( self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None @@ -91,9 +114,6 @@ def _run( # Replay batch steps hold more than `window` tokens: no deferred MoE finalize. assert isinstance(hidden_states, torch.Tensor) rows = replay_batch.rows - num_rows = rows.shape[0] - for buf in self.row_buffers: - buf[:num_rows].copy_(buf.index_select(0, rows)) inputs = (hidden_states, *states) with override_forward_context(replay_batch.forward_context): if replay_batch.run_graph is not None: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 94fa0b4b8b67..18f8f27843ab 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -623,6 +623,40 @@ def forward( torch.cuda.current_stream().wait_stream(mhc_stream) return x, residual, post_mix, res_mix, ffn_pre, previous_aux + def write_kv( + self, + x: torch.Tensor | MoEOutput, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> None: + """Write the KV ``forward`` would for the replay layers' inputs, skipping + the rest of the layer; the replay batch then reruns its rows.""" + # The replay batch rules out Engram and sequence parallel here. + assert self.engram is None and not self.use_sequence_parallel + *_, x, _, _ = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + reduce_results=self.fuse_mhc_all_reduce, + ) + self.attn.forward_kv(positions, x) + class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -765,10 +799,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): and self._decoder_replay_supported(vllm_config, cut) and self.layers[cut].attn.swa_cache_layer.bounded_replay ): - self.decoder_replay_start = cut + 1 + assert self.layers[cut].attn.compress_ratio <= 1 + # The last KV source writes every row's KV, then replays. + self.decoder_replay_start = cut # The attention metadata keys the replay layers read. metadata_prefixes = set() - for layer in islice(self.layers, cut + 1, self.end_layer): + for layer in islice(self.layers, cut, self.end_layer): attn = typing.cast(DeepseekV4DecoderLayer, layer).attn metadata_prefixes.add(attn.swa_cache_layer.prefix) if attn.compressed_cache_prefix is not None: @@ -778,16 +814,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.decoder_replay_layers = DecoderReplayLayers( config.sliding_window, self._run_replay_layers, - [ - buf - for buf in (self.topk_indices_buffer, self.candidate_block_buffer) - if buf is not None - ], + self.layers[cut].write_kv, metadata_prefixes, + self.layers[cut].attn.swa_cache_layer.prefix, ) logger.info_once( - "Decoder SWA bounded replay: in eager prefill steps, layers " + "Decoder SWA bounded replay: in eager prefill steps, layer %d " + "writes its KV for every token, and the rest of it and layers " "%d-%d run on each request's last %d tokens only.", + cut, cut + 1, self.end_layer - 1, config.sliding_window, @@ -1120,9 +1155,8 @@ def _run_replay_layers( res_mix: torch.Tensor, residual: torch.Tensor, ) -> tuple[torch.Tensor, ...]: - """The layers past the last KV source, on whatever rows they are given; - returns their output, the last FFN's pre-mix and the aux hidden states - they capture.""" + """Layers ``decoder_replay_start``.. on the given rows; returns their output, + the last FFN's pre-mix and the aux hidden states they capture.""" aux_hidden_by_layer: dict[int, torch.Tensor] = {} hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( range(self.decoder_replay_start, self.end_layer), @@ -1154,7 +1188,7 @@ def _run_replay_layers( ) def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: - """Whether this rank may trim the layers after ``cut``; warns when not.""" + """Whether this rank may replay from layer ``cut`` on; warns when not.""" parallel_config = vllm_config.parallel_config spec_config = vllm_config.speculative_config draft_config = spec_config.draft_model_config if spec_config else None @@ -1176,8 +1210,8 @@ def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: "the replay-layer batch shrinks per rank, which sequence and " "prefill-context parallelism and microbatching cannot follow" ) - elif any(i > cut for i in getattr(self.config, "engram_layer_ids", ())): - reason = "an Engram layer sits after the last KV source layer" + elif any(i >= cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits at or after the last KV source layer" elif draft_config is not None and ( draft_window is None or draft_window > window diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index c09e4582ea34..cc3ade80b8ad 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -238,12 +238,13 @@ def __init__( self._kept_kv_start = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=device ) - self._replay_attn_groups: list[list[AttentionGroup]] | None = None + # The replay layers' and the first one's SWA attention groups. + self._replay_attn_groups: ( + tuple[list[list[AttentionGroup]], list[list[AttentionGroup]]] | None + ) = None layers = self.decoder_replay_layers compilation_config = vllm_config.compilation_config self.replay_cudagraphs: DecoderReplayCudaGraphManager | None = None - # PIECEWISE graphs of this many tokens or more break out to the replay batch. - self._trim_threshold: int | None = None if ( layers is not None and compilation_config.cudagraph_decoder_replay @@ -263,7 +264,7 @@ def __init__( self.replay_cudagraphs = DecoderReplayCudaGraphManager( vllm_config, device, layers ) - self._trim_threshold = threshold + layers.trim_threshold = threshold def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) @@ -454,8 +455,7 @@ def _prepare_replay_batch( layers.replay_batch = None force_replay_batch = capture_desc is not None or ( cudagraph_mode == CUDAGraphMode.PIECEWISE - and self._trim_threshold is not None - and input_batch.num_tokens_after_padding >= self._trim_threshold + and layers.graph_uses_replay_batch(input_batch.num_tokens_after_padding) ) if cudagraph_mode != CUDAGraphMode.NONE and not force_replay_batch: return @@ -509,17 +509,35 @@ def _prepare_replay_batch( kept_batch, kept_slot_mappings = self._kept_input_batch( input_batch, slot_mappings, replay_start, kept_lens, num_padded ) + if self._replay_attn_groups is None: + self._replay_attn_groups = ( + self._copy_attn_groups(attn_groups, layers.metadata_prefixes), + self._copy_attn_groups(attn_groups, {layers.first_swa_prefix}), + ) + replay_groups, first_swa_groups = self._replay_attn_groups attn_metadata = super().prepare_attn( kept_batch, CUDAGraphMode.NONE, block_tables, kept_slot_mappings, - self._replay_groups(attn_groups), + replay_groups, kv_cache_config, model_specific_attn_metadata=ReplayAttnMetadata( self._kept_kv_start[:num_reqs] ), ) + # The first replay layer wrote every row's window KV: its window starts + # are the batch's. + first_swa = layers.first_swa_prefix + attn_metadata[first_swa] = super().prepare_attn( + kept_batch, + CUDAGraphMode.NONE, + block_tables, + kept_slot_mappings, + first_swa_groups, + kv_cache_config, + model_specific_attn_metadata=ReplayAttnMetadata(replay_start), + )[first_swa] layers.replay_batch = ReplayBatch( self._kept_rows[:num_tokens], create_forward_context( @@ -636,27 +654,23 @@ def _agree_across_dp( return False, None return True, agreed[1] - def _replay_groups( - self, attn_groups: list[list[AttentionGroup]] + def _copy_attn_groups( + self, attn_groups: list[list[AttentionGroup]], prefixes: set[str] ) -> list[list[AttentionGroup]]: - """The attention groups with metadata builders of their own, like each - microbatch's: a builder keeps the metadata it built, and the runner's - hold the batch's. Only the groups the replay layers read are built.""" - if self._replay_attn_groups is None: - assert self.decoder_replay_layers is not None - prefixes = self.decoder_replay_layers.metadata_prefixes - self._replay_attn_groups = [] - for groups in attn_groups: - replay_groups = [] - for group in groups: - if prefixes.isdisjoint(group.layer_names): - continue - replay_group = replace(group, metadata_builders=[]) - replay_group.create_metadata_builders( - self.vllm_config, - self.device, - group.metadata_builders[0].kernel_block_size, - ) - replay_groups.append(replay_group) - self._replay_attn_groups.append(replay_groups) - return self._replay_attn_groups + """Copies of the groups holding any of ``prefixes``, with builders of their + own like a microbatch's: a builder keeps the metadata it built.""" + own_groups = [] + for groups in attn_groups: + copies = [] + for group in groups: + if prefixes.isdisjoint(group.layer_names): + continue + copy = replace(group, metadata_builders=[]) + copy.create_metadata_builders( + self.vllm_config, + self.device, + group.metadata_builders[0].kernel_block_size, + ) + copies.append(copy) + own_groups.append(copies) + return own_groups From f0a5f111f205b5c73bb34fbd41f8f0d9936b543f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Ma=C5=84czak?= Date: Fri, 9 Oct 2026 08:29:29 +0200 Subject: [PATCH 263/952] [XPU] Add tuned Mamba SSU configs for Intel Arc Pro B60 (#56765) Signed-off-by: pmanczak Co-authored-by: Kunshang Ji --- ..._Pro_B60_Graphics,cache_dtype=float16.json | 51 +++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 51 +++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float16.json | 63 +++++++++++++++++++ ..._Pro_B60_Graphics,cache_dtype=float32.json | 63 +++++++++++++++++++ 8 files changed, 480 insertions(+) create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json create mode 100644 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..64087ce74c28 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..0a378592c522 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..91c6e28fa794 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "16": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "64": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "128": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "8192": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "16384": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "32768": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..c30784ce7762 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..5f29bcd76c22 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..8f815a1b337d --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..2df252554195 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + } +} diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..f943b523c441 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=64,device_name=Intel(R)_Arc(TM)_Pro_B60_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.8.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} From dd30786fef8b15bd83205833c7b3348b0ba0ebc7 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Fri, 9 Oct 2026 12:10:21 +0530 Subject: [PATCH 264/952] [CI] Enable Zen CPU image build on AMD CI (#60016) Signed-off-by: Chinmay Kulkarni Co-authored-by: Cursor Co-authored-by: Micah Williamson --- .buildkite/hardware_tests/cpu.yaml | 54 +++++++++++++ .buildkite/image_build/image_build_zen_cpu.sh | 75 +++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100755 .buildkite/image_build/image_build_zen_cpu.sh diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index d54688b57491..a29b44ed3c71 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -232,3 +232,57 @@ steps: "$$BUILDKITE_PARALLEL_JOB" "$REGISTRY/$REPO:$BUILDKITE_COMMIT-arm64-cpu" parallelism: 3 + +- label: ":docker: Build Zen CPU image" + key: zen-cpu-image-build + # Reuse the shared `:-cpu` image that image-build-cpu publishes + # to public ECR, so Dockerfile.zen just layers on top instead of recompiling + # vLLM from source. If the pull fails, the build script still falls back to + # building the CPU base from source. + depends_on: + - image-build-cpu + soft_fail: true + agents: + queue: amd-cpu + source_file_dependencies: + - .buildkite/hardware_tests/cpu.yaml + - .buildkite/image_build/image_build_zen_cpu.sh + - docker/Dockerfile.zen + - docker/Dockerfile.cpu + - requirements/cpu.txt + - requirements/build/cpu.txt + - requirements/test/cpu.txt + - cmake/cpu_extension.cmake + - vllm/model_executor/layers/utils.py + - vllm/model_executor/kernels/linear/zentorch_utils.py + - vllm/model_executor/kernels/linear/scaled_mm/zentorch.py + - vllm/model_executor/kernels/linear/mixed_precision/zentorch.py + - vllm/platforms/cpu.py + - vllm/platforms/zen_cpu.py + - vllm/platforms/__init__.py + - tests/model_executor/test_cpu_unquantized_gemm_dispatch.py + - tests/test_zen_cpu_platform_detection.py + - tests/kernels/moe/test_zen_cpu_fused_moe.py + - tests/kernels/moe/test_zen_cpu_int8_moe.py + - vllm/v1/attention/backends/zentorch_sdpa.py + - tests/test_zentorch_version_pin.py + - vllm/v1/attention/backends/mla/prefill/zen_cpu_sdpa.py + - vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py + - tests/kernels/quantization/test_zen_da8w4.py + commands: + - .buildkite/image_build/image_build_zen_cpu.sh "$REGISTRY" "$REPO" "$BUILDKITE_COMMIT" + env: + DOCKER_BUILDKIT: "1" + # Install zentorch directly (if-branch in Dockerfile.zen) instead of + # `uv pip install "vllm[zen]"`, which forces uv to re-resolve vLLM's full + # dependency closure and chokes on its transitive triton-cpu URL dependency. + # Keep this pin in sync with the `zen` extra in setup.py. + ZENTORCH_VERSION: "2.13.0.1" + retry: + automatic: + - exit_status: 1 # Transient Docker/BuildKit failure + limit: 1 + - exit_status: -1 # Agent was lost + limit: 2 + - exit_status: -10 # Agent was lost + limit: 2 diff --git a/.buildkite/image_build/image_build_zen_cpu.sh b/.buildkite/image_build/image_build_zen_cpu.sh new file mode 100755 index 000000000000..3e42e38a8d75 --- /dev/null +++ b/.buildkite/image_build/image_build_zen_cpu.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Build the AMD Zen CPU image (vLLM + zentorch) by layering docker/Dockerfile.zen +# on top of a CPU base image that already has vLLM installed. +# +# To avoid recompiling vLLM from source, prefer reusing the shared CPU image that +# the image-build-cpu step already builds and publishes for this commit +# (`:-cpu`). It lives in public ECR, which allows anonymous pulls, +# so no registry credentials are required. If the pull fails (e.g. the image +# isn't published yet), fall back to building the base from docker/Dockerfile.cpu. +# +# The zen image is not pushed to a registry. +# +# See docker/Dockerfile.zen for the build workflow this mirrors. +set -euo pipefail + +if [[ $# -lt 3 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 + +image_repo() { + if [[ -n "${REGISTRY}" ]]; then + printf '%s/%s' "${REGISTRY}" "${REPO}" + else + printf '%s' "${REPO}" + fi +} + +IMAGE_REPO="$(image_repo)" + +# The shared CPU image published by image-build-cpu (public ECR, anonymous pull). +SHARED_CPU_IMAGE="" +if [[ -n "${REGISTRY}" ]]; then + SHARED_CPU_IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu" +fi +FALLBACK_BASE_IMAGE="zen-cpu-base:$BUILDKITE_COMMIT" +IMAGE="$IMAGE_REPO:$BUILDKITE_COMMIT-zen-cpu" + +# ZENTORCH_VERSION is optional; when unset the Dockerfile falls back to +# installing zentorch via `vllm[zen]`. +ZENTORCH_VERSION=${ZENTORCH_VERSION:-} + +# Step 1: obtain the CPU base image that Dockerfile.zen layers on. Prefer pulling +# the published `-cpu` image; fall back to building it from source if unavailable. +if [[ -n "${SHARED_CPU_IMAGE}" ]] && docker pull "$SHARED_CPU_IMAGE"; then + echo "--- :docker: Using published CPU image as base: $SHARED_CPU_IMAGE" + BASE_IMAGE="$SHARED_CPU_IMAGE" +else + echo "--- :docker: Published CPU image unavailable; building base from source" + docker build --file docker/Dockerfile.cpu \ + --platform linux/amd64 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg VLLM_CPU_X86=true \ + --tag "$FALLBACK_BASE_IMAGE" \ + --target vllm-openai \ + --progress plain . + BASE_IMAGE="$FALLBACK_BASE_IMAGE" +fi + +# Step 2: build the zen test image on top of the CPU base. +echo "--- :docker: Building Zen test image" +docker build --file docker/Dockerfile.zen \ + --platform linux/amd64 \ + --build-arg BASE_IMAGE="$BASE_IMAGE" \ + ${ZENTORCH_VERSION:+--build-arg ZENTORCH_VERSION="$ZENTORCH_VERSION"} \ + --tag "$IMAGE" \ + --target vllm-zen-test \ + --progress plain . From 19c0bed584bde3fe785296d2923247ab9d678483 Mon Sep 17 00:00:00 2001 From: Jakub Zakrzewski Date: Fri, 9 Oct 2026 09:12:03 +0200 Subject: [PATCH 265/952] [Fusion] Support 3D activation quantization fusion and enable it for Pixtral (#59612) Signed-off-by: Jakub Zakrzewski Co-authored-by: Codex Co-authored-by: Misha Goin --- .../fusion/test_quant_activation_contract.py | 2 + .../test_flashinfer_nvfp4_scaled_mm.py | 52 +++++++++++++++++++ tests/kernels/test_fused_quant_activation.py | 32 +++++++----- .../kernels/linear/nvfp4/flashinfer.py | 31 +++++++---- .../layers/fusion/fused_act_quant.py | 8 +-- vllm/model_executor/models/pixtral.py | 5 +- 6 files changed, 101 insertions(+), 29 deletions(-) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index 15b9835d6cf5..da507e5caf50 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -21,6 +21,7 @@ NvFp4LinearLayerConfig, ) from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) @@ -59,6 +60,7 @@ SUPPORTING = { CutlassFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, + FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, AiterHipbMMPerTokenFp8ScaledMMLinearKernel, diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index 54dcaa0b97ea..47c25d046d2c 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -13,11 +13,15 @@ from vllm import _custom_ops as ops from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, FlashInferCuteDslNvFp4W4A16LinearKernel, + FlashInferCutlassNvFp4LinearKernel, ) +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_to_dtype, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.flashinfer import ( flashinfer_scaled_fp4_mm, @@ -47,6 +51,54 @@ CUDA_DEVICES = ["cuda:0"] +@pytest.mark.parametrize("token_shape", [(17,), (1, 17), (2, 17)]) +@pytest.mark.parametrize( + "kernel_cls", + [FlashInferCutlassNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel], +) +@torch.inference_mode() +def test_flashinfer_nvfp4_preserves_quantized_input_shape( + token_shape: tuple[int, ...], + kernel_cls: type[FlashInferCutlassNvFp4LinearKernel] + | type[FlashInferCuteDslNvFp4LinearKernel], +) -> None: + """Pre-quantized inputs must match the tensor path and retain token dimensions.""" + supported, reason = kernel_cls.is_supported() + if not supported: + pytest.skip(reason) + + set_random_seed(42) + hidden_size, output_size = 128, 128 + x = torch.randn(*token_shape, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(output_size, hidden_size, dtype=x.dtype, device=x.device) + input_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / x.abs().max()).float() + weight_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).float() + x_fp4, x_blockscale = ops.scaled_fp4_quant(x.reshape(-1, hidden_size), input_scale) + weight_fp4, weight_blockscale = ops.scaled_fp4_quant(weight, weight_scale) + quantized_input = QuantizedActivation( + data=x_fp4.reshape(*token_shape, hidden_size // 2), + scale=x_blockscale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kNvfp4Dynamic, + ) + layer = torch.nn.Module() + layer.input_size_per_partition = hidden_size + layer.output_size_per_partition = output_size + layer.weight = weight_fp4 + layer.weight_scale = weight_blockscale + layer.input_global_scale_inv = input_scale + layer.alpha = (input_scale * weight_scale).reciprocal() + kernel = kernel_cls(NvFp4LinearLayerConfig()) + + expected = kernel.apply_weights(layer, x.reshape(-1, hidden_size)) + actual = kernel.apply_weights(layer, quantized_input) + + assert actual.shape == (*token_shape, output_size) + assert actual.dtype == x.dtype + torch.testing.assert_close(actual.reshape(expected.shape), expected) + + def get_ref_results( a_fp4, b_fp4, diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 321cb6978e3d..90309dbd1d9b 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -115,13 +115,13 @@ class MockLinearRequiresUnquantized(MockLinearFp8Static): requires_unquantized_input = True -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_static( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -133,26 +133,29 @@ def test_maybe_fused_act_quant_fp8_static( scale = torch.tensor([0.5], device=device, dtype=torch.float32) linear = MockLinearFp8Static(scale) - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8StaticTensorSym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape - ref_out = ref_impl(act_fn, x, scale) + ref_out = ref_impl(act_fn, x.reshape(-1, hidden_size * 2), scale).reshape( + result.orig_shape + ) torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32)) -@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("token_shape", [(1,), (16,), (128,), (1, 17), (2, 17)]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @torch.inference_mode() def test_maybe_fused_act_quant_fp8_dynamic_block( default_vllm_config, - num_tokens: int, + token_shape: tuple[int, ...], hidden_size: int, dtype: torch.dtype, ) -> None: @@ -166,31 +169,34 @@ def test_maybe_fused_act_quant_fp8_dynamic_block( linear = MockLinearFp8Dynamic128() scale = 1 / hidden_size - x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale + x = torch.randn(*token_shape, hidden_size * 2, dtype=dtype, device=device) * scale result = maybe_fused_act_quant(act_fn, x, linear) assert isinstance(result, QuantizedActivation) assert result.quant_key == kFp8Dynamic128Sym assert result.data.dtype == current_platform.fp8_dtype() assert result.orig_dtype == dtype - assert result.orig_shape == (num_tokens, hidden_size) + assert result.orig_shape == (*token_shape, hidden_size) + assert result.data.shape == result.orig_shape num_groups = hidden_size // group_size - assert result.scale.shape == (num_tokens, num_groups) + assert result.scale.shape == (*token_shape, num_groups) gate, up = x.split(hidden_size, dim=-1) silu_out = F.silu(gate) * up ref_out, ref_scales = per_token_group_quant_fp8( - silu_out, group_size=group_size, use_ue8m0=False + silu_out.reshape(-1, hidden_size), group_size=group_size, use_ue8m0=False ) + ref_out = ref_out.reshape(result.orig_shape) + ref_scales = ref_scales.reshape(*token_shape, num_groups) torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5) ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave( - group_size, dim=1 + group_size, dim=-1 ) torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index a555ca17bf89..2a69a3be2100 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -112,6 +112,9 @@ def apply_weights( class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel): """NVFP4 GEMM via FlashInfer's cutedsl backend.""" + def input_quant_key(self) -> QuantKey | None: + return kNvfp4Dynamic + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -137,19 +140,26 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: output_size = layer.output_size_per_partition - output_dtype = x.dtype - output_shape = [*x.shape[:-1], output_size] - - x_fp4, x_blockscale = scaled_fp4_quant( - x, - layer.input_global_scale_inv, - is_sf_swizzled_layout=True, - backend="flashinfer-cutedsl", - ) + qa = as_quantized_activation(x, self.input_quant_key()) + if qa is not None: + x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) + output_dtype = qa.orig_dtype + output_shape = [*qa.orig_shape[:-1], output_size] + else: + assert isinstance(x, torch.Tensor) + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-cutedsl", + ) x_fp4 = pad_nvfp4_activation_for_cutlass( x_fp4, nvfp4_weight_padding_bytes(layer) @@ -219,6 +229,7 @@ def apply_weights( qa = as_quantized_activation(x, self.input_quant_key()) if qa is not None: x_fp4, x_blockscale = qa.data, qa.scale + x_fp4 = x_fp4.reshape(-1, x_fp4.shape[-1]) x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes) output_dtype = qa.orig_dtype output_shape = [*qa.orig_shape[:-1], output_size] diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index a9b14c617000..6261a1bdd019 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -63,10 +63,10 @@ def _silu_and_mul_fp8_dynamic_block( x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey ) -> QuantizedActivation: """SiluAndMul + FP8 dynamic per-block quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] num_groups = d // group_size @@ -102,10 +102,10 @@ def _silu_and_mul_nvfp4_dynamic( x: torch.Tensor, linear: LinearBase ) -> QuantizedActivation: """SiluAndMul + NVFP4 dynamic quantization.""" - assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" - + assert x.ndim in (2, 3), f"Input must be 2D or 3D, got {x.shape}" d = x.shape[-1] // 2 out_shape = x.shape[:-1] + (d,) + x = x.reshape(-1, x.shape[-1]) num_tokens = x.shape[0] # NVFP4 packs 2 values into 1 byte diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index fb10026ad19a..4f96c0acfcc0 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.activation import SiluAndMul, get_act_and_mul_fn from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.conv import Conv2dLayer +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -758,7 +759,7 @@ def __init__( def forward(self, x): x, _ = self.gate_up_proj(x) - x = self.act_fn(x) + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) x, _ = self.down_proj(x) return x @@ -1259,7 +1260,7 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) - x = self.act_and_mul(gate_up) + x = maybe_fused_act_quant(self.act_and_mul, gate_up, self.down_proj) x, _ = self.down_proj(x) return x From 82e89e99440661f3404c58c9c60d712657305cac Mon Sep 17 00:00:00 2001 From: Zheng Cai Date: Fri, 9 Oct 2026 15:18:50 +0800 Subject: [PATCH 266/952] [Perf][Qwen4Exp] Retune H200 M=4 merged QSA LL-GEMM plan (#60776) Signed-off-by: zigzagcai Co-authored-by: zigzagcai Co-authored-by: Claude Code --- vllm/models/qwen4_exp/nvidia/low_latency_gemm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index acc5b3e6bd98..2ee1c7cb5a54 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -79,7 +79,7 @@ (4224, 2560): { 1: SkinnyGemmConfig(1, 128, 2, k_unroll=6, vector_width=2, static_k=2560), 2: SkinnyGemmConfig(2, 128, 2, k_unroll=3, vector_width=2, static_k=2560), - 4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4), + 4: SkinnyGemmConfig(4, 64, 4, k_unroll=1, vector_width=8), }, # GDN and QSA output projections, TP=4. (2560, 1536): { From efaca9547dbdfd2d5269f274aaf3798c1f730435 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Fri, 9 Oct 2026 00:30:18 -0700 Subject: [PATCH 267/952] [Perf][DSv4.1] Exchange packed FP8 Engram rows across DP and feed wkv MXFP8 (#59923) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_engram_dp_shard.py | 73 ++++++ tests/kernels/test_engram.py | 13 +- vllm/models/deepseek_v41/common/engram.py | 223 +++++++++++++++++- .../deepseek_v41/common/ops/query_quant.py | 24 +- vllm/models/deepseek_v41/nvidia/model.py | 30 ++- .../deepseek_v41/nvidia/ops/engram_peer.py | 173 ++++++++++++++ 6 files changed, 505 insertions(+), 31 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/engram_peer.py diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 637d179a5d21..8b7d3e8cec77 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -39,9 +39,16 @@ initialize_model_parallel, ) from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + convert_swizzled_to_linear, +) from vllm.models.deepseek_v41.common import engram as engram_ops from vllm.models.deepseek_v41.common.engram import ( Engram, + EngramBatch, EngramLayout, ParallelEngramEmbedding, engram_head_shard_rank, @@ -207,10 +214,75 @@ def _worker(rank: int, tp_size: int, port: int) -> None: n_heads, sequence_parallel, ) + if tp_size == 1: + _check_fp8_batch(vllm_config) finally: cleanup_dist_env_and_memory() +def _check_fp8_batch(vllm_config): + """MXFP8 staging matches the tables across peer/all-to-all slots, idle + replicas and graph replays.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + dp_rank = vllm_config.parallel_config.data_parallel_rank + layers = [] + for index in range(2): + weight, scale = _full_table(sum(HEAD_SIZES)) + weight = (weight.float() * (1 - 2 * index)).to(torch.float8_e4m3fn) + table = ParallelEngramEmbedding( + sum(HEAD_SIZES), DIM, HEAD_SIZES, cpu_offload=True + ) + table.weight.weight_loader(table.weight, weight) + table.weight_scale_inv.weight_loader(table.weight_scale_inv, scale) + staged = torch.empty(0, device="cuda") + reference = (weight.cuda(), scale.cuda()) + layers.append(SimpleNamespace(embed_tokens=table, staged_rows=staged)) + layers[-1].layer_hash_index, layers[-1].reference = index, reference + batch = EngramBatch(layers, 257) + + def check(ids): + tokens = ids.shape[0] + for layer in layers: + x = batch.wkv_input(layer, tokens) + sf = convert_swizzled_to_linear(x.scale, tokens, x.data.shape[1], BLOCK) + actual = dequant_mxfp8_to_bf16(x.data, sf).unflatten(-1, (-1, DIM)) + weight, scale = layer.reference + expected = _reference(weight, scale, ids[:, layer.layer_hash_index]) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + for counts in ((7, 0, 5, 1), (127,) * 4, (128,) * 4, (257, 0, 129, 1)): + tokens = counts[dp_rank] + seed = 700 + 10 * dp_rank + ids = torch.stack([_make_ids(HEAD_SIZES, tokens, seed + i) for i in (0, 1)], 1) + with set_forward_context( + None, + vllm_config, + num_tokens=tokens, + num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32), + ): + batch.prepare_embeddings(ids) + check(ids) + for capture in ("full", "breakable") if len(set(counts)) == 1 else (): + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + if capture == "full": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=warmup): + batch.prepare_embeddings(ids) + else: + graph = BreakableCUDAGraphCapture() + with torch.cuda.stream(warmup), graph: + batch.prepare_embeddings(ids) + torch.cuda.current_stream().wait_stream(warmup) + for step in (2, 3): + replay_ids = _make_ids(HEAD_SIZES, tokens, seed + step) + ids.copy_(torch.stack([replay_ids] * 2, 1)) + graph.replay() + check(ids) + del graph + + def _check_shared_storage_lifetime_and_failures(failure): """One peer's failure must roll back every mapping/registration, without hangs.""" import gc @@ -516,6 +588,7 @@ def __call__(self, hidden, positions, input_ids, *args, **kwargs): model.use_native_mega_moe = False model.use_sequence_parallel = False model.fuse_mhc_all_reduce = False + model.engram_batch = None model.engram_hash = state model.engram_swa_prefix = "swa" model.engram_dp_shared_memory = dp_shared_memory diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 15bdc9162cf6..3663108d9d1a 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -7,6 +7,9 @@ import pytest import torch +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, +) from vllm.models.deepseek_v41.common import engram as engram_ops from vllm.models.deepseek_v41.common.engram import ( Engram, @@ -714,8 +717,9 @@ def test_engram_lookup_reuses_jit_across_token_shapes(): ) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize("packed", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 7, 256]) -def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): +def test_engram_lookup_matches_torch(cpu_offload, background, packed, num_tokens): """The fused gather must be bit-exact with the torch dequant path it replaces, from HBM and from pinned host memory alike, and must contribute zeros for rows another TP rank owns.""" @@ -729,8 +733,13 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): layer.vocab_start_idx, layer.vocab_end_idx, ) - out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda") + width = layer.dim + layer.dim // 32 if packed else layer.dim + dtype = torch.uint8 if packed else torch.bfloat16 + out = torch.empty(num_tokens, cols, width, dtype=dtype, device="cuda") layer.lookup(ids, out, background=background) + if packed: + values = out[..., : layer.dim].view(torch.float8_e4m3fn) + out = dequant_mxfp8_to_bf16(values, out[..., layer.dim :]) assert torch.equal(out, expected) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index d4810ee4ea9e..0e160c8b06cf 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -62,9 +62,16 @@ from vllm.distributed.parallel_state import GroupCoordinator from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v41.common.ops.query_quant import ( + can_fuse_query_quant, + mxfp8_scale_bytes, + mxfp8_scale_offset, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available @@ -72,6 +79,10 @@ logger = init_logger(__name__) +# Per-rank token slot from which the Engram DP exchange switches from peer +# reads to all-to-all. +ENGRAM_A2A_MIN_SLOT = 128 + # Cache value for tokens that take no part in an n-gram (image spans). DEAD_ID = -1 @@ -644,8 +655,10 @@ def _engram_lookup_kernel( BLOCK_R: tl.constexpr, GRID, SORTED: tl.constexpr if _SORTED_IS_CONSTEXPR else None, # type: ignore[valid-type] + PACKED: tl.constexpr, + OUT_STRIDE: tl.constexpr, ): - """Gather fp8 rows, apply their ue8m0 block scales, write bf16. + """Gather FP8 rows as packed values/scales or dequantized BF16. Only this rank's heads are read; padded heads write zeros for all-gather. `weight`/`scales` may address pinned host memory through UVA. If SORTED, @@ -684,13 +697,28 @@ def _engram_lookup_kernel( mask=owned[:, None], other=0, ) - # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. - scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) - tl.store( - out + dst[:, None] * DIM + cols[None, :], - (values.to(tl.float32) * scale).to(tl.bfloat16), - mask=valid[:, None], - ) + if PACKED: + width: tl.constexpr = DIM + DIM // QUANT_BLOCK + dst = dst.to(tl.int64) + dst = dst // LOCAL_HEADS * OUT_STRIDE + dst % LOCAL_HEADS * width + tl.store( + out + dst[:, None] + cols[None, :], + values.to(tl.uint8, bitcast=True), + mask=valid[:, None], + ) + tl.store( + out + dst[:, None] + DIM + scale_cols[None, :], + scale, + mask=valid[:, None] & (cols[None, :] % QUANT_BLOCK == 0), + ) + else: + # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. + scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) + tl.store( + out + dst[:, None] * DIM + cols[None, :], + (values.to(tl.float32) * scale).to(tl.bfloat16), + mask=valid[:, None], + ) def engram_head_shard_rank() -> int: @@ -1082,7 +1110,7 @@ def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: def lookup( self, indices: torch.Tensor, out: torch.Tensor, background: bool = False ) -> None: - """Look up local heads of [T, heads] into [T, local_heads, dim] bf16. + """Look up local heads into BF16, or packed FP8 bytes followed by scales. `background` limits the grid to leave SMs for concurrent work. """ @@ -1131,6 +1159,8 @@ def lookup( BLOCK_R=16, GRID=grid, SORTED=sort, + PACKED=out.dtype == torch.uint8, + OUT_STRIDE=out.stride(0), ) def forward(self, indices: torch.Tensor) -> torch.Tensor: @@ -1311,6 +1341,175 @@ def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor: return rows +@triton.jit(do_not_specialize=["num_tokens"]) +def _engram_unpack_fp8_kernel( + packed, + values, + scales, + num_tokens, + values_stride, + scales_stride, + HEADS: tl.constexpr, + LOCAL_HEADS: tl.constexpr, + LAYERS: tl.constexpr, + RANK: tl.constexpr, + DIM: tl.constexpr, + BLOCK: tl.constexpr, + PEER: tl.constexpr, +): + """Split `[rank][token][layer][local head]` packed rows into MXFP8 values and + F8_128x4 scales, one layer per grid column; `packed` holds per-rank + pointers if PEER. Scales past `num_tokens` zero-pad the 128-row tile. + """ + row = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + layer = tl.program_id(1).to(tl.int64) + token, head = row // HEADS, row % HEADS + if PEER: + src = tl.load(packed + head // LOCAL_HEADS) + src = tl.multiple_of(src, DIM // 32).to(tl.pointer_type(tl.uint8)) + source = RANK * num_tokens + token + else: + src = packed + source = head // LOCAL_HEADS * num_tokens + token + source = (source * LAYERS + layer) * LOCAL_HEADS + head % LOCAL_HEADS + # Packed rows are only scale-width aligned. + src += tl.multiple_of(source * (DIM + DIM // 32), DIM // 32) + valid = token < num_tokens + cols = tl.arange(0, DIM) + data = tl.load(src[:, None] + cols[None, :], mask=valid[:, None], other=0) + values += layer * values_stride + row * DIM + tl.store(values[:, None] + cols[None, :], data, mask=valid[:, None]) + scale_cols = tl.arange(0, DIM // 32) + scale = tl.load( + src[:, None] + DIM + scale_cols[None, :], mask=valid[:, None], other=0 + ) + group = head[:, None] * (DIM // 32) + scale_cols[None, :] + offset = mxfp8_scale_offset(token[:, None], group, HEADS * DIM // 32) + tl.store(scales + layer * scales_stride + offset, scale) + + +class EngramBatch(nn.Module): + """Stage every Engram layer's WKV input as MXFP8 with one exchange. + + A module so the model owns the peer exchange's pointer buffers. + """ + + def __init__(self, layers: list["Engram"], max_tokens: int) -> None: + super().__init__() + first = layers[0].layer_hash_index + hash_indices = [layer.layer_hash_index for layer in layers] + assert hash_indices == list(range(first, first + len(layers))) + self.layers = layers + table = layers[0].embed_tokens + device = layers[0].staged_rows.device + self.packed_rows = torch.empty( + max_tokens * table.dp_size, + len(layers), + table.part_n_hash_cols, + table.dim + table.dim // 32, + dtype=torch.uint8, + device=device, + ) + width = table.n_hash_cols * table.dim + self.fp8_values = torch.empty( + len(layers), max_tokens, width, dtype=torch.uint8, device=device + ) + self.fp8_scales = torch.empty( + len(layers), + mxfp8_scale_bytes(max_tokens, width), + dtype=torch.uint8, + device=device, + ) + max_peer_slot = min(max_tokens, ENGRAM_A2A_MIN_SLOT - 1) + self.peer_exchange = None + if table.cpu_offload: + from vllm.models.deepseek_v41.nvidia.ops.engram_peer import ( + EngramPeerExchange, + ) + + self.peer_exchange = EngramPeerExchange.create(layers, max_peer_slot) + for layer in layers: + # Owned by the model; keep it out of each layer's submodules. + object.__setattr__(layer, "_batch", self) + # The layers read this batch's staging instead of their BF16 rows. + layer.staged_rows = layer.staged_rows.new_empty( + (0, *layer.staged_rows.shape[1:]) + ) + + @classmethod + def create(cls, layers: list["Engram"], max_tokens: int) -> "EngramBatch | None": + """Batch DP-sharded tables whose WKV takes FlashInfer MXFP8 input.""" + if not layers: + return None + table = layers[0].embed_tokens + if ( + table.dp_size == 1 + or table.tp_size > 1 + or table.dim % 128 + or get_current_vllm_config().lora_config is not None + or not can_fuse_query_quant([layer.wkv for layer in layers]) + ): + return None + return cls(layers, max_tokens) + + def prepare_embeddings(self, hashes: torch.Tensor) -> None: + """Stage every layer's WKV input before the decoder layers. + + DP-sharded tables decide on the group's slot, so every rank agrees. + """ + slot = engram_gathered_num_tokens() + peer = self.peer_exchange is not None and slot < ENGRAM_A2A_MIN_SLOT + self._stage(hashes if peer else gather_engram_hashes(hashes), slot, peer) + + @eager_break_during_capture + def _stage(self, hashes: torch.Tensor, slot: int, peer: bool) -> None: + if slot == 0: + return + if peer: + assert self.peer_exchange is not None + rows = self.peer_exchange.exchange(hashes, slot) + rank = self.peer_exchange.dp_rank + else: + rows, rank = self.packed_rows[: hashes.shape[0]], 0 + for index, layer in enumerate(self.layers): + layer.embed_tokens.lookup( + hashes[:, layer.layer_hash_index], rows[:, index] + ) + group = get_engram_dp_group() + assert group is not None + sent, rows = rows, torch.empty_like(rows) + torch.distributed.all_to_all_single(rows, sent, group=group.device_group) + table = self.layers[0].embed_tokens + heads = table.n_hash_cols + grid = (triton.cdiv(slot, 128) * 128 * heads // 16, len(self.layers)) + _engram_unpack_fp8_kernel[grid]( + rows, + self.fp8_values, + self.fp8_scales, + slot, + self.fp8_values.stride(0), + self.fp8_scales.stride(0), + HEADS=heads, + LOCAL_HEADS=table.part_n_hash_cols, + LAYERS=len(self.layers), + RANK=rank, + DIM=table.dim, + BLOCK=16, + PEER=peer, + ) + + def wkv_input(self, layer: "Engram", num_tokens: int) -> QuantizedActivation: + index = layer.layer_hash_index - self.layers[0].layer_hash_index + values = self.fp8_values[index, :num_tokens].view(torch.float8_e4m3fn) + return QuantizedActivation( + values, + self.fp8_scales[index, : mxfp8_scale_bytes(num_tokens, values.shape[1])], + torch.bfloat16, + values.shape, + kMxfp8Dynamic, + ) + + class Engram(nn.Module): """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. @@ -1323,6 +1522,7 @@ class Engram(nn.Module): _prefetch_stream: torch.cuda.Stream | None = None _prefetch_done: torch.cuda.Event | None = None + _batch: EngramBatch | None = None def __init__( self, @@ -1482,7 +1682,10 @@ def forward( """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all tokens, pre sequence-parallel shard); token_mask: [T], False shuts the gate so those positions pass through untouched.""" - kv = self.wkv(self.embed(hash_ids).flatten(-2)) + if self._batch is not None: + kv = self.wkv(self._batch.wkv_input(self, hash_ids.shape[0])) + else: + kv = self.wkv(self.embed(hash_ids).flatten(-2)) num_kv_tokens = hash_ids.shape[0] assert token_mask is None or token_mask.shape == (num_kv_tokens,) if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py index 6094d1980cea..97714e29f151 100644 --- a/vllm/models/deepseek_v41/common/ops/query_quant.py +++ b/vllm/models/deepseek_v41/common/ops/query_quant.py @@ -21,6 +21,18 @@ def _norm_row(x, weight, row, stride, eps, SIZE: tl.constexpr, BLOCK: tl.constex return values * rrms * w +@triton.jit +def mxfp8_scale_offset(row, group, GROUPS: tl.constexpr): + """F8_128x4: [row/128, group/4, row%32, row%128/32, group%4].""" + offset = (row // 128 * triton.cdiv(GROUPS, 4) + group // 4) * 512 + return offset + row % 32 * 16 + row % 128 // 32 * 4 + group % 4 + + +def mxfp8_scale_bytes(tokens: int, width: int) -> int: + """Bytes of F8_128x4 scales for `tokens` rows of `width` MXFP8 values.""" + return triton.cdiv(tokens, 128) * 128 * triton.cdiv(width // 32, 4) * 4 + + @triton.jit(do_not_specialize=["num_tokens"]) def _q_kv_norm_quant_kernel( q, @@ -68,14 +80,7 @@ def _q_kv_norm_quant_kernel( sf = tl.full((BLOCK // 32,), 0, tl.uint32) padded_groups: tl.constexpr = triton.cdiv(Q_SIZE // 32, 4) * 4 sf = tl.where(groups < Q_SIZE // 32, sf, 0) - # F8_128x4: [row/128, group/4, row%32, row%128/32, group%4]. - offsets = ( - row // 128 * (128 * padded_groups) - + groups // 4 * 512 - + row % 32 * 16 - + row % 128 // 32 * 4 - + groups % 4 - ) + offsets = mxfp8_scale_offset(row, groups, Q_SIZE // 32) tl.store(scales + offsets, sf, groups < padded_groups) elif row < num_tokens: y = _norm_row(kv, kvw, row, kv_stride, eps, KV_SIZE, BLOCK) @@ -99,9 +104,8 @@ def fused_q_kv_rmsnorm_quant( qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device) kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device) padded_tokens = triton.cdiv(tokens, 128) * 128 - padded_groups = triton.cdiv(q_size // 32, 4) * 4 scales = torch.empty( - padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device + mxfp8_scale_bytes(tokens, q_size), dtype=torch.uint8, device=qr.device ) if tokens: block = triton.next_power_of_2(max(q_size, kv_size)) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 18f8f27843ab..8ec0155f736a 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -100,6 +100,7 @@ from ..common.engram import ( Engram, + EngramBatch, EngramLayout, NgramHashState, can_share_engram_tables, @@ -777,6 +778,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) + self.engram_batch = EngramBatch.create( + [ + layer.engram + for layer in islice(self.layers, self.start_layer, self.end_layer) + if getattr(layer, "engram", None) is not None + ], + vllm_config.scheduler_config.max_num_batched_tokens, + ) if self.fuse_mhc_all_reduce: # A MoE's top-k finalize folds into the next layer's first mHC # boundary, which the last local layer lacks and an engram layer @@ -971,15 +980,18 @@ def forward( if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. # One gather feeds every layer sharing the DP-split table. - gathered_hashes = gather_engram_hashes( - engram_hashes, dp_shared_memory=self.engram_dp_shared_memory - ) - for layer in islice(self.layers, self.start_layer, self.end_layer): - engram = getattr(layer, "engram", None) - if engram is not None: - engram.prepare_embeddings( - gathered_hashes[:, engram.layer_hash_index] - ) + if self.engram_batch is not None: + self.engram_batch.prepare_embeddings(engram_hashes) + else: + gathered_hashes = gather_engram_hashes( + engram_hashes, dp_shared_memory=self.engram_dp_shared_memory + ) + for layer in islice(self.layers, self.start_layer, self.end_layer): + engram = getattr(layer, "engram", None) + if engram is not None: + engram.prepare_embeddings( + gathered_hashes[:, engram.layer_hash_index] + ) full_num_tokens = positions.shape[0] if self.use_sequence_parallel: diff --git a/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py b/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py new file mode 100644 index 000000000000..4f9423bb1e8b --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Symmetric-memory Engram exchange for slots below ENGRAM_A2A_MIN_SLOT.""" + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +from torch import nn + +from vllm.distributed import get_engram_dp_group +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.common.engram import Engram + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _peer_lookup_kernel( + metadata, + id_ptrs, + output, + num_tokens, + HEADS: tl.constexpr, + LAYERS: tl.constexpr, + DP: tl.constexpr, + RANK: tl.constexpr, + DIM: tl.constexpr, + BLOCK: tl.constexpr, +): + """Look up this rank's heads for every peer's tokens into packed rows.""" + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + local_heads: tl.constexpr = HEADS // DP + valid = row < num_tokens * LAYERS * HEADS + token = row // (LAYERS * local_heads) + layer, head = row // local_heads % LAYERS, row % local_heads + ids = tl.load(id_ptrs + token // num_tokens, valid, other=0).to( + tl.pointer_type(tl.int32) + ) + index = tl.load( + ids + (token % num_tokens * LAYERS + layer) * HEADS + RANK * local_heads + head, + valid, + other=-1, + ).to(tl.int64) + # Per-layer metadata: weight pointer, scale pointer, vocab start, vocab end. + table = metadata + layer * 4 + start, end = tl.load(table + 2), tl.load(table + 3) + owned = valid & (index >= start) & (index < end) + local = tl.where(owned, index - start, 0) + weight = tl.multiple_of(tl.load(table), DIM).to(tl.pointer_type(tl.uint8)) + scales = tl.load(table + 1) + scales = tl.multiple_of(scales, DIM // 32).to(tl.pointer_type(tl.uint8)) + col = tl.arange(0, DIM) + value = tl.load( + weight[:, None] + local[:, None] * DIM + col[None, :], owned[:, None], other=0 + ) + dst = row.to(tl.int64) * (DIM + DIM // 32) + tl.store(output + dst[:, None] + col[None, :], value, valid[:, None]) + scale_cols = tl.arange(0, DIM // 32) + scale = tl.load( + scales[:, None] + local[:, None] * (DIM // 32) + scale_cols[None, :], + owned[:, None], + other=0, + ) + tl.store(output + dst[:, None] + DIM + scale_cols[None, :], scale, valid[:, None]) + + +class EngramPeerExchange(nn.Module): + """Two phase barriers fence reuse of each other's buffers across forwards. + + Pointer tables are buffers so level-2 sleep restores them on wake-up. + """ + + def __init__(self, layers: list["Engram"], max_tokens: int): + import torch.distributed._symmetric_memory as symm_mem + + super().__init__() + group = get_engram_dp_group() + assert group is not None + self.dp_size, self.dp_rank = group.world_size, group.rank_in_group + device = layers[0].staged_rows.device + self.num_layers = len(layers) + self.hash_start = layers[0].layer_hash_index + table = layers[0].embed_tokens + self.heads, self.dim = table.n_hash_cols, table.dim + self.ids = symm_mem.empty( + (max_tokens, self.num_layers, self.heads), + dtype=torch.int32, + device=device, + ) + self.rows = symm_mem.empty( + ( + max_tokens * self.dp_size, + self.num_layers, + self.heads // self.dp_size, + self.dim + self.dim // 32, + ), + dtype=torch.uint8, + device=device, + ) + self.id_handle = symm_mem.rendezvous(self.ids, group.device_group) + self.row_handle = symm_mem.rendezvous(self.rows, group.device_group) + metadata = [] + for layer in layers: + table = layer.embed_tokens + weight, scale = table._storage() + metadata.append( + ( + weight.data_ptr(), + scale.data_ptr(), + table.vocab_start_idx, + table.vocab_end_idx, + ) + ) + for name, data in ( + ("metadata", metadata), + ("id_ptrs", self.id_handle.buffer_ptrs), + ("row_ptrs", self.row_handle.buffer_ptrs), + ): + self.register_buffer( + name, + torch.tensor(data, dtype=torch.int64, device=device), + persistent=False, + ) + + @classmethod + def create( + cls, layers: list["Engram"], max_tokens: int + ) -> "EngramPeerExchange | None": + # The Engram DP group is node-local and every rank agrees on these. + group = get_engram_dp_group() + assert group is not None + if not ( + current_platform.is_device_capability_family(100) + and group.world_size in (2, 4, 8) + and layers[0].embed_tokens.n_hash_cols % group.world_size == 0 + ): + return None + physical_id = current_platform.visible_device_id_to_physical_device_id( + torch.accelerator.current_device_index() + ) + physical_ids = [None] * group.world_size + dist.all_gather_object(physical_ids, physical_id, group=group.cpu_group) + if not current_platform.is_fully_connected(physical_ids): + return None + return cls(layers, max_tokens) + + def exchange(self, hashes: torch.Tensor, slot: int) -> torch.Tensor: + """Look up this rank's heads for every peer; return per-rank row pointers.""" + tokens = hashes.shape[0] + assert 0 < slot <= self.ids.shape[0] and tokens <= slot + hash_end = self.hash_start + self.num_layers + self.ids[:tokens].copy_(hashes[:, self.hash_start : hash_end]) + if tokens < slot: + self.ids[tokens:slot].fill_(-1) + # The preceding rows barrier completed every peer's hash reads. This + # barrier completes their previous output reads before rows are reused. + self.id_handle.barrier() + _peer_lookup_kernel[(triton.cdiv(slot * self.num_layers * self.heads, 16),)]( + self.metadata, + self.id_ptrs, + self.rows, + slot, + HEADS=self.heads, + LAYERS=self.num_layers, + DP=self.dp_size, + RANK=self.dp_rank, + DIM=self.dim, + BLOCK=16, + ) + self.row_handle.barrier() + return self.row_ptrs From f43e4c39e30120abf6250a209a42e5cd4849ca23 Mon Sep 17 00:00:00 2001 From: Denis Ermakov Date: Fri, 9 Oct 2026 09:49:47 +0200 Subject: [PATCH 268/952] [Structured Output] Allow propertyNames with additionalProperties on xgrammar (#57594) Signed-off-by: errmakov Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar.py | 136 ++++++++++++------ vllm/v1/structured_output/backend_xgrammar.py | 101 +++++++------ 2 files changed, 152 insertions(+), 85 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar.py b/tests/v1/structured_output/test_backend_xgrammar.py index bc29112077cc..9024b614dd77 100644 --- a/tests/v1/structured_output/test_backend_xgrammar.py +++ b/tests/v1/structured_output/test_backend_xgrammar.py @@ -31,21 +31,6 @@ def grammar_accepts(schema: dict, text: str) -> bool: # ================================================ -@pytest.fixture -def unsupported_string_schemas(): - return [ - {"type": "string", "format": "non_existing_format"}, - # TODO(arpera): - # pattern/format is compiled but length bounds are silently dropped, - # so the combination must be rejected instead of producing quietly - # wrong output - # https://github.com/mlc-ai/xgrammar/issues/749 - {"type": "string", "pattern": "^a+$", "maxLength": 2}, - {"type": "string", "pattern": "^a+$", "minLength": 3}, - {"type": "string", "format": "email", "maxLength": 10}, - ] - - @pytest.fixture def unsupported_multipleOf_schemas(): return [ @@ -59,6 +44,8 @@ def unsupported_multipleOf_schemas(): @pytest.fixture def unsupported_array_schemas(): return [ + # array + some constraints is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/968 {"type": "array", "uniqueItems": True}, {"type": "array", "contains": {"type": "string"}}, {"type": "array", "minContains": 1}, @@ -67,74 +54,113 @@ def unsupported_array_schemas(): @pytest.fixture -def unsupported_property_names_combinations(): +def unsupported_string_schemas(): + return [ + # ======================================================== + # string + format is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/967 + {"type": "string", "format": "non_existing_format"}, + # ======================================================== + # + # ======================================================== + # string + format/pattern + length constraint is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/966 + # + # pattern/format is compiled but length bounds are silently dropped, + # so the combination must be rejected instead of producing quietly + # wrong output + {"type": "string", "pattern": "^a+$", "maxLength": 2}, + {"type": "string", "pattern": "^a+$", "minLength": 3}, + {"type": "string", "format": "email", "maxLength": 10}, + # ======================================================== + ] + + +@pytest.fixture +def unsupported_propertyNames_combinations(): return [ - # TODO(arpera): - # this case is not covered by xgrammar, so we should report this bug to xgrammar + # ======================================================== + # propertyNames + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/959 + # # xgrammar drops propertyNames whenever patternProperties is present, # so "grade_12" matches the pattern and escapes the name constraint { "type": "object", - "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, "propertyNames": {"pattern": "^grade_[0-9]$"}, # does NOT match "grade_12" + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, }, - # propertyNames makes xgrammar discard the sibling additionalProperties - # value schema: https://github.com/mlc-ai/xgrammar/issues/826 - { - "type": "object", - "propertyNames": {"pattern": "^[a-z]+$"}, - "additionalProperties": {"type": "integer"}, - }, - # propertyNames is a string schema that conventionally omits "type", so - # it escapes the string check while xgrammar drops its length bound all - # the same: https://github.com/mlc-ai/xgrammar/issues/749 + # ======================================================== + # + # ======================================================== + # propertyNames + maxLength is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/960 { "type": "object", "propertyNames": {"pattern": "^a+$", "maxLength": 2}, }, - # xgrammar emits named `properties` separately from `propertyNames` and - # only applies `propertyNames` in its additional-properties branch, so a - # key declared in `properties` escapes the name constraint regardless of - # what `propertyNames` contains. + # ======================================================== + # + # ======================================================== + # propertyNames + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/958 { "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, "properties": {"Bad": {"type": "integer"}}, "required": ["Bad"], - "propertyNames": {"pattern": "^[a-z]+$"}, }, { "type": "object", - "properties": {"Bad": {"type": "integer"}}, "propertyNames": {"enum": ["good"]}, + "properties": {"Bad": {"type": "integer"}}, }, + # ======================================================== + # + # ======================================================== + # propertyNames + unevaluatedProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/961 + # # xgrammar's propertyNames-only branch falls back to an unconstrained # value type, so it drops unevaluatedProperties whether it restricts # values (a schema) or forbids them (false). { "type": "object", "propertyNames": {"pattern": "^[a-z]+$"}, - "unevaluatedProperties": {"type": "integer"}, + "unevaluatedProperties": False, }, + # also keep in mind that unevaluatedProperties can be a schema { "type": "object", "propertyNames": {"pattern": "^[a-z]+$"}, - "unevaluatedProperties": False, + "unevaluatedProperties": {"type": "integer"}, }, + # ======================================================== ] @pytest.fixture -def unsupported_pattern_properties_combinations(): +def unsupported_patternProperties_combinations(): return [ + # ======================================================== + # patternProperties + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/964 + # # JSON Schema requires a property matched by both `properties` and # `patternProperties` to satisfy both (conjunction), but xgrammar # compiles them as alternative branches, so satisfying either one is # enough. { "type": "object", - "properties": {"x": {"type": "string"}}, "patternProperties": {"^x$": {"type": "integer"}}, + "properties": {"x": {"type": "string"}}, }, + # ======================================================== + # + # ======================================================== + # patternProperties + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/965 + # # The same alternative-branches problem applies to overlapping # patternProperties patterns: a key matching both patterns only has to # satisfy one of them. @@ -145,6 +171,7 @@ def unsupported_pattern_properties_combinations(): "^[a-z]*z$": {"type": "integer"}, }, }, + # ======================================================== ] @@ -273,6 +300,24 @@ def property_names_schema(): return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} +@pytest.fixture +def property_names_with_additional_properties_schema(): + """This combination propertyNames + additionalProperties + was unsupported some time ago by xgrammar. + So, this test case was in unsupported category before. + But this xgrammar's PR + https://github.com/mlc-ai/xgrammar/pull/836 + added support of propertyNames + additionalProperties + This fix was released in xgrammar v0.2.7 + which vLLM pins in PR #57272 + """ + return { + "type": "object", + "propertyNames": {"pattern": "^[a-z_]+$"}, + "additionalProperties": {"type": "integer"}, + } + + @pytest.fixture def supported_allof_anyof_and_oneof(): """Xgrammar DO support some variants of anyOf, allOf, oneOf @@ -334,8 +379,8 @@ class TestHasXGrammarUnsupportedJsonFeatures: "unsupported_string_schemas", "unsupported_multipleOf_schemas", "unsupported_array_schemas", - "unsupported_property_names_combinations", - "unsupported_pattern_properties_combinations", + "unsupported_propertyNames_combinations", + "unsupported_patternProperties_combinations", # vLLM issue #56556 "unsupported_multibranch_allof", "unsupported_vllm_issue_56556_schema", @@ -354,6 +399,7 @@ def test_unsupported_json_features(self, schema_type, request): "supported_frankenstein_schema", "pattern_properties_schema", "property_names_schema", + "property_names_with_additional_properties_schema", # Additional test cases implemented during work on # vLLM issue #56556 "supported_allof_anyof_and_oneof", @@ -421,6 +467,14 @@ def test_property_names_constrains_keys(self, property_names_schema): assert not grammar_accepts(property_names_schema, '{"Score": 5}') assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + def test_property_names_keeps_additional_properties_value_schema( + self, property_names_with_additional_properties_schema + ): + schema = property_names_with_additional_properties_schema + assert grammar_accepts(schema, '{"score": 5}') + assert not grammar_accepts(schema, '{"Score": 5}') + assert not grammar_accepts(schema, '{"score": "seven"}') + def test_pattern_properties_constrains_keys_and_values( self, pattern_properties_schema ): diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index b675a9a3f193..0199c5f75d6f 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -296,18 +296,29 @@ def check_object(obj: dict[str, Any]) -> bool: schema_types = _schema_types(obj) - # Check for numeric ranges + # integer/number + multipleOf is unsupported by xgrammar + # This is known behavior and xgrammar emits warning in logs: + # [21:18:08] /project/cpp/json_schema_converter.cc:1053: + # Warning: multipleOf is not supported for type:number; ignoring multipleOf + # This warning was added in PR + # https://github.com/mlc-ai/xgrammar/pull/670 + # So, no need to track progress on this if (schema_types & {"integer", "number"}) and ("multipleOf" in obj): return True - # Check for array unsupported keywords + # array + some constraints is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/968 if "array" in schema_types and any( key in obj for key in ("uniqueItems", "contains", "minContains", "maxContains") ): return True - # Unsupported keywords for strings + # string + format with unsupported keywords + # is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/967 + # See tests on this in test_backend_xgrammar.py + # unsupported_string_schemas if ( "string" in schema_types and "format" in obj @@ -315,48 +326,51 @@ def check_object(obj: dict[str, Any]) -> bool: ): return True - # A string mixing a generative constraint (pattern or format) with - # explicit length bounds. xgrammar compiles the pattern/format side - # and silently drops minLength/maxLength from the grammar, so output - # can violate the bound without any error surfacing. Verified against - # the compiled EBNF: pattern/format grammars come out byte-identical - # with and without the length keywords, while maxLength alone lowers - # to {0, N} correctly. + # string + format/pattern + length constraint is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/966 + # See tests on this in test_backend_xgrammar.py + # unsupported_string_schemas if "string" in schema_types and _has_pattern_and_length_bounds(obj): return True - # propertyNames validates names, so it is a string schema even when it - # omits "type", which is the form that escapes the check above. - if ( - "object" in schema_types - and isinstance(obj.get("propertyNames"), dict) - and _has_pattern_and_length_bounds(obj["propertyNames"]) - ): - return True - - # FIXME: propertyNames conflicts with properties/patternProperties/ - # additionalProperties/unevaluatedProperties under xgrammar. - # https://github.com/mlc-ai/xgrammar/issues/826 - if ( - "object" in schema_types - and "propertyNames" in obj - and ( - "properties" in obj - or "patternProperties" in obj - or isinstance(obj.get("additionalProperties"), dict) - or obj.get("unevaluatedProperties", True) is not True - ) - ): - return True - - # FIXME: multiple patternProperties, or patternProperties alongside - # properties, conflict under xgrammar. - if ( - "object" in schema_types - and isinstance(obj.get("patternProperties"), dict) - and ("properties" in obj or len(obj["patternProperties"]) > 1) - ): - return True + # propertyNames is not supported in pair with some constraints + # in xgrammar + # See tests on this in test_backend_xgrammar.py + # unsupported_propertyNames_combinations + if "object" in schema_types and "propertyNames" in obj: + # propertyNames + maxLength is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/960 + property_names = obj.get("propertyNames") + if isinstance(property_names, dict) and _has_pattern_and_length_bounds( + property_names + ): + return True + # propertyNames + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/959 + if "patternProperties" in obj: + return True + # propertyNames + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/958 + if "properties" in obj: + return True + # propertyNames + unevaluatedProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/961 + if obj.get("unevaluatedProperties", True) is not True: + return True + + # patternProperties is not supported in pair with some constraints + # in xgrammar + # See tests on this in test_backend_xgrammar.py + # unsupported_patternProperties_combinations + if "object" in schema_types and isinstance(obj.get("patternProperties"), dict): + # patternProperties + properties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/964 + if "properties" in obj: + return True + # patternProperties + patternProperties is unsupported by xgrammar + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/965 + if len(obj["patternProperties"]) > 1: + return True # Note(arpera): # Xgrammar lacks support of multi-branch allOf @@ -370,8 +384,7 @@ def check_object(obj: dict[str, Any]) -> bool: # would accept any kind of json, such as # "maybe", "", 42, {}, [], {"a": 1}, etc. # which is NOT what is expected. - # Reported this issue to xgrammar team to track progress on resolving: - # https://github.com/mlc-ai/xgrammar/issues/937 + # Tracking issue: https://github.com/mlc-ai/xgrammar/issues/937 allof = obj.get("allOf") if isinstance(allof, list) and len(allof) >= 2: return True From 8a03ccf7b3a26dc0c6ccf2a3ba12113475824db3 Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Fri, 9 Oct 2026 11:14:13 +0300 Subject: [PATCH 269/952] [Frontend] Add return_mm_kwargs to render requests (#60728) Signed-off-by: Maroon Ayoub Co-authored-by: Claude --- docs/serving/online_serving/renderer.md | 7 +++++++ .../token_in_token_out/test_mm_serde.py | 14 ++++++++++++++ vllm/entrypoints/anthropic/protocol.py | 11 +++++++++++ .../openai/chat_completion/protocol.py | 12 ++++++++++++ vllm/entrypoints/openai/completion/protocol.py | 12 ++++++++++++ vllm/entrypoints/openai/responses/protocol.py | 11 +++++++++++ vllm/entrypoints/scale_out/render/serving.py | 16 +++++++++++++--- .../scale_out/token_in_token_out/mm_features.py | 16 ++++++++++------ 8 files changed, 90 insertions(+), 9 deletions(-) diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 79d0ade85b4f..93ecf34f4f36 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -153,6 +153,13 @@ should split those fields: - Legacy clients that ignore `mm_metadata` and keep sending `kwargs_data` continue to work. +Callers that need only the token layout and item hashes, such as cache-aware +routers, can set `"return_mm_kwargs": false` on the render request. The response +then keeps `mm_hashes` and `mm_placeholders` and sets `kwargs_data` and +`mm_metadata` to null, so the processed tensors are neither serialized nor sent. +Do not forward such a response to `/inference/v1/generate`, which reads a null +`kwargs_data` as every item being cached. + ## Example The example below shows how a disaggregated encode / prefill coordinator can diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index 5cd60a959fb3..62223f1d7cf2 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -168,6 +168,20 @@ def test_render_extracts_metadata_fields_separately(): assert torch.equal(full_data["pixel_values"].data, pixel_values.data) +def test_extract_without_kwargs_keeps_layout(): + """Layout-only callers get hashes and placeholders, no serialized tensors.""" + engine_input, _pixel_values, _image_grid_thw = _image_engine_input() + + full = extract_mm_features(engine_input) + layout = extract_mm_features(engine_input, include_mm_kwargs=False) + + assert full is not None and layout is not None + assert layout.kwargs_data is None + assert layout.mm_metadata is None + assert layout.mm_hashes == full.mm_hashes + assert layout.mm_placeholders == full.mm_placeholders + + def test_extract_includes_declared_placeholder_metadata_fields(): """EC placeholder metadata is kept even when keep_on_cpu is unset.""" engine_input, _pixel_values, image_grid_thw = _image_engine_input( diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 5997d2e3d2e6..b2d29e0b34fe 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -229,6 +229,17 @@ class AnthropicMessagesRequest(BaseModel): top_p: float | None = None # vLLM-specific fields that are not in Anthropic spec + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 15905b8a3da6..dee8f283cf68 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -458,6 +458,18 @@ class ChatCompletionRequest(OpenAIBaseModel): ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) + cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 78722b45ec7b..f26dea2e1f9e 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -197,6 +197,18 @@ class CompletionRequest(OpenAIBaseModel): ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) + cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index d1a4ea4f6c21..406ab2d0818f 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -330,6 +330,17 @@ class ResponsesRequest(OpenAIBaseModel): "if the served model does not use priority scheduling." ), ) + return_mm_kwargs: bool = Field( + default=True, + description=( + "If false, the render response's `features` set `kwargs_data` " + "and `mm_metadata` to null, for callers that need only the token " + "layout and item hashes, such as cache-aware routers. Do not send " + "such a response to `/inference/v1/generate`, which reads a null " + "`kwargs_data` as every item being cached. Only supported on the " + "render endpoints; ignored on regular generation endpoints." + ), + ) cache_salt: str | None = Field( default=None, min_length=1, diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index bddbe90f286f..84771c5f7390 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -144,7 +144,9 @@ async def render_chat_request( return GenerateRequest( request_id=request_id, token_ids=token_ids, - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, reasoning_ended=reasoning_ended, @@ -169,6 +171,7 @@ async def render_messages_request( chat_req = AnthropicServingMessages.to_chat_completion_request( request, merge_inline_system=self._merge_inline_system ) + chat_req.return_mm_kwargs = request.return_mm_kwargs return await self.render_chat_request(chat_req) async def render_completion_request( @@ -217,7 +220,9 @@ async def render_completion_request( GenerateRequest( request_id=request_id, token_ids=token_ids, - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, stream=bool(request.stream), @@ -277,7 +282,9 @@ async def render_responses_request( return GenerateRequest( request_id=request.request_id, token_ids=list(token_ids), - features=self._extract_mm_features(engine_input), + features=self._extract_mm_features( + engine_input, include_mm_kwargs=request.return_mm_kwargs + ), sampling_params=params, model=request.model, reasoning_parser_kwargs=self._reasoning_parser_kwargs(request), @@ -328,8 +335,11 @@ def _placeholder_metadata_fields(self, modality: str) -> set[str]: def _extract_mm_features( self, engine_input: EngineInput, + *, + include_mm_kwargs: bool = True, ) -> MultiModalFeatures | None: return extract_mm_features( engine_input, metadata_fields_for=self._placeholder_metadata_fields, + include_mm_kwargs=include_mm_kwargs, ) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py index 62d006783a2e..a0cb0d4b5eb6 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/mm_features.py @@ -139,13 +139,15 @@ def extract_mm_features( engine_input: EngineInput, *, metadata_fields_for: Callable[[str], Collection[str]] | None = None, + include_mm_kwargs: bool = True, ) -> MultiModalFeatures | None: """Extract multimodal features from a rendered engine prompt. Returns ``None`` for text-only prompts. ``mm_metadata`` keeps the intersection of processed kwargs that prefill needs after EC transfer: fields declared as embedding metadata, plus fields marked - ``keep_on_cpu`` (for example M-RoPE grid dims). + ``keep_on_cpu`` (for example M-RoPE grid dims). With + ``include_mm_kwargs=False`` neither is serialized. """ if engine_input.get("type") != "multimodal": return None @@ -157,11 +159,13 @@ def extract_mm_features( kwargs_data: dict[str, list[str | None]] | None = None mm_metadata: dict[str, list[str | None]] | None = None - if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"): - kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata( - raw_mm_kwargs, - metadata_fields_for=metadata_fields_for, - ) + if include_mm_kwargs: + raw_mm_kwargs = mm_engine_input.get("mm_kwargs") + if raw_mm_kwargs: + kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata( + raw_mm_kwargs, + metadata_fields_for=metadata_fields_for, + ) return MultiModalFeatures( mm_hashes=mm_hashes, From 560fae64cdaf84468590fd362916969f18bb4d9b Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Fri, 9 Oct 2026 11:16:20 +0300 Subject: [PATCH 270/952] [ROCm][Bugfix] Pass per-sequence context lengths to the AITER paged MQA-logits kernels (#59420) Signed-off-by: Sriram Kumar --- .../test_rocm_paged_mqa_context_lens.py | 32 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 16 +++++++--- 2 files changed, 44 insertions(+), 4 deletions(-) create mode 100644 tests/v1/attention/test_rocm_paged_mqa_context_lens.py diff --git a/tests/v1/attention/test_rocm_paged_mqa_context_lens.py b/tests/v1/attention/test_rocm_paged_mqa_context_lens.py new file mode 100644 index 000000000000..3122cc915a93 --- /dev/null +++ b/tests/v1/attention/test_rocm_paged_mqa_context_lens.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.platforms import current_platform + +if not current_platform.is_rocm() or not is_aiter_found_and_supported(): + pytest.skip("Requires ROCm with AITER.", allow_module_level=True) + +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _per_sequence_context_lens + + +def test_per_row_table_uses_last_column(): + seq_lens = torch.tensor([8192, 1500, 3000, 700], dtype=torch.int32) + next_n = 3 + offsets = torch.arange(next_n, dtype=torch.int32) - next_n + 1 + per_row = seq_lens.unsqueeze(1) + offsets + + out = _per_sequence_context_lens(per_row) + + assert out.shape == (4,) + assert out.is_contiguous() + assert torch.equal(out, seq_lens) + + +@pytest.mark.parametrize("shape", [(4,), (4, 1)]) +def test_per_sequence_input_is_unchanged(shape): + context_lens = torch.arange(1, 5, dtype=torch.int32).reshape(shape) + assert _per_sequence_context_lens(context_lens) is context_lens diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 90a73c101c0d..6a93d2151127 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -707,6 +707,14 @@ def paged_mqa_logits_module(): return None +def _per_sequence_context_lens(context_lens: torch.Tensor) -> torch.Tensor: + # AITER takes one length per sequence: the last column of a (B, next_n) table. + # TODO: remove once AITER accepts (B, next_n) context_lens (ROCm/aiter#6153). + if context_lens.dim() == 2 and context_lens.shape[1] > 1: + return context_lens[:, -1].contiguous() + return context_lens + + def rocm_fp8_paged_mqa_logits( q_fp8: torch.Tensor, kv_cache_fp8: torch.Tensor, @@ -726,8 +734,8 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. - context_lens: Tensor of shape [B], dtype int32; effective context length - for each batch element. + context_lens: Tensor of shape [B] or [B, next_n], dtype int32; effective + context length for each batch element, or for each of its Q rows. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical block indices to physical blocks in the paged cache. schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; @@ -779,7 +787,7 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8, weights, out_logits, - context_lens, + _per_sequence_context_lens(context_lens), block_tables, max_model_len, ChunkK=256, @@ -801,7 +809,7 @@ def rocm_fp8_paged_mqa_logits( kv_cache_fp8, weights, out_qk, - context_lens, + _per_sequence_context_lens(context_lens), block_tables, max_model_len, ChunkQ=heads, From bacbbe187885db62859f5a4a1443f8ec3006e987 Mon Sep 17 00:00:00 2001 From: The Anh Nguyen Date: Fri, 9 Oct 2026 15:22:18 +0700 Subject: [PATCH 271/952] [Bugfix][NIXL] Guard record_transfer against None telemetry (#57734) Signed-off-by: The Anh Nguyen --- vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index 26c4275efbe1..645525bc6e2b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -44,8 +44,10 @@ def reset(self): "num_notifications_after_expiry": [], } - def record_transfer(self, res: "nixlXferTelemetry"): + def record_transfer(self, res: "nixlXferTelemetry | None"): # Keep metrics units consistent with rest of the code: time us->s + if res is None: + return self.data["transfer_duration"].append(res.xferDuration / 1e6) self.data["post_duration"].append(res.postDuration / 1e6) self.data["bytes_transferred"].append(res.totalBytes) From a6f43e57d6bdcbe1a42cc07fd5298ae33052180d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Fri, 9 Oct 2026 16:29:55 +0800 Subject: [PATCH 272/952] [ROCm] Enable the cuMem CUDA-graph pool offload on ROCm (#59523) Signed-off-by: aoshen02 Signed-off-by: indianspeedster Co-authored-by: Claude Fable 5.1 Co-authored-by: indianspeedster --- csrc/cumem_allocator.cpp | 22 ------- docs/features/sleep_mode.md | 2 +- .../memory/cumem/test_cumem.py | 11 ++-- .../memory/sleep_mode/test_sleep_mode.py | 62 +++++++++++++++++-- vllm/_aiter_ops.py | 6 +- vllm/config/vllm.py | 2 +- vllm/device_allocator/cumem.py | 8 ++- .../aiter_custom_all_reduce.py | 7 ++- .../device_communicators/cuda_communicator.py | 2 + 9 files changed, 84 insertions(+), 38 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 6965d2124a65..efb1f879d2ae 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -226,28 +226,6 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, return; } -#ifdef USE_ROCM - // ROCm-only: zero the freshly (re)mapped region. - // - // On some amdgpu driver versions the physical VRAM backing a buffer is not - // scrubbed on release, so the pages handed back here (e.g. on wake_up after a - // sleep that discarded a tag) can still contain stale data from a previously - // freed allocation -- possibly from another process. Callers of the device - // allocator expect newly mapped memory to be well defined; in particular - // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent - // state cache, and a new sequence that reads an unzeroed slot will propagate - // garbage/NaN, collapsing generation to token id 0 (decoded as "!"). - // - // This is scoped to the ROCm path on purpose: the CUDA path is unchanged. - // Use CUDA_CHECK so a failure sets the global error_code; the Python caller - // (python_create_and_map) only inspects error_code, so returning without - // setting it would make a failed memset look successful and still hand back - // unzeroed memory. - CUDA_CHECK((CUresult)hipMemset(reinterpret_cast(d_mem), 0, size)); - if (error_code != 0) { - return; - } -#endif // std::cout << "create_and_map: device=" << device << ", size=" << size << ", // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; } diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index a8aa15af7b52..0db82de9d3c5 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -107,7 +107,7 @@ By default, CUDA graph memory stays on the GPU while asleep. With `sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured into a cuMem pool that sleep backs up to CPU memory at both levels and any wake restores in place, so graphs are replayed, not recaptured. It needs the default -`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect. +`cumem` backend, CUDA or ROCm, and CUDA graphs; otherwise it has no effect. ```python llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index df8c6048824c..c7ebfd2609d6 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -359,7 +359,7 @@ def test_runtime_state_survives_sleep(level): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") def test_level2_discards_ordinary_tensor_with_weights_tag(): - """Reproduce the level-2 variant for an ordinary tensor in weights.""" + """Discarded weights-tag memory is remapped; ROCm zeroes it over stale pages.""" allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): @@ -374,8 +374,9 @@ def test_level2_discards_ordinary_tensor_with_weights_tag(): torch.accelerator.synchronize() assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers - assert torch.all(fake_weight == 0xA5) - assert torch.all(ordinary_tensor == 0xA5) + expected = 0 if current_platform.is_rocm() else 0xA5 + assert torch.all(fake_weight == expected) + assert torch.all(ordinary_tensor == expected) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @@ -427,7 +428,9 @@ def model(x): @pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"]) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_cudagraph_pool_sleep(level): """Routing, and the graph pool backed up at both sleep levels and restored in place by any wake.""" diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 9f6180335ab1..d897d710015b 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -209,8 +209,20 @@ def _cudagraph_bytes(worker) -> int: def _custom_ar_active(worker) -> bool: from vllm.distributed.parallel_state import get_tp_group - ca_comm = get_tp_group().device_communicator.ca_comm - return ca_comm is not None and not ca_comm.disabled + comm = get_tp_group().device_communicator + return any( + c is not None and not c.disabled for c in (comm.ca_comm, comm.aiter_ar_comm) + ) + + +def _capture_registers_buffers(worker) -> bool: + from vllm.distributed.parallel_state import get_tp_group + + comm = get_tp_group().device_communicator + ca, aiter = comm.ca_comm, comm.aiter_ar_comm + return (ca is not None and ca._capture_registered) or ( + aiter is not None and aiter.aiter_ca.enable_register_for_capturing + ) @pytest.mark.parametrize( @@ -227,7 +239,9 @@ def _custom_ar_active(worker) -> bool: ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"], ) @create_new_process_for_each_test() -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): """Each capture site uses the pool, exact across sleeps; TP=2 needs unregistered custom AR; off changes nothing.""" @@ -255,8 +269,10 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert all(graph_bytes) if offload else not any(graph_bytes) if not offload: return - if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)): - pytest.skip("Custom allreduce is unavailable on these GPUs") + if tp > 1: + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("Custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) expected = llm.generate(prompt, params)[0].outputs[0].text for level in (1, 2): @@ -268,6 +284,42 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert llm.generate(prompt, params)[0].outputs[0].text == expected +@multi_gpu_test(num_gpus=2) +@create_new_process_for_each_test() +@pytest.mark.skipif(not current_platform.is_rocm(), reason="AITER fused allreduce") +def test_sleep_cudagraph_pool_aiter_fused_ar_group_quant(monkeypatch): + """The fused AITER AR + RMSNorm + per-group FP8 quant captures into the + cuMem pool without IPC-registering its input, across sleeps.""" + for name, value in [ + ("VLLM_USE_V2_MODEL_RUNNER", "1"), + ("VLLM_ROCM_USE_AITER", "1"), + ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"), + ]: + monkeypatch.setenv(name, value) + monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False) + llm = LLM( + "Qwen/Qwen3-0.6B-FP8", + enable_sleep_mode=True, + sleep_mode_offload_cudagraph=True, + tensor_parallel_size=2, + max_model_len=1024, + compilation_config={"pass_config": {"fuse_allreduce_rms": True}}, + ) + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("AITER custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) + assert all(llm.collective_rpc(_cudagraph_bytes)) + prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) + assert llm.generate(prompt, params)[0].outputs[0].text + for level in (1, 2): + llm.sleep(level=level) + llm.wake_up(tags=["weights"]) + if level == 2: + llm.collective_rpc("reload_weights") + llm.wake_up(tags=["kv_cache"]) + assert llm.generate(prompt, params)[0].outputs[0].text + + @create_new_process_for_each_test() def test_deep_sleep_async(gpu_memory_cleared): async def test(): diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 44b89aeaf12c..0ed84906185c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1167,7 +1167,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, ) assert result is not None @@ -1219,7 +1220,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, emit_bf16=True, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5d38b6765308..bbda00af40bb 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -702,7 +702,7 @@ def use_cumem_cudagraph_pool(self) -> bool: and model_config.enable_sleep_mode and model_config.sleep_mode_backend == "cumem" and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and current_platform.is_cuda() + and current_platform.is_cuda_alike() ) @property diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index b85be08d28e3..a2e582f85b07 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -98,7 +98,7 @@ class CuMemAllocator: offloaded to CPU memory, and the rest of the tensors will be discarded. When we call `wake_up`, all tensors that are previously offloaded will be loaded back to GPU memory, and the rest of the tensors will - have empty memory. + have empty memory (zeroed on ROCm). Why it needs to be a singleton? When allocated tensors are garbage collected, PyTorch will call @@ -331,7 +331,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None: def wake_up(self, tags: list[str] | None = None) -> None: """Wake up the allocator from sleep mode. All data that is previously offloaded will be loaded back to GPU - memory, and the rest of the data will have empty memory. + memory, and the rest will have empty memory (zeroed on ROCm). Args: tags: The deferrable tags (weights, kv_cache) to load back to GPU @@ -358,6 +358,10 @@ def wake_up(self, tags: list[str] | None = None) -> None: cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + elif current_platform.is_rocm(): + # amdgpu <6.14 may return stale VRAM (#44972); drop with its support + libcudart.cudaMemset(ptr, 0, handle[1]) + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 6fbffcaaf73d..9a5970c565d7 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -30,6 +30,8 @@ def __init__( group: ProcessGroup, device: int | str | torch.device, max_size: int | None = None, + *, + register_graph_buffers: bool = True, ): from aiter.dist.device_communicators.custom_all_reduce import ( CustomAllreduce as _AiterCustomAllreduce, @@ -38,7 +40,10 @@ def __init__( if max_size is None: max_size = self.MAX_SIZE - self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + kwargs = {} + if not register_graph_buffers: + kwargs["enable_register_for_capturing"] = False + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs) @property def aiter_ca(self): diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index d4d588bf450d..992d07e15dad 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -151,6 +151,7 @@ def __init__( self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: @@ -174,6 +175,7 @@ def __init__( self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if self.aiter_ar_comm.disabled: self.aiter_ar_comm = None From 18c19eaf239f09bc2543637c73f3b2d05f9bf5c5 Mon Sep 17 00:00:00 2001 From: Taewoong Kim <116135174+ultramancode@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:53:50 +0900 Subject: [PATCH 273/952] [Frontend] Add cache salt and tracing to structured decisions (#60786) Signed-off-by: Taewoong Kim Co-authored-by: Chauncey --- .../online_serving/structured_decisions.md | 5 + docs/usage/security.md | 2 +- .../structured_decisions/test_strategies.py | 135 +++++++++++++++++- .../test_non_object_body_validation.py | 15 ++ .../generate/structured_decisions/protocol.py | 22 ++- .../generate/structured_decisions/serving.py | 23 +++ .../structured_decisions/strategies.py | 15 +- 7 files changed, 207 insertions(+), 10 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 2194b54e47b4..a1f58a3769b9 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -135,6 +135,11 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | `questions` | question id to `{type, instructions, criteria}`, asked in this order | | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | +| `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. + +When [tracing](../../../examples/observability/opentelemetry/README.md) is enabled +with `--otlp-traces-endpoint`, the HTTP `traceparent` and `tracestate` headers are +forwarded to every question's generation request. diff --git a/docs/usage/security.md b/docs/usage/security.md index c0d2c79700e6..a74ed1d6ebe9 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details. -`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, and by the Anthropic `/v1/messages` endpoint. +`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint. #### Usage with the OpenAI Python client diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 3d9538e4dc8b..5dcb8edb3061 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -2,24 +2,46 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import pytest -from fastapi import FastAPI +from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.structured_decisions.api_router import ( register_structured_decisions_api_router, ) +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, + StructuredDecisionResponse, +) from vllm.entrypoints.generate.structured_decisions.question_types import LABELS -from vllm.entrypoints.generate.structured_decisions.serving import state_text +from vllm.entrypoints.generate.structured_decisions.serving import ( + ServingStructuredDecisions, + state_text, +) from vllm.entrypoints.generate.structured_decisions.strategies import ( NextTokenStrategy, + ReadContext, reply_label_ids, select_read_strategy, ) +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.exception_handling.register import ( init_exception_handler, ) +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.renderers.hf import HfRenderer +from vllm.renderers.online_renderer import OnlineRenderer + +TRACE_HEADERS = { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", +} def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: @@ -69,9 +91,9 @@ def test_structured_state_keeps_unicode_in_prompt(): @pytest.fixture(scope="module") def qwen(): - from transformers import AutoTokenizer + from vllm.tokenizers import get_tokenizer - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + tokenizer = get_tokenizer("Qwen/Qwen3-0.6B") prompt_ids = tokenizer.apply_chat_template( [{"role": "user", "content": "state"}], add_generation_prompt=True, @@ -93,3 +115,108 @@ def test_labels_start_the_reply(qwen): # After a colon, Qwen writes ":A" as one token, so "A" is not one token. with pytest.raises(ValueError, match="not one distinct token"): reply_label_ids(tokenizer, tokenizer.encode("team:")) + + +@pytest.fixture +def serving_and_engine(qwen): + tokenizer, _ = qwen + model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32") + renderer = HfRenderer( + SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(_api_process_rank=0), + ), + tokenizer, + ) + engine = MagicMock(spec=EngineClient) + engine.errored = False + engine.model_config = model_config + engine.renderer = renderer + engine.input_processor = MagicMock() + online_renderer = OnlineRenderer( + model_config, + renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + ) + + async def generate(prompt, params, request_id, **kwargs): + label_ids = params.logprob_token_ids + yield RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=prompt["prompt_token_ids"], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="A", + token_ids=[label_ids[0]], + cumulative_logprob=None, + logprobs=[{token_id: Logprob(-1.0) for token_id in label_ids}], + finish_reason="length", + ) + ], + finished=True, + ) + + engine.generate.side_effect = generate + strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {})) + models = OpenAIServingModels( + engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)] + ) + yield ServingStructuredDecisions(models, strategy), engine + renderer.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cache_salt,tracing_enabled,request_headers", + [ + ("tenant-a", True, {"authorization": "Bearer test", **TRACE_HEADERS}), + ("tenant-b", False, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {}), + ("tenant-direct", True, None), + ], +) +async def test_request_metadata_reaches_every_question( + serving_and_engine, cache_salt, tracing_enabled, request_headers +): + serving, engine = serving_and_engine + engine.is_tracing_enabled.return_value = tracing_enabled + raw_request = None + if request_headers is not None: + raw_request = Request( + { + "type": "http", + "headers": [ + (k.encode(), v.encode()) for k, v in request_headers.items() + ], + } + ) + request = StructuredDecisionRequest( + state={"team": "billing", "language": "English"}, + questions={ + "team": { + "type": "choice", + "instructions": "Which team does the message name?", + "criteria": {"billing": "payments", "shipping": "deliveries"}, + }, + "lang": { + "type": "choice", + "instructions": "Which language does the message name?", + "criteria": {"English": "English text", "French": "French text"}, + }, + }, + cache_salt=cache_salt, + ) + response = await serving.create_decision(request, raw_request) + assert isinstance(response, StructuredDecisionResponse) + assert engine.generate.call_count == 2 + for call in engine.generate.call_args_list: + assert call.args[0].get("cache_salt") == cache_salt + assert call.kwargs["trace_headers"] == ( + TRACE_HEADERS if tracing_enabled and request_headers else None + ) diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py index 5bdc1d57873b..dd0e3eafd0a4 100644 --- a/tests/entrypoints/unit_tests/test_non_object_body_validation.py +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -13,6 +13,9 @@ import pytest from pydantic import ValidationError +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, +) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -74,6 +77,18 @@ def test_tokenize_chat_request_still_validates_dict_bodies(): CACHE_SALT_MODELS = [ + ( + StructuredDecisionRequest, + { + "state": "hello", + "questions": { + "lang": { + "type": "choice", + "criteria": {"English": "Text written in English."}, + } + }, + }, + ), (CompletionRequest, {"prompt": "hello"}), (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}), (ResponsesRequest, {"input": "hello"}), diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index 3a193e203bad..b15fee171fca 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -5,10 +5,11 @@ import time from typing import Any, Literal -from pydantic import Field +from pydantic import Field, model_validator from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs from vllm.utils import random_uuid @@ -30,9 +31,27 @@ class StructuredDecisionRequest(OpenAIBaseModel): default=None, description="Context placed ahead of the questions." ) chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = Field( + default=None, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1) request_id: str = Field(default_factory=random_uuid) + @model_validator(mode="before") + @classmethod + def check_cache_salt_support(cls, data: Any) -> Any: + if isinstance(data, dict): + validate_cache_salt(data.get("cache_salt")) + return data + class DecisionUsage(OpenAIBaseModel): input_tokens: int @@ -65,6 +84,7 @@ class ReadPromptRequest(OpenAIBaseModel): would follow a thought rather than start the reply.""" chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = None def build_chat_params( self, diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 4921bfcf3574..e6e990bd46df 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json +from collections.abc import Mapping from typing import Any from fastapi import Request @@ -11,6 +12,11 @@ from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger +from vllm.tracing import ( + contains_trace_headers, + extract_trace_headers, + log_tracing_disabled_warning, +) from .protocol import ( DecisionUsage, @@ -71,6 +77,16 @@ def __init__( self.strategy = strategy self.limits = strategy.limits() + async def _get_trace_headers( + self, headers: Mapping[str, str] + ) -> Mapping[str, str] | None: + if not contains_trace_headers(headers): + return None + if not await self.strategy.context.engine_client.is_tracing_enabled(): + log_tracing_disabled_warning() + return None + return extract_trace_headers(headers) + async def create_decision( self, request: StructuredDecisionRequest, @@ -88,6 +104,11 @@ async def create_decision( questions = parse_questions(request, self.limits) lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] engine_client.check_admission(len(questions)) + trace_headers = ( + None + if raw_request is None + else await self._get_trace_headers(raw_request.headers) + ) reads = await self.strategy.read( questions, request.instructions, @@ -96,6 +117,8 @@ async def create_decision( chat_template_kwargs=request.chat_template_kwargs, lora_request=lora_request, priority=request.priority, + cache_salt=request.cache_salt, + trace_headers=trace_headers, ) except StructuredDecisionError as e: return self.create_error_response(e) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index af588e13a289..65b5a9dd7b05 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -7,7 +7,7 @@ import math from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any @@ -74,6 +74,8 @@ async def read( chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: """One read per question, in order.""" @@ -136,13 +138,15 @@ def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) def _read_request( - self, chat_template_kwargs: dict[str, Any] | None + self, chat_template_kwargs: dict[str, Any] | None, cache_salt: str | None = None ) -> ReadPromptRequest: if (chat_template_kwargs or {}).get("enable_thinking"): raise StructuredDecisionError( "a read is the reply's first token, so thinking must be off" ) - return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs) + return ReadPromptRequest( + chat_template_kwargs=chat_template_kwargs, cache_salt=cache_salt + ) async def _render( self, read_request: ReadPromptRequest, messages: list[dict[str, Any]] @@ -170,9 +174,11 @@ async def read( chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: ctx = self.context - read_request = self._read_request(chat_template_kwargs) + read_request = self._read_request(chat_template_kwargs, cache_salt) slots, engine_inputs = [], [] for q in questions: @@ -197,6 +203,7 @@ async def read( ], request_id, lora_request=lora_request, + trace_headers=trace_headers, priority=priority, ) From bc713cf8cc861c0d0ccfe46c9ce322ab81993ace Mon Sep 17 00:00:00 2001 From: Chao Lei Date: Fri, 9 Oct 2026 17:19:33 +0800 Subject: [PATCH 274/952] [Bugfix]MooncakeConnect add query timeout and retry (#55923) --- docs/features/mooncake_connector_usage.md | 5 ++-- .../v1/mooncake/mooncake_connector.py | 28 +++++++++++++++++-- vllm/models/deepseek_v41/common/engram.py | 2 +- 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/docs/features/mooncake_connector_usage.md b/docs/features/mooncake_connector_usage.md index 866d6d6bd48c..00bfdeabc14c 100644 --- a/docs/features/mooncake_connector_usage.md +++ b/docs/features/mooncake_connector_usage.md @@ -49,12 +49,13 @@ cache), see [ECMooncakeConnector Usage Guide](mooncake_ec_connector_usage.md). - For headless instances, must be the same as the master instance - Each instance needs a unique port on its host; using the same port number across different hosts is fine -- `VLLM_MOONCAKE_CONNECTOR_TIMEOUT`: Per-request timeout (in seconds) for MooncakeConnector HTTP calls to the bootstrap server, including prefiller worker registration at startup. (Optional) +- `VLLM_MOONCAKE_CONNECTOR_TIMEOUT`: Per-request timeout (in seconds) for MooncakeConnector HTTP calls to the bootstrap server, including prefiller worker registration at startup and metadata queries. (Optional) - Default: 30.0 - Global rank 0 hosts the bootstrap server in the same process that mounts the Mooncake transfer engine's host segment, so a large `MOONCAKE_GLOBAL_SEGMENT_SIZE` can delay responses for several seconds during startup - Raise this if bootstrap server timeout warnings appear on hosts with very large host-memory segments - Bootstrap registration uses up to 3 attempts internally, with exponential backoff capped at 10 seconds; connection errors and timeouts are both retried - - Once the attempts are exhausted the caller raises instead of blocking, so failures surface rather than hanging + - Once registration attempts are exhausted the caller raises instead of blocking, so failures surface rather than hanging + - Metadata queries reuse the same timeout and 3-attempt limit, with a fixed 100ms delay between attempts. If all attempts fail, the error is logged and waiting queries are notified - `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM. - Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 8300cd950b0d..9d5194f33e27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -2438,9 +2438,30 @@ def process_pulling_result( async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): url = remote_bootstrap_addr + "/query" + max_attempts = _BOOTSTRAP_MAX_ATTEMPTS try: - async with httpx.AsyncClient() as client: - response = await client.get(url) + async with httpx.AsyncClient( + timeout=envs.VLLM_MOONCAKE_CONNECTOR_TIMEOUT + ) as client: + retry_delay = 0.1 + for attempt in range(1, max_attempts + 1): + try: + response = await client.get(url) + break + except httpx.RequestError as e: + if attempt == max_attempts: + raise + logger.warning( + "Bootstrap query to %s failed on attempt %d/%d " + "(%s: %s); retrying in %.1f seconds", + remote_bootstrap_addr, + attempt, + max_attempts, + type(e).__name__, + e, + retry_delay, + ) + await asyncio.sleep(retry_delay) response.raise_for_status() data: dict = response.json() for _, dp_entry in data.items(): @@ -2455,8 +2476,9 @@ async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): self._tp_size[remote_engine_id] = len(dp_entry["worker_addr"]) except Exception as e: logger.error( - "Failed to connect to bootstrap server %s: %s", + "Failed to connect to bootstrap server %s (%s): %s", remote_bootstrap_addr, + type(e).__name__, e, ) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index 0e160c8b06cf..236235dfa82f 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -1305,7 +1305,7 @@ def _engram_select_rows( ) -> None: """Copy one token window out of a rank-major gathered buffer. - Both gathers land rank-major ([rank][token][local width]); this walks the + Both gathers land rank-major (`[rank][token][local width]`); this walks the window the rank keeps and lays its ranks out side by side as width. """ if output.numel() == 0: From 2c99ee9333821030daf84718b2a76ea28d94ba3c Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:29:03 +0100 Subject: [PATCH 275/952] [Chore] Remove vendored code made redundant by Transformers 5.16.1 (#60623) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .agents/skills/transformers-upgrade/SKILL.md | 178 +++++++++++ .../transformers-upgrade/agents/openai.yaml | 4 + .../scripts/diff_dumps.py | 51 ++++ .../scripts/dump_configs.py | 30 ++ .../scripts/dump_model.py | 114 +++++++ .../scripts/find_upstreamed.py | 147 +++++++++ .buildkite/test-amd.yaml | 10 - .buildkite/test_areas/lm_eval.yaml | 7 - .buildkite/test_areas/misc.yaml | 1 - .buildkite/test_areas/models_language.yaml | 2 - .claude/skills/transformers-upgrade | 1 + tests/config/test_model_arch_config.py | 16 +- .../token_in_token_out/test_serving_tokens.py | 10 +- .../layers/test_fused_shared_expert.py | 2 +- .../multimodal/generation/test_common.py | 5 +- .../processing/test_transformers_image.py | 6 +- tests/models/registry.py | 12 +- .../bitsandbytes/test_bitsandbytes.py | 5 +- tests/tokenizers_/test_registry.py | 3 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 2 +- vllm/model_executor/models/AXK1.py | 2 +- vllm/model_executor/models/aimv2.py | 20 +- vllm/model_executor/models/cohere2_moe.py | 14 - vllm/model_executor/models/cosmos3.py | 19 +- vllm/model_executor/models/cosmos3_edge.py | 3 +- vllm/model_executor/models/diffusion_gemma.py | 27 +- .../model_executor/models/embedding_gemma2.py | 7 +- vllm/model_executor/models/gemma3.py | 15 +- vllm/model_executor/models/gemma3n.py | 12 +- vllm/model_executor/models/gemma4.py | 7 +- vllm/model_executor/models/gemma4_dspark.py | 3 +- vllm/model_executor/models/gemma4_mm.py | 16 +- vllm/model_executor/models/gemma4_mtp.py | 3 +- vllm/model_executor/models/granite4_vision.py | 96 +++--- vllm/model_executor/models/granite_speech.py | 7 +- .../model_executor/models/granitemoehybrid.py | 4 - vllm/model_executor/models/hy_v3.py | 3 +- vllm/model_executor/models/hyperclovax.py | 17 +- vllm/model_executor/models/kimi_k25.py | 17 +- vllm/model_executor/models/kimi_k25_vit.py | 44 +-- vllm/model_executor/models/lfm2_moe.py | 2 +- vllm/model_executor/models/modernbert.py | 28 +- vllm/model_executor/models/nemotron.py | 2 +- vllm/model_executor/models/nemotron_h.py | 2 +- vllm/model_executor/models/nemotron_h_mtp.py | 2 +- vllm/model_executor/models/ovis.py | 2 +- vllm/model_executor/models/qwen3_5.py | 11 +- vllm/model_executor/models/qwen3_5_mtp.py | 3 +- vllm/model_executor/models/qwen3_next.py | 2 +- vllm/model_executor/models/qwen3_next_mtp.py | 2 +- vllm/models/kimi_k3/amd/model.py | 1 + vllm/models/kimi_k3/nvidia/model.py | 1 + vllm/transformers_utils/config.py | 47 +-- vllm/transformers_utils/configs/AXK1.py | 216 ------------- vllm/transformers_utils/configs/__init__.py | 40 +-- vllm/transformers_utils/configs/afmoe.py | 87 ------ vllm/transformers_utils/configs/cosmos3.py | 7 - .../configs/cosmos3_edge.py | 11 +- .../configs/deepseek_vl2.py | 5 - .../configs/diffusion_gemma.py | 44 --- vllm/transformers_utils/configs/gemma4.py | 34 --- .../configs/granite4_vision.py | 94 ------ vllm/transformers_utils/configs/hy_v3.py | 185 ----------- .../transformers_utils/configs/hyperclovax.py | 277 ----------------- vllm/transformers_utils/configs/kimi_k25.py | 131 -------- vllm/transformers_utils/configs/lfm2_moe.py | 164 ---------- vllm/transformers_utils/configs/nemotron.py | 219 ------------- vllm/transformers_utils/configs/nemotron_h.py | 288 ------------------ .../transformers_utils/configs/olmo_hybrid.py | 285 ----------------- vllm/transformers_utils/configs/ovis.py | 72 +---- vllm/transformers_utils/configs/qwen3_5.py | 193 ------------ .../transformers_utils/configs/qwen3_5_moe.py | 205 ------------- vllm/transformers_utils/configs/qwen3_next.py | 277 ----------------- .../model_arch_config_convertor.py | 24 +- .../transformers_utils/processors/__init__.py | 4 - .../processors/granite4_vision.py | 78 ----- 76 files changed, 719 insertions(+), 3268 deletions(-) create mode 100644 .agents/skills/transformers-upgrade/SKILL.md create mode 100644 .agents/skills/transformers-upgrade/agents/openai.yaml create mode 100644 .agents/skills/transformers-upgrade/scripts/diff_dumps.py create mode 100644 .agents/skills/transformers-upgrade/scripts/dump_configs.py create mode 100644 .agents/skills/transformers-upgrade/scripts/dump_model.py create mode 100644 .agents/skills/transformers-upgrade/scripts/find_upstreamed.py create mode 120000 .claude/skills/transformers-upgrade delete mode 100644 vllm/transformers_utils/configs/AXK1.py delete mode 100644 vllm/transformers_utils/configs/afmoe.py delete mode 100644 vllm/transformers_utils/configs/cosmos3.py delete mode 100644 vllm/transformers_utils/configs/diffusion_gemma.py delete mode 100644 vllm/transformers_utils/configs/gemma4.py delete mode 100644 vllm/transformers_utils/configs/granite4_vision.py delete mode 100644 vllm/transformers_utils/configs/hy_v3.py delete mode 100644 vllm/transformers_utils/configs/hyperclovax.py delete mode 100644 vllm/transformers_utils/configs/kimi_k25.py delete mode 100644 vllm/transformers_utils/configs/lfm2_moe.py delete mode 100644 vllm/transformers_utils/configs/nemotron.py delete mode 100644 vllm/transformers_utils/configs/nemotron_h.py delete mode 100644 vllm/transformers_utils/configs/olmo_hybrid.py delete mode 100644 vllm/transformers_utils/configs/qwen3_5.py delete mode 100644 vllm/transformers_utils/configs/qwen3_5_moe.py delete mode 100644 vllm/transformers_utils/configs/qwen3_next.py delete mode 100644 vllm/transformers_utils/processors/granite4_vision.py diff --git a/.agents/skills/transformers-upgrade/SKILL.md b/.agents/skills/transformers-upgrade/SKILL.md new file mode 100644 index 000000000000..de9065b6fa81 --- /dev/null +++ b/.agents/skills/transformers-upgrade/SKILL.md @@ -0,0 +1,178 @@ +--- +name: transformers-upgrade +description: "Playbook for changing vLLM's Transformers version bounds. Use when bumping the pinned or maximum Transformers version, raising the minimum Transformers version, or deleting vendored configs, processors, or compat shims that upstream Transformers now provides." +--- + +# Upgrading Transformers + +The supported range lives in `requirements/common.txt`. CI pins one exact +version in `requirements/test/*.in` and the compiled `requirements/test/*.txt`. +A version change is one of two jobs, and they usually come as a pair: + +- **Bump the pin**: a new Transformers release becomes the CI version and + the upper bound. vLLM grows compat code for it. +- **Raise the floor**: the minimum moves up. That compat code, and every + vendored config or processor upstream now ships, becomes dead weight. + Every floor raise ends with an **unvendoring pass**. + +## Bump the pin + +1. Edit the pin line in each `requirements/test/*.in` and `*.txt` and the + upper bound in `requirements/common.txt`. Past bumps edit only the pin + line in the compiled files; see `git log -- requirements/common.txt`. +2. Read the release notes for config validation changes, renamed auto-mapping + formats, and removed attributes; these are what break vLLM. +3. Search `tests/` for `min_transformers_version` and skip markers naming the + new version, and un-gate them. +4. Run the suites in [Testing](#testing) and fix what breaks. Mark every + shim with a `# Transformers >= X` comment so the next floor raise finds it. + +Done when the [Testing](#testing) suites pass on the new pin and every test +waiting on this version is un-gated. + +## Raise the floor (unvendoring pass) + +### 1. Inventory + +Inventory candidates with the floor version installed (no vLLM install needed): + +```bash +uv run --no-project --with transformers== \ + python .agents/skills/transformers-upgrade/scripts/find_upstreamed.py +``` + +It lists: + +- vendored configs whose `model_type` upstream now owns; +- `_CONFIG_REGISTRY` overrides of upstream types; +- registered processors that collide with upstream class names or models; +- version-gated code at or below the floor; +- code that reads the installed Transformers version, whatever its threshold; +- patches citing a Transformers issue or PR (with whether a referenced PR is + in the floor release) or a future minimum version. + +It finds candidates by name only. Every hit still needs a verdict. + +### 2. Port each candidate to upstream + +Every hit is deleted by default. A vendored copy is a liability: it drifts +from upstream, hides upstream fixes, and shadows the native class. Port the +vLLM model code to the upstream schema instead, renaming attribute reads and +deriving values from upstream fields. + +Fan candidates out to parallel subagents. Each one compares the vendored and +upstream classes, then greps every attribute the vLLM model code reads. +[Upstreaming hazards](#upstreaming-hazards) lists the usual fixes. + +Keep a vendored class only when upstream cannot load a checkpoint vLLM must +keep serving, and record why in the PR description. A different upstream +schema, a model-code rewrite, or a needed eval is work for the PR, not a +reason to keep. + +### 3. Prove it on real checkpoints + +For every checkpoint in `tests/models/registry.py` that uses a ported +config, dump what vLLM builds before and after the change, then diff the two. +Run both dumps; each catches what the other misses: + +- `dump_configs.py`: the resolved config, attribute by attribute. +- `dump_model.py`: the model built on the meta device with no GPU or weights. + It records ModelConfig sizes, every parameter shape, and every module's + scalar attributes (`head_dim`, `num_kv_heads`, `scaling`, ...). Missing + fields that are read with defaults only show up here. + +```bash +git worktree add --detach /vllm-base upstream/main +S=.agents/skills/transformers-upgrade/scripts +for dump in dump_configs dump_model; do + PYTHONPATH=/vllm-base .venv/bin/python $S/$dump.py before.jsonl + PYTHONPATH=$PWD .venv/bin/python $S/$dump.py after.jsonl + .venv/bin/python $S/diff_dumps.py before.jsonl after.jsonl +done +``` + +Add `--trust-remote-code` for checkpoints whose `config.json` has an +`auto_map`, and dump both ways. + +How to read the diffs: + +- **A `` config field the model reads:** this crashes or silently + falls back to a default. Port the model code to the upstream field. +- **New upstream-only defaults** (`attention_bias=False`, + `pad_token_id=None`): usually benign. A model diff of 0 confirms it. +- **Any model diff:** this is a numerics change, and the PR needs a model + eval. For example, A.X-K1 gained the checkpoint's YaRN rope scaling, which + the vendored class had been dropping; it shows up as + `layers.*.self_attn.scaling (x61)`. + +### 4. Delete + +For each deleted config: + +- Delete the file. +- Remove its entries from `_CLASS_TO_MODULE` and `__all__` in + `vllm/transformers_utils/configs/__init__.py`. +- Remove its entries from `_CONFIG_REGISTRY` in + `vllm/transformers_utils/config.py`. +- Switch every importer, including tests, to `from transformers import ...`. +- Remove the `source_file_dependencies` lines naming the file from + `.buildkite/`. + +Deleted processors follow the same steps in +`vllm/transformers_utils/processors/__init__.py`. + +Done when every line of the inventory output is deleted, or kept with a +reason that meets the bar in step 2. The PR also carries the config diff +results and evals for any changed values. + +## Upstreaming hazards + +Upstream configs are `@strict` dataclasses. Checkpoint keys upstream does not +declare survive as plain attributes. Defaults that only vLLM sets, and values +vLLM derives, vanish. + +| Hazard | Signal | Fix | +| --- | --- | --- | +| Upstream consumes legacy keys | `__post_init__` pops keys the model reads (`deepseek_v4` `compress_ratios`) | Read the upstream replacement (`layer_types`, `mlp_layer_types`, per-layer `rope_parameters`) | +| Renamed or dropped fields | Diff shows `` for a field the model reads | Read the upstream name; derive dropped flags from what implies them (`enable_moe_block` from `num_experts`) | +| Type check on the config | `ctx.get_hf_config(SomeConfig)` relies on the vendored base class | Accept the upstream class too, e.g. a `Cosmos3ProcessingInfo` overriding `get_hf_config` | +| Remote code takes over | Checkpoint `config.json` has `auto_map`; with `--trust-remote-code`, the hub class replaces the native one once vLLM stops registering its own | Check the remote class loads; if not, register the upstream class for that `model_type` | +| Read-only upstream fields | Properties or derived fields break subclasses and `hf_overrides` (`nemotron_h` `hybrid_override_pattern`) | Override the source field (`layers_block_type`) instead | +| Behavioural rewrite | Upstream derives fields differently (`diffusion_gemma` `sliding_window`) | Upstream matches the reference implementation; verify with an eval | +| Original checkpoint format | Upstream targets `-hf` conversions; vLLM serves the originals (`qwen3_asr`, `qianfan_ocr`) | Keep only if upstream cannot load the original at all | +| Processor output contract | Model code consumes keys upstream does not emit | Port the model to upstream's keys | +| Name collision | A registered vendored processor shares an upstream class name and hijacks `get_processor` (`InternVLProcessor`) | Drop the registry entry | + +Version gates from the inventory need a decision: + +- **Code** only reachable below the floor is deleted. This covers: + - comparisons against `TRANSFORMERS_VERSION`; + - branches labelled `# Transformers < X`; + - fallbacks for fields that upstream configs now always set, e.g. + `mlp_layer_types`, `per_layer_config` or legacy `layer_types` names. + + A `skipif` whose condition is now always true becomes an unconditional + `skip` with the same reason. +- **`min_transformers_version`** at or below the floor is removed. +- **`max_transformers_version`** below the floor means CI never tests that + model; report it rather than deleting it silently. +- **Patches citing a Transformers issue or PR:** + - Delete the patch when its fix is in the floor release. + - For an issue, find the PR that closed it and check that PR. A closed + issue is not a fixed one: #43329 was closed while the bug was still in + 5.16.1. + - Re-run the patched behaviour before deleting. A fix in the release does + not always cover vLLM's path: #47924 made Emu3's processor add BOS, but + vLLM tokenizes without the processor. + - Links that only explain a design choice (e.g. a float32 RoPE note) are + not patches; leave them. + +## Testing + +```bash +.venv/bin/python -m pytest tests/transformers_utils tests/config tests/tokenizers_ +.venv/bin/python -m pytest tests/models/multimodal/processing/test_common.py -k "" +``` + +Add `tests/models/test_initialization.py -k ""` where a GPU is +available, and a model eval for any model whose resolved config changed. diff --git a/.agents/skills/transformers-upgrade/agents/openai.yaml b/.agents/skills/transformers-upgrade/agents/openai.yaml new file mode 100644 index 000000000000..b881c5d7ea27 --- /dev/null +++ b/.agents/skills/transformers-upgrade/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Transformers Upgrade" + short_description: "Bump vLLM's Transformers version and delete vendored code upstream now provides." + default_prompt: "Use $transformers-upgrade to raise the minimum Transformers version and remove vendored configs, processors, and shims it makes redundant." diff --git a/.agents/skills/transformers-upgrade/scripts/diff_dumps.py b/.agents/skills/transformers-upgrade/scripts/diff_dumps.py new file mode 100644 index 000000000000..76fb2038a151 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/diff_dumps.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diff two dump_configs.py or dump_model.py outputs, per repo. + +python diff_dumps.py before.jsonl after.jsonl +""" + +import json +import sys +from collections import Counter + +import regex as re + +IGNORED = {"transformers_version", "_name_or_path", "architectures", "auto_map"} + + +def load(path: str) -> dict[str, dict]: + with open(path) as f: + return {record["repo"]: record for record in map(json.loads, f)} + + +def diff(before: dict, after: dict, prefix: str = ""): + for key in sorted(set(before) | set(after)): + if key in IGNORED: + continue + old, new = before.get(key, ""), after.get(key, "") + if isinstance(old, dict) and isinstance(new, dict) and key != "rope_parameters": + yield from diff(old, new, f"{prefix}{key}.") + elif old != new: + yield f"{prefix}{key}", json.dumps(old), json.dumps(new) + + +def collapse(diffs) -> Counter: + """Merge diffs that differ only by layer index, e.g. `layers.*.scaling`.""" + return Counter( + (re.sub(r"\.\d+\.", ".*.", path), old, new) for path, old, new in diffs + ) + + +before, after = load(sys.argv[1]), load(sys.argv[2]) +for repo, old in before.items(): + new = after.get(repo, {"error": "missing from after"}) + if "error" in old or "error" in new: + print(f"== {repo}\n before: {old.get('error', 'ok')}") + print(f" after: {new.get('error', 'ok')}") + continue + diffs = collapse(diff(old["cfg"], new["cfg"])) + print(f"== {repo}: {old['cls']} -> {new['cls']} ({sum(diffs.values())} diffs)") + for (path, old_value, new_value), count in diffs.items(): + times = f" (x{count})" if count > 1 else "" + print(f" {path}{times}: before={old_value} after={new_value}") diff --git a/.agents/skills/transformers-upgrade/scripts/dump_configs.py b/.agents/skills/transformers-upgrade/scripts/dump_configs.py new file mode 100644 index 000000000000..6d6913fb5178 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/dump_configs.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dump the config vLLM resolves for each checkpoint, one JSON line per repo. + +Run once per vLLM tree (before and after a change), then compare the two +outputs with diff_dumps.py: + + PYTHONPATH= .venv/bin/python dump_configs.py out.jsonl ... + +Pass --trust-remote-code to resolve checkpoints the way users who set it do. +""" + +import json +import sys + +from vllm.transformers_utils.config import get_config + +out, *args = sys.argv[1:] +trust_remote_code = "--trust-remote-code" in args +with open(out, "w") as f: + for repo in (arg for arg in args if arg != "--trust-remote-code"): + try: + config = get_config(repo, trust_remote_code=trust_remote_code) + cls = f"{type(config).__module__}.{type(config).__name__}" + record = {"repo": repo, "cls": cls, "cfg": config.to_dict()} + except Exception as e: + error = f"{type(e).__name__}: {e}".splitlines()[0] + record = {"repo": repo, "error": error} + f.write(json.dumps(record, default=str, sort_keys=True) + "\n") + f.flush() diff --git a/.agents/skills/transformers-upgrade/scripts/dump_model.py b/.agents/skills/transformers-upgrade/scripts/dump_model.py new file mode 100644 index 000000000000..47efbe873f20 --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/dump_model.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build each model on the meta device and dump what vLLM constructs. + +Records ModelConfig-derived sizes, every parameter shape, and scalar +attributes of every module (head_dim, num_kv_heads, sliding_window, ...). +Needs no GPU and no weights. Run once per vLLM tree, then compare the two +outputs with diff_dumps.py: + + PYTHONPATH= .venv/bin/python dump_model.py out.jsonl ... + +Pass --trust-remote-code to resolve checkpoints the way users who set it do. +""" + +import json +import socket +import sys + +import torch + +from vllm.config import set_current_vllm_config +from vllm.config.vllm import VllmConfig +from vllm.distributed import ( + init_distributed_environment, + initialize_model_parallel, + model_parallel_is_initialized, +) +from vllm.engine.arg_utils import EngineArgs +from vllm.model_executor.model_loader.utils import initialize_model +from vllm.utils.torch_utils import set_default_torch_dtype + +# Construction only: skip model runner checks that need a GPU or Triton +VllmConfig._get_v1_model_runner_unsupported_features = lambda self: [] # type: ignore[method-assign] + +# Model code that moves submodules to a device must leave them on meta +_module_to = torch.nn.Module.to + + +def _meta_safe_to(self, *args, **kwargs): + kwargs.pop("device", None) + args = tuple(a for a in args if not isinstance(a, (torch.device, str, int))) + return _module_to(self, *args, **kwargs) + + +torch.nn.Module.to = _meta_safe_to # type: ignore[method-assign] + +SCALARS = (bool, int, float, str) + + +def module_attrs(model: torch.nn.Module) -> dict[str, dict]: + attrs = {} + for name, module in model.named_modules(): + scalars = { + key: value + for key, value in vars(module).items() + if not key.startswith("_") and isinstance(value, SCALARS) + } + if scalars: + attrs[name or ""] = scalars + return attrs + + +def dump(repo: str, trust_remote_code: bool) -> dict: + args = EngineArgs( + model=repo, + load_format="dummy", + enforce_eager=True, + max_model_len=4096, + trust_remote_code=trust_remote_code, + ) + vllm_config = args.create_engine_config() + model_config = vllm_config.model_config + with set_current_vllm_config(vllm_config): + if not model_parallel_is_initialized(): + with socket.socket() as sock: + sock.bind(("", 0)) + port = sock.getsockname()[1] + init_distributed_environment(1, 0, f"tcp://127.0.0.1:{port}", 0, "gloo") + initialize_model_parallel(1, 1) + with set_default_torch_dtype(model_config.dtype), torch.device("meta"): + model = initialize_model(vllm_config=vllm_config) + return { + "repo": repo, + "cls": type(model).__name__, + "cfg": { + "model_config": { + "head_size": model_config.get_head_size(), + "total_num_kv_heads": model_config.get_total_num_kv_heads(), + "sliding_window": model_config.get_sliding_window(), + "num_layers": model_config.get_num_layers(vllm_config.parallel_config), + }, + "params": {n: list(p.shape) for n, p in model.named_parameters()}, + "modules": module_attrs(model), + }, + } + + +def main() -> None: + out, *args = sys.argv[1:] + trust_remote_code = "--trust-remote-code" in args + repos = [arg for arg in args if arg != "--trust-remote-code"] + with open(out, "w") as f: + for repo in repos: + try: + record = dump(repo, trust_remote_code) + except Exception as e: + error = f"{type(e).__name__}: {e}".splitlines()[0] + record = {"repo": repo, "error": error} + f.write(json.dumps(record, default=str, sort_keys=True) + "\n") + f.flush() + + +if __name__ == "__main__": + main() diff --git a/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py b/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py new file mode 100644 index 000000000000..4cd73c29cbac --- /dev/null +++ b/.agents/skills/transformers-upgrade/scripts/find_upstreamed.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""List vendored vLLM code that the installed Transformers may make redundant. + +Run with the Transformers version under test, from the vLLM repo root: + + uv run --no-project --with transformers== \ + python .agents/skills/transformers-upgrade/scripts/find_upstreamed.py + +Imports nothing from vLLM, so no vLLM install is needed. Uses the `gh` CLI, if +available, to check whether referenced Transformers PRs are in this release. +""" + +import json +import subprocess +import sys +from pathlib import Path + +import regex as re +import transformers +from packaging.version import Version +from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES +from transformers.models.auto.processing_auto import PROCESSOR_MAPPING_NAMES + +ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else ".") +FLOOR = Version(transformers.__version__) +UPSTREAM_NAMES = set(dir(transformers)) + +MODEL_TYPE = re.compile(r'model_type\s*(?::\s*str\s*)?=\s*"([^"]+)"') +REGISTRY_KEY = re.compile(r'^\s+(?:\*\*\{)?"?([\w-]+)"?(?:=|": )"(\w+)"', re.M) +CLASS_KEY = re.compile(r'^\s+"(\w+)": ', re.M) +UPSTREAM_REF = re.compile( + r"huggingface/transformers(?:/(?:issues|pull)/|#)(\d+)" + r"|(?:once|until|when) [Tt]ransformers [\d.]+ is the minimum" +) +VERSIONED = re.compile( + r"(?:min_transformers_version|max_transformers_version|check_version|Version)" + r'\(?\s*=?\s*"([45]\.\d+(?:\.\d+)?)[^"]*"' + r"|[Tt]ransformers\s*(?:v|[<>]=?\s*)([45]\.\d+(?:\.\d+)?|4)\b" +) +READS_VERSION = re.compile( + r"transformers import __version__|transformers\.__version__" + r"|version\([\"']transformers[\"']\)|TRANSFORMERS_VERSION|check_version\(" +) + + +def section(title: str) -> None: + print(f"\n## {title}") + + +def vendored_configs() -> None: + section(f"Vendored configs whose model_type is upstream in {FLOOR}") + files = sorted(ROOT.glob("vllm/transformers_utils/configs/*.py")) + files += sorted(ROOT.glob("vllm/models/*/config*.py")) + files += sorted(ROOT.glob("vllm/models/*/configs/*.py")) + for path in files: + types = sorted(set(MODEL_TYPE.findall(path.read_text()))) + hits = [t for t in types if t in CONFIG_MAPPING_NAMES] + if hits: + upstream = ", ".join(f"{t} -> {CONFIG_MAPPING_NAMES[t]}" for t in hits) + print(f"- {path.relative_to(ROOT)}: {upstream}") + + registry = (ROOT / "vllm/transformers_utils/config.py").read_text() + start = registry.index("_CONFIG_REGISTRY") + body = registry[start : registry.index("\n)\n", start)] + section("_CONFIG_REGISTRY entries that override an upstream model_type") + for model_type, cls in REGISTRY_KEY.findall(body): + if model_type in CONFIG_MAPPING_NAMES: + print(f"- {model_type}: vLLM {cls} over {CONFIG_MAPPING_NAMES[model_type]}") + + +def vendored_processors() -> None: + init = (ROOT / "vllm/transformers_utils/processors/__init__.py").read_text() + names = sorted(set(CLASS_KEY.findall(init))) + section("Registered processor names that collide with upstream classes") + for name in names: + if name in UPSTREAM_NAMES: + print(f"- {name}") + section("Processor files named after a model_type with an upstream processor") + for path in sorted(ROOT.glob("vllm/transformers_utils/processors/*.py")): + if path.stem in PROCESSOR_MAPPING_NAMES: + upstream = PROCESSOR_MAPPING_NAMES[path.stem] + print(f"- {path.relative_to(ROOT)} -> {upstream}") + + +def version_gates() -> None: + files = sorted([*ROOT.glob("vllm/**/*.py"), *ROOT.glob("tests/**/*.py")]) + section(f"Version-gated code at or below {FLOOR} (review each)") + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + for match in VERSIONED.finditer(line): + version = Version(match.group(1) or match.group(2)) + if version <= FLOOR: + print(f"- {path.relative_to(ROOT)}:{lineno}: {line.strip()}") + break + section("Code that reads the installed Transformers version (review each)") + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + if READS_VERSION.search(line): + print(f"- {path.relative_to(ROOT)}:{lineno}: {line.strip()}") + + +def upstream_status(number: str) -> str: + """Whether Transformers issue/PR `number` is resolved in this release.""" + + def gh(path: str) -> dict | None: + try: + out = subprocess.run(["gh", "api", path], capture_output=True, text=True) + except FileNotFoundError: + return None + return json.loads(out.stdout) if out.returncode == 0 else None + + repo = "repos/huggingface/transformers" + if (issue := gh(f"{repo}/issues/{number}")) is None: + return "status unknown" + if "pull_request" not in issue: + return f"issue {issue['state']}, check which PR fixed it" + pr = gh(f"{repo}/pulls/{number}") + if pr is None or not pr["merged_at"]: + return "PR not merged" + compare = gh(f"{repo}/compare/v{FLOOR}...{pr['merge_commit_sha']}") + included = compare and compare["status"] in ("behind", "identical") + return f"PR merged, {'in' if included else 'not in'} v{FLOOR}" + + +def upstream_refs() -> None: + section("Patches referencing Transformers issues/PRs or a future minimum") + files = sorted([*ROOT.glob("vllm/**/*.py"), *ROOT.glob("tests/**/*.py")]) + statuses: dict[str, str] = {} + for path in files: + for lineno, line in enumerate(path.read_text().splitlines(), 1): + if (match := UPSTREAM_REF.search(line)) is None: + continue + note = "" + if number := match.group(1): + if number not in statuses: + statuses[number] = upstream_status(number) + note = f" [#{number}: {statuses[number]}]" + print(f"- {path.relative_to(ROOT)}:{lineno}:{note} {line.strip()}") + + +if __name__ == "__main__": + print(f"# Transformers {FLOOR}") + vendored_configs() + vendored_processors() + version_gates() + upstream_refs() diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 4ae92315b4c9..a3f7dbd83874 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -237,8 +237,6 @@ steps: - tests/models/language/generation_ppl_test - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py @@ -2605,7 +2603,6 @@ steps: - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py - - vllm/transformers_utils/configs/qwen3_5.py - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - vllm/v1/kv_cache_interface.py @@ -3914,8 +3911,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py @@ -3957,8 +3952,6 @@ steps: - vllm/model_executor/models/qwen3_next_mtp.py - vllm/model_executor/models/nemotron_h.py - vllm/model_executor/models/nemotron_h_mtp.py - - vllm/transformers_utils/configs/qwen3_next.py - - vllm/transformers_utils/configs/nemotron_h.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/layers/fused_moe/ @@ -4156,7 +4149,6 @@ steps: - tests/evals/gsm8k/ - tests/utils.py - vllm/distributed/ - - vllm/transformers_utils/configs/diffusion_gemma.py - vllm/v1/attention/ops/ - tests/evals/gsm8k/configs/models-small-tp-mi355.txt commands: @@ -5048,8 +5040,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d1a115a1e1b0..49eea0b78f7c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -208,7 +208,6 @@ steps: - tests/evals/gsm8k/ - tests/utils.py - vllm/distributed/ - - vllm/transformers_utils/configs/diffusion_gemma.py - vllm/v1/attention/ops/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt @@ -417,8 +416,6 @@ steps: - vllm/model_executor/models/qwen3_next_mtp.py - vllm/model_executor/models/nemotron_h.py - vllm/model_executor/models/nemotron_h_mtp.py - - vllm/transformers_utils/configs/qwen3_next.py - - vllm/transformers_utils/configs/nemotron_h.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/layers/fused_moe/ @@ -445,8 +442,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/qwen3_next_mtp.py - vllm/third_party/flash_linear_attention/ops/ @@ -463,8 +458,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index e1490139c052..52ff1d65b3ee 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -296,7 +296,6 @@ steps: - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py - - vllm/transformers_utils/configs/qwen3_5.py - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - vllm/v1/kv_cache_interface.py diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index fa8013ec54e8..f123f2c57b13 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -184,8 +184,6 @@ steps: source_file_dependencies: - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - - vllm/transformers_utils/configs/qwen3_5.py - - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen2.py - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py diff --git a/.claude/skills/transformers-upgrade b/.claude/skills/transformers-upgrade new file mode 120000 index 000000000000..d71bb1472413 --- /dev/null +++ b/.claude/skills/transformers-upgrade @@ -0,0 +1 @@ +../../.agents/skills/transformers-upgrade \ No newline at end of file diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index ae4c7220f24c..8ec16ee7f9fb 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -10,13 +10,11 @@ from unittest.mock import Mock import pytest -from transformers import PreTrainedConfig +from transformers import DiffusionGemmaConfig, PreTrainedConfig from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.config.model_arch import ModelArchitectureConfig -from vllm.transformers_utils.configs.diffusion_gemma import DiffusionGemmaConfig -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.transformers_utils.model_arch_config_convertor import ( MODEL_ARCH_CONFIG_CONVERTORS, Gemma4ModelArchConfigConvertor, @@ -318,11 +316,7 @@ def _gemma4_text_config(**overrides) -> Gemma4TextConfig: def test_gemma4_head_dims_vary_by_layer_type(): - """Gemma4's full attention layers are wider than its sliding ones. - - Transformers >= 5.15.0 says so in the config; this exercises the convertor - building the same per-layer view from the flat attributes used before that. - """ + """Gemma4's full attention layers are wider than its sliding ones.""" text_config = _gemma4_text_config( num_global_key_value_heads=8, attention_k_eq_v=True ) @@ -332,9 +326,9 @@ def test_gemma4_head_dims_vary_by_layer_type(): assert (arch.head_size, arch.total_num_kv_heads) == (32, 8) assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32] assert [arch[i].total_num_kv_heads for i in range(6)] == [4] * 5 + [8] - # The model files resolve each layer through the same helper, so the KV cache - # vLLM allocates and the projections the model builds cannot disagree. - assert [gemma4_layer_config(text_config, i).head_dim for i in range(6)] == [ + # The model files read the same per-layer configs, so the KV cache vLLM + # allocates and the projections the model builds cannot disagree. + assert [text_config.per_layer_config[i].head_dim for i in range(6)] == [ arch[i].head_size for i in range(6) ] diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 0919431abac6..48b9dd509977 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -14,7 +14,6 @@ import vllm.envs as envs from tests.utils import RemoteOpenAIServer from vllm.config import ModelConfig -from vllm.config.utils import getattr_iter from vllm.v1.engine.detokenizer import check_stop_strings MODEL_NAME = "Qwen/Qwen3-0.6B" @@ -402,14 +401,7 @@ async def test_same_response_as_chat_completions(client, tokenizer, messages): # Post-EOS generation is undefined and may differ eos_tokens = { tokenizer.eos_token_id, - *getattr_iter( - tokenizer, - [ - "extra_special_tokens_ids", # Transformers v5 - "additional_special_tokens_ids", # Transformers v4 - ], - [], - ), + *tokenizer.extra_special_tokens_ids, } # Find first EOS in generated tokens eos_pos = None diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 78aa563925ff..0fcc2540e252 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -12,6 +12,7 @@ import pytest import torch from torch import nn +from transformers import Qwen3_5MoeTextConfig import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config @@ -34,7 +35,6 @@ from vllm.models.minimax_m3.amd import model as minimax_m3_model from vllm.platforms import current_platform from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig pytestmark = pytest.mark.skipif( current_platform.is_xpu(), diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 32a9adda8822..9ea2cc43cd6e 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -9,14 +9,12 @@ from pathlib import PosixPath import pytest -from packaging.version import Version from transformers import ( AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, AutoModelForTextToWaveform, ) -from transformers import __version__ as TRANSFORMERS_VERSION from vllm.platforms import current_platform from vllm.utils.func_utils import identity @@ -818,8 +816,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): patch_hf_runner=model_utils.paddleocr_vl_patch_hf_runner, image_size_factors=[(0.25,)], marks=[ - pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), + pytest.mark.skip( reason="Model's custom code uses ROPE_INIT_FUNCTIONS" "['default'] which was removed in transformers v5", ), diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index 26a3a592ca1c..d45912a4b141 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -249,9 +249,9 @@ def process(): "BAAI/Emu3-Chat-hf", " and more text", marks=pytest.mark.xfail( - reason="Emu3Processor prepends its BOS token only when images are " - "passed, so the unexpanded prompt vLLM tokenizes never " - "gets one. Fixed by huggingface/transformers#47924, unreleased.", + reason="Emu3Processor prepends the BOS token itself because the " + "tokenizer doesn't, so the unexpanded prompt vLLM tokenizes " + "never gets one.", strict=False, ), ), diff --git a/tests/models/registry.py b/tests/models/registry.py index 4255d439cbdc..31c3c5433aec 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -775,10 +775,7 @@ def check_available_online( # [Decoder-only] "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"), "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo( - "nvidia/audio-flamingo-3-hf", - transformers_version_reason={ - "vllm": "Needs https://github.com/huggingface/transformers/pull/43538" - }, + "nvidia/audio-flamingo-3-hf" ), "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"), "BailingMoeV3VLForConditionalGeneration": _HfExamplesInfo( @@ -817,12 +814,7 @@ def check_available_online( "nvidia/Cosmos3-Edge", max_model_len=4096, use_original_num_layers=True, - hf_overrides={ - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "*-", - } - }, + hf_overrides={"text_config": {"layers_block_type": ["full_attention", "mlp"]}}, ), "DeepseekVLV2ForCausalLM": _HfExamplesInfo( "deepseek-ai/deepseek-vl2-tiny", diff --git a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py index f5845efb0466..92803d93151d 100644 --- a/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py +++ b/tests/plugins_tests/bitsandbytes/test_bitsandbytes.py @@ -6,9 +6,7 @@ from unittest.mock import MagicMock, patch import pytest -from packaging.version import Version from transformers import BitsAndBytesConfig -from transformers import __version__ as TRANSFORMERS_VERSION from vllm_bnb_plugin import bitsandbytes_loader as bnb from vllm.platforms import current_platform @@ -174,8 +172,7 @@ def test_load_pp_4bit_bnb_model( ) -@pytest.mark.skipif( - Version(TRANSFORMERS_VERSION) >= Version("5.0.0"), +@pytest.mark.skip( reason="Need to add support for quantizing MoE experts with bnb in " "transformers v5. See https://github.com/bitsandbytes-foundation/" "bitsandbytes/issues/1849", diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index 312b0c03a921..b9d9c732f333 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -6,7 +6,7 @@ from unittest.mock import patch import pytest -from transformers import AutoConfig +from transformers import AutoConfig, Qwen3_5MoeConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING from vllm.tokenizers import TokenizerLike @@ -18,7 +18,6 @@ get_tokenizer, resolve_tokenizer_args, ) -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeConfig class TestTokenizer(TokenizerLike): diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 1e7883864027..516ef132365b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -8,6 +8,7 @@ import torch from einops import rearrange from torch import nn +from transformers import Qwen3NextConfig from vllm import _custom_ops as ops from vllm import envs @@ -58,7 +59,6 @@ ) from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index d8ceb5083c2f..83859946a255 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -29,6 +29,7 @@ import torch from torch import nn +from transformers import AXK1Config from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile @@ -78,7 +79,6 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.AXK1 import AXK1Config from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index 4e31b99944f1..c9e5dcfdcdd1 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/vllm/model_executor/models/aimv2.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from transformers import Aimv2VisionConfig from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.utils import divide @@ -21,20 +22,19 @@ ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper -from vllm.transformers_utils.configs.ovis import AIMv2Config class AIMv2SwiGLUFFN(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): super().__init__() hidden_features = config.intermediate_size in_features = config.hidden_size - bias = config.use_bias + bias = config.mlp_bias self.fc13 = MergedColumnParallelLinear( in_features, @@ -60,7 +60,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2PatchEmbed(nn.Module): - def __init__(self, config: AIMv2Config): + def __init__(self, config: Aimv2VisionConfig): super().__init__() self.proj = Conv2dLayer( config.num_channels, @@ -77,7 +77,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2ViTPreprocessor(nn.Module): - def __init__(self, config: AIMv2Config): + def __init__(self, config: Aimv2VisionConfig): super().__init__() num_patches = (config.image_size // config.patch_size) ** 2 @@ -95,7 +95,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Attention(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): @@ -124,7 +124,7 @@ def __init__( self.proj = RowParallelLinear( input_size=self.embed_dim, output_size=self.embed_dim, - bias=config.use_bias, + bias=config.mlp_bias, quant_config=quant_config, prefix=f"{prefix}.proj", ) @@ -151,7 +151,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Block(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, prefix: str, ): @@ -174,7 +174,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class AIMv2Transformer(nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, @@ -214,7 +214,7 @@ class AIMv2Model(torch.nn.Module): def __init__( self, - config: AIMv2Config, + config: Aimv2VisionConfig, quant_config: QuantizationConfig | None, *, require_post_norm: bool | None = None, diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 7c17edc2a1db..3a4d8ef039f7 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -403,20 +403,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vocab_size, config.hidden_size ) - # Decoder layers read per-layer MLP layout from config.mlp_layer_types - # (dense MLP vs MoE) and use it for weight loading. Transformers >=5.10 - # populates this field; older versions only expose first_k_dense_replace. - # Normalize here so layer construction below sees a consistent layout. - if getattr(config, "mlp_layer_types", None) is None: - first_k_dense_replace = getattr(config, "first_k_dense_replace", None) - n = config.num_hidden_layers - if first_k_dense_replace is not None: - config.mlp_layer_types = ["dense"] * first_k_dense_replace + [ - "sparse" - ] * (n - first_k_dense_replace) - else: - config.mlp_layer_types = ["sparse"] * n - self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: Cohere2MoeDecoderLayer( diff --git a/vllm/model_executor/models/cosmos3.py b/vllm/model_executor/models/cosmos3.py index 7e4e82745772..5023266351e5 100644 --- a/vllm/model_executor/models/cosmos3.py +++ b/vllm/model_executor/models/cosmos3.py @@ -2,11 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import regex +from transformers import Cosmos3OmniConfig -from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) from vllm.model_executor.models.utils import WeightsMapper +from vllm.multimodal import MULTIMODAL_REGISTRY +class Cosmos3ProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Cosmos3OmniConfig: + return self.ctx.get_hf_config(Cosmos3OmniConfig) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Cosmos3ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration): # Cosmos3 unified checkpoints store a Qwen3-VL understanding tower # alongside a generation tower in a flat key layout. This mapper drops diff --git a/vllm/model_executor/models/cosmos3_edge.py b/vllm/model_executor/models/cosmos3_edge.py index ae65e25c901d..c10fafd0b192 100644 --- a/vllm/model_executor/models/cosmos3_edge.py +++ b/vllm/model_executor/models/cosmos3_edge.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn -from transformers import ProcessorMixin +from transformers import NemotronHConfig, ProcessorMixin from vllm.config import VllmConfig from vllm.model_executor.layers.layernorm import RMSNorm @@ -20,7 +20,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig from vllm.utils.torch_utils import async_tensor_h2d from .interfaces import ( diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 1b27df68abb0..7c64980b57bb 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Sequence from types import SimpleNamespace from typing import Any @@ -25,7 +25,7 @@ import torch._dynamo from torch import nn from torch.nn import functional as F -from transformers import AutoModel +from transformers import AutoModel, DiffusionGemmaConfig from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode @@ -113,26 +113,13 @@ def forward( class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo): - """Processing info for DiffusionGemma. + """Processing info for DiffusionGemma, which has no audio tower.""" - Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig`` - (which inherits from ``PreTrainedConfig``, not ``Gemma4Config``). - Supports image and video modalities. - """ - - def get_hf_config(self): - # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we - # accept any PreTrainedConfig here. - return self.ctx.get_hf_config() - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - # DiffusionGemma supports image and video inputs. - return {"image": None, "video": None} + def get_hf_config(self) -> DiffusionGemmaConfig: + return self.ctx.get_hf_config(DiffusionGemmaConfig) - def get_mm_max_tokens_per_item( - self, seq_len: int, mm_counts: Mapping[str, int] - ) -> Mapping[str, int] | None: - return super().get_mm_max_tokens_per_item(seq_len, mm_counts) + def has_audio(self) -> bool: + return False @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py index 150128505d5e..0f92aeaa3716 100644 --- a/vllm/model_executor/models/embedding_gemma2.py +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -36,10 +36,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) -from vllm.model_executor.models.gemma4 import ( - Gemma4MLP, - gemma4_layer_config, -) +from vllm.model_executor.models.gemma4 import Gemma4MLP from vllm.model_executor.models.gemma4_mm import ( _SUPPORTED_SOFT_TOKENS, Gemma4DummyInputsBuilder, @@ -244,7 +241,7 @@ def __init__( prefix: str, ): super().__init__() - lc = gemma4_layer_config(config, idx) + lc = config.per_layer_config[idx] self.self_attn = EmbeddingGemma2Attention( config, lc, diff --git a/vllm/model_executor/models/gemma3.py b/vllm/model_executor/models/gemma3.py index e886c2ccdd07..4646ebeca462 100644 --- a/vllm/model_executor/models/gemma3.py +++ b/vllm/model_executor/models/gemma3.py @@ -154,20 +154,7 @@ def __init__( self.is_sliding = layer_type == "sliding_attention" sliding_window = config.sliding_window if self.is_sliding else None - # Initialize the rotary embedding. - if layer_type in config.rope_parameters: - # Transformers v5 rope config. - rope_parameters = config.rope_parameters[layer_type] - else: - # Transformers v4 rope config. - # Global attention. Use the values in config.json. - rope_parameters = config.rope_parameters - # Local attention. Override the values in config.json. - if self.is_sliding: - rope_parameters = dict( - rope_type="default", rope_theta=config.rope_local_base_freq - ) - + rope_parameters = config.rope_parameters[layer_type] self.rotary_emb = get_rope( self.head_dim, max_position=max_position_embeddings, diff --git a/vllm/model_executor/models/gemma3n.py b/vllm/model_executor/models/gemma3n.py index 3c9df39e674c..9e7c26bf59b9 100644 --- a/vllm/model_executor/models/gemma3n.py +++ b/vllm/model_executor/models/gemma3n.py @@ -350,17 +350,7 @@ def __init__( is_sliding = layer_type == "sliding_attention" self.sliding_window = config.sliding_window if is_sliding else None - # Initialize the rotary embedding. - if layer_type in config.rope_parameters: - # Transformers v5 rope config. - rope_parameters = config.rope_parameters[layer_type] - else: - # Transformers v4 rope config. - # Global attention. Use the values in config.json. - rope_parameters = config.rope_parameters.copy() - # Local attention. Override the values in config.json. - if is_sliding: - rope_parameters["rope_theta"] = config.rope_local_base_freq + rope_parameters = config.rope_parameters[layer_type] kv_sharing_target_layer_name = None if self.is_kv_shared_layer: diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 17b8251e9a45..d1498f16ddb4 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -59,7 +59,6 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata @@ -576,7 +575,7 @@ def __init__( self.layer_idx = layer_idx # Gemma4 uses different head dimensions for sliding vs full attention - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads @@ -628,9 +627,7 @@ def __init__( ) # MoE (Mixture of Experts) — router + expert block parallel to MLP - self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr( - config, "use_second_mlp_block", False - ) + self.enable_moe_block = getattr(config, "num_experts", None) is not None self.router: Gemma4Router | None self.experts: MoERunner | None self.post_feedforward_layernorm_1: RMSNorm | None diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py index 33b1a1aada23..13c868e59b38 100644 --- a/vllm/model_executor/models/gemma4_dspark.py +++ b/vllm/model_executor/models/gemma4_dspark.py @@ -20,7 +20,6 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal @@ -48,7 +47,7 @@ def __init__( config, "attention_k_eq_v", False ) - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads super().__init__( diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 796dfbb6e8fa..204db6ce8d51 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -222,12 +222,11 @@ def get_hf_processor(self, **kwargs: object) -> Gemma4Processor: **kwargs, ) + def has_audio(self) -> bool: + return self.get_hf_config().audio_config is not None + def validate_num_items(self, modality: str, num_items: int) -> None: - if ( - modality == "audio" - and num_items > 0 - and self.get_hf_config().audio_config is None - ): + if modality == "audio" and num_items > 0 and not self.has_audio(): model_config = self.ctx.model_config model = get_served_model_name( model_config.model, model_config.served_model_name @@ -243,7 +242,7 @@ def validate_num_items(self, modality: str, num_items: int) -> None: def get_supported_mm_limits(self) -> Mapping[str, int | None]: limits: dict[str, int | None] = {"image": None} - if self.get_hf_config().audio_config is not None: + if self.has_audio(): limits["audio"] = None limits["video"] = None return limits @@ -261,7 +260,7 @@ def get_mm_max_tokens_per_item( if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: tokens_per_image = val tokens: dict[str, int] = {"image": tokens_per_image} - if config.audio_config is not None: + if self.has_audio(): # Audio max tokens from the processor's audio_seq_length. processor = self.get_hf_processor() tokens["audio"] = processor.audio_seq_length @@ -275,9 +274,8 @@ def get_mm_max_tokens_per_item( return tokens def get_data_parser(self) -> MultiModalDataParser: - config = self.get_hf_config() kwargs: dict[str, Any] = {"video_needs_metadata": True} - if getattr(config, "audio_config", None) is not None: + if self.has_audio(): processor = self.get_hf_processor() kwargs["target_sr"] = processor.feature_extractor.sampling_rate return MultiModalDataParser(**kwargs) diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 01f226ff07a5..f9931ddafc0a 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -45,7 +45,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from .gemma4 import Gemma4MLP, _get_text_config from .utils import ( @@ -265,7 +264,7 @@ def __init__( self.hidden_size = config.hidden_size layer_idx = extract_layer_index(prefix) - layer_config = gemma4_layer_config(config, layer_idx) + layer_config = config.per_layer_config[layer_idx] head_dim = layer_config.head_dim num_kv_heads = layer_config.num_key_value_heads diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index b62a86543414..d738c665c2ca 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -493,13 +493,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "vision_tower"), ) - # image_newline parameter - if config.use_image_newline_parameter: - self.image_newline = nn.Parameter( - torch.empty(config.text_config.hidden_size) - ) - else: - self.image_newline = None + self.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size) + ) cache_config = vllm_config.cache_config @@ -517,20 +513,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) # Spatial projectors: 4 offset groups - self.spatial_projectors = None - if config.use_spatial_sampling: - self.spatial_projectors = nn.ModuleList( - [ - WindowQFormerDownsampler( - config, - quant_config=quant_config, - cache_config=cache_config, - spatial_offset=i, - prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), - ) - for i in range(4) - ] - ) + self.spatial_projectors = nn.ModuleList( + [ + WindowQFormerDownsampler( + config, + quant_config=quant_config, + cache_config=cache_config, + spatial_offset=i, + prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), + ) + for i in range(4) + ] + ) # ----- Language model (marked as LM) ----- with self._mark_language_model(vllm_config): @@ -545,12 +539,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Store config values we need self._deepstack_layer_map = config.deepstack_layer_map # [[-19, 9], ...] - self._use_spatial_sampling = getattr(config, "use_spatial_sampling", False) - self._spatial_vision_layer = getattr(config, "spatial_vision_layer", -1) - self._spatial_target_layers = getattr(config, "spatial_target_layers", []) - self._vision_feature_select_strategy = getattr( - config, "vision_feature_select_strategy", "full" - ) + self._spatial_vision_layer = config.spatial_vision_layer + self._spatial_target_layers = config.spatial_target_layers + self._vision_feature_select_strategy = config.vision_feature_select_strategy self._downsample_rate = Fraction(config.downsample_rate) # Ordered list of LLM layer indices for each deepstack level. @@ -558,7 +549,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # (before any embed_multimodal call). self._ds_layer_indices: list[int] = [ llm_layer for _, llm_layer in config.deepstack_layer_map - ] + list(getattr(config, "spatial_target_layers", [])) + ] + list(config.spatial_target_layers) # Share ds_layer_indices with the LLM causal model so # make_empty_intermediate_tensors includes the correct keys @@ -571,7 +562,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # the updated values written just before each prefill. # Shape: (max_num_batched_tokens, lm_hidden_size) per level. n_layerwise = len(config.deepstack_layer_map) - n_spatial = len(getattr(config, "spatial_target_layers", [])) + n_spatial = len(config.spatial_target_layers) num_ds_levels = n_layerwise + n_spatial lm_hidden = config.text_config.hidden_size max_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -644,26 +635,24 @@ def _pack_and_unpad_image_features( ) image_feature = unpad_image(image_feature, image_sizes[image_idx]) - if self.image_newline is not None: - image_feature = torch.cat( - ( - image_feature, - self.image_newline[:, None, None] - .expand(*image_feature.shape[:-1], 1) - .to(image_feature.device, image_feature.dtype), - ), - dim=-1, - ) + image_feature = torch.cat( + ( + image_feature, + self.image_newline[:, None, None] + .expand(*image_feature.shape[:-1], 1) + .to(image_feature.device, image_feature.dtype), + ), + dim=-1, + ) image_feature = image_feature.flatten(1, 2).transpose(0, 1) image_feature = torch.cat((base_image_feature, image_feature), dim=0) else: image_feature = image_feature[0] - if self.image_newline is not None: - image_feature = torch.cat( - (image_feature, self.image_newline[None].to(image_feature)), - dim=0, - ) + image_feature = torch.cat( + (image_feature, self.image_newline[None].to(image_feature)), + dim=0, + ) new_image_features.append(image_feature) @@ -718,16 +707,15 @@ def _get_all_layer_features( ) levels.append((llm_layer, per_image)) - if self._use_spatial_sampling and self.spatial_projectors is not None: - spatial_hidden = all_hidden_states[self._spatial_vision_layer] - if select_strategy == "default": - spatial_hidden = spatial_hidden[:, 1:] - for group_idx, llm_layer in enumerate(self._spatial_target_layers): - projected = self.spatial_projectors[group_idx](spatial_hidden) - per_image = self._pack_and_unpad_image_features( - torch.split(projected, image_num_patches, dim=0), image_sizes - ) - levels.append((llm_layer, per_image)) + spatial_hidden = all_hidden_states[self._spatial_vision_layer] + if select_strategy == "default": + spatial_hidden = spatial_hidden[:, 1:] + for group_idx, llm_layer in enumerate(self._spatial_target_layers): + projected = self.spatial_projectors[group_idx](spatial_hidden) + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes + ) + levels.append((llm_layer, per_image)) llm_layer_indices = [llm_layer for llm_layer, _ in levels] num_images = len(image_sizes) diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 3ae85cfaac69..2380bba10a65 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -158,9 +158,7 @@ def _get_prompt_updates( feature_extractor = processor.audio_processor vocab = tokenizer.get_vocab() - # Use getattr with default to be compatible with transformers<4.48 - audio_token = getattr(processor, "audio_token", "<|audio|>") - audio_token_id = vocab[audio_token] + audio_token_id = vocab[processor.audio_token] def get_replacement(item_idx: int): audios = mm_items.get_items("audio", AudioProcessorItems) @@ -221,8 +219,7 @@ def get_dummy_mm_data( def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_audios = mm_counts.get("audio", 0) hf_processor = self.info.get_hf_processor() - audio_token = getattr(hf_processor, "audio_token", "<|audio|>") - return audio_token * num_audios + return hf_processor.audio_token * num_audios ### QFormer Projector diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 84a981e6bba2..979d272c13d1 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -320,10 +320,6 @@ def forward( ALL_DECODER_LAYER_TYPES = { - # Transformers < 5.13.0 - "attention": GraniteMoeHybridAttentionDecoderLayer, - "mamba": GraniteMoeHybridMambaDecoderLayer, - # Transformers >= 5.13.0 "full_attention": GraniteMoeHybridAttentionDecoderLayer, "linear_attention": GraniteMoeHybridMambaDecoderLayer, } diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 2e24cf65e8f3..80fc16e23202 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -31,7 +31,7 @@ import torch from torch import nn -from transformers import PreTrainedConfig +from transformers import HYV3Config, PreTrainedConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config @@ -67,7 +67,6 @@ maybe_remap_kv_scale_name, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.hy_v3 import HYV3Config from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/hyperclovax.py b/vllm/model_executor/models/hyperclovax.py index 66c3b4a270a3..e86260bca96f 100644 --- a/vllm/model_executor/models/hyperclovax.py +++ b/vllm/model_executor/models/hyperclovax.py @@ -31,6 +31,7 @@ import torch from torch import nn +from transformers import HyperCLOVAXConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -51,7 +52,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.hyperclovax import HyperCLOVAXConfig from .interfaces import SupportsLoRA, SupportsPP from .utils import ( @@ -137,9 +137,7 @@ def __init__( # the KV heads across multiple tensor parallel GPUs. assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) + self.head_dim = config.head_dim self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = config.attention_multiplier @@ -166,7 +164,7 @@ def __init__( self.head_dim, max_position=max_position_embeddings, is_neox_style=True, - rope_parameters=getattr(config, "rope_parameters", None), + rope_parameters=config.rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, ) @@ -217,18 +215,15 @@ def __init__( "dual_chunk_attention_config", None, ) - attention_bias = getattr(config, "attention_bias", False) self.self_attn = HyperCLOVAXAttention( config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, - num_kv_heads=getattr( - config, "num_key_value_heads", config.num_attention_heads - ), + num_kv_heads=config.num_key_value_heads, max_position_embeddings=max_position_embeddings, quant_config=quant_config, - bias=attention_bias, + bias=config.attention_bias, cache_config=cache_config, prefix=f"{prefix}.self_attn", dual_chunk_attention_config=dual_chunk_attention_config, @@ -238,7 +233,7 @@ def __init__( intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, - bias=getattr(config, "mlp_bias", False), + bias=config.mlp_bias, prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 8cd354b903af..e555b6368d4a 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -11,7 +11,7 @@ import torch from torch import nn -from transformers import BatchFeature +from transformers import BatchFeature, Kimi_K25Config from vllm.config import VllmConfig from vllm.config.multimodal import MultiModalDummyOptions @@ -63,7 +63,6 @@ ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config from vllm.transformers_utils.processor import cached_get_image_processor from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( @@ -135,7 +134,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: # Resolve token ID from the tokenizer because transformers v5 # may remap token IDs vs config.json. - config_token_id = hf_config.media_placeholder_token_id + config_token_id = hf_config.image_token_id resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") unk_token_id = getattr(tokenizer, "unk_token_id", None) is_valid_resolved = isinstance(resolved_token_id, int) and ( @@ -143,7 +142,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: ) if is_valid_resolved and resolved_token_id != config_token_id: logger.warning_once( - "Kimi-K2.5 config.media_placeholder_token_id (%d) disagrees " + "Kimi-K2.5 config.image_token_id (%d) disagrees " "with tokenizer mapping for <|media_pad|> (%d). " "Using tokenizer value.", config_token_id, @@ -151,7 +150,7 @@ def __init__(self, ctx: InputProcessingContext) -> None: ) media_token_id = resolved_token_id # Patch config so downstream code also sees the correct ID. - hf_config.media_placeholder_token_id = resolved_token_id + hf_config.image_token_id = resolved_token_id else: media_token_id = config_token_id @@ -170,7 +169,7 @@ def get_hf_processor(self, **kwargs: object): return self.hf_processor def get_hf_config(self): - return self.ctx.get_hf_config(KimiK25Config) + return self.ctx.get_hf_config(Kimi_K25Config) def get_supported_mm_limits(self) -> Mapping[str, int | None]: # None means unlimited @@ -338,7 +337,7 @@ def __init__( ) -> None: super().__init__() model_config = vllm_config.model_config - config: KimiK25Config = model_config.hf_config + config: Kimi_K25Config = model_config.hf_config self.config = config quant_config = vllm_config.quant_config @@ -347,7 +346,6 @@ def __init__( ) self.hidden_size = config.text_config.hidden_size self.device = current_platform.current_device() - # Build vision tower directly with KimiK25VisionConfig with self._mark_tower_model(vllm_config, "vision_chunk"): self.vision_tower = MoonViT3dPretrainedModel( config.vision_config, @@ -364,6 +362,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), @@ -383,7 +382,7 @@ def __init__( self.make_empty_intermediate_tensors = ( self.language_model.make_empty_intermediate_tensors ) - self.media_placeholder: int = self.config.media_placeholder_token_id + self.media_placeholder: int = self.config.image_token_id def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig | None): if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index 44e8498d212b..bb1be1ba32f8 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -15,6 +15,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from transformers import Kimi_K25VisionConfig from transformers.activations import GELUActivation from vllm.distributed import divide, get_tensor_model_parallel_world_size @@ -40,13 +41,21 @@ run_dp_sharded_mrope_vision_model, ) from vllm.platforms import current_platform -from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig +from vllm.transformers_utils.configs.kimi_k3 import KimiK3VisionConfig from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import async_tensor_h2d logger = init_logger(__name__) +def _get_pos_emb_size( + config: Kimi_K25VisionConfig | KimiK3VisionConfig, dim: str +) -> int: + # Legacy checkpoints and Kimi-K3 use `init_pos_emb_*` + legacy = getattr(config, f"init_pos_emb_{dim}", None) + return legacy if legacy is not None else getattr(config, f"pos_emb_{dim}") + + def _apply_rope_input_validation(x, freqs_cis): assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) @@ -698,14 +707,11 @@ def tpool_patch_merger_packed( class MoonViT3dPretrainedModel(nn.Module): - """Main vision tower model. - - Uses KimiK25VisionConfig directly from transformers_utils/configs/kimi_k25.py. - """ + """Main vision tower model.""" def __init__( self, - config: KimiK25VisionConfig, + config: Kimi_K25VisionConfig | KimiK3VisionConfig, quant_config: QuantizationConfig | None = None, input_norm: nn.Module | None = None, prefix: str = "", @@ -713,17 +719,18 @@ def __init__( super().__init__() config = deepcopy(config) self.config = config # Required for run_dp_sharded_mrope_vision_model - self.merge_kernel_size = config.merge_kernel_size + merge_h, merge_w = config.merge_kernel_size + self.merge_kernel_size = (merge_h, merge_w) self.patch_size = config.patch_size - self.merge_type = config.merge_type + self.merge_type = getattr(config, "merge_type", "sd2_tpool") self.patch_embed = MoonVision3dPatchEmbed( out_dim=config.hidden_size, patch_size=config.patch_size, - pos_emb_height=config.init_pos_emb_height, - pos_emb_width=config.init_pos_emb_width, - pos_emb_time=config.init_pos_emb_time, - pos_emb_type=config.pos_emb_type, + pos_emb_height=_get_pos_emb_size(config, "height"), + pos_emb_width=_get_pos_emb_size(config, "width"), + pos_emb_time=_get_pos_emb_size(config, "time"), + pos_emb_type=getattr(config, "pos_emb_type", "divided_fixed"), patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), pos_emb_interpolation_mode=getattr( config, "pos_emb_interpolation_mode", "bicubic" @@ -740,14 +747,14 @@ def __init__( "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), "mlp_dim": config.intermediate_size, "activation": get_act_fn( - getattr(config, "activation_func", "gelu_pytorch_tanh") + getattr(config, "activation_func", None) or config.hidden_act ), "attn_bias": getattr(config, "attn_bias", True), "norm_type": getattr(config, "norm_type", "layernorm"), "mlp_type": getattr(config, "mlp_type", "mlp2"), "linear_bias": getattr(config, "linear_bias", True), }, - video_attn_type=config.video_attn_type, + video_attn_type=getattr(config, "video_attn_type", "spatial_temporal"), quant_config=quant_config, prefix=maybe_prefix(prefix, "encoder"), ) @@ -886,7 +893,8 @@ class KimiK25MultiModalProjector(nn.Module): def __init__( self, - config: KimiK25VisionConfig, + config: Kimi_K25VisionConfig | KimiK3VisionConfig, + out_hidden_size: int, use_data_parallel: bool = False, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -909,13 +917,13 @@ def __init__( ) self.linear_2 = ReplicatedLinear( self.hidden_size, - getattr(config, "text_hidden_size", config.mm_hidden_size), + out_hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.linear_2", ) self.post_norm = torch.nn.RMSNorm( - getattr(config, "text_hidden_size", config.mm_hidden_size), + out_hidden_size, eps=config.projector_ln_eps, ) self.act = GELUActivation() @@ -931,7 +939,7 @@ def __init__( ) self.linear_2 = ReplicatedLinear( self.hidden_size, - config.mm_hidden_size, + out_hidden_size, bias=True, quant_config=quant_config, prefix=f"{prefix}.linear_2", diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index f6e819c2639e..66c9520bc935 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/vllm/model_executor/models/lfm2_moe.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +from transformers import Lfm2MoeConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config @@ -40,7 +41,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig from .interfaces import ( HasInnerState, diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 3a8abee67ec3..0c1909445b48 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -94,28 +94,12 @@ def __init__( prefix=f"{prefix}.Wqkv", ) - if layer_types := getattr(config, "layer_types", None): - # Transformers v5 - layer_type = layer_types[layer_id] - rope_parameters = config.rope_parameters[layer_type] - sliding_window: int | None = None - if layer_type == "sliding_attention": - # Treats the local attention boundary as inclusive - sliding_window = config.sliding_window + 1 - else: - # Transformers v4 - sliding_window = None - if layer_id % config.global_attn_every_n_layers != 0: - # ModernBertConfig does not expose sliding_window - sliding_window = config.local_attention // 2 - rope_theta = ( - config.local_rope_theta - if config.local_rope_theta is not None - else config.global_rope_theta - ) - else: - rope_theta = config.global_rope_theta - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} + layer_type = config.layer_types[layer_id] + rope_parameters = config.rope_parameters[layer_type] + sliding_window: int | None = None + if layer_type == "sliding_attention": + # Treats the local attention boundary as inclusive + sliding_window = config.sliding_window + 1 self.rotary_emb = get_rope( head_size=self.head_dim, diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index a271432813a3..55218576969d 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -29,6 +29,7 @@ import torch from torch import nn +from transformers import NemotronConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -49,7 +50,6 @@ VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron import NemotronConfig from .interfaces import SupportsLoRA, SupportsPP from .utils import ( diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e477c65302b6..ae653950ac48 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -23,6 +23,7 @@ import torch from torch import nn +from transformers import NemotronHConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import ( @@ -87,7 +88,6 @@ sequence_parallel_chunk, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig class NemotronHMLP(nn.Module): diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index af93cc62ae25..1e393296bc97 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from transformers import NemotronHConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig @@ -33,7 +34,6 @@ maybe_prefix, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig from .interfaces import SupportsPP, SupportsQuant from .nemotron_h import ( diff --git a/vllm/model_executor/models/ovis.py b/vllm/model_executor/models/ovis.py index 7c1668ddabcd..9e7778da1589 100644 --- a/vllm/model_executor/models/ovis.py +++ b/vllm/model_executor/models/ovis.py @@ -125,7 +125,7 @@ def _init_backbone( prefix: str = "", ) -> nn.Module: model_type = config.backbone_config.model_type - if model_type == "aimv2": + if model_type == "aimv2_vision_model": # No post rms_norm in Ovis2's AIMv2 ViT. return AIMv2Model( config=config.backbone_config, diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 32660be5150f..0d110dd58be1 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -29,6 +29,12 @@ import torch from torch import nn +from transformers import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, +) from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -59,11 +65,6 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors from vllm.tokenizers.registry import cached_tokenizer_from_config -from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import ( - Qwen3_5MoeConfig, - Qwen3_5MoeTextConfig, -) from .interfaces import ( HasInnerState, diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index d7fb1d6e09ea..2acac6b5b7c5 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -6,6 +6,7 @@ import torch from torch import nn +from transformers import Qwen3_5MoeTextConfig, Qwen3_5TextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -32,8 +33,6 @@ ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig -from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig from .interfaces import ( MultiModalEmbeddings, diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 9580894e22ee..b0bf0fecd19c 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -7,6 +7,7 @@ import torch from torch import nn +from transformers import Qwen3NextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig @@ -58,7 +59,6 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.v1.attention.backend import AttentionType from .interfaces import ( diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index ef08878eaa7f..4ae8f21840a1 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -6,6 +6,7 @@ import torch from torch import nn +from transformers import Qwen3NextConfig from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -29,7 +30,6 @@ ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from .utils import ( AutoWeightsLoader, diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py index 227f6dba2487..124fca3e337a 100644 --- a/vllm/models/kimi_k3/amd/model.py +++ b/vllm/models/kimi_k3/amd/model.py @@ -117,6 +117,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 22d2c4a1a9a1..eb2ac88aa436 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1825,6 +1825,7 @@ def __init__( self.mm_projector = KimiK25MultiModalProjector( config=config.vision_config, + out_hidden_size=config.text_config.hidden_size, use_data_parallel=self.use_data_parallel, quant_config=self._maybe_ignore_quant_config(quant_config), prefix=maybe_prefix(prefix, "mm_projector"), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c22f6fc3c2ea..e046c301d23a 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -79,8 +79,6 @@ def __getitem__(self, key): _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( - afmoe="AfmoeConfig", - axk1="AXK1Config", bagel="BagelConfig", bailing_moe_v3_vl="BailingMoeV3VLConfig", chatglm="ChatGLMConfig", @@ -89,24 +87,20 @@ def __getitem__(self, key): colqwen3="ColQwen3Config", ops_colqwen3="OpsColQwen3Config", qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig", - cosmos3_omni="Cosmos3Config", cosmos3_edge="Cosmos3EdgeConfig", - diffusion_gemma="DiffusionGemmaConfig", deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", dots3_note="Dots3NoteConfig", k3_dspark="K3DSparkConfig", funaudiochat="FunAudioChatConfig", - granite4_vision="Granite4VisionConfig", - hyperclovax="HyperCLOVAXConfig", - hy_v3="HYV3Config", hy_v4="HYV4Config", + hyperclovax="HyperCLOVAXConfig", # Upstream class, hub remote code is broken isaac="IsaacConfig", kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3 kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", - kimi_k25="KimiK25Config", + kimi_k25="Kimi_K25Config", # Upstream class, hub remote code uses old schema muse_glimmer="MuseGlimmerConfig", muse_glimmer_text="MuseGlimmerTextConfig", muse_glimmer_vision="MuseGlimmerVisionConfig", @@ -123,8 +117,6 @@ def __getitem__(self, key): moss_transcribe_diarize="MossTranscribeDiarizeConfig", eagle="EAGLEConfig", speculators="SpeculatorsConfig", - nemotron="NemotronConfig", - olmo_hybrid="OlmoHybridConfig", openvla="OpenVLAConfig", ovis="OvisConfig", ultravox="UltravoxConfig", @@ -133,12 +125,6 @@ def __getitem__(self, key): step3p5="Step3p5Config", qianfan_ocr="QianfanOCRConfig", qwen3_asr="Qwen3ASRConfig", - qwen3_next="Qwen3NextConfig", - qwen3_5="Qwen3_5Config", - qwen3_5_text="Qwen3_5TextConfig", - qwen3_5_moe="Qwen3_5MoeConfig", - qwen3_5_moe_text="Qwen3_5MoeTextConfig", - lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, **{"deepseek_v41": "DeepseekV41Config"}, inkling_mm_model="InklingMMConfig", @@ -162,14 +148,6 @@ def __getitem__(self, key): # `gemma4_text` ships them in the checkpoint itself. _PATCH_HF_NESTED_ROPE_VALIDATION: set[str] = {"laguna", "gemma4_text"} -# Model types whose checkpoints declare `layer_types` entries that upstream -# transformers has not added to `ALLOWED_LAYER_TYPES` yet, so its strict config -# validation rejects them (e.g. GLM-5.2 `glm_moe_dsa` use -# `deepseek_sparse_attention`). Extend the allowed set for these model types. -_PATCH_HF_ALLOWED_LAYER_TYPES: dict[str, tuple[str, ...]] = { - "glm_moe_dsa": ("deepseek_sparse_attention",), -} - _CONFIG_ATTRS_MAPPING: dict[str, str] = { "llm_config": "text_config", } @@ -314,24 +292,6 @@ def patched_validate_rope(self, *args, **kwargs): ) -def _patch_hf_transformers_allowed_layer_types( - extra_layer_types: tuple[str, ...], -) -> None: - """Extend transformers' ``ALLOWED_LAYER_TYPES`` so its strict config - validation accepts layer types (e.g. ``deepseek_sparse_attention``) that a - checkpoint declares but upstream transformers has not registered yet. - """ - import transformers.configuration_utils as hf_configuration_utils - - missing = tuple( - layer_type - for layer_type in extra_layer_types - if layer_type not in hf_configuration_utils.ALLOWED_LAYER_TYPES - ) - if missing: - hf_configuration_utils.ALLOWED_LAYER_TYPES += missing - - class HFConfigParser(ConfigParserBase): def parse( self, @@ -379,9 +339,6 @@ def parse( if model_type in _PATCH_HF_NESTED_ROPE_VALIDATION: _patch_hf_transformers_nested_rope_validation() - if extra_layer_types := _PATCH_HF_ALLOWED_LAYER_TYPES.get(model_type): - _patch_hf_transformers_allowed_layer_types(extra_layer_types) - if model_type == "vlm": # HyperCLOVAX remote code registers this alias in a bare # `try/except` that fails silently since transformers 5.18 diff --git a/vllm/transformers_utils/configs/AXK1.py b/vllm/transformers_utils/configs/AXK1.py deleted file mode 100644 index 50bc891573c8..000000000000 --- a/vllm/transformers_utils/configs/AXK1.py +++ /dev/null @@ -1,216 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers import PreTrainedConfig - - -class AXK1Config(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`AXK1Model`]. - It is used to instantiate an A.X model according to the specified arguments, - defining the model architecture. Instantiating a configuration with the defaults - will yield a similar configuration to that of the A.X K1. - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control - the model outputs. Read the documentation from [`PreTrainedConfig`] for more - information. - - Args: - vocab_size (`int`, *optional*, defaults to 163840): - Vocabulary size of the A.X K1 model. Defines the number of different - tokens that can be represented by the `inputs_ids` passed when calling - [`AXK1Model`] - hidden_size (`int`, *optional*, defaults to 7168): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 18432): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 2048): - Dimension of the MoE representations. - num_hidden_layers (`int`, *optional*, defaults to 61): - Number of hidden layers in the Transformer decoder. - num_nextn_predict_layers (`int`, *optional*, defaults to 1): - Number of nextn predict layers in the AXK1 Model. - num_attention_heads (`int`, *optional*, defaults to 64): - Number of attention heads for each attention layer in the Transformer - decoder. - n_shared_experts (`int`, *optional*, defaults to 1): - Number of shared experts, None means dense model. - n_routed_experts (`int`, *optional*, defaults to 192): - Number of routed experts, None means dense model. - routed_scaling_factor (`float`, *optional*, defaults to 2.5): - Scaling factor or routed experts. - topk_method (`str`, *optional*, defaults to `noaux_tc`): - Topk method used in routed gate. - n_group (`int`, *optional*, defaults to 8): - Number of groups for routed experts. - topk_group (`int`, *optional*, defaults to 4): - Number of selected groups for each token(for each token, ensuring the - selected experts is only within `topk_group` groups). - num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of selected experts, None means dense model. - moe_layer_freq (`int`, *optional*, defaults to 1): - The frequency of the MoE layer: one expert layer for every - `moe_layer_freq - 1` dense layers. - first_k_dense_replace (`int`, *optional*, defaults to 1): - Number of dense layers in shallow layers - (embed->dense->dense->...->dense->moe->moe...->lm_head). - \--k dense layers--/ - norm_topk_prob (`bool`, *optional*, defaults to True): - Whether to normalize the weights of the routed experts. - scoring_func (`str`, *optional*, defaults to 'sigmoid'): - Method of computing expert weights. - aux_loss_alpha (`float`, *optional*, defaults to 0.0001): - Auxiliary loss weight coefficient. - seq_aux = (`bool`, *optional*, defaults to True): - Whether to compute the auxiliary loss for each individual sample. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to implement - Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, - the model will use Multi Head Attention (MHA), if `num_key_value_heads=1 - the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and - value head should be constructed by meanpooling all the original heads - within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). - If it is not specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 131072): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions - (not used by all models). Only relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 163691): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 163691): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during pretraining. - Please refer to - [this document](https://huggingface.co/docs/transformers/parallelism) - to understand more about it. This value is necessary to ensure exact - reproducibility of the pretraining results. Please refer to - [this issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. - Currently supports two scaling strategies: linear and dynamic. - Their scaling factor must be a float greater than 1. The expected format - is `{"type": strategy name, "factor": scaling factor}`. When using this - flag, don't update `max_position_embeddings` to the expected new maximum. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection - layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - - """ - - model_type = "axk1" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size: int = 163840, - hidden_size: int = 7168, - intermediate_size: int = 18432, - moe_intermediate_size: int = 2048, - num_hidden_layers: int = 61, - num_nextn_predict_layers: int | None = 1, - num_attention_heads: int = 64, - num_key_value_heads: int = 64, - n_shared_experts: int | None = 1, - n_routed_experts: int | None = 192, - ep_size: int | None = 8, ## Ignored - Expert parallel size - routed_scaling_factor: float | None = 2.5, - kv_lora_rank: int | None = 512, - q_lora_rank: int | None = 1536, - qk_rope_head_dim: int | None = 64, - v_head_dim: int | None = 128, - qk_nope_head_dim: int | None = 128, - topk_method: str | None = "noaux_tc", - n_group: int | None = 8, - topk_group: int | None = 4, - num_experts_per_tok: int | None = 8, - moe_layer_freq: int | None = 1, - first_k_dense_replace: int = 1, - norm_topk_prob: bool = True, - scoring_func: str | None = "sigmoid", - aux_loss_alpha: float | None = 0.0001, - seq_aux: float | None = True, - hidden_act: str | None = "silu", - max_position_embeddings: int | None = 131072, - initializer_range: float | None = 0.02, - rms_norm_eps: float = 1e-6, - use_cache: bool | None = True, - pad_token_id: int | None = None, - bos_token_id: int | None = 163691, - eos_token_id: int | None = 163691, - pretraining_tp: int | None = 1, - tie_word_embeddings: bool | None = False, - rope_theta: float | None = 10000.0, - rope_scaling: dict[str, Any] | None = None, - rope_parameters: dict[str, Any] | None = None, - attention_bias: bool | None = False, - attention_dropout: float | None = 0.0, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_nextn_predict_layers = num_nextn_predict_layers - self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.ep_size = ep_size - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.topk_method = topk_method - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.moe_layer_freq = moe_layer_freq - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.scoring_func = scoring_func - self.aux_loss_alpha = aux_loss_alpha - self.seq_aux = seq_aux - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.rope_parameters = rope_parameters - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5bd47fb9482c..336b96c1a316 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -14,8 +14,6 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { - "AfmoeConfig": "vllm.transformers_utils.configs.afmoe", - "AXK1Config": "vllm.transformers_utils.configs.AXK1", "BagelConfig": "vllm.transformers_utils.configs.bagel", "BailingMoeV3TextConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VisionConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", @@ -26,13 +24,10 @@ "ColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3", - "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3", "Cosmos3EdgeConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeProjectorConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeTextConfig": "vllm.transformers_utils.configs.cosmos3_edge", "Cosmos3EdgeVisionConfig": "vllm.transformers_utils.configs.cosmos3_edge", - "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma", - "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", "DeepseekV41Config": "vllm.transformers_utils.configs.deepseek_v41", @@ -42,16 +37,12 @@ "EAGLEConfig": "vllm.transformers_utils.configs.eagle", "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat", "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat", - "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision", - "HYV3Config": "vllm.transformers_utils.configs.hy_v3", "HYV4Config": "vllm.transformers_utils.configs.hy_v4", - "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax", "IsaacConfig": "vllm.transformers_utils.configs.isaac", # RWConfig is for the original tiiuae/falcon-40b(-instruct) and # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the # `FalconConfig` class from the official HuggingFace transformers library. "RWConfig": "vllm.transformers_utils.configs.falcon", - "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe", "MedusaConfig": "vllm.transformers_utils.configs.medusa", "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm", "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3", @@ -67,16 +58,12 @@ "MoonViTConfig": "vllm.transformers_utils.configs.moonvit", "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", - "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", "MuseGlimmerConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerTextConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerVisionConfig": "vllm.transformers_utils.configs.muse_glimmer", "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer", "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", - "NemotronConfig": "vllm.transformers_utils.configs.nemotron", - "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h", - "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid", "OpenVLAConfig": "vllm.transformers_utils.configs.openvla", "OvisConfig": "vllm.transformers_utils.configs.ovis", "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac", @@ -91,22 +78,17 @@ "QianfanOCRConfig": "vllm.transformers_utils.configs.qianfan_ocr", "QianfanOCRVisionConfig": "vllm.transformers_utils.configs.qianfan_ocr", "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr", - "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next", - "Qwen3_5Config": "vllm.transformers_utils.configs.qwen3_5", - "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5", - "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe", - "Qwen3_5MoeTextConfig": "vllm.transformers_utils.configs.qwen3_5_moe", "InklingModelConfig": "vllm.models.inkling.configs", "InklingAudioConfig": "vllm.models.inkling.configs", "InklingVisionConfig": "vllm.models.inkling.configs", "InklingMMConfig": "vllm.models.inkling.configs", - # Special case: DeepseekV3Config is from HuggingFace Transformers + # Upstream Transformers classes registered in _CONFIG_REGISTRY "DeepseekV3Config": "transformers", + "HyperCLOVAXConfig": "transformers", + "Kimi_K25Config": "transformers", } __all__ = [ - "AfmoeConfig", - "AXK1Config", "BagelConfig", "BailingMoeV3TextConfig", "BailingMoeV3VisionConfig", @@ -117,13 +99,10 @@ "ColQwen3Config", "OpsColQwen3Config", "Qwen3VLNemotronEmbedConfig", - "Cosmos3Config", "Cosmos3EdgeConfig", "Cosmos3EdgeProjectorConfig", "Cosmos3EdgeTextConfig", "Cosmos3EdgeVisionConfig", - "DiffusionGemmaConfig", - "DiffusionGemmaTextConfig", "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", @@ -133,13 +112,10 @@ "EAGLEConfig", "FunAudioChatConfig", "FunAudioChatAudioEncoderConfig", - "Granite4VisionConfig", - "HYV3Config", "HYV4Config", "HyperCLOVAXConfig", "IsaacConfig", "RWConfig", - "Lfm2MoeConfig", "MedusaConfig", "MiDashengLMConfig", "MiniMaxM3Config", @@ -153,16 +129,13 @@ "MoonViTConfig", "KimiLinearConfig", "KimiVLConfig", - "KimiK25Config", + "Kimi_K25Config", "MuseGlimmerConfig", "MuseGlimmerTextConfig", "MuseGlimmerVisionConfig", "MuseGlimmerAssistantConfig", "KimiK3Config", "KimiK3VisionConfig", - "NemotronConfig", - "NemotronHConfig", - "OlmoHybridConfig", "OpenVLAConfig", "OvisConfig", "PixelShuffleSiglip2VisionConfig", @@ -177,11 +150,6 @@ "QianfanOCRConfig", "QianfanOCRVisionConfig", "Qwen3ASRConfig", - "Qwen3NextConfig", - "Qwen3_5Config", - "Qwen3_5TextConfig", - "Qwen3_5MoeConfig", - "Qwen3_5MoeTextConfig", "InklingModelConfig", "InklingAudioConfig", "InklingVisionConfig", diff --git a/vllm/transformers_utils/configs/afmoe.py b/vllm/transformers_utils/configs/afmoe.py deleted file mode 100644 index b3d10dc5f274..000000000000 --- a/vllm/transformers_utils/configs/afmoe.py +++ /dev/null @@ -1,87 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from transformers.configuration_utils import PreTrainedConfig - - -class AfmoeConfig(PreTrainedConfig): - model_type = "afmoe" - - def __init__( - self, - vocab_size: int = 200_192, - hidden_size: int = 2048, - intermediate_size: int = 6144, - moe_intermediate_size: int = 1408, - num_hidden_layers: int = 32, - num_dense_layers: int = 1, - num_attention_heads: int = 16, - num_key_value_heads: int | None = None, - head_dim: int = 128, - hidden_act: str = "silu", - max_position_embeddings: int = 131072, - initializer_range: float = 0.02, - rms_norm_eps: float = 1e-5, - use_cache: bool = True, - tie_word_embeddings: bool = False, - rope_parameters: dict | None = None, - rope_scaling: dict | None = None, - num_experts: int = 64, - num_experts_per_tok: int = 6, - num_shared_experts: int = 2, - num_expert_groups: int = 1, - num_limited_groups: int = 1, - score_func: str = "sigmoid", - route_norm: bool = True, - route_scale: float = 1.0, - global_attn_every_n_layers: int = 4, - sliding_window: int = 2048, - layer_types: list[str] | None = None, - attention_dropout: float = 0.0, - mup_enabled: bool = False, - n_group: int = 1, - topk_group: int = 1, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_dense_layers = num_dense_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads or num_attention_heads - self.head_dim = head_dim - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - rope_theta = kwargs.pop("rope_theta", 10000.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.rope_scaling = rope_scaling - - self.moe_intermediate_size = moe_intermediate_size - self.num_experts = num_experts - self.num_experts_per_tok = num_experts_per_tok - self.num_shared_experts = num_shared_experts - self.num_expert_groups = num_expert_groups - self.num_limited_groups = num_limited_groups - self.score_func = score_func - self.route_norm = route_norm - self.route_scale = route_scale - - self.global_attn_every_n_layers = global_attn_every_n_layers - self.sliding_window = sliding_window - self.layer_types = layer_types - self.attention_dropout = attention_dropout - - self.mup_enabled = mup_enabled - self.n_group = n_group - self.topk_group = topk_group - - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["AfmoeConfig"] diff --git a/vllm/transformers_utils/configs/cosmos3.py b/vllm/transformers_utils/configs/cosmos3.py deleted file mode 100644 index 965801888a4b..000000000000 --- a/vllm/transformers_utils/configs/cosmos3.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig - - -class Cosmos3Config(Qwen3VLConfig): - model_type = "cosmos3_omni" diff --git a/vllm/transformers_utils/configs/cosmos3_edge.py b/vllm/transformers_utils/configs/cosmos3_edge.py index 447b64cc99ed..ddf17933fd50 100644 --- a/vllm/transformers_utils/configs/cosmos3_edge.py +++ b/vllm/transformers_utils/configs/cosmos3_edge.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import NemotronHConfig from transformers.configuration_utils import PreTrainedConfig from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig -from .nemotron_h import NemotronHConfig - def _normalize_mrope_parameters(config: NemotronHConfig) -> None: """Normalize the checkpoint's mRoPE configuration for vLLM.""" @@ -31,11 +30,14 @@ def __init__( num_hidden_layers: int = 28, hidden_act: str = "relu2", rms_norm_eps: float = 1e-5, + layers_block_type: list[str] | None = None, **kwargs, ) -> None: + # Each decoder block is an attention layer followed by an MLP layer + if layers_block_type is None: + layers_block_type = ["full_attention", "mlp"] * num_hidden_layers super().__init__( - num_hidden_layers=2 * num_hidden_layers, - hybrid_override_pattern="*-" * num_hidden_layers, + layers_block_type=layers_block_type, mlp_hidden_act=hidden_act, layer_norm_epsilon=rms_norm_eps, **kwargs, @@ -46,7 +48,6 @@ def to_dict(self) -> dict: config_dict["num_hidden_layers"] = self.num_hidden_layers // 2 config_dict["hidden_act"] = config_dict.pop("mlp_hidden_act") config_dict["rms_norm_eps"] = config_dict.pop("layer_norm_epsilon") - config_dict.pop("hybrid_override_pattern", None) return config_dict diff --git a/vllm/transformers_utils/configs/deepseek_vl2.py b/vllm/transformers_utils/configs/deepseek_vl2.py index 3543db833a4d..a819f12f3e4e 100644 --- a/vllm/transformers_utils/configs/deepseek_vl2.py +++ b/vllm/transformers_utils/configs/deepseek_vl2.py @@ -90,11 +90,6 @@ def __init__( @strict class DeepseekVLV2TextConfig(DeepseekV2Config): - # DeepSeek-VL2 checkpoints (e.g. deepseek-vl2-small) omit ``vocab_size`` and - # rely on the original DeepSeek-VL2 default of 102400; the generic - # Transformers ``DeepseekV2Config`` incorrectly defaults it to 32000. - # Upstream fix: https://github.com/huggingface/transformers/pull/48159 - vocab_size: int = 102400 kv_lora_rank: int | None = DeepseekV2Config.kv_lora_rank diff --git a/vllm/transformers_utils/configs/diffusion_gemma.py b/vllm/transformers_utils/configs/diffusion_gemma.py deleted file mode 100644 index 884d176cd0b4..000000000000 --- a/vllm/transformers_utils/configs/diffusion_gemma.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers import PreTrainedConfig -from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig - - -def _init_text_config(self: PreTrainedConfig, **kwargs: Any) -> None: - PreTrainedConfig.__init__(self, **kwargs) - # DiffusionGemma always uses MoE and K=V sharing for full_attention - # layers. The HF reference removed these config fields entirely. - if getattr(self, "num_experts", None): - self.enable_moe_block = True - self.attention_k_eq_v = True - - -class DiffusionGemmaTextConfig(PreTrainedConfig): - model_type = "diffusion_gemma_text" - - def __init__(self, **kwargs: Any): - _init_text_config(self, **kwargs) - - -class DiffusionGemmaConfig(PreTrainedConfig): - model_type = "diffusion_gemma" - - def __init__( - self, - text_config: dict[str, Any] | None = None, - canvas_length: int = 256, - self_conditioning_size: int | None = None, - **kwargs: Any, - ): - self.text_config = DiffusionGemmaTextConfig(**(text_config or {})) - self.canvas_length = canvas_length - self.self_conditioning_size = self_conditioning_size - vision_config = kwargs.pop("vision_config", None) - if isinstance(vision_config, dict): - self.vision_config = Gemma4VisionConfig(**vision_config) - else: - self.vision_config = vision_config - self.audio_config = None - PreTrainedConfig.__init__(self, **kwargs) diff --git a/vllm/transformers_utils/configs/gemma4.py b/vllm/transformers_utils/configs/gemma4.py deleted file mode 100644 index 69cbc06d0482..000000000000 --- a/vllm/transformers_utils/configs/gemma4.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Per-layer config resolution for Gemma4 and its variants.""" - -from copy import copy - -from transformers import PreTrainedConfig - - -def gemma4_layer_config( - text_config: PreTrainedConfig, layer_idx: int -) -> PreTrainedConfig: - """The Gemma4 text config as it applies to one layer. - - Gemma4 uses a larger head dimension on its full attention layers than on its - sliding ones, and with `attention_k_eq_v` it uses more KV heads there too. - Transformers >= 5.15.0 says so in the config itself; before that the values - are flat attributes that `layer_types` picks between, so resolve them here. - The result is homogeneous either way, so callers read `head_dim` and - `num_key_value_heads` off it without caring which version they are on. - """ - if getattr(text_config, "is_heterogeneous", False): - return text_config.per_layer_config[layer_idx] - - layer = copy(text_config) - if text_config.layer_types[layer_idx] == "full_attention": - global_head_dim = getattr(text_config, "global_head_dim", None) - layer.head_dim = global_head_dim or text_config.head_dim - if getattr(text_config, "attention_k_eq_v", False): - global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) - layer.num_key_value_heads = ( - global_kv_heads or text_config.num_key_value_heads - ) - return layer diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py deleted file mode 100644 index 8522bdd47a3e..000000000000 --- a/vllm/transformers_utils/configs/granite4_vision.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any - -import transformers - - -class Granite4VisionConfig(transformers.PreTrainedConfig): - """Configuration for Granite 4 Vision model. - - This config is needed because the granite4_vision model type is not yet - in the transformers version pinned by vLLM. Once transformers adds native - support, this file can be removed and the _CONFIG_REGISTRY entry dropped. - """ - - model_type = "granite4_vision" - is_composition = False - - def __init__( - self, - vision_config: dict[str, Any] | None = None, - text_config: dict[str, Any] | None = None, - image_token_index: int = 100352, - image_seq_length: int = 576, - image_grid_pinpoints: list[list[int]] | None = None, - vision_feature_select_strategy: str = "full", - vision_feature_layer: int | list[int] = -2, - projector_hidden_act: str = "gelu", - projector_dropout: float = 0.1, - downsample_rate: str | None = None, - use_image_newline_parameter: bool = True, - deepstack_layer_map: list[list[int]] | None = None, - use_spatial_sampling: bool = False, - spatial_stride: int = 2, - spatial_vision_layer: int = -1, - spatial_target_layers: list[int] | None = None, - # Hub aliases — base model config uses different field names - vision_layer_to_llm_layer: list[list[int]] | None = None, - use_checkerboard_sampling: bool | None = None, - checkerboard_stride: int | None = None, - checkerboard_vision_layer: int | None = None, - checkerboard_llm_layers: list[int] | None = None, - **kwargs: Any, - ): - self.image_token_index = image_token_index - self.image_seq_length = image_seq_length - self.image_grid_pinpoints = image_grid_pinpoints or [] - self.vision_feature_select_strategy = vision_feature_select_strategy - self.vision_feature_layer = vision_feature_layer - self.projector_hidden_act = projector_hidden_act - self.projector_dropout = projector_dropout - self.downsample_rate = downsample_rate - self.use_image_newline_parameter = use_image_newline_parameter - self.deepstack_layer_map = deepstack_layer_map or vision_layer_to_llm_layer - self.use_spatial_sampling = ( - use_spatial_sampling - if use_checkerboard_sampling is None - else use_checkerboard_sampling - ) - self.spatial_stride = ( - spatial_stride if checkerboard_stride is None else checkerboard_stride - ) - self.spatial_vision_layer = ( - spatial_vision_layer - if checkerboard_vision_layer is None - else checkerboard_vision_layer - ) - self.spatial_target_layers = ( - spatial_target_layers or checkerboard_llm_layers or [0, 10, 20, 30] - ) - - if vision_config is None: - vision_config = {} - if text_config is None: - text_config = {} - - vision_model_type = vision_config.get("model_type", "siglip_vision_model") - if vision_model_type in transformers.CONFIG_MAPPING: - self.vision_config = transformers.CONFIG_MAPPING[vision_model_type]( - **vision_config - ) - else: - self.vision_config = transformers.PreTrainedConfig(**vision_config) - - text_model_type = text_config.get("model_type", "granite") - if text_model_type in transformers.CONFIG_MAPPING: - self.text_config = transformers.CONFIG_MAPPING[text_model_type]( - **text_config - ) - else: - self.text_config = transformers.PreTrainedConfig(**text_config) - - super().__init__(**kwargs) diff --git a/vllm/transformers_utils/configs/hy_v3.py b/vllm/transformers_utils/configs/hy_v3.py deleted file mode 100644 index 939dd5d78ceb..000000000000 --- a/vllm/transformers_utils/configs/hy_v3.py +++ /dev/null @@ -1,185 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers.configuration_utils import PreTrainedConfig - - -class HYV3Config(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`HYV3Model`]. - It is used to instantiate a HYV3 model (HY V3 MoE language model) according to - the specified arguments. - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to - control the model outputs. Read the documentation from [`PreTrainedConfig`] - for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 120832): - Vocabulary size of the model. - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 13312): - Dimension of the dense FFN intermediate representations. - num_hidden_layers (`int`, *optional*, defaults to 80): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 64): - Number of attention heads for each attention layer. - num_key_value_heads (`int`, *optional*, defaults to 8): - Number of key-value heads for grouped-query attention. - head_dim (`int`, *optional*, defaults to 128): - Dimension per attention head. - hidden_act (`str`, *optional*, defaults to `"silu"`): - Activation function used in FFN layers. - max_position_embeddings (`int`, *optional*, defaults to 131072): - Maximum sequence length supported by the model. - initializer_range (`float`, *optional*, defaults to 0.006): - Standard deviation of the truncated normal initializer for weight - initialization. - rms_norm_eps (`float`, *optional*, defaults to 1e-5): - Epsilon for RMS normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether to use KV cache for decoding. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*): - Beginning-of-sequence token id. - eos_token_id (`int` or `List[int]`, *optional*): - End-of-sequence token id(s). - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. - qk_norm (`bool`, *optional*, defaults to `True`): - Whether to apply RMSNorm to query and key states before attention. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie input and output embedding weights. - enable_attention_fp32_softmax (`bool`, *optional*, defaults to `False`): - Whether to upcast attention softmax to float32. Note: the eager attention - path always computes softmax in float32 regardless of this setting; this - flag is reserved for future use with custom attention backends. - enable_lm_head_fp32 (`bool`, *optional*, defaults to `True`): - Whether to upcast the LM head computation to float32. - num_experts (`int`, *optional*, defaults to 192): - Total number of MoE experts. - num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of experts selected per token (top-k routing). - num_shared_experts (`int`, *optional*, defaults to 1): - Number of always-active shared experts combined into a single MLP. - expert_hidden_dim (`int`, *optional*, defaults to 1536): - Intermediate dimension of each individual MoE expert. - moe_router_enable_expert_bias (`bool`, *optional*, defaults to `True`): - Whether to use per-expert load-balancing bias in the router. - moe_router_use_sigmoid (`bool`, *optional*, defaults to `True`): - Whether to use sigmoid (instead of softmax) for router scoring. - route_norm (`bool`, *optional*, defaults to `True`): - Whether to normalize routing scores when using sigmoid routing. - router_scaling_factor (`float`, *optional*): - Optional multiplicative scaling factor applied to routing scores. - use_grouped_mm (`bool`, *optional*, defaults to `False`): - Whether to use grouped GEMM for expert computation (not yet implemented). - enable_moe_fp32_combine (`bool`, *optional*, defaults to `False`): - Whether to accumulate expert outputs in float32. - first_k_dense_replace (`int`, *optional*, defaults to 1): - Number of initial decoder layers that use a dense FFN instead of MoE. - output_router_logits (`bool`, *optional*, defaults to `False`): - Whether to output router logits from each MoE layer. Useful for computing - auxiliary load-balancing loss during training. Disabled by default to avoid - the memory overhead of storing per-layer router tensors during inference. - - Example: - ```python - >>> from transformers import HYV3Config, HYV3Model - - >>> config = HYV3Config() - >>> model = HYV3Model(config) - ``` - - """ - - model_type = "hy_v3" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=120832, - hidden_size=4096, - intermediate_size=13312, - num_hidden_layers=80, - num_attention_heads=64, - num_key_value_heads=8, - head_dim=128, - hidden_act="silu", - max_position_embeddings=131072, - initializer_range=0.006, - rms_norm_eps=1e-5, - use_cache=True, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - rope_parameters: dict[str, Any] | None = None, - qk_norm=True, - tie_word_embeddings=False, - enable_attention_fp32_softmax=False, - enable_lm_head_fp32=True, - # MoE specific - num_experts=192, - num_experts_per_tok=8, - num_shared_experts=1, - expert_hidden_dim=1536, - moe_router_enable_expert_bias=True, - moe_router_use_sigmoid=True, - route_norm=True, - router_scaling_factor=None, - use_grouped_mm=False, - enable_moe_fp32_combine=False, - # Dense/MoE layer control - first_k_dense_replace=1, - output_router_logits=False, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - rope_theta = kwargs.pop("rope_theta", 11158840.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.qk_norm = qk_norm - self.tie_word_embeddings = tie_word_embeddings - self.enable_lm_head_fp32 = enable_lm_head_fp32 - self.enable_attention_fp32_softmax = enable_attention_fp32_softmax - - # MoE specific - self.num_experts = num_experts - self.num_experts_per_tok = num_experts_per_tok - self.num_shared_experts = num_shared_experts - self.expert_hidden_dim = expert_hidden_dim - self.moe_router_enable_expert_bias = moe_router_enable_expert_bias - self.moe_router_use_sigmoid = moe_router_use_sigmoid - self.route_norm = route_norm - self.use_grouped_mm = use_grouped_mm - self.router_scaling_factor = router_scaling_factor - self.enable_moe_fp32_combine = enable_moe_fp32_combine - - # Dense/MoE layer control - self.first_k_dense_replace = first_k_dense_replace - self.output_router_logits = output_router_logits - - if eos_token_id is not None and isinstance(eos_token_id, int): - eos_token_id = [eos_token_id] - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py deleted file mode 100644 index 5470408b5803..000000000000 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ /dev/null @@ -1,277 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Copyright 2025 NAVER Cloud HyperCLOVA team -# -# Copyright 2025 NAVER Cloud HyperCLOVA team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""HyperCLOVA X model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class HyperCLOVAXConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`HyperCLOVAXModel`]. It is used to instantiate a HyperCLOVAX model - according to the specified arguments, defining the model architecture. - Configuration objects inherit from [`PreTrainedConfig`] and can be used - to control the model outputs. Read the documentation from - [`PreTrainedConfig`] for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 32000): - Vocabulary size of the HyperCLOVAX model. Defines the number of - different tokens that can be represented by the `input_ids` - passed when calling [`HyperCLOVAXModel`] - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 11008): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the - Transformer decoder. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if `num_key_value_heads=1` the model - will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each - group key and value head should be constructed by meanpooling all - the original heads within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not - specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the - decoder. - max_position_embeddings (`int`, *optional*, defaults to 2048): - The maximum sequence length that this model might ever be used - with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - pretraining_tp (`int`, *optional*, defaults to 1): - Experimental feature. Tensor parallelism rank used during - pretraining. Please refer to [this document](https://huggingface. - co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism) - to understand more about it. This value is necessary to ensure - exact reproducibility of the pretraining results. Please refer to - [this issue](https://github.com/pytorch/pytorch/issues/76232). - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_theta (`float`, *optional*, defaults to 10000.0): - The base period of the RoPE embeddings. - rope_scaling (`Dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE - embeddings. NOTE: if you apply new rope type and you expect the - model to work on longer `max_position_embeddings`, we recommend - you to update this value accordingly. - Expected contents: - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', - 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with - 'default' being the original RoPE implementation. - `factor` (`float`, *optional*): - Used with all rope types except 'default'. The scaling - factor to apply to the RoPE embeddings. In most scaling - types, a `factor` of x will enable the model to handle - sequences of length x * original maximum pre-trained - length. - `original_max_position_embeddings` (`int`, *optional*): - Used with 'dynamic', 'longrope' and 'llama3'. The - original max position embeddings used during pretraining. - `attention_factor` (`float`, *optional*): - Used with 'yarn' and 'longrope'. The scaling factor to be - applied on the attention computation. If unspecified, it - defaults to value recommended by the implementation, using - the `factor` field to infer the suggested value. - `beta_fast` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for - extrapolation (only) in the linear ramp function. If - unspecified, it defaults to 32. - `beta_slow` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for - interpolation (only) in the linear ramp function. If - unspecified, it defaults to 1. - `short_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be - applied to short contexts (< - `original_max_position_embeddings`). Must be a list of - numbers with the same length as the hidden size divided - by the number of attention heads divided by 2 - `long_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be - applied to long contexts (< - `original_max_position_embeddings`). Must be a list of - numbers with the same length as the hidden size divided - by the number of attention heads divided by 2 - `low_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to low - frequency components of the RoPE - `high_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to high - frequency components of the RoPE - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output - projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in up_proj, down_proj and gate_proj layers - in the MLP layers. - head_dim (`int`, *optional*): - The attention head dimension. If None, it will default to - hidden_size // num_heads - embedding_multiplier (`float`, *optional*, defaults to `None`): - Multiplier applied to the embedding weights. If `None`, it is - equivalent to `1.0`. - logits_scaling (`float`, *optional*, defaults to `None`): - Scaling factor for logits. If `None`, it is equivalent to `1.0`. - attention_multiplier (`float`, *optional*, defaults to `None`): - Multiplier applied to the attention weights. If `None`, it is - equivalent to `self.head_dim ** -0.5`. - residual_multiplier (`float`, *optional*, defaults to `None`): - Scaling factor for residual connections. If `None`, it is - equivalent to `1.0`. - use_post_norm (`bool`, *optional*, defaults to `True`): - Determines whether to apply Peri-Layer Normalization. Set to - False to disable this feature. - rope_parameters (`dict`, *optional*): - Dictionary containing the RoPE parameters used by vLLM's - `get_rope`. When provided, takes precedence over `rope_theta` - and `rope_scaling`. If `None`, it is derived from `rope_theta` - and `rope_scaling` automatically. - - """ - - model_type = "hyperclovax" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=32000, - hidden_size=4096, - intermediate_size=11008, - num_hidden_layers=32, - num_attention_heads=32, - num_key_value_heads=None, - hidden_act="silu", - max_position_embeddings=2048, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=None, - bos_token_id=1, - eos_token_id=2, - pretraining_tp=1, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - mlp_bias=False, - head_dim=None, - embedding_multiplier=None, # mup - logits_scaling=None, # mup - attention_multiplier=None, # mup - residual_multiplier=None, # mup - use_post_norm=True, # post-norm(peri-LN) - rope_parameters=None, - auto_map=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - self.head_dim = ( - head_dim - if head_dim is not None - else self.hidden_size // self.num_attention_heads - ) - # Derive rope_parameters for vLLM's get_rope() from rope_theta / - # rope_scaling, unless the caller already provided rope_parameters. - if rope_parameters is None: - if rope_scaling is not None: - # Shallow-copy to avoid mutating the caller's dict. - rope_parameters = dict(rope_scaling) - # BC: 'type' field -> 'rope_type', remove stale key. - if "type" in rope_parameters: - rope_parameters.setdefault("rope_type", rope_parameters.pop("type")) - else: - rope_parameters = {"rope_type": "default"} - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - self.rope_parameters = rope_parameters - - # BC: keep self.rope_scaling consistent for HF serialization. - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - - # mup - self.embedding_multiplier = ( - embedding_multiplier if embedding_multiplier is not None else 1.0 - ) - self.logits_scaling = logits_scaling if logits_scaling is not None else 1.0 - self.attention_multiplier = ( - attention_multiplier - if attention_multiplier is not None - else self.head_dim**-0.5 - ) - self.residual_multiplier = ( - residual_multiplier if residual_multiplier is not None else 1.0 - ) - - # post-norm (Peri-LN) - self.use_post_norm = use_post_norm - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - auto_map=auto_map, - **kwargs, - ) diff --git a/vllm/transformers_utils/configs/kimi_k25.py b/vllm/transformers_utils/configs/kimi_k25.py deleted file mode 100644 index ed3d05d7cc21..000000000000 --- a/vllm/transformers_utils/configs/kimi_k25.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Kimi-K2.5 Model Configuration. - -This configuration supports video-chunk as an internal modality type. -A video-chunk is the smallest independently processable unit of video. -""" - -from transformers import DeepseekV3Config -from transformers.configuration_utils import PreTrainedConfig - - -class KimiK25VisionConfig(PreTrainedConfig): - model_type = "kimi_k25_vision" - - def __init__( - self, - # Vision Tower - patch_size: int = 14, - init_pos_emb_height: int = 64, - init_pos_emb_width: int = 64, - init_pos_emb_time: int = 4, - pos_emb_type: str = "divided_fixed", - num_attention_heads: int = 16, - num_hidden_layers: int = 27, - hidden_size: int = 1152, - intermediate_size: int = 4304, - merge_kernel_size: tuple[int, int] = (2, 2), - video_attn_type: str = "spatial_temporal", - merge_type: str = "sd2_tpool", - # MM Projector - mm_projector_type: str = "patchmerger", - mm_hidden_size: int | None = None, - projector_hidden_act: str = "gelu", - projector_ln_eps: float = 1e-5, - **kwargs, - ): - super().__init__(**kwargs) - # Vision Tower - self.patch_size = patch_size - self.init_pos_emb_height = init_pos_emb_height - self.init_pos_emb_width = init_pos_emb_width - self.init_pos_emb_time = init_pos_emb_time - self.pos_emb_type = pos_emb_type - self.num_attention_heads = num_attention_heads - self.num_hidden_layers = num_hidden_layers - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.merge_kernel_size = merge_kernel_size - self.video_attn_type = video_attn_type - self.merge_type = merge_type - # MM Projector - self.mm_projector_type = mm_projector_type - if mm_hidden_size is not None: - self.mm_hidden_size = mm_hidden_size - else: - self.mm_hidden_size = hidden_size - self.projector_hidden_act = projector_hidden_act - self.projector_ln_eps = projector_ln_eps - - -class KimiK25Config(PreTrainedConfig): - """Kimi-K2.5 model configuration. - - Kimi-K2.5 extends Kimi-K2 with vision support using video-chunks. - A video-chunk consists of multiple consecutive frames - that are processed together with temporal pooling. - - Args: - vision_config: Configuration for the vision tower and projector. - text_config: Configuration for the text model (DeepseekV3). - ignore_index: The ignore index for the loss function. - media_placeholder_token_id: The token ID for media placeholders. - pad_token_id: The token ID for padding. - - """ - - model_type = "kimi_k25" - - def __init__( - self, - vision_config: dict | KimiK25VisionConfig | None = None, - text_config: dict | DeepseekV3Config | None = None, - ignore_index: int = -100, - media_placeholder_token_id: int = 163605, - pad_token_id: int = 0, - use_unified_vision_chunk: bool = False, - video_placeholder: str = "<|kimi_k25_video_placeholder|>", - **kwargs, - ): - # Vision config - if vision_config is None: - self.vision_config = KimiK25VisionConfig() - elif isinstance(vision_config, dict): - self.vision_config = KimiK25VisionConfig(**vision_config) - else: - self.vision_config = vision_config - - # Text config - if text_config is None: - self.text_config = DeepseekV3Config() - elif isinstance(text_config, dict): - self.text_config = DeepseekV3Config(**text_config) - else: - self.text_config = text_config - - # Set mm_hidden_size to text hidden size if not explicitly set - if self.vision_config.mm_hidden_size == self.vision_config.hidden_size: - self.vision_config.mm_hidden_size = self.text_config.hidden_size - - # Other config - self.ignore_index = ignore_index - self.media_placeholder_token_id = media_placeholder_token_id - self.use_unified_vision_chunk = use_unified_vision_chunk - self.video_placeholder = video_placeholder - - # Propagate quantization config from text model - if getattr(self.text_config, "quantization_config", None) is not None: - self.quantization_config = self.text_config.quantization_config - - super().__init__(pad_token_id=pad_token_id, **kwargs) - - @property - def hidden_size(self) -> int: - """Get hidden size from text config for compatibility.""" - return self.text_config.hidden_size - - @property - def vocab_size(self) -> int: - """Get vocab size from text config for compatibility.""" - return self.text_config.vocab_size diff --git a/vllm/transformers_utils/configs/lfm2_moe.py b/vllm/transformers_utils/configs/lfm2_moe.py deleted file mode 100644 index 903ae9d54700..000000000000 --- a/vllm/transformers_utils/configs/lfm2_moe.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any - -from transformers.configuration_utils import PreTrainedConfig - - -class Lfm2MoeConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`Lfm2MoeModel`]. It is used to instantiate a LFM2 Moe - model according to the specified arguments, defining the model architecture. Instantiating a configuration with the - defaults will yield a similar configuration to that of the LFM2-8B-A1B model. - e.g. [LiquidAI/LFM2-8B-A1B](https://huggingface.co/LiquidAI/LFM2-8B-A1B) - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 65536): - Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`Lfm2Model`] - hidden_size (`int`, *optional*, defaults to 2048): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 7168): - Dimension of the MLP representations. - moe_intermediate_size (`int`, *optional*, defaults to 1792): - Intermediate size of the routed expert. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - pad_token_id (`int`, *optional*, defaults to 0): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 1): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 2): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, defaults to `True`): - Whether to tie weight embeddings - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. - max_position_embeddings (`int`, *optional*, defaults to 128000): - The maximum sequence length that this model might ever be used with. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - norm_eps (`float`, *optional*, defaults to 1e-05): - The epsilon used by the rms normalization layers. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the Transformer decoder. - num_key_value_heads (`int`, *optional*, defaults to 8): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details, check out [this - paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to - `num_attention_heads`. - conv_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the conv layers. - conv_L_cache (`int`, *optional*, defaults to 3): - L_cache dim in the conv layers. - num_dense_layers (`int`, *optional*, defaults to 2): - Number of dense Lfm2MoeMLP layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). - num_experts_per_tok (`int`, *optional*, defaults to 4): - Number of selected experts. - num_experts (`int`, *optional*, defaults to 32): - Number of routed experts. - use_expert_bias (`bool`, *optional*, defaults to `True`): - Whether to use the expert bias on the routing weights. - routed_scaling_factor (`float`, *optional*, defaults to 1.0): - Scaling factor for routed experts in MoE models. - norm_topk_prob (`bool`, *optional*, defaults to `True`): - Whether to normalize the topk probabilities. - layer_types (`Optional`, *optional*): - Type of each layers. - - ```python - >>> from transformers import Lfm2MoeModel, Lfm2MoeConfig - - >>> # Initializing a LFM2 Moe model - >>> configuration = Lfm2MoeConfig() - - >>> # Initializing a model from the LFM2-8B-A1B style configuration - >>> model = Lfm2MoeModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ # noqa: E501 - - model_type = "lfm2_moe" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size: int = 65536, - hidden_size: int = 2048, - intermediate_size: int = 7168, - moe_intermediate_size: int = 1792, - num_hidden_layers: int = 32, - pad_token_id: int = 0, - bos_token_id: int = 1, - eos_token_id: int = 2, - tie_word_embeddings: bool = True, - rope_parameters: dict[str, Any] | None = None, - max_position_embeddings: int = 128_000, - use_cache: bool = True, - norm_eps: float = 0.00001, - num_attention_heads: int = 32, - num_key_value_heads: int = 8, - conv_bias: bool = False, - conv_L_cache: int = 3, - num_dense_layers: int = 2, - num_experts_per_tok: int = 4, - num_experts: int = 32, - use_expert_bias: bool = True, - routed_scaling_factor: float = 1.0, - norm_topk_prob: bool = True, - layer_types: list[str] | None = None, - **kwargs, - ): - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - rope_theta = kwargs.pop("rope_theta", 1000000.0) - if rope_parameters is None: - rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} - self.rope_parameters = rope_parameters - self.max_position_embeddings = max_position_embeddings - self.use_cache = use_cache - self.norm_eps = norm_eps - - # attn operator config - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - - # custom operator config - self.conv_bias = conv_bias - self.conv_L_cache = conv_L_cache - - # moe config - self.num_dense_layers = num_dense_layers - self.moe_intermediate_size = moe_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.use_expert_bias = use_expert_bias - self.routed_scaling_factor = routed_scaling_factor - self.norm_topk_prob = norm_topk_prob - self.layer_types = layer_types - - tie_word_embeddings = kwargs.get( - "tie_embedding", tie_word_embeddings - ) # to fit original config keys - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - -__all__ = ["Lfm2MoeConfig"] diff --git a/vllm/transformers_utils/configs/nemotron.py b/vllm/transformers_utils/configs/nemotron.py deleted file mode 100644 index 135a0e1596dc..000000000000 --- a/vllm/transformers_utils/configs/nemotron.py +++ /dev/null @@ -1,219 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2024 HuggingFace Inc. team. All rights reserved. -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Nemotron model configuration.""" - -from transformers import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class NemotronConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`NemotronModel`]. It is used to instantiate a Nemotron model - according to the specified arguments, defining the model architecture. - Instantiating a configuration with the defaults will yield a similar - configuration to that of the Nemotron-8B. - - Configuration objects inherit from [`PreTrainedConfig`] and can be - used to control the model outputs. Read the documentation from - [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 256000): - Vocabulary size of the Nemotron model. Defines the number of - different tokens that can be represented by the - `inputs_ids` passed when calling [`NemotronModel`] - hidden_size (`int`, *optional*, defaults to 6144): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 24576): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 32): - Number of hidden layers in the Transformer decoder. - num_attention_heads (`int`, *optional*, defaults to 48): - Number of attention heads for each attention layer in the - Transformer decoder. - head_dim (`int`, *optional*): - Projection weights dimension in multi-head attention. Set to - hidden_size // num_attention_heads if None - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention - (MQA) otherwise GQA is used. When converting a multi-head - checkpoint to a GQA checkpoint, each group key and value - head should be constructed by meanpooling all the original - heads within that group. For more details checkout - [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it - is not specified, will default to `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, defaults to `"relu2"`): - The non-linear activation function (function or string) in the - decoder. - max_position_embeddings (`int`, *optional*, defaults to 4096): - The maximum sequence length that this model might ever be used - with. - initializer_range (`float`, *optional*, defaults to 0.0134): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - norm_eps (`float`, *optional*, defaults to 1e-05): - The epsilon used by the normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - pad_token_id (`int`, *optional*): - Padding token id. - bos_token_id (`int`, *optional*, defaults to 2): - Beginning of stream token id. - eos_token_id (`int`, *optional*, defaults to 3): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings - rope_parameters (`dict`, *optional*): - The parameters of the RoPE embeddings. Expected contents: - `rope_theta` (`float`): The base period of the RoPE embeddings. - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', 'linear', - 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the - original RoPE implementation. - `partial_rotary_factor` (`float`, *optional*, defaults to 0.5): - Percentage of the query and keys which will have rotary embedding. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output - projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in up_proj and down_proj layers in the MLP - layers. - - ```python - >>> from transformers import NemotronModel, NemotronConfig - >>> # Initializing a Nemotron nemotron-15b style configuration - >>> configuration = NemotronConfig() - >>> # Initializing a model from the nemotron-15b style configuration - >>> model = NemotronModel(configuration) - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ - - model_type = "nemotron" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=256000, - hidden_size=6144, - intermediate_size=24576, - num_hidden_layers=32, - num_attention_heads=48, - head_dim=None, - num_key_value_heads=None, - hidden_act="relu2", - max_position_embeddings=4096, - initializer_range=0.0134, - norm_eps=1e-5, - use_cache=True, - pad_token_id=None, - bos_token_id=2, - eos_token_id=3, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - mlp_bias=False, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - head_dim = head_dim or kwargs.get("kv_channels") - self.head_dim = ( - head_dim if head_dim is not None else (hidden_size // num_attention_heads) - ) - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.norm_eps = norm_eps - self.use_cache = use_cache - # Try to set `rope_scaling` if available, otherwise use `rope_parameters` - rope_scaling = kwargs.pop("rope_scaling", None) - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"} - rope_theta = kwargs.pop("rope_theta", 10000.0) - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - # for backward compatibility - partial_rotary_factor = ( - kwargs.get("rope_percent") - or kwargs.get("rope_percentage") - or kwargs.get("partial_rotary_factor") - or 0.5 - ) - if "partial_rotary_factor" not in rope_parameters: - rope_parameters["partial_rotary_factor"] = partial_rotary_factor - self.rope_parameters = rope_parameters - self._rope_parameters_validation() - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - def _rope_parameters_validation(self): - """Validate the `rope_parameters` configuration.""" - if self.rope_parameters is None: - return - - rope_type: str | None = self.rope_parameters.get("rope_type", None) - factor: float | None = self.rope_parameters.get("factor", None) - - if rope_type not in {"default", "linear", "dynamic"}: - raise ValueError( - "`rope_type` must be one of ['default', 'linear', 'dynamic'], " - f"got {rope_type}" - ) - if rope_type != "default": - if factor is None: - raise ValueError( - "If `rope_type` is not 'default', `rope_parameters` " - "must include a `factor` field. Got `None`." - ) - if not isinstance(factor, float) or factor <= 1.0: - raise ValueError( - "`rope_parameters`'s factor field must be a float > 1, got " - f"{factor}" - ) diff --git a/vllm/transformers_utils/configs/nemotron_h.py b/vllm/transformers_utils/configs/nemotron_h.py deleted file mode 100644 index c25a5d54e3d7..000000000000 --- a/vllm/transformers_utils/configs/nemotron_h.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2024 HuggingFace Inc. team. All rights reserved. -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""NemotronH model configuration.""" - -import regex as re -from transformers.configuration_utils import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class NemotronHConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a - [`NemotronHModel`]. It is used to instantiate a NemotronH model according - to the specified arguments, defining the model architecture. Instantiating - a configuration with the defaults will yield a similar configuration to - that of the NemotronH-v0.1 model. - - Args: - vocab_size (`int`, *optional*, defaults to 131072): - Vocabulary size of the NemotronH model. Defines the number of - different tokens that can be represented by the `inputs_ids` - passed when calling [`NemotronHModel`] - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether the model's input and output word embeddings should be - tied. Note that this is only relevant if the model has an output - word embedding layer. - hidden_size (`int`, *optional*, defaults to 4096): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 21504): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 52): - Number of hidden layers in the Transformer encoder. - hybrid_override_pattern (`str`, *optional*, defaults to - `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`): - The pattern of the hybrid model. The pattern is a string of - characters where each character represents - M: Mamba2, *: Attention, -: MLP - mtp_hybrid_override_pattern (`str`, *optional*, defaults to `"*E"`): - The pattern of the MTP layers. - num_attention_heads (`int`, *optional*, defaults to 32): - Number of attention heads for each attention layer in the - Transformer encoder. - attention_head_dim (`int`, *optional*, defaults to 128): - Dimension of each attention head. - num_key_value_heads (`int`, *optional*, defaults to 8): - This is the number of key_value heads that should be used to - implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use - Multi Head Attention (MHA), if `num_key_value_heads=1` the model - will use Multi Query Attention (MQA) otherwise GQA is used. - mlp_hidden_act (`str`, *optional*, defaults to "relu2"): - The non-linear activation function in the MLP layers. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in attention layers. - mlp_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in MLP layers. - use_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the model. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): - The epsilon used by the layer normalization layers. - residual_in_fp32 (`bool`, *optional*, defaults to `False`): - Whether or not residuals should be in `float32`. If set to `False` - residuals will keep the same `dtype` as the rest of the model. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). Only relevant if - `config.is_decoder=True`. - num_logits_to_keep (`int` or `None`, *optional*, defaults to 1): - Number of prompt logits to calculate during generation. If `None`, - all logits will be calculated. If an integer value, only last - `num_logits_to_keep` logits will be calculated. - pad_token_id (`int`, *optional*, defaults to 0): - The id of the padding token. - bos_token_id (`int`, *optional*, defaults to 1): - The id of the "beginning-of-sequence" token. - eos_token_id (`int`, *optional*, defaults to 2): - The id of the "end-of-sequence" token. - sliding_window (`int`, *optional*, defaults to None): - Sliding window attention window size. - max_position_embeddings (`int`, *optional*, defaults to 4096): - The maximum sequence length that this model might ever be used - with. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - hidden_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the hidden states. - use_mamba_kernels (`bool`, *optional*, defaults to `True`): - Flag indicating whether or not to use the fast mamba kernels. - These are available only if `mamba-ssm` and `causal-conv1d` - are installed, and the mamba modules are running on a CUDA device. - ssm_state_size (`int`, *optional*, defaults to 128): - The dimension of the mamba state space latents. - mamba_num_heads (`int`, *optional*, defaults to 128): - Number of heads in Mamba layers. - mamba_n_groups (`int`, *optional*, defaults to 8): - Number of groups in Mamba layers. - mamba_head_dim (`int`, *optional*, defaults to 64): - Dimension of each Mamba head. - mamba_d_conv (`int`, *optional*, defaults to 4): - The size of the mamba convolution kernel. - mamba_expand (`int`, *optional*, defaults to 2): - Expanding factor used to determine the mamba intermediate size. - mamba_hidden_act (`str`, *optional*, defaults to "silu"): - The non-linear activation function in the Mamba layers. - mamba_dt_min (`float`, *optional*, defaults to 0.001): - Minimum value for the time step in Mamba. - mamba_dt_max (`float`, *optional*, defaults to 0.1): - Maximum value for the time step in Mamba. - mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))): - Limits for the time step in Mamba. - mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4): - Floor value for time step initialization in Mamba. - mamba_conv_bias (`bool`, *optional*, defaults to `True`): - Whether to use bias in the convolution layer of the mamba mixer - block. - mamba_proj_bias (`bool`, *optional*, defaults to `False`): - Whether to use bias in the input and output projections of the - mamba mixer block. - mamba_chunk_size (`int`, *optional*, defaults to 256): - Size of chunks for Mamba processing. - rescale_prenorm_residual (`bool`, *optional*, defaults to `True`): - Whether to rescale the pre-normalization residual connections. - - """ - - model_type = "nemotron_h" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=131072, - tie_word_embeddings=False, - hidden_size=4096, - intermediate_size=21504, - num_hidden_layers=52, - hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-", - mtp_hybrid_override_pattern="*E", - num_attention_heads=32, - head_dim=128, - num_key_value_heads=8, # nemo: num_query_groups - mlp_hidden_act="relu2", - attention_bias=False, - mlp_bias=False, - use_bias=False, - initializer_range=0.02, # nemo: init_method_std - layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon - residual_in_fp32=False, # Megatron Core default value - use_cache=True, - num_logits_to_keep=1, - pad_token_id=0, - bos_token_id=1, - eos_token_id=2, - sliding_window=None, - max_position_embeddings=4096, - attention_dropout=0.0, - hidden_dropout=0.0, # * ADDED - use_mamba_kernels=True, - ssm_state_size=128, # mamba_state_size - mamba_num_heads=128, - mamba_n_groups=8, # nemo: mamba_ssm_ngroups = num_heads - mamba_head_dim=64, - mamba_d_conv=4, - mamba_expand=2, - mamba_hidden_act="silu", - mamba_dt_min=0.001, - mamba_dt_max=0.1, - mamba_dt_limit=(0.0, float("inf")), - mamba_dt_init_floor=1e-4, - mamba_conv_bias=True, - mamba_proj_bias=False, - mamba_chunk_size=256, - rescale_prenorm_residual=True, - n_routed_experts=8, - n_shared_experts=1, - moe_intermediate_size=7688, - moe_shared_expert_intermediate_size=7688, - moe_latent_size=None, - num_experts_per_tok=2, - routed_scaling_factor=1.0, - n_group=1, - topk_group=1, - norm_topk_prob=True, - **kwargs, - ): - self.vocab_size = vocab_size - self.tie_word_embeddings = tie_word_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.hybrid_override_pattern = hybrid_override_pattern - self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern - self.num_attention_heads = num_attention_heads - self.head_dim = head_dim - self.sliding_window = sliding_window - self.max_position_embeddings = max_position_embeddings - self.attention_dropout = attention_dropout - self.hidden_dropout = hidden_dropout - - # Validate hybrid_override_pattern - # M: Mamba2, *: Attention, -: MLP - assert len(self.hybrid_override_pattern) == self.num_hidden_layers, ( - "hybrid_override_pattern must have same length as num_hidden_layers" - ) - assert re.match(r"^[*-ME]+$", self.hybrid_override_pattern), ( - "hybrid_override_pattern must only contain characters 'M', '*', '-', or 'E'" - ) - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.mlp_hidden_act = mlp_hidden_act - self.attention_bias = attention_bias - self.mlp_bias = mlp_bias - self.use_bias = use_bias - self.initializer_range = initializer_range - self.layer_norm_epsilon = layer_norm_epsilon - self.residual_in_fp32 = residual_in_fp32 - - self.use_cache = use_cache - self.num_logits_to_keep = num_logits_to_keep - - self.use_mamba_kernels = use_mamba_kernels - self.n_groups = mamba_n_groups - self.mamba_head_dim = mamba_head_dim - self.ssm_state_size = ssm_state_size - self.mamba_num_heads = mamba_num_heads - self.conv_kernel = mamba_d_conv - self.expand = mamba_expand - self.mamba_hidden_act = mamba_hidden_act - self.time_step_min = mamba_dt_min - self.time_step_max = mamba_dt_max - self.time_step_limit = mamba_dt_limit - self.time_step_floor = mamba_dt_init_floor - self.use_conv_bias = mamba_conv_bias - self.mamba_proj_bias = mamba_proj_bias - self.chunk_size = mamba_chunk_size - self.rescale_prenorm_residual = rescale_prenorm_residual - self.n_routed_experts = n_routed_experts - self.n_shared_experts = n_shared_experts - self.moe_intermediate_size = moe_intermediate_size - self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size # noqa: E501 - self.moe_latent_size = moe_latent_size - self.num_experts_per_tok = num_experts_per_tok - self.routed_scaling_factor = routed_scaling_factor - self.n_group = n_group - self.topk_group = topk_group - self.norm_topk_prob = norm_topk_prob - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - - @property - def layers_block_type(self): - return [ - "mamba" - if self.hybrid_override_pattern[i] == "M" - else "attention" - if self.hybrid_override_pattern[i] == "*" - else "mlp" - if self.hybrid_override_pattern[i] == "-" - else "moe" - for i in range(self.num_hidden_layers) - ] diff --git a/vllm/transformers_utils/configs/olmo_hybrid.py b/vllm/transformers_utils/configs/olmo_hybrid.py deleted file mode 100644 index 62d3c4e0d5d1..000000000000 --- a/vllm/transformers_utils/configs/olmo_hybrid.py +++ /dev/null @@ -1,285 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -from transformers.configuration_utils import PreTrainedConfig - - -class OlmoHybridConfig(PreTrainedConfig): - r"""Configuration class for [`OlmoHybridModel`]. It is used to - instantiate an OLMo Hybrid model according to the specified - arguments, defining the model architecture. Instantiating a - configuration with the defaults will yield a similar - configuration to that of the - [allenai/Olmo-Hybrid-7B](https://huggingface.co/allenai/Olmo-Hybrid-7B) - model. - - Configuration objects inherit from [`PreTrainedConfig`] and - can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - Args: - vocab_size (`int`, *optional*, defaults to 100352): - Vocabulary size of the OlmoHybrid model. Defines - the number of different tokens that can be - represented by the `inputs_ids` passed when - calling [`OlmoHybridModel`]. - hidden_size (`int`, *optional*, defaults to 3840): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, - defaults to 11008): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, - defaults to 32): - Number of hidden layers in the Transformer - decoder. - num_attention_heads (`int`, *optional*, - defaults to 30): - Number of attention heads for each attention - layer in the Transformer decoder. - num_key_value_heads (`int`, *optional*): - This is the number of key_value heads that - should be used to implement Grouped Query - Attention. If - `num_key_value_heads=num_attention_heads`, - the model will use Multi Head Attention (MHA), - if `num_key_value_heads=1` the model will use - Multi Query Attention (MQA) otherwise GQA is - used. When converting a multi-head checkpoint - to a GQA checkpoint, each group key and value - head should be constructed by meanpooling all - the original heads within that group. For more - details, check out - [this paper](https://huggingface.co/papers/2305.13245). - If it is not specified, will default to - `num_attention_heads`. - hidden_act (`str` or `function`, *optional*, - defaults to `"silu"`): - The non-linear activation function (function - or string) in the decoder. - max_position_embeddings (`int`, *optional*, - defaults to 65536): - The maximum sequence length that this model - might ever be used with. - initializer_range (`float`, *optional*, - defaults to 0.02): - The standard deviation of the - truncated_normal_initializer for initializing - all weight matrices. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last - key/values attentions (not used by all models). - Only relevant if `config.is_decoder=True`. - pad_token_id (`int`, *optional*, - defaults to 100277): - Padding token id. - bos_token_id (`int`, *optional*): - Beginning of stream token id. - eos_token_id (`int`, *optional*, - defaults to 100257): - End of stream token id. - tie_word_embeddings (`bool`, *optional*, - defaults to `False`): - Whether to tie weight embeddings. - rope_parameters (`RopeParameters`, *optional*): - Dictionary containing the configuration - parameters for the RoPE embeddings. Can be - `None` to disable RoPE. - attention_bias (`bool`, *optional*, - defaults to `False`): - Whether to use a bias in the query, key, value - and output projection layers during - self-attention. - attention_dropout (`float`, *optional*, - defaults to 0.0): - The dropout ratio for the attention - probabilities. - rms_norm_eps (`float`, *optional*, - defaults to 1e-06): - The epsilon used by the rms normalization - layers. - layer_types (`list`, *optional*): - Attention pattern for each layer. Can contain - `"full_attention"` or `"linear_attention"`. - Defaults to linear attention for most layers - with full attention for every 4th layer. - linear_num_key_heads (`int`, *optional*): - Number of key heads for the linear attention - layers. Defaults to `num_attention_heads`. - linear_num_value_heads (`int`, *optional*): - Number of value heads for the linear attention - layers. Defaults to `num_attention_heads`. - linear_key_head_dim (`int`, *optional*): - Dimension of each key head in linear attention - layers. Defaults to - `0.75 * hidden_size / linear_num_key_heads`. - linear_value_head_dim (`int`, *optional*): - Dimension of each value head in linear - attention layers. Defaults to - `2 * linear_key_head_dim`. - linear_a_log_min (`float`, *optional*, - defaults to 0.0): - Minimum value for uniform initialization of - A_log in GatedDeltaNet layers. - linear_a_log_max (`float`, *optional*, - defaults to 16.0): - Maximum value for uniform initialization of - A_log in GatedDeltaNet layers. - linear_dt_min (`float`, *optional*, - defaults to 0.001): - Minimum value for dt initialization in - GatedDeltaNet layers. - linear_dt_max (`float`, *optional*, - defaults to 0.1): - Maximum value for dt initialization in - GatedDeltaNet layers. - linear_dt_init_floor (`float`, *optional*, - defaults to 0.0001): - Floor value for clamping dt during - initialization in GatedDeltaNet layers. - linear_conv_kernel_dim (`int`, *optional*, - defaults to 4): - Kernel size for the short convolution applied - to queries, keys, and values in linear - attention layers. - linear_allow_neg_eigval (`bool`, *optional*, - defaults to `True`): - Whether to allow negative eigenvalues in the - GatedDeltaNet recurrence. When `True`, the - beta parameter is scaled by 2.0 to allow - values in range [0, 2] instead of [0, 1]. - ```python - >>> from transformers import ( - ... OlmoHybridModel, - ... OlmoHybridConfig, - ... ) - - >>> configuration = OlmoHybridConfig() - >>> model = OlmoHybridModel(configuration) - >>> configuration = model.config - ``` - - """ - - model_type = "olmo_hybrid" - keys_to_ignore_at_inference = ["past_key_values"] - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise_gather_output", - "layers.*.self_attn.k_proj": "colwise_gather_output", - "layers.*.self_attn.v_proj": "colwise_gather_output", - "layers.*.self_attn.o_proj": "rowwise_split_input", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - def __init__( - self, - vocab_size: int | None = 100352, - hidden_size: int | None = 3840, - intermediate_size: int | None = 11008, - num_hidden_layers: int | None = 32, - num_attention_heads: int | None = 30, - num_key_value_heads: int | None = None, - hidden_act: str | None = "silu", - max_position_embeddings: int | None = 65536, - initializer_range: float | None = 0.02, - use_cache: bool | None = True, - pad_token_id: int | None = 100277, - bos_token_id: int | None = None, - eos_token_id: int | None = 100257, - tie_word_embeddings: bool | None = False, - rope_parameters=None, - attention_bias: bool | None = False, - attention_dropout: float | None = 0.0, - rms_norm_eps: float | None = 1e-06, - layer_types: list[str] | None = None, - linear_num_key_heads: int | None = None, - linear_num_value_heads: int | None = None, - linear_key_head_dim: int | None = None, - linear_value_head_dim: int | None = None, - linear_a_log_min: float = 0.0, - linear_a_log_max: float = 16.0, - linear_dt_min: float = 0.001, - linear_dt_max: float = 0.1, - linear_dt_init_floor: float = 1e-4, - linear_conv_kernel_dim: int = 4, - linear_allow_neg_eigval: bool = True, - **kwargs, - ): - super().__init__(**kwargs) - - assert num_hidden_layers is not None - assert hidden_size is not None - assert num_attention_heads is not None - - if layer_types is None: - # Default: linear attention for most layers, full attention every 4th layer - layer_types = ["linear_attention"] * int(num_hidden_layers) - for i in range(int(num_hidden_layers)): - if i % 4 == 3: - layer_types[i] = "full_attention" - # Ensure at least one full attention layer for small num_hidden_layers - if "full_attention" not in layer_types: - layer_types[-1] = "full_attention" - - self.layer_types = layer_types - self.validate_layer_type() - if "linear_attention" not in layer_types: - raise ValueError( - "OLMoHybrid expects at least one 'linear_attention' layer." - ) - if all(t == "linear_attention" for t in layer_types): - raise ValueError("OLMoHybrid expects at least one attention layer.") - - self.layer_types = layer_types - - if linear_num_key_heads is None: - linear_num_key_heads = num_attention_heads - if linear_num_value_heads is None: - linear_num_value_heads = num_attention_heads - if linear_key_head_dim is None: - linear_key_head_dim = int(0.75 * hidden_size / linear_num_key_heads) - if linear_value_head_dim is None: - linear_value_head_dim = 2 * linear_key_head_dim - - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_a_log_min = linear_a_log_min - self.linear_a_log_max = linear_a_log_max - self.linear_dt_min = linear_dt_min - self.linear_dt_max = linear_dt_max - self.linear_dt_init_floor = linear_dt_init_floor - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_allow_neg_eigval = linear_allow_neg_eigval - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.rope_parameters = rope_parameters - - self.tie_word_embeddings = tie_word_embeddings - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id diff --git a/vllm/transformers_utils/configs/ovis.py b/vllm/transformers_utils/configs/ovis.py index bc84bfe812fc..d5662d7ef542 100644 --- a/vllm/transformers_utils/configs/ovis.py +++ b/vllm/transformers_utils/configs/ovis.py @@ -2,69 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 -# adapted from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_aimv2.py -# and https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py -# Ovis Config with AimV2 config registration removed for Transformers compatibility +# adapted from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py from typing import Any -from transformers import AutoConfig, PreTrainedConfig - - -class AIMv2Config(PreTrainedConfig): - """This is the configuration class to store the configuration of an [`AIMv2Model`]. - Instantiating a configuration with the defaults will yield a similar configuration - to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224). - - Args: - hidden_size: Dimension of the hidden representations. - intermediate_size: Dimension of the SwiGLU representations. - num_hidden_layers: Number of hidden layers in the Transformer. - num_attention_heads: Number of attention heads for each attention layer - in the Transformer. - num_channels: Number of input channels. - image_size: Image size. - patch_size: Patch size. - rms_norm_eps: Epsilon value used for the RMS normalization layer. - attention_dropout: Dropout ratio for attention probabilities. - projection_dropout: Dropout ratio for the projection layer after the attention. - qkv_bias: Whether to add a bias to the queries, keys and values. - use_bias: Whether to add a bias in the feed-forward and projection layers. - kwargs: Keyword arguments for the [`PreTrainedConfig`]. - - """ - - model_type: str = "aimv2" - - def __init__( - self, - hidden_size: int = 1024, - intermediate_size: int = 2816, - num_hidden_layers: int = 24, - num_attention_heads: int = 8, - num_channels: int = 3, - image_size: int = 224, - patch_size: int = 14, - rms_norm_eps: float = 1e-5, - attention_dropout: float = 0.0, - projection_dropout: float = 0.0, - qkv_bias: bool = False, - use_bias: bool = False, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_channels = num_channels - self.patch_size = patch_size - self.image_size = image_size - self.attention_dropout = attention_dropout - self.rms_norm_eps = rms_norm_eps - - self.projection_dropout = projection_dropout - self.qkv_bias = qkv_bias - self.use_bias = use_bias +from transformers import Aimv2VisionConfig, AutoConfig, PreTrainedConfig # ---------------------------------------------------------------------- @@ -96,14 +37,13 @@ def __init__( f"expect `backbone_config` to be instance of PreTrainedConfig or dict, but got {type(backbone_config)} type" ) if not isinstance(backbone_config, PreTrainedConfig): - model_type = backbone_config["model_type"] - if model_type != "aimv2": - backbone_config.pop("model_type") + model_type = backbone_config.pop("model_type") + if model_type == "aimv2": + backbone_config = Aimv2VisionConfig(**backbone_config) + else: backbone_config = AutoConfig.for_model( model_type, **backbone_config ) - else: - backbone_config = AIMv2Config(**backbone_config) self.backbone_config = backbone_config self.hidden_stride = hidden_stride diff --git a/vllm/transformers_utils/configs/qwen3_5.py b/vllm/transformers_utils/configs/qwen3_5.py deleted file mode 100644 index 57f64f0f7814..000000000000 --- a/vllm/transformers_utils/configs/qwen3_5.py +++ /dev/null @@ -1,193 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3.5 model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class Qwen3_5TextConfig(PreTrainedConfig): - model_type = "qwen3_5_text" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - base_config_key = "text_config" - - def __init__( - self, - vocab_size=248320, - hidden_size=4096, - intermediate_size=12288, - num_hidden_layers=32, - num_attention_heads=16, - num_key_value_heads=4, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" - if bool((i + 1) % interval_pattern) - else "full_attention" - for i in range(self.num_hidden_layers) - ] - kwargs["ignore_keys_at_rope_validation"] = { - "mrope_section", - "mrope_interleaved", - } - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - super().__init__(**kwargs) - # Set these AFTER super().__init__() because transformers v4's - # PreTrainedConfig.__init__ has these as explicit params with different - # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - - -class Qwen3_5VisionConfig(PreTrainedConfig): - model_type = "qwen3_5" - base_config_key = "vision_config" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5Config(PreTrainedConfig): - model_type = "qwen3_5" - sub_configs = { - "vision_config": Qwen3_5VisionConfig, - "text_config": Qwen3_5TextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, - tie_word_embeddings=False, - **kwargs, - ): - if isinstance(vision_config, dict): - self.vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - self.vision_config = self.sub_configs["vision_config"]() - - if isinstance(text_config, dict): - self.text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - self.text_config = self.sub_configs["text_config"]() - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id - self.vision_end_token_id = vision_end_token_id - super().__init__(**kwargs) - # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite - self.tie_word_embeddings = tie_word_embeddings - - -__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig"] diff --git a/vllm/transformers_utils/configs/qwen3_5_moe.py b/vllm/transformers_utils/configs/qwen3_5_moe.py deleted file mode 100644 index 34e66c825ac5..000000000000 --- a/vllm/transformers_utils/configs/qwen3_5_moe.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3.5-MoE model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig - - -class Qwen3_5MoeTextConfig(PreTrainedConfig): - model_type = "qwen3_5_moe_text" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.shared_expert.gate_proj": "colwise", - "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - base_config_key = "text_config" - - def __init__( - self, - vocab_size=248320, - hidden_size=2048, - num_hidden_layers=40, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=8, - num_experts=256, - output_router_logits=False, - router_aux_loss_coef=0.001, - layer_types=None, - pad_token_id=None, - bos_token_id=None, - eos_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.rope_parameters = rope_parameters - kwargs.setdefault("partial_rotary_factor", 0.25) - - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.get("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" - if bool((i + 1) % interval_pattern) - else "full_attention" - for i in range(self.num_hidden_layers) - ] - kwargs["ignore_keys_at_rope_validation"] = { - "mrope_section", - "mrope_interleaved", - } - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - super().__init__(**kwargs) - # Set these AFTER super().__init__() because transformers v4's - # PreTrainedConfig.__init__ has these as explicit params with different - # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings - - -class Qwen3_5MoeVisionConfig(PreTrainedConfig): - model_type = "qwen3_5_moe" - base_config_key = "vision_config" - - def __init__( - self, - depth=27, - hidden_size=1152, - hidden_act="gelu_pytorch_tanh", - intermediate_size=4304, - num_heads=16, - in_channels=3, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, - out_hidden_size=3584, - num_position_embeddings=2304, - initializer_range=0.02, - **kwargs, - ): - super().__init__(**kwargs) - - self.depth = depth - self.hidden_size = hidden_size - self.hidden_act = hidden_act - self.intermediate_size = intermediate_size - self.num_heads = num_heads - self.in_channels = in_channels - self.patch_size = patch_size - self.spatial_merge_size = spatial_merge_size - self.temporal_patch_size = temporal_patch_size - self.out_hidden_size = out_hidden_size - self.num_position_embeddings = num_position_embeddings - self.initializer_range = initializer_range - - -class Qwen3_5MoeConfig(PreTrainedConfig): - model_type = "qwen3_5_moe" - sub_configs = { - "vision_config": Qwen3_5MoeVisionConfig, - "text_config": Qwen3_5MoeTextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config=None, - vision_config=None, - image_token_id=248056, - video_token_id=248057, - vision_start_token_id=248053, - vision_end_token_id=248054, - tie_word_embeddings=False, - **kwargs, - ): - if isinstance(vision_config, dict): - self.vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - self.vision_config = self.sub_configs["vision_config"]() - - if isinstance(text_config, dict): - self.text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - self.text_config = self.sub_configs["text_config"]() - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.vision_start_token_id = vision_start_token_id - self.vision_end_token_id = vision_end_token_id - super().__init__(**kwargs) - # Set after super().__init__() to avoid v4 PreTrainedConfig overwrite - self.tie_word_embeddings = tie_word_embeddings - - -__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] diff --git a/vllm/transformers_utils/configs/qwen3_next.py b/vllm/transformers_utils/configs/qwen3_next.py deleted file mode 100644 index 8b083ff9633d..000000000000 --- a/vllm/transformers_utils/configs/qwen3_next.py +++ /dev/null @@ -1,277 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. -# All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3-Next model configuration.""" - -from transformers.configuration_utils import PreTrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class Qwen3NextConfig(PreTrainedConfig): - r"""This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a - Qwen3-Next model according to the specified arguments, defining the model architecture. - Instantiating a configuration with the defaults will yield a similar configuration to that of - Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct). - - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 151936): - Vocabulary size of the model. Defines the number of different tokens that can be represented by the - `inputs_ids`. - hidden_size (`int`, *optional*, defaults to 2048): - Dimension of the hidden representations. - intermediate_size (`int`, *optional*, defaults to 5632): - Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 48): - Number of hidden layers in the Transformer encoder. - num_attention_heads (`int`, *optional*, defaults to 16): - Number of attention heads for each attention layer in the Transformer encoder. - num_key_value_heads (`int`, *optional*, defaults to 2): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details checkout [this - paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. - hidden_act (`str`, *optional*, defaults to `"silu"`): - The non-linear activation function in the decoder. - max_position_embeddings (`int`, *optional*, defaults to 32768): - The maximum sequence length that this model might ever be used with. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for initializing all weight matrices. - rms_norm_eps (`float`, *optional*, defaults to 1e-06): - The epsilon used by the rms normalization layers. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. - tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether the model's input and output word embeddings should be tied. - rope_parameters (`dict`, *optional*): - Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type - and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value - accordingly. - Expected contents: - `rope_theta` (`float`): The base period of the RoPE embeddings. - `rope_type` (`str`): - The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', - 'llama3'], with 'default' being the original RoPE implementation. - `factor` (`float`, *optional*): - Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In - most scaling types, a `factor` of x will enable the model to handle sequences of length x * - original maximum pre-trained length. - `original_max_position_embeddings` (`int`, *optional*): - Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during - pretraining. - `attention_factor` (`float`, *optional*): - Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention - computation. If unspecified, it defaults to value recommended by the implementation, using the - `factor` field to infer the suggested value. - `beta_fast` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear - ramp function. If unspecified, it defaults to 32. - `beta_slow` (`float`, *optional*): - Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear - ramp function. If unspecified, it defaults to 1. - `short_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be applied to short contexts (< - `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden - size divided by the number of attention heads divided by 2 - `long_factor` (`List[float]`, *optional*): - Only used with 'longrope'. The scaling factor to be applied to long contexts (< - `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden - size divided by the number of attention heads divided by 2 - `low_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE - `high_freq_factor` (`float`, *optional*): - Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE - `partial_rotary_factor` (`float`, *optional*, defaults to 0.25): - Percentage of the query and keys which will have rotary embedding. - attention_bias (`bool`, *optional*, defaults to `False`): - Whether to use a bias in the query, key, value and output projection layers during self-attention. - attention_dropout (`float`, *optional*, defaults to 0.0): - The dropout ratio for the attention probabilities. - head_dim (`int`, *optional*, defaults to 256): - Projection weights dimension in multi-head attention. - linear_conv_kernel_dim (`int`, *optional*, defaults to 4): - Kernel size of the convolution used in linear attention layers. - linear_key_head_dim (`int`, *optional*, defaults to 128): - Dimension of each key head in linear attention. - linear_value_head_dim (`int`, *optional*, defaults to 128): - Dimension of each value head in linear attention. - linear_num_key_heads (`int`, *optional*, defaults to 16): - Number of key heads used in linear attention layers. - linear_num_value_heads (`int`, *optional*, defaults to 32): - Number of value heads used in linear attention layers. - decoder_sparse_step (`int`, *optional*, defaults to 1): - The frequency of the MoE layer. - moe_intermediate_size (`int`, *optional*, defaults to 512): - Intermediate size of the routed expert. - shared_expert_intermediate_size (`int`, *optional*, defaults to 512): - Intermediate size of the shared expert. - num_experts_per_tok (`int`, *optional*, defaults to 10): - Number of selected experts. - num_experts (`int`, *optional*, defaults to 512): - Number of routed experts. - norm_topk_prob (`bool`, *optional*, defaults to `True`): - Whether to normalize the topk probabilities. - output_router_logits (`bool`, *optional*, defaults to `False`): - Whether or not the router logits should be returned by the model. Enabling this will also - allow the model to output the auxiliary loss, including load balancing loss and router z-loss. - router_aux_loss_coef (`float`, *optional*, defaults to 0.001): - The aux loss factor for the total loss. - mlp_only_layers (`list[int]`, *optional*, defaults to `[]`): - Indicate which layers use Qwen3NextMLP rather than Qwen3NextSparseMoeBlock - The list contains layer index, from 0 to num_layers-1 if we have num_layers layers - If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. - layer_types (`list[str]`, *optional*): - Types of each layer (attention or linear). - - ```python - >>> from transformers import Qwen3NextModel, Qwen3NextConfig - - >>> # Initializing a Qwen3Next style configuration - >>> configuration = Qwen3NextConfig() - - >>> # Initializing a model from the Qwen3-Next-80B-A3B style configuration - >>> model = Qwen3NextModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ``` - - """ # noqa: E501 - - model_type = "qwen3_next" - keys_to_ignore_at_inference = ["past_key_values"] - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.*.gate_proj": "colwise", - "layers.*.mlp.experts.*.up_proj": "colwise", - "layers.*.mlp.experts.*.down_proj": "rowwise", - "layers.*.mlp.shared_experts.gate_proj": "colwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - def __init__( - self, - vocab_size=151936, - hidden_size=2048, - intermediate_size=5632, - num_hidden_layers=48, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_parameters=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - decoder_sparse_step=1, - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=10, - num_experts=512, - norm_topk_prob=True, - output_router_logits=False, - router_aux_loss_coef=0.001, - mlp_only_layers=None, - layer_types=None, - **kwargs, - ): - if mlp_only_layers is None: - mlp_only_layers = [] - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - # Try to set `rope_scaling` if available, otherwise use `rope_parameters` - rope_scaling = kwargs.pop("rope_scaling", None) - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"} - rope_theta = kwargs.pop("rope_theta", 10000.0) - if "rope_theta" not in rope_parameters: - rope_parameters["rope_theta"] = rope_theta - partial_rotary_factor = kwargs.pop("partial_rotary_factor", 0.25) - if "partial_rotary_factor" not in rope_parameters: - rope_parameters["partial_rotary_factor"] = partial_rotary_factor - self.rope_parameters = rope_parameters - self.partial_rotary_factor = partial_rotary_factor - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - - self.layer_types = layer_types - if self.layer_types is None: - self.layer_types = [ - "linear_attention" if bool((i + 1) % 4) else "full_attention" - for i in range(self.num_hidden_layers) - ] - self.validate_layer_type() - - # linear attention part - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - - # MoE arguments - self.decoder_sparse_step = decoder_sparse_step - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.norm_topk_prob = norm_topk_prob - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - self.mlp_only_layers = mlp_only_layers - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["Qwen3NextConfig"] diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 9dee4af924ce..d79c1275b6e7 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -17,7 +17,6 @@ ConfigFormat, get_safetensors_params_metadata, ) -from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config from vllm.utils.torch_utils import common_broadcastable_dtype logger = init_logger(__name__) @@ -41,7 +40,7 @@ def get_per_layer_hf_configs( This is the only place that decides whether a checkpoint is heterogeneous and how many layers it has, so a convertor can declare variation that - Transformers does not express (see `Gemma4ModelArchConfigConvertor`). + Transformers does not express. Each pair is converted independently and the results are diffed, so a convertor that overrides this must not also collapse a varying field in @@ -745,27 +744,6 @@ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: == "vision" ) - def get_per_layer_hf_configs( - self, - ) -> list[tuple[PreTrainedConfig, PreTrainedConfig]] | None: - # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its - # full attention layers than on its sliding ones. Transformers >= 5.15.0 - # says so in the config; before that the values are flat attributes picked - # apart by `layer_types`, so build the per-layer configs here instead. Both - # then reach the base convertor as ordinary homogeneous configs. - if getattr(self.hf_text_config, "is_heterogeneous", False): - return super().get_per_layer_hf_configs() - text_config = self.hf_text_config - if not getattr(text_config, "layer_types", None): - return None - # `num_hidden_layers` is what the stack is built from, and it is what - # Transformers sizes `per_layer_config` by. `layer_types` can be longer, - # e.g. when a test truncates the layer count to build a small model. - return [ - (self.hf_config, gemma4_layer_config(text_config, layer_idx)) - for layer_idx in range(self.get_num_hidden_layers()) - ] - class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase): def _language_config(self) -> PreTrainedConfig: diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index b913545377d0..f2af4901e2e9 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -17,10 +17,8 @@ "FireRedASR2Processor", "FunASRProcessor", "GLM4VProcessor", - "Granite4VisionProcessor", "H2OVLProcessor", "Moondream3Processor", - "InternVLProcessor", "IsaacProcessor", "KimiAudioProcessor", "KimiK25Processor", @@ -55,9 +53,7 @@ "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2", "FunASRProcessor": "vllm.transformers_utils.processors.funasr", "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v", - "Granite4VisionProcessor": "vllm.transformers_utils.processors.granite4_vision", "H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl", - "InternVLProcessor": "vllm.transformers_utils.processors.internvl", "IsaacProcessor": "vllm.transformers_utils.processors.isaac", "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio", "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25", diff --git a/vllm/transformers_utils/processors/granite4_vision.py b/vllm/transformers_utils/processors/granite4_vision.py deleted file mode 100644 index 7972015adf9d..000000000000 --- a/vllm/transformers_utils/processors/granite4_vision.py +++ /dev/null @@ -1,78 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from fractions import Fraction - -from transformers import LlavaNextProcessor -from transformers.image_processing_utils import select_best_resolution - - -class Granite4VisionProcessor(LlavaNextProcessor): - """Processor for Granite 4 Vision. - - Extends LlavaNextProcessor to account for the Window Q-Former - downsampling when computing the number of image features. - - This processor is needed because the granite4_vision processor type - is not yet in the transformers version pinned by vLLM. - """ - - model_type = "granite4_vision" - - def __init__( - self, - image_processor=None, - tokenizer=None, - patch_size=None, - vision_feature_select_strategy=None, - chat_template=None, - image_token="", - num_additional_image_tokens=0, - downsample_rate=None, - **kwargs, - ): - super().__init__( - image_processor=image_processor, - tokenizer=tokenizer, - patch_size=patch_size, - vision_feature_select_strategy=vision_feature_select_strategy, - chat_template=chat_template, - image_token=image_token, - num_additional_image_tokens=num_additional_image_tokens, - ) - self.downsample_rate = downsample_rate - - def _get_number_of_features( - self, - orig_height: int, - orig_width: int, - height: int, - width: int, - ) -> int: - image_grid_pinpoints = self.image_processor.image_grid_pinpoints - - height_best_resolution, width_best_resolution = select_best_resolution( - [orig_height, orig_width], image_grid_pinpoints - ) - scale_height = height_best_resolution // height - scale_width = width_best_resolution // width - - patches_height = height // self.patch_size - patches_width = width // self.patch_size - if self.downsample_rate is not None: - ds_rate = Fraction(self.downsample_rate) - patches_height = int(patches_height * ds_rate) - patches_width = int(patches_width * ds_rate) - - unpadded_features, newline_features = self._get_unpadded_features( - orig_height, - orig_width, - patches_height, - patches_width, - scale_height, - scale_width, - ) - base_features = ( - patches_height * patches_width + self.num_additional_image_tokens - ) - return unpadded_features + newline_features + base_features From 2186b89a9a02fac11fa0e274f3201b9812dbf65c Mon Sep 17 00:00:00 2001 From: Damien Laine Date: Fri, 9 Oct 2026 11:40:29 +0200 Subject: [PATCH 276/952] [Bugfix][Model] Voxtral realtime: fix boot OOM and engine crash at max_model_len (#45022) Signed-off-by: Damien Laine Co-authored-by: Cyrus Leung --- .../generation/test_voxtral_realtime.py | 84 ++++++++++++++++++- .../speech_to_text/realtime/connection.py | 29 ++++++- .../model_executor/models/voxtral_realtime.py | 43 +++++++++- 3 files changed, 149 insertions(+), 7 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 09432b4494cb..e05930b0f36b 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -106,7 +106,8 @@ def tokenizer() -> MistralTokenizer: @pytest_asyncio.fixture -async def async_engine(): +async def async_engine(request): + engine_config = {**ENGINE_CONFIG, **getattr(request, "param", {})} gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9) from vllm.platforms import current_platform @@ -115,7 +116,7 @@ async def async_engine(): wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) - engine_args = AsyncEngineArgs(**ENGINE_CONFIG) + engine_args = AsyncEngineArgs(**engine_config) llm = AsyncLLM.from_engine_args(engine_args) try: yield llm @@ -275,3 +276,82 @@ async def input_stream(audio_buffer): f" got: {got!r}\n" f" expected: {expected!r}" ) + + +@pytest.mark.asyncio +async def test_voxtral_realtime_buffer_stops_at_max_model_len(tokenizer): + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + audio = Audio.from_file(AudioAsset("winning_call").get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + + max_model_len = 64 + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + # Unlimited feedback tokens so only the length cap can stop the stream. + await buffer.append_tokens([0] * max_model_len) + + n_tokens = 0 + async for prompt in buffer.get_input_stream(): + n_tokens += len(prompt["prompt_token_ids"]) + + assert len(audio_enc.tokens) < n_tokens < max_model_len + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_engine", [{"max_model_len": 128}], indirect=True) +async def test_voxtral_realtime_generator_max_model_len( + audio_assets, tokenizer, async_engine +): + # Lazy import to avoid CUDA-reinitialization error + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + max_model_len = async_engine.model_config.max_model_len + sampling_params = SamplingParams(temperature=0.0, max_tokens=1) + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + + async def input_stream(audio_buffer): + async for prompt in audio_buffer.get_input_stream(): + parsed_prompt = parse_model_prompt(async_engine.model_config, prompt) + (engine_input,) = await async_engine.renderer.render_cmpl_async( + [parsed_prompt] + ) + yield StreamingInput(prompt=engine_input) + + # winning_call is ~300 tokens long, well past max_model_len. + audio = Audio.from_file(audio_assets[1].get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + + output_tokens = [] + async for resp in async_engine.generate( + prompt=input_stream(buffer), + sampling_params=sampling_params, + request_id="session-max-len", + ): + tokens = resp.outputs[0].token_ids[-1:] + output_tokens.extend(tokens) + await buffer.append_tokens(tokens) + + assert resp.finished + assert 0 < len(output_tokens) < max_model_len + # The engine is still healthy after the session hit the cap. + assert not async_engine.errored diff --git a/vllm/entrypoints/speech_to_text/realtime/connection.py b/vllm/entrypoints/speech_to_text/realtime/connection.py index ed8ec7f862c5..4a5e6cef9a57 100644 --- a/vllm/entrypoints/speech_to_text/realtime/connection.py +++ b/vllm/entrypoints/speech_to_text/realtime/connection.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import contextlib import json from collections.abc import AsyncGenerator from http import HTTPStatus @@ -48,6 +49,7 @@ def __init__(self, websocket: WebSocket, serving: OpenAIServingRealtime): self.serving = serving self.audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue() self.generation_task: asyncio.Task | None = None + self.request_id: str | None = None self._is_connected = False self._is_model_validated = False @@ -204,6 +206,7 @@ async def _run_generation( 6. Cleans up the audio queue """ request_id = f"rt-{self.connection_id}-{uuid4()}" + self.request_id = request_id full_text = "" prompt_token_ids_len: int = 0 @@ -263,7 +266,9 @@ async def _run_generation( except Exception as e: logger.exception("Error in generation: %s", e) - await self.send_error(sanitize_message(str(e)), "processing_error") + # The socket may already be closed; nothing more we can do. + with contextlib.suppress(Exception): + await self.send_error(sanitize_message(str(e)), "processing_error") async def send( self, event: SessionCreated | TranscriptionDelta | TranscriptionDone @@ -282,8 +287,28 @@ async def cleanup(self): # Signal audio stream to stop self.audio_queue.put_nowait(None) - # Cancel generation task if running + # Abort the engine request. Cancelling the asyncio task alone does NOT + # abort the underlying engine request: a request parked in the + # VAD-waiting state keeps occupying a slot and generating into a now + # dead websocket (an orphaned request that never frees and starves + # other streams). Aborting also unblocks the generation task, which is + # otherwise stuck awaiting the next engine output. + if self.request_id is not None: + try: + await self.serving.engine_client.abort(self.request_id) + except Exception: + logger.exception("Failed to abort request %s", self.request_id) + self.request_id = None + + # Cancel the generation task if still running, and await it so the + # cancellation (and any generator cleanup) actually completes. if self.generation_task and not self.generation_task.done(): self.generation_task.cancel() + try: + await self.generation_task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Error awaiting cancelled generation task") logger.debug("Connection cleanup complete: %s", self.connection_id) diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 8531d7d18f39..72318976f1d4 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -120,9 +120,35 @@ def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor: return (base.unsqueeze(1) + offsets).view(-1) +class VoxtralRealtimeProcessingInfo(VoxtralProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # Each streaming update carries a single audio frame. + return {"audio": 1} + + def get_max_audio_array_len(self) -> int: + # Profile with the largest streaming chunk, the first one: it covers + # the prompt prefix (BOS + left padding + delay), later chunks are + # ~1 token. get_max_audio_tokens keeps max_model_len so the encoder + # cache stays as large as before: each paused session keeps its last + # chunk cached until it resumes, and a cache sized to one chunk + # deadlocks a few concurrent sessions. + tokenizer = self.get_tokenizer() + audio_encoder = tokenizer.instruct.audio_encoder + prompt_tokens = ( + tokenizer.instruct.start() + audio_encoder.encode_streaming_tokens() + ) + return len(prompt_tokens) * audio_encoder.audio_config.raw_audio_length_per_tok + + class VoxtralRealtimeBuffer: - def __init__(self, config: AudioConfig, prompt_tokens: list[int]) -> None: + def __init__( + self, + config: AudioConfig, + prompt_tokens: list[int], + max_model_len: int | None = None, + ) -> None: self._config = config + self._max_model_len = max_model_len _look_ahead_in_ms = self._config.streaming_look_ahead_ms _look_back_in_ms = self._config.streaming_look_back_ms @@ -165,7 +191,16 @@ async def append_tokens(self, tokens: Iterable[int]) -> None: await self._token_queue.put(token) async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]: + n_tokens = 0 for frame_size, num_tokens in self._generate_frame_size_and_num_tokens(): + n_tokens += num_tokens + if self._max_model_len is not None and n_tokens >= self._max_model_len: + logger.warning( + "Stopping realtime audio stream: reached max_model_len (%d).", + self._max_model_len, + ) + return + next_tokens = [await self._token_queue.get() for _ in range(num_tokens)] audio_arrays: list[np.ndarray] = ( @@ -198,7 +233,7 @@ async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]: @MULTIMODAL_REGISTRY.register_processor( VoxtralRealtimeMultiModalProcessor, - info=VoxtralProcessingInfo, + info=VoxtralRealtimeProcessingInfo, dummy_inputs=VoxtralDummyInputsBuilder, ) @support_torch_compile @@ -248,7 +283,9 @@ async def buffer_realtime_audio( # Get left/right padding audio left_pad, right_pad = audio_encoder.get_padding_audio() - buffer = VoxtralRealtimeBuffer(config, prompt_tokens) + buffer = VoxtralRealtimeBuffer( + config, prompt_tokens, max_model_len=model_config.max_model_len + ) # Feed audio with padding into buffer in background async def feed_audio(): From 59ffd73e8597d860730f4a8c4f778f27c75a0dfa Mon Sep 17 00:00:00 2001 From: Warren Deng <35615489+warrendeng@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:46:41 -0700 Subject: [PATCH 277/952] [Bugfix][Kernel] Restore FP fusion for indexer Q quantization (#60716) Signed-off-by: Warren Deng --- vllm/models/deepseek_v4/common/ops/fused_indexer_q.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index 669186edbb1b..54b1b9be7e89 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -305,6 +305,7 @@ def __call__( FP8_MAX=fp8_max, USE_FNUZ=use_fnuz, USE_EXPLICIT_FMA=current_platform.is_rocm(), + enable_fp_fusion=True, num_warps=1, ) @@ -540,6 +541,7 @@ def __call__( MXFP4_BLOCK=MXFP4_BLOCK_SIZE, index_weights_stride=index_weights.stride(0), index_weights_out_stride=index_weights_out.stride(0), + enable_fp_fusion=True, num_warps=1, ) From b8bd3357548ac07b1877053b6c5ba415763bad36 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Fri, 9 Oct 2026 06:26:27 -0400 Subject: [PATCH 278/952] [CI] Enroll five more steps in automatic sharding (#60662) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/entrypoints.yaml | 3 +++ .buildkite/test_areas/pytorch.yaml | 1 + 3 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 53499d4a986a..29aec98b07ee 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -108,6 +108,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) E2E Core" device: h200_18gb key: e2e-core-1-gpu + automatic_shard: true timeout_in_minutes: 65 source_file_dependencies: - vllm/v1/ diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 34a190e810b9..d31c66a8cbd8 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -26,6 +26,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Unit" device: h200_35gb key: entrypoints-unit-tests + automatic_shard: true timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -168,6 +169,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Responses API)" device: h200_35gb key: entrypoints-integration-responses-api + automatic_shard: true timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -188,6 +190,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (Speech to Text)" device: h200_35gb key: entrypoints-integration-speech_to_text + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 0ba608e72580..60061af41488 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -123,6 +123,7 @@ steps: - label: ":nvidia: (L4) PyTorch Compilation Passes" key: pytorch-compilation-passes-unit-tests + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 45 From efd141047d8e879536806a86e3f93c02e6164add Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Fri, 9 Oct 2026 06:26:47 -0400 Subject: [PATCH 279/952] [CI] Enroll a third batch of five steps in automatic sharding (#60757) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/lm_eval.yaml | 1 + .buildkite/test_areas/misc.yaml | 1 + .buildkite/test_areas/quantization.yaml | 1 + 5 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 45b3dc768e6d..b3bfa9f09708 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -33,6 +33,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) CUDAGraph" device: h200_18gb key: cudagraph + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: - tests/v1/cudagraph diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 29aec98b07ee..26485e0a172d 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -38,6 +38,7 @@ steps: - label: ":nvidia: (L4) V1 Engine" key: engine-1-gpu + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 45 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 49eea0b78f7c..06607951057b 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) LM Eval Small Models" device: h200_35gb key: lm-eval-small-models + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: - csrc/ diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 52ff1d65b3ee..9763b2ce7de4 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) V1 Sample" key: v1-sample + automatic_shard: true timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: &v1-sample-logits-deps diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 429b01f27320..55cef9fbd33a 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -80,6 +80,7 @@ steps: - label: ":nvidia: (B200) Quantized MoE" key: quantized-moe-test-b200 + automatic_shard: true timeout_in_minutes: 120 working_dir: "/vllm-workspace/" device: b200-k8s From 4e318ac56319476be1bc81f44acaf99516891afb Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 9 Oct 2026 18:38:01 +0800 Subject: [PATCH 280/952] [XPU] fix online fp8_per_channel quantization (#56027) Signed-off-by: Yan Ma --- .buildkite/intel_jobs/quantization.yaml | 4 +++- tests/quantization/test_fp8_per_channel.py | 4 +++- vllm/platforms/xpu.py | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.buildkite/intel_jobs/quantization.yaml b/.buildkite/intel_jobs/quantization.yaml index 8f1dfebe5001..af1b23b8a004 100644 --- a/.buildkite/intel_jobs/quantization.yaml +++ b/.buildkite/intel_jobs/quantization.yaml @@ -25,5 +25,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]" - && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py' + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_fp8_per_channel.py + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/models/quantization/test_fp8_per_channel.py --deselect="tests/models/quantization/test_fp8_per_channel.py::test_fp8_per_channel_logprobs[moe]"' diff --git a/tests/quantization/test_fp8_per_channel.py b/tests/quantization/test_fp8_per_channel.py index 2c3fddcb0fe4..d4f2fe264e6c 100644 --- a/tests/quantization/test_fp8_per_channel.py +++ b/tests/quantization/test_fp8_per_channel.py @@ -61,7 +61,9 @@ def test_scaled_fp8_quant_per_channel_shape() -> None: weight to `ops.scaled_fp8_quant` with `use_per_token_if_dynamic=True` yields one scale per output row -- a [N, 1] fp32 tensor. """ - x = (torch.randn(size=(96, 256), device="cuda") * 13).to(torch.bfloat16) + x = (torch.randn(size=(96, 256), device=current_platform.device_type) * 13).to( + torch.bfloat16 + ) y, s = ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True) assert y.shape == (96, 256) assert y.dtype == current_platform.fp8_dtype() diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 355a15094db0..992bfae0534c 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -123,6 +123,7 @@ class XPUPlatform(Platform): "mxfp8", "fp8_per_tensor", "fp8_per_block", + "fp8_per_channel", "online", "gpt_oss_mxfp4", "modelopt", From 5cadabfafde31dcccf4e1f773d00cc6e8a134cfb Mon Sep 17 00:00:00 2001 From: yuvalluria <131433771+yuvalluria@users.noreply.github.com> Date: Fri, 9 Oct 2026 13:41:29 +0300 Subject: [PATCH 281/952] [Docs] Add Qwen2.5-Coder-7B-Instruct to batch invariance tested models (#46368) Signed-off-by: Yuval Luria Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/features/batch_invariance.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index b40ff4da8cc6..03d07896cce4 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -126,7 +126,7 @@ Batch invariance has been tested and verified on the following models: - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ` - **Qwen3-VL (Vision-Language)**: `Qwen/Qwen3-VL-2B-Instruct`, `Qwen/Qwen3-VL-4B-Instruct` (single image and video inputs) - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8` -- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` +- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`, `Qwen/Qwen2.5-Coder-7B-Instruct` - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-1B-Instruct`, `meta-llama/Llama-3.2-3B-Instruct` for example - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` - **Mistral**: `mistralai/Mistral-7B-v0.3` From b027ac831253e83ff605a2ff5d4aa7c00619802c Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Fri, 9 Oct 2026 18:58:07 +0800 Subject: [PATCH 282/952] [XPU] Support register KV offload mmap region as pinned host memory on XPU (#51956) Signed-off-by: Chaojun Zhang --- .../cpu/test_shared_offload_region.py | 81 +++++++--------- tests/v1/kv_offload/test_factory.py | 3 + vllm/v1/kv_offload/cpu/gpu_worker.py | 45 ++------- vllm/v1/kv_offload/cpu/host_register.py | 97 +++++++++++++++++++ .../kv_offload/cpu/shared_offload_region.py | 21 ++-- vllm/v1/kv_offload/cpu/spec.py | 4 +- 6 files changed, 151 insertions(+), 100 deletions(-) create mode 100644 vllm/v1/kv_offload/cpu/host_register.py diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index f84e1bd7296c..b7259e6fd621 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -772,17 +772,14 @@ def test_cleanup_idempotent(iid): def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): """cleanup() must release every chunk from a split host registration.""" r = _make_region(iid) - cudart = MagicMock() - cudart.cudaHostUnregister.return_value = MagicMock(value=0) - monkeypatch.setattr(region_module, "current_platform", MagicMock()) - region_module.current_platform.is_cuda_alike.return_value = True - monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart) + unregister = MagicMock() + monkeypatch.setattr(region_module, "host_unregister", unregister) r.pinned_addresses = [0x100000, 0x200000, 0x300000] r.is_pinned = True r.cleanup() - assert cudart.cudaHostUnregister.call_args_list == [ + assert unregister.call_args_list == [ call(0x300000), call(0x200000), call(0x100000), @@ -791,40 +788,35 @@ def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch): @pytest.fixture def host_register(monkeypatch): - """Pretend to run on CUDA with host registrations capped at seven pages. - - Registration goes to a mock CudaRTLibrary and the region's cleanup to a - separate mock torch runtime. Both stay installed through the cleanup that - ``_region`` runs on exit. + """Mock the cross-platform host_register/host_unregister helpers with + registrations capped at seven pages. The same unregister mock backs both + pin_mmap_region's rollback path and the region's cleanup() on exit from + ``_region``. """ - cudart = MagicMock() - cudart.cudaHostRegister.return_value = 0 - cudart.cudaHostUnregister.return_value = 0 - cudart.cudaGetLastError.return_value = 2 - torch_cudart = MagicMock() - torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0) - monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True) + register = MagicMock(return_value=True) + unregister = MagicMock() monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE) - monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart) - monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart) - return cudart, torch_cudart + monkeypatch.setattr(gpu_worker, "host_register", register) + monkeypatch.setattr(gpu_worker, "host_unregister", unregister) + monkeypatch.setattr(region_module, "host_unregister", unregister) + return register, unregister def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register): """Chunks end on row boundaries: 3-page rows under a 7-page cap register as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1.""" - cudart, torch_cudart = host_register + register, unregister = host_register with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: gpu_worker.pin_mmap_region(region) base = region._base.data_ptr() - assert cudart.mock_calls == [ - call.cudaHostRegister(base, 6 * PAGE_SIZE), - call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE), - call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE), + assert register.call_args_list == [ + call(base, 6 * PAGE_SIZE), + call(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE), + call(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE), ] assert region.is_pinned - assert torch_cudart.cudaHostUnregister.call_args_list == [ + assert unregister.call_args_list == [ call(base + 12 * PAGE_SIZE), call(base + 6 * PAGE_SIZE), call(base), @@ -833,10 +825,10 @@ def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register): @pytest.mark.parametrize("fail_at", [0, 1, 2]) def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at): - """A failed chunk drains its pending error and unregisters the chunks - before it on the same handle, so the region is pinned whole or not at all.""" - cudart, torch_cudart = host_register - cudart.cudaHostRegister.side_effect = [0] * fail_at + [2] + """A failed chunk unregisters the chunks before it on the same range, so + the region is pinned whole or not at all.""" + register, unregister = host_register + register.side_effect = [True] * fail_at + [False] with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region: gpu_worker.pin_mmap_region(region) base = region._base.data_ptr() @@ -845,25 +837,22 @@ def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)] failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE - assert cudart.mock_calls == [ - *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered), - call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size), - call.cudaGetLastError(), - *(call.cudaHostUnregister(address) for address in reversed(registered)), + assert register.call_args_list == [ + *(call(address, 6 * PAGE_SIZE) for address in registered), + call(base + fail_at * 6 * PAGE_SIZE, failed_size), + ] + assert unregister.call_args_list == [ + call(address) for address in reversed(registered) ] - torch_cudart.cudaHostUnregister.assert_not_called() -def test_pin_mmap_region_without_cuda_runtime_stays_pageable( - iid, host_register, monkeypatch +def test_pin_mmap_region_when_host_register_unsupported_stays_pageable( + iid, monkeypatch ): - """A CUDA runtime that cannot be loaded leaves the region pageable instead - of failing worker startup.""" - - def missing_runtime(): - raise AssertionError("libcudart is not loaded in the current process") - - monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime) + """A platform without host registration support (host_register() + returning False) leaves the region pageable instead of failing worker + startup.""" + monkeypatch.setattr(gpu_worker, "host_register", lambda ptr, num_bytes: False) with _region(iid) as region: gpu_worker.pin_mmap_region(region) assert not region.is_pinned diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 01ef756cae53..84b9da36e06b 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -307,6 +307,7 @@ def test_cpu_spec_replicated_disabled_without_shared_region( # those private buffers empty and corrupting subsequent loads. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT @@ -341,6 +342,7 @@ def test_cpu_spec_replicated_layout_truth_matrix( # actually allocates on the shared region. import vllm.v1.kv_offload.cpu.spec as cpu_spec_module + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr( cpu_spec_module.current_platform, "is_cuda_alike", lambda: cuda_alike ) @@ -417,6 +419,7 @@ def fake_worker_ctor(**kwargs): worker_calls.append(kwargs) return MagicMock() + monkeypatch.setattr(cpu_spec_module.current_platform, "is_xpu", lambda: False) monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: rocm) monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: rocm) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index a56febf702ed..f6053480a7ab 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -11,7 +11,6 @@ import torch from vllm import _custom_ops as ops -from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, triton @@ -27,6 +26,7 @@ OffloadingWorker, TransferResult, ) +from vllm.v1.kv_offload.cpu.host_register import host_register, host_unregister from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.swap_blocks_triton import ( THRESHOLD_BYTES, @@ -200,26 +200,7 @@ def _canonical_block_sizes( def pin_mmap_region(region: SharedOffloadRegion) -> None: """Register row-aligned chunks, rolling back on failure.""" - if not current_platform.is_cuda_alike(): - logger.info( - "Skipping mmap host registration on %s; cudaHostRegister is only " - "available on CUDA/ROCm.", - current_platform.device_name, - ) - return - rank = region.rank - try: - cudart = CudaRTLibrary() - except (AssertionError, AttributeError, OSError): - logger.warning( - "Could not load the CUDA runtime for host registration on rank=%d; " - "the offload region stays pageable", - rank, - exc_info=True, - ) - return - base_ptr = region._base.data_ptr() total_size = region.total_size_bytes # Chunks end on block-row boundaries, which are page aligned, so neither the @@ -228,42 +209,30 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None: rows_per_chunk = max(MAX_HOST_REGISTER_CHUNK_BYTES // region._row_stride, 1) chunk_size = rows_per_chunk * region._row_stride - # Register, drain and roll back through the same runtime handle, so a - # failed chunk leaves neither a pending error nor a partly pinned region. + # Register, drain and roll back through the same helper, so a failed + # chunk leaves neither a pending error nor a partly pinned region. addresses: list[int] = [] for offset in range(0, total_size, chunk_size): address = base_ptr + offset size = min(chunk_size, total_size - offset) - result = cudart.cudaHostRegister(address, size) - if result == 0: + if host_register(address, size): addresses.append(address) continue - cudart.cudaGetLastError() logger.warning( - "cudaHostRegister failed for rank=%d at %.2f of %.2f GB (code=%d); " + "host_register failed for rank=%d at %.2f of %.2f GB; " "the offload region stays pageable", rank, offset / 1e9, total_size / 1e9, - result, ) for registered in reversed(addresses): - unregister_result = cudart.cudaHostUnregister(registered) - if unregister_result != 0: - cudart.cudaGetLastError() - logger.warning( - "cudaHostUnregister failed for rank=%d at %#x (code=%d); " - "that chunk stays registered until the process exits", - rank, - registered, - unregister_result, - ) + host_unregister(registered) return region.pinned_addresses.extend(addresses) region.is_pinned = True logger.debug( - "cudaHostRegister rank=%d %.2f GB in %d chunk(s)", + "Host-registered mmap region rank=%d %.2f GB in %d chunk(s)", rank, total_size / 1e9, len(addresses), diff --git a/vllm/v1/kv_offload/cpu/host_register.py b/vllm/v1/kv_offload/cpu/host_register.py new file mode 100644 index 000000000000..37448857409f --- /dev/null +++ b/vllm/v1/kv_offload/cpu/host_register.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Registration of existing host memory for pinned DMA transfers. + +Host allocations that the offloading workers share across processes (the +`SharedOffloadRegion` mmap) cannot be allocated through the accelerator's +pinned allocator, so they must be registered after the fact. Pinning is a +pure bandwidth optimization: when registration is unavailable the transfers +still run correctly, just as pageable (slower) copies. +""" + +import torch + +from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def _load_cudart() -> CudaRTLibrary | None: + try: + return CudaRTLibrary() + except (AssertionError, AttributeError, OSError): + logger.warning( + "Could not load the CUDA runtime for host registration; " + "the offload region stays pageable", + exc_info=True, + ) + return None + + +def host_register(ptr: int, num_bytes: int) -> bool: + """Register a host range as device-accessible pinned memory. + + Args: + ptr: Base address of the host range. + num_bytes: Size of the host range in bytes. + + Returns: + True if the range was registered and later needs `host_unregister`. + + """ + if not ptr or num_bytes <= 0: + return False + + if current_platform.is_cuda_alike(): + cudart = _load_cudart() + if cudart is None: + return False + result = cudart.cudaHostRegister(ptr, num_bytes) + if result != 0: + cudart.cudaGetLastError() + logger.warning( + "cudaHostRegister failed (code=%d) — transfers will still work " + "but may be slower (unpinned DMA)", + result, + ) + return False + return True + + if current_platform.is_xpu(): + if not torch.ops._C.xpu_host_register(ptr, num_bytes): + logger.warning( + "xpu_host_register failed — transfers will still work " + "but may be slower (unpinned DMA)" + ) + return False + return True + + logger.info_once( + "Host registration is not implemented on %s; KV offload transfers will " + "use unpinned host memory.", + current_platform.device_name, + ) + return False + + +def host_unregister(ptr: int) -> None: + """Release a host range previously registered by `host_register`.""" + if not ptr: + return + + if current_platform.is_cuda_alike(): + cudart = _load_cudart() + if cudart is None: + return + result = cudart.cudaHostUnregister(ptr) + if result != 0: + cudart.cudaGetLastError() + logger.warning("cudaHostUnregister failed (code=%d)", result) + return + + if current_platform.is_xpu(): + if not torch.ops._C.xpu_host_unregister(ptr): + logger.warning("xpu_host_unregister failed") + return diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index af183fbc4f06..7e0e97aa1219 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -14,7 +14,7 @@ check_shm_free_space, ) from vllm.logger import init_logger -from vllm.platforms import current_platform +from vllm.v1.kv_offload.cpu.host_register import host_unregister logger = init_logger(__name__) @@ -343,20 +343,11 @@ def create_kv_memoryview(self) -> memoryview: def cleanup(self) -> None: if self.is_pinned and self._base is not None: - if current_platform.is_cuda_alike(): - base_ptr = self._base.data_ptr() - addresses = self.pinned_addresses or [base_ptr] - for address in reversed(addresses): - result = torch.cuda.cudart().cudaHostUnregister(address) - if result.value != 0: - logger.warning( - "cudaHostUnregister failed for rank=%d, " - "address=%#x (code=%d)", - self.rank, - address, - result.value, - ) - self.pinned_addresses.clear() + base_ptr = self._base.data_ptr() + addresses = self.pinned_addresses or [base_ptr] + for address in reversed(addresses): + host_unregister(address) + self.pinned_addresses.clear() self.is_pinned = False # Release views before _base: each view holds a _base reference and a # direct StorageImpl reference. Freeing views first lets both refcounts diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index fddf56608590..bf289d37a65b 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -162,7 +162,9 @@ def get_manager(self) -> OffloadingManager: def _uses_shared_region(self) -> bool: """Whether the worker CPU buffer is the shared mmap region (vs a private per-rank tensor); replicated-layout dedup is gated on this being True.""" - return current_platform.is_cuda_alike() and not current_platform.is_rocm() + return ( + current_platform.is_cuda_alike() and not current_platform.is_rocm() + ) or current_platform.is_xpu() def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: mmap_region: SharedOffloadRegion | None = None From 28d76f3bad3df9270019b850a5c29f241b3d7bbf Mon Sep 17 00:00:00 2001 From: Shreyansh Jain Date: Fri, 9 Oct 2026 16:46:13 +0530 Subject: [PATCH 283/952] [Bugfix][Structured Output] Reject unsupported regex inside groups (#60359) Signed-off-by: Shreyansh Jain Signed-off-by: Artem Perevedentsev Co-authored-by: Artem Perevedentsev --- .../test_backend_outlines.py | 73 +++++++++++++++++++ vllm/v1/structured_output/backend_outlines.py | 8 +- 2 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 tests/v1/structured_output/test_backend_outlines.py diff --git a/tests/v1/structured_output/test_backend_outlines.py b/tests/v1/structured_output/test_backend_outlines.py new file mode 100644 index 000000000000..abff1e7f51b6 --- /dev/null +++ b/tests/v1/structured_output/test_backend_outlines.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.exceptions import VLLMValidationError +from vllm.v1.structured_output.backend_outlines import validate_regex_is_buildable + +pytestmark = pytest.mark.cpu_test + + +# ================================================ +# Test validate_regex_is_buildable functionality +# ================================================ + + +class TestValidateRegexIsBuildable: + class TestPR60359Regressions: + """This class implements regression tests for PR #60359. + Initial bug report issue #60344. + validate_regex_is_buildable misses \b and backref inside groups + So, we check here that these features are rejected by us. + """ + + @pytest.mark.parametrize( + "regex", + [ + # top-level word boundary \b + r"a\b", + # \b inside capturing group + r"(a\b)", + # word boundary \b in nested group + r"((a\b))", + # top-level backref + r"(a)\1", + # backref in nested group + r"((a)\2)", + # positive lookahead + r"(a(?=b))", + # inline flag group + word boundary + r"(?i:a\b)", + # counted repeat over a group containing \b + r"(?:(a\b)){2}", + # lazy repeat + word boundary + r"(?:a\b)*?", + # possessive repeat + word boundary + r"(?:a\b)*+", + # atomic group + word boundary + r"(?>a\b)", + ], + ) + def test_unsupported_regex(self, regex): + with pytest.raises(VLLMValidationError): + validate_regex_is_buildable(regex) + + @pytest.mark.parametrize( + "regex", + [ + # [\b] is character class, not word-boundary \b + r"[\b]", + r"(a[\b]c)", + # escaped backslash + "b", not word-boundary \b + r"a\\b", + r"(a\\bc)", + # no \b or backref here, must be ok + r"(abc)", + # BRANCH and MAX_REPEAT must be ok + r"(a|b)+c", + r"(?:ab)*?c", + ], + ) + def test_supported_regex(self, regex): + validate_regex_is_buildable(regex) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 655b8b14631f..1cfc6c8f5d1a 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -327,9 +327,11 @@ def _check_unsupported(parsed) -> None: for branch in tval[1]: _check_unsupported(branch) - # tval is (min, max, subpattern) - elif ttype == _parser.MAX_REPEAT: - _check_unsupported(tval[2]) + else: + nested = tval if isinstance(tval, tuple) else (tval,) + for sub in nested: + if isinstance(sub, _parser.SubPattern): + _check_unsupported(sub) def validate_regex_is_buildable(pattern: str) -> None: From 83ed0c7c0944e9d124589d8fa3f82a019343a726 Mon Sep 17 00:00:00 2001 From: Luca Motz Date: Fri, 9 Oct 2026 13:30:11 +0200 Subject: [PATCH 284/952] [Bugfix][NIXL] Fix receive post-process for different P/D block sizes (#58860) Signed-off-by: Luca Motz Co-authored-by: Claude Opus 5.5 --- .../kv_connector/unit/test_nixl_connector.py | 16 ++++- .../unit/test_nixl_connector_hma.py | 66 ++++++++++++------- .../kv_transfer/kv_connector/utils.py | 44 ++----------- .../kv_connector/v1/nixl/base_worker.py | 54 ++++++--------- 4 files changed, 83 insertions(+), 97 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 52ac746391ff..d918c38fbf8c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1128,17 +1128,24 @@ def test_concurrent_load_kv( return raise TimeoutError("Took too long to complete async handshake.") + @pytest.mark.parametrize( + ("kv_cache_layout", "enable_permute_local_kv"), + # enable_permute_local_kv only covers a local LBNHC cache. + [("LBHNC", False), ("BLHNC", True)], + ) @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_fails_on_kv_cache_layout_mismatch( - self, default_vllm_config, dist_init + self, default_vllm_config, dist_init, kv_cache_layout, enable_permute_local_kv ): """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ - vllm_config = create_vllm_config() + vllm_config = create_vllm_config( + enable_permute_local_kv=enable_permute_local_kv + ) # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 @@ -1151,7 +1158,10 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) ) connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 + vllm_config, + connector.engine_id, + hand_shake_latency=0, + kv_cache_layout=kv_cache_layout, ) worker = connector.connector_worker diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index f016b26475f0..c665ce0c5b88 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1404,46 +1404,66 @@ def test_map_block_ids_for_block_size_ratio_hybrid(): @pytest.mark.cpu_test -def test_post_process_zeroes_untransferred_tail(): - """The untransferred sub-blocks of the last local block are zeroed on - receive; mamba state caches are untouched by the attention permute.""" +@pytest.mark.parametrize( + ("kv_cache_layout", "enable_permute_local_kv"), + [("LBHNC", False), ("LBNHC", False), ("LBNHC", True)], +) +def test_post_process_zeroes_untransferred_tail( + kv_cache_layout, enable_permute_local_kv +): + """Received remote sub-blocks are regrouped per head and the untransferred + sub-blocks of the last local block are zeroed on receive, once per block + although two attention groups alias the tensor.""" from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) - from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec ratio = 4 block_tokens = 8 # 2 tokens per remote sub-block + num_kv_heads = 2 # fewer than ratio worker = MagicMock(spec=NixlConnectorWorker) - worker._group_spec_types = (FullAttentionSpec, MambaSpec) worker.transfer_topo = MagicMock() worker.device_type = "cpu" - worker.enable_permute_local_kv = False - attn_cache = torch.ones(6, block_tokens, 2, 4) - mamba_cache = torch.ones(6, 16) - worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + worker.enable_permute_local_kv = enable_permute_local_kv + worker.kv_cache_layout = kv_cache_layout + # Attention caches are [B, H, N, C]; distinct values per head and token. + expected = torch.arange(6 * num_kv_heads * block_tokens * 4).view( + 6, num_kv_heads, block_tokens, 4 + ) + # Blocks [2, 3] as received: `ratio` head-major remote sub-blocks each. + received = ( + expected[2:4] + .unflatten(2, (ratio, -1)) + .transpose(1, 2) + .reshape(2, num_kv_heads, block_tokens, 4) + ) + if kv_cache_layout == "LBNHC": + attn_cache = expected.transpose(1, 2).contiguous().transpose(1, 2) + if enable_permute_local_kv: + # The remote is LBHNC: its bytes land in token-major memory. + attn_cache.transpose(1, 2)[2:4] = received.view(2, block_tokens, -1, 4) + # Otherwise token-major blocks receive the sub-blocks in token order. + else: + attn_cache = expected.clone() + attn_cache[2:4] = received + worker.device_kv_caches = {"attn.0": attn_cache, "swa.0": attn_cache} fa_group = MagicMock(layer_names=["attn.0"]) - ssm_group = MagicMock(layer_names=["mamba.0"]) - worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, ssm_group]) - # The cached property filters mamba layers out of the permuted caches. - attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) - assert len(attn_caches) == 1 and attn_caches[0] is attn_cache - worker._attention_kv_caches = attn_caches + swa_group = MagicMock(layer_names=["swa.0"]) + worker.kv_cache_config = MagicMock(transfer_groups=[fa_group, swa_group]) _bind_worker_method(worker, "post_process_device_kv_on_receive") - # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. - worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + # Group 0 request in blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [(0, [2, 3], 6)]) # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). - assert torch.all(attn_cache[2] == 1) - assert torch.all(attn_cache[3, :4] == 1) - assert torch.all(attn_cache[3, 4:] == 0) - # Untouched blocks and the mamba cache keep their content. - assert torch.all(attn_cache[4] == 1) - assert torch.all(mamba_cache == 1) + assert torch.equal(attn_cache[2], expected[2]) + assert torch.equal(attn_cache[3, :, :4], expected[3, :, :4]) + assert torch.all(attn_cache[3, :, 4:] == 0) + # Untouched blocks keep their content. + assert torch.equal(attn_cache[4], expected[4]) @pytest.mark.cpu_test diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index a0ead27c2ff1..d314d6107ad7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -256,8 +256,6 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): """ blocks_to_update = cache.index_select(0, indices) - # use physical order - blocks_to_update = blocks_to_update.permute(0, 2, 1, 3) n_kv_heads, block_size, head_size = blocks_to_update.shape[1:] remote_block_size = block_size // block_size_ratio n_blocks = block_size_ratio @@ -267,56 +265,28 @@ def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio): .permute(0, 2, 1, 3, 4) .flatten(2, 3) ) - permuted_blocks = permuted_blocks.permute(0, 2, 1, 3) - cache.index_copy_(0, indices, permuted_blocks) - - -def kv_postprocess_layout_on_receive(cache, indices): - """Transforms the layout of received KV cache blocks to the local format. - - This method corrects layout mismatches from direct memory copies by - permuting the tensor dimensions. - - 4D cache: - - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]` - - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]` - 5D cache: - - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]` - - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]` - - Implementation: - - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout - - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size - - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back - - """ - blocks_to_update = cache.index_select(0, indices) - target_shape = list(blocks_to_update.shape) - target_shape[0] = -1 - inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3] - src_shape = tuple(target_shape[i] for i in inv_order) - blocks_to_update = cache.index_select(0, indices) - permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order) cache.index_copy_(0, indices, permuted_blocks) def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio): """Transforms the layout of received KV cache to the local block_size and LBHNC. - (Only works for local blocksize > remote blocksize) + (Only works for local blocksize >= remote blocksize) prefill is LBHNC, smaller block_size decode(local) is LBNHC, larger block_size """ blocks_to_update = cache.index_select(0, indices) - block_size, n_kv_heads, head_size = blocks_to_update.shape[1:] + n_kv_heads, block_size, head_size = blocks_to_update.shape[1:] remote_block_size = block_size // block_size_ratio n_blocks = block_size_ratio + # View the received bytes in memory order as the remote head-major blocks. permuted_blocks = ( - blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) - .permute(0, 1, 3, 2, 4) - .flatten(1, 2) + blocks_to_update.transpose(1, 2) + .reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size) + .permute(0, 2, 1, 3, 4) + .flatten(2, 3) ) cache.index_copy_(0, indices, permuted_blocks) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7ce6e93ebb46..38a656954eee 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -14,7 +14,6 @@ from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor -from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -31,7 +30,6 @@ get_current_attn_backends, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, - kv_postprocess_layout_on_receive, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( CopyBlocksOp, @@ -2596,6 +2594,7 @@ def _validate_remote_agent_handshake( if ( self.kv_transfer_config.enable_permute_local_kv and nixl_agent_meta.kv_cache_layout == "LBHNC" + and kv_cache_layout == "LBNHC" ): logger.info( "Remote is LBHNC and local is LBNHC, enabled additional permute " @@ -2607,9 +2606,11 @@ def _validate_remote_agent_handshake( self.enable_permute_local_kv = True else: raise RuntimeError( - "Heterogeneous TP expects same kv_cache_layout. " - "Or enable experimental feature to use HND to NHD support by " - "setting 'enable_permute_local_kv'=True in --kv-transfer-config." + "Heterogeneous TP expects same kv_cache_layout (local " + f"{kv_cache_layout}, remote {nixl_agent_meta.kv_cache_layout}). " + "Or enable experimental feature to receive LBHNC into a local " + "LBNHC cache by setting 'enable_permute_local_kv'=True " + "in --kv-transfer-config." ) # if remote_agent used attn is not same as local, # hint heterogenuous attn post process @@ -2795,26 +2796,10 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): "d2h", ) - @cached_property - def _attention_kv_caches(self) -> list[torch.Tensor]: - """Device KV caches of attention layers (mamba states excluded), - as consumed by the receive post-process.""" - assert self.device_kv_caches, ( - "_attention_kv_caches accessed before register_kv_caches" - ) - mamba_layers = { - name - for g, group in enumerate(self.kv_cache_config.transfer_groups) - if _is_ssm_spec(self._group_spec_types[g]) - for name in group.layer_names - } - kv_caches = self.device_kv_caches - return [cache for name, cache in kv_caches.items() if name not in mamba_layers] - def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[tuple[list[int], int]], + block_ids_list: list[tuple[int, list[int], int]], convert: bool = True, ): """Post process device kv cache after receiving from remote. @@ -2863,9 +2848,11 @@ def post_process_device_kv_on_receive( block_size_ratio, ) - attn_caches = self._attention_kv_caches - device = attn_caches[0].device - for block_ids, covered_sub_blocks in block_ids_list: + for group, block_ids, covered_sub_blocks in block_ids_list: + attn_caches = [ + self.device_kv_caches[name] + for name in self.kv_cache_config.transfer_groups[group].layer_names + ] # Blocks the transfer didn't write: the token tail of the last # partially covered block, then everything beyond it. covered_blocks, sub_blocks_in_last = divmod( @@ -2875,17 +2862,16 @@ def post_process_device_kv_on_receive( has_stale = first_stale < len(block_ids) indices = None if convert or has_stale: - indices = async_tensor_h2d(block_ids, device, torch.long) + indices = async_tensor_h2d(block_ids, attn_caches[0].device, torch.long) if convert: for cache in attn_caches: - if self.enable_permute_local_kv and block_size_ratio > 1: + if self.enable_permute_local_kv: kv_postprocess_blksize_and_layout_on_receive( cache, indices, block_size_ratio ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: + elif KVCacheLayout[self.kv_cache_layout].is_block_contiguous: + # Token-major blocks get the sub-blocks in token order. kv_postprocess_blksize_on_receive( cache, indices, block_size_ratio ) @@ -2893,10 +2879,10 @@ def post_process_device_kv_on_receive( if sub_blocks_in_last: last_block_id = block_ids[covered_blocks] for cache in attn_caches: - # Both post-processed layouts leave tokens on dim 1. - sub_block_tokens = cache.shape[1] // block_size_ratio + # Attention caches are [B, H, N, C]: tokens on dim 2. + sub_block_tokens = cache.shape[2] // block_size_ratio zero_from = sub_blocks_in_last * sub_block_tokens - cache[last_block_id, zero_from:].zero_() + cache[last_block_id, :, zero_from:].zero_() if has_stale: assert indices is not None stale_ids = indices[first_stale:] @@ -3068,7 +3054,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults: len(meta.remote.block_ids[g]), ) block_ids_for_blocksize_post_process[block_size_ratio].append( - (local_group, covered_sub_blocks) + (g, local_group, covered_sub_blocks) ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: From 7368fadebd6c3f0c470d0e64b62cf6eea3b0987d Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Fri, 9 Oct 2026 19:34:38 +0800 Subject: [PATCH 285/952] [CI Failure] Move MTEB tests to h200_35gb (#60817) Signed-off-by: wang.yuqi --- .buildkite/test_areas/models_language.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index f123f2c57b13..002fad8ac6c3 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -220,10 +220,10 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 18GB) Language Models (MTEB)" +- label: ":nvidia: (H200 MIG 35GB) Language Models (MTEB)" key: language-models-test-mteb timeout_in_minutes: 68 - device: h200_18gb + device: h200_35gb optional: true source_file_dependencies: - vllm/ From d095e61c0b389753ae86ccd7b2b7d1f2e1f5cd58 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Fri, 9 Oct 2026 07:35:48 -0400 Subject: [PATCH 286/952] [Bugfix][Spec Decode] Fix GLM-5.3-Flash DFlash2/EAGLE3 aux hidden-state capture with the upstream config (#60763) Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- vllm/models/glm5next/common/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index fd7c0263bdda..e35887ca7c54 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -788,7 +788,7 @@ def _aux_hidden_state( return hidden_states assert residual is not None and comb is not None completed = self._aux_post_op(hidden_states, residual, post, comb) - return hc_contract(completed, self.config.mhc_num_residual_streams) + return hc_contract(completed, self.config.hc_mult) def forward( self, From 1c47ebfc57d9e2415486e6080027ba319c73362a Mon Sep 17 00:00:00 2001 From: Praneeth Reddy A <72988597+0xViking@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:13:55 +0530 Subject: [PATCH 287/952] [Docs] Fix VLLM_CPU_KVCACHE_SPACE default and units (#60053) Signed-off-by: Praneeth Co-authored-by: Li, Jiang --- docs/configuration/conserving_memory.md | 2 +- docs/getting_started/installation/cpu.md | 2 +- vllm/platforms/cpu.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/configuration/conserving_memory.md b/docs/configuration/conserving_memory.md index 96a903bc31d3..b397283a623f 100644 --- a/docs/configuration/conserving_memory.md +++ b/docs/configuration/conserving_memory.md @@ -80,7 +80,7 @@ llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", enforce_eager=True) If you run out of CPU RAM, try the following options: - (Multi-modal models only) you can set the size of multi-modal cache by setting `mm_processor_cache_gb` engine argument (default 4 GiB). -- (CPU backend only) you can set the size of KV cache using `VLLM_CPU_KVCACHE_SPACE` environment variable (default 4 GiB). +- (CPU backend only) you can set the size of KV cache using `VLLM_CPU_KVCACHE_SPACE` environment variable, in GiB (default 0, meaning vLLM doesn't override the KV cache size through this variable). ## Multi-modal input limits diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index d27753413823..1e58a68fcf89 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -293,7 +293,7 @@ See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, support ### How to decide `VLLM_CPU_KVCACHE_SPACE`? -This value is 4GB by default. Larger space can support more concurrent requests, longer context length. However, users should take care of memory capacity of each NUMA node. The memory usage of each TP rank is the sum of `weight shard size` and `VLLM_CPU_KVCACHE_SPACE`, if it exceeds the capacity of a single NUMA node, the TP worker will be killed with `exitcode 9` due to out-of-memory. +This value is 0 by default, which means vLLM doesn't override the KV cache size through this variable. When set, the value is in GiB (e.g, `VLLM_CPU_KVCACHE_SPACE`=40 means 40 GiB). Larger space can support more concurrent requests, longer context length. However, users should take care of memory capacity of each NUMA node. The memory usage of each TP rank is the sum of `weight shard size` and `VLLM_CPU_KVCACHE_SPACE`, if it exceeds the capacity of a single NUMA node, the TP worker will be killed with `exitcode 9` due to out-of-memory. This is a legacy setting; use `--kv-cache-memory-bytes` instead. ### How to do performance tuning for vLLM CPU? diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index e0257acb8007..43d472f22e67 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -323,7 +323,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: mamba_backend, ) - # Lagecy setting + # Legacy setting env_key = "VLLM_CPU_KVCACHE_SPACE" if env_key in os.environ and os.environ[env_key] != "": kv_cache_space = int(os.environ[env_key]) From 196483e25b43511e13aa179ed52f50574ecd9fee Mon Sep 17 00:00:00 2001 From: Louie Tsai Date: Fri, 9 Oct 2026 04:44:09 -0700 Subject: [PATCH 288/952] [Tools] Recipes sweep workflow (#57307) Signed-off-by: louie-tsai --- tools/recipes/README.md | 89 +- tools/recipes/RUNTIME_TUNING.md | 82 +- tools/recipes/SWEEP_TUNING.md | 662 ++++++- tools/recipes/VISUALIZATION.md | 73 + tools/recipes/hardware_detection.py | 85 +- tools/recipes/recipe_json_to_vllm_config.py | 260 ++- tools/recipes/requirements.txt | 4 + tools/recipes/sweep/__init__.py | 4 + tools/recipes/sweep/adaptive_concurrency.py | 261 +++ .../sweep/concurrency_recommendation.py | 330 ++++ tools/recipes/{ => sweep}/runtime_tuning.py | 77 +- .../recipes/sweep/scheduler_sweep_prepare.py | 225 +++ tools/recipes/sweep/sweep_generation.py | 1531 +++++++++++++++++ tools/recipes/sweep/sweep_html_report.py | 756 ++++++++ .../{ => sweep}/sweep_recommendation.py | 261 ++- tools/recipes/sweep/sweep_visualization.py | 226 +++ tools/recipes/sweep_generation.py | 358 ---- 17 files changed, 4720 insertions(+), 564 deletions(-) create mode 100644 tools/recipes/VISUALIZATION.md create mode 100644 tools/recipes/requirements.txt create mode 100644 tools/recipes/sweep/__init__.py create mode 100755 tools/recipes/sweep/adaptive_concurrency.py create mode 100644 tools/recipes/sweep/concurrency_recommendation.py rename tools/recipes/{ => sweep}/runtime_tuning.py (81%) create mode 100644 tools/recipes/sweep/scheduler_sweep_prepare.py create mode 100644 tools/recipes/sweep/sweep_generation.py create mode 100644 tools/recipes/sweep/sweep_html_report.py rename tools/recipes/{ => sweep}/sweep_recommendation.py (62%) create mode 100644 tools/recipes/sweep/sweep_visualization.py delete mode 100644 tools/recipes/sweep_generation.py diff --git a/tools/recipes/README.md b/tools/recipes/README.md index 90ad33b3eb10..37486ae5487e 100644 --- a/tools/recipes/README.md +++ b/tools/recipes/README.md @@ -19,7 +19,7 @@ flowchart LR ``` The recipe is the baseline. Hardware and workload information can optionally -refine the initial configuration. Sweep tuning is an optional validation step. +refine the initial configuration. Sweep tuning is optional. ## Getting Started @@ -54,7 +54,9 @@ For `xeon6`, the script enables hardware detection automatically. Generated `config.yml` and `env.sh` files remain in the directory where the script was invoked. -## 1. vLLM Recipes Only +## Detailed Deployment Workflow + +### 1. vLLM Recipes Only Use the converter directly when the recipe already contains the deployment settings you need. This path requires only PyYAML; the vLLM Python package is @@ -64,15 +66,7 @@ not required unless optional runtime tuning or sweep generation is requested. pip install pyyaml ``` -Choose whichever recipe-selection method fits the workflow: - -**Interactive discovery** — search models, then choose hardware and strategy: - -```bash -python3 tools/recipes/recipe_json_to_vllm_config.py -``` - -**Non-interactive discovery** — provide model and hardware and use the +For non-interactive discovery, provide the model and hardware and use the Recipes-recommended strategy: ```bash @@ -81,7 +75,17 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --hardware xeon6 ``` -**Direct JSON input** — use a Recipes JSON URL or a local JSON file: +#### 1.1 Interactive Discovery + +Search models, then choose hardware and strategy interactively: + +```bash +python3 tools/recipes/recipe_json_to_vllm_config.py +``` + +#### 1.2 Direct JSON Input + +Use a Recipes JSON URL or a local JSON file: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -90,15 +94,27 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ python3 tools/recipes/recipe_json_to_vllm_config.py recipe.json ``` -All paths generate `config.yml` and `env.sh`. See -[REFERENCE.md](REFERENCE.md) for recipe discovery, strategy selection, direct -JSON input, custom output files, and deployment scope. +#### 1.3 Test a Preview Recipe API -## 2. Hardware Information (Optional) +A Recipes pull request can expose the same JSON API through its Vercel preview. +Use `--api-base` to validate it before the recipe is available in production: -Add `--detect-hardware` when the target host's effective CPU/NUMA/memory -resources should refine deployment-sensitive values such as -`tensor-parallel-size` and `gpu-memory-utilization`. +```bash +PREVIEW=https://vllm-recipes-git-fork-intel-ai-tce-dockerin-f4c148-inferact-inc.vercel.app + +python3 tools/recipes/recipe_json_to_vllm_config.py \ + --api-base "$PREVIEW" \ + --model meta-llama/Llama-3.2-1B-Instruct \ + --hardware xeon6 +``` + +See [REFERENCE.md](REFERENCE.md) for recipe discovery, strategy selection, +custom output files, and deployment scope. + +### 2. Hardware Information (Optional) + +Add `--detect-hardware` when using the converter directly and the target host's +effective CPU/NUMA/memory resources should refine deployment-sensitive values: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -107,14 +123,13 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --detect-hardware ``` -Hardware detection is optional and uses vLLM CPU resource utilities only when -requested. See [RUNTIME_TUNING.md](RUNTIME_TUNING.md#hardware-information). +`serve_with_recipe.sh` enables this automatically for `xeon6`. See +[RUNTIME_TUNING.md](RUNTIME_TUNING.md#hardware-information) for details. -## 3. Workload Information (Optional) +### 3. Workload Information (Optional) -Workload hints can refine scheduler settings for one initial deployment -suggestion. Inputs include token lengths, concurrency, and optional latency or -capacity objectives. +Add workload hints when token lengths, concurrency, or latency objectives should +seed optional benchmark tuning: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -127,17 +142,13 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 ``` -Hardware detection and workload information are independent optional inputs; -they can also be supplied together. See -[RUNTIME_TUNING.md](RUNTIME_TUNING.md#workload-information) for the supported -inputs and how runtime parameters are calculated. +See [RUNTIME_TUNING.md](RUNTIME_TUNING.md#workload-information) for supported +inputs and runtime calculations. -## 4. Sweep Tuning (Optional) +### 4. Sweep Tuning (Optional) -Use `--generate-sweep` when the initial scheduler suggestion should be validated -with `vllm bench sweep serve`. The sweep benchmarks nearby scheduler values and -`recommend.py` produces one measured `recommended-config.yml` plus the -selection evidence in `recommendation.json`. +Use `--generate-full-sweep` for benchmark-backed tuning of TP/DP, concurrency, +and scheduler settings: ```bash python3 tools/recipes/recipe_json_to_vllm_config.py \ @@ -147,13 +158,15 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --input-tokens 128 \ --output-tokens 128 \ --concurrency 32 \ - --generate-sweep + --generate-full-sweep ``` -See [SWEEP_TUNING.md](SWEEP_TUNING.md) for the benchmark, recommendation, and -vLLM CPU Docker-shell workflow. +See [SWEEP_TUNING.md](SWEEP_TUNING.md) for targeted sweep modes, NUMA binding, +failure handling, recommendations, reporting, and visualization. + +### 5. Start vLLM -## Start vLLM +When using the converter directly: ```bash source env.sh diff --git a/tools/recipes/RUNTIME_TUNING.md b/tools/recipes/RUNTIME_TUNING.md index d0b8ee929974..cf442a34d605 100644 --- a/tools/recipes/RUNTIME_TUNING.md +++ b/tools/recipes/RUNTIME_TUNING.md @@ -64,12 +64,10 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 ``` -`--output-tokens` is used to estimate steady-state request turnover for -`max-num-batched-tokens`. When both `--target-qps` and `--tpot-sla-ms` are -available, they are also used to estimate prompt arrival pressure per scheduler -step. `--ttft-sla-ms` is collected but is not yet converted directly into a -batch-size formula because the relationship between TTFT and scheduler budget is -model- and hardware-dependent. +Scheduler calculation is opt-in with `--tune-scheduler`; workload hints do +not rewrite `max-num-seqs` or `max-num-batched-tokens` by default. The detailed +scheduler and staged-sweep behavior is documented in +[SWEEP_TUNING.md](SWEEP_TUNING.md). ## Deployment-Time Parameters @@ -81,11 +79,30 @@ on a validation environment that differs from the user's target deployment. | --- | --- | --- | --- | | `tensor-parallel-size` | Available topology and model head counts constrain TP. | Hardware topology + model configuration | Use the largest head-compatible power-of-two TP within the effective NUMA-node count, for DP=1 and PP=1. | | `gpu-memory-utilization` | Available memory can differ by machine size, container limits, and other memory use. The vLLM option name is also used by the CPU backend. | Hardware memory + recipe baseline | Calculate a conservative fraction from the most constrained NUMA node. | -| `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active at the same time. | Workload concurrency | Set `max-num-seqs` to `--concurrency` when supplied. | -| `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Calculate decode budget + expected prefill demand, with vLLM scheduler constraints as floors. | +| `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active on each DP replica. | Workload concurrency + DP | Keep the recipe/vLLM value in `config.yml`; use `ceil(--concurrency / data-parallel-size)` as an explicit sweep seed. | +| `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Keep the recipe/vLLM value in `config.yml`; use decode budget + expected prefill demand as an explicit sweep seed. | | `data-parallel-size` | The required replica count depends on the requested throughput and the capacity of one replica. | Capacity target | Keep the recipe value today because per-replica SLA capacity is not known. | -## How Each Runtime Parameter Is Calculated +For `max-num-seqs` and `max-num-batched-tokens`, workload formulas now create +benchmark-only seeds instead of overriding the directly deployable +`config.yml`. The optional benchmark sweep measures those explicit seeds +against vLLM's normal serving defaults. Default-reference runs leave the +selected scheduler option unset at engine-configuration time so vLLM can +resolve its platform-, world-size-, model-, and usage-context-aware default +rather than assuming a fixed numeric value. See +[SWEEP_TUNING.md](SWEEP_TUNING.md). + +This serving-default comparison is distinct from the `SchedulerConfig` fallback +constants used inside the initial tuning heuristics below. + +## How Initial Runtime Tuning Seeds Are Calculated + +The hardware formulas below produce directly deployable starting values. The +scheduler formulas produce explicit benchmark seeds and are not claimed to be +universal performance optima. This separation avoids baking workload guesses +into `config.yml` before measurements exist. The Tune All workflow benchmarks +TP/DP first, then concurrency, and finally scheduler settings before writing a +final recommendation. See [SWEEP_TUNING.md](SWEEP_TUNING.md). ### `tensor-parallel-size` @@ -147,14 +164,20 @@ candidate = min(recipe value or 0.80, safe_fraction) ### `max-num-seqs` -When `--concurrency` is supplied, it directly represents the requested maximum -number of simultaneously active requests: +When `--concurrency` is supplied, it represents global client concurrency. Each +DP replica needs capacity for its expected share of those requests: ```text -max-num-seqs = concurrency +per_replica_concurrency = ceil(concurrency / data_parallel_size) +max-num-seqs = per_replica_concurrency ``` -If concurrency is not supplied, the converter does not override this parameter. +For DP=1 this reduces to `max-num-seqs = concurrency`. This value seeds the +scheduler benchmark; it is not added to `config.yml`. If concurrency is not +supplied, no seed is calculated. Benchmark results show +that setting `max-num-seqs` below per-replica concurrency can create request +queueing and severe TTFT degradation; scheduler sweeps therefore do not use +below-concurrency candidates. ### `max-num-batched-tokens` @@ -165,7 +188,7 @@ First, the policy determines the active sequence count: ```text active_sequences = - --concurrency + ceil(--concurrency / data_parallel_size) or recipe max-num-seqs or vLLM default max_num_seqs (128) ``` @@ -181,17 +204,19 @@ The policy then estimates how many new prompts need prefill work per scheduler step: ```text +per_replica_qps = target_qps / data_parallel_size + prefills_per_step = max( 1, active_sequences / output_tokens, - target_qps * tpot_sla_ms / 1000 + per_replica_qps * tpot_sla_ms / 1000 ) prefills_per_step = min(active_sequences, prefills_per_step) prefill_budget = ceil(input_tokens * prefills_per_step) ``` -The final scheduler budget is: +The final explicit scheduler sweep seed is: ```text max-num-batched-tokens = max( @@ -201,7 +226,7 @@ max-num-batched-tokens = max( ) ``` -For example, with 128 input tokens, 128 output tokens, and concurrency 32: +For example, with 128 input tokens, 128 output tokens, concurrency 32, and DP=1: ```text decode_budget = 32 @@ -210,6 +235,12 @@ prefill_budget = 128 candidate = max(2048, 32, 32 + 128) = 2048 ``` +The calculated token budget is a scheduler sweep seed and is not added to +`config.yml`. If P99 TTFT still +misses its objective while P99 TPOT has substantial headroom, Tune All should +reduce workload concurrency or select a different TP/DP layout rather than +reducing `max-num-seqs` below per-replica concurrency. + If chunked prefill is explicitly disabled and `max-model-len` is available, the policy also ensures: @@ -226,8 +257,10 @@ also needs measured per-replica throughput that still satisfies TTFT/TPOT: DP = ceil(target_qps / qps_per_replica_at_SLO) ``` -Because the converter does not have that measured capacity yet, it intentionally -keeps the recipe DP value rather than guessing. +Because the converter does not have that measured capacity yet, its initial +configuration intentionally keeps the recipe DP value rather than guessing. +Tune All subsequently benchmarks supported TP/DP layouts and replaces this seed +with a measured layout before concurrency and scheduler tuning. ## Precedence @@ -235,8 +268,11 @@ keeps the recipe DP value rather than guessing. vLLM defaults -> vLLM Recipes baseline -> hardware refinement (optional) - -> workload / SLO refinement (optional) - -> config.yml + env.sh + -> config.yml + env.sh + +workload / SLO hints (optional) + -> explicit scheduler sweep seed + -> benchmark-backed recommended-config.yml ``` ## Runtime-Tuning Hardware Scope @@ -312,10 +348,10 @@ with [SWEEP_TUNING.md](SWEEP_TUNING.md). ## Implementation - `hardware_detection.py` collects effective CPU/NUMA/memory information. -- `runtime_tuning.py` owns hardware-policy selection and independent parameter +- `sweep/runtime_tuning.py` owns hardware-policy selection and independent parameter tuning functions. - `recipe_json_to_vllm_config.py` resolves the recipe, collects optional inputs, applies the selected policy, and generates `config.yml` and `env.sh`. -For benchmark-backed validation of the initial scheduler suggestion, see +For benchmark-backed validation of the explicit scheduler seed, see [SWEEP_TUNING.md](SWEEP_TUNING.md). diff --git a/tools/recipes/SWEEP_TUNING.md b/tools/recipes/SWEEP_TUNING.md index e7282e122d78..e39444aef983 100644 --- a/tools/recipes/SWEEP_TUNING.md +++ b/tools/recipes/SWEEP_TUNING.md @@ -1,13 +1,133 @@ # Sweep Tuning -Sweep tuning is optional. The converter always creates one initial `config.yml` -first. Users can deploy it immediately or benchmark nearby scheduler settings -and generate one measured recommendation. +Sweep tuning is optional. The converter always creates an initial `config.yml` +that can be deployed directly. Scheduler calculation is disabled by default: +TP/DP and concurrency stages remove explicit `max-num-seqs` and +`max-num-batched-tokens` so vLLM resolves its normal scheduler defaults. -## Generate +Use `--tune-scheduler` only when you explicitly want workload-derived scheduler +values before the dedicated scheduler stage. For benchmark-backed tuning, the +recommended workflow is **Tune All**, which measures the serving stack in +dependency order: + +```text +TP/DP -> concurrency -> scheduler +``` + +The examples below assume tuning is run inside the target vLLM CPU container so +hardware detection sees the same CPU, NUMA, memory, and cgroup limits as the +deployment. + +## 1. Start the vLLM CPU Docker Shell + +From the vLLM source tree: ```bash -python3 tools/recipes/recipe_json_to_vllm_config.py \ +mkdir -p recipe-output + +docker run --rm -it \ + --entrypoint bash \ + --security-opt seccomp=unconfined \ + --cap-add SYS_NICE \ + --shm-size=4g \ + -p 8000:8000 \ + -v "$PWD/tools/recipes:/recipes:ro" \ + -v "$PWD/recipe-output:/output" \ + -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \ + -w /output \ + vllm/vllm-openai-cpu:latest-x86_64 +``` + +`/recipes` is mounted read-only from the source tree and `/output` is writable. +Run the converter, generated sweep scripts, and hardware detection inside this +container. + +For additional container details, see +[RUNTIME_TUNING.md](RUNTIME_TUNING.md#vllm-cpu-docker-shell). + +## 2. Recommended: Tune All + +Use `--generate-full-sweep` by default when you want benchmark-backed runtime +tuning. It measures TP/DP first, then the highest useful concurrency, then +scheduler parameters using the selected layout and concurrency. + +```mermaid +flowchart TD + I[["Input
Recipe + detected hardware
fixed input/output workload shape"]] + + A1("1. Tune TP/DP
tensor-parallel-size + data-parallel-size") + O1[/"parallel-layout-config.yml
selected TP/DP"/] + + A2("2. Tune concurrency
max_concurrency") + O2[/"concurrency-recommendation.json
recommended max_concurrency"/] + + A3("3. Tune scheduler
max-num-seqs + max-num-batched-tokens") + O3[/"recommended-config.yml
final vLLM server config"/] + + I --> A1 + A1 --> O1 + O1 --> A2 + A2 --> O2 + O2 --> A3 + A3 --> O3 +``` + +### Workflow options + +The table is ordered by the tuning sequence. **Tune All** is the recommended +default; the other modes are useful when only one part of the serving stack +needs to be re-measured. + +| Goal | Option | Tuning sequence | +| --- | --- | --- | +| **Tune all (recommended)** | `--generate-full-sweep` | **1. TP/DP -> 2. concurrency -> 3. scheduler** | +| Tune TP/DP only | `--generate-parallel-layout-sweep` | 1. TP/DP | +| Tune concurrency only | `--generate-concurrency-sweep` | 2. concurrency | +| Tune scheduler only | `--generate-scheduler-sweep` (`--generate-sweep` alias) | 3. scheduler | + +All workflows keep one fixed input/output workload shape and require +`--input-tokens`, `--output-tokens`, and `--concurrency`. + +For **Tune All**, the supplied `--concurrency` is the representative load used +for the initial TP/DP comparison and to size the benchmark request set. The +concurrency stage then measures the final SLA-feasible `max_concurrency`. + +### Scheduler defaults and `--tune-scheduler` + +By default, `--concurrency`, token lengths, and SLA values are workload hints; +they do not cause the converter to calculate `max-num-seqs` or +`max-num-batched-tokens`. This keeps scheduler capacity independent from the +seed concurrency while TP/DP and maximum concurrency are being measured. + +The default Tune All policy is: + +```text +Stage 1: TP/DP -> vLLM scheduler defaults +Stage 2: concurrency -> vLLM scheduler defaults +Stage 3: scheduler -> calculate and benchmark scheduler candidates +``` + +For a directly deployable heuristic configuration, or to intentionally carry +workload-derived scheduler values into the earlier sweep stages, opt in with: + +```bash +--tune-scheduler +``` + +With that flag, the existing workload formulas are enabled. In particular, +`max-num-seqs` uses the per-replica workload concurrency and +`max-num-batched-tokens` uses the estimated decode and prefill pressure. + +A scheduler-only sweep is also an explicit opt-in to scheduler calculation: it +uses the workload formulas as benchmark seeds and compares them with vLLM +defaults. + +### Generate the Tune All package + +Inside the container: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ --model meta-llama/Llama-3.1-8B-Instruct \ --hardware xeon6 \ --detect-hardware \ @@ -16,80 +136,311 @@ python3 tools/recipes/recipe_json_to_vllm_config.py \ --concurrency 32 \ --ttft-sla-ms 3000 \ --tpot-sla-ms 100 \ - --generate-sweep + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-full-sweep \ + --sweep-out-dir /output/sweep +``` + +The initial configuration remains directly deployable: + +```bash +source /output/env.sh +vllm serve --config /output/config.yml +``` + +Before measurements, `config.yml` retains the recipe/vLLM scheduler defaults. +The sweep package separately carries explicit workload-derived values for +`max-num-seqs` and `max-num-batched-tokens`, so the benchmark can compare both +policies without making the workload heuristic the deployment default. + +Stop a manually started server before running the sweep because the generated +sweep scripts start and stop their own vLLM servers. + +Run all tuning stages: + +```bash +/output/sweep/run_full_sweep.sh ``` -The optional package is: +The end-to-end flow produces intermediate recommendations for TP/DP and +concurrency and finishes with: ```text -sweep/ -├── serve_params.json -├── bench_params.json -├── run_sweep.sh -├── recommend.py -└── SWEEP.md +/output/sweep/recommended-config.yml +/output/sweep/recommendation.json +/output/sweep/sweep-report.html ``` -The sweep varies only `max-num-seqs` and `max-num-batched-tokens`. It uses an -eight-point directed design: a batch-budget curve at the initial sequence count -plus lower/higher batch interactions at three-quarters and one-half of that -count. This gives broader coverage than a one-parameter-at-a-time sweep without -the cost of a full Cartesian grid. +Inspect the final recommendation: -## Run and Recommend +```bash +cat /output/sweep/recommendation.json +``` + +## 3. Tune All Stage Details + +### Stage 1: Tune TP/DP + +The parallel-layout stage measures NUMA-aware combinations of: + +- `tensor-parallel-size` +- `data-parallel-size` + +Primary candidates use all effective NUMA nodes: + +```text +tensor-parallel-size * data-parallel-size = effective NUMA nodes +``` + +The effective NUMA-node count comes from hardware detection; it is not fixed by +the sweep generator. The table below shows representative detected topologies. + +TP is restricted to the supported values `1`, `2`, `4`, and `8`. The sweep also +includes the largest supported TP size that does not exceed the effective +NUMA-node count, even if that layout leaves some NUMA nodes idle. + +| Effective NUMA nodes | Generated layouts | +| ---: | --- | +| 2 | `TP=2, DP=1`; `TP=1, DP=2` | +| 4 | `TP=4, DP=1`; `TP=2, DP=2`; `TP=1, DP=4` | +| 6 | `TP=4, DP=1` (4 of 6 nodes); `TP=2, DP=3`; `TP=1, DP=6` | +| 8 | `TP=8, DP=1`; `TP=4, DP=2`; `TP=2, DP=4`; `TP=1, DP=8` | + +The generator then applies model- and runtime-aware safety policy before any +server is started: + +1. For an MoE model, if the CPU communicator does not override the variable-size + `all_gatherv` and `reduce_scatterv` collectives used by MoE data parallelism, + all generated candidates use `DP=1`. Every supported TP value is retained, + so a 6-NUMA-node host tests `TP=4`, `TP=2`, and `TP=1` rather than only `TP=4`. +2. When MoE dimensions are known and expert parallelism is disabled, the fused + MoE path flattens TP, DP, and PCP. A candidate is valid only when + `moe_intermediate_size % (TP * DP * PCP) == 0`. +3. Failures not covered by a preflight rule remain isolated by + `--continue-on-error` when the installed CLI supports it, with bounded + `--resume` retries as the compatibility fallback. + +Skipped candidates and reasons are saved in +`sweep/parallel_layout_skips.json`. If model metadata cannot be loaded, the +generator keeps the generic layouts and relies on the third policy. + +By default, parallel-layout candidates change only TP and DP. The generated +benchmark config removes explicit `max-num-seqs` and +`max-num-batched-tokens`, so every layout uses vLLM-resolved scheduler +defaults. This prevents a scheduler value derived from the seed concurrency +from becoming part of the TP/DP selection. + +With `--tune-scheduler`, the workload-derived per-replica scheduler baseline is +restored for this stage. + +The stage writes: + +```text +sweep/parallel-layout-config.yml +sweep/parallel-layout-recommendation.json +``` + +The selected TP/DP layout becomes the fixed server layout for Stage 2. + +### Temporary TP/DP NUMA-binding workaround + +Xeon TP/DP and Tune All sweeps currently enable a temporary explicit CPU-binding +workaround by default: + +```text +--tp-dp-numa-bind-workaround +``` + +Hardware detection builds one `VLLM_CPU_OMP_THREADS_BIND` CPU list per effective +NUMA node. On x86 it selects one logical CPU per physical core, matching vLLM's +auto-binding SMT policy, and excludes one physical core from each NUMA node for +non-OMP work. The resulting value is written to `env.sh`, for example: ```bash -sweep/run_sweep.sh --dry-run -sweep/run_sweep.sh -sweep/recommend.py +export VLLM_CPU_OMP_THREADS_BIND='48-62|64-78|80-94|96-110' ``` -Resume an interrupted sweep with: +The exact CPU IDs depend on the effective container/cgroup cpuset and NUMA +topology. Discontinuous CPU IDs are preserved as comma-separated ranges. + +`VLLM_CPU_NUM_OF_RESERVED_CPU` is not used by this workaround because vLLM only +applies its reserved-CPU logic in automatic binding mode. With an explicit +`VLLM_CPU_OMP_THREADS_BIND`, reserved cores must already be omitted from the +generated lists. + +After the vLLM CPU DP NUMA-binding issue is fixed, disable the workaround +without removing the implementation: ```bash -sweep/run_sweep.sh --resume +--no-tp-dp-numa-bind-workaround +``` + +An explicit non-`auto` `VLLM_CPU_OMP_THREADS_BIND` supplied by the recipe is +preserved. + +### Stage 2: Tune Concurrency + +The concurrency stage keeps the selected TP/DP layout fixed while finding the +largest useful client concurrency. By default the generated benchmark config +removes explicit `max-num-seqs` and `max-num-batched-tokens`, so the seed +concurrency cannot become an artificial scheduler admission limit. + +With `--tune-scheduler`, explicit scheduler values are preserved as an opt-in +behavior. + +When TTFT and/or TPOT objectives are supplied, Recipes uses an **adaptive SLA +boundary search** instead of uniformly sampling the entire workload range. + +The supplied `--concurrency` is the seed. The controller grows upward until it +finds the first SLA failure (or shrinks downward when the seed fails), then +binary-searches the integer PASS/FAIL bracket. One vLLM server remains running +throughout this refinement, so midpoint probes do not reload the model. + +For example: + +```text +32 PASS, 64 FAIL + -> 48 + -> 40 or 56 + -> ... + -> adjacent PASS / FAIL boundary +``` + +A point is SLA-feasible only when all of the following hold: + +1. Every measured run completes without failed requests. +2. Median P99 TTFT is at or below the supplied TTFT objective. +3. Median P99 TPOT is at or below the supplied TPOT objective. +4. Duration-weighted combined request compliance is at least the configured + minimum (`0.99`, or 99%, by default). + +The combined-compliance check is intentionally stricter than looking at the +TTFT and TPOT P99 columns independently. `vllm bench serve --goodput` counts +requests that satisfy all supplied latency objectives, and the recommender +estimates the combined compliant-request fraction across repeated runs as: + +```text +sum(request_goodput * duration) / sum(completed_requests) +``` + +Therefore a candidate can have P99 TTFT below 3000 ms and P99 TPOT below +100 ms but still be SLA-ineligible if fewer than 99% of requests satisfy both +objectives together. For example, a reported combined compliance of `98.96%` +fails the default `99.00%` requirement even when both displayed P99 metrics +pass. The HTML report shows the combined compliance, required threshold, and +the specific SLA failure reason for each candidate. + +Adaptive concurrency uses a `1800` second server-readiness timeout by default, +which is intentionally longer than the generic benchmark helper default for +CPU model loading and initialization. `run_full_sweep.sh +--server-ready-timeout VALUE` also exports the selected timeout explicitly to +the generated concurrency runner, in addition to forwarding the CLI option. +This avoids falling back to a shorter adaptive-helper timeout when the +concurrency stage is invoked through the full staged workflow. + +The default growth factor is `2` and the default hard cap is `1000`. +Generated runners accept `--growth-factor VALUE` and +`--max-concurrency-cap VALUE`. + +Without TTFT/TPOT objectives there is no SLA boundary. In that mode the +generated runner keeps vLLM Workload Explorer: + +```text +vllm bench sweep serve_workload --workload-var max_concurrency +``` + +The stage writes: + +```text +sweep/concurrency-recommendation.json +``` + +`max_concurrency` is benchmark/deployment metadata. It is **not** written as a +`vllm serve` configuration parameter. + +### Stage 3: Tune Scheduler + +Before the scheduler sweep, `prepare_scheduler.py` recalculates the scheduler +baseline from: + +```text +selected TP/DP + + +selected max_concurrency + + +fixed input/output token shape +``` + +The scheduler stage keeps an explicit `max-num-seqs` at the per-replica +concurrency lower bound while tuning `max-num-batched-tokens`. It also compares +the explicit sequence limit with the vLLM default. + +The per-replica sequence baseline is recalculated as: + +```text +max-num-seqs = ceil(selected max_concurrency / selected data-parallel-size) ``` -`recommend.py` writes: +The directed scheduler sweep does not test `max-num-seqs` below this value; +doing so limits immediately active requests and can turn scheduler queueing into +severe TTFT degradation. It keeps the existing batch-budget curve around the +baseline and also measures three vLLM-default reference candidates: + +| Reference candidate | `max-num-seqs` | `max-num-batched-tokens` | +| --- | --- | --- | +| `vllm_default_max_num_seqs` | vLLM default | generated value | +| `vllm_default_max_num_batched_tokens` | generated value | vLLM default | +| `vllm_defaults` | vLLM default | vLLM default | + +The default references do not hard-code values such as 128 or 2048. The sweep +omits the selected scheduler argument and lets the installed vLLM resolve its +normal platform-, world-size-, model-, and usage-context-aware default. + +If a default-reference candidate wins, the corresponding scheduler key is +omitted from `recommended-config.yml`. + +## 4. Benchmark Size, Warmup, and Failure Handling + +Measured request count scales with concurrency using: ```text -sweep/recommended-config.yml -sweep/recommendation.json +num_prompts = min(1000, max(100, concurrency * 10)) ``` -With TTFT/TPOT objectives, the benchmark uses vLLM `--goodput`. The recommender -calculates duration-weighted combined compliance across repeated runs and -requires both median P99 latency objectives and the minimum compliance ratio -(default `0.99`). Among eligible configurations it selects the highest mean -output-token throughput. Use `recommend.py --minimum-compliance VALUE` to -change the compliance requirement. +This targets about ten concurrency turnovers when neither bound applies. The +100-prompt floor avoids very small samples for P99/compliance measurements, and +the 1000-prompt cap bounds sweep runtime. + +Before a measured parameter combination, generated scripts use one unmeasured +warmup containing one full concurrency window. Warmup output is retained as +`warmup.json` for auditing but is excluded from recommendation statistics. -If no configuration is eligible, `recommend.py` records the highest-goodput -configuration as `best_effort` in `recommendation.json`, does not write a -deployable `recommended-config.yml`, and exits with status 2. Without latency -objectives it selects highest mean output-token throughput. Failed benchmark -configurations are excluded in either mode. +Generated scripts also handle sweep-run failures without discarding completed +work: -When an eligible configuration exists, `recommended-config.yml` changes only -the two swept scheduler parameters. `recommendation.json` records mean, median, -and worst-run P99 values, combined compliance, and the measured evidence for -every candidate. +- If the installed vLLM CLI supports `--continue-on-error`, the generated script + enables it. +- If the installed vLLM does not expose that option, the script omits it and + automatically retries an interrupted sweep with `--resume`. +- The default is two automatic resume retries. Override it with + `VLLM_RECIPE_SWEEP_RETRIES`. -Deploy: +For example: ```bash -source env.sh -vllm serve --config sweep/recommended-config.yml +export VLLM_RECIPE_SWEEP_RETRIES=3 +/output/sweep/run_full_sweep.sh ``` -## vLLM CPU Docker Shell +## 5. Targeted Tuning Workflows + +Use these modes when a full re-tune is unnecessary. -Use the common CPU container setup in -[RUNTIME_TUNING.md](RUNTIME_TUNING.md#vllm-cpu-docker-shell). It ensures -hardware detection runs against the same effective CPU, NUMA, memory, and -cgroup limits used by the deployment. +### Tune TP/DP Only -Inside that container, generate the initial suggestion plus sweep package: +Use this when only the parallel layout needs to be measured. This mode stops +after selecting `tensor-parallel-size` and `data-parallel-size`: ```bash python3 /recipes/recipe_json_to_vllm_config.py \ @@ -103,35 +454,224 @@ python3 /recipes/recipe_json_to_vllm_config.py \ --tpot-sla-ms 100 \ --config-out /output/config.yml \ --env-out /output/env.sh \ - --generate-sweep \ + --generate-parallel-layout-sweep \ --sweep-out-dir /output/sweep ``` -The initial configuration can be deployed directly: +Run: ```bash -source /output/env.sh -vllm serve --config /output/config.yml +/output/sweep/run_parallel_layout_sweep.sh +/output/sweep/recommend_parallel_layout.py +``` + +The result is `/output/sweep/parallel-layout-config.yml`. Use Tune All when +concurrency and scheduler tuning should follow the selected layout. + +### Tune Concurrency Only + +Use this when TP/DP and scheduler configuration are already fixed and only +concurrent serving capacity needs to be measured: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 \ + --detect-hardware \ + --input-tokens 128 \ + --output-tokens 128 \ + --concurrency 32 \ + --ttft-sla-ms 3000 \ + --tpot-sla-ms 100 \ + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-concurrency-sweep \ + --sweep-out-dir /output/sweep +``` + +Run: + +```bash +/output/sweep/run_concurrency_sweep.sh +/output/sweep/recommend_concurrency.py +``` + +### Tune Scheduler Only + +Use this when TP/DP and target concurrency are already known: + +```bash +python3 /recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --hardware xeon6 \ + --detect-hardware \ + --input-tokens 128 \ + --output-tokens 128 \ + --concurrency 32 \ + --ttft-sla-ms 3000 \ + --tpot-sla-ms 100 \ + --config-out /output/config.yml \ + --env-out /output/env.sh \ + --generate-scheduler-sweep \ + --sweep-out-dir /output/sweep ``` -Stop that manually started server before running the sweep, because -`run_sweep.sh` starts and stops its own vLLM servers. +`--generate-sweep` remains a backward-compatible alias for +`--generate-scheduler-sweep`. + +Run: ```bash -/output/sweep/run_sweep.sh --dry-run /output/sweep/run_sweep.sh /output/sweep/recommend.py ``` -Inspect: +## 6. Recommendation Policy + +With TTFT/TPOT objectives, benchmark commands use vLLM `--goodput`. + +For TP/DP and scheduler selection, the recommender: + +1. Excludes invalid configurations and failed benchmark measurements. +2. Calculates duration-weighted combined request compliance across successful + runs from request goodput and run duration. +3. Requires the median P99 TTFT and P99 TPOT objectives independently. +4. Also requires the combined request-compliance ratio to meet + `--minimum-compliance`, which defaults to `0.99`. +5. Establishes the highest mean output-token throughput among eligible + configurations, then applies the selection policy below. + +The P99 checks and combined-compliance check answer different questions. P99 +verifies the tail of each latency metric independently; combined compliance +verifies that the required fraction of requests satisfy all configured +objectives together. A candidate must pass both forms of validation to be +marked SLA eligible. + +For a scheduler-only comparison with one fixed TP/DP layout, candidates within +1% of the highest eligible output-token throughput are treated as practically +equivalent. Within that set, the recommender prefers the candidate with fewer +explicit scheduler overrides, followed by compliance, goodput, and throughput. +This avoids hard-coding a scheduler value for a difference that is likely within +normal run-to-run variation. Change or disable the equivalence range with: ```bash -cat /output/sweep/recommendation.json +/output/sweep/recommend.py --throughput-equivalence-percent VALUE +``` + +Use `0` to restore exact highest-throughput selection. TP/DP comparisons still +use exact highest-throughput selection because changing a parallel layout is a +material deployment decision. + +For concurrency selection, the policy instead selects the **highest +SLA-feasible `max_concurrency`**, with throughput/compliance metrics used to +break ties. + +Change the compliance threshold with: + +```bash +/output/sweep/recommend.py --minimum-compliance VALUE +``` + +If no configuration is SLA-feasible, the recommendation JSON records a +best-effort candidate. A deployable final configuration is produced only when +the required recommendation stages succeed. + +## 7. Post-Benchmark Reporting + +After Tune All finishes, `run_full_sweep.sh` automatically runs `report.py` and +writes a self-contained summary: + +```text +/output/sweep/sweep-report.html +``` + +The HTML report includes the selected configuration, a stage summary, SLA +metrics, candidate tables, and automatic coverage observations. It uses only +Python's standard library. A report-generation error is printed as a warning +and does not change the success of the completed sweep. + +The TTFT and TPOT objectives supplied to `recipe_json_to_vllm_config.py` are +embedded in every generated recommender and in `report.py`. +`run_full_sweep.sh` also passes those values explicitly to every recommendation +stage and the report. Report generation stops with an error if a stale or +manually generated recommendation JSON contains different SLA values. + +Regenerate the report manually when recommendation files change: + +```bash +cd /output/sweep +./report.py +``` + +### Optional figures + +Every generated sweep package includes a standalone analysis helper. It is not +called by `run_full_sweep.sh` or any stage runner, so figure generation cannot +start, resume, or alter a sweep. + +After one or more benchmark stages finish, install the plotting dependencies: + +```bash +cd /output/sweep +python3 -m pip install -r requirements.txt ``` -Deploy the measured result: +Preview the planned figures, then generate them from the completed +`summary.json` files: + +```bash +./visualize.py --dry-run +./visualize.py +``` + +Use `--stage parallel-layout`, `--stage concurrency-tuning`, or +`--stage runtime-tuning` to analyze selected stages. The helper writes PNG files +under each stage's `figures/` directory. See the generated `VISUALIZATION.md` +for the figure set and additional examples. + +## 8. Deploy the Tuned Configuration + +After Tune All completes: ```bash source /output/env.sh vllm serve --config /output/sweep/recommended-config.yml ``` + +## Select sweep stages + +Sweep implementation and reporting helpers live in `tools/recipes/sweep/`. +The converter remains `tools/recipes/recipe_json_to_vllm_config.py`. + +Use `--sweep-stage concurrency` to generate only a concurrency sweep against +the supplied recipe's fixed TP/DP layout, for example: + +```bash +python3 tools/recipes/recipe_json_to_vllm_config.py \ + --model meta-llama/Llama-3.1-8B-Instruct --hardware xeon6 \ + --input-tokens 128 --output-tokens 128 --concurrency 32 \ + --ttft-sla-ms 3000 --tpot-sla-ms 100 \ + --sweep-stage concurrency --sweep-out-dir sweep +./sweep/run_concurrency_sweep.sh +./sweep/recommend_concurrency.py +``` + +`--sweep-stage` accepts `parallel-layout`, `concurrency`, `scheduler`, or `all`. +The existing `--generate-*-sweep` flags remain supported; do not combine them +with `--sweep-stage`. `parallel-layout` and `all` require `--detect-hardware`. + +In a package generated with `--sweep-stage all` or `--generate-full-sweep`, +select the stages to execute: + +```bash +./sweep/run_full_sweep.sh --stages concurrency +./sweep/run_full_sweep.sh --stages concurrency,scheduler +./sweep/run_full_sweep.sh --stages scheduler --resume +``` + +The default is all three stages. Selected stages run in dependency order. +Skipping parallel-layout reuses `parallel-layout-config.yml` if present; +otherwise the runner copies the generated `initial-config.yml` as the fixed +layout. Scheduler-only execution requires `concurrency-recommendation.json` +from an earlier concurrency run. Remaining arguments go to the selected sweep +runners. Use separate output directories for unrelated experiments. diff --git a/tools/recipes/VISUALIZATION.md b/tools/recipes/VISUALIZATION.md new file mode 100644 index 000000000000..4c109102e7ab --- /dev/null +++ b/tools/recipes/VISUALIZATION.md @@ -0,0 +1,73 @@ +# Post-Benchmark Reporting + +## HTML summary + +`run_full_sweep.sh` automatically generates a self-contained HTML report after +the final recommendation: + +```text +sweep-report.html +``` + +The report uses Python's standard library and does not require plotting +dependencies. It summarizes selected settings, stage results, SLA metrics, and +all measured candidates. Report-generation failures are warnings and do not +change sweep success. + +The generated `report.py` uses the same TTFT and TPOT objectives supplied when +the sweep package was created. It displays both objectives and validates them +against every available recommendation JSON before writing the report. + +Regenerate it manually or select a different output path: + +```bash +./report.py +./report.py --output results/my-sweep-report.html +``` + +## Optional figures + +Visualization is an independent analysis step. It reads completed sweep +results and does not start a server, run a benchmark, or change a tuning +recommendation. + +Install the plotting dependencies in the same environment as vLLM: + +```bash +python3 -m pip install -r requirements.txt +``` + +Preview the figures that would be generated: + +```bash +./visualize.py --dry-run +``` + +Generate figures for every completed stage: + +```bash +./visualize.py +``` + +Select one or more stages when needed: + +```bash +./visualize.py --stage concurrency-tuning +./visualize.py --stage parallel-layout --stage runtime-tuning +``` + +Figures are written beside the source results: + +```text +results/parallel-layout/figures/ +results/concurrency-tuning/figures/ +results/runtime-tuning/figures/ +``` + +The analysis produces output-throughput, P99 TTFT, and P99 TPOT figures. The +concurrency stage also produces a throughput/TTFT tradeoff figure and a goodput +figure when the benchmark results contain that metric. + +The native vLLM Pareto command is intentionally not used because its efficiency +axis is expressed as tokens per second per GPU, which is not an accurate label +for Xeon CPU layouts. diff --git a/tools/recipes/hardware_detection.py b/tools/recipes/hardware_detection.py index 2b90dffac492..f5b428dd020d 100644 --- a/tools/recipes/hardware_detection.py +++ b/tools/recipes/hardware_detection.py @@ -4,7 +4,7 @@ """CPU hardware discovery for optional vLLM Recipes runtime tuning. Keep this module focused on facts exposed to the current process/container. -Tuning decisions belong in runtime_tuning.py. +Tuning decisions belong in sweep/runtime_tuning.py. """ from __future__ import annotations @@ -46,6 +46,89 @@ def to_dict(self) -> dict[str, object]: return asdict(self) +def _compress_cpu_ids(cpu_ids: list[int] | tuple[int, ...]) -> str: + """Return a compact Linux CPU-list string while preserving CPU IDs.""" + values = sorted(set(cpu_ids)) + if not values: + return "" + + ranges: list[str] = [] + start = previous = values[0] + for cpu_id in values[1:]: + if cpu_id == previous + 1: + previous = cpu_id + continue + ranges.append(str(start) if start == previous else f"{start}-{previous}") + start = previous = cpu_id + ranges.append(str(start) if start == previous else f"{start}-{previous}") + return ",".join(ranges) + + +def build_numa_omp_threads_bind( + hardware: HardwareInfo, + *, + reserved_cores_per_numa: int = 1, +) -> str: + """Build explicit x86 OMP CPU lists for the effective NUMA nodes. + + Temporary TP/DP sweep workaround: + - select one logical CPU per physical core, matching vLLM x86 auto-binding; + - reserve cores by omitting them from each explicit NUMA CPU list; + - emit one pipe-separated list per effective NUMA node. + + Explicit VLLM_CPU_OMP_THREADS_BIND lists bypass vLLM's auto-binding reserve + logic, so the reserved cores must be excluded here. + """ + if reserved_cores_per_numa < 0: + raise ValueError("reserved_cores_per_numa must be at least 0") + + effective_nodes = {node.node_id for node in hardware.numa_nodes} + allowed_cpus = get_allowed_cpu_list() + + groups: list[str] = [] + for node in hardware.numa_nodes: + node_cpus = [ + cpu + for cpu in allowed_cpus + if cpu.numa_node == node.node_id and cpu.numa_node in effective_nodes + ] + + core_to_logical_ids: dict[int, list[int]] = {} + for cpu in node_cpus: + if cpu.physical_core < 0: + continue + core_to_logical_ids.setdefault(cpu.physical_core, []).append(cpu.id) + + if not core_to_logical_ids: + raise RuntimeError( + f"NUMA node {node.node_id} has no physical-core topology " + "available for explicit OMP binding." + ) + + # Match vLLM's x86 auto-binding selector: for SMT siblings on a physical + # core, use the highest logical CPU ID. + selected_cpu_ids = sorted( + max(logical_ids) for _, logical_ids in sorted(core_to_logical_ids.items()) + ) + + if len(selected_cpu_ids) <= reserved_cores_per_numa: + raise RuntimeError( + f"NUMA node {node.node_id} has only {len(selected_cpu_ids)} " + "usable physical cores, which is not enough after reserving " + f"{reserved_cores_per_numa} core(s)." + ) + + if reserved_cores_per_numa: + selected_cpu_ids = selected_cpu_ids[:-reserved_cores_per_numa] + + groups.append(_compress_cpu_ids(selected_cpu_ids)) + + if not groups: + raise RuntimeError("No effective NUMA CPU lists are available for OMP binding.") + + return "|".join(groups) + + def _read_socket_id(cpu_id: int) -> int | None: path = Path(f"/sys/devices/system/cpu/cpu{cpu_id}/topology/physical_package_id") try: diff --git a/tools/recipes/recipe_json_to_vllm_config.py b/tools/recipes/recipe_json_to_vllm_config.py index bd947b07451c..c5089fbe9bbf 100644 --- a/tools/recipes/recipe_json_to_vllm_config.py +++ b/tools/recipes/recipe_json_to_vllm_config.py @@ -189,17 +189,67 @@ def parse_args() -> argparse.Namespace: type=float, help="Optional capacity target for future DP/capacity tuning.", ) + tuning.add_argument( + "--tune-scheduler", + action="store_true", + help=( + "Opt in to workload-derived max-num-seqs and " + "max-num-batched-tokens. By default scheduler parameters remain " + "at recipe/vLLM defaults." + ), + ) sweep = p.add_argument_group( "optional performance sweep", - ("Generate benchmark files after creating one initial runtime suggestion."), + ("Generate benchmark files alongside the directly deployable runtime config."), + ) + sweep.add_argument( + "--sweep-stage", + choices=("parallel-layout", "concurrency", "scheduler", "all"), + help="Generate only the selected stage, or all stages in dependency order.", ) sweep.add_argument( "--generate-sweep", action="store_true", + help="Backward-compatible alias for --generate-scheduler-sweep.", + ) + sweep.add_argument( + "--generate-scheduler-sweep", + action="store_true", + help=("Generate the max-num-seqs/max-num-batched-tokens scheduler sweep."), + ) + sweep.add_argument( + "--generate-parallel-layout-sweep", + action="store_true", + help=( + "Generate a standalone NUMA-aware TP/DP sweep. Requires --detect-hardware." + ), + ) + sweep.add_argument( + "--generate-concurrency-sweep", + action="store_true", + help=( + "Generate a max_concurrency workload sweep using " + "vllm bench sweep serve_workload." + ), + ) + sweep.add_argument( + "--generate-full-sweep", + action="store_true", + help=( + "Generate the end-to-end TP/DP -> max_concurrency -> scheduler " + "tuning pipeline. Requires --detect-hardware." + ), + ) + sweep.add_argument( + "--tp-dp-numa-bind-workaround", + action=argparse.BooleanOptionalAction, + default=True, help=( - "Generate an optional vllm bench sweep package around the single " - "initial runtime suggestion." + "Temporary Xeon TP/DP NUMA-binding workaround. For parallel-layout " + "and full sweeps, generate an explicit VLLM_CPU_OMP_THREADS_BIND " + "from detected NUMA topology. Enabled by default; disable with " + "--no-tp-dp-numa-bind-workaround after the vLLM DP binding fix." ), ) sweep.add_argument( @@ -766,10 +816,22 @@ def write_config( Path(path).write_text("\n".join(metadata) + "\n" + body, encoding="utf-8") -def write_env(path: str, source: str, recipe: dict[str, Any]) -> None: - env = recipe.get("env") or {} - if not isinstance(env, dict): - raise ValueError(f"Recipe `env` must be an object, got {type(env).__name__}") +def write_env( + path: str, + source: str, + recipe: dict[str, Any], + *, + env_overrides: dict[str, object] | None = None, +) -> None: + recipe_env = recipe.get("env") or {} + if not isinstance(recipe_env, dict): + raise ValueError( + f"Recipe `env` must be an object, got {type(recipe_env).__name__}" + ) + + env = dict(recipe_env) + overrides = env_overrides or {} + env.update(overrides) lines = [ "#!/usr/bin/env bash", @@ -778,6 +840,9 @@ def write_env(path: str, source: str, recipe: dict[str, Any]) -> None: "", ] + if overrides: + lines.append("# Temporary runtime overrides derived from detected hardware.") + if env: for key, value in env.items(): lines.append(f"export {key}={shlex.quote(str(value))}") @@ -793,6 +858,34 @@ def main() -> int: args = parse_args() try: + sweep_modes = { + "--generate-sweep": args.generate_sweep, + "--generate-scheduler-sweep": args.generate_scheduler_sweep, + "--generate-parallel-layout-sweep": args.generate_parallel_layout_sweep, + "--generate-concurrency-sweep": args.generate_concurrency_sweep, + "--generate-full-sweep": args.generate_full_sweep, + } + selected_sweep_modes = [ + name for name, enabled in sweep_modes.items() if enabled + ] + if len(selected_sweep_modes) > 1: + raise ValueError( + "Choose exactly one sweep-generation mode: " + + ", ".join(selected_sweep_modes) + ) + if args.sweep_stage is not None: + if selected_sweep_modes: + raise ValueError( + "Do not combine --sweep-stage with --generate-*-sweep options." + ) + stage_flags = { + "parallel-layout": "generate_parallel_layout_sweep", + "concurrency": "generate_concurrency_sweep", + "scheduler": "generate_scheduler_sweep", + "all": "generate_full_sweep", + } + setattr(args, stage_flags[args.sweep_stage], True) + selected_sweep_modes = ["--sweep-stage=" + args.sweep_stage] source = args.source if source is None: source = discover_recipe_source( @@ -813,7 +906,8 @@ def main() -> int: tuning_requested = ( args.detect_hardware - or args.generate_sweep + or args.tune_scheduler + or bool(selected_sweep_modes) or any( value is not None for value in ( @@ -828,12 +922,15 @@ def main() -> int: ) tuning = None + sweep_config = None workload = None sweep_writer = None + hardware = None + env_overrides: dict[str, object] = {} if tuning_requested: # Keep plain Recipes conversion lightweight. vLLM-specific modules # are imported only for optional runtime tuning or sweep generation. - from runtime_tuning import ( + from sweep.runtime_tuning import ( WorkloadHints, finetune_runtime_config, get_runtime_tuning_policies, @@ -848,24 +945,60 @@ def main() -> int: target_qps=args.target_qps, ) - if args.generate_sweep: - from sweep_generation import ( - validate_sweep_workload, - write_sweep_files, - ) + if selected_sweep_modes: + from sweep.sweep_generation import validate_sweep_workload validate_sweep_workload(workload) + + if args.generate_sweep or args.generate_scheduler_sweep: + from sweep.sweep_generation import write_sweep_files + sweep_writer = write_sweep_files recipe_hardware = recipe.get("hardware") - policies = get_runtime_tuning_policies(recipe_hardware) + scheduler_tuning_requested = bool( + args.tune_scheduler + or args.generate_sweep + or args.generate_scheduler_sweep + ) + policies = get_runtime_tuning_policies( + recipe_hardware, + tune_scheduler=scheduler_tuning_requested, + ) - hardware = None if args.detect_hardware: - from hardware_detection import detect_hardware + from hardware_detection import ( + build_numa_omp_threads_bind, + detect_hardware, + ) hardware = detect_hardware() + if ( + args.tp_dp_numa_bind_workaround + and str(recipe_hardware).lower() == "xeon6" + and ( + args.generate_parallel_layout_sweep or args.generate_full_sweep + ) + ): + recipe_env = recipe.get("env") or {} + current_binding = ( + recipe_env.get("VLLM_CPU_OMP_THREADS_BIND") + if isinstance(recipe_env, dict) + else None + ) + + # Preserve an explicit recipe-provided manual binding. + # Replace an unset or "auto" value with the temporary + # topology-derived binding. + if current_binding in (None, "auto"): + env_overrides["VLLM_CPU_OMP_THREADS_BIND"] = ( + build_numa_omp_threads_bind( + hardware, + reserved_cores_per_numa=1, + ) + ) + tuning = finetune_runtime_config( config, hardware=hardware, @@ -873,19 +1006,73 @@ def main() -> int: policies=policies, ) config.update(tuning.overrides) + if args.tune_scheduler: + config.update(tuning.sweep_overrides) + sweep_config = dict(config) + sweep_config.update(tuning.sweep_overrides) write_config(args.config_out, source, recipe, config) - write_env(args.env_out, source, recipe) + write_env( + args.env_out, + source, + recipe, + env_overrides=env_overrides, + ) sweep_files: list[Path] = [] - if args.generate_sweep: + if args.generate_full_sweep: + if not args.detect_hardware or hardware is None: + raise ValueError("--generate-full-sweep requires --detect-hardware.") + assert workload is not None + from sweep.sweep_generation import write_full_sweep_files + + sweep_files = write_full_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + numa_node_count=hardware.numa_node_count, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_parallel_layout_sweep: + if not args.detect_hardware or hardware is None: + raise ValueError( + "--generate-parallel-layout-sweep requires --detect-hardware." + ) + assert workload is not None + from sweep.sweep_generation import write_parallel_layout_sweep_files + + sweep_files = write_parallel_layout_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + numa_node_count=hardware.numa_node_count, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_concurrency_sweep: + assert workload is not None + from sweep.sweep_generation import write_concurrency_sweep_files + + sweep_files = write_concurrency_sweep_files( + args.sweep_out_dir, + config_path=args.config_out, + env_path=args.env_out, + config=config, + workload=workload, + tune_scheduler=args.tune_scheduler, + ) + elif args.generate_sweep or args.generate_scheduler_sweep: assert workload is not None assert sweep_writer is not None + assert sweep_config is not None sweep_files = sweep_writer( args.sweep_out_dir, config_path=args.config_out, env_path=args.env_out, - config=config, + config=sweep_config, workload=workload, ) @@ -894,6 +1081,13 @@ def main() -> int: print("Initial runtime suggestion:") for key, value in tuning.overrides.items(): print(f" {key}: {value}") + if tuning.sweep_overrides: + if args.tune_scheduler: + print("Workload-derived scheduler tuning:") + elif selected_sweep_modes: + print("Explicit scheduler sweep seed:") + for key, value in tuning.sweep_overrides.items(): + print(f" {key}: {value}") for note in tuning.notes: print(f" tuning: {note}") except Exception as exc: @@ -910,15 +1104,31 @@ def main() -> int: print(f" vllm serve --config {shlex.quote(args.config_out)}") if sweep_files: sweep_dir = Path(args.sweep_out_dir) + run_parallel = sweep_dir / "run_parallel_layout_sweep.sh" + recommend_parallel = sweep_dir / "recommend_parallel_layout.py" + run_concurrency = sweep_dir / "run_concurrency_sweep.sh" + recommend_concurrency = sweep_dir / "recommend_concurrency.py" run_sweep = sweep_dir / "run_sweep.sh" recommend = sweep_dir / "recommend.py" + run_full = sweep_dir / "run_full_sweep.sh" print() print("Optional performance sweep:") - print(f" {shlex.quote(str(run_sweep))} --dry-run") - print(f" {shlex.quote(str(run_sweep))}") - print() - print("After the sweep:") - print(f" {shlex.quote(str(recommend))}") + if args.generate_full_sweep: + print(f" {shlex.quote(str(run_full))}") + elif args.generate_concurrency_sweep: + print(f" {shlex.quote(str(run_concurrency))} --dry-run") + print(f" {shlex.quote(str(run_concurrency))}") + print(f" {shlex.quote(str(recommend_concurrency))}") + elif args.generate_parallel_layout_sweep: + print(f" {shlex.quote(str(run_parallel))} --dry-run") + print(f" {shlex.quote(str(run_parallel))}") + print(f" {shlex.quote(str(recommend_parallel))}") + else: + print(f" {shlex.quote(str(run_sweep))} --dry-run") + print(f" {shlex.quote(str(run_sweep))}") + print() + print("After the sweep:") + print(f" {shlex.quote(str(recommend))}") return 0 diff --git a/tools/recipes/requirements.txt b/tools/recipes/requirements.txt new file mode 100644 index 000000000000..36a96c66b4e5 --- /dev/null +++ b/tools/recipes/requirements.txt @@ -0,0 +1,4 @@ +# Dependencies used by `vllm bench sweep plot`. +matplotlib +pandas +seaborn diff --git a/tools/recipes/sweep/__init__.py b/tools/recipes/sweep/__init__.py new file mode 100644 index 000000000000..c6a8d5b25b68 --- /dev/null +++ b/tools/recipes/sweep/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Runtime tuning and benchmark sweep generation helpers.""" diff --git a/tools/recipes/sweep/adaptive_concurrency.py b/tools/recipes/sweep/adaptive_concurrency.py new file mode 100755 index 000000000000..a0fdf36a24a5 --- /dev/null +++ b/tools/recipes/sweep/adaptive_concurrency.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Find the maximum SLA-feasible serving concurrency adaptively.""" + +from __future__ import annotations + +import argparse +import shlex +from pathlib import Path +from statistics import median +from typing import Any + +from vllm.benchmarks.sweep.param_sweep import ParameterSweepItem +from vllm.benchmarks.sweep.serve import run_server +from vllm.benchmarks.sweep.serve_workload import run_comb_workload + +DEFAULT_MINIMUM_COMPLIANCE = 0.99 +DEFAULT_MAX_CONCURRENCY_CAP = 1000 +DEFAULT_GROWTH_FACTOR = 2.0 + + +def _number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _num_prompts(concurrency: int) -> int: + return min(1000, max(100, concurrency * 10)) + + +def _sla_eligible( + rows: list[dict[str, Any]], + *, + expected_runs: int, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float, +) -> bool: + if len(rows) < expected_runs: + return False + if any(int(_number(row.get("failed")) or 0) for row in rows): + return False + + completed = sum(int(_number(row.get("completed")) or 0) for row in rows) + if completed <= 0: + return False + + compliant = sum( + (_number(row.get("request_goodput")) or 0.0) + * (_number(row.get("duration")) or 0.0) + for row in rows + ) + if compliant / completed < minimum_compliance: + return False + + if ttft_sla_ms is not None: + values = [ + value + for row in rows + if (value := _number(row.get("p99_ttft_ms"))) is not None + ] + if not values or median(values) > ttft_sla_ms: + return False + + if tpot_sla_ms is not None: + values = [ + value + for row in rows + if (value := _number(row.get("p99_tpot_ms"))) is not None + ] + if not values or median(values) > tpot_sla_ms: + return False + + return True + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Find the highest SLA-feasible max_concurrency using exponential " + "bracketing followed by integer binary search." + ) + ) + parser.add_argument("--serve-cmd", required=True) + parser.add_argument("--bench-cmd", required=True) + parser.add_argument("--results-dir", required=True) + parser.add_argument("--seed-concurrency", type=int, required=True) + parser.add_argument("--input-tokens", type=int, required=True) + parser.add_argument("--output-tokens", type=int, required=True) + parser.add_argument("--ttft-sla-ms", type=float) + parser.add_argument("--tpot-sla-ms", type=float) + parser.add_argument( + "--minimum-compliance", + type=float, + default=DEFAULT_MINIMUM_COMPLIANCE, + ) + parser.add_argument( + "--max-concurrency-cap", + type=int, + default=DEFAULT_MAX_CONCURRENCY_CAP, + ) + parser.add_argument( + "--growth-factor", + type=float, + default=DEFAULT_GROWTH_FACTOR, + ) + parser.add_argument("--num-runs", type=int, default=3) + parser.add_argument("--server-ready-timeout", type=int, default=1800) + parser.add_argument("--show-stdout", action="store_true") + parser.add_argument("--continue-on-error", action="store_true") + parser.add_argument("--resume", action="store_true") + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + + if args.seed_concurrency <= 0: + raise ValueError("--seed-concurrency must be positive.") + if not 0.0 < args.minimum_compliance <= 1.0: + raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + if args.growth_factor <= 1.0: + raise ValueError("--growth-factor must be greater than 1.") + if args.max_concurrency_cap <= 0: + raise ValueError("--max-concurrency-cap must be positive.") + if args.num_runs <= 0: + raise ValueError("--num-runs must be positive.") + if args.ttft_sla_ms is None and args.tpot_sla_ms is None: + raise ValueError("Adaptive search requires TTFT and/or TPOT SLA.") + + seed = args.seed_concurrency + cap = max(seed, args.max_concurrency_cap) + results_dir = Path(args.results_dir) + results_dir.mkdir(parents=True, exist_ok=True) + + if args.dry_run: + first_upper = min( + cap, + max(seed + 1, int(round(seed * args.growth_factor))), + ) + print("Adaptive concurrency search plan:") + print(f" seed: {seed}") + print(f" first upper probe: {first_upper}") + print(f" hard cap: {cap}") + print(" strategy: bracket SLA boundary, then integer binary search") + return 0 + + serve_cmd = shlex.split(args.serve_cmd) + bench_cmd = shlex.split(args.bench_cmd) + serve_comb = ParameterSweepItem() + + def run_point(server: object, concurrency: int) -> bool: + bench_comb = ParameterSweepItem( + { + "_benchmark_name": "user_workload", + "random_input_len": args.input_tokens, + "random_output_len": args.output_tokens, + "num_prompts": _num_prompts(concurrency), + } + ) + print(f"[ADAPTIVE CONCURRENCY] testing max_concurrency={concurrency}") + rows = ( + run_comb_workload( + server, + bench_cmd, + serve_comb=serve_comb, + bench_comb=bench_comb, + link_vars=[], + experiment_dir=results_dir, + num_runs=args.num_runs, + dry_run=False, + warmup_num_prompts=min(concurrency, 1000), + continue_on_error=args.continue_on_error, + workload_var="max_concurrency", + workload_value=concurrency, + ) + or [] + ) + + eligible = _sla_eligible( + rows, + expected_runs=args.num_runs, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + minimum_compliance=args.minimum_compliance, + ) + print( + "[ADAPTIVE CONCURRENCY] " + f"max_concurrency={concurrency}: {'PASS' if eligible else 'FAIL'}" + ) + return eligible + + with run_server( + serve_cmd, + [], + show_stdout=args.show_stdout, + serve_overrides=serve_comb, + dry_run=False, + server_ready_timeout=args.server_ready_timeout, + ) as server: + seed_passes = run_point(server, seed) + + if seed_passes: + low = seed + high: int | None = None + while low < cap: + probe = min( + cap, + max(low + 1, int(round(low * args.growth_factor))), + ) + if run_point(server, probe): + low = probe + if low == cap: + print( + "[ADAPTIVE CONCURRENCY] reached cap without an SLA failure" + ) + return 0 + else: + high = probe + break + if high is None: + return 0 + else: + high = seed + low: int | None = None + probe = seed + while probe > 1: + probe = max(1, probe // 2) + if run_point(server, probe): + low = probe + break + high = probe + if low is None: + print( + "[ADAPTIVE CONCURRENCY] no SLA-feasible concurrency found, " + "including max_concurrency=1" + ) + return 0 + + assert low is not None + assert high is not None + print(f"[ADAPTIVE CONCURRENCY] refining bracket PASS={low}, FAIL={high}") + + while high - low > 1: + midpoint = (low + high) // 2 + if run_point(server, midpoint): + low = midpoint + else: + high = midpoint + + print(f"[ADAPTIVE CONCURRENCY] boundary resolved: PASS={low}, FAIL={high}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep/concurrency_recommendation.py b/tools/recipes/sweep/concurrency_recommendation.py new file mode 100644 index 000000000000..908f8ce1da8b --- /dev/null +++ b/tools/recipes/sweep/concurrency_recommendation.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Select one max-concurrency recommendation from a workload sweep.""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path +from statistics import mean, median +from typing import Any + +DEFAULT_TTFT_SLA_MS: float | None = None +DEFAULT_TPOT_SLA_MS: float | None = None +DEFAULT_SEED_CONCURRENCY: int | None = None +DEFAULT_MINIMUM_COMPLIANCE: float = 0.99 +DEFAULT_RESULTS_DIR = "results/concurrency-tuning" +DEFAULT_OUTPUT_JSON = "concurrency-recommendation.json" + + +def _number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _mean(rows: list[dict[str, Any]], key: str) -> float | None: + values = [_number(row.get(key)) for row in rows] + numbers = [value for value in values if value is not None] + return mean(numbers) if numbers else None + + +def _percentile_summary( + rows: list[dict[str, Any]], key: str +) -> tuple[float | None, float | None, float | None]: + values = [_number(row.get(key)) for row in rows] + numbers = [value for value in values if value is not None] + if not numbers: + return None, None, None + return mean(numbers), median(numbers), max(numbers) + + +def _load_runs(results_dir: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for path in sorted(results_dir.rglob("summary.json")): + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, list): + continue + for row in data: + if isinstance(row, dict): + copied = dict(row) + copied["_summary_path"] = str(path) + rows.append(copied) + return rows + + +def _aggregate_candidates( + rows: list[dict[str, Any]], + *, + use_goodput: bool, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float, +) -> list[dict[str, Any]]: + grouped: dict[int, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + concurrency = row.get("max_concurrency") + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency <= 0 + ): + continue + grouped[concurrency].append(row) + + candidates: list[dict[str, Any]] = [] + for concurrency, runs in sorted(grouped.items()): + failed_requests = sum(int(_number(run.get("failed")) or 0) for run in runs) + output_throughput = _mean(runs, "output_throughput") + request_throughput = _mean(runs, "request_throughput") + request_goodput = _mean(runs, "request_goodput") + mean_p99_ttft_ms, median_p99_ttft_ms, worst_p99_ttft_ms = _percentile_summary( + runs, "p99_ttft_ms" + ) + mean_p99_tpot_ms, median_p99_tpot_ms, worst_p99_tpot_ms = _percentile_summary( + runs, "p99_tpot_ms" + ) + + completed_requests = sum( + int(_number(run.get("completed")) or 0) for run in runs + ) + estimated_compliant_requests = sum( + (_number(run.get("request_goodput")) or 0.0) + * (_number(run.get("duration")) or 0.0) + for run in runs + ) + combined_compliance_ratio = ( + estimated_compliant_requests / completed_requests + if completed_requests > 0 + else None + ) + + valid = failed_requests == 0 and output_throughput is not None + reason = None + if failed_requests: + valid = False + reason = f"{failed_requests} failed request(s) across successful runs" + elif use_goodput and request_goodput is None: + valid = False + reason = "request_goodput is missing; rerun with SLO goodput enabled" + + p99_sla_eligible = valid + if ttft_sla_ms is not None: + p99_sla_eligible = bool( + p99_sla_eligible + and median_p99_ttft_ms is not None + and median_p99_ttft_ms <= ttft_sla_ms + ) + if tpot_sla_ms is not None: + p99_sla_eligible = bool( + p99_sla_eligible + and median_p99_tpot_ms is not None + and median_p99_tpot_ms <= tpot_sla_ms + ) + compliance_eligible = bool( + valid + and combined_compliance_ratio is not None + and combined_compliance_ratio >= minimum_compliance + ) + + candidates.append( + { + "max_concurrency": concurrency, + "run_count": len(runs), + "failed_requests": failed_requests, + "mean_request_goodput": request_goodput, + "mean_request_throughput": request_throughput, + "mean_output_throughput": output_throughput, + "mean_p99_ttft_ms": mean_p99_ttft_ms, + "median_p99_ttft_ms": median_p99_ttft_ms, + "worst_p99_ttft_ms": worst_p99_ttft_ms, + "mean_p99_tpot_ms": mean_p99_tpot_ms, + "median_p99_tpot_ms": median_p99_tpot_ms, + "worst_p99_tpot_ms": worst_p99_tpot_ms, + "estimated_compliant_requests": estimated_compliant_requests, + "completed_requests": completed_requests, + "combined_compliance_ratio": combined_compliance_ratio, + "combined_compliance_percent": ( + combined_compliance_ratio * 100 + if combined_compliance_ratio is not None + else None + ), + "benchmark_valid": valid, + "p99_sla_eligible": p99_sla_eligible, + "compliance_eligible": compliance_eligible, + "sla_eligible": p99_sla_eligible and compliance_eligible, + "valid": valid, + "invalid_reason": reason, + "summary_files": sorted({str(run["_summary_path"]) for run in runs}), + } + ) + + return candidates + + +def _select_candidate( + candidates: list[dict[str, Any]], *, use_goodput: bool +) -> tuple[dict[str, Any] | None, dict[str, Any], str]: + valid = [candidate for candidate in candidates if candidate["valid"]] + if not valid: + raise ValueError( + "No valid workload-sweep point was found. Inspect failed runs and metrics." + ) + + if use_goodput: + objective = "highest_sla_feasible_max_concurrency" + eligible = [candidate for candidate in valid if candidate["sla_eligible"]] + winner = ( + max( + eligible, + key=lambda candidate: ( + candidate["max_concurrency"], + candidate["mean_output_throughput"], + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_request_goodput"] or 0.0, + ), + ) + if eligible + else None + ) + best_effort = max( + valid, + key=lambda candidate: ( + candidate["mean_request_goodput"] or 0.0, + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_output_throughput"], + candidate["max_concurrency"], + ), + ) + else: + objective = "highest_mean_output_throughput" + winner = max( + valid, + key=lambda candidate: ( + candidate["mean_output_throughput"], + candidate["max_concurrency"], + ), + ) + best_effort = winner + + return winner, best_effort, objective + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Select max_concurrency from vLLM Workload Explorer results." + ) + parser.add_argument("--results-dir", default=DEFAULT_RESULTS_DIR) + parser.add_argument("--ttft-sla-ms", type=float, default=DEFAULT_TTFT_SLA_MS) + parser.add_argument("--tpot-sla-ms", type=float, default=DEFAULT_TPOT_SLA_MS) + parser.add_argument( + "--seed-concurrency", type=int, default=DEFAULT_SEED_CONCURRENCY + ) + parser.add_argument( + "--minimum-compliance", + type=float, + default=DEFAULT_MINIMUM_COMPLIANCE, + ) + parser.add_argument("--output-json", default=DEFAULT_OUTPUT_JSON) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not 0.0 < args.minimum_compliance <= 1.0: + raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + + script_dir = Path(__file__).resolve().parent + results_dir = Path(args.results_dir) + if not results_dir.is_absolute(): + results_dir = script_dir / results_dir + output_json = Path(args.output_json) + if not output_json.is_absolute(): + output_json = script_dir / output_json + + if not results_dir.exists(): + raise ValueError(f"Workload-sweep results were not found at {results_dir}.") + + rows = _load_runs(results_dir) + if not rows: + raise ValueError(f"No summary.json results found under {results_dir}.") + + use_goodput = args.ttft_sla_ms is not None or args.tpot_sla_ms is not None + candidates = _aggregate_candidates( + rows, + use_goodput=use_goodput, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + minimum_compliance=args.minimum_compliance, + ) + winner, best_effort, objective = _select_candidate( + candidates, use_goodput=use_goodput + ) + + recommendation = { + "status": ( + "sla_feasible" if winner is not None else "no_sla_feasible_concurrency" + ), + "selection_objective": objective, + "seed_concurrency": args.seed_concurrency, + "minimum_compliance_ratio": args.minimum_compliance, + "slo": { + "ttft_ms": args.ttft_sla_ms, + "tpot_ms": args.tpot_sla_ms, + }, + "recommended": ( + {"max_concurrency": winner["max_concurrency"]} + if winner is not None + else None + ), + "measured": ( + { + key: winner[key] + for key in ( + "run_count", + "mean_request_throughput", + "mean_request_goodput", + "mean_output_throughput", + "combined_compliance_ratio", + "combined_compliance_percent", + "mean_p99_ttft_ms", + "median_p99_ttft_ms", + "worst_p99_ttft_ms", + "mean_p99_tpot_ms", + "median_p99_tpot_ms", + "worst_p99_tpot_ms", + ) + } + if winner is not None + else None + ), + "best_effort": { + "max_concurrency": best_effort["max_concurrency"], + "mean_request_goodput": best_effort["mean_request_goodput"], + "mean_output_throughput": best_effort["mean_output_throughput"], + "combined_compliance_ratio": best_effort["combined_compliance_ratio"], + "p99_sla_eligible": best_effort["p99_sla_eligible"], + }, + "candidates": candidates, + } + output_json.write_text( + json.dumps(recommendation, indent=2) + "\n", encoding="utf-8" + ) + + if winner is None: + print("No SLA-feasible max_concurrency was found.") + print(f"Wrote diagnostic evidence to {output_json}") + return 2 + + print(f"Recommended max_concurrency: {winner['max_concurrency']}") + print(f"Selection objective: {objective}") + print(f"Wrote {output_json}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/runtime_tuning.py b/tools/recipes/sweep/runtime_tuning.py similarity index 81% rename from tools/recipes/runtime_tuning.py rename to tools/recipes/sweep/runtime_tuning.py index b6b2d298d24d..f23e6f76ad6b 100644 --- a/tools/recipes/runtime_tuning.py +++ b/tools/recipes/sweep/runtime_tuning.py @@ -47,7 +47,10 @@ def validate(self) -> None: @dataclass class TuningResult: + """Runtime overrides and benchmark-only seeds produced by tuning policies.""" + overrides: dict[str, Any] = field(default_factory=dict) + sweep_overrides: dict[str, Any] = field(default_factory=dict) notes: list[str] = field(default_factory=list) @@ -67,13 +70,24 @@ def _record_override( result.notes.append(f"{key}={value!r}: {reason}") +def _record_sweep_override( + result: TuningResult, + key: str, + value: Any, + reason: str, +) -> None: + """Record an explicit benchmark seed without changing deployment config.""" + result.sweep_overrides[key] = value + result.notes.append(f"sweep {key}={value!r}: {reason}") + + def _get_active_sequence_count( config: dict[str, Any], workload: WorkloadHints, ) -> int: - """Return the sequence count used to size one scheduler iteration.""" + """Return the per-replica sequence count for one scheduler iteration.""" if workload.concurrency is not None: - return workload.concurrency + return ceil(workload.concurrency / _get_data_parallel_size(config)) configured = config.get("max-num-seqs") if ( @@ -86,9 +100,22 @@ def _get_active_sequence_count( return SchedulerConfig.DEFAULT_MAX_NUM_SEQS +def _get_data_parallel_size(config: dict[str, Any]) -> int: + """Return a valid configured DP size, falling back to one replica.""" + configured = config.get("data-parallel-size", 1) + if ( + isinstance(configured, int) + and not isinstance(configured, bool) + and configured > 0 + ): + return configured + return 1 + + def _estimate_prefills_per_step( workload: WorkloadHints, active_sequences: int, + data_parallel_size: int, ) -> float: """Estimate steady-state prompt arrivals per decode scheduler step.""" # Reserve enough prefill budget for at least one representative prompt. @@ -105,9 +132,10 @@ def _estimate_prefills_per_step( # With target QPS + TPOT, approximate how many prompts arrive during one # decode scheduler-step interval. if workload.target_qps is not None and workload.tpot_sla_ms is not None: + per_replica_qps = workload.target_qps / data_parallel_size prefills_per_step = max( prefills_per_step, - workload.target_qps * workload.tpot_sla_ms / 1000.0, + per_replica_qps * workload.tpot_sla_ms / 1000.0, ) # A step cannot replace more requests than are active. @@ -213,15 +241,22 @@ def _resolve_max_num_seqs( workload: WorkloadHints, result: TuningResult, ) -> None: - del config, hardware + del hardware if workload.concurrency is None: return - _record_override( + data_parallel_size = _get_data_parallel_size(config) + per_replica_concurrency = ceil(workload.concurrency / data_parallel_size) + + _record_sweep_override( result, "max-num-seqs", - workload.concurrency, - "derived from the optional target concurrency", + per_replica_concurrency, + ( + "benchmark seed derived from the optional global target concurrency " + f"and data_parallel_size={data_parallel_size}; config.yml keeps the " + "recipe/vLLM scheduler default" + ), ) @@ -235,12 +270,17 @@ def _resolve_max_num_batched_tokens( if workload.input_tokens is None: return + data_parallel_size = _get_data_parallel_size(config) active_sequences = _get_active_sequence_count(config, workload) # Decode requests consume roughly one token per active sequence in one # scheduler iteration. Prefills share the remaining scheduler token budget. decode_budget = active_sequences - prefills_per_step = _estimate_prefills_per_step(workload, active_sequences) + prefills_per_step = _estimate_prefills_per_step( + workload, + active_sequences, + data_parallel_size, + ) prefill_budget = ceil(workload.input_tokens * prefills_per_step) # Use vLLM's scheduler defaults/constraints as floors rather than local @@ -263,16 +303,18 @@ def _resolve_max_num_batched_tokens( ): candidate = max(candidate, max_model_len) - _record_override( + _record_sweep_override( result, "max-num-batched-tokens", candidate, ( - "workload-derived scheduler budget " - f"(active_sequences={active_sequences}, " + "workload-derived benchmark seed " + f"(per_replica_active_sequences={active_sequences}, " + f"data_parallel_size={data_parallel_size}, " f"prefills_per_step={prefills_per_step:.2f}, " f"decode_budget={decode_budget}, " - f"prefill_budget={prefill_budget})" + f"prefill_budget={prefill_budget}); config.yml keeps the " + "recipe/vLLM scheduler default" ), ) @@ -296,9 +338,12 @@ def _resolve_data_parallel_size( DEFAULT_POLICIES: tuple[Policy, ...] = ( _resolve_tensor_parallel_size, _resolve_gpu_memory_utilization, + _resolve_data_parallel_size, +) + +SCHEDULER_TUNING_POLICIES: tuple[Policy, ...] = ( _resolve_max_num_seqs, _resolve_max_num_batched_tokens, - _resolve_data_parallel_size, ) # Select tuning behavior from the hardware requested by the Recipes rendering, @@ -312,6 +357,8 @@ def _resolve_data_parallel_size( def get_runtime_tuning_policies( recipe_hardware: object, + *, + tune_scheduler: bool = False, ) -> tuple[Policy, ...]: if not isinstance(recipe_hardware, str) or not recipe_hardware: raise ValueError( @@ -326,6 +373,8 @@ def get_runtime_tuning_policies( "Runtime tuning is not supported for recipe hardware " f"{recipe_hardware!r}. Currently supported: {supported}." ) + if tune_scheduler: + policies = (*policies, *SCHEDULER_TUNING_POLICIES) return policies @@ -336,7 +385,7 @@ def finetune_runtime_config( workload: WorkloadHints | None = None, policies: tuple[Policy, ...] = DEFAULT_POLICIES, ) -> TuningResult: - """Return overrides; never mutate the recipe-derived config in place.""" + """Return deployment overrides and sweep seeds without mutating config.""" hints = workload or WorkloadHints() hints.validate() diff --git a/tools/recipes/sweep/scheduler_sweep_prepare.py b/tools/recipes/sweep/scheduler_sweep_prepare.py new file mode 100644 index 000000000000..59709a0fc56d --- /dev/null +++ b/tools/recipes/sweep/scheduler_sweep_prepare.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Prepare scheduler sweep inputs after TP/DP and concurrency selection.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +try: + import yaml +except ImportError as exc: + raise SystemExit("PyYAML is required. Install it with: pip install pyyaml") from exc + +PROMPTS_PER_CONCURRENCY = 10 +MIN_NUM_PROMPTS = 100 +MAX_NUM_PROMPTS = 1000 + + +def _positive_int(config: dict[str, Any], key: str) -> int: + value = config.get(key) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"Expected positive integer {key!r}, got {value!r}.") + return value + + +def _num_prompts_for_concurrency(concurrency: int) -> int: + return min( + MAX_NUM_PROMPTS, + max(MIN_NUM_PROMPTS, concurrency * PROMPTS_PER_CONCURRENCY), + ) + + +def _strict_lower_power_of_two(value: int) -> int: + if value <= 1: + return 1 + return 1 << ((value - 1).bit_length() - 1) + + +def _strict_upper_power_of_two(value: int) -> int: + return 1 << value.bit_length() + + +def _scheduler_baseline( + config: dict[str, Any], + *, + concurrency: int, + input_tokens: int, + output_tokens: int, + tpot_sla_ms: float | None, + target_qps: float | None, +) -> tuple[int, int]: + data_parallel_size = int(config.get("data-parallel-size", 1)) + if data_parallel_size <= 0: + raise ValueError("data-parallel-size must be positive.") + + per_replica_concurrency = math.ceil(concurrency / data_parallel_size) + prefills_per_step = max(1.0, per_replica_concurrency / max(output_tokens, 1)) + if target_qps is not None and tpot_sla_ms is not None: + per_replica_qps = target_qps / data_parallel_size + prefills_per_step = max( + prefills_per_step, + per_replica_qps * tpot_sla_ms / 1000.0, + ) + prefills_per_step = min(float(per_replica_concurrency), prefills_per_step) + + batch = max( + 2048, + per_replica_concurrency, + per_replica_concurrency + math.ceil(input_tokens * prefills_per_step), + ) + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool): + batch = max(batch, max_model_len) + + return per_replica_concurrency, batch + + +def _build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: + initial_seqs = _positive_int(config, "max-num-seqs") + initial_batch = _positive_int(config, "max-num-batched-tokens") + + minimum_batch = initial_seqs + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if ( + isinstance(max_model_len, int) + and not isinstance(max_model_len, bool) + and max_model_len > 0 + ): + minimum_batch = max(minimum_batch, max_model_len) + + lower_batch = max(minimum_batch, _strict_lower_power_of_two(initial_batch)) + smaller_batch = max(minimum_batch, _strict_lower_power_of_two(lower_batch)) + higher_batch = max(minimum_batch, _strict_upper_power_of_two(initial_batch)) + + candidates: list[dict[str, Any]] = [] + seen: set[tuple[int | None, int | None]] = set() + + def add( + name: str, + max_num_seqs: int | None, + max_num_batched_tokens: int | None, + ) -> None: + if max_num_seqs is not None and max_num_batched_tokens is not None: + max_num_batched_tokens = max(max_num_batched_tokens, max_num_seqs) + signature = (max_num_seqs, max_num_batched_tokens) + if signature in seen: + return + seen.add(signature) + + candidate: dict[str, Any] = {"_benchmark_name": name} + if max_num_seqs is not None: + candidate["max_num_seqs"] = max_num_seqs + if max_num_batched_tokens is not None: + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + add("initial", initial_seqs, initial_batch) + add("smaller_batch_budget", initial_seqs, smaller_batch) + add("lower_batch_budget", initial_seqs, lower_batch) + add("higher_batch_budget", initial_seqs, higher_batch) + add("vllm_default_max_num_seqs", None, initial_batch) + add("vllm_default_max_num_batched_tokens", initial_seqs, None) + add("vllm_defaults", None, None) + return candidates + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--concurrency-recommendation", required=True) + parser.add_argument("--input-tokens", type=int, required=True) + parser.add_argument("--output-tokens", type=int, required=True) + parser.add_argument("--tpot-sla-ms", type=float) + parser.add_argument("--target-qps", type=float) + parser.add_argument("--output-config", required=True) + parser.add_argument("--serve-params", required=True) + parser.add_argument("--bench-params", required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + config_path = Path(args.config) + recommendation_path = Path(args.concurrency_recommendation) + + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError(f"{config_path} does not contain a YAML object.") + + recommendation = json.loads(recommendation_path.read_text(encoding="utf-8")) + selected = recommendation.get("recommended") + if not isinstance(selected, dict): + raise ValueError( + f"{recommendation_path} has no SLA-feasible concurrency recommendation." + ) + concurrency = selected.get("max_concurrency") + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency <= 0 + ): + raise ValueError("Recommended max_concurrency is missing or invalid.") + + seqs, batch = _scheduler_baseline( + config, + concurrency=concurrency, + input_tokens=args.input_tokens, + output_tokens=args.output_tokens, + tpot_sla_ms=args.tpot_sla_ms, + target_qps=args.target_qps, + ) + + scheduler_seed = dict(config) + scheduler_seed["max-num-seqs"] = seqs + scheduler_seed["max-num-batched-tokens"] = batch + + scheduler_config = dict(config) + scheduler_config.pop("max-num-seqs", None) + scheduler_config.pop("max-num-batched-tokens", None) + Path(args.output_config).write_text( + yaml.safe_dump( + scheduler_config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ), + encoding="utf-8", + ) + + Path(args.serve_params).write_text( + json.dumps(_build_serve_params(scheduler_seed), indent=2) + "\n", + encoding="utf-8", + ) + Path(args.bench_params).write_text( + json.dumps( + [ + { + "random_input_len": args.input_tokens, + "random_output_len": args.output_tokens, + "max_concurrency": concurrency, + "num_prompts": _num_prompts_for_concurrency(concurrency), + } + ], + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + print(f"Prepared scheduler sweep for max_concurrency={concurrency}") + print(f"Scheduler seed max-num-seqs={seqs}") + print(f"Scheduler seed max-num-batched-tokens={batch}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep/sweep_generation.py b/tools/recipes/sweep/sweep_generation.py new file mode 100644 index 000000000000..cd3a131f8951 --- /dev/null +++ b/tools/recipes/sweep/sweep_generation.py @@ -0,0 +1,1531 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate optional vLLM benchmark sweep files from one initial suggestion.""" + +from __future__ import annotations + +import json +import math +import os +import shlex +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + +from .runtime_tuning import WorkloadHints + +SUPPORTED_TENSOR_PARALLEL_SIZES = frozenset({1, 2, 4, 8}) + +MOE_EXPERT_COUNT_KEYS = ( + "num_experts", + "num_local_experts", + "num_sparse_experts", + "num_total_experts", + "n_experts", + "n_routed_experts", + "num_routed_experts", +) +MOE_INTERMEDIATE_SIZE_KEYS = ( + "moe_intermediate_size", + "expert_intermediate_size", + "moe_ffn_hidden_size", +) + +PROMPTS_PER_CONCURRENCY = 10 +MIN_NUM_PROMPTS = 100 +MAX_NUM_PROMPTS = 1000 + + +@dataclass(frozen=True) +class ModelParallelMetadata: + """Model dimensions needed to reject invalid MoE layouts before startup.""" + + is_moe: bool | None + num_experts: int | None = None + moe_intermediate_size: int | None = None + + +UNKNOWN_MODEL_PARALLEL_METADATA = ModelParallelMetadata(is_moe=None) + + +def _num_prompts_for_concurrency(concurrency: int) -> int: + return min( + MAX_NUM_PROMPTS, + max(MIN_NUM_PROMPTS, concurrency * PROMPTS_PER_CONCURRENCY), + ) + + +def _positive_int(config: dict[str, Any], key: str) -> int: + value = config.get(key) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"Sweep generation requires a positive initial {key!r} value.") + return value + + +def _strict_lower_power_of_two(value: int) -> int: + if value <= 1: + return 1 + return 1 << ((value - 1).bit_length() - 1) + + +def _strict_upper_power_of_two(value: int) -> int: + return 1 << value.bit_length() + + +def validate_sweep_workload(workload: WorkloadHints) -> None: + required = { + "--input-tokens": workload.input_tokens, + "--output-tokens": workload.output_tokens, + "--concurrency": workload.concurrency, + } + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError("--generate-sweep requires " + ", ".join(missing) + ".") + + +def build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: + """Build a bounded batch-budget sweep at the sequence-count lower bound. + + Explicit candidates keep the initial per-replica sequence count. Three + reference points compare those tuned values with vLLM-resolved defaults. + """ + initial_seqs = _positive_int(config, "max-num-seqs") + initial_batch = _positive_int(config, "max-num-batched-tokens") + + minimum_batch = initial_seqs + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if ( + isinstance(max_model_len, int) + and not isinstance(max_model_len, bool) + and max_model_len > 0 + ): + minimum_batch = max(minimum_batch, max_model_len) + + lower_batch = max( + minimum_batch, + _strict_lower_power_of_two(initial_batch), + ) + smaller_batch = max( + minimum_batch, + _strict_lower_power_of_two(lower_batch), + ) + higher_batch = max( + minimum_batch, + _strict_upper_power_of_two(initial_batch), + ) + + candidates: list[dict[str, Any]] = [] + seen: set[tuple[int | None, int | None]] = set() + + def add( + name: str, + max_num_seqs: int | None, + max_num_batched_tokens: int | None, + ) -> None: + if max_num_seqs is not None and max_num_batched_tokens is not None: + max_num_batched_tokens = max(max_num_batched_tokens, max_num_seqs) + signature = (max_num_seqs, max_num_batched_tokens) + if signature in seen: + return + seen.add(signature) + + candidate: dict[str, Any] = {"_benchmark_name": name} + if max_num_seqs is not None: + candidate["max_num_seqs"] = max_num_seqs + if max_num_batched_tokens is not None: + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + # Keep the exact initial suggestion as the measured baseline. Additional + # values exist only in the optional sweep package. + add("initial", initial_seqs, initial_batch) + add("smaller_batch_budget", initial_seqs, smaller_batch) + add("lower_batch_budget", initial_seqs, lower_batch) + add("higher_batch_budget", initial_seqs, higher_batch) + # Default-reference candidates intentionally omit one or both scheduler + # keys. The sweep server starts from sweep_config.yml, where both keys are + # removed, so omission lets vLLM resolve its normal runtime default. + add("vllm_default_max_num_seqs", None, initial_batch) + add("vllm_default_max_num_batched_tokens", initial_seqs, None) + add("vllm_defaults", None, None) + + return candidates + + +def _scheduler_baseline_for_dp( + config: dict[str, Any], workload: WorkloadHints, data_parallel_size: int +) -> tuple[int, int]: + """Return per-replica scheduler settings for a DP layout.""" + assert workload.concurrency is not None + assert workload.input_tokens is not None + per_replica_concurrency = math.ceil(workload.concurrency / data_parallel_size) + output_tokens = workload.output_tokens or 1 + prefills_per_step = max(1.0, per_replica_concurrency / output_tokens) + if workload.target_qps is not None and workload.tpot_sla_ms is not None: + per_replica_qps = workload.target_qps / data_parallel_size + prefills_per_step = max( + prefills_per_step, + per_replica_qps * workload.tpot_sla_ms / 1000.0, + ) + prefills_per_step = min(float(per_replica_concurrency), prefills_per_step) + batch = max( + 2048, + per_replica_concurrency, + per_replica_concurrency + math.ceil(workload.input_tokens * prefills_per_step), + ) + if config.get("enable-chunked-prefill") is False: + max_model_len = config.get("max-model-len") + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool): + batch = max(batch, max_model_len) + return per_replica_concurrency, batch + + +def _model_config_value(view: Any, key: str) -> Any: + if isinstance(view, dict): + return view.get(key) + return getattr(view, key, None) + + +def _config_views(hf_config: Any) -> list[Any]: + """Return common nested model config views without duplicates.""" + views = [hf_config] + for view in views: + for key in ( + "text_config", + "language_config", + "llm_config", + "ffn_config", + "moe_config", + ): + value = _model_config_value(view, key) + if value is not None and all(value is not item for item in views): + views.append(value) + return views + + +def _first_positive_model_int(views: list[Any], keys: tuple[str, ...]) -> int | None: + for view in views: + for key in keys: + value = _model_config_value(view, key) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + return None + + +def _config_identifies_moe(views: list[Any]) -> bool: + for view in views: + model_type = _model_config_value(view, "model_type") + if isinstance(model_type, str) and "moe" in model_type.lower(): + return True + architectures = _model_config_value(view, "architectures") + if isinstance(architectures, list) and any( + isinstance(name, str) and "moe" in name.lower() for name in architectures + ): + return True + return False + + +def model_parallel_metadata_from_config(hf_config: Any) -> ModelParallelMetadata: + """Extract MoE markers from a loaded Transformers-compatible config.""" + views = _config_views(hf_config) + num_experts = _first_positive_model_int(views, MOE_EXPERT_COUNT_KEYS) + moe_intermediate_size = _first_positive_model_int(views, MOE_INTERMEDIATE_SIZE_KEYS) + is_moe = ( + num_experts is not None + or moe_intermediate_size is not None + or _config_identifies_moe(views) + ) + if is_moe and moe_intermediate_size is None: + # Some MoE configs use the generic name for each expert's FFN width. + moe_intermediate_size = _first_positive_model_int(views, ("intermediate_size",)) + return ModelParallelMetadata( + is_moe=is_moe, + num_experts=num_experts, + moe_intermediate_size=moe_intermediate_size, + ) + + +def inspect_model_parallel_metadata(config: dict[str, Any]) -> ModelParallelMetadata: + """Load the model config and extract portable MoE topology markers.""" + model = config.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Parallel-layout sweep requires a model in config.yml.") + + from vllm.transformers_utils.config import get_config + + hf_config = get_config( + model, + trust_remote_code=config.get("trust-remote-code") is True, + revision=config.get("revision"), + code_revision=config.get("code-revision"), + config_format=config.get("config-format", "auto"), + ) + return model_parallel_metadata_from_config(hf_config) + + +def cpu_supports_moe_dp_collectives() -> bool: + """Whether CPU overrides the variable-size collectives used by MoE DP.""" + try: + from vllm.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase, + ) + from vllm.distributed.device_communicators.cpu_communicator import ( + CpuCommunicator, + ) + except Exception: + return False + + return all( + getattr(CpuCommunicator, name) is not getattr(DeviceCommunicatorBase, name) + for name in ("all_gatherv", "reduce_scatterv") + ) + + +def _parallel_layout_name( + tensor_parallel_size: int, data_parallel_size: int, numa_node_count: int +) -> str: + name = f"tp{tensor_parallel_size}_dp{data_parallel_size}" + used_nodes = tensor_parallel_size * data_parallel_size + if used_nodes != numa_node_count: + name += f"_numa{used_nodes}of{numa_node_count}" + return name + + +def _default_parallel_layout_sizes(numa_node_count: int) -> list[tuple[int, int]]: + max_supported_tp = max( + size for size in SUPPORTED_TENSOR_PARALLEL_SIZES if size <= numa_node_count + ) + layouts = [] + for tensor_parallel_size in range(numa_node_count, 0, -1): + if tensor_parallel_size not in SUPPORTED_TENSOR_PARALLEL_SIZES: + continue + uses_all_numa_nodes = numa_node_count % tensor_parallel_size == 0 + if not uses_all_numa_nodes and tensor_parallel_size != max_supported_tp: + continue + layouts.append( + (tensor_parallel_size, max(1, numa_node_count // tensor_parallel_size)) + ) + return layouts + + +def build_parallel_layout_plan( + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + *, + model_metadata: ModelParallelMetadata = UNKNOWN_MODEL_PARALLEL_METADATA, + cpu_moe_dp_supported: bool = True, + tune_scheduler: bool = False, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Build candidates and diagnostics after applying parallel safety policy.""" + validate_sweep_workload(workload) + if numa_node_count <= 1: + raise ValueError("Parallel-layout sweep requires at least two NUMA nodes.") + + layouts = _default_parallel_layout_sizes(numa_node_count) + skipped: list[dict[str, Any]] = [] + if model_metadata.is_moe and not cpu_moe_dp_supported: + for tensor_parallel_size, data_parallel_size in layouts: + if data_parallel_size <= 1: + continue + skipped.append( + { + "_benchmark_name": _parallel_layout_name( + tensor_parallel_size, data_parallel_size, numa_node_count + ), + "policy": "cpu_moe_dp_collectives", + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + "reason": ( + "CPU communicator does not implement the variable-size " + "collectives required by MoE data parallelism" + ), + } + ) + # Keep TP coverage instead of collapsing the sweep to whichever DP=1 + # layout happened to exist in the full-NUMA candidate set. + layouts = [ + (tensor_parallel_size, 1) + for tensor_parallel_size in sorted( + SUPPORTED_TENSOR_PARALLEL_SIZES, reverse=True + ) + if tensor_parallel_size <= numa_node_count + ] + + pcp_size = config.get("prefill-context-parallel-size", 1) + if not isinstance(pcp_size, int) or isinstance(pcp_size, bool) or pcp_size <= 0: + raise ValueError("prefill-context-parallel-size must be a positive integer.") + uses_expert_parallel = config.get("enable-expert-parallel") is True + + candidates = [] + for tensor_parallel_size, data_parallel_size in layouts: + name = _parallel_layout_name( + tensor_parallel_size, data_parallel_size, numa_node_count + ) + effective_parallel_size = tensor_parallel_size * data_parallel_size * pcp_size + moe_width = model_metadata.moe_intermediate_size + if ( + model_metadata.is_moe + and not uses_expert_parallel + and moe_width is not None + and moe_width % effective_parallel_size != 0 + ): + skipped.append( + { + "_benchmark_name": name, + "policy": "moe_effective_parallel_size", + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + "prefill_context_parallel_size": pcp_size, + "effective_parallel_size": effective_parallel_size, + "moe_intermediate_size": moe_width, + "reason": ( + "moe_intermediate_size is not divisible by the effective " + "TP x DP x PCP size while expert parallelism is disabled" + ), + } + ) + continue + + candidate = { + "_benchmark_name": name, + "tensor_parallel_size": tensor_parallel_size, + "data_parallel_size": data_parallel_size, + } + if tune_scheduler: + max_num_seqs, max_num_batched_tokens = _scheduler_baseline_for_dp( + config, workload, data_parallel_size + ) + candidate["max_num_seqs"] = max_num_seqs + candidate["max_num_batched_tokens"] = max_num_batched_tokens + candidates.append(candidate) + + if not candidates: + raise ValueError( + "No valid parallel-layout candidates remain after applying model and " + "runtime capability checks. See parallel_layout_skips.json." + ) + return candidates, skipped + + +def build_parallel_layout_params( + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + *, + model_metadata: ModelParallelMetadata = UNKNOWN_MODEL_PARALLEL_METADATA, + cpu_moe_dp_supported: bool = True, + tune_scheduler: bool = False, +) -> list[dict[str, Any]]: + """Build full-NUMA layouts plus the largest supported TP layout.""" + candidates, _ = build_parallel_layout_plan( + config, + workload, + numa_node_count, + model_metadata=model_metadata, + cpu_moe_dp_supported=cpu_moe_dp_supported, + tune_scheduler=tune_scheduler, + ) + return candidates + + +def build_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: + validate_sweep_workload(workload) + assert workload.input_tokens is not None + assert workload.output_tokens is not None + assert workload.concurrency is not None + + return [ + { + "random_input_len": workload.input_tokens, + "random_output_len": workload.output_tokens, + "max_concurrency": workload.concurrency, + "num_prompts": _num_prompts_for_concurrency(workload.concurrency), + } + ] + + +def _benchmark_models(config: dict[str, Any]) -> tuple[str, str]: + model = config.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Sweep generation requires a model in config.yml.") + + served_model_name = config.get("served-model-name") + if isinstance(served_model_name, str) and served_model_name: + return served_model_name, model + if ( + isinstance(served_model_name, list) + and served_model_name + and isinstance(served_model_name[0], str) + ): + return served_model_name[0], model + + return model, model + + +def _relative_to(directory: Path, target: str) -> str: + return os.path.relpath(Path(target).resolve(), directory.resolve()) + + +def _write_json(path: Path, value: object) -> None: + path.write_text( + json.dumps(value, indent=2) + "\n", + encoding="utf-8", + ) + + +def _write_sweep_config(path: Path, config: dict[str, Any]) -> None: + sweep_config = dict(config) + sweep_config.pop("max-num-seqs", None) + sweep_config.pop("max-num-batched-tokens", None) + + body = yaml.safe_dump( + sweep_config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ) + header = ( + "# Generated for runtime sweep only.\n" + "# Scheduler values are supplied per candidate; omitted values use " + "vLLM defaults.\n" + ) + path.write_text(header + body, encoding="utf-8") + + +def _write_run_script( + path: Path, + *, + config_rel: str, + env_rel: str, + request_model: str, + tokenizer: str, + workload: WorkloadHints, + serve_params_name: str = "serve_params.json", + experiment_name: str = "runtime-tuning", + prepare_config_rel: str | None = None, +) -> None: + bench_parts = [ + "vllm bench serve", + "--backend vllm", + f"--model {shlex.quote(request_model)}", + f"--tokenizer {shlex.quote(tokenizer)}", + "--dataset-name random", + "--request-rate inf", + "--ignore-eos", + "--metric-percentiles 99", + ] + + goodput_pairs: list[str] = [] + if workload.ttft_sla_ms is not None: + goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") + if goodput_pairs: + bench_parts.append( + "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) + ) + + bench_cmd = " ".join(bench_parts) + prepare_config = "" + if prepare_config_rel is not None: + prepare_config = f""" +SOURCE_CONFIG_PATH="${{SCRIPT_DIR}}/{prepare_config_rel}" +python3 - "${{SOURCE_CONFIG_PATH}}" "${{CONFIG_PATH}}" <<'PY' +import sys +from pathlib import Path + +import yaml + +source = Path(sys.argv[1]) +target = Path(sys.argv[2]) +config = yaml.safe_load(source.read_text(encoding="utf-8")) +if not isinstance(config, dict): + raise SystemExit(str(source) + " does not contain a YAML configuration object.") +config.pop("max-num-seqs", None) +config.pop("max-num-batched-tokens", None) +target.write_text( + yaml.safe_dump( + config, + sort_keys=False, + default_flow_style=False, + allow_unicode=True, + ), + encoding="utf-8", +) +PY + +""" + + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" +ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" + +source "${{ENV_PATH}}" + +# Released/nightly vLLM images keep recipe-relative assets such as +# examples/*.jinja under /vllm-workspace. Launch sweep server subprocesses +# from there when available so relative recipe assets resolve like they do +# in serve_with_recipe.sh. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +# recipe tools may be newer than the installed vLLM package. Only pass the +# resilience flag when this vLLM CLI actually supports it. +CONTINUE_ON_ERROR_ARG="" +if vllm bench sweep serve --help 2>&1 | grep -q -- "--continue-on-error"; then + CONTINUE_ON_ERROR_ARG="--continue-on-error" +fi + +# On older vLLM versions, a failed benchmark aborts the sweep process. Retry the +# interrupted sweep with --resume so one transient run failure does not discard +# completed work or force a manual restart. +SWEEP_RETRY_COUNT="${{VLLM_RECIPE_SWEEP_RETRY_COUNT:-0}}" +MAX_SWEEP_RETRIES="${{VLLM_RECIPE_SWEEP_RETRIES:-2}}" + +{prepare_config}vllm bench sweep serve \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --serve-params "${{SCRIPT_DIR}}/{serve_params_name}" \ + --bench-params "${{SCRIPT_DIR}}/bench_params.json" \ + --output-dir "${{SCRIPT_DIR}}/results" \ + --experiment-name {experiment_name} \ + --warmup-num-prompts {workload.concurrency} \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" || {{ + if [[ "${{SWEEP_RETRY_COUNT}}" -ge "${{MAX_SWEEP_RETRIES}}" ]]; then + echo "Sweep failed after ${{MAX_SWEEP_RETRIES}} automatic retries." >&2 + exit 1 + fi + + NEXT_RETRY=$((SWEEP_RETRY_COUNT + 1)) + RETRY_ARGS=() + for arg in "$@"; do + [[ "$arg" == "--resume" ]] && continue + RETRY_ARGS+=("$arg") + done + + EXPERIMENT_DIR="${{SCRIPT_DIR}}/results/{experiment_name}" + if [[ -d "${{EXPERIMENT_DIR}}" ]]; then + echo "Sweep command failed;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} with --resume" + RETRY_ARGS=(--resume "${{RETRY_ARGS[@]}}") + else + echo "Sweep command failed before resumable state was created;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} from the start" + fi + + VLLM_RECIPE_SWEEP_RETRY_COUNT="${{NEXT_RETRY}}" exec "$0" "${{RETRY_ARGS[@]}}" + }} +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_recommend_script( + path: Path, + *, + config_rel: str, + env_rel: str, + workload: WorkloadHints, + results_dir: str = "results/runtime-tuning", + output_config: str = "recommended-config.yml", + output_json: str = "recommendation.json", +) -> None: + template_path = Path(__file__).with_name("sweep_recommendation.py") + source = template_path.read_text(encoding="utf-8") + + replacements = { + "DEFAULT_CONFIG_PATH: str | None = None": ( + f"DEFAULT_CONFIG_PATH: str | None = {config_rel!r}" + ), + "DEFAULT_ENV_PATH: str | None = None": ( + f"DEFAULT_ENV_PATH: str | None = {env_rel!r}" + ), + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + 'DEFAULT_RESULTS_DIR = "results/runtime-tuning"': ( + f"DEFAULT_RESULTS_DIR = {results_dir!r}" + ), + 'DEFAULT_OUTPUT_CONFIG = "recommended-config.yml"': ( + f"DEFAULT_OUTPUT_CONFIG = {output_config!r}" + ), + 'DEFAULT_OUTPUT_JSON = "recommendation.json"': ( + f"DEFAULT_OUTPUT_JSON = {output_json!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"Recommender template marker not found: {marker}") + source = source.replace(marker, replacement, 1) + + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_guide(path: Path, workload: WorkloadHints) -> None: + sla_lines: list[str] = [] + if workload.ttft_sla_ms is not None: + sla_lines.append(f"- TTFT objective: `{workload.ttft_sla_ms:g} ms`") + if workload.tpot_sla_ms is not None: + sla_lines.append(f"- TPOT objective: `{workload.tpot_sla_ms:g} ms`") + + sla_text = "\n".join(sla_lines) + if sla_text: + sla_text = ( + "\n## Supplied latency objectives\n\n" + + sla_text + + "\n\nThe generated benchmark uses these values with vLLM " + "`--goodput`.\n" + ) + + content = f"""# Optional Runtime Tuning Sweep + +The generated `config.yml` is the **single initial suggestion** and can be +deployed directly. The scheduler sweep keeps explicit `max-num-seqs` at the +per-replica concurrency lower bound while benchmarking nearby values for: + +- `max-num-batched-tokens` + +It does not test explicit `max-num-seqs` values below per-replica concurrency, +because those values can introduce scheduler queueing and severe TTFT +degradation. The vLLM-default references still test whether an explicit +`max-num-seqs` override is needed. + +In addition to nearby tuned values, the sweep includes three vLLM-default +references: + +- vLLM default `max-num-seqs` with the initial batch-token budget +- initial `max-num-seqs` with vLLM default `max-num-batched-tokens` +- vLLM defaults for both parameters + +The sweep runs every server from a generated `sweep_config.yml` that copies the +initial `config.yml` but removes these two scheduler keys. Tuned candidates add +explicit CLI values; a default-reference candidate simply omits the selected +key, allowing `vllm serve` to resolve its normal runtime default. This avoids +passing the literal string `None` and keeps defaults platform-, world-size-, +model-, and usage-context-aware. + +## Run + +```bash +./run_sweep.sh --dry-run +./run_sweep.sh +./recommend.py +``` + +For a quick one-run experiment: + +```bash +rm -rf results/runtime-tuning +./run_sweep.sh --num-runs 1 +./recommend.py +``` + +Resume an interrupted sweep: + +```bash +./run_sweep.sh --resume +``` + +Before measurement, the generated script runs one unmeasured warmup containing +one full concurrency window ({workload.concurrency} prompts). It is saved as +`warmup.json` for auditing but excluded from recommendation statistics. By +default, vLLM then benchmarks each parameter combination three times. + +The benchmark request count scales with the supplied workload concurrency using +the same model as vLLM's `PROMPTS_PER_CONCURRENCY` performance-benchmark +control: + +```text +num_prompts = min(1000, max(100, concurrency * 10)) +``` + +This targets about ten concurrency turnovers when neither bound applies, keeps +at least 100 requests for percentile/compliance measurements, and caps each run +at 1000 requests. All scheduler candidates use the same request count. +{sla_text} +## Outputs + +`recommend.py` writes: + +```text +recommended-config.yml +recommendation.json +``` + +With TTFT/TPOT objectives it requires duration-weighted combined compliance of +at least 99% plus median P99 TTFT/TPOT compliance. Scheduler candidates within +1% of the highest eligible output-token throughput are treated as practically +equivalent, and the recommender prefers fewer explicit scheduler overrides +within that set. Use `--throughput-equivalence-percent 0` for exact +highest-throughput selection. If no candidate qualifies, it records a +best-effort candidate but does not write `recommended-config.yml`. Without +latency objectives it applies the same scheduler equivalence policy. +Configurations with failed requests are excluded. TP/DP comparisons continue +to use exact highest-throughput selection. + +`recommended-config.yml` copies the initial configuration and changes only +`max-num-seqs` and `max-num-batched-tokens`. If a vLLM-default reference wins, +the corresponding key is removed from `recommended-config.yml`, allowing +`vllm serve` to resolve that parameter from its normal runtime default policy. +""" + path.write_text(content, encoding="utf-8") + + +def _write_report_script(path: Path, *, workload: WorkloadHints) -> None: + template_path = Path(__file__).with_name("sweep_html_report.py") + source = template_path.read_text(encoding="utf-8") + replacements = { + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"HTML report template marker not found: {marker}") + source = source.replace(marker, replacement, 1) + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_post_benchmark_analysis_files( + directory: Path, *, workload: WorkloadHints +) -> list[Path]: + """Copy standalone reporting helpers into a generated sweep package.""" + source_dir = Path(__file__).resolve().parent + generated: list[Path] = [] + for source_name, target_name, executable in ( + ("sweep_visualization.py", "visualize.py", True), + ("requirements.txt", "requirements.txt", False), + ("VISUALIZATION.md", "VISUALIZATION.md", False), + ): + source = ( + source_dir / source_name + if source_name.endswith(".py") + else source_dir.parent / source_name + ) + target = directory / target_name + target.write_text(source.read_text(encoding="utf-8"), encoding="utf-8") + if executable: + target.chmod(target.stat().st_mode | 0o111) + generated.append(target) + report = directory / "report.py" + _write_report_script(report, workload=workload) + generated.append(report) + return generated + + +def write_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, +) -> list[Path]: + """Write an optional sweep package around the initial config suggestion.""" + validate_sweep_workload(workload) + + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + sweep_config = directory / "sweep_config.yml" + serve_params = directory / "serve_params.json" + bench_params = directory / "bench_params.json" + run_script = directory / "run_sweep.sh" + recommend_script = directory / "recommend.py" + guide = directory / "SWEEP.md" + + config_rel = _relative_to(directory, config_path) + sweep_config_rel = _relative_to(directory, str(sweep_config)) + env_rel = _relative_to(directory, env_path) + + _write_sweep_config(sweep_config, config) + _write_json(serve_params, build_serve_params(config)) + _write_json(bench_params, build_bench_params(workload)) + request_model, tokenizer = _benchmark_models(config) + _write_run_script( + run_script, + config_rel=sweep_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_recommend_script( + recommend_script, + config_rel=config_rel, + env_rel=env_rel, + workload=workload, + ) + _write_guide(guide, workload) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + + return [ + sweep_config, + serve_params, + bench_params, + run_script, + recommend_script, + guide, + *analysis_files, + ] + + +def write_parallel_layout_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + tune_scheduler: bool = False, +) -> list[Path]: + """Write a standalone NUMA-aware TP/DP sweep package.""" + validate_sweep_workload(workload) + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + parallel_params = directory / "parallel_layout_serve_params.json" + parallel_skips = directory / "parallel_layout_skips.json" + parallel_config = directory / "parallel-layout-benchmark-config.yml" + bench_params = directory / "bench_params.json" + run_parallel = directory / "run_parallel_layout_sweep.sh" + recommend_parallel = directory / "recommend_parallel_layout.py" + guide = directory / "SWEEP.md" + + initial_config_rel = _relative_to(directory, config_path) + parallel_config_rel = _relative_to(directory, str(parallel_config)) + selected_config_rel = "parallel-layout-config.yml" + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_sweep_config(parallel_config, config) + benchmark_config_rel = initial_config_rel if tune_scheduler else parallel_config_rel + + try: + model_metadata = inspect_model_parallel_metadata(config) + except Exception as exc: + print( + "Warning: could not inspect model parallel metadata; retaining the " + f"generic layouts and relying on sweep error isolation: {exc}", + file=sys.stderr, + ) + model_metadata = UNKNOWN_MODEL_PARALLEL_METADATA + cpu_moe_dp_supported = ( + cpu_supports_moe_dp_collectives() if model_metadata.is_moe else True + ) + parallel_candidates, skipped_candidates = build_parallel_layout_plan( + config, + workload, + numa_node_count, + model_metadata=model_metadata, + cpu_moe_dp_supported=cpu_moe_dp_supported, + tune_scheduler=tune_scheduler, + ) + _write_json(parallel_params, parallel_candidates) + _write_json(parallel_skips, skipped_candidates) + _write_json(bench_params, build_bench_params(workload)) + _write_run_script( + run_parallel, + config_rel=benchmark_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + serve_params_name=parallel_params.name, + experiment_name="parallel-layout", + ) + _write_recommend_script( + recommend_parallel, + config_rel=initial_config_rel, + env_rel=env_rel, + workload=workload, + results_dir="results/parallel-layout", + output_config=selected_config_rel, + output_json="parallel-layout-recommendation.json", + ) + + guide.write_text( + f"""# Parallel-Layout Sweep + +This package tunes only the NUMA-aware TP/DP layout. + +Primary candidates use every effective NUMA node: + +```text +tensor_parallel_size * data_parallel_size = {numa_node_count} +``` + +Tensor parallelism is restricted to `1`, `2`, `4`, or `8`. The largest +supported TP not exceeding the NUMA-node count is also included, even when it +leaves some NUMA nodes idle. + +Before benchmarking, the generator applies two MoE safety checks: + +- When the CPU communicator lacks variable-size MoE collectives, MoE candidates + use `DP=1`; every supported TP size is retained as a replacement candidate. +- With expert parallelism disabled and a known MoE intermediate size, candidates + must satisfy `moe_intermediate_size % (TP * DP * PCP) == 0`. + +Rejected layouts and their reasons are written to +`parallel_layout_skips.json`. If model metadata cannot be inspected, generic +layouts are retained and the runner's per-combination error isolation applies. + +Run: + +```bash +./run_parallel_layout_sweep.sh --dry-run +./run_parallel_layout_sweep.sh +./recommend_parallel_layout.py +``` + +The recommender writes: + +```text +parallel-layout-config.yml +parallel-layout-recommendation.json +``` + +This mode stops after TP/DP selection. Use `--generate-full-sweep` when +concurrency and scheduler tuning should follow automatically. +""", + encoding="utf-8", + ) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + + return [ + parallel_params, + parallel_skips, + bench_params, + run_parallel, + recommend_parallel, + guide, + *analysis_files, + ] + + +def build_concurrency_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: + """Build the fixed workload shape for concurrency tuning.""" + validate_sweep_workload(workload) + assert workload.input_tokens is not None + assert workload.output_tokens is not None + assert workload.concurrency is not None + return [ + { + "_benchmark_name": "user_workload", + "random_input_len": workload.input_tokens, + "random_output_len": workload.output_tokens, + "num_prompts": _num_prompts_for_concurrency(workload.concurrency), + } + ] + + +def _write_adaptive_concurrency_script(path: Path) -> None: + template_path = Path(__file__).with_name("adaptive_concurrency.py") + path.write_text(template_path.read_text(encoding="utf-8"), encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_concurrency_run_script( + path: Path, + *, + config_rel: str, + env_rel: str, + request_model: str, + tokenizer: str, + workload: WorkloadHints, +) -> None: + bench_parts = [ + "vllm bench serve", + "--backend vllm", + f"--model {shlex.quote(request_model)}", + f"--tokenizer {shlex.quote(tokenizer)}", + "--dataset-name random", + "--request-rate inf", + "--ignore-eos", + "--metric-percentiles 99", + ] + goodput_pairs: list[str] = [] + if workload.ttft_sla_ms is not None: + goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") + if goodput_pairs: + bench_parts.append( + "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) + ) + + bench_cmd = " ".join(bench_parts) + use_adaptive_sla_search = bool( + workload.ttft_sla_ms is not None or workload.tpot_sla_ms is not None + ) + adaptive_args = [ + f"--seed-concurrency {workload.concurrency}", + f"--input-tokens {workload.input_tokens}", + f"--output-tokens {workload.output_tokens}", + ] + if workload.ttft_sla_ms is not None: + adaptive_args.append(f"--ttft-sla-ms {workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + adaptive_args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + adaptive_args_text = " \\\n ".join(adaptive_args) + + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" +ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" + +source "${{ENV_PATH}}" + +# Released/nightly vLLM images keep recipe-relative assets such as +# examples/*.jinja under /vllm-workspace. Launch sweep server subprocesses +# from there when available so relative recipe assets resolve like they do +# in serve_with_recipe.sh. +if [[ -d /vllm-workspace/examples ]]; then + cd /vllm-workspace +fi + +# recipe tools may be newer than the installed vLLM package. Only pass the +# resilience flag when this vLLM CLI actually supports it. +CONTINUE_ON_ERROR_ARG="" +if [[ "{int(use_adaptive_sla_search)}" == "1" ]] || \ + vllm bench sweep serve_workload --help 2>&1 | grep -q -- "--continue-on-error"; then + CONTINUE_ON_ERROR_ARG="--continue-on-error" +fi + +# On older vLLM versions, a failed benchmark aborts the sweep process. Retry the +# interrupted sweep with --resume so one transient run failure does not discard +# completed work or force a manual restart. +SWEEP_RETRY_COUNT="${{VLLM_RECIPE_SWEEP_RETRY_COUNT:-0}}" +MAX_SWEEP_RETRIES="${{VLLM_RECIPE_SWEEP_RETRIES:-2}}" +SERVER_READY_TIMEOUT="${{VLLM_RECIPE_SERVER_READY_TIMEOUT:-1800}}" + +run_concurrency_sweep() {{ + if [[ "{int(use_adaptive_sla_search)}" == "1" ]]; then + python3 "${{SCRIPT_DIR}}/adaptive_concurrency.py" \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --results-dir "${{SCRIPT_DIR}}/results/concurrency-tuning" \ + {adaptive_args_text} \ + --server-ready-timeout "${{SERVER_READY_TIMEOUT}}" \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" + else + vllm bench sweep serve_workload \ + --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ + --bench-cmd "{bench_cmd}" \ + --bench-params "${{SCRIPT_DIR}}/concurrency_bench_params.json" \ + --output-dir "${{SCRIPT_DIR}}/results" \ + --experiment-name concurrency-tuning \ + --workload-var max_concurrency \ + --workload-iters 10 \ + --warmup-num-prompts {workload.concurrency} \ + ${{CONTINUE_ON_ERROR_ARG}} \ + "$@" + fi +}} + +run_concurrency_sweep || {{ + if [[ "${{SWEEP_RETRY_COUNT}}" -ge "${{MAX_SWEEP_RETRIES}}" ]]; then + echo "Sweep failed after ${{MAX_SWEEP_RETRIES}} automatic retries." >&2 + exit 1 + fi + + NEXT_RETRY=$((SWEEP_RETRY_COUNT + 1)) + RETRY_ARGS=() + for arg in "$@"; do + [[ "$arg" == "--resume" ]] && continue + RETRY_ARGS+=("$arg") + done + + EXPERIMENT_DIR="${{SCRIPT_DIR}}/results/concurrency-tuning" + if [[ -d "${{EXPERIMENT_DIR}}" ]]; then + echo "Sweep command failed;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} with --resume" + RETRY_ARGS=(--resume "${{RETRY_ARGS[@]}}") + else + echo "Sweep command failed before resumable state was created;" \ + "retry ${{NEXT_RETRY}}/${{MAX_SWEEP_RETRIES}} from the start" + fi + + VLLM_RECIPE_SWEEP_RETRY_COUNT="${{NEXT_RETRY}}" exec "$0" "${{RETRY_ARGS[@]}}" + }} +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_concurrency_recommend_script(path: Path, *, workload: WorkloadHints) -> None: + template_path = Path(__file__).with_name("concurrency_recommendation.py") + source = template_path.read_text(encoding="utf-8") + replacements = { + "DEFAULT_TTFT_SLA_MS: float | None = None": ( + f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" + ), + "DEFAULT_TPOT_SLA_MS: float | None = None": ( + f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" + ), + "DEFAULT_SEED_CONCURRENCY: int | None = None": ( + f"DEFAULT_SEED_CONCURRENCY: int | None = {workload.concurrency!r}" + ), + } + for marker, replacement in replacements.items(): + if marker not in source: + raise ValueError(f"Concurrency recommender marker not found: {marker}") + source = source.replace(marker, replacement, 1) + path.write_text(source, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_scheduler_prepare_script(path: Path) -> None: + template_path = Path(__file__).with_name("scheduler_sweep_prepare.py") + path.write_text(template_path.read_text(encoding="utf-8"), encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_scheduler_wrapper(path: Path, *, workload: WorkloadHints) -> None: + assert workload.input_tokens is not None + assert workload.output_tokens is not None + args = [ + 'python3 "${SCRIPT_DIR}/prepare_scheduler.py"', + '--config "${SCRIPT_DIR}/parallel-layout-config.yml"', + '--concurrency-recommendation "${SCRIPT_DIR}/concurrency-recommendation.json"', + f"--input-tokens {workload.input_tokens}", + f"--output-tokens {workload.output_tokens}", + ] + if workload.tpot_sla_ms is not None: + args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + if workload.target_qps is not None: + args.append(f"--target-qps {workload.target_qps:g}") + args.extend( + [ + '--output-config "${SCRIPT_DIR}/parallel-layout-sweep-config.yml"', + '--serve-params "${SCRIPT_DIR}/serve_params.json"', + '--bench-params "${SCRIPT_DIR}/bench_params.json"', + ] + ) + prepare_cmd = " \\\n ".join(args) + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" + +{prepare_cmd} + +exec "${{SCRIPT_DIR}}/run_sweep.sh" "$@" +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _recommend_sla_args(workload: WorkloadHints) -> list[str]: + args = [] + if workload.ttft_sla_ms is not None: + args.append(f"--ttft-sla-ms {workload.ttft_sla_ms:g}") + if workload.tpot_sla_ms is not None: + args.append(f"--tpot-sla-ms {workload.tpot_sla_ms:g}") + return args + + +def _write_full_run_script(path: Path, *, workload: WorkloadHints) -> None: + recommend_sla_args = " ".join(_recommend_sla_args(workload)) + script = f"""#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" +STAGES=all +RUN_ARGS=() +SERVER_READY_TIMEOUT="${{VLLM_RECIPE_SERVER_READY_TIMEOUT:-1800}}" +RECOMMEND_SLA_ARGS=({recommend_sla_args}) +while (($#)); do + case "$1" in + --stages) + if (($# < 2)); then + echo "--stages requires a comma-separated stage list." >&2 + exit 2 + fi + STAGES="$2" + shift 2 + ;; + --stages=*) STAGES="${{1#*=}}"; shift ;; + --server-ready-timeout) + if (($# < 2)); then + echo "--server-ready-timeout requires a value." >&2 + exit 2 + fi + SERVER_READY_TIMEOUT="$2" + RUN_ARGS+=("$1" "$2") + shift 2 + ;; + --server-ready-timeout=*) + SERVER_READY_TIMEOUT="${{1#*=}}" + RUN_ARGS+=("$1") + shift + ;; + --help|-h) + echo "Usage: $0 [--stages all|parallel-layout,concurrency,scheduler] [sweep args]" + echo "Selected stages always run in dependency order. Default: all." + exit 0 + ;; + *) RUN_ARGS+=("$1"); shift ;; + esac +done +case "$STAGES" in + all) STAGES=parallel-layout,concurrency,scheduler ;; + ""|,*|*,|*,,*) echo "Invalid stage list: $STAGES" >&2; exit 2 ;; +esac +IFS=, read -r -a SELECTED_STAGES <<< "$STAGES" +for stage in "${{SELECTED_STAGES[@]}}"; do + case "$stage" in + parallel-layout|concurrency|scheduler) ;; + *) echo "Unknown sweep stage: $stage" >&2; exit 2 ;; + esac +done +has_stage() {{ [[ ",$STAGES," == *",$1,"* ]]; }} + +# Keep the concurrency runner's server-start timeout explicit even when sweep +# arguments pass through multiple generated wrapper scripts. +export VLLM_RECIPE_SERVER_READY_TIMEOUT="${{SERVER_READY_TIMEOUT}}" + +# Check dependencies before starting any benchmarks. +if has_stage scheduler && ! has_stage concurrency && + [[ ! -f "${{SCRIPT_DIR}}/concurrency-recommendation.json" ]]; then + echo "Scheduler requires concurrency-recommendation.json; run concurrency first." >&2 + exit 2 +fi +if ! has_stage parallel-layout && + [[ ! -f "${{SCRIPT_DIR}}/parallel-layout-config.yml" ]]; then + cp "${{SCRIPT_DIR}}/initial-config.yml" "${{SCRIPT_DIR}}/parallel-layout-config.yml" +fi + +if has_stage parallel-layout; then + "${{SCRIPT_DIR}}/run_parallel_layout_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend_parallel_layout.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi +if has_stage concurrency; then + "${{SCRIPT_DIR}}/run_concurrency_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend_concurrency.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi +if has_stage scheduler; then + "${{SCRIPT_DIR}}/run_scheduler_sweep.sh" "${{RUN_ARGS[@]}}" + "${{SCRIPT_DIR}}/recommend.py" "${{RECOMMEND_SLA_ARGS[@]}}" +fi + +if ! "${{SCRIPT_DIR}}/report.py" "${{RECOMMEND_SLA_ARGS[@]}}"; then + echo "Warning: sweep completed, but HTML report generation failed." >&2 +fi +""" + path.write_text(script, encoding="utf-8") + path.chmod(path.stat().st_mode | 0o111) + + +def _write_full_guide(path: Path, workload: WorkloadHints) -> None: + content = f"""# End-to-End Runtime Tuning + +This package keeps one fixed workload shape and tunes in dependency order: + +```text +TP/DP -> max_concurrency -> scheduler +``` + +The supplied `--concurrency={workload.concurrency}` is the representative load +used for TP/DP selection and to size the Workload Explorer dataset. Stage 2 +selects the measured SLA-feasible `max_concurrency`. Stage 3 recalculates its +scheduler baseline from that concurrency and the selected DP layout. + +Run all stages: + +```bash +./run_full_sweep.sh +``` + +After the final recommendation, the full runner writes `sweep-report.html`. +Report-generation failures are warnings and do not change sweep success. + +Run selected stages (always executed in dependency order): + +```bash +./run_full_sweep.sh --stages concurrency +./run_full_sweep.sh --stages concurrency,scheduler +./run_full_sweep.sh --stages scheduler --resume +``` + +When parallel-layout tuning is skipped, an existing `parallel-layout-config.yml` +is reused; otherwise `initial-config.yml` supplies the fixed TP/DP layout. +Scheduler-only runs require an existing `concurrency-recommendation.json`. +Other arguments are forwarded to the selected benchmark runners. + +Or run each stage explicitly: + +```bash +./run_parallel_layout_sweep.sh +./recommend_parallel_layout.py +./run_concurrency_sweep.sh +./recommend_concurrency.py +./run_scheduler_sweep.sh +./recommend.py +``` + +Generated recipe sweeps use `--continue-on-error`. An isolated warmup or +benchmark invocation failure is recorded and skipped. A combination with no +successful measured runs remains retryable with `--resume`. +""" + path.write_text(content, encoding="utf-8") + + +def write_concurrency_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + tune_scheduler: bool = False, +) -> list[Path]: + """Write a standalone max_concurrency tuning package.""" + validate_sweep_workload(workload) + directory = Path(output_dir) + directory.mkdir(parents=True, exist_ok=True) + + bench_params = directory / "concurrency_bench_params.json" + concurrency_config = directory / "concurrency-benchmark-config.yml" + adaptive_script = directory / "adaptive_concurrency.py" + run_script = directory / "run_concurrency_sweep.sh" + recommend_script = directory / "recommend_concurrency.py" + + config_rel = _relative_to(directory, config_path) + concurrency_config_rel = _relative_to(directory, str(concurrency_config)) + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_sweep_config(concurrency_config, config) + benchmark_config_rel = config_rel if tune_scheduler else concurrency_config_rel + + _write_json(bench_params, build_concurrency_bench_params(workload)) + _write_adaptive_concurrency_script(adaptive_script) + _write_concurrency_run_script( + run_script, + config_rel=benchmark_config_rel, + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_concurrency_recommend_script(recommend_script, workload=workload) + analysis_files = _write_post_benchmark_analysis_files(directory, workload=workload) + return [ + bench_params, + concurrency_config, + adaptive_script, + run_script, + recommend_script, + *analysis_files, + ] + + +def write_full_sweep_files( + output_dir: str, + *, + config_path: str, + env_path: str, + config: dict[str, Any], + workload: WorkloadHints, + numa_node_count: int, + tune_scheduler: bool = False, +) -> list[Path]: + """Write TP/DP -> concurrency -> scheduler end-to-end tuning artifacts.""" + files = write_parallel_layout_sweep_files( + output_dir, + config_path=config_path, + env_path=env_path, + config=config, + workload=workload, + numa_node_count=numa_node_count, + tune_scheduler=tune_scheduler, + ) + + directory = Path(output_dir) + concurrency_bench = directory / "concurrency_bench_params.json" + adaptive_concurrency = directory / "adaptive_concurrency.py" + run_concurrency = directory / "run_concurrency_sweep.sh" + recommend_concurrency = directory / "recommend_concurrency.py" + prepare_scheduler = directory / "prepare_scheduler.py" + scheduler_runner = directory / "run_sweep.sh" + recommend_scheduler = directory / "recommend.py" + run_scheduler = directory / "run_scheduler_sweep.sh" + run_full = directory / "run_full_sweep.sh" + guide = directory / "SWEEP.md" + + env_rel = _relative_to(directory, env_path) + request_model, tokenizer = _benchmark_models(config) + + _write_json(concurrency_bench, build_concurrency_bench_params(workload)) + _write_adaptive_concurrency_script(adaptive_concurrency) + _write_concurrency_run_script( + run_concurrency, + config_rel="parallel-layout-config.yml", + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_concurrency_recommend_script(recommend_concurrency, workload=workload) + + # Full tuning owns the scheduler stage. The standalone parallel-layout mode + # intentionally stops after TP/DP selection. + _write_scheduler_prepare_script(prepare_scheduler) + _write_run_script( + scheduler_runner, + config_rel="parallel-layout-sweep-config.yml", + env_rel=env_rel, + request_model=request_model, + tokenizer=tokenizer, + workload=workload, + ) + _write_recommend_script( + recommend_scheduler, + config_rel="parallel-layout-config.yml", + env_rel=env_rel, + workload=workload, + ) + _write_scheduler_wrapper(run_scheduler, workload=workload) + + initial_config = directory / "initial-config.yml" + initial_config.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") + _write_full_run_script(run_full, workload=workload) + _write_full_guide(guide, workload) + + return [ + *files, + concurrency_bench, + adaptive_concurrency, + run_concurrency, + recommend_concurrency, + prepare_scheduler, + scheduler_runner, + recommend_scheduler, + run_scheduler, + run_full, + initial_config, + ] diff --git a/tools/recipes/sweep/sweep_html_report.py b/tools/recipes/sweep/sweep_html_report.py new file mode 100644 index 000000000000..6d2271360862 --- /dev/null +++ b/tools/recipes/sweep/sweep_html_report.py @@ -0,0 +1,756 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate a self-contained HTML summary from completed recipe sweeps.""" + +from __future__ import annotations + +import argparse +import json +from datetime import datetime, timezone +from html import escape +from pathlib import Path +from typing import Any + +STAGES = ( + ( + "parallel-layout", + "Parallel layout", + "parallel-layout-recommendation.json", + ), + ( + "concurrency-tuning", + "Concurrency", + "concurrency-recommendation.json", + ), + ("runtime-tuning", "Scheduler", "recommendation.json"), +) + +DEFAULT_TTFT_SLA_MS: float | None = None +DEFAULT_TPOT_SLA_MS: float | None = None + + +def _read_object(path: Path) -> dict[str, Any] | None: + if not path.is_file(): + return None + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"{path} does not contain a JSON object.") + return value + + +def _summary_runs(stage_dir: Path) -> list[dict[str, Any]]: + runs: list[dict[str, Any]] = [] + for path in sorted(stage_dir.glob("**/summary.json")): + value = json.loads(path.read_text(encoding="utf-8")) + if isinstance(value, list): + runs.extend(item for item in value if isinstance(item, dict)) + return runs + + +def _number(value: Any) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return float(value) + + +def _fmt(value: Any, digits: int = 1) -> str: + number = _number(value) + if number is None: + return "-" + return f"{number:.{digits}f}" + + +def _fmt_int(value: Any) -> str: + number = _number(value) + return "-" if number is None else f"{number:.0f}" + + +def _fmt_percent(value: Any) -> str: + number = _number(value) + return "-" if number is None else f"{number:.1f}%" + + +def _text(value: Any) -> str: + if value is None: + return "-" + return escape(str(value)) + + +def _resolve_sla( + reports: dict[str, dict[str, Any] | None], + key: str, + expected: float | None, +) -> float | None: + values: dict[str, float | None] = {} + for stage, label, _ in STAGES: + report = reports[stage] + if report is not None: + values[label] = _number((report.get("slo") or {}).get(key)) + + non_null = {value for value in values.values() if value is not None} + if expected is not None: + mismatched = { + label: value for label, value in values.items() if value != expected + } + if mismatched: + details = ", ".join( + f"{label}={value}" for label, value in mismatched.items() + ) + raise ValueError( + f"Recommendation {key} SLA does not match the command-line " + f"value {expected}: {details}" + ) + return expected + + if len(non_null) > 1: + details = ", ".join(f"{label}={value}" for label, value in values.items()) + raise ValueError(f"Recommendation {key} SLA values disagree: {details}") + return next(iter(non_null), None) + + +def _resolve_minimum_compliance( + reports: dict[str, dict[str, Any] | None], +) -> float | None: + values = { + _number(report.get("minimum_compliance_ratio")) + for report in reports.values() + if report is not None + } + non_null = {value for value in values if value is not None} + if len(non_null) > 1: + raise ValueError( + "Recommendation minimum-compliance values disagree: " + + ", ".join(str(value) for value in sorted(non_null)) + ) + return next(iter(non_null), None) + + +def _table(headers: list[str], rows: list[list[Any]]) -> str: + if not rows: + return '

No completed candidates were found.

' + head = "".join(f"{escape(header)}" for header in headers) + body = [] + for row in rows: + cells = "".join(f"{_text(cell)}" for cell in row) + body.append("" + cells + "") + return ( + '
' + + head + + "" + + "".join(body) + + "
" + ) + + +def _status_badge(value: Any) -> str: + status = str(value or "unknown").replace("_", " ") + css_class = "pass" if value == "sla_feasible" else "warn" + return f'{escape(status)}' + + +def _sla_label(candidate: dict[str, Any]) -> str: + if candidate.get("sla_eligible") is True: + return "PASS" + if candidate.get("sla_eligible") is False: + return "FAIL" + return "-" + + +def _sla_reason( + candidate: dict[str, Any], + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> str: + if candidate.get("valid") is False or candidate.get("benchmark_valid") is False: + return str(candidate.get("invalid_reason") or "benchmark invalid") + + reasons: list[str] = [] + ttft = _number(candidate.get("median_p99_ttft_ms")) + tpot = _number(candidate.get("median_p99_tpot_ms")) + compliance = _number(candidate.get("combined_compliance_ratio")) + + if ttft_sla_ms is not None and (ttft is None or ttft > ttft_sla_ms): + reasons.append( + "TTFT " + + ("missing" if ttft is None else f"{ttft:.1f} ms > {ttft_sla_ms:.1f} ms") + ) + if tpot_sla_ms is not None and (tpot is None or tpot > tpot_sla_ms): + reasons.append( + "TPOT " + + ("missing" if tpot is None else f"{tpot:.1f} ms > {tpot_sla_ms:.1f} ms") + ) + if minimum_compliance is not None and ( + compliance is None or compliance < minimum_compliance + ): + reasons.append( + "combined compliance " + + ( + "missing" + if compliance is None + else f"{compliance * 100:.2f}% < {minimum_compliance * 100:.2f}%" + ) + ) + + return "PASS" if not reasons else "; ".join(reasons) + + +def _scheduler_value(candidate: dict[str, Any], key: str) -> str: + value = candidate.get(key) + return "vLLM default" if value is None else str(value) + + +def _selected_settings( + parallel: dict[str, Any] | None, + concurrency: dict[str, Any] | None, + scheduler: dict[str, Any] | None, +) -> list[list[Any]]: + parallel_rec = (parallel or {}).get("recommended") or {} + concurrency_rec = (concurrency or {}).get("recommended") or {} + scheduler_rec = (scheduler or {}).get("recommended") or {} + max_num_seqs = ( + _scheduler_value(scheduler_rec, "max_num_seqs") + if scheduler is not None + else "-" + ) + max_num_batched_tokens = ( + _scheduler_value(scheduler_rec, "max_num_batched_tokens") + if scheduler is not None + else "-" + ) + return [ + ["tensor-parallel-size", parallel_rec.get("tensor_parallel_size")], + ["data-parallel-size", parallel_rec.get("data_parallel_size")], + ["max concurrency", concurrency_rec.get("max_concurrency")], + ["max-num-seqs", max_num_seqs], + ["max-num-batched-tokens", max_num_batched_tokens], + ] + + +def _stage_summary( + reports: dict[str, dict[str, Any] | None], +) -> list[list[Any]]: + rows = [] + for stage, label, _ in STAGES: + report = reports[stage] + if report is None: + rows.append([label, "not available", "-", "-", "-", "-"]) + continue + measured = report.get("measured") or {} + rows.append( + [ + label, + str(report.get("status", "unknown")).replace("_", " "), + _fmt(measured.get("mean_output_throughput")), + _fmt(measured.get("mean_request_goodput"), 3), + _fmt_int(measured.get("mean_p99_ttft_ms")), + _fmt(measured.get("mean_p99_tpot_ms")), + ] + ) + return rows + + +def _parallel_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + rows = [] + for candidate in (report or {}).get("candidates", []): + rows.append( + [ + candidate.get("tensor_parallel_size"), + candidate.get("data_parallel_size"), + candidate.get("max_num_seqs"), + candidate.get("max_num_batched_tokens"), + _fmt(candidate.get("mean_output_throughput")), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _fmt_percent(candidate.get("combined_compliance_percent")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _concurrency_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + candidates = (report or {}).get("candidates", []) + candidates = sorted(candidates, key=lambda item: item.get("max_concurrency", 0)) + rows = [] + for candidate in candidates: + rows.append( + [ + candidate.get("max_concurrency"), + _fmt(candidate.get("mean_output_throughput")), + _fmt(candidate.get("mean_request_throughput"), 3), + _fmt(candidate.get("mean_request_goodput"), 3), + _fmt_percent(candidate.get("combined_compliance_percent")), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _scheduler_rows( + report: dict[str, Any] | None, + *, + ttft_sla_ms: float | None, + tpot_sla_ms: float | None, + minimum_compliance: float | None, +) -> list[list[Any]]: + rows = [] + for candidate in (report or {}).get("candidates", []): + rows.append( + [ + _scheduler_value(candidate, "max_num_seqs"), + _scheduler_value(candidate, "max_num_batched_tokens"), + _fmt(candidate.get("mean_output_throughput")), + _fmt(candidate.get("mean_request_goodput"), 3), + _fmt_int(candidate.get("mean_p99_ttft_ms")), + _fmt(candidate.get("mean_p99_tpot_ms")), + _fmt_percent(candidate.get("combined_compliance_percent")), + _sla_label(candidate), + _sla_reason( + candidate, + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ] + ) + return rows + + +def _metadata(runs: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: + for stage, _, _ in STAGES: + if runs[stage]: + first = runs[stage][0] + return { + "model": first.get("model_id", "unknown"), + "input_tokens": first.get("random_input_len"), + "output_tokens": first.get("random_output_len"), + "request_rate": first.get("request_rate"), + } + return { + "model": "unknown", + "input_tokens": None, + "output_tokens": None, + "request_rate": None, + } + + +def _observations( + reports: dict[str, dict[str, Any] | None], + runs: dict[str, list[dict[str, Any]]], +) -> list[tuple[str, str]]: + observations: list[tuple[str, str]] = [] + parallel = reports["parallel-layout"] + concurrency = reports["concurrency-tuning"] + scheduler = reports["runtime-tuning"] + + parallel_candidates = (parallel or {}).get("candidates", []) + if len(parallel_candidates) == 1: + observations.append( + ( + "warning", + ( + "Only one parallel-layout candidate completed. The result " + "validates that layout but does not establish an optimum." + ), + ) + ) + + if parallel and concurrency and runs["parallel-layout"]: + representative = runs["parallel-layout"][0].get("max_concurrency") + concurrency_values = { + item.get("max_concurrency") for item in concurrency.get("candidates", []) + } + parallel_passed = parallel_candidates and parallel_candidates[0].get( + "p99_sla_eligible" + ) + if representative not in concurrency_values and parallel_passed: + observations.append( + ( + "warning", + ( + f"Concurrency {representative} passed in the " + "parallel-layout stage but was not included in the " + "concurrency grid. Treat the automated concurrency " + "choice as conservative." + ), + ) + ) + + concurrency_candidates = (concurrency or {}).get("candidates", []) + first_failure = next( + ( + item + for item in sorted( + concurrency_candidates, + key=lambda candidate: candidate.get("max_concurrency", 0), + ) + if item.get("p99_sla_eligible") is False + ), + None, + ) + if first_failure and concurrency: + slo = concurrency.get("slo") or {} + ttft = _number(first_failure.get("mean_p99_ttft_ms")) + tpot = _number(first_failure.get("mean_p99_tpot_ms")) + ttft_slo = _number(slo.get("ttft_ms")) + tpot_slo = _number(slo.get("tpot_ms")) + if ( + ttft + and ttft_slo + and ttft > ttft_slo + and tpot is not None + and (tpot_slo is None or tpot <= tpot_slo) + ): + observations.append( + ( + "info", + ( + "TTFT is the first limiting SLA; TPOT remains within " + "its objective at the first failing concurrency." + ), + ) + ) + + scheduler_candidates = (scheduler or {}).get("candidates", []) + scheduler_throughputs = [ + number + for number in ( + _number(item.get("mean_output_throughput")) for item in scheduler_candidates + ) + if number is not None + ] + if len(scheduler_throughputs) > 1 and min(scheduler_throughputs) > 0: + spread = (max(scheduler_throughputs) / min(scheduler_throughputs) - 1) * 100 + if spread <= 3: + observations.append( + ( + "info", + ( + "Scheduler output throughput is tightly grouped: " + f"{spread:.2f}% total spread across measured candidates." + ), + ) + ) + + if not observations: + observations.append( + ("info", "No automatic coverage or SLA warnings were detected.") + ) + return observations + + +def _observation_html(observations: list[tuple[str, str]]) -> str: + return "".join( + f'
{escape(message)}
' + for level, message in observations + ) + + +def generate_report( + root: Path, + output: Path, + title: str, + ttft_sla_ms: float | None = None, + tpot_sla_ms: float | None = None, +) -> None: + reports = {stage: _read_object(root / filename) for stage, _, filename in STAGES} + if not any(reports.values()): + names = ", ".join(filename for _, _, filename in STAGES) + raise ValueError(f"No recommendation files found under {root}: {names}") + + ttft_sla_ms = _resolve_sla(reports, "ttft_ms", ttft_sla_ms) + tpot_sla_ms = _resolve_sla(reports, "tpot_ms", tpot_sla_ms) + minimum_compliance = _resolve_minimum_compliance(reports) + + runs = {stage: _summary_runs(root / "results" / stage) for stage, _, _ in STAGES} + meta = _metadata(runs) + observations = _observations(reports, runs) + concurrency = reports["concurrency-tuning"] or {} + scheduler = reports["runtime-tuning"] or {} + measured = scheduler.get("measured") or concurrency.get("measured") or {} + recommended_concurrency = (concurrency.get("recommended") or {}).get( + "max_concurrency" + ) + generated = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + + styles = """ +:root { color-scheme: light; --ink: #17324d; --blue: #146c94; + --teal: #2a9d8f; --green: #2e7d32; --amber: #d97706; + --red: #c0392b; --line: #d8dee4; --soft: #f3f5f7; } +* { box-sizing: border-box; } +body { margin: 0; background: #eef2f5; color: var(--ink); + font: 15px/1.45 Arial, sans-serif; } +main { max-width: 1180px; margin: 28px auto; padding: 0 20px 40px; } +header, section { background: white; border: 1px solid var(--line); + border-radius: 12px; margin-bottom: 18px; padding: 24px; } +h1 { margin: 0 0 6px; font-size: 30px; } +h2 { margin: 0 0 14px; font-size: 21px; } +h3 { margin: 20px 0 10px; font-size: 16px; color: var(--blue); } +.sub, .muted { color: #607486; } +.cards { display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); + gap: 12px; margin-top: 20px; } +.card { background: var(--soft); border-radius: 9px; padding: 14px; } +.card strong { display: block; font-size: 22px; } +.card span { color: #607486; font-size: 12px; } +.badge { border-radius: 99px; display: inline-block; font-size: 12px; + font-weight: 700; padding: 4px 9px; text-transform: uppercase; } +.badge.pass { background: #eaf5ec; color: var(--green); } +.badge.warn { background: #fff4e5; color: var(--amber); } +.notice { border-left: 5px solid var(--blue); margin: 9px 0; + padding: 11px 14px; background: #eaf4f8; } +.notice.warning { border-color: var(--amber); background: #fff4e5; } +.table-wrap { overflow-x: auto; } +table { width: 100%; border-collapse: collapse; font-size: 13px; } +th { background: var(--ink); color: white; text-align: left; } +th, td { border: 1px solid var(--line); padding: 8px 9px; white-space: nowrap; } +tbody tr:nth-child(even) { background: var(--soft); } +footer { color: #607486; font-size: 12px; text-align: center; } +@media (max-width: 760px) { .cards { grid-template-columns: 1fr 1fr; } } +@media print { body { background: white; } main { margin: 0; max-width: none; } + header, section { border: 0; break-inside: avoid; padding: 12px 0; } + .table-wrap { overflow: visible; } } +""" + + body = f""" + + + + +{escape(title)} + + +
+
+
vLLM recipe sweep report
+

{escape(title)}

+
{_text(meta["model"])} | {_text(meta["input_tokens"])} + input / {_text(meta["output_tokens"])} output tokens | + request rate {_text(meta["request_rate"])}
+
+
{_text(recommended_concurrency)} + recommended concurrency
+
{_fmt(measured.get("mean_output_throughput"))} + output tokens/s
+
{_fmt_int(measured.get("mean_p99_ttft_ms"))} ms + mean p99 TTFT
+
{_fmt(measured.get("mean_p99_tpot_ms"))} ms + mean p99 TPOT
+
{_fmt_int(ttft_sla_ms)} ms + TTFT SLA
+
{_fmt(tpot_sla_ms)} ms + TPOT SLA
+
{ + _fmt_percent( + minimum_compliance * 100 if minimum_compliance is not None else None + ) + } + required combined compliance
+
+
+
+

Recommended configuration

+ { + _table( + ["Setting", "Value"], + _selected_settings( + reports["parallel-layout"], + reports["concurrency-tuning"], + reports["runtime-tuning"], + ), + ) + } +
+
+

Stage summary

+ { + _table( + [ + "Stage", + "Status", + "Output tok/s", + "Goodput req/s", + "p99 TTFT ms", + "p99 TPOT ms", + ], + _stage_summary(reports), + ) + } +

Automatic observations

+ {_observation_html(observations)} +
+
+

Parallel-layout candidates { + _status_badge((reports["parallel-layout"] or {}).get("status")) + }

+ { + _table( + [ + "TP", + "DP", + "Max seqs", + "Batch tokens", + "Output tok/s", + "p99 TTFT ms", + "p99 TPOT ms", + "Compliance", + "SLA", + "SLA reason", + ], + _parallel_rows( + reports["parallel-layout"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
+

Concurrency candidates { + _status_badge((reports["concurrency-tuning"] or {}).get("status")) + }

+ { + _table( + [ + "Concurrency", + "Output tok/s", + "Req/s", + "Goodput req/s", + "Compliance", + "p99 TTFT ms", + "p99 TPOT ms", + "SLA", + "SLA reason", + ], + _concurrency_rows( + reports["concurrency-tuning"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
+

Scheduler candidates { + _status_badge((reports["runtime-tuning"] or {}).get("status")) + }

+ { + _table( + [ + "Max seqs", + "Batch tokens", + "Output tok/s", + "Goodput req/s", + "p99 TTFT ms", + "p99 TPOT ms", + "Compliance", + "SLA", + "SLA reason", + ], + _scheduler_rows( + reports["runtime-tuning"], + ttft_sla_ms=ttft_sla_ms, + tpot_sla_ms=tpot_sla_ms, + minimum_compliance=minimum_compliance, + ), + ) + } +
+
Generated {escape(generated)} from completed recommendation and +summary JSON files. Warmup results are excluded by the recommenders.
+
+ +""" + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(body, encoding="utf-8") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate a self-contained HTML report from recipe sweeps." + ) + parser.add_argument( + "--sweep-dir", + default=None, + help="Sweep package directory (default: beside this script).", + ) + parser.add_argument( + "--output", + default=None, + help="Output HTML path (default: /sweep-report.html).", + ) + parser.add_argument( + "--title", + default="vLLM Recipe Sweep", + help="Report title.", + ) + parser.add_argument( + "--ttft-sla-ms", + type=float, + default=DEFAULT_TTFT_SLA_MS, + help="Expected TTFT objective; must match every recommendation file.", + ) + parser.add_argument( + "--tpot-sla-ms", + type=float, + default=DEFAULT_TPOT_SLA_MS, + help="Expected TPOT objective; must match every recommendation file.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + for name in ("ttft_sla_ms", "tpot_sla_ms"): + value = getattr(args, name) + if value is not None and value <= 0: + option = "--" + name.replace("_", "-") + raise ValueError(f"{option} must be greater than zero.") + + script_dir = Path(__file__).absolute().parent + root = Path(args.sweep_dir).resolve() if args.sweep_dir else script_dir + output = Path(args.output).resolve() if args.output else root / "sweep-report.html" + generate_report( + root, + output, + args.title, + ttft_sla_ms=args.ttft_sla_ms, + tpot_sla_ms=args.tpot_sla_ms, + ) + print(f"Wrote HTML sweep report: {output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep_recommendation.py b/tools/recipes/sweep/sweep_recommendation.py similarity index 62% rename from tools/recipes/sweep_recommendation.py rename to tools/recipes/sweep/sweep_recommendation.py index 70efaad62523..1eebf754c33b 100755 --- a/tools/recipes/sweep_recommendation.py +++ b/tools/recipes/sweep/sweep_recommendation.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Select one measured scheduler recommendation from vLLM sweep results.""" +"""Select one measured parallel-layout or scheduler sweep recommendation.""" from __future__ import annotations @@ -24,6 +24,10 @@ DEFAULT_TTFT_SLA_MS: float | None = None DEFAULT_TPOT_SLA_MS: float | None = None DEFAULT_MINIMUM_COMPLIANCE: float = 0.99 +DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT: float = 1.0 +DEFAULT_RESULTS_DIR = "results/runtime-tuning" +DEFAULT_OUTPUT_CONFIG = "recommended-config.yml" +DEFAULT_OUTPUT_JSON = "recommendation.json" def _number(value: object) -> float | None: @@ -48,11 +52,75 @@ def _percentile_summary( return mean(numbers), median(numbers), max(numbers) -def _positive_int(row: dict[str, Any], key: str) -> int | None: - value = row.get(key) +def _scheduler_value(row: dict[str, Any], key: str) -> tuple[bool, int | None]: + if key not in row: + # Default-reference sweep candidates omit scheduler keys entirely. + return True, None + value = row[key] + if value is None: + return True, None if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - return None - return value + return False, None + return True, value + + +def _parallel_value(row: dict[str, Any], key: str, default: int) -> tuple[bool, int]: + if key not in row: + return True, default + value = row[key] + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return False, default + return True, value + + +def _scheduler_sort_value(value: int | None) -> int: + return -1 if value is None else value + + +def _scheduler_preference_key(value: object) -> tuple[int, int]: + """Prefer vLLM default, then the smaller explicit value on exact ties.""" + if value is None: + return (1, 0) + if isinstance(value, int) and not isinstance(value, bool): + return (0, -value) + return (0, 0) + + +def _has_fixed_parallel_layout(candidates: list[dict[str, Any]]) -> bool: + """Return whether candidates belong to one scheduler-only comparison.""" + layouts = { + (candidate["tensor_parallel_size"], candidate["data_parallel_size"]) + for candidate in candidates + } + return len(layouts) == 1 + + +def _select_near_equivalent_scheduler_candidate( + candidates: list[dict[str, Any]], + *, + throughput_equivalence_percent: float, +) -> dict[str, Any]: + """Prefer fewer scheduler overrides among near-equal configurations.""" + best_throughput = max( + candidate["mean_output_throughput"] for candidate in candidates + ) + throughput_floor = best_throughput * (1.0 - throughput_equivalence_percent / 100.0) + near_equivalent = [ + candidate + for candidate in candidates + if candidate["mean_output_throughput"] >= throughput_floor + ] + return max( + near_equivalent, + key=lambda candidate: ( + len(candidate["vllm_default_parameters"]), + candidate["combined_compliance_ratio"] or 0.0, + candidate["mean_request_goodput"] or 0.0, + candidate["mean_output_throughput"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), + ), + ) def _resolve(script_dir: Path, value: str | None) -> Path | None: @@ -80,21 +148,40 @@ def _aggregate_candidates( rows: list[dict[str, Any]], *, use_goodput: bool, + default_tensor_parallel_size: int = 1, + default_data_parallel_size: int = 1, ttft_sla_ms: float | None = None, tpot_sla_ms: float | None = None, minimum_compliance: float = DEFAULT_MINIMUM_COMPLIANCE, ) -> list[dict[str, Any]]: - grouped: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + grouped: dict[tuple[int, int, int | None, int | None], list[dict[str, Any]]] = ( + defaultdict(list) + ) for row in rows: - seqs = _positive_int(row, "max_num_seqs") - batch = _positive_int(row, "max_num_batched_tokens") - if seqs is None or batch is None: + seqs_valid, seqs = _scheduler_value(row, "max_num_seqs") + batch_valid, batch = _scheduler_value(row, "max_num_batched_tokens") + tp_valid, tp = _parallel_value( + row, "tensor_parallel_size", default_tensor_parallel_size + ) + dp_valid, dp = _parallel_value( + row, "data_parallel_size", default_data_parallel_size + ) + if not seqs_valid or not batch_valid or not tp_valid or not dp_valid: continue - grouped[(seqs, batch)].append(row) + grouped[(tp, dp, seqs, batch)].append(row) candidates: list[dict[str, Any]] = [] - for (seqs, batch), runs in sorted(grouped.items()): + sorted_groups = sorted( + grouped.items(), + key=lambda item: ( + item[0][0], + item[0][1], + _scheduler_sort_value(item[0][2]), + _scheduler_sort_value(item[0][3]), + ), + ) + for (tp, dp, seqs, batch), runs in sorted_groups: failed_requests = sum(int(_number(run.get("failed")) or 0) for run in runs) output_throughput = _mean(runs, "output_throughput") request_throughput = _mean(runs, "request_throughput") @@ -150,9 +237,19 @@ def _aggregate_candidates( candidates.append( { + "tensor_parallel_size": tp, + "data_parallel_size": dp, "max_num_seqs": seqs, "max_num_batched_tokens": batch, "run_count": len(runs), + "vllm_default_parameters": [ + name + for name, value in ( + ("max_num_seqs", seqs), + ("max_num_batched_tokens", batch), + ) + if value is None + ], "failed_requests": failed_requests, "mean_request_goodput": request_goodput, "mean_request_throughput": request_throughput, @@ -188,6 +285,7 @@ def _select_candidate( candidates: list[dict[str, Any]], *, use_goodput: bool, + throughput_equivalence_percent: float = (DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT), ) -> tuple[dict[str, Any] | None, dict[str, Any], str]: valid = [candidate for candidate in candidates if candidate["valid"]] if not valid: @@ -204,8 +302,10 @@ def _select_candidate( candidate["mean_request_goodput"], candidate["combined_compliance_ratio"] or 0.0, candidate["mean_output_throughput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), ), ) if best_effort["mean_request_goodput"] <= 0: @@ -214,30 +314,46 @@ def _select_candidate( "the supplied TTFT/TPOT objectives." ) eligible = [candidate for candidate in valid if candidate["sla_eligible"]] - winner = ( - max( + if eligible and _has_fixed_parallel_layout(valid): + winner = _select_near_equivalent_scheduler_candidate( eligible, + throughput_equivalence_percent=throughput_equivalence_percent, + ) + else: + winner = ( + max( + eligible, + key=lambda candidate: ( + candidate["mean_output_throughput"], + candidate["combined_compliance_ratio"], + candidate["mean_request_goodput"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), + ), + ) + if eligible + else None + ) + else: + objective = "highest_mean_output_throughput" + if _has_fixed_parallel_layout(valid): + winner = _select_near_equivalent_scheduler_candidate( + valid, + throughput_equivalence_percent=throughput_equivalence_percent, + ) + else: + winner = max( + valid, key=lambda candidate: ( candidate["mean_output_throughput"], - candidate["combined_compliance_ratio"], - candidate["mean_request_goodput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], + -candidate["data_parallel_size"], + -candidate["tensor_parallel_size"], + _scheduler_preference_key(candidate["max_num_batched_tokens"]), + _scheduler_preference_key(candidate["max_num_seqs"]), ), ) - if eligible - else None - ) - else: - objective = "highest_mean_output_throughput" - winner = max( - valid, - key=lambda candidate: ( - candidate["mean_output_throughput"], - -candidate["max_num_batched_tokens"], - -candidate["max_num_seqs"], - ), - ) best_effort = winner return winner, best_effort, objective @@ -250,6 +366,31 @@ def _load_config(path: Path) -> dict[str, Any]: return data +def _build_recommended_config( + initial_config: dict[str, Any], + winner: dict[str, Any], +) -> dict[str, Any]: + recommended_config = dict(initial_config) + recommended_config["tensor-parallel-size"] = winner["tensor_parallel_size"] + recommended_config["data-parallel-size"] = winner["data_parallel_size"] + for config_key, result_key in ( + ("max-num-seqs", "max_num_seqs"), + ("max-num-batched-tokens", "max_num_batched_tokens"), + ): + value = winner[result_key] + if value is None: + recommended_config.pop(config_key, None) + else: + recommended_config[config_key] = value + return recommended_config + + +def _format_scheduler_value(value: object) -> str: + if value is None: + return "vLLM default" + return str(value) + + def _write_config( path: Path, *, @@ -274,13 +415,13 @@ def _write_config( def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( - "Select one recommended max-num-seqs/max-num-batched-tokens pair " - "from generated vLLM sweep results." + "Select one recommended TP/DP layout and scheduler setting from " + "generated vLLM sweep results." ) ) parser.add_argument( "--results-dir", - default="results/runtime-tuning", + default=DEFAULT_RESULTS_DIR, help="Sweep experiment directory (default: results/runtime-tuning).", ) parser.add_argument( @@ -314,14 +455,24 @@ def parse_args() -> argparse.Namespace: "meet all supplied objectives (default: 0.99)." ), ) + parser.add_argument( + "--throughput-equivalence-percent", + type=float, + default=DEFAULT_THROUGHPUT_EQUIVALENCE_PERCENT, + help=( + "For a scheduler-only comparison, prefer fewer explicit scheduler " + "overrides when throughput is within this percentage of the best " + "eligible result (default: 1.0)." + ), + ) parser.add_argument( "--output-config", - default="recommended-config.yml", + default=DEFAULT_OUTPUT_CONFIG, help="Recommended config output path.", ) parser.add_argument( "--output-json", - default="recommendation.json", + default=DEFAULT_OUTPUT_JSON, help="Recommendation evidence output path.", ) return parser.parse_args() @@ -331,6 +482,10 @@ def main() -> int: args = parse_args() if not 0.0 < args.minimum_compliance <= 1.0: raise ValueError("--minimum-compliance must be greater than 0 and at most 1.") + if not 0.0 <= args.throughput_equivalence_percent < 100.0: + raise ValueError( + "--throughput-equivalence-percent must be at least 0 and less than 100." + ) script_dir = Path(__file__).resolve().parent config_path = _resolve(script_dir, args.config) @@ -350,6 +505,7 @@ def main() -> int: f"Sweep results were not found at {results_dir}. Run the sweep first." ) + initial_config = _load_config(config_path) rows = _load_runs(results_dir) if not rows: raise ValueError(f"No summary.json sweep results found under {results_dir}.") @@ -358,19 +514,20 @@ def main() -> int: candidates = _aggregate_candidates( rows, use_goodput=use_goodput, + default_tensor_parallel_size=int(initial_config.get("tensor-parallel-size", 1)), + default_data_parallel_size=int(initial_config.get("data-parallel-size", 1)), ttft_sla_ms=args.ttft_sla_ms, tpot_sla_ms=args.tpot_sla_ms, minimum_compliance=args.minimum_compliance, ) winner, best_effort, objective = _select_candidate( - candidates, use_goodput=use_goodput + candidates, + use_goodput=use_goodput, + throughput_equivalence_percent=args.throughput_equivalence_percent, ) - initial_config = _load_config(config_path) if winner is not None: - recommended_config = dict(initial_config) - recommended_config["max-num-seqs"] = winner["max_num_seqs"] - recommended_config["max-num-batched-tokens"] = winner["max_num_batched_tokens"] + recommended_config = _build_recommended_config(initial_config, winner) _write_config( output_config, source_path=config_path, @@ -384,8 +541,11 @@ def main() -> int: recommended = None if winner is not None: recommended = { + "tensor_parallel_size": winner["tensor_parallel_size"], + "data_parallel_size": winner["data_parallel_size"], "max_num_seqs": winner["max_num_seqs"], "max_num_batched_tokens": winner["max_num_batched_tokens"], + "vllm_default_parameters": winner["vllm_default_parameters"], } measured = { key: winner[key] @@ -413,19 +573,25 @@ def main() -> int: ), "selection_objective": objective, "minimum_compliance_ratio": args.minimum_compliance, + "throughput_equivalence_percent": args.throughput_equivalence_percent, "slo": { "ttft_ms": args.ttft_sla_ms, "tpot_ms": args.tpot_sla_ms, }, "initial": { + "tensor_parallel_size": initial_config.get("tensor-parallel-size"), + "data_parallel_size": initial_config.get("data-parallel-size", 1), "max_num_seqs": initial_config.get("max-num-seqs"), "max_num_batched_tokens": initial_config.get("max-num-batched-tokens"), }, "recommended": recommended, "measured": measured, "best_effort": { + "tensor_parallel_size": best_effort["tensor_parallel_size"], + "data_parallel_size": best_effort["data_parallel_size"], "max_num_seqs": best_effort["max_num_seqs"], "max_num_batched_tokens": best_effort["max_num_batched_tokens"], + "vllm_default_parameters": best_effort["vllm_default_parameters"], "mean_request_goodput": best_effort["mean_request_goodput"], "combined_compliance_ratio": best_effort["combined_compliance_ratio"], "p99_sla_eligible": best_effort["p99_sla_eligible"], @@ -445,8 +611,15 @@ def main() -> int: print("Recommended runtime configuration") print() - print(f" max-num-seqs: {winner['max_num_seqs']}") - print(f" max-num-batched-tokens: {winner['max_num_batched_tokens']}") + print(f" tensor-parallel-size: {winner['tensor_parallel_size']}") + print(f" data-parallel-size: {winner['data_parallel_size']}") + print( + " max-num-seqs: " + _format_scheduler_value(winner["max_num_seqs"]) + ) + print( + " max-num-batched-tokens: " + + _format_scheduler_value(winner["max_num_batched_tokens"]) + ) print() print(f"Selection objective: {objective}") if use_goodput: diff --git a/tools/recipes/sweep/sweep_visualization.py b/tools/recipes/sweep/sweep_visualization.py new file mode 100644 index 000000000000..4681ae5f7f37 --- /dev/null +++ b/tools/recipes/sweep/sweep_visualization.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate post-benchmark figures from completed recipe sweep results.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class PlotSpec: + var_x: str + var_y: str + fig_name: str + + +STAGE_PLOTS: dict[str, tuple[PlotSpec, ...]] = { + "parallel-layout": ( + PlotSpec("recipe_candidate", "output_throughput", "parallel_throughput"), + PlotSpec("recipe_candidate", "p99_ttft_ms", "parallel_p99_ttft"), + PlotSpec("recipe_candidate", "p99_tpot_ms", "parallel_p99_tpot"), + ), + "concurrency-tuning": ( + PlotSpec("max_concurrency", "output_throughput", "concurrency_throughput"), + PlotSpec("max_concurrency", "p99_ttft_ms", "concurrency_p99_ttft"), + PlotSpec("max_concurrency", "p99_tpot_ms", "concurrency_p99_tpot"), + PlotSpec("output_throughput", "p99_ttft_ms", "throughput_p99_ttft"), + ), + "runtime-tuning": ( + PlotSpec("recipe_candidate", "output_throughput", "scheduler_throughput"), + PlotSpec("recipe_candidate", "p99_ttft_ms", "scheduler_p99_ttft"), + PlotSpec("recipe_candidate", "p99_tpot_ms", "scheduler_p99_tpot"), + ), +} + + +def _serve_candidate(summary_path: Path, record: dict[str, Any]) -> str: + directory_name = summary_path.parent.name + prefix = "SERVE--" + separator = "-BENCH--" + if directory_name.startswith(prefix) and separator in directory_name: + return directory_name[len(prefix) :].split(separator, 1)[0] + + benchmark_name = record.get("_benchmark_name") + if isinstance(benchmark_name, str) and benchmark_name: + return benchmark_name + return "fixed_configuration" + + +def _load_stage_records(stage_dir: Path) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for summary_path in sorted(stage_dir.rglob("summary.json")): + data = json.loads(summary_path.read_text(encoding="utf-8")) + if not isinstance(data, list): + continue + for record in data: + if not isinstance(record, dict): + continue + normalized = dict(record) + normalized["recipe_candidate"] = _serve_candidate( + summary_path, + normalized, + ) + records.append(normalized) + return records + + +def _has_numeric_metric(records: list[dict[str, Any]], key: str) -> bool: + return any( + isinstance(record.get(key), (int, float)) + and not isinstance(record.get(key), bool) + for record in records + ) + + +def _run_plot( + normalized_dir: Path, + figure_dir: Path, + spec: PlotSpec, + *, + dry_run: bool, +) -> None: + command = [ + "vllm", + "bench", + "sweep", + "plot", + str(normalized_dir), + "--fig-dir", + str(figure_dir.resolve()), + "--var-x", + spec.var_x, + "--var-y", + spec.var_y, + "--fig-name", + spec.fig_name, + ] + if dry_run: + command.append("--dry-run") + + environment = dict(os.environ) + environment.setdefault("MPLBACKEND", "Agg") + subprocess.run(command, check=True, env=environment) + + +def generate_figures( + results_dir: Path, + *, + selected_stages: list[str] | None = None, + dry_run: bool = False, +) -> list[Path]: + if shutil.which("vllm") is None: + raise RuntimeError( + "The vllm executable was not found. Run this analysis inside the " + "vLLM benchmark environment." + ) + + stages = selected_stages or list(STAGE_PLOTS) + discovered: list[Path] = [] + with tempfile.TemporaryDirectory(prefix="vllm-recipe-plot-") as temp_dir: + temporary_root = Path(temp_dir) + for stage in stages: + stage_dir = results_dir / stage + if not stage_dir.is_dir(): + continue + + records = _load_stage_records(stage_dir) + if not records: + continue + discovered.append(stage_dir) + + normalized_dir = temporary_root / stage + normalized_dir.mkdir(parents=True, exist_ok=True) + (normalized_dir / "summary.json").write_text( + json.dumps(records, indent=2) + "\n", + encoding="utf-8", + ) + + figure_dir = stage_dir / "figures" + plots = list(STAGE_PLOTS[stage]) + if stage == "concurrency-tuning" and _has_numeric_metric( + records, "request_goodput" + ): + plots.append( + PlotSpec( + "max_concurrency", + "request_goodput", + "concurrency_goodput", + ) + ) + + for spec in plots: + _run_plot( + normalized_dir, + figure_dir, + spec, + dry_run=dry_run, + ) + + if not discovered: + requested = ", ".join(stages) + raise ValueError( + f"No completed sweep summaries found under {results_dir} " + f"for stages: {requested}." + ) + return discovered + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Plot completed recipe sweep results without running a sweep." + ) + parser.add_argument( + "--results-dir", + default=None, + help="Results root (default: results beside this script).", + ) + parser.add_argument( + "--stage", + action="append", + choices=tuple(STAGE_PLOTS), + help="Plot only this stage; repeat to select multiple stages.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Preview vLLM plot operations without writing figures.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + script_dir = Path(__file__).resolve().parent + results_dir = ( + Path(args.results_dir).resolve() + if args.results_dir is not None + else script_dir / "results" + ) + stages = generate_figures( + results_dir, + selected_stages=args.stage, + dry_run=args.dry_run, + ) + print() + if args.dry_run: + print("Visualization dry run completed for:") + else: + print("Wrote post-benchmark figures under:") + for stage_dir in stages: + print(f" {stage_dir / 'figures'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/recipes/sweep_generation.py b/tools/recipes/sweep_generation.py deleted file mode 100644 index 33a5244a725e..000000000000 --- a/tools/recipes/sweep_generation.py +++ /dev/null @@ -1,358 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Generate optional vLLM benchmark sweep files from one initial suggestion.""" - -from __future__ import annotations - -import json -import os -import shlex -from pathlib import Path -from typing import Any - -from runtime_tuning import WorkloadHints - - -def _positive_int(config: dict[str, Any], key: str) -> int: - value = config.get(key) - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - raise ValueError(f"Sweep generation requires a positive initial {key!r} value.") - return value - - -def _strict_lower_power_of_two(value: int) -> int: - if value <= 1: - return 1 - return 1 << ((value - 1).bit_length() - 1) - - -def _strict_upper_power_of_two(value: int) -> int: - return 1 << value.bit_length() - - -def validate_sweep_workload(workload: WorkloadHints) -> None: - required = { - "--input-tokens": workload.input_tokens, - "--output-tokens": workload.output_tokens, - "--concurrency": workload.concurrency, - } - missing = [name for name, value in required.items() if value is None] - if missing: - raise ValueError("--generate-sweep requires " + ", ".join(missing) + ".") - - -def build_serve_params(config: dict[str, Any]) -> list[dict[str, Any]]: - """Build a diverse, bounded sweep around the initial suggestion. - - The initial five-point sweep changed only one parameter at a time and - skipped the useful middle sequence count. Eight directed points cover the - batch-budget curve at full concurrency plus interactions at 3/4 and 1/2 of - the initial scheduler concurrency without paying for a full Cartesian grid. - """ - initial_seqs = _positive_int(config, "max-num-seqs") - initial_batch = _positive_int(config, "max-num-batched-tokens") - - minimum_batch = initial_seqs - if config.get("enable-chunked-prefill") is False: - max_model_len = config.get("max-model-len") - if ( - isinstance(max_model_len, int) - and not isinstance(max_model_len, bool) - and max_model_len > 0 - ): - minimum_batch = max(minimum_batch, max_model_len) - - lower_seqs = max(1, (initial_seqs + 1) // 2) - middle_seqs = max(lower_seqs, (3 * initial_seqs + 3) // 4) - lower_batch = max( - minimum_batch, - _strict_lower_power_of_two(initial_batch), - ) - smaller_batch = max( - minimum_batch, - _strict_lower_power_of_two(lower_batch), - ) - higher_batch = max( - minimum_batch, - _strict_upper_power_of_two(initial_batch), - ) - - candidates: list[dict[str, Any]] = [] - seen: set[tuple[int, int]] = set() - - def add(name: str, max_num_seqs: int, max_num_batched_tokens: int) -> None: - max_num_batched_tokens = max( - max_num_batched_tokens, - max_num_seqs, - ) - signature = (max_num_seqs, max_num_batched_tokens) - if signature in seen: - return - seen.add(signature) - candidates.append( - { - "_benchmark_name": name, - "max_num_seqs": max_num_seqs, - "max_num_batched_tokens": max_num_batched_tokens, - } - ) - - # Keep the exact initial suggestion as the measured baseline. Additional - # values exist only in the optional sweep package. - add("initial", initial_seqs, initial_batch) - add("smaller_batch_budget", initial_seqs, smaller_batch) - add("lower_batch_budget", initial_seqs, lower_batch) - add("higher_batch_budget", initial_seqs, higher_batch) - add("middle_seqs_lower_batch", middle_seqs, lower_batch) - add("middle_seqs_higher_batch", middle_seqs, higher_batch) - add("lower_seqs_lower_batch", lower_seqs, lower_batch) - add("lower_seqs_higher_batch", lower_seqs, higher_batch) - - return candidates - - -def build_bench_params(workload: WorkloadHints) -> list[dict[str, Any]]: - validate_sweep_workload(workload) - assert workload.input_tokens is not None - assert workload.output_tokens is not None - assert workload.concurrency is not None - - return [ - { - "_benchmark_name": "user_workload", - "random_input_len": workload.input_tokens, - "random_output_len": workload.output_tokens, - "max_concurrency": workload.concurrency, - } - ] - - -def _benchmark_models(config: dict[str, Any]) -> tuple[str, str]: - model = config.get("model") - if not isinstance(model, str) or not model: - raise ValueError("Sweep generation requires a model in config.yml.") - - served_model_name = config.get("served-model-name") - if isinstance(served_model_name, str) and served_model_name: - return served_model_name, model - if ( - isinstance(served_model_name, list) - and served_model_name - and isinstance(served_model_name[0], str) - ): - return served_model_name[0], model - - return model, model - - -def _relative_to(directory: Path, target: str) -> str: - return os.path.relpath(Path(target).resolve(), directory.resolve()) - - -def _write_json(path: Path, value: object) -> None: - path.write_text( - json.dumps(value, indent=2) + "\n", - encoding="utf-8", - ) - - -def _write_run_script( - path: Path, - *, - config_rel: str, - env_rel: str, - request_model: str, - tokenizer: str, - workload: WorkloadHints, -) -> None: - bench_parts = [ - "vllm bench serve", - "--backend vllm", - f"--model {shlex.quote(request_model)}", - f"--tokenizer {shlex.quote(tokenizer)}", - "--dataset-name random", - "--request-rate inf", - "--ignore-eos", - "--metric-percentiles 99", - ] - - goodput_pairs: list[str] = [] - if workload.ttft_sla_ms is not None: - goodput_pairs.append(f"ttft:{workload.ttft_sla_ms:g}") - if workload.tpot_sla_ms is not None: - goodput_pairs.append(f"tpot:{workload.tpot_sla_ms:g}") - if goodput_pairs: - bench_parts.append( - "--goodput " + " ".join(shlex.quote(pair) for pair in goodput_pairs) - ) - - bench_cmd = " ".join(bench_parts) - script = f"""#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${{BASH_SOURCE[0]}}")" && pwd)" -CONFIG_PATH="${{SCRIPT_DIR}}/{config_rel}" -ENV_PATH="${{SCRIPT_DIR}}/{env_rel}" - -source "${{ENV_PATH}}" - -vllm bench sweep serve \ - --serve-cmd "vllm serve --config '${{CONFIG_PATH}}'" \ - --bench-cmd "{bench_cmd}" \ - --serve-params "${{SCRIPT_DIR}}/serve_params.json" \ - --bench-params "${{SCRIPT_DIR}}/bench_params.json" \ - --output-dir "${{SCRIPT_DIR}}/results" \ - --experiment-name runtime-tuning \ - "$@" -""" - path.write_text(script, encoding="utf-8") - path.chmod(path.stat().st_mode | 0o111) - - -def _write_recommend_script( - path: Path, - *, - config_rel: str, - env_rel: str, - workload: WorkloadHints, -) -> None: - template_path = Path(__file__).with_name("sweep_recommendation.py") - source = template_path.read_text(encoding="utf-8") - - replacements = { - "DEFAULT_CONFIG_PATH: str | None = None": ( - f"DEFAULT_CONFIG_PATH: str | None = {config_rel!r}" - ), - "DEFAULT_ENV_PATH: str | None = None": ( - f"DEFAULT_ENV_PATH: str | None = {env_rel!r}" - ), - "DEFAULT_TTFT_SLA_MS: float | None = None": ( - f"DEFAULT_TTFT_SLA_MS: float | None = {workload.ttft_sla_ms!r}" - ), - "DEFAULT_TPOT_SLA_MS: float | None = None": ( - f"DEFAULT_TPOT_SLA_MS: float | None = {workload.tpot_sla_ms!r}" - ), - } - for marker, replacement in replacements.items(): - if marker not in source: - raise ValueError(f"Recommender template marker not found: {marker}") - source = source.replace(marker, replacement, 1) - - path.write_text(source, encoding="utf-8") - path.chmod(path.stat().st_mode | 0o111) - - -def _write_guide(path: Path, workload: WorkloadHints) -> None: - sla_lines: list[str] = [] - if workload.ttft_sla_ms is not None: - sla_lines.append(f"- TTFT objective: `{workload.ttft_sla_ms:g} ms`") - if workload.tpot_sla_ms is not None: - sla_lines.append(f"- TPOT objective: `{workload.tpot_sla_ms:g} ms`") - - sla_text = "\n".join(sla_lines) - if sla_text: - sla_text = ( - "\n## Supplied latency objectives\n\n" - + sla_text - + "\n\nThe generated benchmark uses these values with vLLM " - "`--goodput`.\n" - ) - - content = f"""# Optional Runtime Tuning Sweep - -The generated `config.yml` is the **single initial suggestion** and can be -deployed directly. The sweep benchmarks nearby values for: - -- `max-num-seqs` -- `max-num-batched-tokens` - -## Run - -```bash -./run_sweep.sh --dry-run -./run_sweep.sh -./recommend.py -``` - -For a quick one-run experiment: - -```bash -rm -rf results/runtime-tuning -./run_sweep.sh --num-runs 1 -./recommend.py -``` - -Resume an interrupted sweep: - -```bash -./run_sweep.sh --resume -``` - -By default, vLLM benchmarks each parameter combination three times. -{sla_text} -## Outputs - -`recommend.py` writes: - -```text -recommended-config.yml -recommendation.json -``` - -With TTFT/TPOT objectives it requires duration-weighted combined compliance of -at least 99% plus median P99 TTFT/TPOT compliance, then selects highest mean -output-token throughput. If no candidate qualifies, it records a best-effort -candidate but does not write `recommended-config.yml`. Without latency -objectives it selects highest mean output-token throughput. Configurations with -failed requests are excluded. - -`recommended-config.yml` copies the initial configuration and changes only -`max-num-seqs` and `max-num-batched-tokens`. -""" - path.write_text(content, encoding="utf-8") - - -def write_sweep_files( - output_dir: str, - *, - config_path: str, - env_path: str, - config: dict[str, Any], - workload: WorkloadHints, -) -> list[Path]: - """Write an optional sweep package around the initial config suggestion.""" - validate_sweep_workload(workload) - - directory = Path(output_dir) - directory.mkdir(parents=True, exist_ok=True) - - serve_params = directory / "serve_params.json" - bench_params = directory / "bench_params.json" - run_script = directory / "run_sweep.sh" - recommend_script = directory / "recommend.py" - guide = directory / "SWEEP.md" - - config_rel = _relative_to(directory, config_path) - env_rel = _relative_to(directory, env_path) - - _write_json(serve_params, build_serve_params(config)) - _write_json(bench_params, build_bench_params(workload)) - request_model, tokenizer = _benchmark_models(config) - _write_run_script( - run_script, - config_rel=config_rel, - env_rel=env_rel, - request_model=request_model, - tokenizer=tokenizer, - workload=workload, - ) - _write_recommend_script( - recommend_script, - config_rel=config_rel, - env_rel=env_rel, - workload=workload, - ) - _write_guide(guide, workload) - - return [serve_params, bench_params, run_script, recommend_script, guide] From 0cc046087024bdb403b2032617ea0f713209a185 Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Fri, 9 Oct 2026 20:01:54 +0800 Subject: [PATCH 289/952] [XPU]Fix the accuracy issue when meet topk_ids=-1 on DP+EP scenarios (#57787) Signed-off-by: yisheng Co-authored-by: Kunshang Ji --- vllm/model_executor/layers/fused_moe/experts/xpu_moe.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index ddd891b1c090..35fb91607830 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -175,6 +175,11 @@ def apply( gemm1_clamp_limit=self.gemm1_clamp_limit, ) assert self.fused_moe_impl is not None + if self.moe_config.ep_size > 1: + # Avoid topk_ids = -1, remap them to an + # expert not owned by this rank so they are skipped. + sentinel = global_num_experts - 1 if self.moe_config.ep_rank == 0 else 0 + topk_ids = torch.where(topk_ids < 0, sentinel, topk_ids) self.fused_moe_impl.apply( output=output, hidden_states=hidden_states, From c17fadcfdc80c9f8cfa97a7eec02a66d823989c4 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Fri, 9 Oct 2026 20:16:35 +0800 Subject: [PATCH 290/952] [XPU] fix fused_input_norm for MM (#59865) Signed-off-by: Yan Ma --- .../intel_jobs/model_executor_intel.yaml | 11 +--------- .../intel_jobs/models_multimodal_intel.yaml | 3 ++- .../layers/fusion/mm_input_norm.py | 20 ++++++++----------- 3 files changed, 11 insertions(+), 23 deletions(-) diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index 5da075e235b6..d641c93e5a9e 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -28,16 +28,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype2-out_dtype2]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype3-out_dtype3]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype4-out_dtype4]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes" - --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"' + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' - label: Tensorizer Model Loader Tests (XPU) key: tensorizer-model-loader-tests-xpu diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 061be1b136ed..89654ec3aedf 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -73,7 +73,8 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && - pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model' + pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model + pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' - label: "Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper diff --git a/vllm/model_executor/layers/fusion/mm_input_norm.py b/vllm/model_executor/layers/fusion/mm_input_norm.py index 3dc6974c84c6..5de95b401df5 100644 --- a/vllm/model_executor/layers/fusion/mm_input_norm.py +++ b/vllm/model_executor/layers/fusion/mm_input_norm.py @@ -313,19 +313,15 @@ def forward_cuda( def forward_xpu( self, pixel_values: torch.Tensor, visual_dtype: torch.dtype ) -> torch.Tensor: - """XPU fused custom kernel path. - - On XPU, fuse the whole rescale + normalise into a single custom - kernel. The eager path materializes an fp32 intermediate and then - casts back, which adds device-side compute that cancels the - bandwidth saving of transferring uint8 pixel_values. The fused - kernel reads uint8 directly and writes ``visual_dtype`` in one pass. + """XPU fused custom kernel path for uint8 and Triton kernel path + for others. """ - # The out-of-tree XPU kernel only supports the uint8 input that - # device-side normalisation guarantees; fail loudly otherwise. - assert pixel_values.dtype == torch.uint8, ( - f"xpu_fused_input_norm requires uint8 input, got {pixel_values.dtype}" - ) + if pixel_values.dtype != torch.uint8: + return self.forward_cuda(pixel_values, visual_dtype) + + # For uint8 input, falls to xpu_fused_input_norm from xpu kernels + import vllm._xpu_ops # noqa: F401 + return torch.ops.vllm.xpu_fused_input_norm( pixel_values, self.weight, self.bias, visual_dtype ) From 1ee69dde9cdfc949cfae896016b44357ca864a50 Mon Sep 17 00:00:00 2001 From: Suraj Sharan Date: Fri, 9 Oct 2026 16:39:58 +0400 Subject: [PATCH 291/952] [Bugfix][LoRA] Reject PEFT adapter features vLLM does not implement (#60225) Signed-off-by: suraj Co-authored-by: linitra24 --- .../serve/lora/test_lora_adapters.py | 1 + tests/lora/test_peft_helper.py | 39 ++++++++++++++++++ tests/lora/test_utils.py | 2 + vllm/lora/peft_helper.py | 40 +++++++++++++++++++ vllm/lora/utils.py | 2 +- 5 files changed, 83 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index 8a860dc6251a..484fb9c0f1e6 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -23,6 +23,7 @@ "is greater than max_lora_rank", ), ("test_dora", {"use_dora": True}, "does not yet support DoRA"), + ("test_pissa", {"init_lora_weights": "pissa"}, "init_lora_weights='pissa'"), ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index 9fb52ea7275c..bb1246144b0d 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -24,6 +24,25 @@ ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), + ("test_lora_bias", {"lora_bias": True}, "does not support LoRA bias"), + ("test_pissa", {"init_lora_weights": "pissa"}, "init_lora_weights='pissa'"), + ( + "test_pissa_niter", + {"init_lora_weights": "pissa_niter_4"}, + "init_lora_weights='pissa_niter_4'", + ), + ("test_olora", {"init_lora_weights": "olora"}, "init_lora_weights='olora'"), + ("test_corda", {"init_lora_weights": "corda"}, "init_lora_weights='corda'"), + ("test_loftq", {"init_lora_weights": "loftq"}, "init_lora_weights='loftq'"), + ( + "test_unknown_init", + {"init_lora_weights": "future_init"}, + "init_lora_weights='future_init'", + ), + ("test_alora", {"alora_invocation_tokens": [1, 2]}, "Activated LoRA"), + ("test_layer_replication", {"layer_replication": [[0, 2]]}, "layer_replication"), + ("test_bdlora", {"use_bdlora": {"nblocks": 2}}, "BD-LoRA"), + ("test_qalora", {"use_qalora": True}, "QALoRA"), ] @@ -114,3 +133,23 @@ def test_peft_helper_invalid_rank_direct(bad_rank: int): """ with pytest.raises(ValueError, match="must be a positive integer"): PEFTHelper(r=bad_rank, lora_alpha=16, target_modules=["q_proj"]) + + +@pytest.mark.parametrize( + "init_lora_weights", + [True, False, "gaussian", "eva", "orthogonal", "mica", "lora_ga"], +) +def test_peft_helper_init_lora_weights_supported(init_lora_weights): + """A saved adapter with these inits loads in PEFT as a plain LoRA.""" + lora_config = LoRAConfig(max_lora_rank=16, max_cpu_loras=3, max_loras=2) + PEFTHelper.from_dict( + { + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + "init_lora_weights": init_lora_weights, + "alora_invocation_tokens": None, + "layer_replication": None, + "use_bdlora": None, + } + ).validate_legal(lora_config) diff --git a/tests/lora/test_utils.py b/tests/lora/test_utils.py index 4a9ba66069d7..0b1be49f7e15 100644 --- a/tests/lora/test_utils.py +++ b/tests/lora/test_utils.py @@ -106,6 +106,8 @@ def test_parse_fine_tuned_lora_name_invalid(): fixture = { "base_model.weight", "base_model.model.weight", + "base_model.model.model.layers.9.mlp.down_proj.lora_A.bias", + "base_model.model.model.layers.9.mlp.down_proj.lora_B.bias", } for name in fixture: with pytest.raises(ValueError, match="unsupported LoRA weight"): diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index d84e56bece21..de05dea1cfc1 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -33,6 +33,14 @@ class PEFTHelper: use_rslora: bool = field(default=False) # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353) use_dora: bool = field(default=False) + # PEFT features that change the adapter's semantics and that vLLM does not + # implement. They are read only to reject such adapters. + lora_bias: bool = field(default=False) + init_lora_weights: bool | str = field(default=True) + alora_invocation_tokens: list[int] | None = field(default=None) + layer_replication: list[list[int]] | None = field(default=None) + use_bdlora: dict | None = field(default=None) + use_qalora: bool = field(default=False) # Extra vllm field, start with 'vllm_' to avoid conflict vllm_lora_scaling_factor: float = field(default=1.0) vllm_max_position_embeddings: int | None = field(default=False) @@ -52,8 +60,40 @@ def _validate_features(self) -> list[str]: ) if self.use_dora: error_msg.append("vLLM does not yet support DoRA.") + if self.lora_bias: + error_msg.append("vLLM does not support LoRA bias (lora_bias).") + if not self._init_is_supported(): + error_msg.append( + f"vLLM does not support init_lora_weights={self.init_lora_weights!r}." + " Inits such as PiSSA, OLoRA, CorDA and LoftQ modify the base model" + " weights when PEFT loads the adapter. Convert it to a regular LoRA" + " adapter with PEFT's path_initial_model_for_weight_conversion, or" + " set init_lora_weights to true if the served base model already" + " contains the modified weights." + ) + if self.alora_invocation_tokens: + error_msg.append("vLLM does not support Activated LoRA (aLoRA).") + if self.layer_replication: + error_msg.append("vLLM does not support layer_replication.") + if self.use_bdlora: + error_msg.append("vLLM does not support block-diagonal LoRA (BD-LoRA).") + if self.use_qalora: + error_msg.append("vLLM does not support QALoRA.") return error_msg + def _init_is_supported(self) -> bool: + # These inits only set the initial adapter weights, so a saved adapter + # loads in PEFT as a plain LoRA. lora_ga modifies the base weights only + # during training, when its gradients are attached. + init = self.init_lora_weights + return not isinstance(init, str) or init.lower() in ( + "gaussian", + "eva", + "orthogonal", + "mica", + "lora_ga", + ) + def __post_init__(self): if self.r <= 0: raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 3a4009d28b23..1766c6de626d 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -205,7 +205,7 @@ def parse_fine_tuned_lora_name( parts = name.split(".") if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2: - if parts[-2] in ["lora_A", "lora_B"]: + if parts[-2] in ["lora_A", "lora_B"] and parts[-1] == "weight": new_name = ".".join(parts[start_index:-2]) return new_name, parts[-2] == "lora_A" # For modules_to_save in classification. From 9adba10f3f4bb3bb756f9de53fbf6820668157db Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Fri, 9 Oct 2026 20:48:32 +0800 Subject: [PATCH 292/952] [Bugfix] Fix pre-commit check (#60840) Signed-off-by: jiang1.li --- tools/recipes/sweep/sweep_html_report.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/recipes/sweep/sweep_html_report.py b/tools/recipes/sweep/sweep_html_report.py index 6d2271360862..bac10f9b4a28 100644 --- a/tools/recipes/sweep/sweep_html_report.py +++ b/tools/recipes/sweep/sweep_html_report.py @@ -8,7 +8,7 @@ import argparse import json -from datetime import datetime, timezone +from datetime import UTC, datetime from html import escape from pathlib import Path from typing import Any @@ -501,7 +501,7 @@ def generate_report( recommended_concurrency = (concurrency.get("recommended") or {}).get( "max_concurrency" ) - generated = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + generated = datetime.now(UTC).strftime("%Y-%m-%d %H:%M UTC") styles = """ :root { color-scheme: light; --ink: #17324d; --blue: #146c94; From 81c993658e0bbdbf9d30c92cfc52baa73853b242 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Fri, 9 Oct 2026 16:14:49 +0300 Subject: [PATCH 293/952] [ROCm][Perf] Fuse main QK-norm/RoPE/gate and KV-cache write into the AMD QSA prepare launch (#59732) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- .buildkite/test_areas/models_basic.yaml | 2 +- tests/models/qwen4_exp/test_qsa_amd.py | 72 ++++++-- ...indexer_amd.py => test_qsa_prepare_amd.py} | 85 ++++++++- vllm/config/compilation.py | 2 + vllm/models/qwen4_exp/amd/indexer_qsa.py | 82 ++++++--- .../{qsa_pre_indexer.py => qsa_prepare.py} | 167 ++++++++++++++++-- vllm/models/qwen4_exp/amd/qsa.py | 114 ++++++++++-- 7 files changed, 455 insertions(+), 69 deletions(-) rename tests/models/qwen4_exp/{test_qsa_pre_indexer_amd.py => test_qsa_prepare_amd.py} (80%) rename vllm/models/qwen4_exp/amd/ops/{qsa_pre_indexer.py => qsa_prepare.py} (75%) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index a25c9a78a86c..fd21324a7ece 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -216,7 +216,7 @@ steps: depends_on: - image-build-amd commands: - - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_pre_indexer_amd.py + - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py index 15c490223f2d..c8ebf0190432 100644 --- a/tests/models/qwen4_exp/test_qsa_amd.py +++ b/tests/models/qwen4_exp/test_qsa_amd.py @@ -21,8 +21,8 @@ apply_qsa_rope, ) from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops -from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( - qsa_pre_indexer, +from vllm.models.qwen4_exp.amd.ops.qsa_prepare import ( + qsa_prepare, supports_fused_pre_indexer, ) from vllm.platforms import current_platform @@ -127,6 +127,18 @@ def test_qsa_rmsnorm_uses_portable_implementation(default_vllm_config) -> None: torch.testing.assert_close(output, norm.forward_native(tensor)) +def _make_attn(*, use_fused: bool) -> SimpleNamespace: + """The main-attention fields the fused QSA prepare reads from its owner.""" + return SimpleNamespace( + use_fused_qsa_prepare=use_fused, + kv_cache=torch.zeros(1, 2, 16, 8), + q_norm=GemmaRMSNorm(4, eps=1e-6), + k_norm=GemmaRMSNorm(4, eps=1e-6), + _k_scale_float=1.0, + _v_scale_float=1.0, + ) + + def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: """A QSAIndexer whose unfused path runs for real, with no projection weights. @@ -184,29 +196,45 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: @pytest.mark.usefixtures("default_vllm_config") -def test_qsa_fused_pre_indexer_replaces_the_unfused_chain( +def test_qsa_fused_prepare_replaces_the_unfused_chain( monkeypatch: pytest.MonkeyPatch, ) -> None: indexer = _make_indexer(use_fused=True) + attn = _make_attn(use_fused=True) calls = [] - monkeypatch.setattr( - indexer_qsa, - "qsa_pre_indexer", - lambda *args, **kwargs: calls.append((args, kwargs)), - ) + main_outputs = (torch.empty(0), torch.empty(0)) + + def fake_prepare(*args, **kwargs): + calls.append((args, kwargs)) + return main_outputs + + monkeypatch.setattr(indexer_qsa, "qsa_prepare", fake_prepare) monkeypatch.setattr( QSAIndexer, "_update_and_compress", lambda *args, **kwargs: pytest.fail("fused path ran the unfused chain"), ) rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) - - actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + qkv = torch.zeros(4, 16) + slot_mapping = torch.arange(4) + gate_out = torch.empty(4, 1, 4) + + actual, actual_main = QSAIndexer.forward( + indexer, + torch.zeros(2, 8), + torch.tensor([7, 8]), + rows, + attn=attn, + qkv=qkv, + slot_mapping=slot_mapping, + gate_out=gate_out, + ) assert actual is rows + assert actual_main is main_outputs assert len(calls) == 1 args, kwargs = calls[0] - passed = inspect.signature(qsa_pre_indexer).bind(*args, **kwargs).arguments + passed = inspect.signature(qsa_prepare).bind(*args, **kwargs).arguments _, compressed_metadata = indexer.test_metadata # The kernel addresses the raw cache at its full width and writes the MRoPE # position tail itself, so it needs kv_cache, not the key_cache view. @@ -215,6 +243,17 @@ def test_qsa_fused_pre_indexer_replaces_the_unfused_chain( assert passed["k_work_metadata"] is compressed_metadata.k_work_metadata assert passed["compress_ratio"] == indexer.compress_ratio assert passed["rope_pos_offset"] == indexer.index_head_dim + # Main attention: only the real tokens, the vLLM cache viewed per token, + # and the gate landing in the caller's buffer. + assert passed["main_qkv"].shape[0] == 2 + assert passed["main_qkv"].data_ptr() == qkv.data_ptr() + assert passed["main_slot_mapping"].shape == (2,) + assert passed["main_kv_cache"].shape == (1, 16, 2, 8) + assert passed["main_kv_cache"].data_ptr() == attn.kv_cache.data_ptr() + assert passed["main_gate_out"].shape[0] == 2 + assert passed["main_gate_out"].data_ptr() == gate_out.data_ptr() + assert passed["main_q_norm_weight"] is attn.q_norm.weight + assert passed["main_k_norm_weight"] is attn.k_norm.weight @pytest.mark.usefixtures("default_vllm_config") @@ -224,7 +263,7 @@ def test_qsa_unsupported_config_keeps_the_unfused_chain( indexer = _make_indexer(use_fused=False) monkeypatch.setattr( indexer_qsa, - "qsa_pre_indexer", + "qsa_prepare", lambda *args, **kwargs: pytest.fail("unfused path ran the fused kernel"), ) updates = [] @@ -235,9 +274,16 @@ def test_qsa_unsupported_config_keeps_the_unfused_chain( ) rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) - actual = QSAIndexer.forward(indexer, torch.zeros(2, 8), torch.tensor([7, 8]), rows) + actual, main_outputs = QSAIndexer.forward( + indexer, + torch.zeros(2, 8), + torch.tensor([7, 8]), + rows, + attn=_make_attn(use_fused=False), + ) assert actual is rows + assert main_outputs is None assert len(updates) == 1 diff --git a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py b/tests/models/qwen4_exp/test_qsa_prepare_amd.py similarity index 80% rename from tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py rename to tests/models/qwen4_exp/test_qsa_prepare_amd.py index 7698198923ce..718929720955 100644 --- a/tests/models/qwen4_exp/test_qsa_pre_indexer_amd.py +++ b/tests/models/qwen4_exp/test_qsa_prepare_amd.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Correctness tests for the AMD fused QSA pre-indexer.""" +"""Correctness tests for the AMD fused QSA prepare kernel.""" import pytest import torch @@ -11,9 +11,7 @@ qsa_compress_groups_with_ratio, qsa_store_cache_rows, ) -from vllm.models.qwen4_exp.amd.ops.qsa_pre_indexer import ( - qsa_pre_indexer, -) +from vllm.models.qwen4_exp.amd.ops.qsa_prepare import qsa_prepare from vllm.models.qwen4_exp.common.qsa_cache import ( canonical_qsa_rope_positions, circular_qsa_slot_mapping, @@ -52,6 +50,7 @@ def _reference_gemma_rmsnorm( RTOL = 1.6e-2 ATOL = 1e-2 MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8]) +MAIN_HQ, MAIN_HK, MAIN_D, MAIN_PAGE = 6, 2, 256, 16 def _make_block_table(block_counts): @@ -78,6 +77,76 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all()) +def _make_main_inputs(num_tokens: int) -> dict: + """Random main-attention arguments with a paged BF16 cache in vLLM's layout.""" + num_blocks = num_tokens // MAIN_PAGE + 2 + cache = torch.zeros( + num_blocks, + MAIN_HK, + MAIN_PAGE, + 2 * MAIN_D, + dtype=torch.bfloat16, + device="cuda", + ).transpose(1, 2) + slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens] + slots[-1] = -1 + norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2 + return dict( + main_qkv=torch.randn( + num_tokens, + 2 * (MAIN_HQ + MAIN_HK) * MAIN_D, + dtype=torch.bfloat16, + device="cuda", + ), + main_q_norm_weight=norm_weights[0], + main_k_norm_weight=norm_weights[1], + main_eps=EPS, + main_kv_cache=cache, + main_slot_mapping=slots, + main_k_scale=1.0, + main_v_scale=1.0, + # The AMD op hands the gate back through a buffer it does not own. + main_gate_out=torch.empty( + num_tokens, MAIN_HQ, MAIN_D, dtype=torch.bfloat16, device="cuda" + ), + ) + + +def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: + """Compare with the AMD layer's unfused split, QK-norm, RoPE and cache write.""" + from vllm._custom_ops import reshape_and_cache_flash + + num_tokens = main["main_qkv"].shape[0] + q_gate, k, v = main["main_qkv"].split( + [2 * MAIN_HQ * MAIN_D, MAIN_HK * MAIN_D, MAIN_HK * MAIN_D], dim=-1 + ) + q, gate = torch.chunk(q_gate.view(num_tokens, MAIN_HQ, -1), 2, dim=-1) + q = _reference_gemma_rmsnorm( + q.reshape(-1, MAIN_D), main["main_q_norm_weight"], EPS + ).view(num_tokens, -1) + k = _reference_gemma_rmsnorm( + k.reshape(-1, MAIN_D), main["main_k_norm_weight"], EPS + ).view(num_tokens, -1) + q, k = rope(positions, q, k) + kv_cache = main["main_kv_cache"] + cache = torch.zeros_like(kv_cache) + key_cache, value_cache = cache.split(MAIN_D, dim=-1) + scale = torch.tensor(1.0, device="cuda") + reshape_and_cache_flash( + k.view(-1, MAIN_HK, MAIN_D), + v.view(-1, MAIN_HK, MAIN_D), + key_cache, + value_cache, + main["main_slot_mapping"], + "auto", + scale, + scale, + ) + torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL) + assert torch.equal(gate_out, gate) + torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL) + + @requires_qsa_kernels @pytest.mark.usefixtures("default_vllm_config") @pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn]) @@ -103,7 +172,7 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"), ], ) -def test_qsa_fused_pre_indexer_matches_unfused( +def test_qsa_fused_prepare_matches_unfused( indexer_dtype, mrope, is_2d_positions, @@ -269,7 +338,8 @@ def test_qsa_fused_pre_indexer_matches_unfused( k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) - qsa_pre_indexer( + main = _make_main_inputs(num_tokens) + main_q_out, main_gate_out = qsa_prepare( projected_qk[:, : HQ * D], projected_qk[:, HQ * D :], positions, @@ -289,7 +359,10 @@ def test_qsa_fused_pre_indexer_matches_unfused( compress_ratio=CR, mrope_section=MROPE_SECTION if mrope else None, rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None, + **main, ) + assert main_gate_out is main["main_gate_out"] + _check_main_outputs(main, main_q_out, main_gate_out, rope, positions) unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D) unfused_query = _reference_gemma_rmsnorm( diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index fc371d94b493..7a16e3d6a046 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -774,6 +774,8 @@ class CompilationConfig: # Qwen4Exp's AMD backend still uses these splitting ops. "vllm::qwen4_exp_ple_short_conv", "vllm::qwen4_exp_qsa_with_output", + # The fused prepare calls this op instead of the one above. + "vllm::qwen4_exp_qsa_prepare_with_output", "vllm::linear_attention", "vllm::qwen_gdn_attention_core", "vllm::qwen_gdn_attention_core_fused_norm_packed", diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 162914f39e08..09cb39acf919 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import cast +from typing import TYPE_CHECKING, cast import torch from torch import nn @@ -23,7 +23,10 @@ QSAKeyStateCache, canonical_qsa_rope_positions, ) -from .ops.qsa_pre_indexer import qsa_pre_indexer, supports_fused_pre_indexer +from .ops.qsa_prepare import qsa_prepare, supports_fused_pre_indexer + +if TYPE_CHECKING: + from .qsa import Qwen4ExpQSAAttention def apply_qsa_rope( @@ -265,26 +268,31 @@ def _update_and_compress( position_rows, ) - def _fused_pre_indexer( + def _fused_prepare( self, projected_q: torch.Tensor, raw_keys: torch.Tensor, positions: torch.Tensor, raw_metadata: QSAForwardMetadata, compressed_metadata: QSAForwardMetadata, - ) -> torch.Tensor: - """Normalize and rotate Q, compress K, and write both caches in one launch. + attn: Qwen4ExpQSAAttention, + qkv: torch.Tensor, + slot_mapping: torch.Tensor, + gate_out: torch.Tensor | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Prepare the indexer and ``attn``'s main attention in one launch. - Replaces ``project_qk``'s norm/RoPE half and the whole of - ``_update_and_compress``; the two must stay numerically interchangeable. + Normalizes and rotates the index Q, compresses K and writes both side + caches, which replaces ``project_qk``'s norm/RoPE half and the whole of + ``_update_and_compress``. The same launch applies ``attn``'s QK-norm and + RoPE, copies its gate and writes its K/V into ``attn.kv_cache``, + replacing ``_project_qkv_gate`` and ``do_kv_cache_update``. Both halves + must stay numerically interchangeable with what they replace. """ + num_tokens = raw_metadata.num_actual_tokens raw_key_cache = self.raw_key_cache - q = projected_q.new_empty( - raw_metadata.num_actual_tokens, - self.index_n_heads, - self.index_head_dim, - ) - qsa_pre_indexer( + q = projected_q.new_empty(num_tokens, self.index_n_heads, self.index_head_dim) + main_outputs = qsa_prepare( projected_q, raw_keys, positions, @@ -308,8 +316,17 @@ def _fused_pre_indexer( if raw_key_cache.cache_rope_positions else None ), + main_qkv=qkv[:num_tokens], + main_q_norm_weight=attn.q_norm.weight, + main_k_norm_weight=attn.k_norm.weight, + main_eps=attn.q_norm.variance_epsilon, + main_kv_cache=attn.kv_cache.transpose(1, 2), + main_slot_mapping=slot_mapping[:num_tokens], + main_k_scale=attn._k_scale_float, + main_v_scale=attn._v_scale_float, + main_gate_out=None if gate_out is None else gate_out[:num_tokens], ) - return q + return q, main_outputs def _select( self, @@ -336,13 +353,25 @@ def forward( hidden_states: torch.Tensor, positions: torch.Tensor, out: torch.Tensor | None = None, - ) -> torch.Tensor: - """Return fixed-width request-relative token indices padded with ``-1``.""" + *, + attn: Qwen4ExpQSAAttention, + qkv: torch.Tensor | None = None, + slot_mapping: torch.Tensor | None = None, + gate_out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + """Return fixed-width request-relative token indices padded with ``-1``. + + With ``attn.use_fused_qsa_prepare``, the same launch also writes + ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares + its Q and gate from ``qkv``, returned as the second element (None + otherwise). The gate is written into ``gate_out`` when given. ``qkv``, + ``slot_mapping`` and ``gate_out`` are only read in that mode. + """ metadata = self._metadata() if metadata is None: # Preserve step-0 indices when later MTP steps reuse the buffer. if self.skip_topk and out is not None: - return out + return out, None result = torch.full( (hidden_states.shape[0], self.output_width), -1, @@ -351,20 +380,27 @@ def forward( ) if out is not None: out.copy_(result) - return out - return result + return out, None + return result, None raw_metadata, compressed_metadata = metadata num_tokens = raw_metadata.num_actual_tokens hidden_states = hidden_states[:num_tokens] positions = positions[..., :num_tokens] - if self.use_fused_pre_indexer: + main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None + if attn.use_fused_qsa_prepare: + if qkv is None or slot_mapping is None: + raise ValueError("fused QSA prepare requires qkv and slot_mapping") projected_q, raw_keys = self.project(hidden_states) - q = self._fused_pre_indexer( + q, main_outputs = self._fused_prepare( projected_q, raw_keys, positions, raw_metadata, compressed_metadata, + attn, + qkv, + slot_mapping, + gate_out, ) else: q, token_k = self.project_qk(hidden_states, positions) @@ -377,8 +413,8 @@ def forward( if self.skip_topk: if out is None: raise RuntimeError("QSA top-k reuse requires an output buffer") - return out - return self._select(q, compressed_metadata, out) + return out, main_outputs + return self._select(q, compressed_metadata, out), main_outputs __all__ = ["QSAIndexer", "apply_qsa_rope"] diff --git a/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py similarity index 75% rename from vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py rename to vllm/models/qwen4_exp/amd/ops/qsa_prepare.py index 6a4873882f2c..bc2fee3b3a47 100644 --- a/vllm/models/qwen4_exp/amd/ops/qsa_pre_indexer.py +++ b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused QSA pre-indexer kernel for the AMD Qwen4Exp path. +"""Fused QSA prepare kernel for the AMD Qwen4Exp path. Started as a copy of the NVIDIA kernel (``nvidia/ops/qsa_prepare.py``) and is kept separate so either vendor can change its own copy without re-validating @@ -119,6 +119,27 @@ def _norm_rope( return tl.reshape(result, (TILE_T, TILE_H, D)) +@triton.jit +def _to_dst_dtype(x, dst, scale): + """Round to BF16 like the unfused path, then scale for an FP8 destination.""" + out_ty = dst.dtype.element_ty + x = x.to(tl.bfloat16) + if out_ty == tl.float8e4nv: + x = x.to(tl.float32) / scale + return x.to(out_ty) + + +@triton.jit +def _store_rotated(dst, y, o1, o2, scale): + """Store a normalized head whose first ``2 * len(o1)`` dims are rotated.""" + HALF: tl.constexpr = o1.shape[0] + dims = tl.arange(0, y.shape[0]) + rot = tl.arange(0, HALF) + tl.store(dst + dims, _to_dst_dtype(y, dst, scale), mask=dims >= 2 * HALF) + tl.store(dst + rot, _to_dst_dtype(o1, dst, scale)) + tl.store(dst + HALF + rot, _to_dst_dtype(o2, dst, scale)) + + @triton.jit( do_not_specialize=[ "num_tokens", @@ -127,7 +148,7 @@ def _norm_rope( "num_k_work", ] ) -def _qsa_pre_indexer_kernel( +def _qsa_prepare_kernel( q_ptr, q_stride_token, k_ptr, @@ -171,8 +192,86 @@ def _qsa_pre_indexer_kernel( CACHE_HAS_ROPE_POS: tl.constexpr, MROPE_H: tl.constexpr, MROPE_W: tl.constexpr, + main_qkv_ptr, + main_q_norm_weight_ptr, + main_k_norm_weight_ptr, + main_eps, + main_q_out_ptr, + main_gate_out_ptr, + main_cache_ptr, + main_cache_stride_block, + main_cache_stride_token, + main_cache_stride_head, + main_slots_ptr, + main_k_scale, + main_v_scale, + MAIN_HQ: tl.constexpr, + MAIN_HK: tl.constexpr, + MAIN_D: tl.constexpr, + MAIN_PAGE_SIZE: tl.constexpr, ): pid = tl.program_id(0) + num_index_work = num_k_work + tl.cdiv(num_tokens, TILE_T_Q) * tl.cdiv(HQ, TILE_H_Q) + if pid >= num_index_work: + # Main attention: one program per (token, Q or KV head) after the + # indexer work. RoPE covers the first D // 2 dims, the width of the + # shared cos/sin table. + main_pid = pid - num_index_work + token = (main_pid // (MAIN_HQ + MAIN_HK)).to(tl.int64) + head = main_pid % (MAIN_HQ + MAIN_HK) + HALF: tl.constexpr = D // 4 + dims = tl.arange(0, MAIN_D) + rot = tl.arange(0, HALF) + row = main_qkv_ptr + token * (2 * (MAIN_HQ + MAIN_HK) * MAIN_D) + is_k = head >= MAIN_HQ + kv_head = head - MAIN_HQ + if is_k: + src = row + (2 * MAIN_HQ + kv_head) * MAIN_D + weight_ptr = main_k_norm_weight_ptr + else: + src = row + 2 * head * MAIN_D + weight_ptr = main_q_norm_weight_ptr + x = tl.load(src + dims).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=0) / MAIN_D + main_eps) + w = tl.load(weight_ptr + dims).to(tl.float32) + 1.0 + y = x * inv_rms * w + x1 = tl.load(src + rot).to(tl.float32) + x2 = tl.load(src + HALF + rot).to(tl.float32) + w1 = tl.load(weight_ptr + rot).to(tl.float32) + 1.0 + w2 = tl.load(weight_ptr + HALF + rot).to(tl.float32) + 1.0 + x1 = (x1 * inv_rms * w1).to(tl.bfloat16).to(tl.float32) + x2 = (x2 * inv_rms * w2).to(tl.bfloat16).to(tl.float32) + pos = tl.load(pos_ptr + token * pos_stride_token).to(tl.int64) + if IS_2D_POSITIONS: + pos_h = tl.load(pos_ptr + pos_stride_axis + token * pos_stride_token) + pos_w = tl.load(pos_ptr + 2 * pos_stride_axis + token * pos_stride_token) + is_h = (rot % 3 == 1) & (rot < 3 * MROPE_H) + is_w = (rot % 3 == 2) & (rot < 3 * MROPE_W) + pos = tl.where( + is_h, pos_h.to(tl.int64), tl.where(is_w, pos_w.to(tl.int64), pos) + ) + cos = tl.load(cos_sin_ptr + pos * (D // 2) + rot).to(tl.float32) + sin = tl.load(cos_sin_ptr + pos * (D // 2) + HALF + rot).to(tl.float32) + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if is_k: + slot = tl.load(main_slots_ptr + token).to(tl.int64) + if slot >= 0: + dst = ( + main_cache_ptr + + (slot // MAIN_PAGE_SIZE) * main_cache_stride_block + + (slot % MAIN_PAGE_SIZE) * main_cache_stride_token + + kv_head * main_cache_stride_head + ) + _store_rotated(dst, y, o1, o2, main_k_scale) + v = tl.load(src + MAIN_HK * MAIN_D + dims) + tl.store(dst + MAIN_D + dims, _to_dst_dtype(v, dst, main_v_scale)) + else: + out = (token * MAIN_HQ + head) * MAIN_D + _store_rotated(main_q_out_ptr + out, y, o1, o2, None) + gate = tl.load(src + MAIN_D + dims) + tl.store(main_gate_out_ptr + out + dims, gate) + return # K work occupies the first programs; the remaining programs tile Q. This # keeps both paths in one launch while leaving their register shapes # independent. @@ -181,7 +280,7 @@ def _qsa_pre_indexer_kernel( num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q) token_tile = q_pid // num_head_tiles head_tile = q_pid % num_head_tiles - tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q) + tokens = (token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)).to(tl.int64) heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q) valid_tokens = tokens < num_tokens valid_heads = heads < HQ @@ -283,7 +382,8 @@ def _qsa_pre_indexer_kernel( & (source_tokens < num_tokens) ) current_base = ( - k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token + k_ptr + + tl.maximum(source_tokens, 0).to(tl.int64)[:, None] * k_stride_token ) cached_base = ( state_cache_ptr @@ -397,7 +497,7 @@ def _qsa_pre_indexer_kernel( # One CTA per request commits only the suffix retained by the ring. num_state_rows = tl.minimum(query_len, STATE_SIZE) for state_offset in tl.range(0, num_state_rows): - token = query_end - num_state_rows + state_offset + token = (query_end - num_state_rows + state_offset).to(tl.int64) valid_token = ( (token >= query_start) & (token < query_end) & (token < num_tokens) ) @@ -447,7 +547,7 @@ def _qsa_pre_indexer_kernel( tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) -def qsa_pre_indexer( +def qsa_prepare( q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor, @@ -468,11 +568,34 @@ def qsa_pre_indexer( compress_ratio: int, mrope_section: tuple[int, int, int] | None, rope_pos_offset: int | None, -) -> None: - """Normalize Q, compress K, then update the circular raw state.""" + main_qkv: torch.Tensor, + main_q_norm_weight: torch.Tensor, + main_k_norm_weight: torch.Tensor, + main_eps: float, + main_kv_cache: torch.Tensor, + main_slot_mapping: torch.Tensor, + main_k_scale: float, + main_v_scale: float, + main_gate_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize Q, compress K, then update the circular raw state. + + Also prepares the main attention (QK-norm/RoPE, gate copy, K/V cache write) + and returns its Q and gate. ``main_gate_out`` lets the gate land in a + caller-owned buffer, e.g. one that outlives an opaque custom op. + """ num_tokens = q.shape[0] + main_head_dim = main_kv_cache.shape[-1] // 2 + num_main_kv_heads = main_kv_cache.shape[2] + num_main_q_heads = main_qkv.shape[1] // (2 * main_head_dim) - num_main_kv_heads + main_q_out = main_qkv.new_empty(num_tokens, num_main_q_heads, main_head_dim) + if main_gate_out is None: + main_gate_out = torch.empty_like(main_q_out) + else: + assert main_gate_out.shape == main_q_out.shape + assert main_gate_out.is_contiguous() if num_tokens == 0: - return + return main_q_out, main_gate_out num_q_heads, head_dim = q_out.shape[1:] assert cos_sin_cache.shape[-1] * 2 == head_dim assert q.shape == (num_tokens, num_q_heads * head_dim) @@ -497,6 +620,9 @@ def qsa_pre_indexer( pos_stride_axis, pos_stride_token = 0, positions.stride(0) section = mrope_section if mrope_section is not None else (0, 0, 0) assert len(section) == 3 + qkv_width = 2 * (num_main_q_heads + num_main_kv_heads) * main_head_dim + assert main_qkv.shape == (num_tokens, qkv_width) and main_qkv.is_contiguous() + assert main_slot_mapping.shape == (num_tokens,) # Swept on MI355X (64-lane wavefronts) over TILE_T_Q 1-16, TILE_H_Q 1-4 and # num_warps 1/2/4, from 4 to 16384 tokens: no config beat these by more @@ -507,7 +633,8 @@ def qsa_pre_indexer( TILE_T_Q, TILE_H_Q = 2, 4 num_k_work = k_work_metadata.shape[0] num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q) - _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)]( + num_main_work = num_tokens * (num_main_q_heads + num_main_kv_heads) + _qsa_prepare_kernel[(num_k_work + num_q_work + num_main_work,)]( q, q.stride(0), k, @@ -551,8 +678,26 @@ def qsa_pre_indexer( CACHE_HAS_ROPE_POS=cache_has_rope_pos, MROPE_H=section[1], MROPE_W=section[2], + main_qkv_ptr=main_qkv, + main_q_norm_weight_ptr=main_q_norm_weight, + main_k_norm_weight_ptr=main_k_norm_weight, + main_eps=main_eps, + main_q_out_ptr=main_q_out, + main_gate_out_ptr=main_gate_out, + main_cache_ptr=main_kv_cache, + main_cache_stride_block=main_kv_cache.stride(0), + main_cache_stride_token=main_kv_cache.stride(1), + main_cache_stride_head=main_kv_cache.stride(2), + main_slots_ptr=main_slot_mapping, + main_k_scale=main_k_scale, + main_v_scale=main_v_scale, + MAIN_HQ=num_main_q_heads, + MAIN_HK=num_main_kv_heads, + MAIN_D=main_head_dim, + MAIN_PAGE_SIZE=main_kv_cache.shape[1], num_warps=1, ) + return main_q_out, main_gate_out -__all__ = ["qsa_pre_indexer", "supports_fused_pre_indexer"] +__all__ = ["qsa_prepare", "supports_fused_pre_indexer"] diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 3f03b293b4f0..ff01490bb190 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -116,8 +116,8 @@ def forward_qsa( self, layer: torch.nn.Module, query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, kv_cache: torch.Tensor, attn_metadata: FlashAttentionMetadata, output: torch.Tensor, @@ -294,6 +294,11 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.indexer", ) + self.use_fused_qsa_prepare = ( + self.attn_output_gate + and self.indexer.use_fused_pre_indexer + and self.head_dim & (self.head_dim - 1) == 0 + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.register_buffer( "topk_indices_buffer", @@ -327,16 +332,22 @@ def _run_qsa( self, hidden_states: torch.Tensor, positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + query: torch.Tensor | None, + key: torch.Tensor | None, + value: torch.Tensor | None, output: torch.Tensor, + qkv: torch.Tensor | None = None, + gate: torch.Tensor | None = None, ) -> None: + # query/key/value are None when the fused prepare runs inside the + # indexer launch, which then writes the gate into ``gate``. metadata = get_forward_context().attn_metadata if isinstance(metadata, list): metadata = metadata[0] if not isinstance(metadata, dict): output.zero_() + if gate is not None: + gate.zero_() return main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name]) if self.kv_cache.numel() == 0: @@ -349,10 +360,14 @@ def _run_qsa( ) if side_metadata.num_actual_tokens != num_tokens: raise RuntimeError("QSA main and side metadata token counts disagree") - selected = self.indexer( + selected, main_outputs = self.indexer( hidden_states, positions, self.topk_indices_buffer[:num_tokens], + attn=self, + qkv=qkv, + slot_mapping=main_metadata.slot_mapping, + gate_out=gate, ) if selected.shape != ( num_tokens, @@ -360,13 +375,22 @@ def _run_qsa( ): raise RuntimeError("QSA indexer returned an invalid selection shape") impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) - impl.do_kv_cache_update( - self, - key, - value, - self.kv_cache, - main_metadata.slot_mapping, - ) + if main_outputs is None: + assert key is not None and value is not None + impl.do_kv_cache_update( + self, + key, + value, + self.kv_cache, + main_metadata.slot_mapping, + ) + else: + query, _ = main_outputs + # The gate is applied to every row outside this op; padded rows + # must stay finite so they keep the zero output forward_qsa gives. + assert gate is not None + gate[num_tokens:].zero_() + assert query is not None impl.forward_qsa( self, query, @@ -384,13 +408,40 @@ def forward( hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) - q, k, v, gate = self._project_qkv_gate(qkv, positions) num_tokens = hidden_states.shape[0] + encoded_layer_name = _encode_layer_name(self.layer_name) + if self.use_fused_qsa_prepare: + # Norm/RoPE/gate and the K/V cache write happen inside the op. + attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim) + gate = torch.empty_like(attn_output) + if current_platform.opaque_attention_op(): + torch.ops.vllm.qwen4_exp_qsa_prepare_with_output( + hidden_states, + positions, + qkv, + attn_output, + gate, + encoded_layer_name, + ) + else: + qwen4_exp_qsa_prepare_with_output( + hidden_states, + positions, + qkv, + attn_output, + gate, + encoded_layer_name, + ) + flat_output = attn_output.view(num_tokens, -1) + flat_output = flat_output * torch.sigmoid(gate.view(num_tokens, -1)) + output, _ = self.o_proj(flat_output) + return output + + q, k, v, gate = self._project_qkv_gate(qkv, positions) query = q.view(num_tokens, self.num_heads, self.head_dim) key = k.view(num_tokens, self.num_kv_heads, self.head_dim) value = v.view(num_tokens, self.num_kv_heads, self.head_dim) attn_output = torch.empty_like(query) - encoded_layer_name = _encode_layer_name(self.layer_name) if current_platform.opaque_attention_op(): torch.ops.vllm.qwen4_exp_qsa_with_output( hidden_states, @@ -449,10 +500,43 @@ def qwen4_exp_qsa_with_output( ) +def qwen4_exp_qsa_prepare_with_output( + hidden_states: torch.Tensor, + positions: torch.Tensor, + qkv: torch.Tensor, + output: torch.Tensor, + gate: torch.Tensor, + layer_name: LayerNameType, +) -> None: + """Run the QSA transaction with the main QK-norm/RoPE/gate fused in.""" + layer_name = _resolve_layer_name(layer_name) + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpQSAAttention): + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner") + layer._run_qsa( + hidden_states, + positions, + None, + None, + None, + output, + qkv=qkv, + gate=gate, + ) + + +direct_register_custom_op( + op_name="qwen4_exp_qsa_prepare_with_output", + op_func=qwen4_exp_qsa_prepare_with_output, + mutates_args=["output", "gate"], +) + + __all__ = [ "QSAIndexer", "Qwen4ExpQSAAttention", "Qwen4ExpQSAFlashAttentionBackend", "Qwen4ExpQSAFlashAttentionImpl", + "qwen4_exp_qsa_prepare_with_output", "qwen4_exp_qsa_with_output", ] From 3d52578e13454de0bb8352cac988347cdcc89928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:23:35 +0200 Subject: [PATCH 294/952] [Security] Do not retain Phi-4 MM audio attention matrices (#60614) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/model_executor/test_phi4mm_attention.py | 28 +++++++++++++++++++ vllm/model_executor/models/phi4mm_audio.py | 8 +++--- vllm/model_executor/models/phi4mm_utils.py | 7 ++--- 3 files changed, 34 insertions(+), 9 deletions(-) create mode 100644 tests/model_executor/test_phi4mm_attention.py diff --git a/tests/model_executor/test_phi4mm_attention.py b/tests/model_executor/test_phi4mm_attention.py new file mode 100644 index 000000000000..51710f84bc47 --- /dev/null +++ b/tests/model_executor/test_phi4mm_attention.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.models.phi4mm_utils import MultiHeadedAttention + + +def test_eager_attention_does_not_retain_probability_matrix(): + module = MultiHeadedAttention( + n_head=2, + n_feat=16, + dropout_rate=0.0, + use_pt_scaled_dot_product_attention=False, + ) + module.eval() + batch, time, feat = 2, 8, 16 + x = torch.randn(batch, time, feat) + with torch.no_grad(): + out = module(x, x, x, None, None, None) + + assert out.shape == (batch, time, feat) + assert getattr(module, "attn", None) is None + + +def test_multiheaded_attention_defaults_to_scaled_dot_product(): + module = MultiHeadedAttention(n_head=2, n_feat=16, dropout_rate=0.0) + assert module.use_pt_scaled_dot_product_attention is True diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 2a16b4ca7522..9daedcf38e7b 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -129,7 +129,7 @@ class ConformerEncoderLayer(nn.Module): default False. use_pt_scaled_dot_product_attention: bool, optional if set to True, use pytorch's scaled dot product attention - implementation in training. + implementation in training. Default: True. attn_group_sizes: int, optional the number of groups to use for attention, default 1 (Multi-Head Attention), @@ -163,7 +163,7 @@ def __init__( attention_glu_type: str = "swish", activation_checkpointing: str = "", export: bool = False, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, attn_group_sizes: int = 1, ) -> None: super().__init__() @@ -808,7 +808,7 @@ class ConformerEncoder(TransformerEncoderBase): default 4 use_pt_scaled_dot_product_attention: whether to use pytorch scaled dot product attention in training. - Default: False + Default: True nemo_conv_settings: dict, optional A dictionary of settings for NeMo Subsampling. default: None @@ -873,7 +873,7 @@ def __init__( # pylint: disable-all activation_checkpointing: str = "", relative_attention_bias_args: dict[str, Any] | None = None, time_reduction: int = 4, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, nemo_conv_settings: dict[str, Any] | None = None, conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none", replication_pad_for_subsample_embedding: bool = False, diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py index af54defd3fb0..db1b83f34388 100644 --- a/vllm/model_executor/models/phi4mm_utils.py +++ b/vllm/model_executor/models/phi4mm_utils.py @@ -1642,7 +1642,7 @@ class MultiHeadedAttention(nn.Module): NOTE: this will NOT be used in ONNX decoding due to a lack of support. In that case, we use the original attention implementation, which shows no regression. - default: False. + default: True. n_value: int, optional if set to values other than -1, use a different dimension for value. With the default value (i.e. -1), it is backward compatible. @@ -1666,7 +1666,7 @@ def __init__( attention_inner_dim: int = -1, glu_type: str = "swish", bias_in_glu: bool = True, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, n_value: int = -1, group_size: int = 1, ) -> None: @@ -1690,7 +1690,6 @@ def __init__( self.linear_v = nn.Linear(n_value, attention_inner_dim // group_size) self.linear_out = nn.Linear(attention_inner_dim // group_size, n_value) - self.attn = torch.jit.Attribute(None, Tensor | None) self.dropout = nn.Dropout(p=dropout_rate) self.dropout_rate = dropout_rate self.use_pt_scaled_dot_product_attention = use_pt_scaled_dot_product_attention @@ -1799,8 +1798,6 @@ def forward( attn = masked_softmax(scores, mask) # (batch, head, time1, time2) - self.attn = attn - p_attn = self.dropout(attn) x = torch.matmul(p_attn.to(v.dtype), v) # (batch, head, time1, d_k) if pos_v is not None: From fc39317392714b571868d7a49239db471f2dee0d Mon Sep 17 00:00:00 2001 From: Warren Deng <35615489+warrendeng@users.noreply.github.com> Date: Fri, 9 Oct 2026 07:25:41 -0700 Subject: [PATCH 295/952] [Bugfix][Kernel] Remove invalid NC from Kimi-K3 KDA autotune keys (#60446) Signed-off-by: Warren Deng Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py | 2 +- vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py index 0c600f3ba33d..bc13318d83cd 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -349,7 +349,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=["BK", "NC", "BT", "HV"], + key=["BK", "BT", "HV"], ) @triton.jit(do_not_specialize=["B", "T"]) def chunk_kda_fwd_kernel_intra_sub_chunk( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 8b9aa86f9ed9..07b4a388936e 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -451,7 +451,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=["BK", "NC", "BT", "HV"], + key=["BK", "BT", "HV"], ) @triton.jit(do_not_specialize=["B", "T"]) def chunk_kda_fwd_kernel_intra_sub_chunk( From a6b39f02620e23422113823810148dfda9876304 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 9 Oct 2026 09:32:47 -0500 Subject: [PATCH 296/952] [Bugfix] Clean up async client sockets after event loop closure (#57757) Signed-off-by: Taneem Ibrahim --- vllm/v1/engine/core_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 74520036a7c8..dccc8cba5368 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -538,10 +538,10 @@ def close_sockets_and_tasks(): with contextlib.suppress(Exception): task.cancel() - if loop is not None: + if loop is not None and not loop.is_closed(): if in_loop(loop): close_sockets_and_tasks() - elif not loop.is_closed(): + else: loop.call_soon_threadsafe(close_sockets_and_tasks) else: # Loop has been closed, try to clean up directly. From 970dc63478da41409dec009690f462e8e33d22f2 Mon Sep 17 00:00:00 2001 From: Abhijit Roy <50175934+Roy214@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:11:38 +0530 Subject: [PATCH 297/952] [Docs] Add Giga-Embeddings bidirectional Qwen3 usage to embed docs (#60854) Signed-off-by: Abhijit --- docs/models/pooling_models/embed.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 13a2e0d38147..4c33126977e9 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -50,7 +50,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | ✅︎ | | | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | -| `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ | +| `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, `ai-sage/Giga-Embeddings-instruct-480M-0826`^, etc. | ✅︎ | ✅︎ | | `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | | `VoyageQwen3BidirectionalEmbedModel`C | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ | | `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | | | @@ -97,6 +97,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `*ForConditionalGeneration`C, `*ForCausalLM`C, etc. | Generative models | \* | N/A | \* | \* | C Automatically converted into an embedding model via `--convert embed`. ([details](./README.md#model-conversion)) +^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM. \* Feature support is the same as that of the original model. If your model is not in the above list, we will try to automatically convert the model using From 90ba2f34b5a29d85b8301672e5735676f78fa00c Mon Sep 17 00:00:00 2001 From: Greg Pereira Date: Fri, 9 Oct 2026 08:05:54 -0700 Subject: [PATCH 298/952] [Build] Find compatible published wheels for precompiled installs (#58871) Signed-off-by: greg pereira Signed-off-by: mgoin Co-authored-by: mgoin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/misc.yaml | 3 + .../installation/gpu.cuda.inc.md | 17 +- setup.py | 109 +++++++++- .../test_precompiled_wheel.py | 195 ++++++++++++++++++ 4 files changed, 306 insertions(+), 18 deletions(-) create mode 100644 tests/standalone_tests/test_precompiled_wheel.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 9763b2ce7de4..93f59b50e52f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -556,6 +556,8 @@ steps: - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py + - tests/standalone_tests/test_precompiled_wheel.py + - setup.py - tests/reasoning - tools/pre_commit/check_test_tethering.py - tools/pre_commit/test_tethering_allowlist.txt @@ -572,6 +574,7 @@ steps: commands: - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none - python3 standalone_tests/lazy_imports.py + - python3 standalone_tests/test_precompiled_wheel.py - pytest -v -s device_allocator - pytest -v -s test_uvicorn_logging.py - pytest -v -s test_attention_backend_registry.py diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 6ad581ac97d1..f125a189fcd2 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -117,18 +117,21 @@ If you need to recompile the `vllm-rs` Rust frontend binary, you can rebuild and This will install the required Rust toolchain if needed, build the binary, and place it in `vllm/vllm-rs`. -In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`. +Wheels are published an hour or two after a commit lands on `main`. If the wheel for +your merge-base with upstream `main` is not available yet, the installer uses the +nearest older commit that has one, searching up to 20 commits back. It stops with an +error if compiled sources or build configuration (`csrc/`, `cmake/`, `CMakeLists.txt`, +`pyproject.toml`, `vllm/_custom_ops.py`) changed between that commit and your +merge-base, and warns if `rust/` or `setup.py` did. This is a source-level check, not +a guarantee of binary compatibility. -```bash -export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly -export VLLM_USE_PRECOMPILED=1 -uv pip install --editable . -``` +Setting `VLLM_PRECOMPILED_WHEEL_LOCATION` or a full SHA in +`VLLM_PRECOMPILED_WHEEL_COMMIT` skips this search and check. There are more environment variables to control the behavior of Python-only build: - `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file path of a pre-compiled wheel to use. All other logic to find the wheel will be skipped. -- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit hash to download the pre-compiled wheel. It can be `nightly` to use the last **already built** commit on the main branch. +- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit to download the pre-compiled wheel from. Must be a full 40-character SHA. - `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cu130`, `cpu`. If not specified, the variant is auto-detected based on your system's CUDA version (from PyTorch or nvidia-smi). You can also set `VLLM_MAIN_CUDA_VERSION` to override auto-detection. You can find more information about vLLM's wheels in [Install the latest code](#install-the-latest-code). diff --git a/setup.py b/setup.py index 16a3d35ecad6..b173b9dca8fd 100644 --- a/setup.py +++ b/setup.py @@ -896,6 +896,93 @@ def determine_wheel_url_rocm() -> tuple[str, str | None]: print(f"Using ROCm precompiled wheel: {wheel_url}") return wheel_url, download_filename + # Publication lags main by roughly 5-10 commits; anything much older is + # unlikely to still match the checkout's compiled sources. + WHEEL_SEARCH_DEPTH = 20 + # Changes here make an older wheel's extensions unusable. + WHEEL_BLOCKING_PATHS = ( + "csrc", + "cmake", + "CMakeLists.txt", + "pyproject.toml", + "vllm/_custom_ops.py", + ) + # Changes here are usually harmless for an older wheel. + WHEEL_WARNING_PATHS = ("rust", "setup.py") + + @staticmethod + def find_published_wheel_commit( + base_commit: str, variant: str | None, arch: str + ) -> str: + """Find the nearest ancestor of base_commit that has a published wheel. + + Raises: + ValueError: If no wheel is published within WHEEL_SEARCH_DEPTH + commits, or compiled sources changed since the nearest one. + + """ + from urllib.error import HTTPError + + utils = precompiled_wheel_utils + candidates = subprocess.check_output( + [ + "git", + "rev-list", + "--first-parent", + f"--max-count={utils.WHEEL_SEARCH_DEPTH}", + base_commit, + ], + text=True, + ).splitlines() + for distance, commit in enumerate(candidates): + try: + wheels, _ = utils.fetch_metadata_for_variant(commit, variant) + except HTTPError as err: + if err.code != 404: + raise + continue + if not any( + wheel.get("package_name") == "vllm" + and arch in wheel.get("platform_tag", "") + for wheel in wheels + ): + continue + if distance == 0: + return commit + + diff = ["git", "diff", "--name-only", commit, base_commit, "--"] + blocking, warning = ( + subprocess.check_output([*diff, *paths], text=True).strip() + for paths in (utils.WHEEL_BLOCKING_PATHS, utils.WHEEL_WARNING_PATHS) + ) + if blocking: + raise ValueError( + f"No precompiled wheel is published yet for {base_commit}, and " + f"the nearest one ({commit}, {distance} commits older) was " + f"built from different compiled sources:\n{blocking}\n" + "Wait for wheel publication, rebase onto a commit with a " + "published wheel, build from source (unset " + "VLLM_USE_PRECOMPILED), or set VLLM_PRECOMPILED_WHEEL_COMMIT " + "to a full commit SHA to skip this check." + ) + if warning: + logger.warning( + "Precompiled wheel %s predates changes to:\n%s\n" + "The Rust frontend and packaged files may be stale.", + commit, + warning, + ) + print( + f"No precompiled wheel is published yet for {base_commit}; using " + f"{commit} ({distance} commits older)" + ) + return commit + raise ValueError( + f"No published precompiled wheel for variant {variant} and " + f"architecture {arch} within {utils.WHEEL_SEARCH_DEPTH} commits of " + f"{base_commit}. Wait for wheel publication or build vLLM from source." + ) + @staticmethod def determine_wheel_url() -> tuple[str, str | None]: """Try to determine the precompiled wheel URL or path to use. @@ -906,8 +993,8 @@ def determine_wheel_url() -> tuple[str, str | None]: or CUDA variant selected from VLLM_MAIN_CUDA_VERSION, torch, or nvidia-smi If downloading from the nightly repo, the commit can be specified via - VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch - is used. + VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, local checkouts search backwards + from their merge-base with upstream main for a compatible published wheel. """ wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None) if wheel_location is not None: @@ -933,6 +1020,14 @@ def determine_wheel_url() -> tuple[str, str | None]: ", trying to fetch base commit in main branch" ) commit = precompiled_wheel_utils.get_base_commit_in_main_branch() + if ( + envs.VLLM_USE_PRECOMPILED + and not envs.VLLM_DOCKER_BUILD_CONTEXT + and commit != "nightly" + ): + commit = precompiled_wheel_utils.find_published_wheel_commit( + commit, variant, arch + ) print(f"Using precompiled wheel commit {commit} with variant {variant}") download_filename = None try: @@ -1149,17 +1244,9 @@ def get_base_commit_in_main_branch() -> str: ["git", "fetch", "https://github.com/vllm-project/vllm", "main"] ) - # Then get the commit hash of the current branch that is the same as - # the upstream main commit. - current_branch = ( - subprocess.check_output(["git", "branch", "--show-current"]) - .decode("utf-8") - .strip() - ) - base_commit = ( subprocess.check_output( - ["git", "merge-base", f"{upstream_main_commit}", current_branch] + ["git", "merge-base", upstream_main_commit, "HEAD"] ) .decode("utf-8") .strip() diff --git a/tests/standalone_tests/test_precompiled_wheel.py b/tests/standalone_tests/test_precompiled_wheel.py new file mode 100644 index 000000000000..883a6cea842f --- /dev/null +++ b/tests/standalone_tests/test_precompiled_wheel.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Exercise setup-time wheel selection without importing torch or running setup.""" + +import ast +import logging +import os +import subprocess +import tempfile +import unittest +from email.message import Message +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch +from urllib.error import HTTPError, URLError + +# setup.py executes setup() on import; load only the wheel utility class. +setup_path = Path(__file__).resolve().parents[2] / "setup.py" +node = next( + node + for node in ast.parse(setup_path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == "precompiled_wheel_utils" +) +namespace: dict[str, Any] = { + "subprocess": subprocess, + "os": os, + "logger": logging.getLogger("setup"), + "envs": SimpleNamespace(VLLM_DOCKER_BUILD_CONTEXT=False, VLLM_USE_PRECOMPILED=True), +} +exec( + compile(ast.Module(body=[node], type_ignores=[]), str(setup_path), "exec"), + namespace, +) +wheel_utils = namespace["precompiled_wheel_utils"] +WHEEL = { + "package_name": "vllm", + "platform_tag": "manylinux_x86_64", + "path": "wheel.whl", + "filename": "wheel.whl", +} +NOT_FOUND = HTTPError("url", 404, "missing", Message(), None) + + +class TestPublishedWheelSearch(unittest.TestCase): + """Walk back from the merge-base in a real Git repo; only the index is mocked.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.repo = Path(tmp.name) + self.addCleanup(os.chdir, os.getcwd()) + os.chdir(self.repo) + self.git("init", "-q") + self.published = {self.commit("vllm/a.py")} + + def git(self, *args: str) -> str: + identity = ["-c", "user.name=t", "-c", "user.email=t@t"] + return subprocess.check_output(["git", *identity, *args], text=True).strip() + + def commit(self, path: str) -> str: + file = self.repo / path + file.parent.mkdir(parents=True, exist_ok=True) + file.write_text(file.read_text() + "x" if file.exists() else "x") + self.git("add", "-A") + self.git("commit", "-q", "-m", path) + return self.git("rev-parse", "HEAD") + + def search(self, base: str, arch: str = "x86_64") -> str: + def fetch(commit, variant): + if commit not in self.published: + raise NOT_FOUND + return [WHEEL], "url" + + with patch.object(wheel_utils, "fetch_metadata_for_variant", fetch): + return wheel_utils.find_published_wheel_commit(base, "cu130", arch) + + def test_published_base_is_used_despite_local_native_edits(self): + base = self.commit("csrc/ops.cpp") + self.published.add(base) + self.commit("csrc/ops.cpp") + (self.repo / "pyproject.toml").write_text("dirty") + self.assertEqual(self.search(base), base) + + def test_publication_lag_uses_nearest_published_ancestor(self): + published = self.commit("vllm/b.py") + self.published.add(published) + self.commit("vllm/c.py") + self.assertEqual(self.search(self.commit("vllm/d.py")), published) + + def test_compiled_changes_since_wheel_are_rejected(self): + (published,) = self.published + for path in wheel_utils.WHEEL_BLOCKING_PATHS: + with self.subTest(path=path): + self.git("reset", "-q", "--hard", published) + file = path if "." in path else f"{path}/file" + with self.assertRaisesRegex(ValueError, file): + self.search(self.commit(file)) + + def test_rust_changes_since_wheel_only_warn(self): + (published,) = self.published + base = self.commit("rust/src/main.rs") + with self.assertLogs("setup", "WARNING") as logs: + self.assertEqual(self.search(base), published) + self.assertIn("rust/src/main.rs", logs.output[0]) + + def test_wrong_architecture_is_not_a_match(self): + with self.assertRaisesRegex(ValueError, "No published precompiled wheel"): + self.search("HEAD", arch="aarch64") + + def test_search_stops_at_depth_limit(self): + with patch.object(wheel_utils, "WHEEL_SEARCH_DEPTH", 2): + self.commit("vllm/b.py") + with self.assertRaisesRegex(ValueError, "within 2 commits"): + self.search(self.commit("vllm/c.py")) + + def test_network_and_server_errors_do_not_trigger_walkback(self): + for error in ( + URLError("offline"), + HTTPError("url", 503, "busy", Message(), None), + ): + with ( + self.subTest(error=error), + patch.object( + wheel_utils, "fetch_metadata_for_variant", side_effect=error + ) as fetch, + self.assertRaises(type(error)), + ): + wheel_utils.find_published_wheel_commit("HEAD", "cu130", "x86_64") + fetch.assert_called_once() + + +class TestDetermineWheelUrl(unittest.TestCase): + def determine(self, base: str, **env: str): + with ( + patch.dict( + os.environ, + {"VLLM_PRECOMPILED_WHEEL_VARIANT": "cu130", **env}, + clear=True, + ), + patch("platform.machine", return_value="x86_64"), + patch.object(wheel_utils, "is_rocm_system", return_value=False), + patch.object( + wheel_utils, "get_base_commit_in_main_branch", return_value=base + ), + patch.object( + wheel_utils, + "fetch_metadata_for_variant", + return_value=([WHEEL], "https://wheels.test/"), + ) as fetch, + patch.object( + wheel_utils, "find_published_wheel_commit", return_value="older" + ) as search, + ): + url = wheel_utils.determine_wheel_url() + return url, fetch, search + + def test_automatic_install_downloads_the_searched_commit(self): + url, fetch, search = self.determine("base") + self.assertEqual(url, ("https://wheels.test/wheel.whl", "wheel.whl")) + search.assert_called_once_with("base", "cu130", "x86_64") + fetch.assert_called_once_with("older", "cu130") + + def test_search_is_skipped_when_it_cannot_apply(self): + cases = { + "nightly fallback": ("nightly", {}, {}), + "explicit commit": ( + "base", + {"VLLM_PRECOMPILED_WHEEL_COMMIT": "a" * 40}, + {}, + ), + "rust only": ("base", {}, {"VLLM_USE_PRECOMPILED": False}), + "docker": ("base", {}, {"VLLM_DOCKER_BUILD_CONTEXT": True}), + } + for name, (base, env, flags) in cases.items(): + with self.subTest(name), patch.dict(vars(namespace["envs"]), flags): + _, fetch, search = self.determine(base, **env) + search.assert_not_called() + commit = env.get("VLLM_PRECOMPILED_WHEEL_COMMIT", base) + fetch.assert_called_once_with(commit, "cu130") + + def test_explicit_wheel_location_bypasses_selection(self): + for location in ("/tmp/local.whl", "https://example.com/custom.whl"): + with self.subTest(location=location): + url, fetch, search = self.determine( + "base", VLLM_PRECOMPILED_WHEEL_LOCATION=location + ) + self.assertEqual(url, (location, None)) + fetch.assert_not_called() + search.assert_not_called() + + +if __name__ == "__main__": + unittest.main() From 75e3b17c055c609f02c64cd4ad3f8692b8170d8e Mon Sep 17 00:00:00 2001 From: Taewoong Kim <116135174+ultramancode@users.noreply.github.com> Date: Sat, 10 Oct 2026 01:07:50 +0900 Subject: [PATCH 299/952] [Frontend] Add cache salt support to generative scoring (#60841) Signed-off-by: Taewoong Kim Co-authored-by: Cyrus Leung --- .../online_serving/generative_scoring.md | 7 +++++ docs/usage/security.md | 2 +- .../test_generative_scoring.py | 28 +++++++++++++++---- .../test_non_object_body_validation.py | 12 ++++++++ .../generate/generative_scoring/serving.py | 28 +++++++++++++++++-- 5 files changed, 68 insertions(+), 9 deletions(-) diff --git a/docs/serving/online_serving/generative_scoring.md b/docs/serving/online_serving/generative_scoring.md index d0c578f2105b..eebe1f965ec6 100644 --- a/docs/serving/online_serving/generative_scoring.md +++ b/docs/serving/online_serving/generative_scoring.md @@ -62,3 +62,10 @@ Here, each item is appended to the query to form prompts like `"Is this city the "usage": {"prompt_tokens": 45, "total_tokens": 48, "completion_tokens": 3} } ``` + +## Cache salting + +The optional `cache_salt` request field is applied to every item's prompt. With +prefix caching enabled, requests with different salts do not share cached prefix +blocks. See [Cache Salting](../../usage/security.md#cache-salting) for how to choose +and manage salts. diff --git a/docs/usage/security.md b/docs/usage/security.md index a74ed1d6ebe9..ab5aa96becc5 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details. -`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint. +`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, the structured decisions `/v1/systemone` endpoint, and the generative scoring `/generative_scoring` endpoint. #### Usage with the OpenAI Python client diff --git a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index 69be10b19d70..ea7c81ea2670 100644 --- a/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -139,6 +139,7 @@ def test_request_and_response_all_fields(self): assert req_str.apply_softmax is True # default assert req_str.item_first is False # default assert req_str.add_special_tokens is True # default + assert req_str.cache_salt is None # Test request with pre-tokenized inputs and custom options req_tok = GenerativeScoringRequest( @@ -296,16 +297,26 @@ class TestGeneration: """Tests for the full generation flow with mocked engine.""" @pytest.mark.asyncio - async def test_successful_generation(self): - """Test successful score generation returns valid response.""" + @pytest.mark.parametrize("cache_salt", [None, "tenant-a"]) + @pytest.mark.parametrize( + "query,items", + [ + ("Is this city the capital of France? ", ["Paris", "London"]), + ([100, 101], [[200, 201], [300, 301]]), + ], + ids=["text", "token_ids"], + ) + async def test_successful_generation(self, cache_salt, query, items): + """Test scoring succeeds and preserves cache_salt for every item.""" mock_engine = _create_mock_engine() serving = _create_serving(mock_engine) mock_logprobs = {1234: -0.5, 5678: -2.0, 100: -3.0} mock_output = _create_mock_request_output(mock_logprobs) - state = SimpleNamespace(sampling_params=None) + state = SimpleNamespace(sampling_params=None, prompts=[]) async def mock_generate(*args, **kwargs): + state.prompts.append(args[0]) state.sampling_params = args[1] yield mock_output @@ -313,15 +324,22 @@ async def mock_generate(*args, **kwargs): request = GenerativeScoringRequest( model=MODEL_NAME, - query="Is Paris the capital?", - items=["Yes", "No"], + query=query, + items=items, label_token_ids=[1234, 5678], + cache_salt=cache_salt, ) result = await serving.create_generative_scoring(request, None) assert isinstance(result, GenerativeScoringResponse) assert state.sampling_params is not None assert state.sampling_params.watermarking is False + assert len(state.prompts) == len(items) + for prompt in state.prompts: + if cache_salt is None: + assert "cache_salt" not in prompt + else: + assert prompt["cache_salt"] == cache_salt assert len(result.data) == 2 for item_result in result.data: assert 0.0 <= item_result.score <= 1.0 diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py index dd0e3eafd0a4..2440c880ce91 100644 --- a/tests/entrypoints/unit_tests/test_non_object_body_validation.py +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -13,6 +13,9 @@ import pytest from pydantic import ValidationError +from vllm.entrypoints.generate.generative_scoring.serving import ( + GenerativeScoringRequest, +) from vllm.entrypoints.generate.structured_decisions.protocol import ( StructuredDecisionRequest, ) @@ -31,6 +34,7 @@ pytestmark = pytest.mark.skip_global_cleanup REQUEST_MODELS = [ + GenerativeScoringRequest, CompletionRequest, ResponsesRequest, EmbeddingChatRequest, @@ -77,6 +81,14 @@ def test_tokenize_chat_request_still_validates_dict_bodies(): CACHE_SALT_MODELS = [ + ( + GenerativeScoringRequest, + { + "query": "Is this city the capital of France?", + "items": ["Paris"], + "label_token_ids": [9454, 2753], + }, + ), ( StructuredDecisionRequest, { diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index bd46158edcab..74dbfd075997 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -12,12 +12,13 @@ import math import time from collections.abc import Mapping -from typing import Literal +from typing import Any, Literal from fastapi import Request -from pydantic import Field +from pydantic import Field, model_validator from vllm.engine.protocol import EngineClient +from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.protocol import ( @@ -59,6 +60,7 @@ class GenerativeScoringRequest(OpenAIBaseModel): the full vocab for those ids (False). item_first: If True, prepend items to query. Otherwise append items to query. add_special_tokens: Whether to add special tokens when tokenizing. + cache_salt: Optional salt for prefix caching. """ @@ -91,6 +93,17 @@ class GenerativeScoringRequest(OpenAIBaseModel): default=True, description="Whether to add special tokens when tokenizing.", ) + cache_salt: str | None = Field( + default=None, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) priority: int = Field( default=0, ge=-(2**63), @@ -104,6 +117,13 @@ class GenerativeScoringRequest(OpenAIBaseModel): description="The request_id related to this request.", ) + @model_validator(mode="before") + @classmethod + def check_cache_salt_support(cls, data: Any) -> Any: + if isinstance(data, dict): + validate_cache_salt(data.get("cache_salt")) + return data + class GenerativeScoringItemResult(OpenAIBaseModel): """Result for a single item in the generative scoring response. @@ -395,7 +415,9 @@ async def _build_prompts( if len(prompt_token_ids) > max_prompt_len: prompt_token_ids = prompt_token_ids[:max_prompt_len] - engine_inputs.append(tokens_input(prompt_token_ids)) + engine_inputs.append( + tokens_input(prompt_token_ids, cache_salt=request.cache_salt) + ) prompt_token_counts.append(len(prompt_token_ids)) return engine_inputs, prompt_token_counts From b74233860d4f2419ac5a763f6c83859bfe742ad0 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Fri, 9 Oct 2026 18:48:43 +0200 Subject: [PATCH 300/952] [CI][ROCm] Wait for VRAM to settle before building LLM in test_full_graph (#60872) Signed-off-by: Stefan Koncarevic --- tests/compile/fullgraph/test_full_graph.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/compile/fullgraph/test_full_graph.py b/tests/compile/fullgraph/test_full_graph.py index 4ce7d6767f3f..fceed6346d1e 100644 --- a/tests/compile/fullgraph/test_full_graph.py +++ b/tests/compile/fullgraph/test_full_graph.py @@ -10,11 +10,17 @@ from tests.quantization.utils import is_quant_method_supported from vllm import LLM, SamplingParams -from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig +from vllm.config import ( + CacheConfig, + CompilationConfig, + CompilationMode, + CUDAGraphMode, + PassConfig, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer -from ...utils import create_new_process_for_each_test +from ...utils import create_new_process_for_each_test, wait_for_memory_to_settle def models_list(*, all: bool = True, keywords: list[str] | None = None): @@ -173,6 +179,14 @@ def run_model( if compilation_config.cudagraph_mode is None: compilation_config.cudagraph_mode = CUDAGraphMode.NONE + # V1 startup requires free_memory >= total * gpu_memory_utilization, so a + # device that is not free yet fails engine init rather than waiting for it. + # VllmRunner guards this; building LLM directly does not. No-op off ROCm. + wait_for_memory_to_settle( + threshold_ratio=1.0 + - model_kwargs.get("gpu_memory_utilization", CacheConfig.gpu_memory_utilization) + ) + llm = LLM( model=model, compilation_config=compilation_config, From 89ac8d971be7d2f598598c1551a21c646bdf8b0d Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Fri, 9 Oct 2026 12:59:17 -0400 Subject: [PATCH 301/952] [CI][Test] Allow 10% ngram acceptance drop under preemption on CUDA in test_async_scheduling_accuracy (#60715) Signed-off-by: Andrey Talman Co-authored-by: Andrey Talman --- tests/v1/e2e/general/test_async_scheduling_accuracy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/v1/e2e/general/test_async_scheduling_accuracy.py b/tests/v1/e2e/general/test_async_scheduling_accuracy.py index aa785b3270ce..3ffcaefdc3f3 100644 --- a/tests/v1/e2e/general/test_async_scheduling_accuracy.py +++ b/tests/v1/e2e/general/test_async_scheduling_accuracy.py @@ -271,9 +271,9 @@ def run_tests( context = f"config=[{config}], params={params}" if spec_config.get("max_model_len") is None: # Keep the original acceptance-quality floor per batch. - relative_drop = ( - 0.10 if current_platform.is_rocm() and test_preemption else 0.05 - ) + # Preemption recomputes requests and loses n-gram drafts, + # so allow a larger drop there. + relative_drop = 0.10 if test_preemption else 0.05 assert actual >= baseline * (1 - relative_drop), ( f"{context}: acceptance={actual}, baseline={baseline}" ) From 1a3a7a8f01a8c1a3184151b8185072efda11daab Mon Sep 17 00:00:00 2001 From: Yuvraj Singh Bhadoria Date: Fri, 9 Oct 2026 22:45:22 +0530 Subject: [PATCH 302/952] [Bugfix] Treat max_tokens as output upper bound, not input-space reservation (#57035) Signed-off-by: YuvrajSinghBhadoria2 <245782664+YuvrajSinghBhadoria2@users.noreply.github.com> Signed-off-by: YuvrajSinghBhadoria2 Signed-off-by: Yuvraj Singh Bhadoria Co-authored-by: YuvrajSinghBhadoria2 <245782664+YuvrajSinghBhadoria2@users.noreply.github.com> Co-authored-by: Cyrus Leung Co-authored-by: Claude Opus 4.8 (1M context) --- .../chat_completion/test_serving_chat.py | 22 +++- .../responses/test_serving_responses.py | 6 +- .../pooling/embed/test_protocol.py | 59 +++++++++ .../pooling/test_truncation_control.py | 2 + tests/renderers/test_completions.py | 118 ++++++++++++++++++ vllm/entrypoints/pooling/base/protocol.py | 6 +- vllm/renderers/params.py | 71 +++++++---- 7 files changed, 252 insertions(+), 32 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index bfaa2d4ca0b2..d22f581f67ec 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1684,7 +1684,14 @@ def test_serving_chat_starts_without_tokenizer(): @pytest.mark.asyncio async def test_serving_chat_mistral_token_ids_prompt_is_validated(): """Regression test: when the Mistral tokenizer path returns token IDs - directly, we must still apply input length + max_tokens validation. + directly, we must still apply input length validation and clamp the + generation budget to the context the prompt leaves over. + + `max_tokens` is an upper bound on generation rather than a reservation off + the context window, so a 95-token prompt against `max_model_len=100` is + accepted even though the client asked for 10 output tokens. Genuine + overflow is covered by + test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected. """ mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False @@ -1696,9 +1703,8 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): MockVllmConfig(mock_engine.model_config, parallel_config=MockParallelConfig()), tokenizer=mock_tokenizer, ) - # Force the Mistral chat template renderer to return token IDs. - # Choose a prompt length that is < max_model_len, but large enough that - # adding max_tokens should exceed the model context window. + # Force the Mistral chat template renderer to return token IDs. The prompt + # fits the context on its own; only the naive prompt+output sum overflows. mock_renderer.render_messages_async = AsyncMock( return_value=( [], @@ -1715,8 +1721,12 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): max_tokens=10, ) - with pytest.raises(VLLMValidationError): - await serving_chat.create_chat_completion(req) + # Must not raise: the prompt alone fits within max_model_len. + await serving_chat.create_chat_completion(req) + + # get_max_tokens() clamps the sampling budget to the 5 tokens left over. + sampling_params = mock_engine.generate.call_args.args[1] + assert 1 <= sampling_params.max_tokens <= 5 @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index acab4cb6e2ea..1622d0e8f8cf 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -515,7 +515,11 @@ async def test_online_renderer_applies_responses_token_budget_to_harmony_prompt( result = await renderer.render_responses(request) assert not isinstance(result, ErrorResponse) - assert result.engine_input["prompt_token_ids"] == [4, 5, 6] + # `max_tokens` is an output upper bound rather than a reservation off the + # context window, so the prompt budget is the full max_model_len minus the + # one token held back for output: 6 tokens truncate to 4, not to + # max_model_len - max_output_tokens. + assert result.engine_input["prompt_token_ids"] == [3, 4, 5, 6] @pytest.mark.asyncio diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 91c853f91098..ba797b4ff111 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -4,12 +4,15 @@ underlying packing helpers, plus Cohere-specific serving helpers.""" import struct +from unittest.mock import Mock import numpy as np import pybase64 as base64 import pytest +from vllm.config import ModelConfig, PoolerConfig from vllm.entrypoints.pooling.embed.protocol import ( + EmbeddingCompletionRequest, build_typed_embeddings, ) @@ -162,3 +165,59 @@ def test_subset_types(self, sample_embeddings: list[list[float]]): def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]): result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"]) assert result.float is not None + + +def _embedding_model_config( + *, + max_model_len: int = 128, + max_embed_len: int | None = None, + enable_chunked_processing: bool = False, +) -> Mock: + model_config = Mock(spec=ModelConfig) + model_config.max_model_len = max_model_len + model_config.encoder_config = None + model_config.pooler_config = PoolerConfig( + seq_pooling_type="CLS", + max_embed_len=max_embed_len, + enable_chunked_processing=enable_chunked_processing, + ) + return model_config + + +class TestMaxEmbedLenCapsInput: + """``build_tok_params`` used to enforce ``max_embed_len`` indirectly, by + setting ``max_output_tokens = max_model_len - max_embed_len`` and relying on + ``max_input_tokens == max_total_tokens - max_output_tokens``. The cap is now + expressed directly as ``max_total_tokens``.""" + + def test_non_chunked_max_embed_len_caps_input(self): + cfg = _embedding_model_config(max_model_len=128, max_embed_len=64) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 64 + assert tok_params.max_output_tokens == 0 + + def test_chunked_max_embed_len_caps_input(self): + cfg = _embedding_model_config( + max_model_len=128, max_embed_len=64, enable_chunked_processing=True + ) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 64 + assert tok_params.max_output_tokens == 0 + + def test_no_max_embed_len_falls_back_to_model_len(self): + cfg = _embedding_model_config(max_model_len=128) + + tok_params = EmbeddingCompletionRequest(model="m", input="hi").build_tok_params( + cfg + ) + + assert tok_params.max_input_tokens == 128 + assert tok_params.max_output_tokens == 0 diff --git a/tests/models/language/pooling/test_truncation_control.py b/tests/models/language/pooling/test_truncation_control.py index c4485195558a..003213410c28 100644 --- a/tests/models/language/pooling/test_truncation_control.py +++ b/tests/models/language/pooling/test_truncation_control.py @@ -55,6 +55,8 @@ def test_max_truncation_size(vllm_model): prompt_tokens = vllm_output[0].prompt_token_ids + # Pooling requests emit an embedding rather than generated tokens, so -1 + # truncates to the full context length. assert len(prompt_tokens) == max_model_len diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 0dce7ef68412..4fac06380750 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -247,6 +247,8 @@ def test_neg_truncation(self): ) assert len(results) == 1 + # max_output_tokens defaults to 0, so there is no output budget to + # reserve and -1 truncates to the full max_total_tokens. assert len(results[0]["prompt_token_ids"]) == 100 # max_total_tokens def test_truncation_left(self): @@ -505,6 +507,122 @@ def test_explicit_side_text_pretokenization_guard(self): assert renderer.tokenizer._captured_text_len <= 100 +class TestMaxTokensNotReservation: + """Regression tests for GitHub issue #42474. + + ``max_tokens`` must be treated as an upper bound on generation length, not + as a hard reservation that is subtracted from the available input space. + """ + + def test_large_max_tokens_does_not_eat_prompt_budget(self): + """Prompt fits in context even when max_tokens is large.""" + renderer = _build_renderer(MockModelConfig()) + + # 50-token prompt with a 100-token context window and max_tokens=60. + # Before the fix: max_input_tokens = 100 - 60 = 40 → 50 > 40 → error. + # After the fix: max_input_tokens = 100 → 50 ≤ 100 → OK. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=60), + ) + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 50 + + def test_max_tokens_zero_still_rejects_oversized_prompt(self): + """Genuine prompt overflow is still caught regardless of max_tokens.""" + renderer = _build_renderer(MockModelConfig()) + + # 150-token prompt with 100-token context → should still fail. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 150) + ) + with pytest.raises(VLLMValidationError, match="maximum context length"): + renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=0), + ) + + def test_negative_truncation_uses_max_total(self): + """truncate_prompt_tokens=-1 reserves one output token.""" + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 200) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + max_output_tokens=40, + truncate_prompt_tokens=-1, + ), + ) + # -1 maps to max_total_tokens (100) rather than the old + # max_total - max_output (60), minus the one token held back so + # truncation cannot leave a zero-token output budget. + assert len(results[0]["prompt_token_ids"]) == 99 + + def test_negative_pad_uses_max_total(self): + """pad_prompt_tokens=-1 pads to max_total_tokens, not max_input.""" + renderer = _build_renderer(MockModelConfig()) + pad_id = renderer.tokenizer.pad_token_id + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 30) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + max_output_tokens=40, + pad_prompt_tokens=-1, + ), + ) + # -1 should map to max_total_tokens (100), not max_total - max_output (60). + assert results[0]["prompt_token_ids"] == list(range(30)) + [pad_id] * 70 + + def test_max_tokens_exceeds_context_is_still_clamped_by_api(self): + """Renderer allows the prompt; API layer clamps max_tokens later.""" + renderer = _build_renderer(MockModelConfig()) + + # Simulates the reported scenario: max_tokens > context - prompt_length. + # Renderer should accept the 40-token prompt; get_max_tokens() clamps + # sampling max_tokens to 60 at the entrypoint layer. + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 40) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams(max_total_tokens=100, max_output_tokens=80), + ) + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 40 + + def test_with_kwargs_preserves_output_budget(self): + """with_kwargs without an explicit max_length keeps the output budget. + + Under the old semantics `with_kwargs()` recomputed + `max_output_tokens = max_total - max_input`, which (after this fix) + would have silently reset a real output budget to zero. + """ + tok_params = TokenizeParams(max_total_tokens=100, max_output_tokens=60) + chained = tok_params.with_kwargs(truncation_side="left") + + assert chained.max_output_tokens == 60 + assert chained.max_input_tokens == 100 + + def test_with_kwargs_explicit_max_length_sets_budget(self): + """An explicit max_length still re-derives the output budget.""" + tok_params = TokenizeParams(max_total_tokens=100, max_output_tokens=60) + chained = tok_params.with_kwargs(max_length=80) + + assert chained.max_output_tokens == 20 + assert chained.max_input_tokens == 100 + + class TestRenderEmbedPrompt: def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes: """Helper to create base64-encoded tensor bytes.""" diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index c1366734f7c1..0c646b969ab5 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -183,15 +183,15 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: if pooler_config.enable_chunked_processing: max_total_tokens = pooler_config.max_embed_len else: - max_embed_len = pooler_config.max_embed_len or default_max_total_tokens - max_output_tokens = default_max_total_tokens - max_embed_len + max_total_tokens = ( + pooler_config.max_embed_len or default_max_total_tokens + ) tok_params = self._build_pooling_tok_params( model_config, add_special_tokens=self.add_special_tokens, max_total_tokens=max_total_tokens, max_output_tokens=max_output_tokens, - max_output_tokens_param="max_model_len - max_embed_len", ) if ( diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 9e62dfbb2d23..456a4d0ac563 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -207,11 +207,38 @@ class TokenizeParams: @property def max_input_tokens(self) -> int | None: - """Maximum allowed number of input tokens.""" + """Maximum allowed number of input tokens. + + ``max_output_tokens`` is *not* subtracted here: the value represents the + actual model context length, not a reservation of output space. + Downstream clamping of the sampling ``max_tokens`` is handled by + ``get_max_tokens()`` in the entrypoint layer. + """ if self.max_total_tokens is None: return None - return self.max_total_tokens - self.max_output_tokens + return self.max_total_tokens + + @property + def max_truncation_tokens(self) -> int | None: + """Maximum prompt length that truncation may reduce a prompt to. + + For generative requests this is `max_input_tokens` less one, holding a + token of context back for output. Truncating to the full + `max_total_tokens` would leave `get_max_tokens()` clamping the sampling + budget to 0, turning an over-long prompt into a silent zero-token + generation. + + Pooling requests set `max_output_tokens = 0` and emit an embedding + rather than generated tokens, so they keep the full input budget. + """ + max_input_tokens = self.max_input_tokens + if max_input_tokens is None: + return None + if self.max_output_tokens == 0: + return max_input_tokens + + return max(0, max_input_tokens - 1) def __post_init__(self) -> None: max_total_tokens = self.max_total_tokens @@ -247,14 +274,15 @@ def __post_init__(self) -> None: ): raise VLLMValidationError( f"{self.truncate_prompt_tokens_param}={truncate_prompt_tokens} " - f"cannot be greater than {self.max_total_tokens_param} - " - f"{self.max_output_tokens_param} = {max_input_tokens}. " + f"cannot be greater than {self.max_total_tokens_param} " + f"= {max_input_tokens}. " f"Please request a smaller truncation size.", parameter=self.truncate_prompt_tokens_param, value=truncate_prompt_tokens, ) def with_kwargs(self, **tokenization_kwargs: Any): + has_max_length = "max_length" in tokenization_kwargs max_length = tokenization_kwargs.pop("max_length", self.max_input_tokens) pad_prompt_tokens = tokenization_kwargs.pop( "pad_prompt_tokens", self.pad_prompt_tokens @@ -301,13 +329,17 @@ def with_kwargs(self, **tokenization_kwargs: Any): max_total_tokens = self.max_total_tokens + if has_max_length and max_total_tokens is not None and max_length is not None: + max_output_tokens = max_total_tokens - max_length + else: + # No explicit ``max_length`` override: preserve the original output + # budget instead of recomputing it from ``max_input_tokens`` (which + # is now the full context length). + max_output_tokens = self.max_output_tokens + return TokenizeParams( max_total_tokens=max_total_tokens, - max_output_tokens=( - 0 - if max_total_tokens is None or max_length is None - else max_total_tokens - max_length - ), + max_output_tokens=max_output_tokens, pad_prompt_tokens=pad_prompt_tokens, truncate_prompt_tokens=truncate_prompt_tokens, truncation_side=truncation_side, @@ -320,7 +352,7 @@ def get_encode_kwargs(self) -> dict[str, Any]: """The arguments to pass to `tokenizer.encode`.""" max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: - max_length = self.max_input_tokens + max_length = self.max_truncation_tokens elif max_length is None and self.max_input_tokens is not None: # This prevents tokenization from taking up more resources than necessary # while still failing `self._token_len_check` as expected by users @@ -355,13 +387,11 @@ def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: if len(text) > max_input_chars: raise VLLMValidationError( f"This model's maximum context length is " - f"{self.max_total_tokens} tokens. However, you requested " - f"{self.max_output_tokens} output tokens and your prompt " + f"{self.max_total_tokens} tokens. However, your prompt " f"contains {len(text)} characters (more than " f"{max_input_chars} characters, which is the upper bound " f"for {max_input_tokens} input tokens). " - f"Please reduce the length of the input prompt or the " - f"number of requested output tokens.", + f"Please reduce the length of the input prompt.", parameter="input_text", value=len(text), ) @@ -459,7 +489,7 @@ def _truncation_slice( """ max_length = self.truncate_prompt_tokens if max_length is not None and max_length < 0: - max_length = self.max_input_tokens + max_length = self.max_truncation_tokens if max_length is None or max_length >= length: return None @@ -494,15 +524,12 @@ def _token_len_check(self, tokenizer: TokenizerLike | None, tokens: _S) -> _S: # max_input_tokens + 1 (see get_encode_kwargs), so the # actual prompt length could be larger. qualifier = "at least " if token_count == max_input_tokens + 1 else "" - total = token_count + self.max_output_tokens raise VLLMValidationError( f"This model's maximum context length is " - f"{self.max_total_tokens} tokens. However, you requested " - f"{self.max_output_tokens} output tokens and your prompt " - f"contains {qualifier}{token_count} input tokens, " - f"for a total of {qualifier}{total} tokens. " - f"Please reduce the length of the input prompt or the " - f"number of requested output tokens.", + f"{self.max_total_tokens} tokens. However, your prompt " + f"contains {qualifier}{token_count} input tokens, which " + f"exceeds this limit even without any output tokens. " + f"Please reduce the length of the input prompt.", parameter="input_tokens", value=token_count, ) From 9a585c62aade0bfd290dfc7a975d54cab8ebfa54 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:01:16 +0800 Subject: [PATCH 303/952] [Rust Frontend] Box rarely-set fields of per-token output types (#60588) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/event.rs | 6 ++-- rust/src/chat/src/output/mod.rs | 6 ++-- rust/src/chat/src/output/structured.rs | 4 +-- rust/src/chat/src/stream.rs | 4 +-- .../engine-core-client/src/protocol/output.rs | 20 +++++++----- rust/src/llm/src/output.rs | 32 ++++++++++--------- rust/src/llm/src/request_metrics.rs | 4 +-- rust/src/llm/tests/generate.rs | 20 ++++++------ rust/src/mock-engine/src/engine.rs | 2 +- rust/src/mock-engine/src/tests.rs | 4 +-- rust/src/server/src/grpc/convert.rs | 4 +-- rust/src/server/src/grpc/inference.rs | 4 +-- .../server/src/routes/inference/generate.rs | 18 ++++++----- rust/src/server/src/routes/tests.rs | 4 +-- rust/src/text/src/output/decoded.rs | 6 ++-- rust/src/text/src/output/mod.rs | 6 ++-- 16 files changed, 80 insertions(+), 64 deletions(-) diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index e57b8bbb86ab..5fe25a9969ef 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -237,9 +237,11 @@ pub enum ChatEvent { usage: ChatTokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - kv_transfer_params: Option, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 8e62f3e8918b..8c1dae37f4de 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -54,10 +54,12 @@ pub(crate) enum AssistantEvent { usage: ChatTokenUsage, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - kv_transfer_params: Option, + /// Boxed, like `ec_transfer_params`, to keep the per-token variants of + /// this enum small. + kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for /// disaggregated serving. - ec_transfer_params: Option, + ec_transfer_params: Option>, }, } diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 5f5bbf32fa8e..80035db4cb1e 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -150,8 +150,8 @@ impl StructuredEventState { &mut self, usage: ChatTokenUsage, finish_reason: FinishReason, - kv_transfer_params: Option, - ec_transfer_params: Option, + kv_transfer_params: Option>, + ec_transfer_params: Option>, ) -> Result> { let mut events = Vec::new(); self.close_open_text_block(&mut events); diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 72b913078c29..509628e8e2af 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -95,8 +95,8 @@ impl ChatEventStream { token_ids, usage, finish_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(|value| *value), + ec_transfer_params: ec_transfer_params.map(|value| *value), }); } ChatEvent::ToolCallEnd { call, .. } => { diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c33062eeb3c1..9339c117a889 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -82,6 +82,10 @@ pub struct EngineCoreEvent { /// /// Original Python definition: /// +/// +/// Fields that are set on few outputs (first or final output, or opt-in +/// features) are boxed: outputs are moved by value through the dispatcher and +/// the request stream once per engine step, so their inline size matters. #[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)] pub struct EngineCoreOutput { pub request_id: String, @@ -95,7 +99,7 @@ pub struct EngineCoreOutput { #[serde(default)] pub new_prompt_logprobs_tensors: Option, #[serde(default)] - pub pooling_output: Option, + pub pooling_output: Option>, #[serde(default)] pub finish_reason: Option, #[serde(default)] @@ -103,17 +107,17 @@ pub struct EngineCoreOutput { #[serde(default)] pub events: Option>, #[serde(default)] - pub kv_transfer_params: Option, + pub kv_transfer_params: Option>, #[serde(default)] - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, #[serde(default)] - pub trace_headers: Option, + pub trace_headers: Option>, /// Breakdown of the scheduled prefill computation, set on the first output /// of a newly scheduled prefill and elided for subsequent decode outputs. #[serde(default)] - pub prefill_stats: Option, + pub prefill_stats: Option>, #[serde(default)] - pub routed_experts: Option, + pub routed_experts: Option>, /// Number of NaNs seen in logits. Values above zero indicate corruption. #[serde(default)] pub num_nans_in_logits: u32, @@ -133,11 +137,11 @@ pub struct EngineCoreOutput { /// Per-request speculative-decoding acceptance metrics, set on the final /// output when `--per-request-spec-decode-metrics` is enabled. #[serde(default)] - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option>, /// Log probabilities of `SamplingParams.prompt_logprob_token_ids`, set on /// the first output of a request that asked for them. #[serde(default)] - pub prompt_token_id_logprobs: Option, + pub prompt_token_id_logprobs: Option>, } /// Raw per-sequence speculative-decoding accumulator. diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4b6d03353991..4dde96f2e20c 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -152,7 +152,7 @@ pub struct GenerateOutput { pub request_id: String, /// One-time prompt metadata emitted only on the first output for this /// request. - pub prompt_info: Option, + pub prompt_info: Option>, /// Newly produced token IDs for this step. pub token_ids: Vec, /// Sample logprobs for the generated positions in this step. @@ -162,14 +162,14 @@ pub struct GenerateOutput { /// Number of prompt tokens served from cache, when reported by prefill stats. pub cached_token_count: usize, /// Connector-specific KV transfer parameters for disaggregated serving. - pub kv_transfer_params: Option, + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned one-to-one with `token_ids`. pub sampling_mask: Option, /// Per-request speculative-decoding metrics, present on terminal outputs. - pub spec_decode_metrics: Option, + pub spec_decode_metrics: Option>, } impl GenerateOutput { @@ -207,10 +207,12 @@ impl GenerateOutput { ) -> Self { Self { request_id: String::new(), - prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { - prompt_token_ids: ids, - prompt_logprobs: None, - prompt_token_id_logprobs: None, + prompt_info: prompt_token_ids.map(|ids| { + Box::new(GeneratePromptInfo { + prompt_token_ids: ids, + prompt_logprobs: None, + prompt_token_id_logprobs: None, + }) }), token_ids, logprobs: None, @@ -231,7 +233,7 @@ impl GenerateOutput { /// with `finish_reason = Abort` before the stream ends. /// - For errors or unexpected engine-side closes, the stream terminates with an error. pub struct GenerateOutputStream { - pending_prompt_info: Option, + pending_prompt_info: Option>, raw_stream: EngineCoreOutputStream, request_metrics: RequestMetricsTracker, /// Removes this request's external→internal tracking edge on drop. Held for @@ -249,11 +251,11 @@ impl GenerateOutputStream { request_guard: RequestGuard, ) -> Self { Self { - pending_prompt_info: Some(GeneratePromptInfo { + pending_prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids, prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), raw_stream, request_metrics, _request_guard: request_guard, @@ -291,7 +293,7 @@ impl Stream for GenerateOutputStream { if let Some(info) = &mut self.pending_prompt_info && info.prompt_token_id_logprobs.is_none() { - info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs; + info.prompt_token_id_logprobs = raw.prompt_token_id_logprobs.map(|logprobs| *logprobs); } let logprobs = raw.new_logprobs.map(|value| value.into_direct().unwrap()); @@ -432,9 +434,9 @@ impl> + Send> T { output_token_count: collected.token_ids.len(), cached_token_count, }; - collected.kv_transfer_params = output.kv_transfer_params; - collected.ec_transfer_params = output.ec_transfer_params; - collected.spec_decode_metrics = output.spec_decode_metrics; + collected.kv_transfer_params = output.kv_transfer_params.map(|value| *value); + collected.ec_transfer_params = output.ec_transfer_params.map(|value| *value); + collected.spec_decode_metrics = output.spec_decode_metrics.map(|value| *value); if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 98f9cadd1a35..61cea120940d 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -404,14 +404,14 @@ mod tests { timestamp: 9.0, }, ]), - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 64, num_computed_tokens: 60, num_cached_tokens: 4, num_local_cached_tokens: 4, num_external_cached_tokens: 0, ..Default::default() - }), + })), ..Default::default() }, ); diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index c596c89ba3ba..813171555bba 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -92,8 +92,8 @@ fn request_output_with_logprobs_and_kv( new_logprobs: new_logprobs.map(MaybeWireLogprobs::Direct), new_prompt_logprobs_tensors: prompt_logprobs.map(MaybeWireLogprobs::Direct), finish_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(Box::new), + ec_transfer_params: ec_transfer_params.map(Box::new), ..Default::default() } } @@ -280,11 +280,11 @@ async fn generate_streams_outputs() { assert_eq!(first.request_id, internal_id); assert_eq!( first.prompt_info, - Some(GeneratePromptInfo { + Some(Box::new(GeneratePromptInfo { prompt_token_ids: vec![11, 22].into(), prompt_logprobs: Some(prompt_logprobs()), prompt_token_id_logprobs: None, - }) + })) ); assert_eq!(first.token_ids, vec![1, 2]); assert_eq!( @@ -330,12 +330,12 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { RequestBatchOutputs { outputs: vec![ EngineCoreOutput { - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_cached_tokens: 1, num_local_cached_tokens: 1, ..Default::default() - }), + })), new_sampling_mask: Some(MaybeWireSamplingMask::Direct( SamplingMask { rows: vec![vec![1, 33, 99]], @@ -409,11 +409,11 @@ async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { async fn collect_output_rejects_partial_sampling_mask() { let output = vllm_llm::GenerateOutput { request_id: "req-partial-mask".to_string(), - prompt_info: Some(GeneratePromptInfo { + prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids: Arc::from([11_u32, 22]), prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), token_ids: vec![33, 44], logprobs: None, finish_reason: Some(FinishReason::Length), @@ -763,11 +763,11 @@ async fn generate_records_request_metrics_in_prometheus_output() { engine_index: 4, timestamp: 10.0, outputs: vec![EngineCoreOutput { - prefill_stats: Some(PrefillStats { + prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_computed_tokens: 2, ..Default::default() - }), + })), ..request_output_with_events( &request.request_id, vec![1], diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index 4844aba12018..8685357b1896 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -201,7 +201,7 @@ impl ActiveRequest { let finished = self.generated >= self.max_tokens; EngineCoreOutput { - prefill_stats, + prefill_stats: prefill_stats.map(Box::new), ..request_output( self.request_id.clone(), new_token_ids, diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 71b4286033ad..d82ceac083ca 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -123,11 +123,11 @@ async fn chunk_size_one_outputs_one_token_per_update() { assert_eq!(first.finish_reason, None); assert_eq!( first.prefill_stats, - Some(PrefillStats { + Some(Box::new(PrefillStats { num_prompt_tokens: 3, num_computed_tokens: 3, ..Default::default() - }) + })) ); let second = stream.next().await.expect("second").expect("second ok"); assert_eq!(second.new_token_ids.len(), 1); diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index d495fea485c7..8e31bb84d2d3 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -406,8 +406,8 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result { if let Some(metrics) = output.spec_decode_metrics { - spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(metrics)); + spec_decode_metrics = Some(SpeculativeDecodingMetrics::from(*metrics)); } if prompt_tokens.is_none() && let Some(info) = output.prompt_info.as_ref() @@ -632,11 +632,11 @@ mod tests { }), Ok(GenerateOutput { request_id: String::new(), - prompt_info: Some(GeneratePromptInfo { + prompt_info: Some(Box::new(GeneratePromptInfo { prompt_token_ids: Arc::from([11_u32, 22_u32]), prompt_logprobs: None, prompt_token_id_logprobs: None, - }), + })), token_ids: vec![33], logprobs: None, finish_reason: Some(FinishReason::stop_eos()), @@ -716,10 +716,12 @@ mod tests { ) -> GenerateOutput { GenerateOutput { request_id: String::new(), - prompt_info: prompt_token_ids.map(|ids| GeneratePromptInfo { - prompt_token_ids: Arc::from(ids), - prompt_logprobs: None, - prompt_token_id_logprobs: None, + prompt_info: prompt_token_ids.map(|ids| { + Box::new(GeneratePromptInfo { + prompt_token_ids: Arc::from(ids), + prompt_logprobs: None, + prompt_token_id_logprobs: None, + }) }), token_ids, logprobs: None, @@ -817,7 +819,7 @@ mod tests { ..Default::default() }; let mut output = stream_output(None, Vec::new(), Some(FinishReason::stop_eos())); - output.spec_decode_metrics = Some(metrics.clone()); + output.spec_decode_metrics = Some(Box::new(metrics.clone())); for include_usage in [false, true] { let chunks: Vec<_> = generate_chunk_stream( diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 0ed2ea4d4590..a9c9831bbd37 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -120,8 +120,8 @@ fn request_output_with_logprobs_and_kv( new_prompt_logprobs_tensors: new_prompt_logprobs_tensors.map(MaybeWireLogprobs::Direct), finish_reason, stop_reason, - kv_transfer_params, - ec_transfer_params, + kv_transfer_params: kv_transfer_params.map(Box::new), + ec_transfer_params: ec_transfer_params.map(Box::new), ..Default::default() } } diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index d3dc42265641..64df0492bfbe 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -47,10 +47,12 @@ pub struct Finished { pub usage: TokenUsage, pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. - pub kv_transfer_params: Option, + /// Kept boxed as received from `vllm-llm`, so the chat events built from + /// this stay small without boxing again. + pub kv_transfer_params: Option>, /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. - pub ec_transfer_params: Option, + pub ec_transfer_params: Option>, /// Sampling support sets aligned with all generated token positions. pub sampling_mask: Option, } diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 1293ac6025ff..8e09f0761870 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -101,8 +101,10 @@ impl T { let mut collected = collected.unwrap(); collected.finish_reason = finished.finish_reason; collected.usage = finished.usage; - collected.kv_transfer_params = finished.kv_transfer_params; - collected.ec_transfer_params = finished.ec_transfer_params; + collected.kv_transfer_params = + finished.kv_transfer_params.map(|value| *value); + collected.ec_transfer_params = + finished.ec_transfer_params.map(|value| *value); collected.sampling_mask = finished.sampling_mask; return Ok(collected); } From ce4603cad676af27f08f39b696a31a4022eb9ca4 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:04:00 +0800 Subject: [PATCH 304/952] [Bugfix][Rust Frontend] Pass through raw characters in byte-level decode (#59861) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/tokenizer/src/byte_level_decode.rs | 54 ++++++++++++++++----- rust/src/tokenizer/src/hf.rs | 44 +++++++++++++++++ 2 files changed, 86 insertions(+), 12 deletions(-) diff --git a/rust/src/tokenizer/src/byte_level_decode.rs b/rust/src/tokenizer/src/byte_level_decode.rs index ad71f6decbbf..ef7a3c1ef674 100644 --- a/rust/src/tokenizer/src/byte_level_decode.rs +++ b/rust/src/tokenizer/src/byte_level_decode.rs @@ -7,14 +7,16 @@ /// Reverse GPT-2 byte-to-unicode mapping: codepoint → original byte. The GPT-2 /// table only emits codepoints in U+0000..U+0143, so a flat array suffices. -const CHAR_TO_BYTE: [u8; 324] = build_char_to_byte(); +/// Codepoints in that range the table never emits, such as raw whitespace in +/// an added token's content, map to `None`. +const CHAR_TO_BYTE: [Option; 324] = build_char_to_byte(); const fn is_nice(b: u8) -> bool { (b >= b'!' && b <= b'~') || (b >= 0xA1 && b <= 0xAC) || b >= 0xAE } -const fn build_char_to_byte() -> [u8; 324] { - let mut table = [0u8; 324]; +const fn build_char_to_byte() -> [Option; 324] { + let mut table = [None; 324]; let mut b: u16 = 0; while b < 256 { let cp = if is_nice(b as u8) { @@ -22,7 +24,7 @@ const fn build_char_to_byte() -> [u8; 324] { } else { 256 + nice_offset(b as u8) }; - table[cp as usize] = b as u8; + table[cp as usize] = Some(b as u8); b += 1; } table @@ -41,20 +43,23 @@ const fn nice_offset(b: u8) -> u32 { } /// Decode byte-level encoded token strings into a single UTF-8 string, -/// matching `fastokens::decoders::ByteLevelDecoder`. +/// matching the HuggingFace `tokenizers` `ByteLevel` decoder. pub fn decode_byte_level<'a, I: IntoIterator>(tokens: I) -> String { let iter = tokens.into_iter(); let (lower, _) = iter.size_hint(); let mut bytes: Vec = Vec::with_capacity(lower.saturating_mul(4)); for token in iter { + let start = bytes.len(); for c in token.chars() { - let cp = c as usize; - if cp < CHAR_TO_BYTE.len() { - bytes.push(CHAR_TO_BYTE[cp]); + if let Some(&Some(b)) = CHAR_TO_BYTE.get(c as usize) { + bytes.push(b); } else { - // Non-GPT2 codepoints (e.g. DeepSeek's U+FF5C, U+2581) pass through. - let mut buf = [0u8; 4]; - bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes()); + // As in HF, a token with any non-GPT2 codepoint is raw text and + // passes through whole: DeepSeek's U+FF5C and U+2581, or raw + // characters in added tokens (e.g. `\n`). + bytes.truncate(start); + bytes.extend_from_slice(token.as_bytes()); + break; } } } @@ -87,7 +92,7 @@ mod tests { for b in 0..=255u8 { let cp = byte_to_char[b as usize] as usize; assert!(cp < CHAR_TO_BYTE.len()); - assert_eq!(CHAR_TO_BYTE[cp], b, "mismatch for byte {b:#x}"); + assert_eq!(CHAR_TO_BYTE[cp], Some(b), "mismatch for byte {b:#x}"); } } @@ -119,4 +124,29 @@ mod tests { let tok = "<\u{FF5C}begin\u{2581}of\u{2581}sentence\u{FF5C}>"; assert_eq!(decode_byte_level([tok]), "<｜begin▁of▁sentence｜>"); } + + #[test] + fn decode_preserves_raw_chars_outside_gpt2_table() { + // Added tokens store raw text, which may contain codepoints below + // U+0144 that the GPT-2 table never emits. + let toks = [ + "", + "\n", + "\t\u{7F}\u{A0}\u{AD}", + ]; + assert_eq!( + decode_byte_level(toks), + "\n\t\u{7F}\u{A0}\u{AD}", + ); + } + + #[test] + fn decode_passes_through_whole_token_with_raw_chars() { + // Characters that the GPT-2 table does emit (`Ġ`, `é`) stay verbatim + // in a raw token instead of decoding to bytes 0x20 and 0xE9. + assert_eq!(decode_byte_level(["\u{120}x\n"]), "\u{120}x\n"); + assert_eq!(decode_byte_level(["\ncafé"]), "\ncafé"); + } } diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index 44ddc6cb27cb..3ae1779b42cc 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -554,6 +554,50 @@ mod tests { } } + /// Added tokens store raw text, including whitespace that the GPT-2 byte + /// table never emits, so a byte-level decode must pass it through verbatim, + /// along with characters the table does emit in the same token (`é`). + #[test] + fn added_tokens_with_raw_whitespace_decode_verbatim() { + const OPEN: &str = ">() + ); + let dir = tempdir().expect("create temp dir"); + let path = write_tokenizer_json(dir.path(), "tokenizer.json", &value); + let text = format!("{OPEN}city\">\nParis{CLOSE}{MIXED}"); + + let fastokens = HuggingFaceTokenizer::new_fastokens(&path).expect("load fastokens wrapper"); + assert!(matches!( + fastokens.backend, + super::Backend::FastokensByteLevel(_) + )); + let hf = HuggingFaceTokenizer::new_hf(&path).expect("load hf wrapper"); + for wrapper in [fastokens, hf] { + let ids = wrapper.encode(&text, false).expect("encode"); + assert!((256..=258).all(|id| ids.contains(&id)), "ids={ids:?}"); + assert_eq!(wrapper.decode(&ids, true).expect("decode"), text); + } + } + #[test] fn hf_vocab_size_counts_added_tokens() { let mut tokenizer = tiny_bpe_tokenizer(); From 5dc9bcd8cd9656b852edfd6c5b4a8c9c46a34530 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 02:04:49 +0800 Subject: [PATCH 305/952] [Rust Frontend] Pass reasoning controls through render -> generate (#60546) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../src/protocol/request.rs | 1 + .../src/routes/inference/generate/convert.rs | 52 ++++++++++++++++++- .../src/routes/inference/generate/types.rs | 8 +++ rust/src/server/src/routes/render.rs | 7 +++ rust/src/server/src/routes/tests.rs | 37 +++++++++---- 5 files changed, 92 insertions(+), 13 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs index 02fc1c30cad6..5390eed16f2c 100644 --- a/rust/src/engine-core-client/src/protocol/request.rs +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -64,6 +64,7 @@ impl EngineCoreRequestType { #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] pub struct ReasoningParserKwargs { /// Effective kwargs visible to the chat template for this request. + #[serde(default)] pub chat_template_kwargs: HashMap, } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 36a497f6ddf3..313e9fa1f955 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -85,8 +85,8 @@ pub(super) fn prepare_generate_request( data_parallel_rank: ctx.data_parallel_rank, session_id: ctx.session_id, kv_hints: None, - reasoning_parser_kwargs: Default::default(), - reasoning_ended: None, + reasoning_parser_kwargs: request.reasoning_parser_kwargs.unwrap_or_default(), + reasoning_ended: request.reasoning_ended, lora_request: lora_resolution.lora_request.clone(), arrival_time: None, }; @@ -224,6 +224,54 @@ mod tests { ); } + #[test] + fn prepare_generate_request_forwards_reasoning_controls() { + for (controls, expected) in [ + (json!({}), (json!({"chat_template_kwargs": {}}), None)), + ( + json!({ + "reasoning_ended": false, + "reasoning_parser_kwargs": { + "chat_template_kwargs": {"enable_thinking": true} + } + }), + ( + json!({"chat_template_kwargs": {"enable_thinking": true}}), + Some(false), + ), + ), + ( + json!({"reasoning_ended": true, "reasoning_parser_kwargs": {}}), + (json!({"chat_template_kwargs": {}}), Some(true)), + ), + ] { + let mut body = json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22, 33], + "sampling_params": {} + }); + body.as_object_mut().unwrap().extend(controls.as_object().unwrap().clone()); + let request: GenerateRequest = serde_json::from_value(body).expect("parse request"); + + let prepared = prepare_generate_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + None, + ) + .expect("prepare"); + + assert_eq!( + ( + serde_json::to_value(&prepared.text_request.reasoning_parser_kwargs).unwrap(), + prepared.text_request.reasoning_ended, + ), + expected, + "controls: {controls}" + ); + } + } + #[test] fn prepare_generate_request_gates_continuous_usage_on_include_usage() { let request: GenerateRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index bd277ef34a0a..0a37640d63a1 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{Normalizable, StreamOptions, Usage}; @@ -48,6 +49,13 @@ pub struct GenerateRequest { /// Raw multimodal input; server resolves media. Mutually exclusive with `features`. pub content_parts: Option>, pub return_token_ids: Option, + /// Whether reasoning has ended before the first generated token, as + /// resolved by `/render`. `true` applies structured outputs from the first + /// token; `None` lets the engine check the prompt with its reasoning parser. + pub reasoning_ended: Option, + /// Set by `/render` so the engine-side reasoning parser agrees with the + /// frontend on flags such as `enable_thinking`. + pub reasoning_parser_kwargs: Option, #[serde(flatten)] pub other: Map, } diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index 2a6f70ec4b6e..aceb28f4993f 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -10,6 +10,7 @@ use axum::extract::State; use axum::http::{HeaderMap, StatusCode}; use axum::routing::{get, post}; use thiserror_ext::AsReport as _; +use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_text::TextRequest; use crate::DEFAULT_REQUEST_BODY_LIMIT_BYTES; @@ -106,6 +107,12 @@ fn lower_render_request( ec_transfer_params: None, content_parts: None, return_token_ids: None, + // Carry what the chat route would pass to the engine, so the engine + // gates structured outputs the same way on `/inference/v1/generate`. + reasoning_ended: text_request.reasoning_ended, + reasoning_parser_kwargs: (text_request.reasoning_parser_kwargs + != ReasoningParserKwargs::default()) + .then_some(text_request.reasoning_parser_kwargs), other: Default::default(), }; validate_generate_request(&request, &state.served_model_names)?; diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index a9c9831bbd37..118e412b45b6 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -542,7 +542,7 @@ impl ChatRenderer for FakeChatBackend { Ok(vllm_chat::RenderedPrompt { prompt: Prompt::Text(prompt), media_order: None, - effective_template_kwargs: Default::default(), + effective_template_kwargs: request.chat_options.template_kwargs.clone(), }) } } @@ -615,7 +615,7 @@ fn qwen_multimodal_model_info_with_limits( async fn test_models_with_engine_outputs_and_backend_inner( engine_id: impl Into, output_specs: Vec<(Vec, Option)>, - expected_prompt_token_ids: Option>, + check_request: impl FnOnce(&EngineCoreRequest) + Send + 'static, backend: Arc, ) -> (ChatLlm, MockEngineTask) { let ipc = IpcNamespace::new().expect("create ipc namespace"); @@ -630,12 +630,7 @@ async fn test_models_with_engine_outputs_and_backend_inner( let add = recv_engine_message(dealer).await; let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).expect("decode request"); - if let Some(expected_prompt_token_ids) = expected_prompt_token_ids { - assert_eq!( - request.prompt_token_ids.as_deref(), - Some(expected_prompt_token_ids.as_slice()) - ); - } + check_request(&request); send_outputs( push, engine_outputs_for_request(&request.request_id, output_specs), @@ -667,7 +662,8 @@ async fn test_models_with_engine_outputs_and_backend( output_specs: Vec<(Vec, Option)>, backend: Arc, ) -> (ChatLlm, MockEngineTask) { - test_models_with_engine_outputs_and_backend_inner(engine_id, output_specs, None, backend).await + test_models_with_engine_outputs_and_backend_inner(engine_id, output_specs, |_| {}, backend) + .await } async fn test_chat_with_engine_outputs( @@ -1440,6 +1436,8 @@ async fn render_completion_returns_generate_request_with_body_request_id() { assert!(json[0].get("stream_options").is_none()); assert!(json[0].get("prompt_token_ids").is_none()); assert!(json[0].get("mm_features").is_none()); + assert!(json[0].get("reasoning_ended").is_none()); + assert!(json[0].get("reasoning_parser_kwargs").is_none()); } #[tokio::test] @@ -1479,7 +1477,8 @@ async fn render_chat_response_can_be_submitted_to_generate_unchanged() { "messages": [{"role": "user", "content": "hello"}], "max_completion_tokens": 8, "bad_words": ["blocked"], - "vllm_xargs": {"custom": 1} + "vllm_xargs": {"custom": 1}, + "chat_template_kwargs": {"enable_thinking": false} }) .to_string(), )) @@ -1502,11 +1501,27 @@ async fn render_chat_response_can_be_submitted_to_generate_unchanged() { json!(["blocked"]) ); assert_eq!(render_json["sampling_params"]["vllm_xargs"]["custom"], 1); + assert_eq!( + render_json["reasoning_parser_kwargs"], + json!({"chat_template_kwargs": {"enable_thinking": false}}) + ); let (chat, engine_task) = test_models_with_engine_outputs_and_backend_inner( b"engine-render-generate-round-trip", default_stream_output_specs(), - Some(token_ids), + move |request| { + assert_eq!( + request.prompt_token_ids.as_deref(), + Some(token_ids.as_slice()) + ); + assert_eq!( + request + .reasoning_parser_kwargs + .as_ref() + .map(|kwargs| &kwargs.chat_template_kwargs["enable_thinking"]), + Some(&json!(false)) + ); + }, Arc::new(FakeChatBackend::new()), ) .await; From c811943137dbfecdf9c7af849bc66270917a9d79 Mon Sep 17 00:00:00 2001 From: premsurawut <92110316+surawut-jirasaktavee@users.noreply.github.com> Date: Sat, 10 Oct 2026 01:06:26 +0700 Subject: [PATCH 306/952] [Bugfix][Frontend] Reject beam search with echo and logprobs in /v1/completions (#60876) Signed-off-by: surawut.j Co-authored-by: Claude Opus 5.5 --- .../openai/completion/test_completion_error.py | 14 ++++++++++++++ vllm/entrypoints/openai/completion/protocol.py | 11 +++++++++++ 2 files changed, 25 insertions(+) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 4056ea49d6d3..bbbb61f9c9c8 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -852,6 +852,20 @@ def test_logprobs_below_minus_one_rejected(): ) +@pytest.mark.parametrize("max_tokens", [5, 0]) +def test_beam_search_echo_logprobs_rejected(max_tokens): + """Beam search does not compute prompt logprobs, so echo + logprobs is rejected.""" + with pytest.raises(VLLMValidationError, match="not supported with beam search"): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=max_tokens, + use_beam_search=True, + echo=True, + logprobs=1, + ) + + class TestCompletionPromptListLimit: """Regression tests for CVE: unbounded prompt list fan-out.""" diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index f26dea2e1f9e..e99de1e28625 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -492,6 +492,17 @@ def check_logprobs(cls, data): parameter="logprob_token_ids", ) + if ( + data.get("use_beam_search") + and data.get("echo") + and data.get("logprobs") is not None + ): + raise VLLMValidationError( + "`echo` with `logprobs` is not supported with beam search " + "because beam search does not compute prompt logprobs.", + parameter="echo", + ) + if ( data.get("logprob_token_ids") and data.get("echo") From a0f28d73a83bbf0491428efb0f11c2b8f9dc9067 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Sat, 10 Oct 2026 02:08:31 +0800 Subject: [PATCH 307/952] [sharded state loader] support pp in sharded state loader (#52110) Signed-off-by: Andy Xie Co-authored-by: Simon Mo --- .../save_sharded_state_offline.py | 7 +- .../model_loader/test_sharded_state_loader.py | 438 ++++++++++++++++++ vllm/distributed/parallel_state.py | 5 + .../model_loader/sharded_state_loader.py | 152 ++++-- 4 files changed, 572 insertions(+), 30 deletions(-) diff --git a/examples/features/sharded_state/save_sharded_state_offline.py b/examples/features/sharded_state/save_sharded_state_offline.py index 7ed0de1f1054..44b68cecf0e7 100644 --- a/examples/features/sharded_state/save_sharded_state_offline.py +++ b/examples/features/sharded_state/save_sharded_state_offline.py @@ -38,8 +38,11 @@ def parse_args(): parser.add_argument( "--file-pattern", type=str, - default=ShardedStateLoader.DEFAULT_PATTERN, - help="string pattern of saved filenames", + help=f"string pattern of saved filenames." + f"{ShardedStateLoader.TP_ONLY_PATTERN!r} for tensor parallel only " + f"(pipeline parallelism is 1). " + f"{ShardedStateLoader.PP_AND_TP_PATTERN!r} for pipeline and " + f"tensor parallel.", ) parser.add_argument( "--max-file-size", diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py index 06354fe3c5ea..6ec5f1ce43ab 100644 --- a/tests/model_executor/model_loader/test_sharded_state_loader.py +++ b/tests/model_executor/model_loader/test_sharded_state_loader.py @@ -6,6 +6,7 @@ import os import shutil from tempfile import TemporaryDirectory +from types import SimpleNamespace import pytest import torch @@ -49,6 +50,318 @@ def test_filter_subtensors(): assert tensor is state_dict[key] +def _mock_pp_group( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + *, + pp_rank: int = 0, + tp_rank: int = 0, + tp_world_size: int = 1, +) -> None: + """Stub the parallel-state accessors used by the sharded state loader.""" + monkeypatch.setattr( + "vllm.distributed.get_pp_group", + lambda: SimpleNamespace(world_size=world_size, rank_in_group=pp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tp_group", + lambda: SimpleNamespace(world_size=tp_world_size, rank_in_group=tp_rank), + ) + monkeypatch.setattr( + "vllm.distributed.get_tensor_model_parallel_rank", + lambda: tp_rank, + ) + monkeypatch.setattr( + "vllm.distributed.get_pipeline_model_parallel_rank", + lambda: pp_rank, + ) + + +@pytest.mark.parametrize( + "world_size,expected", + [ + (1, ShardedStateLoader.TP_ONLY_PATTERN), + (2, ShardedStateLoader.PP_AND_TP_PATTERN), + ], +) +def test_get_default_pattern(monkeypatch, world_size, expected): + """The default pattern is chosen from the pipeline parallel group size.""" + _mock_pp_group(monkeypatch, world_size) + assert ShardedStateLoader._get_default_pattern() == expected + + +def test_render_default_pattern_default_tp(monkeypatch): + """Without a custom pattern and PP=1, legacy TP-only names are rendered.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_default_pp(monkeypatch): + """Without a custom pattern and PP>1, pp/tp-aware names are rendered.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern(part=0) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +def test_render_default_pattern_custom_tp_only(monkeypatch): + """A custom TP-only pattern is honored when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + == "model-rank-0-part-0.safetensors" + ) + + +def test_render_default_pattern_tp_only_with_pp_fails(monkeypatch): + """TP-only filenames would collide across pipeline stages; fail fast.""" + _mock_pp_group(monkeypatch, 2) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.TP_ONLY_PATTERN, part=0 + ) + + +def test_render_default_pattern_pp_pattern_without_pp_fails(monkeypatch): + """A pp/tp-aware pattern cannot be rendered when pipeline parallelism is off.""" + _mock_pp_group(monkeypatch, 1) + with pytest.raises(ValueError, match="requires pipeline parallelism"): + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + + +def test_render_default_pattern_custom_pp(monkeypatch): + """A custom pp/tp-aware pattern is honored when pipeline parallelism is on.""" + _mock_pp_group(monkeypatch, 2) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == "model-pp-0-tp-0-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # tp field present but pp field missing: pipeline stages with the + # same tp rank would write the same file. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # pp field present but tp field missing: TP workers would overwrite + # each other's shards. + ( + "model-pp-{pp_rank}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_incomplete_custom_patterns_fail( + monkeypatch, pattern, match +): + """A custom pattern must distinguish every active rank (pp and tp).""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +def test_render_default_pattern_no_rank_field_with_tp_fails(monkeypatch): + """Without any rank field, TP workers would collide even when PP is off.""" + _mock_pp_group(monkeypatch, 1, tp_world_size=2) + with pytest.raises(ValueError, match="does not embed the tensor parallel"): + ShardedStateLoader._render_default_pattern( + "model-part-{part}.safetensors", part=0 + ) + + +def test_render_default_pattern_custom_pp_tp1_allows_pp_only(monkeypatch): + """With TP=1 a pp-only custom pattern is safe and must be accepted.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=1) + assert ( + ShardedStateLoader._render_default_pattern( + "model-pp-{pp_rank}-part-{part}.safetensors", part=0 + ) + == "model-pp-1-part-0.safetensors" + ) + + +@pytest.mark.parametrize( + "pattern,match", + [ + # {{pp_rank}} renders as literal text: both stages would write the + # same file, so it must not count as a pp field. + ( + "model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + "does not embed the pipeline parallel", + ), + # Same for an escaped tp field: TP workers would overwrite each + # other's shards. + ( + "model-pp-{pp_rank}-tp-{{tp_rank}}-part-{part}.safetensors", + "does not embed the tensor parallel", + ), + ], +) +def test_render_default_pattern_escaped_field_not_rank_field( + monkeypatch, pattern, match +): + """Escaped braces are literal text and cannot pose as rank fields.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + with pytest.raises(ValueError, match=match): + ShardedStateLoader._render_default_pattern(pattern, part=0) + + +@pytest.mark.parametrize( + "pattern,world_size,tp_world_size,pp_rank,tp_rank,expected", + [ + # Mixed tp alias works when pipeline parallelism is on. + ( + "model-pp-{pp_rank}-rank-{rank}-part-{part}.safetensors", + 2, + 2, + 1, + 0, + "model-pp-1-rank-0-part-0.safetensors", + ), + # tp_rank alias works for TP-only checkpoints too. + ( + "model-tp-{tp_rank}-part-{part}.safetensors", + 1, + 2, + 0, + 3, + "model-tp-3-part-0.safetensors", + ), + ], +) +def test_render_default_pattern_tp_aliases_accepted( + monkeypatch, pattern, world_size, tp_world_size, pp_rank, tp_rank, expected +): + """Both tp aliases are supplied at render time, so either works.""" + _mock_pp_group( + monkeypatch, + world_size, + pp_rank=pp_rank, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + assert ShardedStateLoader._render_default_pattern(pattern, part=0) == expected + + +def test_render_default_pattern_unsupported_field_fails(monkeypatch): + """Unknown placeholders are rejected instead of KeyError-ing later.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_world_size=2) + with pytest.raises(ValueError, match="unsupported placeholder"): + ShardedStateLoader._render_default_pattern( + "model-{pp_rank}-{tp_rank}-{epoch}-part-{part}.safetensors", part=0 + ) + + +@pytest.mark.parametrize( + "pp_rank,tp_rank,expected", + [ + (0, 0, "model-pp-0-tp-0-part-0.safetensors"), + (1, 2, "model-pp-1-tp-2-part-0.safetensors"), + (3, 5, "model-pp-3-tp-5-part-0.safetensors"), + ], +) +def test_render_default_pattern_non_zero_ranks(monkeypatch, pp_rank, tp_rank, expected): + """The rendered filename must embed the actual pp/tp ranks, not zeros.""" + _mock_pp_group(monkeypatch, 2, pp_rank=pp_rank, tp_rank=tp_rank) + assert ( + ShardedStateLoader._render_default_pattern( + ShardedStateLoader.PP_AND_TP_PATTERN, part=0 + ) + == expected + ) + + +def test_render_default_pattern_part_wildcard(monkeypatch): + """The glob wildcard '*' passes through as the part token for file lookup.""" + _mock_pp_group(monkeypatch, 1, tp_rank=4) + assert ( + ShardedStateLoader._render_default_pattern(part="*") + == "model-rank-4-part-*.safetensors" + ) + + +def test_save_model_pp_aware_filenames(monkeypatch, tmp_path): + """save_model writes pp/tp-aware filenames for pipeline-parallel models.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=2) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + ShardedStateLoader.save_model(model, str(tmp_path)) + assert saved == ["model-pp-1-tp-2-part-0.safetensors"] + + +def test_save_model_part_split(monkeypatch, tmp_path): + """save_model rolls over to a new part when max_size is exceeded.""" + _mock_pp_group(monkeypatch, 1, tp_rank=0) + + saved: list[str] = [] + + def _capture(tensors, path, **kwargs): + saved.append(os.path.basename(path)) + + monkeypatch.setattr("safetensors.torch.save_file", _capture) + + # Two float32[4] tensors (16 B each); max_size=20 fits one but not two. + model = SimpleNamespace( + state_dict=lambda: { + "a": torch.zeros(4, dtype=torch.float32), + "b": torch.zeros(4, dtype=torch.float32), + } + ) + ShardedStateLoader.save_model(model, str(tmp_path), max_size=20) + assert saved == [ + "model-rank-0-part-0.safetensors", + "model-rank-0-part-1.safetensors", + ] + + +def test_save_model_rejects_bad_pattern_before_overwrite(monkeypatch, tmp_path): + """An unsafe pattern is rejected before any shard file is written.""" + _mock_pp_group(monkeypatch, 2, pp_rank=1, tp_rank=0, tp_world_size=2) + + # With the escaped pp field this pattern renders both stages to the + # same literal name; the validation must fire before any write. + colliding = tmp_path / "model-{pp_rank}-tp-0-part-0.safetensors" + colliding.write_bytes(b"existing shard") + before = sorted(p.name for p in tmp_path.iterdir()) + + model = SimpleNamespace( + state_dict=lambda: {"w": torch.zeros(4, dtype=torch.float32)} + ) + with pytest.raises(ValueError, match="does not embed the pipeline parallel"): + ShardedStateLoader.save_model( + model, + str(tmp_path), + pattern="model-{{pp_rank}}-tp-{tp_rank}-part-{part}.safetensors", + ) + + # The pre-existing shard is untouched and nothing new was written. + assert colliding.read_bytes() == b"existing shard" + assert sorted(p.name for p in tmp_path.iterdir()) == before + + @pytest.fixture(scope="module") def llama_3p2_1b_files(): input_dir = hf_api().snapshot_download( @@ -163,3 +476,128 @@ def test_sharded_state_loader( queue.join_thread() assert out_before == out_after + + +def _run_sharded_state_round_trip( + input_dir: str, + *, + tp_size: int, + pp_size: int, + num_gpus_available: int, +) -> None: + """Save a TP/PP sharded checkpoint, reload it, and compare outputs. + + Args: + input_dir: directory containing the original (un-sharded) model. + tp_size: tensor parallel size used for both save and load. + pp_size: pipeline parallel size used for both save and load. + num_gpus_available: number of GPUs on the current machine. + + """ + if num_gpus_available < tp_size * pp_size: + pytest.skip(f"Not enough GPUs for TP={tp_size}, PP={pp_size}") + + weights_patterns = ("*.safetensors",) + gpu_memory_utilization = 0.8 + ctx = mp.get_context("spawn") + + # Keep batching deterministic: this test compares exact greedy outputs + # across separate engine processes, whose schedulers may otherwise form + # different batches depending on prompt-processing timing. + platform_args = {"max_num_seqs": 1} + parallel_args = dict( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + ) + + with TemporaryDirectory() as output_dir: + p = ctx.Process( + target=_run_writer, + args=(input_dir, output_dir, weights_patterns), + kwargs=dict( + **parallel_args, + gpu_memory_utilization=gpu_memory_utilization, + enforce_eager=True, + **platform_args, + ), + ) + p.start() + p.join() + assert p.exitcode == 0, f"writer process failed with exit code {p.exitcode}" + + if pp_size > 1: + # PP checkpoints must use the pp/tp-aware naming scheme so that + # different pipeline stages write distinct files. + sharded_files = [ + f for f in os.listdir(output_dir) if f.endswith(".safetensors") + ] + assert sharded_files, "no sharded checkpoint files were written" + assert any( + fnmatch.fnmatch(f, "model-pp-*-tp-*-part-*.safetensors") + for f in sharded_files + ), f"expected PP-aware filenames, got {sharded_files}" + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(input_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + **parallel_args, + **platform_args, + ), + ) + p.start() + # Call queue.get() before p.join() to prevent deadlock: + # If p.join() is called before queue.get() and the queue is full, + # the child process may block while writing to the queue and never + # terminate, causing the parent to wait indefinitely on p.join(). + # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814 + out_before = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + queue = ctx.Queue() + p = ctx.Process( + target=_run_generate, + args=(output_dir, queue), + kwargs=dict( + enable_lora=False, + gpu_memory_utilization=gpu_memory_utilization, + load_format="sharded_state", + **parallel_args, + **platform_args, + ), + ) + p.start() + out_after = queue.get() + p.join() + assert p.exitcode == 0, f"generate process failed with exit code {p.exitcode}" + queue.close() + queue.join_thread() + + assert out_before == out_after + + +def test_sharded_state_loader_pp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline parallelism (PP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=1, + pp_size=2, + num_gpus_available=num_gpus_available, + ) + + +def test_sharded_state_loader_pp_tp(num_gpus_available, llama_3p2_1b_files): + """Round-trip sharded state save/load with pipeline + tensor parallelism + (PP=2, TP=2).""" + _run_sharded_state_round_trip( + llama_3p2_1b_files, + tp_size=2, + pp_size=2, + num_gpus_available=num_gpus_available, + ) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index feb8f25d1eeb..fea0c4e3c783 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -2444,6 +2444,11 @@ def get_tensor_model_parallel_rank() -> int: return get_tp_group().rank_in_group +def get_pipeline_model_parallel_rank() -> int: + """Return my rank for the pipeline model parallel group.""" + return get_pp_group().rank_in_group + + def get_node_count() -> int: """Return the total number of nodes in the distributed environment.""" assert _NODE_COUNT is not None, "distributed environment is not initialized" diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index d17bfdc6d495..97a31e8e879a 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -4,6 +4,7 @@ import collections import glob import os +import string import time from collections.abc import Generator from copy import copy @@ -28,13 +29,14 @@ class ShardedStateLoader(BaseModelLoader): """Model loader that directly loads each worker's model state dict, which - enables a fast load path for large tensor-parallel models where each worker - only needs to read its own shard rather than the entire checkpoint. See - `examples/features/sharded_state/save_sharded_state_offline.py` for creating - a sharded checkpoint. + enables a fast load path for large tensor-parallel/pipeline-parallel models + where each worker only needs to read its own shard rather than the entire + checkpoint. See `examples/features/sharded_state/save_sharded_state_offline.py` + for creating a sharded checkpoint. """ - DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + TP_ONLY_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + PP_AND_TP_PATTERN = "model-pp-{pp_rank}-tp-{tp_rank}-part-{part}.safetensors" def __init__(self, load_config: LoadConfig): super().__init__(load_config) @@ -44,14 +46,94 @@ def __init__(self, load_config: LoadConfig): if load_config.model_loader_extra_config is None else copy(load_config.model_loader_extra_config) ) - self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN) + # A custom shard filename pattern can be supplied via + # model_loader_extra_config; when omitted (None), the default pattern + # is selected at render time based on the parallel configuration. + self.pattern = extra_config.pop("pattern", None) if extra_config: raise ValueError( f"Unexpected extra config keys for load format " - f"{load_config.load_format}: " - f"{load_config.model_loader_extra_config.keys()}" + f"{load_config.load_format}: {extra_config.keys()}" ) + @staticmethod + def _get_default_pattern() -> str: + from vllm.distributed import get_pp_group + + return ( + ShardedStateLoader.PP_AND_TP_PATTERN + if get_pp_group().world_size > 1 + else ShardedStateLoader.TP_ONLY_PATTERN + ) + + @staticmethod + def _render_default_pattern(pattern: str | None = None, *, part: int | str) -> str: + from vllm.distributed import ( + get_pipeline_model_parallel_rank, + get_pp_group, + get_tensor_model_parallel_rank, + get_tp_group, + ) + + # _get_default_pattern is called here because _render_default_pattern + # can be invoked from save_model without instantiating a + # ShardedStateLoader instance. + if pattern is None: + pattern = ShardedStateLoader._get_default_pattern() + + # Parse the real format fields instead of matching placeholder + # text: escaped braces like {{pp_rank}} render as literal text + # and must not count as a rank field. + fields = { + field + for _, field, _, _ in string.Formatter().parse(pattern) + if field is not None + } + unsupported = fields - {"pp_rank", "tp_rank", "rank", "part"} + if unsupported: + raise ValueError( + f"Pattern '{pattern}' uses unsupported placeholder(s) " + f"{sorted(unsupported)}; supported fields are pp_rank, " + "tp_rank, rank and part" + ) + tp_rank = get_tensor_model_parallel_rank() + pp_rank = get_pipeline_model_parallel_rank() + pp_world_size = get_pp_group().world_size + tp_world_size = get_tp_group().world_size + # The pattern must carry a rank field for every active + # parallelism dimension, otherwise shards from different ranks + # collide on the same filename and silently overwrite each other. + has_pp_field = "pp_rank" in fields + has_tp_field = "tp_rank" in fields or "rank" in fields + if pp_world_size > 1: + if not has_pp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the pipeline parallel " + "rank; with pipeline parallelism every stage would write the " + f"same file. Use '{ShardedStateLoader.PP_AND_TP_PATTERN}' " + "instead." + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel " + "rank; TP workers would overwrite each other's shards." + ) + # Supply both tp aliases so mixed-alias patterns render. + return pattern.format( + pp_rank=pp_rank, tp_rank=tp_rank, rank=tp_rank, part=part + ) + if has_pp_field: + raise ValueError( + f"Pattern '{pattern}' requires pipeline parallelism, but the " + "pipeline parallel size is 1" + ) + if tp_world_size > 1 and not has_tp_field: + raise ValueError( + f"Pattern '{pattern}' does not embed the tensor parallel rank; " + "TP workers would overwrite each other's shards." + ) + return pattern.format(tp_rank=tp_rank, rank=tp_rank, part=part) + @staticmethod def _filter_subtensors( tensors: dict[str, torch.Tensor], @@ -106,30 +188,36 @@ def download_model(self, model_config: ModelConfig) -> None: self._prepare_weights(model_config.model, model_config.revision) def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: - from vllm.distributed import get_tensor_model_parallel_rank - model_weights = model_config.model if model_weights_override := model_config.model_weights: model_weights = model_weights_override local_model_path = model_weights - rank = get_tensor_model_parallel_rank() - pattern = os.path.join( - local_model_path, - self.pattern.format(rank=rank, part="*"), + rendered_pattern = ShardedStateLoader._render_default_pattern( + self.pattern, part="*" ) - filepaths = [] if is_s3(local_model_path): - file_pattern = f"*{self.pattern.format(rank=rank, part='*')}" + file_pattern = f"*{rendered_pattern}" filepaths = s3_glob(path=local_model_path, allow_pattern=[file_pattern]) else: - filepaths = glob.glob(pattern) + filepaths = glob.glob( + os.path.join( + local_model_path, + rendered_pattern, + ) + ) if not filepaths: # TODO: support un-sharded checkpoints too raise ValueError( - f"Could not find checkpoint files '{pattern}', only " - f"pre-sharded checkpoints are currently supported!" + f"Could not find checkpoint files " + f"'{ + os.path.join( + local_model_path, + rendered_pattern, + ) + }', " + f"only pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) counter_before_loading_weights = time.perf_counter() @@ -137,8 +225,17 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # If loading with LoRA enabled, additional padding may # be added to certain parameters. We only load into a # narrowed view of the parameter data. - param_data = state_dict[key].data - param_shape = state_dict[key].shape + param = state_dict[key] + if param.dtype != tensor.dtype: + logger.warning( + "Attempted to load weight %s with dtype %s into " + "parameter with dtype %s", + key, + tensor.dtype, + param.dtype, + ) + param_data = param.data + param_shape = param.shape for dim, size in enumerate(tensor.shape): if size < param_shape[dim]: param_data = param_data.narrow(dim, 0, size) @@ -182,11 +279,6 @@ def save_model( ) -> None: from safetensors.torch import save_file - from vllm.distributed import get_tensor_model_parallel_rank - - if pattern is None: - pattern = ShardedStateLoader.DEFAULT_PATTERN - rank = get_tensor_model_parallel_rank() part_idx = 0 total_size = 0 state_dict = ShardedStateLoader._filter_subtensors(model.state_dict()) @@ -194,7 +286,9 @@ def save_model( for key, tensor in state_dict.items(): param_size = tensor.nelement() * tensor.element_size() if max_size is not None and total_size + param_size > max_size: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename), @@ -205,7 +299,9 @@ def save_model( state_dict_part[key] = tensor total_size += param_size if len(state_dict_part) > 0: - filename = pattern.format(rank=rank, part=part_idx) + filename = ShardedStateLoader._render_default_pattern( + pattern, part=part_idx + ) save_file( state_dict_part, os.path.join(path, filename), From 1dc60dd36734f4436dbc438265bbd9c24164cbad Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 9 Oct 2026 14:42:41 -0400 Subject: [PATCH 308/952] [CI] Pin pydantic in the mypy pre-commit hook (#60894) Signed-off-by: mgoin Co-authored-by: Claude --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 92963cf69a22..8e257edccc12 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -205,7 +205,7 @@ repos: language: python types_or: [python, pyi] require_serial: true - additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic] + additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, "pydantic==2.13.5"] - id: png-lint name: Lint PNG exports from excalidraw entry: tools/pre_commit/png-lint.sh From ad8a6065f4bcbdfafd38828a4b3f1c6fac4dbf89 Mon Sep 17 00:00:00 2001 From: Willow Lopez <100782273+Oxygen56@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:59:02 +0800 Subject: [PATCH 309/952] [Bugfix] Validate guidance tokenizer compatibility before EngineCore (#52298) Signed-off-by: Oxygen56 Signed-off-by: Artem Perevedentsev Co-authored-by: Artem Perevedentsev --- tests/v1/structured_output/test_validation.py | 97 ++++++++++++++++++- vllm/sampling_params.py | 19 +++- vllm/v1/structured_output/backend_guidance.py | 44 ++++++++- 3 files changed, 152 insertions(+), 8 deletions(-) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index cd5526c7b760..276056b501cd 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,6 +5,7 @@ import json import pytest +from transformers import MistralCommonBackend, TokenizersBackend from vllm.config import StructuredOutputsConfig from vllm.exceptions import VLLMClientError, VLLMValidationError @@ -14,6 +15,8 @@ SamplingParams, StructuredOutputsParams, ) +from vllm.tokenizers import mistral as mistral_tokenizers +from vllm.v1.structured_output import backend_guidance pytestmark = pytest.mark.cpu_test @@ -278,10 +281,102 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): def test_auto_backend_falls_back_on_unsupported_schema(schema, expected_backend): """`auto` falls back on rejection, so it must catch what the validators raise.""" params = SamplingParams(structured_outputs=StructuredOutputsParams(json=schema)) + # trick to create tokenizer that guidance backend accepts + tokenizer = object.__new__(TokenizersBackend) params._validate_structured_outputs( _StubModelConfig(is_diffusion=False), StructuredOutputsConfig(backend="auto"), - tokenizer=object(), + tokenizer=tokenizer, ) assert params.structured_outputs is not None assert params.structured_outputs._backend == expected_backend + + +# ================================================ +# Test validate_structured_outputs functionality +# ================================================ + + +class TestValidateStructuredOutputs: + class TestPR52298Regression: + class _StubSlowTokenizer: + is_fast = False + + class _StubVllmMistralTokenizer: + IS_MISTRAL_TOKENIZER = True + + def __init__(self, is_tekken: bool): + self.is_tekken = is_tekken + self.llg_tokenizer = object() + + @staticmethod + def _validate_guidance(tokenizer): + params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend="guidance"), + tokenizer=tokenizer, + ) + return params + + def _make_tokenizer(self, tokenizer_id: str): + if tokenizer_id == "slow-hf": + return self._StubSlowTokenizer() + if tokenizer_id == "plain-object": + return object() + if tokenizer_id == "raw-mistral-common": + return object.__new__(MistralCommonBackend) + if tokenizer_id == "vllm-mistral-non-tekken": + return self._StubVllmMistralTokenizer(is_tekken=False) + if tokenizer_id == "fast-hf": + return object.__new__(TokenizersBackend) + if tokenizer_id == "vllm-mistral-tekken": + return self._StubVllmMistralTokenizer(is_tekken=True) + raise ValueError(tokenizer_id) + + @pytest.mark.parametrize( + "tokenizer_id, patch_mistral", + [ + pytest.param("slow-hf", False, id="slow-hf"), + pytest.param("plain-object", False, id="plain-object"), + pytest.param("raw-mistral-common", False, id="raw-mistral-common"), + pytest.param( + "vllm-mistral-non-tekken", + True, + id="vllm-mistral-non-tekken", + ), + ], + ) + def test_unsupported_tokenizer(self, monkeypatch, tokenizer_id, patch_mistral): + if patch_mistral: + monkeypatch.setattr( + mistral_tokenizers, + "MistralTokenizer", + self._StubVllmMistralTokenizer, + ) + with pytest.raises(VLLMValidationError): + self._validate_guidance(self._make_tokenizer(tokenizer_id)) + + @pytest.mark.parametrize( + "tokenizer_id, patch_mistral", + [ + pytest.param("fast-hf", False, id="fast-hf"), + pytest.param("vllm-mistral-tekken", True, id="vllm-mistral-tekken"), + ], + ) + def test_supported_tokenizer(self, monkeypatch, tokenizer_id, patch_mistral): + if patch_mistral: + monkeypatch.setattr( + mistral_tokenizers, + "MistralTokenizer", + self._StubVllmMistralTokenizer, + ) + monkeypatch.setattr( + backend_guidance, + "validate_guidance_grammar", + lambda *_args, **_kwargs: None, + ) + params = self._validate_guidance(self._make_tokenizer(tokenizer_id)) + assert params.structured_outputs._backend == "guidance" diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 8a37c5a3fd17..8331a1707b7d 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1317,6 +1317,7 @@ def _validate_structured_outputs( ) from vllm.v1.structured_output.backend_guidance import ( + GuidanceBackend, has_guidance_unsupported_json_features, validate_guidance_grammar, ) @@ -1340,6 +1341,13 @@ def _validate_structured_outputs( "Mistral model, the ['xgrammar', 'outlines'] " "backends or tokenizer_mode='hf' instead." ) + if not GuidanceBackend.is_tokenizer_supported(tokenizer): + raise VLLMValidationError( + "The 'guidance' structured output backend only supports fast " + "Hugging Face tokenizers and Tekken-based Mistral tokenizers. " + "Please use the 'xgrammar' backend or configure a fast " + "tokenizer instead." + ) # TODO: ideally we would have the LLTokenizer here as Lark syntax # allows <|special_token|> and similar, see # https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens @@ -1408,12 +1416,19 @@ def _validate_structured_outputs( skip_guidance = has_guidance_unsupported_json_features(schema) if skip_guidance: - # Fall back to outlines if the tokenizer is non-tekken Mistral or - # the schema contains features unsupported by guidance + # Fall back to outlines for non-Tekken Mistral tokenizers or + # schemas containing features unsupported by guidance. validate_structured_output_request_outlines(self) self.structured_outputs._backend = "outlines" else: # Fall back to guidance by default. + if not GuidanceBackend.is_tokenizer_supported(tokenizer): + raise VLLMValidationError( + "No compatible structured output backend was found for " + "this request. XGrammar rejected the constraint, and " + "Guidance does not support this tokenizer. Configure a " + "fast tokenizer or select a compatible backend explicitly." + ) from None validate_guidance_grammar( self, tokenizer=_get_llg_tokenizer(tokenizer), diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 6dc2dc91ff76..4af5ca2da7ac 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -8,10 +8,12 @@ from typing import TYPE_CHECKING, Any import torch +from transformers import TokenizersBackend from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams +from vllm.tokenizers import TokenizerLike from vllm.utils.import_utils import LazyLoader from vllm.utils.mistral import is_mistral_tokenizer from vllm.v1.structured_output.backend_types import ( @@ -128,6 +130,37 @@ def process_for_additional_properties( @dataclass class GuidanceBackend(StructuredOutputBackend): + @staticmethod + def is_tokenizer_supported(tokenizer: TokenizerLike) -> bool: + """Check if tokenizer is supported by guidance backend. + + TODO(arpera): + guidance backend does NOT support some types of tokenizers + so, we MUST verify tokenizer before using guidance + The problem is that guidance API does not provide + a way to validate tokenizer before building LLTokenizer + So, we have to do it manually here in vLLM. + This is ofcourse very bad, we duplicate complex logic + that can change arbitrary in future. + Our check here is based on guidance source code + on function `from_tokenizer` implementation: + https://github.com/guidance-ai/llguidance/blob/main/python/llguidance/hf.py#L29-L47 + If you want to contribute to vLLM, please, add to llguidance + API function to implement this check there. + In that case we will remove this check from vLLM. + That kind of work is really appreciated! + """ + if is_mistral_tokenizer(tokenizer): + return tokenizer.is_tekken + # PreTrainedTokenizerFast is the same as TokenizersBackend + return isinstance(tokenizer, TokenizersBackend) + + @staticmethod + def _build_ll_tokenizer(tokenizer: TokenizerLike, vocab_size: int) -> Any: + if is_mistral_tokenizer(tokenizer): + return tokenizer.llg_tokenizer + return llguidance_hf.from_tokenizer(tokenizer, max(vocab_size, len(tokenizer))) + def __post_init__(self): self.disable_any_whitespace = ( self.vllm_config.structured_outputs_config.disable_any_whitespace @@ -136,12 +169,13 @@ def __post_init__(self): self.vllm_config.structured_outputs_config.disable_additional_properties ) - if is_mistral_tokenizer(self.tokenizer): - self.ll_tokenizer = self.tokenizer.llg_tokenizer - else: - self.ll_tokenizer = llguidance_hf.from_tokenizer( - self.tokenizer, max(self.vocab_size, len(self.tokenizer)) + if not self.is_tokenizer_supported(self.tokenizer): + raise ValueError( + "Only fast tokenizers are supported." + "This error should NOT happen," + "if you see this, please report an issue." ) + self.ll_tokenizer = self._build_ll_tokenizer(self.tokenizer, self.vocab_size) def compile_grammar( self, From 04462732a40a8329ad60593d754b576c54c5fe5d Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:15:10 +0400 Subject: [PATCH 310/952] [Bugfix] Reject max_num_scheduled_tokens=0 at config time instead of hanging (#60714) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Hui Li <40770106+untrall@users.noreply.github.com> --- tests/test_config.py | 6 ++++++ vllm/config/scheduler.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_config.py b/tests/test_config.py index ebf76c76a0c6..425d55f29346 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3486,6 +3486,12 @@ def test_scheduler_config_init(): print(SchedulerConfig.default_factory().max_model_len) +def test_scheduler_config_rejects_zero_max_num_scheduled_tokens(): + """A zero token budget never schedules anything, so generation would hang.""" + with pytest.raises(ValidationError, match="max_num_scheduled_tokens"): + SchedulerConfig.default_factory(max_num_scheduled_tokens=0) + + @pytest.mark.parametrize( ( "model_id", diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 0d10103ed4ec..1dbe4b72eb7c 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -52,7 +52,7 @@ class SchedulerConfig: In real usage, this should be set in `EngineArgs.create_engine_config`. """ - max_num_scheduled_tokens: int | None = Field(default=None, ge=0) + max_num_scheduled_tokens: int | None = Field(default=None, ge=1) """Maximum number of tokens that the scheduler may issue in a single iteration. This is usually equal to max_num_batched_tokens, but can be smaller in cases From 3766a67a32f7ea8130d84383f72883cee5e8b006 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 9 Oct 2026 15:34:58 -0400 Subject: [PATCH 311/952] [WideEP] Explain how to fix a too-old NCCL for DeepEP v2 (#60449) Signed-off-by: mgoin Co-authored-by: Claude --- tools/ep_kernels/README.md | 28 ++++++++----- tools/ep_kernels/install_python_libraries.sh | 26 ++++++++++-- .../device_communicators/all2all.py | 6 +-- vllm/utils/import_utils.py | 41 ++++++++++++------- 4 files changed, 70 insertions(+), 31 deletions(-) diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index 04bc3ef27680..286a6db7bcf9 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -11,12 +11,16 @@ Step 2 is necessary for multi-node deployment. All scripts accept a positional argument as workspace path for staging the build, defaulting to `$(pwd)/ep_kernels_workspace`. -## NCCL version requirement (CUDA 13+) +## NCCL version requirement for DeepEPv2 -DeepEPv2 uses the NCCL GIN (GPU-Initiated Networking) backend, which requires -NCCL >= 2.30.4 at both compile time and runtime. PyTorch 2.11 pins -`nvidia-nccl-cu13==2.28.9` as a transitive dependency, so you need to -override it. +DeepEPv2 (`--all2all-backend deepep_v2`) uses the NCCL GIN (GPU-Initiated +Networking) backend, which requires NCCL >= 2.30.4 both when DeepEP is built +and at runtime. PyTorch pins an older release as a dependency (PyTorch 2.13 +pins `nvidia-nccl-cu13==2.29.7`), so a plain `pip install` or `uv pip install` +of vLLM is not enough. The vLLM Docker images already override it. + +Upgrade NCCL before running `install_python_libraries.sh`. Use +`nvidia-nccl-cu12` on CUDA 12. **With uv** (recommended): @@ -29,18 +33,22 @@ export UV_OVERRIDE=/tmp/nccl-override.txt uv pip install vllm ``` -**With pip**: +Keep `UV_OVERRIDE` set for later installs into the same environment. Without +it, any `uv pip install` that resolves PyTorch again restores the pinned NCCL. + +**With pip**, or to fix an existing environment: ```bash -pip install vllm pip install "nvidia-nccl-cu13>=2.30.4" --no-deps ``` -The override / reinstall must happen before building DeepEP (for GIN device -headers) and must remain in place at runtime. You can verify with: +`install_python_libraries.sh` warns when the installed NCCL is too old. If +DeepEP was built against an older NCCL, upgrade NCCL and run the script again. +When NCCL is too old at runtime, vLLM logs the reason at startup and +`deepep_v2` fails with the same message. You can check with: ```bash -python -c "from vllm.utils.import_utils import has_deep_ep_v2; print(has_deep_ep_v2())" +python -c "from vllm.utils.import_utils import deep_ep_v2_unavailable_reason; print(deep_ep_v2_unavailable_reason() or 'deepep_v2 is available')" ``` ## Usage diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 9ca6a8d6a99e..fdd7aa20e957 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -78,10 +78,25 @@ mkdir -p "$WHEEL_DIR" pushd "$WORKSPACE" # install dependencies if not installed +UV_PIP_TARGET=() if [ -z "$VIRTUAL_ENV" ]; then - uv pip install --system cmake torch ninja -else - uv pip install cmake torch ninja + UV_PIP_TARGET=(--system) +fi +uv pip install "${UV_PIP_TARGET[@]}" cmake torch ninja + +# DeepEPv2 needs NCCL >= 2.30.4 at both build time and runtime, but PyTorch +# pins an older release. +DEEPEP_V2_MIN_NCCL="2.30.4" +NCCL_PACKAGE="nvidia-nccl-cu${CUDA_VERSION_MAJOR}" +NCCL_WARNING="" +NCCL_INSTALLED=$(uv pip show "${UV_PIP_TARGET[@]}" "$NCCL_PACKAGE" 2>/dev/null | sed -n 's/^Version: //p') +if [ -n "$NCCL_INSTALLED" ] && \ + [ "$(printf '%s\n' "$DEEPEP_V2_MIN_NCCL" "$NCCL_INSTALLED" | sort -V | head -n1)" != "$DEEPEP_V2_MIN_NCCL" ]; then + NCCL_WARNING="WARNING: ${NCCL_PACKAGE} ${NCCL_INSTALLED} is installed, but the deepep_v2 backend +requires NCCL >= ${DEEPEP_V2_MIN_NCCL} when DeepEP is built and when it runs. To use deepep_v2, run + uv pip install \"${NCCL_PACKAGE}>=${DEEPEP_V2_MIN_NCCL}\" --no-deps +and then run this script again. See tools/ep_kernels/README.md." + echo "$NCCL_WARNING" >&2 fi # fetch nvshmem @@ -236,3 +251,8 @@ if [ "$MODE" = "wheel" ]; then echo "All wheels written to $WHEEL_DIR" ls -l "$WHEEL_DIR" fi + +# Repeat after the verbose build output, where the first copy is easy to miss. +if [ -n "$NCCL_WARNING" ]; then + echo "$NCCL_WARNING" >&2 +fi diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index ae31c0b1d3cd..cff34139f968 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -23,6 +23,7 @@ from vllm.utils.func_utils import supports_kw from vllm.utils.import_utils import ( check_moonep_system_support, + deep_ep_v2_unavailable_reason, has_deep_ep, has_deep_ep_v2, has_mori, @@ -1104,10 +1105,7 @@ class DeepEPV2All2AllManager(All2AllManagerBase): """ def __init__(self, cpu_group, tcp_store_group=None, device_group=None): - assert has_deep_ep_v2(), ( - "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 " - "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4." - ) + assert has_deep_ep_v2(), deep_ep_v2_unavailable_reason() super().__init__(cpu_group, tcp_store_group) self._device_group = device_group self.handle_cache = Cache() diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 04a7ebec2eab..277f987d4800 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -469,31 +469,44 @@ def _format_nccl_raw_version(raw: int) -> str: return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}" -def has_deep_ep_v2() -> bool: - """Whether deep_ep with ElasticBuffer (v2 API) is available. +def deep_ep_v2_unavailable_reason() -> str | None: + """Why the deep_ep v2 (ElasticBuffer) API cannot be used, or None if it can. Requires both the ElasticBuffer class in the deep_ep module and NCCL >= 2.30.4 (GIN backend), checked against the runtime library. """ + install_doc = "See tools/ep_kernels/README.md." if not _has_module("deep_ep"): - return False + return f"DeepEP is not installed. {install_doc}" import deep_ep # type: ignore[import-not-found] if not hasattr(deep_ep, "ElasticBuffer"): - return False + return f"The installed DeepEP predates v2 (no ElasticBuffer). {install_doc}" + nccl_ver = _get_runtime_nccl_version() + if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW: + import torch + + min_version = _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW) + cuda_major = (torch.version.cuda or "13").split(".")[0] + return ( + f"DeepEP v2 requires NCCL >= {min_version} but found " + f"{_format_nccl_raw_version(nccl_ver) if nccl_ver else 'unknown'}. " + "PyTorch pins an older NCCL; upgrade it with " + f'`pip install "nvidia-nccl-cu{cuda_major}>={min_version}" --no-deps` ' + f"and rebuild DeepEP against it. {install_doc}" + ) + return None + + +def has_deep_ep_v2() -> bool: + """Whether deep_ep with ElasticBuffer (v2 API) is available.""" try: - nccl_ver = _get_runtime_nccl_version() - if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW: - logger.info_once( - "DeepEP v2 requires NCCL >= %s but found %s. " - "deepep_v2 backend will not be available.", - _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW), - _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown", - ) - return False + reason = deep_ep_v2_unavailable_reason() except Exception: return False - return True + if reason is not None and _has_module("deep_ep"): + logger.info_once("deepep_v2 backend will not be available: %s", reason) + return reason is None def has_deep_gemm() -> bool: From 02e55d2af38adab51ec5ffb985063d355b20805b Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:48:37 +0800 Subject: [PATCH 312/952] [KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots (#57169) Signed-off-by: Yifan Qiao Signed-off-by: Dakai An Signed-off-by: Lucas Wilkinson Co-authored-by: Yifan Qiao Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson --- .../attention/test_deepgemm_attention.py | 46 ++ .../attention/test_rocm_aiter_mla_sink.py | 4 +- .../test_jit_warmup_triton_launcher.py | 3 - .../test_dspark_noncausal_sparse_mla.py | 2 +- .../test_flashinfer_mla_sparse_sm90.py | 18 +- ...st_flashmla_nope_sm90_backend_selection.py | 97 +++ .../test_glm5next_indexer_block_sizes.py | 40 ++ .../test_indexer_deepseek_v4_slot_mapping.py | 19 +- .../test_kpool_indexer_block_sizes.py | 54 +- .../attention/test_kpool_tail_slot_mapping.py | 45 +- tests/v1/attention/test_mm_prefix.py | 1 - .../v1/attention/test_rocm_glm5next_sparse.py | 6 + .../v1/attention/test_sparse_mla_backends.py | 69 +- .../test_sparse_mla_kv_cache_layout.py | 21 +- .../test_partial_prefix_cache_hits.py | 8 +- tests/v1/core/test_contiguous_kv_packing.py | 42 +- tests/v1/core/test_kv_cache_utils.py | 600 ++---------------- .../core/test_single_type_kv_cache_manager.py | 14 +- tests/v1/kv_connector/unit/conftest.py | 2 +- .../test_mooncake_connector_hybrid_mamba.py | 9 +- .../unit/test_mooncake_store_hma_e2e.py | 4 +- .../unit/test_nixl_desc_geometry.py | 187 +----- tests/v1/simple_kv_offload/test_scheduler.py | 4 +- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- tests/v1/spec_decode/test_eagle.py | 6 + .../v1/spec_decode/test_llm_base_proposer.py | 24 +- .../test_gpu_model_runner_streaming.py | 1 - tests/v1/worker/test_attn_utils.py | 112 +++- tests/v1/worker/test_gpu_block_table.py | 31 +- tests/v1/worker/test_gpu_input_batch.py | 9 - .../test_gpu_kpool_tail_slot_mapping.py | 8 +- tests/v1/worker/test_gpu_model_runner.py | 108 +--- ...odel_runner_prefill_tail_classification.py | 3 - tests/v1/worker/test_gpu_model_runner_v2.py | 1 - tests/v1/worker/test_gpu_pcp_manager.py | 3 +- ...test_gpu_target_dependent_ar_speculator.py | 51 ++ tests/v1/worker/test_jit_warmup_migration.py | 18 +- .../worker/test_kv_cache_allocation_scope.py | 2 + tests/v1/worker/test_mamba_utils.py | 1 - tests/v1/worker/test_utils.py | 37 ++ .../kv_transfer/kv_connector/utils.py | 30 + .../v1/mooncake/mooncake_connector.py | 28 +- .../kv_connector/v1/nixl/base_worker.py | 124 +--- .../layers/attention/mla_attention.py | 8 - .../models/deepseek_v41/nvidia/model_state.py | 2 +- vllm/models/glm5next/common/attention.py | 64 +- .../glm5next/nvidia/ops/kpool_compress.py | 5 - vllm/platforms/cuda.py | 20 - vllm/platforms/interface.py | 39 +- vllm/utils/deep_gemm.py | 9 +- vllm/v1/attention/backend.py | 8 +- vllm/v1/attention/backends/composite.py | 5 - .../backends/mla/compressor_utils.py | 4 +- .../backends/mla/flashattn_mla_sparse.py | 7 +- .../backends/mla/flashinfer_mla_sparse.py | 22 +- .../mla/flashinfer_mla_sparse_sm120.py | 8 +- .../mla/flashinfer_mla_sparse_sm90.py | 25 +- .../attention/backends/mla/flashmla_sparse.py | 51 +- vllm/v1/attention/backends/mla/indexer.py | 66 +- .../backends/mla/rocm_aiter_mla_sparse.py | 25 +- .../v1/attention/backends/mla/sparse_utils.py | 9 +- .../attention/backends/mla/xpu_mla_sparse.py | 5 + vllm/v1/core/kv_cache_coordinator.py | 8 +- vllm/v1/core/kv_cache_utils.py | 354 ++--------- vllm/v1/core/single_type_kv_cache_manager.py | 13 - vllm/v1/engine/core.py | 4 +- vllm/v1/hisparse/binding.py | 8 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/kv_cache_interface.py | 56 +- vllm/v1/spec_decode/gemma4.py | 8 +- vllm/v1/spec_decode/llm_base_proposer.py | 17 +- vllm/v1/spec_decode/step3p5.py | 8 +- vllm/v1/worker/block_table.py | 129 +--- vllm/v1/worker/cp_utils.py | 6 +- vllm/v1/worker/gpu/attn_utils.py | 24 +- vllm/v1/worker/gpu/block_table.py | 26 +- vllm/v1/worker/gpu/input_batch.py | 17 +- vllm/v1/worker/gpu/model_runner.py | 8 +- .../gpu/spec_decode/dflash/speculator.py | 2 +- .../target_dependent_ar/speculator.py | 5 + vllm/v1/worker/gpu_input_batch.py | 2 - vllm/v1/worker/gpu_model_runner.py | 41 +- vllm/v1/worker/utils.py | 252 ++++++-- 83 files changed, 1281 insertions(+), 1987 deletions(-) create mode 100644 tests/v1/attention/test_glm5next_indexer_block_sizes.py diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index 5d11063c7a2b..f332fb0fea21 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -313,3 +313,49 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): ref_logits = ref_logits.masked_fill(~mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") +@pytest.mark.skipif( + not current_platform.has_device_capability(90), reason="SM90 and SM100 only" +) +def test_deepgemm_paged_mqa_packed_manager_block_stride(): + torch.manual_seed(0) + num_blocks, block_size, packed_rows = 4, 256, 320 + num_heads, head_dim, row_bytes = 64, 128, 132 + page_size, pages_per_block = 64, block_size // 64 + dense_stride, packed_stride = pages_per_block, packed_rows // page_size + compact_pages = kv_cache_cast_to_fp8( + torch.randn( + num_blocks, block_size, 1, head_dim, device="cuda", dtype=torch.bfloat16 + ).view(num_blocks * pages_per_block, page_size, 1, head_dim) + ) + packed_backing = torch.zeros( + num_blocks, + packed_rows, + row_bytes, + device="cuda", + dtype=torch.uint8, + ) + packed_manager = packed_backing[:, :block_size] + packed_manager.copy_(compact_pages.squeeze(2).view(num_blocks, block_size, -1)) + packed_pages = packed_backing.view(-1, page_size, row_bytes).unsqueeze(2) + offsets = torch.arange(pages_per_block, device="cuda", dtype=torch.int32) + tables = [ + (stride + offsets).unsqueeze(0) for stride in (dense_stride, packed_stride) + ] + + q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to( + torch.float8_e4m3fn + ) + weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32) + schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms()) + logits = [ + fp8_fp4_paged_mqa_logits( + (q, None), cache, weights, context_lens, table, schedule, block_size, False + ) + for cache, table in zip((compact_pages, packed_pages), tables) + ] + torch.testing.assert_close(*logits, rtol=0, atol=0) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index 0b007dbc875c..f66d408c8495 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -463,7 +463,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): set_random_seed(412) num_tokens, num_heads, block_size = 2, 8, 16 q = torch.randn(num_tokens, num_heads, Q_HEAD_DIM, device="cuda").to(dtype) - kv = torch.randn(2 * block_size, 1, Q_HEAD_DIM, device="cuda").to(dtype) + kv = torch.randn(2, 2, block_size, Q_HEAD_DIM, device="cuda").to(dtype)[:, 1] selected = torch.tensor([[1, 7, 18], [0, 3, 20]], dtype=torch.int32, device="cuda") sinks = torch.linspace(-3.0, 5.0, num_heads, device="cuda") impl = object.__new__(ROCMAiterMLASparseImpl) @@ -504,7 +504,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): SimpleNamespace(_q_scale=None, _k_scale=None), ) references = [ - _sink_reference(q[i], kv[:, 0][selected[i].long()], sinks, SM_SCALE) + _sink_reference(q[i], kv.flatten(0, 1)[selected[i].long()], sinks, SM_SCALE) for i in range(num_tokens) ] expected, _, value_scale, score_error = ( diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index a92303dcf288..c19804eae4b4 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -553,12 +553,10 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None: owner = ComputeSlotMappingKernel() compile_key = owner.CompileKey( kv_cache_block_size=16, - blocks_per_kv_block=1, total_cp_world_size=2, total_cp_rank=1, cp_kv_cache_interleave_size=1, block_table_stride=128, - block_size=16, ) launches: list[tuple[Any, ...]] = [] @@ -574,7 +572,6 @@ def launch(launch_spec: Any, inputs: Any) -> None: assert inputs["block_table_stride"] == 128 assert kwargs == { "KV_CACHE_BLOCK_SIZE": 16, - "BLOCKS_PER_KV_BLOCK": 1, "TOTAL_CP_WORLD_SIZE": 2, "TOTAL_CP_RANK": 1, "CP_KV_CACHE_INTERLEAVE_SIZE": 1, diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 7cd7863b0659..0942af4f4570 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py @@ -425,7 +425,7 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i and kv_cache_dtype != "fp8_ds_mla" ): pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype") - if block_size not in backend_cls.get_supported_kernel_block_sizes(): + if not backend_cls.supports_block_size(block_size): pytest.skip( f"{backend_cls.get_name()} does not support block_size={block_size}" ) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 735748ae943e..66113aca0917 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -29,7 +29,7 @@ TOPK = 128 # triton convert requires width % 128 == 0 -def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_): +def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **kw): out = torch.full_like(token_indices, -1) counts = torch.zeros(token_indices.shape[0], dtype=torch.int32) for t in range(token_indices.shape[0]): @@ -41,7 +41,7 @@ def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_): blk = int(block_table[int(req_id[t]), pos // BLOCK_SIZE]) if blk < 0: continue - vals.append(blk * BLOCK_SIZE + pos % BLOCK_SIZE) + vals.append(blk * kw["BLOCK_STRIDE_ROWS"] + pos % BLOCK_SIZE) out[t, : len(vals)] = torch.tensor(vals, dtype=out.dtype) counts[t] = len(vals) return out, counts @@ -133,11 +133,8 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): meta.state = state q_nope = torch.randn(rows, impl.num_heads, HEAD) q_rope = torch.randn(rows, impl.num_heads, qk_rope) - cache = torch.zeros( - 8 * BLOCK_SIZE, - impl.head_size, - dtype=torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16, - ) + dtype = torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16 + cache = torch.zeros(8, 2, BLOCK_SIZE, impl.head_size, dtype=dtype)[:, 0] out, lse = impl.forward_mqa( (q_nope, q_rope), cache, meta, SimpleNamespace(_k_scale_float=0.5) @@ -147,7 +144,10 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): # Reserved buffers carry this step's slots; lengths are NOT refreshed # here (the builder plans them host-side before capture/replay). ref_slots, ref_counts = ref_convert( - meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer + meta.req_id_per_token, + meta.block_table, + impl.topk_indices_buffer, + BLOCK_STRIDE_ROWS=2 * BLOCK_SIZE, ) offset = 0 for t in range(rows): @@ -161,7 +161,7 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert state.wrapper.run_args is not None q_pe, ckv, kpe, kwargs = state.wrapper.run_args[1:] assert q_pe.shape == (rows, impl.num_heads, qk_rope) - assert ckv.shape == (8 * BLOCK_SIZE, 1, HEAD) + assert ckv.shape == ((8 - 1) * 2 * BLOCK_SIZE + BLOCK_SIZE, 1, HEAD) assert kpe.shape[-1] == qk_rope if impl.use_fp8_kv_cache: assert kwargs["ckv_scale"] == 0.5 and kwargs["kpe_scale"] == 1.0 diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 652358876949..8532d2f43aee 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -24,6 +24,21 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import ( QUANTIZED_DS_MLA_CACHE_FORMATS, FlashMLASparseBackend, + _fp8_kv_pages, +) +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheTensor, + KVQuantMode, + MLAAttentionSpec, +) +from vllm.v1.worker.utils import ( + AttentionGroup, + allocate_kv_cache, + prepare_kernel_block_sizes, + select_common_block_size, ) SM90 = DeviceCapability(major=9, minor=0) @@ -83,6 +98,88 @@ def test_supported_head_sizes_include_512(): assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512] +def test_flashmla_bf16_nope_accepts_packed_manager_blocks(): + spec = MLAAttentionSpec( + block_size=1152, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + cache_dtype_str="bfloat16", + block_stride_alignment=1024, + ) + num_blocks = 2 + layers = ["layer.0", "layer.1"] + page_size = spec.page_size_bytes + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * len(layers) * page_size, + layers=layers, + layer_stride=page_size, + block_stride=len(layers) * page_size, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(layers, spec)], + ) + groups = [[AttentionGroup(FlashMLASparseBackend, layers, spec, 0)]] + + kernel_block_sizes = prepare_kernel_block_sizes(config, groups) + caches = allocate_kv_cache( + config, + torch.device("cpu"), + KVCacheLayout.BLHNC, + kernel_block_sizes, + ) + + assert kernel_block_sizes == [1152] + assert caches["layer.0"].shape == (num_blocks, 1, 1152, 512) + + +@pytest.mark.parametrize("sm100", [False, True]) +def test_flashmla_quantized_cache_kernel_pages(monkeypatch, sm100): + """Quantized blocks are read as 64-token pages; SM100 only row-aligns them + (TMA) and keeps 64-token blocks.""" + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda family, device_id=0: sm100 and family == 100, + ) + spec = MLAAttentionSpec( + block_size=1152, + num_kv_heads=1, + head_size=576, + dtype=torch.float8_e4m3fn, + cache_dtype_str="fp8_ds_mla", + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, + state_content_bytes=656, + ) + + kernel_block_size = 64 if sm100 else 1152 + assert ( + select_common_block_size(1152, [FlashMLASparseBackend], [spec]) + == kernel_block_size + ) + + +def test_fp8_kv_pages_address_packed_blocks(): + """Token t of block b is row b * block_stride_rows + t of the 64-row pages; + 64-token blocks are passed through.""" + row, block_size, stride_rows = 656, 128, 192 + backing = torch.arange(3 * stride_rows).repeat_interleave(row) + cache = backing.view(3, stride_rows, row)[:, :block_size] + + pages, block_stride_rows = _fp8_kv_pages(cache, block_size) + assert block_stride_rows == stride_rows + assert pages.shape[1:] == (64, row) + slot = 2 * block_stride_rows + 100 + assert torch.equal(pages[slot // 64, slot % 64], cache[2, 100]) + + pages, block_stride_rows = _fp8_kv_pages(cache[:, :64], 64) + assert block_stride_rows is None + assert torch.equal(pages, cache[:, :64]) + + def test_quantized_ds_mla_formats_are_the_envelope_set(): assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS diff --git a/tests/v1/attention/test_glm5next_indexer_block_sizes.py b/tests/v1/attention/test_glm5next_indexer_block_sizes.py new file mode 100644 index 000000000000..9f27d2994636 --- /dev/null +++ b/tests/v1/attention/test_glm5next_indexer_block_sizes.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash indexer kernel block sizes are DeepGEMM pool pages.""" + +from types import SimpleNamespace + +import pytest + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla import indexer +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend + +KPOOL = 4 + + +@pytest.mark.parametrize( + "rocm,sm120,expected", + [ + (False, False, [KPOOL * 32, KPOOL * 64]), + # SM120 DeepGEMM takes only 64-entry pages for the fp8 cache. + (False, True, [KPOOL * 64]), + # ROCm's AITER paged-MQA logits take the same pool pages. + (True, False, [KPOOL * 32, KPOOL * 64]), + ], +) +def test_glm5next_indexer_kernel_block_sizes(monkeypatch, rocm, sm120, expected): + monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm) + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda family: sm120 and family == 120, + ) + model_config = SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=KPOOL)) + config = SimpleNamespace(model_config=model_config) + monkeypatch.setattr(indexer, "get_current_vllm_config", lambda: config) + + spec = SimpleNamespace(tokens_per_state=KPOOL) + assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes(spec) == expected + # Platform block-size selection queries without a spec. + assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes() == expected diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 17dea5d03c19..464fce66f212 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -122,7 +122,6 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = 16 - builder.kernel_block_size = None builder.num_sms = num_sms builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) builder.kv_cache_spec = SimpleNamespace(num_states=64) @@ -170,30 +169,24 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize( - "block_size,kernel_block_size,compress_ratio", [(64, 64, 1), (256, 64, 4)] -) -def test_indexer_draft_decode_refresh_follows_new_batch( - monkeypatch, block_size, kernel_block_size, compress_ratio -): +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_draft_decode_refresh_follows_new_batch(monkeypatch, compress_ratio): """Without a prior build(), the refresh must not keep the last batch's state.""" monkeypatch.setattr(indexer_module, "has_deep_gemm", lambda: False) device = torch.device("cuda") num_reqs, width = 3, 8 - factor = block_size // kernel_block_size builder = object.__new__(DeepseekV32IndexerMetadataBuilder) builder.dcp_world_size = 1 builder.compress_ratio = compress_ratio - builder.kernel_block_size = kernel_block_size builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) - builder.kv_cache_spec = SimpleNamespace(block_size=block_size, num_states=64) + builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64) def stale(*shape: int) -> torch.Tensor: return torch.full(shape, -7, dtype=torch.int32, device=device) decode = DeepSeekV32IndexerDecodeMetadata( - block_table=stale(num_reqs, width // factor), + block_table=stale(num_reqs, width), seq_lens=stale(num_reqs, 1), decode_lens=stale(num_reqs), requires_padding=False, @@ -219,8 +212,7 @@ def stale(*shape: int) -> torch.Tensor: builder.update_draft_decode_metadata(metadata) - expected_block_table = request_block_table[:, ::factor] // factor - torch.testing.assert_close(decode.block_table, expected_block_table) + torch.testing.assert_close(decode.block_table, request_block_table) torch.testing.assert_close(decode.indices, builder.arange_buffer[:num_reqs]) assert torch.all(decode.per_req_decode_lens == 1) assert torch.all(decode.decode_lens == 1) @@ -471,7 +463,6 @@ def test_zero_token_pcp_rank_participates_in_compressed_mapping_gather(monkeypat builder.pcp_rank = 0 builder.dcp_world_size = 1 builder.cp_kv_cache_interleave_size = 1 - builder.kernel_block_size = None builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64) builder.compressed_slot_mapping_buffer = torch.zeros(8, dtype=torch.int64) diff --git a/tests/v1/attention/test_kpool_indexer_block_sizes.py b/tests/v1/attention/test_kpool_indexer_block_sizes.py index a737988dde66..86e196d539b4 100644 --- a/tests/v1/attention/test_kpool_indexer_block_sizes.py +++ b/tests/v1/attention/test_kpool_indexer_block_sizes.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU tests for storage-block kernel-block selection in prepare_kernel_block_sizes.""" +"""CPU tests for the kpool indexer's kernel block selection.""" from types import SimpleNamespace @@ -11,20 +11,14 @@ from vllm.models.glm5next.common.attention import Glm5NextIndexerCache from vllm.v1.attention.backend import AttentionBackend, MultipleOf from vllm.v1.attention.backends.mla import indexer as indexer_mod -from vllm.v1.attention.backends.mla.indexer import ( - DeepseekV32IndexerBackend, - KpoolTailBackend, -) -from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( - ROCMAiterMLASparseBackend, -) +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend from vllm.v1.kv_cache_interface import MLAAttentionSpec -from vllm.v1.worker.utils import prepare_kernel_block_sizes +from vllm.v1.worker.utils import prepare_kernel_block_sizes, select_common_block_size pytestmark = pytest.mark.cpu_test # ``index_kpool`` of zai-org/GLM-5.3-Flash: one indexer state pools 4 -# compressed states (hence storage_block_size = page * 4). +# compressed states (hence a kernel block is page * 4 tokens). INDEX_KPOOL = 4 # The hybrid KDA/mamba manager floors at TP8..TP1; 640 is the geometry of the @@ -59,14 +53,17 @@ def _prepare( spec: MLAAttentionSpec, backends: list[type[AttentionBackend]] ) -> list[int]: kv_cache_config = SimpleNamespace( - kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)] + kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec, layer_names=[])], + kv_cache_tensors=[], ) - attn_groups = [[SimpleNamespace(backend=backend) for backend in backends]] + attn_groups = [ + [SimpleNamespace(backend=backend, kv_cache_spec=spec) for backend in backends] + ] return prepare_kernel_block_sizes(kv_cache_config, attn_groups) @pytest.mark.parametrize("manager_block", TP_FLOOR_MANAGER_BLOCKS) -def test_prepare_uses_storage_block_for_the_kpool_group( +def test_kpool_group_maps_pool_pages( monkeypatch: pytest.MonkeyPatch, manager_block: int ): """The bug: on ROCm, kpool groups must get 128/256 - not the 640-4352 @@ -78,54 +75,43 @@ def test_prepare_uses_storage_block_for_the_kpool_group( """ _mock_rocm_platform(monkeypatch) spec = _kpool_storage_spec(manager_block) - expected_storage = 256 if manager_block % 256 == 0 else 128 - assert spec.storage_block_size == expected_storage + expected = 256 if manager_block % 256 == 0 else 128 - selected = _prepare( - spec, - [ROCMAiterMLASparseBackend, DeepseekV32IndexerBackend, KpoolTailBackend], - ) - assert selected == [expected_storage] + # The kpool group's AttentionGroup selects its kernel block this way. + selected = [ + select_common_block_size(manager_block, [Glm5NextIndexerBackend], [spec]) + ] + assert selected == [expected] assert selected != [manager_block] # The hybrid block-table split stays integral. assert manager_block % selected[0] == 0 -def test_prepare_falls_back_when_storage_block_is_not_supported(): - """No storage block, or a storage block the backends reject (e.g. the - CUDA indexer's exact ``[64]`` vs a 128-token storage block): selection is - `select_common_block_size`'s answer, unchanged from no-fix behavior. - """ +def test_prepare_uses_the_backend_vote(): + """Unpacked groups take `select_common_block_size`'s answer.""" class Fixed64Backend: @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return [64] @staticmethod def get_name() -> str: return "FIXED64" - # Storage block the group's backend does not accept -> backend vote (64). - spec = _kpool_storage_spec(640) - assert spec.storage_block_size == 128 - assert _prepare(spec, [Fixed64Backend]) == [64] - - # No storage block at all (any non-kpool MLA model) -> backend vote. plain_spec = MLAAttentionSpec( block_size=640, num_kv_heads=1, head_size=128, dtype=torch.bfloat16, ) - assert plain_spec.storage_block_size is None assert _prepare(plain_spec, [Fixed64Backend]) == [64] # ...and with backends that accept the manager block, the manager wins, # exactly as on main. class AcceptAllBackend: @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return [MultipleOf(1)] @staticmethod diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 183c63ec2ffd..062fd783e0b3 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -3,7 +3,7 @@ """CPU tests for the kpool tail slot mapping (no GPU required). The kpool tail cache is a 1-block-per-request circular ring addressed by -``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block +``pos % kpool`` (``CircularBufferSpec`` / ``CircularBufferManager``: exactly one block allocated per request, never grown, so only column 0 of its block table is ever written; the rest stays zero-initialized). @@ -32,27 +32,25 @@ KpoolTailMetadataBuilder, compute_kpool_tail_slot_mapping, ) -from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides +from vllm.v1.kv_cache_interface import CircularBufferSpec, compute_layout_strides from vllm.v1.kv_cache_layout import KVCacheLayout -from vllm.v1.worker.block_table import get_block_table_width KPOOL = 4 def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): (layout,) = KpoolTailBackend.supported_kv_cache_layouts() - spec = KpoolTailSpec( + spec = CircularBufferSpec( block_size=KPOOL, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=KPOOL, ) strides = compute_layout_strides(spec, num_blocks=8, num_layers=3, layout=layout) _, _, head_stride, state_stride, content_stride = strides - assert layout is KVCacheLayout.LBHNC + assert layout is KVCacheLayout.BLHNC assert head_stride == KPOOL * 128 * torch.bfloat16.itemsize assert state_stride == 128 * torch.bfloat16.itemsize assert content_stride == 1 @@ -77,37 +75,14 @@ def test_tail_ring_divides_the_attention_block(num_speculative_tokens, ring): SimpleNamespace(num_speculative_tokens=num_speculative_tokens) ) - assert spec.block_size == spec.sliding_window == ring + assert isinstance(spec, CircularBufferSpec) and spec.block_size == ring assert ring >= KPOOL + num_speculative_tokens assert 640 % ring == 0 -def test_tail_spec_opts_out_of_generic_slot_mapping(): - """The tail row is one block wide (padded to the block-table alignment), so - the generic kernel's ``pos // kpool`` column index runs off the end of the - allocation for long prompts. The spec must opt out of it entirely.""" - spec = KpoolTailSpec( - block_size=KPOOL, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - sliding_window=KPOOL, - ) - max_len = 1 << 20 - width = get_block_table_width( - spec.max_num_blocks_per_req(None, max_len), - spec.block_size, - token_alignment=spec.block_table_token_alignment, - ) - - assert width * KPOOL < max_len - assert spec.uses_slot_mapping is False - - def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the - request's single KpoolTailManager block, the remaining columns are never + request's single CircularBufferManager block, the remaining columns are never written and stay zero.""" bt = torch.zeros(len(own_blocks), width, dtype=torch.int32) bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32) @@ -447,6 +422,7 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len) tokens between the last request boundary and num_actual_tokens (mapped to the last request), untouched padding beyond num_actual, and PAD for requests on the null block.""" + ring_size = 2 * KPOOL positions, qsl, slot_mapping, _, num_reqs = make_batch( per_req, padded_len=padded_len ) @@ -455,14 +431,17 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len) slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000 bt = make_tail_block_table(own_blocks) - ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs) - got = circular_tail_slots( + ref = compute_kpool_tail_slot_mapping( + slot_mapping, bt, qsl, positions, num_actual, num_reqs, ring_size + ) + got = compute_kpool_tail_slot_mapping( slot_mapping.cuda(), bt.cuda(), qsl.cuda().to(torch.int32), positions.cuda(), num_actual, num_reqs, + ring_size, ) torch.testing.assert_close(got.cpu(), ref) for req, blk in enumerate(own_blocks): diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index f878ede4865f..b8085f0bc584 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/tests/v1/attention/test_mm_prefix.py @@ -805,7 +805,6 @@ def variant_uses_mm_prefix(variant): ) layer.mm_prefix_clamp_sliding_window = True builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE) - builder.set_kernel_block_size(block_size) # Match production, where unwritten page tails are zero-initialized. cache = torch.zeros( num_blocks, diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index e69686f5005d..6077b9eb2fb6 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -11,6 +11,7 @@ from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, _use_rocm_sparse_triton, fit_kpool_indices_to_aiter, ) @@ -19,6 +20,7 @@ _validate_dsv4_sparse_dims, _validate_sparse_dims, ) +from vllm.v1.kv_cache_interface import KVCacheLayout @triton.jit @@ -104,6 +106,10 @@ def test_fit_kpool_indices_rejects_narrow_input(): ) +def test_rocm_sparse_mla_supports_glm_packed_layout(): + assert KVCacheLayout.BLHNC in ROCMAiterMLASparseBackend.supported_kv_cache_layouts() + + @pytest.mark.parametrize( ( "kv_cache_dtype", diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 973a9e03e88b..160b766da077 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -140,7 +140,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): metadata = SimpleNamespace( req_id_per_token=torch.zeros(1, dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), - block_size=1, + block_size=32, ) recorded_scale = None @@ -178,7 +178,7 @@ def fake_flashinfer(**kwargs): ) impl.forward_mqa( torch.zeros(1, 1, kv_lora_rank), - torch.zeros(1, 1, kv_lora_rank), + torch.zeros(1, 32, kv_lora_rank), metadata, SimpleNamespace(), ) @@ -437,6 +437,7 @@ def test_sparse_backend_decode_correctness( qk_rope_head_dim: int, v_head_dim: int, monkeypatch, + block_stride_rows: int | None = None, ): if qk_rope_head_dim == 0 and ( backend_cls != FlashMLASparseBackend @@ -478,8 +479,7 @@ def test_sparse_backend_decode_correctness( if device_capability is None or device_capability.major != 10: pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x") - supported_block_sizes = backend_cls.get_supported_kernel_block_sizes() - if block_size not in supported_block_sizes: + if not backend_cls.supports_block_size(block_size): pytest.skip( f"{backend_cls.get_name()} does not support block_size={block_size}" ) @@ -719,6 +719,11 @@ def test_sparse_backend_decode_correctness( kv_cache_dtype=kv_cache_dtype, scale=kv_cache_scale, ) + if block_stride_rows is not None: + pad = (0, 0, 0, block_stride_rows - block_size) + kv_cache = torch.nn.functional.pad(kv_cache, pad)[:, :, :block_size] + slot = int(common_attn_metadata.slot_mapping[-1]) + query_row = kv_cache[slot // block_size, 0, slot % block_size].zero_() # The sparse builder clones the layer's dense-MHA prefill backend from # static_forward_context; register a mock layer carrying one. @@ -814,6 +819,12 @@ def test_sparse_backend_decode_correctness( out_buffer, ) + if block_stride_rows is not None and kv_cache_dtype == "auto": + expected_row = torch.cat((kv_c_vllm[-1], k_pe_vllm[-1, 0])) + torch.testing.assert_close(query_row, expected_row, rtol=0, atol=0) + elif block_stride_rows is not None: + assert query_row.any() + assert backend_output.shape == sdpa_reference.shape assert backend_output.dtype == sdpa_reference.dtype assert torch.isfinite(backend_output).all() @@ -828,6 +839,56 @@ def test_sparse_backend_decode_correctness( torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01) +@pytest.mark.parametrize( + ("backend_cls", "kv_cache_dtype", "block_size", "block_stride_rows", "head_dims"), + [ + pytest.param( + FlashInferMLASparseTRTLLMBackend, + "auto", + 256, + 320, + (128, 128, 64, 128), + id="FlashInferTRTLLM", + ), + # GLM-5.3 on SM90: the fp8 kernel's 64-token pages in larger blocks. + pytest.param( + FlashMLASparseBackend, + "fp8_ds_mla", + 128, + 192, + (64, 256, 0, 256), + id="FlashMLA-fp8_ds_mla", + ), + ], +) +def test_sparse_mla_packed_stride( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch, + backend_cls, + kv_cache_dtype, + block_size, + block_stride_rows, + head_dims, +): + test_sparse_backend_decode_correctness( + default_vllm_config, + dist_init, + backend_cls, + "mixed_small", + kv_cache_dtype, + 1, + block_size, + workspace_init, + 1.0, + 1.0, + *head_dims, + monkeypatch, + block_stride_rows=block_stride_rows, + ) + + def _triton_convert_reference_impl( req_ids: torch.Tensor, block_table: torch.Tensor, diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index cd3c9b218147..d869141df81a 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform +from vllm.v1.attention.backend import MultipleOf from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseTRTLLMBackend, ) @@ -25,7 +26,7 @@ from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs from vllm.v1.kv_cache_layout import KVCacheLayout -from vllm.v1.worker.utils import allocate_kv_cache +from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec pytestmark = pytest.mark.skip_global_cleanup @@ -48,10 +49,10 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576), + (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152), + (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576), (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, None), + (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)), (FlashMLASparseBackend, "fp8_ds_mla", True, 656), (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), ], @@ -79,7 +80,9 @@ def test_allocation_and_warmup_follow_addressing_mode( non_causal_multi_token_decode=False, ) layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer) - spec = MLAAttention.get_kv_cache_spec(layer, config) + spec = customize_attention_spec( + backend, MLAAttention.get_kv_cache_spec(layer, config) + ) indexer = SimpleNamespace( cache_config=config.cache_config, head_dim=132, dtype=torch.uint8 ) @@ -102,10 +105,8 @@ def test_allocation_and_warmup_follow_addressing_mode( MLAAttention.bind_kv_cache(layer, views["mla"]) DeepseekV32IndexerCache.bind_kv_cache(indexer, views["indexer"]) assert indexer.kv_cache.stride(0) == views["indexer"].stride(0) - if stride_alignment is not None: - assert ( - views["mla"].stride(0) * views["mla"].element_size() % stride_alignment == 0 - ) + alignment = spec.get_block_stride_alignment() + assert views["mla"].stride(0) * views["mla"].element_size() % alignment == 0 if not layer._uses_flat_kv_cache(): register.assert_not_called() else: @@ -113,7 +114,7 @@ def test_allocation_and_warmup_follow_addressing_mode( register.assert_called_once_with(config, block_stride_rows=stride) layer.kv_cache[1, 0].fill_(3) torch.testing.assert_close(rows[stride], layer.kv_cache[1, 0]) - if stride_alignment is None: + if alignment == 1: assert cache.kv_cache_tensors[0].size == cache.num_blocks * sum( item.page_size_bytes for item in specs.values() ) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 51317deb8d01..4bffe91a6534 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -26,8 +26,8 @@ ) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, MambaSpec, @@ -2308,12 +2308,11 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits(): ), KVCacheGroupSpec( ["tail"], - KpoolTailSpec( + CircularBufferSpec( block_size=mamba_block_size, num_kv_heads=1, head_size=1, dtype=torch.float32, - sliding_window=mamba_block_size, ), ), ], @@ -2366,12 +2365,11 @@ def test_kpool_tail_supports_128_token_partial_hash_hits(): ), KVCacheGroupSpec( ["tail"], - KpoolTailSpec( + CircularBufferSpec( block_size=kpool_block_size, num_kv_heads=1, head_size=1, dtype=torch.float32, - sliding_window=kpool_block_size, ), ), ], diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 638deeedc579..6a705b6e5e8e 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -88,7 +88,7 @@ def test_packed_alignment_preserves_hot_pages_and_generic_stride(): hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2) groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)] alignment = lcm(1152, hot.page_size_bytes) - stride = _pool_bytes_per_block(groups) + stride = _pool_bytes_per_block(_mock_vllm_config("BLHNC"), groups) assert stride == alignment layout = KVCacheLayout.BLHNC tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride) @@ -254,7 +254,7 @@ def test_roles_land_in_separate_groups(self): def test_every_group_fits_one_packed_block(self): config = _shared_layout_config() groups = get_kv_cache_groups(config, _make_csa_linear_specs()) - bytes_per_block = _get_kv_cache_bytes_per_block(groups) + bytes_per_block = _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) pages = _pages(groups) for group in groups: @@ -418,7 +418,7 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid groups = _get_packed_kv_cache_groups(config, specs) gdn = [g for g in groups if g.layer_names[0].startswith("gdn.")] - assert _get_kv_cache_bytes_per_block(groups) == sum( + assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == sum( specs[name].page_size_bytes for name in specs if name.startswith("wide.") ) # That block holds every GDN state at once, so the repeat pattern alone @@ -465,7 +465,7 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): pages = _pages(groups) main = next(g for g in groups if "layers.2.attn" in g.layer_names) main_bytes = sum(pages[n] for n in main.layer_names) - assert _get_kv_cache_bytes_per_block(groups) == main_bytes + assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == main_bytes swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")] per_group = main_bytes // pages["layers.0.swa"] assert len(swa_groups) == -(-43 // per_group) @@ -479,9 +479,9 @@ def test_sliding_window_bucket_is_capped_at_the_main_page(self): class TestDensePacking: def test_bytes_per_block_is_largest_group(self): groups, g1, g2 = _mixed_page_groups() - assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block( - groups - ) + assert _get_kv_cache_bytes_per_block( + groups, KVCacheLayout.BLHNC + ) == _expected_bytes_per_block(groups) config = get_kv_cache_config_from_groups( _mock_vllm_config("BLHNC"), groups, MEMORY @@ -554,6 +554,20 @@ def test_single_group_mixed_pages_follows_layout(self, layout): assert idx_tensor.offset == 2 * _mla(512).page_size_bytes assert mla_tensor.block_stride == block_stride + def test_layer_outer_pads_only_unaligned_pages(self): + """Layer-outer block strides are pages, so only an unaligned page pads.""" + mla_page = _mla(512).page_size_bytes + idx = replace(_mla(128), block_stride_alignment=12288) + groups = [_uniform_group({"mla.0": _mla(512), "idx.0": idx})] + config = get_kv_cache_config_from_groups( + _mock_vllm_config("LBNHC"), groups, MEMORY + ) + mla_tensor, idx_tensor = config.kv_cache_tensors + assert config.num_blocks == MEMORY // (mla_page + 12288) + assert (mla_tensor.block_stride, idx_tensor.block_stride) == (mla_page, 12288) + assert idx_tensor.offset == mla_page * config.num_blocks + _bind(config, "LBNHC") + def test_overlaid_groups_alias_and_stay_isolated(self): groups, g1, g2 = _mixed_page_groups() # Overlay models resolve to a block-outer layout at backend selection (the @@ -562,7 +576,9 @@ def test_overlaid_groups_alias_and_stay_isolated(self): _mock_vllm_config("BLNHC"), groups, MEMORY ) assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups) - assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups) + assert _pool_bytes_per_block( + _mock_vllm_config("BLNHC"), groups + ) == _expected_bytes_per_block(groups) views = _bind(config, "BLNHC") assert set(views) == set(g1) | set(g2) @@ -667,7 +683,10 @@ def test_ring_beside_paged_groups_keeps_the_kv_block_as_scheduler_block(self): assert len(ring_groups) == 1 assert not ring_groups[0].kv_cache_spec.prefix_cacheable kv_cache_config = get_kv_cache_config_from_groups( - config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + config, + groups, + available_memory=64 + * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC), ) assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64) manager = KVCacheManager( @@ -716,7 +735,10 @@ def test_non_cacheable_swa_leaves_hits_to_the_paged_group(self): for g in groups ) == [(False, 128)] * 4 + [(True, 0)] kv_cache_config = get_kv_cache_config_from_groups( - config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups) + config, + groups, + available_memory=64 + * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC), ) manager = KVCacheManager( generate_scheduler_kv_cache_config([kv_cache_config]), diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 53833ba83711..b06479f1caf1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -10,7 +10,7 @@ from dataclasses import replace from pathlib import Path from types import SimpleNamespace -from typing import Any, cast +from typing import Any import pytest import torch @@ -19,7 +19,6 @@ import vllm.v1.hisparse.runtime as hisparse_runtime_module from vllm.config import ( CacheConfig, - KVTransferConfig, ModelConfig, SchedulerConfig, VllmConfig, @@ -69,7 +68,6 @@ HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, @@ -1942,13 +1940,12 @@ def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch): config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) config.parallel_config.decode_context_parallel_size = 4 config.cache_config.kv_cache_layout = None - tail = KpoolTailSpec( + tail = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=4, dcp_sharded=False, ) layout = resolve_kv_cache_layout(config, [["LBHNC"]], [new_mla_spec(), tail]) @@ -2906,564 +2903,77 @@ def test_generate_scheduler_kv_cache_config(): ) -def _glm5_like_kv_cache_spec( - mamba_spec_factory=new_mamba_spec, -) -> tuple[dict[str, KVCacheSpec], list[str]]: - """(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba.""" +def test_packed_groups_glm5_like_hybrid(): + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + vllm_config.cache_config.kv_cache_layout = "BLHNC" + kpool, block_size = 4, 1024 kv_cache_spec: dict[str, KVCacheSpec] = {} - mamba_layers = [] - for i in range(45): - if i % 4 == 3: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=16, - ) - else: - name = f"layers.{i}.linear_attn" - kv_cache_spec[name] = mamba_spec_factory() - mamba_layers.append(name) - return kv_cache_spec, mamba_layers - - -def _glm5_like_kv_cache_spec_with_tail( - mamba_spec_factory=new_mamba_spec, -) -> dict[str, KVCacheSpec]: - """Production kpool=4 proportions: (mamba*3, MLA + indexer + tail) * 11. - - The tail's logical page (2 * kpool * 2*indexer_head_dim * bf16) must fit - inside the indexer page it parasitizes (block_size//kpool * 132 B), which - is why the fixture drops the tokens_per_state=16 shape used above. - """ - kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory) - for i in range(3, 45, 4): - kpool = 4 - kv_cache_spec[f"layers.{i}.indexer"] = replace( - cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]), + for i in range(11): + kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + block_stride_alignment=32 * 576 * 2, + ) + kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, tokens_per_state=kpool, + block_stride_alignment=64 * 132, ) - kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec( + kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( block_size=kpool, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=kpool, ) - return kv_cache_spec - - -def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]: - return { - layer_name: tensor - for tensor in kv_cache_config.kv_cache_tensors - for layer_name in tensor.layers - } - - -def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int: - return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride - - -def test_get_kv_cache_config_balanced_mamba_hybrid(): - """Hybrid slot sharing: mamba layers co-own the MLA slot tensors.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - uniform_groups = [ - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - ] - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) - ] - - # Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer - # gets an MLA slot. - assert len(groups) == 5 - assert len(uniform_groups) == 1 - assert len(uniform_groups[0].layer_names) == 22 - assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11 - assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] - for k, name in enumerate(mamba_layers): - assert name in mamba_groups[k % 4].layer_names - # Mamba pages are padded up to the MLA page. - for group in mamba_groups: - assert group.kv_cache_spec.page_size_padded == mla_page - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - # Every block id is charged the full per-block sum. - attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config) - mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks - blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group - assert ( - kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) - == blocks_per_request * bytes_per_block - ) - - available_memory = bytes_per_block * 100 + 1 - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, available_memory - ) - assert kv_cache_config.num_blocks == 100 - - # Every logical layer has a view into one backing allocation. Mamba views - # alias their corresponding MLA slot by using the same byte offset. - assert len(kv_cache_config.kv_cache_tensors) == 56 - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - tensors = _tensor_by_layer(kv_cache_config) - mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)] - for i, mla_name in enumerate(mla_layer_names): - mla_tensor = tensors[mla_name] - assert mla_tensor.block_stride == mla_page - for group in mamba_groups: - if i < len(group.layer_names): - assert tensors[group.layer_names[i]].offset == mla_tensor.offset - for i in range(11): - assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page - - total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors})) - assert total_allocated == bytes_per_block * 100 - assert 0 <= available_memory - total_allocated < bytes_per_block - - assert get_max_concurrency_for_kv_cache_config( - vllm_config, kv_cache_config - ) == pytest.approx(100 / blocks_per_request) - - -def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor(): - """The kpool tail parasitizes the indexer tensors instead of getting its - own: sibling idx/tail tensors paired by layer order, zero standalone tail - bytes, one shared block per request, and no prefix-caching leakage from - the kpool-sized scratch group.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec = _glm5_like_kv_cache_spec_with_tail() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes - assert tail_logical_page < idx_page - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - tail_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - and all( - isinstance(spec, KpoolTailSpec) - for spec in group.kv_cache_spec.kv_cache_specs.values() + for i in range(34): + kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( + shapes=((3, 6144), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + block_size=block_size, ) - ) - # The tail never prefix-caches; its page is padded up to the indexer page - # so the runner's strided view rides the indexer storage. - assert not tail_group.kv_cache_spec.prefix_cacheable - tail_inner = cast( - KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"] - ) - assert tail_inner.page_size_padded == idx_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - # Tail layers get logical views but no additional storage: each view aliases - # its sibling indexer at the same offset in the shared allocation. - assert len(kv_cache_config.kv_cache_tensors) == 67 - tensors = _tensor_by_layer(kv_cache_config) - for i in range(11): - idx_name = f"layers.{4 * i + 3}.indexer" - tail_name = f"layers.{4 * i + 3}.tail" - assert tensors[idx_name].offset == tensors[tail_name].offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - # The layout detector surfaces the sibling names in layer order for the - # accounting and connector paths. - layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups) - assert layout is not None - assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)] - assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)] - - # Accounting charges exactly one extra shared block per request for the - # tail, not a per-sequence allocation. - attn_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - and group is not tail_group - ) - attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config) - mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks - blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1 - assert ( - kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups) - == blocks_per_request * bytes_per_block - ) - - -def test_glm5_kpool_tail_does_not_drag_hash_block_size(): - """The tail's kpool-sized scratch block (4 tokens) must not constrain the - prefix-cache hash granularity: participating groups alone decide it.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - def align_mamba(): - return new_mamba_spec(mamba_cache_mode="align") - - groups = kv_cache_utils.get_kv_cache_groups( - vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba) - ) - kv_cache_config = KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=groups, - ) - hash_vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ) - # gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is - # the lcm including the tail (1024 % 4 == 0, so it coincides). - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, hash_vllm_config - ) == (1024, 16) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible(): - """Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - # (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page. - def big_mamba_spec(): - return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,)) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert big_mamba_spec().page_size_bytes > mla_page - - with pytest.raises(ValueError, match="does not fit the MLA page"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer(): - """Use the generic layout error when no kpool indexer is present.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBHNC" - - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(27): - if i % 4 == 3 or i == 26: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec( - shapes=((512, 1024),), dtypes=(torch.float32,) - ) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page - - with pytest.raises(NotImplementedError, match="page size"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba(): - """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not - disable slot sharing; the layout re-pads them to the MLA page.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - def prepadded_mamba_spec(): - return new_mamba_spec(page_size_padded=294_912) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - assert prepadded_mamba_spec().page_size_bytes < mla_page groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) ] - assert len(groups) == 5 - assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8] - for group in mamba_groups: - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - assert len(kv_cache_config.kv_cache_tensors) == 56 - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection(): - """Round-robin mamba grouping keeps practical PP splits balanced: every - stage's largest projected mamba group slice fits its projected MLA - layers, so per-stage slot tensors all have an MLA owner.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - - # PP=2 split at transformer layer 22/23. - # Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer) - # layers, projected mamba groups [5, 5, 4, 4]. - worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22} - groups = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) ] - assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4] - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 5 * mla_page + 5 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - tensors = _tensor_by_layer(kv_cache_config) - # Every projected Mamba slot aliases an MLA owner. - assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset - assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset - # The last slot only hosts the two groups with a 5th projected layer. - assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset - assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch): - """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA - layers but a round-robin slice of 3 under the minimum 4 mamba groups; - grouping bumps to 5 groups so every stage's projection keeps sharing - on instead of silently falling back on stage 0.""" - from vllm.config import ParallelConfig - from vllm.distributed.utils import get_pp_indices - - monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=8192), - parallel_config=ParallelConfig(pipeline_parallel_size=4), - ) - - kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec() - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)] - assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6] - for k, name in enumerate(mamba_layers): - assert name in mamba_groups[k % 5].layer_names - - for rank in range(4): - start, end = get_pp_indices(45, rank, 4) - worker_spec = { - n: s - for n, s in kv_cache_spec.items() - if start <= int(n.split(".")[1]) < end - } - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - groups, worker_spec + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec ) - assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None - - -def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch): - """Reject PP stages whose Mamba layers have no MLA slot to share.""" - from vllm.config import ParallelConfig - - monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45) - monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42") - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=8192), - parallel_config=ParallelConfig(pipeline_parallel_size=2), - ) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"): - kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - -def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag(): - """Hybrid slot sharing must take precedence over the experimental - enable_cross_layers_blocks packed layout: generic packing would give the - MLA slots a strided view, breaking the contiguous virtual split.""" - model_config = ModelConfig(max_model_len=8192) - kv_transfer_config = KVTransferConfig( - kv_connector="NixlConnector", - kv_role="kv_both", - kv_connector_extra_config={"enable_cross_layers_blocks": "true"}, - ) - vllm_config = VllmConfig( - model_config=model_config, kv_transfer_config=kv_transfer_config - ) - - kv_cache_spec, _ = _glm5_like_kv_cache_spec() - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 11 * mla_page + 11 * idx_page - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - assert len(kv_cache_config.kv_cache_tensors) == 56 - assert all( - tensor.block_stride in (mla_page, idx_page) - for tensor in kv_cache_config.kv_cache_tensors - ) - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer(): - """Kimi-Linear-like: MLA without indexer layers, idx_stride == 0.""" - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6]. - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(27): - if i % 4 == 3 or i == 26: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() - mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - mamba_groups = [ - group for group in groups if isinstance(group.kv_cache_spec, MambaSpec) ] - assert len(groups) == 4 - assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6] - for group in mamba_groups: - assert group.kv_cache_spec.page_size_bytes == mla_page - - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - assert bytes_per_block == 7 * mla_page + assert len(attn) == 1 and len(attn[0].layer_names) == 22 + assert len(rings) == 1 and len(rings[0].layer_names) == 11 + # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. + assert len(mamba) == 3 + assert sum(len(g.layer_names) for g in mamba) == 34 kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - assert kv_cache_config.num_blocks == 100 - # Each group has one tensor descriptor and all groups alias the same slots. - assert len(kv_cache_config.kv_cache_tensors) == 4 - tensors = _tensor_by_layer(kv_cache_config) - mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)] - for index, mla_name in enumerate(mla_names): - mla_offset = _layer_offset(tensors[mla_name], mla_name) - for group in mamba_groups: - if index < len(group.layer_names): - mamba_name = group.layer_names[index] - assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset - assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100} - - -def test_get_kv_cache_capacity_after_scheduler_unwrap(): - """max_concurrency must survive the scheduler-config unwrap. - - Regression for the balanced-mamba hybrid layout: - ``generate_scheduler_kv_cache_config`` flattens the MLA - ``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the - scheduler config holds an MLA spec (page size A) next to ``MambaSpec`` - groups (page size B). ``EngineCore._initialize_kv_caches`` calls - ``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size - path used to assert-fail (``assert len(page_sizes) == 1``) on this topology. - """ - model_config = ModelConfig(max_model_len=8192) - vllm_config = VllmConfig(model_config=model_config) - - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(45): - if i % 4 == 3: - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=1024, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=16, - ) - else: - kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec() - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups) - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, bytes_per_block * 100 + 1 - ) - - scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config]) - # Confirm we are exercising the regression path: the MLA group is no longer - # UniformTypeKVCacheSpecs, so sibling specs have differing page sizes. - assert not any( - isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - for g in scheduler_config.kv_cache_groups - ) - - unwrapped_groups = scheduler_config.kv_cache_groups - expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups( - vllm_config, unwrapped_groups - ) - expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups) - expected_blocks_per_request = ( - expected_max_mem + expected_pool - 1 - ) // expected_pool - - _, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config) - assert max_concurrency > 0 - assert max_concurrency == pytest.approx( - scheduler_config.num_blocks / expected_blocks_per_request + vllm_config, groups, available_memory=2 * 1024**3 ) + strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} + assert len(strides) == 1 + (block_stride,) = strides + # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages + # apart so both caches can be re-paged for their kernels. + assert block_stride % (32 * 576 * 2) == 0 + assert block_stride % (64 * 132) == 0 def new_mla_spec(cache_dtype_str=None, block_size: int = 16): @@ -4483,6 +3993,7 @@ def test_kv_cache_reserves_null_block_for_max_model_len(use_override): block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. @@ -4528,6 +4039,7 @@ def test_check_enough_kv_cache_memory_reserves_null_block(): block_size = 16 max_model_len = 512 # needs 512 / 16 = 32 blocks vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" spec = new_kv_cache_spec(block_size=block_size) # 32 blocks -> only 31 usable after the null block: one short -> reject. diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 415609aefb8d..f926f496514a 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -20,7 +20,6 @@ ChunkedLocalAttentionManager, CircularBufferManager, FullAttentionManager, - KpoolTailManager, MambaManager, RSWAManager, SlidingWindowManager, @@ -29,7 +28,6 @@ ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, MambaSpec, RSWASpec, SlidingWindowSpec, @@ -406,15 +404,15 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): @pytest.mark.parametrize("record_for_zeroing", [True, False]) def test_external_kpool_tail_zeroing(record_for_zeroing): - """A transferred ring at index zero skips zeroing even with a local hit.""" - spec = KpoolTailSpec( + """The kpool tail ring is never zeroed, transferred or not: a pool only + reads ring slots its own request wrote.""" + spec = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16, - sliding_window=4, ) - manager = KpoolTailManager( + manager = CircularBufferManager( spec, block_pool=BlockPool(10, enable_caching=True, hash_block_size=4), enable_caching=True, @@ -428,9 +426,7 @@ def test_external_kpool_tail_zeroing(record_for_zeroing): ) blocks = manager.req_to_blocks["request"] assert len(blocks) == 1 - assert manager.take_new_block_ids() == ( - [blocks[0].block_id] if record_for_zeroing else [] - ) + assert manager.take_new_block_ids() == [] assert manager.allocate_new_blocks("request", 20, 20) == [] assert manager.take_new_block_ids() == [] diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py index 58c6e632d3d9..85862da685a7 100644 --- a/tests/v1/kv_connector/unit/conftest.py +++ b/tests/v1/kv_connector/unit/conftest.py @@ -23,5 +23,5 @@ def _mock_cpu_backend_block_sizes(monkeypatch): monkeypatch.setattr( CPUAttentionBackend, "get_supported_kernel_block_sizes", - staticmethod(lambda: [MultipleOf(16)]), + staticmethod(lambda kv_cache_spec=None: [MultipleOf(16)]), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 23856b4541aa..969bf65f8c54 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -164,13 +164,12 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8) - gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8) + gdn_block_stride = gdn_spec.page_size_bytes + 64 + gdn_raw = torch.empty(num_blocks, gdn_block_stride, dtype=torch.int8) (fa_cache,) = dense_kv_cache_views( fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC ) - (gdn_cache,) = dense_kv_cache_views( - gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC - ) + gdn_cache = gdn_raw[:, : gdn_spec.page_size_bytes] worker.register_kv_caches( { @@ -184,7 +183,7 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): "model.layers.0.self_attn", "model.layers.1.linear_attn", ] - assert worker.block_len_per_layer == [ + assert worker.kv_block_len_per_layer == [ fa_spec.page_size_bytes, gdn_spec.page_size_bytes, ] diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6c2cff0a542c..a8354390d71e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -31,7 +31,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -711,13 +710,12 @@ def test_worker_setup_tolerates_finer_scratch_group(): mamba_cache_mode="align", ) # Scratch block 4 is not divisible by the hash unit 8 below. - scratch = KpoolTailSpec( + scratch = CircularBufferSpec( block_size=4, num_kv_heads=2, head_size=64, head_size_v=0, dtype=torch.bfloat16, - sliding_window=4, ) cfg = KVCacheConfig( num_blocks=4, diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 9f6bc56deeb5..4e7a7b843143 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -641,164 +641,43 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo @pytest.mark.cpu_test -@pytest.mark.parametrize("logical_block_size", [1152, 640]) -@pytest.mark.parametrize("tail_first", [False, True]) -def test_register_compressed_indexer_uses_virtual_transfer_pages( - logical_block_size, tail_first -): - """Compressed indexer rows must split into contiguous NIXL transfer pages.""" - from unittest.mock import MagicMock - - from vllm.config import set_current_vllm_config - from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( - base_worker as bw, - ) - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( - NixlConnectorWorker, - ) - from vllm.v1.kv_cache_interface import ( - KpoolTailSpec, - KVCacheConfig, - KVCacheGroupSpec, - KVCacheLayout, - KVCacheTensor, - MLAAttentionSpec, - UniformTypeKVCacheSpecs, - create_kv_cache_views, - ) - - num_logical_blocks = 3 - transfer_block_size = 64 - kernel_block_size = 128 - tokens_per_state = 4 - state_content_bytes = 132 - - indexer_spec = MLAAttentionSpec( - block_size=logical_block_size, - num_kv_heads=1, - head_size=128, - head_size_v=0, - dtype=torch.uint8, - state_content_bytes=state_content_bytes, - tokens_per_state=tokens_per_state, - ) - indexer_page_size = indexer_spec.page_size_bytes - tail_spec = KpoolTailSpec( - block_size=tokens_per_state, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - page_size_padded=indexer_page_size, - sliding_window=tokens_per_state, - ) - - allocation_size = num_logical_blocks * indexer_page_size - indexer_tensor = KVCacheTensor( - size=allocation_size, - layers=["indexer"], - layer_stride=allocation_size, - block_stride=indexer_page_size, - ) - tail_tensor = KVCacheTensor( - size=allocation_size, - layers=["tail"], - layer_stride=allocation_size, - block_stride=indexer_page_size, - ) - kv_cache_config = KVCacheConfig( - num_blocks=num_logical_blocks, - kv_cache_tensors=[indexer_tensor, tail_tensor], - kv_cache_groups=[ - KVCacheGroupSpec( - ["indexer"], - UniformTypeKVCacheSpecs( - block_size=logical_block_size, - kv_cache_specs={"indexer": indexer_spec}, - ), - ), - KVCacheGroupSpec( - ["tail"], - UniformTypeKVCacheSpecs( - block_size=tokens_per_state, - kv_cache_specs={"tail": tail_spec}, - ), - ), - ], - ) - - raw = torch.zeros(allocation_size, dtype=torch.int8) - (indexer_cache,) = create_kv_cache_views( - raw, - indexer_spec, - num_logical_blocks, - KVCacheLayout.LBHNC, - indexer_tensor, - kernel_block_size=kernel_block_size, +def test_nixl_keeps_packed_sparse_mla_manager_block(): + """Like the worker, NIXL keeps packed sparse MLA + indexer manager blocks.""" + from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend + from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend + from vllm.v1.kv_cache_interface import UniformTypeKVCacheSpecs + + specs = { + "mla": MLAAttentionSpec( + 1152, num_kv_heads=1, head_size=512, dtype=torch.bfloat16 + ), + "indexer": MLAAttentionSpec( + 1152, num_kv_heads=1, head_size=128, dtype=torch.uint8, tokens_per_state=4 + ), + } + backends = {"mla": FlashMLASparseBackend, "indexer": Glm5NextIndexerBackend} + layers = { + name: MagicMock(spec=AttentionLayerBase, **{"get_attn_backend.return_value": b}) + for name, b in backends.items() + } + worker = object.__new__(bw.NixlBaseConnectorWorker) + worker.block_size, worker.num_blocks, worker.attn_backends = 1152, 2, [] + worker._physical_blocks_per_logical_kv_block = 1 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC + ), + compilation_config=SimpleNamespace(static_forward_context=layers), ) - (tail_cache,) = create_kv_cache_views( - raw, - tail_spec, - num_logical_blocks, - KVCacheLayout.LBHNC, - tail_tensor, + worker.kv_cache_config = KVCacheConfig( + 2, [], [KVCacheGroupSpec(list(specs), UniformTypeKVCacheSpecs(1152, specs))] ) - assert indexer_cache.data_ptr() == tail_cache.data_ptr() == raw.data_ptr() - vllm_config = create_vllm_config(block_size=logical_block_size) - vllm_config.cache_config.kv_cache_layout = "LBHNC" - vllm_config.kv_transfer_config.kv_buffer_device = "cuda" - fake_backend = MagicMock() - fake_backend.get_supported_kernel_block_sizes.return_value = [transfer_block_size] - fake_backend.get_name.return_value = "DEEPSEEK_V32_INDEXER" - fake_backend.full_cls_name.return_value = "fake.DEEPSEEK_V32_INDEXER" - fake_platform = MagicMock() - fake_platform.device_type = "cuda" - fake_platform.get_nixl_memory_type.return_value = "VRAM" - - caches = [("indexer", indexer_cache), ("tail", tail_cache)] - if tail_first: - caches.reverse() + worker._sync_block_size_with_kernel() - with ( - patch.object(bw, "NixlWrapper", _RecordingNixl), - patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), - patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), - patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), - patch.object(bw, "current_platform", fake_platform), - set_current_vllm_config(vllm_config), - ): - worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) - worker.use_mla = True - worker.register_kv_caches(dict(caches)) - - transfer_page_size = transfer_block_size // tokens_per_state * state_content_bytes - num_transfer_blocks = num_logical_blocks * ( - logical_block_size // transfer_block_size - ) - expected_descs = np.asarray( - [ - [ - raw.data_ptr() + block_idx * transfer_page_size, - transfer_page_size, - 0, - ] - for block_idx in range(num_transfer_blocks) - ], - dtype=np.uint64, - ) - - assert worker.block_size == transfer_block_size - assert worker.num_regions == 1 - assert worker.block_len_per_layer == [transfer_page_size] - assert worker.block_stride_per_layer == [transfer_page_size] - assert worker._region_is_mla == [True] - assert worker.kv_caches_base_addr[worker.engine_id][0] == [raw.data_ptr()] - assert worker._registered_descs[0] == [(raw.data_ptr(), raw.nbytes, 0, "")] - np.testing.assert_array_equal(worker.src_blocks_data, expected_descs) - assert expected_descs[-1, 0] + expected_descs[-1, 1] == ( - raw.data_ptr() + raw.nbytes - ) + assert worker.block_size == 1152 + assert worker._physical_blocks_per_logical_kv_block == 1 def _make_remote_meta( diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 08759bc1e34b..6c0cc310c729 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -47,7 +47,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, @@ -145,13 +144,12 @@ def _make_scratch_kv_cache_config( register_all_kvcache_specs(vllm_config=None) fa_config = _make_kv_cache_config(num_blocks, num_groups=1) scratch_layers = ["layer_scratch"] - scratch_spec = KpoolTailSpec( + scratch_spec = CircularBufferSpec( block_size=scratch_block_size, num_kv_heads=2, head_size=HEAD_SIZE, head_size_v=0, dtype=DTYPE, - sliding_window=scratch_block_size, ) assert not scratch_spec.prefix_cacheable scratch_bytes = scratch_spec.page_size_bytes * num_blocks diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 1d004c36677a..3fc5caec2e08 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -102,7 +102,7 @@ class InputsPrepared(Exception): block_tables=SimpleNamespace( slot_mappings=query_slot_mapping.unsqueeze(0), input_block_tables=[block_table], - kernel_block_sizes=[4], + block_sizes=[4], cp_rank=cp_rank, cp_size=cp_size, cp_interleave=cp_interleave, diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 2ddd0a977481..62a9a1635c1f 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -963,6 +963,9 @@ def create_deterministic_logits(token_ids): # Mock draft_attn_groups for attention metadata building. mock_attn_group = mock.MagicMock() mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder + mock_attn_group.build_metadata_for_drafting = ( + attn_metadata_builder.build_for_drafting + ) mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec proposer.draft_attn_groups = [mock_attn_group] @@ -1066,6 +1069,9 @@ def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel): ) mock_attn_group = mock.MagicMock() mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder + mock_attn_group.build_metadata_for_drafting = ( + attn_metadata_builder.build_for_drafting + ) mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec proposer.draft_attn_groups = [mock_attn_group] diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py index fb1ecdd2df12..086fd467acd4 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer.py +++ b/tests/v1/spec_decode/test_llm_base_proposer.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for SpecDecodeBaseProposer.initialize_attn_backend. -Block tables are stored at kernel-block granularity, so the proposer's -``block_size`` (used for slot-mapping math) must be the kernel block size, -not the KV cache manager's block size — the two differ when manager blocks -are split for the attention kernel. The value must also be deterministic: +Block tables are stored at manager-block granularity, so the proposer's +``block_size`` (used for slot-mapping math) must be the KV cache manager's +block size. Attention groups map to kernel blocks for their metadata builders. +The value must also be deterministic: ``_draft_attn_layer_names`` is a set, whose iteration order varies across processes, so anything derived from iteration order must not leak into ``block_size``. @@ -68,9 +68,9 @@ def _make_kv_cache_config(layer_names: set[str]) -> KVCacheConfig: return SimpleNamespace(kv_cache_groups=[group]) # type: ignore[return-value] -def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): - """The proposer's slot-mapping math runs against the kernel-granularity - block table, so block_size must come from kernel_block_sizes.""" +def test_block_size_uses_manager_block_size(monkeypatch: pytest.MonkeyPatch): + """The proposer's slot-mapping math runs against the manager-granularity + block table; attention groups map kernel blocks for metadata builders.""" layer_names = {"draft.0.self_attn.attn"} proposer = _make_proposer(monkeypatch, layer_names) @@ -79,12 +79,8 @@ def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): kernel_block_sizes=[KERNEL_BLOCK_SIZE], ) - assert proposer.block_size == KERNEL_BLOCK_SIZE - assert proposer.block_size != SCHEDULER_BLOCK_SIZE - # The metadata builder keeps receiving the kernel block size as well. - group = proposer.draft_attn_groups[0] - assert isinstance(group, _FakeAttentionGroup) - assert group.kernel_block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size == SCHEDULER_BLOCK_SIZE + assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch): @@ -112,4 +108,4 @@ def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch) ) assert len(proposer.draft_attn_groups) == 1 assert proposer.draft_attn_groups[0].layer_names == expected_order - assert proposer.block_size == KERNEL_BLOCK_SIZE + assert proposer.block_size == SCHEDULER_BLOCK_SIZE diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index 9457c66050db..7002176dcc57 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -38,7 +38,6 @@ def mock_model_runner_with_input_batch(): device="cpu", vocab_size=32000, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], logitsprocs=None, is_pooling_model=False, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9c9c26c3c2c0..8d30403f30fc 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -15,7 +15,11 @@ import torch import vllm.v1.hisparse.binding as attn_utils_module -from tests.v1.attention.utils import dense_kv_cache_views +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + dense_kv_cache_views, +) from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -42,6 +46,8 @@ AttentionGroup, allocate_kv_cache, copy_kv_cache_blocks_inplace, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, ) @@ -79,7 +85,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( backend = SimpleNamespace( get_name=lambda name=name: name, customize_spec=AttentionBackend.customize_spec, - get_supported_kernel_block_sizes=lambda sizes=sizes: sizes, + get_supported_kernel_block_sizes=lambda _=None, sizes=sizes: sizes, ) layers[name] = SimpleNamespace( get_kv_cache_spec=lambda _, spec=specs[name]: spec, @@ -99,7 +105,8 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( resolved = attn_utils.get_kv_cache_spec(config) assert resolved["main"].block_size == resolved["indexer"].block_size == expected - assert resolved["dense"] is specs["dense"] + dense_backend = layers["dense"].get_attn_backend() + assert resolved["dense"] == customize_attention_spec(dense_backend, specs["dense"]) assert all(spec.block_size == block_size for spec in specs.values()) @@ -495,20 +502,17 @@ def test_reshape_padded_kv_cache_strides_by_padded_page(): @pytest.mark.parametrize( ( "kernel_block_sizes", - "storage_block_size", "expected_num_blocks", "expected_num_states", ), [ - (None, None, 4, 64), - ([256], None, 4, 64), - ([64], None, 16, 16), - ([64], 256, 4, 64), + (None, 4, 64), + ([256], 4, 64), + ([64], 16, 16), ], ) def test_allocate_compressed_mla_cache( kernel_block_sizes: list[int] | None, - storage_block_size: int | None, expected_num_blocks: int, expected_num_states: int, ): @@ -518,7 +522,6 @@ def test_allocate_compressed_mla_cache( head_size=128, dtype=torch.bfloat16, tokens_per_state=4, - storage_block_size=storage_block_size, ) num_pages = 4 config = KVCacheConfig( @@ -744,3 +747,92 @@ def host_allocator(size: int) -> torch.Tensor: backing.untyped_storage().data_ptr() == caches["indexer"].untyped_storage().data_ptr() ) + + +class _TableBuilder: + requires_block_table_width = False + supports_update_block_table = False + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device): + self.kv_cache_spec = kv_cache_spec + + def build(self, common_prefix_len, common_attn_metadata, fast_build=False): + return common_attn_metadata.block_table_tensor + + +_VLLM_CONFIG = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=8192), + scheduler_config=SimpleNamespace(max_num_seqs=4), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), +) + + +def _packed_group(tokens_per_state, row_bytes, supported_sizes): + spec = MLAAttentionSpec( + block_size=1024, + num_kv_heads=1, + head_size=128, + dtype=torch.uint8, + state_content_bytes=row_bytes, + tokens_per_state=tokens_per_state, + ) + backend = SimpleNamespace( + get_builder_cls=lambda: _TableBuilder, + get_supported_kernel_block_sizes=lambda _=None: supported_sizes, + get_name=lambda: "FAKE", + customize_spec=lambda spec: spec, + ) + return AttentionGroup(backend, ["layer.0"], spec, 0) + + +def test_packed_compressed_group_gets_kernel_block_table(): + """A packed indexer cache whose kernel needs 64-state pages gets a table and + a bound view in kernel blocks that address the same packed-block rows.""" + row, stride_rows = 132, 5 * 64 + group = _packed_group(4, row, [128, 256]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024) + assert group.get_metadata_builder().kv_cache_spec.block_size == 256 + + backing = torch.arange(3 * stride_rows).repeat_interleave(row) + manager = backing.view(3, 1, stride_rows, row)[:, :, :256] + pages = map_kv_caches_to_kernel_blocks({"layer.0": manager}, [group])["layer.0"] + metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu") + metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32) + table = group.build_metadata(metadata) + assert table.tolist() == [[0, 1, 2, 3, 10, 11, 12, 13]] + assert torch.equal(pages[table[0, 4:], 0].flatten(0, 1), manager[2, 0]) + + +def test_split_group_maps_manager_block_table(): + """Groups split in a dense layout keep manager block tables; builders get + b * blocks_per_kv_block + j against the view allocated in kernel blocks.""" + group = _packed_group(4, 132, [128, 256]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=256) + kernel_view = torch.zeros(12, 1, 64, 132) + bound = map_kv_caches_to_kernel_blocks({"layer.0": kernel_view}, [group]) + assert bound["layer.0"] is kernel_view + + metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu") + metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32) + table = group.build_metadata(metadata) + assert table.tolist() == [[0, 1, 2, 3, 8, 9, 10, 11]] + + +@pytest.mark.parametrize( + ("block_size", "page_states"), [(1024, 64), (384, 32), (256, 0)] +) +def test_split_blocks_align_to_kernel_blocks(block_size, page_states): + """Backends taking only fixed kernel blocks align packed blocks to the kernel + block the final block size is split into, and not at all when unsplit.""" + group = _packed_group(4, 132, [128, 256]) + spec = customize_attention_spec(group.backend, group.kv_cache_spec) + spec = spec.copy_with_new_block_size(block_size) + assert spec.get_block_stride_alignment() == max(page_states * 132, 1) + + +def test_packed_token_cache_cannot_be_split(): + group = _packed_group(1, 1152, [64]) + group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024) + packed = torch.zeros(3, 1, 2048, 1152, dtype=torch.uint8)[:, :, :1024] + with pytest.raises(ValueError, match="cannot be split"): + map_kv_caches_to_kernel_blocks({"layer.0": packed}, [group]) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5c2231a1f267..dac9f12ae69a 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -22,7 +22,6 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch): max_num_batched_tokens=64, max_num_blocks_per_group=[8, 8, 8], device=device, - kernel_block_sizes=[16, 16, 8], ) def fail_if_apply_write_called(): @@ -48,33 +47,31 @@ def fail_if_apply_write_called(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) - # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two - # kernel block IDs. assert torch.equal( - block_tables.block_tables[1].gpu[0, :4], - torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[0, :2], + torch.tensor([10, 11], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[0].gpu[1, :1], torch.tensor([3], dtype=torch.int32, device=device), ) assert torch.equal( - block_tables.block_tables[1].gpu[1, :2], - torch.tensor([24, 25], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[1, :1], + torch.tensor([12], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[2].gpu[1, :2], torch.tensor([5, 6], dtype=torch.int32, device=device), ) assert block_tables.num_blocks.np[0, 0] == 2 - assert block_tables.num_blocks.np[1, 0] == 4 + assert block_tables.num_blocks.np[1, 0] == 2 assert block_tables.num_blocks.np[2, 0] == 0 assert block_tables.num_blocks.np[0, 1] == 1 - assert block_tables.num_blocks.np[1, 1] == 2 + assert block_tables.num_blocks.np[1, 1] == 1 assert block_tables.num_blocks.np[2, 1] == 2 assert torch.equal( block_tables.num_blocks.gpu[:, :2], - torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device), + torch.tensor([[2, 1], [2, 1], [0, 2]], dtype=torch.int32, device=device), ) for block_table in block_tables.block_tables: @@ -96,15 +93,15 @@ def fail_if_apply_write_called(): torch.tensor([1, 2, 7], dtype=torch.int32, device=device), ) assert torch.equal( - block_tables.block_tables[1].gpu[0, :6], - torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device), + block_tables.block_tables[1].gpu[0, :3], + torch.tensor([10, 11, 13], dtype=torch.int32, device=device), ) assert torch.equal( block_tables.block_tables[2].gpu[0, :1], torch.tensor([8], dtype=torch.int32, device=device), ) assert block_tables.num_blocks.np[0, 0] == 3 - assert block_tables.num_blocks.np[1, 0] == 6 + assert block_tables.num_blocks.np[1, 0] == 3 assert block_tables.num_blocks.np[2, 0] == 1 @@ -116,7 +113,6 @@ def test_block_tables_apply_staged_writes_single_group(): max_num_batched_tokens=16, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[16], ) block_tables.append_block_ids( @@ -141,7 +137,6 @@ def test_block_tables_skip_custom_slot_mapping_groups(): max_num_batched_tokens=4, max_num_blocks_per_group=[1, 1], device=device, - kernel_block_sizes=[8, 262144], slot_mapping_enabled=[False, True], ) block_tables.append_block_ids( @@ -170,7 +165,7 @@ def test_block_tables_skip_custom_slot_mapping_groups(): @pytest.mark.parametrize("cp_rank", range(4)) -def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): +def test_dcp_slot_mapping_only_shards_sharded_groups(cp_rank: int): """Only sharded groups use DCP interleave in logical-block coordinates.""" device = torch.device("cuda") block_tables = BlockTables( @@ -179,7 +174,6 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): max_num_batched_tokens=1024, max_num_blocks_per_group=[2, 8], device=device, - kernel_block_sizes=[64, 64], dcp_sharded=[True, False], cp_size=4, cp_rank=cp_rank, @@ -230,7 +224,6 @@ def test_v1_block_table_move_row_clears_vacated_row(): max_num_batched_tokens=64, pin_memory=False, device=torch.device("cuda"), - kernel_block_size=16, cp_kv_cache_interleave_size=1, ) block_table.add_row([7, 8, 9], row_idx=0) @@ -260,7 +253,6 @@ def test_get_dummy_block_tables_returns_zeroed_rows(): max_num_batched_tokens=64, max_num_blocks_per_group=[8], device=device, - kernel_block_sizes=[16], ) # Simulate a real step: stage a request's blocks and gather them into # the persistent input block tables. @@ -292,7 +284,6 @@ def test_dummy_request_slot_mapping_is_pad(): max_num_batched_tokens=16, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[4], ) block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) block_tables.apply_staged_writes() diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index f37e824e6c63..21157aae33b2 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -233,7 +233,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int): device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) reqs: list[CachedRequestState] = [] @@ -327,7 +326,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) ref_input_batch: InputBatch = InputBatch( @@ -337,7 +335,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis device=torch.device(device), vocab_size=1024, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -388,7 +385,6 @@ def test_condense_clears_stale_allowed_token_ids_mask(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -446,7 +442,6 @@ def test_swap_states_preserves_allowed_token_ids_mask(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[1], - kernel_block_sizes=[1], max_num_blocks_per_req=[1024], ) @@ -516,7 +511,6 @@ def test_pooling_prompt_lens_not_aliased(device: str): device=torch.device(device), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=True, ) @@ -552,7 +546,6 @@ def test_placeholder_spec_token_ids_written_verbatim(): device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[1], ) req = CachedRequestState( @@ -599,7 +592,6 @@ def test_pooling_metadata_token_id_buffers( device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=True, ) @@ -630,7 +622,6 @@ def _make_input_batch(is_pooling_model: bool = False) -> InputBatch: device=torch.device("cpu"), vocab_size=VOCAB_SIZE, block_sizes=[16], - kernel_block_sizes=[16], max_num_blocks_per_req=[64], is_pooling_model=is_pooling_model, ) diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py index 8a7445f7d49a..d9d5c4de63af 100644 --- a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py +++ b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py @@ -5,7 +5,7 @@ The kpool tail cache holds one circular block per request, so its block table row is orders of magnitude shorter than the sequence it serves. `_compute_slot_mappings_kernel` indexes block tables by -`position // kernel_block_size` with an effectively unmasked load, so +`position // block_size` with an effectively unmasked load, so this leads to a read of neighboring requests' rows (and for the later requests, outside the buffer entirely). """ @@ -15,7 +15,7 @@ from vllm.platforms import current_platform from vllm.v1.attention.backends.utils import PAD_SLOT_ID -from vllm.v1.kv_cache_interface import KpoolTailSpec +from vllm.v1.kv_cache_interface import CircularBufferSpec from vllm.v1.worker.block_table import get_block_table_width from vllm.v1.worker.gpu.block_table import BlockTables @@ -32,13 +32,12 @@ def make_tail_spec(): - return KpoolTailSpec( + return CircularBufferSpec( block_size=KPOOL, num_kv_heads=2, head_size=128, head_size_v=0, dtype=torch.bfloat16, - sliding_window=KPOOL, ) @@ -67,7 +66,6 @@ def test_kpool_tail_group_never_position_indexes_its_block_table(): max_num_batched_tokens=len(positions), max_num_blocks_per_group=[mla_width, tail_width], device=device, - kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size], slot_mapping_enabled=[True, spec.uses_slot_mapping], ) mla_blocks = list(range(1, mla_width + 1)) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 3851d5dad6bd..158befc15027 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -111,9 +111,6 @@ def initialize_kv_cache(runner: GPUModelRunner): device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[ - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size - ], max_num_blocks_per_req=[NUM_BLOCKS], ) runner.initialize_attn_backend(kv_cache_config) @@ -303,7 +300,7 @@ def get_name(): return "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return supported_sizes return _MockBackend() @@ -441,7 +438,7 @@ def get_name() -> str: return "MOCK_EXACT" if exact else "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return list(supported) if exact: @@ -1625,7 +1622,6 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) runner.device = torch.device("cpu") runner.is_pooling_model = False runner._init_block_sizes = [16] - runner._init_kernel_block_sizes = [16] runner._init_max_num_blocks = [4] runner._init_slot_mapping_modes = [ gpu_model_runner_module.SlotMappingMode.TOKEN_TO_KV_SLOT @@ -1653,7 +1649,7 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) lambda _: gpu_model_runner_module.KVCacheSpecKind.FULL_ATTENTION, ) - runner.may_reinitialize_input_batch(kv_cache_config, [16]) + runner.may_reinitialize_input_batch(kv_cache_config) assert input_batch_cls.call_count == 1 assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16 @@ -1681,68 +1677,10 @@ class StopInitialization(Exception): assert runner.cp_interleave == 16 -def test_hybrid_block_table_initialization(): - """Test hybrid block table with different kernel and kvcache_manager block - sizes.""" - from vllm.v1.worker.block_table import BlockTable - - # Test configuration: kvcache_manager block size = 32, - # kernel block size = 16 - block_size = 32 - kernel_block_sizes = [16] - max_num_reqs = 10 - max_num_blocks_per_req = 20 - max_num_batched_tokens = 512 - cp_kv_cache_interleave_size = 8 - - block_table = BlockTable( - block_size=block_size, - max_num_reqs=max_num_reqs, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max_num_batched_tokens, - pin_memory=False, - device=torch.device(DEVICE_TYPE), - kernel_block_size=kernel_block_sizes[0], - cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, - ) - - # Verify hybrid block configuration - assert block_table.use_hybrid_blocks is True - assert block_table.block_size == kernel_block_sizes[0] - assert block_table.blocks_per_kv_block == ( - block_size // kernel_block_sizes[0] - ) # Changed to use first element - - # Test block table conversion logic - # One kvcache_manager block should map to multiple kernel blocks - kvcache_manager_blocks = [0, 1, 2] - - # Verify that kvcache_manager blocks can be converted to kernel blocks - # and that block table operations work correctly. - req_index = 0 - block_table.append_row(kvcache_manager_blocks, req_index) - # Get expected kernel blocks from the implementation for verification. - expected_kernel_blocks = block_table.map_to_kernel_blocks( - np.array(kvcache_manager_blocks), - block_table.blocks_per_kv_block, - block_table._kernel_block_arange, - ) - # Verify block table state - assert block_table.num_blocks_per_row[req_index] == len(expected_kernel_blocks) - assert np.array_equal( - block_table.block_table.np[req_index, : len(expected_kernel_blocks)], - expected_kernel_blocks, - ) - - def test_get_block_table_width_aligns_to_128_tokens(): assert get_block_table_width(1875, 64) == 1876 -def test_get_block_table_width_splits_virtual_blocks(): - assert get_block_table_width(235, 256, 64) == 940 - - def test_mamba_state_table_width_is_not_aligned(): block_tables = MultiGroupBlockTable( max_num_reqs=1, @@ -1750,7 +1688,6 @@ def test_mamba_state_table_width_is_not_aligned(): pin_memory=False, device=torch.device("cpu"), block_sizes=[39664], - kernel_block_sizes=[39664], max_num_blocks=[1], slot_mapping_modes=[SlotMappingMode.NONE], ) @@ -1758,42 +1695,6 @@ def test_mamba_state_table_width_is_not_aligned(): assert block_tables[0].max_num_blocks_per_req == 1 -def test_input_batch_with_kernel_block_sizes(): - """Test InputBatch initialization with kernel_block_sizes parameter.""" - max_num_reqs = 10 - max_model_len = 512 - max_num_batched_tokens = 512 - device = torch.device(DEVICE_TYPE) - vocab_size = 50272 - - # Test with different kernel block sizes - block_sizes = [32, 64] - kernel_block_sizes = [16, 32] - - input_batch = InputBatch( - max_num_reqs=max_num_reqs, - max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, - device=device, - vocab_size=vocab_size, - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks_per_req=[16, 8], - ) - - # Verify that block tables were created with kernel block sizes - assert len(input_batch.block_table.block_tables) == len(block_sizes) - - for i, (kv_size, kernel_size) in enumerate(zip(block_sizes, kernel_block_sizes)): - block_table = input_batch.block_table.block_tables[i] - if kv_size != kernel_size: - assert block_table.use_hybrid_blocks is True - assert block_table.block_size == kernel_size - else: - assert block_table.use_hybrid_blocks is False - assert block_table.block_size == kernel_size - - def test_hybrid_cache_integration(default_vllm_config, dist_init): """Test hybrid cache architecture integration with GPUModelRunner.""" # Create a new model runner with hybrid cache configuration @@ -1843,9 +1744,8 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[16], max_num_blocks_per_req=[NUM_BLOCKS], - ) # Use kernel block size + ) runner.initialize_attn_backend(kv_cache_config) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index 92207aa38f6b..dd9046de612d 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -119,9 +119,6 @@ def _initialize_kv_cache(runner: GPUModelRunner) -> None: device=runner.device, vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], - kernel_block_sizes=[ - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size - ], max_num_blocks_per_req=[NUM_BLOCKS], ) runner.initialize_attn_backend(kv_cache_config) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index dec382725bc1..47173da08155 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -237,7 +237,6 @@ def stage_write(self, *_args): block_tables = BlockTables.__new__(BlockTables) block_tables.num_kv_cache_groups = 1 - block_tables.blocks_per_kv_block = [1] block_tables.block_tables = [_BlockTable()] block_tables.num_blocks = SimpleNamespace( np=torch.tensor([[0, 3]], dtype=torch.int32) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index d4b464e01463..fcf171a89ca3 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -47,8 +47,7 @@ def _make_capture_manager(block_table: torch.Tensor): block_tables = SimpleNamespace( input_block_tables=(block_table,), num_kv_cache_groups=1, - kernel_block_sizes=(2,), - blocks_per_kv_block=(1,), + block_sizes=(2,), ) manager = PCPManager( pcp_world_size=2, diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index b63bd7a2ff05..a3ec5f4eb7f4 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -20,6 +20,7 @@ from vllm.v1.attention.backends import flash_attn as flash_attn_module from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import ( @@ -30,6 +31,7 @@ from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( TargetDependentARSpeculator, ) +from vllm.v1.worker.utils import AttentionGroup class _TestSpeculator(TargetDependentARSpeculator): @@ -102,6 +104,7 @@ def build(common_prefix_len, common_attn_metadata): get_metadata_builder=lambda _: SimpleNamespace( build=build, supports_update_block_table=False ), + build_metadata=lambda cm, ubatch_idx=0, **_: build(0, cm), layer_names=["draft"], ) ] @@ -427,6 +430,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected( query_start_loc=torch.arange(3), ) speculator.idx_mapping = torch.arange(2) + speculator.attn_groups = [] + speculator.block_tables = SimpleNamespace(input_block_tables=[]) generate_draft = Mock() speculator._generate_draft = generate_draft run_fullgraph = Mock() @@ -450,6 +455,52 @@ def test_multi_step_decode_replays_captured_graph_as_expected( assert run_fullgraph.call_count == expected_graph_replays +def test_fused_full_graph_maps_split_group_block_tables(): + """The fused FULL replay skips build(), so split groups' tables are mapped first.""" + spec = MLAAttentionSpec( + block_size=1024, num_kv_heads=1, head_size=128, dtype=torch.uint8 + ) + group = AttentionGroup( + SimpleNamespace(), + ["layer.0"], + spec, + 0, + kernel_block_size=256, + kernel_block_table=torch.zeros(1, 4, 32, dtype=torch.int32), + kernel_block_offsets=torch.arange(4, dtype=torch.int32), + ) + speculator = object.__new__(_TestSpeculator) + speculator.input_buffers = SimpleNamespace( + positions=torch.arange(2), query_start_loc=torch.arange(3) + ) + speculator.idx_mapping = torch.arange(2) + speculator.attn_groups = [[group]] + speculator.block_tables = SimpleNamespace( + input_block_tables=[torch.tensor([[1, 2], [0, 3]], dtype=torch.int32)] + ) + replayed_tables = [] + speculator.decode_cudagraph_manager = SimpleNamespace( + run_fullgraph=lambda _: replayed_tables.append( + group.kernel_block_table[0, :2, :8].tolist() + ) + ) + + speculator._fused_multi_step_decode( + num_reqs=2, + skip_attn=False, + batch_desc=BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, num_tokens=2, num_reqs=2 + ), + num_tokens_across_dp=None, + seq_lens_cpu_upper_bound=None, + num_speculative_steps=4, + ) + + assert replayed_tables == [ + [[4, 5, 6, 7, 8, 9, 10, 11], [0, 1, 2, 3, 12, 13, 14, 15]] + ] + + def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata( monkeypatch, ): diff --git a/tests/v1/worker/test_jit_warmup_migration.py b/tests/v1/worker/test_jit_warmup_migration.py index ef887c5351f6..1875e605e53a 100644 --- a/tests/v1/worker/test_jit_warmup_migration.py +++ b/tests/v1/worker/test_jit_warmup_migration.py @@ -12,40 +12,24 @@ from vllm.v1.worker.block_table import ComputeSlotMappingKernel -@pytest.mark.parametrize( - ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"), - [ - (256, 1, 256, 16), - (256, 4, 64, 16), - (64, 1, 64, 16), - (8, 1, 8, 2), - (4, 1, 4, 2), - ], -) +@pytest.mark.parametrize("kv_cache_block_size", [256, 64, 8, 4]) def test_compute_slot_mapping_warmup_matches_runtime_specializations( kv_cache_block_size: int, - blocks_per_kv_block: int, - block_size: int, - block_size_rep: int, ) -> None: kernel = ComputeSlotMappingKernel() kwargs = dict( kv_cache_block_size=kv_cache_block_size, - blocks_per_kv_block=blocks_per_kv_block, total_cp_world_size=1, total_cp_rank=0, cp_kv_cache_interleave_size=1, block_table_stride=32768, - block_size=block_size, ) expected = kernel.CompileKey( kv_cache_block_size=kv_cache_block_size, - blocks_per_kv_block=blocks_per_kv_block, total_cp_world_size=1, total_cp_rank=0, cp_kv_cache_interleave_size=1, block_table_stride=16, - block_size=block_size_rep, ) assert kernel.dispatch(**kwargs) == expected diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 7f4bea52c227..52cf2a374172 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -112,6 +112,8 @@ def bind(*args, **kwargs): model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")), compilation_config=SimpleNamespace(static_forward_context={}), kv_caches=[], + attn_groups=[], + speculative_config=None, ) result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors( runner, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 4407d6fc5430..7e588ccde2e5 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -659,7 +659,6 @@ def test_v2_align_ctx_binds_request_slot_tables_for_state_copies(): max_num_batched_tokens=64, max_num_blocks_per_group=[4], device=device, - kernel_block_sizes=[cfg.block_size], ) state = object.__new__(MambaHybridModelState) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index f6f7227954bf..e17f214871f3 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -443,6 +443,43 @@ def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): assert (backing[9] == -1).all() +def test_bind_kv_cache_keeps_scheduler_block_views_for_runner(): + """Block copies index scheduler blocks, so the runner keeps the unmapped + views while a packed layer is bound with kernel blocks.""" + + class _Layer: + def bind_kv_cache(self, kv_cache): + self.kv_cache = kv_cache + + num_blocks, row = 4, 64 + storage = torch.zeros(num_blocks * row) + # Two layers packed into one block row: an 8x4 page, then a 4x2 page + # split into two 2-row kernel blocks. + mla = storage.as_strided((num_blocks, 8, 4), (row, 4, 1)) + indexer = storage.as_strided((num_blocks, 1, 4, 2), (row, 8, 2, 1), 40) + # The indexer in 2-row kernel blocks: block b's kernel block j is b * 16 + j. + mapped = indexer.as_strided((50, 1, 2, 2), (4, 8, 2, 1)) + kv_caches = {"layers.0.attn": mla, "layers.1.indexer": indexer} + ctx = {name: _Layer() for name in kv_caches} + runner_kv_caches: list[torch.Tensor] = [] + + bind_kv_cache( + kv_caches, + ctx, + runner_kv_caches, + layer_kv_caches={**kv_caches, "layers.1.indexer": mapped}, + ) + + assert ctx["layers.1.indexer"].kv_cache is mapped + assert runner_kv_caches == [mla, indexer] + copies = [KVCacheBlockCopy(1, 3)] + with pytest.raises(AssertionError, match="not divisible"): + copy_kv_cache_blocks_inplace([mapped], num_blocks, copies) + storage[row : 2 * row] = 7 + copy_kv_cache_blocks_inplace(runner_kv_caches, num_blocks, copies) + assert (mla[3] == 7).all() and (indexer[3] == 7).all() + + @pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)]) @pytest.mark.parametrize("has_copies", [False, True]) def test_hisparse_host_copy_waits_for_writes_on_writer_only( diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index d314d6107ad7..d09e5da71b0e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -20,6 +20,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: @@ -351,6 +352,35 @@ def get_current_attn_backends( ] +def get_current_attn_backends_and_specs( + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + fallback: list[type[AttentionBackend]], +) -> tuple[list[type[AttentionBackend]], list[AttentionSpec | None]]: + """Distinct (backend, spec) pairs of the transfer layers, else ``fallback``. + + Compressed caches are skipped: their kernel pages never split transferred + blocks. + """ + pairs: dict[tuple[type[AttentionBackend], AttentionSpec | None], None] = {} + layer_type = cast(type[Any], AttentionLayerBase) + for group in kv_cache_config.transfer_groups: + specs = getattr(group.kv_cache_spec, "kv_cache_specs", {}) + layers = get_layers_from_vllm_config(vllm_config, layer_type, group.layer_names) + for name, layer in layers.items(): + spec = specs.get(name, group.kv_cache_spec) + if isinstance(spec, AttentionSpec) and spec.tokens_per_state > 1: + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + assert layout.is_block_outermost, "Compressed caches must be packed" + continue + if not isinstance(spec, AttentionSpec): + spec = None + pairs[layer.get_attn_backend(), spec] = None + if not pairs: + return fallback, [None] * len(fallback) + return [backend for backend, _ in pairs], [spec for _, spec in pairs] + + def get_current_attn_backend( vllm_config: VllmConfig, layer_names: list[str] | None = None ) -> type[AttentionBackend]: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 9d5194f33e27..41c1337ecd85 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -26,6 +26,7 @@ EngineId, TransferTopology, get_current_attn_backends, + get_current_attn_backends_and_specs, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -59,8 +60,8 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( AttentionSpec, + CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheSpec, MambaSpec, MLAAttentionSpec, @@ -69,7 +70,6 @@ ) from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.request import RequestStatus -from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size logger = init_logger(__name__) @@ -1350,7 +1350,10 @@ def _sync_block_size_with_kernel(self) -> None: # physical block sizes. Pick the common (smallest) block size so that # KV-cache registration and transfer work correctly for both models. backends = get_current_attn_backends(self.vllm_config) - kernel_block_size = select_common_block_size(self.block_size, backends) + backends, specs = get_current_attn_backends_and_specs( + self.vllm_config, self.kv_cache_config, backends + ) + kernel_block_size = select_common_block_size(self.block_size, backends, specs) if self.block_size != kernel_block_size: logger.info_once( "User-specified logical block size (%s) does not match" @@ -1727,11 +1730,12 @@ def _logical_to_kernel_block_ids( ) group_specs = self.kv_cache_config.transfer_groups return [ - BlockTable.map_to_kernel_blocks( - np.array(group), - self._physical_blocks_per_logical_kv_block, - block_arange, - ).tolist() + ( + np.array(group)[:, None] * self._physical_blocks_per_logical_kv_block + + block_arange + ) + .reshape(-1) + .tolist() if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) else group for i, group in enumerate(block_ids) @@ -2054,7 +2058,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # hetero PP can align by name (#56033). A non-contiguous row cannot # name a sub-span, so it is one region anchored at the allocation. use_packed_row = ( - not isinstance(layer_spec, (MambaSpec, KpoolTailSpec)) + not isinstance(layer_spec, (MambaSpec, CircularBufferSpec)) and storage_is_block_major and block_stride > 0 and (is_mla_region or not hnc_contiguous) @@ -2121,9 +2125,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): block_len = region_cache.stride(0) * region_cache.element_size() region_base_addresses.append(base_addr) - if isinstance(layer_spec, KpoolTailSpec): - kv_block_len = layer_spec.unpadded_page_size_bytes // 2 - elif isinstance(layer_spec, AttentionSpec) and block_is_contiguous: + if isinstance(layer_spec, AttentionSpec) and block_is_contiguous: assert ( layer_spec.page_size_bytes % self._physical_blocks_per_logical_kv_block @@ -2133,6 +2135,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): layer_spec.page_size_bytes // self._physical_blocks_per_logical_kv_block ) + elif isinstance(layer_spec, MambaSpec) and block_is_contiguous: + kv_block_len = layer_spec.page_size_bytes else: kv_block_len = block_len if kv_block_len > block_len: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 38a656954eee..778eb0f7560d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -28,6 +28,7 @@ EngineTransferInfo, TransferTopology, get_current_attn_backends, + get_current_attn_backends_and_specs, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, ) @@ -79,7 +80,6 @@ from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, - KpoolTailSpec, KVCacheLayout, KVCacheSpec, MambaSpec, @@ -88,7 +88,6 @@ UniformTypeKVCacheSpecs, iter_layer_specs, ) -from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size if TYPE_CHECKING: @@ -120,46 +119,6 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool: return len(block_strides) == len(storage_ptrs) == 1 -def _tensor_byte_span_end(cache: torch.Tensor) -> int: - """Return the exclusive end address touched by a nonnegative-stride view.""" - if cache.numel() == 0: - return cache.data_ptr() - if any(stride < 0 for stride in cache.stride()): - raise ValueError("NIXL cache views must have nonnegative strides") - max_element_offset = sum( - (size - 1) * stride for size, stride in zip(cache.shape, cache.stride()) - ) - return cache.data_ptr() + (max_element_offset + 1) * cache.element_size() - - -def _uses_dense_virtual_transfer_pages( - layer_spec: KVCacheSpec, - cache: torch.Tensor, - physical_page_size: int, - num_blocks: int, -) -> bool: - """Return whether a compressed kernel view can be split into NIXL pages.""" - if not ( - isinstance(layer_spec, MLAAttentionSpec) - and layer_spec.tokens_per_state > 1 - and cache.ndim == 4 - and cache.shape[1] == 1 - and cache.is_contiguous() - and physical_page_size > 0 - and layer_spec.state_content_size_bytes > 0 - ): - return False - - block_stride = cache.stride(0) * cache.element_size() - return ( - block_stride > physical_page_size - and block_stride % physical_page_size == 0 - and physical_page_size % layer_spec.state_content_size_bytes == 0 - and cache.shape[0] * (block_stride // physical_page_size) == num_blocks - and cache.nbytes == num_blocks * physical_page_size - ) - - class NixlBaseConnectorWorker: """Base implementation of Worker side methods shared by pull and push.""" @@ -985,8 +944,10 @@ def _validate_remote_parallel_config( ) def _sync_block_size_with_kernel(self) -> None: - backends = get_current_attn_backends(self.vllm_config) - kernel_block_size = select_common_block_size(self.block_size, backends) + backends, specs = get_current_attn_backends_and_specs( + self.vllm_config, self.kv_cache_config, self.attn_backends + ) + kernel_block_size = select_common_block_size(self.block_size, backends, specs) # Number of blocks not accounting for kernel block mismatches self._logical_num_blocks = self.num_blocks if self.block_size != kernel_block_size: @@ -1425,7 +1386,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): if self._supports_pp_hma and self._has_packed_cache: # PP can reorder the same backends across stages of a packed model. backend_name = ",".join( - sorted(backend.get_name() for backend in self.attn_backends) + sorted({backend.get_name() for backend in self.attn_backends}) ) self.compat_hash = compute_nixl_compatibility_hash( self.vllm_config, @@ -1478,31 +1439,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ): raise NotImplementedError("PP push with packed KV caches requires MLA") - layer_specs: dict[str, KVCacheSpec] = {} - compressed_region_owners: dict[int, torch.Tensor] = {} - for layer_name, cache in xfer_buffers.items(): - layer_spec = self._layer_specs.get(layer_name) - if isinstance(layer_spec, UniformTypeKVCacheSpecs): - layer_spec = layer_spec.kv_cache_specs[layer_name] - if layer_spec is None: - continue - layer_specs[layer_name] = layer_spec - physical_page_size = ( - layer_spec.page_size_bytes - if isinstance(layer_spec, MambaSpec) - else layer_spec.page_size_bytes - // self._physical_blocks_per_logical_kv_block - ) - num_blocks = ( - self._logical_num_blocks - if isinstance(layer_spec, MambaSpec) - else self.num_blocks - ) - if _uses_dense_virtual_transfer_pages( - layer_spec, cache, physical_page_size, num_blocks - ): - compressed_region_owners.setdefault(cache.data_ptr(), cache) - track_region_layers = self._tracks_region_layers() region_layers: list[list[str]] = [] packed_member_layouts: dict[str, tuple[int, int]] = {} @@ -1527,7 +1463,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # kernel requires a specific block size. This leads to SSM and FA layers # having different num_blocks. # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. - layer_spec = layer_specs.get(layer_name) + layer_spec = self._layer_specs.get(layer_name) if layer_spec is None: logger.debug( "Skipping layer %s as no KVCache spec is present. " @@ -1587,32 +1523,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_device_id, ) - if isinstance(layer_spec, KpoolTailSpec): - compressed_owner = compressed_region_owners.get(cache.data_ptr()) - if compressed_owner is not None: - owner_storage = compressed_owner.untyped_storage() - owner_end = compressed_owner.data_ptr() + compressed_owner.nbytes - tail_is_covered = ( - compressed_owner.is_contiguous() - and owner_storage.data_ptr() == storage_addr - and _tensor_byte_span_end(cache) <= owner_end - ) - if not tail_is_covered: - raise AssertionError( - "Kpool tail cache is not fully covered by its compressed " - f"indexer region: layer={layer_name}, " - f"tail_shape={tuple(cache.shape)}, " - f"tail_stride={tuple(cache.stride())}, " - f"owner_shape={tuple(compressed_owner.shape)}, " - f"owner_stride={tuple(compressed_owner.stride())}" - ) - logger.debug( - "Skipping layer %s because its compressed indexer region " - "covers the same storage", - layer_name, - ) - continue - if isinstance(layer_spec, MambaSpec): physical_ratio = self._physical_blocks_per_logical_kv_block block_len = physical_page_size // physical_ratio @@ -1647,9 +1557,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): and cache[0].is_contiguous() and cache[0].nbytes == physical_page_size ) - virtual_transfer_pages = _uses_dense_virtual_transfer_pages( - layer_spec, cache, physical_page_size, num_blocks - ) # Push workers address the pages of packed MLA rows by layer. packed_mla_push = ( track_region_layers @@ -1658,12 +1565,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): and storage_is_block_major and is_mla_region ) - if virtual_transfer_pages: - # A compressed kernel row can contain multiple NIXL transfer pages. - region_specs = [ - (cache.data_ptr(), physical_page_size, physical_page_size) - ] - elif packed_mla_push and use_layer_name_routing: + if packed_mla_push and use_layer_name_routing: # PP>1 prefill stage: register only this layer's page of each row. region_specs = [ (cache.data_ptr(), physical_page_size, block_stride) @@ -1671,7 +1573,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): elif packed_mla_push: # PP=1 decode or prefill: register whole rows, and advertise where # this layer's page sits so a PP>1 prefill can write into it. - storage_block_len = registration_len // num_blocks + storage_block_len = block_stride region_specs = [ (registration_base, storage_block_len, storage_block_len) ] @@ -3388,11 +3290,9 @@ def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> Block physical_block_ids.append(group) else: physical_block_ids.append( - BlockTable.map_to_kernel_blocks( - np.array(group), - ratio, - block_arange, - ).tolist() + (np.array(group)[:, None] * ratio + block_arange) + .reshape(-1) + .tolist() ) return physical_block_ids diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c98..861c894c383e 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1376,14 +1376,6 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: is_index_group_leader=self.indexer is not None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) - # SM100 FlashMLA paged kernels also express TMA coordinates in token rows. - uses_tma_rows = ( - self.attn_backend.get_name() == "FLASHMLA_SPARSE" - and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") - and current_platform.is_device_capability_family(100) - ) - if self._uses_flat_kv_cache() or uses_tma_rows: - spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes) return spec def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index cc3ade80b8ad..7fc0bd25ff75 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -669,7 +669,7 @@ def _copy_attn_groups( copy.create_metadata_builders( self.vllm_config, self.device, - group.metadata_builders[0].kernel_block_size, + group.kernel_block_size, ) copies.append(copy) own_groups.append(copies) diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index a264717f0d78..b398a614c070 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -33,9 +33,9 @@ from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8 from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool from vllm.platforms import current_platform -from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES from vllm.utils.math_utils import cdiv, next_power_of_2 -from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec +from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend +from vllm.v1.kv_cache_interface import CircularBufferSpec, MLAAttentionSpec logger = init_logger(__name__) @@ -89,15 +89,9 @@ class Glm5NextIndexerCache(DeepseekV32IndexerCache): The indexer shares one block with the co-located MLA (a single ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide - ``cache_config.block_size``. DeepGEMM's paged-MQA kernel - (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or - 64, so the storage block is virtually split into pool pages of the - largest such size that tiles it (``storage_kernel_block_size``); this - needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for - ``index_kpool = 16``). A smaller block (e.g. the default 64) silently - collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at - the opaque C++ assert; ``get_kv_cache_spec`` guards this up front - instead. + ``cache_config.block_size``. DeepGEMM's paged-MQA kernel takes 32- or + 64-state pages (only 64 on SM120), which ``Glm5NextIndexerBackend`` + declares as its kernel block sizes. """ def __init__( @@ -113,12 +107,6 @@ def __init__( head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config ) assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1" - # Keep chunked-prefill boundaries aligned to complete pools. - assert cache_config.block_size % index_kpool == 0, ( - "Glm5NextIndexerCache: cache_config.block_size " - f"({cache_config.block_size}) must be a multiple of index_kpool " - f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned." - ) self._index_kpool = index_kpool def get_kv_cache_spec(self, vllm_config: VllmConfig): @@ -128,33 +116,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig): # ``tokens_per_state`` is the KV-spec representation of kpool # compression in the current cache-layout API. assert isinstance(spec, MLAAttentionSpec) - spec = replace(spec, tokens_per_state=self._index_kpool) - - # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block - # (= block_size // index_kpool) is virtually split into pool pages of - # the largest such size that tiles it, so it must be a multiple of 32. - storage_block_size = spec.block_size // self._index_kpool - assert ( - spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0 - ), ( - "Glm5NextIndexerCache: kpool indexer requires cache block_size to " - f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so " - "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the " - f"storage block, got block_size={spec.block_size} -> " - f"storage_block_size={storage_block_size}." - ) - max_page_size = max(PAGED_MQA_PAGE_SIZES) - min_page_size = min(PAGED_MQA_PAGE_SIZES) - if storage_block_size <= max_page_size: - page_size = storage_block_size - elif storage_block_size % max_page_size == 0: - page_size = max_page_size - else: - page_size = min_page_size - return replace( - spec, - storage_block_size=page_size * self._index_kpool, - ) + return replace(spec, tokens_per_state=self._index_kpool) + + def get_attn_backend(self): + return Glm5NextIndexerBackend class Glm5NextTailCache(DeepseekV32IndexerCache): @@ -164,9 +129,9 @@ class Glm5NextTailCache(DeepseekV32IndexerCache): ``ring`` slots per request, overwritten in place by ``pos % ring`` as decode/spec-decode advances. Prefill seeds it (instead of discarding the tail raw K+gate); the connector transfers it across PD; decode reads it to - compress the boundary pool correctly. ``KpoolTailSpec`` / - ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets - the in-progress pool survive across steps and across transfer. + compress the boundary pool correctly. ``CircularBufferSpec`` provides the + no-prune, 1-block/req allocation that lets the in-progress pool survive + across steps and across transfer. Stores raw bf16 K (``head_dim``) as the "K" half of each block and the bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed @@ -202,13 +167,12 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig): f"({self.cache_config.block_size}) must be a multiple of the " f"tail ring ({ring})" ) - return KpoolTailSpec( + return CircularBufferSpec( block_size=ring, num_kv_heads=2, head_size=self.head_dim, head_size_v=0, dtype=torch.bfloat16, - sliding_window=ring, # Tiny per-request ring; every DCP rank keeps the full tail so the # rank owning a pool can compress it. dcp_sharded=False, @@ -303,7 +267,7 @@ def __init__( # Paged tail cache (in-progress pool's raw K + gate score). Written by # prefill (seeds the boundary pool) and decode (per-step stash); read by # the decode kernel to compress the boundary pool. Transferred across PD - # so the decode side sees the prefill tail. See KpoolTailSpec/Manager. + # so the decode side sees the prefill tail. See CircularBufferManager. self.tail_cache = Glm5NextTailCache( head_dim=self.head_dim, dtype=torch.bfloat16, diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index db496d300b46..eb289c65cf68 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -396,11 +396,6 @@ def _kpool_tail_seed_kernel( ahead belongs to a different tail block (or is past the batch / padding, slot < 0). ``tslot = block * RING + pos % RING``; the destination is ``tail[block, {0:K, 1:score}, pos % RING, :]``. - - The tail cache aliases the indexer cache with the indexer's (padded) block - stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` / - ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense - ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array. """ i = tl.program_id(0) t = tl.load(tslot_ptr + i).to(tl.int64) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 68c3aa7851cd..401a98f66feb 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -414,26 +414,6 @@ def get_valid_backends( return valid_backends_priorities, invalid_reasons - @classmethod - def _get_indexer_block_alignment(cls, vllm_config: VllmConfig) -> int | None: - index_kpool = getattr( - vllm_config.model_config.hf_text_config, "index_kpool", None - ) - if not index_kpool or index_kpool <= 1: - return None - from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES - - # kpool paged-MQA indexer: the storage block (block_size / - # index_kpool) is virtually split into pool pages, so block_size - # must be a multiple of index_kpool times a legal pool page. - page = min(PAGED_MQA_PAGE_SIZES) - if cls.is_device_capability_family(120): - # On sm120 the DeepGEMM paged-MQA kernel only accepts block_kv - # 64 for the fp8 indexer cache, so align to the largest pool - # page here to make the page split land on 64 not the min 32. - page = max(PAGED_MQA_PAGE_SIZES) - return index_kpool * page - @classmethod def get_attn_backend_cls( cls, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index dc3200759c82..3df44a927531 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -737,14 +737,14 @@ def _preferred_block_size_for_backends( for candidate in candidates: if all(b.supports_block_size(candidate) for b in backend_classes): return candidate - raise ValueError( - "The attention backends share no supported KV cache block size (" - + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" - for b in backend_classes + raise ValueError( + "The attention backends share no supported KV cache block size (" + + "; ".join( + f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" + for b in backend_classes + ) + + ")." ) - + ")." - ) @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @@ -935,17 +935,6 @@ def per_token_page_bytes(dtype: "torch.dtype", cache_dtype: str) -> int: if cache_config.mamba_page_size_padded is not None: cache_config.mamba_page_size_padded = shared_page - @classmethod - def _get_indexer_block_alignment(cls, vllm_config: "VllmConfig") -> int | None: - """Extra ``block_size`` multiple a sparse indexer needs, else ``None``. - - The CUDA kpool paged-MQA indexer virtually splits each storage block - into pool pages, so ``block_size`` must be a multiple of - ``index_kpool * min(PAGED_MQA_PAGE_SIZES)`` — implemented in the CUDA - platform override. Other platforms impose no extra constraint. - """ - return None - @classmethod def _align_hybrid_block_size( cls, @@ -1077,6 +1066,14 @@ def _align_hybrid_block_size( ), cache_config.block_size, ) + # Also a multiple of every backend's smallest kernel block. + backend_block_alignment = lcm( + *( + min(s.base if isinstance(s, MultipleOf) else s for s in sizes) + for b in cls._find_non_ssm_backends(vllm_config) + if (sizes := b.get_supported_kernel_block_sizes()) + ) + ) if model_config.use_mla: # TRTLLM/FlashInfer MLA decode kernels require the physical # number of kernel blocks to be aligned to 128 / kernel_block_size. @@ -1090,9 +1087,9 @@ def _align_hybrid_block_size( mamba_page_size, kernel_block_alignment_size * attn_page_size_1_token, ) - indexer_align = cls._get_indexer_block_alignment(vllm_config) - if indexer_align: - attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align) + attn_block_size = backend_block_alignment * cdiv( + attn_block_size, backend_block_alignment + ) if cache_config.block_size < attn_block_size: cache_config.block_size = attn_block_size diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 0eb979ce05a5..04c4508f0ba9 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -29,10 +29,11 @@ "qwen3_5_moe_text", } -# KV page sizes (in cache entries) supported by the paged-MQA logits kernels -# (fp8_fp4_paged_mqa_logits / get_paged_mqa_logits_metadata). Larger storage -# blocks must be virtually split into one of these page sizes. -PAGED_MQA_PAGE_SIZES = (32, 64) + +def get_paged_mqa_page_sizes() -> tuple[int, ...]: + """KV page sizes (in cache entries) the paged-MQA logits kernels take (only + 64 on SM120); larger storage blocks are split into one of these.""" + return (64,) if current_platform.is_device_capability_family(120) else (32, 64) def should_auto_disable_deep_gemm(model_type: str | None) -> bool: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 3387a31794c9..68bb0b3bdd63 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -48,12 +48,10 @@ class AttentionType(str, Enum): """Attention between dec. Q and enc. K/V for encoder-decoder.""" +@dataclass(frozen=True) class MultipleOf: base: int - def __init__(self, base: int): - self.base = base - def __repr__(self) -> str: return f"MultipleOf({self.base})" @@ -622,10 +620,6 @@ def __init__( self.layer_names = layer_names self.vllm_config = vllm_config self.device = device - self.kernel_block_size: int | None = None - - def set_kernel_block_size(self, kernel_block_size: int) -> None: - self.kernel_block_size = kernel_block_size @classmethod def get_cudagraph_support( diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index 0ba40e99f20a..f77e5e6d9561 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/vllm/v1/attention/backends/composite.py @@ -297,11 +297,6 @@ def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs): ) self.routing = routing_policy(layer_names, vllm_config) - def set_kernel_block_size(self, kernel_block_size): - super().set_kernel_block_size(kernel_block_size) - for builder in self._builders: - builder.set_kernel_block_size(kernel_block_size) - @classmethod def get_cudagraph_support(cls, vllm_config, kv_cache_spec): return min( diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 41d973f2963b..bf939dc04990 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -95,7 +95,9 @@ def kernel( block_table_ptr + batch_idx * block_table_stride + block_ids, mask=mask & is_valid, ) - slot_ids = block_numbers * block_size + pos_after_compress % block_size + slot_ids = block_numbers.to(tl.int64) * block_size + ( + pos_after_compress % block_size + ) # NOTE slot_ids = tl.where(is_valid, slot_ids, PAD_ID) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index aabd68224093..0016eafedc8a 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -24,6 +24,7 @@ from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -42,7 +43,11 @@ class FlashAttnMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64] + return [MultipleOf(64)] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) @staticmethod def get_name() -> str: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 4ebda50178cb..21eb66aae490 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """FlashInfer sparse MLA attention backend.""" +import math from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar @@ -30,6 +31,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, prepare_sparse_mla_safe_lengths, triton_convert_req_index_to_global_index, @@ -69,6 +71,13 @@ def is_mla(cls) -> bool: def is_sparse(cls) -> bool: return True + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + # The kernels read the rows as pages of their kernel block size. + (page,) = cls.get_supported_kernel_block_sizes(spec) + assert isinstance(page, MultipleOf) + return align_blocks_to_rows(spec, page.base) + class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): """FlashInfer sparse MLA backend using the TRTLLM-gen launcher.""" @@ -84,7 +93,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [32, 64] + return [MultipleOf(32)] @staticmethod def get_impl_cls() -> type[MLAAttentionImpl]: @@ -173,7 +182,7 @@ def get_name() -> str: @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64, 256] + return [MultipleOf(64)] @staticmethod def get_impl_cls() -> type[MLAAttentionImpl]: @@ -575,7 +584,7 @@ def forward_mqa( assert prefill_lse is not None return output, torch.cat((decode_lse, prefill_lse)) - _, block_stride_rows = flat_kv_row_view( + kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) @@ -602,7 +611,12 @@ def forward_mqa( return self._run_mqa_kernel( q, - kv_c_and_k_pe_cache, + # Block sizes and strides have been aligned to 32 rows. + kv_rows.view( + -1, + math.gcd(block_stride_rows, attn_metadata.block_size, 64), + kv_rows.shape[-1], + ), topk_indices_physical, seq_lens, ) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 2c918c0e4acd..2bfcee05c400 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -16,6 +16,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -184,6 +185,9 @@ def forward_mqa( output = torch.cat(outputs) if len(outputs) > 1 else outputs[0] return output, None + kv_rows, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) topk_indices_physical = cast( torch.Tensor, triton_convert_req_index_to_global_index( @@ -191,13 +195,15 @@ def forward_mqa( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], ), ) return ( self._run_mqa_kernel( q, - kv_c_and_k_pe_cache, + # The SM120 kernels are built for 64-row pages only. + kv_rows.view(-1, 64, kv_rows.shape[-1]), topk_indices_physical, ), None, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 16d7c12e6720..9e4e85761a67 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -26,7 +26,7 @@ reserved buffers by kernels inside the captured forward, and captured runs read the refreshed plan buffers on replay. -KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]`` +KV cache format: E4M3 ``[num_blocks, block_size, 512]`` (any block stride) (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the kernel but not wired yet (it needs a group-quantizing cache-write op). @@ -57,9 +57,11 @@ FlashInferMLASparseMetadataBuilder, ) from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, + flat_kv_row_view, triton_convert_req_index_to_global_index, ) -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout +from vllm.v1.kv_cache_interface import AttentionSpec _FP8_KV_DTYPES = ("fp8", "fp8_e4m3") _WORKSPACE_BYTES = 128 * 1024 * 1024 @@ -99,7 +101,12 @@ def _pack_topk_indices( end = tl.load(indptr + row + 1) mask = (cols < end - start) & (cols < WIDTH) values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0) - tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask) + first_slot = tl.load(slots + row * WIDTH) + tl.store( + indices + start + cols, + tl.where(values >= 0, values, tl.maximum(first_slot, 0)), + mask=mask, + ) @triton.jit @@ -149,6 +156,10 @@ class FlashInferMLASparseSM90Backend(AttentionBackend): def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [MultipleOf(64)] + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE_SM90" @@ -222,10 +233,6 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: return (num_blocks, block_size, head_size) - @classmethod - def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - return (KVCacheLayout.LBHNC,) - class _SM90State: """Builder-owned wrapper, capture-stable buffers, and plan parameters. @@ -602,6 +609,9 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_tokens] + kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) # return_valid_counts=True keeps the compacted-prefix layout: valid # entries at [0, valid_count), -1 past it — exactly the prefix the # planned per-row lengths address. @@ -610,6 +620,7 @@ def forward_mqa( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d441934c022a..0ff34f8ac5e0 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, ClassVar import numpy as np @@ -36,6 +36,7 @@ ) from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, request_row_bounds, triton_convert_req_index_to_global_index, @@ -53,7 +54,7 @@ flash_mla_with_kvcache, get_mla_metadata, ) -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -132,6 +133,17 @@ FP8_DS_MLA_ROPE_DIM = 64 +def _fp8_kv_pages( + kv_cache: torch.Tensor, block_size: int +) -> tuple[torch.Tensor, int | None]: + """The cache as the fp8 kernel's 64-token pages, and the block stride in + rows when a block spans several pages.""" + if block_size == 64: + return kv_cache, None + rows, block_stride_rows = flat_kv_row_view(kv_cache, block_size) + return rows.view(-1, 64, rows.shape[-1]), block_stride_rows + + class FlashMLASparseBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -144,7 +156,23 @@ class FlashMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: - return [64] + if ( + isinstance(kv_cache_spec, AttentionSpec) + and kv_cache_spec.kv_quant_mode != KVQuantMode.NONE + and current_platform.is_device_capability_family(100) + ): + return [64] + return [MultipleOf(64)] + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + # Off SM100 the quantized kernel reads 64-token pages, only needed + # when blocks span several. + if spec.kv_quant_mode != KVQuantMode.NONE and ( + not current_platform.is_device_capability_family(100) + ): + return replace(spec, block_stride_alignment=MultipleOf(64)) + return align_blocks_to_rows(spec) @staticmethod def get_name() -> str: @@ -998,13 +1026,16 @@ def _forward_fp8_kv_separate_prefill_decode( # and the prefill rows per chunk, against the gathered workspace. lse: torch.Tensor | None = None topk_length: torch.Tensor | None = None + kv_pages, block_stride_rows = _fp8_kv_pages( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) if self.pcp_dcp_kv_gather: pass elif num_prefill_tokens == 0 and not uses_host_cache: topk_indices, topk_length = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, - block_stride_rows=None, + block_stride_rows=block_stride_rows, return_valid_counts=True, ) elif num_prefill_tokens > 0: @@ -1013,6 +1044,7 @@ def _forward_fp8_kv_separate_prefill_decode( attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], HAS_PREFILL_WORKSPACE=has_prefill_workspace, prefill_workspace_request_ids=prefill_request_ids, @@ -1056,7 +1088,7 @@ def _fp8_decode( topk_indices = topk_indices.view(num_decodes, seq_len, -1) attn_out, _ = self._fp8_flash_mla_kernel( q=q, - kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, + kv_c_and_k_pe_cache=kv_pages, topk_indices=topk_indices, kernel_metadata=fp8_metadata.decode.kernel_metadata, ) @@ -1173,6 +1205,9 @@ def _forward_fp8_kv_mixed_batch( # req_id_per_token covers the whole batch; slice it to the MQA tokens # (q may exclude prefill tokens routed to dense MHA). req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] + kv_pages, block_stride_rows = _fp8_kv_pages( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) if self.dcp_world_size > 1: # The indexer emits global token ids; keep this rank's shard and # convert to local slots. compact_valid_to_front=False keeps the @@ -1187,6 +1222,7 @@ def _forward_fp8_kv_mixed_batch( dcp_rank=self.dcp_rank, cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], compact_valid_to_front=False, ) @@ -1202,7 +1238,7 @@ def _forward_fp8_kv_mixed_batch( topk_indices = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, - block_stride_rows=None, + block_stride_rows=block_stride_rows, return_valid_counts=False, ) else: @@ -1211,12 +1247,13 @@ def _forward_fp8_kv_mixed_batch( block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=topk_indices.shape[1], ) _attn_out, _lse = self._fp8_flash_mla_kernel( q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) - kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, + kv_c_and_k_pe_cache=kv_pages, topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk) kernel_metadata=fp8_metadata, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1005b8e4246a..e776057b05b7 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -7,7 +7,7 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs @@ -22,6 +22,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, + get_paged_mqa_page_sizes, has_deep_gemm, native_next_n_supported, ) @@ -234,12 +235,25 @@ def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]: return DeepseekV32IndexerMetadataBuilder +class Glm5NextIndexerBackend(DeepseekV32IndexerBackend): + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + if kv_cache_spec is not None: + # Worker get_kv_cache_spec: runs without the current config. + index_kpool = int(kv_cache_spec.tokens_per_state) + else: + # Platform block-size selection: no spec yet, under the current config. + model_config = get_current_vllm_config().model_config + index_kpool = model_config.hf_text_config.index_kpool + return [index_kpool * n for n in get_paged_mqa_page_sizes()] + + class KpoolTailBackend(DeepseekV32IndexerBackend): """Storage-only backend for the GLM-5.3-Flash kpool tail cache.""" @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - return (KVCacheLayout.LBHNC,) + return (KVCacheLayout.BLHNC,) @staticmethod def get_name() -> str: @@ -1065,8 +1079,6 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: dtype=torch.int32, device=self.device, ) - self.indexer_decode_block_table_buffer: torch.Tensor | None = None - self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens # DCP not supported yet self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 @@ -1353,16 +1365,7 @@ def build( dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens compressed_slot_mapping = slot_mapping - indexer_block_table = block_table if self.compress_ratio > 1: - kernel_block_size = self.kernel_block_size - if ( - kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - indexer_block_table = (block_table[:, ::factor] // factor).contiguous() padded_num_tokens = num_tokens local_slot_mapping = slot_mapping if self.use_pcp: @@ -1378,7 +1381,7 @@ def build( local_slot_mapping, query_start_loc, seq_lens, - indexer_block_table, + block_table, self.kv_cache_spec.num_states, self.compress_ratio, out=self.compressed_slot_mapping_buffer, @@ -1497,7 +1500,7 @@ def build( seq_lens, compressed_seq_lens, compressed_seq_lens_cpu, - indexer_block_table, + block_table, self.compress_ratio, query_slice=query_slice, skip_kv_gather=query_slice.start > 0, @@ -1632,27 +1635,6 @@ def build( ) ) - if self.compress_ratio > 1: - kernel_block_size = self.kernel_block_size - if ( - kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - compressed = block_table[:, ::factor] // factor - rows, cols = compressed.shape - if self.indexer_decode_block_table_buffer is None: - self.indexer_decode_block_table_buffer = torch.zeros( - (self._max_num_batched_tokens, cols), - dtype=torch.int32, - device=self.device, - ) - self.indexer_decode_block_table_buffer[:rows, :cols].copy_( - compressed - ) - block_table = self.indexer_decode_block_table_buffer[:rows, :cols] - # Flattening always returns a buffer view, including single-token # batches. Keep its address stable across varlen graph replays. seq_lens_is_buffer_view = not use_native or next_n > 1 @@ -1741,17 +1723,7 @@ def update_draft_decode_metadata( # One query token per request, so the per-token decode block table is # the request block table (at the indexer kernel's page size). build() # may copy it into a builder buffer, so re-derive it for this batch. - block_table = metadata.block_table[: metadata.num_decode_tokens] - kernel_block_size = self.kernel_block_size - if ( - self.compress_ratio > 1 - and kernel_block_size is not None - and self.kv_cache_spec.block_size != kernel_block_size - and self.kv_cache_spec.block_size % kernel_block_size == 0 - ): - factor = self.kv_cache_spec.block_size // kernel_block_size - block_table = block_table[:, ::factor] // factor - decode.block_table.copy_(block_table) + decode.block_table.copy_(metadata.block_table[: metadata.num_decode_tokens]) if decode.indices is not None: decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens]) if decode.per_req_decode_lens is not None: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e3d982e882bb..e30fec715deb 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -35,6 +35,10 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( AiterMLAHelper, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, + flat_kv_row_view, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode_bf16, @@ -160,6 +164,7 @@ def _convert_req_index_to_global_index_kernel( # shapes (compile-time where possible) max_num_blocks_per_req: tl.constexpr, BLOCK_SIZE: tl.constexpr, + BLOCK_STRIDE_ROWS: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) bt_stride0, @@ -201,9 +206,8 @@ def _convert_req_index_to_global_index_kernel( bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 base = tl.load(bt_ptr, mask=valid_block, other=0) - # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset out_val = tl.where( - is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + is_invalid_tok | (~valid_block), 0, base * BLOCK_STRIDE_ROWS + inblock_off ) out_ptr_ij = out_ptr + seq_start + indice_id out_ptr_ij_mask = (seq_start + indice_id) < seq_end @@ -218,13 +222,14 @@ def triton_convert_req_index_to_global_index( token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] cu_seqlens: torch.Tensor, # int32 [num_tokens + 1] paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer + BLOCK_STRIDE_ROWS: int, BLOCK_SIZE: int = 64, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 128, # tile width along columns ): """out[token_id, indice_id] = block_table[req_id[token_id], - token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_STRIDE_ROWS + token_indices[token_id, indice_id] % BLOCK_SIZE Only when token_indices[token_id, indice_id] == -1 do we output -1. @@ -264,6 +269,7 @@ def triton_convert_req_index_to_global_index( # shapes / constexprs max_num_blocks_per_req, BLOCK_SIZE, + BLOCK_STRIDE_ROWS, BLOCK_N, # strides bt_stride0, @@ -345,6 +351,10 @@ class ROCMAiterMLASparseBackend(AttentionBackend): def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [1, MultipleOf(16)] + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_name() -> str: return "ROCM_AITER_MLA_SPARSE" @@ -371,8 +381,7 @@ def is_sparse(cls) -> bool: @classmethod def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: - # Global index conversion assumes contiguous pages within each layer. - return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC) + return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC, KVCacheLayout.BLHNC) @classmethod def supports_sink(cls) -> bool: @@ -1155,6 +1164,9 @@ def forward_mqa( self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens ) + kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( + kv_c_and_k_pe_cache, attn_metadata.block_size + ) triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token, attn_metadata.block_table, @@ -1162,6 +1174,7 @@ def forward_mqa( attn_metadata.paged_kv_indptr, attn_metadata.paged_kv_indices, BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=attn_metadata.topk_tokens, ) @@ -1179,5 +1192,5 @@ def forward_mqa( return output, None mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q) return self._forward_mla( - layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata + layer, mla_padded_q, kv_c_and_k_pe_cache.unsqueeze(1), attn_metadata ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 2ec10218d99f..2447df5ceb71 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for sparse MLA backends.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import numpy as np @@ -22,6 +22,7 @@ ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.block_table import get_block_table_width @@ -40,6 +41,12 @@ def request_row_bounds(req_idx: np.ndarray) -> np.ndarray: return bounds +def align_blocks_to_rows(spec: AttentionSpec, rows: int = 1) -> AttentionSpec: + """Sparse kernels read the cache as flat rows, viewed as ``rows``-row pages, + so blocks must be whole pages apart.""" + return replace(spec, block_stride_alignment=rows * spec.state_content_size_bytes) + + def flat_kv_row_view( kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] block_size: int, diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 036dfe01730b..87800ecd1727 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -27,6 +27,7 @@ MLAAttentionImpl, ) from vllm.v1.attention.backends.mla.sparse_utils import ( + align_blocks_to_rows, flat_kv_row_view, triton_convert_req_index_to_global_index, ) @@ -50,6 +51,10 @@ class XPUMLASparseBackend(AttentionBackend): def get_name() -> str: return "XPU_MLA_SPARSE" + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + return align_blocks_to_rows(spec) + @staticmethod def get_metadata_cls() -> type["XPUMLASparseMetadata"]: return XPUMLASparseMetadata diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c32173a63554..019941a1bf03 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -661,8 +661,8 @@ def __init__( self.hash_block_size = hash_block_size self.dcp_world_size = dcp_world_size # Only groups that participate in prefix caching must satisfy the - # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool - # tail, block_size=kpool) are scratch buffers and excluded. + # divisibility constraint; groups that opt out (e.g. circular buffers + # such as the GLM-5.3-Flash indexer tail) are scratch and excluded. cacheable_block_sizes = [ manager.block_size for manager, group in zip( @@ -724,8 +724,8 @@ def verify_and_split_kv_cache_groups(self) -> None: """ self.attention_groups: list[SpecGroup] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): - # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash - # kpool tail): their blocks are per-request scratch, never + # Skip groups that opt out of prefix caching (e.g. circular + # buffers): their blocks are per-request scratch, never # shareable, so they must not participate in hit lookup (their # manager-level hooks already no-op). Their slot in the per-group # hit tuple stays empty. diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 53a01a5cabe2..919a7b4b3f52 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace from functools import partial, reduce -from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs from vllm.config import VllmConfig @@ -32,7 +32,6 @@ FullAttentionSpec, HiddenStateCacheSpec, HiSparseHotSpec, - KpoolTailSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, @@ -1126,7 +1125,7 @@ def check_enough_kv_cache_memory( # of the specs since grouping may unify them in-place. groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) check_memory = ( - available_memory - _pool_bytes_per_block(groups) + available_memory - _pool_bytes_per_block(vllm_config, groups) if groups else available_memory ) @@ -1237,14 +1236,17 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: return num_blocks -def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: """Bytes consumed by one block in the worker's shared KV cache pool, mirroring the divisor used by `get_kv_cache_config_from_groups` to convert `available_memory` into `num_blocks`. Used to compute the effective KV cache capacity once `num_gpu_blocks_override` is applied. """ return _get_kv_cache_bytes_per_block( - [group for group in kv_cache_groups if not group.host_resident] + [group for group in kv_cache_groups if not group.host_resident], + vllm_config.cache_config.get_resolved_kv_cache_layout(), ) @@ -1287,205 +1289,6 @@ def _get_kv_cache_groups_uniform_type( return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] -def _pp_balanced_mamba_group_count( - vllm_config: VllmConfig, - mamba_layer_names: list[str], - mla_layer_names: list[str], -) -> int | None: - """Return a Mamba group count whose PP projections fit the MLA slots.""" - num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names)) - pp_size = vllm_config.parallel_config.pipeline_parallel_size - if pp_size == 1: - return num_groups - - from vllm.distributed.utils import get_pp_indices - from vllm.model_executor.models.utils import extract_layer_index - - total_layers = vllm_config.model_config.get_total_num_hidden_layers() - mamba_indices = [extract_layer_index(name) for name in mamba_layer_names] - mla_indices = [extract_layer_index(name) for name in mla_layer_names] - for rank in range(pp_size): - start, end = get_pp_indices(total_layers, rank, pp_size) - num_mamba = sum(start <= index < end for index in mamba_indices) - num_mla = sum(start <= index < end for index in mla_indices) - if not num_mamba: - continue - if not num_mla: - return None - num_groups = max(num_groups, cdiv(num_mamba, num_mla)) - return num_groups - - -def _get_kv_cache_groups_glm5_next( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing.""" - mamba_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if isinstance(spec, MambaSpec) - } - tail_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if isinstance(spec, KpoolTailSpec) - } - attn_specs = { - name: spec - for name, spec in kv_cache_spec.items() - if not isinstance(spec, (MambaSpec, KpoolTailSpec)) - } - if not mamba_specs or not all( - type(spec) is MLAAttentionSpec for spec in attn_specs.values() - ): - return None - - mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs) - idx_pages = { - spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1 - } - if not idx_pages: - return None - - assert all(spec.page_size_padded is None for spec in mla_specs.values()) - assert len(idx_pages) == 1 - mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1] - mla_pages = {mla_specs[name].page_size_bytes for name in mla_names} - assert len(mla_pages) == 1 - mla_page = mla_pages.pop() - uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs) - assert uniform_spec is not None - - tail_group: KVCacheGroupSpec | None = None - if tail_specs: - idx_page = next(iter(idx_pages)) - padded_tail_specs: dict[str, KVCacheSpec] = { - name: replace(spec, page_size_padded=idx_page) - for name, spec in tail_specs.items() - } - tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs) - assert tail_uniform is not None - tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform) - - any_mamba = next(iter(mamba_specs.values())) - assert all(spec == any_mamba for spec in mamba_specs.values()) - if any_mamba.real_page_size_bytes > mla_page: - raise ValueError( - f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) " - f"does not fit the MLA page ({mla_page} bytes); increase tensor " - "parallelism or use a wider KV cache dtype" - ) - padded_specs: dict[str, KVCacheSpec] = { - name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs - } - num_groups = _pp_balanced_mamba_group_count( - vllm_config, list(mamba_specs), mla_names - ) - if num_groups is None: - raise ValueError( - "a pipeline stage has mamba layers but no MLA layer to share " - "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)" - ) - mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)] - for index, name in enumerate(mamba_specs): - mamba_grouped_names[index % num_groups].append(name) - - return ( - [KVCacheGroupSpec(list(attn_specs), uniform_spec)] - + ([tail_group] if tail_group is not None else []) - + create_kv_cache_group_specs(padded_specs, mamba_grouped_names) - ) - - -def _glm5_next_tensor_layout( - kv_cache_groups: list[KVCacheGroupSpec], -) -> ( - tuple[ - KVCacheGroupSpec, - list[KVCacheGroupSpec], - list[str], - list[str], - int, - int, - list[str], - int, - ] - | None -): - """Recognize the GLM-5.3-Flash grouping after optional PP projection.""" - uniform_groups = [ - group - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - ] - mamba_groups = [ - group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec) - ] - attn_group: KVCacheGroupSpec | None = None - tail_group: KVCacheGroupSpec | None = None - for group in uniform_groups: - inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs - if all(type(spec) is MLAAttentionSpec for spec in inner.values()): - attn_group = group - elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()): - tail_group = group - if attn_group is None or not mamba_groups: - return None - if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups): - return None - - attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) - mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs) - if not all( - type(spec) is MLAAttentionSpec and spec.page_size_padded is None - for spec in mla_inner.values() - ): - return None - mla_names = [ - name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1 - ] - idx_names = [ - name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1 - ] - mla_pages = {mla_inner[name].page_size_bytes for name in mla_names} - idx_pages = {mla_inner[name].page_size_bytes for name in idx_names} - if len(mla_pages) != 1 or len(idx_pages) != 1: - return None - mla_page = mla_pages.pop() - idx_page = idx_pages.pop() - if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups): - return None - - tail_names: list[str] = [] - tail_page = 0 - if tail_group is not None: - tail_names = list(tail_group.layer_names) - tail_inner = cast( - UniformTypeKVCacheSpecs, tail_group.kv_cache_spec - ).kv_cache_specs - tail_pages = { - cast(KpoolTailSpec, spec).unpadded_page_size_bytes - for spec in tail_inner.values() - } - if len(tail_pages) != 1 or len(tail_names) != len(idx_names): - return None - tail_page = tail_pages.pop() - if tail_page > idx_page: - return None - - return ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - tail_page, - ) - - def unify_kv_cache_spec_page_size( kv_cache_spec: dict[str, KVCacheSpec], ) -> dict[str, KVCacheSpec]: @@ -1709,14 +1512,24 @@ def _get_per_layer_spec( return spec +def _get_layer_block_stride(spec: KVCacheSpec) -> int: + """A layer-outer block stride: the page padded to its alignment.""" + return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) + + def _get_kv_cache_bytes_per_block( - kv_cache_groups: list[KVCacheGroupSpec], + kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout ) -> int: """Return the largest cache group's bytes per block.""" - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout - return len(mla_names) * mla_page + len(idx_names) * idx_page - + if not layout.is_block_outermost: + # Each layer's block stride is its own padded page. + return max( + sum( + _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) bytes_per_block = max( sum( _get_per_layer_spec(group, layer_name).page_size_bytes @@ -1731,7 +1544,7 @@ def _get_kv_cache_bytes_per_block( if isinstance(group.kv_cache_spec, HiSparseHotSpec) ] alignments.extend( - _get_per_layer_spec(group, layer_name).block_stride_alignment or 1 + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() for group in kv_cache_groups for layer_name in group.layer_names ) @@ -1806,72 +1619,9 @@ def get_kv_cache_config_from_groups( vllm_config, kv_cache_groups, available_memory, host_budget ) - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - _, - ) = glm5_layout - bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page - num_blocks = may_override_num_blocks( - vllm_config, available_memory // bytes_per_block - ) - size = bytes_per_block * num_blocks - attn_specs = cast( - UniformTypeKVCacheSpecs, attn_group.kv_cache_spec - ).kv_cache_specs - - kv_cache_tensors: list[KVCacheTensor] = [] - - def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=[layer_name], - layer_stride=spec.page_size_bytes * num_blocks, - block_stride=spec.page_size_bytes, - offset=offset, - ) - ) - - for index, mla_name in enumerate(mla_names): - offset = index * mla_page * num_blocks - add_tensor(mla_name, attn_specs[mla_name], offset) - for group in mamba_groups: - if index < len(group.layer_names): - add_tensor(group.layer_names[index], group.kv_cache_spec, offset) - - idx_base = len(mla_names) * mla_page * num_blocks - for index, idx_name in enumerate(idx_names): - offset = idx_base + index * idx_page * num_blocks - add_tensor(idx_name, attn_specs[idx_name], offset) - if tail_names: - tail_name = tail_names[index] - tail_group = next( - group for group in kv_cache_groups if tail_name in group.layer_names - ) - tail_specs = cast( - UniformTypeKVCacheSpecs, tail_group.kv_cache_spec - ).kv_cache_specs - add_tensor(tail_name, tail_specs[tail_name], offset) - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None num_blocks = available_memory // bytes_per_block @@ -1901,17 +1651,17 @@ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: byte_offset = 0 for spec, layer_names in layers_by_spec.items(): + block_stride = interleaved_block_stride + padded_page = _get_layer_block_stride(spec) + if block_stride is None and padded_page != spec.page_size_bytes: + assert layout.is_block_compact, f"{layout.name} cannot pad pages" + block_stride = padded_page layer_stride, block_stride, _, _, _ = compute_layout_strides( spec, num_blocks, len(layer_names), layout, - fixed_strides=(None, interleaved_block_stride, None, None, None), - ) - offset = ( - byte_offset - * max(layer_stride, spec.page_size_bytes) - // spec.page_size_bytes + fixed_strides=(None, block_stride, None, None, None), ) kv_cache_tensors.append( KVCacheTensor( @@ -1919,10 +1669,13 @@ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: layers=layer_names, layer_stride=layer_stride, block_stride=block_stride, - offset=offset, + offset=byte_offset, ) ) - byte_offset += len(layer_names) * spec.page_size_bytes + # Block-outer layers share each block; layer-outer ones own regions. + byte_offset += len(layer_names) * ( + spec.page_size_bytes if layout.is_block_outermost else layer_stride + ) return KVCacheConfig( num_blocks=num_blocks, @@ -2438,9 +2191,6 @@ def get_kv_cache_groups( # full attention, or all layers are sliding window attention with the # same window size). Put all layers into one group. return _get_kv_cache_groups_uniform_type(uniform_spec) - elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec): - return glm5_groups - # Hidden-state layers use their own block table and must not be absorbed # into a compatible attention bucket. hidden_specs = { @@ -2594,31 +2344,7 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 - if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - ( - attn_group, - mamba_groups, - mla_names, - idx_names, - mla_page, - idx_page, - tail_names, - _, - ) = glm5_layout - uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) - total_blocks = uniform_spec.max_memory_usage_pages(vllm_config) - total_blocks += sum( - cdiv( - group.kv_cache_spec.max_memory_usage_bytes(vllm_config), - group.kv_cache_spec.page_size_bytes, - ) - for group in mamba_groups - ) - if tail_names: - total_blocks += 1 - return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page) - - bytes_per_block = _pool_bytes_per_block(kv_cache_groups) + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) total_blocks = 0 for group in kv_cache_groups: spec = group.kv_cache_spec @@ -2874,7 +2600,7 @@ def get_kv_cache_configs( if not groups: adjusted_memory.append(avail_mem) continue - bytes_per_block = _pool_bytes_per_block(groups) + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) logger.info( "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", avail_mem // bytes_per_block, @@ -2887,7 +2613,7 @@ def get_kv_cache_configs( # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. check_memory = [ - avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem for groups, avail_mem in zip(projected_groups_per_worker, available_memory) ] @@ -2931,7 +2657,9 @@ def get_kv_cache_configs( # strides and offsets stay consistent with the shrunken allocation. groups = kv_cache_config.kv_cache_groups kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups) + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(vllm_config, groups), ) for kv_cache_config in kv_cache_configs: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 3e7ae1b7063a..e9f5f2f3c231 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, ClassVar from vllm.distributed.kv_events import MEDIUM_CPU -from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( @@ -27,7 +26,6 @@ HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, - KpoolTailSpec, KVCacheGroupRole, KVCacheSpec, MambaSpec, @@ -45,8 +43,6 @@ if TYPE_CHECKING: from vllm.v1.hisparse.coordinator import HiSparseCoordinator -logger = init_logger(__name__) - class SingleTypeKVCacheManager(ABC): """An abstract base class for a manager that handle the kv cache management @@ -1299,10 +1295,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: return 0 -class KpoolTailManager(CircularBufferManager): - """One-block circular scratch manager for ``KpoolTailSpec``.""" - - class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) @@ -2687,11 +2679,6 @@ def register_all_kvcache_specs(vllm_config): SlidingWindowManager, uniform_type_base_spec=SlidingWindowMLASpec, ) - KVCacheSpecRegistry.register( - KpoolTailSpec, - KpoolTailManager, - uniform_type_base_spec=KpoolTailSpec, - ) KVCacheSpecRegistry.register( MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bfbaa38d9478..d0dadae27e9b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -335,8 +335,8 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups if kv_cache_groups: - # Exclude groups that opt out of prefix caching (e.g. GLM-5.3-Flash - # kpool tail, a 1-block/req scratch buffer with block_size=kpool): + # Exclude groups that opt out of prefix caching (e.g. a circular + # buffer, a 1-block/req scratch buffer with a small ring block): # their small block_size would otherwise drag the global block_size # below the real allocator block size and desync it from mamba. participating = [ diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index b8c0d5dff269..60c60a3ff504 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -55,8 +55,9 @@ def resolve_hisparse_specs( if len(block_sizes) != 1: raise ValueError("HiSparse requires one scheduler block size.") backends = [attn_layers[name].get_attn_backend() for name in mla_specs] + specs = list(mla_specs.values()) try: - block_size = select_common_block_size(block_sizes.pop(), backends) + block_size = select_common_block_size(block_sizes.pop(), backends, specs) except ValueError as error: raise ValueError( "HiSparse requires a GPU block size supported by every sparse " @@ -104,16 +105,13 @@ def allocate_hisparse_kv_caches( if isinstance(host_spec, UniformTypeKVCacheSpecs) else host_spec ) - kernel_block_size = kernel_block_sizes[host_group_id] - if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: - kernel_block_size = spec.storage_block_size views = create_kv_cache_views( backing, spec, kv_cache_config.num_blocks_of(tensor), layout, tensor, - kernel_block_size=kernel_block_size, + kernel_block_size=kernel_block_sizes[host_group_id], ) kv_caches.update(zip(tensor.layers, views)) return kv_caches diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 98636866f8ac..96d93650ad29 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -333,7 +333,7 @@ def get_hisparse_kv_cache_config( device_groups = hisparse_layout.device_groups layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, device_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(device_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(device_groups, layout) num_blocks = may_override_num_blocks( vllm_config, available_memory // bytes_per_block ) @@ -345,7 +345,7 @@ def get_hisparse_kv_cache_config( host_groups = [hisparse_layout.source_group] host_layout = KVCacheLayout.LBNHC validate_kv_cache_layout(host_layout, host_groups) - host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups) + host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups, host_layout) host_size = host_bytes_per_block * hisparse_layout.host_num_blocks kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors( host_groups, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index ca3f485917eb..79e88ca91d8f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy +import math from collections import Counter from collections.abc import Collection, Sequence from dataclasses import dataclass, field, fields, replace @@ -24,6 +25,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backend import MultipleOf logger = init_logger(__name__) @@ -168,13 +170,28 @@ class KVCacheSpec: dcp_sharded: bool = field(default=False, kw_only=True) """Whether DCP shards this cache's token positions across ranks.""" - block_stride_alignment: int | None = field(default=None, kw_only=True) - """Required byte alignment between physical blocks, including packed layers.""" + block_stride_alignment: int | MultipleOf | None = field(default=None, kw_only=True) + """Required byte alignment between physical blocks, including packed layers. + ``MultipleOf(n)``: whole kernel blocks, when blocks are split into kernel + blocks of up to ``n`` tokens.""" def __post_init__(self): - if self.block_stride_alignment is not None and self.block_stride_alignment <= 0: + alignment = self.block_stride_alignment + if isinstance(alignment, int) and alignment <= 0: raise ValueError("block_stride_alignment must be positive") + def get_block_stride_alignment(self) -> int: + """Bytes the stride between physical blocks must be a multiple of.""" + from vllm.v1.attention.backend import MultipleOf + + alignment = self.block_stride_alignment + if not isinstance(alignment, MultipleOf): + return alignment or 1 + kernel_block_size = math.gcd(alignment.base, self.block_size) + if kernel_block_size == self.block_size: + return 1 + return self.copy_with_new_block_size(kernel_block_size).page_size_bytes + @property def prefix_cacheable(self) -> bool: """Whether this spec's group participates in prefix caching.""" @@ -674,8 +691,6 @@ class MLAAttentionSpec(FullAttentionSpec): model_version: str | None = None cache_role: SparseCacheRole = SparseCacheRole.SPARSE is_index_group_leader: bool = False - storage_block_size: int | None = None - """Token width used to view storage when it differs from the kernel block.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -696,7 +711,6 @@ def merge(cls, specs: list[Self]) -> Self: model_version_set = set(spec.model_version for spec in specs) cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} - storage_block_size_set = set(spec.storage_block_size for spec in specs) block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 @@ -704,12 +718,11 @@ def merge(cls, specs: list[Self]) -> Self: and len(model_version_set) == 1 and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 - and len(storage_block_size_set) == 1 and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, storage block size and block stride alignment." + "index-sharing role, and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -727,7 +740,6 @@ def merge(cls, specs: list[Self]) -> Self: model_version=model_version_set.pop(), cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), - storage_block_size=storage_block_size_set.pop(), block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs @@ -1019,32 +1031,6 @@ def is_uniform_with_collection( ) -@dataclass(frozen=True, kw_only=True) -class KpoolTailSpec(SlidingWindowSpec): - """One-block circular scratch cache for a kpool indexer's raw tail.""" - - def max_admission_blocks_per_request( - self, max_in_flight_tokens: int, max_model_len: int - ) -> int: - return 1 - - def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: - return 1 - - def is_uniform_with_collection( - self, kv_cache_specs: dict[str, KVCacheSpec] - ) -> bool: - return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values()) - - @property - def prefix_cacheable(self) -> bool: - return False - - @property - def uses_slot_mapping(self) -> bool: - return False - - @dataclass(frozen=True) class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index bea677dde85a..af504b0a1108 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -95,7 +95,7 @@ def build_per_group_and_layer_attn_metadata( cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size] else: cm = common_attn_metadata - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=cm, draft_index=draft_index ) per_group_attn_metadata.append(attn_metadata) @@ -291,11 +291,7 @@ def initialize_attn_backend( self.draft_attn_groups = list(attention_groups.values()) if self.draft_attn_groups: self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size else: self.kv_cache_gid = 0 self.block_size = kv_cache_config.kv_cache_groups[ diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index e4718bb3d7d1..79be2e8a4d86 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -991,7 +991,7 @@ def build_per_group_and_layer_attn_metadata( per_group_attn_metadata: list[object] = [] per_layer_attn_metadata: dict[str, object] = {} for attn_group in self.draft_attn_groups: - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=common_attn_metadata, draft_index=draft_index ) per_group_attn_metadata.append(attn_metadata) @@ -1765,20 +1765,7 @@ def initialize_attn_backend( attention_groups[backend_key].layer_names.append(layer_name) self.draft_attn_groups = list(attention_groups.values()) - if kernel_block_sizes is not None and 0 <= self.kv_cache_gid < len( - kernel_block_sizes - ): - # Slot mappings are computed against the block table, which is - # stored at kernel-block granularity. Use the kernel block size - # rather than the KV cache manager's block size; the two differ - # when manager blocks are split for the attention kernel. - self.block_size = kernel_block_sizes[self.kv_cache_gid] - else: - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size logger.debug("Using block size %d for drafting layers", self.block_size) def _determine_batch_execution_and_padding( diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 7a73473adc2f..0853067f3906 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -144,7 +144,7 @@ def build_per_group_and_layer_attn_metadata( cm.slot_mapping = sm else: cm = common_attn_metadata - attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + attn_metadata = attn_group.build_metadata_for_drafting( common_attn_metadata=cm, draft_index=draft_index, ) @@ -244,11 +244,7 @@ def initialize_attn_backend( self.draft_attn_groups = list(attention_groups.values()) if self.draft_attn_groups: self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id - self.block_size = ( - self.draft_attn_groups[0] - .get_metadata_builder() - .kv_cache_spec.block_size - ) + self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size else: self.kv_cache_gid = 0 self.block_size = kv_cache_config.kv_cache_groups[ diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index c63908dad0d0..2beec249bff0 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -29,24 +29,16 @@ def get_block_table_width( max_num_blocks: int, block_size: int, - kernel_block_size: int | None = None, *, token_alignment: int | None = 128, ) -> int: - """Return the width after optional alignment and virtual block splitting.""" - if kernel_block_size is None: - kernel_block_size = block_size - if block_size % kernel_block_size != 0: - raise ValueError( - f"kernel_block_size {kernel_block_size} must divide " - f"block_size {block_size} evenly" - ) + """Return the width after optional alignment.""" if token_alignment is not None: if token_alignment <= 0: raise ValueError("token_alignment must be positive") block_alignment = token_alignment // math.gcd(token_alignment, block_size) max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment - return max_num_blocks * block_size // kernel_block_size + return max_num_blocks class SlotMappingMode(Enum): @@ -63,7 +55,6 @@ def __init__( max_num_batched_tokens: int, pin_memory: bool, device: torch.device, - kernel_block_size: int, cp_kv_cache_interleave_size: int, slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT, ): @@ -74,9 +65,6 @@ def __init__( max_num_batched_tokens: Maximum number of tokens in a batch. pin_memory: Whether to pin memory for faster GPU transfers. device: Target device for the block table. - kernel_block_size: The block_size of underlying attention kernel. - Will be the same as `block_size` if `block_size` is supported - by the attention kernel. slot_mapping_mode: How this cache group maps scheduled tokens to cache slots. Mamba-like state caches do not use token slot mappings and should use SlotMappingMode.NONE. @@ -86,30 +74,8 @@ def __init__( self.max_num_batched_tokens = max_num_batched_tokens self.pin_memory = pin_memory self.device = device - self.kv_cache_block_size = block_size - - if kernel_block_size == block_size: - # Standard case: allocation and computation use same block size - # No block splitting needed, direct mapping - self.block_size = block_size - self.blocks_per_kv_block = 1 - self.use_hybrid_blocks = False - else: - # Hybrid case: allocation block size differs from kernel block size - # Memory blocks are subdivided to match kernel requirements - # Example: 32-token memory blocks with 16-token kernel blocks - # → Each memory block corresponds to 2 kernel blocks - if block_size % kernel_block_size != 0: - raise ValueError( - f"kernel_block_size {kernel_block_size} must divide " - f"kv_manager_block_size size {block_size} evenly" - ) - - self.block_size = kernel_block_size - self.blocks_per_kv_block = block_size // kernel_block_size - self.use_hybrid_blocks = True - - self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block + self.block_size = block_size + self.max_num_blocks_per_req = max_num_blocks_per_req self.block_table = self._make_buffer( self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32 @@ -120,13 +86,6 @@ def __init__( self.max_num_batched_tokens, dtype=torch.int64 ) - if self.use_hybrid_blocks: - self._kernel_block_arange = np.arange(0, self.blocks_per_kv_block).reshape( - 1, -1 - ) - else: - self._kernel_block_arange = None - try: self.pcp_world_size = get_pcp_group().world_size self.pcp_rank = get_pcp_group().rank_in_group @@ -145,13 +104,11 @@ def __init__( self.slot_mapping_mode = slot_mapping_mode if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT: _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup( - kv_cache_block_size=self.kv_cache_block_size, - blocks_per_kv_block=self.blocks_per_kv_block, + kv_cache_block_size=self.block_size, total_cp_world_size=self.dcp_world_size, total_cp_rank=self.dcp_rank, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, block_table_stride=self.block_table.gpu.stride(0), - block_size=self.block_size, ) def append_row( @@ -162,11 +119,6 @@ def append_row( if not block_ids: return - if self.use_hybrid_blocks: - block_ids = self.map_to_kernel_blocks( - np.array(block_ids), self.blocks_per_kv_block, self._kernel_block_arange - ) - num_blocks = len(block_ids) start = self.num_blocks_per_row[row_idx] self.num_blocks_per_row[row_idx] += num_blocks @@ -219,10 +171,8 @@ def compute_slot_mapping( positions, self.block_table.gpu, self.block_table.gpu.stride(0), - self.block_size, self.slot_mapping.gpu, - self.kv_cache_block_size, - self.blocks_per_kv_block, + self.block_size, self.dcp_world_size, self.dcp_rank, self.cp_kv_cache_interleave_size, @@ -235,37 +185,6 @@ def clear(self) -> None: self.block_table.gpu.fill_(0) self.block_table.cpu.fill_(0) - @staticmethod - def map_to_kernel_blocks( - kv_manager_block_ids: np.ndarray, - blocks_per_kv_block: int, - kernel_block_arange: np.ndarray, - ) -> np.ndarray: - """Convert kv_manager_block_id IDs to kernel block IDs. - - Example: - # kv_manager_block_ids: 32 tokens, - # Kernel block size: 16 tokens - # blocks_per_kv_block = 2 - >>> kv_manager_block_ids = np.array([0, 1, 2]) - >>> Result: [0, 1, 2, 3, 4, 5] - - # Each kv_manager_block_id maps to 2 kernel block id: - # kv_manager_block_id 0 → kernel block id [0, 1] - # kv_manager_block_id 1 → kernel block id [2, 3] - # kv_manager_block_id 2 → kernel block id [4, 5] - - """ - if blocks_per_kv_block == 1: - return kv_manager_block_ids - - kernel_block_ids = ( - kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block - + kernel_block_arange - ) - - return kernel_block_ids.reshape(-1) - def get_device_tensor(self, num_reqs: int) -> torch.Tensor: """Returns the device tensor of the block table.""" return self.block_table.gpu[:num_reqs] @@ -296,16 +215,10 @@ def __init__( pin_memory: bool, device: torch.device, block_sizes: list[int], - kernel_block_sizes: list[int], max_num_blocks: list[int], cp_kv_cache_interleave_size: int = 1, slot_mapping_modes: list[SlotMappingMode] | None = None, ) -> None: - if len(kernel_block_sizes) != len(block_sizes): - raise ValueError( - f"kernel_block_sizes length ({len(kernel_block_sizes)}) " - f"must match block_sizes length ({len(block_sizes)})" - ) if slot_mapping_modes is None: slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes) if len(slot_mapping_modes) != len(block_sizes): @@ -339,17 +252,11 @@ def __init__( max_num_batched_tokens, pin_memory, device, - kernel_block_size, cp_kv_cache_interleave_size, slot_mapping_mode=slot_mapping_mode, ) - for ( - block_size, - kernel_block_size, - max_num_blocks_per_req, - slot_mapping_mode, - ) in zip( - block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes + for block_size, max_num_blocks_per_req, slot_mapping_mode in zip( + block_sizes, max_num_blocks, slot_mapping_modes ) ] @@ -403,12 +310,10 @@ class ComputeSlotMappingKernel( @dataclass(frozen=True) class CompileKey: kv_cache_block_size: int - blocks_per_kv_block: int total_cp_world_size: int total_cp_rank: int cp_kv_cache_interleave_size: int block_table_stride: int - block_size: int @staticmethod @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) @@ -419,10 +324,8 @@ def kernel( positions_ptr, # [num_tokens], int64 block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) block_table_stride, # max_num_blocks_per_req - block_size, slot_mapping_ptr, # [max_num_tokens], int64 KV_CACHE_BLOCK_SIZE: tl.constexpr, - BLOCKS_PER_KV_BLOCK: tl.constexpr, TOTAL_CP_WORLD_SIZE: tl.constexpr, TOTAL_CP_RANK: tl.constexpr, CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, @@ -463,17 +366,12 @@ def kernel( virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE ) - block_indices = ( - virtual_block_indices * BLOCKS_PER_KV_BLOCK - + local_block_offsets // block_size - ) block_numbers = tl.load( - block_table_ptr + row_offset + block_indices, + block_table_ptr + row_offset + virtual_block_indices, mask=mask & is_local, other=0, ).to(tl.int64) - slot_offsets = local_block_offsets % block_size - slot_ids = block_numbers * block_size + slot_offsets + slot_ids = block_numbers * KV_CACHE_BLOCK_SIZE + local_block_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) @@ -481,13 +379,11 @@ def dispatch( # type: ignore[override] self, *, block_table_stride: int, - block_size: int, **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, block_table_stride=triton_scalar_specialization_rep(block_table_stride), - block_size=triton_scalar_specialization_rep(block_size), ) def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]: @@ -507,10 +403,8 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: shape=(1, compile_key.block_table_stride), ), block_table_stride=compile_key.block_table_stride, - block_size=compile_key.block_size, slot_mapping=int64_ptr, kv_cache_block_size=compile_key.kv_cache_block_size, - blocks_per_kv_block=compile_key.blocks_per_kv_block, total_cp_world_size=compile_key.total_cp_world_size, total_cp_rank=compile_key.total_cp_rank, cp_kv_cache_interleave_size=compile_key.cp_kv_cache_interleave_size, @@ -526,17 +420,14 @@ def __call__( positions: torch.Tensor, block_table: torch.Tensor, block_table_stride: int, - block_size: int, slot_mapping: torch.Tensor, kv_cache_block_size: int, - blocks_per_kv_block: int, total_cp_world_size: int, total_cp_rank: int, cp_kv_cache_interleave_size: int, ) -> LaunchSpec: return (num_reqs + 1,), dict( KV_CACHE_BLOCK_SIZE=kv_cache_block_size, - BLOCKS_PER_KV_BLOCK=blocks_per_kv_block, TOTAL_CP_WORLD_SIZE=total_cp_world_size, TOTAL_CP_RANK=total_cp_rank, CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size, diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index bcf56dc0bf78..49dcccf44318 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -95,11 +95,9 @@ def prepare_dcp_dummy_context_metadata( max_valid_block_id = kv_cache_config.num_blocks - 1 assert max_valid_block_id > 0 for blk_table in input_batch.block_table.block_tables: - max_row_blocks = ( - blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block - ) block_ids = [ - (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks) + (block_idx % max_valid_block_id) + 1 + for block_idx in range(blk_table.max_num_blocks_per_req) ] for req_idx in range(num_reqs): blk_table.add_row(block_ids, req_idx) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8aa9ba3cb8b2..d1feaef031d9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -40,6 +40,8 @@ add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, bind_kv_cache_to_layers, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, prepare_kernel_block_sizes, ) @@ -141,7 +143,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): if isinstance(spec, AttentionSpec): - spec = attn_module.get_attn_backend().customize_spec(spec) + spec = customize_attention_spec(attn_module.get_attn_backend(), spec) kv_cache_spec[layer_name] = spec if vllm_config.attention_config.hisparse_config is not None: kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers) @@ -382,6 +384,7 @@ def init_kv_cache( kv_cache_allocation_context: AbstractContextManager | None = None, *, block_tables: "BlockTables | None" = None, + attn_groups: Iterable[AttentionGroup] = (), ) -> dict[str, Any]: allocation_context = kv_cache_allocation_context or nullcontext() with allocation_context: @@ -409,9 +412,10 @@ def init_kv_cache( num_attn_module = ( 2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1 ) - bindable_caches = { - name: cache for name, cache in kv_caches.items() if name in forward_context - } + bindable_caches = map_kv_caches_to_kernel_blocks( + {name: cache for name, cache in kv_caches.items() if name in forward_context}, + attn_groups, + ) bind_kv_cache_to_layers( bindable_caches, forward_context, @@ -536,8 +540,8 @@ def build_attn_metadata( cached_metadata[reuse_key], block_table, slot_mapping ) elif for_cudagraph_capture: - metadata = attn_metadata_builder.build_for_cudagraph_capture( - common_attn_metadata + metadata = attn_group.build_metadata_for_cudagraph_capture( + common_attn_metadata, ubatch_idx ) else: attn_metadata_extra_kwargs = ( @@ -548,10 +552,8 @@ def build_attn_metadata( if model_specific_attn_metadata is not None else {} ) - metadata = attn_metadata_builder.build( - common_prefix_len=0, - common_attn_metadata=common_attn_metadata, - **attn_metadata_extra_kwargs, + metadata = attn_group.build_metadata( + common_attn_metadata, ubatch_idx, **attn_metadata_extra_kwargs ) if reuse_key is not None and reuse_key not in cached_metadata: cached_metadata[reuse_key] = metadata diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index c7665c46b876..1ab78864c31e 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -25,7 +25,6 @@ def __init__( max_num_batched_tokens: int, max_num_blocks_per_group: list[int], device: torch.device, - kernel_block_sizes: list[int], cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, @@ -33,7 +32,6 @@ def __init__( dcp_sharded: list[bool] | None = None, ): self.block_sizes = block_sizes - self.kernel_block_sizes = kernel_block_sizes self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens self.device = device @@ -53,14 +51,9 @@ def __init__( assert len(dcp_sharded) == self.num_kv_cache_groups self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device) - self.blocks_per_kv_block = [ - bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) - ] - # num_kv_cache_groups x [max_num_reqs, max_num_blocks] self.block_tables: list[StagedWriteTensor] = [] - for i in range(self.num_kv_cache_groups): - max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i] + for max_num_blocks in max_num_blocks_per_group: block_table = StagedWriteTensor( (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device ) @@ -109,9 +102,6 @@ def init_block_table_layout_tensors(self) -> None: self.block_sizes_tensor = torch.tensor( self.block_sizes, dtype=torch.int32, device=self.device ) - self.kernel_block_sizes_tensor = torch.tensor( - self.kernel_block_sizes, dtype=torch.int32, device=self.device - ) self.slot_mapping_enabled = torch.tensor( self._slot_mapping_enabled, dtype=torch.bool, device=self.device ) @@ -126,9 +116,6 @@ def append_block_ids( for i in range(self.num_kv_cache_groups): start = self.num_blocks.np[i, req_index] if not overwrite else 0 block_ids = new_block_ids[i] - bpk = self.blocks_per_kv_block[i] - if bpk > 1: - block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] if self.redirect_writes_to_null_block: block_ids = [0] * len(block_ids) end = start + len(block_ids) @@ -218,7 +205,6 @@ def compute_slot_mappings( self.block_table_ptrs, self.block_table_strides, self.block_sizes_tensor, - self.kernel_block_sizes_tensor, self.slot_mapping_enabled, self.dcp_sharded, slot_mappings, @@ -292,7 +278,6 @@ def _compute_slot_mappings_kernel( block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] block_sizes, # [num_kv_cache_groups] - kernel_block_sizes, # [num_kv_cache_groups] slot_mapping_enabled, # [num_kv_cache_groups] dcp_sharded, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] @@ -322,7 +307,6 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) kv_block_size = tl.load(block_sizes + group_id) - kernel_block_size = tl.load(kernel_block_sizes + group_id) mapping_enabled = tl.load(slot_mapping_enabled + group_id) if CP_SIZE != 1: sharded = tl.load(dcp_sharded + group_id) @@ -354,16 +338,14 @@ def _compute_slot_mappings_kernel( local_positions = tl.where(sharded, local_positions, positions) is_local = ~sharded | is_local - block_indices = tl.where( - mapping_enabled, local_positions // kernel_block_size, 0 - ) - block_offsets = local_positions % kernel_block_size + block_indices = tl.where(mapping_enabled, local_positions // kv_block_size, 0) + block_offsets = local_positions % kv_block_size block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices, mask=is_local & is_real_req, other=0, ) - slot_ids = block_numbers * kernel_block_size + block_offsets + slot_ids = block_numbers * kv_block_size + block_offsets if CP_SIZE != 1: slot_ids = tl.where(is_local, slot_ids, PAD_ID) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index e2f68ad59865..ea1c925dadc0 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -252,17 +252,18 @@ def set_dummy_context( input_batch.positions.copy_(torch.from_numpy(local_pos + context_len)) seq_len = context_len + query_len - for block_table, block_size, bpk in zip( - input_block_tables, - block_tables.kernel_block_sizes, - block_tables.blocks_per_kv_block, - ): + for block_table, block_size in zip(input_block_tables, block_tables.block_sizes): num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1]) # Spans are disjoint until the pool runs out, then they wrap and share # blocks: profiling only needs the reads to be realistic, not distinct. - block_ids = torch.arange( - num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device - ) % (num_kv_blocks * bpk) + block_ids = ( + torch.arange( + num_reqs * num_blocks, + dtype=block_table.dtype, + device=block_table.device, + ) + % num_kv_blocks + ) block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 61d6197ceb9d..134b4bfb0e69 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -688,7 +688,6 @@ def initialize_kv_cache( max_num_batched_tokens=self.max_num_tokens, max_num_blocks_per_group=max_num_blocks_per_group, device=self.device, - kernel_block_sizes=self.kernel_block_sizes, slot_mapping_enabled=slot_mapping_enabled, dcp_sharded=dcp_sharded, cp_size=self.dcp_size, @@ -773,6 +772,13 @@ def initialize_kv_cache( self.vllm_config, kv_cache_allocation_context=kv_cache_allocation_context, block_tables=self.block_tables, + # Draft groups build their own metadata, so they map the views too. + attn_groups=( + g + for groups in self.attn_groups + + getattr(self.speculator, "attn_groups", []) + for g in groups + ), ) self.kv_caches = [ cache for cache in kv_caches_dict.values() if cache.device == self.device diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index d59d87ae7e64..ec46039533d7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -432,7 +432,7 @@ def propose( temperature, seeds, self.block_tables.input_block_tables[gid], - self.block_tables.kernel_block_sizes[gid], + self.block_tables.block_sizes[gid], self.block_tables.cp_rank, self.dcp_size, self.block_tables.cp_interleave, diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 9a8ed7b8b4ca..4894ce77e962 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -577,6 +577,11 @@ def _fused_multi_step_decode( idx_mapping = self.idx_mapping[:num_reqs] if batch_desc.cg_mode == CUDAGraphMode.FULL: + for groups in self.attn_groups: + for group in groups: + group.map_to_kernel_block_table( + self.block_tables.input_block_tables[group.kv_cache_group_id] + ) assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.run_fullgraph(batch_desc) return diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 4707e4f2ca95..73a8c0eb9a59 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -96,7 +96,6 @@ def __init__( device: torch.device, vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group - kernel_block_sizes: list[int], max_num_blocks_per_req: list[int], logitsprocs: LogitsProcessors | None = None, logitsprocs_need_output_token_ids: bool = False, @@ -187,7 +186,6 @@ def __init__( pin_memory=PIN_MEMORY, device=device, block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, max_num_blocks=max_num_blocks_per_req, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, slot_mapping_modes=slot_mapping_modes, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c02bdc74e22b..22ad662b9018 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -233,6 +233,8 @@ allocate_kv_cache, bind_kv_cache, copy_kv_cache_blocks_inplace, + customize_attention_spec, + map_kv_caches_to_kernel_blocks, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -718,7 +720,6 @@ def __init__( max(self.max_model_len, self.max_encoder_len), placeholder_block_size ) self._init_block_sizes = [placeholder_block_size] - self._init_kernel_block_sizes = [placeholder_block_size] self._init_max_num_blocks = [placeholder_max_num_blocks] self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] self.cp_kv_cache_interleave_size = ( @@ -734,7 +735,6 @@ def __init__( device=self.device, vocab_size=self.model_config.get_vocab_size(), block_sizes=[placeholder_block_size], - kernel_block_sizes=[placeholder_block_size], max_num_blocks_per_req=[placeholder_max_num_blocks], num_spec_tokens=self.num_spec_tokens, logitsprocs=build_logitsprocs( @@ -2508,8 +2508,8 @@ def _build_attn_group_metadata( ) if for_cudagraph_capture: - attn_metadata_i = builder.build_for_cudagraph_capture( - common_attn_metadata + attn_metadata_i = attn_group.build_metadata_for_cudagraph_capture( + common_attn_metadata, ubid or 0 ) elif ( cache_key in cached_attn_metadata @@ -2517,13 +2517,16 @@ def _build_attn_group_metadata( ): attn_metadata_i = builder.update_block_table( cached_attn_metadata[cache_key], - common_attn_metadata.block_table_tensor, + attn_group.map_to_kernel_block_table( + common_attn_metadata.block_table_tensor, ubid or 0 + ), common_attn_metadata.slot_mapping, ) else: - attn_metadata_i = builder.build( + attn_metadata_i = attn_group.build_metadata( + common_attn_metadata, + ubid or 0, common_prefix_len=cascade_attn_prefix_len, - common_attn_metadata=common_attn_metadata, **extra_attn_metadata_args, ) if builder.supports_update_block_table: @@ -7143,9 +7146,7 @@ def calculate_reorder_batch_threshold(self) -> None: return self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment] - def may_reinitialize_input_batch( - self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] - ) -> None: + def may_reinitialize_input_batch(self, kv_cache_config: KVCacheConfig) -> None: """Re-initialize the input batch if the block sizes are different from what it was originally created with. This happens when the final block size (determined after model loading) differs from the @@ -7154,7 +7155,6 @@ def may_reinitialize_input_batch( Args: kv_cache_config: The KV cache configuration. - kernel_block_sizes: The kernel block sizes for each KV cache group. """ block_sizes = [] @@ -7179,14 +7179,12 @@ def may_reinitialize_input_batch( if ( block_sizes != self._init_block_sizes - or kernel_block_sizes != self._init_kernel_block_sizes or max_num_blocks != self._init_max_num_blocks or slot_mapping_modes != self._init_slot_mapping_modes or self.cp_kv_cache_interleave_size != self.parallel_config.cp_kv_cache_interleave_size ): self._init_block_sizes = block_sizes - self._init_kernel_block_sizes = kernel_block_sizes self._init_max_num_blocks = max_num_blocks self._init_slot_mapping_modes = slot_mapping_modes self.cp_kv_cache_interleave_size = ( @@ -7201,7 +7199,6 @@ def may_reinitialize_input_batch( device=self.device, vocab_size=self.model_config.get_vocab_size(), block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, max_num_blocks_per_req=max_num_blocks, num_spec_tokens=self.num_spec_tokens, logitsprocs=self.input_batch.logitsprocs, @@ -7217,10 +7214,6 @@ def may_reinitialize_input_batch( f"InputBatch block_sizes {self._init_block_sizes} != " f"kv_cache block_sizes {block_sizes}" ) - assert self._init_kernel_block_sizes == kernel_block_sizes, ( - f"InputBatch kernel_block_sizes {self._init_kernel_block_sizes} " - f"!= kv_cache kernel_block_sizes {kernel_block_sizes}" - ) def _attn_group_iterator(self) -> Iterator[AttentionGroup]: return itertools.chain.from_iterable(self.attn_groups) @@ -7267,12 +7260,16 @@ def initialize_kv_cache_tensors( num_attn_module = ( 2 if self.model_config.hf_config.model_type == "longcat_flash" else 1 ) + attn_groups = [g for groups in self.attn_groups for g in groups] + if self.speculative_config: + attn_groups += getattr(self.drafter, "draft_attn_groups", []) bind_kv_cache( kv_caches, self.compilation_config.static_forward_context, self.kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, + layer_kv_caches=map_kv_caches_to_kernel_blocks(kv_caches, attn_groups), ) return kv_caches @@ -7336,6 +7333,8 @@ def initialize_kv_cache( # backends for that group only supports block_size 64, we will return # kernel_block_size 64 and split the 256-token-block to 4 blocks with 64 # tokens each. + # Block-outermost packed groups are not split; their attention groups + # map kernel blocks themselves. kernel_block_sizes = prepare_kernel_block_sizes( kv_cache_config, self.attn_groups ) @@ -7345,7 +7344,7 @@ def initialize_kv_cache( self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes) # Reinitialize need to after initialize_attn_backend - self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) + self.may_reinitialize_input_batch(kv_cache_config) # Capture warmup providers that depend on allocated KV-cache strides. with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] kv_caches = self.initialize_kv_cache_tensors( @@ -7422,7 +7421,9 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(self.vllm_config): if isinstance(spec, AttentionSpec): - spec = attn_module.get_attn_backend().customize_spec(spec) + spec = customize_attention_spec( + attn_module.get_attn_backend(), spec + ) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index db0f81b81c36..d10cdcec6453 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -3,7 +3,7 @@ import math from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import product as iprod from typing import Any @@ -24,6 +24,7 @@ from vllm.v1.attention.backend import ( AttentionBackend, AttentionMetadataBuilder, + CommonAttentionMetadata, MultipleOf, ) from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy @@ -35,7 +36,6 @@ KVCacheLayout, KVCacheSpec, MambaSpec, - MLAAttentionSpec, UniformTypeKVCacheSpecs, create_kv_cache_views, ) @@ -166,6 +166,9 @@ def __init__( kv = static_forward_context[layer_name].kv_cache if not isinstance(kv, torch.Tensor): continue + if group.kernel_block_stride: + # Bound views address kernel blocks; zero by manager block. + kv = group.map_kv_cache_to_manager_blocks(kv, num_blocks) if kv.device.type != self.device.type: continue dp = kv.data_ptr() @@ -255,12 +258,22 @@ def warmup(self, num_kv_blocks: int) -> None: self.zero_block_ids([0]) -@dataclass +@dataclass(eq=False) class AttentionGroup: backend: type[AttentionBackend] layer_names: list[str] kv_cache_spec: KVCacheSpec kv_cache_group_id: int + # Set when the backend needs blocks smaller than the KV cache manager's: its + # block tables and KV cache views then address kernel blocks, kernel block j + # of block b being b * kernel_block_stride + j. The stride is set when a + # packed KV cache is bound (other layers' pages sit between its blocks); + # otherwise it stays 0 and the kernel blocks per block are used. + kernel_block_size: int | None = None + kernel_block_stride: int = 0 + # Persistent per ubatch: CUDA graphs capture the mapped block table. + kernel_block_table: torch.Tensor | None = None + kernel_block_offsets: torch.Tensor | None = None # When ubatching is enabled we will have a metadata builder for each ubatch # so that if they use internal persistent buffers for cudagraphs, and they # won't have to worry about conflicting with the other ubatches. @@ -275,27 +288,44 @@ def create_metadata_builders( kernel_block_size: int | None = None, num_metadata_builders: int = 1, ): + spec = self.kv_cache_spec + if kernel_block_size is not None and isinstance(spec, AttentionSpec): + # A backend in a packed KV cache group may need smaller blocks than + # the group is split into. + kernel_block_size = select_common_block_size( + kernel_block_size, [self.backend], [spec] + ) + if kernel_block_size != spec.block_size: + self.kernel_block_size = kernel_block_size if kernel_block_size is None: kv_cache_spec_builder = self.kv_cache_spec - elif ( - isinstance(self.kv_cache_spec, MLAAttentionSpec) - and self.kv_cache_spec.storage_block_size is not None - ): - kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( - self.kv_cache_spec.storage_block_size - ) else: kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size( kernel_block_size ) builder_cls = self.backend.get_builder_cls() builder_kwargs = {} - if builder_cls.requires_block_table_width: + if builder_cls.requires_block_table_width or self.kernel_block_size is not None: max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req( vllm_config, vllm_config.model_config.max_model_len ) - builder_kwargs["block_table_width"] = get_block_table_width( - max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size + width = get_block_table_width(max_num_blocks, self.kv_cache_spec.block_size) + if kernel_block_size is not None: + width *= self.kv_cache_spec.block_size // kernel_block_size + if builder_cls.requires_block_table_width: + builder_kwargs["block_table_width"] = width + if self.kernel_block_size is not None: + self.kernel_block_table = torch.zeros( + num_metadata_builders, + vllm_config.scheduler_config.max_num_seqs, + width, + dtype=torch.int32, + device=device, + ) + self.kernel_block_offsets = torch.arange( + spec.block_size // self.kernel_block_size, + dtype=torch.int32, + device=device, ) self.metadata_builders = [ builder_cls( @@ -307,9 +337,109 @@ def create_metadata_builders( ) for _ in range(num_metadata_builders) ] - if kernel_block_size is not None: - for builder in self.metadata_builders: - builder.set_kernel_block_size(kernel_block_size) + + def build_metadata( + self, + common_attn_metadata: CommonAttentionMetadata, + ubatch_idx: int = 0, + common_prefix_len: int = 0, + **kwargs, + ) -> Any: + return self.get_metadata_builder(ubatch_idx).build( + common_prefix_len, + self._with_kernel_block_table(common_attn_metadata, ubatch_idx), + **kwargs, + ) + + def build_metadata_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0 + ) -> Any: + return self.get_metadata_builder(ubatch_idx).build_for_cudagraph_capture( + self._with_kernel_block_table(common_attn_metadata, ubatch_idx) + ) + + def build_metadata_for_drafting( + self, common_attn_metadata: CommonAttentionMetadata, draft_index: int + ) -> Any: + return self.get_metadata_builder().build_for_drafting( + self._with_kernel_block_table(common_attn_metadata), draft_index + ) + + def _with_kernel_block_table( + self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0 + ) -> CommonAttentionMetadata: + if self.kernel_block_size is None: + return common_attn_metadata + return replace( + common_attn_metadata, + block_table_tensor=self.map_to_kernel_block_table( + common_attn_metadata.block_table_tensor, ubatch_idx + ), + ) + + def map_to_kernel_block_table( + self, block_table: torch.Tensor, ubatch_idx: int = 0 + ) -> torch.Tensor: + if self.kernel_block_size is None: + return block_table + assert self.kernel_block_table is not None + blocks_per_kv_block = self.kv_cache_spec.block_size // self.kernel_block_size + rows, cols = block_table.shape + out = self.kernel_block_table[ubatch_idx, :rows, : cols * blocks_per_kv_block] + torch.add( + self.kernel_block_offsets, + block_table.unsqueeze(-1), + alpha=self.kernel_block_stride or blocks_per_kv_block, + out=out.unflatten(1, (cols, -1)), + ) + return out + + def map_kv_cache_to_kernel_blocks(self, kv_cache: torch.Tensor) -> torch.Tensor: + """Views a packed cache in manager blocks as kernel blocks.""" + assert self.kernel_block_size is not None + spec = self.kv_cache_spec + kernel_rows = spec.get_num_kernel_states(self.kernel_block_size) + if kv_cache.shape[-2] == kernel_rows: + return kv_cache + # Token and pooled caches are written through manager-block slot + # mappings, so only compressed caches (written from their block table) + # can be mapped. + if ( + spec.tokens_per_state <= 1 + or kv_cache.shape[1] != 1 + or kv_cache.stride(0) % (kernel_rows * kv_cache.stride(-2)) + ): + raise ValueError( + f"{self.backend.get_name()} needs {self.kernel_block_size}-token " + f"kernel blocks, which packed {spec.block_size}-token KV cache " + f"blocks cannot be split into. Reduce --block-size to " + f"{self.kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a " + "layer-compact layout (e.g. LBNHC)." + ) + # Kernel block j of block b is b * kernel_block_stride + j. The view + # spans the other layers' pages between blocks (never addressing them), + # which ``view`` cannot express. + page_stride = kernel_rows * kv_cache.stride(-2) + self.kernel_block_stride = kv_cache.stride(0) // page_stride + num_pages = (kv_cache.shape[0] - 1) * self.kernel_block_stride + ( + kv_cache.shape[-2] // kernel_rows + ) + return kv_cache.as_strided( + (num_pages, *kv_cache.shape[1:-2], kernel_rows, kv_cache.shape[-1]), + (page_stride, *kv_cache.stride()[1:]), + ) + + def map_kv_cache_to_manager_blocks( + self, kv_cache: torch.Tensor, num_blocks: int + ) -> torch.Tensor: + """Views a packed cache in kernel blocks as manager blocks.""" + assert self.kernel_block_stride + spec = self.kv_cache_spec + rows = spec.get_num_kernel_states(spec.block_size) + return kv_cache.as_strided( + (num_blocks, *kv_cache.shape[1:-2], rows, kv_cache.shape[-1]), + (kv_cache.stride(0) * self.kernel_block_stride, *kv_cache.stride()[1:]), + ) def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder: assert len(self.metadata_builders) > ubatch_id @@ -327,17 +457,48 @@ def update_draft_decode_metadata( self.get_metadata_builder().update_draft_decode_metadata(metadata) +def map_kv_caches_to_kernel_blocks( + kv_caches: dict[str, torch.Tensor], attn_groups: Iterable[AttentionGroup] +) -> dict[str, torch.Tensor]: + """The views to bind to layers. The manager views in ``kv_caches`` stay for + KV connectors.""" + mapped = dict(kv_caches) + for group in attn_groups: + if group.kernel_block_size is not None: + mapped.update( + (name, group.map_kv_cache_to_kernel_blocks(kv_caches[name])) + for name in group.layer_names + if name in kv_caches + ) + return mapped + + +def customize_attention_spec( + backend: type[AttentionBackend], spec: AttentionSpec +) -> AttentionSpec: + """The backend's spec. Blocks larger than every kernel block size a backend + takes are split into kernel blocks, which packed blocks must keep whole.""" + spec = backend.customize_spec(spec) + sizes = backend.get_supported_kernel_block_sizes(spec) + fixed_sizes = [s for s in sizes if isinstance(s, int)] + if spec.block_stride_alignment is None and sizes and fixed_sizes == sizes: + spec = replace(spec, block_stride_alignment=MultipleOf(max(fixed_sizes))) + return spec + + def _block_size_is_supported( - backends: list[type[AttentionBackend]], block_size: int + backends: list[type[AttentionBackend]], + block_size: int, + specs: Sequence[KVCacheSpec | None], ) -> bool: """Check if the block size is supported by all backends. An exact ``int`` declaration must match exactly; a ``MultipleOf`` declaration accepts any multiple of its base. """ - for backend in backends: + for backend, spec in zip(backends, specs, strict=True): is_supported = False - for supported_size in backend.get_supported_kernel_block_sizes(): + for supported_size in backend.get_supported_kernel_block_sizes(spec): if isinstance(supported_size, int): if block_size == supported_size: is_supported = True @@ -354,6 +515,7 @@ def _block_size_is_supported( def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], + kv_cache_specs: Sequence[KVCacheSpec | None] | None = None, ) -> int: """Select a block size that is supported by all backends and is a factor of kv_manager_block_size. @@ -364,6 +526,7 @@ def select_common_block_size( Args: kv_manager_block_size: Block size of KV cache. backends: List of attention backend classes. + kv_cache_specs: Per-backend cache specs for spec-specific constraints. Returns: The selected block size. @@ -372,25 +535,27 @@ def select_common_block_size( ValueError: If no valid block size found. """ - if _block_size_is_supported(backends, kv_manager_block_size): + specs = kv_cache_specs or [None] * len(backends) + if _block_size_is_supported(backends, kv_manager_block_size, specs): return kv_manager_block_size # MultipleOf constraints also accept the manager size if they accept a divisor. # Any remaining candidate must therefore be an explicit size from a backend. candidates = { size - for backend in backends - for size in backend.get_supported_kernel_block_sizes() + for backend, spec in zip(backends, specs, strict=True) + for size in backend.get_supported_kernel_block_sizes(spec) if isinstance(size, int) and kv_manager_block_size % size == 0 } for size in sorted(candidates, reverse=True): - if _block_size_is_supported(backends, size): + if _block_size_is_supported(backends, size, specs): return size raise ValueError( f"No common block size for {kv_manager_block_size} (" + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" for b in backends + f"{b.get_name()}: {b.get_supported_kernel_block_sizes(s)}" + for b, s in zip(backends, specs, strict=True) ) + ")." ) @@ -450,9 +615,6 @@ def allocate_kv_cache( kernel_block_size = None if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[group_id] - if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: - kernel_block_size = spec.storage_block_size - views = create_kv_cache_views( buf, spec, @@ -497,24 +659,20 @@ def prepare_kernel_block_sizes( # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size group_backends = [g.backend for g in attn_groups[kv_cache_gid]] - storage_block_size = ( - kv_cache_spec.storage_block_size - if isinstance(kv_cache_spec, MLAAttentionSpec) - else None - ) - if storage_block_size is not None and _block_size_is_supported( - group_backends, storage_block_size - ): - # Storage-block specs (e.g. the GLM-5.3-Flash kpool indexer - # cache) address the cache in pool pages, and every other - # consumer (cache views, metadata builders, hisparse) already - # uses storage_block_size as the kernel block. Fall back to the - # backend vote when the group's backends do not accept it. - selected_kernel_size = storage_block_size - else: - selected_kernel_size = select_common_block_size( - kv_manager_block_size, group_backends + group_specs = [g.kv_cache_spec for g in attn_groups[kv_cache_gid]] + # Block-outermost packing puts other layers' pages between the + # blocks, so they are never split group-wide. + selected_kernel_size = ( + kv_manager_block_size + if any( + t.block_stride > t.layer_stride + for t in kv_cache_config.kv_cache_tensors + if not set(t.layers).isdisjoint(kv_cache_group.layer_names) ) + else select_common_block_size( + kv_manager_block_size, group_backends, group_specs + ) + ) kernel_block_sizes.append(selected_kernel_size) elif isinstance(kv_cache_spec, MambaSpec): # This is likely Mamba or other non-attention cache, no splitting. @@ -656,6 +814,7 @@ def bind_kv_cache( runner_kv_caches: list[torch.Tensor], num_attn_module: int = 1, kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, + layer_kv_caches: dict[str, torch.Tensor] | None = None, ) -> None: """Bind the allocated KV cache to both ModelRunner and forward context so that the KV cache can be used in the forward pass. @@ -674,6 +833,9 @@ def bind_kv_cache( runner_kv_caches: The kv_cache declared by ModelRunner. kv_cache_groups: The KV cache groups of the model, used to resolve layers that share a KV cache. + layer_kv_caches: The views to bind to layers when they differ from + kv_caches (e.g. kernel-block views). The runner keeps kv_caches, + which are addressed by scheduler block IDs. """ # Bind kv_caches to ModelRunner @@ -699,7 +861,7 @@ def bind_kv_cache( runner_kv_caches.append(kv_caches[layer_name]) bind_kv_cache_to_layers( - kv_caches, forward_context, num_attn_module, kv_cache_groups + layer_kv_caches or kv_caches, forward_context, num_attn_module, kv_cache_groups ) From e0cb3efce05388e21454556193d4069984b3154a Mon Sep 17 00:00:00 2001 From: Kai Hoang <56827800+hpdkhoa@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:52:08 +0700 Subject: [PATCH 313/952] [Bugfix][Structured Output] Decode choice specs as JSON in outlines and LMFE (#60748) Signed-off-by: hpdkhoa Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structured_output_manager.py | 18 +++++++++++++++++- .../backend_lm_format_enforcer.py | 3 +-- vllm/v1/structured_output/backend_outlines.py | 3 +-- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index fbceb3bd07e4..67c15825580c 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -16,7 +16,8 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_outlines import OutlinesGrammar +from vllm.v1.structured_output.backend_outlines import OutlinesBackend, OutlinesGrammar +from vllm.v1.structured_output.request import get_structured_output_key TOKENIZER = "gpt2" THINK_END = "\n" # reasoning-end marker (single GPT-2 token) @@ -560,3 +561,18 @@ def test_outlines_termination(tokenizer): assert grammar.validate_tokens([eos]) == [] assert grammar.accept_tokens(request.request_id, [one, eos, one]) assert grammar.is_terminated() + + +def test_outlines_choice_with_non_bmp_characters(tokenizer): + """The choice spec is JSON, which encodes emoji as surrogate pairs.""" + choices = ["😀 yes", "no"] + backend = OutlinesBackend( + VllmConfig(), tokenizer=tokenizer, vocab_size=len(tokenizer) + ) + request_type, grammar_spec = get_structured_output_key( + StructuredOutputsParams(choice=choices) + ) + + grammar = backend.compile_grammar(request_type, grammar_spec) + + assert grammar.accept_tokens("", tokenizer.encode(choices[0])) diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 36eb5eb159e1..9a16ddc1fe38 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from dataclasses import dataclass, field from functools import lru_cache @@ -117,7 +116,7 @@ def compile_grammar( grammar_spec, ) elif request_type == StructuredOutputOptions.CHOICE: - choices = ast.literal_eval(grammar_spec) + choices = json.loads(grammar_spec) character_level_parser = lmformatenforcer.UnionParser( [lmformatenforcer.StringParser(choice) for choice in choices] ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 1cfc6c8f5d1a..176d1a46e684 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -3,7 +3,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations -import ast import json from dataclasses import dataclass, field from re import _constants, _parser # type: ignore[attr-defined] @@ -68,7 +67,7 @@ def compile_grammar( elif request_type == StructuredOutputOptions.REGEX: regex = grammar_spec elif request_type == StructuredOutputOptions.CHOICE: - choices = ast.literal_eval(grammar_spec) + choices = json.loads(grammar_spec) choices = [regex_escape(c) for c in choices] regex = "(" + "|".join(choices) + ")" else: From ffdde24f546c503c7e47c5264b180c16fd668888 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Fri, 9 Oct 2026 13:54:40 -0700 Subject: [PATCH 314/952] [Docs] Document adaptive verification without a confidence head (#60913) Signed-off-by: Giancarlo Delfin --- .../adaptive_verification.md | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/docs/features/speculative_decoding/adaptive_verification.md b/docs/features/speculative_decoding/adaptive_verification.md index e9314c0e09f0..d560b5af3e4f 100644 --- a/docs/features/speculative_decoding/adaptive_verification.md +++ b/docs/features/speculative_decoding/adaptive_verification.md @@ -4,7 +4,7 @@ Speculative decoding buys fewer decode steps with more compute. At batch size 1 That matters because per-position acceptance decays fast. While the GPU is memory-bound that slot is effectively free and worth the gamble; once it saturates the gamble has a real throughput cost. The crossover moves with load and with workload-dependent acceptance rates, so no static `num_speculative_tokens` is right across concurrencies. -Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position confidences, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two. +Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position acceptance estimates, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two. The budget itself comes from a cost model profiled at startup. vLLM measures what a step costs at each shape, then picks the token count that maximizes expected accepted tokens per second. @@ -12,11 +12,14 @@ The practical effect is that one configuration holds up across the whole load ra ## Support -Adaptive verification needs per-position acceptance estimates, so today it is only supported for DSpark with a **confidence head**. +Adaptive verification is supported for model-based drafters (EAGLE, MTP, DFlash, DSpark and standalone draft models). The per-position acceptance estimates come from one of two sources: + +- **Confidence head**: DSpark checkpoints that ship a trained confidence head use its predictions directly. +- **Online acceptance estimator**: every other drafter predicts acceptance from its draft logits with a small logistic model, recalibrated during serving against the target's verification results. ## Usage -It is off by default. Enable it in the speculative config: +It is off by default. Enable it in the speculative config with `enable_adaptive_verification`. For example, with DSpark: ```bash vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ @@ -30,6 +33,21 @@ vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ }' ``` +Or with DFlash on MiMo-V2.5-Pro, whose drafter ships in the checkpoint's `dflash/` subdirectory: + +```bash +hf download XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash --local-dir MiMo-V2.5-Pro-FP4-DFlash + +vllm serve MiMo-V2.5-Pro-FP4-DFlash \ + --tensor-parallel-size 4 --trust-remote-code --max-model-len auto \ + --speculative-config '{ + "method": "dflash", + "model": "MiMo-V2.5-Pro-FP4-DFlash/dflash", + "num_speculative_tokens": 8, + "enable_adaptive_verification": true + }' +``` + Set `enable_adaptive_verification: false` to verify the full block for every request. ## Requirements and limitations @@ -37,6 +55,7 @@ Set `enable_adaptive_verification: false` to verify the full block for every req - The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend. - Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup. - Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries) or pipeline parallelism (cost curves and confidences exist only on the last rank). +- Without a confidence head, `use_local_argmax_reduction` is rejected. The estimator needs the full draft logits, which local argmax reduction never materializes. ## Tuning the cost profile From 35e920d671539d7066b648dc82058f8ef5348b9f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 10 Oct 2026 05:10:42 +0800 Subject: [PATCH 315/952] [Rust Frontend] Support `--stream-interval` and per-request `stream_interval` (#60589) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/cmd/src/cli.rs | 14 + rust/src/cmd/src/cli/tests.rs | 34 ++ rust/src/cmd/src/cli/unsupported.rs | 4 - .../examples/external_engine_logprobs.rs | 3 +- rust/src/engine-core-client/src/client.rs | 21 +- rust/src/engine-core-client/src/client/imp.rs | 73 ++-- .../src/client/output_channel.rs | 321 ++++++++++++++++++ .../engine-core-client/src/client/state.rs | 286 ++++++++++++---- .../engine-core-client/src/client/stream.rs | 83 ++++- rust/src/engine-core-client/src/lib.rs | 2 +- .../src/protocol/sampling.rs | 7 + .../engine-core-client/src/tests/client.rs | 43 +-- rust/src/llm/src/lib.rs | 23 +- rust/src/llm/src/output.rs | 79 ++++- rust/src/llm/tests/generate.rs | 106 ++++++ rust/src/mock-engine/src/tests.rs | 10 +- .../examples/external_engine_openai_qwen.rs | 2 + rust/src/server/src/config.rs | 5 + rust/src/server/src/lib.rs | 4 +- .../routes/openai/chat_completions/convert.rs | 13 + .../routes/openai/chat_completions/types.rs | 8 + .../src/routes/openai/completions/convert.rs | 10 +- .../src/routes/openai/completions/types.rs | 7 + rust/src/text/src/lower.rs | 29 ++ rust/src/text/src/request.rs | 7 + 25 files changed, 1040 insertions(+), 154 deletions(-) create mode 100644 rust/src/engine-core-client/src/client/output_channel.rs diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index d3e369bbe388..14b949c87de4 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -12,6 +12,7 @@ mod unsupported; use std::collections::HashMap; use std::ffi::{OsStr, OsString}; +use std::num::NonZeroU32; use std::path::PathBuf; use std::time::Duration; @@ -268,6 +269,13 @@ pub struct SharedRuntimeArgs { #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] #[serde(default)] pub max_logprobs: Option, + /// The interval (or buffer size) for streaming in terms of token length. + /// A smaller value (1) makes streaming smoother by sending each token + /// immediately, while a larger value (e.g., 10) reduces host overhead and + /// may increase throughput by batching multiple tokens before sending. + #[arg(long, default_value_t = NonZeroU32::MIN)] + #[serde(default = "default_stream_interval")] + pub stream_interval: NonZeroU32, /// Maximum time to wait for active requests to drain during shutdown. #[arg(long, default_value_t = 0)] #[serde(default)] @@ -547,6 +555,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -609,6 +618,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -646,6 +656,10 @@ fn default_engine_ready_timeout_secs() -> u64 { 600 } +fn default_stream_interval() -> NonZeroU32 { + NonZeroU32::MIN +} + fn default_cors_wildcard() -> JsonStringList { JsonStringList(vec!["*".to_string()]) } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index aa164bc618f8..57780027deac 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -189,6 +189,7 @@ fn serve_args_forward_python_flags_with_separator() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -373,6 +374,34 @@ fn serve_args_forward_profiler_config_to_managed_engine() { ); } +#[test] +fn serve_args_forward_stream_interval_to_frontend_config() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "4", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string()); + assert_eq!(frontend_config.stream_interval.get(), 4); + + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "0", + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); +} + #[test] fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() { let cli = Cli::try_parse_from([ @@ -1030,6 +1059,7 @@ fn frontend_args_accept_json() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1702,6 +1732,7 @@ fn serve_args_accept_handshake_aliases() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1866,6 +1897,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1961,6 +1993,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -2100,6 +2133,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index f65f766e847c..0a14d1b841c4 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -359,10 +359,6 @@ pub struct EngineUnsupportedArgs { #[arg(long)] pub collect_detailed_traces: Option, - /// The interval (or buffer size) for streaming in terms of token length. - #[arg(long)] - pub stream_interval: Option, - /// Maximum number of requests that can be in-flight (waiting or running) /// at the same time. When the limit is reached, new requests are rejected /// with HTTP 503 so the client can retry on another instance. diff --git a/rust/src/engine-core-client/examples/external_engine_logprobs.rs b/rust/src/engine-core-client/examples/external_engine_logprobs.rs index d4ec116f173f..e0da15623ca3 100644 --- a/rust/src/engine-core-client/examples/external_engine_logprobs.rs +++ b/rust/src/engine-core-client/examples/external_engine_logprobs.rs @@ -87,8 +87,9 @@ fn build_request( } async fn wait_for_final_output( - mut stream: vllm_engine_core_client::EngineCoreOutputStream, + stream: vllm_engine_core_client::EngineCoreOutputStream, ) -> Result { + let mut stream = stream.into_outputs(); while let Some(output) = stream.next().await { let output = output.context("failed to receive engine-core output")?; if output.finished() { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index b3516936fef2..d8b436c6466c 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeMap; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::Duration; @@ -25,10 +26,11 @@ use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime}; use crate::transport::{self, ConnectedEngine}; pub(crate) mod imp; +mod output_channel; mod state; mod stream; -pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput}; +pub use stream::{EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput}; /// How the frontend acquires its request/response transport with Python /// `EngineCoreProc`s. @@ -614,6 +616,12 @@ fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> { impl EngineCoreClient { /// Add a new request to the engine and return a per-request raw output /// stream. + /// + /// With `sampling_params.stream_interval` above one, outputs after the + /// first are delivered to the stream in batches of at least that many new + /// tokens, and the terminal output is delivered immediately together with + /// any held-back outputs. The stream still yields every raw output in + /// order; batching only reduces how often the consuming task is woken. pub async fn call(&self, mut req: EngineCoreRequest) -> Result { req.client_index = self.config.client_index; req.validate()?; @@ -628,8 +636,15 @@ impl EngineCoreClient { let request_id = req.request_id.clone(); let lora_name = req.lora_request.as_ref().map(|lora| lora.lora_name.clone()); let data_parallel_rank = req.data_parallel_rank; - let (engine_id, rx) = - self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?; + let stream_interval = (req.sampling_params.as_ref()) + .and_then(|params| params.stream_interval) + .unwrap_or(NonZeroU32::MIN); + let (engine_id, rx) = self.inner.register_request( + request_id.clone(), + lora_name, + data_parallel_rank, + stream_interval, + )?; // Construct the output stream first before actually sending the request to the engine. // This ensures that cancelling the future will properly clean up the request registration diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index c8abb95bbdf8..0a4af1a5ee5f 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet}; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; @@ -15,7 +16,9 @@ use tracing::{debug, info, trace, warn}; use vllm_metrics::METRICS; use zeromq::RouterSendHalf; -use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, UtilityRegistry}; +use crate::client::state::{ + OutputReceiver, OutputSender, RequestRegistry, UtilityReceiver, UtilityRegistry, +}; use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; @@ -111,12 +114,13 @@ impl ClientInner { request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { let mut registry = self.request_reg.lock(); if registry.is_closed() { return Err(self.closed_error()); } - registry.register(request_id, lora_name, data_parallel_rank) + registry.register(request_id, lora_name, data_parallel_rank, stream_interval) } /// Allocate the next utility `call_id` and register its waiting receiver. @@ -159,13 +163,15 @@ impl ClientInner { Ok(registry.abortable_request_ids(request_ids)) } - /// Obtain stream senders for a whole engine output batch with one registry - /// lock acquisition. - pub fn take_senders_for_outputs<'a>( + /// Send a whole engine output batch to the request streams with one + /// registry lock acquisition. Returns the outputs whose requests are no + /// longer tracked. + pub fn send_outputs( &self, - outputs: impl IntoIterator, - ) -> Vec>>> { - self.request_reg.lock().senders_for_outputs(outputs) + outputs: impl IntoIterator, + ) -> Vec { + let mut registry = self.request_reg.lock(); + outputs.into_iter().filter_map(|output| registry.send_output(output)).collect() } /// Remove a batch of requests that have finished or aborted, returning @@ -173,7 +179,7 @@ impl ClientInner { pub fn finish_requests<'a>( &self, request_ids: impl IntoIterator, - ) -> Vec>> { + ) -> Vec { self.request_reg.lock().finish_many(request_ids) } @@ -214,7 +220,7 @@ impl ClientInner { // Notify all ongoing requests that the client is closed. for sender in request_senders { - let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); + sender.send_error(Error::Shared(persistent_error.clone())); } for sender in utility_senders { let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); @@ -413,6 +419,16 @@ pub(crate) async fn run_abort_loop( } } +/// Number of tokens one output accounts for in `vllm:iteration_tokens_total`. +fn output_iteration_tokens(output: &EngineCoreOutput) -> u64 { + let mut tokens = output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + tokens += u64::from(prefill.num_computed_tokens); + } + tokens +} + /// Background loop that listens for engine-core outputs and dispatches them to /// the corresponding request streams based on their `request_id`. pub(crate) async fn run_output_dispatcher_loop( @@ -435,29 +451,20 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { let has_outputs = !batch.outputs.is_empty(); - let mut iteration_tokens = 0_u64; - let senders = inner.take_senders_for_outputs(&batch.outputs); - for (output, sender) in batch.outputs.into_iter().zip(senders) { - let request_id = output.request_id.clone(); - let Some(sender) = sender else { - debug!(request_id, "dropping output for inactive request"); - continue; - }; - - iteration_tokens += output.new_token_ids.len() as u64; - if let Some(prefill) = &output.prefill_stats { - // The engine emits prefill_stats once, on the first output. - iteration_tokens += u64::from(prefill.num_computed_tokens); - } - - let wrapped_output = EngineCoreStreamOutput { - engine_index: batch.engine_index, - timestamp: batch.timestamp, - output, - }; - if sender.send(Ok(wrapped_output)).is_err() { - debug!(request_id, "request output stream receiver dropped"); - } + let mut iteration_tokens: u64 = + batch.outputs.iter().map(output_iteration_tokens).sum(); + let (engine_index, timestamp) = (batch.engine_index, batch.timestamp); + let outputs = batch.outputs.into_iter().map(|output| EngineCoreStreamOutput { + engine_index, + timestamp, + output, + }); + for output in inner.send_outputs(outputs) { + debug!( + request_id = output.request_id, + "dropping output for inactive request" + ); + iteration_tokens -= output_iteration_tokens(&output); } if has_outputs diff --git a/rust/src/engine-core-client/src/client/output_channel.rs b/rust/src/engine-core-client/src/client/output_channel.rs new file mode 100644 index 000000000000..7b58d3dc2f1f --- /dev/null +++ b/rust/src/engine-core-client/src/client/output_channel.rs @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Per-request output channel between the output dispatcher and +//! [`EngineCoreOutputStream`](crate::EngineCoreOutputStream). +//! +//! With a stream interval above one, the sender holds back the outputs after +//! the first one until they carry at least that many new tokens, then delivers +//! them in one channel message, which the stream yields as one +//! [`EngineCoreStreamDelivery`](crate::EngineCoreStreamDelivery). The request +//! task is then woken once per delivery instead of once per engine step. +//! +//! Outputs are never merged, reordered, or dropped here; the consumer merges a +//! delivery into one output (see `vllm-llm`). Merging here would hide the raw +//! outputs that request metrics observe one by one (per-step timestamps for +//! inter-token latency, prefill stats, lifecycle events), and would have to +//! combine wire-level fields such as undecoded logprobs tensors and opaque +//! values. + +use std::num::NonZeroU32; + +use tokio::sync::mpsc; +use tracing::{debug, warn}; + +use crate::client::stream::{EngineCoreStreamDelivery as Delivery, EngineCoreStreamOutput}; +use crate::error::{Error, Result}; + +/// Create the output channel for one request. +pub(crate) fn output_channel(stream_interval: NonZeroU32) -> (OutputSender, OutputReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + let sender = OutputSender { + tx, + stream_interval: stream_interval.get() as usize, + delivered_first: false, + buffered: Vec::new(), + buffered_tokens: 0, + }; + (sender, rx) +} + +/// Receiving half of a request's output channel. +pub(crate) type OutputReceiver = mpsc::UnboundedReceiver>; + +/// Sending half of a request's output channel, owned by its registry entry. +#[derive(Debug)] +pub(crate) struct OutputSender { + tx: mpsc::UnboundedSender>, + /// Minimum number of new tokens per delivery after the first one. + stream_interval: usize, + /// Whether the first output has been delivered. + delivered_first: bool, + /// Outputs held back since the last delivery. + buffered: Vec, + /// Number of new tokens in `buffered`. + buffered_tokens: usize, +} + +impl OutputSender { + /// Send one output, or hold it back until the held-back outputs reach the + /// stream interval. The first output and the terminal output are always + /// delivered immediately, the latter together with any held-back outputs. + pub fn send(&mut self, output: EngineCoreStreamOutput) { + let finished = output.finished(); + if self.buffered.is_empty() + && (finished || !self.delivered_first || self.stream_interval == 1) + { + self.delivered_first = true; + self.deliver(Delivery::One(output)); + return; + } + + self.buffered_tokens += output.new_token_ids.len(); + self.buffered.push(output); + if finished || self.buffered_tokens >= self.stream_interval { + self.flush(); + } + } + + /// Deliver the held-back outputs, if any. + pub fn flush(&mut self) { + if self.buffered.is_empty() { + return; + } + self.buffered_tokens = 0; + let outputs = std::mem::take(&mut self.buffered); + self.deliver(Delivery::Many(outputs)); + } + + /// Deliver the held-back outputs, then terminate the stream with `error`. + pub fn send_error(mut self, error: Error) { + self.flush(); + let _ = self.tx.send(Err(error)); + } + + fn deliver(&self, delivery: Delivery) { + if let Err(mpsc::error::SendError(Ok(delivery))) = self.tx.send(Ok(delivery)) { + debug!( + request_id = delivery.first().request_id, + "request output stream receiver dropped" + ); + } + } +} + +impl Drop for OutputSender { + fn drop(&mut self) { + // Every path that ends a live request delivers held-back outputs first + // (terminal output, `flush`, or `send_error`). Outputs left here are lost. + if !self.buffered.is_empty() && !self.tx.is_closed() { + warn!( + request_id = self.buffered[0].request_id, + count = self.buffered.len(), + "dropping request output sender with held-back outputs" + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use expect_test::expect; + use tokio::sync::mpsc::error::TryRecvError; + + use super::*; + use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput}; + + fn output(token_ids: &[u32]) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output: EngineCoreOutput { + request_id: "req".to_string(), + new_token_ids: token_ids.to_vec(), + ..EngineCoreOutput::default() + }, + } + } + + fn finished(token_ids: &[u32]) -> EngineCoreStreamOutput { + let mut output = output(token_ids); + output.output.finish_reason = Some(EngineCoreFinishReason::Length); + output + } + + fn channel(stream_interval: u32) -> (OutputSender, OutputReceiver) { + output_channel(NonZeroU32::new(stream_interval).unwrap()) + } + + /// Render a delivery as the token ids of its outputs. + fn token_ids(delivery: Delivery) -> Vec> { + delivery.into_iter().map(|output| output.output.new_token_ids).collect() + } + + /// Drain all pending channel messages, rendering each delivery as the + /// token ids of its outputs. + fn deliveries(receiver: &mut OutputReceiver) -> Vec>> { + let mut deliveries = Vec::new(); + loop { + match receiver.try_recv() { + Ok(Ok(delivery)) => deliveries.push(token_ids(delivery)), + Ok(Err(error)) => panic!("unexpected error delivery: {error:?}"), + Err(TryRecvError::Empty | TryRecvError::Disconnected) => return deliveries, + } + } + } + + #[test] + fn stream_interval_one_delivers_every_output() { + let (mut sender, mut receiver) = channel(1); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[4])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [ + 4, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn stream_interval_batches_outputs_after_the_first() { + let (mut sender, mut receiver) = channel(3); + // The first output is delivered immediately; later ones are held back + // until they carry at least 3 tokens, counted across multi-token + // outputs. The terminal output flushes whatever is held back. + for token_ids in [&[1][..], &[2], &[3], &[4, 5], &[6], &[7, 8, 9, 10], &[11]] { + sender.send(output(token_ids)); + } + sender.send(finished(&[12])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + [ + 3, + ], + [ + 4, + 5, + ], + ], + [ + [ + 6, + ], + [ + 7, + 8, + 9, + 10, + ], + ], + [ + [ + 11, + ], + [ + 12, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn terminal_output_without_held_back_outputs_is_delivered_alone() { + let (mut sender, mut receiver) = channel(2); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn flush_delivers_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.flush(); + sender.flush(); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn error_follows_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.send_error(Error::EngineCoreDead); + + let mut items = Vec::new(); + while let Ok(item) = receiver.try_recv() { + items.push(item.map(|delivery| token_ids(delivery).len()).ok()); + } + // One delivery for the first output, one for the held-back output, + // then the error. + assert_eq!(items, [Some(1), Some(1), None]); + } +} diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index a8f0fa3c24ed..9c183b87e4e0 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -2,12 +2,15 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::num::NonZeroU32; use std::sync::atomic::{AtomicU64, Ordering}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::oneshot; use tracing::trace; use crate::EngineId; +use crate::client::output_channel::output_channel; +pub(crate) use crate::client::output_channel::{OutputReceiver, OutputSender}; use crate::client::stream::EngineCoreStreamOutput; use crate::error::{Error, Result}; use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput}; @@ -15,8 +18,6 @@ use crate::protocol::stats::SchedulerStats; use crate::protocol::utility::UtilityOutput; use crate::transport::ConnectedEngine; -pub type OutputSender = mpsc::UnboundedSender>; -pub type OutputReceiver = mpsc::UnboundedReceiver>; pub type UtilitySender = oneshot::Sender>; pub type UtilityReceiver = oneshot::Receiver>; @@ -137,18 +138,22 @@ impl RequestRegistry { /// When `data_parallel_rank` is provided, the request is routed directly to /// the engine at that rank index, bypassing load balancing. Otherwise /// the engine with the lowest routing score is chosen. + /// + /// `stream_interval` sets how many new tokens the channel batches into one + /// delivery after the first output; see [`OutputSender::send`]. pub fn register( &mut self, request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { if self.requests.contains_key(&request_id) { return Err(Error::DuplicateRequestId { request_id }); } let engine_id = self.choose_engine_for_request(data_parallel_rank)?; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = output_channel(stream_interval); let lora = lora_name.map(|adapter_name| LoraRequestState { adapter_name, phase: LoraPhase::Waiting, @@ -213,17 +218,30 @@ impl RequestRegistry { by_engine } - /// Obtain the stream sender for one output. If it indicates the request is - /// finished, it will be removed from the registry. - pub fn sender_for_output(&mut self, output: &EngineCoreOutput) -> Option { - self.apply_lora_events(output); + /// Send one output to its request stream. If it indicates the request is + /// finished, the request will be removed from the registry. Returns the + /// output back if its request is no longer tracked. + /// + /// The send happens under the registry lock: the sender is stateful (see + /// [`OutputSender::send`]) and owned by the registry entry, and sending on + /// its unbounded channel never blocks. + pub fn send_output( + &mut self, + output: EngineCoreStreamOutput, + ) -> Option { + self.apply_lora_events(&output.output); if output.finished() { - self.remove(output.request_id.as_str()).map(|tracked| tracked.0) + match self.remove(output.request_id.as_str()) { + Some((mut sender, _)) => sender.send(output), + None => return Some(output), + } } else { - self.requests - .get(output.request_id.as_str()) - .map(|tracked| tracked.sender.clone()) + match self.requests.get_mut(output.request_id.as_str()) { + Some(tracked) => tracked.sender.send(output), + None => return Some(output), + } } + None } /// Advance the request's LoRA scheduling phase from the engine-core events @@ -267,17 +285,8 @@ impl RequestRegistry { (running, waiting) } - /// Obtain stream senders for a whole engine output batch under one - /// registry lock. Finished outputs are removed before returning. - pub fn senders_for_outputs<'a>( - &mut self, - outputs: impl IntoIterator, - ) -> Vec> { - outputs.into_iter().map(|output| self.sender_for_output(output)).collect() - } - /// Remove a batch of requests that have finished or aborted, returning - /// their stream senders. + /// their stream senders after delivering any held-back outputs. pub fn finish_many<'a>( &mut self, request_ids: impl IntoIterator, @@ -285,6 +294,10 @@ impl RequestRegistry { request_ids .into_iter() .filter_map(|request_id| self.remove(request_id.as_str()).map(|tracked| tracked.0)) + .map(|mut sender| { + sender.flush(); + sender + }) .collect() } @@ -325,7 +338,7 @@ impl RequestRegistry { ) -> Vec { let mut aborted = Vec::new(); for request_id in request_ids { - let Some((sender, engine_id)) = self.remove(request_id) else { + let Some((mut sender, engine_id)) = self.remove(request_id) else { continue; }; let output = EngineCoreStreamOutput { @@ -337,7 +350,7 @@ impl RequestRegistry { ..EngineCoreOutput::default() }, }; - let _ = sender.send(Ok(output)); + sender.send(output); aborted.push(request_id.clone()); } aborted @@ -469,11 +482,13 @@ impl UtilityRegistry { #[cfg(test)] mod tests { use std::collections::BTreeSet; + use std::num::NonZeroU32; use crate::EngineId; use crate::client::state::{ - EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry, + EngineLoadSnapshot, EngineRoutingState, OutputReceiver, RequestRegistry, UtilityRegistry, }; + use crate::client::stream::EngineCoreStreamOutput; use crate::mock_engine::default_ready_response; use crate::protocol::output::{ EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, @@ -508,6 +523,14 @@ mod tests { } } + fn stream_output(output: EngineCoreOutput) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output, + } + } + fn adapter_names(values: &[&str]) -> BTreeSet { values.iter().map(|name| (*name).to_string()).collect() } @@ -515,8 +538,9 @@ mod tests { #[test] fn registry_rejects_duplicate_request_ids() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - let error = registry.register("req-1".to_string(), None, None).unwrap_err(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let error = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap_err(); assert!(matches!( error, crate::error::Error::DuplicateRequestId { request_id } if request_id == "req-1" @@ -526,25 +550,99 @@ mod tests { #[test] fn registry_removes_finished_request_on_output() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); - let sender = registry.sender_for_output(&EngineCoreOutput { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { request_id: "req-1".to_string(), finish_reason: Some(EngineCoreFinishReason::Length), ..Default::default() - }); + })); - assert!(sender.is_some()); + assert!(unrouted.is_none()); assert!(!registry.contains("req-1")); } + /// Drain the outputs that are ready on `receiver` without waiting, + /// rendering each as its token ids and finish reason. + fn ready_outputs( + receiver: &mut OutputReceiver, + ) -> Vec<(Vec, Option)> { + let mut outputs = Vec::new(); + while let Ok(delivery) = receiver.try_recv() { + for output in delivery.unwrap() { + outputs.push((output.output.new_token_ids, output.output.finish_reason)); + } + } + outputs + } + + #[test] + fn registry_abort_delivers_held_back_outputs_before_terminal_output() { + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); + let (_, mut receiver) = registry + .register("req-1".to_string(), None, None, NonZeroU32::new(4).unwrap()) + .unwrap(); + + for token_id in 1..=3 { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { + request_id: "req-1".to_string(), + new_token_ids: vec![token_id], + ..Default::default() + })); + assert!(unrouted.is_none()); + } + // Only the first output is delivered; the next two are held back + // below the stream interval. + expect_test::expect![[r#" + [ + ( + [ + 1, + ], + None, + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + + registry.abort_many(&["req-1".to_string()], 0.0); + expect_test::expect![[r#" + [ + ( + [ + 2, + ], + None, + ), + ( + [ + 3, + ], + None, + ), + ( + [], + Some( + Abort, + ), + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + } + #[test] fn registry_tracks_lora_phases_from_engine_events() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Registered but not yet scheduled: counted as waiting. The non-LoRA // request never shows up. @@ -554,7 +652,7 @@ mod tests { ); // Queued then scheduled in one output: running. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Queued, EngineCoreEventType::Scheduled], None, @@ -565,7 +663,7 @@ mod tests { ); // Preempted: back to waiting. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Preempted], None, @@ -576,7 +674,7 @@ mod tests { ); // Finished: dropped from tracking entirely. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Scheduled], Some(EngineCoreFinishReason::Stop), @@ -591,18 +689,33 @@ mod tests { fn registry_unions_lora_adapters_across_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-a1".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a1".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-a2".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a2".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-b".to_string(), Some("adapter-b".to_string()), None) + .register( + "req-b".to_string(), + Some("adapter-b".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); // One of adapter-a's requests starts running while the other waits: // the adapter appears in both sets. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-a1", &[EngineCoreEventType::Scheduled], None, @@ -620,7 +733,7 @@ mod tests { fn registry_counts_only_active_lora_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(registry.active_lora_requests(), 0); assert_eq!( registry.lora_adapter_states(), @@ -632,6 +745,7 @@ mod tests { "req-lora-a".to_string(), Some("adapter-a".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); registry @@ -639,6 +753,7 @@ mod tests { "req-lora-b".to_string(), Some("adapter-b".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); assert_eq!(registry.active_lora_requests(), 2); @@ -661,7 +776,12 @@ mod tests { fn registry_clears_lora_count_on_close() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); assert_eq!(registry.active_lora_requests(), 1); @@ -677,7 +797,12 @@ mod tests { fn registry_drops_lora_tracking_on_abort() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); drop(registry.finish_many(&["req-lora".to_string()])); @@ -691,8 +816,8 @@ mod tests { #[test] fn registry_closes_all_requests_on_failure() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - registry.register("req-2".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); let senders = registry.close(); @@ -708,9 +833,12 @@ mod tests { connected_engine(engine_0.clone()), connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -737,9 +865,12 @@ mod tests { connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -809,7 +940,8 @@ mod tests { } )); - let (chosen, _) = registry.register("req-stats".to_string(), None, None).unwrap(); + let (chosen, _) = + registry.register("req-stats".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -824,7 +956,9 @@ mod tests { for (rank, running) in counts.into_iter().enumerate() { for index in 0..running { let request_id = format!("old-{rank}-{index}"); - registry.register(request_id.clone(), None, Some(rank as u32)).unwrap(); + registry + .register(request_id.clone(), None, Some(rank as u32), NonZeroU32::MIN) + .unwrap(); old_requests.push(request_id); } registry.apply_scheduler_counts( @@ -839,7 +973,8 @@ mod tests { let mut distribution = [0; 2]; for index in 0..14 { - let (engine, _) = registry.register(format!("new-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("new-{index}"), None, None, NonZeroU32::MIN).unwrap(); distribution[engine.engine_index().unwrap() as usize] += 1; } distributions.push(distribution); @@ -865,7 +1000,8 @@ mod tests { ); } let request_id = format!("req-{index}"); - let (engine, _) = registry.register(request_id.clone(), None, None).unwrap(); + let (engine, _) = + registry.register(request_id.clone(), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); drop(registry.finish_many(&[request_id])); } @@ -890,12 +1026,15 @@ mod tests { }, ); for index in 0..4 { - registry.register(format!("pinned-{index}"), None, Some(0)).unwrap(); + registry + .register(format!("pinned-{index}"), None, Some(0), NonZeroU32::MIN) + .unwrap(); } // Even unchanged counts replace the optimistic estimate for that engine. registry.apply_scheduler_counts(0, snapshot); - let (after_refresh, _) = registry.register("after-refresh".into(), None, None).unwrap(); + let (after_refresh, _) = + registry.register("after-refresh".into(), None, None, NonZeroU32::MIN).unwrap(); // An older, empty snapshot must not erase the five admitted requests. for rank in [0, 1] { registry.apply_scheduler_counts( @@ -906,7 +1045,8 @@ mod tests { }, ); } - let (after_empty, _) = registry.register("after-empty".into(), None, None).unwrap(); + let (after_empty, _) = + registry.register("after-empty".into(), None, None, NonZeroU32::MIN).unwrap(); expect_test::expect!["(Some(0), Some(1))"].assert_eq(&format!( "{:?}", (after_refresh.engine_index(), after_empty.engine_index()) @@ -927,10 +1067,11 @@ mod tests { }, ); } - registry.register("pinned".into(), None, Some(9)).unwrap(); + registry.register("pinned".into(), None, Some(9), NonZeroU32::MIN).unwrap(); let mut chosen = Vec::new(); for index in 0..3 { - let (engine, _) = registry.register(format!("auto-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("auto-{index}"), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); } expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}")); @@ -948,15 +1089,18 @@ mod tests { ]); // Explicitly target rank 2 (third engine). - let (chosen, _) = registry.register("req-1".to_string(), None, Some(2)).unwrap(); + let (chosen, _) = + registry.register("req-1".to_string(), None, Some(2), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_2); // Explicitly target rank 0 (first engine). - let (chosen, _) = registry.register("req-2".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-2".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); // Explicitly target rank 1. - let (chosen, _) = registry.register("req-3".to_string(), None, Some(1)).unwrap(); + let (chosen, _) = + registry.register("req-3".to_string(), None, Some(1), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -970,11 +1114,12 @@ mod tests { ]); // Load-balance: first two go to engine_0 and engine_1. - registry.register("req-lb-0".to_string(), None, None).unwrap(); + registry.register("req-lb-0".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Now engine_0 has 1 in-flight. Without dp_rank, next would go to engine_1. // But with dp_rank=0, it should still go to engine_0. - let (chosen, _) = registry.register("req-dp".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-dp".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); } @@ -985,7 +1130,9 @@ mod tests { connected_engine(EngineId::from_engine_index(1)), ]); - let error = registry.register("req-1".to_string(), None, Some(2)).unwrap_err(); + let error = registry + .register("req-1".to_string(), None, Some(2), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { @@ -1000,10 +1147,13 @@ mod tests { let engine_3 = EngineId::from_engine_index(3); let mut registry = RequestRegistry::new(&[connected_engine(engine_3.clone())]); - let (chosen, _) = registry.register("req-ok".to_string(), None, Some(3)).unwrap(); + let (chosen, _) = + registry.register("req-ok".to_string(), None, Some(3), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_3); - let error = registry.register("req-bad".to_string(), None, Some(0)).unwrap_err(); + let error = registry + .register("req-bad".to_string(), None, Some(0), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index ba408f2c8788..6f363f61017d 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -5,8 +5,9 @@ use std::ops::Deref; use std::pin::Pin; use std::task::{Context, Poll}; -use futures::Stream; use futures::stream::FusedStream; +use futures::{Stream, StreamExt as _}; +use itertools::Either; use thiserror_ext::AsReport as _; use tokio::sync::mpsc; use tracing::{debug, error, warn}; @@ -40,6 +41,68 @@ impl Deref for EngineCoreStreamOutput { } } +/// Consecutive outputs of one request, delivered to its stream together and in +/// order. Outputs are grouped by the request's `stream_interval` (see +/// [`EngineCoreClient::call`]) but not merged; consumers merge them after +/// observing each raw output. +/// +/// [`EngineCoreClient::call`]: crate::EngineCoreClient::call +#[derive(Debug)] +#[allow( + clippy::large_enum_variant, + reason = "single outputs stay inline, avoiding an allocation per engine step" +)] +pub enum EngineCoreStreamDelivery { + /// A single output, sent without being held back. + One(EngineCoreStreamOutput), + /// Outputs held back by the stream interval, in order. Never empty. + Many(Vec), +} + +impl EngineCoreStreamDelivery { + /// Return the first output. + pub fn first(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.first().expect("deliveries are never empty"), + } + } + + /// Return the last output, the only one that can be terminal. + pub fn last(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.last().expect("deliveries are never empty"), + } + } + + /// Iterate over the outputs in order. + pub fn iter( + &self, + ) -> Either< + std::iter::Once<&EngineCoreStreamOutput>, + std::slice::Iter<'_, EngineCoreStreamOutput>, + > { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.iter()), + } + } +} + +impl IntoIterator for EngineCoreStreamDelivery { + type Item = EngineCoreStreamOutput; + type IntoIter = + Either, std::vec::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.into_iter()), + } + } +} + /// Stream of raw engine-core outputs for one request. /// /// The stream yields only [`EngineCoreStreamOutput`] values whose embedded @@ -79,10 +142,21 @@ impl EngineCoreOutputStream { pub fn engine_index(&self) -> u32 { self.engine_index } + + /// Flatten the stream into one item per raw output, for consumers that do + /// not care how outputs are grouped into deliveries. + pub fn into_outputs(self) -> impl Stream> + Send + Unpin { + self.flat_map(|item| { + futures::stream::iter(match item { + Ok(delivery) => Either::Left(delivery.into_iter().map(Ok)), + Err(error) => Either::Right(std::iter::once(Err(error))), + }) + }) + } } impl Stream for EngineCoreOutputStream { - type Item = Result; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.is_terminated() { @@ -93,10 +167,11 @@ impl Stream for EngineCoreOutputStream { Poll::Pending => Poll::Pending, Poll::Ready(Some(item)) => { match &item { - Ok(output) => { + Ok(delivery) => { // If the output indicates the request is finished, mark the stream as // terminated with cleanly-finished state and expect no more outputs to - // come. + // come. Only the last output of a delivery can be terminal. + let output = delivery.last(); if output.finished() { if output.finish_reason == Some(EngineCoreFinishReason::Error) { error!( diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs index 7b60206194ce..8f79e227de11 100644 --- a/rust/src/engine-core-client/src/lib.rs +++ b/rust/src/engine-core-client/src/lib.rs @@ -14,7 +14,7 @@ mod transport; pub use client::{ AbortCause, CoordinatorMode, EngineCoreClient, EngineCoreClientConfig, EngineCoreOutputStream, - EngineCoreStreamOutput, TransportMode, + EngineCoreStreamDelivery, EngineCoreStreamOutput, TransportMode, }; pub use error::{Error, Result}; pub use transport::{ENGINE_CORE_DEAD_SENTINEL, EngineId}; diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 577dcb651f20..51e80d0f66eb 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeSet, HashMap}; +use std::num::NonZeroU32; use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; @@ -161,6 +162,11 @@ pub struct EngineCoreSamplingParams { /// Number of prompt tokens to skip from returned routed-expert data. /// A value of zero returns routing data for the entire prompt. pub routed_experts_prompt_start: u32, + /// Number of newly generated tokens to batch into each streamed output. + /// The first and final outputs are always emitted immediately. Engine-core + /// ignores it; [`EngineCoreClient::call`](crate::EngineCoreClient::call) + /// batches output deliveries by it. + pub stream_interval: Option, } impl EngineCoreSamplingParams { @@ -195,6 +201,7 @@ impl EngineCoreSamplingParams { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index dbd3fa69ceb2..0ad58b027154 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -755,8 +755,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let final_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -778,7 +778,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { tokio::time::sleep(Duration::from_millis(100)).await; - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); let final_3 = timeout(Duration::from_secs(1), stream_3.next()) .await .unwrap() @@ -933,7 +933,7 @@ async fn coordinator_accepts_stats_only_outputs() { ) .await; - let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap(); + let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap().into_outputs(); let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(final_output.request_id, "req-stats"); @@ -1026,8 +1026,8 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() { assert!(client.ready_responses()[0].max_model_len > 0); assert_eq!(client.vllm_version(), "test-vllm-version"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_2 = timeout(Duration::from_secs(1), stream_2.next()) .await @@ -1103,7 +1103,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = match client.call(sample_request()).await { Ok(_) => panic!("expected duplicate request error"), Err(error) => error, @@ -1170,7 +1170,7 @@ async fn finished_requests_without_final_output_is_treated_as_unexpected_close() ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = timeout(Duration::from_secs(1), stream.next()) .await .unwrap() @@ -1233,7 +1233,7 @@ async fn dropping_a_live_stream_triggers_abort() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(first.new_token_ids, vec![99]); drop(stream); @@ -1316,7 +1316,7 @@ async fn dropping_multiple_live_streams_aborts_all_in_a_burst() { // emits outputs, then drain the first token from each stream. let mut streams = Vec::new(); for id in request_ids { - streams.push(client.call(sample_request_with_id(id)).await.unwrap()); + streams.push(client.call(sample_request_with_id(id)).await.unwrap().into_outputs()); } for stream in streams.iter_mut() { let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); @@ -1363,8 +1363,8 @@ async fn dispatcher_failure_propagates_to_streams_and_future_calls() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -1862,7 +1862,7 @@ async fn client_decodes_multipart_logprob_outputs() { ) .await; - let stream = client.call(sample_request()).await.unwrap(); + let stream = client.call(sample_request()).await.unwrap().into_outputs(); let outputs = stream.collect::>().await; assert_eq!(outputs.len(), 1); @@ -1950,7 +1950,7 @@ async fn client_sends_large_multimodal_tensor_as_aux_frame() { ) .await; - let outputs = client.call(request).await.unwrap().collect::>().await; + let outputs = client.call(request).await.unwrap().into_outputs().collect::>().await; assert_eq!(outputs.len(), 1); assert!(outputs[0].is_ok()); @@ -2079,8 +2079,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(client.ready_responses().len(), 2); assert_eq!(client.engine_identities()[0], b"engine-0"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-1" @@ -2100,7 +2100,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(final_1.new_token_ids, vec![10]); assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Length)); - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-3" @@ -2271,8 +2271,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() { assert!(client.is_sleeping().await.unwrap()); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); client .abort(&[ @@ -2803,6 +2803,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, }, ); @@ -3266,7 +3267,7 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends .await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let wakeup = timeout( Duration::from_secs(1), @@ -3352,7 +3353,7 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() { send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); assert!( timeout( diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 015624b409c8..7e9b5e0ba53f 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::num::NonZeroU32; + use tracing::Span; use vllm_engine_core_client::EngineCoreClient; @@ -34,6 +36,7 @@ use crate::request_metrics::RequestMetricsTracker; pub struct Llm { client: EngineCoreClient, randomize_request_id: bool, + stream_interval: NonZeroU32, stats_logger: Option, inflight: InflightRequests, } @@ -45,6 +48,7 @@ impl Llm { Self { client, randomize_request_id: true, + stream_interval: NonZeroU32::MIN, stats_logger: None, inflight: InflightRequests::new(), } @@ -71,6 +75,14 @@ impl Llm { self } + /// Set the frontend-level stream interval: the minimum number of newly + /// generated tokens batched into each streamed output after the first one. + /// A request's own `stream_interval` can only raise it. + pub fn with_stream_interval(mut self, stream_interval: NonZeroU32) -> Self { + self.stream_interval = stream_interval; + self + } + /// Expose the underlying engine-core client for low-level utility/admin /// calls. pub fn engine_core_client(&self) -> &EngineCoreClient { @@ -79,7 +91,16 @@ impl Llm { /// Submit one tokenized generate request and return a per-request output /// stream. - pub async fn generate(&self, req: GenerateRequest) -> Result { + pub async fn generate(&self, mut req: GenerateRequest) -> Result { + // Clamp the request's stream interval to the frontend-level one, like + // Python. Engine-core ignores it; the client batches deliveries by it. + let stream_interval = (req.sampling_params.stream_interval) + .map_or(self.stream_interval, |interval| { + interval.max(self.stream_interval) + }); + req.sampling_params.stream_interval = + (stream_interval > NonZeroU32::MIN).then_some(stream_interval); + let prepared = req.prepare(self.randomize_request_id)?; let prompt_token_ids = prepared.prompt_token_ids().into(); let external_request_id = prepared diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4dde96f2e20c..b6ccb7717508 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -16,7 +16,9 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_engine_core_client::protocol::tensor::WireNdArray; -use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; +use vllm_engine_core_client::{ + AbortCause, EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput, +}; use crate::error::Result; use crate::inflight::RequestGuard; @@ -195,6 +197,35 @@ impl GenerateOutput { pub fn finished(&self) -> bool { self.finish_reason.is_some() } + + /// Append a later output of the same request, as if both were produced in + /// one step. Per-token data is concatenated; per-output fields take the + /// later output's value, like Python's `RequestState.make_request_output()` + /// does for the output that flushes a stream interval. + fn merge(&mut self, next: GenerateOutput) { + debug_assert!( + next.prompt_info.is_none(), + "only the first output carries prompt info" + ); + self.token_ids.extend(next.token_ids); + if let Some(next_logprobs) = next.logprobs { + match &mut self.logprobs { + Some(logprobs) => logprobs.positions.extend(next_logprobs.positions), + None => self.logprobs = Some(next_logprobs), + } + } + if let Some(next_mask) = next.sampling_mask { + match &mut self.sampling_mask { + Some(mask) => mask.rows.extend(next_mask.rows), + None => self.sampling_mask = Some(next_mask), + } + } + self.cached_token_count = self.cached_token_count.max(next.cached_token_count); + self.finish_reason = next.finish_reason; + self.kv_transfer_params = next.kv_transfer_params; + self.ec_transfer_params = next.ec_transfer_params; + self.spec_decode_metrics = next.spec_decode_metrics; + } } #[cfg(any(test, feature = "test-util"))] @@ -266,18 +297,27 @@ impl GenerateOutputStream { pub fn request_id(&self) -> &str { self.raw_stream.request_id() } -} - -impl Stream for GenerateOutputStream { - type Item = Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let raw = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { - Some(Ok(raw)) => raw, - Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), - None => return Poll::Ready(None), + /// Merge the outputs of one delivery into one output, like Python + /// `RequestState.make_request_output()` does for outputs held back by + /// `stream_interval`. Request metrics still observe every raw output. + fn merge_delivery(&mut self, delivery: EngineCoreStreamDelivery) -> Result { + // Match rather than iterate, so the common single output is converted + // without first being moved into an iterator. + let mut outputs = match delivery { + EngineCoreStreamDelivery::One(output) => return self.convert_output(output), + EngineCoreStreamDelivery::Many(outputs) => outputs.into_iter(), }; + let first = outputs.next().expect("deliveries are never empty"); + let mut merged = self.convert_output(first)?; + for raw in outputs { + merged.merge(self.convert_output(raw)?); + } + Ok(merged) + } + /// Convert one raw engine-core output. + fn convert_output(&mut self, raw: EngineCoreStreamOutput) -> Result { let received_at = current_unix_timestamp_secs(); self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); @@ -301,11 +341,11 @@ impl Stream for GenerateOutputStream { if let Some(mask) = sampling_mask.as_ref() && mask.rows.len() != raw.new_token_ids.len() { - return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { + return Err(crate::Error::SamplingMaskTokenCountMismatch { request_id: raw.request_id, token_count: raw.new_token_ids.len(), row_count: mask.rows.len(), - }))); + }); } let cached_token_count = raw .prefill_stats @@ -331,7 +371,20 @@ impl Stream for GenerateOutputStream { spec_decode_metrics: raw.spec_decode_metrics, }; - Poll::Ready(Some(Ok(output))) + Ok(output) + } +} + +impl Stream for GenerateOutputStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let delivery = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { + Some(Ok(delivery)) => delivery, + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), + None => return Poll::Ready(None), + }; + Poll::Ready(Some(self.merge_delivery(delivery))) } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 813171555bba..e1205ede9961 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; +use std::num::NonZeroU32; use std::sync::{Arc, Once}; use std::time::Duration; @@ -308,6 +309,111 @@ async fn generate_streams_outputs() { llm.shutdown().await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn generate_merges_outputs_by_stream_interval() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-stream-interval".to_vec(); + + // One logprobs position per token, so merged outputs can be checked + // against the logprobs of their concatenated tokens. + fn token_logprobs(token_ids: &[u32]) -> Logprobs { + Logprobs { + positions: token_ids + .iter() + .map(|&token_id| PositionLogprobs { + entries: vec![TokenLogprob { + token_id, + logprob: -0.5, + rank: 1, + }], + }) + .collect(), + } + } + + let (shutdown_tx, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + |dealer, push| { + Box::pin(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + + // One engine step per batch. + let steps: [(&[u32], Option); 6] = [ + (&[1], None), + (&[2], None), + (&[3], None), + (&[4, 5], None), + (&[6], None), + (&[7], Some(EngineCoreFinishReason::Length)), + ]; + for (token_ids, finish_reason) in steps { + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output_with_logprobs( + &request.request_id, + token_ids.to_vec(), + finish_reason, + Some(token_logprobs(token_ids)), + None, + )], + ..Default::default() + } + .into(), + ) + .await; + } + }) + }, + ); + + // The request's interval raises the frontend-level one. + let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc) + .await + .with_stream_interval(NonZeroU32::new(2).unwrap()); + let mut request = sample_generate_request("req-stream-interval", 7); + request.sampling_params.stream_interval = NonZeroU32::new(3); + let outputs: Vec<_> = llm + .generate(request) + .await + .unwrap() + .map(|output| output.unwrap()) + .collect() + .await; + + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + llm.shutdown().await.unwrap(); + + // The first output is yielded immediately, later ones once they carry at + // least 3 tokens, and the terminal one flushes the rest. + let shapes: Vec<_> = outputs + .iter() + .map(|output| { + ( + output.prompt_info.is_some(), + output.token_ids.clone(), + output.finish_reason.clone(), + ) + }) + .collect(); + assert_eq!( + shapes, + [ + (true, vec![1], None), + (false, vec![2, 3, 4, 5], None), + (false, vec![6, 7], Some(FinishReason::Length)), + ] + ); + for output in &outputs { + assert_eq!(output.logprobs, Some(token_logprobs(&output.token_ids))); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { init_tracing(); diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index d82ceac083ca..0b053cbbb16a 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -116,7 +116,7 @@ async fn chunk_size_one_outputs_one_token_per_update() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-1", 3)).await.expect("call"); + let mut stream = client.call(sample_request("req-1", 3)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); @@ -147,7 +147,7 @@ async fn chunk_size_clips_final_output_to_max_tokens() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await; - let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call"); + let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 4); @@ -165,7 +165,11 @@ async fn abort_cancels_active_request_and_emits_terminal_output() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call"); + let mut stream = client + .call(sample_request("req-abort", 1_000_000)) + .await + .expect("call") + .into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.finish_reason, None); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 43192ab78f47..a391f56caeaa 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -81,6 +82,7 @@ async fn main() -> Result<()> { lora_modules: Vec::new(), chat_template_content_format: ChatTemplateContentFormatOption::Auto, max_logprobs: None, + stream_interval: NonZeroU32::MIN, api_server_options: ApiServerOptions::default(), cors: CorsConfig::default(), tls: None, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 8e51f1440b3e..3f95b5891cb2 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::str::FromStr; use std::time::Duration; @@ -258,6 +259,10 @@ pub struct Config { /// Optional maximum number of top log probabilities accepted by the /// frontend. `None` delegates to the text layer default. pub max_logprobs: Option, + /// Minimum number of newly generated tokens batched into each streamed + /// output after the first one. Requests can raise it with their own + /// `stream_interval`. + pub stream_interval: NonZeroU32, /// HTTP/API-server behavior switches. pub api_server_options: ApiServerOptions, /// CORS settings applied to every HTTP response. diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index f54f8ce3e423..d5999cd3e78b 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -123,7 +123,9 @@ async fn build_state(config: &Config) -> Result> { .await .context("failed to connect to engine core")?; - let llm = Llm::new(client).with_log_stats(!config.disable_log_stats); + let llm = Llm::new(client) + .with_log_stats(!config.disable_log_stats) + .with_stream_interval(config.stream_interval); let text = TextLlm::new(llm, text_backend).with_max_logprobs(config.max_logprobs); let chat = ChatLlm::new(text, chat_backend) diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 4f80dda6c7b0..46aaafbc69e2 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -179,6 +179,7 @@ pub(super) fn prepare_chat_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, chat_options: ChatOptions { generation_prompt_mode, @@ -443,6 +444,7 @@ fn convert_tool_choice(tool_choice: &ToolChoice) -> Result(request).is_err()); + } + #[test] fn chat_http_request_defaults_missing_or_null_model() { for model in [None, Some(serde_json::Value::Null)] { @@ -1004,6 +1015,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..base_request() }; @@ -1019,6 +1031,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..VllmSamplingParams::default() }; assert_eq!(prepared.chat_request.sampling_params, expected); diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 13c7c7233f14..204f3bab5aab 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -256,6 +257,12 @@ pub struct ChatCompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, } impl Default for ChatCompletionRequest { @@ -323,6 +330,7 @@ impl Default for ChatCompletionRequest { ec_transfer_params: None, vllm_xargs: None, repetition_detection: None, + stream_interval: None, } } } diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 0b37f84057f5..33fc6bbe0a07 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, decode_options: TextDecodeOptions { skip_special_tokens: request.skip_special_tokens, @@ -211,6 +212,8 @@ fn completion_echo_text( #[cfg(test)] mod tests { + use std::num::NonZeroU32; + use axum::http::HeaderMap; use serde_json::json; use validator::Validate; @@ -400,7 +403,8 @@ mod tests { "presence_penalty": 0.3, "repetition_penalty": 1.1, "ignore_eos": true, - "skip_special_tokens": false + "skip_special_tokens": false, + "stream_interval": 4 })) .expect("parse request"); @@ -436,6 +440,10 @@ mod tests { ); assert!(prepared.text_request.sampling_params.ignore_eos); assert!(!prepared.text_request.decode_options.skip_special_tokens); + assert_eq!( + prepared.text_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); } #[test] diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 2a6901370682..b10f8a23d701 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -114,6 +115,12 @@ pub struct CompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, + /// Length penalty for beam search pub length_penalty: Option, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 2b09070049df..9fa8c78f4f91 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -124,6 +124,7 @@ pub fn lower_sampling_params( structured_outputs, skip_reading_prefix_cache, vllm_xargs, + stream_interval, } = sampling_params; validate_logprobs( @@ -214,6 +215,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, routed_experts_prompt_start: 0, + stream_interval, }; validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; @@ -341,6 +343,7 @@ fn merge_unique_token_ids( #[cfg(test)] mod tests { use std::collections::{BTreeSet, HashMap}; + use std::num::NonZeroU32; use serial_test::file_serial; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; @@ -693,6 +696,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -746,6 +750,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -921,6 +926,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -992,6 +998,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1056,6 +1063,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1401,6 +1409,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1430,6 +1439,26 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_carries_stream_interval_in_sampling_params() { + let mut request = sample_request(); + request.sampling_params.stream_interval = NonZeroU32::new(4); + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!( + prepared.generate_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); + } + #[test] fn lower_text_request_passes_kv_hints_through() { let hints = KvHintsEnvelope { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 08758748ac79..b00d8f3b2f12 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; @@ -147,6 +148,11 @@ pub struct SamplingParams { pub skip_reading_prefix_cache: Option, /// Additional request parameters for custom extensions. pub vllm_xargs: Option>, + /// Number of newly generated tokens to batch into each streamed output. + /// Raises the interval above the frontend-level `--stream-interval`; + /// values below it are clamped up to it. The first and final outputs are + /// always emitted immediately. + pub stream_interval: Option, } #[allow(clippy::derivable_impls)] // more explicit @@ -180,6 +186,7 @@ impl Default for SamplingParams { structured_outputs: None, skip_reading_prefix_cache: None, vllm_xargs: None, + stream_interval: None, } } } From 51144dd83ec557afa6fc6032bbf03d7f50a55b89 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sat, 10 Oct 2026 05:22:31 +0800 Subject: [PATCH 316/952] [Bugfix][Spec Decode] Trim token_ids/logprobs left past a stop string under speculative decoding (#47616) Signed-off-by: Ting Sun Signed-off-by: Nick Hill Co-authored-by: Nick Hill Co-authored-by: Claude Opus 5.5 --- ...stop_string_while_stop_model_terminates.py | 113 +++++++++++++++++- vllm/v1/engine/detokenizer.py | 24 +++- vllm/v1/engine/output_processor.py | 34 +++++- 3 files changed, 158 insertions(+), 13 deletions(-) diff --git a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py index 3b72d612dd1c..38d860969f49 100644 --- a/tests/detokenizer/test_stop_string_while_stop_model_terminates.py +++ b/tests/detokenizer/test_stop_string_while_stop_model_terminates.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest from vllm.sampling_params import SamplingParams -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import BaseIncrementalDetokenizer +from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.outputs import LogprobsLists, SamplingMaskLists @pytest.fixture(params=[True, False]) @@ -82,7 +85,7 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo # Simulate that the last token ('Z') is a stop token (stop_terminated=True). result = detok.update(new_token_ids=token_ids, stop_terminated=True) - # The update should not report a stop string + # The update should report the matched stop string. assert result == stop_string # Output text should reflect stop-string handling: @@ -91,9 +94,111 @@ def test_stop_string_while_stop_token_terminates(include_stop_str_in_output: boo expected_text = "abcd" if include_stop_str_in_output else "ab" assert detok.output_text == expected_text - # The skipped final token should still be recorded in token_ids. - assert detok.output_token_ids == token_ids + # Tokens after the stop string, including the stop token, are dropped. + assert detok.output_token_ids == token_ids[:4] # get_next_output_text should return the full text when finished=True. # (Buffering only applies during streaming when finished=False.) assert detok.get_next_output_text(finished=True, delta=False) == expected_text + + +@pytest.mark.parametrize( + "min_tokens,num_kept", + [(0, 2), (1, 2), (2, 4), (3, 4), (4, None)], +) +@pytest.mark.parametrize("num_prev_tokens", [0, 1]) +def test_stop_string_min_tokens_within_step( + include_stop_str_in_output: bool, min_tokens, num_kept, num_prev_tokens +): + """Stop strings completing within min_tokens are ignored, even when the + min_tokens boundary falls in the middle of a multi-token step.""" + token_ids = [ord(c) for c in "cdcdx"] + req = _make_request( + stop=["cd"], + include_stop_str_in_output=include_stop_str_in_output, + min_tokens=min_tokens, + ) + detok = _DummyDetokenizer(req) + + if num_prev_tokens: + assert detok.update(token_ids[:num_prev_tokens], False) is None + result = detok.update(token_ids[num_prev_tokens:], False) + + if num_kept is None: + assert result is None + assert detok.output_text == "cdcdx" + assert detok.output_token_ids == token_ids + return + assert result == "cd" + text_len = num_kept if include_stop_str_in_output else num_kept - 2 + assert detok.output_text == "cdcdx"[:text_len] + assert detok.output_token_ids == token_ids[:num_kept] + + +@pytest.mark.parametrize( + "decoded_tokens,keep", + [ + (["a", "b", "c", "d", "e", "f"], 4), + (["ab", "cdx", "ef"], 2), + (["ab", "", "c", "", "d", "", "ef"], 5), + ], +) +@pytest.mark.parametrize("stop_terminated", [False, True]) +def test_stop_string_trims_speculative_overflow( + include_stop_str_in_output: bool, + decoded_tokens, + keep, + stop_terminated: bool, + monkeypatch, +): + if stop_terminated: + # The engine also stopped on a stop token at the end of the batch. + decoded_tokens = decoded_tokens + [""] + token_ids = list(range(len(decoded_tokens))) + stop_string = "cd" + expected_token_ids = token_ids[:keep] + + req = _make_request( + stop=[stop_string], include_stop_str_in_output=include_stop_str_in_output + ) + req.external_req_id = req.request_id + assert req.sampling_params is not None + req.sampling_params.logprobs = 0 + detok = _DummyDetokenizer(req) + monkeypatch.setattr(detok, "decode_next", decoded_tokens.__getitem__) + processor = OutputProcessor(tokenizer=None, log_stats=False) + processor.add_request(req, prompt=None) + processor.request_states[req.request_id].detokenizer = detok + + outputs = processor.process_outputs( + [ + EngineCoreOutput( + request_id=req.request_id, + new_token_ids=token_ids, + finish_reason=FinishReason.STOP if stop_terminated else None, + new_logprobs=LogprobsLists( + np.array(token_ids).reshape(-1, 1), + np.zeros((len(token_ids), 1)), + np.ones(len(token_ids), dtype=int), + None, + ), + new_sampling_mask=SamplingMaskLists( + np.array(token_ids), np.arange(len(token_ids) + 1) + ), + routed_experts=np.array(token_ids).reshape(-1, 1, 1), + ) + ] + ) + result = outputs.request_outputs[0].outputs[0] + + assert result.stop_reason == stop_string + expected_text = "abcd" if include_stop_str_in_output else "ab" + assert detok.output_text == expected_text + assert detok.output_token_ids == expected_token_ids + assert result.token_ids == expected_token_ids + assert result.logprobs is not None + assert len(result.logprobs) == len(expected_token_ids) + assert result.sampling_mask is not None + assert result.sampling_mask.token_ids == [[token] for token in expected_token_ids] + assert result.routed_experts is not None + assert result.routed_experts.ravel().tolist() == expected_token_ids diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index c27e01895cac..adc0be31584b 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import sys from abc import ABC, abstractmethod +from bisect import bisect_left import tokenizers.decoders from tokenizers import Tokenizer @@ -112,21 +113,27 @@ def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: skipped_stop_token_id = None # 1) Detokenize the new token ids incrementally. - stop_check_offset = len(self.output_text) + num_prev_tokens = self.num_output_tokens() + first_new_token_index = len(self.token_ids) + # Output text offsets at the boundaries of the new tokens. + token_offsets = [len(self.output_text)] if self.stop else None for new_token_id in new_token_ids: self.token_ids.append(new_token_id) self.output_text += self.decode_next(new_token_id) - # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 - if self.min_tokens and self.num_output_tokens() <= self.min_tokens: - stop_check_offset = len(self.output_text) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) if skipped_stop_token_id is not None: # Cleanup after skipping detokenization. self.token_ids.append(skipped_stop_token_id) + if token_offsets is not None: + token_offsets.append(len(self.output_text)) # 2) Evaluate stop strings. stop_string = None - if self.stop and self.num_output_tokens() > self.min_tokens: + if token_offsets is not None and self.num_output_tokens() > self.min_tokens: + # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014 + stop_check_offset = token_offsets[max(self.min_tokens - num_prev_tokens, 0)] stop = check_stop_strings( output_text=self.output_text, new_char_count=len(self.output_text) - stop_check_offset, @@ -135,6 +142,13 @@ def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None: ) if stop is not None: stop_string, truncate_to = stop + stop_end = len(self.output_text) if truncate_to == -1 else truncate_to + if not self.include_stop_str_in_output: + stop_end += len(stop_string) + # Drop tokens generated after the one completing the stop string + # (possible when multiple tokens are appended in one step). + num_kept = bisect_left(token_offsets, stop_end, lo=1) + del self.token_ids[first_new_token_index + num_kept :] if truncate_to != -1: self.output_text = self.output_text[:truncate_to] diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index c22c5474c7c0..abb45f7239be 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -736,17 +736,22 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None - if engine_core_output.new_sampling_mask is not None: - req_state.sampling_mask_chunks.append( - engine_core_output.new_sampling_mask - ) # 2) Detokenize the token ids into text and perform stop checks. + num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( new_token_ids, finish_reason == FinishReason.STOP ) if stop_string: finish_reason = FinishReason.STOP stop_reason = stop_string + self._trim_surplus_tokens( + req_state, engine_core_output, num_prev_tokens + ) + + if engine_core_output.new_sampling_mask is not None: + req_state.sampling_mask_chunks.append( + engine_core_output.new_sampling_mask + ) # 3) Compute sample and prompt logprobs for request, # if required. @@ -799,6 +804,27 @@ def process_outputs( reqs_to_abort=reqs_to_abort, ) + @staticmethod + def _trim_surplus_tokens( + req_state: RequestState, output: EngineCoreOutput, num_prev_tokens: int + ) -> None: + """Drop new tokens generated past a stop string matched by the + detokenizer, e.g. when multiple tokens are accepted in one step.""" + assert req_state.detokenizer is not None + num_kept = req_state.detokenizer.num_output_tokens() - num_prev_tokens + num_dropped = len(output.new_token_ids) - num_kept + if num_dropped <= 0: + return + del output.new_token_ids[num_kept:] + if output.new_logprobs is not None: + output.new_logprobs = output.new_logprobs.slice_request(0, num_kept) + if output.new_sampling_mask is not None: + output.new_sampling_mask = output.new_sampling_mask.slice_request( + 0, num_kept + ) + if output.routed_experts is not None: + req_state.routed_experts_chunks[-1] = output.routed_experts[:-num_dropped] + def _finish_request(self, req_state: RequestState) -> None: req_id = req_state.request_id self.request_states.pop(req_id) From ed7ae705a1f366c2008964af562a7543e91d934e Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 14:38:11 -0700 Subject: [PATCH 317/952] [Bugfix][MRV2] Use int64 index mappings to fix row-offset overflow (#60736) Signed-off-by: Nick Hill Co-authored-by: vllmellm Co-authored-by: Claude Opus 5.5 (1M context) --- tests/v1/worker/test_gpu_row_offsets_i64.py | 203 ++++++++++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 2 +- vllm/v1/worker/gpu/input_batch.py | 6 +- vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/pcp_manager.py | 4 +- vllm/v1/worker/gpu/sample/batch_shard.py | 20 +- vllm/v1/worker/gpu/sample/penalties.py | 2 +- .../gpu/spec_decode/dflash/speculator.py | 2 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 2 +- .../spec_decode/standalone_ar/speculator.py | 2 +- 10 files changed, 221 insertions(+), 24 deletions(-) create mode 100644 tests/v1/worker/test_gpu_row_offsets_i64.py diff --git a/tests/v1/worker/test_gpu_row_offsets_i64.py b/tests/v1/worker/test_gpu_row_offsets_i64.py new file mode 100644 index 000000000000..374332e2dafa --- /dev/null +++ b/tests/v1/worker/test_gpu_row_offsets_i64.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""int64 row offsets in the MRV2 staged-write and idx-mapping kernels. + +With int32 request-slot indices, `slot * stride` wraps once the product +crosses 2**31 (issue #57030): at max_model_len=1M from row 2148, and from +row 716 for the 3-dim mrope positions. Each test runs a control row below +and an overflow row above the threshold in the same kernel launch. +Requires CUDA/ROCm and ~11 GiB free device memory: the tensors must exceed +2**31 elements by construction. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor +from vllm.v1.worker.gpu.input_batch import post_update +from vllm.v1.worker.gpu.mm.rope import RopeState +from vllm.v1.worker.gpu.sample.penalties import bincount + +STRIDE = 1_000_000 # max_model_len +ROW = 2559 # ROW * STRIDE > 2**31 (first overflowing row at 1M: 2148) +NUM_ROWS = ROW + 1 # (2560, 1M) int32 = 10.24 GB +NEEDED_BYTES = NUM_ROWS * STRIDE * 4 +VOCAB_SIZE = 8192 +PROMPT_LEN = 1024 +PREFILL_LEN = 2048 + +# mrope: flat row = NUM_DIMS * req + j and stride0 = NUM_DIMS * STRIDE, so +# the overflow threshold drops to row 716 of a 1M-row buffer. +NUM_DIMS = 3 +ROPE_ROW = 716 # NUM_DIMS * ROPE_ROW * STRIDE > 2**31 + +pytestmark = [ + pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires CUDA or ROCm" + ), + pytest.mark.skipif( + not current_platform.is_cuda_alike() + or torch.accelerator.get_memory_info(0)[0] + < NEEDED_BYTES + 2**30, # slack for the outputs + reason=f"requires >= {NEEDED_BYTES / 2**30:.0f} GiB free memory", + ), +] + + +def test_staged_write_high_row(): + """Apply-write past the int32 offset limit must land on the target row.""" + device = torch.device("cuda:0") + content = list(range(9000, 9016)) + staged = StagedWriteTensor((NUM_ROWS, STRIDE), torch.int32, device) + staged.stage_write(ROW, 0, content) + staged.apply_write() + torch.accelerator.synchronize() + + assert staged.gpu[ROW, : len(content)].tolist() == content + # A wrapped offset must not clobber row 0 (silent-corruption variant). + assert staged.gpu[0, : len(content)].tolist() == [0] * len(content) + + +def test_bincount_high_row(): + """penalties.bincount past the int32 offset limit must hit the right row.""" + device = torch.device("cuda:0") + tokens = torch.arange(PREFILL_LEN, dtype=torch.int32, device=device) + + all_token_ids = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + all_token_ids[0, :PREFILL_LEN] = tokens + all_token_ids[ROW, :PREFILL_LEN] = tokens + + prompt_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + prefill_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + prompt_len[[0, ROW]] = PROMPT_LEN + prefill_len[[0, ROW]] = PREFILL_LEN + + idx_mapping = torch.tensor([0, ROW], dtype=torch.int64, device=device) + prompt_bin_mask = torch.zeros( + NUM_ROWS, cdiv(VOCAB_SIZE, 32), dtype=torch.int32, device=device + ) + output_bin_counts = torch.zeros( + NUM_ROWS, VOCAB_SIZE, dtype=torch.int32, device=device + ) + + bincount( + idx_mapping, + all_token_ids, + prompt_len, + prefill_len, + prompt_bin_mask, + output_bin_counts, + max_prefill_len=PREFILL_LEN, + ) + torch.accelerator.synchronize() + + assert torch.equal(prompt_bin_mask[ROW], prompt_bin_mask[0]) + assert torch.equal(output_bin_counts[ROW], output_bin_counts[0]) + # Sanity: the control row got prompt bits and output counts. + assert prompt_bin_mask[0, : cdiv(PROMPT_LEN, 32)].any().item() + assert ( + output_bin_counts[0, PROMPT_LEN:PREFILL_LEN].sum().item() + == PREFILL_LEN - PROMPT_LEN + ) + + +def test_post_update_high_row(): + """post_update (runs every sampled token) must store on the mapped row.""" + device = torch.device("cuda:0") + # Big buffer first so a wrapped (negative) offset cannot land in-bounds. + all_token_ids = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + num_computed = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + last_sampled = torch.full((NUM_ROWS, 1), -1, dtype=torch.int64, device=device) + total_len = torch.zeros(NUM_ROWS, dtype=torch.int32, device=device) + total_len[[0, ROW]] = torch.tensor([50, 100], dtype=torch.int32, device=device) + + idx_mapping = torch.tensor([0, ROW], dtype=torch.int64, device=device) + sampled_tokens = torch.tensor([[111], [222]], dtype=torch.int64, device=device) + num_sampled = torch.tensor([1, 1], dtype=torch.int32, device=device) + num_rejected = torch.tensor([0, 0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) + + post_update( + idx_mapping, + num_computed, + last_sampled, + None, # output_bin_counts: a second product of the same idx load + sampled_tokens, + num_sampled, + num_rejected, + query_start_loc, + all_token_ids, + total_len, + ) + torch.accelerator.synchronize() + + assert int(all_token_ids[0, 50]) == 111 + assert int(all_token_ids[ROW, 100]) == 222 + # Each mapped row holds exactly its own token; unmapped rows stay zero. + assert int(all_token_ids[0].sum()) == 111 + assert int(all_token_ids[ROW].sum()) == 222 + assert int(all_token_ids[1].sum()) == 0 + assert int(last_sampled[0, 0]) == 111 + assert int(last_sampled[ROW, 0]) == 222 + assert int(total_len[0]) == 51 + assert int(total_len[ROW]) == 101 + assert int(num_computed[0]) == 1 + assert int(num_computed[ROW]) == 1 + + +def test_rope_positions_high_row(): + """M-RoPE prepare_positions must read staged rows past the int32 limit. + + Threshold is 3x lower than the other tests (stride0 = num_dims * 1M), so + this is the variant that trips at default max_num_seqs on 1M context. + """ + device = torch.device("cuda:0") + max_num_reqs = NUM_ROWS // NUM_DIMS + base = torch.arange(PREFILL_LEN, dtype=torch.int32, device=device) + # (num_dims * max_num_reqs, max_model_len), logically [req][dim][pos]. + prefill_positions = torch.zeros(NUM_ROWS, STRIDE, dtype=torch.int32, device=device) + # Distinct per-(dim, request) patterns so mixed-up rows are detectable. + for j in range(NUM_DIMS): + prefill_positions[j, :PREFILL_LEN] = 1_000 * (j + 1) + base + target_row = NUM_DIMS * ROPE_ROW + j + prefill_positions[target_row, :PREFILL_LEN] = 7_000 * (j + 1) + base + + # Skeletal RopeState: the real __init__ pins num_dims * max_num_reqs host + # memory, which the stride arithmetic does not depend on. + rope = object.__new__(RopeState) + rope.num_dims = NUM_DIMS + rope.max_num_reqs = max_num_reqs + rope.max_model_len = STRIDE + rope.prefill_positions = SimpleNamespace(gpu=prefill_positions) + rope.positions = torch.zeros( + (NUM_DIMS, 2 * PREFILL_LEN + 1), dtype=torch.int64, device=device + ) + rope.prefill_delta = SimpleNamespace( + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + ) + + idx_mapping = torch.tensor([0, ROPE_ROW], dtype=torch.int64, device=device) + query_start_loc = torch.tensor( + [0, PREFILL_LEN, 2 * PREFILL_LEN], dtype=torch.int32, device=device + ) + prefill_len = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + prefill_len[[0, ROPE_ROW]] = PREFILL_LEN + num_computed = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + + rope.prepare_positions(idx_mapping, query_start_loc, prefill_len, num_computed) + torch.accelerator.synchronize() + + expect_control = torch.stack([1_000 * (j + 1) + base for j in range(NUM_DIMS)]).to( + torch.int64 + ) + expect_target = torch.stack([7_000 * (j + 1) + base for j in range(NUM_DIMS)]).to( + torch.int64 + ) + assert torch.equal(rope.positions[:, :PREFILL_LEN], expect_control) + assert torch.equal(rope.positions[:, PREFILL_LEN : 2 * PREFILL_LEN], expect_target) + # Padding column stays zero. + assert bool((rope.positions[:, 2 * PREFILL_LEN :] == 0).all()) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 32d997358c13..7364156848bf 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -310,7 +310,7 @@ def _apply_write_kernel( MULTI_GROUP: tl.constexpr, ): pid = tl.program_id(0) - row_idx = tl.load(write_indices_ptr + pid) + row_idx = tl.load(write_indices_ptr + pid).to(tl.int64) start_idx = tl.load(write_starts_ptr + pid) cu_start = tl.load(write_cu_lens_ptr + pid - 1) if pid > 0 else 0 diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index ea1c925dadc0..bd55a43f6412 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -137,7 +137,7 @@ def make_dummy( req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)] idx_mapping_np = np.arange(num_reqs, dtype=np.intp) - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + idx_mapping = torch.arange(num_reqs, dtype=torch.int64, device=device) expanded_idx_mapping = idx_mapping expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) @@ -715,7 +715,9 @@ def expand_idx_mapping( ) -> tuple[torch.Tensor, torch.Tensor]: num_reqs = idx_mapping.shape[0] expanded_idx_mapping = idx_mapping.new_empty(total_num_logits) - expanded_local_pos = idx_mapping.new_empty(total_num_logits) + expanded_local_pos = torch.empty( + total_num_logits, dtype=torch.int32, device=idx_mapping.device + ) _expand_idx_mapping_kernel[(num_reqs,)]( idx_mapping, expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 134b4bfb0e69..2e9925f1ad1c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1361,7 +1361,7 @@ def prepare_inputs( num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens idx_mapping_np = batch_req_state.idx_mapping_np idx_mapping = async_tensor_h2d( - idx_mapping_np, device=self.device, dtype=torch.int32 + idx_mapping_np, device=self.device, dtype=torch.int64 ) num_reqs = len(req_ids) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 6aec96bfe14b..aa245fa267c4 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -653,9 +653,7 @@ def partition_batch( idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, - expanded_local_pos=torch.zeros( - num_local_reqs, dtype=torch.int32, device=self.device - ), + expanded_local_pos=input_batch.expanded_local_pos.new_zeros(num_local_reqs), num_scheduled_tokens=local_num_scheduled_tokens, num_tokens=num_local_tokens, num_tokens_after_padding=num_local_tokens_padded, diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index a7e369b1886a..695b1d543563 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -214,21 +214,15 @@ def shard_sampler_inputs( local_logits_indices = torch.empty( num_local_logits, dtype=torch.int64, device=self.device ) - local_idx_mapping = torch.empty( - num_local_reqs, dtype=torch.int32, device=self.device + local_idx_mapping = input_batch.idx_mapping.new_empty(num_local_reqs) + local_cu_num_logits = input_batch.cu_num_logits.new_empty(num_local_reqs + 1) + local_expanded_idx_mapping = input_batch.expanded_idx_mapping.new_empty( + num_local_logits ) - local_cu_num_logits = torch.empty( - num_local_reqs + 1, dtype=torch.int32, device=self.device - ) - local_expanded_idx_mapping = torch.empty( - num_local_logits, dtype=torch.int32, device=self.device - ) - local_expanded_local_pos = torch.empty( - num_local_logits, dtype=torch.int32, device=self.device - ) - local_seq_lens = torch.empty( - num_local_reqs, dtype=torch.int32, device=self.device + local_expanded_local_pos = input_batch.expanded_local_pos.new_empty( + num_local_logits ) + local_seq_lens = input_batch.seq_lens.new_empty(num_local_reqs) if num_reqs > 0: _build_shard_plan_kernel[(num_reqs,)]( input_batch.idx_mapping, diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index eb751e8be355..5df7c60afd56 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -62,7 +62,7 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: def apply_staged_writes(self) -> None: if self._new_penalties_reqs: idx_mapping = async_tensor_h2d( - self._new_penalties_reqs, dtype=torch.int32, device=self.device + self._new_penalties_reqs, dtype=torch.int64, device=self.device ) prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index ec46039533d7..9538dcb4d4bf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -102,7 +102,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # full buffer before a real batch has populated it, so zero would make # every padding row scatter into request slot 0. self.sample_idx_mapping = torch.full( - (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device + (max_num_sampled_tokens,), -1, dtype=torch.int64, device=device ) # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into # draft_logits[req, step, :]. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d07c043de915..cbd94ed0542d 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -152,7 +152,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): device=device, ) self.idx_mapping = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=device + self.max_num_reqs, dtype=torch.int64, device=device ) self.temperature = torch.zeros( self.max_num_reqs, dtype=torch.float32, device=device diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 8bdfb8530be6..898de1e3f886 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -402,7 +402,7 @@ def _prepare_prefill_inputs_kernel( target_input_ids_ptr, # [src_tokens] int32 target_positions_ptr, # [src_tokens] int64 last_sampled_ptr, # [max_num_reqs] int32 - idx_mapping_ptr, # [num_reqs] int32 + idx_mapping_ptr, # [num_reqs] int64 out_input_ids_ptr, # [src_tokens + num_reqs] int32 out_positions_ptr, # [src_tokens + num_reqs] int64 last_token_indices_ptr, # [max_num_reqs] int64 From 783ac12071ec319d7d3e5af29d039034115d8934 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 9 Oct 2026 15:43:37 -0700 Subject: [PATCH 318/952] [CI] Accept whitespace around /ci commands and reply instead of going silent (#60925) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .github/workflows/run-ci-command.yml | 29 +++++++------- .github/workflows/scripts/run_ci_command.py | 34 ++++++++++++++++- .../workflows/scripts/test_run_ci_command.py | 38 ++++++++++++++++++- 3 files changed, 82 insertions(+), 19 deletions(-) diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 3435148ea4e3..b0764bfa0683 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -18,22 +18,8 @@ jobs: if: >- github.repository_owner == 'vllm-project' && github.event.issue.pull_request && - (github.event.comment.body == '/ci run' || - github.event.comment.body == '/ci run all' || - github.event.comment.body == '/ci run nightly' || - github.event.comment.body == '/ci run --allow-stale' || - github.event.comment.body == '/ci run all --allow-stale' || - github.event.comment.body == '/ci run nightly --allow-stale' || - github.event.comment.body == '/ci retry' || - github.event.comment.body == '/ci cancel' || - github.event.comment.body == '/amd-ci run' || - github.event.comment.body == '/amd-ci run all' || - github.event.comment.body == '/amd-ci run nightly' || - github.event.comment.body == '/amd-ci run --allow-stale' || - github.event.comment.body == '/amd-ci run all --allow-stale' || - github.event.comment.body == '/amd-ci run nightly --allow-stale' || - github.event.comment.body == '/amd-ci retry' || - github.event.comment.body == '/amd-ci cancel') + (startsWith(github.event.comment.body, '/ci') || + startsWith(github.event.comment.body, '/amd-ci')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -54,3 +40,14 @@ jobs: BUILDKITE_PIPELINE: ci CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} GH_TOKEN: ${{ github.token }} + - name: Report failure on the PR + if: failure() + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + with: + script: | + const runUrl = `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`; + await github.rest.issues.createComment({ + ...context.repo, + issue_number: context.issue.number, + body: `❌ @${context.payload.comment.user.login}, your CI command failed with an unexpected error, so CI may not have started. See the [workflow run](${runUrl}) for details.`, + }); diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 530bd697d008..537efa83c4a5 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -58,6 +58,7 @@ RETRY_COMMANDS = frozenset({COMMAND_RETRY_FAILED, COMMAND_RETRY_AMD_FAILED}) CANCEL_COMMANDS = frozenset({COMMAND_CANCEL_CI, COMMAND_CANCEL_AMD_CI}) ALL_CI_COMMANDS = UPSTREAM_CI_COMMANDS | AMD_CI_COMMANDS +CI_COMMAND_PREFIXES = frozenset({"/ci", "/amd-ci"}) CI_AUTHORIZED_COMMENT_MARKER = "" READY_LABELS = {"ready", "ready-run-all-tests"} TRUSTED_PERMISSIONS = {"admin", "maintain", "write"} @@ -461,11 +462,17 @@ def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: def parse_command(body: str) -> str | None: + body = body.strip() if body in ALL_CI_COMMANDS: return body return None +def is_ci_command_attempt(body: str) -> bool: + words = body.split(maxsplit=1) + return bool(words) and words[0].casefold() in CI_COMMAND_PREFIXES + + def pipeline_for_command( command: str, *, @@ -1086,6 +1093,28 @@ def handle_cancel_ci( ) +def reply_unrecognized_command( + event: Mapping[str, Any], + github: GitHubClient, +) -> None: + if "pull_request" not in event["issue"]: + return + issue_number = event["issue"]["number"] + comment_id = event["comment"]["id"] + actor = event["comment"]["user"]["login"] + if is_already_handled(github, issue_number, comment_id): + print(f"Comment {comment_id} was already handled.") + return + add_reaction_safely(github, comment_id, "-1") + commands = ", ".join(f"`{command}`" for command in sorted(ALL_CI_COMMANDS)) + github.add_comment( + issue_number, + f"❌ @{actor}, that is not a recognized CI command, so no CI was " + "started. Post a new comment containing only one of: " + f"{commands}.\n\n{command_comment_marker(comment_id)}", + ) + + def run( event: Mapping[str, Any], github: GitHubClient, @@ -1217,8 +1246,11 @@ def main() -> None: ) return - command = parse_command(event["comment"]["body"]) + body = event["comment"]["body"] + command = parse_command(body) if not command: + if is_ci_command_attempt(body): + reply_unrecognized_command(event, github) return pipeline = pipeline_for_command( diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index 9750cefaa394..623b41247641 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -32,10 +32,12 @@ has_trusted_approval, is_active_build, is_build_for_pr, + is_ci_command_attempt, notify_authorized, parse_command, parse_trusted_users, pipeline_for_command, + reply_unrecognized_command, resolve_workflow_run_pr, run, select_latest_build, @@ -340,14 +342,46 @@ def test_only_exact_ci_commands_are_accepted(self) -> None: self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) self.assertIsNone(parse_command("/ci cancel please")) - self.assertIsNone(parse_command(" /ci run")) self.assertIsNone(parse_command("/amd-ci run please")) - self.assertIsNone(parse_command("/amd-ci retry ")) self.assertIsNone(parse_command("/AMD-CI run")) self.assertIsNone(parse_command("/amdci run")) self.assertIsNone(parse_command("/ci run --allow-stale all")) self.assertIsNone(parse_command("/amd-ci run --allow-stale please")) + def test_surrounding_whitespace_is_ignored(self) -> None: + self.assertEqual(parse_command("/ci run\r\n"), COMMAND_RUN_CI) + self.assertEqual(parse_command(" /ci run "), COMMAND_RUN_CI) + self.assertEqual(parse_command("/amd-ci retry\n"), COMMAND_RETRY_AMD_FAILED) + + def test_ci_command_attempts_are_detected(self) -> None: + for body in ("/ci", "/ci rn", "/CI run", "/amd-ci run please", " /ci\nrun"): + with self.subTest(body=body): + self.assertTrue(is_ci_command_attempt(body)) + for body in ("", "LGTM", "/cifoo", "/amdci run", "please /ci run"): + with self.subTest(body=body): + self.assertFalse(is_ci_command_attempt(body)) + + def test_unrecognized_command_gets_feedback(self) -> None: + github = FakeGitHub() + + reply_unrecognized_command(make_event("/ci rn", actor="author"), github) + + self.assertEqual(github.reactions, ["-1"]) + self.assertEqual(len(github.comments), 1) + self.assertIn( + "@author, that is not a recognized CI command", github.comments[0] + ) + self.assertIn(f"`{COMMAND_RUN_CI}`", github.comments[0]) + self.assertIn("", github.comments[0]) + + def test_unrecognized_command_feedback_is_not_repeated(self) -> None: + github = FakeGitHub(comments=["❌ earlier\n\n"]) + + reply_unrecognized_command(make_event("/ci rn"), github) + + self.assertEqual(len(github.comments), 1) + self.assertEqual(github.reactions, []) + def test_commands_select_only_their_configured_pipeline(self) -> None: cases = ( (COMMAND_RUN_CI, "ci"), From 4c43254d52785d7a69300fdfb6c79987200c277c Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 9 Oct 2026 18:45:34 -0400 Subject: [PATCH 319/952] [HiSparse] Report host-pool block residency as HiSparse metrics (#59485) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- docs/usage/metrics.md | 5 +- .../unit/test_hisparse_connector.py | 44 ++++++++++++- .../kv_connector/unit/test_hisparse_stats.py | 31 +++++++++- tests/v1/worker/test_utils.py | 3 + .../kv_connector/v1/hisparse/connector.py | 2 + .../kv_connector/v1/hisparse/stats.py | 61 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 7 ++- 7 files changed, 147 insertions(+), 6 deletions(-) diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index b87ff1c28fee..ef75ce11c1fc 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -50,7 +50,10 @@ The following metrics are exposed: These metrics are exposed when the `HiSparseConnector` KV connector is configured. The host-pool gauges are reported by the scheduler once per engine step, for the host pool shared by all tensor-parallel ranks; evictable cached -blocks count as free. +blocks count as free. With `--kv-cache-metrics`, the +`vllm:hisparse_host_block_*_seconds` histograms sample host-pool block +residency, like the device `vllm:kv_block_*_seconds` histograms, which cover +only the GPU pool. --8<-- "gen:metrics-hisparse" diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index f08982a5e924..f9598a24f82e 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -3,7 +3,7 @@ import contextlib import ctypes from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -25,6 +25,7 @@ HiSparseConnectorWorker, ) from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.hisparse import runtime as runtime_module from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.hisparse.runtime import HiSparseCacheHandle @@ -33,6 +34,7 @@ SparseKVPageTransfer, SparseKVRowMirror, ) +from vllm.v1.metrics.stats import KVCacheEvictionEvent from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector @@ -100,6 +102,46 @@ def test_scheduler_stats_report_host_pool_usage(): assert stats.data["host_cache_usage_perc"] == [expected] +def test_hisparse_host_residency_reported_as_hisparse_metrics(): + """Host blocks sample into HiSparse's own residency metrics.""" + from tests.v1.core.test_prefix_caching import ( + make_hisparse_kv_cache_config, + make_kv_cache_manager, + ) + from tests.v1.core.utils import create_requests + + collector = KVCacheMetricsCollector(sample_rate=1.0) + manager = make_kv_cache_manager( + make_hisparse_kv_cache_config(2, 2), + max_model_len=128, + enable_caching=True, + hash_block_size=16, + metrics_collector=collector, + ) + coordinator = get_hisparse_coordinator(manager) + assert coordinator.host_manager is not None + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(coordinator) + device = manager.block_pool + host = coordinator.host_manager.block_pool + request = create_requests(1, num_tokens=16)[0] + for pool, birth in ((device, 1), (host, 2)): + with patch("time.monotonic_ns", return_value=birth * 10**9): + block = pool.get_new_blocks(1)[0] + assert block.block_id == 1 + pool.cache_full_blocks(request, [block], 0, 1, 16, 0) + with patch("time.monotonic_ns", return_value=4_000_000_000): + host.free_blocks([host.blocks[1]]) + host.evict_blocks({1}) + with patch("time.monotonic_ns", return_value=6_000_000_000): + device.evict_blocks({1}) + + assert collector.drain_events() == [KVCacheEvictionEvent(5.0, 5.0, ())] + stats = scheduler.get_kv_connector_stats() + assert stats.data["host_block_lifetime_seconds"] == [2.0] + assert stats.data["host_block_idle_before_evict_seconds"] == [2.0] + + def test_no_forward_enqueues_deferred_hisparse_transfers(): """A zero-token step must still enqueue deferred post-forward transfers.""" connector = object.__new__(ActiveKVConnector) diff --git a/tests/v1/kv_connector/unit/test_hisparse_stats.py b/tests/v1/kv_connector/unit/test_hisparse_stats.py index 91d00b1e2260..f0b09ecf365b 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_stats.py +++ b/tests/v1/kv_connector/unit/test_hisparse_stats.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( HiSparseConnector, ) @@ -82,7 +84,8 @@ def test_build_kv_connector_stats_round_trip(): assert rebuilt.reduce() == stats.reduce() -def test_prom_metrics_observe_host_usage_gauges(): +@pytest.mark.parametrize("kv_cache_metrics", [True, False]) +def test_prom_metrics_observe_host_metrics(kv_cache_metrics: bool): from types import SimpleNamespace from typing import Any @@ -91,12 +94,14 @@ def test_prom_metrics_observe_host_usage_gauges(): from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( HiSparsePromMetrics, ) + from vllm.v1.metrics.stats import KVCacheEvictionEvent class _FakeMetric: def __init__(self, **kwargs: Any): self.kwargs = kwargs self.increments: list[int | float] = [] self.set_values: list[int | float] = [] + self.observed: list[float] = [] def labels(self, *labelvalues: object) -> "_FakeMetric": return self @@ -107,6 +112,9 @@ def inc(self, value: int | float) -> None: def set(self, value: int | float) -> None: self.set_values.append(value) + def observe(self, value: float) -> None: + self.observed.append(value) + created: dict[str, _FakeMetric] = {} class _NamedFake(_FakeMetric): @@ -115,7 +123,12 @@ def __init__(self, **kwargs: Any): created[kwargs["name"]] = self prom = HiSparsePromMetrics( - vllm_config=SimpleNamespace(kv_transfer_config=None), + vllm_config=SimpleNamespace( + kv_transfer_config=None, + observability_config=SimpleNamespace( + kv_cache_metrics=kv_cache_metrics, custom_histogram_buckets=None + ), + ), metric_types={Gauge: _NamedFake, Counter: _NamedFake, Histogram: _NamedFake}, labelnames=["model_name"], per_engine_labelvalues={0: ["model"]}, @@ -125,9 +138,21 @@ def __init__(self, **kwargs: Any): stats.record_snapshot(hits=3, misses=2, host_to_device_bytes=32) stats.record_host_usage(usage=0.5, pending_page_transfers=1) stats.record_host_usage(usage=0.75, pending_page_transfers=3) - prom.observe(stats.to_dict()) + stats.record_host_evictions([KVCacheEvictionEvent(5.0, 2.0, (1.0, 2.0))]) + later = HiSparseKVConnectorStats() + later.record_host_evictions([KVCacheEvictionEvent(3.0, 3.0, ())]) + prom.observe(stats.aggregate(later).to_dict()) assert created["vllm:hisparse_cache_hits"].increments == [3] assert created["vllm:hisparse_host_to_device_bytes"].increments == [32] assert created["vllm:hisparse_host_cache_usage_perc"].set_values == [0.75] assert created["vllm:hisparse_pending_page_transfers"].set_values == [3] + if not kv_cache_metrics: + assert not any("host_block" in name for name in created) + return + assert created["vllm:hisparse_host_block_lifetime_seconds"].observed == [5.0, 3.0] + assert created["vllm:hisparse_host_block_idle_before_evict_seconds"].observed == [ + 2.0, + 3.0, + ] + assert created["vllm:hisparse_host_block_reuse_gap_seconds"].observed == [1.0, 2.0] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index e17f214871f3..98395e90c4dd 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -94,6 +94,9 @@ def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeyp "host_to_device_bytes": [64], "host_cache_usage_perc": [], "pending_page_transfers": [], + "host_block_lifetime_seconds": [], + "host_block_idle_before_evict_seconds": [], + "host_block_reuse_gap_seconds": [], } diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 7e8960f644a9..863f52e02514 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -112,6 +112,8 @@ def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None: stats.record_host_usage( host_pool.get_usage(), len(self.coordinator.pending_spills) ) + if host_pool.metrics_collector is not None: + stats.record_host_evictions(host_pool.metrics_collector.drain_events()) return stats def build_connector_meta( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py index 6f6cfced9ede..9c7e8750976a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/stats.py @@ -12,6 +12,8 @@ PromMetric, PromMetricT, ) +from vllm.v1.metrics.buckets import histogram_buckets +from vllm.v1.metrics.stats import KVCacheEvictionEvent from vllm.v1.metrics.utils import create_metric_per_engine _HISPARSE_LEVEL_KEYS = frozenset({"host_cache_usage_perc", "pending_page_transfers"}) @@ -40,6 +42,9 @@ def reset(self): "host_to_device_bytes": [], "host_cache_usage_perc": [], "pending_page_transfers": [], + "host_block_lifetime_seconds": [], + "host_block_idle_before_evict_seconds": [], + "host_block_reuse_gap_seconds": [], } def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): @@ -51,6 +56,12 @@ def record_host_usage(self, usage: float, pending_page_transfers: int) -> None: self.data["host_cache_usage_perc"].append(usage) self.data["pending_page_transfers"].append(pending_page_transfers) + def record_host_evictions(self, events: list[KVCacheEvictionEvent]) -> None: + for event in events: + self.data["host_block_lifetime_seconds"].append(event.lifetime_seconds) + self.data["host_block_idle_before_evict_seconds"].append(event.idle_seconds) + self.data["host_block_reuse_gap_seconds"].extend(event.reuse_gaps_seconds) + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: if not other.is_empty(): for k, v in other.data.items(): @@ -143,6 +154,53 @@ def __init__( ), } + self.hisparse_histograms: dict[str, Any] = {} + if vllm_config.observability_config.kv_cache_metrics: + residency_buckets = histogram_buckets( + "kv_cache_residency", + overrides=vllm_config.observability_config.custom_histogram_buckets, + ) + histogram_host_block_lifetime = self._histogram_cls( + name="vllm:hisparse_host_block_lifetime_seconds", + documentation=( + "Histogram of HiSparse host KV block lifetime from allocation " + "to eviction. Sampled metrics (controlled by " + "--kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + histogram_host_block_idle_before_evict = self._histogram_cls( + name="vllm:hisparse_host_block_idle_before_evict_seconds", + documentation=( + "Histogram of HiSparse host KV block idle time before eviction. " + "Sampled metrics (controlled by --kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + histogram_host_block_reuse_gap = self._histogram_cls( + name="vllm:hisparse_host_block_reuse_gap_seconds", + documentation=( + "Histogram of time gaps between consecutive HiSparse host KV " + "block accesses. Sampled metrics (controlled by " + "--kv-cache-metrics-sample)." + ), + buckets=residency_buckets, + labelnames=labelnames, + ) + self.hisparse_histograms = { + "host_block_lifetime_seconds": create_metric_per_engine( + histogram_host_block_lifetime, self.per_engine_labelvalues + ), + "host_block_idle_before_evict_seconds": create_metric_per_engine( + histogram_host_block_idle_before_evict, self.per_engine_labelvalues + ), + "host_block_reuse_gap_seconds": create_metric_per_engine( + histogram_host_block_reuse_gap, self.per_engine_labelvalues + ), + } + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): for name, counter in self.hisparse_counters.items(): for value in transfer_stats_data.get(name, []): @@ -151,3 +209,6 @@ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): values = transfer_stats_data.get(name, []) if values: gauge[engine_idx].set(values[-1]) + for name, histogram in self.hisparse_histograms.items(): + for value in transfer_stats_data.get(name, []): + histogram[engine_idx].observe(value) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e9f5f2f3c231..eed15427b2dd 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -9,6 +9,7 @@ from vllm.distributed.kv_events import MEDIUM_CPU from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHashList, BlockHashListWithBlockSize, @@ -2272,7 +2273,11 @@ def bind_host_pool(self, num_blocks: int) -> None: enable_caching=self.enable_caching and self.kv_cache_spec.prefix_cacheable, hash_block_size=device_pool.hash_block_size, enable_kv_cache_events=device_pool.enable_kv_cache_events, - metrics_collector=None, + metrics_collector=( + KVCacheMetricsCollector(device_pool.metrics_collector.sample_rate) + if device_pool.metrics_collector is not None + else None + ), medium=MEDIUM_CPU, event_owner=device_pool, ) From fa94ad41431d7942abcd1adfffd3262edaed975a Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 9 Oct 2026 17:54:06 -0500 Subject: [PATCH 320/952] [Test] Allow one bf16 ulp in the fused mHC post/pre residual check (#60724) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/kernels/mhc/test_mhc_kernels.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 00e9834d6215..4afb45f503ef 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -439,13 +439,16 @@ def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carrie return # The post mapping and the collapse are what the unfused kernels produce. + # The two kernels may round a near-midpoint residual differently (FMA + # contraction is up to the compiler), so allow one bf16 ulp there and + # check the collapse against the fused kernel's own residual. residual_ref = torch.ops.vllm.mhc_post_tilelang( x, residual, post_layer_mix, comb_res_mix ) layer_input_ref = mhc_pre_delayed_tilelang( - residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 + residual_cur, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 )[2] - torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0) + torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=2**-7) torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0) post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref( From 04b02fc8a56fc0140d8559b3a07d84b2fdc46ae8 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Fri, 9 Oct 2026 19:13:24 -0400 Subject: [PATCH 321/952] [Core] shm tensor arena for efficient cpu->gpu worker broadcast (#51207) Signed-off-by: Bolin Sun Signed-off-by: Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- docs/design/shm_tensor_arena.md | 302 ++++++++++ tests/distributed/test_shm_broadcast.py | 522 +++++++++++++++++- tools/pre_commit/check_forbidden_imports.py | 1 + vllm/config/parallel.py | 10 + .../device_communicators/shm_broadcast.py | 59 +- .../device_communicators/shm_tensor_arena.py | 397 +++++++++++++ vllm/engine/arg_utils.py | 6 + vllm/v1/executor/multiproc_executor.py | 1 + 8 files changed, 1290 insertions(+), 8 deletions(-) create mode 100644 docs/design/shm_tensor_arena.md create mode 100644 vllm/distributed/device_communicators/shm_tensor_arena.py diff --git a/docs/design/shm_tensor_arena.md b/docs/design/shm_tensor_arena.md new file mode 100644 index 000000000000..c37ac89e363a --- /dev/null +++ b/docs/design/shm_tensor_arena.md @@ -0,0 +1,302 @@ +# Zero-Copy Shared-Memory Tensor Arena for `MessageQueue` + +This document describes an optimization to vLLM's engine→worker IPC path for +multimodal serving under tensor parallelism, implemented in +`vllm/distributed/device_communicators/shm_tensor_arena.py` and wired into +`shm_broadcast.py`'s `MessageQueue`. It **builds on the existing out-of-band +tensor channel** (`_reduce_tensor` / `_rebuild_tensor`, still in +`shm_broadcast.py`), which already keeps CPU-tensor bytes out of the pickle +stream. + +> **Status: experimental, opt-in.** This lands as an increment over the +> existing out-of-band channel, not a replacement for it. vLLM also has an +> existing multimodal-tensor shm caching mechanism; reconciling the two is +> intentionally left as follow-up work rather than blocking this PR. + +## TL;DR + +`MessageQueue.enqueue` already avoids copying CPU tensor bytes into the pickle +stream: `_reduce_tensor` emits each tensor as a protocol-5 out-of-band +`PickleBuffer`, and large payloads are published to the local readers over a +ZMQ socket. That removed the dominant cost of the older *in-band* path — the +per-tensor `pickle` serialize (see [History](#history-the-in-band-pickle-era)). + +Two costs remain for a large multimodal `pixel_values` tensor on a TP=N worker: + +1. **N transport copies.** On the overflow (large-payload) path the out-of-band + buffer is published to each of the N node-local readers via `send_multipart`, + so the tensor bytes are moved once **per reader** by the transport — ~N×200 MB + of host data motion for a payload the model needs exactly once. +2. **A pageable H2D.** Each reader zero-copies a view over its received ZMQ + frame, but that frame is ordinary pageable memory, so the reader's subsequent + `cudaMemcpyAsync` to the GPU pays pageable-copy staging plus first-touch + faults. + +This change adds a slotted shared-memory **arena**: the writer does **one** +memcpy of the tensor into a free slot; every reader takes a **zero-copy view** +of the same slot (no per-reader transport copy); and the mapping is +`cudaHostRegister`-pinned so the H2D is a true DMA. Net for a 200 MB image on +TP=4: **1 copy + a pinned DMA** instead of **4 transport copies + a pageable +H2D**. Small tensors, non-contiguous tensors, and arena exhaustion all fall +through to the existing `_reduce_tensor` path unchanged. + +## Scope of the win (and what still needs measuring) + +The arena is an **incremental** improvement layered on the out-of-band channel, +not a from-scratch fix — the out-of-band channel already eliminated the single +largest cost (the serialize copy). The benefit that remains is: + +- **N→1 copies**, which grows with TP degree (no gain at TP=1), and +- the **pinned vs. pageable H2D**, which shows up as reader-side H2D latency. + +> **The benchmarks in [§4](#4-validation) were measured against the *in-band +> pickle* baseline that predates `_reduce_tensor`**, where the dominant cost was +> the serialize copy. They are retained as motivation and as an upper bound. The +> incremental improvement of the arena *over the current out-of-band path on +> `main`* is smaller and has **not yet been re-measured**; that A/B (arena vs. +> oob, same seeded workload) is the outstanding validation item before this is +> merge-ready. + +## 1. Where it bites + +Any multimodal model served with `--tensor-parallel-size > 1` through the +multiproc executor. The `rpc_broadcast_mq` `MessageQueue` +(`v1/executor/multiproc_executor.py`) carries scheduler output — including the +raw preprocessed `pixel_values` — from the EngineCore to every TP worker each +step. Images at native resolution reach tens to hundreds of MB (a ~4k×4k image +is ~200 MB as bf16 patches). The out-of-band channel keeps those bytes out of +the pickle stream, but on the overflow path they are still published **once per +node-local reader**, and each reader's H2D is from **pageable** memory. The +per-reader copy cost scales with TP degree; the pinning cost is a fixed per-image +tax on the reader H2D. Latency-sensitive (interactive) serving makes the residual +visible in TTFT and, via head-of-line blocking, in the end-to-end latency of +*other* requests. + +### History: the in-band-pickle era + +Before `_reduce_tensor`, `MessageQueue.enqueue` pickled the whole payload — +including tensor bytes — because `torch.Tensor.__reduce_ex__` does not emit +protocol-5 out-of-band buffers. A ~200 MB tensor was byte-copied into the pickle +stream on the engine (`THPStorage_writeFileRaw`, ~6.5 ms/MB → ~1.3 s), *inside* +the EngineCore step loop, so all GPUs on the worker went idle for **0.9–1.3 s per +large image** (the vision-encoder compute for the same image is 50–90 ms → +transport dominated compute ~15:1); each reader then paid a full `pickle.loads` +byte-copy. Per 200 MB image on TP=4 that was ~1.8 GB of host data motion. + +The `_reduce_tensor` out-of-band channel removed the serialize copy (and the +per-reader `pickle.loads` copy on the zero-copy ZMQ-frame path). Note that +merely raising `VLLM_MQ_MAX_CHUNK_BYTES_MB` (ring vs. socket path) did **not** +help in the in-band era: the serialize ran before the ring-vs-socket branch, so +it was common to both — a negative result that pinned the real cost on the +byte-copies, not the transport choice. + +## 2. The arena + +The arena itself (`ShmTensorArena`, `_ArenaPickler`, `_rebuild_arena_tensor`) +lives in `shm_tensor_arena.py`; `shm_broadcast.py`'s `MessageQueue` creates +and wires it in but otherwise treats it as an internal implementation detail. + +### 2.1 `ShmTensorArena` + +A second shared-memory region created alongside the existing `ShmRingBuffer`: + +- **N slots × slot_bytes** (default 8 × 256 MB, internal constants -- see + §3), plus per-slot metadata `[written_flag, reader0_done … readerN_done]`. +- Concurrency uses the **same lock-free single-writer/N-reader protocol as + `ShmRingBuffer`** (memory fences, per-reader done flags), so the model is one + the codebase already trusts. +- Created by the queue **writer** in `MessageQueue.__init__` when all readers are + node-local; readers attach via a new `tensor_arena_handle` field on the queue + `Handle`. Queues with remote readers get no arena and keep today's behavior. +- Pages are allocated lazily by the kernel on first write, so arenas on queues + that never carry big tensors (e.g. worker→engine response queues) cost + approximately nothing beyond the (virtual) reservation. + +### 2.2 Writer path — arena diversion, composing with `_reduce_tensor` + +`enqueue` builds one `dispatch_table` that routes CPU tensors through +`_reduce_tensor` (out-of-band `PickleBuffer`). When a node-local arena exists, +the pickler is an `_ArenaPickler` whose `reducer_override` *additionally* diverts +**large contiguous** CPU tensors into the arena: + +```python +class _ArenaPickler(pickle.Pickler): + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) # {Tensor: ..., Parameter: ...} + if (rebuild_fn is not None and obj.device.type == "cpu" + and obj.layout is torch.strided and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= MIN_BYTES): + idx = self.arena.write_tensor(obj) # ONE memcpy into a free slot + if idx is not None: + return (rebuild_fn, + (self.arena.shared_memory.name, idx, nbytes, dtype_str, shape)) + return NotImplemented # fall through to dispatch_table → _reduce_tensor +``` + +Note the exact-type dict lookup (not `isinstance`), matching how +`dispatch_table` itself dispatches: an unrecognized `Tensor` subclass must +decline here too, or it would come back as a plain `Tensor`, silently +losing its subclass identity. `torch.nn.Parameter` is the one subclass this +*does* divert, via its own rebuild function (§2.3). + +`reducer_override` is consulted before an object's normal reduction, and +returning `NotImplemented` falls through to the `dispatch_table` — so a diverted +tensor becomes a ~100-byte rebuild stub, and **everything the arena declines** +(too small, non-contiguous, or arena full) is handled by `_reduce_tensor` +exactly as on `main`. The single arena copy — +`torch.frombuffer(slot).copy_(t.view(torch.uint8))` — is a multithreaded memcpy +at memory bandwidth (~50–70 ms for 200 MB), and, unlike the oob path, the *same* +slot is then read zero-copy by all N readers rather than transported to each. + +**The writer never blocks.** If no slot is free or the tensor exceeds the slot +size, `write_tensor` returns `None`; the pickler falls through to `_reduce_tensor` +— the worst case is exactly the current out-of-band behavior, and deadlock is +structurally impossible. + +### 2.3 Reader path — zero copies + +The stub unpickles through a module-level rebuild function: + +```python +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + arena = _TENSOR_ARENAS[arena_name] # this process's registry of attached arenas + # get_tensor: a view over the mapped slot — zero bytes copied + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) +``` + +`_rebuild_arena_parameter` wraps the same tensor in +`torch.nn.Parameter(t, requires_grad=False)`, which shares its storage. +`_ARENA_REBUILD_FNS = {torch.Tensor: _rebuild_arena_tensor, torch.nn. +Parameter: _rebuild_arena_parameter}` is the full set of types the arena +diverts; anything else (any other `Tensor` subclass) declines in +`reducer_override` and takes the out-of-band path. + +The registry is keyed by arena shm name (`_TENSOR_ARENAS: +weakref.WeakValueDictionary[str, ShmTensorArena]`) rather than a single +per-process slot, since nothing guarantees a process only ever attaches to +one arena-bearing queue. It holds only a *weak* reference — a strong one +would keep the arena (and its pinned mapping) alive for the process +lifetime — and needs no explicit removal on `MessageQueue.shutdown()` +(unsafe anyway, since shutdown can run on a different thread than an +in-flight `dequeue()`); an entry just disappears once nothing else +references that arena. + +The rebuilt tensor *is* the shared memory — no transport copy and no deserialize +on any rank. + +**Slot lifecycle.** The rebuilt tensor is the *source* of an async H2D while +the worker executes that step, so the reader must not release the slot at +unpickle time. Release is tied to **the lifetime of the slot's tensor +storage**, not a fixed "next dequeue" schedule — needed for callers like +`prompt_embeds` that retain the tensor across many `dequeue` calls during +chunked prefill, where a fixed schedule would let the writer reclaim (and +corrupt) the slot while the worker was still reading it. `get_tensor` builds +the tensor via `torch.from_numpy` over a numpy array wrapping the slot; the +tensor's storage holds the only reference to that array, so a +`weakref.finalize` on it fires only once every tensor, view/slice and wrapper +(e.g. `Parameter`) sharing the storage has been freed. Once queued, the +pinned fast path (§2.4) adds a second gate: the H2D is a true async DMA that +can outlive `execute_model`, so "the storage was freed" alone isn't +sufficient. +`flush_releases` records a CUDA event on the compute stream after each +step's H2D and only marks a slot done once that event completes +(non-blocking `event.query()`; not-yet-done just waits one more dequeue). +Unpinned, `cudaMemcpyAsync` from pageable memory already stages +synchronously, so the slot releases immediately. Either way, the writer +requires every reader's done flag before reusing a slot. + +> **Note.** Assumes the multimodal H2D is issued on the worker's +> current/default compute stream (true today: mm inputs are copied eagerly, outside the decode CUDA +> graph). If a future vLLM issues that copy on a dedicated side stream, the event +> must be recorded at the copy site rather than at `flush_releases`. + +### 2.4 Pinning — the zero-copy trap + +A zero-copy view over shared memory the reader process has never touched makes +the subsequent `cudaMemcpyAsync` pay **first-touch page faults on ~50k pages plus +pageable-copy staging** — the same tax the current oob path pays on its pageable +ZMQ frame. Fix: each reader lazily `cudaHostRegister`s the whole arena mapping +once (~1 s, first use), after which every H2D from the arena is a **pinned-memory +DMA** (~10 ms for 192 MB). Processes without a CUDA context skip registration +silently. (Because the arena mapping is stable and reused, this one-time pin +amortizes across all future images — an option the transient per-message ZMQ +frame does not have.) + +## 3. Configuration + +A single CLI flag toggles the arena; the slot count/size and divert threshold +are internal constants (no env vars). + +| Flag (`ParallelConfig` field) | Default | Meaning | +| --- | --- | --- | +| `--enable-shm-tensor-arena` / `--no-enable-shm-tensor-arena` (`enable_shm_tensor_arena`) | **off** | Opt in to route large CPU tensors through the arena (reserves slots in `/dev/shm`). Off = the out-of-band `_reduce_tensor` path only, identical to stock behavior. | + +Internal constants in `shm_tensor_arena.py`: **8 slots × 256 MB**; tensors larger +than a slot, or smaller than the **8 MB** divert threshold, take the out-of-band +`_reduce_tensor` path. + +## 4. Validation + +> **Baseline caveat.** The A/B numbers below were collected against the *in-band +> pickle* path that predates `_reduce_tensor`. They quantify the arena vs. that +> older baseline (an upper bound), **not** the arena vs. the current out-of-band +> path. The current-`main` oob baseline would land between the "stock" and +> "arena + pinning" rows; measuring exactly where is the open item. + +1. **Unit tests** (`tests/distributed/test_shm_broadcast.py`): byte-exact + zero-copy round-trips (incl. bf16/fp8), slot lifecycle (no reuse until every + reader releases; exhaustion / oversize / non-contiguous fall back to the + out-of-band path), event-gated release on the pinned path, `_ArenaPickler` ⊕ + `_reduce_tensor` composition, and an end-to-end `MessageQueue` broadcast + through forked processes. (Historical microbenchmark: enqueue of a 199 MB + payload took **66.7 ms** vs ~1275 ms for the old in-band serialize.) +2. **Same-seed A/B** (baseline = in-band pickle, pre-`_reduce_tensor`), interactive + multimodal workload, uncapped input images, TP=4 ×2 workers on one 8-GPU node, + low qps so individual images decompose cleanly (~1.2k aligned requests): + + | TTFT (ms) | p50 | p90 | p99 | max | >1 s | >1.5 s | + | --- | ---: | ---: | ---: | ---: | ---: | ---: | + | in-band pickle (old baseline) | 93 | 401 | 1321 | 2139 | 30 | 9 | + | + bigger MQ chunk (config lever) | 91 | 368 | 1287 | 2027 | 22 | 5 | + | **arena + pinning** | **89** | **241** | **862** | **1375** | **7** | **0** | + + p50 unchanged → `reducer_override` adds no measurable overhead when no large + tensor is present. +3. **Per-image stall** (CUDA-level all-GPU idle gap around the pixel H2D), vs the + in-band baseline: 192 MB image 1318 ms → 681 ms (arena only) → **351 ms class** + (arena + pinning); vision-encoder compute for the same image is ~80 ms. +4. **Outstanding — arena vs. out-of-band on `main`.** Re-run the same seeded A/B + with the baseline set to current `main` (i.e. `_reduce_tensor` on, arena off + via `--no-enable-shm-tensor-arena`) vs. arena on, to quantify the incremental N→1 + copy + pinning gain at TP=4 (and ideally TP=2/8 to show the copy-count scaling). + +## 5. Limitations and future work + +- **Value is TP- and size-dependent.** The arena helps only when a queue has + ≥2 node-local readers *and* carries tensors ≥ `MIN_MB`; at TP=1, or for small + payloads, it adds nothing and everything takes the out-of-band path. +- **Shared-memory reservation.** The arena reserves `slots × slot_bytes` of + `/dev/shm` (default 2 GB, lazily paged). Like `ShmRingBuffer`, creation is + guarded by a free-space check (`check_shm_free_space`), so an undersized + `/dev/shm` fails fast with a clear error instead of a `SIGBUS` the first + time a tensor is copied into pages beyond tmpfs capacity. +- **Relationship to `_reduce_tensor`.** This is strictly additive: the arena is + an opt-in fast path for the large-multimodal-tensor case; declining it (or + disabling via env) reverts to the merged out-of-band behavior. +- **Relationship to existing mm tensor shm caching.** vLLM already has a + separate shared-memory caching path for multimodal tensors; this arena is + not yet unified with it (see the experimental note in §1). Reconciling the + two is follow-up work. +- **Slot release granularity**: releases are tied to the slot's tensor + storage being freed and, on the pinned path, further gated on a per-slot + CUDA event recorded after the consuming H2D (§2.3), which closes both the + multi-step-retention and the async-DMA reuse windows. Bursts deeper than + the slot count safely fall back to the out-of-band path when the arena is + exhausted. +- **Fallback observability**: arena exhaustion (no free slot) is a + rate-limited log line today; a counter metric would be better. An + oversize tensor (bigger than a slot) falls back silently, with no log at + all -- worth adding if that path turns out to matter in practice. +- **Scope**: the arena activates only when every queue reader is node-local. + Remote readers (multi-node PP/TP) keep the existing socket path. diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index d28732261a13..4984a2480a61 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import gc import io import pickle @@ -17,7 +18,7 @@ import torch import torch.distributed as dist -from vllm.distributed.device_communicators import shm_broadcast +from vllm.distributed.device_communicators import shm_broadcast, shm_tensor_arena from vllm.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, @@ -25,6 +26,10 @@ _reduce_tensor, check_shm_free_space, ) +from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + _ArenaPickler, +) from vllm.distributed.utils import StatelessProcessGroup from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -38,15 +43,18 @@ def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]: return [np.random.randint(1, 100, i) for i in sizes] -def distributed_run(fn, world_size, timeout=60): +def distributed_run(fn, world_size, timeout=60, mp_context=None): """Run a function in multiple processes with proper error handling. Args: fn: Function to run in each process world_size: Number of processes to spawn timeout: Maximum time in seconds to wait for processes (default: 60) + mp_context: Optional multiprocess context (e.g. ``mp.get_context("spawn")``) + for workers that must not inherit the parent's thread/CUDA state. """ + ctx = mp_context or mp number_of_processes = world_size processes = [] for i in range(number_of_processes): @@ -57,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60): env["LOCAL_WORLD_SIZE"] = str(number_of_processes) env["MASTER_ADDR"] = "localhost" env["MASTER_PORT"] = "12345" - p = mp.Process(target=fn, args=(env,)) + p = ctx.Process(target=fn, args=(env,)) processes.append(p) p.start() @@ -765,6 +773,20 @@ def test_shm_ring_buffer_creation_checks_free_space(): ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) +def test_shm_tensor_arena_creation_checks_free_space(): + """Like ShmRingBuffer, arena creation must fail fast with a clear error + on an undersized /dev/shm instead of discovering the shortfall via a + SIGBUS the first time a tensor is copied into pages past tmpfs capacity.""" + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) + ), + mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + ShmTensorArena(n_reader=1, slot_bytes=256 << 20, n_slots=8) + + def test_remote_subscribe_addr_unique_concurrent_writers( monkeypatch: pytest.MonkeyPatch, ): @@ -828,3 +850,497 @@ def reader_main(): for q in queues: q.remote_socket.close(linger=0) + + +# --------------- ShmTensorArena (zero-copy tensor arena) tests --------------- + + +def _make_arena(n_reader: int = 2, slot_bytes: int = 4 << 20, n_slots: int = 3): + """Writer arena plus attached per-reader arenas (same process).""" + writer = ShmTensorArena(n_reader, slot_bytes, n_slots) + readers = [ShmTensorArena(*writer.handle(), reader_rank=i) for i in range(n_reader)] + return writer, readers + + +def _get_view(reader: ShmTensorArena, idx: int, ref: torch.Tensor) -> torch.Tensor: + return reader.get_tensor( + idx, ref.numel() * ref.element_size(), ref.dtype, tuple(ref.shape) + ) + + +def _drain(reader: ShmTensorArena) -> None: + """Flush a reader's releases to completion. On the pinned path the first + flush defers the release behind an H2D-completion event; the second flush + retires it once the event has fired.""" + reader.flush_releases() + if reader._deferred_releases: + torch.accelerator.synchronize() + reader.flush_releases() + assert not reader._pending_release + assert not reader._deferred_releases + + +def test_arena_zero_copy_roundtrip(): + writer, (r0, r1) = _make_arena(n_reader=2) + src = torch.randn(400, 512) + idx = writer.write_tensor(src) + assert idx is not None + a = _get_view(r0, idx, src) + b = _get_view(r1, idx, src) + assert torch.equal(a, src) + assert torch.equal(b, src) + assert a.dtype == src.dtype and a.shape == src.shape + # Readers alias one shared mapping: a write through one reader's view is + # visible through the other's (this is what "zero-copy" means here). + a[0, 0] = 12345.0 + assert b[0, 0].item() == 12345.0 + del a, b + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_arena_dtype_roundtrip(dtype): + """Dtypes numpy doesn't recognize traverse the arena's uint8 view path.""" + writer, (reader,) = _make_arena(n_reader=1) + src = torch.randn(256, 256).to(dtype) + idx = writer.write_tensor(src) + assert idx is not None + got = _get_view(reader, idx, src) + assert got.dtype == dtype + assert torch.equal(got.view(torch.uint8), src.view(torch.uint8)) + del got + + +def test_arena_slot_lifecycle(): + """A slot must not be reusable until EVERY reader has released it — + the writer overwriting a slot a reader still consumes would corrupt data. + Release is tied to the returned tensor's storage lifetime (`get_tensor`): + a reader that still holds its view keeps the slot reserved even across + repeated `flush_releases()` calls, not just until the "next" one.""" + writer, readers = _make_arena(n_reader=2, n_slots=3) + t = torch.ones(1000) + first = writer.write_tensor(t) + views = [_get_view(r, first, t) for r in readers] + # Fill the remaining slots; the arena is now exhausted. + assert all(writer.write_tensor(t) is not None for _ in range(2)) + assert writer.write_tensor(t) is None # exhausted -> caller falls back + + # Flushing while both readers still hold their view releases nothing: + # the slot is only queued once the returned tensor's storage is freed. + readers[0].flush_releases() + readers[1].flush_releases() + assert writer.write_tensor(t) is None + + # Dropping one reader's view queues its release; still not enough. + views[0] = None + _drain(readers[0]) + assert writer.write_tensor(t) is None + # Once every reader has dropped (and released) its view, it's reused. + views[1] = None + _drain(readers[1]) + assert writer.write_tensor(t) == first + + +def test_arena_slot_not_released_while_tensor_retained_across_multiple_flushes(): + """A caller retaining the tensor across many `flush_releases()` calls + (e.g. `prompt_embeds` across chunked-prefill steps) must keep its slot + reserved the whole time, not just until the reader's next flush.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + idx = writer.write_tensor(t) + retained = _get_view(reader, idx, t) + # Several flush cycles pass while the caller still holds `retained`. + for _ in range(5): + reader.flush_releases() + assert reader._pending_release == [] + # The slot is still reserved: only the arena's other slot is available... + assert writer.write_tensor(t) is not None + assert writer.write_tensor(t) is None # ...and now the arena is exhausted. + # ...and the bytes the caller is reading are unchanged throughout. + assert torch.equal(retained, t) + # Only once the caller drops its reference is the slot released. + del retained + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_slot_not_released_while_view_retained(): + """A view/slice (or wrapper) sharing the returned tensor's storage keeps + the slot reserved even after the original tensor object is dropped.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=1) + t = torch.arange(1000, dtype=torch.float32) + idx = writer.write_tensor(t) + view = _get_view(reader, idx, t)[10:20].unsqueeze(0) + param = torch.nn.Parameter(view, requires_grad=False) + del view + reader.flush_releases() + assert reader._pending_release == [] + assert writer.write_tensor(t) is None + assert torch.equal(param[0], t[10:20]) + del param + assert reader._pending_release == [idx] + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_release_survives_intervening_flush_from_another_tensor(): + """`flush_releases()` must clear `_pending_release` in place, never + rebind it: a finalizer registered before a rebind would append into the + now-orphaned old list and permanently leak its slot. Exercises the + interleaving a long-lived `prompt_embeds` tensor sees once a + shorter-lived tensor through the same reader triggers a flush first.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + # `a` is long-lived: its finalizer is registered first, while + # _pending_release is still the arena's original (empty) list object. + idx_a = writer.write_tensor(t) + a = _get_view(reader, idx_a, t) + # `b` is short-lived and its finalizer is registered against that SAME + # list object (nothing has flushed yet). + idx_b = writer.write_tensor(t) + b = _get_view(reader, idx_b, t) + del b + assert reader._pending_release == [idx_b] + # This flush must not orphan `a`'s still-pending finalizer. `_drain`, + # not a bare `flush_releases()` call: on the pinned fast path (real + # CUDA present) the release is gated behind an event and may need a + # second flush once it fires -- see `flush_releases`. + _drain(reader) + assert writer.write_tensor(t) == idx_b # b's slot correctly reused. + # `a` is dropped well after that intervening flush. + del a + assert reader._pending_release == [idx_a] + _drain(reader) + assert writer.write_tensor(t) == idx_a # a's slot must NOT be leaked. + + +def test_arena_oversize_falls_back(): + writer, _ = _make_arena(n_reader=1, slot_bytes=1 << 20, n_slots=2) + big = torch.empty((1 << 20) + 4096, dtype=torch.uint8) + assert writer.write_tensor(big) is None + + +def _dumps_arena(obj, arena: ShmTensorArena) -> tuple[bytes, list, list[int]]: + """Pickle `obj` the same way `MessageQueue.enqueue` does when an arena is + attached: arena diversion first (reducer_override), then the tensor + dispatch table, with out-of-band buffers >= 1MiB. Also returns the slot + indices the arena actually wrote to, in write order.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + written: list[int] = [] + orig_write_tensor = arena.write_tensor + + def _spy_write_tensor(t): + idx = orig_write_tensor(t) + if idx is not None: + written.append(idx) + return idx + + bio = io.BytesIO() + pickler = _ArenaPickler(bio, arena, buffer_callback=callback) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + with mock.patch.object(arena, "write_tensor", side_effect=_spy_write_tensor): + pickler.dump(obj) + return bio.getvalue(), buffers, written + + +def test_arena_pickler_composes(monkeypatch): + """Large contiguous tensors are diverted into the arena; everything the + arena declines falls through to `_reduce_tensor` unchanged.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + big = torch.randn(1024, 1024) # 4MiB -> diverted into the arena + small = torch.randn(16, 16) # 1KiB -> falls through to _reduce_tensor + data, buffers, written = _dumps_arena({"big": big, "small": small}, writer) + # The diverted tensor's bytes are in the arena, not the pickle stream. + assert len(data) + sum(b.nbytes for b in buffers) < big.numel() * 4 + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out["big"], big) + assert torch.equal(out["small"], small) + # "big" is a zero-copy view of the reader's slot; "small" is not. + (idx,) = written + nbytes = big.numel() * big.element_size() + slot_ptr = torch.frombuffer(reader._slot(idx, nbytes), dtype=torch.uint8).data_ptr() + assert out["big"].data_ptr() == slot_ptr + assert out["small"].data_ptr() != slot_ptr + # Not yet queued for release: `out["big"]` is still alive (see get_tensor). + assert reader._pending_release == [] + del out + assert reader._pending_release == [idx] + _drain(reader) + + +def test_arena_pickler_noncontig_falls_through(monkeypatch): + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + nc = torch.randn(2048, 1024)[:, ::2] # non-contiguous, above threshold + assert not nc.is_contiguous() + data, buffers, written = _dumps_arena(nc, writer) + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, nc) + # The arena never saw it: no slot consumed on the reader. + assert written == [] + assert reader._pending_release == [] + del out + + +@pytest.mark.parametrize("case", ["requires_grad", "conj"]) +def test_arena_pickler_excludes_requires_grad_and_conj(monkeypatch, case: str): + """Tensors `_reduce_tensor` excludes (requires_grad, or an exotic + conjugate-bit view that can't alias as uint8) must fall through the + arena the same way, even when every other criterion qualifies.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + if case == "requires_grad": + tensor = torch.randn(1024, 1024, requires_grad=True) # 4MiB, above threshold + else: + tensor = torch.randn(600, 600, dtype=torch.complex64).conj() # ~2.9MiB + assert tensor.is_contiguous() + assert tensor.numel() * tensor.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(tensor, writer) + # Declined by the arena -- no slot consumed -- and handled by + # _reduce_tensor's identical exclusion instead. + assert written == [] + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, tensor) + assert out.requires_grad == tensor.requires_grad + + +def test_arena_pickler_excludes_tensor_subclasses(monkeypatch): + """`torch.nn.Parameter` is diverted via its own `_rebuild_arena_parameter` + hook, which wraps an intermediate tensor it keeps no Python reference to. + The slot must stay reserved while the Parameter is alive -- the + `_pending_release` assertions below check this even though + `type(out) is torch.nn.Parameter` alone would not.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + # requires_grad=False to isolate the subclass check from the (already + # separately tested) requires_grad exclusion. + param = torch.nn.Parameter(torch.randn(1024, 1024), requires_grad=False) + assert type(param) is not torch.Tensor + assert param.is_contiguous() + assert param.numel() * param.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(param, writer) + # The Parameter itself is diverted directly (_ARENA_REBUILD_FNS routes + # it to _rebuild_arena_parameter) -- Parameter.__reduce_ex__ never runs, + # so there's no separate "inner" tensor diversion to speak of. + assert written != [] + out = pickle.loads(data, buffers=buffers) + # The outer Parameter's type identity survives. + assert type(out) is torch.nn.Parameter + assert torch.equal(out, param) + assert out.requires_grad == param.requires_grad + # Not yet queued for release: `out` still shares the slot's storage. + assert reader._pending_release == [] + del out + assert reader._pending_release == written + _drain(reader) + + +def test_arena_unregisters_pinned_memory_on_del(monkeypatch): + """A reader that cudaHostRegister-ed its mapping must cudaHostUnregister + the same pointer before the mapping is closed; an unpinned arena must not + unregister memory it never registered.""" + writer, (reader,) = _make_arena(n_reader=1) + cudart = mock.MagicMock() + cudart.cudaHostRegister.return_value = 0 + cudart.cudaHostUnregister.return_value = 0 + monkeypatch.setattr( + shm_tensor_arena, + "current_platform", + SimpleNamespace(is_cuda_alike=lambda: True, cudart=lambda: cudart), + ) + reader._ensure_pinned() + assert reader._pinned + ptr = cudart.cudaHostRegister.call_args.args[0] + del reader + cudart.cudaHostUnregister.assert_called_once_with(ptr) + del writer # never pinned -> no unregister + cudart.cudaHostUnregister.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_arena_event_gated_release(): + """On the pinned path a slot release is gated on an H2D-completion CUDA + event: the writer must not see the reader's done flag until the async DMA + sourced from the slot has been retired. Release is first queued once the + caller drops the CPU-side view -- the realistic pattern: the view is only + needed as the H2D source, while the GPU tensor is what's retained for the + rest of the step.""" + writer, (reader,) = _make_arena(n_reader=1) + reader._ensure_pinned() + if not reader._pinned: + pytest.skip("cudaHostRegister unavailable in this environment") + src = torch.randn(512, 1024) + idx = writer.write_tensor(src) + view = _get_view(reader, idx, src) + dev = view.to("cuda", non_blocking=True) + # Not yet queued: this test still holds `view`. + assert reader._pending_release == [] + del view + assert reader._pending_release == [idx] + reader.flush_releases() + # Deferred behind the event, not applied eagerly. + assert len(reader._deferred_releases) == 1 + with reader._meta(idx) as meta: + assert meta[1] == 0 + torch.accelerator.synchronize() + reader.flush_releases() + assert reader._deferred_releases == [] + with reader._meta(idx) as meta: + assert meta[1] == 1 + assert torch.equal(dev.cpu(), src) + del dev + + +def test_tensor_arena_registry_is_weak_and_survives_shutdown(): + """`_TENSOR_ARENAS` must not keep an arena alive past its owning + MessageQueue, but also must not be cleared by `shutdown()` -- which can + run on a different thread than an in-flight `dequeue()` still + reconstructing a tensor from that arena. A weak-value registry gets + both: the entry survives `shutdown()` and vanishes only once + unreferenced.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + reader_mq = MessageQueue.create_from_handle(handle, rank=0) + + name = reader_mq.tensor_arena.shared_memory.name + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # shutdown() must NOT remove the entry: that would race with a + # concurrent dequeue() still reconstructing a tensor from this arena. + reader_mq.shutdown() + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # Once nothing else references the arena, the registry entry vanishes + # on its own. + reader_mq.tensor_arena = None + assert name not in shm_tensor_arena._TENSOR_ARENAS + + writer_mq.shutdown() + + +def test_create_from_handle_tolerates_arena_attach_failure(): + """If a reader's `ShmTensorArena.__init__` can't attach to the writer's + segment (`FileNotFoundError`, e.g. deserialized on a different node), + `create_from_handle` must not crash -- same deferred-failure convention + as the sibling `ShmRingBuffer` path.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + n_reader, slot_bytes, n_slots, _real_name = handle.tensor_arena_handle + bogus_handle = dataclasses.replace( + handle, + tensor_arena_handle=( + n_reader, + slot_bytes, + n_slots, + "vllm-arena-does-not-exist", + ), + ) + + reader_mq = MessageQueue.create_from_handle(bogus_handle, rank=0) + + # Attach failed silently: the arena object exists but never got its + # shared_memory attribute, and was never registered. + assert reader_mq.tensor_arena is not None + assert not hasattr(reader_mq.tensor_arena, "shared_memory") + assert "vllm-arena-does-not-exist" not in shm_tensor_arena._TENSOR_ARENAS + + reader_mq.shutdown() + writer_mq.shutdown() + + +@worker_fn_wrapper +def worker_fn_arena_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + if rank == writer_rank: + message_queue = MessageQueue( + 1, + 1, + local_reader_ranks=[1], + max_chunk_bytes=8 * 1024 * 1024, + enable_shm_tensor_arena=True, + ) + handles = [message_queue.export_handle()] + dist.broadcast_object_list(handles, src=writer_rank) + else: + handles = [None] + dist.broadcast_object_list(handles, src=writer_rank) + message_queue = MessageQueue.create_from_handle(handles[0], rank) + message_queue.wait_until_ready() + + torch.manual_seed(42) + payload = { + # 16MiB: above the arena divert threshold (8MiB) -> arena slot. + "huge": torch.randn(2048, 2048), + # 2MiB: declined by the arena -> out-of-band _reduce_tensor path. + "mid": torch.randn(1024, 512), + } + + if rank == writer_rank: + assert message_queue.tensor_arena is not None + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + # "huge" was diverted into the arena before the dispatch table was + # consulted; only "mid" went through _reduce_tensor. + assert wrapped_reduce.call_count == 1 + else: + received = message_queue.dequeue(timeout=30) + assert torch.equal(received["huge"], payload["huge"]) + assert torch.equal(received["mid"], payload["mid"]) + # The huge tensor is a zero-copy view into the arena's shared memory + # (not a copy received over the transport). Its slot is not yet + # queued for release while `received` still holds the reference. + arena = message_queue.tensor_arena + assert arena is not None + assert arena._pending_release == [] + huge_ptr = received["huge"].data_ptr() + arena_base = torch.frombuffer( + arena.shared_memory.buf, dtype=torch.uint8 + ).data_ptr() + assert arena_base <= huge_ptr < arena_base + arena.total_bytes + # Dropping the last reference to the slot's storage queues it for + # release (see get_tensor). + del received + assert len(arena._pending_release) == 1 + + dist.barrier() + print(f"arena broadcast passed the test! Rank {rank}") + + +def test_arena_broadcast(): + # Spawn (not fork): by the time this test runs, the pytest process may be + # multi-threaded / CUDA-initialized (earlier GPU tests), and a forked + # child running the arena's slot memcpy can deadlock on inherited lock + # state. Spawned workers start clean; the larger timeout absorbs their + # interpreter + import startup. + distributed_run( + worker_fn_arena_broadcast, + 2, + timeout=180, + mp_context=mp.get_context("spawn"), + ) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 531f6fb8a2b4..d80a8e1fd26d 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -56,6 +56,7 @@ class ForbiddenImport: "vllm/distributed/device_communicators/all_reduce_utils.py", "vllm/distributed/device_communicators/shm_broadcast.py", "vllm/distributed/device_communicators/shm_object_storage.py", + "vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", "vllm/model_executor/model_loader/weight_cache/protocol.py", diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 37affbaf9ebf..6960661e0c7e 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -235,6 +235,15 @@ class ParallelConfig: disable_custom_all_reduce: bool = False """Disable the custom all-reduce kernel and fall back to NCCL.""" + enable_shm_tensor_arena: bool = False + """Route large CPU tensors (e.g. multimodal ``pixel_values``) in the + engine→worker broadcast through a zero-copy shared-memory arena instead of + sending a copy to every reader. **Experimental**: may be reconciled with + vLLM's existing multimodal shm tensor caching in a future release. Opt-in; + most beneficial for multimodal serving with large images and tensor + parallelism. Active only when all queue readers are node-local; reserves + slots in ``/dev/shm``.""" + enable_elastic_ep: bool = False """Enable elastic expert parallelism with stateless NCCL groups for DP/EP.""" elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment] @@ -901,6 +910,7 @@ def compute_hash(self): "nnodes", "max_parallel_loading_workers", "disable_custom_all_reduce", + "enable_shm_tensor_arena", "ray_workers_use_nsight", "ray_runtime_env", "placement_group", diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index da2288f77b54..5656e8aaa153 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -61,6 +61,10 @@ if TYPE_CHECKING: from _typeshed import SizedBuffer + from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + ) + VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL # Cap on how long an idle reader parks before re-reading the authoritative SHM # written-flag. Bounds lost-notify recovery latency to ~5s while the periodic @@ -471,6 +475,7 @@ class Handle: local_notify_addr: str | None = None remote_subscribe_addr: str | None = None remote_addr_ipv6: bool = False + tensor_arena_handle: tuple[int, int, int, str] | None = None def _close_zmq(context: zmq.Context, sockets: list[zmq.Socket]) -> None: @@ -490,6 +495,7 @@ def __init__( max_chunk_bytes: int = 1024 * 1024 * 24, max_chunks: int = 10, connect_ip: str | None = None, + enable_shm_tensor_arena: bool = False, ): if local_reader_ranks is None: local_reader_ranks = list(range(n_local_reader)) @@ -501,12 +507,22 @@ def __init__( self.shutting_down = False context = Context() + self.tensor_arena: ShmTensorArena | None = None if n_local_reader > 0: # for local readers, we will: # 1. create a shared memory ring buffer to communicate small data # 2. create a publish-subscribe socket to communicate large data self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) + # Zero-copy arena for large CPU tensors (see ShmTensorArena). + # Local readers only: remote readers receive the pickled bytes + # over a socket and cannot map the arena, so the substitution + # would break them. + if enable_shm_tensor_arena and n_remote_reader == 0: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena(n_local_reader) + # XPUB is very similar to PUB, # except that it can receive subscription messages # to confirm the number of subscribers @@ -567,6 +583,9 @@ def __init__( local_notify_addr=local_notify_addr, remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, + tensor_arena_handle=( + self.tensor_arena.handle() if self.tensor_arena is not None else None + ), ) self._close_zmq_on_collect(context) @@ -599,6 +618,7 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": context = Context() + self.tensor_arena = None if rank in handle.local_reader_ranks: assert handle.buffer_handle is not None self.buffer = ShmRingBuffer(*handle.buffer_handle) @@ -607,6 +627,23 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._is_local_reader = True self._is_remote_reader = False + arena_handle = handle.tensor_arena_handle + if arena_handle is not None: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena( + *arena_handle, reader_rank=self.local_reader_rank + ) + # `shared_memory` is unset (not just None) if attaching to + # the writer's segment raised FileNotFoundError -- see + # ShmTensorArena.__init__'s reader branch. Registering a + # broken arena is a no-op deferred failure (same convention + # ShmRingBuffer uses above), not a crash here. + if hasattr(self.tensor_arena, "shared_memory"): + shm_tensor_arena._TENSOR_ARENAS[ + self.tensor_arena.shared_memory.name + ] = self.tensor_arena + self.local_socket = context.socket(SUB) self.local_socket.setsockopt_string(SUBSCRIBE, "") socket_addr = handle.local_subscribe_addr @@ -880,12 +917,21 @@ def oob_callback(buf: PickleBuffer) -> bool: # dispatch table would otherwise shadow them. dispatch_table = dict(copyreg.dispatch_table) dispatch_table[torch.Tensor] = _reduce_tensor + arena = self.tensor_arena with io.BytesIO() as bio: - pickler = pickle.Pickler( - bio, - protocol=pickle.HIGHEST_PROTOCOL, - buffer_callback=oob_callback, - ) + pickler: pickle.Pickler + if arena is not None: + from vllm.distributed.device_communicators.shm_tensor_arena import ( + _ArenaPickler, + ) + + pickler = _ArenaPickler(bio, arena, buffer_callback=oob_callback) + else: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) pickler.dispatch_table = dispatch_table pickler.dump(obj) all_buffers[0] = bio.getvalue() @@ -922,6 +968,9 @@ def dequeue( ): """Read from message queue with optional timeout (in seconds)""" if self._is_local_reader: + # See ShmTensorArena.flush_releases for what this retires and why. + if self.tensor_arena is not None: + self.tensor_arena.flush_releases() with self.acquire_read(timeout, indefinite) as buf: overflow = buf[0] == 1 if not overflow: diff --git a/vllm/distributed/device_communicators/shm_tensor_arena.py b/vllm/distributed/device_communicators/shm_tensor_arena.py new file mode 100644 index 000000000000..257cce06b2b4 --- /dev/null +++ b/vllm/distributed/device_communicators/shm_tensor_arena.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zero-copy shared-memory arena for large CPU tensors (multimodal fast path). + +Problem: `MessageQueue.enqueue` (see `shm_broadcast.py`) pickles the whole +payload, and torch's default tensor pickling copies bytes into the pickle +stream on the writer and back out on every local reader. For a large +`pixel_values` tensor (up to hundreds of MB) this blocks the EngineCore step +loop for ~1s with all GPUs idle. + +Fast path: `_ArenaPickler.reducer_override` intercepts qualifying CPU +tensors during pickling, memcpys them ONCE into a free arena slot, and +pickles a tiny `(arena_name, slot, nbytes, dtype, shape)` stub. Readers +rebuild a zero-copy view of the mapped slot (`torch.frombuffer`). + +Slot lifecycle: single writer, n_reader readers, per-slot metadata +[written_flag, reader0_done, ..., readerN_done] (same protocol as +ShmRingBuffer). A slot is not released when rebuilt: the rebuilt tensor is +the SOURCE of an async H2D copy, and callers like chunked-prefill +`prompt_embeds` retain it across many steps. Release is queued only once the +tensor's storage is freed, i.e. once no tensor or view of the slot remains +(see `get_tensor`), and is further gated on a CUDA event on the pinned fast +path, since the H2D can outlive the step that issued it (see +`flush_releases`). + +The writer NEVER blocks: no free slot (or an oversized tensor) falls back to +the default pickle path, so worst case matches today's behavior. + +Physical memory: slots are paged in lazily, so an arena on a queue that +never carries big tensors costs ~0. Creation is guarded by +`check_shm_free_space` (like `ShmRingBuffer`) so an undersized `/dev/shm` +fails fast instead of a `SIGBUS` mid-copy. +""" + +import pickle +import weakref +from contextlib import suppress +from multiprocessing import shared_memory +from unittest.mock import patch + +import numpy as np +import torch + +from vllm.distributed.device_communicators.shm_broadcast import ( + check_shm_free_space, + memory_fence, +) +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +# Arena sizing — internal constants (kept off the env/CLI surface). +_ARENA_SLOTS = 8 +_ARENA_SLOT_BYTES = 256 * 1024 * 1024 +# Minimum contiguous CPU-tensor size to divert into the arena; smaller tensors +# take the out-of-band pickle-buffer path. +_ARENA_MIN_BYTES = 8 * 1024 * 1024 + +# Reader-side registry: arena shm name -> attached ShmTensorArena, keyed by +# name rather than a single slot since a process isn't guaranteed to attach +# to only one. Populated by `MessageQueue.create_from_handle`; consumed by +# the `_rebuild_arena_*` unpickle hooks. WEAK-value: a strong reference here +# would keep the arena (and its pinned mapping) alive for the process +# lifetime, and would need `MessageQueue.shutdown` to remove it explicitly -- +# unsafe since shutdown() can run on a different thread than an in-flight +# `dequeue`. Instead an entry just disappears once nothing else references +# that arena. +_TENSOR_ARENAS: weakref.WeakValueDictionary[str, "ShmTensorArena"] = ( + weakref.WeakValueDictionary() +) + + +class ShmTensorArena: + """Slotted shared-memory arena: one writer, n_reader zero-copy readers. + + Memory layout: [slot0 | slot1 | ... | slotN-1 | meta0 | meta1 | ... ] + where each meta is (1 + n_reader) bytes: [written_flag, reader_done...]. + Slot states mirror ShmRingBuffer: + written=0 -> free (never written / being written) + written=1, some reader_done=0 -> in use, cannot reuse + written=1, all reader_done=1 -> consumed, can reuse + """ + + def __init__( + self, + n_reader: int, + slot_bytes: int = _ARENA_SLOT_BYTES, + n_slots: int = _ARENA_SLOTS, + name: str | None = None, + reader_rank: int = -1, + ): + self.n_reader = n_reader + self.slot_bytes = slot_bytes + self.n_slots = n_slots + self.reader_rank = reader_rank # -1 for the writer + self.metadata_size = 1 + n_reader + self.metadata_offset = slot_bytes * n_slots + self.total_bytes = (slot_bytes + self.metadata_size) * n_slots + self._next_slot = 0 + self._fallbacks = 0 + self._pin_attempted = False + self._pinned = False + self._pinned_ptr = 0 + # slots queued for release once the storage of the tensor `get_tensor` + # handed out is freed; flushed at the next dequeue. + self._pending_release: list[int] = [] + # slots whose release is gated on an H2D-completion CUDA event: each + # entry is (event, [slot_idx, ...]). Only populated on the pinned fast + # path, where the source-side DMA is asynchronous (see flush_releases). + self._deferred_releases: list[tuple[torch.cuda.Event, list[int]]] = [] + + if name is None: + self.is_creator = True + # Same guard as ShmRingBuffer: SharedMemory/ftruncate can succeed + # on an undersized /dev/shm, so check up front instead of SIGBUS + # on first write past tmpfs capacity. + check_shm_free_space(self.total_bytes, allocation_name="ShmTensorArena") + self.shared_memory = shared_memory.SharedMemory( + create=True, size=self.total_bytes + ) + assert self.shared_memory.buf is not None + with self.shared_memory.buf[self.metadata_offset :] as meta: + torch.frombuffer(meta, dtype=torch.uint8).fill_(0) + else: + self.is_creator = False + # same resource_tracker workaround as ShmRingBuffer + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + try: + self.shared_memory = shared_memory.SharedMemory(name=name) + assert self.shared_memory.size >= self.total_bytes + except FileNotFoundError: + # deserialized on a different node; arena unused there + pass + + def handle(self): + return (self.n_reader, self.slot_bytes, self.n_slots, self.shared_memory.name) + + def _meta(self, idx: int) -> memoryview: + start = self.metadata_offset + idx * self.metadata_size + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + self.metadata_size] + + def _slot(self, idx: int, nbytes: int) -> memoryview: + start = idx * self.slot_bytes + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + nbytes] + + # ---- writer side ---- + + def write_tensor(self, t: torch.Tensor) -> int | None: + """Copy tensor bytes into a free slot; return slot idx or None + (caller must then fall back to the default pickle path).""" + nbytes = t.numel() * t.element_size() + if nbytes > self.slot_bytes: + return None + try: + # Computed before claiming a slot: exotic tensors (e.g. conjugate + # bit set) raise RuntimeError here, same as `_reduce_tensor`'s + # identical view -- fail before claiming, not after. + src = t.detach().reshape(-1).view(torch.uint8) + except RuntimeError: + return None + memory_fence() + for probe in range(self.n_slots): + idx = (self._next_slot + probe) % self.n_slots + with self._meta(idx) as meta: + free = meta[0] == 0 or sum(meta[1:]) == self.n_reader + if not free: + continue + meta[0] = 0 # claim + slot_mv = self._slot(idx, nbytes) + try: + dst = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + dst.copy_(src) + finally: + slot_mv.release() + with self._meta(idx) as meta: + for i in range(1, self.n_reader + 1): + meta[i] = 0 + memory_fence() + meta[0] = 1 + memory_fence() + self._next_slot = (idx + 1) % self.n_slots + return idx + self._fallbacks += 1 + if self._fallbacks == 1 or self._fallbacks % 100 == 0: + logger.info( + "ShmTensorArena: no free slot (%d bytes, %d fallbacks so far); " + "falling back to the out-of-band pickle path.", + nbytes, + self._fallbacks, + ) + return None + + # ---- reader side ---- + + def _ensure_pinned(self): + """Pin the whole arena mapping via cudaHostRegister in THIS process + (lazy, once). Without it the HtoD of a zero-copy tensor pays + first-touch page faults on the tmpfs mapping plus pageable staging + (~hundreds of ms for 200MB); registration allocates+pins the pages + once, making every later HtoD a true DMA. Failure (no CUDA in this + process, etc.) is fine — the copy still works, just slower.""" + if self._pin_attempted: + return + self._pin_attempted = True + try: + if not current_platform.is_cuda_alike(): + return + import ctypes + + buf = self.shared_memory.buf + assert buf is not None + ptr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) + ret = current_platform.cudart().cudaHostRegister(ptr, self.total_bytes, 0) + self._pinned = int(ret) == 0 + if self._pinned: + self._pinned_ptr = ptr + logger.info( + "ShmTensorArena: cudaHostRegister(%d MB) -> %s", + self.total_bytes >> 20, + "pinned" if self._pinned else f"error {int(ret)}", + ) + except Exception as e: + logger.info("ShmTensorArena: host-register skipped: %s", e) + + def _unpin(self): + """Undo the pin from _ensure_pinned via cudaHostUnregister; must run + before the mapping is closed. Failures are ignored — at interpreter + shutdown the CUDA context may already be gone, and the registration + dies with the process anyway.""" + if not self._pinned: + return + self._pinned = False + with suppress(Exception): + current_platform.cudart().cudaHostUnregister(self._pinned_ptr) + + def get_tensor( + self, idx: int, nbytes: int, dtype: torch.dtype, shape: tuple[int, ...] + ) -> torch.Tensor: + """Zero-copy view of a slot as a tensor. The slot is queued for + release once the returned tensor's storage is freed, i.e. once the + caller has dropped it along with any views/slices or wrappers (such + as `torch.nn.Parameter`) sharing that storage.""" + self._ensure_pinned() + # The tensor's storage holds the only reference to `holder` and drops + # it when freed, firing the finalizer. (`weakref.finalize` can't + # target the tensor itself, which views and wrappers don't keep alive, + # nor the memoryview, which doesn't support weakrefs.) + holder = np.frombuffer(self._slot(idx, nbytes), dtype=np.uint8) + # Binds to the current list OBJECT: `flush_releases` must mutate it in + # place, never rebind `self._pending_release`. + weakref.finalize(holder, self._pending_release.append, idx) + return torch.from_numpy(holder).view(dtype).view(shape) + + def _mark_released(self, idxs: list[int]): + """Set THIS reader's done flag on the given slots; a slot becomes + reusable once every reader has done so.""" + for idx in idxs: + with self._meta(idx) as meta: + meta[1 + self.reader_rank] = 1 + memory_fence() + + def _record_release_event(self): + """Record a CUDA event on the current (compute) stream, or return None + if this process has no usable CUDA context. Called from flush_releases, + which runs after the previous step's H2D was enqueued on that same + stream, so the event is ordered strictly after that H2D.""" + try: + event = current_platform.Event() + event.record() + return event + except Exception: + return None + + def flush_releases(self): + """Retire this reader's arena slots that are ready for writer reuse. + + Called on every `MessageQueue.dequeue`. `_pending_release` holds + slots whose tensor storage was just freed (see `get_tensor`). On + the pinned fast path that tensor was the source of an async + `non_blocking=True` H2D that can outlive the collection, so instead + of marking released immediately we record a CUDA event (ordered + after that H2D on the same stream) and move the slot to + `_deferred_releases` until `event.query()` says it's done -- a + not-yet-complete slot just waits one more `dequeue`. Unpinned (or no + CUDA context), the H2D already staged synchronously, so the slot is + safe to release right away. + + A slot only becomes free for the writer (`write_tensor`) once every + reader has marked it released (`_mark_released`). + """ + if self._deferred_releases: + still_pending = [] + for deferred_event, idxs in self._deferred_releases: + if deferred_event.query(): + self._mark_released(idxs) + else: + still_pending.append((deferred_event, idxs)) + self._deferred_releases = still_pending + + if not self._pending_release: + return + # Drain in place with atomic pops, never rebind (see get_tensor): a + # finalizer may append from another thread concurrently. + idxs = [] + while self._pending_release: + idxs.append(self._pending_release.pop()) + + event = self._record_release_event() if self._pinned else None + if event is None: + self._mark_released(idxs) + else: + self._deferred_releases.append((event, idxs)) + + def __del__(self): + if hasattr(self, "shared_memory"): + self._unpin() + try: + self.shared_memory.close() + if self.is_creator: + self.shared_memory.unlink() + except BufferError: + # zero-copy tensor views may still hold exported pointers at + # interpreter shutdown; the mapping dies with the process. + pass + + +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook: rebuild a tensor as a zero-copy view of an arena slot.""" + arena = _TENSOR_ARENAS[arena_name] + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + + +def _rebuild_arena_parameter(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook for `torch.nn.Parameter` (shares the slot's storage).""" + t = _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape) + return torch.nn.Parameter(t, requires_grad=False) + + +# Types `_ArenaPickler.reducer_override` diverts, mapped to the rebuild hook +# that reconstructs them. Any other type -- including other `torch.Tensor` +# subclasses -- falls through to the normal pickle path, since an arbitrary +# subclass can't be faithfully rebuilt from just the slot's bytes. +_ARENA_REBUILD_FNS = { + torch.Tensor: _rebuild_arena_tensor, + torch.nn.Parameter: _rebuild_arena_parameter, +} + + +class _ArenaPickler(pickle.Pickler): + """Pickler that diverts large contiguous CPU tensors into the arena. + + `reducer_override` is consulted before an object's normal reduction. For + a qualifying tensor (exact type in `_ARENA_REBUILD_FNS`, CPU, strided, + contiguous, not `requires_grad`, above the divert-size threshold -- the + same criteria `_reduce_tensor` uses, plus size) it does one memcpy into a + free slot and returns a tiny rebuild stub. Anything that doesn't qualify, + or finds the arena full, returns `NotImplemented` and falls through to + the normal out-of-band pickle path, exactly as if no arena were attached. + """ + + def __init__(self, file, arena: ShmTensorArena, buffer_callback=None): + super().__init__( + file, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=buffer_callback, + ) + self.arena = arena + + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) + if ( + rebuild_fn is not None + and obj.device.type == "cpu" + and obj.layout == torch.strided + and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= _ARENA_MIN_BYTES + ): + idx = self.arena.write_tensor(obj) + if idx is not None: + return ( + rebuild_fn, + ( + self.arena.shared_memory.name, + idx, + obj.numel() * obj.element_size(), + str(obj.dtype).removeprefix("torch."), + tuple(obj.shape), + ), + ) + return NotImplemented diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 5f32ec8deb9f..31777e266f8b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -622,6 +622,7 @@ class EngineArgs: allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization enforce_eager: bool = ModelConfig.enforce_eager disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce + enable_shm_tensor_arena: bool = ParallelConfig.enable_shm_tensor_arena language_model_only: bool = MultiModalConfig.language_model_only limit_mm_per_prompt: dict[str, int | dict[str, int]] = get_field( MultiModalConfig, "limit_per_prompt" @@ -1274,6 +1275,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--all2all-backend", **parallel_kwargs["all2all_backend"] ) parallel_group.add_argument("--enable-dbo", **parallel_kwargs["enable_dbo"]) + parallel_group.add_argument( + "--enable-shm-tensor-arena", + **parallel_kwargs["enable_shm_tensor_arena"], + ) parallel_group.add_argument( "--ubatch-size", **parallel_kwargs["ubatch_size"], @@ -2534,6 +2539,7 @@ def create_engine_config( enable_elastic_ep=self.enable_elastic_ep, elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, enable_dbo=self.enable_dbo, + enable_shm_tensor_arena=self.enable_shm_tensor_arena, ubatch_size=self.ubatch_size, dbo_decode_token_threshold=self.dbo_decode_token_threshold, dp_sync_interval=self.dp_sync_interval, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index cad34368546d..749f46b42872 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -166,6 +166,7 @@ def _init_executor(self) -> None: self.local_world_size, max_chunk_bytes=max_chunk_bytes, connect_ip=mq_connect_ip, + enable_shm_tensor_arena=(self.parallel_config.enable_shm_tensor_arena), ) scheduler_output_handle = self.rpc_broadcast_mq.export_handle() # Create workers From 03a7d3776fc2f520a70c8d2da1724eab6df6d0e1 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:33:47 -0400 Subject: [PATCH 322/952] [Perf] Support minimum KV splits configuration for TokenSpeed MLA decode (#60345) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 Co-authored-by: Codex --- tests/engine/test_arg_utils.py | 6 ++ tests/v1/attention/test_mla_backends.py | 27 ++++++-- vllm/config/attention.py | 5 ++ .../attention/backends/mla/tokenspeed_mla.py | 67 ++++++++++++++----- 4 files changed, 83 insertions(+), 22 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index fc9f0229c46c..3e3a82f80fbd 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -503,6 +503,12 @@ def test_attention_config(): engine_args = EngineArgs.from_cli_args(args) assert engine_args.attention_config == AttentionConfig() + args = parser.parse_args( + ["--attention-config", '{"tokenspeed_mla_min_split_kv": 8}'] + ) + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.attention_config.tokenspeed_mla_min_split_kv == 8 + # set backend via dot notation args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"]) assert args is not None diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9c504f31ef7e..91bd9008f078 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1098,6 +1098,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): ], ) @pytest.mark.parametrize("cp_interleave_size", [1, 16, 896]) +@pytest.mark.parametrize("min_split_kv", [1, 8]) def test_tokenspeed_mla_decode_contract( monkeypatch, causal, @@ -1105,6 +1106,7 @@ def test_tokenspeed_mla_decode_contract( dcp_world_size, dcp_rank, cp_interleave_size, + min_split_kv, ): decode_call = None num_decodes = 2 @@ -1148,7 +1150,12 @@ def fake_decode(**kwargs): impl.scale = 1.0 impl.softmax_scale = 1.0 impl.output_scale = 1.0 - impl._workspace_buffer = torch.empty(1, dtype=torch.int8) + impl._min_split_kv = min_split_kv + impl._max_decode_tokens = num_decode_tokens + workspace = torch.empty(1, dtype=torch.int8) + monkeypatch.setattr( + tokenspeed_mla_module, "_get_workspace", lambda *args: workspace + ) metadata = SimpleNamespace( num_decodes=num_decodes, @@ -1183,6 +1190,8 @@ def fake_decode(**kwargs): assert lse.shape == (num_decode_tokens, num_heads) assert decode_call is not None + assert decode_call["workspace_buffer"] is workspace + assert decode_call["min_split_kv"] == min_split_kv assert decode_call["query"].shape == ( num_decodes, tokens_per_decode, @@ -1217,16 +1226,23 @@ def fake_decode(**kwargs): pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"), ], ) -def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len): +def test_tokenspeed_mla_dcp_matches_unsharded_decode( + monkeypatch, cp_interleave_size, seq_len +): from tokenspeed_mla import tokenspeed_mla_decode from vllm.v1.attention.ops.dcp import ( _lse_weighted_combine, mask_dcp_empty_shards_, ) + from vllm.v1.worker.workspace import WorkspaceManager torch.manual_seed(7) device = torch.device("cuda") + workspace_manager = WorkspaceManager(device) + monkeypatch.setattr( + tokenspeed_mla_module, "current_workspace_manager", lambda: workspace_manager + ) dcp_world_size = 8 kernel_block_size = 64 num_heads = 128 @@ -1257,7 +1273,9 @@ def make_paged_cache(tokens: torch.Tensor): global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to( torch.float8_e4m3fn ) - workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank) + workspace = tokenspeed_mla_module._get_workspace( + device, num_heads, kv_lora_rank, 1, query.shape[1] + ) global_cache, global_block_table = make_paged_cache(global_tokens) global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device) reference, _ = tokenspeed_mla_decode( @@ -1301,7 +1319,8 @@ def make_paged_cache(tokens: torch.Tensor): impl.scale = head_size**-0.5 impl.softmax_scale = None impl.output_scale = None - impl._workspace_buffer = workspace + impl._min_split_kv = 1 + impl._max_decode_tokens = query.shape[1] metadata = SimpleNamespace( num_decodes=1, num_decode_tokens=query.shape[1], diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 8a9be6336c5e..80548f58f2ca 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -61,6 +61,11 @@ class AttentionConfig: flash_attn_max_num_splits_for_cuda_graph: int = 32 """Flash Attention max number splits for cuda graph decode.""" + tokenspeed_mla_min_split_kv: int = Field(default=1, ge=1, le=256) + """Minimum KV splits for TokenSpeed MLA decode. The default of 1 preserves + automatic scheduling. Larger values request a split floor and may increase + workspace memory. Only applies to the TokenSpeed MLA backend.""" + tq_max_kv_splits_for_cuda_graph: int = 32 """TurboQuant max NUM_KV_SPLITS for cuda graph decode. Fixes the split count so grid dimensions are constant across captures, diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 96859d7f9884..8d41cd178ed8 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -24,6 +24,7 @@ AttentionType, MultipleOf, ) +from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: from vllm.config import VllmConfig @@ -38,21 +39,38 @@ # own buffer check. _TOKENSPEED_MAX_Q_LEN = 8 -_g_workspace: dict[torch.device, torch.Tensor] = {} - def _get_workspace( - device: torch.device, num_heads: int, kv_lora_rank: int + device: torch.device, + num_heads: int, + kv_lora_rank: int, + min_split_kv: int, + max_decode_tokens: int, ) -> torch.Tensor: from tokenspeed_mla import get_num_sm + # TokenSpeed reduction scratch: B * H * Q * split_kv * (D + 1) * 4 bytes, + # using the effective query layout and FP32 partial outputs plus LSE. + # Automatic splitting lowers split_kv as the batch grows; split_kv=1 + # needs no reduction scratch, so an SM-based bound covers all batch sizes. needed = ( get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4 ) - existing = _g_workspace.get(device) - if existing is None or existing.numel() < needed: - _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device) - return _g_workspace[device] + if min_split_kv > 1: + # A forced floor breaks that bound: reserve for the maximum batch. + # Packed queries pad head rows to 128; reserve one tile per query token. + needed = max( + needed, + max_decode_tokens + * max(num_heads, 128) + * min_split_kv + * (kv_lora_rank + 1) + * 4, + ) + logger.info_once( + "TokenSpeed MLA decode workspace: %.2f GiB", needed / (1 << 30) + ) + return current_workspace_manager().get_simultaneous(((needed,), torch.int8))[0] class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): @@ -212,15 +230,25 @@ def __init__( f"got kv_cache_dtype={self.kv_cache_dtype!r}." ) - # Allocate (or fetch the cached) workspace lazily on first forward — - # __init__ runs before the device is necessarily set on the worker; - # we know it for sure at forward time when we see the input tensor. - self._workspace_buffer: torch.Tensor | None = None + config = get_current_vllm_config() + self._min_split_kv = config.attention_config.tokenspeed_mla_min_split_kv + self._max_decode_tokens = 0 + if self._min_split_kv > 1: + spec = config.speculative_config + query_len = 1 + (spec.num_speculative_tokens if spec else 0) + self._max_decode_tokens = max( + min( + config.scheduler_config.max_num_batched_tokens, + config.scheduler_config.max_num_seqs * query_len, + ), + config.compilation_config.max_cudagraph_capture_size or 0, + ) + self.softmax_scale: float | None = None self.output_scale: float | None = None # NIXL resolves interleaving after model construction; retain the config # rather than caching its initial interleave size. - self._parallel_config = get_current_vllm_config().parallel_config + self._parallel_config = config.parallel_config # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always # runs when tokenspeed is selected, prefill backend may not (user can @@ -290,11 +318,13 @@ def forward_mqa( ) self.output_scale = layer._k_scale_float - if self._workspace_buffer is None: - # Parallelism can change the runtime query head count. - self._workspace_buffer = _get_workspace( - q.device, q.shape[-2], self.kv_lora_rank - ) + workspace = _get_workspace( + q.device, + q.shape[-2], + self.kv_lora_rank, + self._min_split_kv, + self._max_decode_tokens, + ) # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size). # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm). @@ -302,7 +332,7 @@ def forward_mqa( kernel_out = tokenspeed_mla_decode( query=q, kv_cache=kv_c_and_k_pe_cache, - workspace_buffer=self._workspace_buffer, + workspace_buffer=workspace, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=block_tables, @@ -317,6 +347,7 @@ def forward_mqa( cp_world=self.dcp_world_size, cp_rank=self.dcp_rank, cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size, + min_split_kv=self._min_split_kv, ) if return_lse: o, lse = kernel_out From f8e6e365aec519285d99961015ee270a3871d871 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Fri, 9 Oct 2026 17:08:08 -0700 Subject: [PATCH 323/952] [Bugfix][Mooncake] Save the EAGLE attention proof with hybrid Mamba partial tails (#60695) Signed-off-by: Dao Le Signed-off-by: Dao007forever Signed-off-by: zjy0516 Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: zjy0516 --- .buildkite/test_areas/models_basic.yaml | 2 + tests/models/kimi_k3/test_prefix_cache.py | 127 +++++- .../unit/test_mooncake_store_coordinator.py | 44 +- .../unit/test_mooncake_store_hma_e2e.py | 95 +++-- .../unit/test_mooncake_store_scheduler.py | 316 ++++++++++++-- .../unit/test_mooncake_store_worker.py | 392 ++++++++++++++++-- .../v1/mooncake/store/coordinator.py | 98 +++-- .../kv_connector/v1/mooncake/store/data.py | 51 ++- .../v1/mooncake/store/scheduler.py | 202 ++++++--- .../kv_connector/v1/mooncake/store/worker.py | 190 ++------- vllm/v1/engine/core.py | 5 +- 11 files changed, 1137 insertions(+), 385 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index fd21324a7ece..1d9a7fe2f3b1 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -165,6 +165,8 @@ steps: - vllm/v1/core/ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ + - vllm/distributed/mooncake_store.py - vllm/v1/simple_kv_offload/ - vllm/v1/spec_decode/ - vllm/model_executor/layers/mamba/ diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 20e5d6670e5c..d6819a507ab8 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -4,7 +4,12 @@ import contextlib import json +import os import random +import shutil +import socket +import subprocess +import time from collections.abc import Iterator from dataclasses import dataclass from math import lcm @@ -60,6 +65,16 @@ } NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"} +MOONCAKE = { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_both", + "kv_load_failure_policy": "recompute", + "kv_connector_extra_config": { + "load_async": True, + "lookup_async": True, + "enable_cross_layers_blocks": False, + }, +} # SimpleCPUOffloadConnector as in the Kimi-K3 recipe. OFFLOAD = { "kv_connector": "SimpleCPUOffloadConnector", @@ -243,10 +258,16 @@ def _tokens_text_logprobs(response: dict): def _reset_gpu_prefix_cache(url: str) -> None: - response = requests.post( - f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 - ) - response.raise_for_status() + deadline = time.monotonic() + 60 + while True: + response = requests.post( + f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 + ) + response.raise_for_status() + if response.json()["success"]: + return + assert time.monotonic() < deadline, "GPU cache reset timed out" + time.sleep(0.1) @dataclass(frozen=True) @@ -354,29 +375,19 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) -@pytest.mark.parametrize( - "prompt_len", - [ - # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: - # ensure lookup reads the cached attention proof at 7040 despite - # the final hit limit of 7039, then applies EAGLE's 128-token drop. - pytest.param(7040, id="eagle-proof-scan"), - # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: - # ensure the Mamba checkpoint is 7168, not 7040 from applying both - # the P - 1 cap and EAGLE's 128-token drop. - pytest.param(7296, id="eagle-double-cap"), - ], -) -def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: +def _check_resend_reuses_prompt_checkpoint( + deployment: Deployment, prompt_len: int +) -> None: """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" mode = Mode(128, True) salt = str(uuid4()) rng = random.Random(0) prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)] - deployment = DEPLOYMENTS["plain"] with _serve(deployment, mode) as servers: url = servers[0].url_root cold = _complete(url, prompt, salt, max_tokens=1) + if deployment.offload: + _reset_gpu_prefix_cache(url) resend = _complete(url, prompt, salt, max_tokens=1) recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) @@ -393,6 +404,84 @@ def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: ) +@pytest.mark.parametrize( + "prompt_len", + [ + # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: + # ensure lookup reads the cached attention proof at 7040 despite + # the final hit limit of 7039, then applies EAGLE's 128-token drop. + pytest.param(7040, id="eagle-proof-scan"), + # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: + # ensure the Mamba checkpoint is 7168, not 7040 from applying both + # the P - 1 cap and EAGLE's 128-token drop. + pytest.param(7296, id="eagle-double-cap"), + ], +) +def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: + _check_resend_reuses_prompt_checkpoint(DEPLOYMENTS["plain"], prompt_len) + + +@pytest.fixture +def mooncake_store(tmp_path, monkeypatch) -> Iterator[None]: + """Run a Mooncake master and point the servers' store config at it.""" + if shutil.which("mooncake_master") is None: + pytest.skip("mooncake_master is not installed") + port = get_open_port() + log_path = tmp_path / "mooncake_master.log" + with open(log_path, "w") as log: + master = subprocess.Popen( + [ + "mooncake_master", + f"--port={port}", + f"--metrics_port={get_open_port()}", + "--enable_metric_reporting=false", + "--default_kv_lease_ttl=60000", + ], + stdout=log, + stderr=subprocess.STDOUT, + ) + try: + deadline = time.monotonic() + 60 + while True: + assert master.poll() is None, log_path.read_text() + with ( + contextlib.suppress(OSError), + socket.create_connection(("127.0.0.1", port)), + ): + break + assert time.monotonic() < deadline, "Mooncake master did not start" + time.sleep(0.1) + config_path = tmp_path / "mooncake.json" + config_path.write_text( + json.dumps( + { + "metadata_server": "P2PHANDSHAKE", + "master_server_address": f"127.0.0.1:{port}", + "protocol": "rdma", + "device_name": os.getenv("MOONCAKE_DEVICE_NAME", "mlx5_12"), + "global_segment_size": "4GB", + "local_buffer_size": "1GB", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(config_path)) + monkeypatch.setenv("MC_MAX_MR_SIZE", str(4 << 30)) + yield + finally: + master.kill() + master.wait() + + +@pytest.mark.usefixtures("mooncake_store") +def test_mooncake_resend_reuses_prompt_checkpoint() -> None: + """Blocks=6144, PMU=128. P=7449 needs Mamba state at 7296 and attention proof + at 7424, both beyond the normal Mooncake save boundary of 6144: ensure they + remain reusable from the store after a GPU cache reset.""" + _check_resend_reuses_prompt_checkpoint( + Deployment(Instance(kv_config=MOONCAKE), offload=True), 7449 + ) + + @pytest.mark.parametrize( "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 8767b1c0d983..ac52ca07ab33 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -30,7 +30,13 @@ def _mamba_align(block_size=32): ) -def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): +def _make_coord( + groups, + hash_block_size, + use_eagle=False, + retention_interval=None, + enable_partial_hash_hits=False, +): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for the test fixtures.""" @@ -42,6 +48,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non hash_block_size=hash_block_size, use_eagle=use_eagle, retention_interval=retention_interval, + enable_partial_hash_hits=enable_partial_hash_hits, ) @@ -219,8 +226,7 @@ def test_coordinator_fine_grained_partial_tail_hit(): KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), ] - coord = _make_coord(groups, hash_block_size=16) - assert coord.enable_partial_hash_hits + coord = _make_coord(groups, hash_block_size=16, enable_partial_hash_hits=True) hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc. # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full. @@ -287,38 +293,6 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): assert masks[1] == [False, False, False, True, False, False, False, True] -def test_store_mask_swa_does_not_double_scale_under_dcp(): - """The coordinator receives DCP-resolved specs (worker.py applies - resolve_dcp_kv_cache_spec) and indexes chunks in that already-scaled block - size. So dcp_world_size must NOT further scale reachable_block_mask; the - mask must be identical for any dcp_world_size given the same specs. - - Regression: previously the coordinator forwarded its dcp_world_size, which - double-scaled the SWA block size and produced a wrong mask when dcp > 1. - """ - full = _full(32) - swa = _swa(block_size=8, sliding_window=8) - groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] - scheduler_block_size = lcm(32, 8) - - def make(dcp_world_size): - return MooncakeStoreCoordinator( - groups, - scheduler_block_size=scheduler_block_size, - hash_block_size=8, - dcp_world_size=dcp_world_size, - ) - - mask_dcp1 = make(1).store_mask(64) - mask_dcp2 = make(2).store_mask(64) - - # dcp must not change the mask; both match the correct single-block-size - # result (chunks ending at 32 and 64 -> blocks 3 and 7). - expected = [False, False, False, True, False, False, False, True] - assert mask_dcp1[1] == expected - assert mask_dcp2[1] == expected - - def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm boundary should now span two SWA tail blocks.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index a8354390d71e..0682f9f465f1 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -18,11 +18,15 @@ MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LoadSpec, ReqMeta, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( # noqa: E501 KVCacheStoreRecvingThread, KVCacheStoreSendingThread, @@ -378,6 +382,15 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): assert out[1][2].hex() == fine_hashes[6].hex() +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def test_sub_block_partial_tail_offload_reads_cow_block(): """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the partial tail is offloaded for both groups under the boundary sub-hash. The @@ -394,7 +407,12 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -441,11 +459,13 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + num_prompt_tokens=13, + prefill_end_tokens=20, + completed_token_len=20, + boundary_puts=[(1, mamba_cow_block, 12)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2]. partial_hash = hs[2] @@ -473,7 +493,12 @@ def test_offload_syncs_event_before_put(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) event = MagicMock() class _FencedStore(_DictStore): @@ -513,12 +538,13 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_hashes=hs, can_save=True, num_prompt_tokens=12, + completed_token_len=12, store_job_id=1, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 8)], ) req.current_event = event - send.add_request(req) + send.add_request(_resolve_partial_tail(send, req)) send._handle_request(send.request_queue.get()) assert send.request_queue.qsize() == 0 assert store._data @@ -532,7 +558,8 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): offload must persist every FA block up to the boundary — the normal save floors to the lcm, so those blocks are otherwise never written and the consumer's per-group lookup would miss. The mamba boundary block still - reads the core-provided CoW target.""" + reads the core-provided CoW target. Without EAGLE, the 12-token prompt's + resend-safe checkpoint is 8, leaving two attention blocks to save.""" full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -544,7 +571,12 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -590,19 +622,19 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): block_hashes=hs, can_save=True, num_prompt_tokens=12, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + completed_token_len=12, + boundary_puts=[(1, mamba_cow_block, 8)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) - # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal + # FA (block 4): full blocks ending at 4 and 8, keyed by their normal # block-end hashes; mamba (block 16): the partial boundary block under # the boundary sub-hash, read from the CoW target. expected = { token_dbs[0].key_for(hs[0]): [1 * 512], token_dbs[0].key_for(hs[1]): [2 * 512], - token_dbs[0].key_for(hs[2]): [3 * 512], - token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512], + token_dbs[1].key_for(hs[1]): [10_000 + mamba_cow_block * 512], } assert store.puts == expected @@ -611,11 +643,10 @@ def test_worker_lookup_hits_sub_block_partial_tail(): """worker.lookup must query sub-block keys when partial hash hits are on. Regression test: ``lookup`` hard-coded ``fine_grained = False``, so the - sub-block keys persisted by ``_sub_block_tail_puts`` were never probed and - partial prefix hits silently returned 0 while the store side kept writing - them. Here the mamba block (16) exceeds the hash unit (4), so - ``enable_partial_hash_hits`` is on and the lookup must find the stored - boundary at 12. + partial tail's sub-block keys were never probed and partial prefix hits + silently returned 0 while the store side kept writing them. Here the + mamba block (16) exceeds the hash unit (4), so ``enable_partial_hash_hits`` + is on and the lookup must find the stored boundary at 8. """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( @@ -644,6 +675,8 @@ def test_worker_lookup_hits_sub_block_partial_tail(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=4, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 4 < mamba block 16 -> partial hash hits are enabled. @@ -672,7 +705,7 @@ def test_worker_lookup_hits_sub_block_partial_tail(): replicate_config=MagicMock(), ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//4-1]). + # A 12-token prompt leaves a resend-safe sub-block checkpoint at 8. hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)] req = ReqMeta( req_id="r0", @@ -680,15 +713,16 @@ def test_worker_lookup_hits_sub_block_partial_tail(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store - # A 13-token prompt sharing the prefix must hit the stored boundary at 12. - assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 12 + # A 13-token prompt sharing the prefix must hit the stored boundary at 8. + assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8 def test_worker_setup_tolerates_finer_scratch_group(): @@ -744,6 +778,8 @@ def test_worker_setup_tolerates_finer_scratch_group(): KVCacheGroupSpec(["L1"], mamba), KVCacheGroupSpec(["L2"], scratch), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=8, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 8 divides the participating groups (16) but not the scratch @@ -777,7 +813,7 @@ def test_worker_setup_tolerates_finer_scratch_group(): group_participates=[True, True], ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]). + # Persist the sub-block partial tail at boundary 8 (keyed by hs[8//8-1]). hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)] req = ReqMeta( req_id="r0", @@ -785,10 +821,11 @@ def test_worker_setup_tolerates_finer_scratch_group(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 7ced38d65a41..47796eb14707 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -7,7 +7,11 @@ import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, LoadSpec, MooncakeLookupResult, MooncakeStoreConnectorMetadata, @@ -17,6 +21,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( MooncakeStoreScheduler, + _partial_tail_non_mamba_puts, ) from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks @@ -62,6 +67,31 @@ def _make_bare_scheduler( scheduler._next_store_job_id = 0 scheduler._pinned_saves = {} scheduler._boundary_state_group_ids = frozenset({1}) + scheduler._store_coord = MooncakeStoreCoordinator( + [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=max(16, hash_block_size), + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=max(16, hash_block_size), + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + max(16, hash_block_size), + hash_block_size, + enable_partial_hash_hits=enable_partial_hash_hits, + ) return scheduler @@ -121,6 +151,7 @@ def _make_new_scheduler_output() -> SimpleNamespace: req_id="req-0", num_computed_tokens=0, prompt_token_ids=list(range(32)), + num_prompt_tokens=32, prefill_token_ids=None, block_ids=([0, 1],), block_hashes=[b"h0", b"h1"], @@ -199,6 +230,7 @@ def _add_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=prefill_end_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -268,6 +300,7 @@ def _make_qsa_hybrid_cache_config(): def test_scheduler_projects_nonprefix_groups_and_mamba_ids(): vllm_config = SimpleNamespace( + speculative_config=None, kv_transfer_config=SimpleNamespace( kv_role="kv_both", kv_connector_extra_config={} ), @@ -528,6 +561,7 @@ def _make_pending_load_unfinished_request( ) -> None: request = SimpleNamespace( num_tokens=num_tokens, + num_prompt_tokens=num_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -594,6 +628,7 @@ def _make_resumed_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=32, block_hashes=block_hashes, num_computed_tokens=num_computed_tokens, num_output_placeholders=0, @@ -796,13 +831,24 @@ def test_from_request_tracker_load_with_can_load_false_still_saves(): assert tracker.num_saved_tokens == 48 -def test_from_request_tracker_no_load_saves_normally(): +@pytest.mark.parametrize("token_len,save_partial_tail", [(48, False), (13, True)]) +def test_from_request_tracker_no_load_saves_normally(token_len, save_partial_tail): tracker = RequestTracker( req_id="req-0", - token_len=48, + token_len=token_len, allocated_block_ids=([0, 1, 2],), num_saved_tokens=0, + prefill_end_tokens=token_len, ) + if save_partial_tail: + tracker.token_len = 8 + assert ( + ReqMeta.from_request_tracker( + tracker, 16, save_partial_tail=True, num_prompt_tokens=token_len + ) + is None + ) + tracker.token_len = token_len req_meta = ReqMeta.from_request_tracker( tracker, @@ -810,12 +856,70 @@ def test_from_request_tracker_no_load_saves_normally(): load_spec=None, skip_save=False, block_hashes=[b"h0", b"h1", b"h2"], + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, ) assert req_meta is not None assert req_meta.can_save is True assert req_meta.load_spec is None - assert tracker.num_saved_tokens == 48 + assert req_meta.token_len_chunk == token_len // 16 * 16 + assert req_meta.completed_token_len == token_len + assert req_meta.boundary_puts is None + assert tracker.num_saved_tokens == token_len // 16 * 16 + assert ( + ReqMeta.from_request_tracker( + tracker, + 16, + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, + ) + is None + ) + + +def test_partial_tail_is_resolved_once_on_the_prompt_completing_save(): + scheduler = _make_bare_scheduler( + hash_block_size=4, enable_partial_hash_hits=True, save_decode_cache=True + ) + scheduler._store_group_ids = (0, 1) + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[bytes([i]) for i in range(12)], + prefill_end_tokens=45, + ) + block_ids = ([0, 1, 2], [50]) + scheduler._request_trackers["req-0"].allocated_block_ids = block_ids + + def step(num_computed_tokens: int, num_scheduled_tokens: int): + return scheduler.build_connector_meta( + SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=["req-0"], + new_block_ids=[()], + num_computed_tokens=[num_computed_tokens], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req-0": num_scheduled_tokens}, + scheduled_spec_decode_tokens={}, + kv_connector_block_state=_make_connector_block_state(block_ids), + ) + ).requests + + # The save that completes the prompt carries the tail for the worker. + (req_meta,) = step(44, 1) + assert req_meta.publish_partial_tail + assert req_meta.boundary_puts == [(0, 2, 44)] + + # A later save does not recompute or resend it. + (req_meta,) = step(45, 3) + assert req_meta.token_len_chunk == 48 + assert not req_meta.publish_partial_tail + assert req_meta.boundary_puts is None class _StubLookupClient: @@ -1018,7 +1122,7 @@ def _add_pending_partial_tail_request( all_token_ids=list(range(num_tokens)), block_hashes=block_hashes, num_output_placeholders=0, - num_prompt_tokens=12, + num_prompt_tokens=13, ) scheduler._unfinished_requests["req-0"] = (request, block_ids) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1057,7 +1161,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) @@ -1069,8 +1173,8 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert req_meta.req_id == "req-0" assert req_meta.can_save is True assert req_meta.token_len_chunk == 0 - assert req_meta.boundary_state_offloads == [(1, 7, 12)] - assert req_meta.num_prompt_tokens == 12 + assert req_meta.boundary_puts == [(1, 7, 12)] + assert req_meta.num_prompt_tokens == 13 assert req_meta.block_ids == ([0],) store_job_id = req_meta.store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7] @@ -1081,34 +1185,71 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert tracker.has_pending_offload is True -def test_finished_partial_tail_is_pre_pinned_as_store_job(): +@pytest.mark.parametrize("attention_block_size", [4, 8, 16]) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("num_workers", [1, 2]) +def test_finished_partial_tail_is_pre_pinned_as_store_job( + attention_block_size, use_eagle, num_workers +): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) scheduler._store_group_ids = (0, 1) scheduler.client = SimpleNamespace(discard=lambda *_: None) + scheduler._num_workers = num_workers + groups = list(scheduler._store_coord.kv_cache_groups) + groups[0] = KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=use_eagle, + ) + scheduler._store_coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle, enable_partial_hash_hits=True + ) + prompt_tokens = 49 if use_eagle else 45 + attention_ids = list( + range(1, (prompt_tokens + attention_block_size - 1) // attention_block_size + 1) + ) + block_ids = (attention_ids, [50]) request = SimpleNamespace( request_id="req-0", - block_hashes=[b"h0", b"h1", b"h2"], + block_hashes=[bytes([i]) for i in range(prompt_tokens // 4)], + num_computed_tokens=prompt_tokens, + num_prompt_tokens=prompt_tokens, ) scheduler._request_trackers["req-0"] = RequestTracker( req_id="req-0", - token_len=12, - allocated_block_ids=([3], [9]), - token_ids=list(range(12)), - prefill_end_tokens=12, + token_len=prompt_tokens, + allocated_block_ids=block_ids, + token_ids=list(range(prompt_tokens)), + prefill_end_tokens=prompt_tokens, ) - block_ids = ([3], [9]) + pool = scheduler._gpu_block_pool + owned_ids = attention_ids + [50] + pool.touch([pool.blocks[i] for i in owned_ids]) delay_free = scheduler.register_finished_partial_tail( request, block_ids, - [(1, 9, 12)], + [(1, 50, 44)], ) # The exact source is pinned immediately, so the request can be freed # before the next connector metadata build. assert delay_free is False - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 - assert scheduler._gpu_block_pool.blocks[3].ref_cnt == 0 + proof_end = 48 if use_eagle else 44 + expected = attention_ids[ + 32 // attention_block_size : (proof_end + attention_block_size - 1) + // attention_block_size + ] + pool.free_blocks([pool.blocks[i] for i in owned_ids]) + assert all( + pool.blocks[i].ref_cnt == (1 if i in expected else 0) for i in attention_ids + ) + assert pool.blocks[50].ref_cnt == 1 out = SimpleNamespace( finished_req_ids={"req-0"}, @@ -1132,13 +1273,116 @@ def test_finished_partial_tail_is_pre_pinned_as_store_job(): assert req_meta.token_len_chunk == 0 assert req_meta.block_ids == block_ids assert req_meta.block_hashes == request.block_hashes - assert req_meta.boundary_state_offloads == [(1, 9, 12)] - assert scheduler._pinned_saves[req_meta.store_job_id][0] == [9] - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 + assert req_meta.boundary_puts == [ + (0, attention_ids[idx], min((idx + 1) * attention_block_size, proof_end)) + for idx in range( + 32 // attention_block_size, + (proof_end + attention_block_size - 1) // attention_block_size, + ) + ] + [(1, 50, 44)] + assert set(scheduler._pinned_saves[req_meta.store_job_id][0]) == {50, *expected} + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) assert scheduler._finished_partial_tail_metas == {} + for _ in range(num_workers - 1): + scheduler.update_connector_output( + _make_worker_output({req_meta.store_job_id: 1}) + ) + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) scheduler.update_connector_output(_make_worker_output({req_meta.store_job_id: 1})) - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 0 + assert all(pool.blocks[i].ref_cnt == 0 for i in owned_ids) + + +def test_partial_tail_boundary_follows_eagle_block_drop_not_group_flags(): + """Core drops the Mamba checkpoint whenever EAGLE block drop is on, even if + only the Mamba group carries the eagle flag.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + is_eagle_group=True, + ), + ] + coord = MooncakeStoreCoordinator(groups, 16, 4, use_eagle=True) + # Core's checkpoint for a 49-token prompt: 48 minus one EAGLE hash unit. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=([1, 2, 3, 4], [50]), + block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=49, + completed_token_len=49, + boundary_puts=[BoundaryPut(1, 50, 44)], + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [16, 16]) == [ + BoundaryPut(0, 3, 44) + ] + + +@pytest.mark.parametrize( + "completed, mamba_tails, publish, expected", + [ + # The chunk ending at the junction (40) hands off its Mamba state. Its + # attention KV comes from the normal saves, so nothing is added here. + (40, [40], False, []), + # The prompt-completing save: blocks from the LCM floor 64 up to the + # checkpoint 68 plus the EAGLE proof unit, whether or not the junction + # hand-off is in the same save. + (75, [68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + (75, [40, 68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + ], +) +def test_shared_prefix_junction_handoff_adds_no_attention_tail( + completed, mamba_tails, publish, expected +): + """``--enable-mamba-shared-prefix-checkpoint`` also hands off a sub-block + Mamba state at the shared-prefix junction, below the prompt checkpoint.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle=True, enable_partial_hash_hits=True + ) + # 75-token prompt: checkpoint 72 - 4 = 68. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=(list(range(1, 20)), [50, 51, 52, 53, 54]), + block_hashes=[bytes([i]) for i in range(18)], + num_prompt_tokens=75, + completed_token_len=completed, + boundary_puts=[BoundaryPut(1, 60 + i, n) for i, n in enumerate(mamba_tails)], + publish_partial_tail=publish, + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [4, 16]) == expected def test_decode_boundary_state_offload_dropped_unclaimed(): @@ -1235,7 +1479,7 @@ def test_boundary_state_group_ids_are_remapped_to_store_projection(): scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta) - assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)] + assert meta.requests[0].boundary_puts == [(1, 7, 800)] def test_resumed_prefill_claims_boundaries_past_prompt_length(): @@ -1245,13 +1489,13 @@ def test_resumed_prefill_claims_boundaries_past_prompt_length(): # missing for boundaries full attention does store and no joint hit could # complete there. scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) - _register_offload_request(scheduler, prefill_end_tokens=20, num_prompt_tokens=12) - out = _make_offload_only_output([(1, 7, 16), (1, 9, 20), (1, 11, 24)]) + _register_offload_request(scheduler, prefill_end_tokens=37, num_prompt_tokens=13) + out = _make_offload_only_output([(1, 7, 16), (1, 9, 32), (1, 11, 48)]) meta = scheduler.build_connector_meta(out) - # 16 and 20 are inside the resumed prefill; 24 is past it. - assert meta.requests[0].boundary_state_offloads == [(1, 7, 16), (1, 9, 20)] + # Aligned states at 16 and 32 are inside the resumed prefill; 48 is past it. + assert meta.requests[0].boundary_puts == [(1, 7, 16), (1, 9, 32)] store_job_id = meta.requests[0].store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7, 9] @@ -1267,7 +1511,7 @@ def test_boundary_state_job_pins_exact_blocks_once(): meta = scheduler.build_connector_meta(out) req_meta = meta.requests[0] - assert req_meta.boundary_state_offloads == [ + assert req_meta.boundary_puts == [ (1, 7, 16), (1, 8, 32), (1, 9, 48), @@ -1294,6 +1538,7 @@ def test_store_job_pins_current_non_null_non_mamba_blocks(): request = SimpleNamespace( all_token_ids=list(range(48)), block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=48, num_output_placeholders=0, ) scheduler._unfinished_requests["req-0"] = (request, ([7, 2, 0], [21, 0, 22])) @@ -1379,7 +1624,7 @@ def test_preemption_and_request_id_reuse_do_not_release_inflight_job(): _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64) second = scheduler.build_connector_meta(_make_offload_only_output([(1, 8, 32)])) second_job_id = second.requests[0].store_job_id - assert second.requests[0].boundary_state_offloads == [(1, 8, 32)] + assert second.requests[0].boundary_puts == [(1, 8, 32)] scheduler.update_connector_output(_make_worker_output({first_job_id: 1})) assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0 @@ -1412,9 +1657,9 @@ def test_resumed_partial_tail_uses_exact_boundary(): meta = scheduler.build_connector_meta(out) assert len(meta.requests) == 1 - assert meta.requests[0].boundary_state_offloads == [(1, 7, 12)] - # Ordinary metadata retains the full resumed prefill range. - assert meta.requests[0].num_prompt_tokens == 20 + assert meta.requests[0].boundary_puts == [(1, 7, 12)] + assert meta.requests[0].num_prompt_tokens == 13 + assert meta.requests[0].prefill_end_tokens == 20 tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 assert tracker.has_pending_offload is True @@ -1427,7 +1672,7 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): all_token_ids=list(range(48)), block_hashes=[b"h0", b"h1", b"h2"], num_output_placeholders=0, - num_prompt_tokens=36, + num_prompt_tokens=37, ) scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1445,8 +1690,9 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): assert len(meta.requests) == 1 assert meta.requests[0].can_save is True - assert meta.requests[0].boundary_state_offloads == [(0, 7, 36)] - assert meta.requests[0].num_prompt_tokens == 48 + assert meta.requests[0].boundary_puts == [(0, 7, 36)] + assert meta.requests[0].num_prompt_tokens == 37 + assert meta.requests[0].prefill_end_tokens == 48 # Ordinary saving still covers the full resumed prefill range. tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 48 @@ -1460,7 +1706,7 @@ def test_partial_tail_cow_block_is_referenced_for_the_job(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index cb6bd8d995be..09d27eb947a0 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -30,6 +30,7 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( BlobBlockHashes, + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LBHNCStoreLayout, @@ -44,6 +45,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( @@ -942,6 +946,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def _make_partial_tail_send_thread( store, *, @@ -954,6 +967,8 @@ def _make_partial_tail_send_thread( hash_block_size=4, lcm_block_size=16, mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, ) db = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0), @@ -987,8 +1002,63 @@ def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta: block_ids=(block_ids, [1]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, + ) + + +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_partial_tail_ignores_handoff_off_prompt_boundary(use_eagle): + """A Mamba hand-off away from the prompt checkpoint (a shared-prefix + junction) gets no attention tail; only its own Mamba put remains.""" + store = MagicMock() + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + req = _make_partial_tail_req([1, 2, 3]) + req.num_prompt_tokens = 17 if use_eagle else 13 + req.boundary_puts = [(1, 7, 8)] + + assert _resolve_partial_tail(thread, req).boundary_puts == [(1, 7, 8)] + + +def test_eagle_attention_proof_published_after_checkpoint_handoff(): + store = MagicMock() + stored = set() + store.batch_is_exist.side_effect = lambda keys: [int(k in stored) for k in keys] + + def put(keys, *args): + stored.update(keys) + return [256] * len(keys) + + store.batch_put_from_multi_buffers.side_effect = put + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = True + thread.coord.eagle_proof_margin_by_group = {0: 4} + metadata = _make_partial_tail_req([1, 2, 3]) + metadata.num_prompt_tokens = 13 + metadata.completed_token_len = 8 + metadata.boundary_puts = [(1, 7, 8)] + + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) ) + mamba_key = thread.token_databases[1].key_for(b"a1") + attention_key = thread.token_databases[0].key_for(b"a2") + assert mamba_key in stored + assert attention_key not in stored + + metadata.completed_token_len = 13 + metadata.boundary_puts = None + metadata.publish_partial_tail = True + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) + ) + assert attention_key in stored + keys, addrs, *_ = store.batch_put_from_multi_buffers.call_args.args + assert addrs[keys.index(attention_key)] == [0x1000 + 3 * 256] + assert mamba_key not in keys def test_partial_tail_offload_skips_null_source_blocks(): @@ -997,7 +1067,9 @@ def test_partial_tail_offload_skips_null_source_blocks(): store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([0, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([0, 2, 3])) + ) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1131,9 +1203,11 @@ def make_db(group_id: int, block_size: int, base_addr: int): block_hashes=[b"a0", b"a1", b"a2"], can_save=True, # The scheduler accepted group 1 and omitted group 2 at boundary 12. - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1156,7 +1230,9 @@ def test_partial_tail_offload_replaces_stale_group_ids_after_filtering(): supports_group_ids=True, ) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([1, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1190,13 +1266,17 @@ def test_partial_tail_put_failure_activates_pressure_gate(): store.batch_put_from_multi_buffers.return_value = [256, -200, 256] thread = _make_partial_tail_send_thread(store) - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert thread._store_pressure_active is True assert thread._skip_store_requests == {"req-a"} assert thread._saved_offset.get("req-a", 0) == 0 - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert store.batch_put_from_multi_buffers.call_count == 1 @@ -1277,9 +1357,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): block_ids=([1, 2, 3], [5]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 7, 32)], + boundary_puts=[(1, 7, 32)], ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args # boundary 32 is block-aligned for the mamba group (block 16): one key, @@ -1289,7 +1369,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): assert addrs == [[0x2000 + 7 * 256]] -def test_mixed_snapshot_and_sub_block_offloads(): +@pytest.mark.parametrize("saved_tokens", [0, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_mixed_snapshot_and_sub_block_offloads(saved_tokens, use_eagle): """A retention snapshot and the prompt-end sub-block CoW tail can arrive in one hand-off; the sub-block path covers FA gap blocks but reads the mamba boundary only from the CoW block, never positionally.""" @@ -1297,39 +1379,150 @@ def test_mixed_snapshot_and_sub_block_offloads(): store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + thread._saved_offset["req-a"] = saved_tokens - hs = [bytes([i + 1]) * 4 for i in range(11)] # 11 hash units = 44 tokens + hs = [bytes([i + 1]) * 4 for i in range(12)] req = ReqMeta( req_id="req-a", token_len_chunk=0, - block_ids=([1, 2, 3], [0, 0, 0]), + block_ids=(list(range(1, 14)), [0, 0, 0]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 9, 32), (1, 7, 44)], + boundary_puts=[(1, 9, 32), (1, 7, 44)], + num_prompt_tokens=49 if use_eagle else 45, + completed_token_len=49 if use_eagle else 45, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args db_full, db_mamba = thread.token_databases + proof_end = 48 if use_eagle else 44 assert keys == [ + # The tail reads only the LCM gap, even when normal saves lag. + *[db_full.key_for(BlockHash(hs[i])) for i in range(8, proof_end // 4)], # Aligned snapshot: boundary 32 from the handed-off block 9. db_mamba.key_for(BlockHash(hs[7])), - # Sub-block boundary 44: FA gap blocks ending at 4, 8, 12. - db_full.key_for(BlockHash(hs[0])), - db_full.key_for(BlockHash(hs[1])), - db_full.key_for(BlockHash(hs[2])), # Mamba boundary block from the CoW hand-off (block 7). db_mamba.key_for(BlockHash(hs[10])), ] assert addrs == [ + *[[0x1000 + (i + 1) * 256] for i in range(8, proof_end // 4)], [0x2000 + 9 * 256], - [0x1000 + 1 * 256], - [0x1000 + 2 * 256], - [0x1000 + 3 * 256], [0x2000 + 7 * 256], ] +@pytest.mark.parametrize("tp_rank", [0, 1]) +def test_partial_tail_with_smaller_mamba_blocks_writes_one_mamba_key(tp_rank): + """Under DCP2 the attention block spans the LCM (16) and Mamba blocks are + half that. The prompt-completing save writes one attention block and one + Mamba key: the boundary state from the CoW block, never the interior + Mamba blocks.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=16, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=8, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba], tp_rank=tp_rank + ) + # Mamba is replicated across two ranks, which split its keys by block. + thread.group_put_steps = [1, 2] + + hs = [bytes([i + 1]) * 4 for i in range(11)] + req = ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], list(range(20, 26))), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 44)], + num_prompt_tokens=45, + completed_token_len=45, + publish_partial_tail=True, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(0, 3, 44), (1, 7, 44)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + # Mamba block cdiv(44, 8) - 1 = 5 belongs to put_step_rank 5 % 2 = 1, + # which is tp_rank 0 for group 1. + mamba_puts = [db_mamba.key_for(BlockHash(hs[10]))] if tp_rank == 0 else [] + assert keys == [db_full.key_for(BlockHash(hs[10])), *mamba_puts] + assert addrs == [[0x1000 + 3 * 256], *([[0x2000 + 7 * 256]] * len(mamba_puts))] + + +def test_shared_prefix_junction_handoff_writes_only_its_mamba_key(): + """A junction hand-off (``--enable-mamba-shared-prefix-checkpoint``) below + the prompt checkpoint stores its Mamba state under the junction hash and + no attention tail.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=True, + eagle_proof_margin_by_group={0: 4}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=4, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=16, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba] + ) + + hs = [bytes([i + 1]) * 4 for i in range(18)] + # 75-token prompt (checkpoint 68), chunk cut at the junction 40. + req = ReqMeta( + req_id="req-b", + token_len_chunk=0, + block_ids=(list(range(1, 11)), [20, 21, 22]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 40)], + num_prompt_tokens=75, + completed_token_len=40, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(1, 7, 40)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + assert keys == [db_mamba.key_for(BlockHash(hs[9]))] + assert addrs == [[0x2000 + 7 * 256]] + + def test_snapshot_offload_skips_null_handoff_block(): """A hand-off the core could not materialize (null block) carries no committed state; persisting it would poison the boundary key.""" @@ -1338,13 +1531,16 @@ def test_snapshot_offload_skips_null_handoff_block(): hs = [bytes([i + 1]) * 4 for i in range(8)] assert thread._maybe_offload_boundary_states( - ReqMeta( - req_id="req-a", - token_len_chunk=0, - block_ids=([1, 2, 3], [5]), - block_hashes=hs, - can_save=True, - boundary_state_offloads=[(1, NULL_BLOCK_ID, 32)], + _resolve_partial_tail( + thread, + ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], [5]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, NULL_BLOCK_ID, 32)], + ), ) ) store.batch_is_exist.assert_not_called() @@ -1398,6 +1594,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): [True, False], ), ) + coord.store_mask = MagicMock(side_effect=coord.store_mask) db_full = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), @@ -1427,6 +1624,8 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, + num_prompt_tokens=33, + prefill_end_tokens=64, ), ) @@ -1436,6 +1635,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): assert full_hashes == [b"a2".hex(), b"a3".hex()] assert masked_hashes == [b"a2".hex()] + coord.store_mask.assert_called_once_with(64, 32, num_prompt_tokens=64) def test_store_sending_thread_prepares_missing_chunks_once_per_group(): @@ -3580,6 +3780,139 @@ def _make_bare_worker( return worker +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("prefix_match_unit", [None, 8]) +@pytest.mark.parametrize("dcp_world_size,mamba_block_size", [(1, 32), (8, 16)]) +def test_mooncake_lookup_reuses_resolved_hash_checkpoint( + use_eagle, prefix_match_unit, dcp_world_size, mamba_block_size +): + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( + coordinator as mooncake_coordinator, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=use_eagle, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=prefix_match_unit, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_world_size, world_size=dcp_world_size + ), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=prefix_match_unit or 16, + ), + ) + assert scheduler.enable_partial_hash_hits + hash_block_size = prefix_match_unit or 16 + assert scheduler._hash_block_size == hash_block_size + coordinator = mooncake_coordinator.MooncakeStoreCoordinator( + groups, + scheduler_block_size=math.lcm(16 * dcp_world_size, mamba_block_size), + hash_block_size=hash_block_size, + enable_partial_hash_hits=True, + ) + hashes = [BlockHash(bytes([i]) * 16) for i in range(512 // hash_block_size)] + pool = mooncake_coordinator.ExternalCachedBlockPool( + hash_block_size=hash_block_size, + exists={(group, bytes(h)) for group in range(2) for h in hashes}, + ) + _, hit = coordinator.find_longest_cache_hit(hashes, 511, pool) + assert hit == 512 - hash_block_size + + +@pytest.mark.parametrize("alignment, expected", [(8, True), (32, False), (None, False)]) +def test_scheduler_follows_core_hit_alignment(alignment, expected): + """The store takes partial hits from core's resolved alignment, not from its + own groups: these Mamba specs alone would allow them, but core may still + keep hits on the block (e.g. for a sliding-window group it never sees).""" + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=32, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=32, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=32, enable_prefix_caching=True, prefix_match_unit=8 + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + cache_hit_alignment_tokens=alignment, + ), + ) + + assert scheduler.enable_partial_hash_hits is expected + assert scheduler._store_coord.enable_partial_hash_hits is expected + + def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker = _make_bare_worker() worker.tp_size = 4 @@ -3873,6 +4206,7 @@ def test_lookup_partial_tail_uses_hash_alignment(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1] @@ -3946,6 +4280,7 @@ def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes(): scheduler_block_size=16, hash_block_size=4, use_eagle=True, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(6)] @@ -4015,6 +4350,7 @@ def test_lookup_plan_recovers_tail_key_after_multi_chunk_convergence(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(12)] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index d05b553b64b9..48a36bc82f8b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -3,7 +3,8 @@ """External-store cache-hit coordinator for MooncakeStoreConnector.""" from collections.abc import Sequence -from typing import NamedTuple, cast +from dataclasses import replace +from typing import TYPE_CHECKING, NamedTuple, cast from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( chunk_hashes_for_block_size, @@ -13,12 +14,15 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + resolve_dcp_kv_cache_spec, ) from vllm.v1.core.single_type_kv_cache_manager import ( SingleTypeKVCacheManager, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, MambaSpec, @@ -26,6 +30,9 @@ ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +if TYPE_CHECKING: + from vllm.config import VllmConfig + class StoreSpecGroup(NamedTuple): spec: KVCacheSpec @@ -82,7 +89,7 @@ def __init__( hash_block_size: int, use_eagle: bool = False, retention_interval: int | None = None, - dcp_world_size: int = 1, + enable_partial_hash_hits: bool = False, ) -> None: # Mirrors core's resolve_kv_cache_block_sizes: the hash unit only has # to divide groups that participate in prefix caching. Non-shareable @@ -110,14 +117,51 @@ def __init__( } self.hash_block_size = hash_block_size self.lcm_block_size = scheduler_block_size - self.enable_partial_hash_hits = partial_hash_hits_enabled( - kv_cache_groups, hash_block_size, dcp_world_size - ) + # Core's decision: whether hits land on the hash unit, below the block. + self.enable_partial_hash_hits = enable_partial_hash_hits self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval self._verify_and_split_kv_cache_groups() + @classmethod + def from_kv_cache_config( + cls, + kv_cache_config: KVCacheConfig, + vllm_config: "VllmConfig", + scheduler_block_size: int, + hash_block_size: int, + ) -> "MooncakeStoreCoordinator": + """Build the coordinator the connector's scheduler and workers share. + + Specs are DCP-resolved, so a block spans the tokens of one block-table + entry. Partial hits follow core's hit alignment, which checks every KV + cache group and its manager, not only the groups the store transfers; + unset (a config the engine core did not resolve) keeps hits on the + block. + """ + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + spec_config = vllm_config.speculative_config + alignment = kv_cache_config.cache_hit_alignment_tokens + return cls( + [ + replace( + group, + kv_cache_spec=resolve_dcp_kv_cache_spec( + group.kv_cache_spec, dcp_world_size + ), + ) + for group in kv_cache_config.prefix_cacheable_groups + ], + scheduler_block_size, + hash_block_size, + use_eagle=spec_config is not None and spec_config.use_eagle_block_drop(), + retention_interval=kv_cache_config.prefix_cache_retention_interval, + enable_partial_hash_hits=( + alignment is not None and alignment < scheduler_block_size + ), + ) + def align_lookup_length(self, length: int) -> int: alignment = ( self.hash_block_size @@ -167,6 +211,17 @@ def _verify_and_split_kv_cache_groups(self) -> None: self.eagle_group_ids = { gid for g in attention_groups if g.use_eagle for gid in g.group_ids } + self.eagle_proof_margin_by_group = { + gid: eagle_proof_margin( + group.spec.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and group.manager_cls.supports_fine_grained_hash_lookup, + ) + for group in attention_groups + if group.use_eagle and not isinstance(group.spec, MambaSpec) + for gid in group.group_ids + } def find_longest_cache_hit( self, @@ -394,14 +449,16 @@ def _find_hit_blocks( # never drops a block, so a widened bound would match past the # attention-verified hit and resume from speculative state (#43559). if drop_eagle_block and not isinstance(spec, MambaSpec): - eagle_margin = ( - self.hash_block_size - if self.enable_partial_hash_hits - and manager_cls.supports_fine_grained_hash_lookup - and spec.block_size > self.hash_block_size - else spec.block_size + eagle_margin = eagle_proof_margin( + spec.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, + ) + _max_length = min( + curr_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, ) - _max_length = min(curr_hit_length + eagle_margin, max_length) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, @@ -449,20 +506,3 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: if isinstance(spec, UniformTypeKVCacheSpecs): return next(iter(spec.kv_cache_specs.values())) return spec - - -def partial_hash_hits_enabled( - kv_cache_groups: Sequence[KVCacheGroupSpec], - hash_block_size: int, - dcp_world_size: int = 1, -) -> bool: - """Match core's DCP-aware Mamba partial-hit condition.""" - return any( - isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec) - and spec.mamba_cache_mode == "align" - and ( - (dcp_world_size == 1 and spec.block_size > hash_block_size) - or (dcp_world_size > 1 and spec.block_size >= hash_block_size) - ) - for g in kv_cache_groups - ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 2510a7ee8e76..4cd0402b7b4f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -7,7 +7,7 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import cast +from typing import NamedTuple, cast import numpy as np import torch @@ -712,6 +712,7 @@ class RequestTracker: # For a fresh request this is len(prompt). For a resumed-from-preemption # request it includes previously-generated tokens, which are re-prefilled. prefill_end_tokens: int = 0 + partial_tail_sent: bool = False def reset(self) -> None: self.token_len = 0 @@ -720,6 +721,7 @@ def reset(self) -> None: self.token_ids = None self.has_pending_offload = False self.prefill_end_tokens = 0 + self.partial_tail_sent = False def update( self, @@ -739,6 +741,14 @@ def update( existing.extend(new) +class BoundaryPut(NamedTuple): + """A block stored under the hash of the prefix ending at ``num_tokens``.""" + + group_id: int + block_id: int + num_tokens: int + + @dataclass class ReqMeta: """Per-request metadata for store put/get operations.""" @@ -756,15 +766,24 @@ class ReqMeta: # Absolute request offset represented by token_ids[0]. token_ids_start: int = 0 num_prompt_tokens: int | None = None + # Total prefill tokens, including generated tokens replayed after preemption. + prefill_end_tokens: int | None = None # Identifies this store job for the engine's lifetime. A request id cannot # serve that purpose: it is reused once a preempted request resumes, so it # would release the wrong job's blocks. store_job_id: int | None = None - # Core-provided (group_id, block_id, boundary_tokens) mamba "align" - # boundary states. A block-aligned entry is a committed boundary snapshot; - # a non-aligned entry is the sub-block CoW tail. The store-job reference - # keeps each exact block alive until every worker rank finishes the job. - boundary_state_offloads: list[tuple[int, int, int]] | None = None + # Blocks stored under an explicit prefix hash rather than by position: + # core-provided mamba "align" states (a block-aligned entry is a committed + # snapshot, a non-aligned one the sub-block CoW prompt tail) and the + # scheduler-resolved non-Mamba blocks of the prompt's partial tail. The + # store-job reference keeps each block alive until every worker rank + # finishes the job. + boundary_puts: list[BoundaryPut] | None = None + # Total computed prefix length at the end of this step. + completed_token_len: int | None = None + # Set on the save that first covers the whole prompt, so the scheduler + # publishes the prompt's partial tail once rather than on every save. + publish_partial_tail: bool = False @staticmethod def from_request_tracker( @@ -773,8 +792,12 @@ def from_request_tracker( load_spec: LoadSpec | None = None, skip_save: bool | None = False, block_hashes: list[BlockHash] | None = None, + save_partial_tail: bool = False, + num_prompt_tokens: int | None = None, ) -> "ReqMeta | None": """Create ReqMeta from a RequestTracker.""" + if save_partial_tail: + assert num_prompt_tokens is not None if block_hashes is None: block_hashes = [] input_token_len = tracker.token_len @@ -783,7 +806,15 @@ def from_request_tracker( chunk_boundary = cdiv(token_ids_start + 1, block_size) * block_size num_tokens_to_save = input_token_len // block_size * block_size - skip_save = skip_save or num_tokens_to_save < chunk_boundary + publish_tail = ( + save_partial_tail + and not tracker.partial_tail_sent + and num_prompt_tokens is not None + and 0 < num_prompt_tokens <= input_token_len + ) + skip_save = skip_save or ( + num_tokens_to_save < chunk_boundary and not publish_tail + ) # A ReqMeta must never carry both a save AND a load. # The save would also be wasted work — the bytes are being looked up # in the store right now. Later cached_reqs steps save new tokens @@ -795,6 +826,7 @@ def from_request_tracker( if not skip_save: tracker.num_saved_tokens = num_tokens_to_save + tracker.partial_tail_sent |= publish_tail token_ids = None if tracker.token_ids and not skip_save: @@ -826,7 +858,10 @@ def from_request_tracker( block_hashes=block_hashes, token_ids=token_ids, token_ids_start=token_ids_start, - num_prompt_tokens=tracker.prefill_end_tokens, + num_prompt_tokens=num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + completed_token_len=input_token_len, + publish_partial_tail=publish_tail and not skip_save, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 8f277d7eea37..239dad736490 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -5,14 +5,17 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Scheduler-side logic for MooncakeStoreConnector.""" +from collections.abc import Sequence + from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 - partial_hash_hits_enabled, +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 + BoundaryPut, LoadSpec, MooncakeStoreConnectorMetadata, MooncakeStoreWorkerMetadata, @@ -23,12 +26,19 @@ LookupKeyClient, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks -from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes +from vllm.v1.core.kv_cache_utils import ( + resolve_kv_cache_block_sizes, +) from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + MambaSpec, + get_mamba_prefill_checkpoint_position, +) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request @@ -49,6 +59,79 @@ def _new_req_prefill_tokens(request: NewRequestData) -> list[int]: return request.prompt_token_ids +def _partial_tail_non_mamba_puts( + coord: MooncakeStoreCoordinator, + req_meta: ReqMeta, + block_sizes: Sequence[int], +) -> list[BoundaryPut]: + """Locate the non-Mamba blocks a partial-tail save publishes. + + A later request resumes at the prompt's Mamba checkpoint (``boundary``) only + if both pieces of KV are stored: + + - Mamba groups: the state at ``boundary``, handed off by the core as a + boundary put of its own, so Mamba groups are not listed here. + - Other groups: the KV from the last LCM-aligned normal save up to + ``boundary`` plus the group's EAGLE proof margin (``proof_end``). Each + block is keyed by the hash at its end, clipped to ``proof_end``. + + Groups whose proof is not computed yet are omitted. For example, with LCM + 16, blocks of 4 tokens, and a checkpoint at 44, the full-attention blocks + are 8-10, keyed at 36, 40 and 44. + + A Mamba state handed off at any other position, i.e. a shared-prefix + junction (``--enable-mamba-shared-prefix-checkpoint``), adds no puts here. + + ``block_sizes`` are the per-group block sizes, indexed like + ``req_meta.block_ids``. + """ + mamba_tails = [ + put.num_tokens + for put in req_meta.boundary_puts or [] + if put.group_id in coord.mamba_group_ids + and put.num_tokens % block_sizes[put.group_id] + ] + # The tail is due on the save that first covers the prompt, or when Mamba + # hands off its checkpoint state. + if not mamba_tails and not req_meta.publish_partial_tail: + return [] + prompt_tokens = req_meta.num_prompt_tokens + completed = req_meta.completed_token_len + assert prompt_tokens is not None and completed is not None + hash_block_size = coord.hash_block_size + # Core drops the Mamba checkpoint whenever EAGLE block drop is on, whichever + # groups carry the eagle flag. + boundary = get_mamba_prefill_checkpoint_position( + prompt_tokens, hash_block_size, coord.use_eagle + ) + # Any other sub-block Mamba hand-off is a shared-prefix junction. Its + # attention KV, EAGLE proof included, ends on an attention block end no later + # than this checkpoint's proof, so the normal and prompt-tail saves store it. + if not req_meta.publish_partial_tail and boundary not in mamba_tails: + return [] + num_hashes = len(req_meta.block_hashes) + if boundary == 0 or boundary // hash_block_size > num_hashes: + return [] + start = boundary // coord.lcm_block_size * coord.lcm_block_size + puts: list[BoundaryPut] = [] + for group_id, block_size in enumerate(block_sizes): + if group_id in coord.mamba_group_ids: + continue + proof_end = boundary + coord.eagle_proof_margin_by_group.get(group_id, 0) + if proof_end > completed or proof_end // hash_block_size > num_hashes: + continue + group_blocks = req_meta.block_ids[group_id] + for block_idx in range(start // block_size, cdiv(proof_end, block_size)): + if block_idx >= len(group_blocks): + break + block_id = group_blocks[block_idx] + if block_id == NULL_BLOCK_ID: + continue + num_tokens = min((block_idx + 1) * block_size, proof_end) + puts.append(BoundaryPut(group_id, block_id, num_tokens)) + return puts + + class MooncakeStoreScheduler: """Scheduler-side component for MooncakeStoreConnector.""" @@ -84,11 +167,6 @@ def __init__( self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) - self.enable_partial_hash_hits = partial_hash_hits_enabled( - store_groups, - self._hash_block_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) mamba_groups = { group_id: group.kv_cache_spec for group_id, group in enumerate(store_groups) @@ -99,6 +177,11 @@ def __init__( ), "MooncakeStoreScheduler requires mamba_cache_mode='align'" self._boundary_state_group_ids = frozenset(mamba_groups) + self._store_coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self._block_size, self._hash_block_size + ) + self.enable_partial_hash_hits = self._store_coord.enable_partial_hash_hits + self._gpu_block_pool: BlockPool | None = None self._num_workers = vllm_config.parallel_config.world_size self._next_store_job_id = 0 @@ -273,6 +356,8 @@ def build_connector_meta( # producer. Loads are still carried by the same metadata. skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: meta.add_request(req_meta) @@ -323,6 +408,8 @@ def build_connector_meta( load_spec=load_spec, skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) else: # Decode/chunked request @@ -364,6 +451,8 @@ def build_connector_meta( load_spec=None, skip_save=False, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: @@ -398,6 +487,7 @@ def build_connector_meta( load_spec=load_spec, skip_save=None, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, ) if req_meta is not None: meta.add_request(req_meta) @@ -455,6 +545,8 @@ def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: itself is freed, until every rank reports the job done. """ pool = self._gpu_block_pool + coord = self._store_coord + block_sizes = [g.kv_cache_spec.block_size for g in coord.kv_cache_groups] for req_meta in meta.requests: if not req_meta.can_save: continue @@ -464,27 +556,27 @@ def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: if req_meta.store_job_id is not None: assert req_meta.store_job_id in self._pinned_saves continue + if tail := _partial_tail_non_mamba_puts(coord, req_meta, block_sizes): + req_meta.boundary_puts = [*tail, *(req_meta.boundary_puts or [])] req_meta.store_job_id = store_job_id = self._next_store_job_id self._next_store_job_id += 1 - block_ids: list[int] = [] - if req_meta.boundary_state_offloads: - block_ids.extend( - block_id for _, block_id, _ in req_meta.boundary_state_offloads - ) + block_ids = [put.block_id for put in req_meta.boundary_puts or []] assert NULL_BLOCK_ID not in block_ids, ( - "A null block cannot back a boundary-state offload" - ) - # Every allocated block is referenced, not just the ones covering - # this job's token range: a rank resumes from its own last - # successful offset, which lags the scheduler's whenever a save was - # skipped or failed, so it may read anywhere below the range. - block_ids.extend( - block_id - for group_id, group in enumerate(req_meta.block_ids) - if group_id not in self._boundary_state_group_ids - for block_id in group - if block_id != NULL_BLOCK_ID + "A null block cannot back a boundary put" ) + if req_meta.token_len_chunk: + # Normal prefix save: pin all non-Mamba sources for retries. + # Every allocated block is referenced, not just the ones covering + # this job's token range: a rank resumes from its own last + # successful offset, which lags the scheduler's whenever a save was + # skipped or failed, so it may read anywhere below the range. + block_ids.extend( + block_id + for group_id, group in enumerate(req_meta.block_ids) + if group_id not in self._boundary_state_group_ids + for block_id in group + if block_id != NULL_BLOCK_ID + ) # An aligned boundary block may also be present in the request's # block table. Take and release exactly one reference per block. block_ids = list(dict.fromkeys(block_ids)) @@ -515,17 +607,39 @@ def register_finished_partial_tail( if boundary_tokens > tracker.prefill_end_tokens: return False - pinned_block_ids: list[int] = [] - remapped_offloads: list[tuple[int, int, int]] = [] + mamba_puts: list[BoundaryPut] = [] for group_id, block_id, boundary in partial_tail_offloads: store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: return False if block_id == NULL_BLOCK_ID: return False - pinned_block_ids.append(block_id) - remapped_offloads.append((store_group_id, block_id, boundary)) - pinned_block_ids = list(dict.fromkeys(pinned_block_ids)) + mamba_puts.append(BoundaryPut(store_group_id, block_id, boundary)) + req_meta = ReqMeta( + req_id=request.request_id, + token_len_chunk=0, + block_ids=tuple( + block_ids[group_id].copy() for group_id in self._store_group_ids + ), + block_hashes=list(request.block_hashes), + can_save=True, + num_prompt_tokens=request.num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=mamba_puts, + completed_token_len=request.num_computed_tokens, + ) + coord = self._store_coord + req_meta.boundary_puts = [ + *_partial_tail_non_mamba_puts( + coord, + req_meta, + [g.kv_cache_spec.block_size for g in coord.kv_cache_groups], + ), + *mamba_puts, + ] + pinned_block_ids = list( + dict.fromkeys(put.block_id for put in req_meta.boundary_puts) + ) pool = self._gpu_block_pool assert pool is not None, ( @@ -537,18 +651,8 @@ def register_finished_partial_tail( self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers) pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids]) - self._finished_partial_tail_metas[request.request_id] = ReqMeta( - req_id=request.request_id, - token_len_chunk=0, - block_ids=tuple( - block_ids[group_id].copy() for group_id in self._store_group_ids - ), - block_hashes=list(request.block_hashes), - can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - store_job_id=store_job_id, - boundary_state_offloads=remapped_offloads, - ) + req_meta.store_job_id = store_job_id + self._finished_partial_tail_metas[request.request_id] = req_meta tracker.has_pending_offload = True # The store job owns exact block refs, so request cleanup need not wait. return False @@ -575,7 +679,7 @@ def _handle_boundary_state_offloads( # going away, so the offload is conservatively dropped. logger.debug("Dropping boundary-state offload for request %s", req_id) continue - accepted: list[tuple[int, int, int]] = [] + accepted: list[BoundaryPut] = [] for group_id, block_id, boundary_tokens in entries: # Every other group stops saving at the end of this prefill, so # a mamba-only key past it can never complete a joint hybrid @@ -589,12 +693,12 @@ def _handle_boundary_state_offloads( store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: continue - accepted.append((store_group_id, block_id, boundary_tokens)) + accepted.append(BoundaryPut(store_group_id, block_id, boundary_tokens)) if not accepted: continue tracker.has_pending_offload = True if (req_meta := save_metas.get(req_id)) is not None: - req_meta.boundary_state_offloads = accepted + req_meta.boundary_puts = accepted continue meta.add_request( ReqMeta( @@ -603,8 +707,10 @@ def _handle_boundary_state_offloads( block_ids=tracker.allocated_block_ids, block_hashes=req_tuple[0].block_hashes, can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - boundary_state_offloads=accepted, + num_prompt_tokens=req_tuple[0].num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=accepted, + completed_token_len=tracker.token_len, ) ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 173aa324e725..004b69286ea3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -74,7 +74,6 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, maybe_convert_block_hash, - resolve_dcp_kv_cache_spec, resolve_kv_cache_block_sizes, ) from vllm.v1.kv_cache_interface import ( @@ -569,40 +568,43 @@ def _clear_store_pressure(self) -> bool: self._skip_store_requests.clear() return True - def _boundary_snapshot_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] + def _boundary_puts( + self, req_meta: ReqMeta ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for committed mamba "align" boundary-state snapshots. - - These are block-aligned boundaries, i.e. exactly what the normal save - would key — but ``store_mask`` masks mamba groups out of it entirely, so - this is their *only* writer. The exclusion is not an optimization: the - normal save resolves a chunk's address as - ``req_meta.block_ids[g][start // block_size]``, and ``block_ids`` is the - connector's append-only mirror of the core's per-group table. An - align-mode table is mutated in place (a superseded state block is freed - and nulled; speculative blocks relocate), and the connector is never - told, so a stale mirror entry is indistinguishable from a live one — a - retry of a failed or pressure-skipped chunk would read a block that now - belongs to another request. - - Each entry's handed-off block *is* the boundary state and is pinned by - the core, so it is uploaded under its boundary-end hash key and never - resolved positionally. + """Puts for ``req_meta.boundary_puts``: each block under the hash of + the prefix ending at its ``num_tokens``. + + The scheduler resolves and pins every block, so none is looked up + positionally here. That matters for mamba "align" states: ``store_mask`` + masks mamba groups out of the normal save, which resolves a chunk's + address as ``req_meta.block_ids[g][start // block_size]``. An align-mode + table is mutated in place (a superseded state block is freed and + nulled; speculative blocks relocate) without telling the connector, so + a retried positional read could hit a block that now belongs to another + request. The handed-off block *is* the state. """ hash_block_size = self.coord.hash_block_size puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for group_id, block_id, boundary in entries: - if boundary == 0 or block_id == NULL_BLOCK_ID: + for group_id, block_id, num_tokens in req_meta.boundary_puts or []: + if ( + not self.group_participates[group_id] + or num_tokens == 0 + or block_id == NULL_BLOCK_ID + ): continue - hash_idx = boundary // hash_block_size - 1 + # A negative index would silently key the block by the last hash. + assert num_tokens % hash_block_size == 0, ( + f"Boundary put at {num_tokens} tokens is not a multiple of the " + f"hash block size {hash_block_size}" + ) + hash_idx = num_tokens // hash_block_size - 1 if hash_idx >= len(req_meta.block_hashes): continue db = self.token_databases[group_id] # Distribute across ranks by the same rule as normal chunks. put_step = self.group_put_steps[group_id] put_step_rank = (self.tp_rank + group_id) % put_step - if (boundary // db.block_size - 1) % put_step != put_step_rank: + if (cdiv(num_tokens, db.block_size) - 1) % put_step != put_step_rank: continue addr, size = db.prepare_value_for_block(block_id) puts.append( @@ -610,115 +612,20 @@ def _boundary_snapshot_puts( ) return puts - def _sub_block_tail_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] - ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for the request's sub-block partial tail (its last prompt hash - boundary), so a later request can hit the sub-block prefix. - - Covers every group's blocks from the normal save's lcm floor to the - boundary: the normal save floors to ``lcm_block_size``, so a - smaller-block group's full blocks in that gap are never persisted - elsewhere, and the consumer's lookup needs every group at every probed - boundary. Full blocks are keyed by their block-end hash and the partial - boundary block by the boundary sub-hash; a mamba "align" group - contributes only its boundary block, from the core-provided CoW block. - """ - boundaries = {boundary for _, _, boundary in entries} - if len(boundaries) != 1: - raise ValueError( - "Sub-block partial-tail offloads for one request must share a boundary" - ) - boundary = boundaries.pop() - hash_block_size = self.coord.hash_block_size - if boundary == 0 or boundary // hash_block_size - 1 >= len( - req_meta.block_hashes - ): - return [] - - mamba_offloads = {group_id: block_id for group_id, block_id, _ in entries} - saved = self._saved_offset.get(req_meta.req_id, 0) - puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for g_idx, db in enumerate(self.token_databases): - if not self.group_participates[g_idx]: - continue - group_blocks = req_meta.block_ids[g_idx] - # Distribute across ranks by the same rule as normal chunks. - put_step = self.group_put_steps[g_idx] - put_step_rank = (self.tp_rank + g_idx) % put_step - # Always include the boundary block: its sub-hash key is written - # only here, even if normal saves already advanced past it. - last_block = cdiv(boundary, db.block_size) - 1 - for block_idx in range( - min(saved // db.block_size, last_block), last_block + 1 - ): - if block_idx % put_step != put_step_rank: - continue - valid_end = min((block_idx + 1) * db.block_size, boundary) - key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1] - if g_idx in mamba_offloads: - if valid_end != boundary: - # Interior align-mode state positions are null or - # stale (the block table is not append-only) and never - # valid gap content; only the boundary block is - # persisted, from the core-provided hand-off. - continue - block_id = mamba_offloads[g_idx] - elif g_idx in self.coord.mamba_group_ids: - continue - elif block_idx < len(group_blocks): - block_id = group_blocks[block_idx] - else: - continue - if block_id == NULL_BLOCK_ID: - logger.debug( - "Skipping unavailable partial-tail source block " - "(req=%s, group=%d, block=%d)", - req_meta.req_id, - g_idx, - block_idx, - ) - continue - addr, size = db.prepare_value_for_block(block_id) - puts.append((db.key_for(key_hash), addr, size, db.metadata)) - return puts - def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool: - """Persist connector-pinned mamba "align" boundary states handed off - for this request, deduped against the store. - - This is every mamba key the connector writes — ``store_mask`` excludes - mamba groups from the positional normal save, aligned boundaries - included (see :meth:`_boundary_snapshot_puts`). + """Persist this job's boundary puts, deduped against the store. - The two entry kinds are keyed and sourced differently, so they are - prepared separately and put in one batch: - - - block-aligned for its group: a committed boundary-state snapshot, - the handed-off block itself; - - not block-aligned: the sub-block CoW partial tail, which also has to - cover the other groups' blocks in the normal save's lcm gap. + These are every mamba key the connector writes (``store_mask`` excludes + mamba groups from the positional normal save) plus the other groups' + partial-tail blocks in the normal save's lcm gap. Returns: True when no put is needed or every put succeeds, False otherwise. """ - offloads = req_meta.boundary_state_offloads - if not offloads or not req_meta.block_hashes: + if not req_meta.boundary_puts or not req_meta.block_hashes: return True - - snapshots: list[tuple[int, int, int]] = [] - sub_block: list[tuple[int, int, int]] = [] - for group_id, block_id, boundary in offloads: - entry = (group_id, block_id, boundary) - if boundary % self.token_databases[group_id].block_size == 0: - snapshots.append(entry) - else: - sub_block.append(entry) - - puts = self._boundary_snapshot_puts(req_meta, snapshots) - if sub_block and self.coord.enable_partial_hash_hits: - puts.extend(self._sub_block_tail_puts(req_meta, sub_block)) + puts = self._boundary_puts(req_meta) if not puts: return True @@ -854,9 +761,7 @@ def _handle_request(self, req_meta: ReqMeta): # Offload the handed-off mamba boundary states (independent of the # normal positional save, which may be skipped this step). - if req_meta.boundary_state_offloads is not None and not ( - self._maybe_offload_boundary_states(req_meta) - ): + if not self._maybe_offload_boundary_states(req_meta): return if token_len == 0: @@ -870,7 +775,7 @@ def _handle_request(self, req_meta: ReqMeta): store_masks = self.coord.store_mask( token_len, save_start, - num_prompt_tokens=req_meta.num_prompt_tokens, + num_prompt_tokens=req_meta.prefill_end_tokens, ) starts: list[int] = [] @@ -1603,31 +1508,10 @@ def __init__( ) return - self._kv_cache_groups = [ - dataclasses.replace( - group, - kv_cache_spec=resolve_dcp_kv_cache_spec( - group.kv_cache_spec, - self.dcp_size, - ), - ) - for group in kv_cache_config.prefix_cacheable_groups - ] - spec_cfg = getattr(vllm_config, "speculative_config", None) - use_eagle_block_drop = bool( - spec_cfg.use_eagle_block_drop() - if spec_cfg is not None - and callable(getattr(spec_cfg, "use_eagle_block_drop", None)) - else False - ) - self.coord = MooncakeStoreCoordinator( - self._kv_cache_groups, - scheduler_block_size=self.block_size, - hash_block_size=self.hash_block_size, - use_eagle=use_eagle_block_drop, - retention_interval=kv_cache_config.prefix_cache_retention_interval, - dcp_world_size=self.dcp_size, + self.coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self.block_size, self.hash_block_size ) + self._kv_cache_groups = self.coord.kv_cache_groups self.store_tp_size, store_namespace, store_layout_cls = ( self._select_store_layout(extra_config) ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d0dadae27e9b..b13764abec3a 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -362,7 +362,10 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: scheduler_block_size, hash_block_size, ) - for kv_cache_config in kv_cache_configs: + # The scheduler's config is a deep copy taken before these are resolved, + # so set them there too: the KV connector's scheduler side reads the same + # fields as its workers. + for kv_cache_config in (*kv_cache_configs, scheduler_kv_cache_config): kv_cache_config.hash_block_size = hash_block_size kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens From 9f1c78388673f90d5b2be3ef77242143e7f7a3f6 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 9 Oct 2026 19:26:53 -0500 Subject: [PATCH 324/952] [Test] Allow one fp8 step on K bytes in the kpool decode fuzz test (#60725) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- .../test_kpool_decode_update_batched.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 1f5693bad92a..4accf527f3d1 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -220,12 +220,19 @@ def _torch_reference( return kv_out, tail_cpu.to(torch.bfloat16).to(device="cuda") -def _assert_eq(r_ref, r_kern): +def _assert_eq(r_ref, r_kern, k_byte_tol=0): kv_ref, tail_ref = r_ref kv_kern, tail_kern = r_kern - assert torch.equal(kv_ref, kv_kern), ( - "kv_cache differs: max diff " - f"{(kv_ref.int() - kv_kern.int()).abs().max().item()}" + # Each page holds the fp8 K bytes followed by the fp32 scales. + k_region = PAGE_SIZE * HEAD_DIM + pages_ref = kv_ref.view(kv_ref.shape[0], -1) + pages_kern = kv_kern.view(kv_kern.shape[0], -1) + k_diff = (pages_ref[:, :k_region].int() - pages_kern[:, :k_region].int()).abs() + assert k_diff.max().item() <= k_byte_tol, ( + f"kv_cache K differs: max diff {k_diff.max().item()}" + ) + assert torch.equal(pages_ref[:, k_region:], pages_kern[:, k_region:]), ( + "kv_cache scales differ" ) assert torch.equal(tail_ref, tail_kern), ( "tail_kv_cache differs: max diff " @@ -652,4 +659,6 @@ def test_batched_matches_reference_fuzz(seed): r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) - _assert_eq(r_ref, r_kern) + # Random inputs can land the pooled value on a bf16 rounding midpoint, + # where the kernel and the CPU reference may round apart by one fp8 step. + _assert_eq(r_ref, r_kern, k_byte_tol=1) From 7d1d8660d620149bc4eadf795cbb80a1f884ddc2 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 17:58:09 -0700 Subject: [PATCH 325/952] [MRV2] Miscellaneous code structure updates (#60754) Signed-off-by: Nick Hill --- tests/v1/spec_decode/test_ngram_gpu.py | 2 +- tests/v1/worker/test_gpu_kv_connector.py | 12 ++-- tests/v1/worker/test_gpu_model_runner_v2.py | 10 +-- .../worker/test_gpu_model_runner_v2_eplb.py | 2 +- tests/v1/worker/test_gpu_ubatch_slicing.py | 10 +-- vllm/compilation/backends.py | 4 +- vllm/v1/worker/gpu/dp_utils.py | 38 +++++------ vllm/v1/worker/gpu/kv_connector.py | 39 +++++++++++- vllm/v1/worker/gpu/model_runner.py | 63 ++++++------------- .../gpu/model_states/encoder_decoder.py | 2 +- .../gpu/spec_decode/dflash/speculator.py | 10 +-- .../gpu/spec_decode/extract_hidden_states.py | 6 +- .../multi_module_mtp/speculator.py | 4 +- .../gpu/spec_decode/ngram/speculator.py | 2 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 10 +-- .../spec_decode/standalone_ar/speculator.py | 4 +- .../target_dependent_ar/speculator.py | 12 ++-- vllm/v1/worker/gpu/spec_decode/utils.py | 19 ++++++ 18 files changed, 140 insertions(+), 109 deletions(-) diff --git a/tests/v1/spec_decode/test_ngram_gpu.py b/tests/v1/spec_decode/test_ngram_gpu.py index 2359f0cdef2e..0170b58ad104 100644 --- a/tests/v1/spec_decode/test_ngram_gpu.py +++ b/tests/v1/spec_decode/test_ngram_gpu.py @@ -147,7 +147,7 @@ def _propose( next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE), temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE), seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE), - dp_sync=None, + dp_sync_state=None, ) return drafts.cpu().tolist() diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index c442c6fc59e2..78fe3c185bd9 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -63,13 +63,11 @@ def test_load_start_phase( request_indices = torch.tensor([3, 1]) request_ids = ["first", "second"] - attn_metadata = {"layer": object()} - connector.pre_forward( - output, - request_state_indices=request_indices, - request_ids=request_ids, - attn_metadata=attn_metadata, + input_batch = SimpleNamespace( + idx_mapping=request_indices, req_ids=request_ids, num_tokens=5 ) + attn_metadata = {"layer": object()} + connector.pre_forward(output, input_batch, attn_metadata=attn_metadata) assert events == ( ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"] ) @@ -80,6 +78,8 @@ def test_load_start_phase( kwargs = connector.kv_connector.start_load_kv.call_args.kwargs assert kwargs["request_state_indices"] is request_indices assert kwargs["request_ids"] is request_ids + assert kwargs["num_tokens"] == 5 + assert kwargs["scheduler_output"] is output assert kwargs["attn_metadata"] is attn_metadata # A subsequent step without a forward must not reuse the prior batch. diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 47173da08155..85c3ddfce425 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner +from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(): @@ -43,7 +44,7 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(): slot_mappings_by_layer=None, hidden_states=None, aux_hidden_states=None, - dp_sync=None, + dp_sync_state=None, finished_req_ids=set(), ec_connector_output=None, cudagraph_stats=None, @@ -342,16 +343,15 @@ def test_capture_model_profile_only_skips_lock(monkeypatch): @pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"]) def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer): """Targets allocate the MTP hidden buffer only for hidden-state drafters.""" - runner = GPUModelRunner.__new__(GPUModelRunner) hidden_states = torch.zeros(4, 8) buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8) if target_buffer == "absent": - runner.model = SimpleNamespace() + model = SimpleNamespace() else: returned = buffer if target_buffer == "tensor" else None - runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) + model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned) - out = runner._get_drafter_hidden_states(hidden_states) + out = get_drafter_hidden_states(model, hidden_states) if target_buffer == "tensor": assert torch.equal(out, buffer[:4]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index d8090de2a359..6b38bfef968d 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -185,7 +185,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): slot_mappings_by_layer=None, hidden_states=None, aux_hidden_states=None, - dp_sync=None, + dp_sync_state=None, finished_req_ids=set(), ec_connector_output=None, routed_experts=None, diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index ebf334e45d4f..127d289afa44 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -350,10 +350,10 @@ def test_decode_threshold_only_applies_when_every_rank_is_a_uniform_decode(): def test_microbatching_pads_all_ranks_to_the_largest(): """Ranks must run the same token count so microbatch sizes line up.""" - desc, dp_sync = _sync_dp([200, 256]) + desc, dp_sync_state = _sync_dp([200, 256]) assert desc.num_tokens == 256 - assert dp_sync is not None - assert dp_sync.num_tokens_across_dp.tolist() == [256, 256] + assert dp_sync_state is not None + assert dp_sync_state.num_tokens_across_dp.tolist() == [256, 256] def test_microbatching_survives_a_rank_that_cannot_fill_it(): @@ -959,7 +959,7 @@ def test_microbatched_graph_needs_every_rank_to_reach_the_split(): assert desc.num_ubatches == 2 # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33. - desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager) + desc, dp_sync_state = _sync_dp([33, 128], uniform, cudagraph_manager=manager) assert desc.cg_mode == CUDAGraphMode.NONE assert desc.num_ubatches == 2 - assert dp_sync is not None and dp_sync.eager + assert dp_sync_state is not None and dp_sync_state.eager diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index d985d8de757a..434c1af47ee8 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -1091,12 +1091,12 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm # torch compile cache errors. The V2 implementation is pure Triton and # does not need the cache disabled. - is_ngram_gpu_enabled = ( + is_v1_ngram_gpu_enabled = ( vllm_config.speculative_config is not None and vllm_config.speculative_config.use_ngram_gpu() and not vllm_config.use_v2_model_runner ) - disable_cache = disable_cache or is_ngram_gpu_enabled + disable_cache = disable_cache or is_v1_ngram_gpu_enabled if disable_cache: logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix()) diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index 15db7c197c96..5c71ebcea018 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -152,7 +152,7 @@ def sync_cudagraph_and_dp_padding( num_reqs=num_reqs, ) - synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) + synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min())) # If any rank wants to run eager, all ranks run eager if synced_cg_mode == CUDAGraphMode.NONE: @@ -175,13 +175,13 @@ def sync_cudagraph_and_dp_padding( "cudagraph_manager should only be None during profile run, " "where synced_cg_mode must be NONE across all DP ranks" ) - synced_num_tokens = int(num_tokens_across_dp.max().item()) + synced_num_tokens = int(num_tokens_across_dp.max()) # Varlen decode graphs are selected by the query-length bound, so ranks must agree # on it or they pad to different token counts below. synced_max_query_len: int | None = None - if bool(torch.all(max_query_lens_across_dp != -1).item()): - synced_max_query_len = int(max_query_lens_across_dp.max().item()) + if bool(torch.all(max_query_lens_across_dp != -1)): + synced_max_query_len = int(max_query_lens_across_dp.max()) # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs # so we don't perform request padding for PIECEWISE graphs. @@ -223,12 +223,12 @@ def dispatch_cg_and_sync_dp( parallel_config: ParallelConfig | None = None, allow_ubatching: bool = False, uniform_decode: bool = False, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks. Runs a collective when dp_size > 1 so every rank dispatches to the same - shape. Pass `dp_sync` from a dispatch already made over this same batch (a + shape. Pass `dp_sync_state` from a dispatch already made over this same batch (a drafter's prefill runs the target's batch shape) to reuse that agreement instead, with no collective. @@ -238,7 +238,7 @@ def dispatch_cg_and_sync_dp( num_reqs: Requests in this rank's batch. num_tokens: Tokens in this rank's batch, already padded by the caller. uniform_token_count: Per-request token count if this rank's batch is a - uniform decode, else None. `dp_sync.uniform_token_count` takes its + uniform decode, else None. `dp_sync_state.uniform_token_count` takes its place when a sync is reused, since that one is agreed across ranks. dp_size: Data-parallel world size. 1 skips all cross-rank work. dp_rank: This rank's index in the DP group. @@ -254,7 +254,7 @@ def dispatch_cg_and_sync_dp( microbatches. Agreed across ranks before it takes effect. uniform_decode: Whether this rank's batch is a uniform decode, used to pick the microbatching split. - dp_sync: Agreement from a prior dispatch over this same batch, to reuse. + dp_sync_state: Agreement from a prior dispatch over this same batch, to reuse. Must come from a batch with this same padded `num_tokens` and the same `uniform_token_count`; `num_reqs` may differ, as neither depends on it. Passing a sync from a different batch is a caller @@ -265,7 +265,7 @@ def dispatch_cg_and_sync_dp( dispatch to reuse. It is None when `dp_size` is 1 or no rank has work. """ - reuse_eager = dp_sync is not None and dp_sync.eager + reuse_eager = dp_sync_state is not None and dp_sync_state.eager if need_eager or reuse_eager: batch_desc = BatchExecutionDescriptor( @@ -282,7 +282,7 @@ def dispatch_cg_and_sync_dp( batch_desc = cudagraph_manager.dispatch( num_reqs, num_tokens, - dp_sync.uniform_token_count if dp_sync is not None else uniform_token_count, + dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count, num_active_loras=num_active_loras, max_query_len=max_query_len, ) @@ -290,24 +290,24 @@ def dispatch_cg_and_sync_dp( if dp_size == 1: return batch_desc, None - if dp_sync is not None: - assert dp_sync.num_tokens_across_dp[dp_rank] == num_tokens, ( + if dp_sync_state is not None: + assert dp_sync_state.num_tokens_across_dp[dp_rank] == num_tokens, ( "reusing a DP sync taken over a different batch" ) assert ( - dp_sync.uniform_token_count is None - or uniform_token_count == dp_sync.uniform_token_count + dp_sync_state.uniform_token_count is None + or uniform_token_count == dp_sync_state.uniform_token_count ), "reusing a DP sync taken over a different batch" - if not dp_sync.eager and batch_desc.num_tokens != num_tokens: + if not dp_sync_state.eager and batch_desc.num_tokens != num_tokens: # Capture sizes can differ between managers, so this one may # pad further. Every rank pads alike, so report what will run. - dp_sync = replace( - dp_sync, + dp_sync_state = replace( + dp_sync_state, num_tokens_across_dp=torch.full_like( - dp_sync.num_tokens_across_dp, batch_desc.num_tokens + dp_sync_state.num_tokens_across_dp, batch_desc.num_tokens ), ) - return batch_desc, dp_sync + return batch_desc, dp_sync_state return sync_cudagraph_and_dp_padding( cudagraph_manager, diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index 1ba75beae271..e9b71dae3e8c 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -24,12 +24,18 @@ if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.worker.gpu.input_batch import InputBatch class KVConnector: """KVConnector interface used by GPUModelRunner.""" - def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + def pre_forward( + self, + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None" = None, + attn_metadata: dict[str, Any] | None = None, + ) -> None: pass def finish_forward(self) -> None: @@ -61,7 +67,12 @@ def __init__( self._pending_load_kwargs: dict[str, Any] | None = None self._disabled = False - def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None: + def pre_forward( + self, + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None" = None, + attn_metadata: dict[str, Any] | None = None, + ) -> None: if self._disabled: return @@ -69,7 +80,9 @@ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> Non assert kv_connector_metadata is not None self.kv_connector.handle_preemptions(kv_connector_metadata) self.kv_connector.bind_connector_metadata(kv_connector_metadata) - self._pending_load_kwargs = kwargs + self._pending_load_kwargs = self._build_load_kwargs( + scheduler_output, input_batch, attn_metadata + ) if scheduler_output.has_sync_kv_loads: # Sync loads need to run before this step's forward. @@ -77,6 +90,26 @@ def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> Non # Otherwise start the async load after forward to keep submission # off the critical path. + @staticmethod + def _build_load_kwargs( + scheduler_output: "SchedulerOutput", + input_batch: "InputBatch | None", + attn_metadata: dict[str, Any] | None, + ) -> dict[str, Any]: + """Keyword args for the connector's start_load_kv().""" + if input_batch is None: + # No forward this step. + return {} + load_kwargs: dict[str, Any] = dict( + scheduler_output=scheduler_output, + request_state_indices=input_batch.idx_mapping, + request_ids=input_batch.req_ids, + num_tokens=input_batch.num_tokens, + ) + if attn_metadata is not None: + load_kwargs["attn_metadata"] = attn_metadata + return load_kwargs + def _start_load_kv(self) -> None: load_kwargs = self._pending_load_kwargs assert load_kwargs is not None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 2e9925f1ad1c..f39fd0a6e6a7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -162,7 +162,10 @@ get_max_chunk_logits, ) from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator -from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler +from vllm.v1.worker.gpu.spec_decode.utils import ( + DraftTokensHandler, + get_drafter_hidden_states, +) from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import ( StructuredOutputsWorker, @@ -578,9 +581,7 @@ def main_stream(self) -> torch.cuda.Stream: def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: encoder_runner = getattr(self.model_state, "encoder_runner", None) - if encoder_runner is None: - return {} - return encoder_runner.get_encoder_timing_stats() + return encoder_runner.get_encoder_timing_stats() if encoder_runner else {} def get_kv_cache_spec(self): kv_cache_spec = get_kv_cache_spec(self.vllm_config) @@ -905,7 +906,7 @@ def _dummy_run( slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states - dp_sync = self.execute_model_state.dp_sync + dp_sync_state = self.execute_model_state.dp_sync_state self.execute_model_state = None self.step_timing.forward_end() @@ -921,7 +922,7 @@ def _dummy_run( ) mm_inputs = [], all_false - spec_hidden_states = self._get_drafter_hidden_states(hidden_states) + spec_hidden_states = get_drafter_hidden_states(self.model, hidden_states) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping @@ -943,7 +944,7 @@ def _dummy_run( next_prefill_tokens=self.req_states.next_prefill_tokens, temperature=self.sampler.sampling_states.temperature.gpu, seeds=self.sampler.sampling_states.seeds.gpu, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, dummy_run=True, skip_attn_for_dummy_run=skip_attn, mm_inputs=mm_inputs, @@ -1771,7 +1772,7 @@ def execute_model( # cross-attention cache with dynamic encoder outputs. skip_compiled = True - batch_desc, dp_sync = dispatch_cg_and_sync_dp( + batch_desc, dp_sync_state = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, num_toks, @@ -1983,13 +1984,6 @@ def execute_model( ) self.step_timing.forward_start() - connector_kwargs = dict( - scheduler_output=scheduler_output, - request_state_indices=input_batch.idx_mapping, - request_ids=input_batch.req_ids, - num_tokens=input_batch.num_tokens, - ) - # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. @@ -1997,7 +1991,7 @@ def execute_model( # because they are already copied to the CUDA graph input buffers. assert self.cudagraph_manager is not None self.kv_connector.pre_forward( - **connector_kwargs, attn_metadata=attn_metadata + scheduler_output, input_batch, attn_metadata=attn_metadata ) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: @@ -2013,16 +2007,15 @@ def execute_model( self.vllm_config, num_tokens=input_batch.num_tokens_after_padding, cudagraph_runtime_mode=batch_desc.cg_mode, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), + num_tokens_across_dp=dp_sync_state + and dp_sync_state.num_tokens_across_dp, batch_descriptor=batch_descriptor, ubatch_slices=ubatch_slices, slot_mapping=slot_mappings_by_layer, skip_compiled=skip_compiled, is_padding=input_batch.is_padding, ): - self.kv_connector.pre_forward(**connector_kwargs) + self.kv_connector.pre_forward(scheduler_output, input_batch) if ubatch_state is not None: assert self.ubatch_runner is not None model_output = self.ubatch_runner.run( @@ -2064,7 +2057,7 @@ def execute_model( slot_mappings_by_layer=slot_mappings_by_layer, hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, finished_req_ids=finished_req_ids, ec_connector_output=ec_connector_output, cudagraph_stats=cudagraph_stats, @@ -2080,24 +2073,6 @@ def execute_model( ) return None - def _get_drafter_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Hidden states fed to the drafter. - - Targets such as DeepSeek V4 expose the pre-hc_head residual through - get_mtp_target_hidden_states(). The buffer is sized at - max_num_batched_tokens and only allocated for drafters that consume - target hidden states, so None means "use the regular hidden states". - """ - get_target_hidden_states = getattr( - self.model, "get_mtp_target_hidden_states", None - ) - if get_target_hidden_states is None: - return hidden_states - target_hidden_states = get_target_hidden_states() - if target_hidden_states is None: - return hidden_states - return target_hidden_states[: hidden_states.shape[0]] - @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -2112,7 +2087,7 @@ def sample_tokens( slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states - dp_sync = self.execute_model_state.dp_sync + dp_sync_state = self.execute_model_state.dp_sync_state finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output cudagraph_stats = self.execute_model_state.cudagraph_stats @@ -2243,7 +2218,9 @@ def sample_tokens( self.speculator.observe_verification( input_batch.idx_mapping, num_sampled, num_rejected ) - spec_hidden_states = self._get_drafter_hidden_states(draft_hidden_states) + spec_hidden_states = get_drafter_hidden_states( + self.model, draft_hidden_states + ) if isinstance(self.sampler, GPUWatermarkSampler): self.speculator.prepare_watermarking( self.sampler, input_batch.idx_mapping @@ -2262,7 +2239,7 @@ def sample_tokens( self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, num_speculative_tokens=num_spec_tokens, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, mm_inputs=mm_inputs, ) if num_spec_tokens < self.num_speculative_steps: @@ -2419,7 +2396,7 @@ class ExecuteModelState(NamedTuple): slot_mappings_by_layer: dict[str, torch.Tensor] | None hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None - dp_sync: DPSyncState | None + dp_sync_state: DPSyncState | None finished_req_ids: set[str] ec_connector_output: ECConnectorOutput | None cudagraph_stats: CUDAGraphStat | None diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index bebb9cd41be2..6363dbad14d3 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -135,7 +135,7 @@ def prepare_attn( if for_capture: max_seq_len = self.max_model_len else: - max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) + max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max()) attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 9538dcb4d4bf..884de5a264be 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -355,7 +355,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -446,8 +446,10 @@ def propose( ) batch_sync, num_batch_tokens = ( - self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req) - if dp_sync is not None + self._build_uniform_batch_dp_sync( + dp_sync_state, num_reqs, self.num_query_per_req + ) + if dp_sync_state is not None else (None, num_query_tokens) ) # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs @@ -459,7 +461,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=batch_sync, + dp_sync_state=batch_sync, ) num_tokens_padded = batch_desc.num_tokens num_tokens_across_dp = ( diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index 5afbe640b2f2..e02d7ee11d70 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -98,7 +98,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -145,9 +145,7 @@ def propose( draft_attn_metadata, self.vllm_config, num_tokens=num_tokens, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), + num_tokens_across_dp=dp_sync_state and dp_sync_state.num_tokens_across_dp, cudagraph_runtime_mode=CUDAGraphMode.NONE, slot_mapping=draft_slot_mappings, is_padding=input_batch.is_padding[:num_tokens], diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 919e939c03cc..a453baf6e3fa 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -150,7 +150,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -201,7 +201,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, ) num_tokens_across_dp = ( batch_sync.num_tokens_across_dp if batch_sync is not None else None diff --git a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py index 7a61e3581a11..b74a0ec4f984 100644 --- a/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/ngram/speculator.py @@ -226,7 +226,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index cbd94ed0542d..c3e8dcbff4d1 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -92,7 +92,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -525,18 +525,18 @@ def _copy_request_inputs( def _build_uniform_batch_dp_sync( self, - target_dp_sync: DPSyncState, + target_dp_sync_state: DPSyncState, num_reqs: int, num_query_per_req: int = 1, ) -> tuple[DPSyncState, int]: - num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req + num_batch_tokens = target_dp_sync_state.num_reqs * num_query_per_req assert num_reqs * num_query_per_req <= num_batch_tokens, ( "reusing a DP sync that does not cover this batch's requests" ) return replace( - target_dp_sync, + target_dp_sync_state, num_tokens_across_dp=torch.full_like( - target_dp_sync.num_tokens_across_dp, num_batch_tokens + target_dp_sync_state.num_tokens_across_dp, num_batch_tokens ), uniform_token_count=num_query_per_req, eager=False, diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 898de1e3f886..04e79debb869 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -350,7 +350,7 @@ def propose( next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -362,7 +362,7 @@ def propose( num_reqs = input_batch.num_reqs skip_attn = dummy_run and skip_attn_for_dummy_run num_tokens_across_dp = ( - dp_sync.num_tokens_across_dp if dp_sync is not None else None + dp_sync_state.num_tokens_across_dp if dp_sync_state is not None else None ) # Copy per-request temperature/seeds/idx_mapping into pre-allocated diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 4894ce77e962..02e28875b823 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -219,7 +219,7 @@ def propose( temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - dp_sync: DPSyncState | None = None, + dp_sync_state: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -299,7 +299,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=dp_sync, + dp_sync_state=dp_sync_state, ) num_tokens_across_dp = ( prefill_batch_sync.num_tokens_across_dp @@ -352,8 +352,10 @@ def propose( ) decode_batch_sync, num_batch_tokens = ( - self._build_uniform_batch_dp_sync(dp_sync, num_reqs, num_query_per_req=1) - if dp_sync is not None + self._build_uniform_batch_dp_sync( + dp_sync_state, num_reqs, num_query_per_req=1 + ) + if dp_sync_state is not None else (None, num_reqs) ) # Each request produces exactly 1 token per draft generation step, @@ -366,7 +368,7 @@ def propose( dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, - dp_sync=decode_batch_sync, + dp_sync_state=decode_batch_sync, ) if self.decode_cudagraph_manager is not None: decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens( diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index f73acb2bdaed..b005f5c3f3ae 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -13,6 +13,25 @@ logger = init_logger(__name__) +def get_drafter_hidden_states( + model: torch.nn.Module, hidden_states: torch.Tensor +) -> torch.Tensor: + """Hidden states fed to the drafter. + + Targets such as DeepSeek V4 expose the pre-hc_head residual through + get_mtp_target_hidden_states(). The buffer is sized at + max_num_batched_tokens and only allocated for drafters that consume + target hidden states, so None means "use the regular hidden states". + """ + get_target_hidden_states = getattr(model, "get_mtp_target_hidden_states", None) + if get_target_hidden_states is None: + return hidden_states + target_hidden_states = get_target_hidden_states() + if target_hidden_states is None: + return hidden_states + return target_hidden_states[: hidden_states.shape[0]] + + def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig: """Avoid collectives that include PP ranks without a draft model.""" if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors": From 3c802af3cf6263d0ef6fb356a7884680ecf15f04 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Fri, 9 Oct 2026 19:33:22 -0600 Subject: [PATCH 326/952] [Feat][Frontend] structured decisions: DiffusionGemma read strategy (#60466) Signed-off-by: Matt Mastracci Signed-off-by: Simon Veitner Co-authored-by: Claude Opus 5.5 Co-authored-by: Simon Veitner Co-authored-by: Misha Goin --- .../online_serving/structured_decisions.md | 5 +- .../test_question_types.py | 2 +- .../structured_decisions/test_strategies.py | 37 +++- vllm/entrypoints/generate/api_router.py | 7 +- .../structured_decisions/question_types.py | 13 +- .../structured_decisions/strategies.py | 169 +++++++++++++++--- 6 files changed, 193 insertions(+), 40 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index a1f58a3769b9..a9ed5a2ee5aa 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -20,7 +20,10 @@ log probabilities. Thinking is off unless `chat_template_kwargs` turns it on. -This endpoint does not serve diffusion models yet. +Probabilities can differ between identical requests. They depend on which +requests share a batch and on whether a read hits the prefix cache. Set +`VLLM_BATCH_INVARIANT=1` for reads that repeat exactly, at a cost in +throughput. Every read of a request starts with the state, so with `--enable-prefix-caching` the state is prefilled once and each further diff --git a/tests/entrypoints/generate/structured_decisions/test_question_types.py b/tests/entrypoints/generate/structured_decisions/test_question_types.py index 3b9ad7ba60a9..26047b5d6c07 100644 --- a/tests/entrypoints/generate/structured_decisions/test_question_types.py +++ b/tests/entrypoints/generate/structured_decisions/test_question_types.py @@ -105,7 +105,7 @@ def test_noul_labels_and_prompt(): assert q.labels == ("yes", "no") assert [a.name for a in q.options] == ["yes", "no"] assert q.type.prompt(q) == ( - "Question: Reply within the hour?\nAnswer with yes or no only." + "Question: Reply within the hour?\nyes\nno\nAnswer with yes or no only." ) described = noul(criteria={"true": "same day", "false": "later"}) assert described.type.prompt(described) == ( diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 5dcb8edb3061..a63e2587521c 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -17,12 +17,16 @@ StructuredDecisionRequest, StructuredDecisionResponse, ) -from vllm.entrypoints.generate.structured_decisions.question_types import LABELS +from vllm.entrypoints.generate.structured_decisions.question_types import ( + LABELS, + StructuredDecisionError, +) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, state_text, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + DiffusionGemmaCanvasStrategy, NextTokenStrategy, ReadContext, reply_label_ids, @@ -52,6 +56,10 @@ def test_strategy_selection(): qwen = "Qwen3ForCausalLM" assert select_read_strategy(model(qwen)) is NextTokenStrategy assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + assert ( + select_read_strategy(model("DiffusionGemmaForBlockDiffusion")) + is DiffusionGemmaCanvasStrategy + ) with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): select_read_strategy(model("LlamaForCausalLM")) with pytest.raises(ValueError, match="not raw_logits"): @@ -220,3 +228,30 @@ async def test_request_metadata_reaches_every_question( assert call.kwargs["trace_headers"] == ( TRACE_HEADERS if tracing_enabled and request_headers else None ) + + +def test_canvas_read(): + strategy = DiffusionGemmaCanvasStrategy.__new__(DiffusionGemmaCanvasStrategy) + strategy.thought, strategy.end, strategy.pad = [10, 11, 12, 13], 106, 0 + strategy.width, strategy.vocab_size = 16, 1000 + strategy.max_model_len = 23 + params = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3]) + assert params.extra_args is not None + canvas = params.extra_args["diffusion_seed_canvas"] + # The label slot as noise, the end of the turn, padding. + assert canvas[1:] == [106] + [0] * 14 + assert params.max_tokens == 2 and params.logprob_token_ids == [65, 66] + again = strategy._sampling_params([65, 66], prompt_ids=[1, 2, 3]) + assert again.extra_args == params.extra_args + read_input = strategy._read_input( + {"type": "token", "prompt_token_ids": [1, 2, 3]}, [1, 2, 3] + ) + assert read_input == { + "type": "token", + "prompt_token_ids": [1, 2, 3, 10, 11, 12, 13], + } + # One token more and the thought and the canvas no longer fit. + with pytest.raises(StructuredDecisionError, match="max_model_len=23"): + strategy._read_input( + {"type": "token", "prompt_token_ids": [1, 2, 3, 4]}, [1, 2, 3, 4] + ) diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 67a256192443..18dfb87a26da 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -5,6 +5,7 @@ from fastapi import FastAPI import vllm.envs as envs +from vllm.logger import init_logger if TYPE_CHECKING: from argparse import Namespace @@ -17,6 +18,8 @@ else: RequestLogger = object +logger = init_logger(__name__) + def register_generate_api_routers(app: FastAPI): from vllm.entrypoints.openai.chat_completion.api_router import ( @@ -243,7 +246,9 @@ async def init_generate_state( default_chat_template_kwargs=default_chat_template_kwargs, ) ) - except ValueError: + except ValueError as e: + # Info, since every model without a read strategy lands here. + logger.info("/v1/systemone is disabled: %s", e) strategy = None state.serving_structured_decisions = ( ServingStructuredDecisions( diff --git a/vllm/entrypoints/generate/structured_decisions/question_types.py b/vllm/entrypoints/generate/structured_decisions/question_types.py index 0e80694b14eb..18e7a6001855 100644 --- a/vllm/entrypoints/generate/structured_decisions/question_types.py +++ b/vllm/entrypoints/generate/structured_decisions/question_types.py @@ -55,8 +55,7 @@ def answer( ``question.labels[i]`` among the labels, and the list sums to 1. ``label_mass`` is the labels' total probability over the vocabulary.""" - def reply_line(self, label: str, option: Option) -> str | None: - """The option's line in the prompt, or None to leave it out.""" + def reply_line(self, label: str, option: Option) -> str: return ( f"{label}: {option.name} - {option.description}" if option.description @@ -67,9 +66,7 @@ def prompt(self, question: Question) -> str: """The question as the model reads it, after the state.""" lines = [f"Question: {question.instructions}"] if question.instructions else [] for label, o in zip(question.labels, question.options): - line = self.reply_line(label, o) - if line is not None: - lines.append(line) + lines.append(self.reply_line(label, o)) lines.append(self.reply_instruction) return "\n".join(lines) @@ -192,8 +189,10 @@ def parse_options(self, qid: str, criteria: Any) -> list[Option]: Option("no", None if false is None else str(false)), ] - def reply_line(self, label: str, option: Option) -> str | None: - return f"{label}: {option.description}" if option.description else None + def reply_line(self, label: str, option: Option) -> str: + # A side without a description still shows its label. Without the + # label in the prompt, models answer "Yes" or "No" instead. + return f"{label}: {option.description}" if option.description else label def answer( self, question: Question, probs: list[float], label_mass: float diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 65b5a9dd7b05..f79375e42a83 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Read strategies: how a model is asked for label probabilities. -Only models in NEXT_TOKEN_ARCHITECTURES are currently supported. +Only models in READ_STRATEGIES are currently supported. """ +import hashlib +import json import math +import random from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -15,7 +18,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.label_reads import next_token_label_reads -from vllm.inputs import EngineInput +from vllm.inputs import EngineInput, tokens_input from vllm.lora.request import LoRARequest from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer @@ -80,6 +83,13 @@ async def read( """One read per question, in order.""" +def prompt_tail(tokenizer: TokenizerLike, prompt_ids: Sequence[int]) -> list[int]: + """The prompt's ids after its last added token.""" + added = set(tokenizer.get_added_vocab().values()) + start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) + return list(prompt_ids[start:]) + + def reply_label_ids( tokenizer: TokenizerLike, prompt_ids: Sequence[int], @@ -88,9 +98,7 @@ def reply_label_ids( """The prompt's ids after its last added token, and the token each of ``labels`` adds after them as the first token of the reply. Raises ValueError if a label is not one distinct token there.""" - added = set(tokenizer.get_added_vocab().values()) - start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0) - tail = list(prompt_ids[start:]) + tail = prompt_tail(tokenizer, prompt_ids) text = tokenizer.decode(tail) ids: list[int] = [] for label in labels: @@ -104,6 +112,18 @@ def reply_label_ids( return tail, ids +def type_label_ids( + tokenizer: TokenizerLike, prompt_ids: Sequence[int] +) -> dict[str, int]: + """The reply's first token for each label of every registered question + type. Raises ValueError as reply_label_ids does.""" + label_ids: dict[str, int] = {} + for qtype in QUESTION_TYPES.values(): + _, ids = reply_label_ids(tokenizer, prompt_ids, qtype.label_set) + label_ids.update(zip(qtype.label_set, ids)) + return label_ids + + class NextTokenStrategy(ReadStrategy): """Autoregressive models. Each question is one request: the state, then the question with its labeled options, and the logprobs of the label tokens as @@ -112,7 +132,18 @@ class NextTokenStrategy(ReadStrategy): def __init__(self, context: ReadContext): super().__init__(context) - tokenizer = context.online_renderer.renderer.get_tokenizer() + # Every prompt ends with the same generation prompt, so a label's token + # is the same for every question. + prompt = self._generation_prompt() + self.tail = prompt_tail(self._tokenizer(), prompt) + self.label_ids = type_label_ids(self._tokenizer(), prompt) + + def _tokenizer(self) -> TokenizerLike: + return self.context.online_renderer.renderer.get_tokenizer() + + def _generation_prompt(self) -> list[int]: + """A chat prompt's ids up to the generation prompt, with thinking off.""" + context, tokenizer = self.context, self._tokenizer() probe = tokenizer.apply_chat_template( [{"role": "user", "content": "x"}], **{ @@ -126,13 +157,20 @@ def __init__(self, context: ReadContext): ) if isinstance(probe, str): probe = tokenizer.encode(probe, add_special_tokens=False) - # Every prompt ends with the same generation prompt, so a label's token - # is the same for every question. - self.tail, _ = reply_label_ids(tokenizer, probe, ()) - self.label_ids: dict[str, int] = {} - for qtype in QUESTION_TYPES.values(): - _, ids = reply_label_ids(tokenizer, probe, qtype.label_set) - self.label_ids.update(zip(qtype.label_set, ids)) + return list(probe) + + def _sampling_params( + self, label_ids: list[int], prompt_ids: list[int] + ) -> SamplingParams: + return SamplingParams( + max_tokens=1, temperature=0.0, logprob_token_ids=label_ids + ) + + def _read_input( + self, engine_input: EngineInput, prompt_ids: list[int] + ) -> EngineInput: + """The engine input for a read of the rendered prompt.""" + return engine_input def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -180,7 +218,7 @@ async def read( ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) - slots, engine_inputs = [], [] + slots, engine_inputs, params = [], [], [] for q in questions: messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] if instructions: @@ -192,15 +230,13 @@ async def read( "labels' tokens are unknown" ) slots.append([self.label_ids[label] for label in q.labels]) - engine_inputs.append(engine_input) + engine_inputs.append(self._read_input(engine_input, prompt_ids)) + params.append(self._sampling_params(slots[-1], prompt_ids)) label_reads = await next_token_label_reads( ctx.engine_client, engine_inputs, - [ - SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids) - for ids in slots - ], + params, request_id, lora_request=lora_request, trace_headers=trace_headers, @@ -223,13 +259,87 @@ async def read( return reads -NEXT_TOKEN_ARCHITECTURES = frozenset( - { - "Qwen3ForCausalLM", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForConditionalGeneration", - } -) +class DiffusionGemmaCanvasStrategy(NextTokenStrategy): + """DiffusionGemma. Each question is one read-only request with one denoising + step.""" + + THOUGHT = "<|channel>thought\n" + END_OF_TURN = "" + CANVAS_STEP = 16 + + def __init__(self, context: ReadContext): + ReadStrategy.__init__(self, context) + tokenizer = self._tokenizer() + self.tail = prompt_tail(tokenizer, self._generation_prompt()) + self.thought = tokenizer.encode(self.THOUGHT, add_special_tokens=False) + self.label_ids = type_label_ids(tokenizer, self.thought) + end = tokenizer.encode(self.END_OF_TURN, add_special_tokens=False) + self.pad = tokenizer.pad_token_id + if self.pad is None or len(end) != 1: + raise ValueError( + "Structured decisions need the model's " + f"{self.END_OF_TURN!r} and pad tokens" + ) + self.end = end[0] + vllm_config = context.engine_client.vllm_config + served = ( + vllm_config.diffusion_config.canvas_length + if vllm_config.diffusion_config + and vllm_config.diffusion_config.canvas_length + else vllm_config.model_config.hf_config.canvas_length + ) + self.width = self.CANVAS_STEP + if self.width > served: + raise ValueError( + f"Structured decisions need a canvas of {self.width}, " + f"and the served canvas is {served}" + ) + self.vocab_size = vllm_config.model_config.get_vocab_size() + self.max_model_len = vllm_config.model_config.max_model_len + + def _sampling_params( + self, label_ids: list[int], prompt_ids: list[int] + ) -> SamplingParams: + # The noise at the label slot is drawn from the prompt, so a repeated + # read gets the same canvas. + rng = random.Random(hashlib.sha256(json.dumps(prompt_ids).encode()).digest()) + canvas = [rng.randrange(self.vocab_size), self.end] + canvas += [self.pad] * (self.width - len(canvas)) + return SamplingParams( + max_tokens=2, + logprob_token_ids=label_ids, + extra_args={ + "diffusion_seed_canvas": canvas, + "diffusion_canvas_length": self.width, + "diffusion_max_steps": 1, + "diffusion_read_only": True, + }, + ) + + def _read_input( + self, engine_input: EngineInput, prompt_ids: list[int] + ) -> EngineInput: + extra = len(self.thought) + self.width + if len(prompt_ids) + extra > self.max_model_len: + raise StructuredDecisionError( + f"the prompt has {len(prompt_ids)} tokens and a canvas read adds " + f"{extra}, more than max_model_len={self.max_model_len}" + ) + salt = engine_input.get("cache_salt") + # The thought block _must_ appear outside of the canvas (far worse results + # otherwise) + return tokens_input( + prompt_ids + self.thought, + cache_salt=salt if isinstance(salt, str) else None, + ) + + +READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "Qwen3ForCausalLM": NextTokenStrategy, + "Qwen3_5ForConditionalGeneration": NextTokenStrategy, + "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, + "DiffusionGemmaForBlockDiffusion": DiffusionGemmaCanvasStrategy, +} # label_mass sums the labels' full-vocabulary probabilities. The logprobs modes @@ -239,15 +349,16 @@ async def read( def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: """Raise ValueError when the model cannot serve structured decisions.""" - if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES: + strategy = READ_STRATEGIES.get(model_config.architecture) + if strategy is None: raise ValueError( "The structured decisions API does not support " f"{model_config.architecture}. Supported architectures: " - f"{sorted(NEXT_TOKEN_ARCHITECTURES)}" + f"{sorted(READ_STRATEGIES)}" ) if model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" ) - return NextTokenStrategy + return strategy From 193922d6ff5f7e213eacadd3cb7dd96612b30ed3 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:53:18 -0400 Subject: [PATCH 327/952] [ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X (#60397) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_basic.yaml | 20 +- tests/models/test_dsv41_mono_decode.py | 317 ++++ tests/models/test_dsv41_mono_layers.py | 335 ++++ tests/models/test_dsv41_mono_numerics.py | 529 ++++++ vllm/envs.py | 7 + vllm/models/deepseek_v41/amd/model.py | 17 + vllm/models/deepseek_v41/amd/mono/README.md | 94 + vllm/models/deepseek_v41/amd/mono/__init__.py | 5 + .../amd/mono/attention/__init__.py | 4 + .../deepseek_v41/amd/mono/attention/back.py | 884 +++++++++ .../deepseek_v41/amd/mono/attention/device.py | 504 +++++ .../deepseek_v41/amd/mono/attention/front.py | 662 +++++++ .../deepseek_v41/amd/mono/attention/gemv.py | 234 +++ .../deepseek_v41/amd/mono/attention/plan.py | 173 ++ .../deepseek_v41/amd/mono/common/__init__.py | 9 + .../models/deepseek_v41/amd/mono/common/mx.py | 78 + .../deepseek_v41/amd/mono/common/ops.py | 300 +++ .../amd/mono/common/peer_memory.py | 104 ++ .../deepseek_v41/amd/mono/common/plan.py | 135 ++ .../deepseek_v41/amd/mono/common/sync.py | 177 ++ vllm/models/deepseek_v41/amd/mono/layer.py | 940 ++++++++++ vllm/models/deepseek_v41/amd/mono/runner.py | 437 +++++ vllm/models/deepseek_v41/amd/mono/sources.py | 8 + .../deepseek_v41/amd/mono/stages/__init__.py | 4 + .../deepseek_v41/amd/mono/stages/dims.py | 60 + .../deepseek_v41/amd/mono/stages/gemv.py | 294 +++ .../deepseek_v41/amd/mono/stages/moe.py | 1635 +++++++++++++++++ .../deepseek_v41/amd/mono/stages/moe_plan.py | 54 + .../deepseek_v41/amd/mono/stages/moe_shape.py | 229 +++ .../deepseek_v41/amd/mono/stages/router.py | 86 + .../deepseek_v41/amd/mono/stages/seam.py | 406 ++++ vllm/models/deepseek_v41/amd/mono_decode.py | 376 ++++ vllm/models/deepseek_v41/amd/rocm.py | 4 +- 33 files changed, 9119 insertions(+), 2 deletions(-) create mode 100644 tests/models/test_dsv41_mono_decode.py create mode 100644 tests/models/test_dsv41_mono_layers.py create mode 100644 tests/models/test_dsv41_mono_numerics.py create mode 100644 vllm/models/deepseek_v41/amd/mono/README.md create mode 100644 vllm/models/deepseek_v41/amd/mono/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/back.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/device.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/front.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/attention/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/mx.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/ops.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/peer_memory.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/common/sync.py create mode 100644 vllm/models/deepseek_v41/amd/mono/layer.py create mode 100644 vllm/models/deepseek_v41/amd/mono/runner.py create mode 100644 vllm/models/deepseek_v41/amd/mono/sources.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/__init__.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/dims.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/gemv.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/router.py create mode 100644 vllm/models/deepseek_v41/amd/mono/stages/seam.py create mode 100644 vllm/models/deepseek_v41/amd/mono_decode.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 1d9a7fe2f3b1..483a9aceac22 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -91,8 +91,9 @@ steps: - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" @@ -259,3 +260,20 @@ steps: device: cpu-small commands: - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py + +- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' + key: dsv41-mono-decode-numerics-amd + dind: false + device: mi355_4 + num_devices: 4 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/deepseek_v41/amd/ + - tests/models/test_dsv41_mono_numerics.py + - tests/models/test_dsv41_mono_layers.py + commands: + - pytest -v -s tests/models/test_dsv41_mono_numerics.py + - pytest -v -s tests/models/test_dsv41_mono_layers.py diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py new file mode 100644 index 000000000000..684e72c00482 --- /dev/null +++ b/tests/models/test_dsv41_mono_decode.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode layer runs only where its kernels are exact: +decode-only steps of at most MAX_ROWS rows with causal SWA windows, eager or in +a FULL graph. A PIECEWISE capture is replayed for mixed batches, so it must +never record the mono launches. Layers outside the whole-layer kernels run the +FFN launch on the same steps, from wo_b's unreduced output.""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.models.deepseek_v41.amd import mono_decode as md + +M = 6 + + +class _Runner: + def __init__(self): + self.calls = [] + + def forward(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return ("mono outputs",) + + def ffn(self, *args): + self.calls.append((args, {})) + return ("ffn outputs",) + + +def _layer(comp="comp"): + cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), + compressed_cache_prefix=comp, + topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), + compress_ratio=1, + _compressed_kv_cache=cache, + ) + return SimpleNamespace(attn=attn) + + +def _swa(**overrides): + fields = dict( + num_prefills=0, + num_decodes=1, + num_decode_tokens=M, + decode_swa_width=md.SWA_WIDTH, + decode_swa_indices=torch.zeros(64, 1, md.SWA_WIDTH, dtype=torch.int32), + decode_swa_lens=torch.zeros(64, dtype=torch.int32), + token_to_req_indices=torch.zeros(64, dtype=torch.int32), + slot_mapping=torch.zeros(64, dtype=torch.int64), + block_size=32, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _inputs(rows=M): + return dict( + x=torch.zeros(rows, 5120, dtype=torch.bfloat16), + positions=torch.zeros(rows, dtype=torch.int64), + residual=torch.zeros(rows, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.zeros(rows, 4, 1), + res_mix=torch.zeros(rows, 4, 4), + pre_mix=torch.zeros(rows, 4), + ) + + +@pytest.fixture +def run(monkeypatch): + """Call the mono path for one layer under a forward context built from + ``mode`` and ``metadata``; returns (result, runner calls).""" + + def _run( + mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs + ): + if metadata is None: + metadata = { + "swa": _swa(), + "comp": SimpleNamespace(block_size=128, block_table=None), + } + fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata) + runner = _Runner() + monkeypatch.setattr(md, "is_forward_context_available", lambda: True) + monkeypatch.setattr(md, "get_forward_context", lambda: fc) + monkeypatch.setattr(md, "_mono_runner", lambda device: runner) + monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) + args = {**_inputs(), **inputs} + mono = md.MonoDecodeLayer(ffn_only) + if (entry or ("ffn" if ffn_only else "layer")) == "ffn": + # x stands for wo_b's unreduced output; the first layer's attention + # has no compressed cache + part, _ = args.pop("x"), args.pop("positions") + return mono.ffn(_layer(comp=None), part, **args), runner.calls + return mono(_layer(), **args), runner.calls + + return _run + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_mono(run, mode): + out, calls = run(mode) + assert out == ("mono outputs",) and len(calls) == 1 + args, kwargs = calls[0] + # the caches reach the kernels as [blocks, block, record] bytes + assert args[8].shape == (4, 32, md.RECORD) and args[8].stride() == ( + 128 * md.RECORD, + md.RECORD, + 1, + ) + assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) + + +def test_padded_decode_step_runs_mono(run): + """A FULL graph pads the step to its capture size: the kernels take it and + skip the pad rows (their slot is -1).""" + swa = _swa(num_decode_tokens=M - 2) + comp = SimpleNamespace(block_size=128, block_table=None) + out, calls = run(metadata={"swa": swa, "comp": comp}) + assert out == ("mono outputs",) and len(calls) == 1 + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict( + metadata={ + "swa": _swa(num_prefills=1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(decode_swa_width=256), + "comp": SimpleNamespace(block_size=128), + } + ), + dict( + metadata={ + "swa": _swa(num_decode_tokens=M + 1), + "comp": SimpleNamespace(block_size=128), + } + ), + dict(metadata={}), # a profile / dummy run: no attention metadata + dict(**_inputs(md.MAX_ROWS + 6)), + dict(residual=None), # the first layer's seam broadcasts the embedding + ], + ids=[ + "piecewise", + "prefill", + "noncausal-window", + "tokens-past-rows", + "no-metadata", + "rows", + "first-layer", + ], +) +def test_other_steps_take_the_original_path(run, case): + out, calls = run(**case) + assert out is None and not calls + + +def _decoder_layer(engram=None, **attn): + attn = ( + dict( + layer_id=21, + compressor=None, + indexer=None, + compress_ratio=1, + kv_mxfp8=False, + kv_cache_dtype="fp8_ds_mla", + ) + | attn + ) + backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + quant_method = SimpleNamespace(mxfp4_backend=backend) + ffn = SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(quant_method=quant_method) + ), + shared_experts=object(), + gate=SimpleNamespace(tid2eid=None), + n_routed_experts=384, + n_activated_experts=6, + routed_scaling_factor=1.5, + swiglu_limit=10.0, + scoring_func="sqrtsoftplus", + renormalize=True, + ) + return SimpleNamespace( + attn=SimpleNamespace(**attn), + ffn=ffn, + engram=engram, + use_sequence_parallel=False, + fuse_seam_norm=True, + ) + + +@pytest.mark.parametrize( + "layer, path", + [ + (_decoder_layer(), "whole"), + (_decoder_layer(compress_ratio=0), "ffn"), + (_decoder_layer(engram=object()), "ffn"), + (_decoder_layer(compressor=object(), indexer=object()), "ffn"), + (_decoder_layer(indexer=object()), "ffn"), + (_decoder_layer(kv_cache_dtype="auto"), "ffn"), + (_decoder_layer(layer_id=40), None), + ], + ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"], +) +def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): + """The whole-layer kernels serve the standard layers; every other backbone + layer keeps vLLM's attention and runs its FFN half as one launch.""" + _deployment(monkeypatch) + mono = md.MonoDecodeLayer.create(layer, _config()) + got = None if mono is None else "ffn" if mono.ffn_only else "whole" + assert got == path + + +@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL]) +def test_decode_step_runs_ffn_launch(run, mode): + out, calls = run(mode, ffn_only=True, metadata={"swa": _swa()}) + assert out == ("ffn outputs",) and len(calls) == 1 + assert calls[0][0][1].shape == (M, 5120) # wo_b's unreduced output + + +def test_ffn_layer_never_runs_whole(run): + out, calls = run(ffn_only=True, entry="layer") + assert out is None and not calls + + +@pytest.mark.parametrize( + "case", + [ + dict(mode=CUDAGraphMode.PIECEWISE), + dict(metadata={"swa": _swa(num_prefills=1)}), + dict(metadata={"swa": _swa(num_decode_tokens=M + 1)}), + dict(metadata={}), + dict(**_inputs(md.MAX_ROWS + 6)), + ], + ids=["piecewise", "prefill", "tokens-past-rows", "no-metadata", "rows"], +) +def test_other_steps_reduce_wo_b_themselves(run, case): + """None: the layer all-reduces wo_b's partial and runs the FFN as before.""" + case.setdefault("metadata", {"swa": _swa()}) + out, calls = run(ffn_only=True, **case) + assert out is None and not calls + + +def _config(**parallel): + parallel = ( + dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) + | parallel + ) + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace(**parallel), + ) + + +def _deployment(monkeypatch, cdna=4, cus=256): + """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" + import vllm.platforms.rocm as rocm + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1") + monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) + monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) + monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(md, "_compute_units", lambda: cus) + name = "vllm.models.deepseek_v41.amd.mono.runner" + runner = ModuleType(name) + runner.BLOCKS = 256 # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, name, runner) + + +@pytest.mark.parametrize("cdna", [3, 5]) +def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): + _deployment(monkeypatch, cdna=cdna) + with pytest.raises(ValueError, match="needs CDNA4"): + md.MonoDecodeLayer.create(_decoder_layer(), _config()) + + +def _other_backend(): + layer = _decoder_layer() + experts = layer.ffn.experts.routed_experts + experts.quant_method.mxfp4_backend = Mxfp4MoeBackend.TRITON + return layer + + +@pytest.mark.parametrize( + "case, match", + [ + (dict(parallel=dict(enable_expert_parallel=True)), "no expert"), + (dict(parallel=dict(data_parallel_size=2)), "no expert"), + (dict(parallel=dict(enable_eplb=True)), "no expert"), + (dict(a4w4=True), "A8W4"), + (dict(layer=_other_backend), "AITER_MXFP4_BF16"), + (dict(cus=128), "256 compute units"), + ], + ids=["expert-parallel", "data-parallel", "eplb", "a4w4", "moe-backend", "cus"], +) +def test_opt_in_outside_the_kernels_moe_raises(monkeypatch, case, match): + """The kernels read all 384 experts on each rank in AITER's A8W4 layout, with + all 256 CTAs resident: any other deployment is an error, not garbage.""" + _deployment(monkeypatch, cus=case.get("cus", 256)) + if case.get("a4w4"): + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", "1") + layer = case.get("layer", _decoder_layer)() + with pytest.raises(ValueError, match=match): + md.MonoDecodeLayer.create(layer, _config(**case.get("parallel", {}))) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py new file mode 100644 index 000000000000..580462a82394 --- /dev/null +++ b/tests/models/test_dsv41_mono_layers.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The DeepSeek-V4.1 mono decode path through the real model: a five-layer +DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``VLLM_ROCM_MONO_DECODE`` on +and off on the same random weights: every decoder layer's outputs at every +decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which +checks the kernels against vLLM's ops one stage at a time, this runs +``DeepseekV4DecoderLayer.forward`` as the engine does -- the model's hooks, +the per-step path choice, the weights taken off the layer, the metadata and +the indexer's top-k -- so a change to the decoder layer that the kernels do +not follow fails here. + +The five layers hold both kinds the mono path takes: +- 0: sliding window only -> the FFN launch +- 1: ratio-2 KV and index source -> the FFN launch +- 2: ratio-2 consumer -> the whole layer (K1 + K2) +- 3: ratio-1 KV, index and candidate source -> the FFN launch +- 4: ratio-1 consumer -> the whole layer +""" + +import zlib + +import pytest +import torch + +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.platforms import current_platform + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +LOAD_FORMAT = "dsv41_mono_random" +TRUNCATED = { + "num_hidden_layers": 5, + # one entry a layer, the three MTP layers' included + "compress_ratios": [0, 2, 2, 1, 1, 0, 0, 0], + "kv_source_layer_ids": [1, 3], + "index_source_layer_ids": [1, 3], + "candidate_source_layer_id": 3, + "engram_layer_ids": [], + "engram_num_embeddings": [], + "dspark_target_layer_ids": [2, 3, 4], +} +WHOLE, FFN = (2, 4), (0, 1, 3) +STEPS = 5 # decode steps after the prefill's token + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +@torch.no_grad() +def _fill(name: str, t: torch.Tensor) -> None: + """A random value for every weight, seeded by its name (alike on every + rank and in every run), in the ranges the checkpoint's formats hold.""" + if t.device.type == "meta": + return # materialized and filled by the dummy loader's own path + g = torch.Generator(device=t.device).manual_seed(zlib.crc32(name.encode())) + + def randn(scale=1.0, shift=0.0): + return torch.randn(t.shape, generator=g, device=t.device) * scale + shift + + if t.dtype == torch.uint8 and "scale" in name: + # E8M0 2^-12 .. 2^-8, alike over each 32 rows as in the checkpoint (the + # loader folds per-row scales into 32 x 32 blocks only then). Larger + # random layers make the stack chaotic: atomics' order alone moves + # vLLM's own logits run to run. + rows = t.shape[-2] if t.dim() > 1 else 1 + shape = (*t.shape[:-2], -(-rows // 32), t.shape[-1]) if t.dim() > 1 else t.shape + e8m0 = torch.randint(115, 120, shape, generator=g, device=t.device) + if t.dim() > 1: + e8m0 = e8m0.repeat_interleave(32, dim=-2)[..., :rows, :] + t.copy_(e8m0) + elif t.dtype in (torch.uint8, torch.float4_e2m1fn_x2): + u8 = t.view(torch.uint8) + u8.copy_(torch.randint(0, 256, u8.shape, generator=g, device=t.device)) + elif t.dtype == torch.float8_e4m3fn: + t.copy_(randn().clamp(-6, 6)) + elif not torch.is_floating_point(t): + return + elif "norm" in name: + t.copy_(randn(0.1, 1.0)) + elif name.endswith(("hc_attn_scale", "hc_ffn_scale")): + t.copy_(randn(0.5).abs() + 0.5) + elif name.endswith(("hc_attn_base", "hc_ffn_base", "attn_sink")): + t.copy_(randn(0.1)) + elif name.endswith("embed_tokens.weight"): + t.copy_(randn()) + else: + t.copy_(randn(0.02)) + + +@register_model_loader(LOAD_FORMAT) +class RandomWeightLoader(DummyModelLoader): + """Dummy loading with weights in the checkpoint's ranges, before post-load + processing, so every processed copy follows (plain dummy weights leave the + E8M0 scales zero, and both paths would agree on zeros).""" + + def load_weights(self, model, model_config) -> None: + from vllm.distributed import get_tensor_model_parallel_world_size + + super().load_weights(model, model_config) + for name, t in model.state_dict().items(): + _fill(name, t) + # the routed experts' intermediate past this rank's share is the + # loader's zero padding (576 -> 640 at TP4): the kernels skip it + inter = model_config.hf_text_config.moe_intermediate_size + inter //= get_tensor_model_parallel_world_size() + for m in model.modules(): + w13, w2 = getattr(m, "w13_weight", None), getattr(m, "w2_weight", None) + if w13 is None or w2 is None or w13.device.type == "meta": + continue + padded = w13.shape[1] // 2 + if padded > inter: + w13, w2 = w13.data.view(torch.uint8), w2.data.view(torch.uint8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + + +class MonoProbe: + """Worker extension (``worker_extension_cls``): it also brings this module, + and so the loader, into the workers.""" + + def _decoder_layers(self) -> dict: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + + model = self.model_runner.model # type: ignore[attr-defined] + return { + m.attn.layer_id: m + for m in model.modules() + if isinstance(m, DeepseekV4DecoderLayer) + } + + def capture_steps(self, lens: list[int]) -> None: + """Keep each decoder layer's outputs (x, residual) for every decode + token taken in a decode-only step (the steps the mono path takes), + keyed by (layer, prompt length, position): a request's step and row + can differ between engines (requests reach the engine core while it + runs), its prompt length and position cannot. ``lens``: the prompt + lengths, all distinct; a row's request is the longest prompt not past + its position.""" + from vllm.forward_context import get_forward_context + + self._captured: dict = {} + lens = sorted(lens) + + def hook(layer_id): + def keep(module, args, kwargs, out): + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(module.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + pos = kwargs.get("positions", args[1] if len(args) > 1 else None) + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + kept = [t[row].float().cpu() for t in out[:2]] + self._captured.setdefault((layer_id, n, p), kept) + + return keep + + def gap_hook(layer_id, layer): + def keep(ffn, args, kwargs): + # vLLM's path only (the mono kernels route inside): each row's + # top-6 margin, sqrt(softplus(logits)) + bias as vLLM routes + md = get_forward_context().attn_metadata + if not isinstance(md, dict): + return + swa = md.get(layer.attn.swa_cache_layer.prefix) + if swa is None or swa.num_prefills != 0: + return + h = args[0] if args else kwargs["hidden_states"] + logits = h.float() @ ffn.gate.weight.float().T + score = torch.nn.functional.softplus(logits).sqrt() + score = score + ffn.gate.e_score_correction_bias.float() + top = score.topk(7, dim=-1).values + gap = (top[:, 5] - top[:, 6]).tolist() + pos = self._positions + for row, p in enumerate(pos.tolist()): + n = max(n for n in lens if n <= p) + self._gaps.setdefault((layer_id, n, p), gap[row]) + + return keep + + def pos_hook(module, args, kwargs): + self._positions = kwargs.get( + "positions", args[1] if len(args) > 1 else None + ) + + self._gaps: dict = {} + for i, layer in self._decoder_layers().items(): + layer.register_forward_hook(hook(i), with_kwargs=True) + layer.register_forward_pre_hook(pos_hook, with_kwargs=True) + gap = gap_hook(i, layer) + layer.ffn.register_forward_pre_hook(gap, with_kwargs=True) + + def save_capture(self, path: str) -> None: + rank = self.rank # type: ignore[attr-defined] + torch.save((self._captured, self._gaps), f"{path}.rank{rank}") + + def mono_layers(self) -> dict: + from vllm.models.deepseek_v41.amd import mono_decode + + layers = { + i: None + if m.mono is None + else ["ffn" if m.mono.ffn_only else "whole", m.mono._weights is not None] + for i, m in self._decoder_layers().items() + } + runner = mono_decode._runner + return { + "layers": layers, + "epoch": -1 if runner is None else int(runner.epoch[0]), + } + + +def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): + """Greedy decode of ``prompts`` with the mono path on or off: (each + prompt's output, the workers' mono layers); each decoder layer's outputs + for the decode tokens of decode-only steps, and vLLM's routing margins, + saved under ``path``.""" + from vllm import SamplingParams + from vllm.inputs import TokensPrompt + + monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1" if mono else "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + with vllm_runner( + MODEL, + tensor_parallel_size=tp, + load_format=LOAD_FORMAT, + hf_overrides=TRUNCATED, + worker_extension_cls=f"{__name__}.MonoProbe", + moe_backend="aiter", + language_model_only=True, + tokenizer_mode="deepseek_v41", + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + block_size=128, + enforce_eager=True, + max_model_len=4096, + max_num_seqs=len(prompts), + # a prompt prefills in one step, so its decode tokens never share one + max_num_batched_tokens=16384, + enable_prefix_caching=False, + gpu_memory_utilization=0.5, + ) as runner: + runner.llm.collective_rpc("capture_steps", args=([len(p) for p in prompts],)) + params = SamplingParams(temperature=0.0, max_tokens=STEPS + 1) + outs = runner.llm.generate( + [TokensPrompt(prompt_token_ids=p) for p in prompts], params + ) + runner.llm.collective_rpc("save_capture", args=(path,)) + probe = runner.llm.collective_rpc("mono_layers") + return [o.outputs[0] for o in outs], probe + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a - b).norm() / b.norm()).item() + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + g = torch.Generator().manual_seed(0) + # contexts past both ratios' 512 compressed entries, so top-k selects + lens = torch.randperm(1700, generator=g)[:8].add(300).tolist() # distinct + prompts = [torch.randint(1000, 60000, (n,), generator=g).tolist() for n in lens] + + on, probe_on = _decode( + vllm_runner, monkeypatch, tp, True, prompts, f"{tmp_path}/on" + ) + off, probe_off = _decode( + vllm_runner, monkeypatch, tp, False, prompts, f"{tmp_path}/off" + ) + + # both layer kinds took the mono path on every rank, at every decode step + launches = len(WHOLE + FFN) # one launch (pair) a layer a step + for rank in probe_on: + assert rank["layers"] == { + i: ["whole" if i in WHOLE else "ffn", True] for i in WHOLE + FFN + }, rank + assert rank["epoch"] >= launches * STEPS, rank + assert rank["epoch"] % launches == 0, rank + assert all(set(r["layers"].values()) == {None} for r in probe_off) + # each decode token's input matches between the runs while the greedy + # tokens before it do: (prompt length, position) pairs to compare + same: set[tuple[int, int]] = set() + for prompt, a, b in zip(prompts, on, off): + for step in range(STEPS): + if a.token_ids[: step + 1] != b.token_ids[: step + 1]: + break + same.add((len(prompt), len(prompt) + step)) + + # each layer's outputs for those tokens, mono against vLLM's path, token by + # token. Layer 0 sees the same inputs on both paths at a request's first + # decode token: there the bound is test_dsv41_mono_numerics' (about 3x the + # kernels' 0.4%); its later tokens attend over KV each path wrote, 3%. Each + # later layer takes the previous one's slightly different output, which + # random weights amplify (2-4% by layer 4): 8%. A token whose top-6 routing + # is a near-tie on vLLM's path (6th / 7th score within 5e-3: a 0.4% input + # difference moves scores near 0.8 by about 3e-3) may pick another expert; + # once it has, its residual stream differs and the later layers are exempt + # for it. The residual stream: 3% (0.3% seen). A layer the kernels no + # longer follow is off by far more, on every token. + for rank in range(tp): + got, _ = torch.load(f"{tmp_path}/on.rank{rank}") + want, gaps = torch.load(f"{tmp_path}/off.rank{rank}") + keys: list[tuple[int, int]] = sorted( + same & {k[1:] for k in got} & {k[1:] for k in want} + ) + assert len(keys) >= len(prompts), keys + # the bound that sees a 3% error: a request's first decode token + assert any(p == n for n, p in keys), keys + flipped: set = set() + for i in sorted(WHOLE + FFN): + for k in keys: + if k in flipped: + continue + (x, res), (x_ref, res_ref) = got[(i, *k)], want[(i, *k)] + err = _rel(res, res_ref) + assert err < 0.03, (rank, i, k, "residual", err) + bound = (0.015 if k[1] == k[0] else 0.03) if i == 0 else 0.08 + err = _rel(x, x_ref) + if err >= bound: + assert gaps[(i, *k)] < 5e-3, (rank, i, k, "x", err) + flipped.add(k) + assert len(flipped) <= len(keys) // 4, (rank, sorted(flipped)) diff --git a/tests/models/test_dsv41_mono_numerics.py b/tests/models/test_dsv41_mono_numerics.py new file mode 100644 index 000000000000..f7cd877caf47 --- /dev/null +++ b/tests/models/test_dsv41_mono_numerics.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Numerics of the DeepSeek-V4.1 mono decode kernels at TP2 on CDNA4 against the +ops vLLM runs for the same decode step: random weights in the checkpoint's +formats, the step's KV caches as vLLM's insert writes them. + +- The whole layer (K1 + K2) against the mHC seams, the ROCm attention op chain, + the TP all-reduce and ``DeepseekV4MoE``. +- The FFN launch against the same FFN half, from wo_b's unreduced output. + +Each stage is compared on its own input, so a failure names the stage; the MoE +is checked on the kernels' own FFN input, where routing cannot differ.""" + +import pytest +import ray +import torch + +from vllm.platforms import current_platform + +from ..utils import ( + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_process_parallel, +) + +MODEL = "deepseek-ai/DeepSeek-V4.1-Flash" +HIDDEN, Q_RANK, HEAD_DIM, NOPE, ROPE = 5120, 1280, 512, 448, 64 +N_HEADS, O_GROUPS, O_RANK = 64, 8, 1024 +WINDOW, TOPK, SWA_BLOCK, CACHE_BLOCK = 128, 512, 32, 128 +RECORD, DATA, ALIGN = 584, 576, 576 +REPLAYS = 300 +HC = 4 + + +def _on_cdna4() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import get_cdna_version + + return get_cdna_version() == 4 + + +def _cos(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.float().reshape(-1), b.float().reshape(-1) + return torch.nn.functional.cosine_similarity(a, b, dim=0).item() + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + """Relative L2 error of ``a`` against ``b``: unlike the cosine, it sees a + wrong scale.""" + a, b = a.float().reshape(-1), b.float().reshape(-1) + return ((a - b).norm() / b.norm()).item() + + +# ---------------------------------------------------------------- caches + + +def _paged_cache(num_blocks: int, block: int, device) -> torch.Tensor: + """[blocks, block, 584] u8 over pages of round_up(block * 584, 576) bytes, + the way vLLM binds a packed V4 cache.""" + page = -(-block * RECORD // ALIGN) * ALIGN + raw = torch.zeros(num_blocks * page, dtype=torch.uint8, device=device) + return torch.as_strided(raw, (num_blocks, block, RECORD), (page, RECORD, 1)) + + +def _rope(kv: torch.Tensor, positions: torch.Tensor, cos_sin: torch.Tensor): + """GPT-J RoPE of the last 64 dims (fp32 math, bf16 out).""" + cs = cos_sin[positions.long()] + cos, sin = cs[:, :32], cs[:, 32:] + x = kv.float() + e, o = x[:, NOPE::2].clone(), x[:, NOPE + 1 :: 2].clone() + x[:, NOPE::2] = e * cos - o * sin + x[:, NOPE + 1 :: 2] = e * sin + o * cos + return x.to(torch.bfloat16) + + +def _write_records(cache, block: int, slots, kv): + """fp8_ds_mla records of roped bf16 kv rows at ``slots``: the 448 nope dims + in e4m3 with a UE8M0 per 64, the rope dims in bf16, as vLLM's insert.""" + nope = kv.float()[:, :NOPE].view(-1, 7, 64) + exp = torch.ceil(torch.log2(nope.abs().amax(-1).clamp_min(1e-4) / 448.0)) + q = (nope * torch.exp2(-exp)[..., None]).clamp(-448, 448).reshape(-1, NOPE) + data = torch.cat( + [ + q.to(torch.float8_e4m3fn).view(torch.uint8), + kv[:, NOPE:].contiguous().view(torch.uint8), + ], + 1, + ) + scales = torch.zeros(kv.shape[0], 8, dtype=torch.uint8, device=kv.device) + scales[:, :7] = (exp + 127).clamp(0, 255).to(torch.uint8) + n = cache.untyped_storage().nbytes() - cache.storage_offset() + flat = torch.as_strided(cache, (n,), (1,)) + slots = slots.to(kv.device).long() + b, o = slots // block, slots % block + d0 = b * cache.stride(0) + o * DATA + s0 = b * cache.stride(0) + block * DATA + o * 8 + ar = torch.arange(DATA, device=kv.device) + flat[(d0[:, None] + ar).reshape(-1)] = data.reshape(-1) + ar = torch.arange(8, device=kv.device) + flat[(s0[:, None] + ar).reshape(-1)] = scales.reshape(-1) + + +def _normed_kv(n: int, kv_norm, g, device): + x = torch.randn(n, HEAD_DIM, generator=g).to(device) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)) * kv_norm.float() + return x.to(torch.bfloat16) + + +def _decode_step(ratio: int, R: int, T: int, kv_norm, cos_sin, seed: int): + """R requests x T tokens past random contexts: the step's metadata as + vLLM's builders lay it out, its window rows and (ratio > 0) the compressed + cache filled, a random top-512 a token.""" + g = torch.Generator().manual_seed(seed) + dev = kv_norm.device + ctx = [int(c) for c in torch.randint(2000, 40000, (R,), generator=g)] + M = R * T + spans = [ + (max(0, c - WINDOW - T) // SWA_BLOCK, (c + T - 1) // SWA_BLOCK) for c in ctx + ] + n_swa = 1 + sum(hi - lo + 1 for lo, hi in spans) + perm = (torch.randperm(n_swa - 1, generator=g) + 1).tolist() + bt = torch.zeros(R, max(hi for _, hi in spans) + 1, dtype=torch.int64) + for r, (lo, hi) in enumerate(spans): + for lb in range(lo, hi + 1): + bt[r, lb] = perm.pop() + swa = _paged_cache(n_swa, SWA_BLOCK, dev) + + def slot(r, p): + return int(bt[r, p // SWA_BLOCK]) * SWA_BLOCK + p % SWA_BLOCK + + pos = [ctx[r] + t for r in range(R) for t in range(T)] + t2r = [r for r in range(R) for _ in range(T)] + swa_idx = torch.full((M, WINDOW), -1, dtype=torch.int32) + swa_len = torch.zeros(M, dtype=torch.int32) + for i in range(M): + ps = range(max(pos[i] - WINDOW + 1, 0), pos[i] + 1) + swa_idx[i, : len(ps)] = torch.tensor([slot(t2r[i], p) for p in ps]) + swa_len[i] = len(ps) + hist = [ + (r, p) for r in range(R) for p in range(max(0, ctx[r] - WINDOW - T), ctx[r]) + ] + hp = torch.tensor([p for _, p in hist], device=dev) + kv = _rope(_normed_kv(len(hist), kv_norm, g, dev), hp, cos_sin) + _write_records(swa, SWA_BLOCK, torch.tensor([slot(r, p) for r, p in hist]), kv) + + step = dict( + M=M, + positions=torch.tensor(pos, dtype=torch.int64, device=dev), + slot_mapping=torch.tensor( + [slot(t2r[i], pos[i]) for i in range(M)], dtype=torch.int64, device=dev + ), + token_to_req=torch.tensor(t2r, dtype=torch.int32, device=dev), + swa_indices=swa_idx.to(dev), + swa_lens=swa_len.to(dev), + swa_cache=swa, + topk=None, + comp_cache=None, + comp_block_table=None, + comp_entries=0, + ) + if ratio: + entries = CACHE_BLOCK // ratio + lens = [(c + T) // ratio for c in ctx] + nblk = [-(-n // entries) for n in lens] + perm = (torch.randperm(sum(nblk), generator=g) + 1).tolist() + cbt = torch.zeros(R, max(nblk), dtype=torch.int32) + for r in range(R): + for b in range(nblk[r]): + cbt[r, b] = perm.pop() + comp = _paged_cache(1 + sum(nblk), entries, dev) + for r in range(R): + idx = torch.arange(lens[r]) + slots = cbt[r, idx // entries].long() * entries + idx % entries + kv = _rope( + _normed_kv(lens[r], kv_norm, g, dev), (idx * ratio).to(dev), cos_sin + ) + _write_records(comp, entries, slots, kv) + topk = torch.full((M, TOPK), -1, dtype=torch.int32) + for i in range(M): + cand = (pos[i] + 1) // ratio + n = min(cand, TOPK) + topk[i, :n] = torch.randperm(cand, generator=g)[:n].to(torch.int32) + step.update( + topk=topk.to(dev), + comp_cache=comp, + comp_block_table=cbt.to(dev), + comp_entries=entries, + ) + return step + + +# ---------------------------------------------------------------- weights + + +def _mxfp8(n: int, k: int, g, device): + """A random MXFP8 weight [n, k] e4m3 and its E8M0 32 x 32 block scales.""" + w = torch.randn(n, k, generator=g).clamp(-6, 6).to(torch.float8_e4m3fn) + s = torch.randint(118, 123, (n // 32, k // 32), generator=g, dtype=torch.uint8) + return w.to(device), s.to(device) + + +def _random_moe(moe, gd, gr, inter: int) -> None: + """Random parameters in the loaded layout, before post-load processing: this + rank's shards (``gd``, a device generator) of the MXFP4 routed experts + (random nibbles, E8M0 scales; the intermediate past ``inter`` zero, as the + loader pads it) and the MXFP8 shared expert (one E8M0 a 32 x 32 block); the + replicated router from ``gr``, seeded alike on every rank.""" + u8 = torch.uint8 + for lin in (moe.shared_experts.gate_up_proj, moe.shared_experts.down_proj): + n = lin.weight.shape[0] + w = torch.randn(lin.weight.shape, generator=gd, device=gd.device) + lin.weight.data.copy_(w.clamp(-6, 6).to(lin.weight.dtype)) + sc = lin.weight_scale.data.view(u8) + rows = -(-n // 32) if sc.shape[0] == n else sc.shape[0] + blocks = torch.randint( + 118, 123, (rows, sc.shape[1]), generator=gd, device=gd.device, dtype=u8 + ) + sc.copy_(blocks.repeat_interleave(32, 0)[:n] if sc.shape[0] == n else blocks) + e = moe.experts.routed_experts + for name, lo, hi in (("weight", 0, 256), ("weight_scale", 118, 122)): + for w in ("w13", "w2"): + t = getattr(e, f"{w}_{name}").data.view(u8) + t.copy_( + torch.randint(lo, hi, t.shape, generator=gd, device=gd.device, dtype=u8) + ) + padded = e.w13_weight.shape[1] // 2 + if padded > inter: + w13, w2 = e.w13_weight.data.view(u8), e.w2_weight.data.view(u8) + for half in (0, padded): + w13[:, half + inter : half + padded] = 0 + w2[..., inter // 2 :] = 0 + for t in (moe.gate.weight, moe.gate.e_score_correction_bias): + t.data.copy_(torch.randn(t.shape, generator=gr, device=gr.device) * 0.02) + + +# ---------------------------------------------------------------- the worker + + +@ray.remote(num_gpus=1, max_calls=1) +def _mono_numerics(monkeypatch, tp_size, pp_size, rank, distributed_init_port): + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.worker.workspace import init_workspace_manager + + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + # the checkpoint's MXFP4 experts run on AITER's MoE; its flags are read + # at import, so refresh them + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + rocm_aiter_ops.refresh_env_variables() + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + init_workspace_manager(device) + _check_layers(rank, tp_size, device) + + +def _check_layers(rank: int, tp: int, device) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce + from vllm.engine.arg_utils import EngineArgs + from vllm.forward_context import set_forward_context + from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( + rocm_mxfp8_block32_gemm, + ) + from vllm.model_executor.layers.mhc import MHCPreDelayedOp + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.model_executor.model_loader.utils import process_weights_after_loading + from vllm.models.common.ops import fused_q_kv_rmsnorm + from vllm.models.deepseek_v41.amd.model import DeepseekV4MoE + from vllm.models.deepseek_v41.amd.mono.layer import scratch_layout + from vllm.models.deepseek_v41.amd.mono.runner import ( + AttnWeights, + DSV41MonoLayer, + MonoLayerWeights, + ) + from vllm.models.deepseek_v41.amd.rocm import ( + compute_global_topk_ragged_indices_and_indptr, + ) + from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope + from vllm.utils.torch_utils import set_default_torch_dtype + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_ragged_indices_from_dense, + rocm_inverse_rope_mxfp8_rows, + ) + + vc = EngineArgs( + model=MODEL, + tensor_parallel_size=tp, + moe_backend="aiter", + load_format="dummy", + max_model_len=65536, + language_model_only=True, + attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, + ).create_engine_config() + cfg = vc.model_config.hf_config + cfg = getattr(cfg, "text_config", cfg) + # replicated tensors alike on every rank; each rank's own TP shards + g = torch.Generator().manual_seed(1000) + gs = torch.Generator().manual_seed(1001 + rank) + H, G = N_HEADS // tp, O_GROUPS // tp + bf = torch.bfloat16 + + with set_current_vllm_config(vc), set_default_torch_dtype(bf), device: + moe = DeepseekV4MoE(vc, prefix="model.layers.21.ffn") + mhc = MHCPreDelayedOp() + _random_moe( + moe, + torch.Generator(device=device).manual_seed(2001 + rank), + torch.Generator(device=device).manual_seed(2000), + inter=2304 // tp, + ) + process_weights_after_loading(moe, vc.model_config, device) + + def f32(*shape, s=1.0, gen=g): + return (torch.randn(*shape, generator=gen) * s).to(device) + + mix = (2 + HC) * HC + seams = { + sub: dict( + fn=f32(mix, HC * HIDDEN, s=0.01), + scale=f32(3, s=0.5).abs() + 0.5, + base=f32(mix, s=0.1), + norm=(1 + f32(HIDDEN, s=0.1)).to(bf), + ) + for sub in ("attn", "ffn") + } + + def seam(sub, residual, x, post, comb, pre): + s = seams[sub] + return mhc( + residual, s["fn"], s["scale"], s["base"], cfg.rms_norm_eps, + cfg.hc_eps, cfg.hc_eps, 2.0, cfg.hc_sinkhorn_iters, + pre_mix=pre, sublayer_out=x, post_layer_mix=post, comb_res_mix=comb, + norm_weight=s["norm"], norm_eps=cfg.rms_norm_eps, + ) # fmt: skip + + def moe_out(x): + with set_forward_context(None, vc): + ids = torch.zeros(x.shape[0], dtype=torch.int64, device=device) + return moe(x.clone(), ids) + + wqkv, wqkv_s = _mxfp8(1792, HIDDEN, g, device) + wq_b, wq_b_s = _mxfp8(H * HEAD_DIM, Q_RANK, gs, device) + wo_a, wo_a_s = _mxfp8(G * O_RANK, H // G * HEAD_DIM, gs, device) + wo_b, wo_b_s = _mxfp8(HIDDEN, G * O_RANK, gs, device) + q_norm, kv_norm = (1 + f32(Q_RANK, s=0.1)).to(bf), (1 + f32(HEAD_DIM, s=0.1)).to(bf) + sink = f32(H, gen=gs) + e, sh = moe.experts.routed_experts, moe.shared_experts + u8 = lambda t: t.view(torch.uint8) # noqa: E731 + runner = DSV41MonoLayer(tp, rank, get_tp_group().cpu_group, device) + gw, bias = moe.gate.weight.float(), moe.gate.e_score_correction_bias.float() + + def top6(x): + """The top-6 experts as vLLM routes: sqrt(softplus(logits)) + bias.""" + sc = torch.nn.functional.softplus(x.float() @ gw.T).sqrt() + bias + return sc.topk(6, dim=1).indices.sort(1).values + + steps = [] # each case's launches, on fixed buffers, for the replay stress + for ratio, R, T in ((1, 1, 6), (2, 8, 6)): + with set_current_vllm_config(VllmConfig()): + rope = build_deepseek_v4_rope( + cfg, head_dim=HEAD_DIM, rope_head_dim=ROPE, + max_position_embeddings=cfg.max_position_embeddings, + compress_ratio=ratio, + ) # fmt: skip + cos_sin = rope.cos_sin_cache.to(device=device, dtype=torch.float32) + s = _decode_step(ratio, R, T, kv_norm, cos_sin, seed=7 + ratio) + M = s["M"] + mw = MonoLayerWeights( + attn=AttnWeights( + layer_id=21, wqkv=wqkv, wqkv_scale=wqkv_s, q_norm=q_norm, + kv_norm=kv_norm, wq_b=wq_b, wq_b_scale=wq_b_s, wo_a=wo_a, + wo_a_scale=wo_a_s, wo_b=wo_b, wo_b_scale=wo_b_s, attn_sink=sink, + cos_sin=cos_sin, ratio=ratio, + ), + hc_attn_fn=seams["attn"]["fn"], hc_attn_scale=seams["attn"]["scale"], + hc_attn_base=seams["attn"]["base"], attn_norm=seams["attn"]["norm"], + hc_ffn_fn=seams["ffn"]["fn"], hc_ffn_scale=seams["ffn"]["scale"], + hc_ffn_base=seams["ffn"]["base"], ffn_norm=seams["ffn"]["norm"], + gate_w=moe.gate.weight, bias=moe.gate.e_score_correction_bias, + w13=e.w13_weight, w13_s=e.w13_weight_scale, w2=e.w2_weight, + w2_s=e.w2_weight_scale, sgu=sh.gate_up_proj.weight, + sgu_s=u8(sh.gate_up_proj.weight_scale), sw2=sh.down_proj.weight, + sw2_s=u8(sh.down_proj.weight_scale), + ) # fmt: skip + + # the layer's inputs at the attention seam (the same on every rank) + gi = torch.Generator().manual_seed(11) + res0 = torch.randn(M, HC, HIDDEN, generator=gi) + res0 = ( + (res0 * torch.rsqrt(res0.pow(2).mean(-1, keepdim=True))).to(bf).to(device) + ) + x0 = (torch.randn(M, HIDDEN, generator=gi) * 0.5).to(bf).to(device) + comb0 = torch.rand(M, HC, HC, generator=gi) + 0.1 + for _ in range(20): + comb0 = comb0 / comb0.sum(-1, keepdim=True) + comb0 = comb0 / comb0.sum(-2, keepdim=True) + comb0 = comb0.to(device) + post0 = (torch.sigmoid(torch.randn(M, HC, 1, generator=gi)) * 2).to(device) + pre0 = (torch.sigmoid(torch.randn(M, HC, generator=gi)) + 1e-6).to(device) + + # vLLM's decode chain + res, post, comb, xa, attn_pre = seam("attn", res0, x0, post0, comb0, pre0) + swa_ind, swa_ptr = build_ragged_indices_from_dense( + s["swa_indices"].reshape(M, -1), s["swa_lens"] + ) + top_ind = top_ptr = None + if ratio: + valid = torch.ones(M, dtype=torch.bool, device=device) + top_ind, top_ptr, _ = compute_global_topk_ragged_indices_and_indptr( + s["topk"], s["token_to_req"], s["comp_block_table"], + s["comp_entries"], valid, + ) # fmt: skip + x8, xs = mxfp8_e4m3_quantize(xa) + qr_kv = rocm_mxfp8_block32_gemm(x8, xs, wqkv, wqkv_s, bf) + qr, kv = fused_q_kv_rmsnorm( + qr_kv[:, :Q_RANK], qr_kv[:, Q_RANK:], q_norm, kv_norm, cfg.rms_norm_eps + ) + q8, qs = mxfp8_e4m3_quantize(qr) + q = rocm_mxfp8_block32_gemm(q8, qs, wq_b, wq_b_s, bf).view(M, H, HEAD_DIM) + swa2d = s["swa_cache"].view(s["swa_cache"].shape[0], -1) + q_roped = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, kv, swa2d, s["slot_mapping"], s["positions"], cos_sin, H, + cfg.rms_norm_eps, SWA_BLOCK, False, False, True, False, + ) # fmt: skip + attn = torch.empty(M, H, HEAD_DIM, dtype=bf, device=device) + rocm_aiter_ops.triton_sparse_mla_fwd( + q_roped, s["swa_cache"], attn, HEAD_DIM**-0.5, swa_ptr, swa_ind, + attn_sink=sink, extra_kv_buffer=s["comp_cache"], + extra_kv_indptr=top_ptr, extra_kv_indices=top_ind, + ) # fmt: skip + a8 = torch.empty(M, H * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + a_s = torch.empty(M, H * HEAD_DIM // 32, dtype=torch.uint8, device=device) + rocm_inverse_rope_mxfp8_rows(attn, s["positions"], cos_sin, ROPE, a8, a_s) + z = torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(a8, a_s, wo_a, wo_a_s, G, O_RANK) + z8, zs = mxfp8_e4m3_quantize(z) + part = rocm_mxfp8_block32_gemm(z8, zs, wo_b, wo_b_s, bf).contiguous() + a = tensor_model_parallel_all_reduce(part.clone()) + res2, post2, comb2, xf, ffn_pre = seam("ffn", res, a, post, comb, attn_pre) + out = moe_out(xf) + + # ---- the whole layer: K1 + K2 + o = runner.forward( + mw, x0, res0, post0, comb0, pre0, s["positions"], s["slot_mapping"], + s["swa_cache"], s["swa_indices"], s["swa_lens"], s["token_to_req"], + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + off = scratch_layout(M, tp)["normed"][0] + normed = runner.scratch(M)[off : off + M * HIDDEN * 2].view(bf).view(M, HIDDEN) + tag = f"ratio {ratio}, {R} x {T}" + assert _cos(runner.res_mid[:M], res) > 0.99999, tag + for got, want in ( + (runner.post_a, post), + (runner.comb_a, comb), + (runner.pre_a, attn_pre), + ): + torch.testing.assert_close( + got[:M].view_as(want), want, rtol=1e-4, atol=1e-5 + ) + same_in = moe_out(normed) + # relative errors about 3x what the kernels show (0.3-0.5%) + assert _cos(o[1], res2) > 0.9999 and _rel(o[1], res2) < 0.01, tag + assert _cos(normed, xf) > 0.999 and _rel(normed, xf) < 0.012, tag + assert _cos(o[0], same_in) > 0.9999 and _rel(o[0], same_in) < 0.015, tag + # per token, where both FFN inputs pick the same experts + same = (top6(xf) == top6(normed)).all(1) + assert same.float().mean() >= 0.5, tag + tok = torch.nn.functional.cosine_similarity(o[0].float(), out.float(), dim=1) + assert tok[same].min() > 0.99, tag + + # ---- the FFN launch, from the same unreduced wo_b output + f = runner.ffn(mw, part, res, post, comb, attn_pre) + torch.testing.assert_close(f[1], res2, rtol=2**-8, atol=1e-6) + torch.testing.assert_close(normed, xf, rtol=2**-7, atol=1e-6) + assert _cos(f[0], out) > 0.9999 and _rel(f[0], out) < 0.015, tag + torch.accelerator.synchronize() + + whole, ffn = runner._outs(M, res0), runner._outs(M, res) + args = (s["positions"], s["slot_mapping"], s["swa_cache"], s["swa_indices"]) + meta = dict( + topk_indices=s["topk"], comp_cache=s["comp_cache"], + comp_block_table=s["comp_block_table"], + ) # fmt: skip + + def step(mw=mw, x0=x0, res0=res0, post0=post0, comb0=comb0, pre0=pre0, + args=args, s=s, meta=meta, whole=whole, part=part, res=res, + post=post, comb=comb, attn_pre=attn_pre, ffn=ffn): # fmt: skip + runner.forward( + mw, x0, res0, post0, comb0, pre0, *args, s["swa_lens"], + s["token_to_req"], **meta, outs=whole, + ) # fmt: skip + runner.ffn(mw, part, res, post, comb, attn_pre, outs=ffn) + return [*whole, *ffn] + + steps.append(step) + + # ---- the hand-off protocol under back-to-back launches: both step widths + # alternating in one graph, every replay bit for bit the first's + outs = [t for step in steps for t in step()] + torch.accelerator.synchronize() + ref = [t.clone() for t in outs] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for step in steps: + step() + for i in range(REPLAYS): + graph.replay() + if i % 50 == 49 or i == REPLAYS - 1: + torch.accelerator.synchronize() + bad = [ + j for j, (t, r) in enumerate(zip(outs, ref)) if not torch.equal(t, r) + ] + assert not bad, f"replay {i}: outputs {bad} differ from the first run" + + +@pytest.mark.skipif(not _on_cdna4(), reason="the mono kernels target CDNA4") +@pytest.mark.parametrize("tp", [2, 4]) +def test_mono_decode_layer_matches_vllm(monkeypatch: pytest.MonkeyPatch, tp: int): + if torch.accelerator.device_count() < tp: + pytest.skip(f"needs {tp} GPUs") + pytest.importorskip("flydsl") + multi_process_parallel(monkeypatch, tp, 1, _mono_numerics) diff --git a/vllm/envs.py b/vllm/envs.py index 851741d5253b..02c282fd8ec2 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ VLLM_ROCM_USE_AITER_FP4BMM: bool = True VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1387,6 +1388,12 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") ), + # Run the eligible layers of decode steps on the model's mono decode + # kernels: persistent launches a layer, the TP all-reduces in-kernel + # (DeepSeek-V4.1 on CDNA4). By default is disabled. + "VLLM_ROCM_MONO_DECODE": lambda: ( + os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") + ), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index cff72f12dc22..5de173a8706c 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -14,6 +14,7 @@ get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, ) from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger @@ -58,6 +59,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) +from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.sequence import IntermediateTensors @@ -255,6 +257,12 @@ def __init__( # attn_norm / ffn_norm into its collapse, so the separate norms are # skipped for the seams it takes. self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM + # Decode steps of an eligible layer may run as the mono decode layer. + self.mono = MonoDecodeLayer.create(self, vllm_config) + if self.mono is not None and self.mono.ffn_only: + # A decode step hands wo_b's TP partial to the FFN launch; forward + # reduces it on the others. + self.attn.wo_b.reduce_results = False @staticmethod def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: @@ -273,6 +281,10 @@ def forward( engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if self.mono is not None: + out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix) + if out is not None: + return out # Layer 0's attention seam projects the 2-D embedding with the folded # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn. # The fused kernel only takes the latter, so that seam keeps the @@ -371,6 +383,11 @@ def forward( x = self.attn(positions, x, None) if self.use_sequence_parallel: x = sp_reduce_scatter(x) + elif self.mono is not None and self.mono.ffn_only: + out = self.mono.ffn(self, x, residual, post_mix, res_mix, attn_pre) + if out is not None: + return out + x = tensor_model_parallel_all_reduce(x) residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md new file mode 100644 index 000000000000..88e40d73dd0a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -0,0 +1,94 @@ +# DeepSeek-V4.1-Flash mono decode layer (gfx950, FlyDSL) + +One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, +the attention, the attention all-reduce, the FFN seam, the MoE and the MoE +all-reduce, in two persistent launches a layer. The other backbone layers keep +vLLM's attention and run the rest -- the attention all-reduce, the FFN seam, +the MoE and its all-reduce -- as one launch (the FFN launch). ATOM's V4.1 mono decode +(ROCm/ATOM `atom/models/deepseek_v41/mono`) is the reference design: the seam +and MoE stages (`stages/`) and the framework under them (`common/`) are adapted +from it, each file naming its origin; the attention stages (`attention/`) are +new, reading vLLM's KV records and weights as loaded. + +## Scope + +- Layers: the whole layer on every backbone layer with the standard seam and no + compressor / indexer: 3-7, 9-13, 15-19 (compress ratio 2) and 21-23, 25-27, + 29-31, 33-35, 37-39 (ratio 1), 30 of 40. The FFN launch on the other 10: + layer 0 (embedding-broadcast seam), layer 1 (Engram at the seam) and the + compressor / indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36, + whose wo_b leaves its TP reduce to it. The DSpark draft layers keep vLLM's + path. The paths interleave freely (same tensors at the layer boundary). +- Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 + DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the + routing bias is `e_score_correction_bias`). +- Contract: a drop-in for `DeepseekV4DecoderLayer.forward`: + in `x` (the previous FFN output, reduced, bf16 [M, 5120]), `residual` + (bf16 [M, 4, 5120]), `post_mix` [M, 4, 1], `res_mix` [M, 4, 4], + `pre_mix` [M, 4] (f32, the previous seam's) + out the MoE output (reduced), the residual after the FFN seam, and that + seam's post / comb / pre mixes. + +## Launches + + K1 seam (stages/seam.py): slice (160 x 32 columns): fold x into the + residual (post x + sum comb r), the collapsed input (pre . R'), partial + mixes (fn . R') and sums of squares -> gate (S): 24 mixes -> pre / post + / comb (Sinkhorn 20) -> norm (S): attn_norm -> vLLM MXFP8 -> X8 + attention front (attention/front.py): wqkv, q / kv norms + KV insert, + wq_b + RoPE -> Q + K2 attention back (attention/back.py): split attention, combine + inverse + RoPE, wo_a, wo_b; the bf16 wo_b partial pushed to every rank + FFN seam: reduce (the ranks' partials summed in rank order, fp32 -> + bf16: the attention output), slice, gate, norm -> the MoE input + MoE (stages/moe.py): xq, router (bf16 GEMV, f32 logits), route + (sqrtsoftplus + bias, top-6, renormalized x 1.5), shared expert (MXFP8 + GEMVs, clamped SiLU), ug (MXFP4 x MXFP8 per expert of the step's + union), down (+ routing weight, top-k order; + shared) -> push, rank-order + sum -> out + FFN K2 without the attention back: push (160 x 32 columns) of vLLM's + unreduced wo_b output, as its bits, to every rank's attention region, + then K2's FFN seam and MoE (the same code, `layer.run_ffn`) + +## Execution model and hand-offs + +- Grid = 256 CTAs (one a CU, all resident), 512 threads a CTA, LDS a per-stage + union. A stage is a task table: CTA `b` runs tasks `b, b + 256, ...`, stages + in order; a consumer only waits on earlier stages, so no deadlock while the + grid is resident. A GEMV task issues its weight loads before it polls its + activations. +- In a launch: tagged pairs (value + tag in one 8 B store) or plain data + + flags, at device scope. The tag is the launch pair's epoch (K1 `2e + 1`, + K2 or the FFN launch `2e + 2`), a device word their CTA 0 moves on once every + CTA has marked the launch done: no clearing between steps, no host step hook, every argument + a fixed device pointer (graph-safe). +- Across ranks: one symmetric uncached buffer a rank (aiter `UncachedIpcHeap`), + every rank holding every peer's address; pushes are system-scope stores of + (value, tag) pairs. Every rank runs the same launches, so the epochs agree; + peer regions alternate by epoch parity, and a rank cannot run two steps ahead + of a peer (each step's reduce waits on every rank's push). + +## Numerics (the vLLM path each stage reproduces) + +- Seams: aiter `mhc_fused_post_pre_delayed_rmsnorm` (post fold order, fn as + bf16 hi + lo, products summed in order, Triton sigmoid / exp / rcp). +- Attention: vLLM's ROCm chain. `mxfp8(x)`: one E8M0 per 32 values, + `code = clamp(ceil(log2(amax / 448)) + 127, 0, 254)`; weights in vLLM's + row-major e4m3 [N, K] with [N/32, K/32] E8M0 scales (one shared copy); KV + records as vLLM's insert writes them (dims 0..447 e4m3 with a UE8M0 per 64, + RoPE dims bf16, 584 B); keys = the sliding window plus the top-k compressed + rows, softmax in base 2 with the attention sink folded in. +- MoE: vLLM's router (bf16 gate GEMV, f32 logits, `topk_hash_softplus_sqrt`) + and AITER a8w4 `fused_moe` (interleaved gate / up, swiglu limit 10, weights + in aiter's (16, 16) shuffle and `shuffle_scale` order); the shared expert as + vLLM's MXFP8 block-32 linears with `silu_and_mul_with_clamp`. +- All-reduces: exact, fp32 in rank order -> bf16. vLLM's decode all-reduces at + M <= 48 are below its INT4 quick-reduce threshold, so vLLM runs them exact too. + +## Validation + +- Per layer against vLLM's own modules in a TP process group (`DeepseekV4MoE`, + the mHC seams, the attention chain, the all-reduce), real checkpoint weights. +- Bench: layers of one type back to back in one HIP graph, cold weights, + us per layer, M = 1..48, TP2 / TP4; then e2e (AgentX, InferenceX#3690 + recipe). diff --git a/vllm/models/deepseek_v41/amd/mono/__init__.py b/vllm/models/deepseek_v41/amd/mono/__init__.py new file mode 100644 index 000000000000..72a6ec2ab196 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: one decoder layer of a decode step in +two persistent FlyDSL launches, its TP all-reduces in-kernel (README.md). The +host side is ``runner``; vLLM docks it in ``..mono_decode``.""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/__init__.py b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py new file mode 100644 index 000000000000..da1bf8c269a0 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's attention stages: ``front`` (K1: wqkv, q / kv norms and the +KV insert, wq_b) and ``back`` (K2: split attention, combine, wo_a, wo_b).""" diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py new file mode 100644 index 000000000000..4526f8eda7ee --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py @@ -0,0 +1,884 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K2 ``back``: one DeepSeek-V4.1 attention layer from the rotated query to the +rank's wo_b partial. + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S rows: + + split (S x chunks of CK keys): flash-decoding over one chunk of a token's + keys (K1's key table: top-k compressed rows, then window rows) for + every head of the rank. The token's q in LDS; each wave one + (head tile, 16-key tile): its fp8_ds_mla rows loaded once, + dequantized into its QK operands and kept raw in LDS; QK (bf16 + MFMA), the chunk's max / sum per head, bf16 p in LDS; PV over the + chunk's raw rows (wave w: dims 64 w ..) -> the chunk's partial + (bf16) and its (max, sum) per head, published + combine (S x H): a (token, head)'s chunks merged with the sink folded in, + bf16, inverse GPT-J RoPE of dims 448.., vLLM's MXFP8 of each 32 + -> XO (published) + wo_a (G x 1024 / 32, two 16-row tiles): the group's XO rows times wo_a + -> bf16 -> MXFP8 of the 32 rows -> X8B (published) + wo_b (320, 16 rows): X8B times wo_b -> bf16 -> ``out`` + +Hand-offs are plain device-scope data behind tagged flags (half the bytes of a +tagged pair per value). Adapted from ATOM's V4.1 ``attn_post`` (K2a), which +reads a bf16 pool and ends in a peer all-reduce. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T, as_ir_value + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + WAVES, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + clamp_fp8, + first_task, + fp8_pack4, + fp8x8_bf16, + fresh, + gload, + gstore, + hw_exp2, + lane_gather, + ld_f32, + ld_i32, + mfma_bf16, + mx_code, + mx_mul, + pow2, + rsrc, + spin_until, + stamp, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + HALF, + HEAD_DIM, + HEAD_TILE, + KEYS, + NOPE, + ROPE, + ROWS, + WOA_ROWS, + Dims, + back_scratch, + cdiv, + front_end, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +NEG = -3.4028234663852886e38 +LOG2E = 1.4426950408889634 +REC = DATA + 8 # a raw record row in LDS: 576 data bytes + 8 scale bytes +Q_ROW = ( + HEAD_DIM // 2 + 4 +) # a q row in LDS (i32 words): 512 bf16 + 16 B, no bank aliasing +P_PAD = 2 # i32 words past a P row: 16 head rows read together hit distinct banks +# wo_a's K order: waves 0..3 / 4..7 take the two 16-row tiles, a contiguous K +# range a wave, summed in wave order +WOA_SPLIT_K = WAVES * ROWS // WOA_ROWS +BACK_POINTS = 16 # start, stage ends, split internals, GEMV internals + + +@dataclass(frozen=True) +class BackBuild: + tokens: int + tp: int + ratio: int # 0: window only; 1 / 2: top-k over a compressed cache + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +def head_groups(d: Dims, s: int) -> int: + """Wave groups a split unit's head tiles are spread over: two (a tile each, + 64-key chunks, twice the units) while the units fit the grid, else one (a + wave's keys loaded and dequantized once for every head tile).""" + return 2 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= BLOCKS else 1 + + +def chunk_keys(d: Dims, s: int) -> int: + """A split unit's keys: a 16-key tile a wave of a head group.""" + return 16 * (WAVES // head_groups(d, s)) + + +def p_row(d: Dims, s: int) -> int: + """A P row (one head's chunk of bf16 p) in LDS, i32 words.""" + return chunk_keys(d, s) // 2 + P_PAD + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- keys + + +def key_rec(c, t, k): + """(data address, scale address, valid) of token t's key k from K1's key + table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for + none (it reads a zero record: its scores are masked, its values zeros).""" + a, ratio = c["args"], c["ratio"] + val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1)) + ok = val != fx.Int32(-1) + slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) + if const_expr(ratio > 0): + is_c = ok & (val < 0) + blk = is_c.select(a["comp_block"], a["swa_block"]) + stride = is_c.select(a["comp_stride"], a["swa_stride"]) + base = is_c.select(a["comp"], a["swa"]) + else: + blk, stride, base = a["swa_block"], a["swa_stride"], a["swa"] + b = slot // blk + off = slot % blk + page = base + fx.Int64(b) * stride + z = a["zrec"] + data = ok.select(page + fx.Int64(off * DATA), z) + sc = ok.select(page + fx.Int64(blk * DATA + off * 8), z + fx.Int64(DATA)) + return data, sc, ok + + +def _scale_byte(scw, blk64): + return (scw[blk64 // 4] >> (blk64 % 4 * 8)) & 0xFF + + +def _fp8_pair_f32(word, scale, hi): + v = rocdl.cvt_scalef32_pk_f32_fp8( + fx.Vector.make_type(2, fx.Float32), + as_ir_value(fx.Int32(word)), + as_ir_value(fx.Float32(scale)), + hi, + ) + v = fx.Vector(v) + return [v[0], v[1]] + + +# ---------------------------------------------------------------- split + + +@traced +def stage_split(c, unit): + """Flash-decoding over chunk ``unit % nchunk`` of token ``unit // nchunk``, + every head of the rank -> the chunk's partial and (max, sum) per head. Wave + (head group hg, key tile kt): keys 16 kt .. of the chunk (loaded and + dequantized once) against the group's head tiles; then dims 64 w .. of the + PV over every key and head.""" + tid, lane, wave, d, s = c["tid"], c["lane"], c["wave"], c["d"], c["S"] + a, sl = c["args"], c["sl"] + nchunk = c["nchunk"] + t = unit // nchunk + ch = unit % nchunk + ck = chunk_keys(d, s) + ngrp = head_groups(d, s) + nkt = WAVES // ngrp # 16-key tiles a chunk + nht = d.head_tiles // ngrp # head tiles a wave + hgi = wave // nkt + kt = wave % nkt + g = lane // 16 + r16 = lane % 16 + k0 = ch * ck + # the key count, this lane's key and q: independent loads, out together + # (K1's key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t) + key = k0 + kt * 16 + r16 + data, sc, okv = key_rec(c, t, key) + qrow = d.heads * HEAD_DIM // 8 + qv = [] + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + qv.append( + fx.Vector( + bo.buffer_load( + rsrc(a["q"]), + (t * d.q_rows) // 2 + fx.min(u, qrow - 1) * 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + if k0 < kv_len: + # a key's 448 e4m3 as 7 x 16 B (lane g: dims 64 c + 16 g ..), its 64 + # bf16 rope dims as 2 x 16 B (32 c + 8 g ..), its 7 scale bytes + scw = gload(sc, words=2) + nope = [ + gload(data + fx.Int64(64 * c4 + 16 * g), words=4) + for c4 in range(NOPE // 64) + ] + rope = [ + gload(data + fx.Int64(NOPE + (cc * 32 + 8 * g) * 2), words=4) + for cc in range(ROPE // 32) + ] + # q of every head into LDS (bf16 rows of 512 + pad), 16 B a thread + for i in range_constexpr(cdiv(qrow, THREADS)): + u = fresh(tid) + THREADS * i + if u < qrow: + fx.ptr_store( + qv[i], + sl["q"] + + (u // (HEAD_DIM // 8)) * Q_ROW + + (u % (HEAD_DIM // 8)) * 4, + ) + # the raw rows to LDS for the PV (head group 0's waves) + if hgi == 0: + rawb = sl["raw"] + (kt * 16 + r16) * (REC // 4) + for c4 in range_constexpr(NOPE // 64): + fx.ptr_store(nope[c4], rawb + (64 * c4 + 16 * g) // 4) + for cc in range_constexpr(ROPE // 32): + fx.ptr_store(rope[cc], rawb + (NOPE + (cc * 32 + 8 * g) * 2) // 4) + if g == 0: + fx.ptr_store(scw, rawb + DATA // 4) + gpu.barrier() + stamp(c, 5) + # QK: A = the keys (dequantized once), B = q of a head tile (LDS); the + # contraction's 32-dim chunk 2 c + h of the nope dims is {64 c + 16 g + + # 8 h ..} (lane g's 16 B load), q read in the same order + scale = fx.Int32(a["qk_scale"]).bitcast(fx.Float32) + st = sl["st"] + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(nht)] + for c4 in range_constexpr(NOPE // 64): + ksc = pow2(_scale_byte(scw, c4)) + for h2 in range_constexpr(2): + kw = fp8x8_bf16(nope[c4][2 * h2], nope[c4][2 * h2 + 1], ksc).bitcast( + fx.BFloat16 + ) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + 32 * c4 + 8 * g + 4 * h2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + for cc in range_constexpr(ROPE // 32): + kw = rope[cc].bitcast(fx.BFloat16) + for ht in range_constexpr(nht): + qb = sl["q"] + ((hgi * nht + ht) * HEAD_TILE + r16) * Q_ROW + qw = fx.Vector( + fx.ptr_load( + qb + (NOPE + cc * 32 + 8 * g) // 2, + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + accs[ht] = mfma_bf16(kw, qw.bitcast(fx.BFloat16), accs[ht]) + # scores of keys 4 g + i (C rows), head r16 of a tile (C column) + ok = (key < kv_len) & okv + okw = ok.select(fx.Int32(1), fx.Int32(0)) + valid = [lane_gather(okw, 4 * g + i) != 0 for i in range(4)] + svs = [] + for ht in range_constexpr(nht): + sv = [ + fx.Float32(valid[i].select(accs[ht][i] * scale, fx.Float32(NEG))) + for i in range(4) + ] + mw = butterfly( + fx.max(fx.max(sv[0], sv[1]), fx.max(sv[2], sv[3])), (32, 16), fx.max + ) + if g == 0: + fx.ptr_store(mw, st + ((hgi * nht + ht) * nkt + kt) * HEAD_TILE + r16) + svs.append(sv) + gpu.barrier() + lst = st + d.heads * nkt + for ht in range_constexpr(nht): + hti = hgi * nht + ht + m = fx.ptr_load(st + (hti * nkt) * HEAD_TILE + r16) + for j in range_constexpr(1, nkt): + m = fx.max(m, fx.ptr_load(st + (hti * nkt + j) * HEAD_TILE + r16)) + sv = svs[ht] + p = [ + fx.Float32(valid[i].select(hw_exp2(sv[i] - m), fx.Float32(0.0))) + for i in range(4) + ] + lw = (p[0] + p[2]) + (p[1] + p[3]) + lw = lw + xshfl(lw, 32) + lw = lw + xshfl(lw, 16) + pw = ( + fx.Vector.from_elements(p, fx.Float32).to(fx.BFloat16).bitcast(fx.Int32) + ) + pbase = ( + sl["p"] + (hti * HEAD_TILE + r16) * p_row(d, s) + (kt * 16 + 4 * g) // 2 + ) + fx.ptr_store(fx.Vector.from_elements([pw[0], pw[1]], fx.Int32), pbase) + if g == 0: + fx.ptr_store(lw, lst + (hti * nkt + kt) * HEAD_TILE + r16) + gpu.barrier() + # chunk stats: (max, sum) per head + if tid < d.heads: + hts = tid // HEAD_TILE + hr = tid % HEAD_TILE + mm = fx.ptr_load(st + (hts * nkt) * HEAD_TILE + hr) + ll = fx.ptr_load(lst + (hts * nkt) * HEAD_TILE + hr) + for j in range_constexpr(1, nkt): + mm = fx.max(mm, fx.ptr_load(st + (hts * nkt + j) * HEAD_TILE + hr)) + ll = ll + fx.ptr_load(lst + (hts * nkt + j) * HEAD_TILE + hr) + bo.buffer_store( + fx.Vector.from_elements([mm, ll], fx.Float32), + rsrc(c["pstat"]), + ((t * nchunk + ch) * d.heads + tid) * 2, + cache_modifier=CM_DEV, + ) + # PV: wave w dims 64 w + 4 r16 + dt (dt = 0..3: one LDS word holds the + # four of a key); the rope wave (dims 448..) reads bf16 pairs. Lane + # (r16, g) takes keys 4 g + j of a 16-key tile: A[dim][key] + is_rope = wave == WAVES - 1 + dbase = 64 * wave + sh = (wave % 4) * 8 + accs = [ + [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(4)] + for _ in range(d.head_tiles) + ] + for kt2 in range_constexpr(nkt): + rb0 = sl["raw"] + (kt2 * 16 + 4 * g) * (REC // 4) + vv = fx.Vector.filled(16, 0.0, fx.Float32) # [dt][j] + if is_rope: + rv = [] + for j in range_constexpr(4): + rw = fx.Vector( + fx.ptr_load( + rb0 + j * (REC // 4) + (NOPE + 8 * r16) // 4, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ) + rv.append([bf_lo(rw[0]), bf_hi(rw[0]), bf_lo(rw[1]), bf_hi(rw[1])]) + vv = fx.Vector.from_elements( + [rv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + else: + nv = [] + for j in range_constexpr(4): + rb = rb0 + j * (REC // 4) + w = fx.Int32(fx.ptr_load(rb + (dbase + 4 * r16) // 4)) + scb = fx.Int32(fx.ptr_load(rb + DATA // 4 + wave // 4)) + scl = pow2((scb >> sh) & 0xFF) + nv.append( + _fp8_pair_f32(w, scl, False) + _fp8_pair_f32(w, scl, True) + ) + vv = fx.Vector.from_elements( + [nv[j][dt] for dt in range(4) for j in range(4)], fx.Float32 + ) + for ht2 in range_constexpr(d.head_tiles): + pb = fx.Vector( + fx.ptr_load( + sl["p"] + + (ht2 * HEAD_TILE + r16) * p_row(d, s) + + (kt2 * 16 + 4 * g) // 2, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ).bitcast(fx.Int16) + for dt in range_constexpr(4): + av = ( + fx.Vector.from_elements( + [vv[dt * 4 + j] for j in range(4)], fx.Float32 + ) + .to(fx.BFloat16) + .bitcast(fx.Int16) + ) + accs[ht2][dt] = fx.Vector( + rocdl.mfma_f32_16x16x16bf16_1k( + T.vec(4, T.f32), [av, pb, accs[ht2][dt]] + ) + ) + stamp(c, 6) + # partial[t][ch][head][dim] bf16: lane (r16, g) holds C rows 4 g + i -> + # dims 64 w + 16 g + 4 i + dt of head ht2 * 16 + r16: 32 contiguous + # bytes, two 16 B stores (4 lanes fill a 128 B line) + for ht2 in range_constexpr(d.head_tiles): + hh = ht2 * HEAD_TILE + r16 + words = [] + for i in range_constexpr(4): + v4 = fx.Vector.from_elements( + [accs[ht2][dt][i] for dt in range(4)], fx.Float32 + ) + w2 = v4.to(fx.BFloat16).bitcast(fx.Int32) + words += [w2[0], w2[1]] + off = (((t * nchunk + ch) * d.heads + hh) * HEAD_DIM + dbase + 16 * g) // 2 + for hlf in range_constexpr(2): + bo.buffer_store( + fx.Vector.from_elements(words[4 * hlf : 4 * hlf + 4], fx.Int32), + rsrc(c["part"]), + off + 4 * hlf, + cache_modifier=CM_DEV, + ) + publish(c, c["srdy"], t * nchunk + ch) + stamp(c, 7) + + +# ---------------------------------------------------------------- combine + +COMB_HEADS = WAVES # heads a combine unit: a wave each + + +@traced +def stage_combine(c, unit): + """(token, 8 heads) ``unit``, a wave a head, 8 dims a lane: the chunks + merged with the sink folded in, bf16, inverse RoPE, MXFP8 of each 32 -> XO, + published (a flag a (token, head group)).""" + lane, wave, d, s = c["lane"], c["wave"], c["d"], c["S"] + a = c["args"] + hg_n = d.heads // COMB_HEADS + t = unit // hg_n + hg = unit % hg_n + h = hg * COMB_HEADS + wave + nchunk = c["nchunk"] + ck = chunk_keys(d, s) + kv_len = ld_i32(a["klen"], t) + nlive = (kv_len + ck - 1) // ck + pos = ld_i32(a["pos"], 2 * t) + sink = ld_f32(a["sink"], h) * fx.Float32(LOG2E) + d0 = lane * 8 # this lane's 8 dims + rope_lane = d0 >= NOPE + jj = fx.max((d0 - NOPE) // 2, 0) # 4 pairs from here + cs = [ + rope_lane.select(ld_f32(a["cos_sin"], pos * ROPE + jj + i), fx.Float32(1.0)) + for i in range(4) + ] + sn = [ + rope_lane.select( + ld_f32(a["cos_sin"], pos * ROPE + HALF + jj + i), fx.Float32(0.0) + ) + for i in range(4) + ] + await_flags(c, c["srdy"], nlive, lambda u: t * nchunk + u) + m = fx.Float32(NEG) + ms, ls, vs = [], [], [] + for cc in range_constexpr(nchunk): + live = cc < nlive + stv = fx.Vector( + bo.buffer_load( + rsrc(c["pstat"]), + ((t * nchunk + cc) * d.heads + h) * 2, + vec_width=2, + dtype=T.f32, + ) + ) + pv = fx.Vector( + bo.buffer_load( + rsrc(c["part"]), + (((t * nchunk + cc) * d.heads + h) * HEAD_DIM + d0) // 2, + vec_width=4, + dtype=T.i32, + ) + ) + mc = fx.Float32(live.select(stv[0], fx.Float32(NEG))) + ms.append(mc) + ls.append(fx.Float32(live.select(stv[1], fx.Float32(0.0)))) + vals = [] + for e in range_constexpr(4): + vals += [bf_lo(pv[e]), bf_hi(pv[e])] + vs.append([fx.Float32(live.select(v, fx.Float32(0.0))) for v in vals]) + m = fx.max(m, mc) + m_final = fx.max(m, sink) + lsum = hw_exp2(sink - m_final) + tot = [fx.Float32(0.0)] * 8 + for cc in range_constexpr(nchunk): + wgt = hw_exp2(ms[cc] - m_final) + lsum = lsum + ls[cc] * wgt + tot = [tot[e] + vs[cc][e] * wgt for e in range(8)] + inv = fx.Float32(1.0) / lsum + o = [bf16_round(tot[e] * inv) for e in range(8)] + # inverse GPT-J RoPE of each pair: e' = e cos + o sin, o' = o cos - e sin + y = [] + for i in range_constexpr(4): + e, od = o[2 * i], o[2 * i + 1] + y += [e * cs[i] + od * sn[i], od * cs[i] - e * sn[i]] + amax = abs(y[0]) + for v in y[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) # 32 dims = 4 lanes + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(y[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(y[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["xo"]), + t * (d.q_rows // 4) + (h * HEAD_DIM + d0) // 4, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["xos"]), + t * (d.q_rows // 32) + (h * HEAD_DIM + d0) // 32, + cache_modifier=CM_DEV, + ) + publish(c, c["crdy"], t * hg_n + hg) + + +# ---------------------------------------------------------------- wo_a / wo_b + + +@traced +def stage_woa(c, task, ops, t0, n): + """Rows 32 task .. of wo_a -> bf16 -> vLLM's MXFP8 of the 32 -> X8B.""" + tid, red, d, lane = c["tid"], c["red"], c["d"], c["lane"] + split = WOA_SPLIT_K + gemv_mfmas(c, d.group_k, c["xl"], c["xsl"], red, ops, split, tiled=True, rows=n) + gpu.barrier() + mine = tid < WOA_ROWS * n + tl = fx.min(tid // WOA_ROWS, n - 1) + t = plus(t0, tl) + row = tid % WOA_ROWS + tile = row // ROWS + y = bf16_round( + row_sum(red, row % ROWS, tl, [tile * split + w for w in range(split)]) + ) + amax = butterfly(abs(y), (1, 2, 4, 8, 16), fx.max) + code = mx_code(amax) + q = clamp_fp8(y * mx_mul(code)) + w = fp8_pack4( + q, + *[ + lane_gather(q.bitcast(fx.Int32), fx.min(lane + k, 63)).bitcast(fx.Float32) + for k in (1, 2, 3) + ], + ) + if mine & (row % 4 == 0): + bo.buffer_store( + w, + rsrc(c["x8b"]), + t * (d.o_rows // 4) + (task * WOA_ROWS + row) // 4, + cache_modifier=CM_DEV, + ) + if mine & (row == 0): + bo.buffer_store( + code, rsrc(c["x8bs"]), t * (d.o_rows // 32) + task, cache_modifier=CM_DEV + ) + gpu.barrier() + + +def woa_loads(c, task): + a = c["args"] + return gemv_loads( + c, + a["woa"], + a["woa_s"], + c["d"].group_k, + task * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + ) + + +def wob_loads(c, task, pred=None): + a = c["args"] + return gemv_loads(c, a["wob"], a["wob_s"], c["d"].o_rows, task * ROWS, pred=pred) + + +def _wob_store(a, t, col, v0, v1): + word = ( + fx.Vector.from_elements([v0, v1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["out"]), (t * a["out_stride"] + col) // 2) + + +@traced +def stage_wob(c, task, ops, task2, ops2, has2, t0, n, mark=False): + """Rows 16 task .. (and 16 task2 .., ``has2``) of wo_b -> bf16 -> out, or + (``c["wob_out"]``, the mono layer) the caller's hook with each bf16 pair.""" + tid, red = c["tid"], c["red"] + o_rows = c["d"].o_rows + a = c["args"] + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red, ops, rows=n) + gemv_mfmas(c, o_rows, c["xl"], c["xsl"], red + RED, ops2, rows=n) + gpu.barrier() + if const_expr(mark): + stamp(c, 15) + per = ROWS * n // 2 # bf16 pairs a task + k = tid // per + emit = c.get("wob_out") or (lambda t, col, v0, v1: _wob_store(a, t, col, v0, v1)) + if (tid < 2 * per) & ((k == 0) | has2): + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + rk = red + k * RED + col = (k == 0).select(task, task2) * ROWS + 2 * rp + v0 = row_sum(rk, 2 * rp, tl) + v1 = row_sum(rk, 2 * rp + 1, tl) + emit(t, col, v0, v1) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def back_plan(key): + """(nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob) of a build.""" + s, d = key.tokens, Dims(key.tp) + ck = chunk_keys(d, s) + nchunk = cdiv(KEYS, ck) + n_split = s * nchunk + n_comb = s * (d.heads // COMB_HEADS) + COMB0 = n_split % BLOCKS + assert d.woa_tasks <= BLOCKS + assert BLOCKS < d.wob_tasks <= 2 * BLOCKS + # the GEMVs past the attention units, so CTAs without attention work have + # their weights in flight from the start + WOA0 = (n_split + n_comb) % BLOCKS + WOB0 = (WOA0 + d.woa_tasks) % BLOCKS + # wo_b's weights before the wo_a stage only while the wo_a CTAs have no + # attention work (loads complete in issue order: else they would hold up + # wo_a's activation rows) + early_wob = n_split + n_comb + d.woa_tasks <= BLOCKS + return nchunk, n_split, n_comb, COMB0, WOA0, WOB0, early_wob + + +def back_lds_members(s: int, d: Dims) -> dict: + """K2's attention LDS structs: the split stage's and the GEMVs' (a union's + members: one at a time).""" + rows = tile_rows(s) + ck = chunk_keys(d, s) + + @fx.struct + class SplitLds: + q: fx.Array[fx.Int32, d.heads * Q_ROW, 16] + raw: fx.Array[fx.Int32, ck * REC // 4, 16] + p: fx.Array[fx.Int32, d.heads * p_row(d, s), 16] + st: fx.Array[fx.Float32, 2 * d.heads * (WAVES // head_groups(d, s)), 16] + + @fx.struct + class GemvLds: + red: fx.Array[fx.Float32, 2 * RED, 16] + xl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 4), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(d.group_k // 32), 16] + + return {"split": SplitLds, "gemv": GemvLds} + + +def back_lds(s: int, d: Dims): + """K2's attention LDS: the union of ``back_lds_members``.""" + m = back_lds_members(s, d) + SplitLds, GemvLds = m["split"], m["gemv"] + + @fx.union + class BackLds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + + return BackLds + + +def back_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The attention back's context: LDS views (``lds``: an allocated + ``back_lds``, or a dict of the members' views), mailbox (tag ``tag``), + scratch regions, arguments by name.""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + if isinstance(lds, dict): + spl, gml = lds["split"], lds["gemv"] + else: + spl, gml = lds.split.peek(), lds.gemv.peek() + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": gml.red.ptr, + "xl": gml.xl.ptr, + "xsl": gml.xsl.ptr, + "d": d, + "ratio": key.ratio, + "mb": Mailbox(tag), + "nchunk": back_plan(key)[0], + "sl": {"q": spl.q.ptr, "raw": spl.raw.ptr, "p": spl.p.ptr, "st": spl.st.ptr}, + "args": args, + } + for region, (off, _) in ( + layout or back_scratch(s, d, start=front_end(s, d)) + ).items(): + c[region] = scratch + fx.Int64(off) + return c + + +EPOCH_MARKS = 4 # word offset of the per-CTA marks in the epoch buffer + + +@traced +def epoch_begin(c, epoch, ep): + """This CTA has read the launch pair's epoch: its own mark word (no + contended atomic) := ep + 1, which ``epoch_end`` checks.""" + if c["tid"] == 0: + bo.buffer_store( + ep + 1, rsrc(epoch), EPOCH_MARKS + c["bid"], cache_modifier=CM_DEV + ) + + +@traced +def epoch_end(c, epoch, ep, reset=None): + """CTA 0, once every CTA marked this epoch (thread u checks CTA u's mark): + ``reset()`` (state later pairs must find cleared), then the epoch moved on. + The next pair's hand-offs carry a tag no earlier pair wrote; later launches' + kernels start after this one ends.""" + if c["bid"] == 0: + tid = c["tid"] + if tid < BLOCKS: + r = rsrc(epoch) + + def load_mark(): + return fx.Int32( + bo.buffer_load( + r, + EPOCH_MARKS + tid, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + + _ = spin_until(load_mark, lambda v: v != ep + 1) + gpu.barrier() + if tid == 0: + if const_expr(reset is not None): + reset() + gstore(epoch, ep + 1, words=1) + + +@traced +def run_back(c, key, bid): + """The attention back's stages on this CTA: split, combine, wo_a, wo_b.""" + s = key.tokens + dd = Dims(key.tp) + # plain ints: a Python object read inside traced control flow is carried + woa_tasks, groups, heads, group_k = dd.woa_tasks, dd.groups, dd.heads, dd.group_k + q_rows, o_rows, wob_tasks = dd.q_rows, dd.o_rows, dd.wob_tasks + a = c["args"] + tiles = token_tiles(s) + _, n_split, n_comb, COMB0, WOA0, WOB0, early_wob = back_plan(key) + stamp(c, 0) + for unit in range(first_task(bid, 0), n_split, BLOCKS): + stage_split(c, unit) + stamp(c, 1) + for unit in range(first_task(bid, COMB0), n_comb, BLOCKS): + stage_combine(c, unit) + gpu.barrier() + stamp(c, 2) + # the GEMV weights, every CTA right after its attention work: its wo_a + # task (if any), then its one or two wo_b tasks + woa_task = first_task(bid, WOA0) + has_woa = woa_task < woa_tasks + woa_t = fx.min(woa_task, woa_tasks - 1) + ops = gemv_loads( + c, + a["woa"], + a["woa_s"], + group_k, + woa_t * WOA_ROWS, + WOA_SPLIT_K, + tiled=True, + pred=has_woa, + ) + wob_task = first_task(bid, WOB0) + wob_second = wob_task + BLOCKS + has_second = wob_second < wob_tasks + + def wob_weights(): + return ( + wob_loads(c, wob_task), + wob_loads(c, fx.min(wob_second, wob_tasks - 1), pred=has_second), + ) + + if const_expr(early_wob): + ops1, ops2 = wob_weights() + if const_expr(c.get("tl") is not None): + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + stamp(c, 8) + # wo_a: every token tile of its group's XO rows, the next tile's in flight + if has_woa: + grp = woa_t // (woa_tasks // groups) + await_flags(c, c["crdy"], s, lambda u: u * (heads // COMB_HEADS) + grp) + stamp(c, 9) + xo = ( + c["xo"], + c["xos"], + q_rows // 4, + grp * group_k // 4, + group_k // 4, + q_rows // 32, + grp * group_k // 32, + group_k // 32, + ) + regs = rows_loads(c, *xo, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, group_k // 4, group_k // 32, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *xo, *tiles[ti + 1]) + stage_woa(c, woa_t, ops, t0, n) + stamp(c, 10) + publish(c, c["ardy"], woa_t) + stamp(c, 3) + # wo_b + if const_expr(not early_wob): + ops1, ops2 = wob_weights() + stamp(c, 11) + await_flags(c, c["ardy"], woa_tasks, lambda u: u) + stamp(c, 12) + x8b = ( + c["x8b"], + c["x8bs"], + o_rows // 4, + 0, + o_rows // 4, + o_rows // 32, + 0, + o_rows // 32, + ) + regs = rows_loads(c, *x8b, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, o_rows // 4, o_rows // 32, c["xl"], c["xsl"], n) + if const_expr(ti == 0): + stamp(c, 14) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *x8b, *tiles[ti + 1]) + stage_wob( + c, + wob_task, + ops1, + fx.min(wob_second, wob_tasks - 1), + ops2, + has_second, + t0, + n, + mark=ti == 0, + ) + if const_expr(ti == 0): + stamp(c, 13) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/device.py b/vllm/models/deepseek_v41/amd/mono/attention/device.py new file mode 100644 index 000000000000..ec230cd7074a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/device.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device primitives of the DeepSeek-V4.1 attention mega kernel. + +Wave reductions and shuffles, hardware math, bf16 / fp8 / MX packing, the +scaled MFMA, buffer and 64-bit global accesses, and the tagged-pair mailbox. +Adapted from ATOM's mono decode framework (``atom/mono/device``), itself ported +from the FlyDSL GLM-5 layer kernel. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +# The execution model: one CTA a CU of an MI355X, all resident. +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16, NT = 2. SC1 alone is +# device scope (past the per-XCD L2s); NT streams data read once. +CM_DEV = 16 +CM_NT = 2 + +FP8_MAX = 448.0 +FP8, FP4 = 0, 4 # f8f6f4 MFMA operand formats +UNIT_SCALE = 127 # E8M0 code of 1.0 +FLT_MIN = 1.1754943508222875e-38 + +POLL_MAX = 12 # mailbox pairs polled a batch + + +def traced(fn): + """FlyDSL's AST rewrite for a helper, so its ``if`` / ``while`` / ``for`` on + traced values lower to scf ops (kernels rewrite only their own body).""" + return ASTRewriter.transform(fn) + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``aiter::_``.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN5aiter{len(name)}{name}E" + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base``, then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ----------------------------------------------------------------------------- memory + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def ld_i32(ptr, i, cm=0): + return fx.Int32( + bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.i32, cache_modifier=cm) + ) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def ld_i64(ptr, i): + """Element i of an i64 buffer (two dwords).""" + w = fx.Vector(bo.buffer_load(rsrc(ptr), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(w[1]) << 32) | fx.Int64(fx.Uint32(w[0])) + + +def _gptr(addr): + return llvm.IntToPtrOp( + ir.Type.parse("!llvm.ptr<1>"), as_ir_value(fx.Int64(addr)) + ).result + + +def gload(addr, words=4, nt=False): + """``words`` i32 (1, 2 or 4) at a 64-bit global address: a cache past 4 GiB is + reachable only this way (a buffer offset is 32 bits).""" + i32 = ir.IntegerType.get_signless(32) + ty = i32 if words == 1 else ir.VectorType.get([words], i32) + v = llvm.LoadOp(ty, _gptr(addr), alignment=4 * words, nontemporal=nt or None).result + return fx.Int32(v) if words == 1 else fx.Vector(v) + + +def atomic_add_dev(addr, value): + """Device-scope atomic add on the global i32 at ``addr``; the old value. + Relaxed: an acquire / release here would write back and invalidate L2.""" + return fx.Int32( + fx.atomic_add( + _gptr(addr), + fx.Int32(value), + syncscope="agent", + ordering=fx.AtomicOrdering.Monotonic, + ) + ) + + +def gstore(addr, value, words=4): + """Store ``words`` i32 (an Int32 or a Vector) at a 64-bit global address.""" + v = as_ir_value(fx.Int32(value) if words == 1 else value) + llvm.StoreOp(v, _gptr(addr), alignment=4 * words) + + +# ----------------------------------------------------------------------------- math + + +def fresh(v): + """``v`` through an opaque move, so the lane math built on it is computed at + its use instead of being hoisted out of a task loop and held in registers.""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def tile_map(tid, width): + """(tid // width, tid % width) for 0 <= tid < THREADS without a division: + shifts for a power of two, else a multiply-shift exact on that range.""" + if width & (width - 1) == 0: + sh = width.bit_length() - 1 + return tid >> sh, tid & (width - 1) + m = -(-(1 << 20) // width) + assert all((t * m) >> 20 == t // width for t in range(THREADS)), width + row = (tid * m) >> 20 + return row, tid - row * width + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def fma(a, b, c): + return fx.Float32(fmath.fma(fx.Float32(a), fx.Float32(b), fx.Float32(c))) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest (IEEE), given ``r = RN(1 / d)``.""" + q = x * r + return fma(fma(-q, d, x), r, q) + + +def permlane_swap(off, x, y): + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def lane_f32(v, src_lane): + return lane_gather(v.bitcast(fx.Int32), src_lane).bitcast(fx.Float32) + + +def readlane(v, src_lane): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def xshfl(v, off): + """Value of lane ``lane ^ off`` (DPP in a row, permlane swaps across).""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the xor ``offsets``, in that order.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +def row_shr(v, lane, k): + """Lane - k of the same 16-lane row, else 0 (DPP row_shr).""" + moved = lane_gather(v.bitcast(fx.Int32), fx.max(lane - k, 0)).bitcast(fx.Float32) + return fx.Float32((lane % 16 >= k).select(moved, fx.Float32(0.0))) + + +# ----------------------------------------------------------------------------- packing + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def fp8_pack4(a, b, c, d): + """Four f32 (in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +def fp8x8_bf16(w0, w1, scale): + """Eight E4M3 of two dwords times an f32 ``scale`` -> eight bf16 as a + 4-dword Vector (exact for a power-of-two scale in range).""" + out = [] + for w in (w0, w1): + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(scale)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32) + + +def clamp_fp8(v): + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def pow2(code): + """The fp32 power of two a biased exponent (1..254) stands for.""" + return (code << 23).bitcast(fx.Float32) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact for code in [1, 253].""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def ceil_exp(v): + """Biased exponent of the smallest power of two >= v (v > 0 finite, normal).""" + bits = v.bitcast(fx.Int32) + return ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + + +def mx_code(amax): + """VLLM's ROCm MXFP8 scale code of a 32-group: clamp(ceil(log2(amax / 448)) + + 127, 0, 254), amax floored at FLT_MIN.""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + return fx.min(fx.max(ceil_exp(q), fx.Int32(0)), fx.Int32(254)) + + +def mx_mul(code): + """2^(127 - code) as fp32, for code in [0, 254] (the quant multiplier).""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +# ----------------------------------------------------------------------------- MFMA + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8): + """16x16x128 scaled MFMA (gfx950 f8f6f4): lane l (row l % 16) holds K bytes + 16 (l // 16) .. +16 and 64 + 16 (l // 16) .. +16 of the 128-step, and the + E8M0 of K block l // 16.""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, 0, sa, 0, sb] + ) + ) + + +def mfma_bf16(a, b, c): + """16x16x32 bf16 MFMA: lane l (row / col l % 16) holds K 8 (l // 16) .. +8.""" + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +# ----------------------------------------------------------------------------- mailbox + + +@traced +def _spin(load_all, pending): + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +def spin_until(load, pending): + """``load()`` again (s_sleep between) while ``pending(value)``; the value.""" + return _spin(load, pending) + + +class Mailbox: + """Tagged-pair hand-off between the CTAs of one launch (device scope). + + Every 32-bit value is stored next to the launch's tag in one 8-byte store; a + consumer polls until every tag matches, so a hand-off is one round trip with + no separate flag. The tag is ``epoch * 64 + layer + 1`` (``epoch`` a device + counter bumped once a step): a pair from an earlier layer or step never + carries a live tag, and nothing is cleared between launches.""" + + def __init__(self, tag): + self.tag = tag + + def put(self, base, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def put_bf(self, base, i, vs): + """Elements i .. i + len(vs) (2 or 4, i aligned) as bf16 pairs, one store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(base), + i, + cache_modifier=CM_DEV, + ) + + def put_words(self, base, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. (fp8 words must come this way).""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(base), + i * 2, + cache_modifier=CM_DEV, + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of ``specs`` = [(base, pair index, npairs in {1, 2})]: all + re-loaded while any tag is stale. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + tag = self.tag + + def load_all(): + words = [] + for b, i, n in specs: + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +LOAD_BATCH = 4 + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over v = tid + THREADS i < n: ``batch`` rounds' loads + issued before their uses; a round past n loads element n - 1, uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def barrier(): + gpu.barrier() + + +# ----------------------------------------------------------------------------- timeline + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock.""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +@traced +def stamp(c, point): + """A timeline build's stamp: thread 0 of the CTA stores the clock at + ``point`` (``c["tl"]`` [BLOCKS, points] i64; None: not a timeline build).""" + if const_expr(c.get("tl") is None): + return + if c["tid"] == 0: + t = memrealtime() + w = fx.Vector.from_elements([fx.Int32(t), fx.Int32(t >> 32)], fx.Int32) + bo.buffer_store(w, rsrc(c["tl"]), (c["bid"] * c["tl_points"] + point) * 2) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py new file mode 100644 index 000000000000..3ac9e3fe9b10 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K1 ``front``: one DeepSeek-V4.1 attention layer from its normed input row to +the rotated query and the step's KV rows in the sliding-window cache. + +One launch a layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, S rows. +Every CTA issues its weight loads (its wqkv part and its wq_b tasks) at the +kernel start, before it waits on anything; the stages: + + xq (S > 16: S x 2 K halves, on the 32 CTAs without a wqkv part): the + bf16 input's half -> vLLM's MXFP8 -> X8 (published, a flag each) + kt (S): the token's attention keys for K2 -- top-k compressed slots + through the kv source's block table, then window slots -> KT + wqkv_a (112 row tiles x 2 K halves, 16 rows): the input's MXFP8 half -- + quantized in the CTA (S <= 16) or X8's -- a token tile at a time + (the next tile's rows in flight) times wqkv's -> fp32 partials -> + QKVP + qkv (S): the halves summed -> bf16; q_norm -> bf16 -> MXFP8 -> QX8 + (published, a flag a token); kv_norm -> bf16 -> GPT-J RoPE -> bf16 + -> the fp8_ds_mla record at the token's slot + wq_b (H x 512 / 16, 16 rows): QX8 times wq_b -> bf16 -> GPT-J RoPE of + each head's dims 448.. -> q (global, bf16) + +Each elementwise step follows vLLM's ROCm path (``mxfp8_quantize``, +``FusedQKVRMSNorm``, ``fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert``); +the GEMV reductions differ from the Triton GEMMs' only in their order. +Adapted from ATOM's V4.1 ``attn_pre`` (K1), less its mHC seam. +""" + +from dataclasses import dataclass + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from .device import ( + BLOCKS, + CM_DEV, + THREADS, + Mailbox, + bf16_round, + bf_hi, + bf_lo, + butterfly, + ceil_exp, + clamp_fp8, + div_rn, + first_task, + fp8_pack4, + fresh, + gstore, + hw_rsq, + ld_f32, + ld_i32, + mx_code, + mx_mul, + rsrc, + stamp, + tile_map, + traced, + xshfl, +) +from .gemv import RED, gemv_loads, gemv_mfmas, lds_row, row_sum, rows_loads, rows_store +from .plan import ( + DATA, + EPS, + HALF, + HEAD_DIM, + HIDDEN, + KEYS, + KV_DIM, + NOPE, + Q_RANK, + QKV_ROWS, + ROPE, + ROWS, + TILE, + TOPK, + WINDOW, + WQKV_KPARTS, + WQKV_TASKS, + Dims, + front_scratch, + tile_rows, + token_tiles, +) + +_STREAM = fx.Stream(None) +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +QX_WORDS = Q_RANK // 4 +KSTEPS = HIDDEN // 128 // WQKV_KPARTS # a wqkv part's 128-steps +KWORDS = KSTEPS * 32 +KGROUPS = KSTEPS * 4 +XH = HIDDEN // WQKV_KPARTS # a K half's values +XH_COLS = XH // 8 // 2 # 16 B items (8 values) a half of a K half +XH_RPR = THREADS // XH_COLS # half-rows a round +XQ_ITEMS = XH // 8 # an xq unit's 16 B items (one a thread) +FRONT_POINTS = 6 # start, kt done, wqkv done, qkv done, wq_b done, (spare) + + +@dataclass(frozen=True) +class FrontBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (0: no top-k keys) + timeline: bool = False + + +def plus(t0, v): + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +@traced +def publish(c, region, i): + """This CTA's stores visible device-wide, then the flag pair (region, i).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] == 0: + c["mb"].put(region, i, fx.Int32(1)) + + +@traced +def await_flags(c, region, n, index): + """Thread u < n polls flag ``index(u)``; then the CTA proceeds.""" + tid = c["tid"] + if tid < n: + c["mb"].poll([(region, index(fx.min(tid, n - 1)), 1)]) + gpu.barrier() + + +# ---------------------------------------------------------------- xq + + +def xq_loads(c, unit, pred): + """(token, K half) ``unit``'s bf16 input: its load only, 8 values a thread + (threads past the half mirror its last item).""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + col = fx.min(c["tid"], XQ_ITEMS - 1) + e = (unit >> 1) * a["x_stride"] + (unit & 1) * XH + col * 8 + return fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32)) + + +@traced +def stage_xq(c, unit, w, live): + """``xq_loads``' values -> vLLM's MXFP8 -> X8 words / X8S codes of the + unit (published by the caller).""" + tid, lane = c["tid"], c["lane"] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + t = unit >> 1 + part = unit & 1 + if live & (tid < XQ_ITEMS): + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + part * KWORDS + tid * 2, + cache_modifier=CM_DEV, + ) + if lane % 4 == 0: + bo.buffer_store( + code, + rsrc(c["x8s"]), + t * X_GROUPS + part * KGROUPS + (tid >> 2), + cache_modifier=CM_DEV, + ) + + +# ---------------------------------------------------------------- kt + + +@traced +def stage_kt(c, t): + """Token t's keys for K2: ``KT[t][k] = slot | 1 << 31`` for its top-k compressed + rows (k < ntopk), the window slot for k in [ntopk, kv_len), -1 for an + invalid or absent key; KLEN[t] = kv_len (0 for a pad row), KLEN[S + t] = + ntopk.""" + tid = c["tid"] + a, ratio = c["args"], c["ratio"] + slot = ld_i32(a["slot"], 2 * t) + live = slot >= 0 + nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) + if const_expr(ratio > 0): + pos = ld_i32(a["pos"], 2 * t) + ntopk = live.select(fx.min((pos + 1) // ratio, fx.Int32(TOPK)), fx.Int32(0)) + req = ld_i32(a["t2r"], t) + else: + ntopk = fx.Int32(0) + kv_len = ntopk + nswa + for i in range_constexpr((KEYS + THREADS - 1) // THREADS): + k = fresh(tid) + THREADS * i + sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1) + s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx) + val = (k < kv_len).select(s_slot, fx.Int32(-1)) + if const_expr(ratio > 0): + local = ld_i32(a["topk"], t * TOPK + fx.min(k, TOPK - 1)) + lc = fx.max(local, 0) + e = a["comp_block"] + blk = ld_i32(a["comp_bt"], req * a["bt_stride"] + lc // e) + c_slot = (blk * e + lc % e) | fx.Int32(-2147483648) + c_val = (local >= 0).select(c_slot, fx.Int32(-1)) + val = (k < ntopk).select(c_val, val) + if k < KEYS: + bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k) + if tid == 0: + bo.buffer_store(kv_len, rsrc(a["klen"]), t) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t) + + +# ---------------------------------------------------------------- wqkv_a + + +def x_half_loads(c, part, pred, t0, n): + """Tokens t0 .. t0 + n's bf16 input, K half ``part``: its loads only (16 B a + lane-item, 8 values; a thread one column of [2 n half-rows, 160]), masked + off on CTAs without a wqkv part.""" + a = c["args"] + xr = rsrc(a["x"], pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0))) + hr0, col = tile_map(c["tid"], XH_COLS) + e0 = part * XH + col * 8 + out = [] + for i in range((2 * n + XH_RPR - 1) // XH_RPR): + hr = fx.min(hr0 + XH_RPR * i, 2 * n - 1) + e = (t0 + (hr >> 1)) * a["x_stride"] + (hr & 1) * (XH // 2) + e0 + out.append(fx.Vector(bo.buffer_load(xr, e >> 1, vec_width=4, dtype=T.i32))) + return out + + +@traced +def quant_x_half(c, regs, n): + """``x_half_loads``' values -> vLLM's MXFP8 rows of the K half in LDS (a + 32-group = 4 lanes).""" + tid, lane = c["tid"], c["lane"] + x_row, s_row = lds_row(KWORDS), lds_row(KGROUPS) + hr0, col = tile_map(tid, XH_COLS) + act = tid < XH_RPR * XH_COLS + for i in range_constexpr(len(regs)): + hr = hr0 + XH_RPR * i + w = regs[i] + vals = [] + for d in range_constexpr(4): + vals += [bf_lo(w[d]), bf_hi(w[d])] + amax = abs(vals[0]) + for v in vals[1:]: + amax = fx.max(amax, abs(v)) + amax = butterfly(amax, (1, 2), fx.max) + code = mx_code(amax) + mul = mx_mul(code) + w0 = fp8_pack4(*[clamp_fp8(vals[q] * mul) for q in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(vals[4 + q] * mul) for q in range(4)]) + tok = hr >> 1 + half = hr & 1 + if act & (hr < 2 * n): + fx.ptr_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + c["xl"] + (tok * x_row + half * (KWORDS // 2) + col * 2), + ) + if lane % 4 == 0: + fx.ptr_store( + code, c["xsl"] + (tok * s_row + half * (KGROUPS // 2) + (col >> 2)) + ) + gpu.barrier() + + +@traced +def stage_wqkv(c, task, part, ops, t0, n, live): + """Rows 16 task .. of wqkv over K half ``part`` against tokens t0 .. t0 + n + -> fp32 partials QKVP[part] (``live``: a CTA with a wqkv part).""" + tid, red = c["tid"], c["red"] + gemv_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, ops, rows=n, nsteps=KSTEPS) + gpu.barrier() + if (tid < ROWS * n) & live: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + v = row_sum(red, r, tl) + c["mb"].put(c["qkvp"], (part * c["S"] + t) * QKV_ROWS + task * ROWS + r, v) + gpu.barrier() + + +# ---------------------------------------------------------------- q / kv norms + + +@traced +def stage_qkv(c, t): + """Token t: the K halves summed -> bf16; q_norm -> bf16 -> MXFP8 (QX8); + kv_norm -> bf16 -> RoPE -> bf16 -> the fp8_ds_mla record at + slot_mapping[t]. 256 threads (8 q columns and 2 kv columns each); threads + 256.. mirror thread 255 and store nothing.""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, s = c["args"], c["S"] + tt = fx.min(tid, 255) + mine = tid < 256 + qc = [fx.min(tt * 8 + j, Q_RANK - 1) for j in range(8)] + # independent loads first: they overlap the partials' wait + qw = [c["ldbf"](a["qn_w"], qc[j]) for j in range(8)] + kw = [c["ldbf"](a["kvn_w"], tt * 2 + j) for j in range(2)] + pos = ld_i32(a["pos"], 2 * t) + slot = ld_i32(a["slot"], 2 * t) + jr = tt - NOPE // 2 + is_rope = jr >= 0 + jj = fx.max(jr, 0) + cs = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + jj), fx.Float32(1.0)) + sn = is_rope.select(ld_f32(a["cos_sin"], pos * ROPE + HALF + jj), fx.Float32(0.0)) + specs = [] + for p in range_constexpr(2): + specs += [(c["qkvp"], (p * s + t) * QKV_ROWS + qc[j], 1) for j in range(8)] + specs += [ + (c["qkvp"], (p * s + t) * QKV_ROWS + Q_RANK + tt * 2 + j, 1) + for j in range(2) + ] + got = c["mb"].poll(specs) + vals = [ + bf16_round(got[j][0].bitcast(fx.Float32) + got[10 + j][0].bitcast(fx.Float32)) + for j in range(10) + ] + xq = [((tt * 8 + j) < Q_RANK).select(vals[j], fx.Float32(0.0)) for j in range(8)] + xk = vals[8:10] + ssq = xq[0] * xq[0] + for j in range_constexpr(1, 8): + ssq = ssq + xq[j] * xq[j] + ssq = butterfly(ssq, (32, 16, 8, 4, 2, 1)) + ssk = butterfly(xk[0] * xk[0] + xk[1] * xk[1], (32, 16, 8, 4, 2, 1)) + if (lane == 0) & mine: + fx.ptr_store(ssq, red + wave) + fx.ptr_store(ssk, red + 4 + wave) + gpu.barrier() + tq = (fx.ptr_load(red + 0) + fx.ptr_load(red + 1)) + ( + fx.ptr_load(red + 2) + fx.ptr_load(red + 3) + ) + tk = (fx.ptr_load(red + 4) + fx.ptr_load(red + 5)) + ( + fx.ptr_load(red + 6) + fx.ptr_load(red + 7) + ) + rq = hw_rsq( + div_rn(tq, fx.Float32(float(Q_RANK)), fx.Float32(1.0 / Q_RANK)) + + fx.Float32(EPS) + ) + rk = hw_rsq( + div_rn(tk, fx.Float32(float(KV_DIM)), fx.Float32(1.0 / KV_DIM)) + + fx.Float32(EPS) + ) + # q: bf16 normed row, then vLLM's MXFP8 (a 32-group = 4 threads) + ys = [bf16_round((xq[j] * rq) * qw[j]) for j in range(8)] + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = mx_code(butterfly(amax, (1, 2), fx.max)) + mul = mx_mul(code) + if mine & (tt * 8 < Q_RANK): + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["qx8"]), + t * QX_WORDS + 2 * tt, + cache_modifier=CM_DEV, + ) + if tt % 4 == 0: + bo.buffer_store( + code, + rsrc(c["qx8s"]), + t * (Q_RANK // 32) + tt // 4, + cache_modifier=CM_DEV, + ) + # kv: this thread's pair (2 tt, 2 tt + 1) + kn = [bf16_round((xk[j] * rk) * kw[j]) for j in range(2)] + e0 = bf16_round(kn[0] * cs - kn[1] * sn) + e1 = bf16_round(kn[0] * sn + kn[1] * cs) + # record quant of the nope dims: a 64-block = 32 threads (half a wave) + amax = butterfly(fx.max(abs(e0), abs(e1)), (1, 2, 4, 8, 16), fx.max) + rcode = ceil_exp( + div_rn( + fx.max(amax, fx.Float32(1e-4)), fx.Float32(448.0), fx.Float32(1.0 / 448.0) + ) + ) + inv = mx_mul(rcode) + q0 = clamp_fp8(e0 * inv) + q1 = clamp_fp8(e1 * inv) + n0 = xshfl(q0, 1) + n1 = xshfl(q1, 1) + word = fp8_pack4(q0, q1, n0, n1) + rope_word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + blk = a["swa_block"] + sblock = slot // blk + soff = slot % blk + base = a["swa"] + fx.Int64(sblock) * fx.Int64(a["swa_stride"]) + rec = base + fx.Int64(soff * DATA) + if mine & (tt % 32 == 0) & (tt < NOPE // 2): + fx.ptr_store(rcode.bitcast(fx.Float32), red + 8 + tt // 32) + gpu.barrier() + if mine & (slot >= 0): + if tt < NOPE // 2: + if tt % 2 == 0: + gstore(rec + fx.Int64(2 * tt), word, words=1) + else: + gstore(rec + fx.Int64(NOPE + 4 * (tt - NOPE // 2)), rope_word, words=1) + if tt == 0: + sb = [fx.ptr_load(red + 8 + b).bitcast(fx.Int32) for b in range(7)] + w0 = sb[0] | (sb[1] << 8) | (sb[2] << 16) | (sb[3] << 24) + w1 = sb[4] | (sb[5] << 8) | (sb[6] << 16) + sc = base + fx.Int64(blk * DATA + soff * 8) + gstore(sc, fx.Vector.from_elements([w0, w1], fx.Int32), words=2) + publish(c, c["qrdy"], t) + + +# ---------------------------------------------------------------- wq_b + + +def wqb_rope(c, task0, pos_t, n): + """(cos, sin) of this thread's pair in ``stage_wqb``'s epilogue over a tile + of n rows (``pos_t``: its token's position; every task of the CTA shares + the pair's dims: tasks BLOCKS apart are whole heads apart).""" + a = c["args"] + rp = c["tid"] % (ROWS // 2) + dim = (task0 * ROWS) % HEAD_DIM + 2 * rp + jj = fx.max((dim - NOPE) >> 1, 0) + return ld_f32(a["cos_sin"], pos_t * ROPE + jj), ld_f32( + a["cos_sin"], pos_t * ROPE + HALF + jj + ) + + +def wqb_pos(c, t0, n): + """This thread's token position for ``stage_wqb``'s epilogue of tile + t0 .. t0 + n.""" + per = ROWS * n // 2 + tl = fx.min((c["tid"] % per) // (ROWS // 2), n - 1) + return ld_i32(c["args"]["pos"], 2 * (t0 + tl)) + + +@traced +def stage_wqb(c, task0, bops, t0, n, rope): + """Rows 16 (task0 + BLOCKS k) .. of the rank's H x 512 query rows, every + task k of the CTA: GEMVs -> bf16 -> GPT-J RoPE of a head's dims 448.. + (fp32) -> bf16 -> q.""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + nt = len(bops) + for k in range_constexpr(nt): + gemv_mfmas(c, Q_RANK, c["qxl"], c["qxsl"], red + k * RED, bops[k], rows=n) + gpu.barrier() + per = ROWS * n // 2 # bf16 pairs a task + if tid < nt * per: + k = tid // per + u = tid % per + rp = u % (ROWS // 2) + tl = u // (ROWS // 2) + t = plus(t0, tl) + r0 = (task0 + BLOCKS * k) * ROWS + 2 * rp + rk = red + k * RED + v0 = bf16_round(row_sum(rk, 2 * rp, tl)) + v1 = bf16_round(row_sum(rk, 2 * rp + 1, tl)) + is_rope = (task0 * ROWS) % HEAD_DIM + 2 * rp >= NOPE + cs, sn = rope + e0 = is_rope.select(v0 * cs - v1 * sn, v0) + e1 = is_rope.select(v0 * sn + v1 * cs, v1) + word = ( + fx.Vector.from_elements([e0, e1], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + ) + bo.buffer_store(word, rsrc(a["q"]), (t * d.q_rows + r0) // 2) + gpu.barrier() + + +# ---------------------------------------------------------------- kernel + + +def front_lds(s: int, d: Dims): + """K1's LDS (the attention front's stages).""" + rows = tile_rows(s) + wqb_per_cta = d.wqb_tasks // BLOCKS + + @fx.struct + class FrontLds: + red: fx.Array[fx.Float32, wqb_per_cta * RED + 16, 16] + qxl: fx.Array[fx.Int32, rows * lds_row(QX_WORDS), 16] + qxsl: fx.Array[fx.Int32, rows * lds_row(Q_RANK // 32), 16] + xl: fx.Array[fx.Int32, rows * lds_row(KWORDS), 16] + xsl: fx.Array[fx.Int32, rows * lds_row(KGROUPS), 16] + + return FrontLds + + +def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: + """The front stages' context: LDS views, mailbox (tag ``tag``), scratch + regions, arguments (``args``: the kernel's, by name).""" + s, d = key.tokens, Dims(key.tp) + tid = fx.thread_idx.x + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "xl": lds.xl.ptr, + "xsl": lds.xsl.ptr, + "qxl": lds.qxl.ptr, + "qxsl": lds.qxsl.ptr, + "mb": Mailbox(tag), + "d": d, + "ratio": key.ratio, + "ldbf": lambda ptr, i: fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ), + "args": args, + } + for region, (off, _) in (layout or front_scratch(s, d)).items(): + c[region] = scratch + fx.Int64(off) + return c + + +@traced +def run_front(c, key, bid, x8_given=False, before_wqkv=None, after_wqkv=None): + """The front's stages on this CTA. ``x8_given``: a stage before this one + (the mono layer's seam norm) publishes X8 / XRDY for every token (no input + quantization here). ``before_wqkv()`` runs ahead of the wqkv weights' loads + (a latency-critical stage of the caller's), ``after_wqkv()`` right after.""" + s = key.tokens + wqb_tasks = Dims(key.tp).wqb_tasks + a = c["args"] + tid = c["tid"] + tiles = token_tiles(s) + wqb_per_cta = wqb_tasks // BLOCKS + assert wqb_tasks % BLOCKS == 0 + n_wqkv = WQKV_TASKS * WQKV_KPARTS # 224 + assert n_wqkv <= BLOCKS + n_free = BLOCKS - n_wqkv # CTAs without a wqkv part: xq, kt + QKV0 = 0 # on wqkv CTAs + WQB0 = 0 + in_cta_quant = s <= TILE and not x8_given + xq_stage = s > TILE and not x8_given + xq_units = WQKV_KPARTS * s + xq_per = (xq_units + n_free - 1) // n_free + stamp(c, 0) + + # ---- Loads complete in issue order, so every CTA issues its latency- + # critical loads before its later stages' weights: a wqkv CTA its first + # input tile and its wqkv part, the next tiles while it quantizes, then + # its wq_b tasks; a CTA with a key-table task those after that work. + # The q / kv norms run on wqkv CTAs, their partials' wait behind the + # wq_b weights that are due by then anyway. + wqkv_unit = first_task(bid, 0) + has_wqkv = wqkv_unit < n_wqkv + task = fx.min(wqkv_unit, n_wqkv - 1) % WQKV_TASKS + part = fx.min(wqkv_unit, n_wqkv - 1) // WQKV_TASKS + wqb_task = first_task(bid, WQB0) + free = bid - n_wqkv + is_free = free >= 0 + if const_expr(in_cta_quant): + xregs = [x_half_loads(c, part, has_wqkv, *tiles[0])] + if const_expr(xq_stage): + xunits = [free + n_free * i for i in range(xq_per)] + xlive = [is_free & (u < xq_units) for u in xunits] + xq_regs = [ + xq_loads(c, fx.min(fx.max(u, 0), xq_units - 1), lv) + for u, lv in zip(xunits, xlive) + ] + if const_expr(before_wqkv is not None): + before_wqkv() + kops = gemv_loads( + c, + a["wqkv"], + a["wqkv_s"], + HIDDEN, + task * ROWS, + kbeg=part * KSTEPS, + nsteps=KSTEPS, + pred=has_wqkv, + ) + if const_expr(after_wqkv is not None): + after_wqkv() + if const_expr(xq_stage): # noqa: SIM102 (a build-time and a run-time test) + if is_free: + for i in range_constexpr(xq_per): + stage_xq( + c, fx.min(fx.max(xunits[i], 0), xq_units - 1), xq_regs[i], xlive[i] + ) + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid < xq_per: + u = free + n_free * tid + if u < xq_units: + c["mb"].put(c["xrdy"], u, fx.Int32(1)) + for t in range(is_free.select(free, fx.Int32(s)), s, n_free): + stage_kt(c, t) + stamp(c, 1) + poss = [wqb_pos(c, t0, n) for t0, n in tiles] + + # ---- wqkv_a: one (row tile, K half) a CTA, every token tile; the wq_b + # weights issued behind the input rows' loads + def wqb_weights(): + return [ + gemv_loads(c, a["wqb"], a["wqb_s"], Q_RANK, (wqb_task + BLOCKS * i) * ROWS) + for i in range(wqb_per_cta) + ] + + if const_expr(in_cta_quant): + if has_wqkv: + quant_x_half(c, xregs[0], tiles[0][1]) + bops = wqb_weights() + if has_wqkv: + stage_wqkv(c, task, part, kops, *tiles[0], has_wqkv) + else: + await_flags(c, c["xrdy"], s, lambda u: WQKV_KPARTS * u + part) + xr = ( + c["x8"], + c["x8s"], + X_WORDS, + part * KWORDS, + KWORDS, + X_GROUPS, + part * KGROUPS, + KGROUPS, + ) + regs = rows_loads(c, *xr, tiles[0][0], tiles[0][1], pred=has_wqkv) + if const_expr(len(tiles) == 1): + bops = wqb_weights() + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, KWORDS, KGROUPS, c["xl"], c["xsl"], n) + if const_expr(ti + 1 < len(tiles)): + t1, n1 = tiles[ti + 1] + regs = rows_loads(c, *xr, t1, n1, pred=has_wqkv) + if const_expr(ti + 2 == len(tiles)): + bops = wqb_weights() + stage_wqkv(c, task, part, kops, t0, n, has_wqkv) + stamp(c, 2) + + # ---- q / kv norms, KV insert + for t in range(first_task(bid, QKV0), s, BLOCKS): + stage_qkv(c, t) + stamp(c, 3) + + # ---- wq_b: every task's rows of a token tile at a time, the next + # tile's rows in flight + ropes = [wqb_rope(c, wqb_task, poss[ti], tiles[ti][1]) for ti in range(len(tiles))] + await_flags(c, c["qrdy"], s, lambda u: u) + qx = (c["qx8"], c["qx8s"], QX_WORDS, 0, QX_WORDS, Q_RANK // 32, 0, Q_RANK // 32) + regs = rows_loads(c, *qx, *tiles[0]) + for ti in range_constexpr(len(tiles)): + t0, n = tiles[ti] + rows_store(c, regs, QX_WORDS, Q_RANK // 32, c["qxl"], c["qxsl"], n) + if const_expr(ti + 1 < len(tiles)): + regs = rows_loads(c, *qx, *tiles[ti + 1]) + stage_wqb(c, wqb_task, bops, t0, n, ropes[ti]) + stamp(c, 4) diff --git a/vllm/models/deepseek_v41/amd/mono/attention/gemv.py b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py new file mode 100644 index 000000000000..9471a1e3f4df --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/gemv.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 GEMV of the mega kernel: 16 rows of a row-major e4m3 weight [N, K] +(vLLM's loaded layout, E8M0 scales in 32x32 blocks [N / 32, K / 32]) against +every token of a tile, its MXFP8 rows in LDS, on the gfx950 scaled MFMA. + +A task may cover a window of K (``kbeg`` .. ``kbeg + nsteps`` 128-steps): a +split-K part, its activations in LDS holding just that window. The weight +operands are loaded first (all in flight, before the task waits on its +activations); wave w takes the window's steps w, w + 8, ... (or, ``split``, a +contiguous range of waves 0 .. split - 1), its 16 x 16 partial in ``red``, +summed in wave order by ``row_sum``. Adapted from ATOM's +``attn_pre.gemv_fp8_*``, which reads aiter's (16, 16)-preshuffled layout. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import T + +from .device import THREADS, WAVES, mfma_scaled, rsrc, tile_map, traced + +# a token's LDS row of MXFP8 words / E8M0 codes, padded by 16 B so 16 rows read +# by one MFMA operand do not share a bank group +LDS_PAD = 4 +RED = WAVES * 64 * 4 # one task's accumulators in ``red`` (f32) + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def _steps(c, nsteps, split, tiled): + """This wave's steps of the window: [(local step, live)].""" + wave = c["wave"] + if const_expr(split is None): + per = (nsteps + WAVES - 1) // WAVES + ragged = nsteps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < nsteps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and nsteps % split == 0 + per = nsteps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def _step_load(c, r_w, r_ws, k, row0, st, pred=None): + """Global step ``st``'s weight operand of rows row0 .. +16: two 16 B chunks a + lane (K bytes 16 j .. and 64 + 16 j .. of the step, row row0 + lane % 16) + and the E8M0 byte of its K block. ``pred`` false: masked (out of range: no + memory traffic, zeros).""" + lane = c["lane"] + j = lane // 16 + row = row0 + lane % 16 + base = row * k + 128 * st + 16 * j + lo = fx.Vector(bo.buffer_load(r_w, base // 4, vec_width=4, dtype=T.i32)) + hi = fx.Vector(bo.buffer_load(r_w, (base + 64) // 4, vec_width=4, dtype=T.i32)) + kg = k // 32 + wsi = (row0 // 32) * kg + st * 4 + j + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def gemv_loads( + c, w, ws, k, row0, split=None, tiled=False, kbeg=0, nsteps=None, pred=None +): + """The weight operands of rows row0 .. +16 (``tiled``: row0 + 16 (wave // + split)) over steps kbeg .. kbeg + nsteps for this wave, all in flight + (``pred`` false: masked off, no memory traffic).""" + nsteps = k // 128 if nsteps is None else nsteps + if const_expr(pred is None): + r_w, r_ws = rsrc(w), rsrc(ws) + else: + # a masked-off CTA reads through a zero-size buffer: zeros, no traffic + nb = pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + r_w, r_ws = rsrc(w, nb), rsrc(ws, nb) + if const_expr(tiled): + row0 = row0 + 16 * (c["wave"] // split) + return [ + _step_load( + c, + r_w, + r_ws, + k, + row0, + kbeg + (st if live is None else fx.min(st, nsteps - 1)), + pred, + ) + for st, live in _steps(c, nsteps, split, tiled) + ] + + +def _step_mfma(c, nsteps, xl, xsl, st, ops, acc, live, rows): + lane = c["lane"] + lo, hi, sa = ops + x_row, s_row = lds_row(nsteps * 32), lds_row(nsteps * 4) + j = lane // 16 + col = fx.min(lane % 16, rows - 1) + if const_expr(live is not None): + st = fx.min(st, nsteps - 1) + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [hi[d] for d in range(4)], fx.Int32 + ) + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + [xhi[d] for d in range(4)], fx.Int32 + ) + sb = fx.ptr_load(xsl + (col * s_row + st * 4 + j)) + if const_expr(live is not None): + # a step past the end adds zero at unit scale + av = fx.Vector.from_elements( + [live.select(av[d], fx.Int32(0)) for d in range(8)], fx.Int32 + ) + sa = live.select(sa, fx.Int32(127)) + sb = live.select(sb, fx.Int32(127)) + return mfma_scaled(av, xv, acc, sa, sb) + + +def gemv_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=16, nsteps=None): + """``gemv_loads``' operands against ``rows`` tokens' MXFP8 rows of the + window in LDS (``xl`` words, ``xsl`` codes, rows ``lds_row`` of the window) + -> ``red[wave][lane]`` (a 16 x 16 C: row 4 (lane // 16) + i, token lane % 16).""" + lane, wave = c["lane"], c["wave"] + nsteps = k // 128 if nsteps is None else nsteps + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_steps(c, nsteps, split, tiled)): + acc = _step_mfma(c, nsteps, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' accumulators, summed in wave order.""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def span_items(width, n): + """A [n, width] span of 16 B items: rows a round of THREADS threads, rounds.""" + wq = width // 4 + rpr = THREADS // wq + assert rpr >= 1, width + return wq, rpr, (n + rpr - 1) // rpr + + +def rows_loads( + c, + words_base, + codes_base, + stride_w, + off_w, + words, + stride_g, + off_g, + groups, + t0, + n, + pred=None, +): + """Tokens t0 .. t0 + n's published MXFP8 row slices -- ``words`` i32 from + ``off_w`` and ``groups`` codes (an i32 each) from ``off_g``, rows + ``stride_w`` / ``stride_g`` i32 apart -- loaded (issued, not waited on), + 16 B a lane-item: a thread one column, rows a round apart. Plain loads: + their flags are seen, and nothing in the launch read these lines before. + ``pred`` false: masked off (zeros, no memory traffic).""" + tid = c["tid"] + nb = None if pred is None else pred.select(fx.Int64(0xFFFFFFFF), fx.Int64(0)) + out = [] + for base, stride, off, width in ( + (words_base, stride_w, off_w, words), + (codes_base, stride_g, off_g, groups), + ): + wq, rpr, rounds = span_items(width, n) + row0, col = tile_map(tid, wq) + a0 = col * 4 + off + t0 * stride + r = rsrc(base, nb) + regs = [] + for i in range(rounds): + row = row0 + rpr * i + if (i + 1) * rpr > n or rpr * wq < THREADS: + row = fx.min(row, n - 1) + regs.append( + fx.Vector( + bo.buffer_load(r, row * stride + a0, vec_width=4, dtype=T.i32) + ) + ) + out.append(regs) + return out + + +@traced +def _store_span(c, regs, width, dst, n): + tid = c["tid"] + wq, rpr, _ = span_items(width, n) + row0, col = tile_map(tid, wq) + lrow = lds_row(width) + base = dst + (row0 * lrow + col * 4) + for i in range_constexpr(len(regs)): + if const_expr((i + 1) * rpr <= n and rpr * wq == THREADS): + fx.ptr_store(regs[i], base + i * rpr * lrow) + else: + if (row0 + rpr * i < n) & (tid < rpr * wq): + fx.ptr_store(regs[i], base + i * rpr * lrow) + + +def rows_store(c, regs, words, groups, xl, xsl, n): + """``rows_loads``' registers into LDS rows 0 .. n (``lds_row`` strides).""" + _store_span(c, regs[0], words, xl, n) + _store_span(c, regs[1], groups, xsl, n) + gpu.barrier() diff --git a/vllm/models/deepseek_v41/amd/mono/attention/plan.py b/vllm/models/deepseek_v41/amd/mono/attention/plan.py new file mode 100644 index 000000000000..58865ac7d01e --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/attention/plan.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Geometry of the DeepSeek-V4.1 attention mega kernel: model widths, a TP +rank's share of them, task counts and the scratch (mailbox) layout.""" + +from dataclasses import dataclass + +HIDDEN = 5120 +Q_RANK = 1280 +KV_DIM = 512 +QKV_ROWS = Q_RANK + KV_DIM # 1792 +HEAD_DIM = 512 +NOPE = 448 +ROPE = 64 +HALF = ROPE // 2 +N_HEADS = 64 +O_GROUPS = 8 +O_RANK = 1024 +WINDOW = 128 +TOPK = 512 +EPS = 1e-20 + +ROWS = 16 # GEMV rows a task +TILE = 16 # tokens a GEMV pass (an MFMA's N) +HEAD_TILE = 16 # heads an attention MFMA takes +WOA_ROWS = 32 # a wo_a task's rows: one MX group of its output +MAX_TOKENS = 48 + +# KV record (V4 fp8_ds_mla): 576 B data rows (448 e4m3 | 64 bf16), then 8 B of +# scales a token (7 UE8M0 + pad) after the block's rows +DATA = 576 +SCALE_BYTES = 8 + +# attention: a split is 16 keys; a token has at most TOPK + WINDOW keys +BK = 16 +SPLITS = 64 # split slots a (token, head tile): 1024 keys >= TOPK + WINDOW +COLS = 32 # PV columns a group (one MX group of the output) +GROUPS = HEAD_DIM // COLS # 16 + +PAIR = 8 # a mailbox (value, tag) pair + + +def cdiv(a, b): + return -(-a // b) + + +def round_up(a, b): + return cdiv(a, b) * b + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + assert self.tp in (2, 4), self.tp + + @property + def heads(self) -> int: + return N_HEADS // self.tp + + @property + def head_tiles(self) -> int: + return self.heads // HEAD_TILE + + @property + def groups(self) -> int: + return O_GROUPS // self.tp + + @property + def group_k(self) -> int: + """wo_a's K a group: its eight heads' outputs.""" + return self.heads // self.groups * HEAD_DIM # 4096 + + @property + def o_rows(self) -> int: + return self.groups * O_RANK + + @property + def q_rows(self) -> int: + return self.heads * HEAD_DIM + + @property + def wqb_tasks(self) -> int: + return self.q_rows // ROWS + + @property + def woa_tasks(self) -> int: + return self.o_rows // WOA_ROWS + + @property + def wob_tasks(self) -> int: + return HIDDEN // ROWS # 320 + + +WQKV_TASKS = QKV_ROWS // ROWS # 112 +WQKV_KPARTS = 2 # wqkv's K split: 224 (row tile, K half) parts +KEYS = TOPK + WINDOW # a token's attention keys at most + + +def token_tiles(s): + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + return min(s, TILE) + + +def pair_layout(items, start=0, align=256): + out, off = {}, start + for name, pairs in items: + off = round_up(off, align) + out[name] = (off, pairs * PAIR) + off += pairs * PAIR + return out + + +def byte_layout(items, start=0, align=256): + out, off = {}, start + for name, nbytes in items: + off = round_up(off, align) + out[name] = (off, nbytes) + off += nbytes + return out + + +def front_scratch(s: int, d: Dims) -> dict: + """K1's scratch regions -> (byte offset, bytes): published data (plain + words behind a flag) and mailbox pairs.""" + return byte_layout( + [ + ("x8", s * HIDDEN), # e4m3 words (S > 16: quantized once, shared) + ("x8s", s * HIDDEN // 32 * 4), # an i32 code a group + ("xrdy", WQKV_KPARTS * s * PAIR), # a flag a (token, K half) + ("qkvp", WQKV_KPARTS * s * QKV_ROWS * PAIR), + ("qx8", s * Q_RANK), # e4m3 words + ("qx8s", s * Q_RANK // 32 * 4), # an i32 code a group + ("qrdy", s * PAIR), + ] + ) + + +def back_scratch(s: int, d: Dims, start: int = 0) -> dict: + """K2's scratch regions, laid out past ``start``: published data (plain) + and flag pairs.""" + ck = ( + 64 if d.head_tiles == 2 and s * cdiv(KEYS, 64) <= 256 else 128 + ) # back.chunk_keys + nchunk = cdiv(KEYS, ck) + return byte_layout( + [ + ("part", s * nchunk * d.heads * HEAD_DIM * 2), # bf16 partials + ("pstat", s * nchunk * d.heads * 8), # (max, sum) f32 + ("srdy", s * nchunk * PAIR), + ("xo", s * d.q_rows), # e4m3 words + ("xos", s * d.q_rows // 32 * 4), # an i32 code a group + ("crdy", s * d.heads // 8 * PAIR), + ("x8b", s * d.o_rows), + ("x8bs", s * d.o_rows // 32 * 4), + ("ardy", d.woa_tasks * PAIR), + ], + start=start, + ) + + +def front_end(s: int, d: Dims) -> int: + """Where K2's regions start: past K1's, which share the launch tag.""" + return round_up(max(o + n for o, n in front_scratch(s, d).values()), 256) + + +def scratch_bytes(s: int, d: Dims) -> int: + b = back_scratch(s, d, start=front_end(s, d)) + return round_up(max(o + n for o, n in b.values()), 4096) diff --git a/vllm/models/deepseek_v41/amd/mono/common/__init__.py b/vllm/models/deepseek_v41/amd/mono/common/__init__.py new file mode 100644 index 000000000000..397642c6b441 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/__init__.py +"""The mono kernels' shared mechanisms: ``ops`` (wave / math primitives, the TP +group's peer buffers), ``sync`` (the tagged mailbox), ``mx`` (FP8 / MX +arithmetic), ``plan`` (execution model, layouts, build keys) and +``peer_memory`` (symmetric peer buffers).""" diff --git a/vllm/models/deepseek_v41/amd/mono/common/mx.py b/vllm/models/deepseek_v41/amd/mono/common/mx.py new file mode 100644 index 000000000000..9c228abdbf27 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/mx.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/mx.py +"""FP8 / MX arithmetic the mono kernels share: the E4M3 range, E8M0 scale +codes, the byte of a scale in aiter's ``shuffle_scale`` layout, an FP4 weight +tile's operand loads and the scaled MFMA that takes them. + +An E8M0 code is a biased exponent; ``pow2(code)`` is the scale it stands for.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import CM_NT + +FP8_MAX = 448.0 +# the f8f6f4 MFMA's operand formats (cbsz / blgp) +FP8, FP4 = 0, 4 +UNIT_SCALE = 127 # the E8M0 code of 1.0 + + +def clamp_fp8(v): + """``v`` clamped to the E4M3 range.""" + return fx.max(fx.min(v, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + +def code_ceil(v): + """Biased exponent of the smallest power of two >= ``v`` (v >= 0, finite), + clamped to [1, 254]: the codes a scale may take.""" + bits = v.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + + +def rcp_pow2(code): + """1 / ``pow2(code)``, exact (code in [1, 253]): x times it is x / scale + bit for bit, a multiply where the division is a dozen instructions.""" + return ((254 - code) << 23).bitcast(fx.Float32) + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def fp4_tile_load(r_w, base, rg, k, st, at): + """A lane's 16 B of an FP4 weight's (16, 16)-shuffled 16-row group ``rg`` + (K = ``k``) at 128-k step ``st``: dword ``base`` + the group's 2 k + 256 a + step + ``at`` (``lane * 4``: 32 fp4, K block lane / 16 of row lane % 16). + Nontemporal: a weight a step reads once.""" + return fx.Vector( + bo.buffer_load( + r_w, + base + rg * (2 * k) + st * 256 + at, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + + +def mfma_scaled(a, b, c, sa, sb, a_fmt=FP8, b_fmt=FP8, sa_byte=0, sb_byte=0): + """One 16x16x128 scaled MFMA (gfx950 f8f6f4): A ``a`` (``a_fmt``) with E8M0 + ``sa`` (its byte ``sa_byte``), B ``b`` (``b_fmt``) with E8M0 ``sb`` (its + byte ``sb_byte``).""" + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [a, b, c, a_fmt, b_fmt, sa_byte, sa, sb_byte, sb] + ) + ) diff --git a/vllm/models/deepseek_v41/amd/mono/common/ops.py b/vllm/models/deepseek_v41/amd/mono/common/ops.py new file mode 100644 index 000000000000..f9f7fb22ec92 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/ops.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/ops.py +# atom/mono/device/ranks.py +# atom/mono/device/stamps.py +"""Device-side primitives of the mono kernels: wave reductions and shuffles, the +approximate hardware math, bf16 / fp8 packing, buffer resources, the +execution-model constants and the kernel symbol naming. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``). Only published flydsl APIs +and aiter's vendored ``buffer_ops`` / ``dpp_utils`` are used. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries kept, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``. + +Also the TP group seen from inside a kernel (every rank's peer buffer, an +all-reduce's rank-ordered sum) and a timeline build's phase stamps.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES + +# ---------------------------------------------------------------- primitives +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 + + +CM_SYS = 17 + + +# NT = 2: nontemporal. A weight a layer reads once streams past the caches it +# would only evict the step's reused data from (the MoE at 48 rows: 143 -> 133 us) +CM_NT = 2 + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def fresh(v): + """``v`` (a per-lane Int32) through an opaque move, here: the lane math built + on it is computed at its use, where LLVM would otherwise hoist it out of a task + loop or CSE it across stages into registers held for the whole kernel (spills, + and a kernel with scratch is refused: ``runtime.widths``).""" + return fx.Int32( + llvm.inline_asm( + T.i32, + [fx.Int32(v).ir_value()], + "v_mov_b32 $0, $1", + "=v,v", + has_side_effects=True, + ) + ) + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS). Powers of two only: any other + offset would take the lane ^ 1 path.""" + assert off in (1, 2, 4, 8, 16, 32), off + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +LOAD_BATCH = 4 # rounds of a strided loop whose loads go out together + + +@traced +def batched_rounds(tid, n, load, use, batch=LOAD_BATCH): + """A CTA's strided loop over elements v = tid + THREADS i < n (n a Python + int): ``load(v)`` of ``batch`` rounds issued before any ``use(v, loaded)``. + A round whose ``use`` stores to global memory keeps the next round's loads + behind it (they may alias), so the rounds' loads would go out in series. A + round past n loads element n - 1 (in range) and uses nothing.""" + rounds = -(-n // THREADS) + for b0 in range_constexpr(0, rounds, batch): + idx = range(b0, min(b0 + batch, rounds)) + got = [load(fx.min(tid + THREADS * i, n - 1)) for i in idx] + for i, g in zip(idx, got): + v = tid + THREADS * i + if v < n: + use(v, g) + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def row_sum(red, r, t, waves=None): + """Row r, token t of the waves' 16x16 MFMA accumulators (wave w's at ``red`` + + (w * 64 + lane) * 4), summed in wave order from 0 (``waves``: which, in + order; all of them by default).""" + tot = fx.Float32(0.0) + for w in waves if waves is not None else range(WAVES): + tot = tot + fx.ptr_load(red + ((w * 64 + 16 * (r // 4) + t) * 4 + r % 4)) + return tot + + +def ballot(pred): + """The wave's lanes where ``pred`` holds, a bit a lane.""" + return fx.Int64(rocdl.ballot(T.i64, as_ir_value(pred))) + + +def popcount(x): + return fx.Int32(fx.ctpop(x)) + + +def lanes_below(bal, lane): + """How many of the lanes in ballot ``bal`` are below ``lane``.""" + return popcount(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + + +def load_ptr64(table, i): + """Entry ``i`` of an i64 table in memory, wave-uniform (each half through + ``uniform``): a peer's buffer base, a kernel argument block's pointer.""" + pw = fx.Vector(bo.buffer_load(rsrc(table), i * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64(fx.Uint32(uniform(pw[0]))) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf_lo(w): + """The f32 of a packed bf16 pair word's low half.""" + return (w << 16).bitcast(fx.Float32) + + +def bf_hi(w): + """The f32 of a packed bf16 pair word's high half.""" + return (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return bf_lo(w), bf_hi(w) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +# ---------------------------------------------------------------- TP ranks +def peer_bases(peers, tp): + """Every rank's peer buffer base from the ``peers`` table (an i64 a rank), + loaded once: a put's address must not wait on a load of its own.""" + return [load_ptr64(peers, p) for p in range(tp)] + + +def sum_partials(poll, own, pair_of, tp): + """The ``tp`` ranks' bf16 pair partials at ``pair_of(src)`` of the mailbox + region ``own``, summed in fp32 in rank order 0 .. tp - 1: one poll batch -> + (lo sum, hi sum). Every rank sums in the same order, so they agree.""" + ws = poll([(own, pair_of(src), 1) for src in range(tp)]) + acc0, acc1 = bf2_f32(ws[0][0]) + for src in range_constexpr(1, tp): + lo, hi = bf2_f32(ws[src][0]) + acc0 = acc0 + lo + acc1 = acc1 + hi + return acc0, acc1 + + +# ---------------------------------------------------------------- timeline stamps +def stamp(on, tls, tid, k): + # compile-time gate outside, traced condition inside: they cannot be one `and` + if const_expr(on): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) diff --git a/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py new file mode 100644 index 000000000000..58c95a567e07 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/peer_memory.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/runtime/consensus.py +# atom/mono/runtime/peer_memory.py +"""Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every +rank must reach before anything collective (a rank that turned the mono layer +off alone would leave the others spinning on the GPU).""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +# ---------------------------------------------------------------- TP agreement +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" + + +def tp_agree(ok: bool, group) -> bool: + """True when every rank of ``group`` passed ``ok`` true. Every rank must call + it at the same point; ``group`` is a CPU (gloo) process group, or None for a + single rank.""" + if group is None: + return ok + flag = torch.tensor([int(ok)], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group) + return bool(flag.item()) + + +# ---------------------------------------------------------------- peer buffers +class _DeviceBytes: + """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a + non-owning torch view.""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address + (``addresses``, int64, indexed by TP rank). What lives where, and why no rank + overwrites a peer's unread data, is the model's layout's to say. + + The constructor is collective over ``group`` (the handle exchange, then an + agreement that every rank mapped every peer): every rank must reach it, so a + runner reaches it only after ``tp_agree``. A failed mapping on any rank raises + ``MonoUnsupported`` on all of them. + """ + + def __init__( + self, + nbytes: int, + group, + rank: int, + npes: int, + device: torch.device, + ): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + # a view, not an owner: close() frees the memory under it + self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device) + self.rank = rank + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses, failure = [], None + try: + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + except RuntimeError as err: # every rank must still reach the agreement + failure = err + if not tp_agree(failure is None, group): + self.close() + raise MonoUnsupported( + "peer memory handshake failed on " + f"{'this' if failure else 'another'}" + " TP rank" + ) from failure + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent. Only once no + kernel of any rank can still touch it (the runner is being dropped).""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + self.bytes = None + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 diff --git a/vllm/models/deepseek_v41/amd/mono/common/plan.py b/vllm/models/deepseek_v41/amd/mono/common/plan.py new file mode 100644 index 000000000000..2b05d7a073e9 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/plan.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/plan/execution.py +# atom/mono/plan/layout.py +# atom/mono/plan/arith.py +# atom/mono/plan/shard.py +# atom/mono/plan/build_key.py +"""Plain-Python planning of the mono kernels: the execution model (one CTA a CU, +all resident), mailbox layouts, arithmetic that runs on ints and traced values +alike, TP shards, and build keys (every build parameter, and this package's +sources, in FlyDSL's JIT cache key).""" + +import enum +import functools +import hashlib +from dataclasses import dataclass, fields +from pathlib import Path + +# ---------------------------------------------------------------- execution model +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + + +def first_task(bid, base): + """CTA ``bid``'s first task of a stage placed from CTA ``base`` (any count of + CTAs past 0: placements that add task counts run past the grid), then every + BLOCKS-th.""" + return (bid - base % BLOCKS + BLOCKS) % BLOCKS + + +# ---------------------------------------------------------------- mailbox layouts +PAIR_BYTES = 8 # a (value, tag) pair + + +def pair_layout(items, align: int = PAIR_BYTES, start: int = 0) -> dict: + """``items`` (name, pairs) laid out in order: name -> (byte offset, bytes), + each region from an ``align``-byte boundary at or past ``start``.""" + out, off = {}, start + for name, pairs in items: + off = (off + align - 1) // align * align + out[name] = (off, pairs * PAIR_BYTES) + off += pairs * PAIR_BYTES + return out + + +# ---------------------------------------------------------------- arithmetic +def sel(pred, a, b): + """A if pred else b: a conditional on Python values, ``select`` on traced + ones (rewrapped in the traced operand's type).""" + if isinstance(pred, bool): + return a if pred else b + out = pred.select(a, b) + traced = [x for x in (a, b) if not isinstance(x, int)] + return type(traced[0])(out) if traced else out + + +def imin(a, b): + return sel(a < b, a, b) + + +def i32(x): + """``x``, asserted inside the device's Int32 range when it is a Python int.""" + if isinstance(x, int): + assert -(2**31) <= x < 2**31, x + return x + + +def cdiv(a, b): + """ceil(a / b) for a >= 0, b > 0.""" + return i32((a + b - 1) // b) + + +# ---------------------------------------------------------------- TP shards +class ShardError(ValueError): + """A width the TP size does not divide, or a TP size the model refuses.""" + + +@dataclass(frozen=True) +class Shard: + tp: int + + def __post_init__(self): + if self.tp < 1: + raise ShardError(f"TP {self.tp}") + + def split(self, full: int, what: str) -> int: + """``full`` / tp, refusing a remainder: a rank's share of ``what``.""" + if full % self.tp: + raise ShardError(f"{what} {full} is not divisible by TP {self.tp}") + return full // self.tp + + +def padded(n: int, align: int) -> int: + return cdiv(n, align) * align + + +# ---------------------------------------------------------------- build keys +_SCALARS = (int, float, bool, str, enum.Enum) +_ROOT = Path(__file__).resolve().parents[1] # the mono package + + +@functools.cache +def source_digest(*paths: str) -> str: + """sha256 of every ``.py`` file under ``paths`` (files or directories, + relative to the mono package), each file's path and contents: one model's + kernel sources, so an edit to another model's rebuilds none of these.""" + h = hashlib.sha256() + for rel in paths: + root = _ROOT / rel + files = [root] if root.is_file() else sorted(root.rglob("*.py")) + if not files: + raise FileNotFoundError(f"no kernel sources at {root}") + for f in files: + h.update(str(f.relative_to(_ROOT)).encode()) + h.update(f.read_bytes()) + return h.hexdigest()[:16] + + +def key_tuple(key, sources: str) -> tuple: + """``key``'s (field name, value) pairs, every value a scalar, then + ("sources", ``sources``): the model's ``source_digest``.""" + out = [] + for f in fields(key): + value = getattr(key, f.name) + if not isinstance(value, _SCALARS): + raise TypeError( + f"{type(key).__name__}.{f.name} = {value!r}: a build parameter must be" + " a scalar to reach the JIT cache key" + ) + out.append((f.name, value)) + out.append(("sources", sources)) + return tuple(out) diff --git a/vllm/models/deepseek_v41/amd/mono/common/sync.py b/vllm/models/deepseek_v41/amd/mono/common/sync.py new file mode 100644 index 000000000000..0220ce12c32a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/common/sync.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/mono/device/sync.py +"""The tagged mailbox: hand-offs between the CTAs of one GPU (device scope) and +between GPUs (system scope), every address a named region (``Addr``).""" + +from dataclasses import dataclass +from enum import Enum + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_SYS, + bf16_pair, + rsrc, + traced, +) + +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(1) + v = load_all() + return v + + +class Space(Enum): + SCRATCH = "scratch" # this GPU's memory: device-scope hand-off between CTAs + PEER = "peer" # a peer-symmetric region: system-scope hand-off between GPUs + + +@dataclass(frozen=True) +class Addr: + """A mailbox region's address: the region's name and space, and the traced + address. The space sets the hand-off's scope: device for SCRATCH, system for + PEER.""" + + region: str + space: Space + value: object + + +def sreg(base, offset, region): + """Scratch region ``region`` at byte ``offset`` of ``base``.""" + return Addr(region, Space.SCRATCH, base + fx.Int64(offset)) + + +def preg(base, offset, region): + """Peer-symmetric region ``region`` at byte ``offset`` of a rank's buffer + ``base``.""" + return Addr(region, Space.PEER, base + fx.Int64(offset)) + + +# the store / load cache policy of a hand-off in each space +_SCOPE_CM = {Space.SCRATCH: CM_DEV, Space.PEER: CM_SYS} + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The mono layer's launches tag their + hand-offs with their epoch (``layer.py``), so a pair from an earlier step + never carries a live tag. Every address is an ``Addr``: its space picks the + scope. + """ + + def __init__(self, tag): + self.tag = tag + + def put(self, a, i, v): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_bf(self, a, i, vs): + """Elements i .. i+len(vs) (2 or 4, i aligned) as packed bf16 pairs, one + store.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a.value), + i, + cache_modifier=_SCOPE_CM[a.space], + ) + + def put_words(self, a, i, words): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(a.value), + i * 2, + cache_modifier=_SCOPE_CM[a.space], + ) + + def poll(self, specs, batch=POLL_MAX): + """Batched poll of pairs ``specs`` = [(Addr, pair index, npairs in {1, 2})], + all of one space. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. Returns the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], batch) + self.poll(specs[batch:], batch) + spaces = {a.space for a, _, _ in specs} + assert len(spaces) == 1, f"one poll batch spans {spaces}" + cm = _SCOPE_CM[spaces.pop()] + bases = [a.value for a, _, _ in specs] + tag = self.tag + + def load_all(): + words = [] + for b, (_, i, n) in zip(bases, specs): + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +@traced +def publish(put, a, i, v, who): + """This CTA's plain stores made readable by every CTA, then (thread ``who``) + the mailbox pair (``a``, i) := v. The stores are drained first: a flag may + not overtake the data it announces, and a workgroup-scope release fence does + not wait for them (LLVM's AMDGPU memory model omits vmcnt outside tgsplit + mode: the workgroup's waves share one CU, the readers here do not).""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if who: + put(a, i, v) diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py new file mode 100644 index 000000000000..c1be9e3428ae --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -0,0 +1,940 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono decode layer's two launches (README): K1 = the attention seam +(``stages.seam``: slice, gate, the norm into vLLM's MXFP8) + the attention front +(``attention.front``); K2 = the attention back (``attention.back``, wo_b pushed +to every rank) + the FFN seam (the TP sum folded in) + the MoE (``stages.moe``, +its all-reduce in-kernel too). A layer whose attention vLLM runs takes the FFN +launch instead: K2 with the attention back replaced by a push of vLLM's +unreduced wo_b output (``seam.stage_push``). + +Tags: K1's hand-offs carry 2 e + 1, K2's 2 e + 2 (e: the launch pair's epoch, +moved on by K2's CTA 0 at its end), so K2's seam may reuse K1's seam regions. +Peer regions alternate by e's parity; every rank runs the same launches, so +the epochs agree across ranks. +""" + +from dataclasses import dataclass + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Int32, Int64, T + +from .attention import back as mb_back +from .attention import front as mb_front +from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc +from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch +from .common.mx import clamp_fp8 +from .common.ops import butterfly, fp8_pack4, peer_bases, traced +from .common.plan import WAVES, first_task, key_tuple +from .common.sync import Mailbox, preg, publish, sreg +from .sources import SOURCES +from .stages import moe, seam +from .stages.dims import Dims as StageDims +from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape + +_STREAM = fx.Stream(None) +MAX_TOKENS = 48 +X_WORDS = HIDDEN // 4 +X_GROUPS = HIDDEN // 32 +SLICES = seam.SLICES # 160 +# K1's MXFP8 x and the MoE's: plain words in both, the same size, so the two +# launches share the MoE's (placed later); any other name is one region, so a +# word is never a tag to one launch and data to another +SHARED_REGIONS = ("x8", "x8s") +COUNTER_WORDS = 2 * 256 # the MoE's ug queue and down counts, a word a slot +COUNTERS = ("ugq", "dq") +# the epoch buffer, in words: the epoch, a mark a CTA (back.EPOCH_MARKS), then on +# lines of their own the MoE's counters +EPOCH_COUNTERS = -(-(mb_back.EPOCH_MARKS + BLOCKS) // 64) * 64 +EPOCH_WORDS = EPOCH_COUNTERS + COUNTER_WORDS + + +@dataclass(frozen=True) +class MonoBuild: + tokens: int + tp: int + ratio: int # the layer's compress ratio (its attention's key sources) + timeline: bool = False + + +def _moe_key(key: MonoBuild) -> MoeBuild: + return MoeBuild(tokens=key.tokens, tp=key.tp) + + +def _front_key(key: MonoBuild): + return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def _back_key(key: MonoBuild): + return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline) + + +def scratch_layout(s: int, tp: int) -> dict: + """Every region of step width s's launches -> (byte offset, bytes), disjoint: + K1's front and seam, K2's back, the MoE's regions, the normed rows and their + flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``): + under another width's layout a mailbox pair's tag word could be plain data, + which a poll may take for a current tag. The MoE's counters, whose slots + outlive a step width, are in the epoch buffer.""" + d = Dims(tp) + out: dict = {} + off = 0 + + def place(regions): + nonlocal off + base = off + for name, (o, n) in regions.items(): + assert name not in out or name in SHARED_REGIONS, name + out[name] = (base + o, n) + off = max(off, base + max(o + n for o, n in regions.values())) + off = -(-off // 256) * 256 + + place(front_scratch(s, d)) + place(seam.scratch_layout(s)) + place(back_scratch(s, d, start=0)) + regions = moe.scratch_layout(_moe_key(MonoBuild(s, tp, 1))).items() + mr = {n: v for n, v in regions if n not in COUNTERS} + lo = min(o for o, _ in mr.values()) + place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) + place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)}) + return out + + +def scratch_bytes(s: int = MAX_TOKENS, tp: int = 2) -> int: + return max(o + n for o, n in scratch_layout(s, tp).values()) + + +def peer_half_bytes(tp: int) -> int: + """One parity's peer regions: the attention's partials, then the MoE's.""" + return seam.attn_peer_bytes(MAX_TOKENS, tp) + moe.peer_bytes(MAX_TOKENS, tp) + + +def _epoch(epoch): + return fx.Int32(bo.buffer_load(rsrc(epoch), 0, vec_width=1, dtype=T.i32)) + + +# ---------------------------------------------------------------- K1 + + +@traced +def store_x8(c, t, col, ys, live): + """The attention seam's norm output (``seam.stage_norm``'s 8 columns at + ``col``) -> vLLM's MXFP8 of each 32 (the attention's input quant) -> the + front's X8 / X8S, plain at device scope.""" + amax = abs(ys[0]) + for y in ys[1:]: + amax = fx.max(amax, abs(y)) + code = moe.vllm_mx_code(butterfly(amax, (1, 2), fx.max)) + mul = ((254 - code) << 23).bitcast(fx.Float32) + w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) + w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) + if live: + bo.buffer_store( + fx.Vector.from_elements([w0, w1], fx.Int32), + rsrc(c["x8"]), + t * X_WORDS + col // 4, + cache_modifier=CM_DEV, + ) + if c["tid"] % 4 == 0: + bo.buffer_store( + code, rsrc(c["x8s"]), t * X_GROUPS + col // 32, cache_modifier=CM_DEV + ) + + +@traced +def publish_x8(c, t): + """Token t's two XRDY flags (one a wqkv K half) once its X8 / X8S landed.""" + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if c["tid"] < 2: + c["mb"].put(c["xrdy"], 2 * t + c["tid"], fx.Int32(1)) + gpu.barrier() + + +def build_mono_k1(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + fkey = _front_key(key) + FrontLds = mb_front.front_lds(s, Dims(tp)) + NORM0 = SLICES # the norms on the CTAs past the slices: the wqkv waits on them + GATE0 = SLICES + s + assert GATE0 + s <= BLOCKS + + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + + @fx.union + class K1Lds: + seam: SeamLds + front: FrontLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k1", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k1( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K1Lds) + sl = lds.seam.peek() + ep = _epoch(epoch) + tag = (ep << 1) + 1 + fargs = { + "x": 0, + "x_stride": 0, + "wqkv": wqkv, + "wqkv_s": wqkv_s, + "qn_w": qn_w, + "kvn_w": kvn_w, + "wqb": wqb, + "wqb_s": wqb_s, + "cos_sin": cos_sin, + "pos": pos, + "slot": slot, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "q": q, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "topk": topk, + "t2r": t2r, + "kt": kt, + "klen": klen, + } + flayout = {n: layout[n] for n in front_scratch(s, Dims(tp))} + c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + amb = Mailbox(tag) + tid = fx.thread_idx.x + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "args": { + "res_in": res_in, + "pend": pend, + "post_in": post_in, + "comb_in": comb_in, + "pre_in": pre_in, + "hc_fn": hc_fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "norm_w": attn_w, + "res_out": res_out, + "post_out": post_out, + "comb_out": comb_out, + "pre_out": pre_out, + }, + "d": StageDims(tp), + } + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_front.FRONT_POINTS + + def seam_slice(): + for task in range(first_task(bid, 0), SLICES, BLOCKS): + seam.stage_slice(ca, task) + gpu.barrier() + + def seam_rest(): + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_x8(c, *v)) + publish_x8(c, t) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + gpu.barrier() + + mb_front.run_front( + c, fkey, bid, x8_given=True, before_wqkv=seam_slice, after_wqkv=seam_rest + ) + + @flyc.jit + def launch( + res_in: Int64, + pend: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + attn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + wqkv: Int64, + wqkv_s: Int64, + qn_w: Int64, + kvn_w: Int64, + wqb: Int64, + wqb_s: Int64, + cos_sin: Int64, + pos: Int64, + slot: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + q: Int64, + swa_idx: Int64, + swa_lens: Int64, + comp_block: Int32, + comp_bt: Int64, + bt_stride: Int32, + topk: Int64, + t2r: Int64, + kt: Int64, + klen: Int64, + scratch: Int64, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k1( + res_in, + pend, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + attn_w, + res_out, + post_out, + comb_out, + pre_out, + wqkv, + wqkv_s, + qn_w, + kvn_w, + wqb, + wqb_s, + cos_sin, + pos, + slot, + swa, + swa_stride, + swa_block, + q, + swa_idx, + swa_lens, + comp_block, + comp_bt, + bt_stride, + topk, + t2r, + kt, + klen, + scratch, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- K2 + + +@traced +def store_normed(normed, t, col, ys, live): + """The FFN seam's norm output -> the MoE input row (bf16, device scope: + this launch reads it).""" + if live: + for j in range_constexpr(8): + bo.buffer_store( + ys[j].to(fx.BFloat16), + rsrc(normed), + t * HIDDEN + col + j, + cache_modifier=CM_DEV, + ) + + +def _seam_lds(s): + @fx.struct + class SeamLds: + red: fx.Array[fx.Float32, WAVES * 2 * 64 * 4, 16] + rl: fx.Array[fx.Float32, s * seam.KT, 16] + fl: fx.Array[fx.Float32, seam.MIX * seam.KT, 16] + pl: fx.Array[fx.Float32, s * seam.COLS, 16] + + return SeamLds + + +@traced +def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None): + """K2's FFN half on this CTA: the FFN seam (the attention's TP partials -- + ``part`` pushed to every rank's ATTN region first -- summed and folded into + the residual), the MoE and its all-reduce into ``a["out"]``, then the + launch pair's epoch moved on.""" + s, tp = key.tokens, key.tp + layout = scratch_layout(s, tp) + mkey = _moe_key(key) + rs = route_shape(mkey) + GATE0 = SLICES + NORM0 = SLICES + s + bid = fx.block_idx.x + tid = fx.thread_idx.x + sl = lds.seam.peek() + scratch = a["scratch"] + seam_args = ( + "res_in", + "post_in", + "comb_in", + "pre_in", + "hc_fn", + "hc_scale", + "hc_base", + "res_out", + "post_out", + "comb_out", + "pre_out", + ) + ca = { + "S": s, + "tid": tid, + "bid": bid, + "lane": tid % 64, + "wave": tid // 64, + "rl": sl.rl.ptr, + "fl": sl.fl.ptr, + "red": sl.red.ptr, + "pend_lds": sl.pl.ptr, + "put": amb.put, + "put_bf": amb.put_bf, + "put_words": amb.put_words, + "poll": amb.poll, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": own, + "args": {n: a[n] for n in seam_args} | {"norm_w": a["ffn_w"]}, + "d": StageDims(tp), + } + normed = scratch + fx.Int64(layout["normed"][0]) + for region in ("lin", "pmix"): + ca[region] = sreg(scratch, layout[region][0], region) + mlayout = {n: layout[n] for n in moe.scratch_layout(mkey) if n not in COUNTERS} + mlayout["xrdy"] = layout["xrdy_moe"] + margs = moe.moe_args( + normed, + a["gate_w"], + a["bias"], + a["w13"], + a["w13_s"], + a["w2"], + a["w2_s"], + a["sgu"], + a["sgu_s"], + a["sw2"], + a["sw2_s"], + a["out"], + ) + cb = moe.moe_context( + s, rs, lds.moe.peek(), amb, bases, mlayout, scratch, margs, rank, own + ) + cb["x_cm"], cb["x_ready"] = CM_DEV, True + cb["tag"] = ep & 255 + for i, region in enumerate(COUNTERS): + cb[region] = sreg(a["epoch"], 4 * (EPOCH_COUNTERS + i * moe.TAGS), region) + for task in range(first_task(bid, 0), SLICES, BLOCKS): + if const_expr(part is not None): + seam.stage_push(ca, task, part) + seam.stage_reduce(ca, task) + seam.stage_slice(ca, task) + for t in range(first_task(bid, GATE0), s, BLOCKS): + seam.stage_gate(ca, t) + for t in range(first_task(bid, NORM0), s, BLOCKS): + seam.stage_norm(ca, t, lambda *v: store_normed(normed, *v)) + publish(cb["put"], cb["xrdy"], t, 1, tid == 0) + gpu.barrier() + + # ---- the MoE, its all-reduce into ``out`` + moe.run_moe(cb, mkey, bid, 0, 0) + + def reset(): + # the MoE counter slot 128 launch pairs ahead: its last use long + # done, its next use far off + slot = (ep + 128) & 255 + for region in COUNTERS: + gstore(cb[region].value + fx.Int64(slot * 4), fx.Int32(0), words=1) + + mb_back.epoch_end({"bid": bid, "tid": tid}, a["epoch"], ep, reset) + + +def build_mono_k2(key: MonoBuild): + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + layout = scratch_layout(s, tp) + bkey = _back_key(key) + mkey = _moe_key(key) + rs = route_shape(mkey) + assert rs.experts // 64 in SORT_NETS + BackM = mb_back.back_lds_members(s, Dims(tp)) + SplitLds, GemvLds = BackM["split"], BackM["gemv"] + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class K2Lds: + split: SplitLds # type: ignore[valid-type] + gemv: GemvLds # type: ignore[valid-type] + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def k2( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + ): + _ = keyed + bid = fx.block_idx.x + lds = fx.SharedAllocator().allocate(K2Lds) + ep = _epoch(epoch) + tag = (ep << 1) + 2 + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + own = sym + par + bargs = { + "q": q, + "swa": swa, + "swa_stride": swa_stride, + "swa_block": swa_block, + "comp": comp, + "comp_stride": comp_stride, + "comp_block": comp_block, + "kt": kt, + "klen": klen, + "pos": pos, + "sink": sink, + "qk_scale": qk_scale, + "cos_sin": cos_sin, + "woa": woa, + "woa_s": woa_s, + "wob": wob, + "wob_s": wob_s, + "out": 0, + "out_stride": 0, + "zrec": zrec, + } + blayout = {n: layout[n] for n in back_scratch(s, Dims(tp), start=0)} + views = {"split": lds.split.peek(), "gemv": lds.gemv.peek()} + c = mb_back.back_context(bkey, views, bargs, scratch, tag, blayout) + amb = Mailbox(tag) + + def push(t, col, v0, v1): + # this rank's bf16 partial pair to every rank's ATTN region + for p in range_constexpr(tp): + dst = preg(bases[p], 0, "attn") + amb.put_bf(dst, 2 * seam.attn_region_pair(rank, t, s, col), [v0, v1]) + + c["wob_out"] = push + if const_expr(key.timeline): + c["tl"], c["tl_points"] = tl, mb_back.BACK_POINTS + mb_back.epoch_begin(c, epoch, ep) + mb_back.run_back(c, bkey, bid) + gpu.barrier() + + ffn = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + run_ffn(key, lds, ffn, amb, bases, own, rank, ep) + + @flyc.jit + def launch( + q: Int64, + swa: Int64, + swa_stride: Int64, + swa_block: Int32, + comp: Int64, + comp_stride: Int64, + comp_block: Int32, + kt: Int64, + klen: Int64, + pos: Int64, + sink: Int64, + qk_scale: Int32, + cos_sin: Int64, + woa: Int64, + woa_s: Int64, + wob: Int64, + wob_s: Int64, + zrec: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + tl: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + k2( + q, + swa, + swa_stride, + swa_block, + comp, + comp_stride, + comp_block, + kt, + klen, + pos, + sink, + qk_scale, + cos_sin, + woa, + woa_s, + wob, + wob_s, + zrec, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch + + +# ---------------------------------------------------------------- the FFN launch + + +def build_mono_ffn(key: MonoBuild): + """The FFN launch of a layer whose attention vLLM ran (``part``: this rank's + unreduced wo_b output, bf16 [S, HIDDEN]): K2's FFN half, the partial pushed + to every rank first. It takes a launch pair's epoch, as K2 does.""" + s, tp = key.tokens, key.tp + assert 1 <= s <= MAX_TOKENS + rs = route_shape(_moe_key(key)) + assert rs.experts // 64 in SORT_NETS + MoeLds = moe.moe_smem(s, rs, False) + half = peer_half_bytes(tp) + assert SLICES + 2 * s <= BLOCKS + SeamLds = _seam_lds(s) + + @fx.union + class FfnLds: + seam: SeamLds # type: ignore[valid-type] + moe: MoeLds # type: ignore[valid-type] + + name = kernel_symbol("dsv41_mono_ffn", s=s, tp=tp) + keyed = key_tuple(key, SOURCES) + + @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) + def kffn( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + ): + _ = keyed + lds = fx.SharedAllocator().allocate(FfnLds) + ep = _epoch(epoch) + par = fx.Int64(ep & 1) * fx.Int64(half) + bases = [b + par for b in peer_bases(peers, tp)] + mb_back.epoch_begin({"bid": fx.block_idx.x, "tid": fx.thread_idx.x}, epoch, ep) + a = dict( + res_in=res_in, + post_in=post_in, + comb_in=comb_in, + pre_in=pre_in, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + ffn_w=ffn_w, + res_out=res_out, + post_out=post_out, + comb_out=comb_out, + pre_out=pre_out, + gate_w=gate_w, + bias=bias, + w13=w13, + w13_s=w13_s, + w2=w2, + w2_s=w2_s, + sgu=sgu, + sgu_s=sgu_s, + sw2=sw2, + sw2_s=sw2_s, + out=out, + scratch=scratch, + epoch=epoch, + ) + amb = Mailbox((ep << 1) + 2) + run_ffn(key, lds, a, amb, bases, sym + par, rank, ep, part=part) + + @flyc.jit + def launch( + part: Int64, + res_in: Int64, + post_in: Int64, + comb_in: Int64, + pre_in: Int64, + hc_fn: Int64, + hc_scale: Int64, + hc_base: Int64, + ffn_w: Int64, + res_out: Int64, + post_out: Int64, + comb_out: Int64, + pre_out: Int64, + gate_w: Int64, + bias: Int64, + w13: Int64, + w13_s: Int64, + w2: Int64, + w2_s: Int64, + sgu: Int64, + sgu_s: Int64, + sw2: Int64, + sw2_s: Int64, + out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + rank: Int32, + epoch: Int64, + stream: fx.Stream = _STREAM, + ): + _ = keyed + kffn( + part, + res_in, + post_in, + comb_in, + pre_in, + hc_fn, + hc_scale, + hc_base, + ffn_w, + res_out, + post_out, + comb_out, + pre_out, + gate_w, + bias, + w13, + w13_s, + w2, + w2_s, + sgu, + sgu_s, + sw2, + sw2_s, + out, + scratch, + sym, + peers, + rank, + epoch, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py new file mode 100644 index 000000000000..eb379902754f --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1-Flash mono decode layer: host side. + +One ``DSV41MonoLayer`` a TP rank runs every mono layer of a decode step (M <= 48 +rows): two persistent launches a layer (``layer``) from the layer's +inputs at the attention seam to its outputs at the next one -- the MoE's +reduced output, the residual after the FFN seam and that seam's mixes -- with +both TP all-reduces inside the kernels (symmetric peer memory). A layer whose +attention vLLM runs (``ffn``) takes one launch from its unreduced attention +output on. Every argument is a device pointer: the launches can be captured in a +HIP graph. The kernels move their mailbox epoch on themselves; every rank runs +the same launches, so the ranks' epochs agree. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path + +import torch + +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Dims +from .common.plan import BLOCKS +from .layer import ( + EPOCH_WORDS, + MAX_TOKENS, + MonoBuild, + build_mono_ffn, + build_mono_k1, + build_mono_k2, + peer_half_bytes, + scratch_bytes, +) +from .stages.dims import Dims as MoeDims +from .stages.moe_shape import EXPERTS + +__all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "DSV41MonoLayer", "MonoLayerWeights"] + +_LOG2E = 1.4426950408889634 +HC = 4 + + +def _ensure_writable_flydsl_cache() -> None: + """Aiter points FlyDSL at its bundled, read-only cache; new kernels need a + writable one.""" + cur = os.environ.get("FLYDSL_RUNTIME_CACHE_DIR") + if cur and os.access(cur, os.W_OK): + return + path = Path.home() / ".flydsl" / "cache" + path.mkdir(parents=True, exist_ok=True) + os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = str(path) + + +def _check_tensors(owner, tag: str, want: dict) -> None: + """Each named tensor of ``owner`` has the (shape, dtype) the kernels read, + contiguous: their buffer loads are unbounded, so a mismatch would read + garbage rather than fault.""" + for name, (shape, dtype) in want.items(): + t = getattr(owner, name) + assert tuple(t.shape) == shape and t.dtype == dtype and t.is_contiguous(), ( + f"{tag} {name}: {tuple(t.shape)} {t.dtype}, want {shape} {dtype}" + ) + + +@dataclass +class AttnWeights: + """One layer's attention tensors on this rank, in vLLM's loaded layout.""" + + layer_id: int + wqkv: torch.Tensor # [1792, 5120] e4m3: [wq_a; wkv] + wqkv_scale: torch.Tensor # [56, 160] uint8 E8M0 (32 x 32 blocks) + q_norm: torch.Tensor # [1280] bf16 + kv_norm: torch.Tensor # [512] bf16 + wq_b: torch.Tensor # [H * 512, 1280] e4m3 + wq_b_scale: torch.Tensor # [H * 16, 40] uint8 + wo_a: torch.Tensor # [G * 1024, 4096] e4m3 + wo_a_scale: torch.Tensor # [G * 32, 128] uint8 + wo_b: torch.Tensor # [5120, G * 1024] e4m3 + wo_b_scale: torch.Tensor # [160, G * 32] uint8 + attn_sink: torch.Tensor # [>= H] f32 + cos_sin: torch.Tensor # [positions, 64] f32 (the layer's rope table) + ratio: int # 0 / 1 / 2: the layer's compress ratio + + def check(self, d: Dims) -> None: + h, g = d.heads, d.groups + want = { + "wqkv": ((1792, HIDDEN), torch.float8_e4m3fn), + "wqkv_scale": ((56, HIDDEN // 32), torch.uint8), + "wq_b": ((h * HEAD_DIM, 1280), torch.float8_e4m3fn), + "wq_b_scale": ((h * HEAD_DIM // 32, 40), torch.uint8), + "wo_a": ((g * 1024, d.group_k), torch.float8_e4m3fn), + "wo_a_scale": ((g * 32, d.group_k // 32), torch.uint8), + "wo_b": ((HIDDEN, g * 1024), torch.float8_e4m3fn), + "wo_b_scale": ((HIDDEN // 32, g * 32), torch.uint8), + } + _check_tensors(self, f"layer {self.layer_id}", want) + assert self.cos_sin.dtype == torch.float32 and self.cos_sin.shape[-1] == 64 + + +@dataclass +class MonoLayerWeights: + """One layer's tensors on this rank, in vLLM's loaded layout (``attn``: + None for a layer whose attention vLLM runs).""" + + attn: AttnWeights | None + hc_attn_fn: torch.Tensor # [24, 4 * 5120] f32 + hc_attn_scale: torch.Tensor # [3] f32 + hc_attn_base: torch.Tensor # [24] f32 + attn_norm: torch.Tensor # [5120] bf16 + hc_ffn_fn: torch.Tensor + hc_ffn_scale: torch.Tensor + hc_ffn_base: torch.Tensor + ffn_norm: torch.Tensor + gate_w: torch.Tensor # [384, 5120] bf16 + bias: torch.Tensor # [384] f32 (e_score_correction_bias) + w13: torch.Tensor # [384, 2 inter, 2560] fp4x2, aiter (16, 16) shuffle + w13_s: torch.Tensor # its scales, aiter shuffle_scale order + w2: torch.Tensor # [384, 5120, inter / 2] fp4x2 + w2_s: torch.Tensor + sgu: torch.Tensor # shared gate_up [2 inter, 5120] e4m3, row-major + sgu_s: torch.Tensor # [2 inter / 32, 160] E8M0 + sw2: torch.Tensor # shared down [5120, inter] e4m3 + sw2_s: torch.Tensor + + def check(self, tp: int) -> None: + """The MoE's tensors: 384 experts, TP-sharded intermediates (the routed + one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0 + blocks for the shared expert.""" + d, fp4, e4m3, u8 = ( + MoeDims(tp), + torch.float4_e2m1fn_x2, + torch.float8_e4m3fn, + torch.uint8, + ) + want = { + "gate_w": ((EXPERTS, HIDDEN), torch.bfloat16), + "bias": ((EXPERTS,), torch.float32), + "w13": ((EXPERTS, 2 * d.inter, HIDDEN // 2), fp4), + "w13_s": ((EXPERTS * 2 * d.inter, HIDDEN // 32), u8), + "w2": ((EXPERTS, HIDDEN, d.inter // 2), fp4), + "w2_s": ((EXPERTS * HIDDEN, d.down_scale_cols), u8), + "sgu": ((2 * d.sh_inter, HIDDEN), e4m3), + "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8), + "sw2": ((HIDDEN, d.sh_inter), e4m3), + "sw2_s": ((HIDDEN // 32, d.sh_inter // 32), u8), + } + _check_tensors(self, "MoE", want) + + +class DSV41MonoLayer: + """The mono decode runner of one TP rank (``group``: its TP group, for the + peer-memory handle exchange; a gloo / CPU group).""" + + def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): + _ensure_writable_flydsl_cache() + from .common.peer_memory import PeerBuffer + + self.tp, self.rank = tp, rank + self.d = Dims(tp) + dev = self.device = torch.device(device) + self._scratch: dict[int, torch.Tensor] = {} + # [epoch, -, -, -, a mark per CTA, the MoE's counters] + self.epoch = torch.zeros(EPOCH_WORDS, dtype=torch.int32, device=dev) + self.q = torch.zeros( + MAX_TOKENS, self.d.heads, HEAD_DIM, dtype=torch.bfloat16, device=dev + ) + self.kt = torch.zeros(MAX_TOKENS * KEYS, dtype=torch.int32, device=dev) + self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev) + self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) + self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # the attention seam's outputs, K1 -> K2 + self.res_mid = torch.zeros( + MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev + ) + self.post_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.comb_a = torch.zeros(MAX_TOKENS, HC, HC, dtype=torch.float32, device=dev) + self.pre_a = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self.peer = PeerBuffer(2 * peer_half_bytes(tp), group, rank, tp, dev) + self.peer.bytes.zero_() + self._qk_scale = ( + torch.tensor([HEAD_DIM**-0.5 * _LOG2E], dtype=torch.float32) + .view(torch.int32) + .item() + ) + self._kernels: dict = {} + + def scratch(self, tokens: int) -> torch.Tensor: + """Step width ``tokens``'s scratch: a width's layout has memory of its own + (``layer.scratch_layout``). Allocated at the width's first step, which is + eager: vLLM runs every graph's batch eagerly before capturing it.""" + buf = self._scratch.get(tokens) + if buf is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"DSv4.1 mono decode: step width {tokens} first reached inside " + "a CUDA graph capture" + ) + buf = torch.zeros( + scratch_bytes(tokens, self.tp), dtype=torch.uint8, device=self.device + ) + self._scratch[tokens] = buf + return buf + + @staticmethod + def supports(tokens: int) -> bool: + return 1 <= tokens <= MAX_TOKENS + + def kernels(self, tokens: int, ratio: int): + key = (tokens, ratio) + if key not in self._kernels: + b = MonoBuild(tokens, self.tp, ratio) + self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) + return self._kernels[key] + + def ffn_kernel(self, tokens: int): + key = (tokens, "ffn") + if key not in self._kernels: + self._kernels[key] = build_mono_ffn(MonoBuild(tokens, self.tp, 0)) + return self._kernels[key] + + @staticmethod + def _outs(M: int, residual: torch.Tensor) -> tuple: + dev = residual.device + return ( + torch.empty(M, HIDDEN, dtype=torch.bfloat16, device=dev), + torch.empty_like(residual), + torch.empty(M, HC, 1, dtype=torch.float32, device=dev), + torch.empty(M, HC, HC, dtype=torch.float32, device=dev), + torch.empty(M, HC, dtype=torch.float32, device=dev), + ) + + def forward( + self, + w: MonoLayerWeights, + x: torch.Tensor, # [M, 5120] bf16: the previous FFN output (reduced) + residual: torch.Tensor, # [M, 4, 5120] bf16 + post_mix: torch.Tensor, # [M, 4, 1] f32 + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + positions: torch.Tensor, # [M] int64 + slot_mapping: torch.Tensor, # [M] int64 + swa_cache: torch.Tensor, # [blocks, block, 584] uint8 + swa_indices: torch.Tensor, # [>= M, 1, 128] int32 + swa_lens: torch.Tensor, # [>= M] int32 + token_to_req: torch.Tensor, # [>= M] int32 + topk_indices: torch.Tensor | None = None, + comp_cache: torch.Tensor | None = None, + comp_block_table: torch.Tensor | None = None, + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs + (``outs``: buffers to write them into).""" + M = x.shape[0] + assert self.supports(M), M + a = w.attn + assert a is not None + ratio = a.ratio + k1, k2 = self.kernels(M, ratio) + st = torch.cuda.current_stream() + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + for t_ in (x, residual, post_mix, res_mix, pre_mix, positions, slot_mapping): + assert t_.is_contiguous() + assert positions.dtype == slot_mapping.dtype == torch.int64 + assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32 + assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8 + assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584 + swa_block = swa_cache.shape[1] + if ratio: + assert topk_indices is not None and comp_cache is not None + assert comp_block_table is not None + assert topk_indices.dtype == torch.int32 and topk_indices.stride(0) == 512 + assert comp_cache.dtype == torch.uint8 and comp_cache.stride(1) == 584 + assert comp_block_table.dtype == torch.int32 + comp, comp_stride, comp_block = ( + comp_cache, + comp_cache.stride(0), + comp_cache.shape[1], + ) + bt, bt_stride, topk = ( + comp_block_table, + comp_block_table.stride(0), + topk_indices, + ) + else: + comp, comp_stride, comp_block = self._dummy, 0, 1 + bt, bt_stride, topk = self._dummy, 0, self._dummy + k1( + residual.data_ptr(), + x.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_attn_fn.data_ptr(), + w.hc_attn_scale.data_ptr(), + w.hc_attn_base.data_ptr(), + w.attn_norm.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + a.wqkv.data_ptr(), + a.wqkv_scale.data_ptr(), + a.q_norm.data_ptr(), + a.kv_norm.data_ptr(), + a.wq_b.data_ptr(), + a.wq_b_scale.data_ptr(), + a.cos_sin.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + self.q.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + comp_block, + bt.data_ptr(), + bt_stride, + topk.data_ptr(), + token_to_req.data_ptr(), + self.kt.data_ptr(), + self.klen.data_ptr(), + self.scratch(M).data_ptr(), + self.epoch.data_ptr(), + 0, + stream=st, + ) + k2( + self.q.data_ptr(), + swa_cache.data_ptr(), + swa_cache.stride(0), + swa_block, + comp.data_ptr(), + comp_stride, + comp_block, + self.kt.data_ptr(), + self.klen.data_ptr(), + positions.data_ptr(), + a.attn_sink.data_ptr(), + self._qk_scale, + a.cos_sin.data_ptr(), + a.wo_a.data_ptr(), + a.wo_a_scale.data_ptr(), + a.wo_b.data_ptr(), + a.wo_b_scale.data_ptr(), + self._zero_rec.data_ptr(), + self.res_mid.data_ptr(), + self.post_a.data_ptr(), + self.comb_a.data_ptr(), + self.pre_a.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + 0, + stream=st, + ) + return outs + + def ffn( + self, + w: MonoLayerWeights, + part: torch.Tensor, # [M, 5120] bf16: this rank's unreduced wo_b output + residual: torch.Tensor, # [M, 4, 5120] bf16, after the attention seam + post_mix: torch.Tensor, # [M, 4, 1] f32, the attention seam's + res_mix: torch.Tensor, # [M, 4, 4] f32 + pre_mix: torch.Tensor, # [M, 4] f32 + outs: tuple | None = None, + ): + """-> (out, residual, post_mix, res_mix, pre_mix) of a layer whose + attention vLLM ran: the attention's TP reduce, the FFN seam, the MoE and + its all-reduce in one launch.""" + M = part.shape[0] + assert self.supports(M), M + for t_ in (part, residual, post_mix, res_mix, pre_mix): + assert t_.is_contiguous() + assert part.dtype == residual.dtype == torch.bfloat16 + if outs is None: + outs = self._outs(M, residual) + out, res_out, post_out, comb_out, pre_out = outs + self.ffn_kernel(M)( + part.data_ptr(), + residual.data_ptr(), + post_mix.data_ptr(), + res_mix.data_ptr(), + pre_mix.data_ptr(), + w.hc_ffn_fn.data_ptr(), + w.hc_ffn_scale.data_ptr(), + w.hc_ffn_base.data_ptr(), + w.ffn_norm.data_ptr(), + res_out.data_ptr(), + post_out.data_ptr(), + comb_out.data_ptr(), + pre_out.data_ptr(), + w.gate_w.data_ptr(), + w.bias.data_ptr(), + w.w13.data_ptr(), + w.w13_s.data_ptr(), + w.w2.data_ptr(), + w.w2_s.data_ptr(), + w.sgu.data_ptr(), + w.sgu_s.data_ptr(), + w.sw2.data_ptr(), + w.sw2_s.data_ptr(), + out.data_ptr(), + self.scratch(M).data_ptr(), + self.peer.local, + self.peer.addresses.data_ptr(), + self.rank, + self.epoch.data_ptr(), + stream=torch.cuda.current_stream(), + ) + return outs diff --git a/vllm/models/deepseek_v41/amd/mono/sources.py b/vllm/models/deepseek_v41/amd/mono/sources.py new file mode 100644 index 000000000000..49f34d8fed67 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/sources.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The mono layer's kernel sources, in every build's JIT cache key +(``common.plan.key_tuple``): an edit anywhere in this package rebuilds.""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import source_digest + +SOURCES = source_digest(".") diff --git a/vllm/models/deepseek_v41/amd/mono/stages/__init__.py b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py new file mode 100644 index 000000000000..ebb582c99b8d --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The V4.1 stages the mono layer adapts from ROCm/ATOM: the mHC seams +(``seam``), the MoE (``moe``) and the GEMV / LDS helpers they share (``gemv``).""" diff --git a/vllm/models/deepseek_v41/amd/mono/stages/dims.py b/vllm/models/deepseek_v41/amd/mono/stages/dims.py new file mode 100644 index 000000000000..6fd119c77206 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/dims.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/dims.py +"""A TP rank's share of V4.1's MoE widths: every constant the MoE stages derive +from the TP size, in one place (a build key's ``tp`` picks it).""" + +from dataclasses import dataclass + +from vllm.models.deepseek_v41.amd.mono.common.plan import Shard, padded + +# the model's full widths +MOE_INTER = 2304 # a routed expert's intermediate width +SHARED_INTER = 2304 # the shared expert's +# the loader pads a routed expert's rank share to a multiple of this (``FusedMoE`` +# ``pad_align``) with zero weights +MOE_PAD = 128 +ROWS = 16 # GEMV rows a task +UG_PART = 32 # intermediate columns an ug task part: one FP4 group + + +@dataclass(frozen=True) +class Dims: + tp: int + + def __post_init__(self): + shard = Shard(self.tp) + shard.split(MOE_INTER, "routed intermediate") + shard.split(SHARED_INTER, "shared intermediate") + + @property + def inter_real(self) -> int: + """A routed expert's real intermediate width on this rank.""" + return MOE_INTER // self.tp + + @property + def inter(self) -> int: + """... padded as the loader pads it (the rest zero weights).""" + return padded(self.inter_real, MOE_PAD) + + @property + def sh_inter(self) -> int: + """The shared expert's intermediate width on this rank.""" + return SHARED_INTER // self.tp + + @property + def shared_tasks(self) -> int: + return self.sh_inter // ROWS + + @property + def mid_words(self) -> int: + """MXFP8 words of a pick's routed intermediate.""" + return self.inter // 4 + + @property + def down_scale_cols(self) -> int: + """w2_s's e8m0 columns: inter / 32, padded to a multiple of 8 + (aiter ``shuffle_scale``).""" + return padded(self.inter // 32, 8) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/gemv.py b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py new file mode 100644 index 000000000000..31a19ea4c72f --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/gemv.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (its GEMV and LDS helpers). +"""The FP8 GEMV tiles of the seam and MoE stages and the helpers they share: +token tiles, LDS rows of MXFP8 words, scalar loads and the mailbox-to-LDS +copy.""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + UNIT_SCALE, + mfma_scaled, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_NT, + fresh, + rsrc, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES +from vllm.models.deepseek_v41.amd.mono.common.sync import POLL_MAX + +ROWS = 16 # GEMV rows a task +TILE = 16 # an MFMA's N: the tokens a GEMV pass takes (``token_tiles``) + + +def ld_f32(ptr, i): + return fx.Float32(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.f32)) + + +def ld_bf(ptr, i): + return fx.Float32( + fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) + ) + + +def token_tiles(s): + """A step's tokens in GEMV passes: [(first token, tokens)]. A GEMV loads its + weights once and runs every pass on them, an x tile in LDS at a time.""" + return [(t0, min(TILE, s - t0)) for t0 in range(0, s, TILE)] + + +def tile_rows(s): + """The LDS rows an x tile holds.""" + return min(s, TILE) + + +# a token's LDS row of MXFP8 words / E8M0 codes, 16 of which a MFMA B operand +# gathers, is padded by 16 B: at a multiple of 64 words they share a bank group +LDS_PAD = 4 + + +def lds_row(words): + return words + LDS_PAD + + +def lds_at(u, row): + """Element u of a span of rows ``row`` wide, at the ``lds_row`` stride.""" + return u + u // row * LDS_PAD + + +def lds_tail(k): + """Words a K = ``k`` row's last K step reads past its ``lds_row`` (selected + out): the last row's buffer must hold them.""" + return max(0, 32 * -(-k // 128) - lds_row(k // 4)) + + +def plus(t0, v): + """Token ``v`` of the tile from ``t0`` (a first tile adds nothing; ``t0`` + traced: a tile a CTA picks at run time).""" + return v if isinstance(t0, int) and t0 == 0 else v + t0 + + +# ---------------------------------------------------------------- GEMV +def _fp8_step_load(c, r_w, r_ws, k, rg, st, row_major=False): + """K step st's global operands of rows 16 rg ..: the two 64-column weight + chunks (the upper clamped into the row) and the raw scale word. The (16, 16) + shuffle lays a row group out in 32-column blocks of 16 rows, a chunk two of + them (lanes 0-31, 32-63); ``k`` % 64 = 32 (a last quarter step): the lanes + of a block past the row read its last block. ``row_major`` (vLLM's dense + layout, k % 64 == 0): the same registers from a row-major [N, K] weight -- + lane l row 16 rg + l % 16, K bytes 128 st + 16 (l // 16) .. and + 64.""" + lane = c["lane"] + j = lane // 16 + kg = k // 32 + if const_expr(row_major): + assert k % 64 == 0 + row = rg * 16 + lane % 16 + lo_at = row * (k // 4) + st * 32 + j * 4 + hi_at = row * (k // 4) + fx.min(2 * st + 1, k // 64 - 1) * 16 + j * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + if const_expr(k % 64 == 0): + lo_at = (rg * (k // 64) + 2 * st) * 256 + lane * 4 + else: + lo_at = rg * (k * 4) + fx.min(4 * st + lane // 32, kg - 1) * 128 + lane % 32 * 4 + lo = fx.Vector( + bo.buffer_load(r_w, lo_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + if const_expr(k % 64 == 0): + hi_at = (rg * (k // 64) + fx.min(2 * st + 1, k // 64 - 1)) * 256 + lane * 4 + else: + hi_at = rg * (k * 4) + fx.min(2 * st + 1, k // 64 - 1) * 256 + lane * 4 + hi = fx.Vector( + bo.buffer_load(r_w, hi_at, vec_width=4, dtype=T.i32, cache_modifier=CM_NT) + ) + wsi = (rg // 2) * kg + fx.min(st * 4 + j, kg - 1) + word = fx.Int32(bo.buffer_load(r_ws, wsi // 4, vec_width=1, dtype=T.i32)) + return lo, hi, (word >> (wsi % 4 * 8)) & 0xFF + + +def _fp8_step_mfma(c, k, xl, xsl, st, ops, acc, live=None, rows=None): + """``acc`` through K step st's scaled MFMA: ``_fp8_step_load``'s operands + against every token's MXFP8 row in LDS. ``live`` false: a step past the end + (its operands those of the last step), zero at unit scale. ``k`` % 64 = 32: + a lane's lower block past the row is zero on both sides.""" + s, lane = c["S"] if rows is None else rows, c["lane"] + lo, hi, sa = ops + kg = k // 32 + x_row, s_row = lds_row(k // 4), lds_row(kg) + j = lane // 16 + col = fx.min(lane % 16, s - 1) + if const_expr(live is not None): + st = fx.min(st, (k + 127) // 128 - 1) + lo = [live.select(lo[d], fx.Int32(0)) for d in range(4)] + past = (2 * st + 1) * 64 >= k + if const_expr(live is not None): + past = past | ~live + av = fx.Vector.from_elements( + [lo[d] for d in range(4)] + [past.select(fx.Int32(0), hi[d]) for d in range(4)], + fx.Int32, + ) + grp = st * 4 + j + xlo = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + xhi = fx.Vector( + fx.ptr_load( + xl + (col * x_row + st * 32 + 16 + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + ) + if const_expr(k % 64 != 0): + over = 4 * st + j // 2 >= kg + av = fx.Vector.from_elements( + [over.select(fx.Int32(0), av[d]) for d in range(4)] + + [av[4 + d] for d in range(4)], + fx.Int32, + ) + xlo = [over.select(fx.Int32(0), xlo[d]) for d in range(4)] + xv = fx.Vector.from_elements( + [xlo[d] for d in range(4)] + + [past.select(fx.Int32(0), xhi[d]) for d in range(4)], + fx.Int32, + ) + sb = fx.ptr_load(xsl + (col * s_row + fx.min(grp, kg - 1))) + unit = grp >= kg + if const_expr(live is not None): + unit = unit | ~live + sa = unit.select(fx.Int32(UNIT_SCALE), sa) + sb = unit.select(fx.Int32(UNIT_SCALE), sb) + return mfma_scaled(av, xv, acc, sa, sb) + + +def _gemv_steps(c, k, split, tiled=False): + """This wave's K steps (a step past the end marked, ``split`` waves past + split repeat the last one's; ``tiled``: wave w takes part w % split of tile + w / split): [(step, live)].""" + wave = c["wave"] + steps = (k + 127) // 128 + if const_expr(split is None): + per = (steps + WAVES - 1) // WAVES + ragged = steps % WAVES != 0 + return [ + (wave + WAVES * i, (wave + WAVES * i < steps) if ragged else None) + for i in range(per) + ] + assert split <= WAVES and steps % split == 0 + per = steps // split + part = wave % split if tiled else fx.min(wave, split - 1) + return [(part * per + i, None) for i in range(per)] + + +def gemv_fp8_loads(c, w, ws, k, rg, split=None, tiled=False, row_major=False): + """The weight operands of ``gemv_fp8``'s rows for this wave, every one in + flight at once (a step's latency each would put them in series): loads + only. A step past the end loads the last step. ``tiled`` (split x tiles = + WAVES): row groups rg .., one a group of ``split`` waves.""" + steps = (k + 127) // 128 + r_w, r_ws = rsrc(w), rsrc(ws) + if const_expr(tiled): + assert WAVES % split == 0 + rg = rg + c["wave"] // split + return [ + _fp8_step_load( + c, + r_w, + r_ws, + k, + rg, + st if live is None else fx.min(st, steps - 1), + row_major, + ) + for st, live in _gemv_steps(c, k, split, tiled) + ] + + +def gemv_fp8_mfmas(c, k, xl, xsl, red, ops, split=None, tiled=False, rows=None): + """``gemv_fp8_loads``' operands against the rows in LDS (``rows``: a token + tile's, every token's by default) -> + ``red[wave][lane]``; the scaled 16x16x128 MFMA, chunk kc in dwords 0-3 and kc + + 1 in 4-7. ``k`` a multiple of 32: a last half step has its upper chunk zeroed + on both sides at unit scale (a NaN byte past the row must not reach the + MFMA), its loads clamped into the row: the last row group's would read past + the tensor; a last quarter step its lower chunk's upper block too. A step + past the end adds zero at unit scale. + + The K order is the original kernel's, which ``row_sum`` completes: + wave w takes steps w, w + 8, ... (``split`` None: aiter's preshuffled group32 + GEMM, P0.5 microbenchmark 1), or waves 0 .. split - 1 one contiguous K range + each, the other waves zero (aiter's split-K bmm, whose last split sums the + partials in split order).""" + lane, wave = c["lane"], c["wave"] + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i, (st, live) in enumerate(_gemv_steps(c, k, split, tiled)): + acc = _fp8_step_mfma(c, k, xl, xsl, st, ops[i], acc, live, rows) + if const_expr(split is not None and not tiled): + acc = fx.Vector.from_elements( + [(wave < split).select(acc[e], fx.Float32(0.0)) for e in range(4)], + fx.Float32, + ) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) + + +# ---------------------------------------------------------------- mailbox -> LDS +@traced +def poll_copy(c, idx, width, spans, live=None): + """Mailbox spans [(region, pair_of, n, dst, row)] into LDS: element u (pair + ``pair_of(u)``) to dst + u, or ``lds_at(u, row)`` for a span of rows + ``row`` words wide. Thread ``idx`` of ``width`` takes u = idx, idx + + width, ..., POLL_MAX a round trip, each batch stored before the next is + polled (holding them all spills past S = 16); a thread past a span's end + re-reads its last element and stores nothing. + + ``live``: a span's elements the reader needs (traced, <= n), a span each; + a batch past them is not polled and an element past them neither polled + nor written (one past a step's rows is never written).""" + idx = fresh(idx) # its elements' indices not hoisted out of a task loop + if const_expr(live is None): + _poll_store_batches( + c, + [ + (region, pair_of, n, (dst, row), idx + width * i) + for region, pair_of, n, dst, row in spans + for i in range((n + width - 1) // width) + ], + ) + return + for (region, pair_of, n, dst, row), need in zip(spans, live): + chunks = [ + (region, pair_of, need, (dst, row), idx + width * i) + for i in range((n + width - 1) // width) + ] + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + if chunks[b0][4] < need: + _poll_store_batches(c, chunks[b0 : b0 + POLL_MAX]) + + +@traced +def _poll_store_batches(c, chunks): + """``poll_copy``'s chunks [(region, pair_of, n, (dst, row), u)], POLL_MAX a + round trip, each batch stored before the next is polled.""" + for b0 in range_constexpr(0, len(chunks), POLL_MAX): + batch = chunks[b0 : b0 + POLL_MAX] + got = c["poll"]([(r, p(fx.min(u, n - 1)), 1) for r, p, n, _, u in batch]) + for (_, _, n, (dst, row), u), v in zip(batch, got): + at = u if row is None else lds_at(u, row) + if u < n: + fx.ptr_store(v[0], dst + at) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe.py b/vllm/models/deepseek_v41/amd/mono/stages/moe.py new file mode 100644 index 000000000000..63f8947166d2 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe.py @@ -0,0 +1,1635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe.py +"""The mono layer's MoE stages: one V4.1 layer's FFN, from the FFN seam's normed +rows to the reduced output (its TP all-reduce in-kernel). + +One launch a layer and rank, ``BLOCKS`` x ``THREADS``, S = ``tokens`` rows, +E routed experts and top-k (``MoeBuild``: the target's 384 / 6, the DSpark +draft's 128 / 3). The rank holds every routed expert's 1/TP of the +intermediate width (TP-sharded MoE) and the shared expert's 1/TP (``Dims``). +Stages, with TP4's task counts: + + xq (S): the row as MXFP8 for the routed experts (aiter + ``fused_mx_quant_moe_sort``: amax seeded 1e-10) and for the shared + one (ATOM ``quantize_fp8``: amax floored 1e-4) -> X8R, X8 (plain + words, then the token's flag X8RDY) + router (E / 16 x 4, 16 rows x a quarter of K): bf16 GEMV -> the part's + f32 sum -> LOGIT + route (S): the parts summed in order -> bf16 logits; aiter + ``topk_gating_kernel_opt`` (sqrtsoftplus + bias, top-k, renormalized + x 1.5), one wave a token -> ROUTE + shared (36, 16 gate + 16 up rows): FP8 GEMV -> bf16 -> aiter silu_and_mul + (gate <= 10, up in [-10, 10], ocml exp) -> bf16 -> SMID + ug (|U| x 9, 2 x (32 gate + 32 up) rows of one expert of the union U of + every token's top-k -- the draft's |U| x 18 of 1 x (32 + 32): + ``RouteShape.ug_groups``; w13 gate / up interleaved a 16-row tile): + the MXFP4 x MXFP8 GEMV, a wave the gate and up 16-row tiles of one + 32-row scale block over a half (the draft's: a quarter) of K (each + B operand from LDS read once for both), the slices summed + pairwise; then its epilogue: + clamp, silu x up, bf16, + and the stage-2 input quant to MXFP8 (amax seeded 1e-10, the ceil + code) -> MID (plain words at device scope, then the task's flag, UGF) + down (320, 16 rows; a CTA's two run together): per expert of U the MXFP4 + x MXFP8 GEMV of its intermediate (x its routing weight -> bf16), + summed in each token's top-k order, rounded to bf16 at every add + (aiter's stage 2 adds bf16 atomically, in no fixed order); the + shared expert's FP8 down projection (quantize_fp8 of SMID) -> bf16; + routed + shared -> bf16 -> pushed to every rank's FFN region + reduce (160, 32 columns): the TP partials in rank order 0 .. TP - 1 + (``sum_partials``), fp32 -> bf16 -> ``out`` +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.mx import ( + FP4, + FP8, + FP8_MAX, + UNIT_SCALE, + clamp_fp8, + code_ceil, + fp4_tile_load, + mfma_scaled, + rcp_pow2, + scale_index, +) +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + CM_DEV, + CM_NT, + _hw_f32, + ballot, + bf16_round, + bf_hi, + bf_lo, + butterfly, + div_rn, + fp8_pack4, + fresh, + hw_exp2, + hw_rcp, + lane_gather, + lanes_below, + popcount, + row_sum, + rsrc, + stamp, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import ( + BLOCKS, + THREADS, + WAVES, + first_task, + pair_layout, +) +from vllm.models.deepseek_v41.amd.mono.common.sync import preg, publish, sreg +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART + +# helpers the kernel calls are imported by name: FlyDSL keys a build by the +# source of the same-directory functions it names, never through a module +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ROWS, + TILE, + gemv_fp8_loads, + gemv_fp8_mfmas, + ld_f32, + lds_row, + lds_tail, + plus, + poll_copy, + tile_rows, + token_tiles, +) +from vllm.models.deepseek_v41.amd.mono.stages.moe_shape import ( + MASK_WORDS, + MoeBuild, + route_shape, +) +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS +from vllm.models.deepseek_v41.amd.mono.stages.router import ( + router_tile_mfma, + router_weight_loads, + router_x_loads, +) +from vllm.models.deepseek_v41.amd.mono.stages.seam import HIDDEN, attn_peer_bytes + +# a pick tile is a token tile's MFMA N: ``gather_x8r`` fills one from the other +assert mplan.TILE == TILE +SCALE = 1.5 +LIMIT = 10.0 # swiglu_limit +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +DOWN_TASKS = HIDDEN // ROWS # 320 +# (task, slot) scale blocks a down wave holds a round (``stage_down``): 3 took +# no less time, 4 spills +DOWN_R_MAX = 2 +# a down slot's K steps on this many independent MFMA accumulators +DOWN_CHAINS = 2 +# a reduce task's columns: at 128 a thread polled 6 pairs at 48 rows, each +# pass a round trip in series (32: 1-2 passes, S=48 -1 us) +RED_COLS = 32 +RED_TASKS = HIDDEN // RED_COLS # 160 +# CTA b polls route copy b % ROUTE_COPIES: every CTA on one small region +# serializes on its lines (S=24: -0.5 us) +ROUTE_COPIES = 8 +X8_WORDS = HIDDEN // 4 +TAGS = 256 # the ug queue's / down count's slots, a word each: ``c["tag"]`` picks one + +XQ0 = 0 +DOWN0 = 0 +RED0 = 0 + + +# ``timeline`` stamps: kernel start, then the end of each stage's task loop +TL_STAGES = ("xq", "router", "route", "shared", "ug", "down", "reduce") +TL_POINTS = 1 + len(TL_STAGES) + + +# the MoE stages' kernel arguments, in launch order +ARGS = ( + "x", + "gate_w", + "bias", + "w13", + "w13_s", + "w2", + "w2_s", + "sgu", + "sgu_s", + "sw2", + "sw2_s", + "out", +) + + +def scratch_layout(key: MoeBuild) -> dict[str, tuple[int, int]]: + tokens, experts, topk = key.tokens, key.experts, key.topk + rs = route_shape(key) + d = rs.dims + pairs = { + # the routed experts' MXFP8 x (aiter's fused_mx_quant_moe_sort) and the + # shared expert's (ATOM quantize_fp8): their amax floors differ + "x8r": tokens * X8_WORDS // 2, + "x8rs": tokens * HIDDEN // 32 // 2, + "x8": tokens * X8_WORDS // 2, + "x8s": tokens * HIDDEN // 32 // 2, + # a token's four MXFP8 regions written (plain words): ``stage_xq`` + "x8rdy": tokens, + "logit": tokens * experts * ROUTER_PARTS, + "route": ROUTE_COPIES * tokens * 2 * topk, + "smid": tokens * d.sh_inter, + # SMID quantize_fp8'd (``smid_quant``): the down stage's MXFP8 rows + "smq": tokens * d.sh_inter // 4, + "smc": tokens * d.sh_inter // 32, + # a pick a row, in the slot order (``RouteShape.order0``) + "mid": rs.picks * d.mid_words, + "mids": rs.picks * d.inter // 32, + # an ug task's done flag (``stage_ug``) + "ugf": (rs.max_tiles if rs.multi else rs.max_u) * rs.ug_parts, + # the ug task queue: a word a mailbox tag (a layer), counting up from the + # step's clear (``take_ug_task``); plain words, 2 a pair + "ugq": TAGS // 2, + # the down stage's arrival count, a word a tag (``take_down_rank``) + "dq": TAGS // 2, + } + return pair_layout(pairs.items()) + + +def peer_bytes(tokens: int, tp: int) -> int: + """FFN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def ffn_offset(tokens: int, tp: int) -> int: + """FFN's byte offset in the peer buffer: past the ATTN region.""" + return attn_peer_bytes(tokens, tp) + + +def region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +# ---------------------------------------------------------------- xq +@traced +def stage_xq(c, t): + """Token t, a group of 32 a thread.""" + tid = c["tid"] + a = c["args"] + # thread g < 160 owns the 32 values of group g; the rest mirror thread 159 + g = fx.min(tid, HIDDEN // 32 - 1) + mine = tid < HIDDEN // 32 + # 32 bf16 as four 16-byte loads (a load an element made the stage issue bound) + words = [ + fx.Vector( + bo.buffer_load( + rsrc(a["x"]), + (t * HIDDEN + g * 32) // 2 + 4 * q, + vec_width=4, + dtype=T.i32, + cache_modifier=c["x_cm"], + ) + ) + for q in range(4) + ] + v = [(bf_lo if i % 2 == 0 else bf_hi)(words[i // 8][i % 8 // 2]) for i in range(32)] + amax0 = abs(v[0]) + for y in v[1:]: + amax0 = fx.max(amax0, abs(y)) + # the routed experts' (aiter fused_mx_quant_moe_sort): amax floored 1e-10; + # the shared expert's: vLLM's mxfp8_e4m3_quantize (its MXFP8 linear) + if mine: + put_mxfp8_group(c, "x8r", "x8rs", t, g, v, fx.max(amax0, fx.Float32(1e-10))) + put_mxfp8_group_vllm(c, "x8", "x8s", t, g, v, amax0) + + +def put_mxfp8_group(c, words_mb, codes_mb, t, g, v, amax): + """Token t's group g of 32 values ``v`` as MXFP8 (the ceil code of + ``amax`` / 448, x / scale clamped, RNE) -> scratch ``words_mb`` / + ``codes_mb`` (plain words: the token's X8RDY announces them).""" + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + inv = rcp_pow2(code) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +FLT_MIN = 1.1754943508222875e-38 + + +def vllm_mx_code(amax): + """VLLM's ``mxfp8_e4m3_quantize`` scale code of a 32-group: the smallest + power of two >= RN(max(amax, FLT_MIN) / 448), clamped to [0, 254].""" + a = fx.max(amax, fx.Float32(FLT_MIN)) + q = div_rn(a, fx.Float32(FP8_MAX), fx.Float32(1.0 / FP8_MAX)) + bits = q.bitcast(fx.Int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).select( + fx.Int32(1), fx.Int32(0) + ) + return fx.min(fx.max(e, fx.Int32(0)), fx.Int32(254)) + + +def put_mxfp8_group_vllm(c, words_mb, codes_mb, t, g, v, amax): + """``put_mxfp8_group`` with vLLM's rule (``vllm_mx_code``; x * 2^(127 - code), + which code 0 keeps finite).""" + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + words = [] + for w4 in range_constexpr(4): + ws = [ + fp8_pack4(*[clamp_fp8(v[8 * w4 + 4 * h + i] * inv) for i in range(4)]) + for h in range(2) + ] + words.append(ws) + for h in range_constexpr(2): + st_dev( + fx.Vector.from_elements(words[2 * h] + words[2 * h + 1], fx.Int32), + c[words_mb], + t * X8_WORDS + g * 8 + 4 * h, + ) + st_dev(fx.Int32(code), c[codes_mb], t * (HIDDEN // 32) + g) + + +# ---------------------------------------------------------------- router +@traced +def stage_router(c, task): + """Router task (token group, row tile, K part): K part ``part`` of rows + 16 ``tile`` .. of the bf16 gate weight [experts, 5120] times the group's + tokens (``RouteShape.router_groups``), a token tile at a time -- the + decode router's tile (``router``), the + same tiles in the same order -> the part's f32 sum (``route`` sums the + parts in order, rounds to bf16). Every tile's x loads issue first: a task + over every token's tiles in series read their x serially (48 rows: a router + CTA 9 us, 8 of them x).""" + s, tid, red, lane, wave = c["S"], c["tid"], c["red"], c["lane"], c["wave"] + a = c["args"] + experts = c["rs"].experts + groups, per = c["rs"].router_groups + per_group = experts // ROWS * ROUTER_PARTS + tile, part = task % per_group // ROUTER_PARTS, task % ROUTER_PARTS + g0 = 0 if groups == 1 else task // per_group * per + # the group's token tiles: (first token, tokens), tokens traced past one group + tiles = [ + ( + plus(g0, j0), + min(TILE, s - j0) + if groups == 1 + else fx.min(fx.Int32(min(TILE, per - j0)), s - plus(g0, j0)), + ) + for j0 in range(0, per, TILE) + ] + wvs = router_weight_loads(a["gate_w"], tile, part, lane, wave, CM_NT) + xs = [router_x_loads(a["x"], part, n, lane, wave, c["x_cm"], t0) for t0, n in tiles] + for (t0, n), xvs in zip(tiles, xs): + router_tile_mfma(wvs, xvs, lane, wave, red) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + v = row_sum(red, r, tl) + c["put"]( + c["logit"], + (plus(t0, tl) * experts + tile * ROWS + r) * ROUTER_PARTS + part, + v, + ) + gpu.barrier() + + +# ---------------------------------------------------------------- route +def _score(x): + """sqrt(log2(1 + exp2(x log2e)) ln2), x itself past 20 (topk_gating).""" + sp = fx.Float32( + (x > fx.Float32(20.0)).select( + x, + _hw_f32( + "llvm.amdgcn.log.f32", fx.Float32(1.0) + hw_exp2(x * fx.Float32(LOG2E)) + ) + * fx.Float32(LN2), + ) + ) + return fx.Float32(fmath.sqrt(sp)) + + +@traced +def stage_route(c, t): + """VLLM's topk_hash_softplus_sqrt for token t on wave 0 (f32 logits).""" + lane, wave = c["lane"], c["wave"] + experts, topk = c["rs"].experts, c["rs"].topk + a = c["args"] + if wave == 0: + ept = experts // 64 + # each logit's router part sums, added in part order (vLLM's gate: f32 out) + assert ROUTER_PARTS == 4 + got = c["poll"]( + [ + (c["logit"], (t * experts + lane + 64 * i) * ROUTER_PARTS + h, 2) + for i in range(ept) + for h in (0, 2) + ] + ) + logits = [] + for i in range_constexpr(ept): + parts = got[2 * i] + got[2 * i + 1] + tot = fx.Float32(0.0) + for p in range_constexpr(ROUTER_PARTS): + tot = tot + parts[p].bitcast(fx.Float32) + logits.append(tot) + vals, orig, idxs = [], [], [] + for i in range_constexpr(ept): + e = lane + 64 * i + sc = _score(logits[i]) + # aiter's NaN guards: a NaN score weighs 0 and never wins + orig.append(fx.Float32(fx.isnan(sc).select(fx.Float32(0.0), sc))) + vals.append(fx.max(sc + ld_f32(a["bias"], e), fx.Float32(float("-inf")))) + idxs.append(fx.Int32(e)) + for p, q in c["rs"].sort_net: + sw = vals[p] < vals[q] + vals[p], vals[q] = sw.select(vals[q], vals[p]), sw.select(vals[p], vals[q]) + orig[p], orig[q] = sw.select(orig[q], orig[p]), sw.select(orig[p], orig[q]) + idxs[p], idxs[q] = sw.select(idxs[q], idxs[p]), sw.select(idxs[p], idxs[q]) + cursor = fx.Int32(0) + total = fx.Float32(0.0) + my_w, my_id = fx.Float32(0.0), fx.Int32(0) + for k in range_constexpr(topk): + cv, ci, co = fx.Float32(float("-inf")), fx.Int32(0), fx.Float32(0.0) + for i in range_constexpr(ept): + at = cursor == i + cv, ci, co = ( + at.select(vals[i], cv), + at.select(idxs[i], ci), + at.select(orig[i], co), + ) + best = butterfly(cv, (32, 16, 8, 4, 2, 1), fx.max) + ball = ballot(cv == best) + win = fx.Int32( + (ball == fx.Int64(0)).select(fx.Int32(0), fx.Int32(fmath.cttz(ball))) + ) + win_id = fx.Int32(rocdl.readlane(T.i32, ci.ir_value(), win.ir_value())) + won = (cursor < ept) & (ci == win_id) + mine_o = fx.Float32(won.select(co, fx.Float32(0.0))) + w = fx.Int32( + rocdl.readlane( + T.i32, mine_o.bitcast(fx.Int32).ir_value(), (win_id & 63).ir_value() + ) + ).bitcast(fx.Float32) + cursor = cursor + won.select(1, 0) + total = total + w + my_w = (lane == k).select(w, my_w) + my_id = (lane == k).select(win_id, my_id) + scale = fx.Float32(SCALE) / fx.max(total, fx.Float32(1e-20)) + if lane < topk: + for r in range_constexpr(ROUTE_COPIES): + at = (r * c["S"] + t) * 2 * topk + lane + c["put"](c["route"], at, my_id) + c["put"](c["route"], at + topk, my_w * scale) + gpu.barrier() + + +@traced +def load_route(c, tab): + """Every token's (ids, weights) into LDS ``tab``, then the union U of the ids + (ascending): tab[2 S topk ..] = |U|, the experts of U, each (t, k)'s slot in + U, each expert's slot, the picks in slot order (a slot's ascending) and each + slot's first position (``RouteShape``). + + A pick a thread, no scan in series: each pick ORs its token's bit into its + expert's token bitmap, so an expert's pick count is its bitmap's popcount, + a slot's first position the counts of the experts below it (a wave's prefix + sum; slots follow the experts ascending), and a pick's rank in its slot the + bitmap's tokens below its own. The bitmaps are the same whatever order the + ORs land in, so every CTA derives the same pick order.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rs = c["rs"] + topk, max_u, experts = rs.topk, rs.max_u, rs.experts + n = s * topk + assert s <= 32 * MASK_WORDS + # a thread's route words in one poll batch: a pass each was a round trip + # each in series past 256 picks (48 rows: 5 us from the last route to the + # table); a word past 2n re-reads the last and is not stored + idxs = [plus(pas, tid) for pas in range(0, 2 * n, THREADS)] + copy_at = c["bid"] % ROUTE_COPIES * (2 * n) + got = c["poll"]([(c["route"], copy_at + fx.min(idx, 2 * n - 1), 1) for idx in idxs]) + for i in range_constexpr(len(idxs)): + if idxs[i] < 2 * n: + fx.ptr_store(got[i][0], tab + idxs[i]) + for pas in range_constexpr(0, MASK_WORDS * experts, THREADS): + idx = plus(pas, tid) + if idx < MASK_WORDS * experts: + fx.ptr_store(fx.Int32(0), tab + (rs.mask0 + idx)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + fx.llvm.atomic_or( + tab + (rs.mask0 + eid * MASK_WORDS + t // 32), + fx.Int32(1) << (t % 32), + syncscope="workgroup", + ) + gpu.barrier() + base = 2 * n + # a chunk of 64 experts a wave (wave 0 alone through 6 chunks in series + # was 3.3 us of the table's 5 at 48 rows): each chunk's own counts and + # prefixes, its totals through LDS ``red`` (free here), then each chunk + # past the totals of the chunks below it -- the serial walk's values + chunks = experts // 64 + assert chunks <= WAVES + sums = c["red"] + e = lane + 64 * fx.min(wave, chunks - 1) + count = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + e * MASK_WORDS + w)) + count = count + popcount(word) + ball = ballot(count != 0) + # this chunk's inclusive prefixes, lanes ascending: picks, and pick tiles + incl = lane_prefix(count, lane) + tiles = mplan.pick_tiles(count) + tincl = lane_prefix(tiles, lane) if rs.multi else tiles + totals = [popcount(ball), lane_gather(incl, 63), lane_gather(tincl, 63)] + if (wave < chunks) & (lane == 0): + for k in range_constexpr(3): + fx.ptr_store(totals[k].bitcast(fx.Float32), sums + (wave * 3 + k)) + gpu.barrier() + below_c = [fx.Int32(0), fx.Int32(0), fx.Int32(0)] # n_u, first, tile0 + for j in range_constexpr(chunks - 1): + for k in range_constexpr(3): + v = fx.ptr_load(sums + (j * 3 + k)).bitcast(fx.Int32) + below_c[k] = below_c[k] + (j < wave).select(v, fx.Int32(0)) + n_u, first, tile0 = below_c + slot = n_u + lanes_below(ball, lane) + # e afresh: hoisted, its address held a VGPR across the prefix sums (spill) + e_at = fresh(e) + if wave < chunks: + fx.ptr_store(slot, tab + (rs.expert_slot0 + e_at)) + if count != 0: + fx.ptr_store(fx.Int32(e), tab + (base + 1 + slot)) + fx.ptr_store(first + incl - count, tab + (rs.first0 + slot)) + if const_expr(rs.multi): + # the slot's pick tiles: their first index, and each one's slot + tf = tile0 + tincl - tiles + if count != 0: + fx.ptr_store(tf, tab + (rs.tile_first0 + slot)) + for j in range_constexpr(rs.slot_tiles): + if j < tiles: + fx.ptr_store(slot, tab + (rs.tile_slot0 + tf + j)) + if (wave == chunks - 1) & (lane == 0): + n_all = n_u + totals[0] + fx.ptr_store(n_all, tab + base) + fx.ptr_store(fx.Int32(n), tab + (rs.first0 + n_all)) + if const_expr(rs.multi): + fx.ptr_store(tile0 + totals[2], tab + (rs.tile_first0 + n_all)) + gpu.barrier() + if tid < n: + t = tid // topk + eid = fx.ptr_load(tab + (t * 2 * topk + tid % topk)) + slot = fx.ptr_load(tab + (rs.expert_slot0 + eid)) + fx.ptr_store(slot, tab + (base + 1 + max_u + tid)) + # the tokens below t that picked this expert: the picks before tid in it + rank = fx.Int32(0) + for w in range_constexpr(MASK_WORDS): + word = fx.ptr_load(tab + (rs.mask0 + eid * MASK_WORDS + w)) + lo = t - 32 * w + keep = (lo >= 32).select( + fx.Int32(-1), + (lo > 0).select((fx.Int32(1) << fx.max(lo, 0)) - 1, fx.Int32(0)), + ) + rank = rank + popcount(word & keep) + pos = fx.ptr_load(tab + (rs.first0 + slot)) + rank + fx.ptr_store(tid, tab + (rs.order0 + pos)) + gpu.barrier() + + +def lane_prefix(v, lane): + """The wave's inclusive prefix sum of v, lanes ascending.""" + incl = v + for d in (1, 2, 4, 8, 16, 32): + moved = lane_gather(incl, fx.max(lane - d, 0)) + incl = incl + (lane >= d).select(moved, fx.Int32(0)) + return incl + + +# ---------------------------------------------------------------- shared ug +def _silu_mul(g, u): + """VLLM's silu_and_mul_with_clamp (the shared expert's): gate min limit, + up clamped; silu = g / (1 + expf(-g)) rounded to bf16, then times up.""" + g = bf16_round(fx.min(g, fx.Float32(LIMIT))) + u = fx.max(fx.min(u, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + act = bf16_round(g / (fx.Float32(1.0) + fx.Float32(fmath.exp(-g)))) + return act * u + + +@traced +def stage_shared(c, task, gops, t0, n): + """Gate rows 16 task .. (weights ``gops``), up rows 576 + 16 task .. against + tokens t0 .. t0 + n (their x in LDS; ``n`` traced: a tile a CTA picks at run + time) -> bf16 -> silu x up -> SMID; an odd task then quantizes its group of + 32 (``smid_quant``).""" + tid, red, d = c["tid"], c["red"], c["d"] + a = c["args"] + rows = None if isinstance(n, int) else n + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, gops, rows=rows) + gpu.barrier() + gate = fx.Float32(0.0) + if tid < ROWS * n: + gate = row_sum(red, tid % ROWS, tid // ROWS) + gpu.barrier() + uops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, d.shared_tasks + task, row_major=True + ) + gemv_fp8_mfmas(c, HIDDEN, c["xl"], c["xsl"], red, uops, rows=rows) + gpu.barrier() + if tid < ROWS * n: + r = tid % ROWS + t = plus(t0, tid // ROWS) + up = row_sum(red, r, tid // ROWS) + v = fx.Float32(_silu_mul(bf16_round(gate), bf16_round(up)).to(fx.BFloat16)) + c["put"](c["smid"], t * d.sh_inter + task * ROWS + r, v) + fx.ptr_store(v, c["pair"] + tid) + gpu.barrier() + if task % 2 == 1: + smid_quant(c, task, t0, n) + gpu.barrier() + + +@traced +def smid_quant(c, task, t0, n): + """Odd shared task ``task``: its 16 columns (LDS ``pair``) and the even + task's (SMID) are a quantize_fp8 group of 32 (amax floored 1e-4, the ceil + code) -> SMQ / SMC, a thread a token. Once here, not in each of the 256 CTAs + whose down stage reads them: that took a down CTA 14 us at 48 rows.""" + tid, sh = c["tid"], c["d"].sh_inter + assert c["d"].shared_tasks % 2 == 0 # tasks pair into groups + if tid < n: + t = plus(t0, tid) + got = c["poll"]( + [(c["smid"], t * sh + (task - 1) * ROWS + r, 1) for r in range(ROWS)] + ) + vals = [got[r][0].bitcast(fx.Float32) for r in range(ROWS)] + [ + fx.ptr_load(c["pair"] + (tid * ROWS + r)) for r in range(ROWS) + ] + amax = abs(vals[0]) + for y in vals[1:]: + amax = fx.max(amax, abs(y)) + # vLLM's down_proj input quant (mxfp8_e4m3_quantize) + code = vllm_mx_code(amax) + inv = ((254 - code) << 23).bitcast(fx.Float32) + g = task // 2 + for w in range_constexpr(8): + word = fp8_pack4(*[clamp_fp8(vals[4 * w + i] * inv) for i in range(4)]) + c["put_words"](c["smq"], t * (sh // 4) + g * 8 + w, [word]) + c["put_words"](c["smc"], t * (sh // 32) + g, [code]) + + +# ---------------------------------------------------------------- routed ug +@traced +def gemv_fp4_pair(c, w, ws, e, n_rows, k, rg0, ks, xl, xsl, scale_cols, col): + """Tiles rg0 and rg0 + 1 (16 rows each, one 32-row scale block) of expert + e's (16, 16)-preshuffled MXFP4 weight [n_rows, k] times 16 MXFP8 rows in + LDS, over this wave's K slice ``ks`` (``ug_rounds``) -> both tiles' partial + C, a Vector of 8 (rows 4 (lane / 16) .., column lane % 16, whose LDS row is + ``col``). Each B operand read from LDS feeds both tiles' MFMAs: the waves of + a task share x, so a wave a tile over all of K read it once a tile. + """ + lane, wave = c["lane"], c["wave"] + sl = c["wsl"] + wave * wsl_words(c["d"]) + slices = c["rs"].ug_k_slices + rounds = ug_rounds(k, ks, slices) + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + # rounds whose loads are in flight past the one computing: a quarter's + # 2-step rounds two (the draft's 20 rows: -0.5 us) + ahead = slices // 2 + ops = [ + fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[q]) + for q in range(min(ahead, len(rounds))) + ] + for r in range_constexpr(len(rounds)): + sc0, st0, n, sc_steps = rounds[r] + sv, avs = ops[r] + fx.ptr_store(sv, sl + lane * (sc_steps // 2)) + sas = [ + [ + lds_u8( + sl, + scale_index( + t * 16 + lane % 16, (st0 + i) * 4 + lane // 16, scale_cols + ) + - 128 * sc0, + ) + for i in range(n) + ] + for t in range(2) + ] + accs = fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col) + if const_expr(r + ahead < len(rounds)): + ops.append( + fp4_pair_loads( + c, w, ws, e, n_rows, k, rg0, scale_cols, *rounds[r + ahead] + ) + ) + return fx.Vector.from_elements( + [accs[t][q] for t in range(2) for q in range(4)], fx.Float32 + ) + + +def ug_rounds(k, ks, slices): + """K slice ``ks`` of ``slices``' rounds (first step of the scales loaded, + first step, steps, steps of scales loaded); ``ks`` traced: the steps are. + Halves: the 8-step scale blocks 2 r + ks, the odd block left split 4 + 4, + so both run k / 256 steps. Quarters: steps 8 i + 2 ks, + 1 of each block + i, a round loading only their scales (the shuffled layout keeps 2 steps' + together: a dword a lane, where the whole block was 4).""" + blocks = k // 128 // 8 + if slices == 4: + return [(fx.Int32(8 * i) + ks * 2,) * 2 + (2, 2) for i in range(blocks)] + assert slices == 2 + rounds = [ + (fx.Int32(16 * r) + ks * 8, fx.Int32(16 * r) + ks * 8, 8, 8) + for r in range(blocks // 2) + ] + if blocks % 2: + last = (blocks - 1) * 8 + rounds.append((fx.Int32(last), fx.Int32(last) + ks * 4, 4, 8)) + return rounds + + +def fp4_pair_loads(c, w, ws, e, n_rows, k, rg0, scale_cols, sc0, st0, n, sc_steps): + """``gemv_fp4_pair``'s loads of a round: the scales of steps sc0 .. sc0 + + sc_steps (of the tiles' 32-row block: 128 bytes a step, sc_steps / 2 dwords + a lane) and both tiles' weight operands of steps st0 .. st0 + n. Loads + only.""" + lane = c["lane"] + block = (e * n_rows + rg0 * 16) // 32 * (scale_cols * 8) + words = sc_steps // 2 + sv = bo.buffer_load( + rsrc(ws), + block + 32 * sc0 + lane * words, + vec_width=words, + dtype=T.i32, + cache_modifier=CM_NT, + ) + sv = fx.Vector(sv) if words > 1 else fx.Int32(sv) + return sv, [fp4_tile_weights(c, w, e, n_rows, k, rg0 + t, st0, n) for t in range(2)] + + +def scale_block_loads(c, ws, first_row, scale_cols): + """The loads of the 32-row e8m0 block holding ``first_row``: in aiter's + shuffle_scale order a block is contiguous (both 16-row halves, all of K), so + a few 16-byte loads a lane fetch it, where a byte load a K step doubled the + stage's load instructions and bound it by their issue (2026-09-29, ug). A + lane past the block re-reads its last 16 bytes.""" + lane = c["lane"] + words = scale_cols * 8 + block = first_row // 32 * words + return [ + fx.Vector( + bo.buffer_load( + rsrc(ws), + block + fx.min(256 * i + lane * 4, words - 4), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_NT, + ) + ) + for i in range((words + 255) // 256) + ] + + +def down_scale_words(d) -> int: + """A down tile's 32-row scale block in LDS: ``scale_block_store`` writes + whole 256-word rounds.""" + return down_lds_words(d.down_scale_cols * 8) + + +def wsl_words(d) -> int: + """A wave's LDS scale words: an ug batch's 256, or a down round's blocks -- + a (task, slot) each, ``DOWN_R_MAX`` at most (``stage_down``).""" + return max(256, DOWN_R_MAX * down_scale_words(d)) + + +def scale_block_store(c, svs, sl): + """``scale_block_loads``' words -> this wave's LDS block ``sl`` + (``down_scale_words``; a lane past the block writes past it, inside it).""" + lane = c["lane"] + for i in range_constexpr(len(svs)): + fx.ptr_store(svs[i], sl + (256 * i + lane * 4)) + + +def lds_scales(sl, local_row, lane, st0, blk, scale_cols): + """K steps st0 .. st0 + blk's scale bytes of this lane's row from a 32-row + scale block in LDS (``local_row`` its row in the block). An even step and + the next are bytes 2 apart of one word (``scale_index``'s d): a read a + pair, not a step (down: 3 reads a slot, not 5).""" + out, word = [], None + for i in range(blk): + idx = scale_index(local_row, (st0 + i) * 4 + lane // 16, scale_cols) + if (st0 + i) % 2 == 0 or word is None: + word = fx.ptr_load(sl + idx // 4) + out.append((word >> (idx % 4 * 8)) & 0xFF) + return out + + +def lds_u8(ptr, i): + return (fx.ptr_load(ptr + i // 4) >> (i % 4 * 8)) & 0xFF + + +def fp4_tile_weights(c, w, e, n_rows, k, rg, st0, blk, k_real=None): + """The weight operands of expert e's 16-row tile rg, K steps st0 .. st0 + + blk: loads only, nothing waits on them here. ``k_real``: a lane whose 32 + columns lie past it (zero weights) reads its row group's first lane's + address instead (no extra traffic) and holds zero.""" + lane = c["lane"] + base = e * n_rows * (k // 2) // 4 # dwords: expert e + out = [] + for i in range_constexpr(blk): + past = k_real is not None and ((st0 + i) * 128 + 96) >= k_real + if const_expr(past): + real = (st0 + i) * 128 + lane // 16 * 32 < k_real + at = real.select(lane * 4, lane % 16 * 4) + else: + at = lane * 4 + v = fp4_tile_load(rsrc(w), base, rg, k, st0 + i, at) + if const_expr(past): + v = fx.Vector.from_elements( + [real.select(v[q], fx.Int32(0)) for q in range(4)], fx.Int32 + ) + out.append(v) + return out + + +def fp4_pair_mfma(c, k, xl, xsl, st0, n, avs, sas, accs, col): + """``accs`` (two tiles') chained through ``n`` K steps from st0 against MXFP8 + rows in LDS (column lane % 16: row ``col``): a step's B operand loaded once + for both tiles' MFMAs.""" + lane = c["lane"] + x_row = lds_row(k // 4) + s_row = lds_row(k // 32) + j = lane // 16 + accs = list(accs) + for i in range_constexpr(n): + bv = fp8_operand( + [ + fx.ptr_load( + xl + (col * x_row + (st0 + i) * 32 + 16 * h + j * 4), + result_type=fx.Vector.make_type(4, fx.Int32), + ) + for h in range(2) + ] + ) + sb = fx.ptr_load(xsl + (col * s_row + (st0 + i) * 4 + j)) + for t in range_constexpr(2): + accs[t] = mfma_scaled(avs[t][i], bv, accs[t], sas[t][i], sb, FP4, FP8) + return accs + + +def fp8_operand(halves): + """A lane's MXFP8 MFMA operand of a K step: its two 16-byte halves (K + 16 j .. and 64 + 16 j .. of the step, j = lane / 16).""" + return fx.Vector.from_elements( + [fx.Vector(h).bitcast(fx.Int32)[d] for h in halves for d in range(4)], + fx.Int32, + ) + + +def down_lds_words(n): + """A down wave's LDS block of n words: ``scale_block_store`` writes whole + 256-word rounds.""" + return -(-n // 256) * 256 + + +def st_dev(v, region, i): + """Word i of scratch ``region`` (plain words, not tagged pairs) at device + scope: past the per-XCD L2, so another CTA reads it after the flag.""" + bo.buffer_store(v, rsrc(region.value), i, cache_modifier=CM_DEV) + + +@traced +def take_ug_task(c): + """This CTA's next ug task past the first round, from the launch's queue + (``ugq``, a word a tag, counting from 0): the CTAs that ran a shared, + router or route task come to it late and so take fewer, where a fixed + placement gave every CTA the same count and down waited on theirs. A + task's result is the same whichever CTA runs it.""" + if c["tid"] == 0: + at = c["ugq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + task = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(task), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +@traced +def stage_ug(c, task): + """One token tile: slot task // ug_parts of U (``RouteShape.ug_parts``), the + ``ug_groups`` FP4 groups of intermediate columns from 32 ug_groups (task % + ug_parts) (``ug_wave_pair``); then silu x up, rounded to bf16, and the + MXFP8 requant of each token's 32 values of a group -> MID (the stage-1 + epilogue and the stage-2 input quant). Past one: ``stage_ug_tile``.""" + s, tid, lane, tab = c["S"], c["tid"], c["lane"], c["tab"] + rs = c["rs"] + ug_parts = rs.ug_parts + u = task // ug_parts + if u < c["rs"].n_union(tab, s): + e = c["rs"].union_expert(tab, s, u) + # the slot's picks, a column each (a column past them repeats the last); + # every token's row is in LDS: a column reads its pick's token's + first = c["rs"].slot_first(tab, u) + count = c["rs"].slot_first(tab, u + 1) - first + pick = c["rs"].pick_at(tab, first + fx.min(lane % 16, count - 1)) + ug_pass(c, e, task % ug_parts * rs.ug_groups, first, count, pick // rs.topk) + publish(c["put"], c["ugf"], task, 1, tid == 0) + gpu.barrier() + + +def ug_wave_pair(c, part0): + """This wave's share of an ug task from group ``part0``: group p of its + ``ug_groups`` (wave / (2 slices)), its 16-column half bb = wave % 2 -- w13's + gate tile 2 b and up tile 2 b + 1 (aiter GUGU interleaves them a 16-row + tile; b = 2 (part0 + p) + bb), one 32-row scale block -- over K slice ks = + wave / 2 % slices (``gemv_fp4_pair``): (p, 2 b, ks, whether the group is in + the real width).""" + wave, slices = c["wave"], c["rs"].ug_k_slices + p = wave // (2 * slices) + b = (part0 + p) * 2 + wave % 2 + return p, 2 * b, wave // 2 % slices, part0 + p < c["d"].inter_real // UG_PART + + +def ug_red_slot(rs, p, bb, ks, t): + """RED's 64-lane slot of an ug task's partial C: group p, half bb, K slice + ks, gate (t 0) or up (t 1).""" + return ((p * 2 + bb) * rs.ug_k_slices + ks) * 2 + t + + +def pairwise_sum(vs): + """Vs summed as a balanced tree, neighbours first: (v0 + v1) + (v2 + v3).""" + while len(vs) > 1: + vs = [vs[i] + vs[i + 1] for i in range(0, len(vs), 2)] + return vs[0] + + +def ug_tile_task(c, task): + """Past one token tile, ug task ``task``: its slot's expert, its first + group, and its pick tile's first pick and the picks from it.""" + tab, rs, s = c["tab"], c["rs"], c["S"] + ug_parts = rs.ug_parts + tile = task // ug_parts + u = rs.tile_slot(tab, tile) + first = mplan.tile_pick0(rs.slot_first(tab, u), tile, rs.tile_first(tab, u)) + count = rs.slot_first(tab, u + 1) - first + return rs.union_expert(tab, s, u), task % ug_parts * rs.ug_groups, first, count + + +@traced +def stage_ug_tile(c, task, total): + """Past one token tile: ug task ``task`` of ``total`` (a pick tile of its + slot), then the next task from the queue (``take_ug_task``). No next task's + weights prefetched under this one: held across its epilogue and gather they + spilled the merged K2 (S = 24 / 48), and without them it ran faster.""" + tid, lane = c["tid"], c["lane"] + e, part0, first, count = ug_tile_task(c, task) + gather_x8r(c, first, count) + ug_pass(c, e, part0, first, count, fx.min(lane % 16, count - 1)) + nxt = (total > BLOCKS).select(take_ug_task(c) + BLOCKS, total) + # the epilogue's MID stores done before the flag + rocdl.s_waitcnt(vmcnt=0) + gpu.barrier() + if tid == 0: + c["put"](c["ugf"], task, 1) + return nxt + + +@traced +def ug_pass(c, e, part0, first, count, col): + """An ug task's GEMV and epilogue for the slot's picks first .. first + + min(count, 16) (column lane % 16's x in LDS row ``col``).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a, d, rs = c["args"], c["d"], c["rs"] + p, rg0, ks, real = ug_wave_pair(c, part0) + # a group past the real width (the padding to inter) has zero weights: + # its GEMV is zero unread, and the epilogue still writes its zero MID + acc = fx.Vector.filled(8, 0.0, fx.Float32) + if real: + acc = gemv_fp4_pair( + c, + a["w13"], + a["w13_s"], + e, + 2 * d.inter, + HIDDEN, + rg0, + ks, + c["xl"], + c["xsl"], + HIDDEN // 32, + col, + ) + for t in range_constexpr(2): + slot = ug_red_slot(rs, p, wave % 2, ks, t) + fx.ptr_store( + fx.Vector.from_elements([acc[4 * t + q] for q in range(4)], fx.Float32), + red + (slot * 64 + lane) * 4, + ) + gpu.barrier() + # thread (p, j, i): group p, the slot's pick j, intermediate column + # 32 part + i, a pass of THREADS at a time; a group's 32 threads stay + # lane-aligned + cols = rs.tile_picks + for pas in range_constexpr(0, rs.ug_groups * cols * UG_PART, THREADS): + ug_epilogue(c, part0, pas + tid, first, count) + + +@traced +def load_x8_rows(c, words_mb, codes_mb, xl, xsl, token_of, live): + """MXFP8 rows ``token_of(j)``, j < ``live`` (<= TILE; traced or int), of + scratch ``words_mb`` / ``codes_mb`` into LDS rows j of ``xl`` / ``xsl``: + the first ``live`` lanes wait on their token's X8RDY, then plain 16-byte + loads (half the bytes of the tagged pairs a poll would read).""" + tid = c["tid"] + if tid < TILE: + c["poll"]([(c["x8rdy"], token_of(fx.min(tid, live - 1)), 1)]) + gpu.barrier() + for mb, width, dst in ((words_mb, X8_WORDS, xl), (codes_mb, HIDDEN // 32, xsl)): + quads = width // 4 + for i in range_constexpr(-(-TILE * quads // THREADS)): + q = tid + THREADS * i + row = q // quads + if (q < TILE * quads) & (row < live): + v = fx.Vector( + bo.buffer_load( + rsrc(c[mb].value), + token_of(row) * width + q % quads * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + fx.ptr_store(v, dst + (row * lds_row(width) + q % quads * 4)) + gpu.barrier() + + +def gather_x8r(c, first, count): + """The X8R rows and codes of the slot's picks first .. first + min(count, 16) + into LDS rows 0 .. (only those rows: an ug pass reads a pick's row, the + columns past them repeat the last one's).""" + tab, topk = c["tab"], c["rs"].topk + + def token(j): + return c["rs"].pick_at(tab, first + fx.min(j, count - 1)) // topk + + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], token, fx.min(count, TILE)) + + +@traced +def ug_epilogue(c, part0, idx, first, count): + """Thread ``idx``'s value of an ug task: silu x up of its (group, pick j, + column), the group's MXFP8 scale, its byte of MID (the pick's row, ``first`` + + j in the pick order).""" + lane, red, d, rs = c["lane"], c["red"], c["d"], c["rs"] + cols = rs.tile_picks + if idx < rs.ug_groups * cols * UG_PART: + p = idx // (cols * UG_PART) + part = part0 + p + t = idx % (cols * UG_PART) // UG_PART + i = idx % UG_PART + tile = i // 16 + r = i % 16 + at = (16 * (r // 4) + t) * 4 + r % 4 + g, up = [ + pairwise_sum( + [ + fx.ptr_load(red + ug_red_slot(rs, p, tile, ks, t) * 64 * 4 + at) + for ks in range(rs.ug_k_slices) + ] + ) + for t in range(2) + ] + gc = fx.min(g, fx.Float32(LIMIT)) + uc = fx.max(fx.min(up, fx.Float32(LIMIT)), fx.Float32(-LIMIT)) + sig = hw_rcp(fx.Float32(1.0) + hw_exp2(gc * fx.Float32(-LOG2E))) + # the stage-1 kernel stores the intermediate as bf16; the stage-2 input + # quant (fused_mx_quant_moe_sort) makes it MXFP8: amax floored 1e-10, + # the ceil code, x / scale clamped, RNE + v = bf16_round(gc * sig * uc) + amax = fx.max(butterfly(abs(v), (1, 2, 4, 8, 16), fx.max), fx.Float32(1e-10)) + code = code_ceil(amax * fx.Float32(1.0 / FP8_MAX)) + q8 = clamp_fp8(v * rcp_pow2(code)) + nb = [ + lane_gather(q8.bitcast(fx.Int32), fx.min(lane + q, 63)).bitcast(fx.Float32) + for q in range(1, 4) + ] + # 4 threads' bytes -> one word + word = fp8_pack4(q8, *nb) + row = first + t + if (i % 4 == 0) & (t < count): + st_dev(word, c["mid"], row * d.mid_words + part * 8 + i // 4) + if (i == 0) & (t < count): + st_dev(code, c["mids"], mid_code_index(d, row, part)) + + +# ---------------------------------------------------------------- down +@traced +def stage_down(c, tasks): + """Rows 16 task .. of the hidden output for each of ``tasks`` (a CTA's one or + two, run together: a second pass would double the CTA's time): routed + experts of U (a wave each, round robin), the shared expert, their sum -> + pushed to every rank.""" + s, wave, tab, d = c["S"], c["wave"], c["tab"], c["d"] + n = len(tasks) + nu = c["rs"].n_union(tab, s) + sl = c["wsl"] + wave * wsl_words(d) + sw = down_scale_words(d) + # slots of U a round: the loads of a round stay the same + slots = DOWN_R_MAX // n + per_round = WAVES * slots + # the shared expert's down projection first: its input is ready long + # before the routed results, which a CTA mostly waits on here; after the + # rounds it sat on the path past the last ug task + down_shared(c, tasks) + for rnd in range(0, (nu + per_round - 1) // per_round, 1): + # the round's slots of U (wave, wave + WAVES, ...), every load of them + # and of every task in flight before the first wait: the weights and + # their scale blocks (no ug result: before the round's ug flags), then + # the slots' picks' intermediate rows, each lane its pick's straight + # from MID (``mid_operands``). A slot past U loads the last slot again + # (L2 hits) and contributes nothing + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + ucs = [fx.min(u, nu - 1) for u in us] + firsts = [c["rs"].slot_first(tab, u) for u in ucs] + svs, avs = down_weights(c, tasks, rnd, slots, nu) + # this round's slots' ug tasks done, not every slot of the wave: the + # round starts while later slots' ug run + c["poll"]([(c["ugf"], f, 1) for f in down_ug_flags(c, ucs)]) + counts = [c["rs"].slot_first(tab, u + 1) - f for u, f in zip(ucs, firsts)] + mids = [mid_operands(c, f, k) for f, k in zip(firsts, counts)] + for q in range_constexpr(len(svs)): + scale_block_store(c, svs[q], sl + q * sw) + for i in range_constexpr(slots): + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + mids[i], + firsts[i], + counts[i], + avs[j * slots + i], + us[i] < nu, + ) + # a slot of more than 16 picks: its next pick tile's rows + if const_expr(s > TILE): + for j0 in range_constexpr(TILE, s, TILE): + if j0 < counts[i]: + more = mid_operands(c, firsts[i] + j0, counts[i] - j0) + for j in range_constexpr(n): + down_mfma( + c, + j, + tasks[j], + sl + (j * slots + i) * sw, + more, + firsts[i] + j0, + counts[i] - j0, + avs[j * slots + i], + us[i] < nu, + ) + gpu.barrier() + for t0, nn in token_tiles(s): + for j in range_constexpr(n): + down_combine(c, tasks[j], j, t0, nn) + + +@traced +def take_down_rank(c): + """This CTA's arrival rank at the down stage (DQ, a word a tag): which + CTAs take the extra tiles. A tile's result is the same whichever CTA + computes it.""" + if c["tid"] == 0: + at = c["dq"].value + fx.Int64(c["tag"] * 4) + ptr = fx.inttoptr(fx.PointerType.get(T.i32, fx.AddressSpace.Global, 4), at) + rank = fx.llvm.atomic_add(ptr, fx.Int32(1), syncscope="agent") + fx.ptr_store(fx.Int32(rank), c["tab"] + c["rs"].queue0) + gpu.barrier() + return fx.ptr_load(c["tab"] + c["rs"].queue0) + + +def down_weights(c, tasks, rnd, slots, nu): + """Round ``rnd``'s weights of ``stage_down``: each (task, slot)'s scale block + and weight operands, this wave's slots of U (a slot past U: the last's). + Loads only.""" + s, wave, tab, a, d = c["S"], c["wave"], c["tab"], c["args"], c["d"] + us = [wave + WAVES * (slots * rnd + i) for i in range(slots)] + es = [c["rs"].union_expert(tab, s, fx.min(u, nu - 1)) for u in us] + svs = [ + scale_block_loads(c, a["w2_s"], e * HIDDEN + task * 16, d.down_scale_cols) + for task in tasks + for e in es + ] + avs = [ + fp4_tile_weights( + c, a["w2"], e, HIDDEN, d.inter, task, 0, d.inter // 128, d.inter_real + ) + for task in tasks + for e in es + ] + return svs, avs + + +def mid_operands(c, first, count): + """The MFMA B operands of picks first .. first + min(count, 16) (column + lane % 16 the pick, past them the last): each K step's two 16-byte halves + of the lane's MID row and its group's code, loaded at device scope straight into + registers -- no LDS block a slot, so a round holds more slots.""" + lane, d = c["lane"], c["d"] + j = lane // 16 + row = first + fx.min(lane % 16, count - 1) + steps = d.inter // 128 + # the step's constant after an opaque lane word folds into the immediate + # offset; summed with j * 4 it was a hoisted VGPR a step (spill) + lane_word = fresh(row * d.mid_words + j * 4) + bvs = [] + for st in range_constexpr(steps): + halves = [] + for h in range_constexpr(2): + v = fx.Vector( + bo.buffer_load( + rsrc(c["mid"].value), + lane_word + (st * 32 + 16 * h), + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ) + if const_expr(st * 128 + 64 * h + 48 >= d.inter_real): + # columns past the rank's real intermediate (the loader pads + # it with zero weights) are never written: stale bytes, which + # a NaN pattern would carry through the zero weights + past = st * 128 + 64 * h + 16 * j >= d.inter_real + v = fx.Vector.from_elements( + [past.select(fx.Int32(0), v[q]) for q in range(4)], fx.Int32 + ) + halves.append(v) + bvs.append(fp8_operand(halves)) + # this lane's codes, one a K step: contiguous in MIDS (``mid_code_index``), + # 4 a load -- a load a step was 5 of a slot's 16 load instructions, and the + # stage is load-issue bound (probe ``harness/down_probe.py``: -8 %) + base = mid_code_index(d, row, j) + sbs = [] + for w0 in range_constexpr(0, steps, 4): + n = min(4, steps - w0) + v = bo.buffer_load( + rsrc(c["mids"].value), + base + w0, + vec_width=n, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + if const_expr(n == 1): + sbs.append(fx.Int32(v)) + else: + v = fx.Vector(v) + sbs += [fx.Int32(v[k]) for k in range(n)] + for st in range_constexpr(steps): + if const_expr((st + 1) * 128 > d.inter_real): + # a group past the real intermediate has no code written either + past = (4 * st + j) * 32 >= d.inter_real + sbs[st] = past.select(fx.Int32(UNIT_SCALE), sbs[st]) + return bvs, sbs + + +def mid_code_index(d, row, group): + """MIDS word of pick row ``row``'s MXFP8 group ``group`` (32 columns): a K + step's 4 groups apart, so the steps of one group % 4 lie together + (``mid_operands``' lane j reads groups j, j + 4, ..).""" + steps = d.inter // 128 + return row * (d.inter // 32) + group % 4 * steps + group // 4 + + +def down_ug_flags(c, ucs): + """This lane's ug flags of a down round's slots ``ucs``: flag f a (slot, + pick tile past one token tile, part), a tile past the slot's repeating its + last; lane l takes flags l, l + 64, .. (TP2's 18 parts a tile, 48 rows: + 108 flags a round).""" + rs = c["rs"] + ug_parts = rs.ug_parts + per_slot = rs.slot_tiles * ug_parts if rs.multi else ug_parts + n = len(ucs) * per_slot + return [ + _down_ug_flag(c, ucs, fx.min(c["lane"] + 64 * j, n - 1), per_slot) + for j in range(-(-n // 64)) + ] + + +def _down_ug_flag(c, ucs, f, per_slot): + rs, tab = c["rs"], c["tab"] + ug_parts = rs.ug_parts + mine = ucs[0] + for i in range_constexpr(1, len(ucs)): + mine = (f // per_slot == i).select(ucs[i], mine) + if const_expr(not rs.multi): + return mplan.ug_flag(mine, f % ug_parts, ug_parts) + tile0 = rs.tile_first(tab, mine) + tiles = rs.tile_first(tab, mine + 1) - tile0 + return mplan.ug_flag( + mplan.polled_tile(tile0, tiles, f % per_slot // ug_parts), + f % ug_parts, + ug_parts, + ) + + +@traced +def down_shared(c, tasks): + """Each of ``tasks``' rows of the shared expert's down projection + (its MXFP8 intermediate, ``load_smid_tile``) for every + token, bf16 -> LDS ``sdl[task][token][16 rows]`` (``down_combine``): + every task's weights in flight first.""" + s, tid, red, d = c["S"], c["tid"], c["red"], c["d"] + a = c["args"] + shops = [ + gemv_fp8_loads(c, a["sw2"], a["sw2_s"], d.sh_inter, task, row_major=True) + for task in tasks + ] + for t0, nn in token_tiles(s): + load_smid_tile(c, t0, nn, d.sh_inter // 32) + for j in range_constexpr(len(tasks)): + gemv_fp8_mfmas( + c, d.sh_inter, c["shared_xl"], c["shared_xsl"], red, shops[j], rows=nn + ) + gpu.barrier() + if tid < ROWS * nn: + v = row_sum(red, tid % ROWS, tid // ROWS) + at = (j * s + plus(t0, tid // ROWS)) * ROWS + tid % ROWS + fx.ptr_store(bf16_round(v).bitcast(fx.Int32), c["sdl"] + at) + gpu.barrier() + + +@traced +def down_combine(c, task, j, t0, n): + """Down task ``task`` (the CTA's j-th) for tokens t0 .. t0 + n: the routed + contributions summed in bf16 in each token's top-k order, plus the shared expert's + (``down_shared``), the sum pushed to every rank.""" + s, tid, d, tab = c["S"], c["tid"], c["d"], c["tab"] + topk = c["rs"].topk + if tid < ROWS * n: + r = tid % ROWS + tl = tid // ROWS + t = plus(t0, tl) + shared = fx.ptr_load(c["sdl"] + ((j * s + t) * ROWS + r)).bitcast(fx.Float32) + # aiter's stage 2 adds each bf16(v * w) into a bf16 buffer: a rounding + # at every add (its order is the atomics'; here top-k order) + routed = fx.Float32(0.0) + for k in range_constexpr(topk): + w = c["rs"].route_w(tab, s, t, k) + v = contrib_load(c, j, (t * topk + k) * 16 + r) + routed = bf16_round(routed + bf16_round(v * w)) + fx.ptr_store(bf16_round(routed + shared), c["pair"] + tid) + gpu.barrier() + if tid < ROWS * n // 2: + rp = tid % (ROWS // 2) + tl = tid // (ROWS // 2) + t = plus(t0, tl) + v0 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp)) + v1 = fx.ptr_load(c["pair"] + (tl * ROWS + 2 * rp + 1)) + col = task * ROWS + 2 * rp + for p in range_constexpr(d.tp): + dst = preg(c["peer_addr"](p), c["ffn_off"], "ffn") + c["put_bf"](dst, 2 * region_pair(c["rank"], t, s, col), [v0, v1]) + gpu.barrier() + + +@traced +def down_mfma(c, j, task, sl, mid, first, count, av, live): + """Down task ``task`` (the CTA's j-th) on a slot of U: the routed MFMA of its + picks first .. first + min(count, 16) (their B operands ``mid``: + ``mid_operands``; the scale block in ``sl``), each pick's contribution.""" + lane, tab, d = c["lane"], c["tab"], c["d"] + sas = lds_scales( + sl, task % 2 * 16 + lane % 16, lane, 0, d.inter // 128, d.down_scale_cols + ) + jj = lane % 16 + col = fx.min(jj, count - 1) + bvs, sbs = mid + # the K steps on DOWN_CHAINS accumulators in turn, summed in order: one + # chain put every FP8-rate MFMA's latency in series + accs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(DOWN_CHAINS)] + for i in range_constexpr(len(av)): + h = i % DOWN_CHAINS + accs[h] = mfma_scaled(av[i], bvs[i], accs[h], sas[i], sbs[i], FP4, FP8) + acc = accs[0] + for h in range_constexpr(1, DOWN_CHAINS): + acc = acc + accs[h] + pick = c["rs"].pick_at(tab, first + col) + for q in range_constexpr(4): + r = 4 * (lane // 16) + q + if live & (jj < count): + contrib_store(c, j, pick * 16 + r, acc[q]) + + +def contrib_store(c, j, i, v): + """Task j's routed contribution word i (``[pick (t, k)][16 rows]`` f32 bits).""" + fx.ptr_store(v.bitcast(fx.Int32), c["contrib"] + (j * c["rs"].picks * 16 + i)) + + +def contrib_load(c, j, i): + return fx.ptr_load(c["contrib"] + (j * c["rs"].picks * 16 + i)).bitcast(fx.Float32) + + +# ---------------------------------------------------------------- loaders + + +@traced +def load_smid_tile(c, t0, n, groups): + """The shared expert's MXFP8 intermediate of tokens t0 .. t0 + n (SMQ / + SMC, ``smid_quant``) into LDS ``shared_xl`` / ``shared_xsl`` rows 0 .., by every + thread.""" + tid, words = c["tid"], c["d"].sh_inter // 4 + poll_copy( + c, + tid, + THREADS, + [ + (c["smq"], lambda u: plus(t0 * words, u), n * words, c["shared_xl"], words), + ( + c["smc"], + lambda u: plus(t0 * groups, u), + n * groups, + c["shared_xsl"], + groups, + ), + ], + ) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """Columns 128 task ..: the TP ranks' partials, this rank's first, then + rank + 1, + 2, .., in fp32 -> bf16 -> ``out``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * RED_COLS // 2, THREADS): + reduce_pair(c, task, pas + tid) + + +@traced +def reduce_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 64.""" + s, tp = c["S"], c["d"].tp + a = c["args"] + if idx < s * RED_COLS // 2: + t = idx // (RED_COLS // 2) + col = task * RED_COLS + 2 * (idx % (RED_COLS // 2)) + own = preg(c["sym"], c["ffn_off"], "ffn") + acc0, acc1 = sum_partials( + c["poll"], own, lambda src: region_pair(src, t, s, col), tp + ) + bo.buffer_store(acc0.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col) + bo.buffer_store(acc1.to(fx.BFloat16), rsrc(a["out"]), t * HIDDEN + col + 1) + + +def moe_smem(s, rs, timeline): + """The MoE stages' LDS (``timeline``: with its own stamp record).""" + + @fx.struct + class Smem: + # an ug task's partial C: 2 tiles a wave (``ug_red_slot``) + red: fx.Array[fx.Float32, 2 * WAVES * 64 * 4, 16] + tab: fx.Array[fx.Int32, rs.tab_words, 16] + pair: fx.Array[fx.Float32, ROWS * s, 16] + # one stage's buffers at a time (``stage_lds``) + stage: fx.Array[fx.Int32, sum(stage_lds(s, rs)["words"].values()), 16] + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] + + return Smem + + +def stage_lds(s, rs): + """Word offsets into ``Smem.stage`` of the buffers only one stage reads: the + shared stage's MXFP8 x tile; the waves' scale blocks (``wsl``, the ug and + down stages'), then the ug stage's routed MXFP8 x tile, or the down stage's + token tile of the shared expert's MXFP8 intermediate (``load_smid_tile``), + each down task's routed contributions (``contrib_store``) and its shared + expert rows (``down_shared``). A barrier ends every stage.""" + d = rs.dims + # the x tile the shared / ug stages hold (every token's up to one tile) + rows = tile_rows(s) + x8 = rows * lds_row(X8_WORDS) + groups = rows * lds_row(HIDDEN // 32) + wsl = WAVES * wsl_words(d) + shared_xl = rows * lds_row(d.sh_inter // 4) + lds_tail(d.sh_inter) + shared_xsl = rows * lds_row(d.sh_inter // 32) + contrib = 2 * rs.picks * 16 # two down tasks' at most + offsets = { + "x8l": 0, + "x8sl": x8, + "wsl": 0, + "x8rl": wsl, + "x8rsl": wsl + x8, + "shared_xl": wsl, + "shared_xsl": wsl + shared_xl, + "contrib": wsl + shared_xl + shared_xsl, + "sdl": wsl + shared_xl + shared_xsl + contrib, + } + words = { + "shared": x8 + groups, + "ug": wsl + x8 + groups, + "down": wsl + shared_xl + shared_xsl + contrib + 2 * s * ROWS, + } + return {"offsets": offsets, "words": {"all": max(words.values())}} + + +def moe_context(s, rs, lds, mb, bases, layout, scratch, args, rank, sym): + """The MoE stages' context: LDS, mailbox, scratch regions, arguments. Its lane + math starts from a ``fresh`` thread id: in K2 it is not the attention and + seam stages', held from theirs on.""" + tid = fresh(fx.thread_idx.x) + c = { + "S": s, + "tid": tid, + "bid": fx.block_idx.x, + "lane": tid % 64, + "wave": tid // 64, + "red": lds.red.ptr, + "tab": lds.tab.ptr, + "pair": lds.pair.ptr, + "put": mb.put, + "put_bf": mb.put_bf, + "put_words": mb.put_words, + "poll": mb.poll, + "tag": mb.tag, + "peer_addr": lambda p: bases[p], + "rank": rank, + "sym": sym, + "ffn_off": ffn_offset(s, rs.tp), + "rs": rs, + "d": rs.dims, + "args": args, + # x (``ffn_normed``) before this launch: plain loads, no readiness wait + "x_cm": 0, + "x_ready": False, + } + for name, off in stage_lds(s, rs)["offsets"].items(): + c[name] = lds.stage.ptr + off + for region, (off, _) in layout.items(): + c[region] = sreg(scratch, off, region) + return c + + +@traced +def run_moe(c, key, bid, tls, tl0): + """The MoE stages, in order; ``timeline`` stamps from point ``tl0`` + 1.""" + s, timeline, tid = key.tokens, key.timeline, c["tid"] + router_tasks = c["rs"].router_tasks + route0, shared0 = c["rs"].route0, c["rs"].shared0 + shared_tasks = c["d"].shared_tasks + ug_tasks = c["rs"].max_u * c["rs"].ug_parts + for t in range(first_task(bid, XQ0), s, BLOCKS): + await_x(c, [t]) + stage_xq(c, t) + publish(c["put"], c["x8rdy"], t, 1, tid == 0) + stamp(timeline, tls, tid, tl0 + 1) + multi = len(token_tiles(s)) > 1 + for task in range(first_task(bid, c["rs"].router0), router_tasks, BLOCKS): + await_x(c, list(range(s))) + stage_router(c, task) + stamp(timeline, tls, tid, tl0 + 2) + for t in range(first_task(bid, route0), s, BLOCKS): + stage_route(c, t) + stamp(timeline, tls, tid, tl0 + 3) + c["xl"], c["xsl"] = c["x8l"], c["x8sl"] + # a CTA's one shared unit (task, token tile) at most: straight-line code, + # whose tile loads LLVM does not hoist into one live set as it would out of + # a task loop (spills). A task's tiles in series on one CTA put three tiles' + # latency on the path down waits on (the weights read again a tile: L2 hits) + assert c["rs"].shared_units <= BLOCKS + unit = first_task(bid, shared0) + if unit < c["rs"].shared_units: + # the gate weights in flight before x's wait + a = c["args"] + gops = gemv_fp8_loads( + c, a["sgu"], a["sgu_s"], HIDDEN, unit % shared_tasks, row_major=True + ) + if const_expr(not multi): + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: j, s) + stage_shared(c, unit, gops, 0, s) + else: + t0 = unit // shared_tasks * TILE + n = fx.min(fx.Int32(TILE), s - t0) + load_x8_rows(c, "x8", "x8s", c["x8l"], c["x8sl"], lambda j: t0 + j, n) + stage_shared(c, unit % shared_tasks, gops, t0, n) + stamp(timeline, tls, tid, tl0 + 4) + c["xl"], c["xsl"] = c["x8rl"], c["x8rsl"] + if const_expr(multi): + # past one token tile a task gathers its pick tile's rows itself. The + # first round placed (the shared units' CTAs take its last tasks); the + # rest, past a round, from the queue: a CTA a shared or router task + # delayed takes fewer of them + load_route(c, c["tab"]) + n_tiles = c["rs"].tile_first(c["tab"], c["rs"].n_union(c["tab"], s)) + ug_total = mplan.ug_tasks(n_tiles, c["rs"].ug_parts) + task = first_task(bid, c["rs"].ug0) + while task < ug_total: + task = stage_ug_tile(c, task, ug_total) + else: + # one token tile: every task placed (its x in LDS once a CTA) + loaded = fx.Int32(0) + for task in range(first_task(bid, c["rs"].ug0), ug_tasks, BLOCKS): + if loaded == 0: + # X8R (ready long before the routing) first: after the routing + # lands only its own poll remains + load_x8_rows(c, "x8r", "x8rs", c["x8rl"], c["x8rsl"], lambda j: j, s) + load_route(c, c["tab"]) + loaded = fx.Int32(1) + stage_ug(c, task) + stamp(timeline, tls, tid, tl0 + 5) + assert BLOCKS < DOWN_TASKS <= 2 * BLOCKS + # the route table is in LDS already past one token tile (every CTA loads it + # for the ug queue) or when every CTA ran an ug iteration (more ug task + # slots than CTAs), whose first loads it; else load it + if const_expr(not multi and ug_tasks <= BLOCKS): + load_route(c, c["tab"]) + # every CTA its tile; the first DOWN_TASKS - BLOCKS to get here an extra + # one too: the CTAs whose ug ended first (by bid, those ending last held + # two tiles and ended the stage) + down_task = first_task(bid, DOWN0) + rank = take_down_rank(c) + if rank < DOWN_TASKS - BLOCKS: + stage_down(c, [down_task, BLOCKS + rank]) + else: + stage_down(c, [down_task]) + stamp(timeline, tls, tid, tl0 + 6) + for task in range(first_task(bid, RED0), RED_TASKS, BLOCKS): + stage_reduce(c, task) + stamp(timeline, tls, tid, tl0 + 7) + + +def await_x(c, tokens): + """In a launch that also writes x (K2: the FFN seam's norm): tokens' rows + ready.""" + if const_expr(c["x_ready"]): + c["poll"]([(c["xrdy"], t, 1) for t in tokens]) + + +def moe_args(*values) -> dict: + """The MoE stages' arguments by name, in ``ARGS`` order.""" + assert len(values) == len(ARGS) + return dict(zip(ARGS, values)) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py new file mode 100644 index 000000000000..5b5aa6fcdded --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_plan.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/moe_plan.py +"""Every per-step decision of which MoE ug flags (UGF) and intermediate rows +(MID) the ug stage writes and the down stage reads, defined once. + +The route table (``stage_route``) gives each expert with picks a slot of U, its +picks a run from ``slot_first``, and, past one token tile, its picks tiles of +``TILE`` from ``tile_first`` (each tile's slot ``tile_slot``). An ug task covers +a unit -- a slot, or past one token tile a pick tile -- and one of ``ug_parts`` parts +of the intermediate; its flag is ``ug_flag``. A down round polls each of its +slots' flags (``polled_tile`` its tiles) before it reads their MID rows. Called +by the kernel on traced Int32 and by the CPU property test on Python ints. +""" + +from vllm.models.deepseek_v41.amd.mono.common.plan import cdiv, imin + +TILE = 16 # picks a pick tile: an MFMA's N + + +def pick_tiles(count): + """A slot's pick tiles: its ``count`` picks, ``TILE`` a tile.""" + return cdiv(count, TILE) + + +def ug_parts(groups, ug_groups): + """A unit's ug parts: the expert's ``groups`` FP4 groups of its real width, + ``ug_groups`` a part. A part wholly in the padding is not run: MID's padded + columns stay the step's zeros, against zero weights in down.""" + return cdiv(groups, ug_groups) + + +def ug_flag(unit, part, ug_parts): + """The UGF pair of ug task (``unit``, ``part``), ``ug_parts`` parts a unit.""" + return unit * ug_parts + part + + +def ug_tasks(units, ug_parts): + """A step's ug tasks: every part of each of its ``units``.""" + return units * ug_parts + + +def tile_pick0(slot_first, tile, tile_first): + """Pick tile ``tile``'s first pick, its slot's picks from ``slot_first`` and + tiles from ``tile_first``.""" + return slot_first + (tile - tile_first) * TILE + + +def polled_tile(tile_first, tiles, j): + """The pick tile a down round polls as a slot's ``j``-th, the slot's + ``tiles`` tiles from ``tile_first``: the last past them.""" + return tile_first + imin(j, tiles - 1) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py new file mode 100644 index 000000000000..a07e88cfb064 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/moe_shape.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/moe_shape.py +"""The MoE stages' build key and route shape: what they are built for (``MoeBuild``), +and the task counts, placements and LDS route table that follow from its +routing (``RouteShape``), with the table's readers. Plain Python but for the +readers; the kernel's stages are in ``moe``.""" + +from dataclasses import dataclass, field + +import flydsl.expr as fx + +from vllm.models.deepseek_v41.amd.mono.common.plan import BLOCKS, WAVES +from vllm.models.deepseek_v41.amd.mono.stages import moe_plan as mplan +from vllm.models.deepseek_v41.amd.mono.stages.dims import UG_PART, Dims +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ROWS, TILE, token_tiles +from vllm.models.deepseek_v41.amd.mono.stages.router import PARTS as ROUTER_PARTS + +EXPERTS, TOPK = 384, 6 # the target's routing (a build's: ``MoeBuild``) +ROUTER0 = 8 # the router's first task: past the xq tasks (a token each) +MASK_WORDS = 2 # an expert's token bitmap in the route table: tokens <= 64 +# a route lane's experts (EXPERTS / 64) sorted descending, ties kept in index +# order: the network by their count +SORT_NETS = { + 6: [ + (0, 1), + (2, 3), + (4, 5), + (0, 2), + (1, 4), + (3, 5), + (0, 1), + (2, 3), + (4, 5), + (1, 2), + (3, 4), + (2, 3), + ], + 2: [(0, 1)], +} + + +@dataclass(frozen=True) +class MoeBuild: + tokens: int + # routed experts and a token's top-k: the target's, or the DSpark draft's + experts: int = EXPERTS + topk: int = TOPK + # the TP size: every width a rank holds (``Dims``) + tp: int = 4 + timeline: bool = field(default=False, metadata={"sym": "tl"}) + # FP4 groups an ug task (``RouteShape.ug_groups``) + ug_groups: int = 2 + + +@dataclass(frozen=True) +class RouteShape: + """A build's routing: ``experts`` routed experts, ``topk`` a token of + ``tokens``, the task counts and placements that follow, and the route table + ``load_route`` lays out in LDS: every token's (ids, weights), |U|, the + experts of U, each (t, k)'s slot in U, each expert's slot, the picks in slot + order and each slot's first position in it (past one token tile also each + slot's first pick tile and each pick tile's slot).""" + + experts: int + topk: int + tokens: int + tp: int + # FP4 groups (32 intermediate columns) an ug task, its waves splitting K in + # halves (2) or quarters (1). 1 is twice the tasks, so the CTAs' ug ends + # closer together: faster at the draft's widths (40 rows -4 us), slower at + # the target's (48 rows +3 us; the runners pick) + ug_groups: int = 2 + + @property + def dims(self) -> Dims: + return Dims(self.tp) + + @property + def ug_parts(self) -> int: + """An expert's ug tasks (``moe_plan.ug_parts``).""" + return mplan.ug_parts(self.dims.inter_real // UG_PART, self.ug_groups) + + @property + def ug_k_slices(self) -> int: + """The K slices of an ug task's waves: a wave a (group, 16-column half, + slice).""" + return WAVES // (2 * self.ug_groups) + + @property + def picks(self) -> int: + return self.tokens * self.topk + + @property + def max_u(self) -> int: + return min(self.picks, self.experts) + + @property + def router_groups(self) -> tuple[int, int]: + """``stage_router``'s token groups and their tokens: a group a token + tile, as many as one round of CTAs holds (a CTA running two serialized + them: the route waited on the 32 that did at 48 rows).""" + per_group = self.experts // ROWS * ROUTER_PARTS + groups = max(1, min(len(token_tiles(self.tokens)), BLOCKS // per_group)) + tokens = -(-self.tokens // groups) + assert (groups - 1) * tokens < self.tokens + return groups, tokens + + @property + def router_tasks(self) -> int: + # (token group, 16-row tile, K part): ``stage_router`` + return self.router_groups[0] * self.experts // ROWS * ROUTER_PARTS + + @property + def router0(self) -> int: + """Past the xq tasks (a token each from CTA 0).""" + return max(ROUTER0, self.tokens) + + @property + def route0(self) -> int: + return self.router0 + self.router_tasks + + @property + def shared0(self) -> int: + return self.route0 + max(8, self.tokens) + + @property + def shared_units(self) -> int: + """The shared expert's (task, token tile) units (``stage_shared``).""" + return self.dims.shared_tasks * len(token_tiles(self.tokens)) + + @property + def ug0(self) -> int: + """The first ug round's placement: the shared units' CTAs last.""" + return self.shared0 + self.shared_units + + @property + def sort_net(self): + return SORT_NETS[self.experts // 64] + + @property + def tab_words(self) -> int: + return self.tile_slot0 + self.max_tiles if self.multi else self.queue0 + 1 + + @property + def multi(self) -> bool: + """More than one token tile: an ug task is a (tile of a slot's picks, + part) (``stage_ug``).""" + return self.tokens > TILE + + @property + def slot_tiles(self) -> int: + """A slot's pick tiles at most (a token picks an expert once).""" + return len(token_tiles(self.tokens)) + + @property + def max_tiles(self) -> int: + """The pick tiles of U at most: a slot's partial tile each, then the + whole tiles, and never more than a tile a pick.""" + return min(self.picks, self.max_u + self.picks // TILE) + + @property + def tile_first0(self) -> int: + """Past one token tile: each slot's first pick tile, then their count.""" + return self.queue0 + 1 + + @property + def tile_slot0(self) -> int: + """Past one token tile: each pick tile's slot of U.""" + return self.tile_first0 + self.max_u + 1 + + @property + def queue0(self) -> int: + """The ug task this CTA took from the queue (``take_ug_task``).""" + return self.mask0 + MASK_WORDS * self.experts + + @property + def expert_slot0(self) -> int: + """The route table's per-expert slots (a slot where the expert is in U).""" + return 3 * self.picks + 1 + self.max_u + + @property + def order0(self) -> int: + """The picks (t topk + k) in slot order, a slot's ascending.""" + return self.expert_slot0 + self.experts + + @property + def first0(self) -> int: + """Each slot's first position in the pick order, then the pick count.""" + return self.order0 + self.picks + + @property + def mask0(self) -> int: + """Each expert's tokens that picked it, a bit a token (``MASK_WORDS``).""" + return self.first0 + self.max_u + 1 + + @property + def tile_picks(self) -> int: + """A slot's picks at most: a token picks an expert once.""" + return min(self.tokens, 16) + + def route_w(self, tab, s, t, k): + return fx.ptr_load(tab + (t * 2 * self.topk + self.topk + k)).bitcast( + fx.Float32 + ) + + def n_union(self, tab, s): + return fx.ptr_load(tab + 2 * s * self.topk) + + def union_expert(self, tab, s, u): + return fx.ptr_load(tab + (2 * s * self.topk + 1 + u)) + + def slot_first(self, tab, u): + return fx.ptr_load(tab + (self.first0 + u)) + + def pick_at(self, tab, i): + return fx.ptr_load(tab + (self.order0 + i)) + + def tile_first(self, tab, u): + return fx.ptr_load(tab + (self.tile_first0 + u)) + + def tile_slot(self, tab, i): + return fx.ptr_load(tab + (self.tile_slot0 + i)) + + +def route_shape(key: MoeBuild) -> RouteShape: + return RouteShape(key.experts, key.topk, key.tokens, key.tp, key.ug_groups) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/router.py b/vllm/models/deepseek_v41/amd/mono/stages/router.py new file mode 100644 index 000000000000..a7d6f411a816 --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/router.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/model_ops/deepseek_v41/router.py +"""V4.1's router GEMV tile for decode: bf16 x [T, 5120] times the gate weight +[384, 5120] -> f32 logits, in one fixed summation order. + +Rows 16 a tile, K in ``PARTS`` contiguous parts, 16x16x32 bf16 MFMA: in part +p wave w takes the part's K steps w, w + 8, ... (32 columns each) in that +order, and the eight wave partials are summed in wave order +(``common.ops.row_sum``); the MoE's route stage sums the parts in part +order. The MoE stages spread a tile's parts over PARTS CTAs, and past 16 tokens +a token tile a CTA too, which the order allows: a CTA alone cannot read a +tile, or every token's x, fast enough. +""" + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import mfma_bf16, rsrc +from vllm.models.deepseek_v41.amd.mono.common.plan import WAVES + +HIDDEN = 5120 +ROWS = 16 # a tile's rows: the MFMA's M +PARTS = 4 # contiguous K parts of a tile, summed in order + + +def _k0s(part, lane, wave): + """The K offsets of this lane's loads in part ``part`` (8 bf16 each).""" + per_wave = HIDDEN // 32 // WAVES // PARTS + assert per_wave * 32 * WAVES * PARTS == HIDDEN + return [ + (part * per_wave * WAVES + wave + WAVES * i) * 32 + 8 * (lane // 16) + for i in range(per_wave) + ] + + +def router_weight_loads(w, task, part, lane, wave, w_cm=0): + """``router_tile``'s weight operands (rows 16 task .., K part ``part``). + ``w_cm``: their load cache policy (nontemporal where nothing else of the + step reuses them).""" + row = task * ROWS + lane % 16 + return [ + fx.Vector( + bo.buffer_load( + rsrc(w), + (row * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=w_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_x_loads(x, part, s, lane, wave, x_cm=0, t0=0): + """``router_tile``'s x operands: the s tokens from ``t0`` (column + lane % 16; ``t0`` traced or an int). ``x_cm``: their load cache policy + (device scope when another CTA of the same launch wrote x).""" + tok = fx.min(lane % 16, s - 1) + if const_expr(not isinstance(t0, int) or t0 != 0): + tok = tok + t0 + return [ + fx.Vector( + bo.buffer_load( + rsrc(x), + (tok * HIDDEN + k0) // 2, + vec_width=4, + dtype=T.i32, + cache_modifier=x_cm, + ) + ) + for k0 in _k0s(part, lane, wave) + ] + + +def router_tile_mfma(wvs, xvs, lane, wave, red): + """``router_tile``'s MFMAs, K steps in order -> ``red[wave][lane]``.""" + acc = fx.Vector.filled(4, 0.0, fx.Float32) + for i in range_constexpr(len(wvs)): + acc = mfma_bf16(wvs[i].bitcast(fx.BFloat16), xvs[i].bitcast(fx.BFloat16), acc) + fx.ptr_store(acc, red + (wave * 64 + lane) * 4) diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py new file mode 100644 index 000000000000..82abee9f46fa --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from ROCm/ATOM (https://github.com/ROCm/ATOM) at 97359d6, MIT License, +# Copyright (C) 2026, Advanced Micro Devices, Inc.: +# atom/models/deepseek_v41/mono/kernels/attn_pre.py (slice, gate, norm) and +# atom/models/deepseek_v41/mono/kernels/attn_post.py (the attention's TP reduce). +"""The mono layer's mHC seam stages, each a task table of a persistent launch +handing off through the tagged mailbox: + + slice (160, 32 hidden columns x 4 streams each): the owed post folded into + the residual (R_new, bf16 out), the collapsed layer input (mailbox + LIN) and the partial mix projection / sum of squares (mailbox PMIX) + gate (S): the 160 partials -> post / comb (Sinkhorn) and the next pre + norm (S): the RMSNorm -> the bf16 row, handed to the caller's store + reduce (160): the FFN seam's attention output, the TP ranks' partials + (pushed to every rank's ATTN region) summed in the all-reduce's + order -> ``pend_lds``, where its slice reads it + push (160): a rank's partial given in memory (an attention vLLM ran) + to every rank's ATTN region, ahead of the slice's reduce + +Every elementwise step follows aiter's ``mhc_fused_post_pre_delayed_rmsnorm`` +rounding; the mix projection and the sums of squares, split-K reductions, +differ only in their split and order. +""" + +import math + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T + +from vllm.models.deepseek_v41.amd.mono.common.ops import ( + batched_rounds, + bf16_round, + bf_hi, + bf_lo, + butterfly, + hw_exp2, + hw_rcp, + hw_rsq, + lane_gather, + rsrc, + sum_partials, + traced, +) +from vllm.models.deepseek_v41.amd.mono.common.plan import THREADS, WAVES, pair_layout +from vllm.models.deepseek_v41.amd.mono.common.sync import preg +from vllm.models.deepseek_v41.amd.mono.stages.gemv import ( + ld_bf, + ld_f32, + plus, + token_tiles, +) + +HC = 4 +HIDDEN = 5120 +K_MIX = HC * HIDDEN +MIX = HC * (HC + 2) # 24 projections: pre 4 | post 4 | comb 16 +EPS = 1e-20 +HC_EPS = 1e-6 +SINKHORN = 20 +POST_MULT = 2.0 +LOG2E = 1.0 / math.log(2.0) + +COLS = 32 # hidden columns a slice task +SLICES = HIDDEN // COLS # 160 +KT = HC * COLS # projection columns a slice task +PART = MIX + 1 # a slice's partials a token: the mixes, then the sum of squares +GATE_BLOCK = 8 # slices a gate thread sums + + +def scratch_layout(tokens: int) -> dict[str, tuple[int, int]]: + """Mailbox region -> (byte offset, bytes) of a seam; 8 B a pair.""" + pairs = { + "lin": tokens * HIDDEN // 2, # bf16 pairs + "pmix": SLICES * tokens * PART, + } + return pair_layout(pairs.items()) + + +def attn_peer_bytes(tokens: int, tp: int) -> int: + """ATTN: ``[source rank][token][hidden / 2]`` bf16 pairs.""" + return tp * tokens * HIDDEN // 2 * 8 + + +def attn_region_pair(src, t, s, col): + return (src * s + t) * (HIDDEN // 2) + col // 2 + + +def fma(a, b, c): + return fx.Float32(fmath.fma(a, b, c)) + + +def sigmoid(z): + """Triton's ``tl.sigmoid``: 1 / (1 + exp2(-z log2(e))).""" + return fx.Float32(1.0) / (fx.Float32(1.0) + hw_exp2(z * fx.Float32(-LOG2E))) + + +def lane_f32(v, src): + return lane_gather(v.bitcast(fx.Int32), src).bitcast(fx.Float32) + + +def pend_value(c, t, col, cc): + """The owed post's sublayer output at (token t, column col): read from + ``pend``, or where a caller staged it (``pend_lds``, ``[token][32]`` of this + task's columns: a reduction it did first).""" + if const_expr(c.get("pend_lds") is None): + return ld_bf(c["args"]["pend"], t * HIDDEN + col) + return fx.ptr_load(c["pend_lds"] + (t * COLS + cc)) + + +# ---------------------------------------------------------------- slice +@traced +def slice_input(c, task, idx): + """Thread ``idx``'s pair of the layer input (token idx / 16, columns + 32 task + 2 (idx % 16) ..).""" + s, rl = c["S"], c["rl"] + a = c["args"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + cc = 2 * (idx % (COLS // 2)) + outs = [] + for d in range_constexpr(2): + acc = fx.Float32(0.0) + for h in range_constexpr(HC): + rb = fx.ptr_load(rl + (t * KT + h * COLS + cc + d)) + acc = acc + ld_f32(a["pre_in"], t * HC + h) * rb + outs.append(acc) + c["put_bf"](c["lin"], t * HIDDEN + c0 + cc, outs) + + +def slice_at(c0, e): + """Element e of a slice task at column c0: (token, stream, hidden column).""" + return e // KT, (e % KT) // COLS, c0 + e % COLS + + +def slice_loads(c, c0, e): + """Element e's loads of ``stage_slice`` (e clamped into the step's).""" + a = c["args"] + t, o, col = slice_at(c0, e) + r = [ld_bf(a["res_in"], (t * HC + h) * HIDDEN + col) for h in range(HC)] + cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] + return r, cm, ld_f32(a["post_in"], t * HC + o) + + +def slice_row(c, c0, e, got): + """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" + a = c["args"] + t, o, col = slice_at(c0, e) + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), (t * HC + o) * HIDDEN + col) + fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) + + +@traced +def stage_slice(c, task): + """R_new, the layer input and the partial projections of hidden columns + 32 task .. +32, all four streams, every token.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl = c["rl"], c["fl"] + a = c["args"] + c0 = task * COLS + batched_rounds( + tid, + s * KT, + lambda e: slice_loads(c, c0, e), + lambda e, got: slice_row(c, c0, e, got), + ) + for i in range_constexpr((MIX * KT + THREADS - 1) // THREADS): + e = tid + THREADS * i + if e < MIX * KT: + j = e // KT + o = (e % KT) // COLS + # the original's two bf16 MFMAs take fn as hi = bf16(fn) and lo = + # bf16(fn - hi): their sum, exact in fp32, is the weight it applies + f = ld_f32(a["hc_fn"], j * K_MIX + o * HIDDEN + c0 + e % COLS) + hi = bf16_round(f) + fx.ptr_store(hi + bf16_round(f - hi), fl + e) + gpu.barrier() + # the layer input: bf16 R_new of the four streams by the incoming pre-mix, + # products rounded, summed in order (the Triton collapse, no contraction); + # two columns a thread, one bf16 pair a mailbox store + for pas in range_constexpr(0, s * COLS // 2, THREADS): + slice_input(c, task, plus(pas, tid)) + # sum of squares: token t on wave t % WAVES, two columns a lane + for t in range_constexpr(s): + if wave == t % WAVES: + x0 = fx.ptr_load(rl + (t * KT + 2 * lane)) + x1 = fx.ptr_load(rl + (t * KT + 2 * lane + 1)) + sq = butterfly(x0 * x0 + x1 * x1, (32, 16, 8, 4, 2, 1)) + if lane == 0: + c["put"](c["pmix"], (task * s + t) * PART + MIX, sq) + # mixes: fp32 MFMA 16x16x4, A = fn rows (two 16-row tiles), B = R_new^T, + # a token tile at a time + for t0, n in token_tiles(s): + slice_mix(c, task, t0, n) + + +@traced +def slice_mix(c, task, t0, n): + """The mix projection partials of tokens t0 .. t0 + n (``stage_slice``).""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + rl, fl, red = c["rl"], c["fl"], c["red"] + tok = plus(t0, fx.min(lane % 16, n - 1)) + kw = KT // WAVES + cs = [fx.Vector.filled(4, 0.0, fx.Float32) for _ in range(2)] + for kk in range_constexpr(kw // 4): + colk = wave * kw + lane // 16 + 4 * kk + b = fx.ptr_load(rl + (tok * KT + colk)) + for h in range_constexpr(2): + row = h * 16 + lane % 16 + av = fx.ptr_load(fl + (fx.min(row, MIX - 1) * KT + colk)) + av = (row < MIX).select(av, fx.Float32(0.0)) + ops = rocdl._split_mfma_operands([av, b, cs[h], 0, 0, 0]) + cs[h] = fx.Vector(rocdl.mfma_f32_16x16x4f32(T.vec(4, T.f32), *ops).result) + for h in range_constexpr(2): + fx.ptr_store(cs[h], red + ((wave * 2 + h) * 64 + lane) * 4) + gpu.barrier() + if tid < MIX * n: + tl = tid // MIX + t = plus(t0, tl) + j = tid % MIX + h = j // 16 + ln = 16 * ((j % 16) // 4) + tl + tot = fx.Float32(0.0) + for w in range_constexpr(WAVES): + tot = tot + fx.ptr_load(red + (((w * 2 + h) * 64 + ln) * 4 + j % 4)) + c["put"](c["pmix"], (task * s + t) * PART + j, tot) + gpu.barrier() + + +# ---------------------------------------------------------------- gate +def _row(v, lane, q): + """Comb lanes 8..23 hold ``comb[r][c]`` at lane 8 + 4 r + c: element q of + the row.""" + return lane_f32(v, (lane - 8) // 4 * 4 + 8 + q) + + +def _col(v, lane, q): + return lane_f32(v, 8 + 4 * q + (lane - 8) % 4) + + +@traced +def stage_gate(c, t): + """Token t's 160 slices, in slice order, -> post / comb / next pre.""" + s, tid, lane, wave = c["S"], c["tid"], c["lane"], c["wave"] + red = c["red"] + a = c["args"] + groups = SLICES // GATE_BLOCK # 20 + if tid < groups * PART: + g = tid // PART + j = tid % PART + vals = c["poll"]( + [ + (c["pmix"], ((g * GATE_BLOCK + k) * s + t) * PART + j, 1) + for k in range(GATE_BLOCK) + ] + ) + acc = vals[0][0].bitcast(fx.Float32) + for k in range_constexpr(1, GATE_BLOCK): + acc = acc + vals[k][0].bitcast(fx.Float32) + fx.ptr_store(acc, red + (g * PART + j)) + gpu.barrier() + if wave == 0: + j = fx.min(lane, MIX - 1) + mix = fx.ptr_load(red + j) + sq = fx.ptr_load(red + MIX) + for g in range_constexpr(1, groups): + mix = mix + fx.ptr_load(red + (g * PART + j)) + sq = sq + fx.ptr_load(red + (g * PART + MIX)) + sc = [ld_f32(a["hc_scale"], i) for i in range(3)] + base = ld_f32(a["hc_base"], j) + # aiter's mhc_fused_post_pre_delayed_rmsnorm reduce: one rstd = + # rsq(fma(sum, 1/K, eps)), then each gate fma(mix rstd, scale, base) + v = mix * hw_rsq(fma(sq, fx.Float32(1.0 / K_MIX), fx.Float32(EPS))) + if lane < HC: + pre = sigmoid(fma(v, sc[0], base)) + bo.buffer_store(pre + fx.Float32(HC_EPS), rsrc(a["pre_out"]), t * HC + lane) + if (lane >= HC) & (lane < 2 * HC): + post = sigmoid(fma(v, sc[1], base)) + bo.buffer_store( + post * fx.Float32(POST_MULT), rsrc(a["post_out"]), t * HC + lane - HC + ) + cv = fma(v, sc[2], base) + row = [_row(cv, lane, q) for q in range(4)] + # Triton's exp: exp2 of (x - max) log2(e) + m = fx.max(fx.max(row[0], row[1]), fx.max(row[2], row[3])) + cv = hw_exp2((cv - m) * fx.Float32(LOG2E)) + # fast_dividef is x rcp(y); its first "+ eps" contracts into the fma + row = [_row(cv, lane, q) for q in range(4)] + cv = fma(cv, hw_rcp((row[0] + row[1]) + (row[2] + row[3])), fx.Float32(HC_EPS)) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp(((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS)) + for _ in range_constexpr(SINKHORN - 1): + row = [_row(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((row[0] + row[1]) + (row[2] + row[3])) + fx.Float32(HC_EPS) + ) + col = [_col(cv, lane, q) for q in range(4)] + cv = cv * hw_rcp( + ((col[0] + col[1]) + (col[2] + col[3])) + fx.Float32(HC_EPS) + ) + if (lane >= 2 * HC) & (lane < MIX): + bo.buffer_store(cv, rsrc(a["comb_out"]), t * HC * HC + lane - 2 * HC) + gpu.barrier() + + +# ---------------------------------------------------------------- norm +@traced +def stage_norm(c, t, store): + """Token t's RMSNorm (rstd = rsq(fma(sum, 1/H, eps)), bf16((x rstd) w)); + 256 threads, 3 chunks of 8 at 8 t + 2048 c, each chunk's bf16 values handed + to ``store(t, col, ys, live)`` by every thread (``live``: a thread's chunk + in the row; threads 256.. mirror thread 255).""" + tid, lane, wave, red = c["tid"], c["lane"], c["wave"], c["red"] + a = c["args"] + tt = fx.min(tid, 255) + mine = tid < 256 + cols = [tt * 8 + 2048 * ch for ch in range(3)] + xs = [] + for ch in range_constexpr(3): + base = t * HIDDEN + fx.min(cols[ch], HIDDEN - 8) + words = c["poll"]([(c["lin"], base // 2 + k, 1) for k in range(4)]) + v = [] + for k in range_constexpr(4): + v += [bf_lo(words[k][0]), bf_hi(words[k][0])] + xs.append([(cols[ch] < HIDDEN).select(x, fx.Float32(0.0)) for x in v]) + acc = fx.Float32(0.0) + for ch in range_constexpr(3): + for x in xs[ch]: + acc = acc + x * x + acc = butterfly(acc, (1, 2, 4, 8, 16, 32)) + if (lane == 0) & mine: + fx.ptr_store(acc, red + wave) + gpu.barrier() + tot = butterfly(fx.ptr_load(red + lane % 4), (1, 2)) + r = hw_rsq(fma(tot, fx.Float32(1.0 / HIDDEN), fx.Float32(EPS))) + for ch in range_constexpr(3): + col = fx.min(cols[ch], HIDDEN - 8) + ys = [ + bf16_round((xs[ch][j] * r) * ld_bf(a["norm_w"], col + j)) for j in range(8) + ] + store(t, col, ys, mine & (cols[ch] < HIDDEN)) + gpu.barrier() + + +# ---------------------------------------------------------------- reduce +@traced +def stage_reduce(c, task): + """The attention output at this slice's 32 columns: the TP ranks' partials + in the all-reduce's order (``sum_partials``), in fp32 -> bf16 -> + ``pend_lds``.""" + s, tid = c["S"], c["tid"] + for pas in range_constexpr(0, s * COLS // 2, THREADS): + reduce_attn_pair(c, task, plus(pas, tid)) + gpu.barrier() + + +@traced +def reduce_attn_pair(c, task, idx): + """Thread ``idx``'s column pair of token idx / 16 (``stage_reduce``).""" + s = c["S"] + c0 = task * COLS + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = c0 + 2 * (idx % (COLS // 2)) + own = preg(c["sym"], 0, "attn") + acc0, acc1 = sum_partials( + c["poll"], + own, + lambda src: attn_region_pair(src, t, s, col), + c["d"].tp, + ) + lds = c["pend_lds"] + (t * COLS + col - c0) + fx.ptr_store(bf16_round(acc0), lds) + fx.ptr_store(bf16_round(acc1), lds + 1) + + +@traced +def stage_push(c, task, part): + """This rank's attention partial (``part``: bf16 [S, HIDDEN], wo_b's + unreduced output) at this slice's 32 columns, every token, to every rank's + ATTN region: the pairs ``stage_reduce`` polls, as the bits vLLM wrote.""" + s, tid, tp = c["S"], c["tid"], c["d"].tp + for pas in range_constexpr(0, s * COLS // 2, THREADS): + idx = plus(pas, tid) + if idx < s * COLS // 2: + t = idx // (COLS // 2) + col = task * COLS + 2 * (idx % (COLS // 2)) + w = bo.buffer_load( + rsrc(part), (t * HIDDEN + col) // 2, vec_width=1, dtype=T.i32 + ) + pair = attn_region_pair(c["rank"], t, s, col) + for p in range_constexpr(tp): + c["put_words"](preg(c["peer_addr"](p), 0, "attn"), pair, [w]) diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py new file mode 100644 index 000000000000..b404a42de37a --- /dev/null +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``VLLM_ROCM_MONO_DECODE=1``). + +A decode step's backbone layer -- the attention seam, the attention, its TP +all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent +FlyDSL launches (``mono``, this directory) on the layer's loaded +weights, the TP reductions inside the kernels. A layer runs whole when it has +the standard seam and no compressor or indexer of its own (compress ratio 1 or +2) over ``fp8_ds_mla`` records. Any other backbone layer (the first, the Engram +layers, the compressor / indexer layers) keeps vLLM's attention seam and +attention and runs the rest as one launch, the FFN launch, from wo_b's +unreduced output. A step takes the kernels when it is decode only, at most +``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); +a whole layer also needs causal SWA windows. Anything else runs the layer as +before; the paths share every tensor at the layer boundary, so they interleave +freely. + +Every TP rank must take the same path at every layer (the kernels wait on each +other's pushes): the conditions read only metadata every rank holds alike. +""" + +from typing import TYPE_CHECKING, cast + +import torch + +import vllm.envs as envs +from vllm.config import CUDAGraphMode +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata + +logger = init_logger(__name__) + +RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes +SWA_WIDTH = 128 # a causal decode token's window slots +MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats +CDNA_VERSIONS = (4,) + +_runner = None + + +def _mono_runner(device: torch.device): + """The process's mono runner (one a TP rank, shared by every layer). Its + peer-memory handshake is collective over the TP group: every rank builds it + at the same layer of the same step.""" + global _runner + if _runner is None: + if torch.cuda.is_current_stream_capturing(): + # its peer memory is allocated and exchanged eagerly; vLLM warms + # every graph up eagerly before capturing it + raise RuntimeError( + "VLLM_ROCM_MONO_DECODE was first reached inside a CUDA graph " + "capture, before an eager decode step." + ) + from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer + + assert MAX_TOKENS >= MAX_ROWS, MAX_TOKENS + _runner = DSV41MonoLayer( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + get_tp_group().cpu_group, + device, + ) + return _runner + + +def _record_view(kv: torch.Tensor, block: int) -> torch.Tensor: + """A layer's fp8_ds_mla cache as [blocks, block, RECORD] bytes (a block's + data rows, then its scale words: the kernels index both from it).""" + kv = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert kv.stride(0) >= block * RECORD, (tuple(kv.shape), kv.stride(), block) + return torch.as_strided(kv, (kv.shape[0], block, RECORD), (kv.stride(0), RECORD, 1)) + + +def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: + return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) + + +def _compute_units() -> int: + return current_platform.num_compute_units(torch.accelerator.current_device_index()) + + +def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None: + """Why the kernels cannot run this deployment's MoE, or None. They read every + one of the 384 experts on each rank, TP-sharded, in AITER's A8W4 MXFP4 + layout.""" + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + + pc = vllm_config.parallel_config + if pc.enable_expert_parallel or pc.enable_eplb or pc.data_parallel_size > 1: + return "needs tensor parallelism only: no expert or data parallelism" + if envs.VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4: + return "needs AITER's A8W4 MoE (VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4 unset)" + experts = getattr(getattr(layer.ffn, "experts", None), "routed_experts", None) + if experts is None: + return None # not a routed MoE: the layer checks decline it + backend = getattr(experts.quant_method, "mxfp4_backend", None) + if backend != Mxfp4MoeBackend.AITER_MXFP4_BF16: + return f"needs the AITER_MXFP4_BF16 MoE backend, not {backend}" + return None + + +class MonoDecodeLayer: + """One decoder layer's mono path (``create``: None when the layer is not + eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN + launch after vLLM's attention (``ffn``). Its weights are taken from the + layer at the first call, after loading.""" + + def __init__(self, ffn_only: bool = False) -> None: + self.ffn_only = ffn_only + self._weights: MonoLayerWeights | None = None + + @staticmethod + def create( + layer: "DeepseekV4DecoderLayer", vllm_config + ) -> "MonoDecodeLayer | None": + if not envs.VLLM_ROCM_MONO_DECODE: + return None + from vllm.platforms.rocm import get_cdna_version + + # the deployment: an explicit opt-in that cannot run is an error + why = None + if get_cdna_version() not in CDNA_VERSIONS: + why = f"needs CDNA{'/'.join(map(str, CDNA_VERSIONS))}" + elif get_tensor_model_parallel_world_size() not in (2, 4): + why = "needs tensor parallel size 2 or 4" + else: + why = _unsupported_moe(layer, vllm_config) + if why is None: + try: + from vllm.models.deepseek_v41.amd.mono.runner import BLOCKS + except ImportError as err: + why = f"needs FlyDSL and AITER's FlyDSL helpers ({err})" + else: + # every CTA of a launch stays resident: a partitioned GPU deadlocks + cus = _compute_units() + if cus < BLOCKS: + why = f"needs {BLOCKS} compute units, the GPU has {cus}" + if why is not None: + raise ValueError(f"VLLM_ROCM_MONO_DECODE {why}.") + # the layer: the kernels serve the backbone's seams and MoE + config = vllm_config.model_config.hf_config + attn, ffn = layer.attn, layer.ffn + if layer.use_sequence_parallel: + why = "sequence-parallel layer" + elif not layer.fuse_seam_norm: + why = "needs aiter's fused mHC seam" + elif attn.layer_id >= config.num_hidden_layers: + why = "draft layer" + elif ( + ffn.shared_experts is None + or ffn.gate.tid2eid is not None + or ffn.n_routed_experts != 384 + or ffn.n_activated_experts != 6 + or ffn.routed_scaling_factor != 1.5 + or ffn.swiglu_limit != 10.0 + or ffn.scoring_func != "sqrtsoftplus" + or not ffn.renormalize + ): + why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + if why is not None: + logger.info_once("DSv4.1 mono decode off for some layers: %s", why) + return None + # the whole layer: the kernels' attention takes the standard layers only + if layer.engram is not None: + why = "Engram" + elif attn.compressor is not None or attn.indexer is not None: + why = "a compressor / indexer" + elif attn.compress_ratio not in (1, 2): + why = f"compress ratio {attn.compress_ratio}" + elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": + why = f"a {attn.kv_cache_dtype} KV cache" + if why is not None: + logger.info_once( + "DSv4.1 mono decode: the FFN launch for layers with %s", why + ) + return MonoDecodeLayer(ffn_only=True) + logger.info_once("DSv4.1 mono decode: the whole standard backbone layers") + return MonoDecodeLayer() + + def weights(self, layer: "DeepseekV4DecoderLayer"): + if self._weights is not None: + return self._weights + from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights + + a, f = layer.attn, layer.ffn + e, sh = f.experts.routed_experts, f.shared_experts + assert sh is not None # ``create`` took only layers with one + self._weights = MonoLayerWeights( + attn=None if self.ffn_only else self._attn_weights(a), + hc_attn_fn=layer.hc_attn_fn, + hc_attn_scale=layer.hc_attn_scale, + hc_attn_base=layer.hc_attn_base, + attn_norm=layer.attn_norm.weight, + hc_ffn_fn=layer.hc_ffn_fn, + hc_ffn_scale=layer.hc_ffn_scale, + hc_ffn_base=layer.hc_ffn_base, + ffn_norm=layer.ffn_norm.weight, + gate_w=f.gate.weight, + bias=f.gate.e_score_correction_bias, + w13=e.w13_weight, + w13_s=e.w13_weight_scale, + w2=e.w2_weight, + w2_s=e.w2_weight_scale, + sgu=sh.gate_up_proj.weight, + sgu_s=_scale_bytes(sh.gate_up_proj.weight_scale), + sw2=sh.down_proj.weight, + sw2_s=_scale_bytes(sh.down_proj.weight_scale), + ) + self._weights.check(get_tensor_model_parallel_world_size()) + return self._weights + + @staticmethod + def _attn_weights(a): + from vllm.models.deepseek_v41.amd.mono.attention.plan import Dims + from vllm.models.deepseek_v41.amd.mono.runner import AttnWeights + + attn = AttnWeights( + layer_id=a.layer_id, + wqkv=a.fused_wqa_wkv.weight, + wqkv_scale=_scale_bytes(a.fused_wqa_wkv.weight_scale), + q_norm=a.q_norm.weight, + kv_norm=a.kv_norm.weight, + wq_b=a.wq_b.weight, + wq_b_scale=_scale_bytes(a.wq_b.weight_scale), + wo_a=a.wo_a.weight, + wo_a_scale=_scale_bytes(a.wo_a.weight_scale), + wo_b=a.wo_b.weight, + wo_b_scale=_scale_bytes(a.wo_b.weight_scale), + attn_sink=a.attn_sink, + cos_sin=a.rotary_emb.cos_sin_cache, + ratio=a.compress_ratio, + ) + attn.check(Dims(get_tensor_model_parallel_world_size())) + return attn + + @staticmethod + def _skip(why: str) -> None: + logger.debug_once("DSv4.1 mono decode skipped a step: %s", why) + return None + + @staticmethod + def _decode_step(layer: "DeepseekV4DecoderLayer", rows: int): + """(why this step takes the original path -- None for a decode-only step + the kernels take -- the layer's SWA metadata, its compressed cache's).""" + if not 1 <= rows <= MAX_ROWS: + return f"{rows} rows", None, None + if not is_forward_context_available(): + return "no forward context", None, None + fc = get_forward_context() + md = fc.attn_metadata + attn = layer.attn + swa, comp = None, None + if isinstance(md, dict): + swa = cast( + "DeepseekV4ROCMAiterSparseSWAMetadata | None", + md.get(attn.swa_cache_layer.prefix), + ) + if attn.compressed_cache_prefix is not None: + comp = cast( + "DeepseekV4FlashMLAMetadata | None", + md.get(attn.compressed_cache_prefix), + ) + why = None + if fc.cudagraph_runtime_mode not in (CUDAGraphMode.NONE, CUDAGraphMode.FULL): + # a PIECEWISE (breakable) capture is replayed for mixed batches too + why = f"{fc.cudagraph_runtime_mode} step" + elif not isinstance(md, dict): # a profile / dummy run, or DBO ubatches + why = "no attention metadata" + elif swa is None: + why = "layer metadata missing" + elif swa.num_prefills != 0 or swa.num_decodes == 0: + why = "not a decode-only step" + elif swa.num_decode_tokens > rows: + why = "more decode tokens than rows" + return why, swa, comp + + def __call__( + self, + layer: "DeepseekV4DecoderLayer", + x: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + pre_mix: torch.Tensor | None, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the + mono kernels, or None when this step takes the original path (always for + an ``ffn_only`` layer).""" + if self.ffn_only: + return None + if residual is None or x.dim() != 2: + return self._skip("first layer") + why, swa, comp = self._decode_step(layer, x.shape[0]) + attn = layer.attn + tensors = (x, residual, post_mix, res_mix, pre_mix) + if why is None: + # the kernels' attention: causal SWA windows, the compressed cache + if comp is None: + why = "layer metadata missing" + elif swa.decode_swa_width != SWA_WIDTH: + why = f"SWA width {swa.decode_swa_width}" + elif ( + swa.decode_swa_indices is None + or swa.decode_swa_lens is None + or swa.token_to_req_indices is None + ): + why = "SWA decode metadata missing" + elif not all(t is not None and t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + assert swa is not None and comp is not None + runner = _mono_runner(x.device) + logger.info_once( + "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS + ) + out = runner.forward( + self.weights(layer), + x, + residual, + post_mix, + res_mix, + pre_mix, + positions, + swa.slot_mapping, + _record_view(attn.swa_cache_layer.kv_cache, swa.block_size), + swa.decode_swa_indices, + swa.decode_swa_lens, + swa.token_to_req_indices, + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio + ), + comp_block_table=comp.block_table, + ) + return out + + def ffn( + self, + layer: "DeepseekV4DecoderLayer", + part: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + pre_mix: torch.Tensor, + ) -> tuple[torch.Tensor, ...] | None: + """The layer's outputs by the FFN launch from ``part``, wo_b's unreduced + output, and the attention seam's outputs; None when this step takes the + original path (the caller then reduces ``part`` itself).""" + why, _, _ = self._decode_step(layer, part.shape[0]) + tensors = (part, residual, post_mix, res_mix, pre_mix) + if why is None and not all(t.is_contiguous() for t in tensors): + why = "non-contiguous layer inputs" + if why is not None: + return self._skip(why) + runner = _mono_runner(part.device) + logger.info_once( + "DSv4.1 mono decode FFN launch active (decode steps of <= %d rows)", + MAX_ROWS, + ) + return runner.ffn(self.weights(layer), *tensors) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3aa54f9b0e97..3ba6eebcbd17 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -929,7 +929,9 @@ def _o_proj( def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor: if self._wo_b_scale is not None and zf.dim() == 2: - return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) + return self._bpre_attn_gemm( + self.wo_b.weight, self._wo_b_scale, zf, self.wo_b.reduce_results + ) return self.wo_b(zf) def forward_mqa( From 1dfe9fbdae734fc2ee8c0494c543c24d25b563d6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 10 Oct 2026 10:26:40 +0800 Subject: [PATCH 328/952] [Core] Release FlashInfer allreduce workspace with --enable-nccl-comm-suspend (#59360) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- tests/distributed/test_comm_ops.py | 71 +++++++++++++++++++ .../device_communicators/cuda_communicator.py | 7 ++ .../flashinfer_all_reduce.py | 24 +++++-- 3 files changed, 98 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 253ec8a188af..78aee1efd18d 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -325,6 +325,77 @@ def test_cuda_communicator_checkpoints_flashinfer_workspaces( workspace.checkpoint_restore.assert_called_once_with(group) +def _fi_ar_communicator( + monkeypatch: pytest.MonkeyPatch, group: object, normal: Mock, quant: Mock +) -> CudaCommunicator: + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal) + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", quant) + monkeypatch.setattr( + flashinfer_all_reduce, + "_fi_ar_workspace_groups", + {id(normal): group, id(quant): group}, + ) + monkeypatch.setattr( + flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False + ) + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.cpu_group = group + communicator.pynccl_comm = Mock() + communicator.all2all_manager = None + return communicator + + +def _symm_mem_workspace() -> Mock: + workspace = Mock() + workspace.checkpoint_prepare.side_effect = NotImplementedError + workspace.checkpoint_restore.side_effect = NotImplementedError + return workspace + + +@pytest.mark.parametrize("symm_mem", [False, True], ids=["mnnvl", "symm-mem+mnnvl"]) +def test_cuda_communicator_suspend_detaches_flashinfer_workspaces( + monkeypatch: pytest.MonkeyPatch, + symm_mem: bool, +) -> None: + group = object() + workspaces = (_symm_mem_workspace() if symm_mem else Mock(), Mock()) + communicator = _fi_ar_communicator(monkeypatch, group, *workspaces) + communicator.suspend() + communicator.resume() + + for workspace in workspaces: + workspace.checkpoint_prepare.assert_called_once_with() + workspace.checkpoint_restore.assert_called_once_with(group) + communicator.pynccl_comm.suspend.assert_called_once_with() + communicator.pynccl_comm.resume.assert_called_once_with() + + +def test_cuda_communicator_checkpoint_rejects_symm_mem_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + communicator = _fi_ar_communicator( + monkeypatch, object(), _symm_mem_workspace(), Mock() + ) + with pytest.raises(NotImplementedError): + communicator.checkpoint_prepare() + with pytest.raises(NotImplementedError): + communicator.checkpoint_restore() + + +def test_cuda_communicator_resume_retries_failed_restore( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = Mock() + workspace.checkpoint_restore.side_effect = [RuntimeError("cuMemCreate"), None] + communicator = _fi_ar_communicator(monkeypatch, object(), workspace, Mock()) + communicator.suspend() + with pytest.raises(RuntimeError): + communicator.resume() + communicator.resume() + + assert workspace.checkpoint_restore.call_count == 2 + + @pytest.mark.parametrize( ("backend", "capability", "world_size", "nodes", "expected"), [ diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 992d07e15dad..29d5bae89d2c 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -814,12 +814,19 @@ def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool: return bd is not None and bd.uniform def suspend(self) -> None: + from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces + + # FlashInfer AR syncs over the gloo cpu_group, so order vs. NCCL is free. + checkpoint_prepare_fi_ar_workspaces(self.cpu_group, skip_unsupported=True) if self.pynccl_comm is not None: self.pynccl_comm.suspend() def resume(self) -> None: + from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces + if self.pynccl_comm is not None: self.pynccl_comm.resume() + checkpoint_restore_fi_ar_workspaces(self.cpu_group, skip_unsupported=True) def checkpoint_prepare(self) -> None: # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index d51e9e0c0e2c..22bcbb775a65 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -360,14 +360,30 @@ def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]: return group_workspaces -def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None: +def checkpoint_prepare_fi_ar_workspaces( + group: ProcessGroup, skip_unsupported: bool = False +) -> None: for workspace in _fi_ar_workspaces_for_group(group): - workspace.checkpoint_prepare() + try: + workspace.checkpoint_prepare() + except NotImplementedError: + if not skip_unsupported: + raise + logger.info_once( + "FlashInfer all-reduce workspace backed by torch symmetric memory " + "stays resident during suspend." + ) -def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None: +def checkpoint_restore_fi_ar_workspaces( + group: ProcessGroup, skip_unsupported: bool = False +) -> None: for workspace in _fi_ar_workspaces_for_group(group): - workspace.checkpoint_restore(TorchDistBackend(group=group)) + try: + workspace.checkpoint_restore(TorchDistBackend(group=group)) + except NotImplementedError: + if not skip_unsupported: + raise atexit.register(destroy_fi_ar_workspace) From 02aa10539f15c8c73aa7fd4a5ee0f691cbf88dcb Mon Sep 17 00:00:00 2001 From: siliangchen-amd Date: Sat, 10 Oct 2026 10:44:06 +0800 Subject: [PATCH 329/952] [Bugfix][ModelLoader] Make enable_weights_track check quantized layers (#57108) Signed-off-by: siliangchen-amd Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> --- .../model_loader/test_registry.py | 82 +++++++++++++++++++ .../layers/fused_moe/runner/moe_runner.py | 5 +- .../model_loader/default_loader.py | 36 +++++++- 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/model_loader/test_registry.py b/tests/model_executor/model_loader/test_registry.py index a795497a11b6..7e4b6c537a52 100644 --- a/tests/model_executor/model_loader/test_registry.py +++ b/tests/model_executor/model_loader/test_registry.py @@ -2,13 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.linear import UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.model_loader import get_model_loader, register_model_loader from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader @register_model_loader("custom_load_format") @@ -88,3 +93,80 @@ def test_default_loader_hf_still_falls_back_to_pt(tmp_path): ) assert use_safetensors is False assert any(f.endswith("model.pt") for f in files) + + +class _QuantMethod: + uses_meta_device = False + + def process_weights_after_loading(self, layer: nn.Module) -> None: + pass + + +def _layer(quant_method=None, **shapes) -> nn.Module: + layer = nn.Module() + for name, shape in shapes.items(): + param = nn.Parameter(torch.zeros(shape), requires_grad=False) + layer.register_parameter(name, param) + layer.quant_method = quant_method + return layer + + +def _track(model: nn.Module, loaded: set[str], quantized: bool) -> str: + loader = DefaultModelLoader(LoadConfig()) + try: + loader.track_weights_loading(model, set(loaded), quantized=quantized) + except ValueError as e: + return str(e) + return "" + + +def test_weights_track_reports_missing_quantized_weights(): + # Quantized checkpoints may only omit KV-cache quantization params, empty + # placeholders, fbgemm_fp8's input_scale_ub and online-quant params. + online = _QuantMethod() + online.uses_meta_device = True + model = nn.Module() + model.proj = _layer(_QuantMethod(), weight=(4, 4), weight_scale_inv=(1, 1)) + model.unquantized = _layer(UnquantizedLinearMethod(), weight=(4, 4), bias=(4,)) + model.gptq = _layer(_QuantMethod(), qweight=(4, 4), qzeros=(0,)) + model.fbgemm = _layer(_QuantMethod(), weight=(4, 4), input_scale_ub=()) + model.attn = _layer(BaseKVCacheMethod(None), k_scale=(), k_zero_point=()) + model.online = _layer(online, weight=(4, 4)) + loaded = {"proj.weight", "unquantized.weight", "gptq.qweight", "fbgemm.weight"} + error = _track(model, loaded, quantized=True) + assert "proj.weight_scale_inv" in error and "unquantized.bias" in error + assert not any(n in error for n in ("gptq", "fbgemm", "attn", "online")) + + +def test_weights_track_keeps_unquantized_models_unchanged(): + # The check is on by default here, so unquantized layers stay exempt. + model = nn.Module() + model.score = _layer(_QuantMethod(), weight=(2, 4)) + model.norm = _layer(weight=(4,)) + assert _track(model, {"norm.weight"}, quantized=False) == "" + assert "norm.weight" in _track(model, set(), quantized=False) + + +def test_moe_runner_reports_registered_expert_names(): + # AutoWeightsLoader prefixes these names, so they must include routed_experts. + class Experts(nn.Module): + def __init__(self): + super().__init__() + self.w13_weight = nn.Parameter(torch.zeros(2, 8, 4), requires_grad=False) + + def load_weights(self, weights): + for _ in weights: + yield "w13_weight" + + class Runner(nn.Module): + load_weights = MoERunner.load_weights + + def __init__(self): + super().__init__() + self.routed_experts = Experts() + + model = nn.Module() + model.experts = Runner() + weights = [("experts.w13_weight", torch.zeros(2, 8, 4))] + loaded = AutoWeightsLoader(model).load_weights(weights) + assert loaded == {"experts.routed_experts.w13_weight"} diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 9cb7cfbb0f3b..d220f8da36bb 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -301,7 +301,10 @@ def __init__( def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[str]: - return self.routed_experts.load_weights(weights) + # routed_experts reports names relative to itself, but callers such as + # AutoWeightsLoader qualify what this returns with the runner's prefix. + for name in self.routed_experts.load_weights(weights): + yield f"routed_experts.{name}" def _select_forward(self) -> Callable: if current_platform.is_tpu(): diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 9f1c25dd2e34..b162f60b0577 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -16,6 +16,7 @@ from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.ep_weight_filter import ( @@ -47,6 +48,9 @@ logger = init_logger(__name__) +# fbgemm_fp8 builds input_scale_ub from its config; checkpoints do not carry it. +_CONFIG_BUILT_PARAMS = ("input_scale_ub",) + class DefaultModelLoader(BaseModelLoader): """Model loader that can load different file types from disk.""" @@ -506,10 +510,15 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: else default_enable_weights_track ) if enable_weights_track: - self.track_weights_loading(model, loaded_weights) + self.track_weights_loading( + model, loaded_weights, quantized=model_config.quantization is not None + ) def track_weights_loading( - self, model: nn.Module, loaded_weights: set[str] | None + self, + model: nn.Module, + loaded_weights: set[str] | None, + quantized: bool = False, ) -> None: weights_to_load = {name for name, _ in model.named_parameters()} if loaded_weights is not None: @@ -523,7 +532,28 @@ def track_weights_loading( # ignore kv_cache scale and online quant scale, # which can be missing in checkpoints if has_online_quant or has_postprocess_quant: - for param_name, _ in module.named_parameters(): + for param_name, param in module.named_parameters(): + # On a quantized model, every parameter has to come from + # the checkpoint except those a checkpoint cannot carry: + # - online quantization (uses_meta_device) quantizes the + # weights while loading, so its parameters do not map + # to checkpoint tensors; + # - KV-cache quantization parameters (BaseKVCacheMethod) + # fall back to defaults when a checkpoint has none; + # - empty placeholders hold no data, e.g. the qzeros + # moe_wna16 registers for symmetric GPTQ; + # - _CONFIG_BUILT_PARAMS come from the quantization config. + # Unquantized models, where this check runs by default, + # keep the module-wide exemption. + if ( + quantized + and not has_online_quant + and not isinstance(quant_method, BaseKVCacheMethod) + and param.numel() > 0 + and param_name.rsplit(".", 1)[-1] + not in _CONFIG_BUILT_PARAMS + ): + continue full_name = f"{name}.{param_name}" if name else param_name loaded_weights.add(full_name) weights_not_loaded = weights_to_load - loaded_weights From 10cc2f6ae2c9ba7cc5841ece95e27d0562aef1bf Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:00:16 +0900 Subject: [PATCH 330/952] [Bugfix] Bound GDN speculative decode query widths (#55504) Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/worker/test_mamba_hybrid_model_state.py | 11 +++++++++++ vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 9 ++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index fa54ccfcedc5..0dcc4725d779 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -314,6 +314,17 @@ def test_padded_rows_keep_the_sentinel(monkeypatch): assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, 0, -1, -1] +def test_rows_wider_than_the_speculative_state_window_keep_the_sentinel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + state = _mamba_hybrid_state(num_speculative_tokens=7) + input_batch = _input_batch([12, 8, 1], None, [False, False, False]) + + metadata = _prepare_attn_metadata(state, monkeypatch, input_batch) + + assert metadata.num_decode_draft_tokens_cpu.tolist() == [-1, 0, 0] + + def test_chunked_prefill_tail_of_two_or_three_tokens_keeps_the_sentinel( monkeypatch, ): diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index c45304efa512..94c57c1968c3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -288,7 +288,14 @@ def prepare_attn( # Test request state, not num_scheduled_tokens == draft_count+1: # adaptive rewrites num_scheduled_tokens to an even split, so that # equality rarely holds and would demote every verify row to decode. - is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0) + is_decode = ( + ~is_prefilling_np + & (input_batch.num_scheduled_tokens > 0) + & ( + input_batch.num_scheduled_tokens + <= self.vllm_config.num_speculative_tokens + 1 + ) + ) num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( is_decode, num_draft_tokens_per_req, -1 ) From 2d69083c4a09dfeb17d761ae96d2e571a628309c Mon Sep 17 00:00:00 2001 From: xiaofei-zheng Date: Sat, 10 Oct 2026 11:14:01 +0800 Subject: [PATCH 331/952] [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 (#55104) Signed-off-by: xiaofei.zheng Co-authored-by: Cursor Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- tests/kernels/moe/test_ocp_mx_moe.py | 76 +++++++++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 13 +++- 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index f20927565ebd..13f29087d951 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -2255,3 +2255,79 @@ def run(quant_config): "w_mxfp4_a_mxfp6_e3m2 output is bit-identical to weight-only w_mxfp4:" " the emulation never fake-quantized the activations" ) + + +# (backend, activation, hidden_in, inter_in, hidden_out, inter_out) +# +# The AITER_MXFP4_MXFP4 (W4A4) rows are the ones this table exists for: the +# per-partition intermediate is aligned to 128, not 256, so a model whose native +# shard is already 128-aligned is allocated at its real width instead of being +# padded up. `hidden_size` deliberately stays on 256 -- the gfx950 stage-1 mxfp4 +# kernels are generated with tile_k 256 only, and stage-1 reduces over model_dim. +# +# The remaining rows pin the neighbouring branches so a future edit to this +# function cannot silently move a backend between them. +ROCM_MXFP4_ROUNDUP_CASES = [ + # W4A4: 3072 / TP8 = 384 stays 384. This is the case the alignment change is + # for; the generic 256 round-up would allocate 512. + ("AITER_MXFP4_MXFP4", "SILU", 6144, 384, 6144, 384), + # W4A4, already 256-aligned (e.g. moe_intermediate 2048 at TP8): identity in + # both the old and the new form, so such models are provably untouched. + ("AITER_MXFP4_MXFP4", "SILU", 6144, 256, 6144, 256), + # W4A4, 1024 / TP8 = 128: the most aggressive case, halved rather than trimmed. + ("AITER_MXFP4_MXFP4", "SILU", 4096, 128, 4096, 128), + # W4A4 ignores the activation -- the branch does not read it. + ("AITER_MXFP4_MXFP4", None, 6144, 384, 6144, 384), + ("AITER_MXFP4_MXFP4", "SWIGLUOAI", 6144, 384, 6144, 384), + # W4A4 hidden_size is still rounded to 256, unlike the intermediate. + ("AITER_MXFP4_MXFP4", "SILU", 384, 384, 512, 384), + # W4A16 CK + SiTU/SILU keeps its own 128/128 carve-out. + ("AITER_MXFP4_BF16", "SITU", 384, 384, 384, 384), + ("AITER_MXFP4_BF16", "SILU", 384, 384, 384, 384), + # ... but only for those activations; anything else takes the generic ROCm arm. + ("AITER_MXFP4_BF16", "SWIGLUOAI", 384, 384, 512, 512), + # Backends that stay on the generic ROCm 256/256 round-up. + ("AITER_MXFP4_FP8", "SILU", 384, 384, 512, 512), + ("AITER_TRITON_MXFP4_BF16", "SILU", 384, 384, 512, 512), + ("TRITON", "SWIGLUOAI", 384, 384, 512, 512), +] + + +@pytest.mark.parametrize( + "backend_name,activation_name,hidden_in,inter_in,hidden_out,inter_out", + ROCM_MXFP4_ROUNDUP_CASES, +) +@pytest.mark.skipif(not ROCM_AVAILABLE, reason="ROCm is required for this test") +def test_rocm_mxfp4_round_up_hidden_and_intermediate( + backend_name: str, + activation_name: str | None, + hidden_in: int, + inter_in: int, + hidden_out: int, + inter_out: int, +): + """Pin the ROCm arm of mxfp4_round_up_hidden_size_and_intermediate_size. + + The allocation this function returns is what every MXFP4 MoE weight buffer, + scale buffer and aiter tuned-config lookup is sized from, so a change here is + invisible locally and expensive everywhere else. Asserting the table keeps + each backend on the alignment its kernels actually require. + """ + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + mxfp4_round_up_hidden_size_and_intermediate_size, + ) + + backend = Mxfp4MoeBackend[backend_name] + activation = MoEActivation[activation_name] if activation_name else None + + hidden, intermediate = mxfp4_round_up_hidden_size_and_intermediate_size( + backend, hidden_in, inter_in, activation + ) + + assert (hidden, intermediate) == (hidden_out, inter_out), ( + f"{backend_name} / {activation_name}: " + f"hidden {hidden_in}->{hidden} (expected {hidden_out}), " + f"intermediate {inter_in}->{intermediate} (expected {inter_out})" + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index f5fccc03436f..7667dfc6aef0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -875,13 +875,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( triton_uses_128 = ( backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4 ) + aiter_w4a4 = backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 - alignment = ( - 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256 + intermediate_alignment = ( + 128 + if aiter_w4a4 or (is_situ_or_silu and (aiter_uses_128 or triton_uses_128)) + else 256 ) + # W4A4 stage 1 still requires a 256-aligned hidden dimension. + hidden_alignment = 256 if aiter_w4a4 else intermediate_alignment - intermediate_size = round_up(intermediate_size, alignment) - hidden_size = round_up(hidden_size, alignment) + intermediate_size = round_up(intermediate_size, intermediate_alignment) + hidden_size = round_up(hidden_size, hidden_alignment) elif backend == Mxfp4MoeBackend.CPU: # CPU AMX kernel uses BLOCK_N=32, align to 32 intermediate_size = round_up(intermediate_size, 32) From f1a3c60558df01b0c1c61af113663c33bb5a84ed Mon Sep 17 00:00:00 2001 From: Giulio De Pasquale Date: Sat, 10 Oct 2026 04:55:45 +0100 Subject: [PATCH 332/952] [Feature][Attention] Cover a whole spec-verify in one DiffKV split-KV program (#59085) Signed-off-by: Giulio De Pasquale --- .../test_triton_unified_attention_diffkv.py | 95 +++++++++++++++++++ .../ops/triton_unified_attention_diffkv.py | 45 ++++++++- 2 files changed, 136 insertions(+), 4 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py index 3b6cf524e2fc..0addd7ccd332 100644 --- a/tests/kernels/attention/test_triton_unified_attention_diffkv.py +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the Triton DiffKV unified-attention kernel.""" +from unittest.mock import patch + import pytest import torch @@ -16,6 +18,7 @@ get_flash_attn_version, is_flash_attn_varlen_func_available, ) +from vllm.v1.attention.ops import triton_unified_attention_diffkv as diffkv_module from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( unified_attention_diffkv, ) @@ -258,3 +261,95 @@ def test_triton_unified_attn_diffkv_fp8_query(seq_lens, seq_threshold_3D): seq_threshold_3D=seq_threshold_3D, quant_query=True, ) + + +def _capture_diffkv_launch( + query_lens: list[int], max_seqlen_q: int, segm_first_dim: int +): + """Run the launcher with both Triton kernels mocked; return grid/kwargs.""" + num_query_heads, num_kv_heads = 8, 2 + head_size_qk, head_size_v = 128, 128 + num_seqs = len(query_lens) + num_tokens = sum(query_lens) + + query = torch.zeros( + num_tokens, num_query_heads, head_size_qk, dtype=torch.bfloat16, device="cpu" + ) + kv_cache = torch.zeros( + NUM_BLOCKS, + BLOCK_SIZES[0], + num_kv_heads, + head_size_qk + head_size_v, + dtype=torch.bfloat16, + device="cpu", + ) + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk:] + out = torch.zeros( + num_tokens, num_query_heads, head_size_v, dtype=torch.bfloat16, device="cpu" + ) + cu_seqlens_q = torch.tensor( + [0] + query_lens, dtype=torch.int32, device="cpu" + ).cumsum(dim=0, dtype=torch.int32) + seqused_k = torch.full((num_seqs,), 512, dtype=torch.int32, device="cpu") + block_tables = torch.zeros(num_seqs, 32, dtype=torch.int32, device="cpu") + segm_output, segm_max, segm_expsum = _alloc_segm_buffers( + segm_first_dim, num_query_heads, head_size_v + ) + + with ( + patch.object(diffkv_module, "kernel_unified_attention_diffkv") as mock_kernel, + patch.object(diffkv_module, "kernel_reduce_segments_diffkv") as mock_reduce, + ): + unified_attention_diffkv( + q=query, + k=key_cache, + v=value_cache, + out=out, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + softmax_scale=head_size_qk**-0.5, + causal=True, + window_size=(-1, -1), + block_table=block_tables, + softcap=0, + max_seqlen_q=max_seqlen_q, + seq_threshold_3D=8, + num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS, + softmax_segm_output=segm_output, + softmax_segm_max=segm_max, + softmax_segm_expsum=segm_expsum, + ) + + grid = mock_kernel.__getitem__.call_args.args[0] + kwargs = mock_kernel.__getitem__.return_value.call_args.kwargs + return grid, kwargs, mock_reduce.__getitem__.called + + +def test_spec_verify_block_m_grouping_and_launch_kwargs() -> None: + """Pin program granularity for verify/decode without launching Triton.""" + verify_grid, verify_kwargs, verify_reduce = _capture_diffkv_launch( + query_lens=[8, 8], max_seqlen_q=8, segm_first_dim=16 + ) + num_queries_per_kv = 4 + expected_block_m = min(128, next_power_of_2(8 * num_queries_per_kv)) + assert verify_grid == (4, 2, NUM_PAR_SOFTMAX_SEGMENTS) + assert verify_kwargs["IS_3D"] is True + assert verify_kwargs["BLOCK_M"] == expected_block_m + assert verify_kwargs["BLOCK_Q"] == expected_block_m // num_queries_per_kv + assert verify_kwargs["TILE_SIZE"] == 16 + assert verify_kwargs["num_warps"] == 4 + assert "num_stages" not in verify_kwargs + assert verify_reduce + + decode_grid, decode_kwargs, decode_reduce = _capture_diffkv_launch( + query_lens=[1], max_seqlen_q=1, segm_first_dim=8 + ) + assert decode_grid == (1, 2, NUM_PAR_SOFTMAX_SEGMENTS) + assert decode_kwargs["IS_3D"] is True + assert decode_kwargs["BLOCK_M"] == 16 + assert decode_kwargs["BLOCK_Q"] == 4 + assert decode_kwargs["TILE_SIZE"] == 16 + assert "num_warps" not in decode_kwargs + assert "num_stages" not in decode_kwargs + assert decode_reduce diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py index 8480754256b2..27b089be0e13 100644 --- a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -46,6 +46,13 @@ is_batch_invariant = envs.VLLM_BATCH_INVARIANT +# Whole-verify split-KV grouping: one program covers all q tokens of a +# spec-decode verify request (BLOCK_M = q_len * GQA group, capped at 128). +_SPEC_3D_MAX_Q = 16 +_SPEC_3D_BLOCK_M = 128 +_SPEC_3D_NUM_WARPS = 8 +_SPEC_3D_TILE = 16 + @triton.jit def kernel_unified_attention_diffkv( @@ -290,6 +297,9 @@ def kernel_unified_attention_diffkv( acc, mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], ) + # A row whose keys are all masked in this segment ends at M=0, L=0; + # -inf makes reduce_segments treat that segment as absent. + M = tl.where(L > 0.0, M, float("-inf")) store_segm_reduce_scalars( segm_max_ptr, segm_expsum_ptr, @@ -365,6 +375,8 @@ def kernel_reduce_segments_diffkv( ) segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) overall_max = tl.max(segm_max) + # Graph-capture dummy lengths can leave a query with no causal keys. + overall_max = tl.where(overall_max > float("-inf"), overall_max, 0.0) segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) @@ -439,29 +451,53 @@ def unified_attention_diffkv( BLOCK_M = ( 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) ) - BLOCK_Q = BLOCK_M // num_queries_per_kv - - total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + launch_kw: dict[str, int] = {} sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0 # Decide between 2D and 3D launch. Mirrors the standard launcher: # 3D requires preallocated softmax buffers, decode-only batches, and # a small number of sequences (otherwise 2D already saturates the SM). + # Spec-decode verify batches (q_len <= _SPEC_3D_MAX_Q, full attention) + # also take 3D: their 2D grid is only q_blocks x kv_heads programs, each + # walking the whole context. Sliding-window layers stay 2D: their loop + # is already window-bounded. + spec_3d = ( + 1 < max_seqlen_q <= _SPEC_3D_MAX_Q + and sliding_window_val == 0 + and softmax_segm_output is not None + and q.shape[0] <= softmax_segm_output.shape[0] + ) use_3d = not ( seq_threshold_3D is None or num_par_softmax_segments is None or softmax_segm_output is None or softmax_segm_max is None or softmax_segm_expsum is None - or max_seqlen_q > 1 + or (max_seqlen_q > 1 and not spec_3d) or num_seqs > seq_threshold_3D or is_batch_invariant ) + spec_tile: int | None = None + if use_3d and spec_3d and _SPEC_3D_BLOCK_M > BLOCK_M: + spec_bm = min( + _SPEC_3D_BLOCK_M, + triton.next_power_of_2(max_seqlen_q * num_queries_per_kv), + ) + if spec_bm > BLOCK_M and spec_bm % num_queries_per_kv == 0: + BLOCK_M = spec_bm + launch_kw["num_warps"] = _SPEC_3D_NUM_WARPS if BLOCK_M >= 128 else 4 + spec_tile = _SPEC_3D_TILE + BLOCK_Q = BLOCK_M // num_queries_per_kv + + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers # smaller tiles to expose more parallelism along the KV dim. tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32) + if spec_tile is not None: + tile_size = spec_tile grid: tuple[Any, ...] if use_3d: @@ -530,6 +566,7 @@ def unified_attention_diffkv( BLOCK_M=BLOCK_M, NUM_SEGMENTS_PER_SEQ=num_segments, IS_3D=use_3d, + **launch_kw, ) if use_3d: From be0823357f2c1c7833805be480470df8995371ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= Date: Sat, 10 Oct 2026 06:09:13 +0200 Subject: [PATCH 333/952] [TEST][XPU][CI] Fix ci uva kernel test instability 2 (#60404) Signed-off-by: microslaw Co-authored-by: Kunshang Ji --- tests/kernels/core/test_uva.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/kernels/core/test_uva.py b/tests/kernels/core/test_uva.py index 9d88eeced3ba..391a6b763cb1 100644 --- a/tests/kernels/core/test_uva.py +++ b/tests/kernels/core/test_uva.py @@ -21,6 +21,9 @@ @pytest.mark.parametrize("device", DEVICES) def test_cpu_write(device): torch.set_default_device(device) + # Initializes the device; otherwise torch.accelerator.synchronize() is a no-op, + # as the UVA view's ops never trigger lazy init. + torch.accelerator.set_device_index(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) assert gpu_view.device.type == DEVICE_TYPE @@ -44,6 +47,9 @@ def test_cpu_write(device): @pytest.mark.parametrize("device", DEVICES) def test_gpu_write(device): torch.set_default_device(device) + # Initializes the device; otherwise torch.accelerator.synchronize() is a no-op, + # as the UVA view's ops never trigger lazy init. + torch.accelerator.set_device_index(device) cpu_tensor = torch.zeros(10, 10, device="cpu", pin_memory=True, dtype=torch.int32) gpu_view = get_accelerator_view_from_cpu_tensor(cpu_tensor) assert gpu_view.device.type == DEVICE_TYPE From 29e4d2d42e1333da6bded6f28dd5bff03573af7c Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Sat, 10 Oct 2026 00:22:35 -0400 Subject: [PATCH 334/952] [Frontend][Build] Add `vllm download-kernels` for FlashInfer precompiled kernels (#58765) Signed-off-by: mgoin Signed-off-by: Misha Goin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/entrypoints.yaml | 1 + README.md | 1 + docker/Dockerfile | 12 +- docker/versions.json | 3 - docs/cli/README.md | 17 ++- docs/contributing/README.md | 1 + .../installation/gpu.cuda.inc.md | 17 ++- docs/getting_started/quickstart.md | 7 +- docs/mkdocs/gen_files/generate_argparse.py | 7 ++ requirements/cuda.txt | 5 - .../unit_tests/test_download_kernels_cli.py | 113 ++++++++++++++++++ vllm/entrypoints/cli/download_kernels.py | 87 ++++++++++++++ vllm/entrypoints/cli/main.py | 2 + vllm/utils/flashinfer.py | 54 ++++++++- vllm/v1/worker/gpu_worker.py | 2 + 15 files changed, 309 insertions(+), 20 deletions(-) create mode 100644 tests/entrypoints/unit_tests/test_download_kernels_cli.py create mode 100644 vllm/entrypoints/cli/download_kernels.py diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index d31c66a8cbd8..bc7fa701b235 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -31,6 +31,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/entrypoints + - vllm/utils/flashinfer.py - tests/entrypoints/unit_tests - tests/entrypoints/launchers commands: diff --git a/README.md b/README.md index ca31a8cce94d..ff92fface442 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,7 @@ Install vLLM with [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip`: ```bash uv pip install vllm +vllm download-kernels # precompiled kernels for faster startup ``` Or [build from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for development. diff --git a/docker/Dockerfile b/docker/Dockerfile index d524d0fddd47..e29142dd963a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -936,13 +936,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt -# Install FlashInfer JIT cache (requires CUDA-version-specific index URL) -# https://docs.flashinfer.ai/installation.html -# From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.7.0.post1 -RUN --mount=type=cache,target=/opt/uv/cache \ - uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ - --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') +# Install FlashInfer's precompiled kernels for all GPU architectures (the default). +# `vllm download-kernels` wraps this command but vLLM is only installed in later +# stages, and running it there would rebuild this large layer for every wheel. +RUN --mount=type=cache,target=/opt/pip/cache \ + PIP_CACHE_DIR=/opt/pip/cache flashinfer download-kernels # ============================================================ # OPENAI API SERVER DEPENDENCIES diff --git a/docker/versions.json b/docker/versions.json index 68f426d0c7b5..12ec728ec8f3 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -73,9 +73,6 @@ "RUN_WHEEL_CHECK": { "default": "true" }, - "FLASHINFER_VERSION": { - "default": "0.7.0.post1" - }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" }, diff --git a/docs/cli/README.md b/docs/cli/README.md index 27425cd92c78..a1eff4c8ae64 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -9,7 +9,7 @@ vllm --help Available Commands: ```bash -vllm {chat,complete,serve,launch,bench,collect-env,run-batch,preload} +vllm {chat,complete,serve,launch,bench,collect-env,download-kernels,run-batch,preload} ``` ## serve @@ -188,6 +188,21 @@ Start collecting environment information. vllm collect-env ``` +## download-kernels + +Install the precompiled kernels that vLLM's dependencies would otherwise download +or compile at startup. Run it after installing or upgrading vLLM. It currently +installs FlashInfer's kernels and does nothing when FlashInfer is not installed. + +```bash +vllm download-kernels + +# Print the install commands without running them +vllm download-kernels --dry-run +``` + +See [vllm download-kernels](./download-kernels.md) for the full reference of all available arguments. + ## run-batch Run batch prompts and write results to file. diff --git a/docs/contributing/README.md b/docs/contributing/README.md index e7a6448ca63e..5acb8d980940 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -41,6 +41,7 @@ If you are only developing vLLM's Python code, install vLLM using: ```bash VLLM_USE_PRECOMPILED=1 uv pip install -e . +vllm download-kernels ``` To rebuild only the Rust frontend binary: diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index f125a189fcd2..6f18681713ef 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -23,14 +23,28 @@ Therefore, it is recommended to install vLLM with a **fresh new** environment. I ```bash uv pip install vllm --torch-backend=auto +vllm download-kernels ``` ??? console "pip" ```bash # Install vLLM with CUDA 12.9. pip install vllm --extra-index-url https://download.pytorch.org/whl/cu129 + vllm download-kernels ``` +`vllm download-kernels` installs [FlashInfer](https://docs.flashinfer.ai/)'s +precompiled kernels for the installed FlashInfer and CUDA versions. Without them, +vLLM downloads and compiles kernels at startup, which can add several minutes on +Hopper and newer GPUs, **especially Blackwell**, and logs a warning. Run it in the +same Python environment after every vLLM install or upgrade, and in custom +container images. The kernels are tied to the FlashInfer version, and FlashInfer +refuses to import when they do not match, so an upgrade needs them refreshed. +With a PyTorch built for CUDA older than 12.9, only `flashinfer-cubin` is +installed, because FlashInfer does not publish `flashinfer-jit-cache` for it. +For CUDA overrides and nightly kernels, use the +[FlashInfer CLI](https://docs.flashinfer.ai/cli.html#download-kernels) directly. + We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu130`), set `--torch-backend=cu130` (or `UV_TORCH_BACKEND=cu130`). If this doesn't work, try running `uv self update` to update `uv` first. !!! note @@ -94,9 +108,10 @@ If you only need to change Python code, you can build and install vLLM without c git clone https://github.com/vllm-project/vllm.git cd vllm VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto +vllm download-kernels ``` -This command will do the following: +The install command will do the following: 1. Look for the current branch in your vLLM clone. 1. Identify the corresponding base commit in the main branch. diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index afa9713341b2..a0939a11c4fd 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -25,8 +25,11 @@ This guide will help you quickly get started with vLLM to perform: uv venv --python 3.12 --seed source .venv/bin/activate uv pip install vllm --torch-backend=auto + vllm download-kernels ``` + `vllm download-kernels` installs precompiled kernels that vLLM otherwise downloads and compiles at startup, which can add several minutes on Hopper and newer GPUs. See the [installation guide](installation/gpu.md) for details. + `uv` can [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu126`), set `--torch-backend=cu126` (or `UV_TORCH_BACKEND=cu126`). Another delightful way is to use `uv run` with `--with [dependency]` option, which allows you to run commands such as `vllm serve` without creating any permanent environment: @@ -329,5 +332,5 @@ Some of the available backend options include: - On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA`, `ROCM_AITER_UNIFIED_ATTN`, `TRITON_MLA`, `ROCM_AITER_MLA` or `ROCM_AITER_TRITON_MLA`. - On Intel XPU: `FLASH_ATTN`, `TRITON_ATTN`, `TRITON_MLA`, `XPU_MLA_SPARSE`, `TORCH_SDPA` or `TURBOQUANT`. -!!! warning - There are no pre-built vllm wheels containing Flash Infer, so you must install it in your environment first. Refer to the [Flash Infer official docs](https://docs.flashinfer.ai/) or see [docker/Dockerfile](../../docker/Dockerfile) for instructions on how to install it. +!!! tip + FlashInfer is installed with vLLM, but its precompiled kernels are not. Run `vllm download-kernels` after installing or upgrading vLLM so they are not downloaded and compiled at startup. diff --git a/docs/mkdocs/gen_files/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py index a45de12f2e4a..ec9511efadf0 100644 --- a/docs/mkdocs/gen_files/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -178,6 +178,9 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION") RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand") PreloadSubcommand = auto_mock("vllm.entrypoints.cli.preload", "PreloadSubcommand") +DownloadKernelsSubcommand = auto_mock( + "vllm.entrypoints.cli.download_kernels", "DownloadKernelsSubcommand" +) sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS") openai_cli_args = auto_mock("vllm.entrypoints.launchers", "cli_args") run_batch = auto_mock("vllm.entrypoints.launchers", "run_batch") @@ -332,6 +335,10 @@ def linkify_docs_urls(text: str) -> str: "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False), "cli/run-batch.md": (create_parser(run_batch.make_arg_parser), True), "cli/preload.md": (create_parser(PreloadSubcommand.add_cli_args), True), + "cli/download-kernels.md": ( + create_parser(DownloadKernelsSubcommand.add_cli_args), + False, + ), "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True), "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True), # URL kept as `mm_processor` for back-compat; command name is `mm-processor` diff --git a/requirements/cuda.txt b/requirements/cuda.txt index e0db38baadd3..1f5a9642b28a 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -10,12 +10,7 @@ torchaudio==2.11.0 torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version torchcodec >= 0.14 PyNvVideoCodec==2.0.4 -# FlashInfer should be updated together with the Dockerfile -# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from -# install_requires so the published wheel does not carry an unresolvable pin ---extra-index-url https://flashinfer.ai/whl/ flashinfer-python==0.7.0.post1 -flashinfer-cubin==0.7.0.post1 apache-tvm-ffi==0.1.11 tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 diff --git a/tests/entrypoints/unit_tests/test_download_kernels_cli.py b/tests/entrypoints/unit_tests/test_download_kernels_cli.py new file mode 100644 index 000000000000..75af038134b0 --- /dev/null +++ b/tests/entrypoints/unit_tests/test_download_kernels_cli.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata +import os +import subprocess +from importlib.util import find_spec +from pathlib import Path +from types import SimpleNamespace + +import pytest +import regex as re + +import vllm.utils.flashinfer as fi +from vllm.entrypoints.cli import download_kernels + + +@pytest.mark.skipif(find_spec("flashinfer") is None, reason="requires FlashInfer") +def test_download_kernels_resolves_kernel_warning( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + capfd: pytest.CaptureFixture[str], + caplog_vllm: pytest.LogCaptureFixture, + disable_log_dedup, +): + """Runs the installed FlashInfer CLI, so a FlashInfer release that renames + its kernel packages, CLI or version check fails here instead of leaving + users with a warning that its own fix cannot clear.""" + # A vLLM upgrade leaves kernels from the previous FlashInfer behind, which + # FlashInfer rejects on import. The command must work in that state. + stale_cubin = tmp_path / "flashinfer_cubin" + stale_cubin.mkdir() + (stale_cubin / "__init__.py").write_text( + f'__version__ = "0"\ndef get_cubin_dir(): return "{stale_cubin}"\n' + ) + monkeypatch.setenv( + "PYTHONPATH", os.pathsep.join([str(tmp_path), os.getenv("PYTHONPATH", "")]) + ) + + assert download_kernels._download_flashinfer_kernels(dry_run=True) == 0 + kernels = dict(re.findall(r"Requirement: ([\w-]+)==(\S+)", capfd.readouterr().out)) + assert kernels + + installed = {"flashinfer-python": importlib.metadata.version("flashinfer-python")} + + def version(name: str) -> str: + if name not in installed: + raise importlib.metadata.PackageNotFoundError(name) + return installed[name] + + def warns() -> bool: + caplog_vllm.clear() + fi.warn_if_flashinfer_kernels_missing() + return "Run `vllm download-kernels`" in caplog_vllm.text + + monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False) + monkeypatch.setattr(importlib.metadata, "version", version) + monkeypatch.setattr(fi, "has_flashinfer", lambda: True) + monkeypatch.setattr( + fi, + "current_platform", + SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True), + ) + + assert warns() + installed.update(dict.fromkeys(kernels, "0")) + assert warns() + installed.update(kernels) + assert not warns() + + +def test_cuda_without_jit_cache_wheels_needs_only_cubin( + monkeypatch: pytest.MonkeyPatch, + caplog_vllm: pytest.LogCaptureFixture, + disable_log_dedup, +): + """FlashInfer publishes no flashinfer-jit-cache below CUDA 12.9, so the + command must not fail on it and the warning must not ask for it.""" + commands = [] + installed = {"flashinfer-python": "1.0", "flashinfer-cubin": "1.0"} + + def version(name: str) -> str: + if name not in installed: + raise importlib.metadata.PackageNotFoundError(name) + return installed[name] + + def run(cmd: list[str], **kwargs) -> SimpleNamespace: + commands.append(cmd) + return SimpleNamespace(returncode=0) + + monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False) + monkeypatch.setattr(fi.torch.version, "cuda", "12.8") + monkeypatch.setattr(download_kernels, "find_spec", lambda name: object()) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(importlib.metadata, "version", version) + monkeypatch.setattr(fi, "has_flashinfer", lambda: True) + monkeypatch.setattr( + fi, + "current_platform", + SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True), + ) + + assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0 + assert commands[0][-1] == "install-cubin-wheel" + fi.warn_if_flashinfer_kernels_missing() + assert "vllm download-kernels" not in caplog_vllm.text + + +def test_download_kernels_skips_without_flashinfer(monkeypatch: pytest.MonkeyPatch): + """Builds without FlashInfer (CPU, ROCm, XPU) have nothing to download.""" + monkeypatch.setattr(download_kernels, "find_spec", lambda name: None) + monkeypatch.setattr(subprocess, "run", pytest.fail) + + assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0 diff --git a/vllm/entrypoints/cli/download_kernels.py b/vllm/entrypoints/cli/download_kernels.py new file mode 100644 index 000000000000..a5d8d35afb75 --- /dev/null +++ b/vllm/entrypoints/cli/download_kernels.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The `download-kernels` subcommand for the vLLM CLI. + +Installs the precompiled kernels that vLLM's dependencies would otherwise +download or compile at startup. Run it after installing or upgrading vLLM: + + vllm download-kernels +""" + +import argparse +import os +import subprocess +import sys +import typing +from importlib.util import find_spec + +from vllm.entrypoints.cli.types import CLISubcommand +from vllm.logger import init_logger + +if typing.TYPE_CHECKING: + from vllm.utils.argparse_utils import FlexibleArgumentParser +else: + FlexibleArgumentParser = argparse.ArgumentParser + +logger = init_logger(__name__) + + +def _download_flashinfer_kernels(dry_run: bool) -> int: + if find_spec("flashinfer") is None: + logger.info("FlashInfer is not installed; skipping its kernels.") + return 0 + from vllm.utils.flashinfer import has_flashinfer_jit_cache_wheels + + cmd = [sys.executable, "-m", "flashinfer"] + if has_flashinfer_jit_cache_wheels(): + cmd.append("download-kernels") + else: + logger.info( + "FlashInfer does not publish flashinfer-jit-cache for this CUDA " + "version; installing flashinfer-cubin only." + ) + cmd.append("install-cubin-wheel") + if dry_run: + cmd.append("--dry-run") + # FlashInfer refuses to import while kernels from another FlashInfer version + # are installed, which would stop its own CLI from replacing them. + env = {**os.environ, "FLASHINFER_DISABLE_VERSION_CHECK": "1"} + return subprocess.run(cmd, env=env, check=False).returncode + + +class DownloadKernelsSubcommand(CLISubcommand): + """The `download-kernels` subcommand for the vLLM CLI.""" + + name = "download-kernels" + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the install commands without running them.", + ) + return parser + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + sys.exit(_download_flashinfer_kernels(args.dry_run)) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + self.name, + help="Install precompiled kernels for faster startup.", + description=( + "Install the precompiled kernels that vLLM's dependencies would " + "otherwise download or compile at startup. Run it after " + "installing or upgrading vLLM." + ), + usage="vllm download-kernels [--dry-run]", + ) + return DownloadKernelsSubcommand.add_cli_args(parser) + + +def cmd_init() -> list[CLISubcommand]: + return [DownloadKernelsSubcommand()] diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index 4e1ae68b008c..c85bec4f1809 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -42,6 +42,7 @@ def main(): import vllm.entrypoints.cli.benchmark.main import vllm.entrypoints.cli.collect_env + import vllm.entrypoints.cli.download_kernels import vllm.entrypoints.cli.launch import vllm.entrypoints.cli.openai import vllm.entrypoints.cli.preload @@ -60,6 +61,7 @@ def main(): vllm.entrypoints.cli.launch, vllm.entrypoints.cli.benchmark.main, vllm.entrypoints.cli.collect_env, + vllm.entrypoints.cli.download_kernels, vllm.entrypoints.cli.preload, vllm.entrypoints.cli.run_batch, vllm.entrypoints.cli.snapshot, diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 9b615cd2fa4b..69401febea96 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -8,6 +8,7 @@ import contextlib import functools import importlib +import importlib.metadata import importlib.util import os import shutil @@ -16,6 +17,7 @@ import requests import torch +from packaging.version import Version import vllm.envs as envs from vllm.logger import init_logger @@ -98,12 +100,60 @@ def has_flashinfer() -> bool: "FlashInfer kernels are disabled: flashinfer-cubin is not installed " "and nvcc (CUDA_HOME, CUDA_PATH, PATH or /usr/local/cuda) or ninja " "(PATH) is missing. Set CUDA_HOME to a CUDA toolkit and put ninja on " - "PATH, or run `flashinfer download-kernels`." + "PATH, or run `vllm download-kernels`." ) return False return True +def _installed_version(distribution: str) -> str | None: + try: + return Version(importlib.metadata.version(distribution)).public + except importlib.metadata.PackageNotFoundError: + return None + + +def has_flashinfer_jit_cache_wheels() -> bool: + """Return whether FlashInfer publishes flashinfer-jit-cache for the installed + PyTorch's CUDA version, which it does from CUDA 12.9.""" + cuda_version = torch.version.cuda + return cuda_version is not None and Version(cuda_version) >= Version("12.9") + + +def warn_if_flashinfer_kernels_missing() -> None: + """Warn on Hopper and newer GPUs when FlashInfer's precompiled kernels are + missing or were installed for another FlashInfer version.""" + if not ( + current_platform.is_cuda() + and current_platform.has_device_capability(90) + and has_flashinfer() + ): + return + flashinfer_version = _installed_version("flashinfer-python") + packages = [] + if not envs.VLLM_HAS_FLASHINFER_CUBIN: + packages.append("flashinfer-cubin") + if has_flashinfer_jit_cache_wheels(): + packages.append("flashinfer-jit-cache") + installed = {name: _installed_version(name) for name in packages} + if flashinfer_version is None or all( + version == flashinfer_version for version in installed.values() + ): + return + logger.warning_once( + "FlashInfer's precompiled kernels are missing or do not match " + "flashinfer-python %s (%s). FlashInfer then downloads and compiles " + "kernels at startup, which can take several minutes, or fails to import " + "mismatched ones. Run `vllm download-kernels` in this Python environment " + "to install them.", + flashinfer_version, + ", ".join( + f"{name} {version or 'not installed'}" + for name, version in installed.items() + ), + ) + + @functools.cache def has_flashinfer_bf16_gemm() -> bool: """Return whether FlashInfer exposes the BF16 dense GEMM API.""" @@ -1243,4 +1293,6 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "should_use_flashinfer_for_blockscale_fp8_gemm", "is_flashinfer_fp8_blockscale_gemm_supported", "is_flashinfer_cudnn_fp8_prefill_attn_supported", + "has_flashinfer_jit_cache_wheels", + "warn_if_flashinfer_kernels_missing", ] diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index f2dfc8a689f0..a01e5dd5f7a3 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -73,6 +73,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.tracing import instrument +from vllm.utils.flashinfer import warn_if_flashinfer_kernels_missing from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes @@ -497,6 +498,7 @@ def init_device(self): if self.use_v2_model_runner: logger.info_once("Using V2 Model Runner") + warn_if_flashinfer_kernels_missing() # Set random seed. set_random_seed(self.model_config.seed) From 9bd074a99631a87095465590362f63dda5997ba1 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 13:24:00 +0900 Subject: [PATCH 335/952] [Attention] Decode an NVFP4 KV cache with XQA on SM12x (#60452) Signed-off-by: lesj0610 Co-authored-by: Misha Goin --- .buildkite/test_areas/attention.yaml | 14 +++ tests/v1/attention/test_attention_backends.py | 96 ++++++++++++++++++- vllm/v1/attention/backends/flashinfer.py | 21 ++-- 3 files changed, 121 insertions(+), 10 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index b9f430a20ee0..83ba90e2cb1b 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -62,3 +62,17 @@ steps: commands: - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":nvidia: (DGX) Spark FlashInfer NVFP4 XQA" + key: flashinfer-nvfp4-xqa-dgx-spark + timeout_in_minutes: 30 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - vllm/v1/attention/backends/flashinfer.py + - tests/v1/attention/test_attention_backends.py + commands: + - pytest -v -s v1/attention/test_attention_backends.py -k test_flashinfer_xqa_nvfp4_decode_correctness diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 77366c02aebf..add611d92b6a 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -31,6 +31,7 @@ STR_DTYPE_TO_TORCH_DTYPE, is_quantized_kv_cache, is_torch_equal_or_newer, + nvfp4_kv_cache_full_dim, set_random_seed, ) from vllm.v1.attention.backend import ( @@ -41,7 +42,7 @@ max_decode_query_len, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode BACKENDS_TO_TEST = [ AttentionBackendEnum.FLASH_ATTN, @@ -1698,6 +1699,99 @@ def causal_mask_mod( ) +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST + or not current_platform.is_device_capability_family(120), + reason="NVFP4 XQA decode requires SM12x.", +) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.parametrize("q_len", [1, 4]) +def test_flashinfer_xqa_nvfp4_decode_correctness( + default_vllm_config, q_len, use_cuda_graph +): + """Single-token and spec-decode queries over an NVFP4 cache decode with + XQA, eager and replayed, and match the fa2 prefill that wrote the cache.""" + import unittest.mock + + from vllm.v1.attention.backends import flashinfer as fi + + seq_lens, bs, d = [40, 300, 1027], 16, 256 + config = create_vllm_config(max_model_len=max(seq_lens), block_size=bs) + # GQA 6 at head_dim 256, as Qwen3.8-27B at TP=2. + config.model_config.model_arch_config = replace( + config.model_config.model_arch_config, + total_num_attention_heads=24, + total_num_kv_heads=4, + head_size=d, + ) + config.cache_config.cache_dtype = "nvfp4" + if q_len > 1: + config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=q_len - 1 + ) + spec = FullAttentionSpec( + block_size=bs, + num_kv_heads=4, + head_size=d, + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.NVFP4, + ) + cg_support = fi.FlashInferMetadataBuilder.get_cudagraph_support(config, spec) + assert cg_support == AttentionCGSupport.UNIFORM_BATCH + + device = torch.device(f"{DEVICE_TYPE}:0") + # Request i owns consecutive pages from starts[i]; page 0 stays null. + pages = [0] + [cdiv(s, bs) for s in seq_lens[:-1]] + starts = 1 + torch.tensor(pages).cumsum(0) + max_pages = cdiv(max(seq_lens), bs) + block_table = (starts[:, None] + torch.arange(max_pages)).int().to(device) + full_dim = nvfp4_kv_cache_full_dim(d) + kv_cache = torch.zeros( + int(starts[-1]) + max_pages, bs, 8, full_dim, dtype=torch.uint8, device=device + ).transpose(1, 2) + q, k, v = ( + torch.randn(sum(seq_lens), h, d).bfloat16().to(device) for h in (24, 4, 4) + ) + + def run(query_lens, use_cuda_graph=False): + req = torch.repeat_interleave(torch.arange(3), torch.tensor(query_lens)) + pos = torch.cat([torch.arange(s - n, s) for s, n in zip(seq_lens, query_lens)]) + rows = pos + torch.tensor([0, *seq_lens]).cumsum(0)[req] + metadata = replace( + create_common_attn_metadata(BatchSpec(seq_lens, query_lens), bs, device), + block_table_tensor=block_table, + slot_mapping=(starts[req] * bs + pos).to(device), + ) + with set_current_vllm_config(config): + return rows, run_attention_backend( + AttentionBackendEnum.FLASHINFER, + spec, + ["placeholder"], + config, + device, + metadata, + q[rows], + k[rows], + v[rows], + kv_cache, + kv_cache_dtype="nvfp4", + use_cuda_graph=use_cuda_graph, + layer_k_scale=0.02, + layer_v_scale=0.03, + ) + + # The full prefill writes every token through the NVFP4 writer; its fa2 + # output for the last q_len tokens of each request is the reference. + _, reference = run(seq_lens) + xqa_fn = fi.flashinfer_xqa_batch_decode_with_kv_cache + with unittest.mock.patch.object( + fi, "flashinfer_xqa_batch_decode_with_kv_cache", wraps=xqa_fn + ) as xqa: + rows, output = run([q_len] * len(seq_lens), use_cuda_graph) + assert xqa.called + torch.testing.assert_close(output, reference[rows], atol=2e-2, rtol=2e-2) + + if current_platform.is_rocm(): # FLASH_ATTN is not supported on ROCm SLIDING_WINDOW_BACKENDS_TO_TEST = [ diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 85b7456dd889..1a4145844705 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -878,12 +878,10 @@ def __init__( # Prefer TRTLLM/XQA for decoding whenever supported. The decode kernel # must be selected statically for FULL cudagraph capture. - can_use_xqa_or_trtllm_gen_decode = ( - can_use_trtllm_attention( - self.num_qo_heads, self.num_kv_heads, is_prefill=False - ) - and not self.nvfp4_fa2 - ) + # XQA reads an NVFP4 cache only on SM12x; SM90 decodes it with fa2. + can_use_xqa_or_trtllm_gen_decode = can_use_trtllm_attention( + self.num_qo_heads, self.num_kv_heads, is_prefill=False + ) and (not self.nvfp4_fa2 or current_platform.is_device_capability_family(120)) # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by # get_supported_kernel_block_sizes). assert self.page_size <= 64 or ( @@ -1078,7 +1076,10 @@ def get_cudagraph_support( ) or ( is_xqa_arch and ( - spec.kv_quant_mode.is_nvfp4 + ( + spec.kv_quant_mode.is_nvfp4 + and not current_platform.is_device_capability_family(120) + ) or not _is_xqa_head_dim_supported(spec.head_size) ) ): @@ -2283,7 +2284,6 @@ def forward( use_dcp = self.dcp_world_size > 1 if decode_with_xqa: assert not use_dcp - assert not self.is_kvcache_nvfp4 assert self.o_sf_scale is None assert output.dtype != FP4_DTYPE @@ -2628,7 +2628,10 @@ def forward( flashinfer_xqa_batch_decode_with_kv_cache( query=decode_query, - kv_cache=kv_cache_tuple, + kv_cache=( + nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple + ), + kv_cache_sf=nvfp4_kv_block_scales, workspace_buffer=workspace_buffer, block_tables=block_tables_decode, seq_lens=seq_lens_decode, From d5f356dae093d2e47b97af7c9cd8d2fc1aedc56f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Sat, 10 Oct 2026 00:25:44 -0400 Subject: [PATCH 336/952] [HiSparse] Own GPU residency instead of editing core block-table rows (#60647) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- csrc/libtorch_stable/hisparse_kernels.cu | 24 ++-- .../v1/attention/test_sparse_mla_backends.py | 63 ++++++++-- tests/v1/core/test_prefix_caching.py | 94 +++++++++++---- .../unit/test_hisparse_connector.py | 52 ++++++-- tests/v1/worker/test_utils.py | 3 + .../kv_connector/v1/hisparse/connector.py | 15 ++- .../kv_connector/v1/hisparse/worker.py | 19 ++- vllm/v1/attention/backends/mla/index_group.py | 12 +- vllm/v1/core/sched/output.py | 3 - vllm/v1/core/single_type_kv_cache_manager.py | 42 ++++++- vllm/v1/hisparse/binding.py | 28 ++++- vllm/v1/hisparse/coordinator.py | 111 +++++++++++------- vllm/v1/hisparse/runtime.py | 76 ++++++++++-- vllm/v1/hisparse/types.py | 9 ++ vllm/v1/worker/gpu/model_runner.py | 6 +- 15 files changed, 429 insertions(+), 128 deletions(-) diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu index 069eddf3b0ca..5c7621c2f598 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/libtorch_stable/hisparse_kernels.cu @@ -194,7 +194,7 @@ constexpr int kHashLoadInverse = 4; // the page is not resident) and returns its host row (-1 when not host-backed). __device__ __forceinline__ int32_t translate_topk_entry( const int32_t token_index, const int32_t request_row, - const int32_t* __restrict__ source_block_table, + const int32_t state_row, const int32_t* __restrict__ source_block_table, const int32_t* __restrict__ resident_block_table, const int64_t host_rows, const int64_t source_bt_stride, const int32_t source_num_reqs, const int32_t source_num_blocks, const int32_t source_block_size, @@ -221,10 +221,10 @@ __device__ __forceinline__ int32_t translate_topk_entry( if (resident_block_table != nullptr) { const int32_t resident_block = token_index >= 0 ? token_index / resident_block_size : -1; - if (request_row >= 0 && request_row < resident_num_reqs && + if (state_row >= 0 && state_row < resident_num_reqs && resident_block >= 0 && resident_block < resident_num_blocks) { const int32_t physical_block = - resident_block_table[static_cast(request_row) * + resident_block_table[static_cast(state_row) * resident_bt_stride + resident_block]; if (physical_block != resident_null_block && physical_block >= 0) { @@ -424,11 +424,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( for (int i = tid; i < top_k; i += blockDim.x) { int32_t resident_row; const int32_t g = translate_topk_entry( - row_topk[i], request_row, source_block_table, resident_block_table, - host_rows, source_bt_stride, source_num_reqs, source_num_blocks, - source_block_size, resident_bt_stride, resident_num_reqs, - resident_num_blocks, resident_block_size, resident_null_block, - resident_row); + row_topk[i], request_row, state_row, source_block_table, + resident_block_table, host_rows, source_bt_stride, source_num_reqs, + source_num_blocks, source_block_size, resident_bt_stride, + resident_num_reqs, resident_num_blocks, resident_block_size, + resident_null_block, resident_row); if (resolved_global_indices != nullptr) { resolved_global_indices[static_cast(row) * top_k + i] = g; } @@ -602,8 +602,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( g = translate_topk_entry( global_indices[static_cast(row) * input_row_stride + position % top_k], - request_row, source_block_table, resident_block_table, host_rows, - source_bt_stride, source_num_reqs, source_num_blocks, + request_row, state_row, source_block_table, resident_block_table, + host_rows, source_bt_stride, source_num_reqs, source_num_blocks, source_block_size, resident_bt_stride, resident_num_reqs, resident_num_blocks, resident_block_size, resident_null_block, resident_row); @@ -708,8 +708,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( const int32_t g = translate_topk_entry( global_indices[static_cast(row) * input_row_stride + position % top_k], - request_row, source_block_table, resident_block_table, host_rows, - source_bt_stride, source_num_reqs, source_num_blocks, + request_row, state_row, source_block_table, resident_block_table, + host_rows, source_bt_stride, source_num_reqs, source_num_blocks, source_block_size, resident_bt_stride, resident_num_reqs, resident_num_blocks, resident_block_size, resident_null_block, resident_row); diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 160b766da077..6eb089ce3ed6 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3241,7 +3241,8 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources(): [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32 ) - plan.ensure_gpu_sources(resident_table, resident_block_size) + state_indices = torch.arange(2, dtype=torch.int32) + plan.ensure_gpu_sources(resident_table, state_indices, resident_block_size) assert plan.gpu_row_ids is not None unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() @@ -3271,12 +3272,51 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources(): check_dtype=False, ) - plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size) + plan.ensure_gpu_sources( + torch.zeros_like(resident_table), state_indices, resident_block_size + ) assert plan.gpu_row_ids is not None assert (plan.gpu_row_ids == -1).all() +def test_hisparse_prefill_resident_sources_follow_each_resident_group(): + """Layers of different resident groups share one batch staging plan. + + Each group's layers must stage from that group's resident blocks. Keying + the plan's cached resident sources on a per-call batch gather let a later + group's gather reuse the freed address of an earlier one, so its layers + read the earlier group's block ids. + """ + block_size = 4 + plan = build_hisparse_prefill_staging_plan( + torch.tensor([[5, 2]], dtype=torch.int32), + torch.tensor([8], dtype=torch.int32), + block_size, + staging_block_capacity=2, + ) + # The prefill runs in request state row 1; one table holds both groups. + # Rows are wide so that a cache keyed on a freed per-call gather of a row + # would see the next gather reuse its address. + state_rows = torch.zeros((4, 2, 1024), dtype=torch.int32) + state_rows[1, 0, :2] = torch.tensor([11, 12]) + state_rows[1, 1, :2] = torch.tensor([21, 22]) + state_indices = torch.tensor([1], dtype=torch.int32) + host_ids = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist() + + for group in (0, 1, 0, 1): + plan.ensure_gpu_sources(state_rows[:, group], state_indices, block_size) + assert plan.gpu_row_ids is not None + gpu_blocks = plan.gpu_row_ids[0].view(-1, block_size)[:, 0] // block_size + staged = { + host_id: gpu_block + for host_id, gpu_block in zip(host_ids, gpu_blocks.tolist()) + if host_id > 0 + } + first_page, second_page = state_rows[1, group, :2].tolist() + assert staged == {5: first_page, 2: second_page} + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_hisparse_gather_prefill_cache_prefers_resident_rows(): """Staged rows come from the resident cache when a shadow page exists.""" @@ -3294,7 +3334,11 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(): dtype=torch.int32, device=device, ) - plan.ensure_gpu_sources(resident_table, resident_block_size) + plan.ensure_gpu_sources( + resident_table, + torch.arange(2, dtype=torch.int32, device=device), + resident_block_size, + ) num_host_blocks, num_res_blocks = 10, 24 host_cache = ( @@ -3677,6 +3721,7 @@ def test_hisparse_resident_prefill_uses_attention_block_stride(): all_context_pages_resident=True, view=SimpleNamespace(block_size=64, attention_block_stride=832), block_table=torch.tensor([[3]], dtype=torch.int32), + batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32), runtime=SimpleNamespace(max_swap_rows=1), ) index_group = object.__new__(HiSparseMLAIndexGroup) @@ -4091,6 +4136,7 @@ def gather(kv_cache, staging_plan, resident_cache=None): cache = SimpleNamespace( runtime=SimpleNamespace( gather_prefill_cache=gather, + request_state_indices=torch.tensor([0], dtype=torch.int32), ), view=SimpleNamespace(cache=resident_cache, block_size=1), block_table=resident_block_table, @@ -4113,10 +4159,6 @@ def gather(kv_cache, staging_plan, resident_cache=None): assert result is staged assert block_table is plan.block_table torch.testing.assert_close(request_ids, metadata.req_id_per_token) - plan.ensure_gpu_sources.assert_called_once() - args = plan.ensure_gpu_sources.call_args.args - torch.testing.assert_close(args[0], resident_block_table) - assert args[1] == 1 assert calls == [(source, plan, resident_cache)] @@ -4134,6 +4176,9 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): cache = SimpleNamespace( view=SimpleNamespace(cache=resident_cache, block_size=4), block_table=torch.tensor([[2, 3]], dtype=torch.int32), + runtime=SimpleNamespace( + request_state_indices=torch.tensor([0], dtype=torch.int32) + ), ) index_group = object.__new__(HiSparseMLAIndexGroup) index_group.caches = [cache] @@ -4178,10 +4223,6 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): assert args.kwargs["host_cache"].data_ptr() == source.data_ptr() assert args.kwargs["host_row_ids"] is plan.row_ids assert args.kwargs["device_row_ids"] is plan.gpu_row_ids - plan.ensure_gpu_sources.assert_called_once() - ensure_args = plan.ensure_gpu_sources.call_args.args - torch.testing.assert_close(ensure_args[0], cache.block_table) - assert ensure_args[1] == 4 def test_sparse_impl_observes_repointed_indexer_buffer(): diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index a89fd085f1fd..836286d127fa 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -279,6 +279,12 @@ def _allocate_scheduled( return blocks +def _resident_block_ids(manager: KVCacheManager, request_id: str) -> list[int]: + """The GPU block HiSparse reads each resident page from; 0 if host-only.""" + resident = get_hisparse_coordinator(manager).resident_managers[0] + return [block.block_id for block in resident.get_residency_row(request_id)] + + def test_hisparse_does_not_write_back_reprefillable_tokens(): """Multi-module MTP re-prefills the last tokens, so they are not final.""" manager = make_hisparse_kv_cache_manager(32, 16, num_prefill_lookahead=3) @@ -319,7 +325,7 @@ def test_hisparse_builds_dma_row_mirrors_across_pages(): def test_hisparse_async_speculation_mirrors_uncertain_position_range(): """Unresolved drafts must not leave gaps in the eager host mirror.""" coordinator = MagicMock(host_group_id=0) - coordinator.take_block_table_updates.return_value = {} + coordinator.take_residency_updates.return_value = {} coordinator.build_offload_command.return_value = None coordinator.build_row_mirrors.return_value = () scheduler = HiSparseConnectorScheduler( @@ -332,7 +338,6 @@ def test_hisparse_async_speculation_mirrors_uncertain_position_range(): connector.connector_scheduler = scheduler connector.update_state_after_alloc(request, None, 0) scheduler_output = SimpleNamespace( - block_table_updates=None, kv_cache_block_copies=None, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace( @@ -370,11 +375,66 @@ def test_hisparse_reports_when_context_is_fully_resident(): hot_manager.require_hot(request.request_id) assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None _publish_hisparse_pages(manager) + coordinator.take_residency_updates([request.request_id]) + core_row = manager.get_block_ids(request.request_id)[2] pool = manager.block_pool pool.get_new_blocks(pool.get_num_free_blocks()) assert not coordinator.all_context_pages_resident(scheduled) - assert coordinator.take_block_table_updates().keys() == {request.request_id} + # The block-table row stays append-only; the worker learns of the lost + # pages from the residency update instead. + assert manager.get_block_ids(request.request_id)[2] == core_row + resident_ids = _resident_block_ids(manager, request.request_id) + lost = [page for page, block_id in enumerate(resident_ids) if block_id == 0] + assert lost + update = coordinator.take_residency_updates([request.request_id]) + assert update[request.request_id].pages == lost + assert update[request.request_id].block_ids[0] == [0] * len(lost) + + +def test_hisparse_residency_update_sends_only_changed_pages(): + """Losing one early page must not resend the pages after it: per-step + residency updates would otherwise grow with the context length.""" + manager = make_hisparse_kv_cache_manager(32, 16) + tokens = list(range(6 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) + coordinator = get_hisparse_coordinator(manager) + for hot_manager in coordinator.hot_managers: + hot_manager.require_hot(request.request_id) + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) + coordinator.take_residency_updates([request.request_id]) + pool = manager.block_pool + num_never_used = pool.get_num_free_blocks() + # Pages unpin as their host copies complete, so page 0 is reused first. + _publish_hisparse_pages(manager) + pool.get_new_blocks(num_never_used + 1) + + update = coordinator.take_residency_updates([request.request_id]) + + assert update[request.request_id].pages == [0] + assert update[request.request_id].block_ids == ([0],) + + +def test_hisparse_readmitted_request_resends_its_full_residency(): + """A preempted request may resume in a worker state row holding another + request's residency, so its first update must cover every page again.""" + manager = make_hisparse_kv_cache_manager(32, 16) + coordinator = get_hisparse_coordinator(manager) + tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) + request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert _allocate_scheduled(manager, request, len(tokens)) is not None + assert request.request_id in coordinator.take_residency_updates( + [request.request_id] + ) + manager.free(request) + + assert _allocate_scheduled(manager, request, len(tokens)) is not None + update = coordinator.take_residency_updates([request.request_id]) + assert update[request.request_id].pages == [0, 1] + assert update[request.request_id].block_ids[0] == _resident_block_ids( + manager, request.request_id + ) def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): @@ -416,9 +476,7 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks assert len(source) == len(indexer) == len(resident) == 4 assert len(hot) == 2 - assert not any(block.is_null for block in resident[:2]) - assert not resident[2].is_null - assert not resident[3].is_null + assert 0 not in _resident_block_ids(manager, resumed.request_id) coordinator = get_hisparse_coordinator(manager) assert not coordinator.build_offload_command().page_transfers coordinator.finish_host_import(resumed.request_id, failed=False) @@ -666,9 +724,10 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach pool.free_blocks(held) assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None assert all(m.has_hot("first") for m in coordinator.hot_managers) + coordinator.take_residency_updates(["first"]) pool.get_new_blocks(pool.get_num_free_blocks()) assert not coordinator.all_context_pages_resident(scheduled) - assert "first" in coordinator.take_block_table_updates() + assert "first" in coordinator.take_residency_updates(["first"]) def test_hisparse_full_pool_keeps_pages_pinned_until_preemption(): @@ -1204,11 +1263,7 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): num_new_computed_tokens=num_computed, new_computed_blocks=computed, ) - resident_blocks = manager.get_blocks("resumed").blocks[2] - assert [block.block_id for block in resident_blocks[:3]] == ( - original_resident_ids[:3] - ) - assert not any(block.is_null for block in resident_blocks[:3]) + assert _resident_block_ids(manager, "resumed")[:3] == original_resident_ids[:3] assert get_hisparse_coordinator(manager).all_context_pages_resident( ((resumed.request_id, num_computed, len(tokens) - num_computed),) ) @@ -1249,17 +1304,13 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( new_computed_blocks=computed, ) - resumed_blocks = manager.get_blocks("resumed").blocks - assert [block.block_id for block in resumed_blocks[2][:3]] == [ + resident_ids = _resident_block_ids(manager, "resumed") + assert resident_ids[:3] == [ copy_ids[page] if page in adopted_pages else pool.null_block.block_id for page in range(3) ] - gpu_ids = [ - block.block_id - for group in resumed_blocks[1:] - for block in group - if not block.is_null - ] + _, indexer, _, hot = manager.get_block_ids("resumed") + gpu_ids = [block_id for block_id in indexer + resident_ids + hot if block_id] assert len(gpu_ids) == len(set(gpu_ids)) @@ -1488,8 +1539,7 @@ def test_hisparse_reused_copy_is_not_adopted(): num_new_computed_tokens=num_computed, new_computed_blocks=computed, ) - resident_blocks = manager.get_blocks("resumed").blocks[2] - assert all(block.is_null for block in resident_blocks[:3]) + assert _resident_block_ids(manager, "resumed")[:3] == [0, 0, 0] def test_hisparse_external_import_uses_hard_gpu_footprint(): diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index f9598a24f82e..9ae0d3bc8eb1 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -28,10 +28,14 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.hisparse import runtime as runtime_module from vllm.v1.hisparse.coordinator import get_hisparse_coordinator -from vllm.v1.hisparse.runtime import HiSparseCacheHandle +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + update_hisparse_residency, +) from vllm.v1.hisparse.types import ( SparseKVOffloadCommand, SparseKVPageTransfer, + SparseKVResidencyUpdate, SparseKVRowMirror, ) from vllm.v1.metrics.stats import KVCacheEvictionEvent @@ -183,7 +187,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model(): connector = object.__new__(HiSparseConnector) connector.connector_worker = worker connector._get_connector_metadata = MagicMock( - return_value=HiSparseConnectorMetadata(None, (), (), {}, True) + return_value=HiSparseConnectorMetadata(None, (), (), {}, True, {}) ) req_id_per_token = torch.tensor([0, 1], dtype=torch.int32) attn_metadata = SimpleNamespace( @@ -321,6 +325,7 @@ def swap_blocks_batch(src, dst, sizes): ) }, True, + {}, ) ) # A verification step: four query tokens of one request fill the page. @@ -354,7 +359,7 @@ def swap_blocks_batch(src, dst, sizes): # The next step's start mirrors the drafter's rows, then hands the page over. compute_stream.wait_event.reset_mock() connector._get_connector_metadata.return_value = HiSparseConnectorMetadata( - None, (), (), {}, True + None, (), (), {}, True, {} ) connector.start_load_kv( SimpleNamespace(), @@ -397,7 +402,8 @@ def swap_blocks_batch(src, dst, sizes): def test_scheduled_prefix_hit_publishes_adopted_copies(): - """Copies adopted after scheduling must reach the worker's block table.""" + """Copies adopted after scheduling reach the worker as a residency update, + leaving the block-table row the scheduler output carries untouched.""" from tests.v1.core.test_prefix_caching import ( HISPARSE_BLOCK_SIZE, _allocate_scheduled, @@ -443,7 +449,39 @@ def test_scheduled_prefix_hit_publishes_adopted_copies(): num_scheduled_tokens={resumed.request_id: num_new_tokens}, ) - scheduler.build_connector_meta(scheduler_output) + core_row = list(scheduler_output.scheduled_new_reqs[0].block_ids[2]) + + metadata = scheduler.build_connector_meta(scheduler_output) + + update = metadata.residency_updates[resumed.request_id] + assert update.pages == [0, 1, 2, 3] + assert update.block_ids[0][:3] == copy_ids[:3] + assert manager.get_block_ids(resumed.request_id)[2] == core_row + assert core_row[:3] == [0, 0, 0] + + +def test_residency_updates_persist_by_state_row(): + """Updates name requests by batch row but are stored by state row. + + An update for a lost page and an appended page must leave the other pages + and the other request's row intact across steps that reorder the batch. + """ + table = torch.zeros((4, 2, 4), dtype=torch.int32) + update_hisparse_residency( + table, + { + "a": SparseKVResidencyUpdate([0, 1, 2], ([1, 2, 3], [11, 12, 13])), + "b": SparseKVResidencyUpdate([0, 1], ([4, 5], [14, 15])), + }, + ["a", "b"], + torch.tensor([2, 0], dtype=torch.int32), + ) + update_hisparse_residency( + table, + {"a": SparseKVResidencyUpdate([1, 3], ([0, 6], [0, 16]))}, + ["b", "a"], + torch.tensor([0, 2], dtype=torch.int32), + ) - resident_ids = scheduler_output.block_table_updates[resumed.request_id][2] - assert resident_ids[:3] == copy_ids[:3] + assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]] + assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 98395e90c4dd..84dd944e1d22 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1068,6 +1068,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) @@ -1079,6 +1080,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) @@ -1152,6 +1154,7 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): row_mirrors={}, all_context_pages_resident=True, row_mirrors_from_resident=False, + residency_updates={}, ), None, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 863f52e02514..4f17f1e737e0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -36,7 +36,11 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.hisparse.coordinator import get_hisparse_coordinator from vllm.v1.hisparse.layout import get_hisparse_steady_state_concurrency -from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror +from vllm.v1.hisparse.types import ( + SparseKVOffloadCommand, + SparseKVResidencyUpdate, + SparseKVRowMirror, +) from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: @@ -56,6 +60,7 @@ class HiSparseConnectorMetadata(KVConnectorMetadata): source_block_ids: tuple[int, ...] row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]] all_context_pages_resident: bool + residency_updates: dict[str, SparseKVResidencyUpdate] @dataclass @@ -143,7 +148,7 @@ def build_connector_meta( scheduler_output.num_scheduled_tokens.items() ) ) - # Runs first: it plans this step's transfers and table updates. + # Runs first: it plans this step's transfers and residency changes. self.coordinator.advance_scheduled( ( request_id, @@ -151,9 +156,6 @@ def build_connector_meta( ) for request_id, start, count in scheduled_requests ) - scheduler_output.block_table_updates = ( - self.coordinator.take_block_table_updates() or None - ) command = self.coordinator.build_offload_command() host_block_copies = self.coordinator.take_host_block_copies() source_group_id = self.coordinator.host_group_id @@ -192,6 +194,9 @@ def build_connector_meta( tuple(source_block_ids), row_mirrors, self.coordinator.all_context_pages_resident(scheduled_requests), + self.coordinator.take_residency_updates( + scheduler_output.num_scheduled_tokens + ), ) def update_connector_output(self, connector_output: KVConnectorOutput) -> None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index b9dd01f05420..bbccecd65784 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -26,7 +26,11 @@ from vllm.utils.torch_utils import current_stream from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX -from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state +from vllm.v1.hisparse.runtime import ( + HiSparseCacheHandle, + release_pinned_state, + update_hisparse_residency, +) from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror from vllm.v1.kv_cache_interface import ( HiSparseHotSpec, @@ -334,6 +338,10 @@ def initialize( "HiSparse request-state mapping does not match max_num_seqs." ) self.hot_backing = hot_backing + self.residency: torch.Tensor | None = next( + (cache.residency for cache in cache_handles if cache.residency is not None), + None, + ) self._pending_invalid_block_ids: list[int] = [] # Destination block ids of host copies this worker has run. self._completed_host_copy_dst_ids: list[int] = [] @@ -437,6 +445,15 @@ def start_step( self._pending_invalid_block_ids.extend(metadata.source_block_ids) if request_state_indices is not None: self.set_request_state_indices(request_state_indices) + if metadata.residency_updates: + assert self.residency is not None + assert request_state_indices is not None and request_ids is not None + update_hisparse_residency( + self.residency, + metadata.residency_updates, + request_ids, + request_state_indices, + ) def _clear_forward_mirror_state(self) -> None: self._per_layer_mirrored.clear() diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 39f205b91d97..0dfda0ded9f0 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -243,7 +243,7 @@ def convert_logical_to_physical_topk( layer_index, logical_topk_indices, req_id_per_token, - leader.block_table, + leader.batch_block_table(), leader.view.block_size, block_stride_rows=leader.view.attention_block_stride, return_valid_counts=return_valid_counts, @@ -268,8 +268,11 @@ def stage_prefill_rows( assert staging_plan is not None resident_cache = None if cache.view is not None and cache.block_table is not None: + state_indices = cache.runtime.request_state_indices + assert state_indices is not None staging_plan.ensure_gpu_sources( - cache.block_table[attn_metadata.num_decodes :], + cache.block_table, + state_indices[attn_metadata.num_decodes :], cache.view.block_size, ) resident_cache = cache.view.cache @@ -291,8 +294,11 @@ def _prefill_gather_plan( plan = prefill.host_staging_plan if prefill is not None else None assert plan is not None assert cache.view is not None and cache.block_table is not None + state_indices = cache.runtime.request_state_indices + assert state_indices is not None plan.ensure_gpu_sources( - cache.block_table[attn_metadata.num_decodes :], + cache.block_table, + state_indices[attn_metadata.num_decodes :], cache.view.block_size, ) assert plan.gpu_row_ids is not None diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 9d16d2f710a3..a5af5a3d5755 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -304,9 +304,6 @@ class SchedulerOutput: # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None - # Complete block-table rows that replace incrementally appended block IDs. - block_table_updates: dict[str, tuple[list[int], ...]] | None = None - # Scheduler-local; always None by the time this reaches a worker. kv_connector_block_state: KVConnectorBlockState | None = None diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index eed15427b2dd..77c6a4a5f750 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2503,6 +2503,10 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): max_admission_blocks_per_request: int + def __init__(self, kv_cache_spec: HiSparseResidentSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + self.residency_changes: dict[str, dict[int, KVCacheBlock]] = {} + def get_num_blocks_to_allocate( self, request_id: str, @@ -2546,7 +2550,7 @@ def allocate_external_computed_blocks( tail_page = (num_tokens - 1) // self.block_size if len(blocks) <= tail_page: blocks.extend([self._null_block] * (tail_page + 1 - len(blocks))) - if blocks[tail_page].is_null: + if self.get_resident_page(request_id, tail_page) is None: blocks[tail_page] = self.block_pool.get_new_blocks(1)[0] def add_local_computed_blocks( @@ -2580,25 +2584,55 @@ def allocate_new_blocks( def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: assert self.coordinator is not None self.coordinator.free(request_id) - return super().pop_blocks_for_free(request_id) + blocks = self.get_residency_row(request_id) + self.residency_changes.pop(request_id, None) + super().pop_blocks_for_free(request_id) + return blocks def adopt_resident_page( self, request_id: str, block_idx: int, block: KVCacheBlock ) -> bool: """Point a null prefix page at a pinned GPU copy of its contents.""" blocks = self.req_to_blocks.get(request_id) - if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null: + if blocks is None or block_idx >= len(blocks): + return False + if self.get_resident_page(request_id, block_idx) is not None: return False - blocks[block_idx] = block + self.residency_changes.setdefault(request_id, {})[block_idx] = block return True + def drop_resident_page( + self, request_id: str, block_idx: int + ) -> KVCacheBlock | None: + """Stop reading a page from the GPU; return the block it was read from.""" + block = self.get_resident_page(request_id, block_idx) + if block is not None: + changes = self.residency_changes.setdefault(request_id, {}) + changes[block_idx] = self._null_block + return block + def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None: blocks = self.req_to_blocks.get(request_id) if blocks is None or block_idx >= len(blocks): return None block = blocks[block_idx] + changes = self.residency_changes.get(request_id) + if changes: + block = changes.get(block_idx, block) return None if block.is_null else block + def get_residency_row(self, request_id: str, start: int = 0) -> list[KVCacheBlock]: + """The request's row from ``start`` with residency changes applied. Unlike + ``get_resident_page``, host-only pages stay as the null block.""" + blocks = self.req_to_blocks.get(request_id, [])[start:] + changes = self.residency_changes.get(request_id) + if not changes: + return blocks + return [ + changes.get(block_idx, block) + for block_idx, block in enumerate(blocks, start) + ] + def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 60c60a3ff504..1420a45396cd 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -214,10 +214,26 @@ def bind_hisparse_kv_caches( for tensor_config in kv_cache_config.kv_cache_tensors for layer_index, name in enumerate(tensor_config.layers) } - resident_source_index = 0 - for group_id, group in enumerate(kv_cache_config.kv_cache_groups): - if not isinstance(group.kv_cache_spec, HiSparseResidentSpec): - continue + resident_groups = [ + (group_id, group) + for group_id, group in enumerate(kv_cache_config.kv_cache_groups) + if isinstance(group.kv_cache_spec, HiSparseResidentSpec) + ] + residency: torch.Tensor | None = None + if resident_groups: + table_shapes = { + block_tables.input_block_tables[group_id].shape + for group_id, _ in resident_groups + } + assert len(table_shapes) == 1 + max_num_reqs, max_num_pages = table_shapes.pop() + residency = torch.zeros( + (max_num_reqs, len(resident_groups), max_num_pages), + dtype=torch.int32, + device=block_tables.input_block_tables[resident_groups[0][0]].device, + ) + for resident_source_index, (group_id, group) in enumerate(resident_groups): + assert residency is not None for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] @@ -230,13 +246,13 @@ def bind_hisparse_kv_caches( block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, - block_table=block_tables.input_block_tables[group_id], + block_table=residency[:, resident_source_index], slot_mapping=block_tables.slot_mappings[group_id], ) + cache_handle.residency = residency assert cache_handle.view is not None kv_caches[cache_name] = cache_handle.view.cache cache_handle.runtime.resident_source_index = resident_source_index - resident_source_index += 1 hot_backing: torch.Tensor | None = None cache_handles: list[HiSparseCacheHandle] = [] diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 390e0d27ced2..374f100bc889 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -22,6 +22,7 @@ ACTIVE_TAIL_PAGES, SparseKVOffloadCommand, SparseKVPageTransfer, + SparseKVResidencyUpdate, SparseKVRowMirror, ) from vllm.v1.kv_cache_interface import ( @@ -65,6 +66,10 @@ class _HiSparseRequestState: unpinned_pages: set[int] = field(default_factory=set) # Prefix pages whose GPU copies are adopted after the admitting allocation. pages_to_adopt: int = 0 + # The worker holds this request's residency for every page below + # synced_pages, except the stale pages changed since it was sent. + synced_pages: int = 0 + stale_pages: set[int] = field(default_factory=set) @dataclass @@ -86,9 +91,10 @@ class HiSparseCoordinator: GPU-resident pages are a write-back cache of the host tier. Once a page's host copy is durable and its owner can read from host (it has, or has asked for, a hot region), the page's allocation reference is released with - ``BlockPool.unpin_blocks``: the block stays in the block table and keeps - being read, but the pool counts it as free and may hand it out, at which - point the owner's page is nulled and its block table republished. + ``BlockPool.unpin_blocks``: the block keeps being read, but the pool counts + it as free and may hand it out, at which point the owner stops reading it. + Residency changes reach the worker through ``take_residency_updates``, not + the block table, which stays append-only. """ def __init__( @@ -167,7 +173,6 @@ def __init__( self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {} # request -> prefix length an external load is still filling in. self._pending_imports: dict[str, int] = {} - self.block_table_updates: set[str] = set() # Requests whose resident pages were all non-null when last scanned. # Pages only become null at admission or import (before the first # scan) and when the pool reuses an unpinned page (_lose_page). @@ -269,7 +274,8 @@ def _adopt_copies( if manager.adopt_resident_page(request_id, host_idx, block): manager.block_pool.touch([block]) state.pinned_clean.add(host_idx) - self.block_table_updates.add(request_id) + if host_idx < state.synced_pages: + state.stale_pages.add(host_idx) def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: """Index the GPU copies of just-published host blocks for later hits.""" @@ -342,14 +348,11 @@ def _resident_page_blocks( ) -> list[KVCacheBlock] | None: blocks: list[KVCacheBlock] = [] for manager in self.resident_managers: - req_blocks = manager.req_to_blocks.get(request_id) - if ( - req_blocks is None - or page_idx >= len(req_blocks) - ACTIVE_TAIL_PAGES - or req_blocks[page_idx].is_null - ): + num_pages = len(manager.req_to_blocks.get(request_id, ())) + block = manager.get_resident_page(request_id, page_idx) + if page_idx >= num_pages - ACTIVE_TAIL_PAGES or block is None: return None - blocks.append(req_blocks[page_idx]) + blocks.append(block) return blocks def _unpin_page( @@ -392,16 +395,15 @@ def _lose_page(self, request_id: str, page_idx: int) -> None: if state is None: return for manager in self.resident_managers: - blocks = manager.req_to_blocks.get(request_id) - if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null: + block = manager.drop_resident_page(request_id, page_idx) + if block is None: continue - block = blocks[page_idx] - blocks[page_idx] = manager._null_block owners = self._owners.get(block.block_id) if owners is not None: owners.discard((request_id, page_idx)) state.unpinned_pages.discard(page_idx) - self.block_table_updates.add(request_id) + if page_idx < state.synced_pages: + state.stale_pages.add(page_idx) for hot_manager in self.hot_managers: hot_manager.require_hot(request_id) @@ -627,11 +629,10 @@ def build_row_mirrors( mirrors: list[SparseKVRowMirror] = [] for request_id, num_computed_tokens, num_scheduled_tokens in requests: host_blocks = self.host_manager.req_to_blocks.get(request_id) - resident_blocks = [ - manager.req_to_blocks.get(request_id) + if host_blocks is None or any( + request_id not in manager.req_to_blocks for manager in self.resident_managers - ] - if host_blocks is None or any(blocks is None for blocks in resident_blocks): + ): continue token_position = num_computed_tokens end_position = token_position + num_scheduled_tokens @@ -644,14 +645,12 @@ def build_row_mirrors( # window only ever moves forward, so dropping them here would # leave their host rows permanently stale. source_starts = [] - for blocks in resident_blocks: - assert blocks is not None - if page_idx >= len(blocks) or blocks[page_idx].is_null: + for manager in self.resident_managers: + block = manager.get_resident_page(request_id, page_idx) + if block is None: break - source_starts.append( - blocks[page_idx].block_id * block_size + row_offset - ) - if len(source_starts) != len(resident_blocks): + source_starts.append(block.block_id * block_size + row_offset) + if len(source_starts) != len(self.resident_managers): token_position += num_rows continue host_block_idx = page_idx @@ -693,21 +692,53 @@ def all_context_pages_resident( if any( block.is_null for manager in self.resident_managers - for block in manager.req_to_blocks.get(request_id, ()) + for block in manager.get_residency_row(request_id) ): return False self._resident_requests.add(request_id) return True - def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]: - updates = { - request_id: tuple( - [block.block_id for block in manager.req_to_blocks.get(request_id, [])] - for manager in self.managers + def take_residency_updates( + self, request_ids: Iterable[str] + ) -> dict[str, SparseKVResidencyUpdate]: + """Resident block ids the worker lacks for the scheduled requests: + stale pages plus pages appended since the last update. + + Requests that are not scheduled keep their changes until they are: + only a scheduled request's residency is read. + """ + updates: dict[str, SparseKVResidencyUpdate] = {} + if not self.resident_managers: + return updates + first = self.resident_managers[0] + for request_id in request_ids: + state = self._get_request_state(request_id) + num_pages = len(first.req_to_blocks.get(request_id, ())) + start = state.synced_pages + if not state.stale_pages and start >= num_pages: + continue + stale_pages = sorted(state.stale_pages) + block_ids: list[list[int]] = [] + for manager in self.resident_managers: + null_block_id = manager.block_pool.null_block.block_id + group_block_ids = [ + null_block_id if block is None else block.block_id + for block in ( + manager.get_resident_page(request_id, page_idx) + for page_idx in stale_pages + ) + ] + group_block_ids.extend( + block.block_id + for block in manager.get_residency_row(request_id, start) + ) + block_ids.append(group_block_ids) + updates[request_id] = SparseKVResidencyUpdate( + pages=[*stale_pages, *range(start, num_pages)], + block_ids=tuple(block_ids), ) - for request_id in self.block_table_updates - } - self.block_table_updates.clear() + state.stale_pages.clear() + state.synced_pages = num_pages return updates def build_offload_command(self) -> SparseKVOffloadCommand | None: @@ -846,9 +877,7 @@ def free(self, request_id: str) -> None: # planned, so its copies do not prove that the KV was computed. state.publication = None for manager in self.resident_managers: - blocks = manager.req_to_blocks.get(request_id) - if not blocks: - continue + blocks = manager.get_residency_row(request_id) # Tail first, so a prefix loses its tail pages before its head. for page_idx in reversed(range(len(blocks))): block = blocks[page_idx] @@ -865,7 +894,7 @@ def free(self, request_id: str) -> None: manager.block_pool.unpin_blocks([block], self._on_block_reused) else: continue - blocks[page_idx] = manager._null_block + manager.drop_resident_page(request_id, page_idx) def get_hisparse_coordinator( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 1f7a5ad5ca3c..f1b0fa688eda 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -7,10 +7,11 @@ import math import mmap import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeAlias +import numpy as np import psutil import torch @@ -21,8 +22,9 @@ from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import get_max_shared_memory_bytes -from vllm.utils.torch_utils import current_stream +from vllm.utils.torch_utils import async_tensor_h2d, current_stream from vllm.v1.attention.backend import max_decode_query_len +from vllm.v1.hisparse.types import SparseKVResidencyUpdate from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor @@ -499,16 +501,19 @@ class HiSparsePrefillStagingPlan: def ensure_gpu_sources( self, - resident_block_table: torch.Tensor, + resident_state_rows: torch.Tensor, + state_indices: torch.Tensor, resident_block_size: int, ) -> None: """Resolve which staged rows can be served from the resident cache. - Computed once per plan (the resident block table is shared by every - layer in the group); non-null resident pages become miss_mask=0 rows - gathered device-to-device by ``gather_prefill_cache``. + ``resident_state_rows`` is a resident group's persistent table by + request state row and ``state_indices`` holds each staged request's + state row. Computed once per plan and resident group (every layer in + the group shares the table); non-null resident pages become + miss_mask=0 rows gathered device-to-device by ``gather_prefill_cache``. """ - source_key = (resident_block_table.data_ptr(), resident_block_size) + source_key = (resident_state_rows.data_ptr(), resident_block_size) if self.gpu_source_key == source_key: return block_size = self.block_size @@ -519,8 +524,11 @@ def ensure_gpu_sources( host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size new_bt = self.block_table.to(torch.int64) num_rows, num_cols = new_bt.shape - if num_rows == 0 or resident_block_table.shape[0] < num_rows: + if num_rows == 0 or state_indices.shape[0] < num_rows: return + resident_block_table = resident_state_rows.index_select( + 0, state_indices[:num_rows].clamp(min=0) + ) # One representative (row, col) per unique host block: any request # referencing the block holds an equivalent (refcounted) resident view. flat_pos = torch.arange(num_rows * num_cols, device=device) @@ -1107,11 +1115,50 @@ def swap_in( return physical_topk_indices +def update_hisparse_residency( + residency: torch.Tensor, + updates: Mapping[str, SparseKVResidencyUpdate], + request_ids: Sequence[str], + request_state_indices: torch.Tensor, +) -> None: + """Apply scheduled requests' residency changes. + + ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of + each page, or block 0 when the page is read from the host. + """ + if not updates: + return + batch_rows = {request_id: row for row, request_id in enumerate(request_ids)} + num_groups = residency.shape[1] + rows: list[int] = [] + pages: list[int] = [] + block_ids: list[list[int]] = [[] for _ in range(num_groups)] + for request_id, update in updates.items(): + rows.append(batch_rows[request_id]) + pages.extend(update.pages) + for group_block_ids, update_block_ids in zip( + block_ids, update.block_ids, strict=True + ): + group_block_ids.extend(update_block_ids) + host_values = np.empty((len(pages), 2 + num_groups), dtype=np.int32) + host_values[:, 0] = np.repeat( + rows, [len(update.pages) for update in updates.values()] + ) + host_values[:, 1] = pages + host_values[:, 2:] = np.asarray(block_ids, dtype=np.int32).T + values = async_tensor_h2d( + host_values, device=request_state_indices.device, dtype=torch.int32 + ) + state_rows = request_state_indices[values[:, 0]] + residency[state_rows, :, values[:, 1]] = values[:, 2:] + + class HiSparseCacheHandle: """Attention-facing handle for resident KV and sparse offload state.""" def __init__(self, runtime: HiSparseRuntime) -> None: self.view: PagedCacheView | None = None + self.residency: torch.Tensor | None = None self.block_table: torch.Tensor | None = None self.source_block_table: torch.Tensor | None = None self.slot_mapping: torch.Tensor | None = None @@ -1128,6 +1175,7 @@ def __init__(self, runtime: HiSparseRuntime) -> None: # Speculator layers write their rows after the target forward. self.draft_layer = False self.index_group_caches: list[HiSparseCacheHandle] = [self] + self._batch_block_table: torch.Tensor | None = None def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: assert self.runtime.is_group_leader @@ -1136,6 +1184,7 @@ def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: def _prepare_for_batch(self, attn_metadata: Any | None) -> None: self.dummy_batch = attn_metadata is None + self._batch_block_table = None self.runtime.begin_forward() self.num_actual_tokens = ( attn_metadata.num_actual_tokens if attn_metadata is not None else 0 @@ -1151,6 +1200,17 @@ def _prepare_for_batch(self, attn_metadata: Any | None) -> None: not self.decode_batch or self.runtime.eager_host_mirror ) + def batch_block_table(self) -> torch.Tensor: + """Resident rows by this step's batch row, for the prefill paths.""" + if self._batch_block_table is None: + assert self.block_table is not None + indices = self.runtime.request_state_indices + assert indices is not None + self._batch_block_table = self.block_table.index_select( + 0, indices.clamp(min=0) + ) + return self._batch_block_table + def write_target( self, num_input_rows: int, num_slot_rows: int ) -> tuple[torch.Tensor, torch.Tensor, int]: diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py index db0ee06e9399..37b25003df25 100644 --- a/vllm/v1/hisparse/types.py +++ b/vllm/v1/hisparse/types.py @@ -28,6 +28,15 @@ class SparseKVRowMirror: num_rows: int +@dataclass(frozen=True) +class SparseKVResidencyUpdate: + """GPU block ids of some of a request's resident pages, per resident group. + A null block id means the page is read from the host.""" + + pages: list[int] + block_ids: tuple[list[int], ...] + + @dataclass class SparseKVOffloadCommand: """Opaque scheduler-to-worker command for one model step.""" diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f39fd0a6e6a7..674cf23527d1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1224,17 +1224,13 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: def update_requests(self, scheduler_output: SchedulerOutput) -> None: # Add new blocks and update num_computed_tokens for the existing requests. reqs = scheduler_output.scheduled_cached_reqs - table_updates = scheduler_output.block_table_updates or {} - for req_id, block_ids in table_updates.items(): - if (idx := self.req_states.req_id_to_index.get(req_id)) is not None: - self.block_tables.append_block_ids(idx, block_ids, overwrite=True) num_computed_tokens_np = self.req_states.num_computed_tokens_np for req_id, num_computed_tokens, req_new_block_ids in zip( reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids ): req_index = self.req_states.req_id_to_index[req_id] num_computed_tokens_np[req_index] = num_computed_tokens - if req_new_block_ids is not None and req_id not in table_updates: + if req_new_block_ids is not None: self.block_tables.append_block_ids( req_index, req_new_block_ids, overwrite=False ) From acec1b1d84c428ced19d53935979107eeab922a2 Mon Sep 17 00:00:00 2001 From: Guancheng Fu <110874468+gc-fu@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:27:35 +0800 Subject: [PATCH 337/952] [Bugfix][XPU] Remove stale W4A8 weight after repacking (#57511) Signed-off-by: Guancheng Fu Co-authored-by: Kunshang Ji --- vllm/model_executor/kernels/linear/mixed_precision/xpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py index c6545f312a91..c82e862e5870 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py @@ -186,7 +186,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Free the original unpacked int8 weight (still registered as "weight") # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory. - layer.register_parameter("weight", None) + delattr(layer, "weight") def apply_weights( self, From 5ca26880ad87fe5f55628cd0d09d8b8bfc0ba59b Mon Sep 17 00:00:00 2001 From: Puneet Sharma Date: Sat, 10 Oct 2026 10:09:58 +0530 Subject: [PATCH 338/952] Optimize ppc64le vLLM runtime image size (#48074) Signed-off-by: puneetsharma21 --- docker/Dockerfile.ppc64le | 29 +-- docker/build_vllm_ppc64le.sh | 434 +++++++++++++++++++++++++++-------- 2 files changed, 346 insertions(+), 117 deletions(-) diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 68e6ab956377..00d24abbbb5c 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -1,5 +1,5 @@ # Base UBI image -ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 +ARG BASE_UBI_IMAGE_TAG=9.8 ############################################################### # BUILDER STAGE # @@ -7,7 +7,7 @@ ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS builder-base -ARG VLLM_VERSION="0.22.1" +ARG VLLM_VERSION="0.0.0.dev" ARG PYTHON_VERSION=3.12 ARG VLLM_TARGET_DEVICE=cpu @@ -45,10 +45,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # copy vllm source code to build cache COPY . . -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - pip install -U uv - # build & install vLLM so that all transitive dependencies are build/downloaded into the uv cache RUN --mount=type=cache,target=/root/.cache/uv \ source /opt/rh/gcc-toolset-14/enable && \ @@ -63,19 +59,24 @@ RUN --mount=type=cache,target=/root/.cache/uv \ export DNNL_ARCH_OPT_FLAGS="-mcpu=power10 -mtune=power10" && \ export C_INCLUDE_PATH=/usr/local/include:$C_INCLUDE_PATH && \ export CPLUS_INCLUDE_PATH=/usr/local/include:$CPLUS_INCLUDE_PATH && \ - uv pip install 'setuptools>=78.1.1' && \ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/OpenBLAS/lib/:/usr/local/lib64:/usr/local/lib && \ export LIBGOMP=/opt/rh/gcc-toolset-14/root/usr/lib/gcc/ppc64le-redhat-linux/14/libgomp.so && \ - export CMAKE_LIBRARY_PATH=$(dirname $LIBGOMP):${CMAKE_LIBRARY_PATH} && \ export LIBRARY_PATH=$(dirname $LIBGOMP):${LIBRARY_PATH} && \ export LD_LIBRARY_PATH=$(dirname $LIBGOMP):${LD_LIBRARY_PATH} && \ - echo "LIBGOMP=${LIBGOMP}" && \ find /root/.cache/uv -name "*.whl" && \ SETUPTOOLS_SCM_PRETEND_VERSION="$VLLM_VERSION" uv build \ --wheel --out-dir ${WHEEL_DIR} --no-build-isolation && \ - uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh + uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh && \ + uv pip install "setuptools==80.10.2" && \ + uv pip uninstall -y pip cmake meson-python scikit-build-core maturin build wheel auditwheel pybind11 nanobind cython setuptools-rust && \ + # Remove packaged test suites + find ${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages \ + -type d \( -name tests -o -name test \) -exec rm -rf {} + && \ + find ${VIRTUAL_ENV} -name "__pycache__" -type d -exec rm -rf {} + && \ + find ${VIRTUAL_ENV} -name "*.pyc" -delete && \ + rm -f ${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/xgrammar/lib/libxgrammar.a ############################################################### # FINAL VLLM IMAGE STAGE # @@ -94,7 +95,6 @@ ENV PKG_CONFIG_PATH=${PCP_DIR}/usr/lib64/pkgconfig:/usr/local/lib/pkgconfig/ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" ENV LD_LIBRARY_PATH=${PCP_DIR}/usr/lib64:${PCP_DIR}/usr/lib:${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/torch/lib:/usr/local/lib:$LD_LIBRARY_PATH:/usr/local/lib64:/usr/lib64:/usr/lib ENV UV_LINK_MODE=copy -ARG VLLM_VERSION="0.22.1" ARG UV_EXTRA_INDEX_URL="https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/" ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY=first-match @@ -103,7 +103,7 @@ ENV UV_INDEX_STRATEGY=first-match RUN --mount=type=cache,target=/root/.cache/uv \ rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ microdnf install --nodocs -y \ - libomp libicu tar autoconf automake libtool findutils openssl numactl numactl-devel \ + libomp libicu tar autoconf automake libtool findutils openssl openblas-devel numactl numactl-devel \ pkgconfig xsimd gcc-toolset-14 libsndfile \ libtiff libjpeg openjpeg2 zlib zeromq \ freetype lcms2 libwebp tcl tk utf8proc \ @@ -113,10 +113,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && microdnf update -y \ && microdnf clean all -# The `lscpu` command was added as a requirement in part of https://github.com/vllm-project/vllm/pull/21032, so installing it. -RUN microdnf install --nodocs -y util-linux && \ - microdnf clean all - COPY --from=builder-base /usr/lib64/libprotobuf.so.25 /usr/lib64/ COPY --from=builder-base /usr/lib64/libprotobuf.so.25.0.0 /usr/lib64/ @@ -154,4 +150,3 @@ USER 2000 ENTRYPOINT ["vllm", "serve"] - diff --git a/docker/build_vllm_ppc64le.sh b/docker/build_vllm_ppc64le.sh index f2fbad8e2199..e49edbd7c687 100644 --- a/docker/build_vllm_ppc64le.sh +++ b/docker/build_vllm_ppc64le.sh @@ -13,7 +13,6 @@ cd "$REPO_ROOT" ######################################## IBM_DEVPI_URL=${IBM_DEVPI_URL:-"https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/"} -RHOAI_INDEX_URL=${RHOAI_INDEX_URL:-"https://console.redhat.com/api/pypi/public-rhai/rhoai/3.4/cpu-ubi9/simple/"} ######################################## # wheel dir @@ -39,16 +38,65 @@ try_install_from_devpi() { ######################################## cd "$REPO_ROOT" TORCH_VERSION=${TORCH_VERSION:-$(grep -E '^torch==.+==\s*"ppc64le"' requirements/cpu.txt | grep -Eo '\b[0-9\.]+\b' || true)} -TORCH_VERSION=${TORCH_VERSION:-2.11.0} - -TORCHVISION_VERSION=${TORCHVISION_VERSION:-0.26.0} -TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION:-${TORCH_VERSION}} +TORCH_VERSION=${TORCH_VERSION:-2.13.0} + +# Parse version parts +TORCH_MAJOR=$(echo "${TORCH_VERSION}" | cut -d. -f1) +TORCH_MINOR=$(echo "${TORCH_VERSION}" | cut -d. -f2) +TORCH_PATCH=$(echo "${TORCH_VERSION}" | cut -d. -f3 | cut -d+ -f1) +TORCH_PATCH=${TORCH_PATCH:-0} + +# 1. Calculate torchvision (always minor + 15) +VISION_MINOR=$((TORCH_MINOR + 15)) +TORCHVISION_VERSION="0.${VISION_MINOR}.${TORCH_PATCH}" + +# 2. Calculate torchaudio (minor - 2 for >= 2.13, otherwise same as torch) +if [ "${TORCH_MAJOR}" -eq 2 ] && [ "${TORCH_MINOR}" -ge 13 ]; then + AUDIO_MINOR=$((TORCH_MINOR - 2)) + TORCHAUDIO_VERSION="${TORCH_MAJOR}.${AUDIO_MINOR}.${TORCH_PATCH}" +else + TORCHAUDIO_VERSION="${TORCH_VERSION}" +fi export TORCH_VERSION export TORCHVISION_VERSION export TORCHAUDIO_VERSION -export OPENCV_VERSION=${OPENCV_VERSION:-4.13.0.92} -export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} + +echo "Resolved PyTorch versions:" +echo " torch: ${TORCH_VERSION}" +echo " torchvision: ${TORCHVISION_VERSION}" +echo " torchaudio: ${TORCHAUDIO_VERSION}" + +OPENCV_REQUIREMENT=$(grep -E '^opencv-python-headless[[:space:]]*(==|>=|<=|~=)' \ + requirements/common.txt | + sed -E 's/^opencv-python-headless[[:space:]]*//' | + sed -E 's/[[:space:]]*#.*$//' | + tr -d ' ' | + head -n1) + +export OPENCV_REQUIREMENT +echo "OpenCV requirement: ${OPENCV_REQUIREMENT}" + + +XGRAMMAR_VERSION=$( + sed -nE 's/^xgrammar[[:space:]]*==[[:space:]]*([^,;[:space:]]+).*/\1/p' requirements/common.txt | + head -n1 +) + +if [ -z "${XGRAMMAR_VERSION}" ]; then + XGRAMMAR_VERSION=$( + sed -nE 's/^xgrammar[[:space:]]*>=[[:space:]]*([^,;[:space:]]+).*/\1/p' requirements/common.txt | + head -n1 + ) +fi + +export XGRAMMAR_VERSION + +if [ -z "${XGRAMMAR_VERSION}" ]; then + echo "xgrammar version not found in requirements/common.txt" + exit 1 +fi +echo "XGRAMMAR_VERSION=${XGRAMMAR_VERSION}" ######################################## # install system dependencies @@ -56,15 +104,20 @@ export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm || true +# Dynamically match openssl-devel version with installed openssl-libs +OPENSSL_DEVEL="openssl-devel-$(rpm -q --queryformat '%{VERSION}-%{RELEASE}' openssl-libs 2>/dev/null || echo 'openssl-devel')" + microdnf install -y \ + "${OPENSSL_DEVEL}" \ python3.12 python3.12-devel python3.12-pip gcc \ git jq gcc-toolset-14 gcc-toolset-14-libatomic-devel \ - automake libtool clang-devel openssl-devel \ + automake libtool clang-devel \ harfbuzz-devel kmod lcms2-devel libimagequant-devel libjpeg-turbo-devel \ llvm15-devel libraqm-devel libtiff-devel libwebp-devel libxcb-devel \ ninja-build openjpeg2-devel pkgconfig \ - tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel - + tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file \ + openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel + rpm -ivh --nodeps \ https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-17.el9.ppc64le.rpm @@ -88,11 +141,9 @@ python --version ######################################## # install build tools (stable uv) ######################################## - -pip install -U pip setuptools-rust pip install uv -pip install "setuptools<70" build wheel cmake auditwheel -uv pip install "setuptools<70" cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation +pip install build cmake +uv pip install setuptools==78.1.1 cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation ######################################## # Rust @@ -115,127 +166,310 @@ export MAX_JOBS=${MAX_JOBS:-$(nproc)} export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 ######################################## -# Install Packages From Devpi +# Install Torch packages from DevPI or build from source ######################################## -uv pip install numpy==2.3.5 pillow==12.2.0 --extra-index-url "$IBM_DEVPI_URL" -try_install_from_devpi "opencv-python-headless==${OPENCV_VERSION}" -try_install_from_devpi "torch==${TORCH_VERSION}" -try_install_from_devpi "torchvision==${TORCHVISION_VERSION}" + +is_available_on_devpi() { + local pkg=$1 + local version=$2 + + echo "Checking DevPI for ${pkg}==${version}..." + + if pip index versions "${pkg}" \ + --index-url "${IBM_DEVPI_URL}" 2>/dev/null | + grep -F "${version}" >/dev/null; then + + echo "${pkg}==${version} is available on DevPI" + return 0 + fi + + echo "${pkg}==${version} is NOT available on DevPI" + return 1 +} + +######################################## +# Common packages from DevPI +######################################## +uv pip install numpy==2.3.5 pillow --extra-index-url "$IBM_DEVPI_URL" ######################################## -# torch audio +# OpenCV ######################################## -TEMP_BUILD_DIR=$(mktemp -d) -cd "${TEMP_BUILD_DIR}" -export BUILD_SOX=1 BUILD_KALDI=1 BUILD_RNNT=1 USE_FFMPEG=0 USE_ROCM=0 USE_CUDA=0 -export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 -git clone --recursive https://github.com/pytorch/audio.git -b "v${TORCHAUDIO_VERSION}" -cd audio -#patching -sed -i ' -s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; -s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; -s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; -' tools/setup_helpers/extension.py -MAX_JOBS=${MAX_JOBS:-$(nproc)} \ -BUILD_VERSION=${TORCHAUDIO_VERSION} \ -uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation -uv pip install "${WHEEL_DIR}"/torchaudio*.whl -cd "${REPO_ROOT}" -rm -rf "${TEMP_BUILD_DIR}" +if [[ "${OPENCV_REQUIREMENT}" == ">=4.13.0" ]]; then + echo "OpenCV requirement is ${OPENCV_REQUIREMENT}" + echo "Installing opencv-python-headless from DevPI" + + OPENCV_VERSION="==4.13.0.92+ppc64le1" +else + echo "Using OpenCV requirement: ${OPENCV_REQUIREMENT}" + + OPENCV_VERSION="${OPENCV_REQUIREMENT}" +fi + +try_install_from_devpi "opencv-python-headless${OPENCV_VERSION}" ######################################## -# Xgrammar +# Check Torch / Torchvision / Torchaudio ######################################## -uv pip install \ - "scikit-build-core==0.11.6" \ - "pyproject-metadata<0.8" \ - pathspec \ - packaging \ - distro \ - "setuptools<70" \ - setuptools_scm \ - cmake \ - ninja \ - pybind11 \ - nanobind -uv pip install apache-tvm-ffi==0.1.12 \ - --no-build-isolation \ - --no-cache - -TEMP_BUILD_DIR=$(mktemp -d) - -pushd "${TEMP_BUILD_DIR}" - -export CFLAGS="-fno-lto -mcpu=power9" -export CXXFLAGS="-fno-lto -mcpu=power9" -export LDFLAGS="-fno-lto" -export PATH=/opt/vllm/bin:$PATH -export Python_EXECUTABLE=/opt/vllm/bin/python3 -export Python3_EXECUTABLE=/opt/vllm/bin/python3 -export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 +TORCH_FROM_DEVPI=false +TORCHVISION_FROM_DEVPI=false +TORCHAUDIO_FROM_DEVPI=false -export Python_ROOT_DIR=/opt/vllm -export Python3_ROOT_DIR=/opt/vllm +if is_available_on_devpi "torch" "${TORCH_VERSION}"; then + TORCH_FROM_DEVPI=true +fi -git clone \ - --recursive \ - https://github.com/mlc-ai/xgrammar \ - -b "v${XGRAMMAR_VERSION}" +if is_available_on_devpi "torchvision" "${TORCHVISION_VERSION}"; then + TORCHVISION_FROM_DEVPI=true +fi -cd xgrammar +if is_available_on_devpi "torchaudio" "${TORCHAUDIO_VERSION}"; then + TORCHAUDIO_FROM_DEVPI=true +fi -cp cmake/config.cmake . -export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} +######################################## +# Torch +######################################## -uv build \ - --wheel \ - --out-dir "${WHEEL_DIR}" \ - --no-build-isolation +if [[ "${TORCH_FROM_DEVPI}" == "true" ]]; then + echo "Installing torch==${TORCH_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torch==${TORCH_VERSION}" +else + echo "Torch wheel not available on DevPI. Building torch==${TORCH_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" -uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + git clone \ + --recursive \ + https://github.com/pytorch/pytorch.git \ + -b "v${TORCH_VERSION}" -popd + cd pytorch -rm -rf "${TEMP_BUILD_DIR}" -cd "${REPO_ROOT}" + export BLAS=OpenBLAS + export USE_OPENMP=1 + export USE_MKLDNN=OFF + export USE_MKLDNN_CBLAS=OFF + export _GLIBCXX_USE_CXX11_ABI=1 + uv pip install maturin + pip install -r requirements.txt --extra-index-url "$IBM_DEVPI_URL" + python setup.py develop + rm -f dist/torch*+git*whl + MAX_JOBS=${MAX_JOBS:-$(nproc)} \ + PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir "${WHEEL_DIR}" + uv pip install "${WHEEL_DIR}"/torch*.whl + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi ######################################## -# RHOAI Binary Downloads +# Torchvision ######################################## -pip download \ - --index-url "${RHOAI_INDEX_URL}" \ - --only-binary=:all: \ - --no-deps \ - llvmlite==0.47.0 \ - -d "${WHEEL_DIR}" -pip download \ - --index-url "${RHOAI_INDEX_URL}" \ - --only-binary=:all: \ - --no-deps \ - Numba==0.65.0 \ - -d "${WHEEL_DIR}" +if [[ "${TORCHVISION_FROM_DEVPI}" == "true" ]]; then + echo "Installing torchvision==${TORCHVISION_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torchvision==${TORCHVISION_VERSION}" +else + echo "Torchvision wheel not available on DevPI. Building torchvision==${TORCHVISION_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" + + export TORCHVISION_USE_NVJPEG=0 TORCHVISION_USE_FFMPEG=0 + git clone \ + --recursive \ + https://github.com/pytorch/vision.git \ + -b "v${TORCHVISION_VERSION}" + + cd vision + uv pip install standard-pkg-resources --no-build-isolation + MAX_JOBS=${MAX_JOBS:-$(nproc)} \ + BUILD_VERSION=${TORCHVISION_VERSION} \ + uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation + + export BUILD_VERSION="${TORCHVISION_VERSION}" + + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/torchvision*.whl + + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi + +######################################## +# Torchaudio +######################################## + +if [[ "${TORCHAUDIO_FROM_DEVPI}" == "true" ]]; then + echo "Installing torchaudio==${TORCHAUDIO_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "torchaudio==${TORCHAUDIO_VERSION}" +else + echo "Torchaudio wheel not available on DevPI. Building torchaudio==${TORCHAUDIO_VERSION} from source." + + TEMP_BUILD_DIR=$(mktemp -d) + cd "${TEMP_BUILD_DIR}" + + export BUILD_SOX=1 + export BUILD_KALDI=1 + export BUILD_RNNT=1 + export USE_FFMPEG=0 + export USE_ROCM=0 + export USE_CUDA=0 + export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 + + git clone \ + --recursive \ + https://github.com/pytorch/audio.git \ + -b "v${TORCHAUDIO_VERSION}" + + cd audio + + # Patches required for newer Python versions + sed -i ' + s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; + s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; + s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; + ' tools/setup_helpers/extension.py + + MAX_JOBS="${MAX_JOBS}" \ + BUILD_VERSION="${TORCHAUDIO_VERSION}" \ + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/torchaudio*.whl + + cd "${REPO_ROOT}" + rm -rf "${TEMP_BUILD_DIR}" +fi + +######################################## +# Xgrammar +######################################## +XGRAMMAR_FROM_DEVPI=false + +if is_available_on_devpi "xgrammar" "${XGRAMMAR_VERSION}"; then + XGRAMMAR_FROM_DEVPI=true +fi + +if [[ "${XGRAMMAR_FROM_DEVPI}" == "true" ]]; then + echo "Installing xgrammar==${XGRAMMAR_VERSION} from DevPI" + + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --only-binary=:all: \ + --no-build-isolation \ + "xgrammar==${XGRAMMAR_VERSION}" +else + echo "Xgrammar wheel not available on DevPI. Building xgrammar==${XGRAMMAR_VERSION} from source." + # Install xgrammar build dependencies only when building from source. + uv pip install \ + "scikit-build-core==0.11.6" \ + "pyproject-metadata<0.8" \ + pathspec \ + packaging \ + distro \ + setuptools_scm \ + cmake \ + ninja \ + pybind11 \ + nanobind + uv pip install apache-tvm-ffi==0.1.12 \ + --no-build-isolation \ + --no-cache + + TEMP_BUILD_DIR=$(mktemp -d) + + pushd "${TEMP_BUILD_DIR}" + + export CFLAGS="-fno-lto -mcpu=power9" + export CXXFLAGS="-fno-lto -mcpu=power9" + export LDFLAGS="-fno-lto" + export PATH=/opt/vllm/bin:$PATH + + export Python_EXECUTABLE=/opt/vllm/bin/python3 + export Python3_EXECUTABLE=/opt/vllm/bin/python3 + export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 + + export Python_ROOT_DIR=/opt/vllm + export Python3_ROOT_DIR=/opt/vllm + + git clone \ + --recursive \ + https://github.com/mlc-ai/xgrammar \ + -b "v${XGRAMMAR_VERSION}" + + cd xgrammar + + cp cmake/config.cmake . + export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} + + uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + + uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + + popd + + rm -rf "${TEMP_BUILD_DIR}" + cd "${REPO_ROOT}" +fi + +NUMBA_REQUIREMENT=$(grep -E '^numba[[:space:]]*(==|>=|<=|~=)' \ + requirements/cpu.txt | + sed -E 's/^numba[[:space:]]*//' | + sed -E 's/[[:space:]]*#.*$//' | + sed -E 's/;.*$//' | + tr -d ' ' | + head -n1) + +echo "Numba requirement: ${NUMBA_REQUIREMENT}" +try_install_from_devpi "numba${NUMBA_REQUIREMENT}" ######################################## # install built wheels ######################################## uv pip install setuptools_scm maturin setuptools-rust ninja scikit-build-core pybind11 nanobind \ --no-build-isolation -uv pip install "${WHEEL_DIR}"/*.whl ######################################## # install remaining deps ######################################## sed -i.bak -e 's/.*torch.*//g' pyproject.toml requirements/*.txt - -uv pip install "setuptools>=78.1.1" --no-build-isolation - export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib64/pkgconfig:/usr/lib64/pkgconfig uv pip install -r requirements/common.txt \ -r requirements/cpu.txt \ -r requirements/build/cpu.txt --index-strategy unsafe-best-match + + From 364bc81a9ecc72d44d8d1486971d8059d3f8932d Mon Sep 17 00:00:00 2001 From: linitra24 Date: Sat, 10 Oct 2026 12:52:51 +0800 Subject: [PATCH 339/952] [Bugfix][LoRA] Fix GPT-OSS gate/up interleaving with mixed MoE LoRA formats (#60812) Signed-off-by: linitra24 --- tests/lora/test_gptoss_tp.py | 38 ++++++++++++++++++- .../fused_moe/experts/lora_experts_mixin.py | 23 +++++++++-- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index 70b94ea2eb2a..baabaedded63 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -64,7 +64,13 @@ def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: +def generate_and_test( + llm: vllm.LLM, + lora_path: str, + lora_id: int, + *, + is_3d_lora_weight: bool = False, +) -> None: prompts = [ PROMPT_TEMPLATE.format( context="Give the average number of working horses on farms with more than 5000 total horses." # noqa: E501 @@ -80,7 +86,11 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: outputs = llm.generate( prompts, sampling_params, - lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None, + lora_request=LoRARequest( + str(lora_id), lora_id, lora_path, is_3d_lora_weight=is_3d_lora_weight + ) + if lora_id + else None, ) # Print the outputs. generated_texts: list[str] = [] @@ -187,3 +197,27 @@ def test_gpt_oss_lora_tp2( generate_and_test(llm, gptoss20b_lora_files, lora_id=1) generate_and_test(llm, gptoss20b_lora_files, lora_id=2) + + +def test_gpt_oss_lora_mixed_moe_format( + gptoss20b_lora_files, + monkeypatch: pytest.MonkeyPatch, +): + _enable_deterministic_lora_shrink(monkeypatch) + + llm = vllm.LLM( + MODEL_PATH, + max_model_len=1024, + enable_lora=True, + enable_mixed_moe_lora_format=True, + max_loras=2, + max_lora_rank=8, + max_num_seqs=2, + max_num_batched_tokens=2048, + compilation_config=vllm.config.CompilationConfig( + cudagraph_specialize_lora=False, + ), + ) + + generate_and_test(llm, gptoss20b_lora_files, lora_id=1, is_3d_lora_weight=True) + generate_and_test(llm, gptoss20b_lora_files, lora_id=2, is_3d_lora_weight=True) diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py index ac75e7bc798b..5d552a88944d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py @@ -3,6 +3,8 @@ import torch +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext @@ -22,6 +24,7 @@ class LoRAExpertsMixin: state is on lora_context or passed as arguments. """ + moe_config: FusedMoEConfig _lora_context: MoELoRAContext | None = None def set_lora_context(self, ctx: MoELoRAContext) -> None: @@ -70,8 +73,13 @@ def apply_w13_lora( # concat swap afterwards. w13_lora_a_stacked = w13_lora_a_stacked[::-1] w13_lora_b_stacked = w13_lora_b_stacked[::-1] - return lora_context.punica_wrapper.add_lora_w13( - y, + interleave_w13 = ( + lora_context.w13_num_slices == 2 + and self.moe_config.activation == MoEActivation.SWIGLUOAI + ) + lora_output = torch.zeros_like(y) if interleave_w13 else y + result = lora_context.punica_wrapper.add_lora_w13( + lora_output, x, w13_lora_a_stacked, w13_lora_b_stacked, @@ -89,9 +97,18 @@ def apply_w13_lora( lora_context.w13_num_slices, lora_context.fully_sharded, lora_context.use_tuned_config, - add_inputs=add_inputs, + add_inputs=add_inputs or interleave_w13, token_lora_mapping=lora_context.local_token_lora_mapping, ) + if interleave_w13: + # Map gate/up slices onto even/odd output channels without a copy. + gate_up_delta = lora_output.unflatten(-1, (2, -1)) + gate_up_output = y.unflatten(-1, (-1, 2)).transpose(-1, -2) + if add_inputs: + gate_up_output.add_(gate_up_delta) + else: + gate_up_output.copy_(gate_up_delta) + return result def apply_w2_lora( self, From 8ea646b82713067406344e315a3c4ae211b0f799 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:59:35 -0700 Subject: [PATCH 340/952] [Rust Frontend][gRPC] Remove `GenerateStream` producer task so disconnects abort requests (#60922) Co-authored-by: Codex Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Signed-off-by: Alec Flowers Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 - rust/src/server/Cargo.toml | 1 - rust/src/server/src/grpc/inference.rs | 130 ++++++++++++++------------ rust/src/server/src/grpc/tests.rs | 64 +++++++++++++ 4 files changed, 132 insertions(+), 64 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index c7b3d66e4c1b..419365580535 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5997,7 +5997,6 @@ dependencies = [ "tls-listener", "tokio", "tokio-openssl", - "tokio-stream", "tokio-util", "tonic", "tonic-health", diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 8b22b6344a0e..b7b00398ced1 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -35,7 +35,6 @@ thiserror-ext.workspace = true tls-listener.workspace = true tokio.workspace = true tokio-openssl.workspace = true -tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true tonic-health.workspace = true diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index a1e631b2e4e8..1484a9f9e375 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -5,16 +5,17 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Instant; -use futures::{Stream, StreamExt as _}; +use asynk_strim_attr::{TryYielder, try_stream}; +use futures::{Stream, StreamExt as _, pin_mut}; use thiserror_ext::AsReport as _; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status}; use tracing::{Span, info, info_span, warn}; use tracing_futures::Instrument as _; use uuid::Uuid; use vllm_llm::current_unix_timestamp_secs; -use vllm_text::{DecodedTextEvent, Prompt, SampledDelta, TextOutputStreamExt as _, TextRequest}; +use vllm_text::{ + DecodedTextEvent, Prompt, SampledDelta, TextOutputStream, TextOutputStreamExt as _, TextRequest, +}; use super::convert::{self, ResponseOpts}; use super::{InferenceServer, pb}; @@ -304,67 +305,72 @@ impl pb::inference_server::Inference for InferenceServiceImpl { .await .map_err(|error| log_text_error(&request_span, started_at, "submission", error))?; - let (tx, rx) = mpsc::channel(32); - - let task_span = request_span.clone(); - tokio::spawn( - async move { - futures::pin_mut!(stream); - while let Some(event) = stream.next().await { - let response = match event { - Err(error) => Err(log_text_error(&task_span, started_at, "stream", error)), - Ok(DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) => { - let prompt_info = convert::to_prompt_info( - &prompt_token_ids, - prompt_logprobs.as_ref(), - &response_opts, - ); - Ok(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: None, - }) - } - Ok(DecodedTextEvent::TextDelta { - decoded, - sampled: - SampledDelta { - token_ids, - logprobs, - }, - finished, - }) => convert::to_sequence_output( - &decoded.text, - &token_ids, - logprobs.as_ref(), - finished.as_deref(), - &response_opts, - ) - .map(|outputs| pb::GenerateResponse { - prompt_info: None, - outputs: Some(outputs), - }), - }; - - let failed = response.is_err(); - if tx.send(response).await.is_err() || failed { - break; - } - } - info!( - parent: &task_span, - elapsed_ms = started_at.elapsed().as_millis() as u64, - "gRPC inference stream closed" + let response_stream = + generate_stream_responses(stream, response_opts, request_span.clone(), started_at) + .instrument(request_span); + Ok(Response::new(Box::pin(response_stream))) + } +} + +/// Convert one decoded text stream into gRPC streaming responses. +/// +/// The response body owns the generation stream, so a client disconnect drops +/// it and aborts the request in the engine, even before any output arrives. +#[try_stream] +async fn generate_stream_responses( + stream: impl TextOutputStream, + response_opts: ResponseOpts, + request_span: Span, + started_at: Instant, + mut y: TryYielder, +) -> Result<(), Status> { + pin_mut!(stream); + + while let Some(event) = stream.next().await { + let response = match event { + Err(error) => Err(log_text_error(&request_span, started_at, "stream", error)), + Ok(DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) => { + let prompt_info = convert::to_prompt_info( + &prompt_token_ids, + prompt_logprobs.as_ref(), + &response_opts, ); + Ok(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: None, + }) } - .instrument(request_span), - ); - - let response_stream = ReceiverStream::new(rx); - Ok(Response::new(Box::pin(response_stream))) + Ok(DecodedTextEvent::TextDelta { + decoded, + sampled: + SampledDelta { + token_ids, + logprobs, + }, + finished, + }) => convert::to_sequence_output( + &decoded.text, + &token_ids, + logprobs.as_ref(), + finished.as_deref(), + &response_opts, + ) + .map(|outputs| pb::GenerateResponse { + prompt_info: None, + outputs: Some(outputs), + }), + }; + y.yield_ok(response?).await; } + info!( + parent: &request_span, + elapsed_ms = started_at.elapsed().as_millis() as u64, + "gRPC inference stream closed" + ); + Ok(()) } fn text_error_to_status(error: vllm_text::Error) -> Status { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 0ddbfb096e43..f46c658b0188 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1827,6 +1827,70 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn dropping_generate_stream_aborts_before_engine_output() { + let (submitted_tx, submitted_rx) = tokio::sync::oneshot::channel(); + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-cancel-pending-output".to_vec(), + default_ready_response(), + Arc::new(FakeTextBackend), + None, + move |dealer, _push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + submitted_tx.send(()).expect("signal submitted request"); + + // A queued request produces no output before cancellation. + let abort = recv_engine_message(dealer).await; + assert_eq!(abort[0].as_ref(), &[0x01]); + let request_ids: Vec = + rmp_serde::from_slice(&abort[1]).expect("decode abort request IDs"); + assert_eq!(request_ids, vec![request.request_id]); + }) + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + let stream = client + .generate_stream(pb::GenerateRequest { + request_id: "test-cancel-pending-output".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect("start generation") + .into_inner(); + tokio::time::timeout(Duration::from_secs(5), submitted_rx) + .await + .expect("timed out waiting for engine submission") + .expect("engine submission signal"); + + drop(stream); + + let result = tokio::time::timeout(Duration::from_secs(5), engine_task).await; + server_task.abort(); + result + .expect("disconnect must abort without waiting for engine output") + .expect("mock engine task"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_abort_resolves_external_id_and_empty_is_noop() { From 7d0b4e57aac4c4323225b0eee9ae960ce75ffffe Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 9 Oct 2026 22:09:19 -0700 Subject: [PATCH 341/952] [ROCm] Remove MRV1 model exceptions (#56360) Signed-off-by: Nick Hill --- tests/test_config.py | 67 ++------------------------------------------ vllm/config/vllm.py | 30 -------------------- 2 files changed, 2 insertions(+), 95 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 425d55f29346..734253edc9f8 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -651,21 +651,14 @@ def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled - MRV1 paths and off breakable cudagraphs by default.""" - from vllm.config.vllm import ( - ROCM_DEFAULT_MRV1_ARCHITECTURES, - default_breakable_cudagraph_architectures, - ) + paths and off breakable cudagraphs by default.""" + from vllm.config.vllm import default_breakable_cudagraph_architectures from vllm.platforms import current_platform monkeypatch.setattr(current_platform, "is_rocm", lambda: True) # The lookup is lru_cached against a fixed platform. default_breakable_cudagraph_architectures.cache_clear() try: - assert "DeepseekV32ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - assert "DeepseekV4ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - assert "GlmMoeDsaForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES - breakable_architectures = default_breakable_cudagraph_architectures() assert "DeepseekV32ForCausalLM" not in breakable_architectures assert "DeepseekV32MTPModel" not in breakable_architectures @@ -674,66 +667,10 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): # supports uniform-batch CUDA graphs, so it must opt in or the # default FULL_AND_PIECEWISE serve path cannot start. assert "DeepseekV41ForCausalLM" in breakable_architectures - - # The carve-out takes effect via the runner-selection property - # (warning_once args must be hashable for its lru_cache). - monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) - config = SimpleNamespace( - model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]), - attention_config=AttentionConfig(), - ) - config._get_v1_model_runner_unsupported_features = lambda: [] - assert VllmConfig.use_v2_model_runner.fget(config) is False finally: default_breakable_cudagraph_architectures.cache_clear() -def test_rocm_mrv1_default_yields_to_v1_unsupported_config(monkeypatch): - """The ROCm V1 default is a speed preference, not a capability claim. - - DSpark runs only on V2, so pinning DeepSeek V4 to V1 would fail config - validation instead of serving it. With nothing V1 refuses, it still holds. - """ - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "is_rocm", lambda: True) - monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) - monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False) - - config = SimpleNamespace( - model_config=SimpleNamespace( - architectures=["DeepseekV4ForCausalLM"], is_diffusion=False - ), - attention_config=AttentionConfig(), - parallel_config=SimpleNamespace( - prefill_context_parallel_size=1, - pipeline_parallel_size=1, - enable_batch_sharded_sampling=False, - ), - scheduler_config=SimpleNamespace(async_scheduling=False), - speculative_config=None, - ) - config._dflash_needs_multi_kv_group = lambda: False - config._is_dflash_candidate_draft = lambda: False - config._get_v2_model_runner_unsupported_features = lambda: [] - # The real predicate, so the test also pins where dspark lands in it. - config._get_v1_model_runner_unsupported_features = lambda: ( - VllmConfig._get_v1_model_runner_unsupported_features(config) - ) - - assert VllmConfig.use_v2_model_runner.fget(config) is False - - config.speculative_config = SimpleNamespace( - method="dspark", enable_adaptive_verification=False - ) - assert VllmConfig.use_v2_model_runner.fget(config) is True - - # Yielding is not the same as selecting V2: the later checks still run, so - # a config neither runner can serve lands on V1 and fails validation there. - config._get_v2_model_runner_unsupported_features = lambda: ["sequence parallelism"] - assert VllmConfig.use_v2_model_runner.fget(config) is False - - @pytest.mark.parametrize( ("model", "architecture"), [ diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bbda00af40bb..b4298cd145b8 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -71,12 +71,6 @@ logger = init_logger(__name__) -# TODO(rocm): These models are either unsupported by MRV2 or slower with -# MRV2 on AMD GPUs. -ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset( - {"DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", "GlmMoeDsaForCausalLM"} -) - DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { "DeepseekV32MTPModel", @@ -727,30 +721,6 @@ def use_v2_model_runner(self) -> bool: if use_v2_model_runner is not None: return use_v2_model_runner - from vllm.platforms import current_platform - - model_config = self.model_config - if model_config is not None and current_platform.is_rocm(): - architectures = getattr(model_config, "architectures", ()) - if any(arch in ROCM_DEFAULT_MRV1_ARCHITECTURES for arch in architectures): - # This default is a speed preference, not a claim that V1 can - # serve the config, so it yields where V1 cannot. It yields by - # falling through to the checks below, not by selecting V2. - v1_unsupported = self._get_v1_model_runner_unsupported_features() - if not v1_unsupported: - logger.warning_once( - "Defaulting to V1 model runner on ROCm for model " - "architectures: %s", - ", ".join(architectures), - ) - return False - logger.warning_once( - "Skipping the ROCm V1 model runner default for %s: V1 does " - "not support %s.", - ", ".join(architectures), - ", ".join(v1_unsupported), - ) - if not HAS_TRITON: logger.warning_once( "Model Runner V2 requires Triton; using the V1 model runner instead." From fa5cfd48d561989d2f5e7fd42788e2f6cfb479c6 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sat, 10 Oct 2026 13:27:58 +0800 Subject: [PATCH 342/952] [Refactor] Deduplicate trace header handling in BaseServing (#60828) Signed-off-by: chaunceyjiang --- .../structured_decisions/test_strategies.py | 3 ++- vllm/entrypoints/generate/base/serving.py | 22 +------------------ .../generate/generative_scoring/serving.py | 20 ----------------- .../generate/structured_decisions/serving.py | 17 +------------- vllm/entrypoints/pooling/base/serving.py | 22 +------------------ vllm/entrypoints/serve/engine/serving.py | 9 ++++++++ vllm/v1/engine/async_llm.py | 6 ++++- 7 files changed, 19 insertions(+), 80 deletions(-) diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index a63e2587521c..41b9a8f01271 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -222,11 +222,12 @@ async def test_request_metadata_reaches_every_question( ) response = await serving.create_decision(request, raw_request) assert isinstance(response, StructuredDecisionResponse) + engine.is_tracing_enabled.assert_not_awaited() assert engine.generate.call_count == 2 for call in engine.generate.call_args_list: assert call.args[0].get("cache_salt") == cache_salt assert call.kwargs["trace_headers"] == ( - TRACE_HEADERS if tracing_enabled and request_headers else None + TRACE_HEADERS if request_headers else None ) diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index d1f3e94b006f..a0eda6d05ff4 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -2,14 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import time -from collections.abc import Awaitable, Mapping +from collections.abc import Awaitable from dataclasses import dataclass, field from http import HTTPStatus from typing import ClassVar, Generic, TypeVar from fastapi import Request from pydantic import ConfigDict -from starlette.datastructures import Headers from vllm import RequestOutput from vllm.engine.protocol import EngineClient @@ -33,11 +32,6 @@ from vllm.lora.request import LoRARequest from vllm.tokenizers import TokenizerLike from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.v1.metrics.stats import RequestStateStats logger = init_logger(__name__) @@ -222,20 +216,6 @@ def _convert_generation_error_to_streaming_response( status_code=e.status_code, ) - async def _get_trace_headers( - self, - headers: Headers, - ) -> Mapping[str, str] | None: - is_tracing_enabled = await self.engine_client.is_tracing_enabled() - - if is_tracing_enabled: - return extract_trace_headers(headers) - - if contains_trace_headers(headers): - log_tracing_disabled_warning() - - return None - @staticmethod def _get_data_parallel_rank(raw_request: Request | None) -> int | None: """Pulls the data parallel rank from a header, if provided.""" diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 74dbfd075997..92c084771643 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -11,7 +11,6 @@ import asyncio import math import time -from collections.abc import Mapping from typing import Any, Literal from fastapi import Request @@ -32,11 +31,6 @@ from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.utils import random_uuid logger = init_logger(__name__) @@ -462,17 +456,3 @@ def _compute_probabilities( token_id: math.exp(logprob) for token_id, logprob in label_logprobs.items() } - - async def _get_trace_headers( - self, - headers: Mapping[str, str], - ) -> Mapping[str, str] | None: - """Extract trace headers from request headers.""" - if not contains_trace_headers(headers): - return None - - if not await self.engine_client.is_tracing_enabled(): - log_tracing_disabled_warning() - return None - - return extract_trace_headers(headers) diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index e6e990bd46df..2633647a0b65 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json -from collections.abc import Mapping from typing import Any from fastapi import Request @@ -12,11 +11,6 @@ from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from .protocol import ( DecisionUsage, @@ -74,19 +68,10 @@ def __init__( super().__init__( models=models, model_config=model_config, request_logger=request_logger ) + self.engine_client = strategy.context.engine_client self.strategy = strategy self.limits = strategy.limits() - async def _get_trace_headers( - self, headers: Mapping[str, str] - ) -> Mapping[str, str] | None: - if not contains_trace_headers(headers): - return None - if not await self.strategy.context.engine_client.is_tracing_enabled(): - log_tracing_disabled_warning() - return None - return extract_trace_headers(headers) - async def create_decision( self, request: StructuredDecisionRequest, diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 92cdfaec252f..2fb0b8e1bc72 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio from abc import ABC, abstractmethod -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator from concurrent.futures import Executor from http import HTTPStatus from typing import ClassVar @@ -10,7 +10,6 @@ import torch from fastapi import Request from fastapi.responses import Response -from starlette.datastructures import Headers from vllm import PoolingRequestOutput, envs from vllm.config import VllmConfig @@ -24,11 +23,6 @@ from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer -from vllm.tracing import ( - contains_trace_headers, - extract_trace_headers, - log_tracing_disabled_warning, -) from vllm.utils.async_utils import make_async, merge_async_iterators from ...serve.engine.protocol import ErrorResponse @@ -262,20 +256,6 @@ def _validate_request(self, ctx: PoolingServeContext) -> None: return None - async def _get_trace_headers( - self, - headers: Headers, - ) -> Mapping[str, str] | None: - is_tracing_enabled = await self.engine_client.is_tracing_enabled() - - if is_tracing_enabled: - return extract_trace_headers(headers) - - if contains_trace_headers(headers): - log_tracing_disabled_warning() - - return None - class PoolingServing(PoolingBaseServing, ABC): def __init__(self, *args, **kwargs): diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 8b4244955295..15f6adc37744 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/vllm/entrypoints/serve/engine/serving.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping from http import HTTPStatus from fastapi import Request @@ -23,6 +24,7 @@ extract_prompt_len, ) from vllm.sampling_params import BeamSearchParams +from vllm.tracing import contains_trace_headers, extract_trace_headers from vllm.utils import random_uuid @@ -37,6 +39,13 @@ def __init__( self.model_config = model_config self.request_logger = request_logger + async def _get_trace_headers( + self, headers: Mapping[str, str] + ) -> Mapping[str, str] | None: + if not contains_trace_headers(headers): + return None + return extract_trace_headers(headers) + async def _check_model( self, request: AnyRequest | AnyPoolingRequest, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7e50aa045123..bdab43e92528 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,7 +44,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tasks import SupportedTask from vllm.tokenizers import TokenizerLike -from vllm.tracing import init_tracer +from vllm.tracing import init_tracer, log_tracing_disabled_warning from vllm.transformers_utils.config import maybe_register_config_serialize_by_value from vllm.usage.usage_lib import UsageContext from vllm.utils.async_utils import cancel_task_threadsafe @@ -403,6 +403,10 @@ async def add_request( if self.errored: raise EngineDeadError() + if trace_headers and self.observability_config.otlp_traces_endpoint is None: + log_tracing_disabled_warning() + trace_headers = None + is_pooling = isinstance(params, PoolingParams) if ( From 46fb84c2d4ef1f89cf472c62b8c964819a9e512a Mon Sep 17 00:00:00 2001 From: devpramod Date: Sat, 10 Oct 2026 01:34:10 -0400 Subject: [PATCH 343/952] [Model] Support EAGLE3/DFlash aux hidden states for Granite (#60855) Signed-off-by: Pai, Pramod Co-authored-by: Claude Opus 5.5 (1M context) --- vllm/model_executor/models/granite.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 81ab62937e1a..24c0851bdf85 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -60,7 +60,13 @@ from vllm.model_executor.utils import set_weight_attrs from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant +from .interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsLoRA, + SupportsPP, + SupportsQuant, +) from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -307,7 +313,7 @@ def forward( @support_torch_compile -class GraniteModel(nn.Module): +class GraniteModel(nn.Module, EagleModelMixin): hf_to_vllm_mapper = WeightsMapper( orig_to_new_stacked={ # weight_name: (param_name, shard_id) @@ -363,7 +369,7 @@ def forward( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -375,8 +381,14 @@ def forward( assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - for layer in islice(self.layers, self.start_layer, self.end_layer): + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, None) + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer) + ): hidden_states = layer(positions, hidden_states) + self._maybe_add_hidden_state( + aux_hidden_states, idx + 1, hidden_states, None + ) if not get_pp_group().is_last_rank: return IntermediateTensors( @@ -386,6 +398,9 @@ def forward( ) hidden_states = self.norm(hidden_states) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -393,7 +408,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) -class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): +class GraniteForCausalLM( + nn.Module, SupportsLoRA, SupportsPP, SupportsQuant, SupportsEagle3 +): hf_to_vllm_mapper = GraniteModel.hf_to_vllm_mapper # LoRA specific attributes packed_modules_mapping: dict[str, list[str]] = { From a205fb317703f54a9b5afee87a03af5c18f95c75 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 10 Oct 2026 00:54:52 -0500 Subject: [PATCH 344/952] [Bugfix][Model] Restore AXK1 config defaults dropped with the vendored config (#60954) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- vllm/model_executor/models/AXK1.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 83859946a255..2d9e5a0dab9b 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -116,7 +116,6 @@ def __init__( self.ep_size = self.ep_group.size() assert config.n_routed_experts is not None assert config.num_experts_per_tok is not None - assert config.scoring_func is not None assert config.hidden_act is not None self.n_routed_experts: int = config.n_routed_experts self.n_shared_experts: int | None = config.n_shared_experts @@ -134,7 +133,7 @@ def __init__( config.n_routed_experts, prefix=f"{prefix}.gate", ) - if config.topk_method == "noaux_tc": + if getattr(config, "topk_method", "noaux_tc") == "noaux_tc": self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32) ) @@ -186,7 +185,7 @@ def __init__( num_expert_group=config.n_group, topk_group=config.topk_group, prefix=f"{prefix}.experts", - scoring_func=config.scoring_func, + scoring_func=getattr(config, "scoring_func", "sigmoid"), # we do scaling outside, set factor to 1.0 to avoid double mul # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, @@ -645,11 +644,10 @@ def __init__( self.routed_scaling_factor = config.routed_scaling_factor def _is_layer_sparse(self) -> bool: - assert self.config.moe_layer_freq is not None return ( self.config.n_routed_experts is not None and self.layer_idx >= self.config.first_k_dense_replace - and self.layer_idx % self.config.moe_layer_freq == 0 + and self.layer_idx % getattr(self.config, "moe_layer_freq", 1) == 0 ) def forward( From ff440c780544abd63bc82d57557adedf15647bfd Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 10 Oct 2026 14:06:57 +0800 Subject: [PATCH 345/952] [BreakableCG] Add eager break for Mamba-family ops (#57881) Signed-off-by: zhenwei-intel Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py | 2 ++ vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py | 2 ++ vllm/model_executor/layers/mamba/mamba_mixer.py | 2 ++ vllm/model_executor/layers/mamba/mamba_mixer2.py | 2 ++ vllm/model_executor/layers/mamba/short_conv.py | 2 ++ 5 files changed, 10 insertions(+) diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index c8f5a9ba64c9..a7bbac74f66d 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -4,6 +4,7 @@ from einops import rearrange from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import ( VllmConfig, get_current_vllm_config, @@ -524,6 +525,7 @@ def weight_loader(param, loaded_weight, loaded_shard_id=None): return weight_loader +@eager_break_during_capture def olmo_hybrid_gdn_full_forward( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py index e7618b29b713..2d88f17e6af3 100644 --- a/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py @@ -8,6 +8,7 @@ import torch.nn.functional as F from einops import rearrange +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.custom_op import PluggableLayer @@ -325,6 +326,7 @@ def _forward( output[:num_actual_tokens], _ = self.out_proj(hidden) +@eager_break_during_capture def linear_attention( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 8f3b8a2638c5..74cac8b63014 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -7,6 +7,7 @@ from torch import nn from torch.nn.parameter import Parameter +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, @@ -465,6 +466,7 @@ def split_batch_to_prefill_and_decode( ) +@eager_break_during_capture def mamba_mixer( hidden_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 3753153fa3f2..bfa1103bb368 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -8,6 +8,7 @@ import torch from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.config.mamba import MambaBackendEnum from vllm.distributed import ( @@ -1214,6 +1215,7 @@ def share_replayssm_ring_trackers( replayssm_mixers[last_layer_name]._updates_replayssm_trackers = True +@eager_break_during_capture def mamba_mixer2( projected_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 37aef2b5bdb0..7a695dfd31c3 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -4,6 +4,7 @@ import torch +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import ForwardContext, get_forward_context @@ -360,6 +361,7 @@ def mamba_type(self) -> MambaAttentionBackendEnum: return MambaAttentionBackendEnum.SHORT_CONV +@eager_break_during_capture def short_conv( hidden_states: torch.Tensor, output: torch.Tensor, From 699c892dd5dd20557d774250df15d584d4d9b318 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 15:47:39 +0900 Subject: [PATCH 346/952] [Perf] Skip the M-RoPE warmup when every holder of the RoPE fuses it (#60246) Signed-off-by: lesj0610 Co-authored-by: Isotr0py --- .../test_qwen_vl_triton_warmup.py | 21 ++++++++++++------- .../warmup/qwen_vl_triton_warmup.py | 6 +++++- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/tests/model_executor/test_qwen_vl_triton_warmup.py b/tests/model_executor/test_qwen_vl_triton_warmup.py index e94f2ae21717..008c77c14f7a 100644 --- a/tests/model_executor/test_qwen_vl_triton_warmup.py +++ b/tests/model_executor/test_qwen_vl_triton_warmup.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch from vllm.model_executor.warmup.qwen_vl_triton_warmup import ( @@ -61,7 +62,10 @@ def forward(self, *args, **kwargs): assert calls == [] -def test_mrope_warmup_launches_triton_shapes() -> None: +@pytest.mark.parametrize( + "owners_fused, skipped", [((), False), ((True,), True), ((True, False), False)] +) +def test_mrope_warmup_launches_triton_shapes(owners_fused, skipped) -> None: from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding launched: list[tuple[torch.Size, torch.Size]] = [] @@ -84,7 +88,13 @@ class FakeMropeModel(torch.nn.Module): def __init__(self) -> None: super().__init__() - self.rotary_emb = FakeRope() + rope = FakeRope() + if not owners_fused: + self.rotary_emb = rope + for i, fused in enumerate(owners_fused): + owner = torch.nn.Module() + owner.rotary_emb, owner.use_fused_qk_norm_rope_gate = rope, fused + self.add_module(f"attn{i}", owner) def get_mrope_input_positions(self, input_tokens, mm_features): raise AssertionError("warmup must not build live M-RoPE positions") @@ -99,11 +109,8 @@ def get_mrope_input_positions(self, input_tokens, mm_features): device=torch.device("cpu"), ) _warm_mrope(runner, FakeMropeModel()) - assert [shape for shape, _ in launched] == [ - torch.Size((3, 1)), - torch.Size((3, 2)), - torch.Size((3, 16)), - ] + expected = [torch.Size((3, n)) for n in (1, 2, 16)] + assert [shape for shape, _ in launched] == ([] if skipped else expected) def test_mrope_warmup_skips_models_without_supports_mrope() -> None: diff --git a/vllm/model_executor/warmup/qwen_vl_triton_warmup.py b/vllm/model_executor/warmup/qwen_vl_triton_warmup.py index 45956fe15cb5..08f6bed2dae2 100644 --- a/vllm/model_executor/warmup/qwen_vl_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_vl_triton_warmup.py @@ -72,7 +72,11 @@ def _warm_mrope(runner: "GPUModelRunner", model: torch.nn.Module) -> None: (module for module in model.modules() if isinstance(module, MRotaryEmbedding)), None, ) - if rope is None: + if rope is None or all( + getattr(module, "use_fused_qk_norm_rope_gate", False) + for module in model.modules() + if any(child is rope for child in module.children()) + ): return model_config = runner.model_config From c1692c85f16272e6b6f3de5ae4157e2f3b61fd4a Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Sat, 10 Oct 2026 14:50:34 +0800 Subject: [PATCH 347/952] [Misc] mv chat_utils.py to renderers. (#60522) Signed-off-by: wang.yuqi --- .github/CODEOWNERS | 2 -- docs/models/generative_models.md | 2 +- .../pooling/embed/vision_embedding_offline.py | 2 +- rust/src/chat/src/multimodal.rs | 2 +- .../multimodal/llm/test_mm_cache_stats.py | 2 +- .../llm/test_mm_processor_kwargs.py | 2 +- .../entrypoints/pooling/embed/test_online.py | 2 +- .../reward/test_token_reward_online.py | 2 +- tests/entrypoints/pooling/scoring/util.py | 6 ++--- .../unit_tests/test_chat_billing_header.py | 2 +- .../multimodal/generation/test_memory_leak.py | 2 +- .../models/multimodal/pooling/test_colpali.py | 6 ++--- .../multimodal/pooling/test_colqwen3.py | 4 +-- .../pooling/test_jinavl_reranker.py | 4 +-- .../pooling/test_llama_nemotron_vl.py | 8 +++--- .../multimodal/processing/test_common.py | 2 +- .../test_chat_utils.py | 16 ++++++------ .../test_chat_utils_prompt_embeds.py | 6 ++--- tests/renderers/test_hf.py | 2 +- tests/renderers/test_warmup.py | 2 +- tests/tokenizers_/test_deepseek_v4.py | 2 +- tests/tokenizers_/test_mistral.py | 2 +- vllm/entrypoints/anthropic/serving.py | 2 +- vllm/entrypoints/cohere/serving.py | 4 +-- vllm/entrypoints/common/offline.py | 8 +++--- vllm/entrypoints/generate/api_router.py | 2 +- vllm/entrypoints/generate/base/protocol.py | 2 +- .../generate/structured_decisions/protocol.py | 2 +- .../structured_decisions/strategies.py | 2 +- .../launchers/api_server/app_state.py | 2 +- vllm/entrypoints/launchers/cli_args.py | 6 ++--- .../entrypoints/launchers/render/app_state.py | 2 +- vllm/entrypoints/llm.py | 10 +++---- .../openai/chat_completion/batch_serving.py | 2 +- .../openai/chat_completion/protocol.py | 10 +++---- .../openai/chat_completion/serving.py | 10 +++---- vllm/entrypoints/openai/responses/context.py | 6 ++--- vllm/entrypoints/openai/responses/protocol.py | 8 +++--- vllm/entrypoints/openai/responses/serving.py | 6 ++--- .../openai/responses/streaming_events.py | 2 +- vllm/entrypoints/openai/responses/utils.py | 2 +- vllm/entrypoints/pooling/base/io_processor.py | 6 ++--- vllm/entrypoints/pooling/base/protocol.py | 8 +++--- vllm/entrypoints/pooling/base/serving.py | 2 +- .../entrypoints/pooling/embed/io_processor.py | 10 +++---- vllm/entrypoints/pooling/embed/protocol.py | 2 +- vllm/entrypoints/pooling/factories.py | 4 +-- vllm/entrypoints/pooling/offline.py | 2 +- .../pooling/scoring/io_processor.py | 2 +- vllm/entrypoints/pooling/scoring/typing.py | 2 +- vllm/entrypoints/pooling/scoring/utils.py | 6 ++--- vllm/entrypoints/pooling/typing.py | 2 +- .../scale_out/token_in_token_out/serving.py | 2 +- vllm/entrypoints/serve/engine/typing.py | 2 +- vllm/entrypoints/serve/tokenize/protocol.py | 8 +++--- vllm/entrypoints/serve/tokenize/serving.py | 2 +- vllm/parser/abstract_parser.py | 8 +++--- vllm/parser/engine/parser_engine.py | 2 +- vllm/parser/harmony.py | 2 +- vllm/parser/response_template.py | 2 +- vllm/parser/utils.py | 2 +- vllm/renderers/base.py | 4 +-- vllm/{entrypoints => renderers}/chat_utils.py | 0 vllm/renderers/cohere.py | 6 ++--- vllm/renderers/deepseek_v32.py | 2 +- vllm/renderers/deepseek_v4.py | 2 +- vllm/renderers/hf.py | 26 +++++++++---------- vllm/renderers/inkling.py | 4 +-- vllm/renderers/inputs/preprocess.py | 2 +- vllm/renderers/kimi_k3.py | 6 ++--- vllm/renderers/mistral.py | 4 +-- vllm/renderers/online_derenderer.py | 2 +- vllm/renderers/online_renderer.py | 12 ++++----- vllm/renderers/params.py | 2 +- vllm/renderers/terratorch.py | 4 +-- vllm/tokenizers/deepseek_v32.py | 2 +- vllm/tokenizers/deepseek_v4.py | 2 +- vllm/tokenizers/deepseek_v41.py | 2 +- vllm/tokenizers/kimi_audio.py | 2 +- vllm/tokenizers/mistral.py | 2 +- vllm/tokenizers/protocol.py | 2 +- vllm/tool_parsers/apertus_tool_parser.py | 2 +- vllm/tool_parsers/deepseekv31_tool_parser.py | 2 +- vllm/tool_parsers/deepseekv3_tool_parser.py | 2 +- vllm/tool_parsers/dots_tool_parser.py | 2 +- .../tool_parsers/functiongemma_tool_parser.py | 2 +- vllm/tool_parsers/gigachat3_tool_parser.py | 2 +- vllm/tool_parsers/granite4_tool_parser.py | 2 +- .../granite_20b_fc_tool_parser.py | 2 +- vllm/tool_parsers/hermes_tool_parser.py | 2 +- vllm/tool_parsers/hy_v3_tool_parser.py | 2 +- vllm/tool_parsers/hy_v4_tool_parser.py | 2 +- vllm/tool_parsers/internlm2_tool_parser.py | 2 +- vllm/tool_parsers/jamba_tool_parser.py | 2 +- vllm/tool_parsers/llama_tool_parser.py | 2 +- vllm/tool_parsers/minicpm5xml_tool_parser.py | 2 +- vllm/tool_parsers/muse_glimmer_tool_parser.py | 2 +- vllm/tool_parsers/phi4mini_tool_parser.py | 2 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 2 +- vllm/tool_parsers/streaming.py | 2 +- vllm/tool_parsers/xlam_tool_parser.py | 2 +- 101 files changed, 185 insertions(+), 187 deletions(-) rename tests/{entrypoints/unit_tests => renderers}/test_chat_utils.py (99%) rename vllm/{entrypoints => renderers}/chat_utils.py (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 65d9454b283b..a92ab6a2c963 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -43,8 +43,6 @@ /vllm/entrypoints/pooling @noooop /vllm/entrypoints/serve/sagemaker @DarkLight1337 /vllm/entrypoints/serve @njhill -/vllm/entrypoints/*.py @njhill -/vllm/entrypoints/chat_utils.py @DarkLight1337 /vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 diff --git a/docs/models/generative_models.md b/docs/models/generative_models.md index 46f8f4096d6a..04d836a8f45e 100644 --- a/docs/models/generative_models.md +++ b/docs/models/generative_models.md @@ -127,7 +127,7 @@ If the model doesn't have a chat template or you want to specify another one, you can explicitly pass a chat template: ```python -from vllm.entrypoints.chat_utils import load_chat_template +from vllm.renderers.chat_utils import load_chat_template # You can find a list of existing chat templates under `examples/` custom_template = load_chat_template(chat_template="") diff --git a/examples/pooling/embed/vision_embedding_offline.py b/examples/pooling/embed/vision_embedding_offline.py index 5175a6cd46ce..02923e9d8c20 100644 --- a/examples/pooling/embed/vision_embedding_offline.py +++ b/examples/pooling/embed/vision_embedding_offline.py @@ -223,7 +223,7 @@ def run_vlm2vec_qwen2vl(seed: int): from peft import PeftConfig, PeftModel from transformers import AutoModelForImageTextToText, AutoProcessor - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B" diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 4138d39bdab3..4b625f3d8b53 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -714,7 +714,7 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result { /// /// Embedding inputs share their base modality's budget rather than getting one /// of their own, matching Python's `modality.replace("_embeds", "")` in -/// `vllm/entrypoints/chat_utils.py`. +/// `vllm/renderers/chat_utils.py`. fn media_part_limit_modality(part: &MediaContentPart) -> Option { match part { MediaContentPart::Text { .. } => None, diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py index ac86788f69a6..f15feb339d93 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py @@ -8,7 +8,7 @@ from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS from vllm import LLM -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.v1.metrics import loggers as stat_loggers from vllm.v1.metrics.reader import Counter, Metric diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index 55990ed1dda0..91d2b0aa434a 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -7,9 +7,9 @@ import pytest from vllm import LLM, SamplingParams -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.inputs import TextPrompt from vllm.outputs import RequestOutput +from vllm.renderers.chat_utils import ChatCompletionMessageParam def _make_mock_llm() -> LLM: diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index 8ac70876f4d2..ec78515de3f1 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -18,7 +18,6 @@ ) from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse from vllm.entrypoints.pooling.utils import ( @@ -27,6 +26,7 @@ decode_pooling_output, ) from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers import get_tokenizer from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor diff --git a/tests/entrypoints/pooling/reward/test_token_reward_online.py b/tests/entrypoints/pooling/reward/test_token_reward_online.py index 9d89f8c42328..0ea0b8daf87b 100644 --- a/tests/entrypoints/pooling/reward/test_token_reward_online.py +++ b/tests/entrypoints/pooling/reward/test_token_reward_online.py @@ -11,13 +11,13 @@ from tests.models.utils import check_embeddings_close from tests.utils import RemoteOpenAIServer -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse from vllm.entrypoints.pooling.utils import ( MetadataItem, build_metadata_items, decode_pooling_output, ) +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers import get_tokenizer from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor diff --git a/tests/entrypoints/pooling/scoring/util.py b/tests/entrypoints/pooling/scoring/util.py index dbf4b3dc3acd..0e989d2d8779 100644 --- a/tests/entrypoints/pooling/scoring/util.py +++ b/tests/entrypoints/pooling/scoring/util.py @@ -11,12 +11,12 @@ from transformers import AutoModel, AutoTokenizer from tests.conftest import HfRunner -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam -from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score from vllm.transformers_utils.repo_utils import hf_api diff --git a/tests/entrypoints/unit_tests/test_chat_billing_header.py b/tests/entrypoints/unit_tests/test_chat_billing_header.py index 31fa97e93915..3ad694759760 100644 --- a/tests/entrypoints/unit_tests/test_chat_billing_header.py +++ b/tests/entrypoints/unit_tests/test_chat_billing_header.py @@ -6,7 +6,7 @@ import pytest -from vllm.entrypoints.chat_utils import _parse_chat_message_content +from vllm.renderers.chat_utils import _parse_chat_message_content HEADER = ( "x-anthropic-billing-header: cc_version=2.1.220.300; " diff --git a/tests/models/multimodal/generation/test_memory_leak.py b/tests/models/multimodal/generation/test_memory_leak.py index f58b7308b51d..555a61a0df33 100644 --- a/tests/models/multimodal/generation/test_memory_leak.py +++ b/tests/models/multimodal/generation/test_memory_leak.py @@ -13,8 +13,8 @@ from tests.models.registry import HF_EXAMPLE_MODELS from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.utils.mem_constants import KiB_bytes, MiB_bytes from vllm.utils.mem_utils import format_mib diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 9e040bf9a1e5..f9f1eef18911 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -14,12 +14,12 @@ import torch from PIL import Image -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam -from vllm.platforms import current_platform from ....conftest import VllmRunner diff --git a/tests/models/multimodal/pooling/test_colqwen3.py b/tests/models/multimodal/pooling/test_colqwen3.py index 992fbf209fed..2ff1d2baeb25 100644 --- a/tests/models/multimodal/pooling/test_colqwen3.py +++ b/tests/models/multimodal/pooling/test_colqwen3.py @@ -14,11 +14,11 @@ import torch from PIL import Image -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import VllmRunner diff --git a/tests/models/multimodal/pooling/test_jinavl_reranker.py b/tests/models/multimodal/pooling/test_jinavl_reranker.py index dca2403e3510..7aebc4620b48 100644 --- a/tests/models/multimodal/pooling/test_jinavl_reranker.py +++ b/tests/models/multimodal/pooling/test_jinavl_reranker.py @@ -5,12 +5,12 @@ from transformers import AutoModel from vllm.assets.base import VLLM_S3_BUCKET_URL -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from ....conftest import HfRunner, VllmRunner diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index e9e58d863541..3cbeda2e49e1 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -18,13 +18,13 @@ from PIL import Image from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor -from vllm.entrypoints.chat_utils import ( - ChatCompletionContentPartImageParam, - ChatCompletionContentPartTextParam, -) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( + ChatCompletionContentPartImageParam, + ChatCompletionContentPartTextParam, +) from vllm.transformers_utils.processors.nemotron_vl import ( SIGLIP_MEAN, SIGLIP_STD, diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 20a5121e6008..820d655e905e 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -23,13 +23,13 @@ PlaceholderMetadataResolver, collect_ec_item_metadata, ) -from vllm.entrypoints.chat_utils import _get_embeds_data from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.platforms import current_platform +from vllm.renderers.chat_utils import _get_embeds_data from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config from vllm.utils.mistral import is_mistral_tokenizer diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/renderers/test_chat_utils.py similarity index 99% rename from tests/entrypoints/unit_tests/test_chat_utils.py rename to tests/renderers/test_chat_utils.py index 75867f04b2e5..d0a9b9e6c7a8 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -14,7 +14,14 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.multimodal.utils import ( + encode_audio_url, + encode_image_url, + encode_video_url, +) +from vllm.renderers.chat_utils import ( MEDIA_CONNECTOR_REGISTRY, AsyncMultiModalItemTracker, ChatCompletionMessageParam, @@ -26,13 +33,6 @@ parse_chat_messages_async, validate_chat_template, ) -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict -from vllm.multimodal.utils import ( - encode_audio_url, - encode_image_url, - encode_video_url, -) from vllm.utils.serial_utils import tensor2base64 KIMI_K2_5_MODEL_ID = "moonshotai/Kimi-K2.5" diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index bbd4dc5c5f2e..f26680793ee2 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -16,7 +16,9 @@ import torch from transformers import AutoTokenizer -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ( _ENABLE_PROMPT_EMBEDS_ERROR, _PROMPT_EMBEDS_MISSING_DATA_ERROR, _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, @@ -26,8 +28,6 @@ parse_chat_messages, parse_chat_messages_async, ) -from vllm.exceptions import VLLMValidationError -from vllm.renderers import TokenizeParams from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, _build_mixed_prompt_embeds, diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index 833656aa3c31..c4419a70313d 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -5,9 +5,9 @@ import pytest from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.hf import ( _consolidate_system_messages, _convert_developer_to_system, diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index b0e3dcb54796..09fe047d56ef 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -28,9 +28,9 @@ import pytest -from vllm.entrypoints.chat_utils import ChatTemplateResolutionError from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.base import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.params import ChatParams diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index 4403c2a877d0..19c40f72c51a 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -8,7 +8,7 @@ import pytest import torch -from vllm.entrypoints.chat_utils import parse_chat_messages +from vllm.renderers.chat_utils import parse_chat_messages from vllm.renderers.registry import RENDERER_REGISTRY from vllm.tokenizers.deepseek_v4 import get_deepseek_v4_tokenizer from vllm.tokenizers.registry import TokenizerRegistry diff --git a/tests/tokenizers_/test_mistral.py b/tests/tokenizers_/test_mistral.py index 24bd8e00f1e0..2480d9b994f2 100644 --- a/tests/tokenizers_/test_mistral.py +++ b/tests/tokenizers_/test_mistral.py @@ -10,9 +10,9 @@ from mistral_common.guidance.grammar_factory import GrammarFactory from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy, SpecialTokens -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.mistral import ( MistralTokenizer, _validate_apply_chat_template_args, diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index f11697fa165e..7b1cf95e99b0 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -35,7 +35,6 @@ AnthropicToolChangeDefinition, AnthropicUsage, ) -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import ( JsonSchemaResponseFormat, ResponseFormat, @@ -54,6 +53,7 @@ from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.hf import HfRenderer, resolve_chat_template from vllm.renderers.online_renderer import OnlineRenderer diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 98242cc55d09..ec6dd89c3ef8 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -45,7 +45,6 @@ from pydantic import BaseModel from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.cohere.cohere_chat_message import ( CitationSource as InternalCitationSource, ) @@ -91,6 +90,7 @@ from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.cohere import ( MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY, @@ -509,7 +509,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # ``TextToolContent`` and ``DocumentToolContent``. The latter # produces ``{"type": "document", "document": {...}}`` blocks, # which vLLM's shared ``_parse_chat_message_content_parts`` in - # ``vllm/entrypoints/chat_utils.py`` rejects because the OpenAI + # ``vllm/renderers/chat_utils.py`` rejects because the OpenAI # ``ChatCompletionContentPartParam`` union only knows text / # image / audio / etc. Emitting them anyway triggers a Pydantic # ``ValidatorIterator`` failure at chat-parsing time, so we diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index 3f75c3033dd6..b33b8c34c5ad 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -15,15 +15,15 @@ SamplingParams, ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.exceptions import VLLMValidationError from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.renderers import BaseRenderer, ChatParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.renderers.inputs.preprocess import ( conversation_to_seq, parse_model_prompt, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 18dfb87a26da..66a37d6b13ee 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -72,7 +72,7 @@ async def init_generate_state( default_chat_template_kwargs: dict[str, Any], ): from vllm.entrypoints.anthropic.serving import AnthropicServingMessages - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template # The Cohere serving handler depends on the optional `cohere` SDK for # its wire-format protocol models, and is additionally gated on the diff --git a/vllm/entrypoints/generate/base/protocol.py b/vllm/entrypoints/generate/base/protocol.py index 89f6816c41cb..9781bf935f33 100644 --- a/vllm/entrypoints/generate/base/protocol.py +++ b/vllm/entrypoints/generate/base/protocol.py @@ -15,10 +15,10 @@ import vllm.envs as envs from vllm.config.utils import replace -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.sampling_params import StructuredOutputsParams logger = init_logger(__name__) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index b15fee171fca..ab89084dff2d 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -8,10 +8,10 @@ from pydantic import Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.utils import random_uuid diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index f79375e42a83..48debdb0b30f 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -16,10 +16,10 @@ from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.label_reads import next_token_label_reads from vllm.inputs import EngineInput, tokens_input from vllm.lora.request import LoRARequest +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py index 4f30beb1a81f..ce68890717e3 100644 --- a/vllm/entrypoints/launchers/api_server/app_state.py +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -7,7 +7,6 @@ from starlette.datastructures import State from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launchers.cli_args import resolve_default_chat_template_kwargs from vllm.entrypoints.mcp.tool_server import init_tool_server from vllm.entrypoints.openai.models.protocol import BaseModelPath @@ -16,6 +15,7 @@ from vllm.entrypoints.serve.utils.api_utils import process_lora_modules from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tasks import FALLBACK_SUPPORTED_TASKS, POOLING_TASKS, SupportedTask diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 28fe79fd691e..5cd58da87eac 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -17,12 +17,12 @@ from vllm.entrypoints.anthropic.protocol import ( AnthropicDisabledThinkingEffortOption, ) -from vllm.entrypoints.chat_utils import ( +from vllm.entrypoints.openai.models.protocol import LoRAModulePath +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatTemplateContentFormatOption, validate_chat_template, ) -from vllm.entrypoints.openai.models.protocol import LoRAModulePath -from vllm.logger import init_logger from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.tool_strict_level import ToolStrictLevelName from vllm.utils.argparse_utils import FlexibleArgumentParser diff --git a/vllm/entrypoints/launchers/render/app_state.py b/vllm/entrypoints/launchers/render/app_state.py index 33aebb9bc39e..d25263f56b5a 100644 --- a/vllm/entrypoints/launchers/render/app_state.py +++ b/vllm/entrypoints/launchers/render/app_state.py @@ -5,7 +5,6 @@ from starlette.datastructures import State from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launchers.cli_args import resolve_default_chat_template_kwargs from vllm.entrypoints.mcp.tool_server import init_tool_server from vllm.entrypoints.openai.models.protocol import BaseModelPath @@ -15,6 +14,7 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.plugins.endpoint_plugins.interface import init_endpoint_plugins_state from vllm.renderers import renderer_from_config +from vllm.renderers.chat_utils import load_chat_template from vllm.renderers.online_derenderer import OnlineDerenderer from vllm.renderers.online_renderer import OnlineRenderer diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index a879b614844e..ff7bd9c57402 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -28,11 +28,6 @@ ) from vllm.config.quantization import QuantizationConfigArgs from vllm.engine.arg_utils import EngineArgs -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - load_chat_template, -) from vllm.entrypoints.common.offline import _O, _R, OfflineInferenceMixin from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin @@ -44,6 +39,11 @@ from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.outputs import PoolingRequestOutput, RequestOutput from vllm.platforms import current_platform +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + load_chat_template, +) from vllm.sampling_params import SamplingParams from vllm.tokenizers import TokenizerLike from vllm.usage.usage_lib import UsageContext diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index d3ea1a1ac3de..fe83b9ddba84 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -8,7 +8,6 @@ from fastapi import Request -from vllm.entrypoints.chat_utils import ConversationMessage from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, @@ -24,6 +23,7 @@ from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ConversationMessage from vllm.tokenizers import TokenizerLike from vllm.utils.async_utils import merge_async_iterators from vllm.utils.collection_utils import as_list diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index dee8f283cf68..d2e09851eba1 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -19,11 +19,6 @@ ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - has_non_text_content, -) from vllm.entrypoints.generate.base.protocol import ( AnyResponseFormat, DeltaMessage, @@ -44,6 +39,11 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + has_non_text_content, +) from vllm.sampling_params import ( BeamSearchParams, RepetitionDetectionParams, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 69b46ba67f61..b291ec6e8aae 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -11,11 +11,6 @@ from fastapi import Request from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, - ConversationMessage, - make_tool_call_id, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, FunctionCall, @@ -61,6 +56,11 @@ from vllm.outputs import RequestOutput from vllm.parser import ParserManager from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, + ConversationMessage, + make_tool_call_id, +) from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index a68df6dd4d65..334e8032e283 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -21,9 +21,6 @@ from openai_harmony import Author, Message, Role, TextContent from vllm import envs -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, -) from vllm.entrypoints.mcp.tool import Tool from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.parser.harmony_utils import render_for_completion @@ -41,6 +38,9 @@ ) from vllm.outputs import RequestOutput from vllm.parser.abstract_parser import Parser +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, +) from vllm.tokenizers import TokenizerLike from vllm.utils import random_uuid diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 406ab2d0818f..373cb35df057 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -54,10 +54,6 @@ ) from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import ( PerRequestMetrics, StopParam, @@ -68,6 +64,10 @@ from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.sampling_params import ( RequestOutputKind, SamplingParams, diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 3a417ca7a566..ba4aa0434768 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -23,9 +23,6 @@ from vllm import envs from vllm.config.utils import replace from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ( - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, RequestResponseMetadata, @@ -86,6 +83,9 @@ from vllm.outputs import CompletionOutput from vllm.parser import Parser, ParserManager from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, +) from vllm.renderers.online_renderer import ( OnlineRenderer, ResponsesPreviousMessages, diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 2f5b151a1b35..b101040a57a2 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -62,7 +62,6 @@ from openai.types.responses.tool import Tool from openai_harmony import Message as HarmonyMessage -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import DeltaMessage, DeltaToolCall from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.parser.harmony_utils import ( @@ -78,6 +77,7 @@ ) from vllm.outputs import CompletionOutput from vllm.parser.harmony import Segment +from vllm.renderers.chat_utils import make_tool_call_id from vllm.utils import random_uuid TOOL_NAME_TO_MCP_SERVER_LABEL: Final[dict[str, str]] = { diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 6c8507840fdf..c4523964e9dd 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -30,7 +30,6 @@ from openai.types.responses.tool import Tool from vllm import envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import FunctionCall, FunctionDefinition from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionMessageParam, @@ -39,6 +38,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponseInputOutputItem from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.utils import ( build_responses_tool_call_name_map, flat_namespace_tool_name, diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 0d9431f022ef..f79841b5d9ee 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -10,12 +10,12 @@ PoolingRequestOutput, ) from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( - ChatTemplateConfig, -) from vllm.exceptions import VLLMValidationError from vllm.lora.request import LoRARequest from vllm.renderers import BaseRenderer, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatTemplateConfig, +) from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq from vllm.utils.async_utils import make_async from vllm.utils.mistral import is_mistral_tokenizer diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index 0c646b969ab5..c6d11488ac23 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -7,14 +7,14 @@ from pydantic import Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.exceptions import VLLMValidationError from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) from vllm.tasks import check_removed_pooling_task from vllm.utils import random_uuid from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 2fb0b8e1bc72..59d4051aa57b 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -14,7 +14,6 @@ from vllm import PoolingRequestOutput, envs from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.serving import BaseServing @@ -23,6 +22,7 @@ from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.utils.async_utils import make_async, merge_async_iterators from ...serve.engine.protocol import ErrorResponse diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 1b23a77ec60d..94f84f09b678 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -12,15 +12,15 @@ from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams -from vllm.entrypoints.chat_utils import ( - ChatCompletionContentPartParam, - ChatCompletionMessageParam, - CustomChatCompletionMessageParam, -) from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput from vllm.renderers import merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionContentPartParam, + ChatCompletionMessageParam, + CustomChatCompletionMessageParam, +) from vllm.renderers.hf import resolve_chat_template from vllm.utils.collection_utils import chunk_list from vllm.utils.mistral import is_mistral_tokenizer diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index a853a8ea00ee..d52871d6a6f7 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -17,8 +17,8 @@ from pydantic import BaseModel, BeforeValidator, Field, model_validator from vllm import PoolingParams -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.utils import random_uuid from ..base.protocol import ( diff --git a/vllm/entrypoints/pooling/factories.py b/vllm/entrypoints/pooling/factories.py index e6d01be5a9e9..52c7e6f42b7e 100644 --- a/vllm/entrypoints/pooling/factories.py +++ b/vllm/entrypoints/pooling/factories.py @@ -6,10 +6,10 @@ from fastapi import FastAPI from vllm.config import ModelConfig, VllmConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.logger import init_logger from vllm.plugins.io_processors import has_io_processor from vllm.renderers import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask from .base.io_processor import PoolingIOProcessor @@ -151,7 +151,7 @@ def init_pooling_state( if model_config is None: return - from vllm.entrypoints.chat_utils import load_chat_template + from vllm.renderers.chat_utils import load_chat_template from vllm.tasks import POOLING_TASKS from .classify.serving import ServingClassification diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 81d78101c429..19be1eddde91 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -6,7 +6,6 @@ from tqdm.auto import tqdm -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.common.offline import OfflineInferenceMixin from vllm.inputs import DataPrompt, PromptType from vllm.logger import init_logger @@ -18,6 +17,7 @@ ScoringRequestOutput, ) from vllm.pooling_params import PoolingParams +from vllm.renderers.chat_utils import ChatTemplateConfig from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask from .base.io_processor import PoolingIOProcessor diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 91c5f762cf2c..1e4c638968f8 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -9,6 +9,7 @@ from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.logger import init_logger from vllm.renderers import TokenizeParams +from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, @@ -18,7 +19,6 @@ from vllm.tasks import PoolingTask from vllm.utils.mistral import is_mistral_tokenizer -from ...chat_utils import ChatTemplateResolutionError from ..base.io_processor import PoolingIOProcessor from ..pooling.protocol import PoolingCompletionRequest from ..typing import ( diff --git a/vllm/entrypoints/pooling/scoring/typing.py b/vllm/entrypoints/pooling/scoring/typing.py index 9afaada6d666..9fa059ce8ea6 100644 --- a/vllm/entrypoints/pooling/scoring/typing.py +++ b/vllm/entrypoints/pooling/scoring/typing.py @@ -5,7 +5,7 @@ from typing_extensions import TypedDict -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionContentPartImageEmbedsParam, ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index d7cc888a1556..8d953cf6ddd4 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -7,7 +7,9 @@ from vllm import PromptType, TextPrompt from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.renderers.chat_utils import ( BaseMultiModalItemTracker, ChatCompletionContentPartParam, ChatCompletionContentPartTextParam, @@ -15,8 +17,6 @@ MultiModalItemTracker, _parse_chat_message_content_parts, ) -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from .typing import ( ScoreContentPartParam, diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 8e45ce33f4b2..a299fc657d54 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -9,10 +9,10 @@ from pydantic import ConfigDict from vllm import PoolingParams, PoolingRequestOutput, PromptType -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest from vllm.renderers import ChatParams, TokenizeParams +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.renderers.inputs import DictPrompt from ...tasks import PoolingTask diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index b6d7bf19ba93..816d70bec000 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -13,7 +13,6 @@ from fastapi import Request from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker from vllm.entrypoints.generate.base.protocol import ( PerRequestMetrics, RequestResponseMetadata, @@ -47,6 +46,7 @@ PlaceholderRange, ) from vllm.outputs import RequestOutput +from vllm.renderers.chat_utils import AsyncMultiModalItemTracker from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike diff --git a/vllm/entrypoints/serve/engine/typing.py b/vllm/entrypoints/serve/engine/typing.py index a6878d500150..59e5642c5784 100644 --- a/vllm/entrypoints/serve/engine/typing.py +++ b/vllm/entrypoints/serve/engine/typing.py @@ -4,7 +4,6 @@ from typing import Protocol, TypeAlias from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, ChatCompletionRequest, @@ -35,6 +34,7 @@ ) from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest from vllm.renderers import ChatParams, TokenizeParams +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption class RendererRequest(Protocol): diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py index 130a1b4b45c4..261f7075f4bf 100644 --- a/vllm/entrypoints/serve/tokenize/protocol.py +++ b/vllm/entrypoints/serve/tokenize/protocol.py @@ -7,14 +7,14 @@ from pydantic import ConfigDict, Field, model_validator from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, -) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.exceptions import VLLMValidationError from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, +) class TokenizeCompletionRequest(OpenAIBaseModel): diff --git a/vllm/entrypoints/serve/tokenize/serving.py b/vllm/entrypoints/serve/tokenize/serving.py index c2232324ce2b..b97737abc938 100644 --- a/vllm/entrypoints/serve/tokenize/serving.py +++ b/vllm/entrypoints/serve/tokenize/serving.py @@ -5,7 +5,6 @@ from fastapi import Request -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.openai.models.serving import ( OpenAIModelRegistry, OpenAIServingModels, @@ -23,6 +22,7 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import TokensPrompt, tokens_input from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers import TokenizerLike diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 4af133c9882c..61826dc8c86c 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -21,10 +21,6 @@ TriggeredTagsFormat, ) -from vllm.entrypoints.chat_utils import ( - get_tool_call_id_type, - make_tool_call_id, -) from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ExtractedToolCallInformation, @@ -42,6 +38,10 @@ from vllm.parser.metrics import record_tool_parser_invocation from vllm.parser.utils import count_history_tool_calls from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.renderers.chat_utils import ( + get_tool_call_id_type, + make_tool_call_id, +) from vllm.sampling_params import ( SamplingParams, StructuredOutputsParams, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index b04fdaa7944d..c74be931d256 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -13,7 +13,6 @@ import regex as re -from vllm.entrypoints.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -27,6 +26,7 @@ from vllm.parser.engine.events import EventType, SemanticEvent from vllm.parser.engine.parser_engine_config import ParserEngineConfig, ParserState from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine +from vllm.renderers.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.tool_parsers.utils import ( coerce_to_schema_type, extract_types_from_schema, diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 68b50250defb..1cc8ccd7ffec 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -23,7 +23,6 @@ TagFormat, ) -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -43,6 +42,7 @@ from vllm.logger import init_logger from vllm.parser.abstract_parser import DelegatingParser, structured_outputs_to_format from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser from vllm.tool_parsers.structural_tag_registry import ( SimplifiedToolChoice, diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index 985b76308a33..85f8b67efbd6 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -11,7 +11,6 @@ from openai.types.responses import ToolChoiceFunction -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -30,6 +29,7 @@ load_response_template, ) from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.utils import iter_response_function_tool_dicts diff --git a/vllm/parser/utils.py b/vllm/parser/utils.py index 9194b7ef1821..0deb227c0bbb 100644 --- a/vllm/parser/utils.py +++ b/vllm/parser/utils.py @@ -5,12 +5,12 @@ from openai.types.responses import ResponseFunctionToolCall -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ( ResponseInputOutputItem, ResponsesRequest, ) +from vllm.renderers.chat_utils import ChatCompletionMessageParam def count_tool_calls(tool_calls: object) -> int: diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 8c1d91e032b1..e04bfa64eb76 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -66,7 +66,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.entrypoints.chat_utils import ( + from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, ) @@ -362,7 +362,7 @@ def warmup(self, chat_params: ChatParams) -> None: no-op if the background warmup — or a prior reset_mm_cache that joined it — already ran it. """ - from vllm.entrypoints.chat_utils import ChatTemplateResolutionError + from vllm.renderers.chat_utils import ChatTemplateResolutionError # prevent MM processor hangs with set_default_torch_num_threads(1): diff --git a/vllm/entrypoints/chat_utils.py b/vllm/renderers/chat_utils.py similarity index 100% rename from vllm/entrypoints/chat_utils.py rename to vllm/renderers/chat_utils.py diff --git a/vllm/renderers/cohere.py b/vllm/renderers/cohere.py index e22ebf1ebf1b..eb7aeffaac31 100644 --- a/vllm/renderers/cohere.py +++ b/vllm/renderers/cohere.py @@ -94,13 +94,13 @@ from typing import Any from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async @@ -165,7 +165,7 @@ class MelodyContentType(str, Enum): # melody-shape content list the renderer emits. # # Motivation: OpenAI's ``ChatCompletionContentPartParam`` union (in -# ``vllm/entrypoints/chat_utils.py``) only allows ``text`` / ``image_url`` +# ``vllm/renderers/chat_utils.py``) only allows ``text`` / ``image_url`` # / ``audio`` / ``video`` / ``file`` content parts on tool messages, so # ``{"type": "document", "document": {...}}`` blocks -- valid in Cohere # v2 -- get rejected by vLLM's shared chat validator before they can diff --git a/vllm/renderers/deepseek_v32.py b/vllm/renderers/deepseek_v32.py index ba2e0d60c1a1..6e4f9fd91bec 100644 --- a/vllm/renderers/deepseek_v32.py +++ b/vllm/renderers/deepseek_v32.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, diff --git a/vllm/renderers/deepseek_v4.py b/vllm/renderers/deepseek_v4.py index 470e21f1eb11..784764bc8b50 100644 --- a/vllm/renderers/deepseek_v4.py +++ b/vllm/renderers/deepseek_v4.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 4a8bcd0c03d7..ec5b2d96ff7f 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -22,13 +22,6 @@ from typing_extensions import override from vllm import envs -from vllm.entrypoints.chat_utils import ( - PROMPT_EMBEDS_PLACEHOLDER_TOKEN, - ChatTemplateResolutionError, - load_chat_template, - parse_chat_messages, - parse_chat_messages_async, -) from vllm.exceptions import VLLMValidationError from vllm.inputs import EmbedsPrompt from vllm.inputs.engine import MultiModalInput @@ -46,6 +39,13 @@ apply_token_matches, find_mm_placeholders, ) +from vllm.renderers.chat_utils import ( + PROMPT_EMBEDS_PLACEHOLDER_TOKEN, + ChatTemplateResolutionError, + load_chat_template, + parse_chat_messages, + parse_chat_messages_async, +) from vllm.tokenizers.hf import HfTokenizer, maybe_make_thread_pool from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path from vllm.transformers_utils.processor import cached_get_processor @@ -59,18 +59,18 @@ from collections.abc import Set from vllm.config import ModelConfig, VllmConfig - from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormat, - ChatTemplateContentFormatOption, - ConversationMessage, - ) from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict, TokensPrompt from vllm.inputs.engine import TokensInput from vllm.multimodal.processing.processor import ( MultiModalPromptUpdates, ResolvedPromptUpdate, ) + from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormat, + ChatTemplateContentFormatOption, + ConversationMessage, + ) from .inputs import DictPrompt from .params import ChatParams diff --git a/vllm/renderers/inkling.py b/vllm/renderers/inkling.py index 87eedfd69de4..6ca82e309912 100644 --- a/vllm/renderers/inkling.py +++ b/vllm/renderers/inkling.py @@ -14,13 +14,13 @@ """ from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/inputs/preprocess.py b/vllm/renderers/inputs/preprocess.py index 82a49f5638cb..9daf38fdbd3b 100644 --- a/vllm/renderers/inputs/preprocess.py +++ b/vllm/renderers/inputs/preprocess.py @@ -21,7 +21,7 @@ import torch from vllm.config import ModelConfig - from vllm.entrypoints.chat_utils import ChatCompletionMessageParam + from vllm.renderers.chat_utils import ChatCompletionMessageParam @overload diff --git a/vllm/renderers/kimi_k3.py b/vllm/renderers/kimi_k3.py index e0f937f2e6ce..0d6e14d072a4 100644 --- a/vllm/renderers/kimi_k3.py +++ b/vllm/renderers/kimi_k3.py @@ -3,14 +3,14 @@ from typing import Any, cast from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.exceptions import VLLMValidationError +from vllm.multimodal.media.connector import merge_media_io_kwargs +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.exceptions import VLLMValidationError -from vllm.multimodal.media.connector import merge_media_io_kwargs from vllm.tokenizers.hf import HfTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/mistral.py b/vllm/renderers/mistral.py index da2275afc6c1..712f2e90dcc8 100644 --- a/vllm/renderers/mistral.py +++ b/vllm/renderers/mistral.py @@ -2,13 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.config import VllmConfig -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.mistral import MistralTokenizer from vllm.utils.async_utils import make_async diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 66b73fb4912c..8e3e4ca25abc 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -4,7 +4,6 @@ from typing import Any from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ToolCall, @@ -42,6 +41,7 @@ from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer +from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.tokenizers import TokenizerLike from vllm.tokenizers.detokenizer_utils import ( convert_prompt_ids_to_tokens, diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index e8220938b26a..b4e3fc6c4666 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -12,12 +12,6 @@ from openai_harmony import ToolNamespaceConfig from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ( - ChatCompletionMessageParam, - ChatTemplateContentFormatOption, - ConversationMessage, - has_non_text_content, -) from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -59,6 +53,12 @@ from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer, ChatParams, TokenizeParams, merge_kwargs +from vllm.renderers.chat_utils import ( + ChatCompletionMessageParam, + ChatTemplateContentFormatOption, + ConversationMessage, + has_non_text_content, +) from vllm.renderers.inputs.preprocess import ( parse_model_prompt, prompt_to_seq, diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 456a4d0ac563..5e4cc58d3d73 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: import torch - from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption + from vllm.renderers.chat_utils import ChatTemplateContentFormatOption else: torch = LazyLoader("torch", globals(), "torch") diff --git a/vllm/renderers/terratorch.py b/vllm/renderers/terratorch.py index 68b4936a352e..5e188bf70757 100644 --- a/vllm/renderers/terratorch.py +++ b/vllm/renderers/terratorch.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.chat_utils import ( +from vllm.logger import init_logger +from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, ConversationMessage, parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from .base import BaseRenderer from .inputs import DictPrompt diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index 3f0fff938ad9..f933e74cf5af 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -5,7 +5,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v32_encoding import encode_messages from .hf import HfTokenizer, get_cached_tokenizer diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index 532ca0df7c9b..471d6fbb4ac9 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -5,7 +5,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v4_encoding import encode_messages from .hf import HfTokenizer, get_cached_tokenizer diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index a1449e59e203..7ff1f5cc3b83 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -6,7 +6,7 @@ from transformers import TokenizersBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.renderers.chat_utils import ChatCompletionMessageParam from .deepseek_v41_encoding import ( IMAGE_PLACEHOLDER, diff --git a/vllm/tokenizers/kimi_audio.py b/vllm/tokenizers/kimi_audio.py index 10f57a7f7ded..32de7abf056b 100644 --- a/vllm/tokenizers/kimi_audio.py +++ b/vllm/tokenizers/kimi_audio.py @@ -13,8 +13,8 @@ from transformers import AddedToken, BatchEncoding from transformers.utils import chat_template_utils as hf_chat_utils -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.protocol import TokenizerLike from vllm.transformers_utils.repo_utils import hf_api diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 4ffefe0db6d9..15cf0322581a 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -32,10 +32,10 @@ from pydantic import ValidationError from transformers.tokenization_mistral_common import MistralCommonBackend -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger +from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.tokenizers.protocol import TokenizerLike if TYPE_CHECKING: diff --git a/vllm/tokenizers/protocol.py b/vllm/tokenizers/protocol.py index 74b32e60d603..b477dd1986e9 100644 --- a/vllm/tokenizers/protocol.py +++ b/vllm/tokenizers/protocol.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from transformers import BatchEncoding - from vllm.entrypoints.chat_utils import ChatCompletionMessageParam + from vllm.renderers.chat_utils import ChatCompletionMessageParam class TokenizerLike(Protocol): diff --git a/vllm/tool_parsers/apertus_tool_parser.py b/vllm/tool_parsers/apertus_tool_parser.py index 6825b3ef39e1..af9654c0cb26 100644 --- a/vllm/tool_parsers/apertus_tool_parser.py +++ b/vllm/tool_parsers/apertus_tool_parser.py @@ -14,7 +14,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -26,6 +25,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/deepseekv31_tool_parser.py b/vllm/tool_parsers/deepseekv31_tool_parser.py index f4c06d467fab..ff5f605f511f 100644 --- a/vllm/tool_parsers/deepseekv31_tool_parser.py +++ b/vllm/tool_parsers/deepseekv31_tool_parser.py @@ -5,7 +5,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -18,6 +17,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/deepseekv3_tool_parser.py b/vllm/tool_parsers/deepseekv3_tool_parser.py index 13fbcb5ed26d..649eded2878f 100644 --- a/vllm/tool_parsers/deepseekv3_tool_parser.py +++ b/vllm/tool_parsers/deepseekv3_tool_parser.py @@ -5,7 +5,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -18,6 +17,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/dots_tool_parser.py b/vllm/tool_parsers/dots_tool_parser.py index 3e8134473939..311e545365f4 100644 --- a/vllm/tool_parsers/dots_tool_parser.py +++ b/vllm/tool_parsers/dots_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ChatCompletionToolsParam, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.utils import is_complete_json, partial_tag_overlap diff --git a/vllm/tool_parsers/functiongemma_tool_parser.py b/vllm/tool_parsers/functiongemma_tool_parser.py index 708ecb1dd997..6fcae6e55220 100644 --- a/vllm/tool_parsers/functiongemma_tool_parser.py +++ b/vllm/tool_parsers/functiongemma_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/gigachat3_tool_parser.py b/vllm/tool_parsers/gigachat3_tool_parser.py index 020d4c1af700..7f9bdf86d6be 100644 --- a/vllm/tool_parsers/gigachat3_tool_parser.py +++ b/vllm/tool_parsers/gigachat3_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser diff --git a/vllm/tool_parsers/granite4_tool_parser.py b/vllm/tool_parsers/granite4_tool_parser.py index 913135f13d70..aabd455c6766 100644 --- a/vllm/tool_parsers/granite4_tool_parser.py +++ b/vllm/tool_parsers/granite4_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/granite_20b_fc_tool_parser.py b/vllm/tool_parsers/granite_20b_fc_tool_parser.py index 4c6c0a1b0efe..8b2ccf58e86b 100644 --- a/vllm/tool_parsers/granite_20b_fc_tool_parser.py +++ b/vllm/tool_parsers/granite_20b_fc_tool_parser.py @@ -9,7 +9,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hermes_tool_parser.py b/vllm/tool_parsers/hermes_tool_parser.py index c5b318cd19fc..1fa19d0321fc 100644 --- a/vllm/tool_parsers/hermes_tool_parser.py +++ b/vllm/tool_parsers/hermes_tool_parser.py @@ -6,7 +6,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -20,6 +19,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py index 0effb377d0a6..9947294ecd52 100644 --- a/vllm/tool_parsers/hy_v3_tool_parser.py +++ b/vllm/tool_parsers/hy_v3_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ChatCompletionToolsParam, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py index 2ce0381a0132..de707c00228e 100644 --- a/vllm/tool_parsers/hy_v4_tool_parser.py +++ b/vllm/tool_parsers/hy_v4_tool_parser.py @@ -11,7 +11,6 @@ import regex as re import vllm.envs as envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -23,6 +22,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.tool_strict_level import ToolStrictLevel diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index 709fd6c8f2a8..9bec016ec907 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -7,7 +7,6 @@ import partial_json_parser from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -21,6 +20,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/jamba_tool_parser.py b/vllm/tool_parsers/jamba_tool_parser.py index 1ba67960b961..83eea1d5800a 100644 --- a/vllm/tool_parsers/jamba_tool_parser.py +++ b/vllm/tool_parsers/jamba_tool_parser.py @@ -8,7 +8,6 @@ import regex as re from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json diff --git a/vllm/tool_parsers/llama_tool_parser.py b/vllm/tool_parsers/llama_tool_parser.py index 93b0980c3775..8696af46994d 100644 --- a/vllm/tool_parsers/llama_tool_parser.py +++ b/vllm/tool_parsers/llama_tool_parser.py @@ -10,7 +10,6 @@ from transformers import PreTrainedTokenizerBase import vllm.envs as envs -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -23,6 +22,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/minicpm5xml_tool_parser.py b/vllm/tool_parsers/minicpm5xml_tool_parser.py index 8b419d36df7f..818b51c47f4f 100644 --- a/vllm/tool_parsers/minicpm5xml_tool_parser.py +++ b/vllm/tool_parsers/minicpm5xml_tool_parser.py @@ -7,7 +7,6 @@ import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -22,6 +21,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/muse_glimmer_tool_parser.py b/vllm/tool_parsers/muse_glimmer_tool_parser.py index 623b3540d7d3..5a0b26cf423a 100644 --- a/vllm/tool_parsers/muse_glimmer_tool_parser.py +++ b/vllm/tool_parsers/muse_glimmer_tool_parser.py @@ -41,7 +41,6 @@ from openai.types.responses import ToolChoiceFunction from transformers import PreTrainedTokenizerBase -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -56,6 +55,7 @@ ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py index 820e797f4d9f..0ef25b35e0f4 100644 --- a/vllm/tool_parsers/phi4mini_tool_parser.py +++ b/vllm/tool_parsers/phi4mini_tool_parser.py @@ -8,7 +8,6 @@ import regex as re from transformers import PreTrainedTokenizerBase -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaMessage, ExtractedToolCallInformation, @@ -19,6 +18,7 @@ ChatCompletionRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index 178036760634..d2dfb7815fbe 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -19,7 +19,6 @@ from openai.types.responses import ToolChoiceFunction from partial_json_parser.core.options import Allow -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -36,6 +35,7 @@ ResponsesRequest, ) from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 07c1a2e031eb..bd5f9a7f3aa0 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -5,12 +5,12 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, ) +from vllm.renderers.chat_utils import make_tool_call_id from vllm.utils.mistral import is_mistral_tokenizer if TYPE_CHECKING: diff --git a/vllm/tool_parsers/xlam_tool_parser.py b/vllm/tool_parsers/xlam_tool_parser.py index 52dc3063556b..dec34fc0b9a1 100644 --- a/vllm/tool_parsers/xlam_tool_parser.py +++ b/vllm/tool_parsers/xlam_tool_parser.py @@ -10,7 +10,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) -from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.renderers.chat_utils import make_tool_call_id from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, DeltaMessage, From a98247ab4db686ee03c66d5feb3c761e52a2f8ab Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Sat, 10 Oct 2026 15:51:52 +0900 Subject: [PATCH 348/952] [Model] Drop the 4x M-RoPE cache headroom for Qwen3-VL, Qwen3.5 and Qwen3.8 (#59853) Signed-off-by: lesj0610 Signed-off-by: lesj0610 Co-authored-by: lesj0610 --- vllm/config/model.py | 4 ++++ .../layers/rotary_embedding/__init__.py | 4 ++++ .../layers/rotary_embedding/mrope.py | 8 ++++--- vllm/model_executor/models/qwen3.py | 10 ++++++++- vllm/model_executor/models/qwen3_5.py | 2 ++ vllm/model_executor/models/qwen3_moe.py | 4 ++++ vllm/model_executor/models/qwen3_next.py | 2 ++ vllm/model_executor/models/qwen3_vl.py | 22 +++++++++++++++++-- vllm/model_executor/models/qwen3_vl_moe.py | 2 ++ vllm/models/qwen4_exp/amd/qsa.py | 2 ++ vllm/models/qwen4_exp/nvidia/qsa.py | 2 ++ 11 files changed, 56 insertions(+), 6 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index cfbb870e7887..502f1bbd8d4f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -508,6 +508,10 @@ def compute_hash(self) -> str: # here early. if self.multimodal_config: factors["language_model_only"] = self.multimodal_config.language_model_only + # Sizes Qwen3.5's M-RoPE cache. + factors["video_pruning"] = ( + self.multimodal_config.is_multimodal_pruning_enabled() + ) return hash_factors(factors) def _update_nested( diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index e23a5128b8a1..85d4cb63e501 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -36,6 +36,7 @@ def get_rope( rope_parameters: dict[str, Any] | None = None, dtype: torch.dtype | None = None, dual_chunk_attention_config: dict[str, Any] | None = None, + mrope_positions_factor: int = 4, ) -> RotaryEmbedding: if dtype is None: dtype = torch.get_default_dtype() @@ -78,6 +79,7 @@ def get_rope( rope_parameters_args, dual_chunk_attention_args, dtype, + mrope_positions_factor, ) if key in _ROPE_DICT: return _ROPE_DICT[key] @@ -118,6 +120,7 @@ def get_rope( dtype, mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), + positions_factor=mrope_positions_factor, ) elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: extra_kwargs = { @@ -263,6 +266,7 @@ def get_rope( dtype, mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), + positions_factor=mrope_positions_factor, scaling_factor=scaling_factor, **extra_kwargs, ) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 11882856e049..9bcc5a95e16d 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -315,8 +315,9 @@ def __init__( dtype: torch.dtype, mrope_section: list[int] | None = None, mrope_interleaved: bool = False, - # YaRN parameters. *, + positions_factor: int = 4, + # YaRN parameters. scaling_factor: float | None = None, beta_fast: int = 32, beta_slow: int = 1, @@ -345,8 +346,9 @@ def __init__( # In Qwen2.5-VL, the maximum index value is related to the duration of # the input video. We enlarge max_position_embeddings to 4 times to get - # a larger the cos and sin cache. - self.cache_max_position_num = max_position_embeddings * 4 + # a larger the cos and sin cache. Models whose positions never exceed + # the token count pass `positions_factor=1`. + self.cache_max_position_num = max_position_embeddings * positions_factor super().__init__( head_size, rotary_dim, diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 7a1d76234f32..31d1a094ebef 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -79,6 +79,7 @@ def __init__( attn_type: str = AttentionType.DECODER, dual_chunk_attention_config: dict[str, Any] | None = None, per_layer_sliding_window: int | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -124,6 +125,7 @@ def __init__( max_position=max_position, rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) attn_cls = ( EncoderOnlyAttention @@ -179,6 +181,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", per_layer_sliding_window: int | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -211,6 +214,7 @@ def __init__( attn_type=attn_type, dual_chunk_attention_config=dual_chunk_attention_config, per_layer_sliding_window=per_layer_sliding_window, + mrope_positions_factor=mrope_positions_factor, ) self.mlp = Qwen3MLP( hidden_size=self.hidden_size, @@ -263,9 +267,13 @@ def forward( } ) class Qwen3Model(Qwen2Model): + decoder_layer_type: type[nn.Module] = Qwen3DecoderLayer + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__( - vllm_config=vllm_config, prefix=prefix, decoder_layer_type=Qwen3DecoderLayer + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=self.decoder_layer_type, ) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 0d110dd58be1..f105948f9283 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -91,6 +91,7 @@ Qwen3VLForConditionalGeneration, Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, + mrope_positions_factor, ) from .utils import ( AutoWeightsLoader, @@ -161,6 +162,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", reduce_results=not self.use_attn_reduce_scatter_for_moe, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) else: raise ValueError(f"Invalid layer_type {self.layer_type}") diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 3f293550ac9a..870a082d3a72 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -256,6 +256,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", dual_chunk_attention_config: dict[str, Any] | None = None, + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -303,6 +304,7 @@ def __init__( max_position=max_position_embeddings, rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) attention_kwargs: dict[str, Any] = {} if dual_chunk_attention_config: @@ -351,6 +353,7 @@ def __init__( vllm_config: VllmConfig, prefix: str = "", is_fused_checkpoint_transposed: bool = False, + mrope_positions_factor: int = 4, ) -> None: super().__init__() @@ -376,6 +379,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", dual_chunk_attention_config=dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) # `mlp_only_layers` in the config. diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b0bf0fecd19c..04d20c574145 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -280,6 +280,7 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, prefix: str = "", + mrope_positions_factor: int = 4, ) -> None: super().__init__() self.config = config @@ -331,6 +332,7 @@ def __init__( max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, dual_chunk_attention_config=self.dual_chunk_attention_config, + mrope_positions_factor=mrope_positions_factor, ) # Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 6648949c7364..77780b3d1d40 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -52,7 +52,7 @@ should_torch_compile_mm_encoder, support_torch_compile, ) -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.multimodal import ( MultiModalConfig, MultiModalDummyOptions, @@ -148,7 +148,7 @@ Qwen2VLProcessingInfo, _create_qwen2vl_field_factory, ) -from .qwen3 import Qwen3ForCausalLM, Qwen3Model +from .qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM, Qwen3Model from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -1748,6 +1748,22 @@ def get_video_repl( return PromptUpdateDetails.from_seq(all_token_ids) +def mrope_positions_factor(vllm_config: VllmConfig) -> int: + # Video pruning keeps the positions of the tokens it drops. A drafter gets + # the target's positions, so it follows the target. + spec = vllm_config.speculative_config + model_config = (spec and spec.target_model_config) or vllm_config.model_config + mm_config = model_config.multimodal_config + pruned = mm_config is not None and mm_config.is_multimodal_pruning_enabled() + return 4 if pruned else 1 + + +class Qwen3VLDecoderLayer(Qwen3DecoderLayer): + def __init__(self, **kwargs) -> None: + factor = mrope_positions_factor(get_current_vllm_config()) + super().__init__(**kwargs, mrope_positions_factor=factor) + + @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, @@ -1761,6 +1777,8 @@ def get_video_repl( } ) class Qwen3LLMModel(Qwen3Model): + decoder_layer_type = Qwen3VLDecoderLayer + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 4d0c7788a6e3..9a1d771ec4e1 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -56,6 +56,7 @@ Qwen3VLForConditionalGeneration, Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, + mrope_positions_factor, ) from .utils import StageMissingLayer, maybe_prefix @@ -73,6 +74,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: vllm_config=vllm_config, prefix=prefix, is_fused_checkpoint_transposed=True, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index ff01490bb190..709fd0b08ef1 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.model_executor.models.qwen3_vl import mrope_positions_factor from vllm.platforms import current_platform from vllm.utils.torch_utils import ( LayerNameType, @@ -248,6 +249,7 @@ def __init__( head_size=self.head_dim, max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 0246db80943d..9f0934d5a0c6 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -29,6 +29,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.model_executor.models.qwen3_vl import mrope_positions_factor from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -427,6 +428,7 @@ def __init__( head_size=self.head_dim, max_position=config.max_position_embeddings, rope_parameters=config.rope_parameters, + mrope_positions_factor=mrope_positions_factor(vllm_config), ) self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) From c41b2639e29c3bc01add1d34bef3032a6d9d8aca Mon Sep 17 00:00:00 2001 From: Kay Yan Date: Sat, 10 Oct 2026 17:08:36 +0800 Subject: [PATCH 349/952] [Bugfix][Responses API] Complete cache-write token accounting (#60973) Signed-off-by: Kay Yan Co-authored-by: Codex --- .../responses/test_parsable_context_unit.py | 10 ++++++++++ tests/entrypoints/unit_tests/test_context.py | 19 ++++++++++++++++++- vllm/entrypoints/openai/responses/context.py | 5 +++++ vllm/entrypoints/openai/responses/serving.py | 2 +- 4 files changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py index 710cd2d04cc9..f3927c9d7152 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py +++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py @@ -169,6 +169,7 @@ def _make_request_output( text: str = "Hello, world!", token_ids: Sequence[int] = (1, 2, 3), finish_reason: str = "stop", + num_cache_creation_tokens: int | None = None, ) -> RequestOutput: return RequestOutput( request_id="test", @@ -186,6 +187,7 @@ def _make_request_output( ) ], finished=True, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -425,6 +427,14 @@ def test_reasoning_tokens_counted_per_round(): assert ctx.num_reasoning_tokens == 3 +def test_cache_creation_tokens_counted_per_round(): + ctx = _make_context(None) + + for count, expected in [(None, 0), (0, 0), (3, 3), (4, 7)]: + ctx.append_output(_make_request_output(num_cache_creation_tokens=count)) + assert ctx.num_cache_creation_tokens == expected + + def test_reasoning_tokens_zero_without_parser(): ctx = _make_context(None) diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index a5b1afe8c50f..f0af86231aaf 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -20,6 +20,7 @@ def create_mock_request_output( output_token_ids=None, num_cached_tokens=0, finished=True, + num_cache_creation_tokens=None, ): """Helper function to create a mock RequestOutput object for testing.""" outputs = [] @@ -44,6 +45,7 @@ def create_mock_request_output( outputs=outputs, finished=finished, num_cached_tokens=num_cached_tokens, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -117,7 +119,8 @@ def make_harmony_context( return context, fake_parser -def test_single_turn_token_counting(): +@pytest.mark.parametrize("num_cache_creation_tokens", [None, 0, 3]) +def test_single_turn_token_counting(num_cache_creation_tokens): """Test token counting behavior for a single turn.""" # Create a context context, _ = make_harmony_context() @@ -127,6 +130,7 @@ def test_single_turn_token_counting(): prompt_token_ids=[1, 2, 3, 4, 5], # 5 prompt tokens output_token_ids=[6, 7, 8], # 3 output tokens num_cached_tokens=2, # 2 cached tokens + num_cache_creation_tokens=num_cache_creation_tokens, ) # Append the output to the context @@ -136,6 +140,7 @@ def test_single_turn_token_counting(): assert context.num_prompt_tokens == 5 assert context.num_output_tokens == 3 assert context.num_cached_tokens == 2 + assert context.num_cache_creation_tokens == (num_cache_creation_tokens or 0) assert context.num_tool_output_tokens == 0 # No tool tokens in first turn # Verify internal state tracking @@ -398,6 +403,7 @@ async def test_streaming_multi_turn_token_counting(): num_prompt_tokens = [3, 8, 13] num_output_tokens = [3, 3, 2] num_cached_tokens = [0, 3, 8] + num_cache_creation_tokens = [3, 4, 2] # Simulate three turns of conversation: # Turn 1: stream tokens one by one, then finish the message @@ -411,6 +417,7 @@ async def test_streaming_multi_turn_token_counting(): prompt_token_ids=[1, 2, 3], # 3 prompt tokens output_token_ids=[101], # Single token num_cached_tokens=num_cached_tokens[0], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, # Not end of message yet ) ) @@ -419,6 +426,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[102], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=False, ) ) @@ -427,6 +435,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[103], + num_cache_creation_tokens=num_cache_creation_tokens[0], finished=True, # End of message ) ) @@ -435,6 +444,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_prompt_tokens == 3 # Initial prompt tokens assert context.num_output_tokens == 3 # Three output tokens assert context.num_cached_tokens == 0 + assert context.num_cache_creation_tokens == 3 assert context.num_tool_output_tokens == 0 # No tool output in first turn assert context.first_tok_of_message is True # Ready for next message @@ -454,6 +464,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 8 tokens (includes previous) output_token_ids=[201], num_cached_tokens=num_cached_tokens[1], # Some tokens cached + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -463,6 +474,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[202], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=False, ) ) @@ -471,6 +483,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[203], + num_cache_creation_tokens=num_cache_creation_tokens[1], finished=True, # End of reasoning message ) ) @@ -480,6 +493,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_output_tokens == 3 + 3 # First turn + second turn assert context.num_reasoning_tokens == 3 # All tokens in analysis channel assert context.num_cached_tokens == 3 # Cached tokens from second turn + assert context.num_cache_creation_tokens == 3 + 4 # Formula: this turn prompt tokens - last turn prompt - last turn output expected_tool_tokens = 8 - 3 - 3 # = 2 @@ -505,6 +519,7 @@ async def test_streaming_multi_turn_token_counting(): ], # 13 tokens output_token_ids=[301], num_cached_tokens=num_cached_tokens[2], # More cached tokens + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=False, ) ) @@ -512,6 +527,7 @@ async def test_streaming_multi_turn_token_counting(): context.append_output( create_mock_request_output( output_token_ids=[302], + num_cache_creation_tokens=num_cache_creation_tokens[2], finished=True, ) ) @@ -523,6 +539,7 @@ async def test_streaming_multi_turn_token_counting(): assert context.num_cached_tokens == sum( num_cached_tokens ) # Accumulated cached tokens + assert context.num_cache_creation_tokens == sum(num_cache_creation_tokens) # Additional tool tokens from third turn # Formula: this turn prompt - last turn prompt - last turn output diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 334e8032e283..9f5a281163da 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -315,6 +315,7 @@ def __init__( self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 # not implemented yet for ParsableContext self.all_turn_metrics: list[TurnMetrics] = [] @@ -347,6 +348,7 @@ def __init__( def append_output(self, output: RequestOutput) -> None: self.num_prompt_tokens = len(output.prompt_token_ids or []) self.num_cached_tokens = output.num_cached_tokens or 0 + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params @@ -651,6 +653,7 @@ def __init__( self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 self.num_reasoning_tokens = 0 self.num_tool_output_tokens = 0 @@ -773,6 +776,8 @@ def _update_prefill_token_usage(self, output: RequestOutput) -> None: self.num_cached_tokens += num_cached_token self.current_turn_metrics.cached_input_tokens = num_cached_token + self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0 + def _update_decode_token_usage(self, output: RequestOutput) -> int: """Update token usage statistics for the decode phase of generation. diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index ba4aa0434768..78a3f1015ff7 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -881,7 +881,7 @@ async def responses_full_generator( output_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, input_tokens_details=InputTokensDetails( - cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0), + cache_write_tokens=context.num_cache_creation_tokens, cached_tokens=num_cached_tokens, input_tokens_per_turn=[ turn.input_tokens for turn in context.all_turn_metrics From 2bea22002d66a6c89d208fa682b40a89538f023c Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Sat, 10 Oct 2026 20:25:11 +0800 Subject: [PATCH 350/952] [BUG] Fix KV offload shared-memory slot selection using worker rank (#58497) Signed-off-by: Tony Lin --- tests/v1/kv_offload/test_factory.py | 33 +++++++++++++---------------- vllm/v1/kv_offload/cpu/spec.py | 9 ++------ vllm/v1/kv_offload/tiering/spec.py | 9 +------- 3 files changed, 18 insertions(+), 33 deletions(-) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 84b9da36e06b..dbc586bcab07 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import pytest +import torch from vllm.v1.kv_offload.base import ( CanonicalKVCaches, @@ -231,9 +232,7 @@ def fake_worker_ctor(**kwargs): monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor) - monkeypatch.setattr( - tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5 - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 5) kv_caches = MagicMock() spec.create_worker(kv_caches) @@ -244,12 +243,13 @@ def fake_worker_ctor(**kwargs): assert worker_calls[0]["mmap_region"] is region -def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch): +def test_tiering_spec_create_worker_uses_worker_rank_for_sharded_layout(monkeypatch): import vllm.v1.kv_offload.tiering.spec as tiering_spec_module spec = _create_spec( spec_name="TieringOffloadingSpec", worker_kv_bytes_per_block=4096, + rank=1, world_size=4, ) assert isinstance(spec, TieringOffloadingSpec) @@ -263,9 +263,9 @@ def fake_region_ctor(**kwargs): monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock()) monkeypatch.setattr( - tiering_spec_module.torch.accelerator, + torch.accelerator, "current_device_index", - lambda: 5, + lambda: 6, ) spec.create_worker(MagicMock()) @@ -386,9 +386,7 @@ def fake_worker_ctor(**kwargs): monkeypatch.setattr(cpu_spec_module.current_platform, "is_rocm", lambda: False) monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor) - monkeypatch.setattr( - cpu_spec_module.torch.accelerator, "current_device_index", lambda: 5 - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 5) kv_caches = MagicMock() spec.create_worker(kv_caches) @@ -464,16 +462,16 @@ def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch): @pytest.mark.parametrize( - ("replicated_layout", "device_index", "world_size", "expected_rank"), + ("replicated_layout", "rank", "device_index", "world_size", "expected_rank"), [ - (True, 5, 4, 0), # replicated: always slot 0 - (True, 0, 4, 0), # replicated: slot 0 regardless of device - (False, 5, 4, 1), # non-replicated: 5 % 4 == 1 - (False, 7, 4, 3), # non-replicated: 7 % 4 == 3 + (True, 1, 5, 4, 0), + (True, 0, 0, 4, 0), + (False, 1, 5, 4, 1), + (False, 3, 6, 4, 3), ], ) def test_cpu_spec_create_worker_rank_assignment( - monkeypatch, replicated_layout, device_index, world_size, expected_rank + monkeypatch, replicated_layout, rank, device_index, world_size, expected_rank ): import vllm.v1.kv_offload.cpu.spec as cpu_spec_module @@ -483,6 +481,7 @@ def test_cpu_spec_create_worker_rank_assignment( spec = _create_spec( cpu_bytes_to_use=worker_kv_bytes_per_block * 8, worker_kv_bytes_per_block=worker_kv_bytes_per_block, + rank=rank, world_size=world_size, replicated_layout=replicated_layout, ) @@ -495,9 +494,7 @@ def fake_region_ctor(**kwargs): monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor) monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", MagicMock()) - monkeypatch.setattr( - cpu_spec_module.torch.accelerator, "current_device_index", lambda: device_index - ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: device_index) spec.create_worker(MagicMock()) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index bf289d37a65b..c26cd32f9733 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any -import torch from typing_extensions import override from vllm.platforms import current_platform @@ -172,12 +171,8 @@ def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: # mmap'd; fall back to the tensor path (empty tensors) as before. if self._uses_shared_region() and self.num_chunks > 0: # Replicated layout puts all ranks on slot 0 (single MLA copy); - # otherwise each rank takes its own slot by physical device index. - if self.replicated_layout: - rank = 0 - else: - world_size = self.config.parallel.world_size - rank = torch.accelerator.current_device_index() % world_size + # otherwise each worker uses its own rank-indexed slot. + rank = 0 if self.replicated_layout else self.config.parallel.rank mmap_region = SharedOffloadRegion( engine_id=self.config.engine_id, num_chunks=self.num_chunks, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 93c6cbabcb85..508361d39261 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -54,7 +54,6 @@ from typing import Any -import torch from typing_extensions import override from vllm.logger import init_logger @@ -436,13 +435,7 @@ def _uses_shared_region(self) -> bool: @override def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: - world_size = self.config.parallel.world_size - if self.replicated_layout: - rank = 0 - else: - # Fold the global physical device index into the replica-local - # [0, world_size) slot range. - rank = torch.accelerator.current_device_index() % world_size + rank = 0 if self.replicated_layout else self.config.parallel.rank worker_mmap = SharedOffloadRegion( engine_id=self._engine_id, num_chunks=self.num_chunks, From bf16c9c13949a10fc8e77189776f95538a598709 Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:05:55 +0900 Subject: [PATCH 351/952] [Bugfix][Model] K-EXAONE: restore loading of pre-2.0 configs and apply RoPE to the MTP layer (#60187) Signed-off-by: godmook Signed-off-by: Mook <68294499+Godmook@users.noreply.github.com> --- tests/models/registry.py | 4 +++ vllm/model_executor/models/exaone_moe.py | 39 ++++++++++++++---------- vllm/model_executor/models/registry.py | 1 + 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 31c3c5433aec..b57f3fcb7277 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -266,6 +266,10 @@ def check_available_online( ), "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"), "ExaoneMoeForCausalLM": _HfExamplesInfo("LGAI-EXAONE/K-EXAONE-236B-A23B"), + "ExaoneMoEForCausalLM": _HfExamplesInfo( + "LGAI-EXAONE/K-EXAONE-236B-A23B", + hf_overrides={"architectures": ["ExaoneMoEForCausalLM"]}, + ), "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"), "FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"), "FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"), diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index d5b8b4066239..ad6b0b3594fe 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -227,24 +227,31 @@ def __init__( layer_idx = extract_layer_index(prefix) - if config.sliding_windows is not None: - self.sliding_window_size = config.sliding_windows[layer_idx] - - if config.layer_types[layer_idx] == "full_attention": - self.sliding_window_size = None - + # Configs published before K-EXAONE 2.0 only have the scalar + # `sliding_window` and no MTP fields; their MTP layer follows the main + # layer with the same index. + layer_types = config.layer_types + sliding_windows = getattr(config, "sliding_windows", None) if is_mtp: - self.sliding_window = ( - config.mtp_layer_types[layer_idx] == "sliding_attention" + layer_types = getattr(config, "mtp_layer_types", None) or layer_types + sliding_windows = ( + getattr(config, "mtp_sliding_windows", None) or sliding_windows ) - if config.mtp_sliding_windows is not None: - self.sliding_window_size = config.mtp_sliding_windows[layer_idx] - - if config.mtp_layer_types[layer_idx] == "full_attention": - self.sliding_window_size = None - # apply rotary embeddings to every layer in full attention models - self.apply_rope_all_layers = "sliding_attention" not in config.layer_types + if layer_types[layer_idx] == "full_attention": + self.sliding_window_size = None + elif sliding_windows is not None: + self.sliding_window_size = sliding_windows[layer_idx] + else: + self.sliding_window_size = config.sliding_window + + # Full-attention layers are NoPE when sliding layers are present, except + # the MTP layer, which is trained with RoPE. + self.use_rope = ( + bool(self.sliding_window_size) + or is_mtp + or "sliding_attention" not in config.layer_types + ) set_default_rope_theta(config, default_theta=1000000) self.rotary_emb = get_rope( @@ -279,7 +286,7 @@ def forward( k = self.k_norm(k) k = k.flatten(-2, -1) - if self.sliding_window_size or self.apply_rope_all_layers: + if self.use_rope: q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e019d674f502..875179a88fae 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -98,6 +98,7 @@ "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), "ExaoneMoeForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), "FalconForCausalLM": ("falcon", "FalconForCausalLM"), "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"), From 838526c7e5e855f05365a939ef83218cd8e2dbbe Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:09:45 -0400 Subject: [PATCH 352/952] [Compile] Fix sampler compile warning (#60680) Signed-off-by: yewentao256 --- csrc/libtorch_stable/sampler.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index 703a748ce686..09151004c1ad 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -642,6 +642,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); } +#ifdef USE_ROCM // Measured native-kernel policy for gfx950, FP32, and k=1024. This // active-column split crossover is independent of the Python AITER/native // backend crossover; the two policies can be retuned separately. At exactly @@ -693,7 +694,6 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } -#ifdef USE_ROCM // The grid and workspace stay fixed while graph replays change device lengths. __device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, int maxBlocks) { @@ -755,7 +755,7 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware } #endif } -#endif +#endif // USE_ROCM } // namespace vllm From 3709632ff2944a5f2ecdacc84ede5cd134b7ae08 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 11:41:10 -0400 Subject: [PATCH 353/952] [Refactor] Remove dead code multiple places (#60120) Signed-off-by: yewentao256 --- vllm/_custom_ops.py | 6 -- vllm/lora/model_manager.py | 4 - .../layers/fused_moe/experts/cpu_moe.py | 28 ------- vllm/model_executor/models/bert.py | 3 - vllm/model_executor/models/blip.py | 7 -- vllm/model_executor/models/diffusion_gemma.py | 11 --- vllm/model_executor/models/keye.py | 19 ----- vllm/model_executor/models/paddleocr_vl.py | 21 ----- vllm/model_executor/models/phi4mm_audio.py | 6 -- vllm/model_executor/models/step3p5.py | 9 --- .../warmup/jit_warmup_triton_helper.py | 81 ------------------- vllm/tokenizers/mistral.py | 8 -- 12 files changed, 203 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index eef1183697f3..e7862073083f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -4464,12 +4464,6 @@ def cpu_gemm_wna16( return out -def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor: - out = torch.empty_like(input) - torch.ops._C.activation_lut_bf16(out, input, activation) - return out - - def cpu_prepack_moe_weight( weight: torch.Tensor, isa: str, diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 8dc7b5605600..a7f9f02e77af 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -40,7 +40,6 @@ from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.models import ( SupportsLoRA, - SupportsMultiModal, is_pooling_model, supports_multimodal, ) @@ -60,9 +59,6 @@ class SupportsLoRAModel(nn.Module, SupportsLoRA): ... -class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ... - - class AdapterLRUCache(LRUCache[int, T]): def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]): super().__init__(capacity) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 14b25ce30e37..54a9026f97b6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -973,34 +973,6 @@ def apply( ) -# =========================================================================== -# INT8 W8A8 MoE -# =========================================================================== - - -def prepare_int8_moe_layer_for_cpu( - w13: torch.Tensor, - w2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Prepack INT8 MoE weights for the current CPU architecture.""" - # SMMLA packing for AArch64 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - return ( - cpu_prepack_moe_weight_int8(w13, "neon"), - cpu_prepack_moe_weight_int8(w2, "neon"), - ) - # VSX packing for POWER - if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: - return ( - cpu_prepack_moe_weight_int8(w13, "vsx"), - cpu_prepack_moe_weight_int8(w2, "vsx"), - ) - # VNNI packing for x86 - packed_w13 = torch.ops._C.convert_weight_packed(w13) - packed_w2 = torch.ops._C.convert_weight_packed(w2) - return packed_w13, packed_w2 - - class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation modular MoE experts.""" diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 8f061b7284ef..cdaa0bce9725 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -564,9 +564,6 @@ def __init__( self.layer_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps) self.decoder = nn.Linear(hidden_size, vocab_size, bias=True) - def tie_weights_with_embeddings(self, embeddings_weight: torch.Tensor): - self.decoder.weight = embeddings_weight - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: x = self.dense(hidden_states) x = self.activation(x) diff --git a/vllm/model_executor/models/blip.py b/vllm/model_executor/models/blip.py index a4e821ec6b12..4f2275ecd21e 100644 --- a/vllm/model_executor/models/blip.py +++ b/vllm/model_executor/models/blip.py @@ -129,13 +129,6 @@ def __init__( prefix=f"{prefix}.attn", ) - def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): - return ( - tensor.view(bsz, seq_len, self.num_heads, self.head_dim) - .transpose(1, 2) - .contiguous() - ) - def forward( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7c64980b57bb..7f4f9ffba99e 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -269,17 +269,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): eps=getattr(text_config, "rms_norm_eps", 1e-6), ) - def compute_self_conditioning( - self, - inputs_embeds: torch.Tensor, - probs: torch.Tensor, - ) -> torch.Tensor: - embed_weight = self.model.embed_tokens.weight - soft_embeds = torch.matmul( - probs.to(embed_weight.dtype), embed_weight - ) * self.model.normalizer.to(inputs_embeds.dtype) - return self.self_conditioning(inputs_embeds, soft_embeds) - # ------------------------------------------------------------------ # # Multimodal: reuse Gemma4's image parsing, processing & embedding # ------------------------------------------------------------------ # diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 60d0d03b659d..83ff14dcf948 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -255,25 +255,6 @@ def interpolate_pos_encoding( patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return patch_pos_embed - def fetch_position_embedding_lfu_cache(self, embeddings, h, w, max_cache: int = 20): - grid = (h, w) - if grid in self.cache_position_embedding: - self.cache_position_count[grid] += 1 - return self.cache_position_embedding[grid] - - if len(self.cache_position_embedding) >= max_cache: - min_hit_grid = min( - self.cache_position_count, - key=self.cache_position_count.__getitem__, - ) - self.cache_position_count.pop(min_hit_grid) - self.cache_position_embedding.pop(min_hit_grid) - - position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True) - self.cache_position_count[grid] = 1 - self.cache_position_embedding[grid] = position_embedding - return position_embedding - def forward( self, pixel_values: torch.FloatTensor, diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 1fde4cd0a11b..1c165c08ef5d 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -455,27 +455,6 @@ def interpolate_pos_encoding( patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) return patch_pos_embed - def fetch_position_embedding_lfu_cache( - self, embeddings: torch.Tensor, h: int, w: int, max_cache: int = 20 - ): - grid = (h, w) - if grid in self.cache_position_embedding: - self.cache_position_count[grid] += 1 - return self.cache_position_embedding[grid] - - if len(self.cache_position_embedding) >= max_cache: - min_hit_grid = min( - self.cache_position_count, - key=self.cache_position_count.__getitem__, - ) - self.cache_position_count.pop(min_hit_grid) - self.cache_position_embedding.pop(min_hit_grid) - - position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True) - self.cache_position_count[grid] = 1 - self.cache_position_embedding[grid] = position_embedding - return position_embedding - def forward( self, pixel_values: torch.FloatTensor, diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 9daedcf38e7b..151fcfc6ecdf 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -1256,12 +1256,6 @@ def __init__(self, config: PreTrainedConfig, **kwargs: Any) -> None: self.input_embeds = None self.audio_embed_sizes = None - def set_audio_embeds(self, input_embeds: torch.Tensor) -> None: - self.input_embeds = input_embeds - - def set_audio_embed_sizes(self, audio_embed_sizes: torch.Tensor) -> None: - self.audio_embed_sizes = audio_embed_sizes - def get_audio_features( self, input_embeds: torch.Tensor, diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index b56d4d3ed8d8..a6dffeb1d40f 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -17,7 +17,6 @@ get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, ) from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul @@ -459,7 +458,6 @@ def __init__( f"Unsupported attention implementation: {config.att_impl_type}" ) self.use_moe = False - self.tp_group = get_tp_group() self.use_fused_all_reduce = ( get_tensor_model_parallel_world_size() > 1 and get_dp_group().world_size == 1 @@ -496,13 +494,6 @@ def __init__( ) self.prefix = prefix - def add_and_maybe_inplace_all_reduce( - self, in1: torch.Tensor, in2: torch.Tensor - ) -> torch.Tensor: - if not self.use_fused_all_reduce: - return in1 + in2 - return self.tp_group.all_reduce(in1 + in2) - def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor ) -> torch.Tensor: diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index f7765cf23293..70ec8cd70fc5 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -718,65 +718,6 @@ def pointer( return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape) -def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]: - if isinstance(node, ast.Constant) and isinstance(node.value, str | int): - return (node.value,) - if isinstance(node, ast.List | ast.Tuple): - refs: list[str | int] = [] - for elt in node.elts: - if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int): - refs.append(elt.value) - else: - raise ValueError( - f"Unsupported Triton specialization ref: {ast.dump(elt)}" - ) - return tuple(refs) - raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}") - - -def _normalize_arg_refs( - refs: tuple[str | int, ...], - arg_names: tuple[str, ...], -) -> frozenset[str]: - names: set[str] = set() - for ref in refs: - if isinstance(ref, int): - names.add(arg_names[ref]) - else: - names.add(ref) - return frozenset(names) - - -def _decorator_keyword_refs( - function_def: ast.FunctionDef, - keyword_name: str, -) -> tuple[str | int, ...]: - for decorator in function_def.decorator_list: - if not isinstance(decorator, ast.Call): - continue - decorator_name = get_ast_full_name(decorator.func) - if decorator_name not in ("triton.jit", "jit"): - continue - for keyword in decorator.keywords: - if keyword.arg == keyword_name: - return _literal_str_refs(keyword.value) - return () - - -def _triton_do_not_specialize_args( - kernel: Callable[..., Any], - function_def: ast.FunctionDef, - arg_names: tuple[str, ...], -) -> frozenset[str]: - refs = getattr(kernel, "do_not_specialize", None) - if refs is not None: - return _normalize_arg_refs(tuple(refs), arg_names) - return _normalize_arg_refs( - _decorator_keyword_refs(function_def, "do_not_specialize"), - arg_names, - ) - - def _triton_constexpr_arg_names( kernel: Callable[..., Any], function_def: ast.FunctionDef, @@ -821,25 +762,3 @@ def _pointer_arg_names( if name in candidate_names: pointer_names.add(name) return frozenset(pointer_names) - - -def trace_triton_kernel_specialization_args( - kernel: Callable[..., Any], -) -> tuple[str, ...]: - function_def = get_function_source_node(kernel) - if not isinstance(function_def, ast.FunctionDef): - raise ValueError("Expected Triton kernel to be defined as a function") - source_fn = getattr(kernel, "fn", kernel) - arg_names = tuple(inspect.signature(source_fn).parameters) - constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names) - do_not_specialize_args = _triton_do_not_specialize_args( - kernel, function_def, arg_names - ) - pointer_args = _pointer_arg_names(function_def, arg_names) - - return tuple( - name - for name in arg_names - if name in constexpr_args - or (name not in pointer_args and name not in do_not_specialize_args) - ) diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 15cf0322581a..de321025a313 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -178,14 +178,6 @@ def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int: return tokenizer.unk_id -def mistral_common_tekkenizer(tokenizer: object) -> "Tekkenizer | None": - """Return the underlying `Tekkenizer` for a `MistralCommonBackend`.""" - mistral = getattr(tokenizer, "tokenizer", None) - instruct = getattr(mistral, "instruct_tokenizer", None) - tekken = getattr(instruct, "tokenizer", None) - return tekken if isinstance(tekken, Tekkenizer) else None - - def tekken_convert_ids_to_tokens( tokenizer: "Tekkenizer", ids: Sequence[int] ) -> list[str | bytes]: From f2376cdc4da852ec48ab5bcb747fdfa65a4dfd4a Mon Sep 17 00:00:00 2001 From: siyu Date: Sat, 10 Oct 2026 23:48:06 +0800 Subject: [PATCH 354/952] [Fast Start] Preload the FlashInfer autotune table on the weight cache daemon (#60085) Signed-off-by: liusy58 Signed-off-by: Isotr0py Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Isotr0py --- docs/features/preload.md | 30 +++++++- .../model_loader/test_weight_cache.py | 73 ++++++++++++++++++- .../test_flashinfer_autotune_warmup.py | 31 ++++++++ vllm/entrypoints/cli/preload.py | 9 ++- .../model_loader/weight_cache/daemon.py | 39 ++++++++++ .../model_loader/weight_cache/utils.py | 35 ++++++++- .../warmup/flashinfer_autotune_cache.py | 16 +++- vllm/model_executor/warmup/kernel_warmup.py | 10 ++- vllm/v1/worker/gpu/model_runner.py | 12 ++- vllm/v1/worker/gpu_model_runner.py | 15 +++- 10 files changed, 253 insertions(+), 17 deletions(-) diff --git a/docs/features/preload.md b/docs/features/preload.md index 74839bf35b38..6e8fad9ad03e 100644 --- a/docs/features/preload.md +++ b/docs/features/preload.md @@ -50,6 +50,26 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \ The daemon itself must load from disk; passing `--load-format ipc_cache` to `vllm preload` is an error. +## Preloading the FlashInfer autotune cache + +FlashInfer has several implementations of each operation and chooses between +them by benchmarking. vLLM runs that pass during kernel warmup and persists the +result per rank in the on-disk FlashInfer autotune cache (under +`VLLM_CACHE_ROOT`, or `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR`), so a later engine +with the same configuration loads the tuned configs instead of profiling again. +Normally the first serving engine pays for the pass, and on a large MoE model it +dominates what is left of the startup time once the weights come from the +daemon. + +Unless FlashInfer autotune is disabled (`--no-enable-flashinfer-autotune`), +`vllm preload` pays it at preload time: each daemon runs the autotune pass in +place — profiling dummy runs against the shard it already holds — after +loading and before binding its socket, then writes the cache: + +```bash +vllm preload --model /path/to/model --tensor-parallel-size 4 +``` + ## How it works 1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its @@ -164,8 +184,14 @@ Socket paths are derived from the GPU UUID, so they are stable regardless of - **Platform**: CUDA and ROCm only; other platforms raise `UnsupportedPlatformForIPCError` even when `fallback` is enabled, since it is a permanent misconfiguration rather than a transient daemon outage. -- **Parallelism**: tensor, expert and data parallelism are supported; - launching the daemon with pipeline parallelism is rejected. +- **Parallelism**: tensor, expert, pipeline and data parallelism are + supported, including across nodes. +- **Autotune preloading**: tuning runs across the daemon's world group, + including across nodes and DP ranks. With model-based speculative decoding + both daemon groups skip it (the target cannot run the drafter's dummy + passes, and a table tuned on the draft model would key on the draft's + config, which no engine looks up), so those deployments' first engine + tunes and fills the cache instead. - **Quantization**: every quantization method in the model must declare support for pre-processed weights (the daemon transfers weights *after* quantization post-processing). Unsupported methods raise diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 7960a237914b..9628cb47e0bb 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -12,12 +12,14 @@ import sys import tempfile import threading +import time import urllib.error import urllib.request from dataclasses import dataclass, field from typing import Any import pytest +import regex as re from vllm import SamplingParams from vllm.assets.image import ImageAsset @@ -81,6 +83,21 @@ def check_health(self) -> None: f"Weight cache daemon is not ready: HTTP {error.code}" ) from error + def wait_until_ready(self, timeout_s: float = 900) -> None: + """Poll until READY (which covers the autotune tuners) or the preload + process dies.""" + deadline = time.monotonic() + timeout_s + while True: + try: + self.check_health() + return + except (AssertionError, urllib.error.URLError): + if time.monotonic() > deadline: + raise + if self._proc is not None and self._proc.poll() is not None: + raise AssertionError("vllm preload exited before ready") from None + time.sleep(2) + def _stop(self) -> None: assert self._proc is not None self._proc.terminate() @@ -99,6 +116,10 @@ class ModelCase: images: list | None = None llm_kwargs: dict[str, Any] = field(default_factory=dict) daemon_args: list[str] = field(default_factory=list) + # Additionally assert the daemons tune in place before serving and write + # the on-disk FlashInfer autotune cache, and the restart engine reuses it + # (needs SM90+/FlashInfer). + check_flashinfer_cache: bool = False def generate( @@ -157,6 +178,9 @@ def generate( max_model_len=4096, ), daemon_args=["--trust-remote-code"], + # Kimi K3 has tunable FlashInfer ops (a dense model like Qwen2.5 tunes + # nothing and upstream never writes the cache file for it). + check_flashinfer_cache=True, ) # Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp" @@ -186,18 +210,38 @@ def generate( [QWEN_CASE, K3_CASE, QWEN_MTP_CASE], ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"], ) -def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): +def test_ipc_cache_cold_start_and_warm_restart( + vllm_runner, case: ModelCase, tmp_path, monkeypatch, capfd +): """Cold start falls back to disk; warm restarts load weights via CUDA IPC. All runs must produce outputs identical to a default-loader baseline. The warm runs disable the disk fallback, so they only pass if the weights really came from the daemon — for the MTP case, both the target's and the draft's daemon groups. + + Cases with ``check_flashinfer_cache`` additionally cover in-daemon + autotune (default-on with `--enable-flashinfer-autotune`): each daemon + tunes before binding its socket and must write the on-disk FlashInfer + autotune cache, and the restart engine must load tactics from it instead + of re-profiling. The daemon's cache file is keyed with the + OPENAI_API_SERVER batch defaults while these engines run under LLM_CLASS, + so they never hit the daemon's file by design; reuse is proven against + the file the earlier in-process engines wrote. """ if not current_platform.is_cuda_alike(): pytest.skip("Weight cache IPC sharing requires CUDA or ROCm") if case is K3_CASE and not current_platform.is_device_capability_family(100): pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend") + if case.check_flashinfer_cache: + from vllm.utils.flashinfer import has_flashinfer + + if not (has_flashinfer() and current_platform.has_device_capability(90)): + pytest.skip("FlashInfer autotune requires FlashInfer and SM90+") + # Isolate the on-disk autotune cache. The engines are in-process + # (envs read lazily); the daemon subprocesses spawn later and + # inherit it. + monkeypatch.setenv("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)) # Baseline: plain disk loading with the default loader. baseline_outputs = generate( @@ -212,15 +256,40 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir: cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True) + if case.check_flashinfer_cache: + # The baseline engine already wrote its own (LLM_CLASS-keyed) table. + engine_cache_files = set(tmp_path.rglob("autotune_configs*.json")) + with WeightCacheDaemon( case.model, tp_size=1, extra_args=case.daemon_args, ) as d: warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) - d.check_health() + d.wait_until_ready() + if case.check_flashinfer_cache: + # Each daemon tunes before its ready report; exactly one new file + # (their OPENAI_API_SERVER-keyed table) must appear. + tuner_files = set(tmp_path.rglob("autotune_configs*.json")) + tuner_files -= engine_cache_files + assert len(tuner_files) == 1, ( + f"expected the tuner's tuned table, got {tuner_files}" + ) # Warm restart: a second engine lifetime against the same daemon. + if case.check_flashinfer_cache: + capfd.readouterr() # drain: the assertion reads the increment restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) + if case.check_flashinfer_cache: + # The restart's EngineCore subprocess inherits pytest's fds; the + # flashinfer autotuner logs "[Autotuner]: Loaded N configs from + # " (INFO, not once-deduplicated) on a disk cache hit. + out = capfd.readouterr() + m = re.search( + r"\[Autotuner\]: Loaded (\d+) configs from", out.out + out.err + ) + assert m and int(m.group(1)) > 0, ( + "restart engine did not load the on-disk autotune cache" + ) assert cold_outputs == baseline_outputs assert warm_outputs == baseline_outputs diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index f7ea6847f26f..59046a5f3e8d 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -386,3 +386,34 @@ def test_rejected_cache_retunes_every_rank_in_its_tuning_group(autotune_run, pp, rerun.assert_collectives_match() assert set(rerun.profile_groups) == set(range(tp)) assert {rank for rank, _, _ in rerun.saves} == set(range(tp)) + + +def test_autotune_cache_key_normalizes_kv_init_fields(): + """mamba_block_size/kv_cache_layout only resolve at engine KV-init time; + they must not feed the autotune cache key, or nothing can compute the key + before startup (the daemon tunes before any KV init). block_size must + stay in the key: the daemon replays the engine's block-size resolution + before tuning, and different block sizes may select different ops.""" + from vllm.config.cache import CacheConfig + from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + _normalize_cache_config_for_hash, + ) + + fresh = CacheConfig() + # Simulate the engine's KV-init write-backs (hybrid model values). + mutated = CacheConfig() + mutated.mamba_block_size = 512 + mutated.kv_cache_layout = "LBHNC" + + # Negative control: the write-backs are hashed factors. + assert mutated.compute_hash() != fresh.compute_hash() + normalized = _normalize_cache_config_for_hash(mutated) + assert normalized.compute_hash() == fresh.compute_hash() + # Factory state is returned untouched. + assert _normalize_cache_config_for_hash(fresh) is fresh + # block_size stays significant (the daemon keys on the resolved value). + resolved = CacheConfig(block_size=64) + assert ( + _normalize_cache_config_for_hash(resolved).compute_hash() + != fresh.compute_hash() + ) diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index 9eab7421f175..90e746fdb686 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/vllm/entrypoints/cli/preload.py @@ -39,6 +39,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( check_ipc_platform_support, ) +from vllm.usage.usage_lib import UsageContext from vllm.utils.network_utils import get_distributed_init_method, get_open_port if typing.TYPE_CHECKING: @@ -157,7 +158,13 @@ def cmd(args: argparse.Namespace) -> None: ) engine_args = EngineArgs.from_cli_args(args) - vllm_config = engine_args.create_engine_config() + # Warmup context for flashinfer autotune. + usage_context = ( + UsageContext.OPENAI_API_SERVER + if engine_args.kernel_config.enable_flashinfer_autotune is not False + else None + ) + vllm_config = engine_args.create_engine_config(usage_context=usage_context) if vllm_config.load_config.load_format == "ipc_cache": raise ValueError( "The weight cache daemon itself must load from disk; use the " diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 2c6944f1eb14..0fc63da7869b 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -112,11 +112,14 @@ verify_peer_is_owner, ) from vllm.model_executor.model_loader.weight_cache.utils import ( + build_warmup_runner, export_model_attrs, format_daemon_role, is_draft_model_cacheable, ) +from vllm.model_executor.warmup.kernel_warmup import flashinfer_autotune from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.worker.workspace import init_workspace_manager @@ -236,6 +239,41 @@ def __init__( is_draft=is_draft, ) + def warmup(self) -> None: + """Run the FlashInfer autotune pass against the cached model.""" + vllm_config = self.vllm_config + # Mirrors the gate in the engine's kernel warmup. + if ( + vllm_config.kernel_config.enable_flashinfer_autotune is False + or not has_flashinfer() + or not current_platform.has_device_capability(90) + or is_draft_model_cacheable(vllm_config.speculative_config) + ): + return + try: + current_platform.update_block_size_for_backend(vllm_config) + runner = build_warmup_runner( + vllm_config, + self.local_rank, + is_draft=self.is_draft, + model_config=self.model_config, + ) + assert self.model is not None, "warmup ran before load_model" + runner.load_model(model=self.model) + # The V2 runner requires skip_attn when no KV cache exists + flashinfer_autotune(runner, skip_attn=vllm_config.use_v2_model_runner) + logger.info( + "Weight cache %s daemon rank %d tuned FlashInfer; the tuned " + "configs are in the on-disk autotune cache", + self.role, + self.global_rank, + ) + finally: + # Free the runner and the dummy-run activations; the daemon's + # weights are untouched. + gc.collect() + torch.accelerator.empty_cache() + def load_model(self) -> None: torch.accelerator.set_device_index(self.local_rank) # Outside the config context so the daemon keeps its own rendezvous @@ -457,6 +495,7 @@ def _run_daemon( pp_rank, ) daemon.load_model() + daemon.warmup() daemon.serve_forever( ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank)) ) diff --git a/vllm/model_executor/model_loader/weight_cache/utils.py b/vllm/model_executor/model_loader/weight_cache/utils.py index 1abbf572f847..bc5de3cc9768 100644 --- a/vllm/model_executor/model_loader/weight_cache/utils.py +++ b/vllm/model_executor/model_loader/weight_cache/utils.py @@ -7,10 +7,17 @@ wire format. """ -from typing import Any +from typing import TYPE_CHECKING, Any, cast -from vllm.config import SpeculativeConfig +import torch + +from vllm.config import SpeculativeConfig, VllmConfig, replace, set_current_vllm_config from vllm.model_executor.models.interfaces import SupportsEagleBase +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.v1.worker.gpu_model_runner import GPUModelRunner # Speculative methods whose draft model the daemon caches in its own group. # Other drafts keep loading from disk in the engine. @@ -46,3 +53,27 @@ def format_daemon_role(is_draft: bool) -> str: def format_socket_role_suffix(is_draft: bool) -> str: """Socket-name suffix keeping the draft group distinct from the target.""" return "_draft" if is_draft else "" + + +def build_warmup_runner( + vllm_config: VllmConfig, + local_rank: int, + *, + is_draft: bool, + model_config: "ModelConfig", +) -> "GPUModelRunner": + """Create a model runner for the daemon's warmup dummy runs.""" + if is_draft: + vllm_config = replace( + vllm_config, model_config=model_config, speculative_config=None + ) + device = torch.device(current_platform.device_type, local_rank) + with set_current_vllm_config(vllm_config): + if vllm_config.use_v2_model_runner: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as Runner + else: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as Runner + runner = Runner(vllm_config, device) + # The two runner classes share the warmup surface; the nominal type + # stays V1's GPUModelRunner. + return cast("GPUModelRunner", runner) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index afc3ef3ac014..ca2c4d1421c8 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -10,14 +10,28 @@ from typing import TYPE_CHECKING import vllm.envs as envs +from vllm.config import CacheConfig, VllmConfig, replace if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig: + """Normalize cache config for hashing""" + if cache_config.mamba_block_size is None and cache_config.kv_cache_layout is None: + return cache_config + normalized = replace(cache_config, mamba_block_size=None) + assert normalized.kv_cache_layout is None + return normalized + + def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: - config_hash = runner.vllm_config.compute_hash(include_version=False) + vllm_config: VllmConfig = runner.vllm_config + cache_config = _normalize_cache_config_for_hash(vllm_config.cache_config) + if cache_config is not vllm_config.cache_config: + vllm_config = replace(vllm_config, cache_config=cache_config) + config_hash = vllm_config.compute_hash(include_version=False) return hashlib.sha256(config_hash.encode()).hexdigest() diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e5751730e85d..685b33a5652a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -406,7 +406,7 @@ def _all_ranks_have_matching_cache(path: Path, group) -> bool: return fingerprint is not None and all(f == fingerprint for f in gathered) -def flashinfer_autotune(runner: "GPUModelRunner") -> None: +def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> None: """Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, autotuning runs benchmarks for each implementation and stores @@ -415,6 +415,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower. + Args: + runner: The model runner to tune. + skip_attn: Additionally skip attention in the tuning dummy runs. + With PP > 1, stages run different layers and may profile different ops, so each stage's TP group tunes separately with its own cache file; otherwise the world group tunes together. Per-tactic timings are @@ -425,6 +429,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: rank. A rank with a cache hit skips the per-tactic reduce the others block in, so ranks keep loaded configs only if every rank in the tuning group has a matching file and successfully loads it. + """ from flashinfer.autotuner import AutoTuner, set_autotune_process_group @@ -485,11 +490,12 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: hisparse_enabled = ( runner.vllm_config.attention_config.hisparse_config is not None ) + skip_attn = skip_attn or hisparse_enabled if hisparse_enabled: # HiSparse hot-buffer attention is bounded by decode batch # size, not the prefill-sized batch used for the full model. autotune_hisparse_flashinfer_attention(runner) - _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled) + _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) replayssm_autotune_warmup(runner) _autotune_kimi_k3_kda_qkvg(runner.get_model()) finally: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 674cf23527d1..e04a381b94e6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -373,9 +373,15 @@ def get_supported_tasks(self) -> tuple[SupportedTask, ...]: tasks.extend(PoolingRunner.get_supported_tasks(self.model)) return tuple(tasks) - def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: + def load_model( + self, + load_dummy_weights: bool = False, + *args, + model: nn.Module | None = None, + **kwargs, + ) -> None: time_before_load = time.perf_counter() - if load_dummy_weights: + if model is None and load_dummy_weights: self.load_config.load_format = "dummy" self.eplb.prepare_load() eplb_models_added = False @@ -385,7 +391,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: # Capture warmup providers selected while constructing the model. with self.jit_warmup_registry.activate(): - self.model = model_loader.load_model( + self.model = model or model_loader.load_model( vllm_config=self.vllm_config, model_config=self.vllm_config.model_config, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 22ad662b9018..90bea22c9d4f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5217,9 +5217,15 @@ def update_config(self, overrides: dict[str, Any]) -> None: setattr(self, config_name, new_config) @instrument(span_name="Loading (GPU)") - def load_model(self, load_dummy_weights: bool = False) -> None: + def load_model( + self, + load_dummy_weights: bool = False, + *, + model: nn.Module | None = None, + ) -> None: """Args: load_dummy_weights: load dummy weights instead of real weights. + model: a pre-loaded model to skip the model loader. """ logger.info_once( @@ -5235,13 +5241,14 @@ def load_model(self, load_dummy_weights: bool = False) -> None: try: with DeviceMemoryProfiler() as m: time_before_load = time.perf_counter() - if load_dummy_weights: + if model is None and load_dummy_weights: self.load_config.load_format = "dummy" model_loader = get_model_loader(self.load_config) # Capture warmup providers selected while constructing the model. with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] - self.model = model_loader.load_model( - vllm_config=self.vllm_config, model_config=self.model_config + self.model = model or model_loader.load_model( + vllm_config=self.vllm_config, + model_config=self.model_config, ) lookback_depth = getattr(self.model, "token_lookback_depth", 0) if lookback_depth > 0: From 187a0eb98aa42341d703f83421d693fa7585581b Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:36:03 +0800 Subject: [PATCH 355/952] [Bugfix][Frontend] Forward --enable-log-outputs to the Anthropic handler (#60982) Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- vllm/entrypoints/anthropic/serving.py | 4 ++++ vllm/entrypoints/generate/api_router.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 7b1cf95e99b0..22319f1c70c7 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -163,6 +163,8 @@ def __init__( tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_log_outputs: bool = False, + enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, disabled_thinking_effort: AnthropicDisabledThinkingEffortOption = "auto", ): @@ -181,6 +183,8 @@ def __init__( tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, + enable_log_outputs=enable_log_outputs, + enable_log_deltas=enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, ) self._merge_inline_system = self._should_merge_inline_system(online_renderer) diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 66a37d6b13ee..6bf8d690265a 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -193,6 +193,8 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, disabled_thinking_effort=args.anthropic_disabled_thinking_effort, ) From 276fbcff2717bd934cfa37c8a2e4c391f3e7237b Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 10 Oct 2026 15:45:33 -0400 Subject: [PATCH 356/952] [Deprecation] Remove scheduled deprecated items (#60923) Signed-off-by: yewentao256 --- rust/src/server/src/routes/pause.rs | 4 --- vllm/benchmarks/throughput.py | 25 ++----------------- vllm/engine/protocol.py | 2 -- .../pooling/pooling/io_processor.py | 21 +--------------- vllm/entrypoints/rl/online/api_router.py | 3 --- vllm/scripts.py | 17 ------------- vllm/v1/engine/async_llm.py | 12 --------- 7 files changed, 3 insertions(+), 81 deletions(-) delete mode 100644 vllm/scripts.py diff --git a/rust/src/server/src/routes/pause.rs b/rust/src/server/src/routes/pause.rs index 1654c85b0200..ed0c175b84c3 100644 --- a/rust/src/server/src/routes/pause.rs +++ b/rust/src/server/src/routes/pause.rs @@ -39,10 +39,6 @@ fn invalid_query(error: QueryRejection) -> ApiError { ApiError::invalid_request(error.body_text(), Some("mode")) } -// TODO: the Python frontend also accepts the deprecated -// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is -// intentionally omitted here in favor of the `mode` parameter. - /// Pause the scheduler so generation can be halted (e.g. for weight updates). pub async fn pause( State(state): State>, diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 6245b20e86c8..76f6a72fd98a 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -625,15 +625,7 @@ def assign_loras(requests, args): def validate_args(args): """Validate command-line arguments.""" - # === Deprecation and Defaulting === - if args.dataset is not None: - warnings.warn( - "The '--dataset' argument will be deprecated in the next release. " - "Please use '--dataset-name' and '--dataset-path' instead.", - stacklevel=2, - ) - args.dataset_path = args.dataset - + # === Defaulting === if not getattr(args, "tokenizer", None): args.tokenizer = args.model @@ -647,11 +639,7 @@ def validate_args(args): raise ValueError("--prequeue-requests is not supported with --async-engine") # === Dataset Configuration === - if ( - not args.dataset - and not args.dataset_path - and args.dataset_name not in {"prefix_repetition"} - ): + if not args.dataset_path and args.dataset_name != "prefix_repetition": print("When dataset path is not set, it will default to random dataset") args.dataset_name = "random" random_input_len = getattr(args, "random_input_len", None) @@ -809,15 +797,6 @@ def add_cli_args(parser: FlexibleArgumentParser): help="Name of the dataset to benchmark on.", default="sharegpt", ) - parser.add_argument( - "--dataset", - type=str, - default=None, - help="Path to the ShareGPT dataset, will be deprecated in\ - the next release. The dataset is expected to " - "be a json in form of list[dict[..., conversations: " - "list[dict[..., value: ]]]]", - ) parser.add_argument( "--dataset-path", type=str, default=None, help="Path to the dataset" ) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 2bb128ba073e..15f6b801bcd6 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -224,7 +224,6 @@ async def pause_generation( self, *, mode: "PauseMode" = "abort", - wait_for_inflight_requests: bool = False, clear_cache: bool = True, ) -> None: """Pause new generation/encoding requests. @@ -236,7 +235,6 @@ async def pause_generation( - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on :meth:`resume_generation`. - wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV and prefix caches after draining. diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index d5132f2a6649..74e7351c258a 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -4,7 +4,6 @@ from typing import Any from vllm import PoolingParams, PoolingRequestOutput -from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq @@ -21,8 +20,6 @@ ) from .protocol import IOProcessorRequest, IOProcessorResponse -logger = init_logger(__name__) - class UnsupportedCombinedTaskIOProcessor(PoolingIOProcessor): name = "embed&token_classify" @@ -117,23 +114,7 @@ def post_process_online( ctx.final_res_batch, request_id=ctx.request_id, ) - - if callable( - output_to_response := getattr(self.io_processor, "output_to_response", None) - ): - logger.warning_once( - "`IOProcessor.output_to_response` is deprecated. To ensure " - "consistency between offline and online APIs, " - "`IOProcessorResponse` will become a transparent wrapper " - "around output data from v0.19 onwards.", - ) - - if hasattr(output, "request_id") and output.request_id is None: - output.request_id = ctx.request_id # type: ignore - - ctx.response = output_to_response(output) # type: ignore - else: - ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) ####################################### # offline APIs diff --git a/vllm/entrypoints/rl/online/api_router.py b/vllm/entrypoints/rl/online/api_router.py index ef9269749f3b..51182fd99683 100644 --- a/vllm/entrypoints/rl/online/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -32,7 +32,6 @@ def engine_client(request: Request) -> EngineClient: async def pause_generation( raw_request: Request, mode: Annotated[PauseMode, Query()] = "abort", - wait_for_inflight_requests: bool = Query(False), clear_cache: Annotated[bool, Query()] = True, ) -> JSONResponse: """Pause generation requests to allow weight updates. @@ -44,7 +43,6 @@ async def pause_generation( - ``"abort"``: Abort all in-flight requests immediately (default). - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on /resume. - wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead. clear_cache: DEPRECATED. Whether to clear KV/prefix caches after draining. Ignored when mode="keep". @@ -55,7 +53,6 @@ async def pause_generation( await engine.pause_generation( mode=mode, clear_cache=clear_cache, - wait_for_inflight_requests=wait_for_inflight_requests, ) return JSONResponse( content={"status": "paused"}, diff --git a/vllm/scripts.py b/vllm/scripts.py deleted file mode 100644 index f158860726be..000000000000 --- a/vllm/scripts.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.entrypoints.cli.main import main as vllm_main -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -# Backwards compatibility for the move from vllm.scripts to -# vllm.entrypoints.cli.main -def main(): - logger.warning( - "vllm.scripts.main() is deprecated. Please re-install " - "vllm or use vllm.entrypoints.cli.main.main() instead." - ) - vllm_main() diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index bdab43e92528..08445d8f913a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -4,7 +4,6 @@ import os import socket import time -import warnings from collections.abc import AsyncGenerator, Iterable, Mapping from copy import copy from typing import Any @@ -944,7 +943,6 @@ async def pause_generation( self, *, mode: PauseMode = "abort", - wait_for_inflight_requests: bool | None = None, clear_cache: bool = True, ) -> None: """Pause generation to allow model weight updates. @@ -960,20 +958,10 @@ async def pause_generation( - ``"wait"``: Wait for in-flight requests to complete. - ``"keep"``: Freeze requests in queue; they resume on :meth:`resume_generation`. - wait_for_inflight_requests: DEPRECATED: use mode argument. clear_cache: Whether to clear KV cache and prefix cache after draining. Set to ``False`` to preserve cache for faster resume. """ - if wait_for_inflight_requests: - warnings.warn( - "The `wait_for_inflight_requests` parameter in " - "`AsyncLLM.pause_generation()` is deprecated. " - "Please use `mode` argument instead.", - DeprecationWarning, - stacklevel=2, - ) - mode = "wait" if clear_cache: await self.renderer.clear_mm_cache_async() await self.engine_core.pause_scheduler_async(mode=mode, clear_cache=clear_cache) From 3c062d391d15cfa41d3f321cad24f3030c892a97 Mon Sep 17 00:00:00 2001 From: Subhash Polisetti <49185065+subhashpolisetti@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:41:06 -0700 Subject: [PATCH 357/952] [Bugfix] Support 0-dim tensors in packed weight transfer (#52235) Signed-off-by: subhash polisetti Signed-off-by: aoshen02 Co-authored-by: aoshen02 Co-authored-by: Codex --- tests/distributed/test_packed_tensor.py | 29 ++++++++++++++++++- .../weight_transfer/packed_tensor.py | 6 ++-- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_packed_tensor.py b/tests/distributed/test_packed_tensor.py index 5b6846af7683..63af7f4b56fc 100644 --- a/tests/distributed/test_packed_tensor.py +++ b/tests/distributed/test_packed_tensor.py @@ -618,6 +618,7 @@ def test_roundtrip_basic(self): """Test basic IPC producer -> consumer roundtrip.""" params = [ ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()), + ("scale", torch.tensor(1.5, dtype=torch.float32).cuda()), ("w2", torch.randn(5, dtype=torch.float16).cuda()), ] gpu_uuid = self._get_gpu_uuid() @@ -634,7 +635,7 @@ def test_roundtrip_basic(self): ) assert num_chunks == 1 - assert len(result) == 2 + assert len(result) == len(params) for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result): assert orig_name == res_name assert res_tensor.shape == orig_tensor.shape @@ -746,3 +747,29 @@ def test_consumer_wrong_uuid_raises(self): tensor_sizes=c.tensor_sizes, device_index=torch.cuda.current_device(), ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_pack_unpack_scalar(dtype): + """Preserve scalar shape, dtype and value alongside a matrix on CPU.""" + params = [ + ("weight", torch.randn(2, 3, dtype=dtype)), + ("scale", torch.tensor(1.5, dtype=dtype)), + ] + + chunk = pack_tensors(iter(params), lambda item: item[1], buffer_size_bytes=1024) + assert chunk is not None + + result = unpack_tensor( + chunk.packed_tensor, + chunk.names, + chunk.shapes, + chunk.dtypes, + chunk.tensor_sizes, + ) + + for (expected_name, expected), (name, actual) in zip(params, result, strict=True): + assert name == expected_name + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index dcd8efc8c4af..2e715383801f 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -50,7 +50,7 @@ def unpack_tensor( unpacked_tensors = packed_tensor.split(tensor_sizes) return [ - (name, tensor.contiguous().view(dtype).view(*shape)) + (name, tensor.contiguous().view(dtype).view(shape)) for name, shape, dtype, tensor in zip(names, shapes, dtypes, unpacked_tensors) ] @@ -105,7 +105,7 @@ def pack_tensors( name, orig_tensor = item # Apply post processing and convert to linearized uint8 tensor - tensor = post_iter_func(item).contiguous().view(torch.uint8).view(-1) + tensor = post_iter_func(item).contiguous().view(-1).view(torch.uint8) if tensor.numel() > buffer_size_bytes: import warnings @@ -350,7 +350,7 @@ def packed_ipc_producer( for name, orig_tensor in iterator: flat = ( - post_iter_func((name, orig_tensor)).contiguous().view(torch.uint8).view(-1) + post_iter_func((name, orig_tensor)).contiguous().view(-1).view(torch.uint8) ) if flat.numel() > buffer_size_bytes: From b89852aad0854cc8df9138bc412b9a75c1fa8120 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:33:29 +0800 Subject: [PATCH 358/952] [RL][Frontend] Return token-level weight-version spans on final outputs (#59582) Signed-off-by: Shuolei Wang Signed-off-by: Shuolei Wang <948904026@qq.com> Co-authored-by: aoshen02 --- .../online_serving/token_in_token_out.md | 15 ++++++ docs/training/async_rl.md | 20 ++++++- examples/rl/rlhf_async_new_apis.py | 43 +++++++-------- .../engine-core-client/src/protocol/output.rs | 3 ++ .../src/tests/python_compat.py | 3 ++ tests/v1/engine/test_output_processor.py | 53 +++++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 5 ++ .../scale_out/token_in_token_out/serving.py | 2 + vllm/outputs.py | 14 +++++ vllm/v1/engine/__init__.py | 3 ++ vllm/v1/engine/async_llm.py | 5 +- vllm/v1/engine/core.py | 7 +++ vllm/v1/engine/llm_engine.py | 1 + vllm/v1/engine/output_processor.py | 23 ++++++++ 14 files changed, 169 insertions(+), 28 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index e22d4906155b..b3e15ae61692 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -55,6 +55,21 @@ The server returns a 400 for: For using `output_mode` with separate prefill and decode pools, see [Disaggregated Prefilling](../../features/disagg_prefill.md#generate-api-output-modes). +## Weight versions + +The final output of every choice carries `weight_versions`, the label of the weights that sampled each range of its output tokens. Ranges are half-open over output-token indexes (prompt excluded), contiguous and in generation order, so a request that crossed one weight update gets two: + +```json +"weight_versions": [ + {"version": "v7", "start": 0, "end": 57}, + {"version": "v8", "start": 57, "end": 128} +] +``` + +- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null`. +- An output with no tokens, such as a request aborted before its first token, gets `[]`. +- The labels are the ones you set with `/update_weight_version` or `/finish_weight_update`, see [Async Reinforcement Learning](../../training/async_rl.md#weight-versions-in-outputs). vLLM doesn't order or check them, and the label is `"default"` until you set one. + ## Aborting requests `POST /inference/v1/abort_requests` aborts in-flight requests. It is registered wherever `/inference/v1/generate` is and requires the API key when `--api-key` is set, like `/inference/v1/generate`. diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index 9e75a24eaa1a..b5ae067e3cb1 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -43,10 +43,11 @@ With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionali - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) +- `POST /update_weight_version` - Set the `weight_version` label without changing weights (send `{"new_version": "v8"}`) - `GET /weight_info` - Return the latest committed `weight_version` !!! note "Data Parallelism" - When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. + When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update, and set the weight version on every instance too. ## Typical Async RL Flow @@ -54,12 +55,27 @@ A typical async RL loop with weight syncing looks like this: 1. Start generating rollouts from the current policy 2. Once trainer has new weights to update to, pause generation with `mode="keep"` -3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) +3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) and set their weight version 4. Resume generation -- in-flight requests continue with the new weights 5. Repeat The key insight is that requests paused with `mode="keep"` will produce tokens from the **old** weights before the pause and tokens from the **new** weights after resume. The `clear_cache` parameter controls whether the KV cache is invalidated during the pause. When `clear_cache=True`, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When `clear_cache=False`, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache). +## Weight Versions in Outputs + +To tell which tokens came from which weights, set a weight-version label while generation is paused, after the new weights are loaded: pass `weight_version` to `finish_weight_update`, or call `update_weight_version` (`POST /update_weight_version` over HTTP). Labels are strings you choose, such as the trainer step; vLLM doesn't order or check them. + +The final output of every sequence then carries `weight_versions`, one span per label over its output-token indexes. A request paused with `mode="keep"` across one update gets two spans: + +```python +completion = final_output.outputs[0] +for span in completion.weight_versions: + # The tokens in [span.start, span.end) were sampled by weights `span.version`. + print(span.version, completion.token_ids[span.start : span.end]) +``` + +The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. `/inference/v1/generate` returns the same list on each choice; see [Tokens In <> Tokens Out API](../serving/online_serving/token_in_token_out.md#weight-versions). + ## Example The [async RLHF example](../../examples/rl/rlhf_async_new_apis.py) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation. diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index 081023156edf..3449d0f17f20 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/examples/rl/rlhf_async_new_apis.py @@ -26,8 +26,8 @@ * Pause generation once any request reaches a token threshold. * Broadcast the training model's weights to the inference engine via the NCCL weight transfer engine, replacing the base weights. -* Resume generation and collect results, noting which tokens were - generated before vs. after the weight swap. +* Label the new weights, resume generation, and use `weight_versions` + to locate the weight swap in each output. * Validate correctness by launching a fresh vLLM instance loaded directly with the training model and comparing its output to the post-swap tokens from the weight-synced engine. @@ -61,6 +61,7 @@ MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base" MODEL_NAME_V2 = "Qwen/Qwen3-1.7B" PAUSE_TOKEN_THRESHOLD = 10 +WEIGHT_VERSION_V2 = "v2" ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN" @@ -77,44 +78,27 @@ def __init__(self, **kwargs): log_requests=engine_args.enable_log_requests, log_stats=not engine_args.disable_log_stats, ) - self._generation_paused = False self._request_pause_flag = False async def do_generate( self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams - ) -> tuple[vllm.RequestOutput, int]: - """Generate a single request, setting the request pause flag once the - token count reaches the threshold. - - Returns (output, pause_token_index). pause_token_index is the number - of tokens generated before the weight change, or -1 if no pause. - """ - pause_token_index = -1 - prev_token_count = 0 + ) -> vllm.RequestOutput: + """Generate a request and set the pause flag at the token threshold.""" async for request_output in self.generate( {"prompt_token_ids": prompt_token_ids}, sampling_params, request_id=str(uuid.uuid4()), ): output = request_output - cur_token_count = len(output.outputs[0].token_ids) - if ( - cur_token_count >= PAUSE_TOKEN_THRESHOLD - and not self._request_pause_flag - ): + if len(output.outputs[0].token_ids) >= PAUSE_TOKEN_THRESHOLD: self._request_pause_flag = True - if self._generation_paused and pause_token_index == -1: - pause_token_index = prev_token_count - prev_token_count = cur_token_count - return output, pause_token_index + return output async def pause_after_n_tokens(self): """Wait for any request to set the pause flag, then pause.""" while not self._request_pause_flag: await asyncio.sleep(0) await super().pause_generation(mode="keep") - await asyncio.sleep(5) - self._generation_paused = True @ray.remote(num_gpus=1) @@ -262,9 +246,18 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: ray.get(llm.pause_after_n_tokens.remote()) ray.get(train_model.broadcast_weights.remote()) +# Label the new weights before resuming generation. +ray.get(llm.update_weight_version.remote(WEIGHT_VERSION_V2)) ray.get(llm.resume_generation.remote()) -results = ray.get(gen_futures) +results = [] +for output in ray.get(gen_futures): + completion = output.outputs[0] + pause_idx = next( + (s.start for s in completion.weight_versions if s.version == WEIGHT_VERSION_V2), + len(completion.token_ids), + ) + results.append((output, pause_idx)) for i, (output, pause_idx) in enumerate(results): all_token_ids = list(output.outputs[0].token_ids) @@ -326,7 +319,7 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]: num_pass = 0 num_total = len(results) -for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)): +for i, ((output, pause_idx), val_output) in enumerate(zip(results, val_results)): expected = list(output.outputs[0].token_ids)[pause_idx:] actual = list(val_output.outputs[0].token_ids) match = actual == expected diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 9339c117a889..8945f8b11020 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -228,6 +228,9 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, + /// Decoded for Python wire compatibility; not exposed by the Rust frontend. + #[serde(default)] + weight_version: Option, } /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`. diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index bd8d40dc32f2..c0ca97dbe638 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -139,6 +139,7 @@ class EngineCoreOutputs( finished_requests: set[str] | None = None wave_complete: int | None = None start_wave: int | None = None + weight_version: str | None = None request = EngineCoreRequest( @@ -240,6 +241,7 @@ class EngineCoreOutputs( ) ], finished_requests={"req-1"}, + weight_version="v7", ) extended_outputs = EngineCoreOutputs( @@ -252,6 +254,7 @@ class EngineCoreOutputs( ) ], finished_requests=outputs.finished_requests, + weight_version=outputs.weight_version, ) extended_outputs_bytes = msgspec.msgpack.encode(extended_outputs) # The ordinary frontend's schema ignores even non-default extension values. diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 44c3ee20b5e4..8eb85765993b 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -24,6 +24,7 @@ PoolingRequestOutput, RequestOutput, SamplingMask, + WeightVersionSpan, ) from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike @@ -87,6 +88,7 @@ def test_completion_output_preserves_each_sampling_mask_position() -> None: ] state.routed_experts_chunks = [] state.spec_decode_metrics = None + state.weight_version_starts = [] output = state._new_completion_output([1, 2, 3], FinishReason.LENGTH, None) @@ -1596,6 +1598,7 @@ def test_sampling_masks_follow_output_kind(output_kind): state.sampling_mask_chunks = [] state.routed_experts_chunks = [] state.spec_decode_metrics = None + state.weight_version_starts = [] supports = [[10, 11], [20], [30, 31]] masks = [] @@ -1725,3 +1728,53 @@ def test_async_pooling_cache_miss_does_not_fail_output_processor(): assert output.error is not None assert output.error.retryable assert not output_processor.has_unfinished_requests() + + +@pytest.mark.parametrize( + "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY] +) +@pytest.mark.parametrize( + "steps,expected", + [ + # Merge adjacent equal labels; A -> B -> A remains three spans. + ( + [("v7", [10, 11]), ("v7", [12]), ("v8", [13, 14]), ("v7", [15])], + [("v7", 0, 3), ("v8", 3, 5), ("v7", 5, 6)], + ), + # An output without tokens has no spans. + ([], []), + ], +) +def test_weight_version_spans(output_kind: RequestOutputKind, steps, expected): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + queue = RequestOutputCollector(output_kind, request_id="request-0-int") + request = EngineCoreRequest( + request_id="request-0-int", + external_req_id="request-0", + prompt_token_ids=[1, 2, 3], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=SamplingParams(detokenize=False, output_kind=output_kind), + pooling_params=None, + ) + output_processor.add_request(request, prompt=None, queue=queue) + + for i, (weight_version, token_ids) in enumerate(steps): + # Copy: merging deltas extends token lists in place. + output = EngineCoreOutput( + request_id="request-0-int", new_token_ids=list(token_ids) + ) + output_processor.process_outputs([output], weight_version=weight_version) + if i == 0 and output_kind == RequestOutputKind.DELTA: + intermediate = queue.get_nowait() + assert intermediate is not None + assert intermediate.outputs[0].weight_versions is None + + # Unread deltas merge into the abort output without losing spans. + output_processor.abort_requests(["request-0"], internal=False) + final = queue.get_nowait() + assert final is not None and final.finished + assert final.outputs[0].weight_versions == [WeightVersionSpan(*s) for s in expected] diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 84a6a2c62388..bccc8d8eff0b 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -31,6 +31,7 @@ ) from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.logprobs import Logprob +from vllm.outputs import WeightVersionSpan from vllm.renderers import TokenizeParams from vllm.sampling_params import SamplingParams from vllm.utils import random_uuid @@ -411,6 +412,8 @@ class GenerateChoiceBase(BaseModel): # or (b) `enable_return_routed_experts` is off server-side. routed_experts: str | None = None sampling_mask: list[list[int]] | None = None + # Contiguous weight version spans over output-token indexes. + weight_versions: list[WeightVersionSpan] | None = None @field_validator("token_ids") @classmethod @@ -443,6 +446,8 @@ class GenerateStreamChoiceBase(BaseModel): token_ids: list[int] | None = None routed_experts: str | None = None sampling_mask: list[list[int]] | None = None + # Set on the final chunk only; covers the whole output. + weight_versions: list[WeightVersionSpan] | None = None class GenerateTokensStreamChoice(GenerateStreamChoiceBase): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 816d70bec000..085ae97e16a7 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -453,6 +453,7 @@ async def serve_tokens_full_generator( token_ids=as_list(output.token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, + weight_versions=output.weight_versions, ) if text_mode: text_choices.append( @@ -625,6 +626,7 @@ async def serve_tokens_stream_generator( token_ids=as_list(delta_token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, + weight_versions=output.weight_versions, ) chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse if text_mode: diff --git a/vllm/outputs.py b/vllm/outputs.py index 8e69ffa91d91..145b7d510f71 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -38,6 +38,15 @@ class SamplingMask: token_ids: list[list[int]] +@dataclass +class WeightVersionSpan: + """Output tokens `[start, end)` were sampled by weights labelled `version`.""" + + version: str + start: int + end: int + + @dataclass class CompletionOutput: """The output data of one completion output of a request. @@ -62,6 +71,8 @@ class CompletionOutput: ``--per-request-spec-decode-metrics`` is enabled; None otherwise. Surfaced in the response as ``metrics.speculative_decoding`` for single-sequence (``n == 1``) requests. + weight_versions: Contiguous weight version spans over output-token indexes. + Set on the final output only. """ @@ -76,6 +87,7 @@ class CompletionOutput: lora_request: LoRARequest | None = None sampling_mask: SamplingMask | None = None spec_decode_metrics: RequestSpecDecodeMetrics | None = None + weight_versions: list[WeightVersionSpan] | None = None def finished(self) -> bool: return self.finish_reason is not None @@ -223,6 +235,8 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None: # aggregation with earlier chunks that have no R3. if next_completion.routed_experts is not None: completion.routed_experts = next_completion.routed_experts + # Like finish_reason, spans come only with the terminal output. + completion.weight_versions = next_completion.weight_versions completion.finish_reason = next_completion.finish_reason completion.stop_reason = next_completion.stop_reason else: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index effe8d2416e4..1b14b08246aa 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -282,6 +282,9 @@ class EngineCoreOutputs( # "old" wave, so the next wave needs to be started in other engines. start_wave: int | None = None + # Weight version used for this step. + weight_version: str | None = None + def __post_init__(self): if self.timestamp == 0.0: self.timestamp = time.monotonic() diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 08445d8f913a..71e30990c578 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -855,7 +855,10 @@ async def output_handler(): outputs_slice = engine_core_outputs[start:end] # 2) Process EngineCoreOutputs. processed_outputs = output_processor.process_outputs( - outputs_slice, outputs.timestamp, iteration_stats + outputs_slice, + outputs.timestamp, + iteration_stats, + weight_version=outputs.weight_version, ) # NOTE: RequestOutputs are pushed to their queues. assert not processed_outputs.request_outputs diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b13764abec3a..bbca64a58ae9 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -639,6 +639,11 @@ def _attach_iteration_details( else: eco.scheduler_stats.iteration_details = iteration_details + def _attach_weight_version(self, outputs: dict[int, EngineCoreOutputs]) -> None: + # Weight versions change only while generation is paused. + for eco in outputs.values(): + eco.weight_version = self._weight_version + def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). Overridden by the DP engine core; never throttles otherwise.""" @@ -672,6 +677,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 @@ -773,6 +779,7 @@ def step_with_batch_queue( scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 2dd70fc17917..330c114a86dd 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -335,6 +335,7 @@ def step(self) -> list[RequestOutput | PoolingRequestOutput]: outputs.outputs, engine_core_timestamp=outputs.timestamp, iteration_stats=iteration_stats, + weight_version=outputs.weight_version, ) mm_processor_cache = self.renderer.mm_processor_cache diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index abb45f7239be..263a28fb3901 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -19,6 +19,7 @@ RequestError, RequestOutput, SamplingMask, + WeightVersionSpan, ) from vllm.sampling_params import RequestOutputKind from vllm.tokenizers import TokenizerLike @@ -190,6 +191,8 @@ def __init__( # Routed experts accumulation (prompt + sample chunks) self.routed_experts_chunks: list[np.ndarray] = [] self.sampling_mask_chunks: list[SamplingMaskLists] = [] + # (weight version, first output-token index) for each span. + self.weight_version_starts: list[tuple[str, int]] = [] # Stream Interval self.stream_interval = stream_interval @@ -471,8 +474,25 @@ def _new_completion_output( finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, spec_decode_metrics=self.spec_decode_metrics if finished else None, + weight_versions=self._weight_version_spans() if finished else None, ) + def track_weight_version(self, weight_version: str) -> None: + """Start a span when the weight version changes.""" + assert self.detokenizer is not None + starts = self.weight_version_starts + if not starts or starts[-1][0] != weight_version: + starts.append((weight_version, self.detokenizer.num_output_tokens())) + + def _weight_version_spans(self) -> list[WeightVersionSpan] | None: + assert self.detokenizer is not None + num_tokens = self.detokenizer.num_output_tokens() + starts = self.weight_version_starts + if not starts and num_tokens: + return None # No label was reported for these tokens. + ends = [start for _, start in starts[1:]] + [num_tokens] + return [WeightVersionSpan(v, s, e) for (v, s), e in zip(starts, ends)] + def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: return PoolingOutput(data=pooling_output) @@ -659,6 +679,7 @@ def process_outputs( engine_core_outputs: list[EngineCoreOutput], engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, + weight_version: str | None = None, ) -> OutputProcessorOutput: """Process the EngineCoreOutputs: 1) Compute stats for logging @@ -736,6 +757,8 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None + if new_token_ids and weight_version is not None: + req_state.track_weight_version(weight_version) # 2) Detokenize the token ids into text and perform stop checks. num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( From 01abe975b5837f9e0722a196a28e57331735e0a6 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:55:41 +0800 Subject: [PATCH 359/952] [Bugfix][KV Connector] Fix speculative block reservations across KV cache groups (#59880) Signed-off-by: bzsuni --- .../unit/test_remote_prefill_lifecycle.py | 133 ++++++++++++++---- tests/v1/kv_connector/unit/utils.py | 2 + vllm/v1/core/sched/scheduler.py | 34 ++--- 3 files changed, 128 insertions(+), 41 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 59fcc678bfb9..b72de811c0f9 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +from dataclasses import replace from unittest.mock import patch import pytest @@ -761,35 +762,54 @@ def test_async_loads_both_admitted_when_pool_fits(): assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS -def test_async_load_reserves_blocks_for_promotion_margin(): - """An async load is not admitted unless the blocks its own promotion will - need are still free. - - A parked load is allocated without lookahead slots, but promotion pads it - to ``1 + num_spec_tokens`` and asks for the lookahead margin on top. Without - reserving that margin, two loads can be admitted that together consume the - whole pool; the head then fails ``allocate_slots`` forever while the load - behind it is never reached, and with nothing running no block is ever freed. +@pytest.mark.parametrize( + ("swa_enabled", "block_size", "num_spec_tokens", "num_blocks"), + [(False, 16, 3, 8), (True, 16, 3, 17), (True, 64, 31, 40)], + ids=["single-group", "multiple-groups", "different-block-sizes"], +) +def test_async_load_reserves_blocks_for_promotion_margin( + swa_enabled, block_size, num_spec_tokens, num_blocks +): + """Async loads must leave enough blocks for promotion in every cache group. - req_a (4 blocks) and req_b (3 blocks) exactly fill the 7 usable blocks, so - req_b must be held back in WAITING holding no blocks, leaving req_a room to - be promoted and run. + Otherwise two loads can fill the pool and prevent the head from being + promoted, with nothing running to free blocks and break the deadlock. """ - vllm_config = create_vllm_config(num_speculative_tokens=3) - BLOCK_SIZE = vllm_config.cache_config.block_size - scheduler = create_scheduler(vllm_config, num_blocks=8) # usable = 7 + vllm_config = create_vllm_config( + block_size=block_size, + num_speculative_tokens=num_spec_tokens, + disable_hybrid_kv_cache_manager=False, + ) + hash_block_size = 16 + kv_cache_config = make_kv_cache_config( + hash_block_size, swa_enabled=swa_enabled, sw_size=4096, num_blocks=num_blocks + ) + full_group = kv_cache_config.kv_cache_groups[0] + full_group.kv_cache_spec = replace(full_group.kv_cache_spec, block_size=block_size) + if swa_enabled: + sliding_group = kv_cache_config.kv_cache_groups[-1] + sliding_group.kv_cache_spec = replace( + sliding_group.kv_cache_spec, + page_size_padded=full_group.kv_cache_spec.page_size_bytes, + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=hash_block_size, + ) req_a = create_request( request_id=1, - block_size=BLOCK_SIZE, - num_tokens=BLOCK_SIZE * 4, + block_size=hash_block_size, + num_tokens=block_size * 4, do_remote_prefill=True, max_tokens=1, ) req_b = create_request( request_id=2, - block_size=BLOCK_SIZE, - num_tokens=BLOCK_SIZE * 3, + block_size=hash_block_size, + num_tokens=block_size * 3, do_remote_prefill=True, max_tokens=1, ) @@ -800,16 +820,13 @@ def test_async_load_reserves_blocks_for_promotion_margin(): with patch.object( scheduler.connector, "get_num_new_matched_tokens", - side_effect=[(BLOCK_SIZE * 4, True), (BLOCK_SIZE * 3, True)], + side_effect=[(block_size * 4, True), (block_size * 3, True)], ): scheduler.schedule() assert req_a.status == RequestStatus.WAITING_FOR_REMOTE_KVS assert req_b.status == RequestStatus.WAITING - req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[ - 0 - ].req_to_blocks - assert req_b.request_id not in req_to_blocks + assert not any(scheduler.kv_cache_manager.get_block_ids(req_b.request_id)) # req_a's load lands: it must be promotable and actually get scheduled. scheduler.update_from_output( @@ -821,6 +838,74 @@ def test_async_load_reserves_blocks_for_promotion_margin(): assert scheduler_output.num_scheduled_tokens[req_a.request_id] > 0 +@pytest.mark.parametrize( + ("mamba_cache_mode", "mamba_block_size", "num_blocks"), + [("none", 128, 18), ("align", 16, 18)], +) +def test_async_load_does_not_double_reserve_mamba_spec_blocks( + mamba_cache_mode, mamba_block_size, num_blocks +): + """Two hybrid loads fit when Mamba's speculative states are counted once.""" + block_size = 16 + num_spec_tokens = 3 + vllm_config = create_vllm_config( + block_size=mamba_block_size, + max_model_len=128, + num_speculative_tokens=num_spec_tokens, + disable_hybrid_kv_cache_manager=False, + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + kv_cache_config = make_kv_cache_config( + block_size, + mamba_enabled=True, + mamba_cache_mode=mamba_cache_mode, + num_blocks=num_blocks, + ) + mamba_group = kv_cache_config.kv_cache_groups[-1] + mamba_group.kv_cache_spec = replace( + mamba_group.kv_cache_spec, + block_size=mamba_block_size, + num_speculative_blocks=num_spec_tokens, + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=block_size, + ) + reqs = [ + create_request( + request_id=i, + block_size=block_size, + num_tokens=num_tokens, + do_remote_prefill=True, + max_tokens=1, + ) + for i, num_tokens in enumerate((64, 48), 1) + ] + for req in reqs: + scheduler.add_request(req) + + with patch.object( + scheduler.connector, + "get_num_new_matched_tokens", + side_effect=[(req.num_tokens, True) for req in reqs], + ): + scheduler.schedule() + + assert all(req.status == RequestStatus.WAITING_FOR_REMOTE_KVS for req in reqs) + scheduler.update_from_output( + scheduler.schedule(), + create_model_runner_output( + [], finished_recving={req.request_id for req in reqs} + ), + ) + output = scheduler.schedule() + for req in reqs: + assert req.status == RequestStatus.RUNNING + assert output.num_scheduled_tokens[req.request_id] > 0 + + def test_deferred_lookup_does_not_block_parked_load(): """A request with a deferred connector lookup must not strand an async load queued behind it. diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index d738e183952f..28bfda419a76 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -177,6 +177,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -207,6 +208,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 0758c6122565..2603fa10a75e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -344,6 +344,17 @@ def __init__( self.cache_config.enable_mamba_shared_prefix_checkpoint ), ) + # Mamba managers already account for their speculative state blocks. + num_spec_decode_tokens = ( + 1 + self.num_spec_tokens + self.num_lookahead_tokens + if self.num_spec_tokens + else 0 + ) + self._spec_decode_step_blocks = sum( + cdiv(num_spec_decode_tokens, manager.block_size) + for manager in self.kv_cache_manager.coordinator.single_type_managers + if not isinstance(manager.kv_cache_spec, MambaSpec) + ) # Bind after construction so connectors can access the cache manager. if self.connector is not None: self.connector.bind_kv_cache_manager(self.kv_cache_manager) @@ -1236,7 +1247,7 @@ def skip_request(from_queue: RequestQueue) -> None: # predictable preemptions. reserved_blocks = ( self._inflight_prefill_reserved_blocks() - + self._spec_decode_step_blocks() + + self._spec_decode_step_blocks ) # Replayed tokens are already counted in the adopted hit; a @@ -3059,9 +3070,12 @@ def _connector_finished( def _request_remaining_blocks(self, request: Request) -> int: """Blocks `request` still needs to hold its full sequence and be promoted.""" full_num_tokens = min(request.num_tokens, self.max_model_len) - num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + num_tokens = full_num_tokens + if self.num_spec_tokens: + num_tokens += 1 + self.num_spec_tokens + self.num_lookahead_tokens + return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, - num_tokens=full_num_tokens, + num_tokens=num_tokens, new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, num_encoder_tokens=0, total_computed_tokens=request.num_computed_tokens, @@ -3070,20 +3084,6 @@ def _request_remaining_blocks(self, request: Request) -> int: apply_admission_cap=True, prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1), ) - return num_blocks + self._spec_decode_step_blocks() - - def _spec_decode_step_blocks(self) -> int: - """Number of blocks for the extra KV slots a spec decode step needs. - - When using async kv load, scheduler must reserve enough blocks for - full sequence + the spec decode step, otherwise request cannot be - able to run after the async_load if we are out of kv blocks. - """ - if not self.num_spec_tokens: - return 0 - return cdiv( - 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size - ) def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None: kv_fetch_stages = self._kv_fetch_stages From f8471164bf9f5190752fd7b903d415a34fc4ba1c Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 10 Oct 2026 21:15:05 -0700 Subject: [PATCH 360/952] [MRV2] Detatch weight-reload logic from MRV1 (#60937) --- tests/v1/worker/test_gpu_model_runner.py | 61 ++++++--- .../worker/test_gpu_worker_weight_transfer.py | 4 +- .../model_loader/reload/inplace.py | 123 ++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 33 ++--- vllm/v1/worker/gpu_model_runner.py | 121 ++--------------- vllm/v1/worker/gpu_worker.py | 23 +++- 6 files changed, 215 insertions(+), 150 deletions(-) create mode 100644 vllm/model_executor/model_loader/reload/inplace.py diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 158befc15027..0578768d68c7 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -10,6 +10,7 @@ import pytest import torch +import vllm.model_executor.model_loader.reload.inplace as reload_inplace_module import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module from vllm.config import ( AttentionConfig, @@ -63,6 +64,7 @@ from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu_input_batch import InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.utils import select_common_block_size BLOCK_SIZE = 16 @@ -1045,13 +1047,22 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent( ) +def _update_config(model_runner: GPUModelRunner, overrides: dict) -> None: + # Apply overrides via the worker without constructing a full worker. + worker = object.__new__(Worker) + worker.vllm_config = model_runner.vllm_config + worker.model_runner = model_runner + worker.update_config(overrides) + + def test_update_config(model_runner): # Simple update - model_runner.update_config({"load_config": {"load_format": "dummy"}}) + _update_config(model_runner, {"load_config": {"load_format": "dummy"}}) assert model_runner.load_config.load_format == "dummy" + assert model_runner.vllm_config.load_config is model_runner.load_config # Raise error on non-existing config with pytest.raises(ValueError, match="do_not_exist_config"): - model_runner.update_config({"do_not_exist_config": "dummy"}) + _update_config(model_runner, {"do_not_exist_config": "dummy"}) def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2): @@ -1059,12 +1070,14 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2): # model_runner_2 loads dummy weights first then load real weights inplace model_runner.load_model() original_load_format = model_runner_2.load_config.load_format - model_runner_2.update_config({"load_config": {"load_format": "dummy"}}) + _update_config(model_runner_2, {"load_config": {"load_format": "dummy"}}) model_runner_2.load_model() # Initial model loading with dummy weights assert str(model_runner.get_model().state_dict()) != str( model_runner_2.get_model().state_dict() ) - model_runner_2.update_config({"load_config": {"load_format": original_load_format}}) + _update_config( + model_runner_2, {"load_config": {"load_format": original_load_format}} + ) model_runner_2.reload_weights() # Load real weights inplace assert str(model_runner.get_model().state_dict()) == str( model_runner_2.get_model().state_dict() @@ -1083,26 +1096,34 @@ def test_reload_weights_path_replaces_object_storage_source(monkeypatch): # silently re-streams the original checkpoint. loader = Mock() loader.get_all_weights.return_value = iter(()) - monkeypatch.setattr(gpu_model_runner_module, "get_model_loader", lambda _: loader) - monkeypatch.setattr(gpu_model_runner_module, "initialize_layerwise_reload", Mock()) - monkeypatch.setattr(gpu_model_runner_module, "finalize_layerwise_reload", Mock()) + monkeypatch.setattr(reload_inplace_module, "get_model_loader", lambda _: loader) + monkeypatch.setattr(reload_inplace_module, "initialize_layerwise_reload", Mock()) + monkeypatch.setattr(reload_inplace_module, "finalize_layerwise_reload", Mock()) - runner = Mock(lora_config=None) - runner.model_config = SimpleNamespace( - model="/tmp/pulled-config-files", - model_weights="s3://bucket/original", - revision="abc123", - quantization=None, + vllm_config = SimpleNamespace( + lora_config=None, + load_config=SimpleNamespace(load_format="runai_streamer"), + model_config=SimpleNamespace( + model="/tmp/pulled-config-files", + model_weights="s3://bucket/original", + revision="abc123", + quantization=None, + ), ) - runner.get_model.return_value.named_parameters.return_value = [] - runner.get_model.return_value.load_weights.return_value = None - - GPUModelRunner.reload_weights(runner, weights_path="org/new-model") + model = Mock() + model.named_parameters.return_value = [] + model.load_weights.return_value = None - loader.get_all_weights.assert_called_once_with( - runner.model_config, runner.get_model.return_value + reload_inplace_module.reload_weights( + vllm_config, + model, + weights_iterator=None, + weights_path="org/new-model", + is_checkpoint_format=True, ) - cfg = runner.model_config + + loader.get_all_weights.assert_called_once_with(vllm_config.model_config, model) + cfg = vllm_config.model_config assert (cfg.model, cfg.model_weights, cfg.revision) == ("org/new-model", "", None) diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index 49ef70789478..f1965cb1fae2 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -13,7 +13,7 @@ from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config from vllm.lora.layers import BaseLayerWithLoRA -from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload +from vllm.model_executor.model_loader.reload.inplace import get_parameter_for_reload from vllm.v1.worker.gpu_worker import Worker @@ -95,7 +95,7 @@ def test_reload_parameter_lookup_preserves_lora_module_names(): assert set(named_parameters) == {"proj.base_layer.qweight"} assert named_parameters["proj.base_layer.qweight"] is qweight assert model.get_parameter("proj.base_layer.qweight") is qweight - assert _get_parameter_for_reload(model, "proj.qweight") is qweight + assert get_parameter_for_reload(model, "proj.qweight") is qweight def test_start_update_finish_delegates_to_engine(): diff --git a/vllm/model_executor/model_loader/reload/inplace.py b/vllm/model_executor/model_loader/reload/inplace.py new file mode 100644 index 000000000000..dfe11ed4e745 --- /dev/null +++ b/vllm/model_executor/model_loader/reload/inplace.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time +from collections.abc import Iterable +from typing import cast + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader import get_model_loader + +from .layerwise import finalize_layerwise_reload, initialize_layerwise_reload + +logger = init_logger(__name__) + +__all__ = ["get_parameter_for_reload", "reload_weights"] + + +def get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: + """Resolve checkpoint names without changing the model's module tree.""" + from vllm.lora.layers import BaseLayerWithLoRA + + module_name, _, parameter_name = name.rpartition(".") + module = model.get_submodule(module_name) + if isinstance(module, BaseLayerWithLoRA): + module = module.base_layer + return module.get_parameter(parameter_name) + + +def reload_weights( + vllm_config: VllmConfig, + model: nn.Module, + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None, + weights_path: str | None, + is_checkpoint_format: bool, +) -> None: + """Reload a model's weights in place from a weights iterator or from disk. + + Callers are responsible for invalidating any state derived from the old + weights (e.g. LoRA, encoder and multimodal caches). + + Args: + vllm_config: config of the model being reloaded. If `weights_path` is + given, `vllm_config.model_config` is updated to point at it. + model: model to load weights into + weights_iterator: weights to load into model + weights_path: path to load weights from if weights_iterator is not + provided. Use path of original model if neither is provided. + is_checkpoint_format: set to False if weights have already been + processed into kernel format (repacking, renaming, etc.) + + """ + # TODO(@kylesayrs): generalize to all runners and loaders + # argument validation + if weights_iterator is None and not is_checkpoint_format: + logger.warning( + "Reloading from disk means that weights will be in checkpoint format. " + "Please use `is_checkpoint_format=True` " + "to avoid weight reloading errors" + ) + + model_config = vllm_config.model_config + load_config = vllm_config.load_config + weights_to_load = { + name.replace(".base_layer.", ".") if vllm_config.lora_config else name + for name, _ in model.named_parameters() + } + counter_before_reloading = time.perf_counter() + + # load weights from disk if none are provided + if weights_iterator is None: + model_loader = get_model_loader(load_config) + if not hasattr(model_loader, "get_all_weights"): + raise NotImplementedError( + f"Model reloading with `{load_config.load_format}` format" + ) + + if weights_path is not None: + # The revision and any object-storage `model_weights` source + # belong to the model we are reloading away from, so they must + # not be carried over to the new path. + model_config.model = weights_path + model_config.model_weights = "" + model_config.revision = None + weights_iterator = model_loader.get_all_weights(model_config, model) + weights_iterator = cast(Iterable[tuple[str, torch.Tensor]], weights_iterator) + + # begin loading weights + logger.info_once("Reloading weights inplace...") + if is_checkpoint_format: + # load weights from checkpoint/ original model format + initialize_layerwise_reload(model) + loaded_weights = model.load_weights(weights_iterator) + finalize_layerwise_reload(model, model_config) + + else: + # load weights from kernel format + logger.warning_once( + "Reloading with `is_checkpoint_format=True` requires that " + "weights be in kernel format and already sharded", + ) + loaded_weights = set() + for name, loaded_weight in weights_iterator: + param = get_parameter_for_reload(model, name) # TODO: buffers? + param.copy_(loaded_weight) + loaded_weights.add(name) + + # logging and validation + counter_after_reloading = time.perf_counter() + diff_seconds = counter_after_reloading - counter_before_reloading + logger.info_once( + "Reloading and processing weights took %.2f seconds", + diff_seconds, + ) + if model_config.quantization is None and loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + logger.warning( + "Following weights were not loaded from checkpoint: %s", + weights_not_loaded, + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e04a381b94e6..20fa4c17bb07 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -19,6 +19,7 @@ import functools import gc import time +from collections.abc import Iterable from contextlib import AbstractContextManager from copy import deepcopy from dataclasses import replace @@ -46,6 +47,7 @@ initialize_mamba_ssu_backend, ) from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.reload.inplace import reload_weights from vllm.model_executor.models.interfaces import requires_raw_input_tokens from vllm.model_executor.offloader import ( create_offloader, @@ -563,22 +565,21 @@ def get_draft_model(self) -> nn.Module | None: return None return speculator.model - def reload_weights(self, *args, **kwargs) -> None: - # TODO(Wentao): Use full version instead of import when fully migrated to v2 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 - - GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type] - - def update_config(self, *args, **kwargs) -> None: - # TODO(Wentao): Use full version instead of import when fully migrated to v2 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 - - GPUModelRunnerV1.update_config(self, *args, **kwargs) # type: ignore[arg-type] - - # v2 reads config via self.vllm_config (e.g. in load_model), so keep it - # in sync with the attributes the v1 helper just replaced. - self.vllm_config.model_config = self.model_config - self.vllm_config.load_config = self.load_config + def reload_weights( + self, + weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None, + weights_path: str | None = None, + is_checkpoint_format: bool = True, + ) -> None: + """Reload weights in place.""" + vllm_config = self.vllm_config + model = self.get_model() + reload_weights( + vllm_config, model, weights_iterator, weights_path, is_checkpoint_format + ) + self.reset_lora_state() + self.reset_encoder_cache() + self.reset_mm_cache() @functools.cached_property def main_stream(self) -> torch.cuda.Stream: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 90bea22c9d4f..4db3f00b91d4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -34,7 +34,6 @@ VllmConfig, get_layers_from_vllm_config, set_current_vllm_config, - update_config, ) from vllm.config.cache import CacheConfig from vllm.config.ec_manager_config import EncoderCacheManagerMetadata @@ -57,7 +56,7 @@ set_forward_context, ) from vllm.logger import init_logger -from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType +from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager @@ -67,10 +66,7 @@ ) from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.model_loader import get_model_loader -from vllm.model_executor.model_loader.reload import ( - finalize_layerwise_reload, - initialize_layerwise_reload, -) +from vllm.model_executor.model_loader.reload.inplace import reload_weights from vllm.model_executor.models.interfaces import ( MixtureOfExperts, MultiModalEmbeddings, @@ -247,15 +243,6 @@ logger = init_logger(__name__) -def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter: - """Resolve checkpoint names without changing the model's module tree.""" - module_name, _, parameter_name = name.rpartition(".") - module = model.get_submodule(module_name) - if isinstance(module, BaseLayerWithLoRA): - module = module.base_layer - return module.get_parameter(parameter_name) - - AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] # list when ubatching is enabled PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict @@ -5203,19 +5190,6 @@ def propose_draft_token_ids( return draft_token_ids - def update_config(self, overrides: dict[str, Any]) -> None: - allowed_config_names = {"load_config", "model_config"} - for config_name, config_overrides in overrides.items(): - if config_name not in allowed_config_names: - allowed = ", ".join(sorted(allowed_config_names)) - raise ValueError( - f"Config override '{config_name}' is not supported. " - f"Supported configs: {allowed}" - ) - config = getattr(self, config_name) - new_config = update_config(config, config_overrides) - setattr(self, config_name, new_config) - @instrument(span_name="Loading (GPU)") def load_model( self, @@ -5472,89 +5446,16 @@ def reload_weights( weights_path: str | None = None, is_checkpoint_format: bool = True, ) -> None: - """Reload weights from a weights iterator or from disk - - Args: - weights_iterator: weights to load into model - weights_path: path to load weights from if weights_iterator is not - provided. Use path of original model if neither is provided. - is_checkpoint_format: set to False if weights have already been - processed into kernel format (repacking, renaming, etc.) - - """ - # TODO(@kylesayrs): generalize to all runners and loaders - # argument validation - if weights_iterator is None and not is_checkpoint_format: - logger.warning( - "Reloading from disk means that weights will be in checkpoint format. " - "Please use `is_checkpoint_format=True` " - "to avoid weight reloading errors" - ) - - model = self.get_model() - weights_to_load = { - name.replace(".base_layer.", ".") if self.lora_config else name - for name, _ in model.named_parameters() - } - counter_before_reloading = time.perf_counter() - - # load weights from disk if none are provided - if weights_iterator is None: - model_loader = get_model_loader(self.load_config) - if not hasattr(model_loader, "get_all_weights"): - raise NotImplementedError( - f"Model reloading with `{self.load_config.load_format}` format" - ) - - if weights_path is not None: - # The revision and any object-storage `model_weights` source - # belong to the model we are reloading away from, so they must - # not be carried over to the new path. - self.model_config.model = weights_path - self.model_config.model_weights = "" - self.model_config.revision = None - weights_iterator = model_loader.get_all_weights(self.model_config, model) - weights_iterator = cast( - Iterable[tuple[str, torch.Tensor]], weights_iterator - ) - - # begin loading weights - logger.info_once("Reloading weights inplace...") - if is_checkpoint_format: - # load weights from checkpoint/ original model format - initialize_layerwise_reload(model) - loaded_weights = model.load_weights(weights_iterator) - finalize_layerwise_reload(model, self.model_config) - - else: - # load weights from kernel format - logger.warning_once( - "Reloading with `is_checkpoint_format=True` requires that " - "weights be in kernel format and already sharded", - ) - loaded_weights = set() - for name, loaded_weight in weights_iterator: - param = _get_parameter_for_reload(model, name) # TODO: buffers? - param.copy_(loaded_weight) - loaded_weights.add(name) - - self.reset_lora_state() - - # logging and validation - counter_after_reloading = time.perf_counter() - diff_seconds = counter_after_reloading - counter_before_reloading - logger.info_once( - "Reloading and processing weights took %.2f seconds", - diff_seconds, + """Reload weights in place; see `reload_weights` in + `vllm.model_executor.model_loader.reload.inplace` for details.""" + reload_weights( + self.vllm_config, + self.get_model(), + weights_iterator=weights_iterator, + weights_path=weights_path, + is_checkpoint_format=is_checkpoint_format, ) - if self.model_config.quantization is None and loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded: - logger.warning( - "Following weights were not loaded from checkpoint: %s", - weights_not_loaded, - ) - + self.reset_lora_state() self.reset_encoder_cache() self.reset_mm_cache() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a01e5dd5f7a3..37fa00bb7efa 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -16,7 +16,12 @@ import torch.nn as nn import vllm.envs as envs -from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config +from vllm.config import ( + CUDAGraphMode, + VllmConfig, + set_current_vllm_config, + update_config, +) from vllm.config.compilation import CompilationMode from vllm.config.profiler import validate_profile_prefix from vllm.device_allocator import get_mem_allocator_instance @@ -591,7 +596,21 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None: set_torch_threads_for_runtime() def update_config(self, overrides: dict[str, Any]) -> None: - self.model_runner.update_config(overrides) + """Apply config overrides to the worker, its model runner and the + shared VllmConfig.""" + allowed_config_names = {"load_config", "model_config"} + for config_name, config_overrides in overrides.items(): + if config_name not in allowed_config_names: + allowed = ", ".join(sorted(allowed_config_names)) + raise ValueError( + f"Config override '{config_name}' is not supported. " + f"Supported configs: {allowed}" + ) + config = getattr(self.vllm_config, config_name) + new_config = update_config(config, config_overrides) + setattr(self.vllm_config, config_name, new_config) + setattr(self.model_runner, config_name, new_config) + setattr(self, config_name, new_config) def reload_weights(self, *args, **kwargs) -> None: with set_current_vllm_config(self.vllm_config): From 2eb491e18b1ce5b8a25e4c88f32dbb5d028d525c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Sat, 10 Oct 2026 22:40:45 -0700 Subject: [PATCH 361/952] [Perf][MiniMax-M3] MSA prefill attends with the E4M3 query; bump fmha_sm100 to 968061a (#60684) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 --- cmake/external_projects/fmha_sm100.cmake | 11 +- ...st_minimax_m3_msa_cutlass_sparse_decode.py | 179 +++++++++++++++--- vllm/models/minimax_m3/nvidia/model.py | 2 +- .../minimax_m3/nvidia/sparse_attention_msa.py | 16 +- 4 files changed, 176 insertions(+), 32 deletions(-) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index f6bb0b8acee7..6b84b35c9049 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG be4a52aea420b6bef73be9af1a649859964790cb + GIT_TAG 968061ab16cb14991a5c7d4c0e8eabf531be0dcf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -36,11 +36,14 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100") install(FILES "${FMHA_SM100_PY_ROOT}/__init__.py" + "${FMHA_SM100_PY_ROOT}/_jit_cache.py" "${FMHA_SM100_PY_ROOT}/api.py" "${FMHA_SM100_PY_ROOT}/bench_utils.py" "${FMHA_SM100_PY_ROOT}/jit.py" + "${FMHA_SM100_PY_ROOT}/msa_warmup.py" "${FMHA_SM100_PY_ROOT}/nvfp4_kv.py" "${FMHA_SM100_PY_ROOT}/q8kv4_decode_adapter.py" + "${FMHA_SM100_PY_ROOT}/q8kv4_prefill_adapter.py" "${FMHA_SM100_PY_ROOT}/sparse.py" "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py" DESTINATION vllm/third_party/fmha_sm100 @@ -65,6 +68,12 @@ install(DIRECTORY "${FMHA_SM100_PY_ROOT}/kvouter/" PATTERN "__pycache__" EXCLUDE PATTERN "*.pyc" EXCLUDE) +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/prefill_q8kv4/" + DESTINATION vllm/third_party/fmha_sm100/prefill_q8kv4 + COMPONENT fmha_sm100 + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE) + install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/" DESTINATION vllm/third_party/fmha_sm100/cute COMPONENT fmha_sm100 diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index 0007f20c5bfc..42e2b86965e0 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -782,6 +782,46 @@ def _make_nvfp4_kv_cache( return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16) +def _e4m3_staged_kv_reference( + kv_reference: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor +) -> torch.Tensor: + """The dequantized NVFP4 cache as an E4M3-query kernel reads it: each + code * block_scale rounded to E4M3, then the global scale.""" + sides = [] + for side, global_scale in zip( + kv_reference.float().split(HEAD_DIM, dim=-1), (k_scale, v_scale) + ): + sides.append( + (side / global_scale).to(torch.float8_e4m3fn).float() * global_scale + ) + return torch.cat(sides, dim=-1).to(torch.bfloat16) + + +def _assert_close_fp8_prefill(actual: torch.Tensor, expected: torch.Tensor) -> None: + """The fp8-q prefill rounds each split's probabilities to E4M3, and the splits + follow the SM count, so a rare element moves past the BF16 tolerance.""" + close = torch.isclose(actual, expected, atol=0.05, rtol=0.05) + assert close.float().mean().item() > 1 - 1e-4 + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.05) + + +def _record_query_dtypes( + monkeypatch: pytest.MonkeyPatch, kernel_name: str +) -> list[torch.dtype]: + """Record the query dtype of every ``fmha_sm100.sparse.`` call.""" + from vllm.third_party.fmha_sm100 import sparse as msa_sparse + + kernel = getattr(msa_sparse, kernel_name) + query_dtypes: list[torch.dtype] = [] + + def record(q, *args, **kwargs): + query_dtypes.append(q.dtype) + return kernel(q, *args, **kwargs) + + monkeypatch.setattr(msa_sparse, kernel_name, record) + return query_dtypes + + def test_nvfp4_kv_cache_views_address_head_slots() -> None: """Slot 2h (K) and 2h + 1 (V) hold head h's data then scales, so each head is one contiguous run of the page: a TP rank's page is a byte slice.""" @@ -898,27 +938,14 @@ def run() -> None: torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) -@pytest.mark.parametrize( - ("num_q_heads", "num_kv_heads"), - [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], -) -def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( - num_q_heads: int, +def _make_prefill_metadata( + q_lens: list[int], + seq_lens_list: list[int], + block_table: torch.Tensor, num_kv_heads: int, -) -> None: - """MSA prefill reads NVFP4 pages as their exact dequantized values, - including chunked prefill over cached context, and takes q from its fp8 - copy (the only one the fused insert writes on this path).""" - torch.manual_seed(0) - q_lens = [300, 128, 77] - seq_lens_list = [300, 640, 1100] - block_table, num_pages = _make_paged_layout(seq_lens_list) - k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") - v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") - kv_cache, kv_reference = _make_nvfp4_kv_cache( - num_pages, num_kv_heads, "nvfp4", k_scale, v_scale - ) - +) -> tuple[MiniMaxM3SparseMetadata, torch.Tensor]: + """Prefill-only metadata (chunks over cached context) and a top-k selecting + every causally visible page (<= TOPK per query), ending with the local one.""" total_q = sum(q_lens) seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") context_lens = seq_lens - torch.tensor(q_lens, dtype=torch.int32, device="cuda") @@ -948,7 +975,6 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( prefill=prefill, ) - # Every causally visible page (<= TOPK per query), ending with the local one. topk = torch.full( (total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device="cuda" ) @@ -959,6 +985,35 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( assert visible <= TOPK topk[token, :, :visible] = torch.arange(visible, device="cuda") token += 1 + return metadata, topk + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """MSA prefill reads NVFP4 pages as their dequantized values, staged to E4M3 + for the E4M3 q, including chunked prefill over cached context; it attends + with the fp8 q (the only one the fused insert writes on this path).""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda") + v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda") + kv_cache, kv_reference = _make_nvfp4_kv_cache( + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale + ) + + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + total_q = sum(q_lens) layer_name = "model.layers.0.self_attn.attn" q_scale = 0.5 @@ -976,6 +1031,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( / q_scale ).to(torch.float8_e4m3fn) query = query_fp8.to(torch.bfloat16) * q_scale + kv_reference = _e4m3_staged_kv_reference(kv_reference, k_scale, v_scale) common = dict( num_heads=num_q_heads, head_size=HEAD_DIM, @@ -984,6 +1040,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( topk_blocks=TOPK, sparse_block_size=BLOCK_SIZE, ) + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_nvfp4_kv_func") forward_context = ForwardContext( no_compile_layers={}, attn_metadata={layer_name: metadata}, @@ -1001,7 +1058,8 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( query_fp8=query_fp8, ) current_platform.synchronize() - torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: @@ -1079,3 +1137,80 @@ def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: ) current_platform.synchronize() torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_reads_query_fp8_over_fp8_cache( + monkeypatch: pytest.MonkeyPatch, + num_q_heads: int, + num_kv_heads: int, +) -> None: + """With CUTLASS decode on, FP8 prefill (incl. chunks over cached context) + attends with the fp8 q the fused insert writes, not a BF16 copy, and + matches Triton on the dequantized q.""" + torch.manual_seed(0) + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] + block_table, num_pages = _make_paged_layout(seq_lens_list) + kv_cache = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + 2 * HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.5 + ).to(torch.float8_e4m3fn) + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + layer_name = "model.layers.0.self_attn.attn" + q_scale = 0.5 + layer = SimpleNamespace( + layer_name=layer_name, + topk_indices_buffer=topk, + _q_scale_float=q_scale, + ) + query_fp8 = ( + torch.randn( + sum(q_lens), num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + / q_scale + ).to(torch.float8_e4m3fn) + query = query_fp8.to(torch.bfloat16) * q_scale + common = dict( + num_heads=num_q_heads, + head_size=HEAD_DIM, + scale=SM_SCALE, + num_kv_heads=num_kv_heads, + kv_cache_dtype="fp8", + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + ) + impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass") + assert impl.use_cutlass_decode + query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_func") + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: metadata}, + slot_mapping={}, + ) + with override_forward_context(forward_context): + expected = MiniMaxM3SparseTritonImpl(**common).forward( + layer, query, kv_cache, torch.empty_like(query) + ) + actual = impl.forward( + layer, + torch.full_like(query, float("nan")), + kv_cache, + torch.zeros_like(query), + query_fp8=query_fp8, + ) + current_platform.synchronize() + assert query_dtypes == [torch.float8_e4m3fn] + _assert_close_fp8_prefill(actual, expected) diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 06c28a7a38b4..d67e9ea19d90 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -614,7 +614,7 @@ def forward( index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) # With query_fp8, q is emitted only in fp8; the attend dequantizes it - # into ``q`` for the tokens that run on bf16-query kernels. + # into ``q`` for the tokens that run on Triton decode. query_fp8 = self._allocate_query_fp8(qkv) # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); # the fused kernel emits fp8 directly when this buffer is e4m3. diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index 5aa23b4ac849..ed8ca2a78474 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -266,7 +266,7 @@ def forward( q = query[:num_tokens].view(-1, self.num_heads, hd) out = output[:num_tokens].view(-1, self.num_heads, hd) # Given query_fp8, it is the only q the fused insert wrote: CUTLASS - # decode reads it directly, and the bf16-query kernels read it + # decode and prefill read it directly, and Triton decode reads it # dequantized into ``query``. q_fp8 = ( None @@ -339,11 +339,11 @@ def forward( # [H, prefill, MK] transposed view; build_k2q_csr consumes the # strided view directly (topK stays innermost-contiguous). prefill_topk = topk[nd:num_tokens].transpose(0, 1) - qp = ( - q[nd:] - if q_fp8 is None - else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:]) - ) + # An E4M3 query attends as is, its scale folded into softmax_scale. + if q_fp8 is None: + qp, softmax_scale = q[nd:], self.scale + else: + qp, softmax_scale = q_fp8[nd:], self.scale * q_scale_float k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr( prefill_topk, p.cu_seqlens_q, @@ -375,7 +375,7 @@ def forward( max_seqlen_k=p.max_seq_len, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, page_table=p.block_table, seqused_k=p.seq_lens, schedule=schedule, @@ -393,7 +393,7 @@ def forward( topK=self.topk_blocks, blk_kv=SPARSE_BLOCK_SIZE, causal=True, - softmax_scale=self.scale, + softmax_scale=softmax_scale, cu_seqlens_q=p.cu_seqlens_q, cu_seqlens_k=p.cu_seqlens_k, max_seqlen_q=p.max_query_len, From 61f084b5f3e58aa3d3eb5c86aeaf4686c460c478 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Sun, 11 Oct 2026 02:21:55 -0400 Subject: [PATCH 362/952] [Bugfix][Multimodal] Fix per-video cache option reuse (#55203) Signed-off-by: bvolpato Co-authored-by: Codex --- .../multimodal/processing/test_qwen3_vl.py | 138 +++++++++++++----- vllm/model_executor/models/qwen3_vl.py | 69 ++++++++- vllm/multimodal/processing/processor.py | 35 +++-- 3 files changed, 197 insertions(+), 45 deletions(-) diff --git a/tests/models/multimodal/processing/test_qwen3_vl.py b/tests/models/multimodal/processing/test_qwen3_vl.py index 24281098154f..0c74a86ff000 100644 --- a/tests/models/multimodal/processing/test_qwen3_vl.py +++ b/tests/models/multimodal/processing/test_qwen3_vl.py @@ -15,6 +15,7 @@ from vllm.config.multimodal import MultiModalDummyOptions from vllm.model_executor.models.qwen3_vl import Qwen3VLProcessingInfo from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache from vllm.multimodal.processing.context import BaseProcessingInfo from ...registry import HF_EXAMPLE_MODELS @@ -336,65 +337,136 @@ def test_processor_kwargs_videos_kwargs_does_not_leak_into_image_budget( ("hf_mm_kwargs", "scalar_kwargs"), [ ({"num_frames": [8, 16]}, [{"num_frames": 8}, {"num_frames": 16}]), - ({"fps": [4.0, 12.0]}, [{"fps": 4.0}, {"fps": 12.0}]), + ({"fps": [2.0, 4.0]}, [{"fps": 2.0}, {"fps": 4.0}]), + ({"num_frames": [8, 16, 32]}, [{"num_frames": 8}, {"num_frames": 16}]), + ({"fps": [2.0, 4.0, 8.0]}, [{"fps": 2.0}, {"fps": 4.0}]), ], ) +@pytest.mark.parametrize("model_kwargs", [False, True]) +@pytest.mark.parametrize("use_uuids", [False, True]) +@pytest.mark.parametrize("scoped", [False, True]) +@pytest.mark.skip_global_cleanup def test_processor_multi_video_list_kwargs( model_id: str, hf_mm_kwargs: dict[str, Any], scalar_kwargs: list[dict[str, Any]], + model_kwargs: bool, + use_uuids: bool, + scoped: bool, ) -> None: - """List-valued video kwargs must be scalarized per item without mutation. - - Each video in the multi-video request must match processing the same input - independently with its corresponding scalar ``fps``/``num_frames`` value. - """ + """Per-video options must survive duplicate items and partial cache hits.""" + if scoped: + hf_mm_kwargs = {"videos_kwargs": hf_mm_kwargs} + hf_mm_kwargs_before = deepcopy(hf_mm_kwargs) ctx = build_model_context( model_id, limit_mm_per_prompt={"image": 0, "video": 2}, + mm_processor_kwargs=hf_mm_kwargs if model_kwargs else None, + mm_processor_cache_gb=1, ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - - prompt = ( - "<|vision_start|><|video_pad|><|vision_end|>" - "<|vision_start|><|video_pad|><|vision_end|>" - ) - mm_data = { - "video": [ - _build_video_mm_data(num_frames=32)["video"][0], - _build_video_mm_data(num_frames=32)["video"][0], - ] - } - - hf_mm_kwargs_before = deepcopy(hf_mm_kwargs) - processed = processor( - prompt, - mm_items=processor.info.parse_mm_data(mm_data), - hf_processor_mm_kwargs=hf_mm_kwargs, - ) + cache = MultiModalProcessorOnlyCache(ctx.model_config) + prompt = "<|vision_start|><|video_pad|><|vision_end|>" * 2 + videos = [ + _build_video_mm_data(num_frames=32, original_fps=8.0)["video"][0] + for _ in range(3) + ] + for i, (frames, _) in enumerate(videos): + frames.fill(i) def process_single(hf_kwargs: dict[str, Any]) -> torch.Tensor: single_prompt = "<|vision_start|><|video_pad|><|vision_end|>" - single_mm_data = _build_video_mm_data(num_frames=32) single = processor( single_prompt, - mm_items=processor.info.parse_mm_data(single_mm_data), + mm_items=processor.info.parse_mm_data({"video": [videos[0]]}), hf_processor_mm_kwargs=hf_kwargs, ) return single["mm_kwargs"].get_data()["video_grid_thw"][0] expected_grids = [process_single(kwargs) for kwargs in scalar_kwargs] - multi_grids = processed["mm_kwargs"].get_data()["video_grid_thw"] - assert not torch.equal(expected_grids[0], expected_grids[1]) - assert torch.equal(multi_grids[0], expected_grids[0]) - assert torch.equal(multi_grids[1], expected_grids[1]) + + def run(indices, *, cache=None, cached_only=False): + return processor( + prompt, + mm_items=processor.info.parse_mm_data( + {"video": [None if cached_only else videos[i] for i in indices]} + ), + mm_uuid_items=( + {"video": [f"video-{i}" for i in indices]} if use_uuids else None + ), + hf_processor_mm_kwargs={} if model_kwargs else hf_mm_kwargs, + cache=cache, + ) + + for indices in ([0, 0], [0, 1], [0, 2], [2, 0]): + expected = run(indices) + actual = run(indices, cache=cache) + assert actual["prompt_token_ids"] == expected["prompt_token_ids"] + assert actual["mm_placeholders"] == expected["mm_placeholders"] + assert actual["mm_kwargs"] == expected["mm_kwargs"] + assert actual["mm_hashes"] == expected["mm_hashes"] + assert len(set(actual["mm_hashes"]["video"])) == 2 + multi_grids = actual["mm_kwargs"].get_data()["video_grid_thw"] + assert torch.equal(multi_grids[0], expected_grids[0]) + assert torch.equal(multi_grids[1], expected_grids[1]) + + if use_uuids: + cached_only = run([2, 0], cache=cache, cached_only=True) + assert cached_only["prompt_token_ids"] == actual["prompt_token_ids"] + assert cached_only["mm_kwargs"] == actual["mm_kwargs"] assert hf_mm_kwargs == hf_mm_kwargs_before - video_phs = processed["mm_placeholders"].get("video", []) - assert len(video_phs) == 2, ( - f"Expected exactly 2 video placeholders, got {len(video_phs)}" + +@pytest.mark.parametrize(("key", "value"), [("fps", 2.0), ("num_frames", 8)]) +@pytest.mark.parametrize("length", [0, 1]) +@pytest.mark.parametrize("use_cache", [False, True]) +@pytest.mark.parametrize("scoped", [False, True]) +def test_processor_rejects_short_per_video_options( + key: str, value: float | int, length: int, use_cache: bool, scoped: bool +) -> None: + """Reject malformed option lists with and without a warm processor cache.""" + ctx = build_model_context( + MODEL_ID, + limit_mm_per_prompt={"image": 0, "video": 2}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + cache = MultiModalProcessorOnlyCache(ctx.model_config) if use_cache else None + prompt = "<|vision_start|><|video_pad|><|vision_end|>" * 2 + videos = _build_video_mm_data(num_frames=32, original_fps=8.0)["video"] * 2 + + def process(values: list[float | int]) -> None: + kwargs = {key: values} + processor( + prompt, + mm_items=processor.info.parse_mm_data({"video": videos}), + hf_processor_mm_kwargs={"videos_kwargs": kwargs} if scoped else kwargs, + cache=cache, + ) + + if use_cache: + process([value, value]) + + with pytest.raises(ValueError, match=f"{key} must have at least 2 entries"): + process([value] * length) + + +@pytest.mark.parametrize("kwargs", [{"fps": [2.0, 4.0]}, {"num_frames": [8, 16]}]) +def test_processor_profiles_fewer_videos_than_configured_options( + kwargs: dict[str, list[float] | list[int]], +) -> None: + """Profiling may use fewer videos than the configured per-video list.""" + ctx = build_model_context( + MODEL_ID, + limit_mm_per_prompt={"image": 0, "video": 2}, + mm_processor_kwargs=kwargs, ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + profiled = processor.get_dummy_mm_inputs({"video": 1}) + + assert len(profiled["mm_placeholders"]["video"]) == 1 + assert profiled["mm_kwargs"].get_data()["video_grid_thw"].shape[0] == 1 def _build_video_embeds_mm_data( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 77780b3d1d40..527e2b0f3e2e 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -60,7 +60,7 @@ VideoPruningMethod, ) from vllm.distributed import get_pp_group, parallel_state -from vllm.inputs import MultiModalDataDict +from vllm.inputs import MultiModalDataDict, MultiModalHashes from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY from vllm.model_executor.layers.attention.mm_encoder_attention import ( @@ -81,6 +81,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.hasher import MultiModalHasher from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -94,6 +95,7 @@ from vllm.multimodal.processing import ( BaseDummyInputsBuilder, BaseMultiModalProcessor, + ProcessorInputs, PromptReplacement, PromptUpdate, PromptUpdateDetails, @@ -1342,6 +1344,67 @@ def _replace_video_token_placeholders( class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]): + def _get_per_video_kwargs( + self, kwargs: Mapping[str, object], num_videos: int | None = None + ) -> dict[str, list[Any]]: + merged = self.info._merge_and_resolve_mm_processor_kwargs(kwargs) + per_video_kwargs = { + key: value + for key, value in merged.get("videos_kwargs", {}).items() + if key in ("fps", "num_frames") and isinstance(value, list) + } + if num_videos: + for key, values in per_video_kwargs.items(): + if len(values) < num_videos: + raise ValueError( + f"{key} must have at least {num_videos} entries " + "(one per video); " + f"got {len(values)}." + ) + return per_video_kwargs + + def _get_mm_hashes(self, inputs: ProcessorInputs) -> MultiModalHashes: + hashes = dict(super()._get_mm_hashes(inputs)) + per_video_kwargs = self._get_per_video_kwargs( + inputs.hf_processor_mm_kwargs, len(hashes.get("video", [])) + ) + if per_video_kwargs and "video" in hashes: + algorithm = self.info.ctx.get_mm_config().mm_hasher_algorithm + hashes["video"] = [ + MultiModalHasher.hash_kwargs( + algorithm, + mm_hash=mm_hash, + **{key: values[i] for key, values in per_video_kwargs.items()}, + ) + for i, mm_hash in enumerate(hashes["video"]) + ] + return hashes + + def _get_cache_missing_processor_kwargs( + self, + inputs: ProcessorInputs, + mm_is_cached: Mapping[str, list[bool]], + ) -> Mapping[str, object]: + per_video_kwargs = self._get_per_video_kwargs( + inputs.hf_processor_mm_kwargs, len(mm_is_cached.get("video", [])) + ) + if not per_video_kwargs or "video" not in mm_is_cached: + return inputs.hf_processor_mm_kwargs + kwargs = self.info._merge_and_resolve_mm_processor_kwargs( + inputs.hf_processor_mm_kwargs + ) + kwargs["videos_kwargs"].update( + { + key: [ + value + for value, hit in zip(values, mm_is_cached["video"]) + if not hit + ] + for key, values in per_video_kwargs.items() + } + ) + return kwargs + @staticmethod def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: """Transformers>=5.10 processors override `replace_video_token` @@ -1412,6 +1475,7 @@ def _apply_hf_processor_main( video_input_ids_lst: list[list[int]] = [] if videos := hf_data.pop("videos", []): assert isinstance(videos, list) + self._get_per_video_kwargs(hf_kwargs, len(videos)) video_grid_thw_lst = [] pixel_values_videos_lst = [] timestamps_per_video = [] @@ -1544,6 +1608,9 @@ def _apply_hf_processor_main( non_video_scoped_kwargs = dict(non_video_mm_kwargs["videos_kwargs"]) non_video_scoped_kwargs.pop("fps", None) non_video_scoped_kwargs.pop("num_frames", None) + # The context merges model defaults again when calling HF. + for key in self._get_per_video_kwargs({}): + non_video_scoped_kwargs[key] = None non_video_mm_kwargs["videos_kwargs"] = non_video_scoped_kwargs processed_data = self.info.ctx.call_hf_processor( self.info.get_hf_processor(**non_video_mm_kwargs), diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index abdb58540255..5b46267dff54 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1393,6 +1393,21 @@ def _recompute_cached_prompt_update( """ return replace(cached_update, item_idx=new_item_idx) + def _get_mm_hashes(self, inputs: ProcessorInputs) -> MultiModalHashes: + """Get each item's processor and encoder cache identity.""" + return inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) + + def _get_cache_missing_processor_kwargs( + self, + inputs: ProcessorInputs, + mm_is_cached: MultiModalIsCached, + ) -> Mapping[str, object]: + """Select processor options for the items missing from the cache.""" + return inputs.hf_processor_mm_kwargs + def _merge_mm_kwargs( self, cache: BaseMultiModalProcessorCache, @@ -1464,10 +1479,7 @@ def _apply_hf_processor( # Use overrides if provided; fallback to data-dependent hashing. with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) + mm_hashes = self._get_mm_hashes(inputs) mm_prompt_updates = self._get_mm_prompt_updates( inputs.mm_data_items, @@ -1500,10 +1512,7 @@ def _cached_apply_hf_processor( return self._apply_hf_processor(inputs, timing_ctx) with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) + mm_hashes = self._get_mm_hashes(inputs) with timing_ctx.record("get_cache_missing_items"): mm_is_cached, mm_missing_data_items = self._get_cache_missing_items( @@ -1512,25 +1521,29 @@ def _cached_apply_hf_processor( mm_hashes=mm_hashes, ) + missing_processor_kwargs = self._get_cache_missing_processor_kwargs( + inputs, mm_is_cached + ) + # NOTE: The prompt does not correspond to `mm_missing_data_items`, # so we can't apply prompt updates until the new multimodal # items are combined with the cached multimodal items with timing_ctx.record("apply_hf_processor"): mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, - hf_kwargs=inputs.hf_processor_mm_kwargs, + hf_kwargs=missing_processor_kwargs, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( mm_missing_processed_data, self._get_mm_fields_config( - mm_missing_processed_data, inputs.hf_processor_mm_kwargs + mm_missing_processed_data, missing_processor_kwargs ), ) mm_missing_prompt_updates = self._get_mm_prompt_updates( mm_missing_data_items, - inputs.hf_processor_mm_kwargs, + missing_processor_kwargs, mm_missing_kwargs, ) From ef42b4385f8512567703692a5bfbb9c0edc59c89 Mon Sep 17 00:00:00 2001 From: debanshd Date: Sun, 11 Oct 2026 00:03:06 -0700 Subject: [PATCH 363/952] [Bugfix] Reclaim cached GPU memory during weight loading (#57516) Signed-off-by: Debanshu Das Signed-off-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Kimi Code --- vllm/model_executor/model_loader/base_loader.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 4c774808bf79..65f64a33267b 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import gc import time from abc import ABC, abstractmethod @@ -86,6 +87,11 @@ def load_model( logger.debug("Loading weights on %s ...", load_device) self.load_weights(model, model_config) + # Release transient tensors held by PyTorch's caching allocator + # before weight post-processing and raw CUDA driver allocations. + gc.collect() + torch.accelerator.empty_cache() + # Log peak GPU memory after loading weights. This is needed # to have test coverage on peak memory for online quantization. if current_platform.is_cuda_alike() or current_platform.is_xpu(): @@ -103,6 +109,9 @@ def load_model( process_weights_after_loading(model, model_config, target_device) log_online_quantization_time(vllm_config) + gc.collect() + torch.accelerator.empty_cache() + return model.eval() From 26a488ad93d12c658449a19dea2eb78f1be38f83 Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Sun, 11 Oct 2026 15:10:04 +0800 Subject: [PATCH 364/952] [Fast Start][Bugfix] Snapshot GPU memory after external weights are ready (#60966) --- vllm/v1/worker/gpu_worker.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 37fa00bb7efa..152def62221c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -512,12 +512,13 @@ def init_device(self): gc.collect() torch.accelerator.empty_cache() - # take current memory snapshot - self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) # Weights from external model loader process external_weight_memory = get_model_loader( self.load_config ).get_external_weight_memory(self.vllm_config) + + # take current memory snapshot after external weights are ready + self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) self.requested_memory = request_memory( init_snapshot, self.cache_config, external_weight_memory ) From fcae5c98870ff2c20659731dce7ba3a3ff0d48aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:26:02 +0200 Subject: [PATCH 365/952] [Security] Restrict default CORS origins (#61009) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Cyrus Leung --- rust/src/cmd/src/cli.rs | 10 ++- rust/src/cmd/src/cli/tests.rs | 28 +++------ rust/src/server/src/config.rs | 4 +- rust/src/server/src/middleware/cors.rs | 2 +- rust/src/server/src/routes/tests.rs | 44 ++++++++++--- tests/entrypoints/launchers/test_cli_args.py | 7 +++ .../serve/middleware/test_cors_middleware.py | 62 +++++++++++++++++++ vllm/entrypoints/launchers/cli_args.py | 4 +- 8 files changed, 124 insertions(+), 37 deletions(-) create mode 100644 tests/entrypoints/serve/middleware/test_cors_middleware.py diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 14b949c87de4..eeb458c86e75 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -401,9 +401,9 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub served_model_name: Vec, - /// CORS allowed origins as a JSON list. `["*"]` allows any origin. - #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = r#"["*"]"#)] - #[serde(default = "default_cors_wildcard")] + /// CORS allowed origins as a JSON list. Defaults to `[]` (no cross-origin access). + #[arg(long, value_parser = parse_json::, value_name = "JSON", default_value = "[]")] + #[serde(default = "default_cors_origins")] pub allowed_origins: JsonStringList, /// CORS allowed methods as a JSON list. `["*"]` allows the standard set. @@ -664,6 +664,10 @@ fn default_cors_wildcard() -> JsonStringList { JsonStringList(vec!["*".to_string()]) } +fn default_cors_origins() -> JsonStringList { + JsonStringList(Vec::new()) +} + fn default_py_bootstrap_parser_selection() -> ParserSelection { ParserSelection::None } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 57780027deac..cc545ce33e80 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -205,9 +205,7 @@ fn serve_args_forward_python_flags_with_separator() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1075,9 +1073,7 @@ fn frontend_args_accept_json() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1382,13 +1378,13 @@ fn serve_args_parse_cors_flags() { } #[test] -fn serve_args_cors_defaults_are_permissive() { +fn serve_args_cors_defaults_restrict_origins() { let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B"]).unwrap(); let Command::Serve(serve) = cli.command else { panic!("expected serve args"); }; - assert_eq!(serve.runtime.allowed_origins.0, ["*"]); + assert!(serve.runtime.allowed_origins.0.is_empty()); assert_eq!(serve.runtime.allowed_methods.0, ["*"]); assert_eq!(serve.runtime.allowed_headers.0, ["*"]); assert!(!serve.runtime.allow_credentials); @@ -1748,9 +1744,7 @@ fn serve_args_accept_handshake_aliases() { disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( - [ - "*", - ], + [], ), allowed_methods: JsonStringList( [ @@ -1906,9 +1900,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], @@ -2002,9 +1994,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], @@ -2142,9 +2132,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present sse_keep_alive_interval: None, }, cors: CorsConfig { - allow_origins: [ - "*", - ], + allow_origins: [], allow_methods: [ "*", ], diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 3f95b5891cb2..ca91ba40cd50 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -65,7 +65,7 @@ pub struct ApiServerOptions { pub sse_keep_alive_interval: Option, } -/// CORS settings mirroring Python's `CORSMiddleware`; the default is permissive. +/// CORS settings mirroring Python's `CORSMiddleware`; origins are denied by default. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct CorsConfig { /// Allowed origins. `["*"]` allows any origin. @@ -81,7 +81,7 @@ pub struct CorsConfig { impl Default for CorsConfig { fn default() -> Self { Self { - allow_origins: vec!["*".to_string()], + allow_origins: Vec::new(), allow_methods: vec!["*".to_string()], allow_headers: vec!["*".to_string()], allow_credentials: false, diff --git a/rust/src/server/src/middleware/cors.rs b/rust/src/server/src/middleware/cors.rs index 89d1b6f3b7bd..bfc030a3a9d9 100644 --- a/rust/src/server/src/middleware/cors.rs +++ b/rust/src/server/src/middleware/cors.rs @@ -105,7 +105,7 @@ pub fn cors_layer(cfg: &CorsConfig) -> CorsLayer { // Emit `Vary: Origin` only when the allow-origin is dynamic (explicit // origins, or credentials); the wildcard + no-credentials case emits no - // `Vary` at all, and an empty list disables the header here. + // `Vary` at all. An empty origin list never permits an origin. let vary: Vec = if !wildcard_origins || cfg.allow_credentials { vec![header::ORIGIN] } else { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 118e412b45b6..7ba4bbf0f416 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1784,7 +1784,7 @@ async fn api_key_auth_allows_unguarded_route_without_token() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn cors_default_simple_request_allows_any_origin() { +async fn cors_default_simple_request_omits_allow_origin() { let (mut app, _engine_task) = test_app_with_cors(CorsConfig::default()).await; let response = app .call( @@ -1799,17 +1799,12 @@ async fn cors_default_simple_request_allows_any_origin() { .expect("call app"); assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - header_value(&response, "access-control-allow-origin"), - Some("*") - ); - // Wildcard origins without credentials emit no `Vary` (Starlette parity). - assert_eq!(header_value(&response, "vary"), None); + assert_eq!(header_value(&response, "access-control-allow-origin"), None); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn cors_default_preflight_returns_explicit_methods_and_max_age() { +async fn cors_default_preflight_omits_allow_origin() { let (mut app, _engine_task) = test_app_with_cors(CorsConfig::default()).await; let response = app .call( @@ -1825,6 +1820,32 @@ async fn cors_default_preflight_returns_explicit_methods_and_max_age() { .await .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(header_value(&response, "access-control-allow-origin"), None); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn cors_wildcard_preflight_returns_explicit_methods_and_max_age() { + let cors = CorsConfig { + allow_origins: vec!["*".to_string()], + ..CorsConfig::default() + }; + let (mut app, _engine_task) = test_app_with_cors(cors).await; + let response = app + .call( + Request::builder() + .method("OPTIONS") + .uri("/v1/chat/completions") + .header("origin", "http://example.com") + .header("access-control-request-method", "POST") + .header("access-control-request-headers", "content-type") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); // `*` methods expand to the explicit method list, matching Starlette // (never the literal `*`). @@ -1917,6 +1938,7 @@ async fn cors_explicit_origin_disallowed_omits_allow_origin() { #[serial] async fn cors_wildcard_with_credentials_reflects_origin_without_panic() { let cors = CorsConfig { + allow_origins: vec!["*".to_string()], allow_credentials: true, ..CorsConfig::default() }; @@ -1971,8 +1993,12 @@ async fn cors_unauthorized_response_has_no_cors_headers() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn cors_preflight_bypasses_auth_and_returns_cors_headers() { + let cors = CorsConfig { + allow_origins: vec!["*".to_string()], + ..CorsConfig::default() + }; let (mut app, _engine_task) = - test_app_with_cors_and_keys(CorsConfig::default(), vec!["secret".to_string()]).await; + test_app_with_cors_and_keys(cors, vec!["secret".to_string()]).await; let response = app .call( Request::builder() diff --git a/tests/entrypoints/launchers/test_cli_args.py b/tests/entrypoints/launchers/test_cli_args.py index 84279e155582..0a4c97141d07 100644 --- a/tests/entrypoints/launchers/test_cli_args.py +++ b/tests/entrypoints/launchers/test_cli_args.py @@ -44,6 +44,13 @@ def serve_parser(): return _build_vllm_parsers()["vllm serve"] +def test_cors_origins_require_explicit_opt_in(serve_parser): + assert serve_parser.parse_args([]).allowed_origins == [] + assert serve_parser.parse_args( + ["--allowed-origins", '["http://trusted.example"]'] + ).allowed_origins == ["http://trusted.example"] + + ### Test config parsing def test_config_arg_parsing(serve_parser, cli_config_file): args = serve_parser.parse_args([]) diff --git a/tests/entrypoints/serve/middleware/test_cors_middleware.py b/tests/entrypoints/serve/middleware/test_cors_middleware.py new file mode 100644 index 000000000000..543c101a9307 --- /dev/null +++ b/tests/entrypoints/serve/middleware/test_cors_middleware.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from fastapi import FastAPI +from starlette.testclient import TestClient + +from vllm.entrypoints.launchers.cli_args import FrontendArgs +from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware + + +def _make_client(allowed_origins: list[str] | None = None) -> TestClient: + args = FrontendArgs() + if allowed_origins is not None: + args.allowed_origins = allowed_origins + + app = FastAPI() + + @app.get("/v1/models") + async def models(): + return {"data": [{"id": "private-model"}]} + + @app.post("/v1/chat/completions") + async def chat(): + return {"message": {"content": "ok"}} + + init_entrypoints_middleware(args, app, ()) + return TestClient(app) + + +def test_default_origins_prevent_cross_origin_reads_and_json_requests(): + client = _make_client() + response = client.get("/v1/models", headers={"Origin": "http://attacker.example"}) + assert response.status_code == 200 + assert "access-control-allow-origin" not in response.headers + + preflight = client.options( + "/v1/chat/completions", + headers={ + "Origin": "http://attacker.example", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "content-type", + }, + ) + assert preflight.status_code == 400 + assert "access-control-allow-origin" not in preflight.headers + + +def test_explicit_origin_allows_browser_requests(): + client = _make_client(["http://trusted.example"]) + response = client.get("/v1/models", headers={"Origin": "http://trusted.example"}) + assert response.headers["access-control-allow-origin"] == "http://trusted.example" + + preflight = client.options( + "/v1/chat/completions", + headers={ + "Origin": "http://trusted.example", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "content-type", + }, + ) + assert preflight.status_code == 200 + assert preflight.headers["access-control-allow-origin"] == "http://trusted.example" diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 5cd58da87eac..e2c2625e9462 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -316,8 +316,8 @@ class FrontendArgs(BaseFrontendArgs): while keeping logs for other endpoints.""" allow_credentials: bool = False """Allow credentials.""" - allowed_origins: list[str] = field(default_factory=lambda: ["*"]) - """Allowed origins.""" + allowed_origins: list[str] = field(default_factory=list) + """Allowed origins. Cross-origin browser access is disabled by default.""" allowed_methods: list[str] = field(default_factory=lambda: ["*"]) """Allowed methods.""" allowed_headers: list[str] = field(default_factory=lambda: ["*"]) From 1789f095881d14966d7d4d9a2647748bf32dd834 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 16:58:07 +0800 Subject: [PATCH 366/952] [Frontend] Explain disabled scale-out endpoints with opt-in hints (#61042) Signed-off-by: aoshen02 Co-authored-by: OpenAI Codex --- docs/usage/security.md | 2 +- rust/src/server/src/routes.rs | 24 ++++++++- rust/src/server/src/routes/tests.rs | 22 +++++--- tests/entrypoints/scale_out/test_factories.py | 52 ++++++++++++++++--- vllm/entrypoints/scale_out/factories.py | 41 ++++++++++++--- 5 files changed, 117 insertions(+), 24 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index ab5aa96becc5..50c4149be81b 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -178,7 +178,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/systemone` - Structured decisions API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) -- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) +- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`; otherwise returns HTTP 403 with instructions to enable it) - `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`) - `/v2/embed` - Cohere Embed API - `/v2/rerank` - Cohere Rerank API diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 28f35573249e..251398611ba1 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -24,17 +24,36 @@ mod world_size; use std::sync::Arc; -use axum::Router; use axum::extract::DefaultBodyLimit; +use axum::http::StatusCode; use axum::middleware::{from_fn, from_fn_with_state}; use axum::routing::{get, post}; +use axum::{Json, Router}; use tower_http::trace::TraceLayer; use tracing::{info, warn}; use crate::DEFAULT_REQUEST_BODY_LIMIT_BYTES; use crate::middleware; +use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse}; use crate::state::AppState; +const SCALE_OUT_DISABLED_MESSAGE: &str = + "Scale-out endpoints are disabled. Restart vllm serve with --enable-scale-out to enable them."; + +async fn disabled_scale_out() -> (StatusCode, Json) { + ( + StatusCode::FORBIDDEN, + Json(ErrorResponse { + error: ErrorDetail { + message: SCALE_OUT_DISABLED_MESSAGE.to_string(), + error_type: "PermissionDeniedError".to_string(), + param: None, + code: Some("403".to_string()), + }, + }), + ) +} + fn server_dev_mode_enabled() -> bool { std::env::var("VLLM_SERVER_DEV_MODE") .ok() @@ -105,7 +124,8 @@ fn build_router_with_options( if scale_out_endpoints_enabled { router = router.route("/inference/v1/generate", post(inference::generate)); } else { - info!("scale-out endpoints are disabled; pass --enable-scale-out to enable them"); + info!("{SCALE_OUT_DISABLED_MESSAGE}"); + router = router.route("/inference/v1/generate", post(disabled_scale_out)); } if runtime_lora_updating_enabled { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 7ba4bbf0f416..d574e02a9215 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -758,7 +758,7 @@ async fn scale_out_generate_route_follows_api_server_options() { .await .expect("call app"); - assert_ne!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -770,10 +770,7 @@ async fn scale_out_generate_route_is_disabled_by_default() { Arc::new(FakeChatBackend::new()), ) .await; - let mut app = build_router_with_scale_out_endpoints( - Arc::new(AppState::new(vec!["model".to_string()], chat)), - false, - ); + let mut app = build_router(Arc::new(AppState::new(vec!["model".to_string()], chat))); let response = app .call( @@ -781,13 +778,22 @@ async fn scale_out_generate_route_is_disabled_by_default() { .method("POST") .uri("/inference/v1/generate") .header("content-type", "application/json") - .body(Body::from("{}")) + .body(Body::from("not JSON")) .expect("build request"), ) .await .expect("call app"); - assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read response"); + let error: serde_json::Value = serde_json::from_slice(&body).expect("decode error"); + assert_eq!(error["error"]["type"], "PermissionDeniedError"); + assert!( + error["error"]["message"] + .as_str() + .expect("error message") + .contains("--enable-scale-out") + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -816,7 +822,7 @@ async fn scale_out_generate_route_can_be_enabled() { .await .expect("call app"); - assert_ne!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); } /// Build a dev-mode router backed by a mock engine using a custom ready diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py index ea96fff876d1..7937c0d45bc0 100644 --- a/tests/entrypoints/scale_out/test_factories.py +++ b/tests/entrypoints/scale_out/test_factories.py @@ -1,13 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging from argparse import Namespace +import pytest from fastapi import FastAPI +from fastapi.testclient import TestClient from vllm.entrypoints.launchers.api_server.routers import register_api_routers +from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers from vllm.entrypoints.scale_out.token_in_token_out.api_router import attach_router -from vllm.entrypoints.serve.middleware.authenticate import GUARDED_PREFIX +from vllm.entrypoints.serve.middleware.authenticate import ( + GUARDED_PREFIX, + AuthenticationMiddleware, +) from vllm.tasks import SupportedTask RENDER_PATHS = { @@ -41,17 +48,50 @@ def registered_paths( return {route.path for route in app.routes} -def test_scale_out_routes_are_disabled_by_default(): - paths = registered_paths(("generate",)) +@pytest.mark.parametrize("path", sorted(SCALE_OUT_PATHS)) +@pytest.mark.parametrize("authorized", [False, True]) +def test_disabled_scale_out_routes_explain_opt_in_without_parsing(path, authorized): + app = FastAPI() + register_scale_out_api_routers(app, ("generate",)) + app.add_middleware(AuthenticationMiddleware, tokens=["secret"]) + headers = {"Authorization": "Bearer secret"} if authorized else {} + client = TestClient(app) + response = client.post(path, content="not JSON", headers=headers) - assert SCALE_OUT_PATHS.isdisjoint(paths) + assert response.status_code == (403 if authorized else 401) + if authorized: + assert "--enable-scale-out" in response.json()["error"]["message"] + assert response.json()["error"]["type"] == "PermissionDeniedError" + assert path not in app.openapi()["paths"] + assert client.post("/unknown", headers=headers).status_code == 404 -def test_disabled_scale_out_routes_are_logged(caplog): +def test_disabled_scale_out_routes_are_logged_at_info(caplog): with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"): registered_paths(("generate",)) - assert any("--enable-scale-out" in record.message for record in caplog.records) + assert [ + record.levelno + for record in caplog.records + if "--enable-scale-out" in record.message + ] == [logging.INFO] + + +@pytest.mark.parametrize("tokens_only", [False, True]) +def test_enabled_generate_routes_use_native_request_validation(tokens_only): + app = FastAPI() + app.state.args = Namespace( + tokens_only=tokens_only, enable_scale_out=not tokens_only + ) + register_scale_out_api_routers(app, ("generate",)) + response = TestClient(app).post( + GENERATE_PATH, + content="not JSON", + headers={"Content-Type": "application/json"}, + ) + + assert response.status_code == 422 + assert response.json()["detail"][0]["type"] == "json_invalid" def test_scale_out_routes_can_be_enabled_for_generate_server(): diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py index bd48c709ddc0..45bb9dbcde6b 100644 --- a/vllm/entrypoints/scale_out/factories.py +++ b/vllm/entrypoints/scale_out/factories.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from argparse import Namespace +from http import HTTPStatus from typing import TYPE_CHECKING from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.routing import APIRoute from vllm.engine.protocol import EngineClient +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.logger import init_logger from vllm.tasks import SupportedTask @@ -66,6 +72,9 @@ def register_scale_out_api_routers( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], ): + from .derender.api_router import router as derender_render + from .render.api_router import router as render_render + args = getattr(app.state, "args", None) # `vllm launch render` and `vllm serve --tokens-only` are dedicated to # serving these endpoints, so they always register them regardless of @@ -73,17 +82,35 @@ def register_scale_out_api_routers( dedicated_mode = "render" in supported_tasks or getattr(args, "tokens_only", False) enabled = dedicated_mode or getattr(args, "enable_scale_out", False) if not enabled: - logger.info( - "Scale-out endpoints are disabled. Set --enable-scale-out to enable them." + message = ( + "Scale-out endpoints are disabled. Restart vllm serve with " + "--enable-scale-out to enable them." ) + logger.info(message) + + async def disabled_scale_out() -> JSONResponse: + error = create_error_response( + message, "PermissionDeniedError", HTTPStatus.FORBIDDEN + ) + return JSONResponse(error.model_dump(), status_code=HTTPStatus.FORBIDDEN) + + from .token_in_token_out.api_router import router as token_router + + for route in ( + *render_render.routes, + *derender_render.routes, + *token_router.routes, + ): + if isinstance(route, APIRoute): + app.add_api_route( + route.path, + disabled_scale_out, + methods=sorted(route.methods), + include_in_schema=False, + ) return - from .render.api_router import router as render_render - app.include_router(render_render) - - from .derender.api_router import router as derender_render - app.include_router(derender_render) if "generate" in supported_tasks: From 2339bd39d661965b821bfbe652db8f536e0b3dcd Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Sun, 11 Oct 2026 02:15:18 -0700 Subject: [PATCH 367/952] [Bugfix][Mooncake] Use the EAGLE replay boundary when retaining sliding-window tails (#61012) Signed-off-by: Stu Cao Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 Co-authored-by: Yifan Qiao Co-authored-by: Codex --- .../unit/test_mooncake_store_coordinator.py | 44 ++++++++++++++++ .../v1/mooncake/store/coordinator.py | 13 +++-- vllm/v1/core/kv_cache_coordinator.py | 51 +++++++++++-------- vllm/v1/core/single_type_kv_cache_manager.py | 4 +- 4 files changed, 85 insertions(+), 27 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index ac52ca07ab33..6896777226cf 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -3,6 +3,7 @@ from math import lcm +import pytest import torch from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 @@ -422,6 +423,49 @@ def test_store_mask_retention_interval_zero_keeps_only_replay_boundary(): assert masks[1] == [i == 11 for i in range(16)] +def test_store_mask_retention_interval_zero_eagle_keeps_lookup_tail(): + """With EAGLE the replay boundary sits one scheduler block lower, as in + core's ``replay_boundaries``: the lookup matches that block and drops + back from it. The retained SWA tail must span the ``need`` contiguous + blocks the lookup requires (cdiv(8 - 1, 8) + 1 = 2 here); the non-EAGLE + boundary leaves only one of them inside the saved prefix.""" + coord = _make_coord( + _retention_groups(), hash_block_size=8, use_eagle=True, retention_interval=0 + ) + # num_prompt=100 -> boundary 64 -> chunks 7, 8 (8 is the EAGLE peek). + masks = coord.store_mask(96, num_prompt_tokens=100) + assert masks[0] is None + assert masks[1] == [i in (7, 8) for i in range(12)] + # Aligned num_prompt=96 -> boundaries 32 (resend) and 64 (longer sibling). + masks = coord.store_mask(96, num_prompt_tokens=96) + assert masks[1] == [i in (3, 4, 7, 8) for i in range(12)] + + +@pytest.mark.parametrize("num_prompt_tokens", [100, 96]) +def test_store_mask_retention_interval_zero_eagle_round_trip(num_prompt_tokens): + """Everything store_mask keeps is in the store; the EAGLE lookup of the + same prompt must then reach the full block below the one it drops (64 of + 96). The unaligned prompt hit 0 before the fix; the aligned one already + hit.""" + groups = _retention_groups() + coord = _make_coord(groups, hash_block_size=8, use_eagle=True, retention_interval=0) + hashes = _hashes(96 // coord.hash_block_size) + exists: set[tuple[int, bytes]] = set() + store_masks = coord.store_mask(96, num_prompt_tokens=num_prompt_tokens) + for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)): + group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec) + for chunk_id, block_hash in enumerate(group_hashes): + if mask is None or mask[chunk_id]: + exists.add((gid, bytes(block_hash))) + + _masks, hit = coord.find_longest_cache_hit( + hashes, + max_length=96, + cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists), + ) + assert hit == 64 + + def test_store_mask_retention_interval_keeps_segment_and_replay_tails(): """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay boundary tail (num_prompt=100 -> chunk 11) coexist.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 48a36bc82f8b..2382948268bb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -11,6 +11,7 @@ ) from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import replay_boundaries from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, @@ -341,6 +342,15 @@ def _reachable_masks( f"aligned_token_len ({aligned_token_len}) must be a multiple of " f"{mask_alignment}" ) + reachable_boundaries = ( + () + if num_prompt_tokens is None + else replay_boundaries( + num_prompt_tokens, + self.lcm_block_size, + use_eagle=bool(self.eagle_group_ids), + ) + ) masks: list[list[bool] | None] = [] for g_idx, g in enumerate(self.kv_cache_groups): spec = _unwrap_spec(g.kv_cache_spec) @@ -352,9 +362,6 @@ def _reachable_masks( manager_cls = KVCacheSpecRegistry.get_manager_class(spec) assert manager_cls is not None use_eagle = g_idx in self.eagle_group_ids - reachable_boundaries = ( - () if num_prompt_tokens is None else (num_prompt_tokens - 1,) - ) mask = manager_cls.reachable_block_mask( start_block=start_chunk, end_block=end_chunk, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 019941a1bf03..b6defc52e847 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -62,6 +62,30 @@ def _validate_prefix_cache_retention_interval( ) +def replay_boundaries( + num_prompt_tokens: int, scheduler_block_size: int, use_eagle: bool +) -> tuple[int, ...]: + """Positions a later request replaying a prompt can resume at. + + A hit is the shortest across all groups, so every group retains state + at each position; EAGLE groups also keep the block above, which they + match and drop back from (see ``reachable_block_mask``). + + Two positions are reachable: a resend of the identical prompt is capped + at ``num_tokens - 1`` (its last token is recomputed for logits), a + longer sibling matches the final aligned block. They differ only on a + block-aligned prompt, where retaining just the higher one collapses the + resend's hit to 0. The alignment is the scheduler block size, not the + finer hash granularity, which would over-estimate the reach. + """ + if not use_eagle: + return (num_prompt_tokens - 1,) + block = scheduler_block_size + resend = (num_prompt_tokens - 1) // block * block + extension = num_prompt_tokens // block * block + return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) + + class KVCacheCoordinator(ABC): """Coordinate the KV cache of different KV cache groups.""" @@ -319,27 +343,6 @@ def allocate_new_blocks( for manager in self.single_type_managers ) - def get_replay_boundaries(self, request: Request) -> tuple[int, ...]: - """Positions a later request replaying this prompt can resume at. - - A hit is the shortest across all groups, so every group retains state - at each position; EAGLE groups also keep the block above, which they - match and drop back from (see ``reachable_block_mask``). - - Two positions are reachable: a resend of the identical prompt is capped - at ``num_tokens - 1`` (its last token is recomputed for logits), a - longer sibling matches the final aligned block. They differ only on a - block-aligned prompt, where retaining just the higher one collapses the - resend's hit to 0. The alignment is the scheduler block size, not the - finer hash granularity, which would over-estimate the reach. - """ - if not self.eagle_group_ids: - return (request.num_prompt_tokens - 1,) - block = self.scheduler_block_size - resend = (request.num_prompt_tokens - 1) // block * block - extension = request.num_prompt_tokens // block * block - return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) - def get_num_cacheable_tokens( self, num_computed_tokens: int, kv_cache_group_id: int ) -> int: @@ -358,7 +361,11 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: (including tokens that are already cached). """ - boundaries = self.get_replay_boundaries(request) + boundaries = replay_boundaries( + request.num_prompt_tokens, + self.scheduler_block_size, + use_eagle=bool(self.eagle_group_ids), + ) for group_id, manager in enumerate(self.single_type_managers): if not manager.enable_caching: continue diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 77c6a4a5f750..218d59f430c9 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -492,7 +492,7 @@ def cache_blocks( boundary; a positive multiple of ``scheduler_block_size`` keeps a tail once per that-sized segment. Only SWA acts on it. replay_boundaries: Positions a later request replaying this prompt - can resume at, from ``get_replay_boundaries``. + can resume at, from ``replay_boundaries``. """ if not self.kv_cache_spec.prefix_cacheable: @@ -505,7 +505,7 @@ def cache_blocks( # Token boundaries whose reachable tail must be retained under sparse # retention: every position a replaying sibling can resume at (see - # ``get_replay_boundaries``) and any detected shared-prefix junction. + # ``replay_boundaries``) and any detected shared-prefix junction. reachable_boundaries = [*replay_boundaries] if request.shared_prefix_boundary: reachable_boundaries.append(request.shared_prefix_boundary) From f425a212081181e3d6ad2ca92daa74ba0f9e83a4 Mon Sep 17 00:00:00 2001 From: Subhash Polisetti <49185065+subhashpolisetti@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:03:33 -0700 Subject: [PATCH 368/952] [Bugfix][Frontend] Support scalar priority and wire priority into LLM.chat() (#57243) Signed-off-by: subhash polisetti --- tests/entrypoints/llm/test_chat.py | 19 +++++++++++++ tests/entrypoints/llm/test_generate.py | 8 ++++++ .../unit_tests/test_offline_utils.py | 6 +++++ vllm/entrypoints/common/offline.py | 17 +++++++----- vllm/entrypoints/llm.py | 27 ++++++++++++++----- 5 files changed, 65 insertions(+), 12 deletions(-) diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 276426e692c7..8fdd41c51119 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -65,6 +65,25 @@ def test_multi_chat(text_llm): assert len(outputs) == 2 +def test_chat_priority(text_llm): + messages = [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "Explain the concept of entropy."}, + ] + + # A single value is applied to every conversation. + outputs = text_llm.chat(messages, priority=1) + assert len(outputs) == 1 + + outputs = text_llm.chat([messages, messages], priority=[0, 1]) + assert len(outputs) == 2 + + # Exception raised, if the length of priority does not match the length of + # conversations + with pytest.raises(VLLMValidationError): + text_llm.chat([messages, messages], priority=[0]) + + def test_llm_chat_tokenization_no_double_bos(text_llm): """LLM.chat() should not add special tokens when using chat templates. Check we get a single BOS token for llama chat. diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py index 4a2c59dc5842..ccff4eb7a0f9 100644 --- a/tests/entrypoints/llm/test_generate.py +++ b/tests/entrypoints/llm/test_generate.py @@ -75,6 +75,10 @@ def test_multiple_priority(llm: LLM): outputs = llm.generate(PROMPTS, sampling_params=None, priority=[0] * len(PROMPTS)) assert len(PROMPTS) == len(outputs) + # Generate works when priority is a single value applied to every prompt + outputs = llm.generate(PROMPTS, sampling_params=None, priority=1) + assert len(PROMPTS) == len(outputs) + # Exception raised, if the length of priority does not match the length of prompts with pytest.raises(VLLMValidationError): outputs = llm.generate( @@ -91,6 +95,10 @@ def test_single_prompt_priority(llm: LLM): outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=[0]) assert len(outputs) == 1 + # Single string prompts should also accept a single priority value. + outputs = llm.generate(PROMPTS[0], sampling_params=None, priority=1) + assert len(outputs) == 1 + def test_max_model_len(vllm_runner): max_model_len = 20 diff --git a/tests/entrypoints/unit_tests/test_offline_utils.py b/tests/entrypoints/unit_tests/test_offline_utils.py index a37969c9b366..8a0308ce3a46 100644 --- a/tests/entrypoints/unit_tests/test_offline_utils.py +++ b/tests/entrypoints/unit_tests/test_offline_utils.py @@ -55,3 +55,9 @@ def test_matching_lengths_pass_through(mixin: OfflineInferenceMixin): ] assert mixin._lora_request_to_seq([lora], num_requests=1) == [lora] assert mixin._priority_to_seq([3], num_requests=1) == [3] + + +def test_scalar_priority_broadcast(mixin: OfflineInferenceMixin): + assert mixin._priority_to_seq(5, num_requests=1) == [5] + assert mixin._priority_to_seq(5, num_requests=3) == [5, 5, 5] + assert mixin._priority_to_seq(None, num_requests=3) == [0, 0, 0] diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index b33b8c34c5ad..1bb0bedccbd0 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -281,10 +281,13 @@ def _lora_request_to_seq( def _priority_to_seq( self, - priority: list[int] | None, + priority: int | Sequence[int] | None, num_requests: int, ) -> Sequence[int]: - if priority is not None: + if priority is None: + priority = 0 + + if isinstance(priority, Sequence): if len(priority) != num_requests: raise VLLMValidationError( f"The lengths of prompts ({num_requests}) " @@ -293,7 +296,7 @@ def _priority_to_seq( return priority - return [0] * num_requests + return [priority] * num_requests def _add_completion_requests( self, @@ -304,7 +307,7 @@ def _add_completion_requests( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[str]: @@ -341,7 +344,7 @@ def _run_completion( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ): @@ -367,6 +370,7 @@ def _run_chat( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -381,6 +385,7 @@ def _run_chat( params=params, use_tqdm=use_tqdm, lora_request=lora_request, + priority=priority, chat_template=chat_template, chat_template_content_format=chat_template_content_format, chat_template_kwargs=chat_template_kwargs, @@ -402,7 +407,7 @@ def _add_chat_requests( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index ff7bd9c57402..6b7f91d08e4c 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -422,7 +422,7 @@ def generate( *, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ) -> list[RequestOutput]: @@ -448,9 +448,9 @@ def generate( lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. Only applicable when priority scheduling policy is enabled. - If provided, must be a list of integers matching the length - of `prompts`, where each priority value corresponds to the prompt - at the same index. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. tokenization_kwargs: Overrides for `tokenizer.encode`. mm_processor_kwargs: Overrides for `processor.__call__`. @@ -486,7 +486,7 @@ def enqueue( prompts: PromptType | Sequence[PromptType], sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, @@ -502,6 +502,10 @@ def enqueue( sampling_params: The sampling parameters for text generation. lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. use_tqdm: If True, shows a tqdm progress bar while adding requests. tokenization_kwargs: Overrides for `tokenizer.encode`. mm_processor_kwargs: Overrides for `processor.__call__`. @@ -617,6 +621,7 @@ def chat( sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -651,6 +656,11 @@ def chat( it is used to create the progress bar. If `False`, no progress bar is created. lora_request: LoRA request to use for generation, if any. + priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. chat_template: The template to use for structuring the chat. If not provided, the model's default chat template will be used. chat_template_content_format: The format to render message content. @@ -695,6 +705,7 @@ def chat( output_type=RequestOutput, use_tqdm=use_tqdm, lora_request=lora_request, + priority=priority, chat_template=chat_template, chat_template_content_format=chat_template_content_format, chat_template_kwargs=chat_template_kwargs, @@ -712,7 +723,7 @@ def enqueue_chat( sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, use_tqdm: bool | Callable[..., tqdm] = True, lora_request: Sequence[LoRARequest] | LoRARequest | None = None, - priority: list[int] | None = None, + priority: int | Sequence[int] | None = None, chat_template: str | None = None, chat_template_content_format: ChatTemplateContentFormatOption = "auto", add_generation_prompt: bool = True, @@ -739,6 +750,10 @@ def enqueue_chat( conversations. lora_request: LoRA request to use for generation, if any. priority: The priority of the requests, if any. + Only applicable when priority scheduling policy is enabled. + When it is a single value, it is applied to every prompt. + When it is a list, the list must have the same length as the + prompts and it is paired one by one with the prompt. chat_template: The template to use for structuring the chat. chat_template_content_format: The format to render message content. add_generation_prompt: If True, adds a generation template From c85616efc6f0eaa50ca1f1a09c32368aedffd3d8 Mon Sep 17 00:00:00 2001 From: tc-mb <157115220+tc-mb@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:21:10 +0800 Subject: [PATCH 369/952] [Model] Support MiniCPM-V 4.7 (#58674) Signed-off-by: tc-mb Co-authored-by: Cyrus Leung --- docs/models/supported_models.md | 2 +- tests/model_executor/test_minicpmv.py | 114 +++ tests/model_executor/test_minicpmv4_7.py | 310 +++++++ tests/models/registry.py | 5 + .../parser/engine/test_minicpmv_reasoning.py | 459 +++++++++++ vllm/model_executor/models/minicpmv4_6.py | 335 ++++++-- vllm/model_executor/models/minicpmv4_7.py | 764 ++++++++++++++++++ vllm/model_executor/models/registry.py | 4 + vllm/parser/minicpmv.py | 405 ++++++++++ vllm/parser/parser_manager.py | 13 + vllm/reasoning/__init__.py | 4 + vllm/reasoning/minicpmv_reasoning_parser.py | 16 + .../chat_templates/registry.py | 1 + 13 files changed, 2368 insertions(+), 64 deletions(-) create mode 100644 tests/model_executor/test_minicpmv4_7.py create mode 100644 tests/parser/engine/test_minicpmv_reasoning.py create mode 100644 vllm/model_executor/models/minicpmv4_7.py create mode 100644 vllm/parser/minicpmv.py create mode 100644 vllm/reasoning/minicpmv_reasoning_parser.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 221ab877eaa8..bbcb94f32cf5 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -592,7 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MiDashengLMModel` | MiDashengLM | T + A+ | `mispeech/midashenglm-7b` | | ✅︎ | | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ | | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ | -| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, `openbmb/MiniCPM-V-4_6`, etc. | ✅︎ | | +| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, `openbmb/MiniCPM-V-4_6`, `openbmb/MiniCPM-V-4_7`, etc. | ✅︎ | | | `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ | | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ | | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ | diff --git a/tests/model_executor/test_minicpmv.py b/tests/model_executor/test_minicpmv.py index 96b80101b414..cffdea260097 100644 --- a/tests/model_executor/test_minicpmv.py +++ b/tests/model_executor/test_minicpmv.py @@ -6,6 +6,9 @@ from vllm.model_executor.models.minicpmv import MiniCPMV2_6 from vllm.model_executor.models.minicpmv4_6 import ( MiniCPMV4_6ForConditionalGeneration, + _expand_flags, + _select_vision_items, + _vision_downsample_mode, ) pytestmark = pytest.mark.skip_global_cleanup @@ -54,3 +57,114 @@ def test_image_and_video_embeds_stay_in_their_own_modality(): assert torch.equal(modalities["images"]["image_embeds"], image_embeds) assert torch.equal(modalities["videos"]["image_embeds"], video_embeds) + + +class TestPerItemMergerFlags: + """The merger flag is per request but items from several requests can be + batched together, so it has to be applied per item.""" + + def test_expand_flags_broadcasts_a_single_value(self): + assert _expand_flags(None, 3) == [None, None, None] + assert _expand_flags(torch.tensor([True]), 3) == [True, True, True] + assert _expand_flags(True, 2) == [True, True] + + def test_expand_flags_keeps_per_item_values(self): + flags = [torch.tensor([True]), torch.tensor([False])] + + assert _expand_flags(flags, 2) == [True, False] + + def test_expand_flags_falls_back_on_a_shape_mismatch(self): + # Too many values to map onto the items: keep one decision for the + # batch rather than failing the request. + assert _expand_flags([True, False, True], 2) == [True, True] + + def test_vision_downsample_mode(self): + assert _vision_downsample_mode(None) is None + assert _vision_downsample_mode(True) == "16x" + assert _vision_downsample_mode(False) == "4x" + + def test_select_vision_items_slices_by_item(self): + # 3 items holding 2, 1 and 2 slices; one channel, one patch, width 4. + pixel_values = [torch.full((1, 1, 4), float(i)) for i in range(5)] + tgt_sizes = torch.arange(10, dtype=torch.long).reshape(5, 2) + image_input = { + "pixel_values": pixel_values, + "tgt_sizes": tgt_sizes, + "num_slices": torch.tensor([2, 1, 2]), + } + + selected = _select_vision_items(image_input, [2], [2, 1, 2]) + + assert len(selected["pixel_values"]) == 2 + assert bool(selected["pixel_values"][0][0, 0, 0] == 3.0) + assert bool(selected["pixel_values"][1][0, 0, 0] == 4.0) + assert selected["num_slices"].tolist() == [2] + assert torch.equal(selected["tgt_sizes"], tgt_sizes[3:5]) + + def test_select_vision_items_preserves_order(self): + pixel_values = [torch.full((1, 1, 4), float(i)) for i in range(3)] + image_input = { + "pixel_values": pixel_values, + "tgt_sizes": torch.zeros(3, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1, 1]), + } + + selected = _select_vision_items(image_input, [2, 0], [1, 1, 1]) + + assert selected["num_slices"].tolist() == [1, 1] + assert bool(selected["pixel_values"][0][0, 0, 0] == 2.0) + assert bool(selected["pixel_values"][1][0, 0, 0] == 0.0) + + def test_process_vision_input_applies_flags_per_item(self): + calls: list[str | None] = [] + + def fake_get_vision_hidden_states(data, downsample_mode=None): + calls.append(downsample_mode) + tag = 1.0 if downsample_mode == "16x" else 0.0 + return [torch.full((1,), tag) for _ in data["pixel_values"]] + + model = object.__new__(MiniCPMV4_6ForConditionalGeneration) + model.get_vision_hidden_states = fake_get_vision_hidden_states + + image_input = { + "type": "pixel_values", + "pixel_values": [torch.zeros(1, 1, 4) for _ in range(3)], + "tgt_sizes": torch.zeros(3, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1, 1]), + } + + result = model._process_vision_input( + image_input, + use_vit_merger=[ + torch.tensor([True]), + torch.tensor([False]), + torch.tensor([True]), + ], + ) + + # Each item keeps its own mode even though they share one call. + assert [float(r[0]) for r in result] == [1.0, 0.0, 1.0] + # Mixed flags are encoded as separate groups, not one collapsed batch. + assert calls == ["16x", "4x"] + + def test_process_vision_input_uses_one_group_for_uniform_flags(self): + calls: list[str | None] = [] + + def fake_get_vision_hidden_states(data, downsample_mode=None): + calls.append(downsample_mode) + return [torch.zeros(1) for _ in data["pixel_values"]] + + model = object.__new__(MiniCPMV4_6ForConditionalGeneration) + model.get_vision_hidden_states = fake_get_vision_hidden_states + + image_input = { + "type": "pixel_values", + "pixel_values": [torch.zeros(1, 1, 4) for _ in range(2)], + "tgt_sizes": torch.zeros(2, 2, dtype=torch.long), + "num_slices": torch.tensor([1, 1]), + } + + model._process_vision_input(image_input, use_vit_merger=torch.tensor([True])) + + # The common case still makes a single encoder call. + assert calls == ["16x"] diff --git a/tests/model_executor/test_minicpmv4_7.py b/tests/model_executor/test_minicpmv4_7.py new file mode 100644 index 000000000000..3d9e8b467830 --- /dev/null +++ b/tests/model_executor/test_minicpmv4_7.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.model_executor.models.minicpmv4_6 import _stack_vit_merger_qkv +from vllm.model_executor.models.minicpmv4_7 import ( + MiniCPMV4_7MultiModalProcessor, + _compute_canvas_single, + build_image_bounds, + canvas_rope_delta, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +def test_canvas_mrope_sizes_survive_prefix_cache_stripping(): + """A cached image must keep `tgt_sizes`, or canvas M-RoPE falls back. + + `strip_covered_mm_data` drops the payload of a prefix-cache-covered item and + keeps only CPU-side metadata. Canvas M-RoPE places slices from `tgt_sizes`, + so 4.7 marks it `keep_on_cpu` while the shared 4.6 config stays as it was. + """ + from vllm.model_executor.models.minicpmv4_6 import _minicpmv4_6_field_config + + shared = _minicpmv4_6_field_config({}) + assert shared["tgt_sizes"].field.keep_on_cpu is False + assert shared["video_tgt_sizes"].field.keep_on_cpu is False + + processor = MiniCPMV4_7MultiModalProcessor.__new__(MiniCPMV4_7MultiModalProcessor) + fields = processor._get_mm_fields_config({}, {}) + + assert fields["tgt_sizes"].field.keep_on_cpu is True + assert fields["video_tgt_sizes"].field.keep_on_cpu is True + # Other fields keep their previous setting. + assert fields["pixel_values"].field.keep_on_cpu is False + assert fields["video_image_sizes"].field.keep_on_cpu is True + + +def test_modality_mm_kwargs_resolution(): + """Flat keys reach both modalities; scopes stay scoped, and win on conflict. + + `_merge_and_resolve_mm_processor_kwargs` routes a flat key into every scope + whose HF processor declares it, so reading is a plain scope lookup. The + routing itself is upstream's, tested here against the MiniCPM-V schema to + pin the precedence rule (`scoped > flat`). + """ + from vllm.model_executor.models.minicpmv4_6 import ( + _flat_processor_kwargs, + _scoped_mm_kwarg, + ) + from vllm.multimodal.processing.context import _resolve_mm_processor_kwargs + + # What MiniCPMV4_6{Image,Video}ProcessorKwargs declare. + processor_keys = {"downsample_mode", "max_slice_nums"} + schema = { + "text_kwargs": set(), + "images_kwargs": set(processor_keys), + "videos_kwargs": set(processor_keys), + "audio_kwargs": set(), + } + + def read(mm_kwargs, modality, key): + merged = _resolve_mm_processor_kwargs(dict(mm_kwargs), schema) + return _scoped_mm_kwarg(merged, modality, key) + + assert read({}, "image", "downsample_mode") is None + # A flat key is routed into both scopes. + assert read({"downsample_mode": "4x"}, "image", "downsample_mode") == "4x" + assert read({"downsample_mode": "4x"}, "video", "downsample_mode") == "4x" + # A scoped key stays in its scope. + assert ( + read({"images_kwargs": {"downsample_mode": "16x"}}, "image", "downsample_mode") + == "16x" + ) + assert ( + read({"images_kwargs": {"downsample_mode": "16x"}}, "video", "downsample_mode") + is None + ) + videos_only = {"videos_kwargs": {"max_slice_nums": 1}} + assert read(videos_only, "video", "max_slice_nums") == 1 + # Scoped to videos, so image must not see it. + assert read(videos_only, "image", "max_slice_nums") is None + # An explicit scope wins over the flat key it is merged with. + assert ( + read( + {"downsample_mode": "16x", "images_kwargs": {"downsample_mode": "4x"}}, + "image", + "downsample_mode", + ) + == "4x" + ) + # A non-mapping scope is ignored rather than raising. + assert read({"images_kwargs": None}, "image", "downsample_mode") is None + + # `_flat_processor_kwargs` drops the scoped keys and pins the mode. + flat = _flat_processor_kwargs( + { + "downsample_mode": "4x", + "images_kwargs": {"a": 1}, + "videos_kwargs": {"b": 2}, + "audio_kwargs": {"c": 3}, + }, + "16x", + ) + assert flat == {"downsample_mode": "16x"} + + +def test_vit_merger_qkv_is_stacked_before_load(): + mapped = list( + _stack_vit_merger_qkv( + [ + ( + "model.vision_tower.vit_merger.self_attn.q_proj.weight", + torch.zeros(1), + ), + ( + "model.vision_tower.vit_merger.self_attn.k_proj.weight", + torch.zeros(1), + ), + ( + "model.vision_tower.vit_merger.self_attn.v_proj.weight", + torch.zeros(1), + ), + ("vit_merger.self_attn.qkv_proj.weight", torch.zeros(1)), + ] + ) + ) + names = [name for name, _ in mapped] + shards = [getattr(tensor, "shard_id", None) for _, tensor in mapped] + assert names == [ + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "model.vision_tower.vit_merger.self_attn.qkv_proj.weight", + "vit_merger.self_attn.qkv_proj.weight", + ] + assert shards == ["q", "k", "v", None] + + +# Canvas M-RoPE goldens, cross-checked against the reference implementation in +# transformers PR 48979 (`modular_minicpmv4_7.py`). Synthetic ids keep the token +# sequences readable: the canvas math only compares ids for equality. +_IM_START = 10 +_IM_END = 11 +_SLICE_START = 12 +_SLICE_END = 13 +_NEWLINE = 14 +_IMAGE_PAD = 20 +_VIDEO_PAD = 21 +_TEXT = 100 + +_CANVAS_IDS = { + "im_start_id": _IM_START, + "im_end_id": _IM_END, + "slice_start_id": _SLICE_START, + "slice_end_id": _SLICE_END, + "newline_id": _NEWLINE, +} + + +def _canvas_positions(tokens, target_sizes): + input_ids = torch.tensor([tokens], dtype=torch.long) + targets = torch.tensor(target_sizes, dtype=torch.long) + bounds = build_image_bounds(input_ids[0], _CANVAS_IDS) + positions = _compute_canvas_single( + input_ids[0], + torch.arange(len(tokens), dtype=torch.long), + bounds, + targets, + _CANVAS_IDS, + ) + return positions.tolist(), canvas_rope_delta(positions, len(tokens)) + + +def test_canvas_mrope_trailing_newline_stays_inside_the_canvas(): + """A newline right after `` belongs to the canvas, not to the text. + + The chat template joins content parts with ``\\n``, so an image followed by + text always leaves one there. The transformers reference closes a frame over + every marker including newlines; leaving it out pushes every later position + one step forward and makes `rope_delta` differ by one. + """ + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _NEWLINE, _TEXT, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 1, 3, 4], + [0, 0, 1, 2, 1, 3, 4], + [0, 0, 1, 2, 1, 3, 4], + ] + assert delta == -2 + + +def test_canvas_mrope_thumbnail_matches_reference(): + # PAD ; 16x turns a 4x4 patch grid into 1 token. + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 3], + [0, 0, 1, 2, 3], + [0, 0, 1, 2, 3], + ] + assert delta == -1 + + +def test_canvas_mrope_slice_row_matches_reference(): + # Thumbnail then a 1x2 slice row; the slices tile the thumbnail's canvas. + tokens = [ + _TEXT, + _IM_START, + _IMAGE_PAD, + _IM_END, + _SLICE_START, + _IMAGE_PAD, + _SLICE_END, + _SLICE_START, + _IMAGE_PAD, + _SLICE_END, + _TEXT, + ] + positions, delta = _canvas_positions(tokens, [[4, 4], [4, 4], [4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 2, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 3, 1, 1, 1, 2, 2, 2, 4], + ] + assert delta == -6 + + +def test_canvas_mrope_resamples_thumbnail_over_canvas(): + # A 8x8 patch grid is 2x2 LLM tokens, so the canvas is resampled (not arange). + tokens = [_TEXT, _IM_START] + [_IMAGE_PAD] * 4 + [_IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[8, 8]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 4], + [0, 0, 1, 1, 2, 2, 3, 4], + [0, 0, 1, 2, 1, 2, 3, 4], + ] + assert delta == -3 + + +def test_canvas_mrope_video_frames_share_one_canvas(): + tokens = [ + _TEXT, + _IM_START, + _VIDEO_PAD, + _IM_END, + _IM_START, + _VIDEO_PAD, + _IM_END, + _TEXT, + ] + positions, delta = _canvas_positions(tokens, [[4, 4], [4, 4]]) + + assert positions == [ + [0, 1, 1, 1, 3, 3, 3, 5], + [0, 0, 1, 2, 2, 3, 4, 5], + [0, 0, 1, 2, 2, 3, 4, 5], + ] + assert delta == -2 + + +def test_canvas_mrope_interleaved_image_and_video(): + tokens = ( + [_TEXT, _IM_START] + + [_IMAGE_PAD] * 4 + + [_IM_END, _IM_START] + + [_VIDEO_PAD] * 4 + + [_IM_END, _TEXT] + ) + positions, delta = _canvas_positions(tokens, [[8, 8], [8, 8]]) + + assert positions == [ + [0, 1, 1, 1, 1, 1, 1, 4, 4, 4, 4, 4, 4, 7], + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 7], + [0, 0, 1, 2, 1, 2, 3, 3, 4, 5, 4, 5, 6, 7], + ] + assert delta == -6 + + +def test_canvas_span_is_shared_by_image_and_video(): + """A video frame and a same-geometry image must place identical offsets. + + Both paths go through `_assign_canvas_span`; this pins that they cannot + drift apart now that the duplicated geometry is gone. + """ + thumb = [_IM_START] + [_IMAGE_PAD] * 4 + [_IM_END] + + img_pos, _ = _canvas_positions([_TEXT, *thumb, _TEXT], [[8, 8]]) + vid_pos, _ = _canvas_positions([_TEXT, *thumb, *thumb, _TEXT], [[8, 8], [8, 8]]) + + def offsets(pos, start): + # The halo is one below the canvas base. + base = pos[1][start] + 1 + return [[c[start + k] - base for k in range(3)] for c in pos] + + assert offsets(img_pos, 0) == offsets(vid_pos, 0) + + +def test_canvas_mrope_delta_leaves_room_for_the_first_decoded_token(): + # First decoded token must land on max_prefill_position + 1, matching + # transformers and the other native M-RoPE models. + tokens = [_TEXT, _IM_START, _IMAGE_PAD, _IM_END, _TEXT] + positions, delta = _canvas_positions(tokens, [[4, 4]]) + + max_position = max(max(channel) for channel in positions) + assert delta == max_position + 1 - len(tokens) diff --git a/tests/models/registry.py b/tests/models/registry.py index b57f3fcb7277..2ed6d4871c27 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1118,6 +1118,11 @@ def check_available_online( "MiniCPMV4_6ForConditionalGeneration": _HfExamplesInfo( "openbmb/MiniCPM-V-4_6", ), + "MiniCPMV4_7ForConditionalGeneration": _HfExamplesInfo( + "openbmb/MiniCPM-V-4_7", + min_transformers_version="5.18.0", + is_available_online=False, + ), "MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo( "MiniMaxAI/MiniMax-M3", trust_remote_code=True, diff --git a/tests/parser/engine/test_minicpmv_reasoning.py b/tests/parser/engine/test_minicpmv_reasoning.py new file mode 100644 index 000000000000..5edf5bfc5da9 --- /dev/null +++ b/tests/parser/engine/test_minicpmv_reasoning.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MiniCPM-V family parser.""" + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from tests.parser.engine.streaming_helpers import simulate_reasoning_streaming +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.parser.minicpmv import ( + EscapedNewlineNormalizer, + MiniCPMVOutputNormalizer, + recover_newlines, +) +from vllm.parser.parser_manager import ParserManager +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager + +_VOCAB = { + "": 10, + "": 11, + "<|im_end|>": 12, + "": 13, + "": 14, +} + +_TOOLS = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + ) +] + + +@pytest.fixture +def mock_tokenizer(): + return make_mock_tokenizer(_VOCAB) + + +@pytest.fixture +def request_obj(): + return ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + +def _make_parser(tokenizer, *, enable_thinking=None, tools=None): + """Build the parser the serving stack resolves for MiniCPM-V. + + Everything is built through ``ParserManager`` so the tests exercise the + ``DelegatingParser`` the server actually drives, not the parser engine on + its own. Tool calling goes through ``qwen3_coder``: MiniCPM-V emits that + XML syntax verbatim and has no tool parser of its own. + """ + parser_cls = ParserManager.get_parser( + tool_parser_name="qwen3_coder" if tools else None, + reasoning_parser_name="minicpmv", + enable_auto_tools=bool(tools), + ) + assert parser_cls is not None + kwargs = {} + if enable_thinking is not None: + kwargs["chat_template_kwargs"] = {"enable_thinking": enable_thinking} + return parser_cls(tokenizer, tools, **kwargs) + + +def test_reasoning_parser_is_registered(): + parser_cls = ReasoningParserManager.get_reasoning_parser("minicpmv") + + assert parser_cls.__name__ == "MiniCPMVParserReasoningAdapter" + + +def test_thinking_defaults_to_the_chat_template(mock_tokenizer): + """Without `enable_thinking`, the grammar must wait for reasoning. + + The chat template ends the assistant turn with `` unless the flag is + explicitly false, so the model reasons by default. This is the flag + `serving.py` latches into the structured-output gate, so a non-thinking + default both disables the grammar and leaks the thinking into `content`. + """ + + def reasoning_pending(**kwargs): + return not _make_parser(mock_tokenizer, **kwargs).is_reasoning_end([1, 2]) + + assert reasoning_pending() is True # default matches the template + assert reasoning_pending(enable_thinking=True) is True + assert reasoning_pending(enable_thinking=False) is False + + +def test_no_minicpmv_tool_parser_is_registered(): + """MiniCPM-V deliberately registers no tool parser. + + The family emits the Qwen3-Coder XML syntax verbatim, so + ``--tool-call-parser qwen3_coder`` already covers it. Registering a + ``minicpmv`` tool parser would duplicate that name for no parsing gain. + """ + with pytest.raises(KeyError): + ToolParserManager.get_tool_parser("minicpmv") + + +def test_composes_with_qwen3_coder_tool_parser(mock_tokenizer): + """`--tool-call-parser qwen3_coder` must work through the serving Parser.""" + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + parser_cls = ParserManager.get_parser( + tool_parser_name="qwen3_coder", + reasoning_parser_name="minicpmv", + enable_auto_tools=True, + ) + assert parser_cls is not None + parser = parser_cls(mock_tokenizer, request.tools) + + _, _, tool_calls = parser.parse( + "" + "SF", + request, + enable_auto_tools=True, + ) + + assert tool_calls is not None + assert tool_calls[0].name == "get_weather" + assert "SF" in tool_calls[0].arguments + + +class TestNonThinking: + @pytest.fixture + def parser(self, mock_tokenizer): + return _make_parser(mock_tokenizer, enable_thinking=False) + + def test_plain_content_is_unchanged(self, parser): + reasoning, content = parser.extract_reasoning("The answer is 42.", None) + + assert reasoning is None + assert content == "The answer is 42." + + def test_reserved_markers_are_removed(self, parser): + output = ( + "reasoning<|reserved_14|>tool call<|reserved_15|>" + ) + + reasoning, content = parser.extract_reasoning(output, None) + + assert reasoning is None + assert content == "reasoningtool call" + + def test_unrelated_reserved_text_is_unchanged(self, parser): + reasoning, content = parser.extract_reasoning( + "content<|reserved_16|>", + None, + ) + + assert reasoning is None + assert content == "content<|reserved_16|>" + + def test_dangling_standard_think_end_is_removed(self, parser): + output = "first part \n\nfinal part" + + reasoning, content = parser.extract_reasoning(output, None) + + assert reasoning is None + assert content == "first part \n\nfinal part" + + def test_truncated_standard_reasoning_does_not_leak(self, parser): + reasoning, content = parser.extract_reasoning( + "prefixprivate reasoning", + None, + ) + + assert reasoning is None + assert content == "prefix" + + def test_streaming_split_standard_markers(self, parser): + reasoning, content = simulate_reasoning_streaming( + parser, + [ + "<", + "think>", + "private reasoning", + "", + "final answer", + ], + ) + + assert reasoning == "" + assert content == "final answer" + + def test_streaming_split_reserved_markers(self, parser): + reasoning, content = simulate_reasoning_streaming( + parser, + [ + "", + "reasoning", + "<|reserved", + "_13|>", + "final answer", + ], + ) + + assert reasoning == "" + assert content == "reasoningfinal answer" + + def test_tool_calls_are_parsed(self, mock_tokenizer): + """Thinking is off by default, and tool calls must still be parsed.""" + parser = _make_parser(mock_tokenizer, enable_thinking=False, tools=_TOOLS) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + + reasoning, content, tool_calls = parser.parse( + "" + "SF", + request, + enable_auto_tools=True, + ) + + assert reasoning is None + assert not content + assert tool_calls is not None + assert tool_calls[0].name == "get_weather" + assert "SF" in tool_calls[0].arguments + + def test_streaming_tool_calls_are_parsed(self, mock_tokenizer): + """Regression: tool calls must stream when thinking is explicitly off. + + The reasoning parser used to report ``reasoning_ended=False`` in + non-thinking mode, which kept the delegating parser in its reasoning + phase and streamed the raw ```` XML to the client as content. + """ + parser = _make_parser(mock_tokenizer, enable_thinking=False, tools=_TOOLS) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + tools=_TOOLS, + ) + + delta = parser.parse_delta( + "" + "SF", + [], + request, + prompt_token_ids=[], + finished=True, + ) + + assert delta is not None + assert delta.tool_calls + assert delta.tool_calls[0].function is not None + assert delta.tool_calls[0].function.name == "get_weather" + assert "SF" in (delta.tool_calls[0].function.arguments or "") + assert not delta.content + + +class TestThinking: + def test_standard_think_tags_are_parsed(self, mock_tokenizer): + parser = _make_parser(mock_tokenizer, enable_thinking=True) + + reasoning, content = parser.extract_reasoning( + "private reasoningfinal answer", + None, + ) + + assert reasoning == "private reasoning" + assert content == "final answer" + + def test_reserved_markers_are_removed(self, mock_tokenizer): + parser = _make_parser(mock_tokenizer, enable_thinking=True) + + reasoning, content = parser.extract_reasoning( + "private reasoning<|reserved_13|>final answer", + None, + ) + + assert reasoning == "private reasoning" + assert content == "final answer" + + def test_streaming_im_end_is_removed(self): + tokenizer = make_mock_tokenizer(_VOCAB, special_tokens=[]) + parser = _make_parser(tokenizer, enable_thinking=True) + + reasoning, content = simulate_reasoning_streaming( + parser, + ["Answer: 111", "<|im_end|>"], + ) + + assert reasoning == "Answer: 111" + assert content == "" + + +class TestNewlineRecovery: + """The tokenizer writes a newline as the two characters ``\\`` and ``n``.""" + + def test_recover_newlines(self): + text = ( + "first\\nsecond\\rthird\\r\\nfourth " + "`code\\nvalue` " + "```python\\nvalue = '\\\\n'\\n``` " + "$x\\ny$ $$x\\ny$$ \\(x\\ny\\) \\[x\\ny\\] " + "escaped \\\\n" + ) + + assert recover_newlines(text) == ( + "first\nsecond\nthird\nfourth " + "`code\\nvalue` " + "```python\\nvalue = '\\\\n'\\n``` " + "$x\\ny$ $$x\\ny$$ \\(x\\ny\\) \\[x\\ny\\] " + "escaped \\\\n" + ) + + def test_streaming_matches_complete_recovery(self): + text = ( + "first\\nsecond `code\\nvalue` $x\\ny$ ```text\\nvalue\\n``` last\\r\\nline" + ) + expected = recover_newlines(text) + + for split in range(len(text) + 1): + normalizer = EscapedNewlineNormalizer() + actual = normalizer.feed(text[:split]) + actual += normalizer.feed(text[split:], final=True) + assert actual == expected + + normalizer = EscapedNewlineNormalizer() + actual = "".join(normalizer.feed(char) for char in text) + actual += normalizer.feed("", final=True) + assert actual == expected + + def test_backslash_runs_stay_consistent_across_splits(self): + # A run of backslashes can be split over chunks. Its parity decides + # whether a following `n` is a newline, so streaming and non-streaming + # must agree for every split point. + for count in range(1, 7): + text = "\\" * count + "n" + expected = recover_newlines(text) + + for split in range(len(text) + 1): + normalizer = EscapedNewlineNormalizer() + actual = normalizer.feed(text[:split]) + actual += normalizer.feed(text[split:], final=True) + assert actual == expected, (count, split) + + def test_streaming_recovers_reasoning_and_content(self): + normalizer = MiniCPMVOutputNormalizer() + + first = normalizer.normalize_delta( + DeltaMessage(reasoning="reason\\"), + finished=False, + ) + second = normalizer.normalize_delta( + DeltaMessage(reasoning="nnext\\", content="answer\\"), + finished=False, + ) + final = normalizer.normalize_delta( + DeltaMessage(content="nend"), + finished=True, + ) + + assert first == DeltaMessage(reasoning="reason") + assert second == DeltaMessage(reasoning="\nnext\\", content="answer") + assert final == DeltaMessage(content="\nend") + + def test_streaming_does_not_modify_tool_arguments(self): + normalizer = MiniCPMVOutputNormalizer() + arguments = '{"value":"first\\\\nsecond"}' + delta = DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=0, + function=DeltaFunctionCall(arguments=arguments), + ) + ] + ) + + result = normalizer.normalize_delta(delta, finished=True) + + assert result is not None + assert result.tool_calls[0].function is not None + assert result.tool_calls[0].function.arguments == arguments + + def test_parse_recovers_content(self, mock_tokenizer, request_obj): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + + reasoning, content, tool_calls = parser.parse( + "first\\nsecond", + request_obj, + ) + + assert reasoning is None + assert content == "first\nsecond" + assert not tool_calls + + def test_parse_delta_recovers_content(self, mock_tokenizer, request_obj): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + + chunk = parser.parse_delta( + "first\\", + [], + request_obj, + finished=False, + ) + rest = parser.parse_delta( + "nsecond", + [], + request_obj, + finished=True, + ) + + # The trailing half of the escape is held back until it is complete. + assert chunk is not None + assert chunk.content == "first" + assert rest is not None + assert rest.content == "\nsecond" + + def test_buffered_delta_is_dropped_not_sent_empty( + self, mock_tokenizer, request_obj + ): + parser = _make_parser(mock_tokenizer, enable_thinking=False) + parser.parse_delta("hello", [], request_obj, finished=False) + + # The lone backtick is held back for the next chunk. Reporting it as an + # all-None DeltaMessage would reach the client as an empty SSE chunk. + assert parser.parse_delta("`", [], request_obj, finished=False) is None + + def test_normalizer_drops_empty_delta(self): + normalizer = MiniCPMVOutputNormalizer() + + assert ( + normalizer.normalize_delta(DeltaMessage(content="`"), finished=False) + is None + ) + + def test_normalizer_flushes_buffered_content_when_finished(self): + normalizer = MiniCPMVOutputNormalizer() + normalizer.normalize_delta(DeltaMessage(content="`"), finished=False) + + final = normalizer.normalize_delta(DeltaMessage(), finished=True) + + assert final is not None + assert final.content == "`" diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 0a96c2d8479d..333ab8eedd66 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -45,6 +45,7 @@ HasInnerState, IsHybrid, MultiModalEmbeddings, + SupportsEagle3, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -87,35 +88,186 @@ def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]): fields["use_vit_merger"] = MultiModalFieldConfig.batched( "image", keep_on_cpu=True ) + if "video_use_vit_merger" in hf_inputs: + fields["video_use_vit_merger"] = MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ) return fields +def _expand_flags(value: object, num_items: int) -> list[bool | None]: + """Broadcast a batched per-item flag to exactly one entry per item.""" + if value is None: + return [None] * num_items + + raw: list[object] + if isinstance(value, torch.Tensor): + raw = value.reshape(-1).tolist() + elif isinstance(value, list | tuple): + raw = [] + for item in value: + if isinstance(item, torch.Tensor): + raw.extend(item.reshape(-1).tolist()) + else: + raw.append(item) + else: + raw = [value] + + if len(raw) == num_items: + return [bool(v) for v in raw] + if len(raw) == 1: + return [bool(raw[0])] * num_items + # Unexpected shape: keep one decision for the whole batch instead of + # failing the request. + return [any(bool(v) for v in raw)] * num_items + + +def _vision_downsample_mode(flag: bool | None) -> str | None: + if flag is None: + return None + return "16x" if flag else "4x" + + +def _select_vision_items( + image_input: MiniCPMVImagePixelInputs, + item_indices: list[int], + num_slices: list[int], +) -> MiniCPMVImagePixelInputs: + """Slice a vision input down to the given items, keeping their order.""" + offsets = [0] + for n in num_slices: + offsets.append(offsets[-1] + n) + + pixel_values: list[torch.Tensor] = [] + tgt_sizes: list[torch.Tensor] = [] + sizes: list[int] = [] + for i in item_indices: + start, end = offsets[i], offsets[i + 1] + pixel_values.extend(image_input["pixel_values"][start:end]) + tgt_sizes.append(image_input["tgt_sizes"][start:end]) + sizes.append(end - start) + + return MiniCPMVImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + tgt_sizes=torch.cat(tgt_sizes, dim=0), + num_slices=torch.tensor(sizes), + ) + + +# HF-style modality-scoped keys inside `mm_processor_kwargs`. They scope a value +# to a single modality and must not be forwarded to the processors, which take +# flat arguments only. +_MODALITY_SCOPED_MM_KWARGS = ("images_kwargs", "videos_kwargs", "audio_kwargs") + + +def _flat_processor_kwargs( + mm_kwargs: Mapping[str, object], + downsample_mode: str, +) -> dict[str, object]: + """Drop the modality-scoped nesting and pin the resolved downsample mode. + + The placeholder count and the vision tower must read the same mode, so the + value resolved for this modality is passed down explicitly rather than + leaving the processor to fall back to its own default. + """ + kwargs = { + key: value + for key, value in mm_kwargs.items() + if key not in _MODALITY_SCOPED_MM_KWARGS + } + kwargs["downsample_mode"] = downsample_mode + return kwargs + + +_MODALITY_TO_MM_KWARGS_SCOPE = { + "image": "images_kwargs", + "video": "videos_kwargs", +} + + +def _scoped_mm_kwarg( + merged: Mapping[str, object], + modality: str, + key: str, +) -> Any: + """Read a processor kwarg for one modality. + + ``merged`` is what ``_merge_and_resolve_mm_processor_kwargs`` returns: it + routes a flat key into every scope whose processor declares it, and a value + scoped to the other modality stays there, so it is not visible here. + """ + scoped = merged.get(_MODALITY_TO_MM_KWARGS_SCOPE[modality]) + if isinstance(scoped, Mapping): + return scoped.get(key) + return None + + class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): def _resolve_downsample_mode( self, mm_kwargs: Mapping[str, object], + modality: str = "image", ) -> str: - ds = mm_kwargs.get("downsample_mode") - if ds is not None: - return str(ds) + # A flat `downsample_mode` applies to both modalities; the nested + # `images_kwargs` / `videos_kwargs` forms scope it to one, matching + # MiniCPMV4_6Processor. info = self.info assert isinstance(info, MiniCPMV4_6ProcessingInfo) + merged = info._merge_and_resolve_mm_processor_kwargs(mm_kwargs) + downsample_mode = _scoped_mm_kwarg(merged, modality, "downsample_mode") + if downsample_mode is not None: + return str(downsample_mode) return info._get_downsample_mode() + def _resolve_max_slice_nums( + self, + mm_kwargs: Mapping[str, object], + modality: str = "image", + ) -> int | None: + # Per-request max_slice_nums. None keeps the processor default. + # Resolved like `downsample_mode` so the nested per-modality form + # (`images_kwargs` / `videos_kwargs`) is honoured too. This must match + # what `process_images` passes to the image processor: the placeholder + # count and the visual token count are derived from it separately, so + # they diverge (and the engine aborts) if the two sides read different + # values. + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + merged = info._merge_and_resolve_mm_processor_kwargs(mm_kwargs) + max_slice = _scoped_mm_kwarg(merged, modality, "max_slice_nums") + if max_slice is None: + return None + return int(max_slice) + def get_image_prompt_texts( self, image_size, image_idx: int = 0, downsample_mode: str | None = None, + max_slice_nums: int | None = None, ) -> str: info = self.info assert isinstance(info, MiniCPMV4_6ProcessingInfo) return info.get_slice_image_placeholder( image_size, image_idx=image_idx, + max_slice_nums=max_slice_nums, downsample_mode=downsample_mode, ) + def _video_local_id_prefix(self, video_idx: int) -> str: + """`{idx}` prefix emitted in front of a video placeholder. + + MiniCPM-V 4.6 numbers each video with a local id, matching + MiniCPMV4_6Processor. 4.7 treats a video as one temporal sequence and is + trained without these ids, so it overrides this to return an empty string. + """ + tokenizer = self.info.get_tokenizer() + id_start = getattr(tokenizer, "image_id_start_token", "") + id_end = getattr(tokenizer, "image_id_end_token", "") + return f"{id_start}{video_idx}{id_end}" + def get_video_prompt_texts( self, image_size, @@ -144,8 +296,6 @@ def get_video_prompt_texts( image_end = getattr(tokenizer, "image_end_token", "") slice_start = getattr(tokenizer, "slice_start_token", "") slice_end = getattr(tokenizer, "slice_end_token", "") - id_start = getattr(tokenizer, "image_id_start_token", "") - id_end = getattr(tokenizer, "image_id_end_token", "") per_frame = image_start + video_token * source_tokens + image_end if grids[0] > 0 and grids[1] > 0 and patch_tokens > 0: @@ -154,7 +304,7 @@ def get_video_prompt_texts( per_frame += "\n".join(rows) body = per_frame * num_frames - return f"{id_start}{video_idx}{id_end}" + body + return self._video_local_id_prefix(video_idx) + body def process_images( self, @@ -179,10 +329,18 @@ def process_images( n_images = len(parsed_images) image_processor = self.info.get_image_processor() patch_size = image_processor.patch_size + ds_mode = self._resolve_downsample_mode(mm_kwargs, modality="image") + image_proc_kwargs = _flat_processor_kwargs(mm_kwargs, ds_mode) + # `_flat_processor_kwargs` drops the nested form, so the value resolved + # from it has to be pinned back for the processor and the prompt text + # to agree on the slice count. + image_max_slice = self._resolve_max_slice_nums(mm_kwargs, modality="image") + if image_max_slice is not None: + image_proc_kwargs["max_slice_nums"] = image_max_slice per_image_pixel_values: list[torch.Tensor] = [] per_image_tgt_sizes: list[torch.Tensor] = [] for image in parsed_images: - ip_out = image_processor([image], **mm_kwargs) + ip_out = image_processor([image], **image_proc_kwargs) pv = ip_out["pixel_values"] # (1, C, P, sum_W) ts = ip_out["target_sizes"] # (n_slices, 2) if pv.ndim == 4 and pv.shape[0] == 1: @@ -210,7 +368,6 @@ def process_images( "tgt_sizes": per_image_tgt_sizes, } - ds_mode = self._resolve_downsample_mode(mm_kwargs) insert_layer_id = getattr( self.info.get_hf_config(), "insert_layer_id", @@ -245,7 +402,11 @@ def process_videos( image_processor = self.info.get_image_processor() patch_size = image_processor.patch_size video_max_slice = self.info.get_video_max_slice_num() - video_mm_kwargs = {**mm_kwargs, "max_slice_nums": video_max_slice} + video_ds_mode = self._resolve_downsample_mode(mm_kwargs, modality="video") + video_mm_kwargs = { + **_flat_processor_kwargs(mm_kwargs, video_ds_mode), + "max_slice_nums": video_max_slice, + } per_video_pixel_values: list[torch.Tensor] = [] per_video_tgt_sizes: list[torch.Tensor] = [] @@ -312,10 +473,17 @@ def process_videos( if not per_video_pixel_values: return {} + insert_layer_id = getattr(self.info.get_hf_config(), "insert_layer_id", -1) + merger_flag = video_ds_mode != "4x" and insert_layer_id >= 0 + return { "video_pixel_values": per_video_pixel_values, "video_tgt_sizes": per_video_tgt_sizes, "video_image_sizes": per_video_image_sizes, + "video_use_vit_merger": [ + torch.tensor([merger_flag], dtype=torch.bool) + for _ in per_video_pixel_values + ], } def _get_prompt_updates( @@ -324,7 +492,13 @@ def _get_prompt_updates( hf_processor_mm_kwargs: Mapping[str, object], out_mm_kwargs, ): - ds_mode = self._resolve_downsample_mode(hf_processor_mm_kwargs) + image_ds_mode = self._resolve_downsample_mode( + hf_processor_mm_kwargs, modality="image" + ) + video_ds_mode = self._resolve_downsample_mode( + hf_processor_mm_kwargs, modality="video" + ) + image_max_slice = self._resolve_max_slice_nums(hf_processor_mm_kwargs) placeholders = [ ("image", self.info.image_pattern), @@ -360,7 +534,8 @@ def get_image_replacement(item_idx: int): self.get_image_prompt_texts( image_size, item_idx, - downsample_mode=ds_mode, + downsample_mode=image_ds_mode, + max_slice_nums=image_max_slice, ), add_special_tokens=False, ), @@ -389,7 +564,7 @@ def get_video_replacement(item_idx: int): self.get_video_prompt_texts( frame_size, num_frames, - downsample_mode=ds_mode, + downsample_mode=video_ds_mode, video_idx=item_idx, ), add_special_tokens=False, @@ -409,7 +584,7 @@ def get_video_replacement(item_idx: int): self.get_video_prompt_texts( frame_size, num_frames, - downsample_mode=ds_mode, + downsample_mode=video_ds_mode, video_idx=item_idx, ), add_special_tokens=False, @@ -437,31 +612,28 @@ def _recompute_cached_prompt_update( new_update = super()._recompute_cached_prompt_update( cached_update, new_item_idx ) - # MiniCPM-V 4.6 prefixes video placeholders with `{idx}` - # (the base class only rewrites the image modality). + # Video placeholders carry an `{idx}` prefix that + # the base class does not rewrite. 4.7 has no prefix, so this is a + # no-op there. if cached_update.modality == "video": - tokenizer = self.info.get_tokenizer() - id_start = getattr(tokenizer, "image_id_start_token", "") - id_end = getattr(tokenizer, "image_id_end_token", "") - video_token = getattr(tokenizer, "video_token", "<|video_pad|>") - - text = tokenizer.decode(cached_update.content.full) - prev_item_idx = cached_update.item_idx - - new_update = new_update.with_content( - PromptUpdateDetails.select_token_ids( - cached_encode( - tokenizer, - text.replace( - f"{id_start}{prev_item_idx}{id_end}", - f"{id_start}{new_item_idx}{id_end}", - 1, + old_prefix = self._video_local_id_prefix(cached_update.item_idx) + new_prefix = self._video_local_id_prefix(new_item_idx) + if old_prefix != new_prefix: + tokenizer = self.info.get_tokenizer() + video_token = getattr(tokenizer, "video_token", "<|video_pad|>") + + text = tokenizer.decode(cached_update.content.full) + + new_update = new_update.with_content( + PromptUpdateDetails.select_token_ids( + cached_encode( + tokenizer, + text.replace(old_prefix, new_prefix, 1), + add_special_tokens=False, ), - add_special_tokens=False, - ), - cached_encode(tokenizer, video_token, add_special_tokens=False), + cached_encode(tokenizer, video_token, add_special_tokens=False), + ) ) - ) return new_update def _get_mm_fields_config( @@ -676,6 +848,31 @@ def get_num_image_tokens( return source_tokens + grids[0] * grids[1] * patch_tokens +def _stack_vit_merger_qkv( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterable[tuple[str, torch.Tensor]]: + """Fold separate vit_merger q/k/v projections into the fused qkv_proj module. + + HF checkpoints store the window-attention projections separately, while + vLLM keeps them fused. The shard_id mechanism on the submodule's + WeightsMapper cannot do this: AutoWeightsLoader still recurses over every + incoming key after calling a submodule's load_weights, so a "k_proj" key + with no matching child raises before the mapping is honoured. + """ + pairs = ( + ("vit_merger.self_attn.q_proj.", "q"), + ("vit_merger.self_attn.k_proj.", "k"), + ("vit_merger.self_attn.v_proj.", "v"), + ) + for name, tensor in weights: + for src, shard in pairs: + if src in name: + name = name.replace(src, "vit_merger.self_attn.qkv_proj.", 1) + tensor.shard_id = shard + break + yield name, tensor + + class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): def __init__( self, @@ -957,6 +1154,7 @@ class MiniCPMV4_6ForConditionalGeneration( SupportsPP, HasInnerState, IsHybrid, + SupportsEagle3, SupportsMRoPE, ): supports_encoder_tp_data = True @@ -1031,7 +1229,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): apply_encoder_attention_mask=True, prefix=maybe_prefix(prefix, "vpm"), ) - if config.drop_vision_last_layer: + if self._drop_vision_last_layer(): self.vpm.encoder.layers = self.vpm.encoder.layers[:-1] self.vit_merger = MiniCPMV4_6ViTWindowAttentionMerger( @@ -1062,6 +1260,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.language_model.make_empty_intermediate_tensors ) + def _drop_vision_last_layer(self) -> bool: + """Whether the last vision encoder layer is dropped. + + MiniCPM-V 4.6 checkpoints always define this. 4.7's config no longer + carries the field, so it overrides this. + """ + return self.config.drop_vision_last_layer + # ----- Multimodal parsing ----- def _parse_and_validate_vision_input( @@ -1174,35 +1380,35 @@ def _process_vision_input(self, image_input, use_vit_merger=None): if image_input["type"] == "image_embeds": return image_input["image_embeds"] - downsample_mode = None - if use_vit_merger is not None: - downsample_mode = "16x" if use_vit_merger else "4x" - image_features = self.get_vision_hidden_states( - image_input, - downsample_mode=downsample_mode, - ) - num_slices = image_input["num_slices"] - results = [] - idx = 0 - for n in num_slices.tolist(): - group = image_features[idx : idx + n] - results.append(torch.cat(group, dim=0)) - idx += n - return results + num_slices = image_input["num_slices"].tolist() + flags = _expand_flags(use_vit_merger, len(num_slices)) + + # The merger flag is per request, but the model runner may batch items + # from several concurrent requests. Encoding such a mixed batch would + # apply one request's downsample mode to another request's items, so + # group by flag and encode each group separately. + encoded: dict[int, torch.Tensor] = {} + for flag in dict.fromkeys(flags): + item_indices = [i for i, f in enumerate(flags) if f == flag] + group_input = _select_vision_items(image_input, item_indices, num_slices) + group_features = self.get_vision_hidden_states( + group_input, + downsample_mode=_vision_downsample_mode(flag), + ) + + idx = 0 + for i in item_indices: + n = num_slices[i] + encoded[i] = torch.cat(group_features[idx : idx + n], dim=0) + idx += n + + return [encoded[i] for i in range(len(num_slices))] # ----- Multimodal embedding interface ----- def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - use_vit_merger_tensors = kwargs.pop("use_vit_merger", None) - use_vit_merger = None - if use_vit_merger_tensors is not None: - if isinstance(use_vit_merger_tensors, torch.Tensor): - use_vit_merger = bool(use_vit_merger_tensors.any().item()) - elif isinstance(use_vit_merger_tensors, list | tuple): - use_vit_merger = any( - bool(t.any().item()) if isinstance(t, torch.Tensor) else bool(t) - for t in use_vit_merger_tensors - ) + use_vit_merger = kwargs.pop("use_vit_merger", None) + video_use_vit_merger = kwargs.pop("video_use_vit_merger", None) image_kwargs = { k: v @@ -1235,7 +1441,7 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: multimodal_embeddings += tuple( self._process_vision_input( video_input, - use_vit_merger=use_vit_merger, + use_vit_merger=video_use_vit_merger, ) ) @@ -1298,7 +1504,10 @@ def load_weights( weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + return loader.load_weights( + _stack_vit_merger_qkv(weights), + mapper=self.hf_to_vllm_mapper, + ) def get_mm_mapping(self) -> MultiModelKeys: return MultiModelKeys.from_string_field( diff --git a/vllm/model_executor/models/minicpmv4_7.py b/vllm/model_executor/models/minicpmv4_7.py new file mode 100644 index 000000000000..3ce860916965 --- /dev/null +++ b/vllm/model_executor/models/minicpmv4_7.py @@ -0,0 +1,764 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MiniCPM-V 4.7 model. + +Same vision/LLM stack as MiniCPM-V 4.6, plus canvas 3D M-RoPE. +""" + +import math +from collections.abc import Mapping +from typing import Any, NamedTuple + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalFieldConfig + +from .minicpmv import MiniCPMVDummyInputsBuilder +from .minicpmv4_6 import ( + MiniCPMV4_6ForConditionalGeneration, + MiniCPMV4_6MultiModalProcessor, + MiniCPMV4_6ProcessingInfo, + _scoped_mm_kwarg, +) + +logger = init_logger(__name__) + + +class MiniCPMV4_7ProcessingInfo(MiniCPMV4_6ProcessingInfo): + def _mm_max_slice_nums(self, **kwargs: object) -> int | None: + # No schema resolution here: it would call back into `get_hf_processor` + # below, which resolves this same value. + merged = self.ctx.get_merged_mm_kwargs(kwargs) + max_slice = _scoped_mm_kwarg(merged, "image", "max_slice_nums") + if max_slice is None: + max_slice = merged.get("max_slice_nums") + if max_slice is None: + return None + return int(max_slice) + + def _configured_max_slice_nums(self) -> int | None: + # Configured (server-level) value only: per-request values are passed + # explicitly to the image processor call, so they must not be written + # back to the shared processor instance. + return self._mm_max_slice_nums() + + def get_hf_processor(self, **kwargs: object): + hf_processor = super().get_hf_processor(**kwargs) + image_processor = getattr(hf_processor, "image_processor", None) + configured = self._configured_max_slice_nums() + if image_processor is not None and configured is not None: + image_processor.max_slice_nums = configured + return hf_processor + + def get_image_max_slice_num(self) -> int: + max_slice = self._configured_max_slice_nums() + if max_slice is not None: + return max_slice + return super().get_image_max_slice_num() + + +class MiniCPMV4_7MultiModalProcessor(MiniCPMV4_6MultiModalProcessor): + def _get_mm_fields_config( + self, + hf_inputs, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields = dict(super()._get_mm_fields_config(hf_inputs, hf_processor_mm_kwargs)) + # Canvas M-RoPE places slices from `tgt_sizes`. Stripping the payload of + # a prefix-cache-covered item only keeps CPU-side metadata, so without + # this a cached image loses its target sizes, `image_bounds` and + # `target_sizes` disagree, and the whole request silently falls back to + # sequential positions. + for key, modality in (("tgt_sizes", "image"), ("video_tgt_sizes", "video")): + if key in fields: + fields[key] = MultiModalFieldConfig.batched(modality, keep_on_cpu=True) + return fields + + def _video_local_id_prefix(self, video_idx: int) -> str: + # MiniCPMV4_7Processor emits no local `` in front of video + # placeholders: a video is one temporal sequence and 4.7 was trained + # without them. + return "" + + def get_image_prompt_texts( + self, + image_size, + image_idx: int = 0, + downsample_mode: str | None = None, + max_slice_nums: int | None = None, + ) -> str: + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + if max_slice_nums is None: + max_slice_nums = info.get_image_max_slice_num() + return info.get_slice_image_placeholder( + image_size, + image_idx=image_idx, + max_slice_nums=max_slice_nums, + downsample_mode=downsample_mode, + ) + + +def _as_hw_pairs(tensor: torch.Tensor) -> torch.Tensor: + data = tensor.to(torch.long) + if data.numel() == 0: + return data.reshape(0, 2) + if data.shape[-1] != 2: + raise ValueError(f"expected last dim 2, got {tuple(data.shape)}") + return data.reshape(-1, 2) + + +def _sequential_positions(seq_len: int) -> torch.Tensor: + """Positions for the fallback where canvas M-RoPE cannot be computed.""" + return torch.arange(seq_len).unsqueeze(0).expand(3, -1) + + +@MULTIMODAL_REGISTRY.register_processor( + MiniCPMV4_7MultiModalProcessor, + info=MiniCPMV4_7ProcessingInfo, + dummy_inputs=MiniCPMVDummyInputsBuilder, +) +class MiniCPMV4_7ForConditionalGeneration(MiniCPMV4_6ForConditionalGeneration): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + self._init_canvas_mrope(vllm_config) + + def _drop_vision_last_layer(self) -> bool: + # MiniCPMV4_7Config no longer carries this field (transformers marks it + # removed), and 4.7 keeps every vision encoder layer. + return False + + def _init_canvas_mrope(self, vllm_config: VllmConfig) -> None: + config = self.config + uses_canvas = bool(getattr(config, "uses_mrope_canvas", False)) + if not uses_canvas: + mrope_mode = getattr(config, "mrope_mode", None) + uses_canvas = mrope_mode is not None and "canvas" in str(mrope_mode).lower() + + self._uses_canvas_mrope = uses_canvas + self._canvas_special_ids: dict[str, int] | None = None + if not uses_canvas: + return + + # MiniCPMV4_7Config pins the structural ids in config.json; the + # conversion script resolves them from the tokenizer. Prefer them so the + # canvas matches training, and fall back to the tokenizer for checkpoints + # whose config predates these fields. + configured_ids: dict[str, int | None] = { + "im_start_id": getattr(config, "image_start_id", None), + "im_end_id": getattr(config, "image_end_id", None), + "slice_start_id": getattr(config, "slice_start_id", None), + "slice_end_id": getattr(config, "slice_end_id", None), + "newline_id": getattr(config, "newline_id", None), + } + if all(value is not None for value in configured_ids.values()): + self._canvas_special_ids = { + name: int(value) + for name, value in configured_ids.items() + if value is not None + } + return + + from vllm.tokenizers.registry import get_tokenizer + + model_config = vllm_config.model_config + tokenizer = get_tokenizer( + model_config.tokenizer, + revision=model_config.tokenizer_revision, + trust_remote_code=model_config.trust_remote_code, + ) + self._canvas_special_ids = { + "im_start_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "image_start_token", "") + ), + "im_end_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "image_end_token", "") + ), + "slice_start_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "slice_start_token", "") + ), + "slice_end_id": tokenizer.convert_tokens_to_ids( + getattr(tokenizer, "slice_end_token", "") + ), + "newline_id": tokenizer.encode("\n", add_special_tokens=False)[0], + } + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + seq_len = len(input_tokens) + if ( + not getattr(self, "_uses_canvas_mrope", False) + or self._canvas_special_ids is None + or not mm_features + ): + return _sequential_positions(seq_len), 0 + + input_ids = torch.tensor(input_tokens, dtype=torch.long) + image_bounds = build_image_bounds(input_ids, self._canvas_special_ids) + + target_sizes_list: list[torch.Tensor] = [] + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + if mm_feature.data is None: + continue + for key in ("tgt_sizes", "video_tgt_sizes"): + tgt = mm_feature.data.get(key) + if tgt is None or not isinstance(tgt.data, torch.Tensor): + continue + try: + target_sizes_list.append(_as_hw_pairs(tgt.data)) + except ValueError: + # Not an (N, 2) size table. The bounds/sizes match below + # turns this into a sequential fallback. + continue + + if not target_sizes_list or image_bounds.numel() == 0: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE missing inputs: " + "tgt_chunks=%s image_bounds=%s; using sequential positions.", + len(target_sizes_list), + tuple(image_bounds.shape) if image_bounds is not None else None, + ) + return _sequential_positions(seq_len), 0 + + target_sizes = torch.cat(target_sizes_list, dim=0) + if target_sizes.ndim != 2 or target_sizes.shape[-1] != 2: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE unexpected tgt_sizes shape %s; " + "using sequential positions.", + tuple(target_sizes.shape), + ) + return _sequential_positions(seq_len), 0 + if image_bounds.shape[0] != target_sizes.shape[0]: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE size mismatch: " + "image_bounds=%s target_sizes=%s; using sequential positions.", + tuple(image_bounds.shape), + tuple(target_sizes.shape), + ) + return _sequential_positions(seq_len), 0 + + try: + pos3d = _compute_canvas_single( + input_ids, + torch.arange(seq_len, dtype=torch.long), + image_bounds, + target_sizes, + self._canvas_special_ids, + ) + except Exception: + logger.warning( + "MiniCPM-V 4.7 canvas M-RoPE failed; " + "falling back to sequential positions.", + exc_info=True, + ) + return _sequential_positions(seq_len), 0 + return pos3d, canvas_rope_delta(pos3d, seq_len) + + +# Canvas 3D M-RoPE: position ids for the slice canvas. +def compute_llm_grid( + n_tokens: int, vision_grid_height: int, vision_grid_width: int +) -> tuple[int, int]: + if n_tokens <= 0 or vision_grid_height <= 0 or vision_grid_width <= 0: + return 0, 0 + vit_total = vision_grid_height * vision_grid_width + if vit_total == n_tokens: + return vision_grid_height, vision_grid_width + if vit_total < n_tokens: + return 0, 0 + factor = int(round(math.sqrt(vit_total / n_tokens))) + if ( + factor > 0 + and vision_grid_height % factor == 0 + and vision_grid_width % factor == 0 + ): + height, width = vision_grid_height // factor, vision_grid_width // factor + if height * width == n_tokens: + return height, width + ratio = vision_grid_height / vision_grid_width + width = max(1, int(round(math.sqrt(n_tokens / ratio)))) + height = n_tokens // width + if height * width == n_tokens and height > 0: + return height, width + return 0, 0 + + +def build_image_bounds( + input_ids: torch.LongTensor, special_token_ids: dict +) -> torch.LongTensor: + start_ids = [ + x + for x in ( + special_token_ids.get("im_start_id"), + special_token_ids.get("slice_start_id"), + ) + if x is not None + ] + end_ids = [ + x + for x in ( + special_token_ids.get("im_end_id"), + special_token_ids.get("slice_end_id"), + ) + if x is not None + ] + if not start_ids or not end_ids: + return torch.zeros(0, 2, dtype=torch.long, device=input_ids.device) + start_cond = torch.zeros_like(input_ids, dtype=torch.bool) + end_cond = torch.zeros_like(input_ids, dtype=torch.bool) + for sid in start_ids: + start_cond |= input_ids == sid + for eid in end_ids: + end_cond |= input_ids == eid + starts = torch.where(start_cond)[0] + 1 + ends = torch.where(end_cond)[0] + n = min(len(starts), len(ends)) + if n == 0: + return torch.zeros(0, 2, dtype=torch.long, device=input_ids.device) + return torch.stack([starts[:n], ends[:n]], dim=-1) + + +def _compute_canvas_single( + input_ids, position_ids_2d, image_bound, target_sizes, special_token_ids +): + # Inference always handles one unpadded sequence. The packed (cu_seqlens) + # form below is kept only to mirror the transformers reference + # implementation of its canvas position computation. + seq_len = input_ids.shape[0] + cu_seqlens = torch.tensor([0, seq_len], device=input_ids.device, dtype=torch.long) + return _compute_canvas_packed( + position_ids_2d.unsqueeze(0), + cu_seqlens, + image_bound, + target_sizes, + input_ids.unsqueeze(0), + special_token_ids, + )[:, 0, :] + + +def _target_hw(target_sizes: torch.Tensor, index: int) -> tuple[int, int]: + """Patch-grid height/width recorded for one image or slice.""" + return int(target_sizes[index, 0]), int(target_sizes[index, 1]) + + +class _CanvasLayout(NamedTuple): + """Grid of one canvas measured in LLM tokens. + + A canvas is a thumbnail plus the slices cut from it, tiled row-major. + """ + + thumb_h: int + thumb_w: int + """Thumbnail grid.""" + slice_h: int + slice_w: int + """Per-slice grid. Zero when the canvas has no slices.""" + cols: int + """Slices sharing a canvas row. Zero when the canvas has no slices.""" + height: int + width: int + """Assembled canvas.""" + + +def _canvas_layout( + thumbnail: tuple[int, int, int], + slices: list[tuple[int, int, int]], + target_sizes: torch.Tensor, +) -> _CanvasLayout: + """Measure the canvas for one ``(thumbnail, slices)`` group. + + ``thumbnail`` and each slice are ``(start, end, target_index)`` token spans. + """ + t_bs, t_be, t_gi = thumbnail + thumb_h, thumb_w = compute_llm_grid(t_be - t_bs, *_target_hw(target_sizes, t_gi)) + if not slices: + return _CanvasLayout( + thumb_h, thumb_w, 0, 0, 0, max(thumb_h, 1), max(thumb_w, 1) + ) + + cols = len(slices) + for k in range(len(slices) - 1): + if slices[k + 1][0] - slices[k][1] > 2: + cols = k + 1 + break + rows = len(slices) // cols if cols > 0 else 1 + if rows * cols != len(slices): + rows, cols = 1, len(slices) + + s0_bs, s0_be, s0_gi = slices[0] + slice_h, slice_w = compute_llm_grid(s0_be - s0_bs, *_target_hw(target_sizes, s0_gi)) + return _CanvasLayout( + thumb_h, thumb_w, slice_h, slice_w, cols, rows * slice_h, cols * slice_w + ) + + +def _assign_canvas_span( + pos3d: torch.Tensor, + *, + thumbnail: tuple[int, int, int], + span_end: int, + slices: list[tuple[int, int, int]], + base: int, + target_sizes: torch.Tensor, +) -> int: + """Place one canvas at ``base``, writing its 3-D positions into ``pos3d``. + + Shared by the video-frame and single-image paths; they differ only in how + the per-canvas base advances, not in the geometry placed here. + + Mutates ``pos3d`` in place, as the packed canvas computation has always + done. Returns the distance to the next canvas base, so each caller keeps + only its own base advancement. + """ + t_bs, t_be, _ = thumbnail + device = pos3d.device + span_start = t_bs - 1 # the token + n_thumb = t_be - t_bs + layout = _canvas_layout(thumbnail, slices, target_sizes) + cols, canvas_H, canvas_W = layout.cols, layout.height, layout.width + + # base coat + pos3d[0, 0, span_start:span_end] = base + pos3d[1, 0, span_start:span_end] = base + pos3d[2, 0, span_start:span_end] = base + + # -> halo + halo_lo = max(base - 1, 0) + pos3d[1, 0, span_start] = halo_lo + pos3d[2, 0, span_start] = halo_lo + + # -> halo (right after thumbnail) + if t_be < span_end: + pos3d[1, 0, t_be] = base + canvas_H + pos3d[2, 0, t_be] = base + canvas_W + + # thumbnail visual tokens + llm_th, llm_tw = layout.thumb_h, layout.thumb_w + if llm_th > 0 and llm_tw > 0 and llm_th * llm_tw == n_thumb: + if slices and canvas_H > 0 and canvas_W > 0: + h_c = torch.linspace(0, canvas_H - 1, llm_th, device=device).round().long() + w_c = torch.linspace(0, canvas_W - 1, llm_tw, device=device).round().long() + else: + h_c = torch.arange(llm_th, device=device) + w_c = torch.arange(llm_tw, device=device) + h_idx = h_c.view(-1, 1).expand(-1, llm_tw).reshape(-1) + w_idx = w_c.view(1, -1).expand(llm_th, -1).reshape(-1) + pos3d[0, 0, t_bs:t_be] = base + pos3d[1, 0, t_bs:t_be] = h_idx + base + pos3d[2, 0, t_bs:t_be] = w_idx + base + else: + fallback = torch.arange(n_thumb, device=device, dtype=torch.long) + base + for d in range(3): + pos3d[d, 0, t_bs:t_be] = fallback + + # slice visual tokens + slice specials + for k, (s_bs, s_be, s_gi) in enumerate(slices): + n_s = s_be - s_bs + sh_k, sw_k = compute_llm_grid(n_s, *_target_hw(target_sizes, s_gi)) + h_off = (k // cols) * layout.slice_h + w_off = (k % cols) * layout.slice_w + + slice_start_pos = s_bs - 1 + if slice_start_pos >= span_start: + pos3d[1, 0, slice_start_pos] = base + h_off + pos3d[2, 0, slice_start_pos] = base + w_off + + slice_end_pos = s_be + if sh_k > 0 and sw_k > 0: + end_h = h_off + sh_k - 1 + end_w = w_off + sw_k - 1 + else: + end_h = h_off + end_w = w_off + if slice_end_pos < span_end: + pos3d[1, 0, slice_end_pos] = base + end_h + pos3d[2, 0, slice_end_pos] = base + end_w + + if sh_k > 0 and sw_k > 0 and sh_k * sw_k == n_s: + h_idx = ( + torch.arange(sh_k, device=device) + .view(-1, 1) + .expand(-1, sw_k) + .reshape(-1) + ) + w_idx = ( + torch.arange(sw_k, device=device) + .view(1, -1) + .expand(sh_k, -1) + .reshape(-1) + ) + pos3d[0, 0, s_bs:s_be] = base + pos3d[1, 0, s_bs:s_be] = h_idx + h_off + base + pos3d[2, 0, s_bs:s_be] = w_idx + w_off + base + else: + fallback = torch.arange(n_s, device=device, dtype=torch.long) + base + for d in range(3): + pos3d[d, 0, s_bs:s_be] = fallback + + # \n between slice rows -> W = right edge + 1. + # H stays within the ended row (does NOT cross to next row): the \n sits at + # the right-outside of the row that just ended, at H = that row's last row + # and W = one column past the rightmost slice. + for k in range(len(slices) - 1): + gap_s = slices[k][1] + gap_e = slices[k + 1][0] + if gap_e - gap_s <= 2: + continue + boundary_h = ((k // cols) + 1) * layout.slice_h - 1 + right_edge_w = cols * layout.slice_w + for nl_pos in range(gap_s + 1, gap_e - 1): + pos3d[1, 0, nl_pos] = base + boundary_h + pos3d[2, 0, nl_pos] = base + right_edge_w + + return max(canvas_H, canvas_W) + 1 + + +def _compute_canvas_packed( + position_ids_2d, + cu_seqlens, + image_bound, + target_sizes, + input_ids, + special_token_ids, +): + if position_ids_2d.ndim == 1: + position_ids_2d = position_ids_2d.unsqueeze(0) + if input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + + has_images = ( + isinstance(image_bound, torch.Tensor) + and image_bound.numel() > 0 + and isinstance(target_sizes, torch.Tensor) + and target_sizes.numel() > 0 + ) + if not has_images: + return position_ids_2d.unsqueeze(0).expand(3, -1, -1).contiguous() + + device = position_ids_2d.device + pos3d = position_ids_2d.unsqueeze(0).expand(3, -1, -1).clone() + ids_flat = input_ids.view(-1) + + im_start_id = special_token_ids.get("im_start_id") + im_end_id = special_token_ids.get("im_end_id") + slice_start_id = special_token_ids.get("slice_start_id") + slice_end_id = special_token_ids.get("slice_end_id") + + structural_ids = { + v + for v in (im_start_id, im_end_id, slice_start_id, slice_end_id) + if v is not None + } + + num_seqs = len(cu_seqlens) - 1 + img_ptr = 0 + num_imgs = len(image_bound) + + newline_id = special_token_ids.get("newline_id") + gap_ok_ids = structural_ids | ({newline_id} if newline_id is not None else set()) + + def _group_end_pos(g, s_end): + """First position after a group's trailing structural tokens. + + Matches the transformers reference, which closes a span over every + marker in ``markers`` -- newlines included. The chat template joins + content parts with ``\\n``, so an image followed by text always ends with + one; leaving it out moves that newline out of the canvas and pushes + every later position one step forward. + """ + le = g["slices"][-1][1] if g["slices"] else g["thumbnail"][1] + ge = le + while ge < s_end and ids_flat[ge].item() in gap_ok_ids: + ge += 1 + return min(ge, s_end) + + def _is_tight_gap(g_prev, g_next, s_end): + """Check if only structural/nl tokens sit between two groups.""" + ge = _group_end_pos(g_prev, s_end) + ns = g_next["thumbnail"][0] - 1 # im_start of next group + return all(ids_flat[p].item() in gap_ok_ids for p in range(ge, ns)) + + for si in range(num_seqs): + s_start = cu_seqlens[si].item() + s_end = cu_seqlens[si + 1].item() + if s_start >= s_end: + continue + + seq_imgs = [] + while img_ptr < num_imgs: + bs_ = image_bound[img_ptr, 0].item() + be_ = image_bound[img_ptr, 1].item() + if bs_ >= s_start and be_ <= s_end: + seq_imgs.append((bs_, be_, img_ptr)) + img_ptr += 1 + else: + break + if not seq_imgs: + continue + + # ---- group images: thumbnail + following slices ---- + # merge consecutive no-slice frames into one video group + raw_groups = [] + cur_group: dict[str, Any] | None = None + for bs_, be_, gi in seq_imgs: + marker_pos = bs_ - 1 + is_slice = ( + slice_start_id is not None + and marker_pos >= s_start + and ids_flat[marker_pos].item() == slice_start_id + ) + if is_slice and cur_group is not None: + cur_group["slices"].append((bs_, be_, gi)) + else: + if cur_group is not None: + raw_groups.append(cur_group) + cur_group = {"thumbnail": (bs_, be_, gi), "slices": []} + if cur_group is not None: + raw_groups.append(cur_group) + + # merge consecutive tightly-adjacent groups into video groups. + # Two groups are "tightly adjacent" if the gap between them + # contains only structural tokens and \n (no real text). + # This handles both no-slice frames and frames-with-slices. + groups = [] + i_g = 0 + while i_g < len(raw_groups): + video_frames = [raw_groups[i_g]] + j_g = i_g + 1 + while j_g < len(raw_groups) and _is_tight_gap( + raw_groups[j_g - 1], raw_groups[j_g], s_end + ): + video_frames.append(raw_groups[j_g]) + j_g += 1 + if len(video_frames) == 1: + groups.append(raw_groups[i_g]) + else: + groups.append({"video_frames": video_frames}) + i_g = j_g + + # ---- assign 3-D positions ---- + pos = 0 + cursor = s_start + + for group in groups: + is_video_group = "video_frames" in group + + if is_video_group: + # === Video group: multiple frames merged === + # Each frame gets its own base_f so different frames have + # distinct spatial positions. Inter-frame gap tokens (\n + # etc.) are pure 1D text. After each gap, pos advances + # by 1 extra so halo_lo of the next frame does not + # collide with the last gap token. + frames = group["video_frames"] + first_frame = frames[0] + last_frame = frames[-1] + + # group_start = of first frame + group_start = first_frame["thumbnail"][0] - 1 + # group_end = after last structural token of last frame + group_end = _group_end_pos(last_frame, s_end) + + # text before video group + text_len = group_start - cursor + if text_len > 0: + t_pos = ( + torch.arange(text_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, cursor:group_start] = t_pos + pos += text_len + + # process each frame with its own per-frame base + frame_cursor = group_start + for vf in frames: + t_bs = vf["thumbnail"][0] + vf_slices = vf["slices"] + im_start_pos = t_bs - 1 + + # --- inter-frame gap: pure 1D text T=H=W=pos --- + gap_len = im_start_pos - frame_cursor + if gap_len > 0: + t_pos = ( + torch.arange(gap_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, frame_cursor:im_start_pos] = t_pos + pos += gap_len + pos += 1 # buffer: prevents halo_lo collision + + # --- per-frame base --- + base_f = pos + frame_end = _group_end_pos(vf, group_end) + pos = base_f + _assign_canvas_span( + pos3d, + thumbnail=vf["thumbnail"], + span_end=frame_end, + slices=vf_slices, + base=base_f, + target_sizes=target_sizes, + ) + frame_cursor = frame_end + + # trailing tokens after last frame (if any) + if frame_cursor < group_end: + trail_len = group_end - frame_cursor + t_pos = ( + torch.arange(trail_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, frame_cursor:group_end] = t_pos + pos += trail_len + + cursor = group_end + + else: + # === Single image group (thumbnail + optional slices) === + slices = group["slices"] + + group_start = group["thumbnail"][0] - 1 + group_end = _group_end_pos(group, s_end) + + # text before group + text_len = group_start - cursor + if text_len > 0: + t_pos = ( + torch.arange(text_len, device=device, dtype=torch.long) + pos + ) + for d in range(3): + pos3d[d, 0, cursor:group_start] = t_pos + pos += text_len + + base = pos + pos = base + _assign_canvas_span( + pos3d, + thumbnail=group["thumbnail"], + span_end=group_end, + slices=slices, + base=base, + target_sizes=target_sizes, + ) + cursor = group_end + + # remaining text + rem = s_end - cursor + if rem > 0: + t_pos = torch.arange(rem, device=device, dtype=torch.long) + pos + for d in range(3): + pos3d[d, 0, cursor:s_end] = t_pos + + return pos3d + + +def canvas_rope_delta(position_ids: torch.Tensor, seq_len: int) -> int: + """Offset applied to the 1-D positions of tokens decoded after the prompt. + + The first decoded token sits one past the last prefill position. + """ + return int(position_ids.amax().item()) + 1 - seq_len diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 875179a88fae..fa2924441314 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -514,6 +514,10 @@ "minicpmv4_6", "MiniCPMV4_6ForConditionalGeneration", ), + "MiniCPMV4_7ForConditionalGeneration": ( + "minicpmv4_7", + "MiniCPMV4_7ForConditionalGeneration", + ), "Mistral3ForConditionalGeneration": ( "mistral3", "Mistral3ForConditionalGeneration", diff --git a/vllm/parser/minicpmv.py b/vllm/parser/minicpmv.py new file mode 100644 index 000000000000..88a6865c40c9 --- /dev/null +++ b/vllm/parser/minicpmv.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniCPM-V family parser. + +Wraps the Qwen3 grammar (```` reasoning + XML tool calls) with two +MiniCPM-V specifics: + +- ``<|im_end|>`` and the reserved marker tokens are passthrough terminals, + so they never leak into ``content``. +- The tokenizer decodes a newline as the two characters ``\\`` and ``n`` + (e.g. token id 1639) instead of ``\\n``, so the decoded text is recovered + before it reaches the client. Code blocks, inline code and LaTeX are left + alone, as are tool call arguments. +""" + +from __future__ import annotations + +import functools +from collections.abc import Sequence +from dataclasses import replace +from typing import TYPE_CHECKING + +from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) +from vllm.parser.qwen3 import Qwen3Parser, qwen3_config + +if TYPE_CHECKING: + from vllm.entrypoints.generate.base.protocol import FunctionCall + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + + +_RESERVED_MARKER_IDS = range(12, 16) +_IM_END = "<|im_end|>" + + +@functools.cache +def minicpmv_config(thinking: bool) -> ParserEngineConfig: + """MiniCPM-V grammar: the qwen3 grammar plus MiniCPM-V markers. + + Both modes keep the qwen3 tool call terminals and transitions, so + ```` output is parsed whether or not thinking is enabled. + """ + name = "minicpmv" if thinking else "minicpmv_no_thinking" + base = qwen3_config(thinking=thinking, name=name) + + transitions = dict(base.transitions) + if not thinking: + # Default mode: the model answers directly, so `` must not open + # a reasoning span that would swallow the answer. A stray + # `...` in the output is still dropped, not printed. + transitions[(ParserState.CONTENT, "THINK_START")] = Transition( + ParserState.REASONING + ) + base = replace(base, strip_trailing_reasoning_whitespace=False) + + terminals = dict(base.terminals) + token_id_terminals = dict(base.token_id_terminals) + terminals["IM_END"] = _IM_END + token_id_terminals["IM_END"] = _IM_END + for state in (ParserState.CONTENT, ParserState.REASONING): + transitions[(state, "IM_END")] = Transition(state) + + for index in _RESERVED_MARKER_IDS: + text_name = f"RESERVED_TEXT_{index}" + token_name = f"RESERVED_TOKEN_{index}" + terminals[text_name] = f"" + terminals[token_name] = f"<|reserved_{index}|>" + token_id_terminals[token_name] = f"<|reserved_{index}|>" + for state in (ParserState.CONTENT, ParserState.REASONING): + transitions[(state, text_name)] = Transition(state) + transitions[(state, token_name)] = Transition(state) + + return replace( + base, + terminals=terminals, + token_id_terminals=token_id_terminals, + transitions=transitions, + ) + + +class EscapedNewlineNormalizer: + """Recover literal escape sequences written by the model as real text. + + Scans character by character so that a backslash inside a code block, + inline code, LaTeX span or a doubled escape is left untouched. A trailing + half of an escape is held in ``_pending`` until the next chunk arrives, + which keeps streaming output identical to the non-streaming result. + """ + + def __init__(self) -> None: + self._closing_marker: str | None = None + self._pending = "" + + def feed(self, text: str, *, final: bool = False) -> str: + text = self._pending + text + self._pending = "" + output: list[str] = [] + index = 0 + + while index < len(text): + if self._closing_marker is not None: + marker = self._closing_marker + remaining = text[index:] + if remaining.startswith(marker): + output.append(marker) + index += len(marker) + self._closing_marker = None + continue + if not final and marker.startswith(remaining): + self._pending = remaining + break + output.append(text[index]) + index += 1 + continue + + remaining = text[index:] + if text[index] == "`": + if not final and remaining in ("`", "``"): + self._pending = remaining + break + if remaining.startswith("```"): + output.append("```") + index += 3 + self._closing_marker = "```" + else: + output.append("`") + index += 1 + self._closing_marker = "`" + continue + + if text[index] == "$": + if not final and remaining == "$": + self._pending = remaining + break + if remaining.startswith("$$"): + output.append("$$") + index += 2 + self._closing_marker = "$$" + else: + output.append("$") + index += 1 + self._closing_marker = "$" + continue + + if text[index] != "\\": + output.append(text[index]) + index += 1 + continue + + slash_end = index + 1 + while slash_end < len(text) and text[slash_end] == "\\": + slash_end += 1 + if slash_end - index > 1: + if not final and slash_end == len(text): + # The run may still grow, and its parity decides whether a + # following `n` is a newline or stays literal. Hold it until + # the run is complete. + self._pending = text[index:slash_end] + break + output.append(text[index:slash_end]) + index = slash_end + continue + + if not final and remaining in ("\\", "\\r", "\\r\\"): + self._pending = remaining + break + if remaining.startswith("\\r\\n"): + output.append("\n") + index += 4 + elif remaining.startswith(("\\n", "\\r")): + output.append("\n") + index += 2 + elif remaining.startswith("\\("): + output.append("\\(") + index += 2 + self._closing_marker = "\\)" + elif remaining.startswith("\\["): + output.append("\\[") + index += 2 + self._closing_marker = "\\]" + else: + output.append("\\") + index += 1 + + if final and self._pending: + output.append(self._pending) + self._pending = "" + return "".join(output) + + +def recover_newlines(text: str | None) -> str | None: + if text is None or "\\" not in text: + return text + return EscapedNewlineNormalizer().feed(text, final=True) + + +class MiniCPMVOutputNormalizer: + """Apply :class:`EscapedNewlineNormalizer` to a whole stream. + + Reasoning and content are tracked separately so that a half-written escape + at the end of a reasoning chunk does not swallow the first characters of + the content that follows. + """ + + def __init__(self) -> None: + self._reasoning = EscapedNewlineNormalizer() + self._content = EscapedNewlineNormalizer() + + def normalize_delta( + self, + delta: DeltaMessage | None, + *, + finished: bool, + ) -> DeltaMessage | None: + reasoning_finished = finished or ( + delta is not None and (delta.content is not None or bool(delta.tool_calls)) + ) + reasoning = self._reasoning.feed( + delta.reasoning if delta and delta.reasoning else "", + final=reasoning_finished, + ) + content = self._content.feed( + delta.content if delta and delta.content else "", + final=finished, + ) + + if delta is None: + if not reasoning and not content: + return None + return DeltaMessage( + reasoning=reasoning or None, + content=content or None, + ) + + if reasoning: + delta.reasoning = reasoning + elif delta.reasoning is not None: + delta.reasoning = None + if content: + delta.content = content + elif delta.content is not None: + delta.content = None + + # A chunk that normalizes to nothing must be dropped, not returned as an + # all-None DeltaMessage: the engine only skips on `None`, so anything + # else reaches the client as an empty delta. + if not delta.content and not delta.reasoning and not delta.tool_calls: + return None + return delta + + +class MiniCPMVParser(Qwen3Parser): + """Parse MiniCPM-V family output in thinking and non-thinking modes.""" + + CONFIG_NAME = "minicpmv" + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} + # The chat template emits `` unless `enable_thinking` is + # explicitly false, so the grammar must default to thinking too, or the + # reasoning span is misread as content. + thinking_enabled = chat_kwargs.get("enable_thinking", True) + kwargs.setdefault("parser_engine_config", minicpmv_config(thinking_enabled)) + super().__init__(tokenizer, tools, **kwargs) + # ``Qwen3Parser.__init__`` reads the same key with a default of True, + # which does not apply here: MiniCPM-V defaults to thinking disabled. + # Keep the attribute in sync with the grammar built above. + self.thinking_enabled = thinking_enabled + + @property + def reasoning_ended(self) -> bool: + # The engine derives this from the presence of THINK terminals, which + # both MiniCPM-V configs keep, so the base property stays False in + # non-thinking mode. There is no reasoning span to end there. + if not self.thinking_enabled: + return True + return super().reasoning_ended + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + if self.thinking_enabled: + return super().extract_reasoning(model_output, request) + _, content = ParserEngine.extract_reasoning(self, model_output, request) + return None, content + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + if self.thinking_enabled: + return super().extract_reasoning_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + + delta = ParserEngine.extract_reasoning_streaming( + self, + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + if delta is None: + return None + delta.reasoning = None + if delta.content is None and not delta.tool_calls: + return None + return delta + + def get_streaming_fallback_content( + self, + text: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> str | None: + if self.thinking_enabled: + return super().get_streaming_fallback_content(text, request) + delta = ParserEngine.finish_streaming(self) + return delta.content if delta is not None else None + + +class MiniCPMVUnifiedParser(DelegatingParser): + """MiniCPM-V serving parser: the adapters plus newline recovery. + + Recovery cannot live in :class:`MiniCPMVParser`: the serving stack drives + the reasoning/tool adapters, never ``ParserEngine.parse`` / + ``ParserEngine.parse_delta``, so a rewrite applied there would only be + visible to callers that construct the engine directly. + """ + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + **kwargs, + ) -> None: + super().__init__(tokenizer, tools, *args, **kwargs) + self._output_normalizer = MiniCPMVOutputNormalizer() + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + reasoning, content, tool_calls = super().parse( + model_output, + request, + enable_auto_tools, + model_output_token_ids, + ) + # Tool call arguments are JSON, where a literal `\n` is meaningful and + # must survive as written. + return ( + recover_newlines(reasoning), + recover_newlines(content), + tool_calls, + ) + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + delta = super().parse_delta( + delta_text, + delta_token_ids, + request, + prompt_token_ids, + finished=finished, + ) + return self._output_normalizer.normalize_delta(delta, finished=finished) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index cf0be908ed63..6766d15e914c 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -242,6 +242,19 @@ class _CohereCommandParser(CohereCommandParser): return _CohereCommandParser + if "minicpmv" in {reasoning_parser_name, tool_parser_name}: + from vllm.parser.minicpmv import MiniCPMVUnifiedParser + + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _MiniCPMVParser(MiniCPMVUnifiedParser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + tool_strict_level = strict_level + + return _MiniCPMVParser + from vllm.parser.abstract_parser import DelegatingParser r_cls = reasoning_parser_cls diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index f9ab1dc9a11f..9e75e9b77ef4 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -116,6 +116,10 @@ "mimo_engine_reasoning_parser", "MiMoParserReasoningAdapter", ), + "minicpmv": ( + "minicpmv_reasoning_parser", + "MiniCPMVParserReasoningAdapter", + ), "minimax_m2": ( "minimax_m2_reasoning_parser", "MiniMaxM2ReasoningParser", diff --git a/vllm/reasoning/minicpmv_reasoning_parser.py b/vllm/reasoning/minicpmv_reasoning_parser.py new file mode 100644 index 000000000000..402c0524cab0 --- /dev/null +++ b/vllm/reasoning/minicpmv_reasoning_parser.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter +from vllm.parser.minicpmv import MiniCPMVParser + + +class MiniCPMVParserReasoningAdapter(ParserEngineReasoningAdapter): + # Reasoning-only on purpose: this family registers no tool parser, because + # MiniCPM-V emits the Qwen3-Coder XML syntax verbatim and `qwen3_coder` + # already covers it. ``make_adapters`` only builds reasoning/tool pairs, so + # the two halves are declared explicitly here. + _parser_engine_cls = MiniCPMVParser + + +__all__ = ["MiniCPMVParserReasoningAdapter"] diff --git a/vllm/transformers_utils/chat_templates/registry.py b/vllm/transformers_utils/chat_templates/registry.py index 422854ab3d4d..f902ecbbaafb 100644 --- a/vllm/transformers_utils/chat_templates/registry.py +++ b/vllm/transformers_utils/chat_templates/registry.py @@ -32,6 +32,7 @@ def _get_minicpmv_chat_template_fallback(tokenizer_name_or_path: str) -> Path | "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja", "minicpmv": _get_minicpmv_chat_template_fallback, "minicpmv4_6": _get_minicpmv_chat_template_fallback, + "minicpmv4_7": _get_minicpmv_chat_template_fallback, "paligemma": CHAT_TEMPLATES_DIR / "template_basic.jinja", "siglip": CHAT_TEMPLATES_DIR / "template_basic.jinja", "siglip2": CHAT_TEMPLATES_DIR / "template_basic.jinja", From bd5bbce6e45e09dbd07ac3e8c394626c059078f7 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 21:15:14 +0800 Subject: [PATCH 370/952] [RL] Release weight-transfer communicators without stopping inference (#61060) Signed-off-by: aoshen02 Signed-off-by: fh3 --- rust/src/engine-core-client/src/client.rs | 12 ++++++++ rust/src/server/src/routes.rs | 4 +++ rust/src/server/src/routes/tests.rs | 10 +++++-- rust/src/server/src/routes/weight_transfer.rs | 16 +++++++++++ tests/distributed/test_weight_transfer.py | 28 +++++++++++++++++++ tests/distributed/test_weight_transfer_m2n.py | 26 ++++++++++++++++- .../test_route_request_bodies.py | 5 ++++ .../rl/weight_transfer/test_offline.py | 7 +++++ .../test_weight_operation_metrics.py | 1 + .../worker/test_gpu_worker_weight_transfer.py | 23 +++++++++++++++ vllm/distributed/weight_transfer/base.py | 4 +-- vllm/distributed/weight_transfer/clients.py | 10 ++++++- .../distributed/weight_transfer/m2n_engine.py | 4 ++- .../weight_transfer/m2n_trainer.py | 4 ++- .../weight_transfer/nccl_engine.py | 6 ++-- .../weight_transfer/sparse_nccl_engine.py | 5 +++- vllm/engine/protocol.py | 4 +++ vllm/entrypoints/rl/offline.py | 4 +++ vllm/entrypoints/rl/online/api_router.py | 7 +++++ vllm/entrypoints/rl/online/metrics.py | 2 +- vllm/v1/engine/async_llm.py | 4 +++ vllm/v1/worker/gpu_worker.py | 9 ++++++ 22 files changed, 181 insertions(+), 14 deletions(-) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index d8b436c6466c..49910d80c740 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -828,6 +828,18 @@ impl EngineCoreClient { Ok(()) } + /// Release the trainer connection without shutting down inference. + pub async fn shutdown_weight_transfer_engine(&self) -> Result<()> { + self.collective_rpc( + "shutdown_weight_transfer_engine", + None, + Vec::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + /// Start a weight update for the base model. pub async fn start_weight_update(&self) -> Result<()> { self.collective_rpc( diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 251398611ba1..23c4244c821a 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -145,6 +145,10 @@ fn build_router_with_options( "/init_weight_transfer_engine", post(weight_transfer::init_weight_transfer_engine), ) + .route( + "/shutdown_weight_transfer_engine", + post(weight_transfer::shutdown_weight_transfer_engine), + ) .route( "/start_weight_update", post(weight_transfer::start_weight_update), diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d574e02a9215..9b8378f89934 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -6708,6 +6708,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { ("POST", "/reset_mm_cache"), ("POST", "/reset_encoder_cache"), ("POST", "/init_weight_transfer_engine"), + ("POST", "/shutdown_weight_transfer_engine"), ("POST", "/start_weight_update"), ("POST", "/start_draft_weight_update"), ("POST", "/update_weights"), @@ -6739,7 +6740,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { boxed_test_future(async move { let mut calls = Vec::new(); - for _ in 0..10 { + for _ in 0..11 { let utility = recv_engine_message(dealer).await; assert_eq!(utility[0].as_ref(), &[0x03]); let payload = decode_value(&utility[1]).expect("decode utility"); @@ -6769,6 +6770,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "collective_rpc [\"finish_weight_update\",null,[],{}]", "set_weight_version [\"step-2\"]", "get_weight_version []", + "collective_rpc [\"shutdown_weight_transfer_engine\",null,[],{}]", ] "#]].assert_debug_eq(&calls); }) @@ -6817,6 +6819,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { Some(json!({"new_version": "step-2"})), ), ("GET", "/weight_info", None), + ("POST", "/shutdown_weight_transfer_engine", None), ] { let mut request = Request::builder().method(method).uri(path); let body = match body { @@ -6838,7 +6841,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { } engine_task.await.expect("mock engine task"); let metrics_after = METRICS.render().unwrap(); - for operation in ["update", "finish"] { + for (operation, count) in [("update", 2.0), ("finish", 2.0), ("shutdown", 1.0)] { assert_eq!( metric_delta( &metrics_before, @@ -6846,7 +6849,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "vllm:rl_weight_update_operation_duration_seconds_count", Some(&format!("operation=\"{operation}\"")), ), - 2.0 + count ); assert_eq!( metric_value( @@ -6868,6 +6871,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { "{\"message\":\"Weight update finished\"}", "{\"success\":true,\"new_version\":\"step-2\"}", "{\"weight_version\":\"step-2\"}", + "{\"message\":\"Weight transfer shut down\"}", ] "#]] .assert_debug_eq(&responses); diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs index 5ccd66e39b68..7e1d6af61c7e 100644 --- a/rust/src/server/src/routes/weight_transfer.rs +++ b/rust/src/server/src/routes/weight_transfer.rs @@ -75,6 +75,22 @@ pub async fn init_weight_transfer_engine( })) } +/// Release the trainer connection without shutting down inference. +pub async fn shutdown_weight_transfer_engine( + State(state): State>, +) -> Result, ApiError> { + let _recorder = METRICS.api_server.record_weight_operation("shutdown"); + state + .engine_core_client() + .shutdown_weight_transfer_engine() + .await + .map_err(|error| utility_call_error("shutdown_weight_transfer_engine", error))?; + + Ok(Json(MessageResponse { + message: "Weight transfer shut down", + })) +} + /// Start a weight update transaction. pub async fn start_weight_update( State(state): State>, diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index d044ef3bb7ac..642193ce882f 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -218,6 +218,27 @@ def test_rejects_negative_num_updates(self): # --- Unit Tests: Engine Parsing --- +@pytest.mark.parametrize( + "engine_type", + [ + NCCLWeightTransferEngine, + NCCLTrainerWeightTransferEngine, + SparseNCCLWeightTransferEngine, + SparseNCCLTrainerWeightTransferEngine, + ], +) +def test_nccl_shutdown_destroys_communicator_once(engine_type): + engine = object.__new__(engine_type) + group = MagicMock() + engine.model_update_group = group + + engine.shutdown() + engine.shutdown() + + group.destroy.assert_called_once_with() + assert engine.model_update_group is None + + class TestNCCLEngineParsing: """Test NCCLWeightTransferEngine parsing methods.""" @@ -1292,6 +1313,9 @@ def test_ray_client_sends_typed_requests(self, monkeypatch): handle.finish_weight_update.remote.assert_called_once_with() handle.update_weight_version.remote.assert_called_once_with("step-42") + client.shutdown_weight_transfer_engine() + handle.shutdown_weight_transfer_engine.remote.assert_called_once_with() + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" captured = {} @@ -1322,6 +1346,7 @@ def test_http_client_passes_through_nccl_update_info(self, monkeypatch): captured = {} def fake_post(self, path, json=None): + captured["path"] = path captured["json"] = json monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post) @@ -1333,6 +1358,9 @@ def fake_post(self, path, json=None): client.finish_weight_update("step-42") assert captured["json"] == {"weight_version": "step-42"} + client.shutdown_weight_transfer_engine() + assert captured == {"path": "shutdown_weight_transfer_engine", "json": None} + class TestModuleSource: """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py index 0591adff4d8d..588ffeeb264b 100644 --- a/tests/distributed/test_weight_transfer_m2n.py +++ b/tests/distributed/test_weight_transfer_m2n.py @@ -46,7 +46,10 @@ mesh_from_tensor, placements_from_tensor, ) -from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo +from vllm.distributed.weight_transfer.m2n_trainer import ( + M2NTrainerInitInfo, + M2NTrainerWeightTransferEngine, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -57,6 +60,27 @@ VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() +@pytest.mark.parametrize( + "engine_type,group_attribute", + [ + (M2NWeightTransferEngine, "model_update_group"), + (M2NTrainerWeightTransferEngine, "group"), + ], +) +def test_shutdown_destroys_communicator_once(engine_type, group_attribute): + engine = object.__new__(engine_type) + engine._handle = None + engine._executor = None + group = Mock() + setattr(engine, group_attribute, group) + + engine.shutdown() + engine.shutdown() + + group.destroy.assert_called_once_with() + assert getattr(engine, group_attribute) is None + + class TestLayout: def test_replicated_splits_nothing(self): """A replicated tensor imposes no divisibility constraint. The shape diff --git a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index b51d08db0758..9e4df28d03b4 100644 --- a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py +++ b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py @@ -89,6 +89,11 @@ def test_invalid_field_is_rejected_before_the_engine(client, path, body): ("/update_weights", {"update_info": {"names": []}}, "update_weights"), ("/update_weights", {"update_info": [{}, {}]}, "update_weights"), ("/collective_rpc", {"method": "m"}, "collective_rpc"), + ( + "/shutdown_weight_transfer_engine", + {}, + "shutdown_weight_transfer_engine", + ), ], ) def test_valid_body_reaches_the_engine(client, path, body, engine_call): diff --git a/tests/entrypoints/rl/weight_transfer/test_offline.py b/tests/entrypoints/rl/weight_transfer/test_offline.py index e4a30f17e61d..0af71a1dee24 100644 --- a/tests/entrypoints/rl/weight_transfer/test_offline.py +++ b/tests/entrypoints/rl/weight_transfer/test_offline.py @@ -169,6 +169,13 @@ def check_init_called(self): assert called, "init_transfer_engine should have been called" assert param == "hello", f"Expected 'hello', got {param}" + llm.shutdown_weight_transfer_engine() + assert all( + llm.collective_rpc( + lambda worker: worker.weight_transfer_engine.shutdown_called + ) + ) + @create_new_process_for_each_test() def test_update_weights_calls_engine(vllm_runner): diff --git a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index 22ec5e611252..a3bb67ed2b60 100644 --- a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py +++ b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py @@ -75,6 +75,7 @@ def test_record_releases_in_flight_and_observes_duration(registry, fail): "operation,route,body", [ ("init", api_router.init_weight_transfer_engine, {"init_info": {}}), + ("shutdown", api_router.shutdown_weight_transfer_engine, None), ("start", api_router.start_weight_update, None), ("start_draft", api_router.start_draft_weight_update, None), ("update", api_router.update_weights, {"update_info": {}}), diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index f1965cb1fae2..ec63c08c70c4 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -24,6 +24,7 @@ def __init__(self, raise_on_update: bool = False): self.raise_on_update = raise_on_update self.started = False self.finished = False + self.closed = False self.reset_count = 0 self.supports_draft_weight_update = False self.update_calls: list[dict] = [] @@ -49,6 +50,9 @@ def finish_weight_update(self) -> None: def reset_weight_update_target(self) -> None: self.reset_count += 1 + def shutdown(self) -> None: + self.closed = True + class _RecordingModelRunner: def __init__(self) -> None: @@ -82,6 +86,25 @@ def test_reload_weights_sets_current_config(): assert model_runner.seen_config is worker.vllm_config +@pytest.mark.parametrize( + "enabled,active", [(True, False), (True, True), (False, False)] +) +def test_shutdown_weight_transfer_engine_preserves_inference(enabled, active): + engine = _RecordingEngine() if enabled else None + worker = _make_worker(engine) + worker._weight_update_active = active + model_runner = worker.model_runner + if active: + with pytest.raises(RuntimeError, match="during a weight update"): + worker.shutdown_weight_transfer_engine() + else: + worker.shutdown_weight_transfer_engine() + if engine is not None: + assert engine.closed is not active + assert worker.weight_transfer_engine is engine + assert worker.model_runner is model_runner + + def test_reload_parameter_lookup_preserves_lora_module_names(): base_layer = nn.Module() qweight = nn.Parameter(torch.ones(1)) diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 3140b70e9656..6eba3fa96f51 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -514,9 +514,7 @@ def receive_weights(self, update_info: TUpdateInfo) -> None: @abstractmethod def shutdown(self) -> None: - """Shutdown the weight transfer engine. - This should be called when the worker is shutting down. - """ + """Release transfer resources without shutting down the inference worker.""" raise NotImplementedError diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 745e67779e0a..0ccce60f973a 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -60,7 +60,7 @@ class HTTPVLLMWeightSyncClient: Mirrors `vllm/entrypoints/rl/online/api_router.py`: `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, - `/finish_weight_update`. + `/finish_weight_update`, `/shutdown_weight_transfer_engine`. """ def __init__(self, base_url: str, timeout: float = 300) -> None: @@ -78,6 +78,9 @@ def _post(self, path: str, json: dict[str, Any] | None = None) -> None: def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: self._post("init_weight_transfer_engine", {"init_info": init_info}) + def shutdown_weight_transfer_engine(self) -> None: + self._post("shutdown_weight_transfer_engine") + def start_weight_update(self) -> None: self._post("start_weight_update") @@ -109,6 +112,11 @@ def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: request = WeightTransferInitRequest(init_info=init_info) ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles]) + def shutdown_weight_transfer_engine(self) -> None: + import ray + + ray.get([h.shutdown_weight_transfer_engine.remote() for h in self.handles]) + def start_weight_update(self) -> None: import ray diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py index c61bac4f2d1a..e480e60eaf91 100644 --- a/vllm/distributed/weight_transfer/m2n_engine.py +++ b/vllm/distributed/weight_transfer/m2n_engine.py @@ -449,5 +449,7 @@ def shutdown(self) -> None: torch.accelerator.synchronize() self._handle.destroy() self._handle = None - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None self._parameter_destinations = [] diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 2d1d9f47ec02..8b9b73f9f9cc 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -410,4 +410,6 @@ def shutdown(self) -> None: if self._executor is not None: self._executor.shutdown() self._executor = None - self.group = None + if self.group is not None: + self.group.destroy() + self.group = None diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 77d8b608727e..331cb5eacc01 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -229,7 +229,7 @@ def state_dict_info_iterator(): def shutdown(self) -> None: if self.model_update_group is not None: - # Clean up the communicator by removing the reference + self.model_update_group.destroy() self.model_update_group = None @@ -456,4 +456,6 @@ def _post_send_sync(self) -> None: torch.cuda.current_stream().synchronize() def shutdown(self) -> None: - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 01b29529deb7..f9233516a56e 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/vllm/distributed/weight_transfer/sparse_nccl_engine.py @@ -186,6 +186,7 @@ def receive_weights(self, update_info: SparseNCCLWeightTransferUpdateInfo) -> No def shutdown(self) -> None: if self.model_update_group is not None: + self.model_update_group.destroy() self.model_update_group = None @@ -371,4 +372,6 @@ def _post_send_sync(self) -> None: torch.cuda.current_stream().synchronize() def shutdown(self) -> None: - self.model_update_group = None + if self.model_update_group is not None: + self.model_update_group.destroy() + self.model_update_group = None diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 15f6b801bcd6..15f8f8b744f0 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -296,6 +296,10 @@ async def init_weight_transfer_engine( """Initialize weight transfer for RL training.""" raise NotImplementedError + async def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + raise NotImplementedError + async def start_weight_update(self) -> None: """Start a new weight update.""" raise NotImplementedError diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py index 1db6a82533dd..be72a9a5540c 100644 --- a/vllm/entrypoints/rl/offline.py +++ b/vllm/entrypoints/rl/offline.py @@ -26,6 +26,10 @@ def init_weight_transfer_engine( "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} ) + def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + self.llm_engine.collective_rpc("shutdown_weight_transfer_engine") + def start_weight_update(self) -> None: """Start a new weight update.""" self.llm_engine.collective_rpc("start_weight_update") diff --git a/vllm/entrypoints/rl/online/api_router.py b/vllm/entrypoints/rl/online/api_router.py index 51182fd99683..847287336ee5 100644 --- a/vllm/entrypoints/rl/online/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -168,6 +168,13 @@ async def init_weight_transfer_engine( return JSONResponse(content={"message": "Weight transfer initialized"}) +@router.post("/shutdown_weight_transfer_engine") +async def shutdown_weight_transfer_engine(raw_request: Request): + with weight_operation_metrics().record("shutdown"): + await engine_client(raw_request).shutdown_weight_transfer_engine() + return JSONResponse(content={"message": "Weight transfer shut down"}) + + @router.post("/start_weight_update") async def start_weight_update(raw_request: Request): with weight_operation_metrics().record("start"): diff --git a/vllm/entrypoints/rl/online/metrics.py b/vllm/entrypoints/rl/online/metrics.py index ed16f9a820b5..b67f479f6dc0 100644 --- a/vllm/entrypoints/rl/online/metrics.py +++ b/vllm/entrypoints/rl/online/metrics.py @@ -10,7 +10,7 @@ from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram -Operation = Literal["init", "start", "start_draft", "update", "finish"] +Operation = Literal["init", "shutdown", "start", "start_draft", "update", "finish"] class WeightOperationMetrics: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 71e30990c578..60588e076097 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1318,6 +1318,10 @@ async def init_weight_transfer_engine( "init_weight_transfer_engine", kwargs={"init_info": request.init_info} ) + async def shutdown_weight_transfer_engine(self) -> None: + """Release the trainer connection without shutting down inference.""" + await self.collective_rpc("shutdown_weight_transfer_engine") + async def start_weight_update(self) -> None: """Start a new weight update.""" await self.collective_rpc("start_weight_update") diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 152def62221c..3efe77a3a4d2 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1482,6 +1482,15 @@ def init_weight_transfer_engine(self, init_info: dict) -> None: typed_init_info = self.weight_transfer_engine.parse_init_info(init_info) self.weight_transfer_engine.init_transfer_engine(typed_init_info) + def shutdown_weight_transfer_engine(self) -> None: + """Release a trainer connection without shutting down inference.""" + if self._weight_update_active: + raise RuntimeError( + "Cannot release a trainer connection during a weight update" + ) + if self.weight_transfer_engine is not None: + self.weight_transfer_engine.shutdown() + def start_weight_update(self) -> None: """Start a new weight update session. From 3e8e1364ff41330fb667b2d76d6645176d61fdd0 Mon Sep 17 00:00:00 2001 From: Zheng Cai Date: Sun, 11 Oct 2026 22:43:37 +0800 Subject: [PATCH 371/952] [Perf][GLM-5.3-Flash] Shard long-context indexer prefill rows across TP (#54951) Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Signed-off-by: Zheng Cai Signed-off-by: Codex Co-authored-by: root Co-authored-by: OpenAI Codex Co-authored-by: zigzagcai Co-authored-by: Claude Code Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .../v1/attention/test_indexer_tp_row_shard.py | 213 ++++++++++++++++++ .../layers/sparse_attn_indexer.py | 124 ++++++++-- vllm/models/glm5next/nvidia/sparse_indexer.py | 63 ++++-- vllm/v1/attention/backends/mla/indexer.py | 74 +++++- 4 files changed, 436 insertions(+), 38 deletions(-) create mode 100644 tests/v1/attention/test_indexer_tp_row_shard.py diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py new file mode 100644 index 000000000000..e65e4789f132 --- /dev/null +++ b/tests/v1/attention/test_indexer_tp_row_shard.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP query-row sharding of the DSA indexer prefill (see PrefillRowShard): +each rank scores a contiguous slice; the group exchanges int32s per row.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +import vllm.v1.attention.backends.mla.indexer as indexer +from vllm.config import CUDAGraphMode +from vllm.model_executor.layers.sparse_attn_indexer import PrefillRowShard + + +def _chunk(start: int, end: int) -> SimpleNamespace: + """Chunk over rows [start, end); bounds derive from the global row.""" + n = end - start + return SimpleNamespace( + token_start=start, + token_end=end, + cu_seqlen_ks=torch.zeros(n, dtype=torch.int32), + cu_seqlen_ke=((torch.arange(n) * 7 + start) % 65).int(), + ) + + +@pytest.mark.parametrize("world", [2, 3, 8]) +def test_shard_narrowing_covers_every_row_once(world: int) -> None: + """Narrowed windows tile the prefill rows exactly once, carry the causal + bounds of those global rows, and skip chunks foreign to this rank.""" + decode, rows = 5, 997 + edges = [decode, decode + 330, decode + 337, decode + rows] + chunks = [_chunk(lo, hi) for lo, hi in zip(edges, edges[1:])] + sizes = [rows // world + int(r < rows % world) for r in range(world)] + covered = [] + for rank in range(world): + shard = PrefillRowShard(decode, sizes, rank) + for chunk in chunks: + window = shard.narrow(chunk) + if window is None: + continue + start, end, ks, ke = window + lo, hi = start - chunk.token_start, end - chunk.token_start + torch.testing.assert_close(ks, chunk.cu_seqlen_ks[lo:hi]) + torch.testing.assert_close(ke, chunk.cu_seqlen_ke[lo:hi]) + covered.append((start, end)) + covered.sort() + # Windows must be disjoint, contiguous and exactly cover the prefill rows. + assert [s for s, _ in covered] == [decode, *[e for _, e in covered][:-1]] + assert covered[-1][1] == decode + rows + + +def test_shard_exchange_reassembles_rows_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Gathered indices land only on real prefill rows: leading decode rows + and trailing CUDA-graph padding stay untouched, and the full k-pool + width (tail columns included) is exchanged.""" + decode, pad, topk, kpool = 5, 3, 8, 4 + width = topk + kpool - 1 + sizes = [17, 9, 21] + buffer = torch.full((decode + sum(sizes) + pad, width), -7, dtype=torch.int32) + gathered = torch.arange(sum(sizes) * width, dtype=torch.int32).reshape(-1, width) + calls = [] + + def fake_group(): + def all_gatherv(local, dim, sizes): + calls.append((local, sizes)) + return gathered + + return SimpleNamespace(all_gatherv=all_gatherv) + + monkeypatch.setattr(sparse_indexer, "get_tp_group", fake_group) + for rank in range(len(sizes)): + PrefillRowShard(decode, sizes, rank).exchange_topk(buffer, width) + assert calls[0][1] == sizes and calls[0][0].is_contiguous() + torch.testing.assert_close(buffer[decode : decode + sum(sizes)], gathered) + assert torch.all(buffer[:decode] == -7), "decode rows stay outside" + assert torch.all(buffer[decode + sum(sizes) :] == -7), "padding untouched" + + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_balanced_row_shard_partitions_and_declines(tp_size: int) -> None: + """The split covers every row once with >= 1 row per rank and declines + below the activation floor (or on a single rank).""" + n = indexer.MIN_TP_SHARD_ROWS_PER_RANK * tp_size + shapes = [ + ([4 * n], [4 * n]), # fresh prompt: the full causal cost ramp + # uneven requests at different context depths + ([2 * n, 7, 4 * n, n // 3 + 5], [2 * n, 9007, 4 * n + 500, n // 3 + 60005]), + ] + ratio = 4 + for query_lens, seq_lens in shapes: + sizes = indexer.balanced_prefill_row_shard( + torch.tensor(seq_lens, dtype=torch.int32), + torch.tensor(query_lens, dtype=torch.int32), + ratio, + tp_size, + ) + assert sizes is not None and len(sizes) == tp_size + assert sum(sizes) == sum(query_lens) and min(sizes) >= 1 + lens = torch.tensor([n - 1], dtype=torch.int32) + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, tp_size) is None + assert indexer.balanced_prefill_row_shard(lens, lens, ratio, 1) is None + + +def test_row_sharding_gate_rejects_unsupported_configurations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The gate is the whole safety envelope; nothing else guards the exchange.""" + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: True) + + def supported(**kwargs): + mode = kwargs.pop("cudagraph_mode", CUDAGraphMode.PIECEWISE) + cfg = SimpleNamespace(compilation_config=SimpleNamespace(cudagraph_mode=mode)) + args = {"dcp_world_size": 1, "use_pcp": False, "tp_size": 4, **kwargs} + return indexer.tp_prefill_row_sharding_supported(cfg, **args) + + assert supported() + for bad in ({"tp_size": 1}, {"dcp_world_size": 2}, {"use_pcp": True}): + assert not supported(**bad) + for name in ( + "VLLM_DISABLE_PYNCCL", + "VLLM_USE_NCCL_SYMM_MEM", + "VLLM_BATCH_INVARIANT", + ): + with monkeypatch.context() as m: + m.setattr(indexer.envs, name, True) + assert not supported() + # FULL captures mixed batches whole, so the exchange must stay out. + assert not supported(cudagraph_mode=CUDAGraphMode.FULL) + assert supported(cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE) + + +def test_profiling_run_reserves_exchange_transient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The dummy run returns before the exchange executes, so its all_gatherv + output plus staging copy must be reserved up front when row sharding is + active -- and only then.""" + rows, topk = 128, 8 + buffer = torch.empty((rows, topk), dtype=torch.int32) + peak = 2 * rows * topk * buffer.element_size() + specs = (((1,), torch.uint8), ((1,), torch.uint8)) + numels: list[int] = [] + real_empty = torch.empty + + def recording_empty(*args, **kwargs): + out = real_empty(*args, **kwargs) + numels.append(out.numel()) + return out + + def run_dummy(tp_world: int) -> None: + numels.clear() + with monkeypatch.context() as m: + m.setattr(torch, "empty", recording_empty) + m.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata=None), + ) + m.setattr(sparse_indexer, "_gather_workspace_shapes", lambda *a: specs) + m.setattr( + sparse_indexer, + "current_workspace_manager", + lambda: SimpleNamespace(get_simultaneous=lambda *s: None), + ) + m.setattr( + sparse_indexer, + "get_tp_group", + lambda: SimpleNamespace(world_size=tp_world), + ) + m.setattr( + sparse_indexer, + "get_current_vllm_config_or_none", + lambda: SimpleNamespace( + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.PIECEWISE + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + ), + ), + ) + m.setattr(sparse_indexer.envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 1) + m.setattr(indexer.current_platform, "is_cuda", lambda: True) + sparse_indexer.sparse_attn_indexer( + hidden_states=torch.empty(4, 4), + k_cache_prefix="k", + kv_cache=torch.empty(4, 4), + q_quant=torch.empty(4, 4), + q_scale=None, + k=None, + weights=torch.empty(4), + quant_block_size=128, + scale_fmt=None, + topk_tokens=topk, + head_dim=64, + max_model_len=1024, + total_seq_lens=512, + topk_indices_buffer=buffer, + skip_k_cache_insert=True, + use_pcp=False, + pcp_shard_decode_requests=False, + dense_mha_metadata_layer_name="", + ) + + run_dummy(tp_world=2) + assert peak in numels + run_dummy(tp_world=1) + assert peak not in numels diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index f8fedf046aca..935e4c0d2d44 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -9,8 +9,17 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CUDAGraphMode, get_current_vllm_config -from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.config import ( + CUDAGraphMode, + get_current_vllm_config, + get_current_vllm_config_or_none, +) +from vllm.distributed import ( + get_dcp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tp_group, +) from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp @@ -43,6 +52,7 @@ ) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, + tp_prefill_row_sharding_supported, ) from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k @@ -317,6 +327,75 @@ def kv_cache_as_quant_view( return kv_cache.unsqueeze(-2) +class PrefillRowShard: + """One TP rank's contiguous slice of the replicated prefill rows. + + Rows are independent, so each rank scores only its slice and the group + exchanges final int32 top-k indices instead of the logits.""" + + def __init__(self, num_decode_tokens: int, sizes: list[int], rank: int): + self.num_decode_tokens = num_decode_tokens + self.sizes = sizes + self.start = num_decode_tokens + sum(sizes[:rank]) + self.stop = self.start + sizes[rank] + + @classmethod + def from_metadata( + cls, prefill_metadata, num_decode_tokens: int + ) -> "PrefillRowShard | None": + sizes = getattr(prefill_metadata, "row_shard_sizes", None) + if sizes is None: + return None + return cls(num_decode_tokens, sizes, get_tensor_model_parallel_rank()) + + def narrow(self, chunk) -> "tuple[int, int, torch.Tensor, torch.Tensor] | None": + """Row window and narrowed causal bounds of ``chunk`` for this rank.""" + start = max(chunk.token_start, self.start) + end = min(chunk.token_end, self.stop) + if start >= end: + return None + lo, hi = start - chunk.token_start, end - chunk.token_start + return start, end, chunk.cu_seqlen_ks[lo:hi], chunk.cu_seqlen_ke[lo:hi] + + def exchange_topk(self, topk_buffer: torch.Tensor, width: int) -> None: + """Gather per-rank final indices back into global row order.""" + # The local input must outlive the async PyNCCL gather's enqueued copy. + local = topk_buffer[self.start : self.stop, :width].contiguous() + end = self.num_decode_tokens + sum(self.sizes) + topk_buffer[self.num_decode_tokens : end, :width] = get_tp_group().all_gatherv( + local, dim=0, sizes=self.sizes + ) + + +def _reserve_row_shard_exchange_transient( + topk_indices_buffer: torch.Tensor, device: torch.device +) -> None: + """Profiling/dummy runs return before ``exchange_topk`` executes; reserve + its worst-case transient so the memory profile accounts for it. + + The exchange stages this rank's contiguous slice and the all_gatherv + wrapper allocates the full gathered tensor; each is bounded by the + buffer's row capacity times its (tile-padded) width. + """ + cfg = get_current_vllm_config_or_none() + if cfg is None: + return + parallel = cfg.parallel_config + if not tp_prefill_row_sharding_supported( + cfg, + parallel.decode_context_parallel_size, + parallel.prefill_context_parallel_size > 1, + get_tp_group().world_size, + ): + return + rows, width = topk_indices_buffer.shape[:2] + _ = torch.empty( + 2 * rows * width * topk_indices_buffer.element_size(), + dtype=torch.uint8, + device=device, + ) + + @eager_break_during_capture def sparse_attn_indexer( hidden_states: torch.Tensor, @@ -388,6 +467,10 @@ def sparse_attn_indexer( max_logits_elems, dtype=torch.uint8, device=hidden_states.device ) + # The dummy run returns before exchange_topk executes, so reserve its + # worst-case transient for the memory profile. + _reserve_row_shard_exchange_transient(topk_indices_buffer, hidden_states.device) + return sparse_attn_indexer_fake( hidden_states, k_cache_prefix, @@ -515,6 +598,12 @@ def sparse_attn_indexer( scales_spec, *gather_specs, ) + # One row window per rank; DSV4.1 candidate rows stay replicated. + row_shard = ( + None + if candidate_blocks is not None + else PrefillRowShard.from_metadata(prefill_metadata, num_decode_tokens) + ) for chunk in prefill_metadata.chunks: cu_seqlen_ks = chunk.cu_seqlen_ks cu_seqlen_ke = chunk.cu_seqlen_ke @@ -530,6 +619,15 @@ def sparse_attn_indexer( chunk.local_cu_seq_lens, ) + # Narrow scoring to this rank; KV gathers remain unconditional for + # continuation chunks that reuse the workspace. + row_start, row_end = chunk.token_start, chunk.token_end + if row_shard is not None: + narrowed = row_shard.narrow(chunk) + if narrowed is None: + continue + row_start, row_end, cu_seqlen_ks, cu_seqlen_ke = narrowed + # PCP + DCP KV all-gather. deinterleave_idx = chunk.pcp_deinterleave_idx if deinterleave_idx is not None: @@ -547,15 +645,9 @@ def sparse_attn_indexer( out_scales, ) - q_slice = q_quant[chunk.token_start : chunk.token_end] - q_scale_slice = ( - q_scale[chunk.token_start : chunk.token_end] - if q_scale is not None - else None - ) - topk_indices = topk_indices_buffer[ - chunk.token_start : chunk.token_end, :topk_tokens - ] + q_slice = q_quant[row_start:row_end] + q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None + topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) @@ -578,7 +670,7 @@ def sparse_attn_indexer( q_slice_cast, k_quant_cast, k_scale_cast, - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, ) @@ -586,7 +678,7 @@ def sparse_attn_indexer( logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], + weights[row_start:row_end], cu_seqlen_ks, cu_seqlen_ke, clean_logits=False, @@ -637,9 +729,13 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=chunk.cu_seqlen_ks, + row_starts=cu_seqlen_ks, ) + if row_shard is not None: + # Rows are already fully ranked, so this is a layout gather. + row_shard.exchange_topk(topk_indices_buffer, topk_tokens) + if has_decode: decode_metadata = attn_metadata_narrowed.decode assert decode_metadata is not None diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 77fa4c409152..154fdcda900e 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -12,7 +12,11 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.indexer_topk import get_indexer_topk -from vllm.model_executor.layers.sparse_attn_indexer import _merge_dcp_topk_global +from vllm.model_executor.layers.sparse_attn_indexer import ( + PrefillRowShard, + _merge_dcp_topk_global, + _reserve_row_shard_exchange_transient, +) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, _build_decode_scatter_indices, @@ -169,6 +173,10 @@ def sparse_attn_indexer_kpool( max_logits_elems, dtype=torch.uint8, device=hidden_states.device ) + # The dummy run returns before exchange_topk executes, so reserve its + # worst-case transient for the memory profile. + _reserve_row_shard_exchange_transient(topk_indices_buffer, hidden_states.device) + return topk_indices_buffer attn_metadata_narrowed = attn_metadata[k_cache_prefix] assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) @@ -282,6 +290,12 @@ def sparse_attn_indexer_kpool( _pos = positions[num_decode_tokens:num_tokens].to(torch.int32) _buf = topk_indices_buffer[num_decode_tokens:num_tokens] _fill_causal_indices(_buf, _pos) + # Short prefills score nothing locally; suppress the shard there. + row_shard = None + if not short_prefill: + row_shard = PrefillRowShard.from_metadata( + prefill_metadata, num_decode_tokens + ) # Get the full shared workspace buffers once (will allocate on first use). # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and @@ -310,12 +324,16 @@ def sparse_attn_indexer_kpool( chunk.local_cu_seq_lens, ) - q_slice = q_quant[chunk.token_start : chunk.token_end] - q_scale_slice = ( - q_scale[chunk.token_start : chunk.token_end] - if q_scale is not None - else None - ) + row_start, row_end = chunk.token_start, chunk.token_end + cu_seqlen_ks, cu_seqlen_ke = chunk.cu_seqlen_ks, chunk.cu_seqlen_ke + if row_shard is not None: + narrowed = row_shard.narrow(chunk) + if narrowed is None: + continue + row_start, row_end, cu_seqlen_ks, cu_seqlen_ke = narrowed + + q_slice = q_quant[row_start:row_end] + q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. if use_fp4_cache: @@ -335,9 +353,9 @@ def sparse_attn_indexer_kpool( logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), (k_quant_cast, k_scale_cast), - weights[chunk.token_start : chunk.token_end], - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, + weights[row_start:row_end], + cu_seqlen_ks, + cu_seqlen_ke, clean_logits=False, ) @@ -351,17 +369,15 @@ def sparse_attn_indexer_kpool( ) topk_dst = pool_topk else: - topk_dst = topk_indices_buffer[ - chunk.token_start : chunk.token_end, :topk_tokens - ] + topk_dst = topk_indices_buffer[row_start:row_end, :topk_tokens] if logits.shape[1] == 0: topk_dst.fill_(-1) else: torch.ops._C.top_k_per_row_prefill( logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, + cu_seqlen_ks, + cu_seqlen_ke, topk_dst, num_rows, logits.stride(0), @@ -376,7 +392,7 @@ def sparse_attn_indexer_kpool( dcp_rank, dcp_world_size, pool_cp_interleave, - row_starts=chunk.cu_seqlen_ks, + row_starts=cu_seqlen_ks, ) if index_kpool > 1: @@ -385,10 +401,7 @@ def sparse_attn_indexer_kpool( # Fused expand-pools + append-tail into one Triton kernel # (replaces ~25 elementwise ops). seq_len is token-granular # (pos+1); the kernel derives pool_len internally. - q_seq = ( - positions[chunk.token_start : chunk.token_end].to(torch.int32) - + 1 - ) + q_seq = positions[row_start:row_end].to(torch.int32) + 1 expanded = kpool_ops.expand_pools_and_append_tail( pool_ids, q_seq, index_kpool ) @@ -397,9 +410,13 @@ def sparse_attn_indexer_kpool( expanded = kpool_ops.expand_pools_to_tokens( pool_ids, valid, topk_tokens, index_kpool ) - topk_indices_buffer[ - chunk.token_start : chunk.token_end, : expanded.shape[-1] - ] = expanded + topk_indices_buffer[row_start:row_end, : expanded.shape[-1]] = expanded + + if row_shard is not None: + # The k-pool expansion appends the request's incomplete tail + # (index_kpool - 1 entries) after the logical top-k history; those + # columns are part of the attention index and are exchanged too. + row_shard.exchange_topk(topk_indices_buffer, topk_tokens + index_kpool - 1) if has_decode: decode_metadata = attn_metadata_narrowed.decode diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index e776057b05b7..97de9ffa9ba2 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -7,7 +7,7 @@ import torch import vllm.envs as envs -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs @@ -204,6 +204,64 @@ def __call__( ) +# Below this many rows per rank, stay replicated (TP4 engages at 64K rows). +MIN_TP_SHARD_ROWS_PER_RANK = 16_384 + + +def balanced_prefill_row_shard( + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + compress_ratio: int, + tp_size: int, +) -> list[int] | None: + """Return contiguous TP row counts balanced by indexer MQA work.""" + num_rows = int(query_lens_cpu.sum()) + if tp_size < 2 or num_rows < MIN_TP_SHARD_ROWS_PER_RANK * tp_size: + return None + + query_lens = query_lens_cpu.to(torch.int64) + first_key = torch.repeat_interleave( + seq_lens_cpu.to(torch.int64) - query_lens + 1, query_lens + ) + row_in_request = torch.arange(num_rows) - torch.repeat_interleave( + torch.cumsum(query_lens, 0) - query_lens, query_lens + ) + cost = torch.cumsum((first_key + row_in_request) // compress_ratio, 0) + total = int(cost[-1]) if cost.numel() else 0 + if total <= 0: + return None + + targets = torch.arange(1, tp_size) * total // tp_size + bounds = [0, *torch.searchsorted(cost, targets).tolist(), num_rows] + # Keep >= 1 row per rank (the floor above leaves ample room); this also + # handles repeated costs from pool compression. + for i in range(1, tp_size): + bounds[i] = max(bounds[i], bounds[i - 1] + 1) + for i in range(tp_size - 1, 0, -1): + bounds[i] = min(bounds[i], bounds[i + 1] - 1) + return [bounds[i + 1] - bounds[i] for i in range(tp_size)] + + +def tp_prefill_row_sharding_supported( + vllm_config: VllmConfig, + dcp_world_size: int, + use_pcp: bool, + tp_size: int, +) -> bool: + """Whether replicated prefill rows may be sharded across TP ranks.""" + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode or CUDAGraphMode.NONE + return ( + current_platform.is_cuda() + and dcp_world_size == 1 + and not use_pcp + and tp_size > 1 + and not envs.VLLM_DISABLE_PYNCCL + and not envs.VLLM_USE_NCCL_SYMM_MEM + and not envs.VLLM_BATCH_INVARIANT + and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL + ) + + class DeepseekV32IndexerBackend(AttentionBackend): @classmethod def supports_pcp(cls) -> bool: @@ -654,6 +712,8 @@ def __call__( class DeepseekV32IndexerPrefillMetadata: chunks: list[DeepseekV32IndexerPrefillChunkMetadata] max_prefill_seq_len: int = -1 + # Per-TP-rank contiguous row counts, or None for the replicated path. + row_shard_sizes: list[int] | None = None @dataclass @@ -941,6 +1001,7 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config self.dcp_world_size = parallel_config.decode_context_parallel_size + self.tp_size = parallel_config.tensor_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.pcp_world_size = parallel_config.prefill_context_parallel_size self.use_pcp = self.pcp_world_size > 1 @@ -1512,6 +1573,16 @@ def build( # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: chunks.append(metadata) + row_shard_sizes = None + if tp_prefill_row_sharding_supported( + self.vllm_config, self.dcp_world_size, self.use_pcp, self.tp_size + ): + row_shard_sizes = balanced_prefill_row_shard( + seq_lens_cpu[num_decodes:], + prefill_query_lens_cpu, + self.compress_ratio, + self.tp_size, + ) prefill_metadata = DeepseekV32IndexerPrefillMetadata( chunks, max_prefill_seq_len=( @@ -1519,6 +1590,7 @@ def build( if num_prefills > 0 else 0 ), + row_shard_sizes=row_shard_sizes, ) decode_metadata = None From 239a3fa3a8ffc6508e867452c115ea7bcee2304b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 11 Oct 2026 23:30:56 +0800 Subject: [PATCH 372/952] [Sleep] Retain draft-private weights with existing parameter selector (#61046) Signed-off-by: aoshen02 Co-authored-by: Codex --- docs/features/sleep_mode.md | 39 ++++++++++++++++++++++++--- tests/v1/worker/test_gpu_worker.py | 42 ++++++++++++++++++++++++++++++ vllm/config/model.py | 3 +++ vllm/v1/worker/gpu_worker.py | 19 +++++++++++--- 4 files changed, 97 insertions(+), 6 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 0db82de9d3c5..bdfef2ecac7d 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -77,8 +77,15 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`) -to runtime parameter-name glob -patterns for weights that stay frozen during training. Each pattern must match +to runtime parameter-name glob patterns for weights that stay frozen during +training. It covers two use cases: + +- **Frozen Engram tables (original use case):** retain table weights and scales + omitted by the trainer, without retransmitting them after every sleep. +- **Frozen draft weights:** retain draft-private parameters when the trainer + updates only the target model. + +Target patterns match `model.named_parameters()`; checkpoint names and `requires_grad` are not used. Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights @@ -87,7 +94,33 @@ copy. Allow enough host memory for backups. Level-1 behavior is unchanged. The trainer must omit these parameters from updates, and reload/post-processing must preserve their values and storage (including avoiding replacement with meta -tensors). This option does not filter updates and applies only to the target model. +tensors). This option does not filter updates. + +For frozen Engram tables omitted by the trainer, select the target table weights +and scales. For example, DeepSeek V4.1 uses these runtime parameter names: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash --enable-sleep-mode \ + --sleep-preserve-parameter-names '*.engram.embed_tokens.weight*' +``` + +CPU-offloaded tables remain on CPU without an extra backup copy. Engram's +[inference-time CPU offload](engram.md#cpu-offload) is configured separately; +this option preserves the frozen tables across level-2 sleep and weight updates. + +To retain a frozen drafter while updating only target weights, use `draft.*`: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.8-27B --enable-sleep-mode \ + --speculative-config '{"method":"mtp","num_speculative_tokens":3}' \ + --sleep-preserve-parameter-names 'draft.*' +``` + +vLLM selects all draft-private parameters internally; you do not need to list +them individually. Use `draft.` to select specific draft parameters. +Parameters shared with the target are excluded from the draft selection, so +target updates remain authoritative. An empty `draft.*` selection is valid. +The Engram and draft patterns can be supplied together. #### Release only KV cache memory diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 6dc20a817794..adb8f52e2e8b 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -19,6 +19,48 @@ ) +def test_sleep_preservation_restores_private_draft_not_shared_target(): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() + draft.private = torch.nn.Linear(2, 2, bias=False) + draft.shared = target.weight + expected = draft.private.weight.detach().clone() + pointer = draft.private.weight.data_ptr() + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + assert set(worker._get_sleep_parameters(target)) == { + "weight", + "draft.private.weight", + } + worker._sleep_saved_parameters = {"draft.private.weight": expected} + with torch.no_grad(): + target.weight.fill_(3.0) + draft.private.weight.zero_() + worker._restore_sleep_parameters(target) + torch.testing.assert_close(draft.private.weight, expected) + assert draft.private.weight.data_ptr() == pointer + assert torch.all(target.weight == 3.0) + assert worker._sleep_saved_parameters == {} + + +@pytest.mark.parametrize("has_shared_draft", [False, True]) +def test_sleep_preservation_allows_empty_draft_but_rejects_target_typo( + has_shared_draft, +): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() if has_shared_draft else None + if draft is not None: + draft.shared = target.weight + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + worker.model_config = SimpleNamespace(sleep_preserve_parameter_names=["draft.*"]) + worker._save_sleep_parameters(target) + assert worker._sleep_saved_parameters == {} + worker.model_config.sleep_preserve_parameter_names = ["missing"] + with pytest.raises(ValueError, match="No parameter matches sleep retention"): + worker._save_sleep_parameters(target) + + def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch): """Profiling must use serving's thread count to keep Dynamo guards valid.""" from torch._dynamo.testing import CompileCounter diff --git a/vllm/config/model.py b/vllm/config/model.py index 502f1bbd8d4f..d4bafad5d6d3 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -372,6 +372,9 @@ class ModelConfig: hip platforms are supported).""" sleep_preserve_parameter_names: list[str] = field(default_factory=list) """Parameter-name globs to preserve across level-2 sleep. + Use target globs for frozen Engram table weights and scales omitted by the + trainer (e.g. ``*.engram.embed_tokens.weight*``). + Use ``draft.`` for draft-private parameters, or ``draft.*`` for all. The sender must omit these parameters; loaders must preserve their storage.""" sleep_mode_backend: str = "cumem" """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"`` diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3efe77a3a4d2..04feb5da025a 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -254,16 +254,28 @@ def sleep_mode_backend(self) -> "SleepModeBackend": ) return self._sleep_mode_backend + def _get_sleep_parameters(self, model: nn.Module) -> dict[str, nn.Parameter]: + parameters = dict(model.named_parameters()) + draft = self.get_draft_model() + if draft is not None: + shared = {id(parameter) for parameter in parameters.values()} + parameters.update( + (f"draft.{name}", parameter) + for name, parameter in draft.named_parameters() + if id(parameter) not in shared + ) + return parameters + def _save_sleep_parameters(self, model: nn.Module) -> None: patterns = self.model_config.sleep_preserve_parameter_names if not patterns: self._sleep_saved_parameters = {} return - parameters = dict(model.named_parameters()) + parameters = self._get_sleep_parameters(model) names: set[str] = set() for pattern in patterns: matches = fnmatch_filter(parameters, pattern) - if not matches: + if not matches and pattern != "draft.*": raise ValueError(f"No parameter matches sleep retention: {pattern}") names.update(matches) self._sleep_saved_parameters = { @@ -274,8 +286,9 @@ def _save_sleep_parameters(self, model: nn.Module) -> None: @torch.no_grad() def _restore_sleep_parameters(self, model: nn.Module) -> None: + parameters = self._get_sleep_parameters(model) for name, value in self._sleep_saved_parameters.items(): - model.get_parameter(name).copy_(value) + parameters[name].copy_(value) self._sleep_saved_parameters.clear() def sleep(self, level: int = 1) -> None: From 93d2cab35e8ae66bbaa1dea1c077986d5a79eb80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:39:11 +0200 Subject: [PATCH 373/952] [Security] Fix P2P ZMQ socket exhaustion via unbounded peer sessions (#51504) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: jperezde Signed-off-by: Juan Pérez de Algaba Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- docs/usage/security.md | 19 ++ tests/test_envs.py | 18 ++ .../v1/kv_offload/tiering/p2p/test_manager.py | 273 ++++++++++++++++++ .../kv_offload/tiering/p2p/test_sessions.py | 138 +++++++++ .../tiering/p2p/test_zmq_transport.py | 44 +++ vllm/envs.py | 18 ++ vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 76 +++-- vllm/v1/kv_offload/tiering/p2p/manager.py | 125 +++++++- .../kv_offload/tiering/p2p/session/session.py | 43 ++- 9 files changed, 712 insertions(+), 42 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 50c4149be81b..f069de0def74 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -303,6 +303,25 @@ Only deployments whose API clients are trusted should start the server with `--trust-request-mm-kwargs` to restore per-request overrides. Do not enable this option on an endpoint exposed to untrusted clients. +### P2P KV-Offload Peer Session Limits + +Deployments that enable the P2P secondary tier (`OffloadingConnector` with `TieringOffloadingSpec` and a `p2p` secondary tier) accept peer addresses from `kv_transfer_params` in API requests. Without limits, an attacker with API access could submit many unique non-listening peer addresses, each consuming ZeroMQ sockets until the context quota is exhausted, crashing the EngineCore. + +vLLM enforces several configurable limits to prevent this: + +| Environment Variable | Default | Description | +| --- | --- | --- | +| `VLLM_P2P_MAX_PEERS` | `128` | Maximum concurrent P2P peer sessions. New peers (both outbound and inbound) are rejected when this limit is reached. | +| `VLLM_P2P_HANDSHAKE_TIMEOUT_S` | `30` | Seconds to wait for a newly connected peer to complete the P2P handshake. Sessions that remain connected but not ready after this deadline are torn down. | +| `VLLM_P2P_IDLE_TIMEOUT_S` | `300` | Seconds a P2P session may remain idle (no pending work, no messages) before it is eligible for eviction. Idle eviction runs only when remaining peer slots are within 8 of `VLLM_P2P_MAX_PEERS`. Set to `0` to disable idle eviction. | + +Connections that fail (ZMQ socket errors, handshake failures) are handled gracefully — the request falls back to local prefill instead of crashing the engine. + +**Recommendations:** + +- **Public-facing deployments with P2P enabled:** Lower `VLLM_P2P_MAX_PEERS` to the expected number of legitimate peers (e.g., `16` or `32`). +- **Sensitive deployments:** Consider deploying behind a reverse proxy that strips or validates `kv_transfer_params` from API requests to prevent client-supplied peer addresses entirely. + ## Tool Server and MCP Security vLLM supports connecting to external tool servers via the `--tool-server` argument. This enables models to call tools through the Responses API (`/v1/responses`). Tool server support works with all models — it is not limited to specific model architectures. diff --git a/tests/test_envs.py b/tests/test_envs.py index dfd08a9ec287..7356539bcdec 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -60,6 +60,24 @@ def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch) assert envs.VLLM_P2P_SIDE_CHANNEL_PORT == 5799 +def test_p2p_resource_limit_defaults_and_override( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("VLLM_P2P_MAX_PEERS", raising=False) + monkeypatch.delenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", raising=False) + monkeypatch.delenv("VLLM_P2P_IDLE_TIMEOUT_S", raising=False) + assert envs.VLLM_P2P_MAX_PEERS == 128 + assert envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S == 30 + assert envs.VLLM_P2P_IDLE_TIMEOUT_S == 300 + + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "64") + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "15") + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "60") + assert envs.VLLM_P2P_MAX_PEERS == 64 + assert envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S == 15 + assert envs.VLLM_P2P_IDLE_TIMEOUT_S == 60 + + def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1") monkeypatch.setenv("VLLM_PORT", "1234") diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index e35b1d022e42..930158c7f6c8 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -647,6 +647,7 @@ def __init__( close_req_ids: list[str] | None = None, close_stores: list[int] | None = None, close_failed_serves: list[ReqContext] | None = None, + last_activity_at: float | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -663,6 +664,9 @@ def __init__( self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] self.finishes: list[str] = [] + self.last_activity_at: float = ( + last_activity_at if last_activity_at is not None else time.monotonic() + ) # Mirror P2PSession._server._inflight (transfer_id → handle) and # P2PSession._client.has_active_loads for the shutdown-drain and # drain_jobs paths. Tests populate _server._inflight when needed. @@ -1844,3 +1848,272 @@ def test_same_host_port_gets_distinct_nixl_names(self, monkeypatch): mgr_b = self._construct(monkeypatch, host="localhost", port=5710) assert mgr_a._local_id == mgr_b._local_id == "localhost:5710" assert mgr_a._nixl_agent_name != mgr_b._nixl_agent_name + + +# --------------------------------------------------------------------------- +# Tests for session capacity limit (VLLM_P2P_MAX_PEERS) +# --------------------------------------------------------------------------- + + +class TestSessionCapLimit: + """_get_or_create_session and _accept_new_peers reject new peers when + the session count reaches VLLM_P2P_MAX_PEERS.""" + + def _make_with_control(self): + mgr = _make_manager() + created = [] + + class FakeControl: + def connect(self, peer_id): + conn = _RecordingConn(peer_id) + created.append(conn) + return conn + + def poll(self): + return [] + + class FakeData: + block_len = 4096 + base_addr = 0x1000 + num_blocks = 16 + config_fingerprint = "" + + def get_agent_metadata(self): + return b"meta" + + def add_remote_peer(self, *a, **k): + pass + + def remove_remote_peer(self, pid): + pass + + mgr._control = FakeControl() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + return mgr, created + + def test_outbound_rejected_at_capacity(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "2") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, created = self._make_with_control() + mgr._sessions["peer-a:1"] = _FakeSession(peer_id="peer-a:1") # type: ignore[assignment] + mgr._sessions["peer-b:2"] = _FakeSession(peer_id="peer-b:2") # type: ignore[assignment] + + result = mgr._get_or_create_session("peer-c:3") + assert result is None + assert "peer-c:3" not in mgr._sessions + assert len(created) == 0 + + def test_on_new_request_records_failure_on_cap(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + mgr._sessions["existing:1"] = _FakeSession(peer_id="existing:1") # type: ignore[assignment] + + ctx = _req_context( + kv_params=_remote_prefiller_kv_params( + remote_host="10.0.0.2", remote_port=9000, kv_request_id="req-cap" + ) + ) + mgr.on_new_request(ctx) + assert "req-cap" in mgr._failed_req_ids + + def test_inbound_rejected_at_capacity(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + mgr._sessions["existing:1"] = _FakeSession(peer_id="existing:1") # type: ignore[assignment] + + conn = _RecordingConn("new-peer:2") + mgr._accept_new_peers([conn]) + assert conn.close_calls == 1 + assert "new-peer:2" not in mgr._sessions + + def test_existing_session_reused_below_cap(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr, _ = self._make_with_control() + existing = _FakeSession(peer_id="existing:1") + mgr._sessions["existing:1"] = existing # type: ignore[assignment] + + result = mgr._get_or_create_session("existing:1") + assert result is existing + + +# --------------------------------------------------------------------------- +# Tests for graceful error boundary in _get_or_create_session +# --------------------------------------------------------------------------- + + +class TestGetOrCreateSessionErrorBoundary: + def test_connect_failure_returns_none(self): + mgr = _make_manager() + + class FailingControl: + def connect(self, peer_id): + raise RuntimeError("socket exhausted") + + def poll(self): + return [] + + mgr._control = FailingControl() # type: ignore[assignment] + result = mgr._get_or_create_session("bad:1") + assert result is None + assert "bad:1" not in mgr._sessions + + def test_session_construction_failure_closes_conn(self): + """connect() succeeding but P2PSession construction failing must + close the conn so it is not left registered without a session.""" + mgr = _make_manager() + created: list[_RecordingConn] = [] + + class RaisingConn(_RecordingConn): + def send(self, msg): + raise RuntimeError("send failed during connect") + + class Control: + def connect(self, peer_id): + conn = RaisingConn(peer_id) + created.append(conn) + return conn + + def poll(self): + return [] + + class FakeData: + block_len = 4096 + base_addr = 0x1000 + num_blocks = 16 + config_fingerprint = "" + + def get_agent_metadata(self): + return b"meta" + + mgr._control = Control() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + result = mgr._get_or_create_session("bad:2") + assert result is None + assert "bad:2" not in mgr._sessions + assert len(created) == 1 + assert created[0].close_calls == 1 + + +# --------------------------------------------------------------------------- +# Tests for idle session eviction +# --------------------------------------------------------------------------- + + +class TestIdleSessionEviction: + def _make_with_fake_data(self): + mgr = _make_manager() + + class FakeControl: + def poll(self): + return [] + + class FakeData: + def remove_remote_peer(self, pid): + pass + + mgr._control = FakeControl() # type: ignore[assignment] + mgr._data = FakeData() # type: ignore[assignment] + return mgr + + def test_idle_session_evicted(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + idle_session = _FakeSession( + peer_id="idle:1", + last_activity_at=time.monotonic() - 20, + ) + mgr._sessions["idle:1"] = idle_session # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "idle:1" not in mgr._sessions + + def test_idle_session_kept_below_headroom(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + idle_session = _FakeSession( + peer_id="idle:1", + last_activity_at=time.monotonic() - 20, + ) + mgr._sessions["idle:1"] = idle_session # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "idle:1" in mgr._sessions + + def test_active_session_kept(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + active = _FakeSession( + peer_id="active:1", + last_activity_at=time.monotonic(), + ) + mgr._sessions["active:1"] = active # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "active:1" in mgr._sessions + + def test_session_with_pending_work_not_evicted(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "1") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + + mgr = self._make_with_fake_data() + busy = _FakeSession( + peer_id="busy:1", + last_activity_at=time.monotonic() - 100, + ) + busy._server._inflight[1] = object() + mgr._sessions["busy:1"] = busy # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "busy:1" in mgr._sessions + + def test_disabled_when_zero(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "0") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + + mgr = self._make_with_fake_data() + old = _FakeSession( + peer_id="old:1", + last_activity_at=time.monotonic() - 999999, + ) + mgr._sessions["old:1"] = old # type: ignore[assignment] + + mgr._reap_idle_sessions() + assert "old:1" in mgr._sessions diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 09739345be20..300cbacd665d 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -2698,3 +2698,141 @@ def test_success_wrong_type(self): } with pytest.raises(ValueError, match="success"): TransferDoneMsg.validate(msg) + + +# --------------------------------------------------------------------------- +# Tests for handshake deadline +# --------------------------------------------------------------------------- + + +class TestHandshakeDeadline: + """A connected-but-not-ready session is marked dead after the + handshake timeout expires.""" + + def test_marks_dead_after_timeout(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="stale:9999") + session = P2PSession( + peer_id="stale:9999", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + + assert session.connected + assert not session.ready + assert session.alive + + session._connected_at = time.monotonic() - 5 + + session.poll() + assert not session.alive + + def test_no_timeout_when_ready(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="good:1234") + session = P2PSession( + peer_id="good:1234", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + + # Complete the handshake by feeding ConnectMsg + ConnectAckMsg. + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "good:1234", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + {TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "good:1234"}, + ] + session.poll() + assert session.ready + + session._connected_at = time.monotonic() - 100 + session.poll() + assert session.alive + + def test_buffered_ack_at_deadline_completes_handshake(self, monkeypatch): + """A ConnectAck already in the inbox at the deadline must complete + the handshake instead of being discarded by mark_dead.""" + monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_HANDSHAKE_TIMEOUT_S", None) + + transport = FakeDataTransport() + conn = FakeConnection(peer_id="late-ack:1") + session = P2PSession( + peer_id="late-ack:1", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + session._connected_at = time.monotonic() - 5 + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "late-ack:1", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + {TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "late-ack:1"}, + ] + session.poll() + assert session.ready + assert session.alive + + def test_activity_timestamp_updated_on_recv(self): + transport = FakeDataTransport() + conn = FakeConnection(peer_id="peer:1") + session = P2PSession( + peer_id="peer:1", + local_id="local:1", + transport=transport, + local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, + conn=conn, + ) + before = session._last_activity_at + + conn._inbox = [ + { + TYPE_KEY: ConnectMsg.TYPE, + ConnectMsg.PEER_ID: "peer:1", + ConnectMsg.AGENT_METADATA: b"meta", + ConnectMsg.BASE_ADDR: 0x1000, + ConnectMsg.NUM_BLOCKS: 16, + ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.CONFIG_FINGERPRINT: "", + ConnectMsg.HASH_SEED: _DEFAULT_HASH_SEED, + }, + ] + session.poll() + assert session._last_activity_at >= before diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index f814a04887c4..90e87bb40742 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -336,3 +336,47 @@ def test_inbound_message_survives_dead_registration(self): finally: transport_a.close() transport_b.close() + + +class TestExpandedMonitorEvents: + """The monitor mask includes EVENT_DISCONNECTED and EVENT_CLOSED so + both are treated as fatal.""" + + def test_fatal_events_include_disconnected_and_closed(self): + from vllm.v1.kv_offload.tiering.p2p.control.zmq import ( + _FATAL_MONITOR_EVENTS, + ) + + assert zmq.EVENT_DISCONNECTED in _FATAL_MONITOR_EVENTS + assert zmq.EVENT_CLOSED in _FATAL_MONITOR_EVENTS + + def test_monitor_mask_covers_fatal_events(self): + from vllm.v1.kv_offload.tiering.p2p.control.zmq import ( + _FATAL_MONITOR_EVENTS, + _MONITOR_EVENT_MASK, + ) + + for ev in _FATAL_MONITOR_EVENTS: + assert _MONITOR_EVENT_MASK & ev == ev + + def test_disconnect_detected_after_peer_shutdown(self): + """When a peer's transport closes, the outbound connection is + eventually marked dead by the monitor.""" + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + + try: + conn_out = transport_b.connect(f"127.0.0.1:{port_a}") + conn_out.send({"type": "hello"}) + _wait_for_inbound(transport_a) + + transport_a.close() + + end = time.monotonic() + 15 + while conn_out.alive and time.monotonic() < end: + transport_b.poll() + time.sleep(0.05) + + assert not conn_out.alive + finally: + transport_b.close() diff --git a/vllm/envs.py b/vllm/envs.py index 02c282fd8ec2..14489a3f2503 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -233,6 +233,9 @@ VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 VLLM_P2P_SIDE_CHANNEL_HOST: str = "localhost" VLLM_P2P_SIDE_CHANNEL_PORT: int = 5710 + VLLM_P2P_MAX_PEERS: int = 128 + VLLM_P2P_HANDSHAKE_TIMEOUT_S: int = 30 + VLLM_P2P_IDLE_TIMEOUT_S: int = 300 VLLM_EC_SIDE_CHANNEL_HOST: str = "localhost" VLLM_EC_SIDE_CHANNEL_PORT: int = 5601 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 @@ -1739,6 +1742,21 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_P2P_SIDE_CHANNEL_PORT": lambda: int( os.getenv("VLLM_P2P_SIDE_CHANNEL_PORT", "5710") ), + # Maximum number of concurrent P2P peer sessions. Limits ZMQ socket + # consumption to prevent denial-of-service via unbounded peer + # creation from client-supplied kv_transfer_params. Default: 128. + "VLLM_P2P_MAX_PEERS": lambda: int(os.getenv("VLLM_P2P_MAX_PEERS", "128")), + # Seconds to wait for a P2P peer handshake to complete (connected + # but not yet ready). Sessions exceeding this deadline are reaped. + # Prevents never-responding peers from accumulating indefinitely. + "VLLM_P2P_HANDSHAKE_TIMEOUT_S": lambda: int( + os.getenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "30") + ), + # Seconds a P2P session may remain idle (no pending work, no + # messages) before it is eligible for eviction. Idle eviction runs + # only when remaining peer slots are within 8 of VLLM_P2P_MAX_PEERS. + # Default: 300. Set to 0 to disable. + "VLLM_P2P_IDLE_TIMEOUT_S": lambda: int(os.getenv("VLLM_P2P_IDLE_TIMEOUT_S", "300")), # IP address used for the EC connector's ZMQ side channel # (producer ROUTER bind, consumer DEALER dial). "VLLM_EC_SIDE_CHANNEL_HOST": lambda: os.getenv( diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index b28ed76dac97..074c3869f4cd 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -27,11 +27,35 @@ _HEARTBEAT_TIMEOUT_MS = 10000 _HEARTBEAT_TTL_MS = 10000 +# Events that indicate a fatal connection failure — the peer will never +# become reachable through this socket. +_FATAL_MONITOR_EVENTS: frozenset[int] = frozenset() # populated after import + # Shared sentinels returned when there is nothing to report. _EMPTY_INBOX: tuple[dict, ...] = () _EMPTY_NEW_CONNECTIONS: tuple[ControlConnection, ...] = () +def _build_fatal_events() -> frozenset[int]: + events = {zmq.EVENT_DISCONNECTED, zmq.EVENT_CLOSED} + for attr in ( + "EVENT_HANDSHAKE_FAILED_AUTH", + "EVENT_HANDSHAKE_FAILED_PROTOCOL", + "EVENT_HANDSHAKE_FAILED_NO_DETAIL", + ): + val = getattr(zmq, attr, None) + if val is not None: + events.add(val) + return frozenset(events) + + +_FATAL_MONITOR_EVENTS = _build_fatal_events() + +_MONITOR_EVENT_MASK: int = 0 +for _ev in _FATAL_MONITOR_EVENTS: + _MONITOR_EVENT_MASK |= _ev + + def _tcp_addr(host: str, port: int | str) -> str: return f"tcp://{host}:{port}" @@ -247,7 +271,7 @@ def _open_connection( safe_id = peer_id.replace(":", "-").replace("/", "-") monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}" self._monitor_seq += 1 - dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) + dealer.monitor(monitor_addr, _MONITOR_EVENT_MASK) monitor_sock = self._zmq_ctx.socket(zmq.PAIR) monitor_sock.connect(monitor_addr) @@ -319,29 +343,33 @@ def _recv_router(self) -> None: self._pending_inbound.append((sender_id, msg)) def _check_monitors(self) -> None: - """Non-blocking: check all monitor sockets for disconnection.""" + """Non-blocking: check all monitor sockets for fatal events.""" for conn in self._connections.values(): if not conn.alive: continue - try: - event = zmq.utils.monitor.recv_monitor_message( - conn.monitor_socket, zmq.NOBLOCK - ) - except zmq.Again: - continue - except zmq.ZMQError as exc: - logger.warning( - "ZmqTransport %s: monitor error for peer %s: %s", - self._local_id, - conn.peer_id, - exc, - ) - continue - - if event["event"] == zmq.EVENT_DISCONNECTED: - logger.debug( - "ZmqTransport %s: peer %s disconnected", - self._local_id, - conn.peer_id, - ) - conn.mark_dead() + while True: + try: + event = zmq.utils.monitor.recv_monitor_message( + conn.monitor_socket, zmq.NOBLOCK + ) + except zmq.Again: + break + except zmq.ZMQError as exc: + logger.warning( + "ZmqTransport %s: monitor error for peer %s: %s", + self._local_id, + conn.peer_id, + exc, + ) + break + + ev = event["event"] + if ev in _FATAL_MONITOR_EVENTS: + logger.debug( + "ZmqTransport %s: peer %s fatal event 0x%x", + self._local_id, + conn.peer_id, + ev, + ) + conn.mark_dead() + break diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index f759a4481e5d..b7268eb420d6 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -69,6 +69,11 @@ # enough to avoid busy-spinning the scheduler thread. _DRAIN_SLEEP_S = 0.001 +# Idle sessions are only reaped when remaining capacity is at most this +# many slots. Session setup is expensive; below this headroom, keep idle +# peers so a later burst does not pay handshake again. +_IDLE_REAP_HEADROOM = 8 + def _remote_prefiller_params(kv_params: dict | None) -> dict | None: """Return the ``remote_prefiller`` sub-dict, or None if absent. @@ -428,7 +433,9 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: _annotate_req_context(req_context) source = req_context.get_state(P2PSourceInfo) if source is not None: - self._get_or_create_session(source.peer_id) + session = self._get_or_create_session(source.peer_id) + if session is None: + self._failed_req_ids.add(source.kv_request_id) dest = req_context.get_state(P2PDestInfo) if dest is not None and dest.kv_request_id: return RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) @@ -685,38 +692,78 @@ def _get_hash_seed(self) -> str: self._hash_seed = get_none_hash_seed() return self._hash_seed - def _get_or_create_session(self, peer_id: str) -> P2PSession: + def _get_or_create_session(self, peer_id: str) -> P2PSession | None: """Return the existing session for peer_id, or open one outbound. Consumer-side helper for on_new_request: when ``remote_prefiller`` - (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer must reach the - producer at peer_id. If we already have a session toward that - peer (from a prior load or a peer-initiated inbound), reuse it; - otherwise open an outbound ControlConnection and build a + (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer + must reach the producer at peer_id. If we already have a session + toward that peer (from a prior load or a peer-initiated inbound), + reuse it; otherwise open an outbound ControlConnection and build a connected session. + + Returns None if the session could not be created (capacity + exceeded or transport error). """ session = self._sessions.get(peer_id) if session is not None: return session - conn = self._control.connect(peer_id) - session = P2PSession( - peer_id=peer_id, - local_id=self._local_id, - transport=self._data, - local_block_len=self._data.block_len, - local_hash_seed=self._get_hash_seed(), - conn=conn, - ) + + max_peers = envs.VLLM_P2P_MAX_PEERS + if len(self._sessions) >= max_peers: + logger.warning( + "P2P %s: rejecting outbound peer %s — session limit " + "reached (%d/%d). Set VLLM_P2P_MAX_PEERS to increase.", + self._local_id, + peer_id, + len(self._sessions), + max_peers, + ) + return None + + conn = None + try: + conn = self._control.connect(peer_id) + session = P2PSession( + peer_id=peer_id, + local_id=self._local_id, + transport=self._data, + local_block_len=self._data.block_len, + local_hash_seed=self._get_hash_seed(), + conn=conn, + ) + except Exception: + logger.exception( + "P2P %s: failed to open connection to %s", + self._local_id, + peer_id, + ) + if conn is not None: + conn.close() + return None + self._sessions[peer_id] = session return session def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> None: + max_peers = envs.VLLM_P2P_MAX_PEERS for conn in new_connections: logger.info( "P2P %s: accepting incoming connection from %s", self._local_id, conn.peer_id, ) + if len(self._sessions) >= max_peers: + logger.warning( + "P2P %s: rejecting inbound peer %s — session limit " + "reached (%d/%d). Set VLLM_P2P_MAX_PEERS to increase.", + self._local_id, + conn.peer_id, + len(self._sessions), + max_peers, + ) + conn.close() + continue try: existing = self._sessions.get(conn.peer_id) if existing is not None: @@ -827,6 +874,53 @@ def _reap_unbound_stores(self) -> None: len(batches), ) + def _reap_idle_sessions(self) -> None: + """Evict idle sessions when remaining peer slots are scarce. + + A session is idle when it has no pending work and its last activity + timestamp exceeds ``VLLM_P2P_IDLE_TIMEOUT_S``. Eviction runs only + when ``len(_sessions) >= VLLM_P2P_MAX_PEERS - _IDLE_REAP_HEADROOM`` + so a quiet but healthy mesh is not torn down. Uses the same + teardown path as ``_reap_dead_sessions``. + """ + idle_timeout = envs.VLLM_P2P_IDLE_TIMEOUT_S + if idle_timeout <= 0: + return + max_peers = envs.VLLM_P2P_MAX_PEERS + headroom = min(_IDLE_REAP_HEADROOM, max(0, max_peers - 1)) + if len(self._sessions) < max_peers - headroom: + return + deadline = time.monotonic() - idle_timeout + idle: list[str] | None = None + for pid, s in self._sessions.items(): + if not s.has_pending_work and s.last_activity_at <= deadline: + if idle is None: + idle = [] + idle.append(pid) + if idle is None: + return + for pid in idle: + session = self._sessions.pop(pid) + stale_kv_ids = [ + kid for kid, s in self._kv_to_session.items() if s is session + ] + for kid in stale_kv_ids: + del self._kv_to_session[kid] + close_result = session.close() + for job_id in close_result.failed_jobs: + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + for job_id in close_result.failed_stores: + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + self._failed_req_ids.update(close_result.failed_req_ids) + self._failed_serve_ctxs.extend(close_result.failed_serves) + self._data.remove_remote_peer(pid) + logger.info( + "P2P %s: evicted idle peer %s (idle > %ds)", + self._local_id, + pid, + idle_timeout, + ) + # ------------------------------------------------------------------ # Polling # ------------------------------------------------------------------ @@ -893,6 +987,7 @@ def _poll_once(self) -> None: self._reap_dead_sessions() self._reap_unbound_stores() + self._reap_idle_sessions() # ------------------------------------------------------------------ # Lifecycle diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index f43e1a5aab5a..ef490304c901 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -18,9 +18,11 @@ from __future__ import annotations import contextlib +import time from collections.abc import Sequence from typing import TYPE_CHECKING, NamedTuple +import vllm.envs as envs from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection @@ -128,6 +130,10 @@ def __init__( # Consecutive non-protocol dispatch errors. Reset on success. self._dispatch_error_count: int = 0 + now = time.monotonic() + self._connected_at: float | None = None + self._last_activity_at: float = now + # kv_request_ids whose FetchMsg arrived during the current poll # tick. Drained and returned in the next poll() result so the # manager can bind kv_request_id → session and replay any @@ -168,6 +174,11 @@ def has_pending_work(self) -> bool: """True while inbound loads or outbound transfers are outstanding.""" return self._client.has_active_loads or self._server.has_inflight_transfers + @property + def last_activity_at(self) -> float: + """Monotonic timestamp of the last meaningful session activity.""" + return self._last_activity_at + # ------------------------------------------------------------------ # Connection lifecycle # ------------------------------------------------------------------ @@ -181,6 +192,9 @@ def attach_connection(self, conn: ControlConnection) -> None: if self._conn is not None: raise ValueError(f"P2PSession {self.peer_id}: already connected") self._conn = conn + now = time.monotonic() + self._connected_at = now + self._last_activity_at = now self._send_connect() # ------------------------------------------------------------------ @@ -249,17 +263,39 @@ def serve_external_requests(self, parent: ParentManager) -> None: def poll(self) -> SessionPollResult: """Process incoming messages, drive transfers, apply timeouts.""" if self._conn is None: - # Pending session — store-job timeouts still apply so buffered - # jobs that never get picked up are surfaced as failures. return SessionPollResult( loads=[], stores=self._server.collect_idle_timeouts(), new_fetch_ids=[], ) - for msg in self._conn.recv(): + # Drain the inbox first so a ConnectAck already buffered at the + # deadline (e.g. after a scheduler stall) can complete the + # handshake instead of being discarded by mark_dead. + messages = self._conn.recv() + if messages: + self._last_activity_at = time.monotonic() + for msg in messages: self._on_message(msg) + # Handshake deadline: a peer that never completes the handshake + # is stuck in connected-but-not-ready state. Mark it dead so the + # manager reaps it instead of accumulating sockets indefinitely. + if ( + not self._send_ready + and self._conn.alive + and self._connected_at is not None + and time.monotonic() - self._connected_at + > envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S + ): + logger.warning( + "P2PSession %s: handshake timed out after %ds " + "— marking connection dead", + self.peer_id, + envs.VLLM_P2P_HANDSHAKE_TIMEOUT_S, + ) + self._conn.mark_dead() + loads = self._client.collect_results() stores = self._server.collect_results() self._server.drain_pending_aborts() @@ -533,6 +569,7 @@ def _do_send(self, msg: dict) -> None: return try: self._conn.send(msg) + self._last_activity_at = time.monotonic() logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY)) except Exception: # A send failure means the connection is broken. Swallowing it From b99fdbe1f896abd7f4b133685a37397c3427f7d5 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 12 Oct 2026 00:27:21 +0800 Subject: [PATCH 374/952] [Bugfix] Respect the configured offline chat template (#61061) Signed-off-by: Ting Sun Co-authored-by: Cyrus Leung --- tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py | 1 + vllm/entrypoints/common/offline.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py index 91d2b0aa434a..a381e0fa3b67 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py @@ -14,6 +14,7 @@ def _make_mock_llm() -> LLM: llm = object.__new__(LLM) + llm.chat_template = None llm.model_config = SimpleNamespace( runner_type="generate", enable_prompt_embeds=False ) diff --git a/vllm/entrypoints/common/offline.py b/vllm/entrypoints/common/offline.py index 1bb0bedccbd0..a5af009f2226 100644 --- a/vllm/entrypoints/common/offline.py +++ b/vllm/entrypoints/common/offline.py @@ -54,6 +54,7 @@ class OfflineInferenceMixin: renderer: BaseRenderer llm_engine: "LLMEngine" model_config: ModelConfig + chat_template: str | None def _resolve_mm_lora( self, @@ -180,7 +181,9 @@ def _preprocess_chat( renderer = self.renderer chat_params = ChatParams( - chat_template=chat_template, + chat_template=( + self.chat_template if chat_template is None else chat_template + ), chat_template_content_format=chat_template_content_format, chat_template_kwargs=merge_kwargs( chat_template_kwargs, From c059bee812842fb76c2b045e240c02f37c0d498d Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sun, 11 Oct 2026 09:57:07 -0700 Subject: [PATCH 375/952] [MRV2] Switch CI tests still running on MRV1 unnecessarily (#61014) Signed-off-by: Nick Hill --- .buildkite/test-amd.yaml | 4 -- .buildkite/test_areas/quantization.yaml | 8 ---- .../memory/sleep_mode/test_sleep_mode.py | 15 ++---- .../correctness_e2e/test_attn_quant.py | 1 - tests/models/kimi_k3/test_kda_metadata.py | 2 +- tests/utils.py | 14 ------ tests/v1/core/test_scheduler.py | 6 +++ tests/v1/core/utils.py | 2 +- tests/v1/e2e/test_replayssm_decode.py | 47 +++++++++++++------ .../unit/offloading_connector/test_config.py | 20 ++++---- .../test_eagle3_aux_hidden_states_pp.py | 2 +- tests/v1/worker/test_gpu_worker.py | 4 +- 12 files changed, 59 insertions(+), 66 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a3f7dbd83874..a14a0a33a4d2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2046,7 +2046,6 @@ steps: - vllm/model_executor/model_loader/ optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #-------------------------------------------------- mi300 · models / transformers ---------------------------------------------------# @@ -2353,7 +2352,6 @@ steps: # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now # we can only upgrade after this is resolved # TODO(jerryzh168): resolve the above comment - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.18.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -5163,7 +5161,6 @@ steps: - vllm/model_executor/model_loader/ optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# @@ -5186,7 +5183,6 @@ steps: - tests/rocm/test_moe_weight_replay.py optional: true commands: - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 55cef9fbd33a..9517b54f8ced 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -7,8 +7,6 @@ steps: key: quantization timeout_in_minutes: 40 parallelism: 4 - env: - VLLM_USE_V2_MODEL_RUNNER: "0" source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -41,8 +39,6 @@ steps: - tests/rocm/test_moe_weight_replay.py - tests/rocm/test_moe_padding.py commands: - # Use the ROCm default instead of the NVIDIA runner override. - - unset VLLM_USE_V2_MODEL_RUNNER - uv pip install --system torchao==0.17.0 - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -130,8 +126,6 @@ steps: key: quantized-models-test automatic_shard: true timeout_in_minutes: 65 - env: - VLLM_USE_V2_MODEL_RUNNER: "0" source_file_dependencies: - vllm/model_executor/layers/quantization - tests/models/quantization @@ -153,6 +147,4 @@ steps: - vllm/platforms/rocm.py - vllm/model_executor/model_loader/ commands: - # Use the ROCm default instead of the NVIDIA runner override. - - unset VLLM_USE_V2_MODEL_RUNNER - pytest -v -s models/quantization diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index d897d710015b..bc0f2756072f 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -411,15 +411,9 @@ def test_deep_sleep_compressed_tensors_kv_scales( @requires_fp8 -def test_deep_sleep_fp8_kvcache_mrv1( - monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared -): - # Regression test for https://github.com/vllm-project/vllm/pull/28783. - # In particular, verify that MRV1 does not rely on post_kv_cache_wake_up() - # to restore correct output after level-2 sleep. - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") - envs.disable_envs_cache() - +def test_deep_sleep_fp8_kvcache(gpu_memory_cleared): + # Regression test for https://github.com/vllm-project/vllm/pull/28783: + # fp8 KV cache output must be restored correctly after level-2 sleep. model = "Qwen/Qwen2-0.5B" used_bytes_baseline = current_platform.get_current_memory_usage() @@ -454,10 +448,9 @@ def test_deep_sleep_fp8_kvcache_mrv1( @requires_fp8 -def test_deep_sleep_fp8_kvcache_mrv1_with_undefined_remap( +def test_deep_sleep_fp8_kvcache_with_undefined_remap( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") envs.disable_envs_cache() diff --git a/tests/compile/correctness_e2e/test_attn_quant.py b/tests/compile/correctness_e2e/test_attn_quant.py index c7fc849ae52f..e13a8f264a6a 100644 --- a/tests/compile/correctness_e2e/test_attn_quant.py +++ b/tests/compile/correctness_e2e/test_attn_quant.py @@ -77,7 +77,6 @@ def settings(enabled): env1=env, env2=env, method="generate", - force_v1_runner=True, ) matches = [ int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 63024634d305..65b29b44bdae 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -219,7 +219,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): scheduler_config=SimpleNamespace(max_num_seqs=4), additional_config={}, num_speculative_tokens=2, - use_v2_model_runner=False, + use_v2_model_runner=True, ) builder = KimiK3KDAMetadataBuilder( kv_cache_spec=kv_cache_spec, diff --git a/tests/utils.py b/tests/utils.py index dd6bf4ae969d..94f62ef8b870 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1165,7 +1165,6 @@ def compare_two_settings( method: str = "generate", max_wait_seconds: float | None = None, include_seeded_sampling: bool = True, - force_v1_runner: bool = False, ) -> None: """Launch API server with two different sets of arguments/environments and compare the results of the API calls. @@ -1178,9 +1177,6 @@ def compare_two_settings( env2: The second set of environment variables to pass to the API server. include_seeded_sampling: Whether to include temperature=1.0 seeded sampling checks in the default generate comparison. - force_v1_runner: Whether to pin all compared settings to the v1 model - runner to avoid mixing model runner differences into correctness - tests. """ compare_all_settings( @@ -1190,7 +1186,6 @@ def compare_two_settings( method=method, max_wait_seconds=max_wait_seconds, include_seeded_sampling=include_seeded_sampling, - force_v1_runner=force_v1_runner, ) @@ -1202,7 +1197,6 @@ def compare_all_settings( method: str = "generate", max_wait_seconds: float | None = None, include_seeded_sampling: bool = True, - force_v1_runner: bool = False, ) -> None: """Launch API server with several different sets of arguments/environments and compare the results of the API calls with the first set of arguments. @@ -1213,16 +1207,8 @@ def compare_all_settings( all_envs: A list of environment dictionaries to pass to the API server. include_seeded_sampling: Whether to include temperature=1.0 seeded sampling checks in the default generate comparison. - force_v1_runner: Whether to pin all compared settings to the v1 model - runner to avoid mixing model runner differences into correctness - tests. """ - if force_v1_runner: - all_envs = [ - {"VLLM_USE_V2_MODEL_RUNNER": "0", **(env or {})} for env in all_envs - ] - trust_remote_code = False for args in all_args: if "--trust-remote-code" in args: diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 615dd82a1c7b..be3ea553d926 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -352,6 +352,12 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders(): _ = scheduler.schedule() assert req.num_output_placeholders > 0 + if scheduler.use_v2_model_runner: + # MRV2 staggers same-request decodes by pp_size steps to match the + # worker-side sampled-token broadcast cadence. + output = scheduler.schedule() + assert req.request_id not in output.num_scheduled_tokens + # before any update_from_output, we still expect the request can be # scheduled again (multi-step in-flight). output = scheduler.schedule() diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index b5e1bf89329e..cf221fc22a34 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -238,7 +238,7 @@ def create_scheduler( structured_output_manager=StructuredOutputManager(vllm_config), ) if use_v2_model_runner is None: - use_v2_model_runner = bool(envs.VLLM_USE_V2_MODEL_RUNNER) + use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER is not False scheduler.use_v2_model_runner = use_v2_model_runner return scheduler diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index ed3c6b0217c8..89797db4cdd7 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -122,7 +122,10 @@ def test_replayssm_flashinfer_decode_matches_baseline( @pytest.mark.parametrize("model_name", FLASHINFER_MODELS) -def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name): +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replayssm_flashinfer_spec_decode_matches_baseline( + vllm_runner, model_name, monkeypatch, use_v2_model_runner +): common = dict( max_model_len=1024, trust_remote_code=True, @@ -134,20 +137,34 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na "prompt_lookup_max": 3, }, ) - with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: - baseline = llm.generate_greedy_logprobs( - FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 - ) - with vllm_runner( - model_name, - use_replayssm=True, - replayssm_buffer_len=16, - mamba_backend="flashinfer", - **common, - ) as llm: - replay = llm.generate_greedy_logprobs( - FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 - ) + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) + envs.disable_envs_cache() + with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: + assert ( + llm.llm.llm_engine.vllm_config.use_v2_model_runner + is use_v2_model_runner + ) + baseline = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, + use_replayssm=True, + replayssm_buffer_len=16, + mamba_backend="flashinfer", + **common, + ) as llm: + assert ( + llm.llm.llm_engine.vllm_config.use_v2_model_runner + is use_v2_model_runner + ) + replay = llm.generate_greedy_logprobs( + FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5 + ) + finally: + envs.disable_envs_cache() check_logprobs_close( outputs_0_lst=baseline, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 237b37600401..07d530f97d1a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -71,7 +71,7 @@ def _make_vllm_config( decode_context_parallel_size=decode_context_parallel_size, ) config.kv_events_config = None - config.use_v2_model_runner = False + config.use_v2_model_runner = True config.kv_transfer_config = KVTransferConfig( kv_connector="OffloadingConnector", kv_role="kv_both", @@ -284,7 +284,7 @@ def _parallelism_agnostic( kv_cache_groups: list[KVCacheGroupSpec], *, canonical: bool = False, - v2: bool = False, + v2: bool = True, ) -> bool: config = _make_vllm_config( extra_config={"canonical_layout": True} if canonical else None @@ -308,7 +308,7 @@ def _replicated_layout( prefill_context_parallel_size: int = 1, decode_context_parallel_size: int = 1, use_mla: bool = True, - use_v2_model_runner: bool = False, + use_v2_model_runner: bool = True, distributed_executor_backend: Any = "mp", nnodes: int = 1, world_size: int | None = None, @@ -1105,7 +1105,10 @@ def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str): def test_parallelism_agnostic_for_single_full_attention_group(): - assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + # The direct (non-canonical) parallelism-agnostic layout is MRV1-only. + assert _parallelism_agnostic( + [KVCacheGroupSpec(["l0"], _full_attention_spec())], v2=False + ) _SWA_SPEC = SlidingWindowSpec( @@ -1147,7 +1150,9 @@ def _uniform_group(*specs: KVCacheSpec) -> KVCacheGroupSpec: ], ) def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): - assert not _parallelism_agnostic(kv_cache_groups) + # The direct layout is never parallelism-agnostic on MRV2, so check the + # MRV1 gate. + assert not _parallelism_agnostic(kv_cache_groups, v2=False) @pytest.mark.parametrize( @@ -1178,7 +1183,7 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): def test_canonical_layout_gate(kv_cache_groups, certified): """The canonical layout certifies portability group by group; none of these shapes are portable in the direct layout.""" - assert not _parallelism_agnostic(kv_cache_groups) + assert not _parallelism_agnostic(kv_cache_groups, v2=False) assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified @@ -1233,7 +1238,7 @@ def test_canonical_layout_certifies_v2_model_runner(): registration, so the static gate must not depend on the model-runner version — the v2 runner is the case the canonical layout exists for.""" groups = _groups(_full_attention_spec()) - assert _parallelism_agnostic(groups, canonical=True, v2=True) + assert _parallelism_agnostic(groups, canonical=True) @pytest.mark.parametrize("tp_size", [1, 2, 4]) @@ -1275,7 +1280,6 @@ def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler): def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() - config.use_v2_model_runner = True kv_cache_config = KVCacheConfig( num_blocks=0, kv_cache_tensors=[], diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py index 7c6acfafe321..f310c5a083ab 100644 --- a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py +++ b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py @@ -105,7 +105,7 @@ def make_stage(start, end, taps, world_size=2): model_config=SimpleNamespace(hf_text_config=config), quant_config=None, speculative_config=None, - use_v2_model_runner=False, + use_v2_model_runner=True, ) ) model._set_aux_hidden_state_layers(taps) diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index adb8f52e2e8b..a34133e2f8d3 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -424,8 +424,8 @@ def run_model(scheduler_output, intermediate_tensors): pipeline_parallel_size=2, distributed_executor_backend="mp" ), ), - use_v2_model_runner=False, - model_runner=SimpleNamespace(execute_model=run_model), + use_v2_model_runner=True, + model_runner=SimpleNamespace(execute_model=run_model, is_pooling_model=False), annotate_profile=lambda scheduler_output: nullcontext(), _pp_send_work=[previous_tensor_send], ) From f26d9f432ced7b1087d02d29bffed2249df1a18e Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 12 Oct 2026 01:55:00 +0800 Subject: [PATCH 376/952] [Bugfix][GLM-5.3-Flash] Honor `--mamba-ssm-cache-dtype` for the KDA recurrent state (#59147) Signed-off-by: Jared Wen Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/models/glm5next/test_kda_state_dtype.py | 41 +++++++++++++++++++ vllm/models/glm5next/common/kda.py | 4 +- vllm/models/glm5next/common/model.py | 4 +- 3 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 tests/models/glm5next/test_kda_state_dtype.py diff --git a/tests/models/glm5next/test_kda_state_dtype.py b/tests/models/glm5next/test_kda_state_dtype.py new file mode 100644 index 000000000000..b112e47cd520 --- /dev/null +++ b/tests/models/glm5next/test_kda_state_dtype.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``--mamba-ssm-cache-dtype`` must reach GLM-5.3-Flash's KDA recurrent state +(it used to be dropped, leaving the state float32 unconditionally).""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.glm5next.common.kda import Glm5NextLinearAttention +from vllm.models.glm5next.common.model import ( + Glm5NextForCausalLM, + Glm5NextForConditionalGeneration, +) + + +@pytest.mark.parametrize( + ("ssm_dtype", "expected"), + [("auto", torch.float32), ("float32", torch.float32), ("bfloat16", torch.bfloat16)], +) +def test_kda_recurrent_state_dtype_follows_ssm_cache_dtype(ssm_dtype, expected): + model_config = SimpleNamespace(dtype=torch.bfloat16) + cache_config = SimpleNamespace( + mamba_cache_dtype="auto", + mamba_ssm_cache_dtype=ssm_dtype, + use_kda_recoverssm=False, + ) + vllm_config = SimpleNamespace(model_config=model_config, cache_config=cache_config) + want = (torch.bfloat16, expected) + + assert Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == want + assert ( + Glm5NextForConditionalGeneration.get_mamba_state_dtype_from_config(vllm_config) + == want + ) + # The layer's own dtype must agree with the model-level one, since the + # KV-cache spec and the kernels read them separately. + layer = Glm5NextLinearAttention.__new__(Glm5NextLinearAttention) + layer.model_config, layer.cache_config = model_config, cache_config + assert layer.get_state_dtype() == want diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 141a8cc7cb82..68d405b40cfc 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -181,7 +181,9 @@ def get_state_dtype( if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") return MambaStateDtypeCalculator.kda_state_dtype( - self.model_config.dtype, self.cache_config.mamba_cache_dtype + self.model_config.dtype, + self.cache_config.mamba_cache_dtype, + self.cache_config.mamba_ssm_cache_dtype, ) def get_state_shape( diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index e35887ca7c54..1e58166054b0 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -1092,7 +1092,9 @@ def get_mamba_state_dtype_from_config( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, torch.dtype]: return MambaStateDtypeCalculator.kda_state_dtype( - vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, ) @classmethod From 3fe917abb13160d60bea5330399f5c09ca282348 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sun, 11 Oct 2026 16:04:41 -0700 Subject: [PATCH 377/952] [Bugfix][MRV2] Fix FlashInfer attention and ReplaySSM warmups on MRV2 (#61013) Signed-off-by: Nick Hill --- tests/model_executor/test_attention_warmup.py | 71 +++++++++++++++++++ .../model_executor/warmup/attention_warmup.py | 42 +++++++++++ vllm/model_executor/warmup/kernel_warmup.py | 13 ++-- .../warmup/minimax_m3_msa_warmup.py | 11 ++- vllm/v1/worker/gpu/model_runner.py | 9 ++- 5 files changed, 129 insertions(+), 17 deletions(-) create mode 100644 tests/model_executor/test_attention_warmup.py create mode 100644 vllm/model_executor/warmup/attention_warmup.py diff --git a/tests/model_executor/test_attention_warmup.py b/tests/model_executor/test_attention_warmup.py new file mode 100644 index 000000000000..861276280e24 --- /dev/null +++ b/tests/model_executor/test_attention_warmup.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest + +from vllm.model_executor.warmup import attention_warmup +from vllm.model_executor.warmup.attention_warmup import mixed_batch_attention_warmup + + +def _worker(*, use_v2: bool, max_num_batched_tokens: int, max_model_len: int = 4096): + calls = [] + worker = SimpleNamespace( + model_runner=SimpleNamespace(_dummy_run=lambda **kw: calls.append(kw)), + use_v2_model_runner=use_v2, + scheduler_config=SimpleNamespace(max_num_batched_tokens=max_num_batched_tokens), + model_config=SimpleNamespace(max_model_len=max_model_len), + execute_model=object(), + sample_tokens=object(), + ) + return worker, calls + + +@pytest.mark.parametrize( + ("max_num_batched_tokens", "max_model_len", "expected"), + [(8192, 4096, 16), (6, 4096, 6), (8192, 8, 8)], +) +def test_v1_runs_mixed_dummy_batch_within_limits( + max_num_batched_tokens, max_model_len, expected +): + worker, calls = _worker( + use_v2=False, + max_num_batched_tokens=max_num_batched_tokens, + max_model_len=max_model_len, + ) + mixed_batch_attention_warmup(worker) + assert calls == [ + dict( + num_tokens=expected, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + ] + + +@pytest.mark.parametrize(("max_num_batched_tokens", "expected"), [(8192, 16), (6, 6)]) +@pytest.mark.parametrize("ran", [True, False]) +def test_v2_runs_mixed_prefill_decode_step( + monkeypatch, max_num_batched_tokens, expected, ran +): + worker, dummy_calls = _worker( + use_v2=True, max_num_batched_tokens=max_num_batched_tokens + ) + mixed_calls = [] + + def fake_mixed_warmup(runner, execute_model, sample_tokens, num_tokens): + mixed_calls.append((runner, execute_model, sample_tokens, num_tokens)) + return ran + + monkeypatch.setattr( + attention_warmup, "run_mixed_prefill_decode_warmup", fake_mixed_warmup + ) + mixed_batch_attention_warmup(worker) + + assert mixed_calls == [ + (worker.model_runner, worker.execute_model, worker.sample_tokens, expected) + ] + # No dummy-run fallback when the mixed step is skipped. + assert dummy_calls == [] diff --git a/vllm/model_executor/warmup/attention_warmup.py b/vllm/model_executor/warmup/attention_warmup.py new file mode 100644 index 000000000000..cae556a49a0f --- /dev/null +++ b/vllm/model_executor/warmup/attention_warmup.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, cast + +from vllm.logger import init_logger +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def mixed_batch_attention_warmup(worker: "Worker", num_tokens: int = 16) -> None: + """Run a mixed prefill+decode batch to warm up both attention paths.""" + num_tokens = min( + num_tokens, + worker.scheduler_config.max_num_batched_tokens, + worker.model_config.max_model_len, + ) + if not worker.use_v2_model_runner: + worker.model_runner._dummy_run( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + return + + if not run_mixed_prefill_decode_warmup( + cast("GPUModelRunner", worker.model_runner), + worker.execute_model, + worker.sample_tokens, + num_tokens, + ): + logger.info_once( + "Skipping mixed prefill+decode attention warmup; attention kernels " + "will be loaded on the first request." + ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 685b33a5652a..b7bf50865fca 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -15,6 +15,9 @@ import vllm.envs as envs from vllm.logger import init_logger +from vllm.model_executor.warmup.attention_warmup import ( + mixed_batch_attention_warmup, +) from vllm.model_executor.warmup.b12x_warmup import b12x_warmup from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup @@ -286,15 +289,7 @@ def _is_flashinfer_backend(backend): ) ): logger.info_once("Warming up FlashInfer attention.") - # Warmup with mixed batch containing both prefill and decode tokens - # This is to warm up both prefill and decode attention kernels - worker.model_runner._dummy_run( - num_tokens=16, - skip_eplb=True, - is_profile=True, - force_attention=True, - create_mixed_batch=True, - ) + mixed_batch_attention_warmup(worker) def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: diff --git a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py index 18bf14249111..2c8996a93341 100644 --- a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py +++ b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py @@ -4,6 +4,9 @@ from typing import TYPE_CHECKING from vllm.logger import init_logger +from vllm.model_executor.warmup.attention_warmup import ( + mixed_batch_attention_warmup, +) from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention from vllm.platforms import current_platform from vllm.tracing import instrument @@ -34,10 +37,4 @@ def minimax_m3_msa_warmup(worker: "Worker") -> None: logger.info("Warming up MiniMax M3 MSA kernels.") # Cover sparse prefill through the normal model path. - worker.model_runner._dummy_run( - num_tokens=16, - skip_eplb=True, - is_profile=True, - force_attention=True, - create_mixed_batch=True, - ) + mixed_batch_attention_warmup(worker) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 20fa4c17bb07..57cef586f81f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -857,6 +857,11 @@ def _dummy_run( - num_reqs * self.speculator.num_extra_query_per_req, ) + assert num_tokens <= self.max_num_tokens, ( + f"Dummy run with {num_tokens} tokens exceeds max_num_batched_tokens " + f"({self.max_num_tokens})" + ) + # Distribute the remainder evenly so no dummy request exceeds # ceil(num_tokens / num_reqs) <= max_model_len tokens. @@ -1898,9 +1903,11 @@ def execute_model( ) assert block_tables is not None attn_groups = self.attn_groups - if dummy_run and is_profile: + if dummy_run and is_profile and not valid_dummy_state_slots: # Mamba layers take a cheap warmup path with no metadata; # attention metadata is still built so those kernels tune. + # valid_dummy_state_slots runs the real Mamba kernels on distinct + # state slots; the caller must restore that state afterwards. attn_groups = [ [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] for groups in attn_groups From cc9b0f097d49abfa6fc94d1d773dca1cc2b9619c Mon Sep 17 00:00:00 2001 From: Laith Sakka Date: Sun, 11 Oct 2026 17:38:56 -0700 Subject: [PATCH 378/952] [Fix] MiniMax: fix CUDA IPC handle exchange on cuda-bindings >= 13.4 (#60917) Signed-off-by: Laith Sakka --- .../minimax_rms_norm/lamport_workspace.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py index 0b2ef3a6aabd..eac2a64a0457 100644 --- a/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py +++ b/vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py @@ -4,6 +4,7 @@ import array import contextlib +import ctypes import struct import sys import threading @@ -16,6 +17,7 @@ from cuda import cudart _ALIGN = 1 << 21 # 2 MiB — CUDA IPC allocation alignment +_IPC_HANDLE_SIZE = 64 # sizeof(cudaIpcMemHandle_t) # --------------------------------------------------------------------------- @@ -30,6 +32,20 @@ def _check(error): raise RuntimeError(f"CUDA runtime error: {error}") +# cuda-bindings >= 13.4 no longer exposes cudaIpcMemHandle_t.reserved, so +# copy the raw handle bytes through getPtr() instead. This is safe because the +# CUDA ABI fixes the handle as an opaque 64-byte struct (char reserved[64]). +def _ipc_handle_to_bytes(handle) -> bytes: + return ctypes.string_at(handle.getPtr(), _IPC_HANDLE_SIZE) + + +def _ipc_handle_from_bytes(raw: bytes): + assert len(raw) == _IPC_HANDLE_SIZE, len(raw) + handle = cudart.cudaIpcMemHandle_t() + ctypes.memmove(handle.getPtr(), raw, _IPC_HANDLE_SIZE) + return handle + + def _cuda_malloc(size: int): aligned = ((size + _ALIGN - 1) >> 21) << 21 err, ptr = cudart.cudaMalloc(aligned) @@ -84,15 +100,16 @@ def __init__(self, rank: int, world_size: int, size: int, process_group=None): all_handles: list[bytes | None] = [None] * world_size torch.distributed.all_gather_object( - all_handles, bytes(local_handle.reserved), group=process_group + all_handles, _ipc_handle_to_bytes(local_handle), group=process_group ) for r in range(world_size): if r == rank: self.peer_ptrs[r] = self.local_ptr else: - handle = cudart.cudaIpcMemHandle_t() - handle.reserved = all_handles[r] + raw = all_handles[r] + assert raw is not None + handle = _ipc_handle_from_bytes(raw) err, ptr = cudart.cudaIpcOpenMemHandle( handle, cudart.cudaIpcMemLazyEnablePeerAccess ) From 03091e28b8ba956634c153beab0e52decafdef31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Ganczarenko?= Date: Mon, 12 Oct 2026 03:25:29 +0200 Subject: [PATCH 379/952] [Bugfix] Read Ling quantization_config from top-level config for Bailing V3 VL (#59026) Signed-off-by: Michal Ganczarenko Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/bailing_moe_v3.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index 9f5c078fac9c..d82e4c1b8b41 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -20,7 +20,13 @@ from transformers.configuration_utils import PreTrainedConfig from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.config import ( + CacheConfig, + ModelConfig, + VllmConfig, + get_current_vllm_config, + get_current_vllm_config_or_none, +) from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -230,6 +236,12 @@ def _configure_ling_fp8_quant_config( quant_config.ignored_layers_match_mode = "suffix" hf_quant_config = getattr(config, "quantization_config", None) + if hf_quant_config is None: + vllm_config = get_current_vllm_config_or_none() + model_config = getattr(vllm_config, "model_config", None) + hf_quant_config = getattr( + getattr(model_config, "hf_config", None), "quantization_config", None + ) if not isinstance(hf_quant_config, dict): return From 124b89d5b7bbbf7599763858c247848e073b024a Mon Sep 17 00:00:00 2001 From: dingdangmao <2756823972@qq.com> Date: Mon, 12 Oct 2026 10:08:14 +0800 Subject: [PATCH 380/952] [Bugfix][RecoverSSM] Correct aligned state indices at exact block boundaries (#59962) Signed-off-by: apex-mochen <2756823972@qq.com> Co-authored-by: Codex --- tests/models/kimi_k3/test_kda.py | 9 ++++++--- tests/v1/worker/test_mamba_hybrid_model_state.py | 16 ++++++++++++---- vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 3 ++- vllm/v1/worker/gpu/model_states/recoverssm.py | 2 +- 4 files changed, 21 insertions(+), 9 deletions(-) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 88acdb6197dc..0b250654450e 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -692,8 +692,10 @@ def test_kda_spec_decode_correctness( pytest.param(True, False, None, True, id="aligned"), ], ) +@pytest.mark.parametrize("accepted", [(2, 8), (3, 4), (5, 8)]) @torch.inference_mode() def test_kda_recoverssm_verify_and_group_commit( + accepted: tuple[int, int], monkeypatch: pytest.MonkeyPatch, lower_bound: float | None, use_request_indices: bool, @@ -738,7 +740,6 @@ def test_kda_recoverssm_verify_and_group_commit( state_indices = torch.tensor( [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE ) - accepted = [2, 8] if use_request_indices: global_num_accepted = torch.tensor( [0, accepted[0], 0, accepted[1]], @@ -891,7 +892,9 @@ def test_kda_recoverssm_verify_and_group_commit( if align_mode: assert block_table is not None row = request_indices[seq_idx] if use_request_indices else seq_idx - final_block = block_table[row, (4 + commit_len) // 8] + final_block = block_table[ + row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len] + ] committed_states[final_block] = committed_state.transpose(-1, -2) if align_mode and 4 + commit_len >= 8: _, boundary_state = naive_recurrent_kda( @@ -944,7 +947,7 @@ def test_kda_recoverssm_verify_and_group_commit( if align_mode: assert block_table is not None row = request_indices[seq_idx] if use_request_indices else seq_idx - block = block_table[row, (4 + commit_len) // 8] + block = block_table[row, {2: 0, 3: 0, 4: 0, 5: 1, 8: 1}[commit_len]] source_block = state_indices[seq_idx] if align_mode else block if conv_state_dim_first: actual_conv = layer.kv_cache[0][block, :, :history_len] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 0dcc4725d779..b52af4224974 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -387,7 +387,13 @@ def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: +@pytest.mark.parametrize( + ("computed", "sampled", "expected_column"), + [(4, 3, 0), (4, 4, 0), (4, 5, 1), (8, 8, 1), (8, 9, 2)], +) +def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias( + computed: int, sampled: int, expected_column: int +) -> None: state = object.__new__(MambaHybridModelState) state._align_mode = True state._mamba_ctx = None @@ -400,11 +406,13 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() - metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata( num_spec_decodes=1, request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"), - num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor( + [6, computed], dtype=torch.int32, device="cuda" + ), block_size=8, block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"), ) - num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + num_sampled = torch.tensor([2, sampled], dtype=torch.int32, device="cuda") idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda") group = SimpleNamespace(layer_names=["layer"]) @@ -412,7 +420,7 @@ def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() - state.postprocess_state(idx_mapping, num_sampled) - expected_state_indices = [-1, 1, -1, -1, -1] + expected_state_indices = [-1, expected_column, -1, -1, -1] assert state._mamba_state_idx_gpu.tolist() == expected_state_indices expected_accepted = [9, 1, 9, 2, 9] assert state.num_accepted_tokens_gpu.tolist() == expected_accepted diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index c603ec7ed1fc..7eda21f1a5ab 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -284,7 +284,8 @@ def _prepare_commit_plan_kernel( ) final_num_computed = num_computed + commit_len final_state_col = tl.minimum( - final_num_computed // mamba_block_size, block_table_width - 1 + tl.maximum(final_num_computed - 1, 0) // mamba_block_size, + block_table_width - 1, ) final_state_idx = tl.load( block_table_ptr diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index cfbe1a4d45d1..985f29079814 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -94,7 +94,7 @@ def _postprocess_recoverssm_align_kernel( tl.store( state_idx_ptr + req_state_idx, tl.minimum( - (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE, BLOCK_TABLE_WIDTH - 1, ), ) From f147949be8b9e576f80133bee01e717ee7c5db0e Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 12 Oct 2026 10:25:32 +0800 Subject: [PATCH 381/952] [XPU] limit FA kernel block size to multipleOf 64 (#57297) Signed-off-by: Yan Ma --- .buildkite/intel_jobs/test-intel.yaml | 2 +- .../unit/test_bidirectional_kv_transfer.py | 6 +- .../unit/test_mooncake_connector.py | 62 ++--- .../unit/test_mooncake_connector_hma.py | 34 ++- .../test_mooncake_connector_hybrid_mamba.py | 16 +- .../kv_connector/unit/test_multi_connector.py | 21 +- .../kv_connector/unit/test_nixl_connector.py | 224 ++++++++++++------ .../unit/test_nixl_connector_hma.py | 3 +- tests/v1/kv_connector/unit/utils.py | 5 + vllm/v1/attention/backends/flash_attn.py | 2 + vllm/v1/worker/gpu/attn_utils.py | 20 +- 11 files changed, 269 insertions(+), 126 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index a2a0e14be18c..1567c0eba724 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -105,7 +105,7 @@ steps: pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector_hma.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 - label: "XPU Qwen3-30B + V1 KV Connector" key: xpu-qwen3-30b-v1-kv-connector diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 3d539f97c8cb..2aae7d71f515 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -51,6 +51,7 @@ create_scheduler, create_vllm_config, make_kv_cache_config, + maybe_update_block_size, ) pytestmark = pytest.mark.cpu_test @@ -87,10 +88,11 @@ def _make_connector_with_fake_worker( hand_shake_latency=0, cycles_before_done=0, do_handshake=True ): """Create a NixlConnector with FakeNixlConnectorWorker.""" + block_size = maybe_update_block_size(32) vllm_config = create_vllm_config( - kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=32 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=block_size ) - kv_cache_config = make_kv_cache_config(block_size=32, num_blocks=2) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=2) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d326615aad8f..f79d7aa316f0 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -57,7 +57,12 @@ ) from vllm.v1.request import RequestStatus -from .utils import create_request, create_scheduler, create_vllm_config +from .utils import ( + create_request, + create_scheduler, + create_vllm_config, + maybe_update_block_size, +) @pytest.mark.parametrize( @@ -148,6 +153,7 @@ def test_region_length_validation_checks_replicated_gqa_heads( def _make_test_kv_cache_config() -> KVCacheConfig: + block_size = maybe_update_block_size(16) return KVCacheConfig( num_blocks=0, kv_cache_tensors=[], @@ -160,7 +166,7 @@ def _make_test_kv_cache_config() -> KVCacheConfig: "model.layers.1.eagle_attn", ], FullAttentionSpec( - block_size=16, + block_size=block_size, num_kv_heads=4, head_size=64, dtype=torch.float16, @@ -170,6 +176,14 @@ def _make_test_kv_cache_config() -> KVCacheConfig: ) +def _make_worker_vllm_config(kv_role: str): + return create_vllm_config( + kv_connector="MooncakeConnector", + kv_role=kv_role, + block_size=maybe_update_block_size(16), + ) + + def _make_packed_mla_kv_cache_config( num_blocks: int, *, num_groups: int = 1 ) -> KVCacheConfig: @@ -429,9 +443,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata( ): """Producer sends its PP layer shard to the matching consumer layer address.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( @@ -1205,7 +1217,7 @@ def _xfer_meta( @contextlib.contextmanager def mooncake_register_worker(kv_cache_config, *, kv_role: str = "kv_consumer"): """Build a worker whose engine registration and sender thread are patched.""" - vllm_config = create_vllm_config(kv_connector="MooncakeConnector", kv_role=kv_role) + vllm_config = _make_worker_vllm_config(kv_role) connector_mod = ( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector" ) @@ -1241,9 +1253,7 @@ async def test_receive_kv_selects_remote_pp_workers( expected_addrs: list[str], ): """Decode workers should not hard-code producer pp_rank 0.""" - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_consumer" - ) + vllm_config = _make_worker_vllm_config("kv_consumer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): decode_connector = MooncakeConnector( @@ -1314,7 +1324,7 @@ def test_resolve_need_send_accounts_for_remote_tp_fanout(): @pytest.mark.asyncio async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers(): """P4/D2 must retain KV until both D stages finish, including a no-op pull.""" - config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer") + config = _make_worker_vllm_config("kv_producer") with ( set_current_vllm_config(config), patch_worker_dependencies(), @@ -1386,9 +1396,7 @@ async def test_kv_producer(monkeypatch): Verifies memory offset calculation: ptr = base_addr + block_id * block_len. """ monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( @@ -1589,9 +1597,7 @@ async def test_kv_consumuer(monkeypatch): Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ. """ - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_consumer" - ) + vllm_config = _make_worker_vllm_config("kv_consumer") with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks: decode_connector = MooncakeConnector( @@ -1661,9 +1667,7 @@ async def test_kv_consumuer(monkeypatch): @pytest.mark.asyncio async def test_worker_get_finished_timeout(monkeypatch): """Tests the cleanup mechanism for requests.""" - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( vllm_config, @@ -1713,7 +1717,7 @@ async def test_worker_get_finished_timeout(monkeypatch): def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool): """Tests the memory registration logic with the underlying Mooncake engine.""" spec = FullAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=4, head_size=64, dtype=torch.float16, @@ -1794,7 +1798,10 @@ def test_register_noncontiguous_packed_allows_matching_pp(): known. Rejecting every local PP>1 layout at startup is too strict. """ spec = FullAttentionSpec( - block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + block_size=maybe_update_block_size(16), + num_kv_heads=4, + head_size=64, + dtype=torch.float16, ) layer_names = ["model.layers.0.self_attn", "model.layers.1.self_attn"] with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, _reg): @@ -1814,8 +1821,9 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg): worker.use_mla = True worker.transfer_topo.is_mla = True - mla_cache = torch.zeros((2, 16, 512), dtype=torch.float16) - eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16) + block_size = maybe_update_block_size(16) + mla_cache = torch.zeros((2, block_size, 512), dtype=torch.float16) + eagle_cache = torch.zeros((2, block_size, 8, 64), dtype=torch.float16) kv_caches = { "model.layers.0.mla_attn": mla_cache, "model.layers.1.eagle_attn": eagle_cache, @@ -2240,7 +2248,7 @@ def _register_sliced_packed_mla( ): """Register column slices of one MLA packed row. Returns worker fields.""" spec = MLAAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=1, head_size=64, dtype=torch.uint8, @@ -2286,7 +2294,7 @@ def test_sliced_packed_mla_registers_each_layer_view(): KVCacheGroupSpec( [_layer_name(i) for i in layer_idxs], MLAAttentionSpec( - block_size=16, + block_size=maybe_update_block_size(16), num_kv_heads=1, head_size=64, dtype=torch.uint8, @@ -2436,9 +2444,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size): remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") - vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) + vllm_config = _make_worker_vllm_config("kv_producer") with set_current_vllm_config(vllm_config), patch_worker_dependencies(): prefill_connector = MooncakeConnector( diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py index a7fb2be02365..87dafebb4c2e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py @@ -30,7 +30,12 @@ ) from .test_mooncake_connector import FakeMooncakeWrapper, patch_worker_dependencies -from .utils import create_request, create_vllm_config, make_kv_cache_config +from .utils import ( + create_request, + create_vllm_config, + make_kv_cache_config, + maybe_update_block_size, +) # --------------------------------------------------------------------------- @@ -224,12 +229,11 @@ def test_metadata_hma_block_ids(): async def test_build_transfer_params_multi_group_trimming(monkeypatch): """_build_transfer_params trims per-group blocks when local > remote.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) - kv_cache_config = make_kv_cache_config( - block_size=vllm_config.cache_config.block_size, swa_enabled=True + kv_connector="MooncakeConnector", kv_role="kv_producer", block_size=block_size ) + kv_cache_config = make_kv_cache_config(block_size=block_size, swa_enabled=True) with set_current_vllm_config(vllm_config), patch_worker_dependencies(): connector = MooncakeConnector( @@ -320,12 +324,11 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch): async def test_build_transfer_params_group_count_mismatch(monkeypatch): """_build_transfer_params reports an error when group counts differ.""" monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - kv_connector="MooncakeConnector", kv_role="kv_producer" - ) - kv_cache_config = make_kv_cache_config( - block_size=vllm_config.cache_config.block_size, swa_enabled=True + kv_connector="MooncakeConnector", kv_role="kv_producer", block_size=block_size ) + kv_cache_config = make_kv_cache_config(block_size=block_size, swa_enabled=True) with set_current_vllm_config(vllm_config), patch_worker_dependencies(): connector = MooncakeConnector( @@ -449,7 +452,7 @@ def test_request_finished_with_hma_groups(): def _make_kv_consumer_worker( swa_enabled: bool = False, kv_cache_config: KVCacheConfig | None = None ): - block_size = 16 + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", @@ -495,6 +498,7 @@ def test_worker_is_hma_required(swa_enabled, expected_is_hma): def test_worker_is_hma_required_multiple_full_attention_groups(): """Two full-attention groups with different block sizes also select request-level failure reporting, matching the scheduler.""" + block_size = maybe_update_block_size(16) kv_cache_config = KVCacheConfig( num_blocks=100, kv_cache_tensors=[], @@ -502,13 +506,19 @@ def test_worker_is_hma_required_multiple_full_attention_groups(): KVCacheGroupSpec( ["layer0"], FullAttentionSpec( - block_size=16, num_kv_heads=4, head_size=16, dtype=torch.float16 + block_size=block_size, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, ), ), KVCacheGroupSpec( ["layer1"], FullAttentionSpec( - block_size=32, num_kv_heads=4, head_size=16, dtype=torch.float16 + block_size=2 * block_size, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, ), ), ], diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 969bf65f8c54..65abd910bbcf 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -36,7 +36,7 @@ ) from .test_mooncake_connector import patch_worker_dependencies -from .utils import create_request, create_vllm_config +from .utils import create_request, create_vllm_config, maybe_update_block_size def noop_shutdown(): @@ -147,6 +147,7 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -203,6 +204,7 @@ def test_register_kv_caches_scales_attention_len_to_kernel_block(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -392,6 +394,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_consumer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -405,8 +408,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch): ) worker = connector.connector_worker - backing = torch.empty((4, 64), dtype=torch.float16) - fa_cache = backing[:2, :16] + fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec + fa_cache_width = ( + fa_spec.page_size_bytes + // torch.tensor([], dtype=torch.float16).element_size() + ) + backing = torch.empty((4, fa_cache_width), dtype=torch.float16) + fa_cache = backing[:2] gdn_cache = backing[:3] with patch.object( @@ -438,6 +446,7 @@ def test_hybrid_gdn_transfer_params_preserve_group_identity(monkeypatch): vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size @@ -567,6 +576,7 @@ def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole( vllm_config = create_vllm_config( kv_connector="MooncakeConnector", kv_role="kv_producer", + block_size=maybe_update_block_size(16), ) kv_cache_config = make_hybrid_gdn_kv_cache_config( vllm_config.cache_config.block_size diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index e576e3d0d9e6..4ac3eaf7c569 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -10,7 +10,7 @@ import pytest import torch -from tests.v1.kv_connector.unit.utils import create_vllm_config +from tests.v1.kv_connector.unit.utils import create_vllm_config, maybe_update_block_size from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory @@ -314,10 +314,13 @@ def test_multi_example_connector_consistency(): }, ) + block_size = maybe_update_block_size(16) + expected_num_blocks = 2 if block_size == 64 else 7 + expected_num_matched_tokens = (expected_num_blocks - 1) * block_size llm = LLM( model=MODEL_NAME, enforce_eager=True, - block_size=16, + block_size=block_size, gpu_memory_utilization=0.5, kv_transfer_config=kv_transfer_config, async_scheduling=False, @@ -372,7 +375,7 @@ def test_multi_example_connector_consistency(): "set_xfer_handshake_metadata_pp_aware", "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] # First three events are from initialization. Layer hooks run before the @@ -399,7 +402,7 @@ def test_multi_example_connector_consistency(): "set_xfer_handshake_metadata_pp_aware", "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] # Reset prefix cache or else we'll just get the tokens back from there. @@ -423,13 +426,14 @@ def test_multi_example_connector_consistency(): assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 96", + "update_state_after_alloc " + f"num_blocks=[{expected_num_blocks}] {expected_num_matched_tokens}", "build_connector_meta", ] assert storage2_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] @@ -457,13 +461,14 @@ def test_multi_example_connector_consistency(): assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 0", + f"update_state_after_alloc num_blocks=[{expected_num_blocks}] 0", "build_connector_meta", ] assert storage2_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", - "update_state_after_alloc num_blocks=[7] 96", + "update_state_after_alloc " + f"num_blocks=[{expected_num_blocks}] {expected_num_matched_tokens}", "build_connector_meta", ] diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d918c38fbf8c..607f5ff3cd5e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -84,6 +84,7 @@ create_scheduler, create_vllm_config, make_kv_cache_config, + maybe_update_block_size, ) @@ -418,7 +419,8 @@ def test_kv_transfer_handshake(dist_init): # Test setup, we creates a scheduler that contains a NixlConnector # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from # all workers of the instance. - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.cache_config.kv_cache_layout = "BLHNC" # in case the test runs on non-GPU machine vllm_config.kv_transfer_config.kv_buffer_device = "cpu" @@ -434,7 +436,7 @@ def test_kv_transfer_handshake(dist_init): KVCacheGroupSpec( ["layer0", "layer1", "layer2"], FullAttentionSpec( - block_size=16, + block_size=block_size, num_kv_heads=4, head_size=16, dtype=torch.float16, @@ -544,7 +546,9 @@ def __init__( **kwargs, ): if kv_cache_config is None: - kv_cache_config = make_kv_cache_config(block_size=16) + kv_cache_config = make_kv_cache_config( + block_size=args[0].cache_config.block_size + ) super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout @@ -674,7 +678,8 @@ def test_pcp_producer_exposes_dcp_shards_or_canonical_replica( """Replicated PCP is canonicalized; PCP-DCP publishes every shard.""" from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - vllm_config = create_vllm_config(kv_role="kv_producer") + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role="kv_producer", block_size=block_size) vllm_config.parallel_config.prefill_context_parallel_size = pcp_size vllm_config.parallel_config.decode_context_parallel_size = dcp_size with ( @@ -692,7 +697,7 @@ def test_pcp_producer_exposes_dcp_shards_or_canonical_replica( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) worker = connector.connector_worker @@ -746,12 +751,13 @@ def test_multi_xfer_one_engine( This test triggers the connector to load remote KV for the same `request_id`. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) request_id = "req_id" # Test worker role in decode server. - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=10) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -837,12 +843,15 @@ def test_async_load_kv( prefill_tp_size, ): """Test that NixlConnector's start_load_kv should be non-blocking.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.parallel_config.tensor_parallel_size = decode_tp_size # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id @@ -899,12 +908,13 @@ def test_prefill_tp_size_greater_than_decode_tp_size( lambda: local_tp_size, ) - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # num_blocks=1 keeps the fake handshake's region_num_blocks consistent # with the minimal single-block registration below (num_descs == 1); # _fa_desc_replicated cross-checks the two. - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=1) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=1) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -972,16 +982,21 @@ def test_prefill_tp_size_greater_than_decode_tp_size_mla( """Verify remote TP > local TP handshake succeeds with different remote configurations for an MLA model. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) d_tp_size = 1 p_tp_size = 2 # Build two separate connectors/workers to emulate P TP=2 ranks. conn_p0 = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) conn_p1 = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) conn_p0.connector_worker = FakeNixlConnectorWorker( vllm_config, conn_p0.engine_id, hand_shake_latency=0 @@ -1075,11 +1090,14 @@ def test_concurrent_load_kv( dist_init, ): """Test that multiple start_load_kv calls should occur concurrently.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id @@ -1143,8 +1161,10 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( - enable_permute_local_kv=enable_permute_local_kv + enable_permute_local_kv=enable_permute_local_kv, + block_size=block_size, ) # Mock TP world size to 2 to force heterogeneous TP when @@ -1155,7 +1175,9 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( ): # Initialize connector and worker (with fake NIXL wrapper) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -1205,7 +1227,10 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( """Verify that adding a remote agent fails if kv_cache_layout differs. This test is only relevant for heterogeneous TP. """ - vllm_config = create_vllm_config(enable_permute_local_kv=True) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + enable_permute_local_kv=True, block_size=block_size + ) # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 @@ -1215,7 +1240,9 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( ): # Initialize connector and worker (with fake NIXL wrapper) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -1261,7 +1288,9 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices( self, default_vllm_config, dist_init ): worker = FakeNixlConnectorWorker( - create_vllm_config(), "engine", hand_shake_latency=0 + create_vllm_config(block_size=maybe_update_block_size(16)), + "engine", + hand_shake_latency=0, ) remote_block_len = 2048 @@ -1316,13 +1345,16 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): heterogeneous TP must NOT raise (previously a NotImplementedError), and the per-region gate must still reject a wrong block_len. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1409,7 +1441,8 @@ def test_handshake_validates_gqa_replicated_block_len( expected local_block_len * tp_ratio and rejected the valid handshake. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 @@ -1418,7 +1451,7 @@ def test_handshake_validates_gqa_replicated_block_len( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1465,7 +1498,8 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( block_lens when GQA replication is NOT in effect (32 KV heads, D_TP=4, P_TP=2: head_ratio=4, both sides have >1 head/rank). """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 @@ -1474,7 +1508,7 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1522,11 +1556,14 @@ def test_handshake_rejects_wrong_block_len_without_gqa_replication( ) def test_kv_connector_stats(default_vllm_config, dist_init): """Test that KV transfer stats are properly recorded and retrieved.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Test worker role in decode server. connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1602,9 +1639,12 @@ def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist remote read pulls another request's data (silent accuracy corruption). The worker must anchor the remaining TTL to its own clock. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -1793,9 +1833,12 @@ def counter_value(name: str) -> float: FakeNixlWrapper, ) def test_notification_after_expiry_is_counted(default_vllm_config, dist_init): - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -2323,7 +2366,10 @@ def test_register_kv_caches( 2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing block layout info """ - vllm_config = create_vllm_config(attention_backend=attn_backend) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + attention_backend=attn_backend, block_size=block_size + ) vllm_config.cache_config.kv_cache_layout = layout # Import the appropriate backend based on the parameter @@ -2348,7 +2394,6 @@ def test_register_kv_caches( patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends, ): mock_get_attn_backends.return_value = [backend_cls] - block_size = 16 num_blocks = 8 num_heads = 4 head_size = 16 @@ -2472,7 +2517,7 @@ def test_register_kv_caches( } assert len(blocks_data) == num_blocks * len(base_addrs) - assert connector.connector_worker.block_size == 16 + assert connector.connector_worker.block_size == block_size def test_register_packed_dsv4_mla_cache_as_single_region( @@ -2573,7 +2618,8 @@ def test_kv_buffer_to_nixl_memory_types( """Test that register_kv_caches() passes the correct memory types from the config to the nixl_wrapper. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) # Override the default memory types in the config vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( @@ -2603,7 +2649,9 @@ def test_kv_buffer_to_nixl_memory_types( ): # noqa: E501 # Create connector and replace its worker with a fake one for isolation connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) # Verify get_reg_descs was called with the correct memory_type @@ -2617,17 +2665,18 @@ def test_kv_buffer_to_nixl_memory_types( ) def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): """Test that shutdown() properly cleans up all resources.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) scheduler = NixlConnectorScheduler( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) worker = NixlConnectorWorker( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) nixl_wrapper = worker.nixl_wrapper @@ -2643,7 +2692,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, block_size): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: @@ -2688,13 +2737,15 @@ def _setup_worker_with_remote_engine( engine_ttl: float = 10.0, ) -> tuple[Any, str]: """Create a worker with one remote engine registered.""" + block_size = maybe_update_block_size(16) vllm_config = create_vllm_config( + block_size=block_size, kv_connector_extra_config={"engine_ttl": engine_ttl}, ) worker = NixlConnectorWorker( vllm_config, vllm_config.kv_transfer_config.engine_id, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) engine_id = "remote-engine-1" @@ -2857,7 +2908,11 @@ def setup(self, monkeypatch): from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - config = create_vllm_config(kv_connector_extra_config={"engine_ttl": 0}) + block_size = maybe_update_block_size(16) + + config = create_vllm_config( + kv_connector_extra_config={"engine_ttl": 0}, block_size=block_size + ) config.kv_transfer_config.kv_buffer_device = "cpu" platform = SimpleNamespace( device_type="cpu", @@ -3258,12 +3313,15 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_ the real scheduler, then simulate an abort (request not in next scheduler iteration) and verify the worker no longer tracks it as in-batch. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) scheduler = create_scheduler(vllm_config) # KVConnector Worker in P connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3398,9 +3456,12 @@ def test_empty_recv_is_reported_only_when_awaited( A failed notification changes neither: the KV is local either way, and the producer frees its own blocks on a timeout. """ - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -3518,12 +3579,13 @@ def test_transfer_failure_logging( """ import logging - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16, swa_enabled=enable_hma), + make_kv_cache_config(block_size=block_size, swa_enabled=enable_hma), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -3648,10 +3710,13 @@ def emit(self, record): ) def test_handshake_failure_returns_finished(default_vllm_config, dist_init): """Test that handshake failures mark blocks invalid and return via get_finished.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0.1 @@ -3704,10 +3769,13 @@ def test_transfer_setup_failure_returns_finished( default_vllm_config, dist_init, is_hma ): """Setup failures report the request; only non-HMA reports block IDs.""" - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3776,7 +3844,9 @@ def release_xfer_handle(self, handle: int) -> None: def _make_split_read_connector(vllm_config, request_id, states): """Seed a request with scripted in-flight xfer handles + recv metadata.""" connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=vllm_config.cache_config.block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -3818,7 +3888,7 @@ def test_split_read_failure_defers_report_until_last_handle( request_id = "split_read_partial_failure" err_handle, live_handle = 11, 22 connector, worker, wrapper = _make_split_read_connector( - create_vllm_config(), + create_vllm_config(block_size=maybe_update_block_size(16)), request_id, {err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]}, ) @@ -3894,10 +3964,15 @@ def test_failed_request_skips_kv_postprocessing( # Use enable_permute_local_kv=True so that # post_process_device_kv_on_receive would be called on the success path, # making the assertion meaningful (not trivially true). - vllm_config = create_vllm_config(enable_permute_local_kv=True) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config( + enable_permute_local_kv=True, block_size=block_size + ) connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -4220,14 +4295,15 @@ def test_compatibility_hash_validation( enforce_handshake_compat: whether to enforce compatibility checking """ + block_size = maybe_update_block_size(16) local_vllm_config = create_vllm_config( model="facebook/opt-125m", - block_size=16, + block_size=block_size, kv_connector_extra_config={ "enforce_handshake_compat": enforce_handshake_compat }, ) - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2) + kv_cache_config = make_kv_cache_config(block_size=block_size, num_blocks=2) decode_connector = NixlConnector( local_vllm_config, KVConnectorRole.WORKER, kv_cache_config ) @@ -4253,7 +4329,7 @@ def test_compatibility_hash_validation( remote_config_params: dict[str, Any] = { "model": "facebook/opt-125m", - "block_size": 16, + "block_size": block_size, **config_overrides, } remote_vllm_config = create_vllm_config(**remote_config_params) @@ -4276,7 +4352,7 @@ def test_compatibility_hash_validation( decode_worker.backend_name, ) - prefill_block_size = config_overrides.get("block_size", 16) + prefill_block_size = remote_vllm_config.cache_config.block_size prefill_block_lens = [4096 * prefill_block_size] prefill_metadata = NixlAgentMetadata( engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, @@ -4352,12 +4428,15 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) - NixlHandshakePayload decoder - NixlAgentMetadata decoder """ + block_size = maybe_update_block_size(16) local_vllm_config = create_vllm_config( model="facebook/opt-125m", - block_size=16, + block_size=block_size, ) decode_connector = NixlConnector( - local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + local_vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) decode_worker = decode_connector.connector_worker @@ -4441,11 +4520,14 @@ def test_mla_broadcast_notif_uses_remote_request_id( decode_tp_size = 1 prefill_tp_size = 4 - vllm_config = create_vllm_config() + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) vllm_config.parallel_config.tensor_parallel_size = decode_tp_size connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), ) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, connector.engine_id, hand_shake_latency=0 @@ -4550,7 +4632,8 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): _print_warning_once.cache_clear() - vllm_config = create_vllm_config(kv_role="kv_both") + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role="kv_both", block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger" @@ -4559,7 +4642,7 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) mock_logger.warning_once.assert_called_once() @@ -4575,14 +4658,15 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init): def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init): """kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning.""" for role in ("kv_consumer", "kv_producer"): - vllm_config = create_vllm_config(kv_role=role) + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(kv_role=role, block_size=block_size) with patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger" ) as mock_logger: NixlConnector( vllm_config, KVConnectorRole.WORKER, - make_kv_cache_config(block_size=16), + make_kv_cache_config(block_size=block_size), ) ( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index c665ce0c5b88..469e0f5d8294 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -39,6 +39,7 @@ create_vllm_config, make_kv_cache_config, make_nixl_scheduler, + maybe_update_block_size, ) @@ -1027,7 +1028,7 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): kv_connector="NixlConnector", kv_role="kv_consumer", ) - block_size = 16 + block_size = maybe_update_block_size(16) llm_kwargs = { "model": model_name, "enforce_eager": True, diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 28bfda419a76..59b2e7ad0349 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -31,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa ExampleConnector, ) +from vllm.platforms import current_platform from vllm.utils.hashing import sha256 from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash @@ -59,6 +60,10 @@ def create_mock_connector() -> Mock: return connector +def maybe_update_block_size(block_size: int) -> int: + return 64 if current_platform.is_xpu() else block_size + + def assert_scheduler_empty(scheduler: Scheduler): """Confirm the scheduler is "empty" - i.e. no leaks.""" # Scheduler Metadata. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 6179b27bb6db..4a336586acd8 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -355,6 +355,8 @@ def _get_fa4_hd256_block_size( def get_supported_kernel_block_sizes( cls, kv_cache_spec: KVCacheSpec | None = None ) -> list[int | MultipleOf]: + if current_platform.is_xpu(): + return [MultipleOf(64)] if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec): return [block_size] if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec): diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index d1feaef031d9..8964c660a1d9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -14,9 +14,11 @@ get_layers_from_vllm_config, ) from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.platforms import current_platform from vllm.v1.attention.backend import ( AttentionCGSupport, CommonAttentionMetadata, @@ -45,6 +47,8 @@ prepare_kernel_block_sizes, ) +logger = init_logger(__name__) + if TYPE_CHECKING: from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -263,10 +267,24 @@ def init_attn_backend( # by all backends within that group. kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) + if current_platform.is_xpu(): + for kv_cache_group_id, selected_kernel_block_size in enumerate( + kernel_block_sizes + ): + if selected_kernel_block_size % 64 != 0: + for attn_group in attn_groups[kv_cache_group_id]: + logger.warning_once( + "On XPU platform, current selected kernel block size is " + "not fully optimized: kernel_block_size=%s, backend=%s, " + "kv_cache_spec=%s", + selected_kernel_block_size, + attn_group.backend, + repr(attn_group.kv_cache_spec), + ) # Phase 3: create metadata builders and determine cudagraph support. attn_backend_workspace: torch.Tensor | None = None for kv_cache_group_id, groups in enumerate(attn_groups): - kernel_block_size = None + kernel_block_size: int | None = None if kv_cache_group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[kv_cache_group_id] for group in groups: From bfa7e5b7afe40ebe48288801a2baf5c1135cd881 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 12 Oct 2026 10:39:40 +0800 Subject: [PATCH 382/952] [XPU][CI] Enable XPU graph in graph tests (#60784) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- .../v1/cudagraph/test_breakable_cudagraph.py | 46 ++++++++++++------- tests/v1/cudagraph/test_cudagraph_dispatch.py | 15 ++++-- tests/v1/cudagraph/test_cudagraph_mode.py | 7 +++ tests/v1/cudagraph/test_encoder_cudagraph.py | 27 ++++++----- 4 files changed, 63 insertions(+), 32 deletions(-) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index b0b49a9567b1..1b9a79a6ef2c 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -11,6 +11,18 @@ import pytest import torch +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + @pytest.fixture(autouse=True) def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch): @@ -96,8 +108,8 @@ def cuda_capture_stream(): stream, so all capture-using tests need to run under ``torch.cuda.stream(...)`` for a separate stream. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA required") + if not HAS_DEVICE_GRAPH: + pytest.skip("CUDA or XPU graph required") from vllm.utils.torch_utils import _current_stream_tls prev_stream = getattr(_current_stream_tls, "value", None) @@ -200,7 +212,7 @@ def worker(): def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) cap = BreakableCUDAGraphCapture() with cap: x.add_(1.0) @@ -212,7 +224,7 @@ def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): def test_add_eager_creates_alternating_graph_eager_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) counter = {"eager_calls": 0} def eager_step(): @@ -253,7 +265,7 @@ def test_capture_replay_matches_eager_simple(cuda_capture_stream): """ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(8, device="cuda") + x = torch.zeros(8, device=DEVICE_TYPE) log: list[str] = [] def eager_break_op(): @@ -269,21 +281,21 @@ def eager_break_op(): # Capture-time: graph kernels were recorded only; eager segment ran # once on x == 0, leaving x == 0. torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(8, device="cuda")) + assert torch.equal(x, torch.zeros(8, device=DEVICE_TYPE)) assert log == ["eager"] # Replay with a fresh input: 10 -> 11 -> 22 -> 27. x.fill_(10.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 27.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 27.0, device=DEVICE_TYPE)) assert log == ["eager", "eager"] # Replay again with another input: 100 -> 101 -> 202 -> 207. x.fill_(100.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 207.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 207.0, device=DEVICE_TYPE)) assert log == ["eager", "eager", "eager"] @@ -301,13 +313,13 @@ def attention_like(t: torch.Tensor) -> None: # In-place double; stands in for "real" attention work. t.mul_(2.0) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) # Outside capture: decorator should just call through. x.fill_(3.0) attention_like(x) torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 6.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 6.0, device=DEVICE_TYPE)) # Inside capture: decorator should split the graph. Only the eager # segment actually mutates state during capture. @@ -318,7 +330,7 @@ def attention_like(t: torch.Tensor) -> None: attention_like(x) # eager: x *= 2 (on x == 0, no-op) x.add_(1.0) # recorded torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(4, device="cuda")) + assert torch.equal(x, torch.zeros(4, device=DEVICE_TYPE)) # 2 graph segments + 1 eager segment, ordered G E G; the arithmetic # equivalence check below verifies the ordering. assert len(cap.segments) == 3 @@ -329,7 +341,7 @@ def attention_like(t: torch.Tensor) -> None: x.fill_(2.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 15.0, device=DEVICE_TYPE)) @pytest.mark.parametrize("keyword", [False, True]) @@ -352,7 +364,7 @@ def consume(x, data_out, scale_out): data_out.copy_(x.data) scale_out.copy_(x.scale) - source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16) + source = torch.ones((128, 128), device=DEVICE_TYPE, dtype=torch.bfloat16) data_out = source.to(torch.float8_e4m3fn) scale_out = source[:, ::32].to(torch.uint8) output = torch.empty_like(source) @@ -391,7 +403,7 @@ def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): ) from vllm.utils.multi_stream_utils import maybe_execute_in_parallel - x = torch.zeros(1024, device="cuda") + x = torch.zeros(1024, device=DEVICE_TYPE) output = torch.empty_like(x) aux_stream = torch.cuda.Stream() main_event = torch.cuda.Event() @@ -440,7 +452,7 @@ def test_replay_invokes_eager_segments_in_order(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture log: list[str] = [] - x = torch.zeros(1, device="cuda") + x = torch.zeros(1, device=DEVICE_TYPE) def make_eager(name): def step(): @@ -502,7 +514,7 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream): eager_break_during_capture, ) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) inner_calls = 0 @eager_break_during_capture @@ -533,7 +545,7 @@ def outer_op(t: torch.Tensor) -> None: cap.replay() torch.accelerator.synchronize() # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113 - assert torch.equal(x, torch.full((4,), 113.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 113.0, device=DEVICE_TYPE)) assert inner_calls == 2 # replay invokes the outer's lambda again diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 40874be5bd59..b0a4587ddfc5 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -24,6 +24,14 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass # Helper MLP for testing @@ -50,6 +58,7 @@ def _create_vllm_config( mock_config.parallel_config = ParallelConfig() mock_config.speculative_config = None # No speculative decoding mock_config.num_speculative_tokens = 0 + mock_config.use_cumem_cudagraph_pool = False if not lora_config: mock_config.lora_config = None else: @@ -268,7 +277,7 @@ def test_get_capture_descs_empty_when_not_initialized(self): assert dispatcher.get_capture_descs() == [] -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") class TestCUDAGraphWrapper: def setup_method(self): self.vllm_config = _create_vllm_config(CompilationConfig()) @@ -415,7 +424,7 @@ def _run_and_monitor_call( @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_capture_replay_bypass_logic(): comp_config = CompilationConfig( mode=CompilationMode.VLLM_COMPILE, @@ -484,7 +493,7 @@ def test_capture_replay_bypass_logic(): @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_nested_wrappers(): """Tests a scenario with a PIECEWISE wrapper inside a FULL one.""" comp_config = CompilationConfig( diff --git a/tests/v1/cudagraph/test_cudagraph_mode.py b/tests/v1/cudagraph/test_cudagraph_mode.py index 0147f90df534..f35ef4d82974 100644 --- a/tests/v1/cudagraph/test_cudagraph_mode.py +++ b/tests/v1/cudagraph/test_cudagraph_mode.py @@ -19,6 +19,13 @@ ("TritonAttn", "FULL", True), ("TritonAttn", "FULL_AND_PIECEWISE", True), ] +elif current_platform.is_xpu(): + combo_cases_1 = [ + ("FA2", "FULL", True), + ("FA2", "FULL_AND_PIECEWISE", True), + ("TritonAttn", "FULL", True), + ("TritonAttn", "FULL_AND_PIECEWISE", True), + ] else: combo_cases_1 = [ ("FA3", "FULL", True), diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 348944387886..2799a3ef4dd6 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -32,6 +32,15 @@ EncoderItemSpec, ) +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -662,12 +671,10 @@ def _make_video_mm_kwargs( # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphCaptureReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTModel().to(self.device).half() self.mgr = _make_manager_for_gpu( @@ -758,9 +765,7 @@ def test_chunking_when_images_exceed_max_batch(self): # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") @pytest.mark.usefixtures("dist_init", "workspace_init") @pytest.mark.parametrize("profile_only", [False, True]) @torch.inference_mode() @@ -776,7 +781,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner from vllm.v1.worker.workspace import lock_workspace - device = torch.device("cuda:0") + device = torch.device(current_platform.device_type, 0) dtype = torch.float16 model = SimpleMockViTModel().to(device).half() manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype) @@ -1047,12 +1052,10 @@ def test_video_model_returns_video_for_video_kwargs(self): _VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphVideoReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTVideoModel().to(self.device).half() self.mgr = _make_manager_for_gpu( From d4e92ccc37e67048a957aa8e08dc5a2bc665f422 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Mon, 12 Oct 2026 11:58:15 +0800 Subject: [PATCH 383/952] [Bugfix][Frontend] Honor endpoint plugin route overrides (#58839) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex Co-authored-by: Cyrus Leung --- docs/design/endpoint_plugins.md | 4 +- tests/plugins_tests/test_endpoint_plugins.py | 160 ++++++++++++++++++- vllm/entrypoints/launchers/app.py | 4 +- vllm/plugins/endpoint_plugins/interface.py | 51 +++++- 4 files changed, 210 insertions(+), 9 deletions(-) diff --git a/docs/design/endpoint_plugins.md b/docs/design/endpoint_plugins.md index 98b755788e2e..361bd7099b46 100644 --- a/docs/design/endpoint_plugins.md +++ b/docs/design/endpoint_plugins.md @@ -126,7 +126,9 @@ Do not expect a single endpoint plugin to also mutate engine/worker state. If yo ## Path-prefix convention -There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). A plugin's `attach_router` can register a path that collides with a core route and routes attached later win. To avoid surprising operators: +Plugin hooks run after the core routers have been registered. HTTP routes added by plugins replace earlier routes with the same path template and HTTP method, including their OpenAPI description. Methods not overridden by the plugin remain available. When multiple plugins register the same operation, the last registered route takes precedence. Different path templates are not treated as equivalent. + +There is currently no rejection of route conflicts (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). To avoid surprising operators: - Namespace your routes under a distinct prefix, e.g. `/plugins//...`, rather than reusing `/v1/...` or other core prefixes - Only register routes under a core prefix (like the worked example's `/v1/admin/scheduler_config`) if you specifically intend to override or extend existing behavior and document that clearly for operators allowlisting your plugin diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py index 2d921539004f..864c915fe083 100644 --- a/tests/plugins_tests/test_endpoint_plugins.py +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -14,7 +14,7 @@ import httpx import pytest -from fastapi import FastAPI +from fastapi import APIRouter, FastAPI from vllm_add_dummy_endpoint_plugin import DummyAdminEndpointPlugin from vllm.entrypoints.launchers.app import build_app @@ -38,6 +38,34 @@ def __init__(self): raise RuntimeError("boom") +class _ShadowEndpointPlugin: + name = "shadow_endpoint_plugin" + required_tasks = None + + def __init__(self, path: str, include_in_schema: bool): + self.path = path + self.include_in_schema = include_in_schema + + def attach_router(self, app: FastAPI) -> None: + assert str(app.url_path_for("show_version")) == "/version" + app.openapi() + + @app.get(self.path, status_code=209, include_in_schema=self.include_in_schema) + async def plugin_version(plugin_arg: str) -> dict[str, str]: + return {"served_by": plugin_arg} + + async def init_state(self, engine_client, state, args) -> None: + pass + + +class _RouterEndpointPlugin(DummyAdminEndpointPlugin): + def __init__(self, router: APIRouter): + self.router = router + + def attach_router(self, app: FastAPI) -> None: + app.include_router(self.router) + + class _FakeEngineClient: """Minimal stand in exercising `collective_rpc`. Not a real engine.""" @@ -162,6 +190,136 @@ def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch) assert app.state.endpoint_plugins == [] +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/version", "/metrics"]) +@pytest.mark.parametrize("include_in_schema", [True, False]) +async def test_endpoint_plugin_override_preserves_hook_and_schema( + monkeypatch: pytest.MonkeyPatch, path: str, include_in_schema: bool +): + """Hooks see core routes; overrides win over both HTTP routes and mounts.""" + args = _build_args() + monkeypatch.setenv("VLLM_PLUGINS", "shadow_endpoint_plugin") + monkeypatch.setattr( + "vllm.plugins.load_plugins_by_group", + _fake_loader( + { + "shadow_endpoint_plugin": lambda: _ShadowEndpointPlugin( + path, include_in_schema + ) + } + ), + ) + app = build_app(args, supported_tasks=()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get(path, params={"plugin_arg": "plugin"}) + schema = (await client.get("/openapi.json")).json() + + assert response.status_code == 209 + assert response.json() == {"served_by": "plugin"} + if include_in_schema: + operation = schema["paths"][path]["get"] + assert "209" in operation["responses"] + assert "200" not in operation["responses"] + assert operation["parameters"][0]["name"] == "plugin_arg" + assert operation["parameters"][0]["required"] is True + response_schema = operation["responses"]["209"]["content"]["application/json"][ + "schema" + ] + assert response_schema["additionalProperties"] == {"type": "string"} + else: + assert path not in schema["paths"] + + +@pytest.mark.asyncio +async def test_endpoint_plugin_override_preserves_other_methods( + monkeypatch: pytest.MonkeyPatch, +): + """Overriding GET preserves POST's validation and the shared core route.""" + core = APIRouter() + + @core.api_route("/shared", methods=["GET", "POST"], status_code=201) + async def original(value: int) -> dict[str, int]: + return {"value": value} + + plugin = APIRouter() + + @plugin.get("/shared", status_code=202) + async def replacement(): + return {"served_by": "plugin"} + + app = FastAPI() + app.include_router(core) + original_route = app.routes[-1] + monkeypatch.setattr( + "vllm.plugins.load_endpoint_plugins", + lambda supported_tasks: [_RouterEndpointPlugin(plugin)], + ) + attach_endpoint_plugins(app, ()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/shared") + post_response = await client.post("/shared", params={"value": 7}) + invalid_response = await client.post("/shared", params={"value": "invalid"}) + + assert response.status_code == 202 + assert response.json() == {"served_by": "plugin"} + assert post_response.status_code == 201 + assert post_response.json() == {"value": 7} + assert invalid_response.status_code == 422 + assert original_route.methods == {"GET", "POST"} + assert core.routes[0].methods == {"GET", "POST"} + operations = app.openapi()["paths"]["/shared"] + assert "202" in operations["get"]["responses"] + assert "201" in operations["post"]["responses"] + + +@pytest.mark.asyncio +async def test_later_endpoint_plugin_overrides_same_operation( + monkeypatch: pytest.MonkeyPatch, +): + """The last registered override wins while unrelated plugin routes survive.""" + first = APIRouter() + + @first.get("/shared", status_code=201) + @first.get("/plugins/first") + async def first_handler(): + return {"served_by": "first"} + + second = APIRouter() + + @second.get("/shared", status_code=202) + async def second_handler(): + return {"served_by": "second"} + + monkeypatch.setattr( + "vllm.plugins.load_endpoint_plugins", + lambda supported_tasks: [ + _RouterEndpointPlugin(first), + _RouterEndpointPlugin(second), + ], + ) + app = FastAPI() + attach_endpoint_plugins(app, ()) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/shared") + unrelated = await client.get("/plugins/first") + unsupported = await client.delete("/shared") + + assert response.status_code == 202 + assert response.json() == {"served_by": "second"} + assert unrelated.status_code == 200 + assert unrelated.json() == {"served_by": "first"} + assert unsupported.status_code == 405 + responses = app.openapi()["paths"]["/shared"]["get"]["responses"] + assert "202" in responses + assert "201" not in responses + + @pytest.mark.asyncio async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch): """`init_app_state` callers that never ran `build_app` (e.g. diff --git a/vllm/entrypoints/launchers/app.py b/vllm/entrypoints/launchers/app.py index ddea70781051..59c4fc37a07a 100644 --- a/vllm/entrypoints/launchers/app.py +++ b/vllm/entrypoints/launchers/app.py @@ -44,8 +44,8 @@ def build_app( register_api_routers(args, app, supported_tasks, model_config) - # Endpoint plugins are attached last so their routes are registered after all core - # routers. This runs even for the CPU only render server. A plugin eligible for + # Endpoint plugin hooks run after core routers so plugins can inspect them. + # This runs even for the CPU only render server. A plugin eligible for # the `render` task still gets its routes registered. It receives # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). attach_endpoint_plugins(app, supported_tasks) diff --git a/vllm/plugins/endpoint_plugins/interface.py b/vllm/plugins/endpoint_plugins/interface.py index 00451217d2f2..0d847895cfa2 100644 --- a/vllm/plugins/endpoint_plugins/interface.py +++ b/vllm/plugins/endpoint_plugins/interface.py @@ -29,16 +29,22 @@ """ from argparse import Namespace +from copy import copy from typing import TYPE_CHECKING, Protocol, runtime_checkable from fastapi import FastAPI +from fastapi.routing import APIRoute from starlette.datastructures import State +from starlette.routing import BaseRoute from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.tasks import SupportedTask +logger = init_logger(__name__) + @runtime_checkable class EndpointPlugin(Protocol): @@ -64,8 +70,8 @@ def attach_router(self, app: FastAPI) -> None: """Register this plugin's routes on `app`. Called once during `build_app()` after all core routers have been - attached. Routes attached here can shadow core routes with the same - path. There is currently no conflict enforcement (see RFC #46565 follow ups). + attached. HTTP routes registered here replace earlier routes with the + same path template and method. Other methods remain available. """ ... @@ -93,15 +99,50 @@ def attach_endpoint_plugins( ) -> None: """Phase A of endpoint plugin wiring: discover, gate and attach routes. - Attached last after all core routers. This is so endpoint plugin routes can - shadow core routes with the same path (see `EndpointPlugin.attach_router` - docstring). No-ops when no plugins are discovered/allowlisted. + Hooks run after all core routers have been attached. Plugin HTTP routes + replace earlier operations with the same path template and method in both + routing and OpenAPI. No-ops when no plugins are discovered/allowlisted. """ from vllm.plugins import load_endpoint_plugins endpoint_plugins = load_endpoint_plugins(supported_tasks) + existing_routes = {id(route): route for route in app.routes} for plugin in endpoint_plugins: plugin.attach_router(app) + + routes: list[BaseRoute] = [] + for route in app.routes: + if id(route) in existing_routes or not isinstance(route, APIRoute): + routes.append(route) + continue + + insert_at = len(routes) + retained: list[BaseRoute] = [] + for previous in routes: + if ( + isinstance(previous, APIRoute) + and previous.path == route.path + and previous.methods & route.methods + ): + logger.info( + "Endpoint plugin overrides %s %s", + ", ".join(sorted(previous.methods & route.methods)), + route.path, + ) + # Keep replacements ahead of mounts, preserving other methods. + insert_at = min(insert_at, len(retained)) + methods = previous.methods - route.methods + if not methods: + continue + previous = copy(previous) + previous.methods = methods + retained.append(previous) + retained.insert(insert_at, route) + routes = retained + + app.router.routes[:] = routes + if endpoint_plugins: + app.openapi_schema = None app.state.endpoint_plugins = endpoint_plugins From fe4df3df0eb703862c686b256e9fe44fa94c8945 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 12 Oct 2026 13:41:52 +0800 Subject: [PATCH 384/952] [Bugfix][GLM-5.3-Flash] Restore MTP index sharing across draft iterations (#61079) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 --- .../test_speculative_draft_hf_overrides.py | 20 ++++++++++++++++++- vllm/config/speculative.py | 11 ++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index e7014faa47d0..39ebd7924bf2 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch import pytest -from transformers import PreTrainedConfig +from transformers import Glm5NextConfig, PreTrainedConfig from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig @@ -270,3 +270,21 @@ def test_mtp_index_share_override( speculative_config.draft_model_config.hf_config.index_share_for_mtp_iteration is expected ) + + +@pytest.mark.cpu_test +def test_glm5_next_mtp_override_promotes_index_share(): + """The flag lives on text_config, but the MTP proposers read the draft config.""" + config = Glm5NextConfig( + architectures=["Glm5NextForConditionalGeneration"], + text_config={ + "num_nextn_predict_layers": 1, + "index_share_for_mtp_iteration": True, + }, + ) + + draft_config = SpeculativeConfig.hf_config_override(config) + + assert draft_config.architectures == ["Glm5NextMTPModel"] + assert draft_config.n_predict == 1 + assert draft_config.index_share_for_mtp_iteration is True diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 4c61ead82dab..b916d7da2bed 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1058,9 +1058,16 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: ) if hf_config.model_type == "glm5_next": hf_config.model_type = "glm5_next_mtp" - n_predict = hf_config.get_text_config().num_nextn_predict_layers + text_config = hf_config.get_text_config() + share_mtp_indices = getattr( + text_config, "index_share_for_mtp_iteration", False + ) hf_config.update( - {"n_predict": n_predict, "architectures": ["Glm5NextMTPModel"]} + { + "n_predict": text_config.num_nextn_predict_layers, + "architectures": ["Glm5NextMTPModel"], + "index_share_for_mtp_iteration": share_mtp_indices, + } ) return hf_config From 74f4404f58c0fb16821b34c16343486eeaf164a5 Mon Sep 17 00:00:00 2001 From: Hubert <163992334+hubert-marek@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:44:12 +0200 Subject: [PATCH 385/952] [Bugfix][Distributed] Own SHM out-of-band buffers before releasing read slots (#61000) Signed-off-by: Hubert Marek <163992334+hubert-marek@users.noreply.github.com> Signed-off-by: Ao Shen Co-authored-by: Codex Co-authored-by: Ao Shen --- tests/distributed/test_shm_broadcast.py | 105 +++++++++++++----- .../device_communicators/shm_broadcast.py | 15 ++- 2 files changed, 87 insertions(+), 33 deletions(-) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 4984a2480a61..0e6efd10429b 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -57,6 +57,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): ctx = mp_context or mp number_of_processes = world_size processes = [] + master_port = str(get_open_port()) for i in range(number_of_processes): env = {} env["RANK"] = str(i) @@ -64,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): env["WORLD_SIZE"] = str(number_of_processes) env["LOCAL_WORLD_SIZE"] = str(number_of_processes) env["MASTER_ADDR"] = "localhost" - env["MASTER_PORT"] = "12345" + env["MASTER_PORT"] = master_port p = ctx.Process(target=fn, args=(env,)) processes.append(p) p.start() @@ -88,11 +89,14 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): break time.sleep(0.1) # Check every 100ms - # Check for timeout if no failures detected yet + # Reap every child; failures can occur after the last polling iteration. + failed_processes.clear() for i, p in enumerate(processes): if p.is_alive(): p.kill() - p.join() + p.join() + if p.exitcode != 0: + failed_processes.append((i, p.exitcode)) # Report failures if failed_processes: @@ -102,6 +106,11 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None): raise AssertionError(error_msg) +def test_distributed_run_reports_timeout(): + with pytest.raises(AssertionError, match="Distributed test failed"): + distributed_run(lambda _: time.sleep(60), world_size=1, timeout=0) + + def worker_fn_wrapper(fn): # `mp.Process` cannot accept environment variables directly # so we need to pass the environment variables as arguments @@ -381,14 +390,10 @@ def worker_fn_tensor_broadcast(): # Both ranks construct the identical reference payload. torch.manual_seed(42) payload = { - # 2MiB: rides the shm ring as an out-of-band buffer (the receiving - # side must copy out of the reusable ring chunk). "mid": torch.randn(1024, 512), - # 16MiB > max_chunk_bytes: overflows to the zmq socket (the - # receiving side aliases the zmq.Frame zero-copy). - "big": torch.randn(4096, 2048, dtype=torch.bfloat16), "nested": ["plain", 123, {"inner": torch.arange(5)}], } + big = torch.randn(4096, 2048, dtype=torch.bfloat16) if rank == writer_rank: with mock.patch( @@ -396,28 +401,23 @@ def worker_fn_tensor_broadcast(): wraps=_reduce_tensor, ) as wrapped_reduce: message_queue.enqueue(payload) + # Overflow is per message, not per tensor: send this separately. + message_queue.enqueue(big) assert wrapped_reduce.call_count == 3 - # Cycle the ring (max_chunks=4) several times over so that aliased - # ring chunks would be overwritten. - for i in range(16): - message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) else: - received = message_queue.dequeue(timeout=30) - for key in ("mid", "big"): - assert torch.equal(received[key], payload[key]), key - assert received[key].dtype == payload[key].dtype, key + with mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv: + received = message_queue.dequeue(timeout=30) + recv.assert_not_called() + received_big = message_queue.dequeue(timeout=30) + recv.assert_called_once() + assert torch.equal(received["mid"], payload["mid"]) + assert received["mid"].dtype == payload["mid"].dtype + assert torch.equal(received_big, big) + assert received_big.dtype == big.dtype assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) - - snapshot = received["mid"].clone() - for i in range(16): - junk = message_queue.dequeue(timeout=30) - assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) - # Tensors received via the shm ring must not alias chunk memory - # that the writer has reused for subsequent messages. - assert torch.equal(received["mid"], snapshot) # Rebuilt tensors must be writable, like regular tensors. received["mid"] += 1.0 - received["big"][0, 0] = 1.0 + received_big[0, 0] = 1.0 dist.barrier() print(f"tensor broadcast passed the test! Rank {rank}") @@ -427,6 +427,61 @@ def test_tensor_broadcast(): distributed_run(worker_fn_tensor_broadcast, 2) +@pytest.mark.parametrize("kind", ["tensor", "numpy", "numpy_readonly"]) +@pytest.mark.parametrize( + "nbytes, overflow, num_oob", + [ + pytest.param(1024 * 1024 - 2, False, 0, id="inline"), + pytest.param(1024 * 1024, False, 1, id="shm_oob"), + pytest.param(2 * 1024 * 1024, True, 1, id="overflow"), + ], +) +def test_broadcast_retains_values_after_ring_reuse(kind, nbytes, overflow, num_oob): + """Received data belongs to its consumer, not the next queue message.""" + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + max_chunk_bytes=2 * 1024 * 1024, + max_chunks=1, + connect_ip="127.0.0.1", + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + writer.wait_until_ready() + reader.wait_until_ready() + array = np.full(nbytes // 2, 17, dtype=np.uint16) + array.flags.writeable = kind != "numpy_readonly" + payload = torch.from_numpy(array) if kind == "tensor" else array + try: + with ( + mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv, + mock.patch.object(pickle, "loads", wraps=pickle.loads) as loads, + ): + writer.enqueue(payload) + received = reader.dequeue(timeout=5) + received_oob = len(loads.call_args.kwargs["buffers"]) + # The spy must not keep the ring's memoryviews alive. + loads.reset_mock() + assert recv.call_count == int(overflow) + assert received_oob == num_oob + assert type(received) is type(payload) + assert received.dtype == payload.dtype + if isinstance(received, np.ndarray): + assert received.flags.writeable == array.flags.writeable + np.testing.assert_array_equal(received, array) + for value in range(4): + # Fit the replacement in SHM even when the original overflowed. + writer.enqueue(np.full(512 * 1024, value, dtype=np.uint16)) + replacement = reader.dequeue(timeout=5) + np.testing.assert_array_equal(replacement, value) + np.testing.assert_array_equal(received, array) + if array.flags.writeable: + received[0] = 42 + assert array[0] == 17 + finally: + reader.shutdown() + writer.shutdown() + + def _dumps_oob(obj) -> tuple[bytes, list]: """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor dispatch table + out-of-band buffers >= 1MiB.""" diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 5656e8aaa153..c5180d68a3c4 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -408,8 +408,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T Counterpart of `_reduce_tensor`. Note that pickle passes the original buffer-providing object from `loads(buffers=...)` straight to this function (no `PickleBuffer` wrapper on the receiving side), so `buf` is - a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` - if the buffer was serialized in-band. + a `zmq.Frame`, an owned `bytearray` from the shared-memory ring, or + `bytes`/`bytearray` if the buffer was serialized in-band. """ dtype = getattr(torch, dtype_str) assert isinstance(dtype, torch.dtype) @@ -423,11 +423,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T except ValueError: # Empty or read-only frame buffer; fall through to the copy path. pass - # Shared-memory ring buffer chunks are reused by the writer once all - # readers have marked them read, so we must copy out of them. bytearray - # (vs bytes) keeps the resulting tensor writable, matching normal tensor - # semantics. - raw = bytearray(buf) + # dequeue already owns SHM buffers. Copy other buffers to keep them writable. + raw = buf if isinstance(buf, bytearray) else bytearray(buf) if not raw: assert 0 in shape return torch.empty(shape, dtype=dtype) @@ -981,7 +978,9 @@ def dequeue( buf_offset = offset + 4 buf_len = from_bytes_big(buf[offset:buf_offset]) offset = buf_offset + buf_len - all_buffers.append(buf[buf_offset:offset]) + buffer = buf[buf_offset:offset] + # Out-of-band arrays may outlive this reusable ring slot. + all_buffers.append(buffer if i == 0 else bytearray(buffer)) obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:]) if overflow: obj = MessageQueue.recv(self.local_socket, timeout) From 6a70913590d8eab63320da4f3ce70976c9192500 Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Mon, 12 Oct 2026 14:47:59 +0800 Subject: [PATCH 386/952] [XPU] bump up auto-round-lib to 0.16.0 (#60976) Signed-off-by: Zhenzhong1 Co-authored-by: Kunshang Ji --- requirements/xpu.txt | 2 +- tests/quantization/test_auto_round.py | 17 +---------------- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/requirements/xpu.txt b/requirements/xpu.txt index d9d0f3080f86..14372071f0b5 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -19,5 +19,5 @@ torchaudio torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib==0.15.0 +auto_round_lib==0.16.0 vllm_xpu_kernels==0.1.15.4 diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index eb8273c6d6ba..267e1a11ba86 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -146,21 +146,6 @@ }, } -# These checkpoints fail during XPU graph capture, so run them eagerly there. -XPU_EAGER_ONLY_MODELS = { - "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", - "Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound", - "Intel/Qwen3-8B-w2g64-for-ut", - "INCModel/Qwen3-30B-A3B-12L-W4A16-test", -} - - -def _runner_kwargs(model: str) -> dict[str, Any]: - kwargs = dict(MODEL_RUNNER_KWARGS.get(model, {})) - if current_platform.is_xpu() and model in XPU_EAGER_ONLY_MODELS: - kwargs["enforce_eager"] = True - return kwargs - @pytest.mark.skipif( not ( @@ -172,7 +157,7 @@ def _runner_kwargs(model: str) -> dict[str, Any]: ) @pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS) def test_auto_round_model(vllm_runner, model): - with vllm_runner(model, **_runner_kwargs(model)) as llm: + with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=8) assert output From 36f397281295f5b58852a7eaf429b90348161079 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:09:10 -0700 Subject: [PATCH 387/952] [Spec Decode] Skip metadata build for dflash/dspark when all model backends support it (#60692) Signed-off-by: Giancarlo Delfin Co-authored-by: Rishi Puri <20074092+puririshi98@users.noreply.github.com> --- .../attention/test_deepseek_v4_swa_visible.py | 54 ++++++ .../test_triton_attn_metadata_builder.py | 154 ++++++++++++++++++ .../v1/attention/test_triton_attn_segments.py | 73 --------- ...test_gpu_target_dependent_ar_speculator.py | 22 +++ .../deepseek_v41/nvidia/flashinfer_sparse.py | 44 +++-- vllm/v1/attention/backends/mla/sparse_swa.py | 55 +++++-- vllm/v1/attention/backends/triton_attn.py | 3 + .../gpu/spec_decode/dflash/cudagraph.py | 9 +- .../gpu/spec_decode/dflash/speculator.py | 45 +++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 21 +++ .../target_dependent_ar/speculator.py | 52 +----- 11 files changed, 360 insertions(+), 172 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_metadata_builder.py delete mode 100644 tests/v1/attention/test_triton_attn_segments.py diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index 8f08929549a8..5ccd72180846 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -828,3 +828,57 @@ def test_v41_flashinfer_mixed_sparse_indices_respect_replay_start(): # Decode row is copied through; prefill rows follow the bounded window. assert bounded[0].cpu().tolist() == plain[0].cpu().tolist() assert bounded[1:].cpu().tolist() == rows[1:] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_dspark_noncausal_draft_decode_update_matches_build(): + """A DSpark draft graph refreshes its non-causal SWA metadata in place + instead of rebuilding it, so the refresh must match a fresh build.""" + from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, + ) + + num_spec = 4 + query_lens = [num_spec] * 3 + + def make_dspark_builder(): + builder = make_builder(vision=False) + builder.is_dspark = True + builder.noncausal_index_width = get_dspark_swa_index_width(WINDOW, num_spec) + builder.decode_threshold = num_spec + return builder + + def common_metadata(seq_lens): + query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch( + seq_lens, query_lens, torch.device("cuda") + ) + return CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc.cpu(), + seq_lens=seq_lens_t, + seq_lens_cpu_upper_bound=seq_lens_t.cpu(), + num_reqs=len(seq_lens), + num_actual_tokens=int(query_start_loc[-1]), + max_query_len=num_spec, + max_seq_len=max(seq_lens), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + causal=False, + ) + + captured_seq_lens = [6, 40, 130] + replay_seq_lens = [130, 9, 90] + builder = make_dspark_builder() + captured = common_metadata(captured_seq_lens) + md = builder.build(0, captured) + replay = common_metadata(replay_seq_lens) + expected = make_dspark_builder().build(0, replay) + + # A replay sees new inputs through the same device buffers. + captured.seq_lens.copy_(replay.seq_lens) + captured.block_table_tensor.copy_(replay.block_table_tensor) + builder.update_draft_decode_metadata(md) + + assert md.decode_swa_width == builder.noncausal_index_width + torch.testing.assert_close(md.decode_swa_lens, expected.decode_swa_lens) + torch.testing.assert_close(md.decode_swa_indices, expected.decode_swa_indices) diff --git a/tests/v1/attention/test_triton_attn_metadata_builder.py b/tests/v1/attention/test_triton_attn_metadata_builder.py new file mode 100644 index 000000000000..9ec814556633 --- /dev/null +++ b/tests/v1/attention/test_triton_attn_metadata_builder.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder +from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVMetadataBuilder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec + + +# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and +# the 64-segment scratch holds every query token (e.g. multi-query verify). +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +@pytest.mark.parametrize( + "capability,num_kv_heads,num_seqs,query_len,segments", + [ + ((12, 0), 1, 11, 1, 64), + ((12, 0), 1, 12, 1, 16), + ((12, 0), 8, 1, 1, 64), + ((12, 0), 8, 2, 1, 16), + ((9, 0), 1, 1, 1, 16), + ((12, 0), 1, 2, 4, 64), + ((12, 0), 1, 10, 4, 16), + ], +) +def test_split_k_segments_follow_sm_occupancy( + monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability", lambda cap: cap == capability + ) + monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: num_kv_heads, + get_head_size=lambda: 128, + rswa_window=None, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16 + ) + builder = builder_cls(spec, ["layer.0"], config, "cpu") + batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs) + metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) + assert metadata.num_par_softmax_segments == segments + assert metadata.softmax_segm_output.shape[2] == segments + assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len + # DiffKV re-allocates the output; it must stay sized like the base buffers. + assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] + assert ( + metadata.softmax_segm_output.data_ptr() + == builder.softmax_segm_output.data_ptr() + ) + + +def _draft_builder(builder_cls, rswa_window: int | None): + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: 1, + get_head_size=lambda: 128, + rswa_window=rswa_window, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ) + return builder_cls(spec, ["layer.0"], config, "cpu") + + +def _draft_decode_batch(seq_lens: list[int]): + batch = BatchSpec(seq_lens=seq_lens, query_lens=[1] * len(seq_lens)) + return create_common_attn_metadata(batch, 16, "cpu") + + +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +@pytest.mark.parametrize("rswa_window", [None, 64]) +def test_draft_decode_update_flag_matches_build_writes(builder_cls, rswa_window): + """The update is a no-op, so the builder may claim support only when + build() writes no builder-owned tensor that a captured graph reads.""" + builder = _draft_builder(builder_cls, rswa_window) + common_metadata = _draft_decode_batch([5, 9]) + common_metadata.rswa_prefix_lens = torch.tensor([3, 7], dtype=torch.int32) + + def owned_bytes(): + # Bytes, not values: uninitialized scratch may hold NaNs. + return { + name: value.view(torch.uint8).clone() + for name, value in vars(builder).items() + if isinstance(value, torch.Tensor) + } + + before = owned_bytes() + builder.build(0, common_metadata) + after = owned_bytes() + + written = [name for name in before if not torch.equal(before[name], after[name])] + assert builder.supports_draft_decode_metadata_update is (not written) + assert builder.supports_draft_decode_metadata_update is (rswa_window is None) + + +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +def test_draft_decode_update_matches_fresh_build(builder_cls): + """A replayed graph reads metadata built at capture; after an in-place + update it must match a fresh build of the replayed batch.""" + torch.manual_seed(0) + builder = _draft_builder(builder_cls, rswa_window=None) + captured = _draft_decode_batch([100, 120]) + metadata = builder.build(0, captured) + + replayed = _draft_decode_batch([128, 30]) + # The replayed batch arrives through the same persistent input buffers. + captured.seq_lens.copy_(replayed.seq_lens) + captured.block_table_tensor.copy_(replayed.block_table_tensor) + captured.slot_mapping.copy_(replayed.slot_mapping) + builder.update_draft_decode_metadata(metadata) + + expected = _draft_builder(builder_cls, rswa_window=None).build(0, replayed) + assert torch.equal(metadata.seq_lens, expected.seq_lens) + assert torch.equal(metadata.block_table, expected.block_table) + assert torch.equal(metadata.slot_mapping, expected.slot_mapping) + assert torch.equal(metadata.query_start_loc, expected.query_start_loc) diff --git a/tests/v1/attention/test_triton_attn_segments.py b/tests/v1/attention/test_triton_attn_segments.py deleted file mode 100644 index 874c61d2379e..000000000000 --- a/tests/v1/attention/test_triton_attn_segments.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch - -from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata -from vllm.config import CUDAGraphMode -from vllm.platforms import current_platform -from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder -from vllm.v1.attention.backends.triton_attn_diffkv import ( - TritonAttentionDiffKVMetadataBuilder, -) -from vllm.v1.kv_cache_interface import FullAttentionSpec - - -# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and -# the 64-segment scratch holds every query token (e.g. multi-query verify). -@pytest.mark.parametrize( - "builder_cls", - [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], -) -@pytest.mark.parametrize( - "capability,num_kv_heads,num_seqs,query_len,segments", - [ - ((12, 0), 1, 11, 1, 64), - ((12, 0), 1, 12, 1, 16), - ((12, 0), 8, 1, 1, 64), - ((12, 0), 8, 2, 1, 16), - ((9, 0), 1, 1, 1, 16), - ((12, 0), 1, 2, 4, 64), - ((12, 0), 1, 10, 4, 16), - ], -) -def test_split_k_segments_follow_sm_occupancy( - monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments -): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) - monkeypatch.setattr( - current_platform, "is_device_capability", lambda cap: cap == capability - ) - monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) - config = SimpleNamespace( - model_config=SimpleNamespace( - get_num_attention_heads=lambda _: 8, - get_num_kv_heads=lambda _: num_kv_heads, - get_head_size=lambda: 128, - rswa_window=None, - ), - parallel_config=None, - scheduler_config=SimpleNamespace(max_num_seqs=16), - speculative_config=None, - compilation_config=SimpleNamespace( - cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} - ), - ) - spec = FullAttentionSpec( - block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16 - ) - builder = builder_cls(spec, ["layer.0"], config, "cpu") - batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs) - metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) - assert metadata.num_par_softmax_segments == segments - assert metadata.softmax_segm_output.shape[2] == segments - assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len - # DiffKV re-allocates the output; it must stay sized like the base buffers. - assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] - assert ( - metadata.softmax_segm_output.data_ptr() - == builder.softmax_segm_output.data_ptr() - ) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index a3ec5f4eb7f4..490b84e77896 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -609,3 +609,25 @@ def fake_get_scheduler_metadata(**kwargs): assert not called assert metadata.scheduler_metadata is None + + +@pytest.mark.parametrize(("num_speculative_steps", "expected"), [(1, False), (3, True)]) +def test_decode_metadata_update_needs_draft_decode_steps( + num_speculative_steps, expected +): + """One speculative step runs no draft decodes, so the in-place metadata + update (and its fallback log) does not apply even on a capable backend.""" + speculator = object.__new__(_TestSpeculator) + speculator.num_speculative_steps = num_speculative_steps + speculator.attn_groups = [ + [ + SimpleNamespace( + supports_draft_decode_metadata_update=True, + backend=SimpleNamespace(get_name=lambda: "CAPABLE"), + ) + ] + ] + + speculator._configure_decode_metadata_update() + + assert speculator.supports_decode_metadata_update is expected diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0e02b486a517..a5ab09283959 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -206,24 +206,36 @@ def build( metadata = super().build( common_prefix_len, common_attn_metadata, fast_build, replay_start ) - num_tokens = metadata.num_decode_tokens - if not common_attn_metadata.causal and num_tokens > 0: - assert metadata.decode_swa_lens is not None - assert metadata.seq_lens is not None - assert metadata.token_to_req_indices is not None - topk_lens = self._decode_topk_lens[:num_tokens] - seq_lens = self._decode_seq_lens[:num_tokens] - torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens) - torch.index_select( - metadata.seq_lens, - 0, - metadata.token_to_req_indices[:num_tokens], - out=seq_lens, - ) - metadata.flashinfer_decode_topk_lens = topk_lens - metadata.flashinfer_decode_seq_lens = seq_lens + if not common_attn_metadata.causal and metadata.num_decode_tokens > 0: + self._fill_noncausal_decode_lens(metadata) return metadata + def update_draft_decode_metadata( + self, metadata: "DeepseekSparseSWAMetadata" + ) -> None: + super().update_draft_decode_metadata(metadata) + if metadata.flashinfer_decode_topk_lens is not None: + self._fill_noncausal_decode_lens(metadata) + + def _fill_noncausal_decode_lens( + self, metadata: "DeepseekSparseSWAMetadata" + ) -> None: + num_tokens = metadata.num_decode_tokens + assert metadata.decode_swa_lens is not None + assert metadata.seq_lens is not None + assert metadata.token_to_req_indices is not None + topk_lens = self._decode_topk_lens[:num_tokens] + seq_lens = self._decode_seq_lens[:num_tokens] + torch.clamp(metadata.decode_swa_lens, min=self.window_size, out=topk_lens) + torch.index_select( + metadata.seq_lens, + 0, + metadata.token_to_req_indices[:num_tokens], + out=seq_lens, + ) + metadata.flashinfer_decode_topk_lens = topk_lens + metadata.flashinfer_decode_seq_lens = seq_lens + class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend): @staticmethod diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c5cbb4ef2822..a691efe8bfb4 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -852,23 +852,44 @@ def update_draft_decode_metadata( num_tokens, ) - _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( - metadata.decode_swa_indices, - metadata.decode_swa_lens, - metadata.decode_swa_indices.shape[-1], - metadata.decode_swa_indices.shape[-1], - metadata.decode_swa_lens, # unused (HAS_IMAGE=False) - metadata.decode_swa_lens, # unused (HAS_IMAGE=False) - metadata.query_start_loc, - metadata.seq_lens, - metadata.token_to_req_indices, - metadata.is_valid_token, - metadata.block_table, - self.block_size, - metadata.replay_start, - num_tokens=metadata.num_decode_tokens, - token_offset=0, - ) + # build() points non-causal (DSpark draft) metadata at its own buffer. + noncausal = self.decode_swa_indices_noncausal + if ( + noncausal is not None + and metadata.decode_swa_indices.data_ptr() == noncausal.data_ptr() + ): + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( + metadata.decode_swa_indices, + metadata.decode_swa_lens, + self.window_size, + self.noncausal_index_width, + metadata.query_start_loc, + metadata.seq_lens, + metadata.token_to_req_indices, + metadata.is_valid_token, + metadata.block_table, + self.block_size, + num_tokens=num_tokens, + token_offset=0, + ) + else: + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( + metadata.decode_swa_indices, + metadata.decode_swa_lens, + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.query_start_loc, + metadata.seq_lens, + metadata.token_to_req_indices, + metadata.is_valid_token, + metadata.block_table, + self.block_size, + metadata.replay_start, + num_tokens=metadata.num_decode_tokens, + token_offset=0, + ) tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens) metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index d8577b61f1c9..095163f14423 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -193,6 +193,9 @@ def __init__( device=device, ) self.rswa_window = model_config.rswa_window + # R-SWA restages prefix lengths in build(), which the no-op + # update_draft_decode_metadata() does not refresh. + self.supports_draft_decode_metadata_update = self.rswa_window is None self.persistent_rswa_prefix_lens: torch.Tensor | None = None if self.rswa_window is not None: self.persistent_rswa_prefix_lens = torch.empty( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index e21ca9f8707e..2e189da0c939 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Mapping +from typing import Any import torch @@ -84,11 +85,11 @@ def capture( kv_cache_config: KVCacheConfig, max_model_len: int, causal: bool | Mapping[int, bool], - precompute_context_kv: Callable[[int], None], + prepare_forward: Callable[[int, dict[str, Any] | None], None], progress_bar_desc: str = "Capturing CUDA graphs", ) -> None: - """``precompute_context_kv(num_reqs)`` is captured ahead of the query - forward.""" + """``prepare_forward(num_reqs, attn_metadata)`` is captured ahead of the + query forward.""" def create_forward_fn( desc: BatchExecutionDescriptor, @@ -115,7 +116,7 @@ def create_forward_fn( attn_metadata, slot_mappings = attn_state def forward(cg_mode: CUDAGraphMode) -> None: - precompute_context_kv(num_reqs) + prepare_forward(num_reqs, attn_metadata) forward_fn( num_reqs, num_tokens, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 884de5a264be..5a13441ee892 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -166,12 +166,18 @@ def capture(self) -> None: self.kv_cache_config, self.max_model_len, causal=self._group_causal, - precompute_context_kv=lambda num_reqs: self._precompute_context_kv( - 0, self._num_graph_context_tokens(num_reqs) - ), + prepare_forward=self._prepare_graph_forward, progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) + def _prepare_graph_forward( + self, num_reqs: int, attn_metadata: dict[str, Any] | None + ) -> None: + self._precompute_context_kv(0, self._num_graph_context_tokens(num_reqs)) + # Recorded in the FULL graph so replays skip the eager metadata build. + if attn_metadata is not None and self.supports_decode_metadata_update: + self._update_draft_decode_metadata(attn_metadata, num_reqs) + def load_draft_model( self, target_model: nn.Module, @@ -486,29 +492,32 @@ def propose( else: self._precompute_context_kv(0, num_target_tokens, dummy_run) - # Rebuild the draft attention metadata even when replaying the FULL - # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_uniform_attn_metadata( - num_reqs=num_reqs, - batch_desc=batch_desc, - num_query_per_req=self.num_query_per_req, - seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, - step=self.num_query_per_req, - causal=self._group_causal, - ) - draft_slot_mappings_by_layer = build_slot_mappings_by_layer( - self.block_tables.slot_mappings[:, :num_tokens_padded], - self.kv_cache_config, - ) + is_full_cudagraph = batch_desc.cg_mode == CUDAGraphMode.FULL + # During FULL cudagraph with a model that supports decode metadata updates, + # the metadata update is captured in the graph, and we can skip the full + # build below. + if not (is_full_cudagraph and self.supports_decode_metadata_update): + draft_attn_metadata = self._build_uniform_attn_metadata( + num_reqs=num_reqs, + batch_desc=batch_desc, + num_query_per_req=self.num_query_per_req, + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, + step=self.num_query_per_req, + causal=self._group_causal, + ) # DFlash processes all speculative tokens in one forward pass, # so the real token count is num_query_tokens. self._prepare_eplb_forward(num_query_tokens) - if batch_desc.cg_mode == CUDAGraphMode.FULL: + if is_full_cudagraph: assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.run_fullgraph(batch_desc) else: + draft_slot_mappings_by_layer = build_slot_mappings_by_layer( + self.block_tables.slot_mappings[:, :num_tokens_padded], + self.kv_cache_config, + ) self._generate_draft( num_reqs, num_tokens_padded, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index c3e8dcbff4d1..d86a0b0351f0 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -207,6 +207,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.supports_mm_inputs = False self.pcp_manager: PCPManager | None = None + # Whether every draft attention builder can refresh its decode metadata + # in place, so the refresh can be recorded in the draft CUDA graph. + self.supports_decode_metadata_update = False @abstractmethod def load_draft_model( @@ -299,6 +302,24 @@ def set_attn( # builders and buffers. self.target_input_buffers = target_input_buffers self.target_attn_groups = target_attn_groups + self._configure_decode_metadata_update() + + def _configure_decode_metadata_update(self) -> None: + unsupported_backends = sorted( + { + attn_group.backend.get_name() + for attn_groups in self.attn_groups + for attn_group in attn_groups + if not attn_group.supports_draft_decode_metadata_update + } + ) + self.supports_decode_metadata_update = not unsupported_backends + if unsupported_backends: + logger.info_once( + "Attention backend(s) %s cannot update draft decode metadata in " + "place; falling back to rebuilding attention metadata eagerly.", + ", ".join(unsupported_backends), + ) def _build_attn_metadata( self, diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 02e28875b823..fdbedd78f6bd 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -10,18 +10,15 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.triton_utils import tl, triton -from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer -from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) -from vllm.v1.worker.utils import AttentionGroup, get_uniform_decode_token_count +from vllm.v1.worker.utils import get_uniform_decode_token_count logger = init_logger(__name__) @@ -51,7 +48,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None - self.use_fused_multi_step_decode = False def load_model(self, target_model: nn.Module) -> None: super().load_model(target_model) @@ -83,44 +79,12 @@ def advance_draft_positions(self) -> bool: """ return True - def set_attn( - self, - model_state: ModelState, - kv_cache_config: KVCacheConfig, - block_tables: BlockTables, - target_input_buffers: InputBuffers, - target_attn_groups: list[list[AttentionGroup]], - ) -> None: - super().set_attn( - model_state, - kv_cache_config, - block_tables, - target_input_buffers, - target_attn_groups, - ) - self._configure_fused_multi_step_decode() - - def _configure_fused_multi_step_decode(self) -> None: + def _configure_decode_metadata_update(self) -> None: + # With one speculative step there are no draft decode steps to update. if self.num_speculative_steps == 1: - self.use_fused_multi_step_decode = False + self.supports_decode_metadata_update = False return - - unsupported_backends = sorted( - { - attn_group.backend.get_name() - for attn_groups in self.attn_groups - for attn_group in attn_groups - if not attn_group.supports_draft_decode_metadata_update - } - ) - self.use_fused_multi_step_decode = not unsupported_backends - if unsupported_backends: - logger.info_once( - "Fused multi-step draft decode is not supported by attention " - "backend(s) %s; falling back to rebuilding attention metadata " - "between draft steps.", - ", ".join(unsupported_backends), - ) + super()._configure_decode_metadata_update() def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: # Initialize cudagraph manager for draft prefill (draft position 0). @@ -181,7 +145,7 @@ def capture(self) -> None: assert self.decode_cudagraph_manager is not None decode_fn = ( self._generate_fused_drafts - if self.use_fused_multi_step_decode + if self.supports_decode_metadata_update else self._generate_draft ) self.decode_cudagraph_manager.capture( @@ -192,7 +156,7 @@ def capture(self) -> None: self.attn_groups, self.kv_cache_config, progress_bar_desc="Capturing decode CUDA graphs", - specialize_spec_tokens=self.use_fused_multi_step_decode + specialize_spec_tokens=self.supports_decode_metadata_update and self.speculative_config.uses_dynamic_speculative_decoding(), ) self.on_multi_step_decode_end(self.max_num_reqs) @@ -384,7 +348,7 @@ def propose( # Generate the remaining num_speculative_steps - 1 draft tokens. decode_fn = ( self._fused_multi_step_decode - if self.use_fused_multi_step_decode + if self.supports_decode_metadata_update else self._multi_step_decode ) decode_fn( From e3bbaf23695f7f30f144a288057b0419ab2b6a8d Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Mon, 12 Oct 2026 15:53:16 +0800 Subject: [PATCH 388/952] [ROCm][DSv4] Fuse the shared expert into AITER MegaMoE (#61004) Signed-off-by: Yanfei Wang Co-authored-by: aoshen02 --- .../models/test_deepseek_v4_aiter_mega_moe.py | 20 +++ .../device_communicators/all2all.py | 6 +- vllm/envs.py | 7 ++ vllm/models/deepseek_v4/amd/mega_moe.py | 117 ++++++++++++++++-- vllm/models/deepseek_v4/amd/model.py | 13 +- 5 files changed, 151 insertions(+), 12 deletions(-) diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index a8f692e07bfb..19b3697811df 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -29,6 +29,7 @@ def _make_experts( intermediate_size=INTER, max_num_batched_tokens=3000, swiglu_limit=None, + fuse_shared_expert=False, ): monkeypatch.setattr( mega_moe_module, @@ -47,6 +48,7 @@ def _make_experts( hidden_size=HIDDEN, intermediate_size=intermediate_size, swiglu_limit=swiglu_limit, + fuse_shared_expert=fuse_shared_expert, ) @@ -108,6 +110,24 @@ def test_experts_load_into_local_slots(monkeypatch): assert torch.all(w2 == 10 * expert + 3) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") +def test_shared_expert_route_appended(monkeypatch): + experts = _make_experts( + monkeypatch, ep_rank=3, ep_size=8, num_experts=384, fuse_shared_expert=True + ) + ids = torch.randint(0, 384, (37, 2), device="cuda", dtype=torch.int32) + ids[::5, 1] = -1 + weights = torch.rand(37, 2, device="cuda") + out_weights, out_ids = experts.append_shared_expert(weights, ids) + + shared = torch.full((37, 1), 3 * 49 + 48, device="cuda", dtype=torch.int32) + expected_ids = torch.cat([torch.where(ids >= 0, ids + ids // 48, ids), shared], 1) + torch.testing.assert_close(out_ids, expected_ids) + torch.testing.assert_close( + out_weights, torch.cat([weights, torch.ones(37, 1, device="cuda")], 1) + ) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") @pytest.mark.parametrize("hash_routing", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 37]) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index cff34139f968..3466700c3334 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1012,7 +1012,11 @@ def __init__(self, cpu_group, all2all_backend: str): if get_current_vllm_config().kernel_config.moe_backend == "aiter_mega_moe": # MegaMoEV2 places its dispatch/combine workspaces on the MoRI # symmetric heap, which defaults to 2 GB. MegaMoEV2 requires > 4GB - heap_size = os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + # (> 8GB with a fused shared expert). + heap_size = os.environ.setdefault( + "MORI_SHMEM_HEAP_SIZE", + "16G" if envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS else "8G", + ) logger.info_once("AITER MegaMoE: MORI_SHMEM_HEAP_SIZE=%s", heap_size) mori.shmem.shmem_torch_process_group_init("mori") diff --git a/vllm/envs.py b/vllm/envs.py index 14489a3f2503..a2b46030772f 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ VLLM_ROCM_USE_AITER_FP4BMM: bool = True VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False + VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True @@ -1391,6 +1392,12 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower() in ("true", "1") ), + # Fuse DeepSeek-V4's shared expert into AITER MegaMoE (requantized to + # MXFP4). By default is disabled. + "VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS": lambda: ( + os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower() + in ("true", "1") + ), # Run the eligible layers of decode steps on the model's mono decode # kernels: persistent launches a layer, the TP all-reduces in-kernel # (DeepSeek-V4.1 on CDNA4). By default is disabled. diff --git a/vllm/models/deepseek_v4/amd/mega_moe.py b/vllm/models/deepseek_v4/amd/mega_moe.py index 2fdff9335145..312c089e5aad 100644 --- a/vllm/models/deepseek_v4/amd/mega_moe.py +++ b/vllm/models/deepseek_v4/amd/mega_moe.py @@ -10,10 +10,51 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import get_ep_group +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + get_fp8_block_weight_scale, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, +) from vllm.model_executor.utils import set_weight_attrs +from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 +@triton.jit +def _append_shared_expert_kernel( + topk_ids_ptr, + topk_weights_ptr, + out_ids_ptr, + out_weights_ptr, + num_tokens, + shared_expert_id, + TOPK: tl.constexpr, + ROUTED_PER_RANK: tl.constexpr, + BLOCK_M: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + mask = rows < num_tokens + for k in tl.static_range(TOPK): + e = tl.load(topk_ids_ptr + rows * TOPK + k, mask=mask, other=0) + w = tl.load(topk_weights_ptr + rows * TOPK + k, mask=mask, other=0.0) + e = tl.where(e >= 0, e + e // ROUTED_PER_RANK, e) + tl.store(out_ids_ptr + rows * (TOPK + 1) + k, e, mask=mask) + tl.store(out_weights_ptr + rows * (TOPK + 1) + k, w, mask=mask) + tl.store( + out_ids_ptr + rows * (TOPK + 1) + TOPK, + tl.zeros_like(rows) + shared_expert_id, + mask=mask, + ) + tl.store( + out_weights_ptr + rows * (TOPK + 1) + TOPK, + tl.full([BLOCK_M], 1.0, tl.float32), + mask=mask, + ) + + class DeepseekV4AiterMegaMoEExperts(nn.Module): # MegaMoEV2 owns large symmetric-heap buffers, so every layer shares one # instance and passes its own weights per call. @@ -28,16 +69,25 @@ def __init__( hidden_size: int, intermediate_size: int, swiglu_limit: float | None, + fuse_shared_expert: bool = False, prefix: str = "", ): super().__init__() ep_group = get_ep_group() self.ep_rank = ep_group.rank_in_group self.ep_size = ep_group.world_size - self.num_experts = num_experts - self.num_local_experts = num_experts // self.ep_size - self.experts_start_idx = self.ep_rank * self.num_local_experts - self.top_k = top_k + self.num_routed_local_experts = num_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.num_routed_local_experts + # A fused shared expert is one extra local expert per rank and one + # extra route per token, to the token's own rank. + self.fuse_shared_expert = fuse_shared_expert + self.num_local_experts = self.num_routed_local_experts + fuse_shared_expert + self.num_experts = self.num_local_experts * self.ep_size + self.routed_top_k = top_k + self.top_k = top_k + fuse_shared_expert + self.shared_expert_id = ( + self.ep_rank * self.num_local_experts + self.num_routed_local_experts + ) self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.swiglu_limit = swiglu_limit or 0.0 @@ -70,7 +120,7 @@ def weight_loader( return_success: bool = False, ) -> bool | None: local_expert_id = expert_id - self.experts_start_idx - if not 0 <= local_expert_id < self.num_local_experts: + if not 0 <= local_expert_id < self.num_routed_local_experts: return False if return_success else None expert_data = param.data[local_expert_id] if shard_id in ("w1", "w3"): @@ -81,6 +131,24 @@ def weight_loader( expert_data.copy_(loaded_weight.view(torch.uint8)) return True if return_success else None + def load_shared_expert(self, shared_experts: nn.Module) -> None: + slot = self.num_routed_local_experts + for linear, weight, weight_scale in ( + (shared_experts.gate_up_proj, self.w13_weight, self.w13_weight_scale), + (shared_experts.down_proj, self.w2_weight, self.w2_weight_scale), + ): + block_scale = get_fp8_block_weight_scale(linear) + assert block_scale is not None, "shared expert must be block-FP8" + group = GroupShape( + linear.weight.shape[0] // block_scale.shape[0], + linear.weight.shape[1] // block_scale.shape[1], + ) + packed, scale = mxfp4_quantize( + scaled_dequantize(linear.weight, block_scale, group, torch.bfloat16) + ) + weight.data[slot].copy_(packed.view(torch.uint8)) + weight_scale.data[slot].copy_(scale.view(torch.uint8)) + def finalize_weights(self) -> None: if self._runtime is not None: return @@ -129,6 +197,33 @@ def finalize_weights(self) -> None: self._runtime_cache[key] = runtime self._runtime = runtime + def append_shared_expert( + self, topk_weights: torch.Tensor, topk_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Remap routed ids to the per-rank (routed + shared) layout and append + a weight-1.0 route to this rank's shared slot.""" + num_tokens = topk_ids.shape[0] + out_ids = torch.empty( + (num_tokens, self.top_k), dtype=topk_ids.dtype, device=topk_ids.device + ) + out_weights = torch.empty( + (num_tokens, self.top_k), dtype=torch.float32, device=topk_ids.device + ) + if num_tokens: + block_m = 64 + _append_shared_expert_kernel[(triton.cdiv(num_tokens, block_m),)]( + topk_ids.contiguous(), + topk_weights.to(torch.float32).contiguous(), + out_ids, + out_weights, + num_tokens, + self.shared_expert_id, + TOPK=self.routed_top_k, + ROUTED_PER_RANK=self.num_routed_local_experts, + BLOCK_M=block_m, + ) + return out_weights, out_ids + def forward( self, hidden_states: torch.Tensor, @@ -137,6 +232,8 @@ def forward( ) -> torch.Tensor: runtime = self._runtime assert runtime is not None + if self.fuse_shared_expert: + topk_weights, topk_ids = self.append_shared_expert(topk_weights, topk_ids) # MegaMoEV2 reads these per launch; rebind to this layer's experts. runtime._s1_w1 = self.w13_weight runtime._s1_w1_scale = self.w13_weight_scale @@ -149,6 +246,10 @@ def forward( def finalize_mega_moe_weights(model: nn.Module) -> None: - for module in model.modules(): - if isinstance(module, DeepseekV4AiterMegaMoEExperts): - module.finalize_weights() + for module in list(model.modules()): + experts = getattr(module, "experts", None) + if isinstance(experts, DeepseekV4AiterMegaMoEExperts): + if experts.fuse_shared_expert: + experts.load_shared_expert(module.shared_experts) + module.shared_experts = None + experts.finalize_weights() diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 8b26fc0060e7..51d22b75e8c1 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -74,7 +74,10 @@ make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.mega_moe import DeepseekV4AiterMegaMoEExperts +from vllm.models.deepseek_v4.amd.mega_moe import ( + DeepseekV4AiterMegaMoEExperts, + finalize_mega_moe_weights, +) from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 @@ -690,6 +693,10 @@ def __init__( hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, swiglu_limit=self.swiglu_limit, + fuse_shared_expert=( + config.n_shared_experts == 1 + and envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS + ), prefix=f"{prefix}.experts", ) return @@ -1524,6 +1531,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def process_weights_after_loading(self) -> None: + # Must precede the gate_up preshuffle of the fused shared experts. + finalize_mega_moe_weights(self) # After per-layer quant finalize, so we preshuffle the final fp8 weights. fused_compressor_layers = 0 for module in self.modules(): @@ -1532,8 +1541,6 @@ def process_weights_after_loading(self) -> None: module.prepare_attn_preshuffle() elif isinstance(module, DeepseekV4MLP): module.prepare_gateup_preshuffle() - elif isinstance(module, DeepseekV4AiterMegaMoEExperts): - module.finalize_weights() if fused_compressor_layers: logger.info( "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers", From f4114e068e5d8e961defeb7747a9449d5eec0659 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 12 Oct 2026 15:55:06 +0800 Subject: [PATCH 389/952] [Bugfix][Frontend] Preserve token IDs in chat derender responses (#61065) Signed-off-by: Ting Sun --- .../scale_out/derender/test_derender_parity.py | 5 +++-- .../scale_out/derender/test_derender_stream.py | 3 ++- tests/entrypoints/scale_out/derender/utils.py | 4 ++++ vllm/renderers/online_derenderer.py | 13 +++++++++++++ 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 4ea3bc246192..f33187387403 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -126,7 +126,7 @@ async def _disagg( ], }, "prompt_tokens": prompt_tokens, - "chat_request": chat_request, + "chat_request": {"return_token_ids": True, **chat_request}, }, ) assert resp.status_code == 200, resp.text @@ -149,6 +149,7 @@ def _assert_parity(coupled: dict, disagg: dict) -> None: assert d["message"].get("reasoning") == c["message"].get("reasoning") assert _tool_sig(d) == _tool_sig(c) assert d["finish_reason"] == c["finish_reason"] + assert d["token_ids"] == c["token_ids"] assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"] assert disagg["usage"]["completion_tokens"] == len(c["token_ids"]) @@ -386,7 +387,7 @@ async def _run_stream_parity_case( [tid for ch in coupled_choices for tid in ch["token_ids"]], [len(ch["token_ids"]) for ch in coupled_choices] + ([0] if finish_only_tail else []), - {"model": MODEL, "messages": messages, **extra}, + {"model": MODEL, "messages": messages, "return_token_ids": True, **extra}, len(prompt_token_ids), prompt_token_ids, finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 0a7d96f3cf90..d97a2f992ac2 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -1311,7 +1311,7 @@ async def test_finish_reason_stop_for_named_tool_choice(self, parsed_derenderer) async def test_include_reasoning_false_suppresses_reasoning( self, parsed_derenderer ): - chat_request = _chat_request(include_reasoning=False) + chat_request = _chat_request(include_reasoning=False, return_token_ids=True) chunk, _ = await parsed_derenderer.derender_chat_stream( model=MODEL_NAME, generate_chunk=_make_stream_chunk( @@ -1320,6 +1320,7 @@ async def test_include_reasoning_false_suppresses_reasoning( chat_request=chat_request, ) delta = chunk.choices[0].delta + assert chunk.choices[0].token_ids is None assert delta.reasoning is None assert delta.content == "c" diff --git a/tests/entrypoints/scale_out/derender/utils.py b/tests/entrypoints/scale_out/derender/utils.py index d8128353bf3c..c249e2053a6a 100644 --- a/tests/entrypoints/scale_out/derender/utils.py +++ b/tests/entrypoints/scale_out/derender/utils.py @@ -95,6 +95,10 @@ async def stream_chat_derender( ) assert resp.status_code == 200, resp.text data = resp.json() + if chat_request.get("return_token_ids") and chat_request.get( + "include_reasoning", True + ): + assert data["chunk"]["choices"][0]["token_ids"] == tids state = data["stream_state"] choices.extend(data["chunk"]["choices"]) if on_chunk is not None: diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 8e3e4ca25abc..3fe50f16c391 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -160,6 +160,11 @@ def _derender_chat( else {} ) + return_token_ids = ( + chat_request is not None + and chat_request.return_token_ids + and (not has_parser or chat_request.include_reasoning) + ) for choice in generate_response.choices: if not choice.token_ids: raise ValueError(f"choice {choice.index} has empty or null token_ids") @@ -201,6 +206,7 @@ def _derender_chat( message=message, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -432,6 +438,7 @@ async def derender_chat_stream( stream_choices: list[ChatCompletionResponseStreamChoice] = [] updated_state = state + return_token_ids = chat_request is not None and chat_request.return_token_ids for choice in generate_chunk.choices: delta_tids = choice.token_ids or [] new_text, updated_state = await self._detokenize_delta_async( @@ -477,6 +484,7 @@ async def derender_chat_stream( delta=delta, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -583,6 +591,10 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: tools_streamed = state.tools_streamed last_tool_call_ids = list(state.last_tool_call_ids) + return_token_ids = ( + chat_request.return_token_ids and chat_request.include_reasoning + ) + # At most one choice: the caller (derender_chat_stream) already # rejects >1 before dispatching here. role_sent/tools_streamed/ # output_token_ids below are updated for that single choice, not @@ -660,6 +672,7 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: index=choice.index, delta=delta_message, finish_reason=finish_reason, + token_ids=choice.token_ids if return_token_ids else None, ) stream_choices.append( maybe_filter_parallel_tool_calls(stream_choice, chat_request) From 4163a80142d41b98f94d9b798506448162a22407 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:56:59 +0800 Subject: [PATCH 390/952] [Rust Frontend] Return weight-version spans on final outputs (#61070) Signed-off-by: Shuolei Wang Signed-off-by: Bugen Zhao Co-authored-by: aoshen02 Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../online_serving/token_in_token_out.md | 2 +- rust/src/engine-core-client/src/client/imp.rs | 4 +- .../src/client/output_channel.rs | 1 + .../engine-core-client/src/client/state.rs | 2 + .../engine-core-client/src/client/stream.rs | 2 + .../engine-core-client/src/protocol/output.rs | 10 +++- .../engine-core-client/src/tests/client.rs | 3 + rust/src/llm/src/lib.rs | 2 +- rust/src/llm/src/output.rs | 55 +++++++++++++++++++ rust/src/llm/tests/generate.rs | 45 ++++++++++++--- .../server/src/routes/inference/generate.rs | 9 +++ .../src/routes/inference/generate/types.rs | 4 ++ 12 files changed, 127 insertions(+), 12 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index b3e15ae61692..eade608bc451 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -66,7 +66,7 @@ The final output of every choice carries `weight_versions`, the label of the wei ] ``` -- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null`. +- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null` or omit the field. - An output with no tokens, such as a request aborted before its first token, gets `[]`. - The labels are the ones you set with `/update_weight_version` or `/finish_weight_update`, see [Async Reinforcement Learning](../../training/async_rl.md#weight-versions-in-outputs). vLLM doesn't order or check them, and the label is `"default"` until you set one. diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 0a4af1a5ee5f..fa82d9591aa8 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -453,10 +453,12 @@ pub(crate) async fn run_output_dispatcher_loop( let has_outputs = !batch.outputs.is_empty(); let mut iteration_tokens: u64 = batch.outputs.iter().map(output_iteration_tokens).sum(); - let (engine_index, timestamp) = (batch.engine_index, batch.timestamp); + let (engine_index, timestamp, weight_version) = + (batch.engine_index, batch.timestamp, batch.weight_version); let outputs = batch.outputs.into_iter().map(|output| EngineCoreStreamOutput { engine_index, timestamp, + weight_version: weight_version.clone(), output, }); for output in inner.send_outputs(outputs) { diff --git a/rust/src/engine-core-client/src/client/output_channel.rs b/rust/src/engine-core-client/src/client/output_channel.rs index 7b58d3dc2f1f..3de4da366d54 100644 --- a/rust/src/engine-core-client/src/client/output_channel.rs +++ b/rust/src/engine-core-client/src/client/output_channel.rs @@ -130,6 +130,7 @@ mod tests { EngineCoreStreamOutput { engine_index: 0, timestamp: 0.0, + weight_version: None, output: EngineCoreOutput { request_id: "req".to_string(), new_token_ids: token_ids.to_vec(), diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 9c183b87e4e0..9f6f8bbf5c2b 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -344,6 +344,7 @@ impl RequestRegistry { let output = EngineCoreStreamOutput { engine_index: engine_id.engine_index().unwrap_or(0), timestamp, + weight_version: None, output: EngineCoreOutput { request_id: request_id.clone(), finish_reason: Some(EngineCoreFinishReason::Abort), @@ -527,6 +528,7 @@ mod tests { EngineCoreStreamOutput { engine_index: 0, timestamp: 0.0, + weight_version: None, output, } } diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index 6f363f61017d..0f54da98c97d 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -3,6 +3,7 @@ use std::ops::Deref; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use futures::stream::FusedStream; @@ -30,6 +31,7 @@ enum State { pub struct EngineCoreStreamOutput { pub engine_index: u32, pub timestamp: f64, + pub weight_version: Option>, pub output: EngineCoreOutput, } diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 8945f8b11020..d269f2491163 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; +use std::sync::Arc; use bytes::Bytes; use enum_as_inner::EnumAsInner; @@ -228,7 +229,8 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, - /// Decoded for Python wire compatibility; not exposed by the Rust frontend. + /// Label of the weights used for this step, as set by the + /// `set_weight_version` utility; `"default"` until set. #[serde(default)] weight_version: Option, } @@ -247,6 +249,9 @@ pub struct RequestBatchOutputs { pub scheduler_stats: Option>, pub timestamp: f64, pub finished_requests: Option>, + /// Label of the weights used for this step. It changes only while + /// generation is paused, so it applies to every output in the batch. + pub weight_version: Option>, } #[derive(Debug, Clone, PartialEq, Default)] @@ -327,6 +332,7 @@ impl TryFrom for EngineCoreOutputs { scheduler_stats: value.scheduler_stats, timestamp: value.timestamp, finished_requests: value.finished_requests, + weight_version: value.weight_version.map(Arc::from), } .into()), (false, Some(_), None, None) => Ok(UtilityCallOutput { @@ -365,6 +371,7 @@ impl From for WireEngineCoreOutputs { scheduler_stats: batch.scheduler_stats, timestamp: batch.timestamp, finished_requests: batch.finished_requests, + weight_version: batch.weight_version.map(|version| version.to_string()), ..Default::default() }, EngineCoreOutputs::Utility(utility) => Self { @@ -558,6 +565,7 @@ mod tests { "req-1", }, ), + weight_version: None, }, ) "#]] diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 0ad58b027154..43eb88db79ed 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2884,6 +2884,9 @@ fn python_msgpack_fixtures_match_rust_encoding() { "req-1", }, ), + weight_version: Some( + "v7", + ), }, ) "#]] diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 7e9b5e0ba53f..c94c1d80f743 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -16,7 +16,7 @@ mod request_metrics; pub use error::{Error, Result}; pub use output::{ CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStream, - GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, + GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, WeightVersionSpan, }; pub use request::GenerateRequest; pub use request_metrics::current_unix_timestamp_secs; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index b6ccb7717508..81dbb3fda53c 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -24,6 +24,14 @@ use crate::error::Result; use crate::inflight::RequestGuard; use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs}; +/// Output tokens `[start, end)` were sampled by weights labelled `version`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct WeightVersionSpan { + pub version: Arc, + pub start: usize, + pub end: usize, +} + /// Token usage metadata for one request. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct TokenUsage { @@ -55,6 +63,7 @@ pub struct CollectedGenerateOutput { pub sampling_mask: Option, /// Per-request speculative-decoding metrics from the terminal output. pub spec_decode_metrics: Option, + pub weight_versions: Option>, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -172,6 +181,8 @@ pub struct GenerateOutput { pub sampling_mask: Option, /// Per-request speculative-decoding metrics, present on terminal outputs. pub spec_decode_metrics: Option>, + /// Weight-version spans over all output tokens, set on terminal outputs only. + pub weight_versions: Option>, } impl GenerateOutput { @@ -225,6 +236,7 @@ impl GenerateOutput { self.kv_transfer_params = next.kv_transfer_params; self.ec_transfer_params = next.ec_transfer_params; self.spec_decode_metrics = next.spec_decode_metrics; + self.weight_versions = next.weight_versions; } } @@ -253,6 +265,7 @@ impl GenerateOutput { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, } } } @@ -267,6 +280,9 @@ pub struct GenerateOutputStream { pending_prompt_info: Option>, raw_stream: EngineCoreOutputStream, request_metrics: RequestMetricsTracker, + /// (weight version, first output-token index) for each span. + weight_version_starts: Vec<(Arc, usize)>, + num_output_tokens: usize, /// Removes this request's external→internal tracking edge on drop. Held for /// its `Drop` side effect only; never read directly. _request_guard: RequestGuard, @@ -289,6 +305,8 @@ impl GenerateOutputStream { })), raw_stream, request_metrics, + weight_version_starts: Vec::new(), + num_output_tokens: 0, _request_guard: request_guard, } } @@ -320,6 +338,15 @@ impl GenerateOutputStream { fn convert_output(&mut self, raw: EngineCoreStreamOutput) -> Result { let received_at = current_unix_timestamp_secs(); self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); + let num_new_tokens = raw.output.new_token_ids.len(); + // Start a span when tokens arrive under a new label. + if num_new_tokens > 0 + && let Some(version) = raw.weight_version + && self.weight_version_starts.last().is_none_or(|(last, _)| *last != version) + { + self.weight_version_starts.push((version, self.num_output_tokens)); + } + self.num_output_tokens += num_new_tokens; let raw = raw.output; @@ -357,6 +384,11 @@ impl GenerateOutputStream { if let Some(finish_reason) = finish_reason.as_ref() { self.request_metrics.record_finished(received_at, finish_reason.clone()); } + let weight_versions = if finish_reason.is_some() { + self.weight_version_spans() + } else { + None + }; let output = GenerateOutput { request_id: raw.request_id, @@ -369,10 +401,31 @@ impl GenerateOutputStream { ec_transfer_params: raw.ec_transfer_params, sampling_mask, spec_decode_metrics: raw.spec_decode_metrics, + weight_versions, }; Ok(output) } + + /// Spans over the output so far; `[]` for no tokens, `None` if no label was reported. + fn weight_version_spans(&self) -> Option> { + let starts = &self.weight_version_starts; + if starts.is_empty() && self.num_output_tokens > 0 { + return None; + } + let ends = starts.iter().skip(1).map(|(_, start)| *start).chain([self.num_output_tokens]); + Some( + starts + .iter() + .zip(ends) + .map(|((version, start), end)| WeightVersionSpan { + version: version.clone(), + start: *start, + end, + }) + .collect(), + ) + } } impl Stream for GenerateOutputStream { @@ -476,6 +529,7 @@ impl> + Send> T { ec_transfer_params: None, sampling_mask, spec_decode_metrics: None, + weight_versions: None, }); } @@ -490,6 +544,7 @@ impl> + Send> T { collected.kv_transfer_params = output.kv_transfer_params.map(|value| *value); collected.ec_transfer_params = output.ec_transfer_params.map(|value| *value); collected.spec_decode_metrics = output.spec_decode_metrics.map(|value| *value); + collected.weight_versions = output.weight_versions; if let Some(mask) = collected.sampling_mask.as_ref() && mask.rows.len() != collected.token_ids.len() { diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index e1205ede9961..8a150d07a44b 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -31,6 +31,7 @@ use vllm_engine_core_client::test_utils::{ use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::{ Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm, + WeightVersionSpan, }; use vllm_metrics::METRICS; use zeromq::prelude::{SocketRecv, SocketSend}; @@ -342,15 +343,16 @@ async fn generate_merges_outputs_by_stream_interval() { let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); // One engine step per batch. - let steps: [(&[u32], Option); 6] = [ - (&[1], None), - (&[2], None), - (&[3], None), - (&[4, 5], None), - (&[6], None), - (&[7], Some(EngineCoreFinishReason::Length)), + // The label changes within the merged output [2, 3, 4, 5]. + let steps: [(&[u32], &str, Option); 6] = [ + (&[1], "v7", None), + (&[2], "v7", None), + (&[3], "v8", None), + (&[4, 5], "v8", None), + (&[6], "v7", None), + (&[7], "v7", Some(EngineCoreFinishReason::Length)), ]; - for (token_ids, finish_reason) in steps { + for (token_ids, weight_version, finish_reason) in steps { send_outputs( push, RequestBatchOutputs { @@ -361,6 +363,7 @@ async fn generate_merges_outputs_by_stream_interval() { Some(token_logprobs(token_ids)), None, )], + weight_version: Some(weight_version.into()), ..Default::default() } .into(), @@ -412,6 +415,21 @@ async fn generate_merges_outputs_by_stream_interval() { for output in &outputs { assert_eq!(output.logprobs, Some(token_logprobs(&output.token_ids))); } + let span = |version: &str, start, end| WeightVersionSpan { + version: version.into(), + start, + end, + }; + let weight_versions: Vec<_> = + outputs.iter().map(|output| output.weight_versions.clone()).collect(); + assert_eq!( + weight_versions, + [ + None, + None, + Some(vec![span("v7", 0, 2), span("v8", 2, 5), span("v7", 5, 7)]) + ] + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -530,6 +548,7 @@ async fn collect_output_rejects_partial_sampling_mask() { rows: vec![vec![1, 33, 99]], }), spec_decode_metrics: None, + weight_versions: None, }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); @@ -731,6 +750,7 @@ async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() push, RequestBatchOutputs { outputs: vec![request_output(&request.request_id, vec![7], None)], + weight_version: Some("v7".into()), ..Default::default() } .into(), @@ -764,6 +784,15 @@ async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() let terminal = stream.next().await.unwrap().unwrap(); assert_eq!(terminal.finish_reason, Some(FinishReason::Abort)); assert!(terminal.token_ids.is_empty()); + // The synthesized abort output reports the spans received so far. + assert_eq!( + terminal.weight_versions, + Some(vec![WeightVersionSpan { + version: "v7".into(), + start: 0, + end: 1 + }]) + ); assert!(stream.next().await.is_none()); let _ = shutdown_tx.send(()); diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 985fe44fb136..a30cd0c85ab4 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -235,6 +235,7 @@ async fn generate_chunk_stream( finish_reason: finish_reason.map(|reason| reason.as_str().to_string()), token_ids, sampling_mask: output.sampling_mask.map(|mask| mask.rows), + weight_versions: output.weight_versions, }], usage: include_continuous_usage .then(|| Usage::from_token_usage(usage, enable_prompt_tokens_details)), @@ -336,6 +337,7 @@ fn collect_generate( finish_reason: Some(finish_reason), token_ids: collected.token_ids, sampling_mask: collected.sampling_mask.map(|mask| mask.rows), + weight_versions: collected.weight_versions, }], prompt_logprobs, prompt_token_id_logprobs: collected @@ -629,6 +631,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }), Ok(GenerateOutput { request_id: String::new(), @@ -645,6 +648,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }), ]); @@ -731,6 +735,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, } } @@ -905,6 +910,7 @@ mod tests { rows: vec![vec![30, 40]], }), spec_decode_metrics: None, + weight_versions: None, }; let response = collect_generate( @@ -959,6 +965,7 @@ mod tests { prompt_token_ids: vec![10], sampling_mask: None, spec_decode_metrics: Some(metrics.clone()), + weight_versions: None, }; let response = collect_generate( @@ -1051,6 +1058,7 @@ mod tests { prompt_token_ids: vec![10, 20], sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, }; let response = collect_generate( @@ -1086,6 +1094,7 @@ mod tests { ec_transfer_params: None, sampling_mask: None, spec_decode_metrics: None, + weight_versions: None, prompt_token_ids, }; diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 0a37640d63a1..6352b187250a 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -9,6 +9,7 @@ use serde_json::{Map, Value}; use validator::Validate; use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; +use vllm_llm::WeightVersionSpan; use vllm_text::SamplingParams; use crate::routes::openai::utils::types::{Normalizable, StreamOptions, Usage}; @@ -73,6 +74,7 @@ pub(super) struct GenerateResponseChoice { pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, + pub weight_versions: Option>, } /// Mirrors the Python vLLM `GenerateResponseStreamChoice` class. @@ -84,6 +86,8 @@ pub(super) struct GenerateResponseStreamChoice { pub finish_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, + /// Set on the final chunk only; covers the whole output. + pub weight_versions: Option>, } /// Mirrors the Python vLLM `GenerateStreamResponse` class. From 5bed289707f3b647c903eb6911a2ca6e20ebf74c Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Mon, 12 Oct 2026 16:06:47 +0800 Subject: [PATCH 391/952] [DX] Update CODEOWNERS for entrypoints (#60984) Signed-off-by: DarkLight1337 --- .github/CODEOWNERS | 51 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a92ab6a2c963..01dab4e76b52 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -34,21 +34,48 @@ /vllm/utils/argparse_utils.py # Entrypoints -/vllm/entrypoints/anthropic @mgoin @DarkLight1337 +/vllm/entrypoints/anthropic @chaunceyjiang @mgoin @DarkLight1337 +/vllm/entrypoints/cohere @DarkLight1337 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 /vllm/entrypoints/mcp @heheda12345 -/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 -/vllm/entrypoints/speech_to_text/realtime @njhill +/vllm/entrypoints/openai/chat_completion @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/completion @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/models @chaunceyjiang @DarkLight1337 +/vllm/entrypoints/openai/parser @aarnphm @chaunceyjiang @sfeng33 +/vllm/entrypoints/openai/responses @aarnphm @chaunceyjiang @sfeng33 +/vllm/entrypoints/pooling @noooop @DarkLight1337 +/vllm/entrypoints/rl @aoshen02 @njhill +/vllm/entrypoints/scale_out @aoshen02 @noooop @DarkLight1337 /vllm/entrypoints/speech_to_text @NickLucche -/vllm/entrypoints/pooling @noooop -/vllm/entrypoints/serve/sagemaker @DarkLight1337 +/vllm/entrypoints/speech_to_text/realtime @njhill @NickLucche /vllm/entrypoints/serve @njhill +/vllm/entrypoints/serve/sagemaker @DarkLight1337 /vllm/entrypoints/common @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 +/tests/entrypoints @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas +/tests/entrypoints/anthropic @mgoin @DarkLight1337 +/tests/entrypoints/cohere @DarkLight1337 +/tests/entrypoints/cli @hmellor @mgoin @DarkLight1337 +/tests/entrypoints/openai/chat_completion @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/completion @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/models @chaunceyjiang @DarkLight1337 +/tests/entrypoints/openai/parser @aarnphm @chaunceyjiang @sfeng33 +/tests/entrypoints/openai/responses @aarnphm @chaunceyjiang @sfeng33 +/tests/entrypoints/pooling @noooop @DarkLight1337 +/tests/entrypoints/rl @aoshen02 @njhill +/tests/entrypoints/scale_out @noooop @DarkLight1337 +/tests/entrypoints/speech_to_text @NickLucche +/tests/entrypoints/speech_to_text/realtime @njhill @NickLucche +/tests/entrypoints/serve @njhill +/tests/entrypoints/serve/sagemaker @DarkLight1337 +/tests/entrypoints/llm @DarkLight1337 +/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera + # Rust Frontend /rust/ @BugenZhao @njhill /rust/src/bench @esmeetu +/rust/src/server/src/routes/inference @aoshen02 @BugenZhao @njhill /tools/build_rust.sh @BugenZhao @njhill /rust-toolchain.toml @BugenZhao @njhill /.buildkite/test_areas/rust* @BugenZhao @njhill @@ -62,6 +89,12 @@ /vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning /vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tokenizers_ @DarkLight1337 @njhill +/tests/renderers @DarkLight1337 @njhill +/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning +/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning + # vLLM V1 /vllm/v1/attention @LucasWilkinson @MatthewBonanni /vllm/v1/attention/backend.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @njhill @@ -99,7 +132,6 @@ /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao -/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas /tests/evals @mgoin @vadiklyutiy @AndreasKaratzas /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa @@ -107,7 +139,6 @@ /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas -/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera /tests/v1/structured_output @mgoin @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/lora @jeejeelee @@ -117,9 +148,6 @@ /tests/v1/kv_offload @ApostaC @orozery /tests/v1/simple_kv_offload @ivanium /tests/v1/determinism @yewentao256 -/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning -/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning -/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor @@ -221,15 +249,12 @@ mkdocs.yaml @hmellor /examples/pooling @noooop /docs/models/pooling_models @noooop /tests/models/*/pooling* @noooop -/tests/entrypoints/pooling @noooop /vllm/config/pooler.py @noooop /vllm/model_executor/layers/pooler @noooop # Reinforcement learning /vllm/distributed/weight_transfer @aoshen02 -/vllm/entrypoints/rl @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 -/tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 /examples/rl @aoshen02 /docs/training/rlhf.md @aoshen02 From 99b44c7a4282e903edc0c4703f50d3a72f4d419b Mon Sep 17 00:00:00 2001 From: Thuongvu Ho Date: Mon, 12 Oct 2026 01:31:07 -0700 Subject: [PATCH 392/952] [Metrics] Add admission rejection metrics (#55812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Thuongvu Ho Co-authored-by: Nicolò Lucchesi --- docs/design/metrics.md | 1 + .../serve/instrumentator/test_metrics.py | 1 + tests/v1/engine/test_admission_control.py | 79 ++++++++++++++++++- vllm/v1/engine/async_llm.py | 11 ++- vllm/v1/metrics/loggers.py | 31 ++++++++ 5 files changed, 119 insertions(+), 4 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 3e070f588890..5d9fa42c11c3 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -30,6 +30,7 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met - `vllm:prompt_tokens_total` (Counter) - Total number of prompt tokens processed. - `vllm:generation_tokens_total` (Counter) - Total number of generated tokens. - `vllm:request_success_total` (Counter) - Number of finished requests (by finish reason). +- `vllm:admission_rejections_total` (Counter) - Admission attempts rejected by a queue limit, by reason. - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts. - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts. - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT). diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index ac4576fbb95d..a685b9efb374 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -188,6 +188,7 @@ async def test_metrics_counts( "vllm:prefix_cache_queries", "vllm:prefix_cache_hits", "vllm:num_preemptions_total", + "vllm:admission_rejections_total", "vllm:prompt_tokens_total", "vllm:generation_tokens_total", "vllm:iteration_tokens_total", diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index 15be166c8c0e..dabd42711ca1 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -6,6 +6,7 @@ These tests cover: - OutputProcessor.get_num_queued_tokens() token counting - AsyncLLM.check_admission() admission control logic +- Admission rejection Prometheus samples - Exception classes (GracefulHTTPError, QueueOverflowError, MaxQueuedTokensError) - create_error_response() mapping GracefulHTTPError to HTTP 503 - SchedulerConfig field defaults and validation @@ -14,6 +15,7 @@ import argparse import asyncio +import logging import multiprocessing from collections.abc import Awaitable, Callable from http import HTTPStatus @@ -21,6 +23,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from prometheus_client import REGISTRY from pydantic import ValidationError from vllm.config.scheduler import SchedulerConfig @@ -40,6 +43,10 @@ from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor +from vllm.v1.metrics.loggers import StatLoggerManager + +_ADMISSION_METRIC_MODEL = "admission-metric-test" +_ADMISSION_LOGGER = "vllm.v1.engine.async_llm" pytestmark = pytest.mark.cpu_test @@ -70,10 +77,43 @@ def _make_async_llm( llm.output_processor = MagicMock() llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + llm.logger_manager = None llm.admission_stats = None return llm +def _attach_prometheus_logger( + llm: AsyncLLM, engine_indexes: list[int] | None = None +) -> None: + """Attach a real Prometheus logger manager to a lightweight AsyncLLM.""" + vllm_config = MagicMock() + vllm_config.observability_config.show_hidden_metrics = False + vllm_config.observability_config.kv_cache_metrics = False + vllm_config.observability_config.custom_histogram_buckets = None + vllm_config.model_config.served_model_name = _ADMISSION_METRIC_MODEL + vllm_config.model_config.max_model_len = 2048 + vllm_config.model_config.is_diffusion = False + vllm_config.speculative_config = None + vllm_config.kv_transfer_config = None + vllm_config.ec_transfer_config = None + vllm_config.lora_config = None + llm.logger_manager = StatLoggerManager( + vllm_config=vllm_config, + engine_idxs=engine_indexes, + enable_default_loggers=False, + ) + + +def _admission_rejection_count(reason: str) -> float: + """Read the Prometheus sample for one admission rejection reason.""" + value = REGISTRY.get_sample_value( + "vllm:admission_rejections_total", + {"model_name": _ADMISSION_METRIC_MODEL, "reason": reason}, + ) + assert value is not None + return value + + def _make_output_processor(**request_states) -> OutputProcessor: op = OutputProcessor.__new__(OutputProcessor) op.request_states = request_states @@ -157,11 +197,13 @@ def test_shared_request_snapshot_enforces_global_limit(): server_0.set_num_requests(2) llm = _make_async_llm(max_num_queued_reqs=2) llm.admission_stats = server_1 + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission() assert server_0.get_num_requests() == 2 + assert _admission_rejection_count("max_num_queued_reqs") == 1 # --------------------------------------------------------------------------- @@ -247,19 +289,35 @@ def test_admission_no_limits_allows_everything(): def test_admission_reqs_allows_when_under_limit(): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=5) + _attach_prometheus_logger(llm) + llm.check_admission() llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 0 + assert _admission_rejection_count("max_num_queued_tokens") == 0 -def test_admission_reqs_rejects_at_limit(): +def test_admission_reqs_rejects_at_limit(caplog): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) - with pytest.raises(QueueOverflowError): + _attach_prometheus_logger(llm, engine_indexes=[0, 1, 2]) + with ( + caplog.at_level(logging.INFO, logger=_ADMISSION_LOGGER), + pytest.raises(QueueOverflowError), + ): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 + assert not any( + "Request queue full" in record.getMessage() for record in caplog.records + ) def test_admission_reqs_rejects_with_n(): llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=8) + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission(3) + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 def test_admission_reqs_allows_n_at_boundary(): @@ -283,8 +341,11 @@ def test_admission_tokens_allows_when_under_limit(): def test_admission_tokens_rejects_at_limit(): llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1000) + _attach_prometheus_logger(llm) with pytest.raises(MaxQueuedTokensError): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 0 + assert _admission_rejection_count("max_num_queued_tokens") == 1 def test_admission_tokens_rejects_over_limit(): @@ -320,8 +381,11 @@ def test_admission_req_limit_checked_before_token_limit(): num_unfinished=10, num_queued_tokens=1000, ) + _attach_prometheus_logger(llm) with pytest.raises(QueueOverflowError): llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 # -- n derived from params at the add_request call site --------------------- @@ -381,6 +445,17 @@ async def add_request_async(_request): assert len(request_states) == 1 +def test_admission_preflight_pass_then_reject_records_once(): + llm = _make_async_llm(max_num_queued_reqs=1, num_unfinished=0) + _attach_prometheus_logger(llm) + llm.check_admission() + llm.output_processor.get_num_unfinished_requests.return_value = 1 + with pytest.raises(QueueOverflowError): + llm.check_admission() + assert _admission_rejection_count("max_num_queued_reqs") == 1 + assert _admission_rejection_count("max_num_queued_tokens") == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("first_n", [1, 3]) async def test_parallel_admission_is_all_or_nothing(first_n: int): diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 60588e076097..f2142af2b3fb 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -348,7 +348,7 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: else self.get_num_unfinished_requests() ) if current_requests + n > max_num_reqs: - logger.info( + logger.debug( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, @@ -356,21 +356,28 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: n, max_num_reqs, ) + self._record_admission_rejection("max_num_queued_reqs") raise QueueOverflowError() max_queued_tokens = self.scheduler_config.max_num_queued_tokens if max_queued_tokens is not None: current_tokens = self.get_num_queued_tokens() if current_tokens >= max_queued_tokens: - logger.info( + logger.debug( "Max queued tokens reached - rejecting request %s " "(current_tokens=%d, max=%d).", request_id, current_tokens, max_queued_tokens, ) + self._record_admission_rejection("max_num_queued_tokens") raise MaxQueuedTokensError() + def _record_admission_rejection(self, reason: str) -> None: + """Increment the frontend admission rejection counter if logging is on.""" + if self.logger_manager is not None: + self.logger_manager.record_admission_rejection(reason) + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: if not hasattr(self, "_supported_tasks"): # Cache the result diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 7ff1d2993cad..c42d17790926 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -43,6 +43,9 @@ WAITING_REASON_CAPACITY = "capacity" WAITING_REASON_DEFERRED = "deferred" +# User-facing reason labels for frontend admission rejections. +ADMISSION_REASONS = ("max_num_queued_reqs", "max_num_queued_tokens") + PerEngineStatLoggerFactory = Callable[[VllmConfig, int], "StatLoggerBase"] AggregateStatLoggerFactory = type["AggregateStatLoggerBase"] StatLoggerFactory = AggregateStatLoggerFactory | PerEngineStatLoggerFactory @@ -102,6 +105,9 @@ class AggregateStatLoggerBase(StatLoggerBase): @abstractmethod def __init__(self, vllm_config: VllmConfig, engine_indexes: list[int]): ... + def record_admission_rejection(self, reason: str) -> None: # noqa + """Record a frontend admission rejection. Optional for custom loggers.""" + class LoggingStatLogger(StatLoggerBase): def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): @@ -816,6 +822,23 @@ def __init__( "request_num_preemptions", overrides=custom_buckets ) + # + # Frontend admission + # + counter_admission_rejections = self._counter_cls( + name="vllm:admission_rejections", + documentation=( + "Number of admission attempts rejected by a queue limit, by reason." + ), + labelnames=["model_name", "reason"], + ) + self.admission_rejections = { + reason: counter_admission_rejections.labels( + model_name=model_name, reason=reason + ) + for reason in ADMISSION_REASONS + } + # # Histograms of counts # @@ -1276,6 +1299,9 @@ def record( finished_request.max_tokens_param ) + def record_admission_rejection(self, reason: str) -> None: + self.admission_rejections[reason].inc() + def record_sleep_state(self, sleep: int = 0, level: int = 0): awake = 1 discard_all = 0 @@ -1383,6 +1409,11 @@ def record_sleep_state(self, sleep: int = 0, level: int = 0): for logger in self.stat_loggers: logger.record_sleep_state(sleep, level) + def record_admission_rejection(self, reason: str) -> None: + """Record one frontend admission rejection on aggregate loggers.""" + for stat_logger in self.stat_loggers: + stat_logger.record_admission_rejection(reason) + def log(self): for logger in self.stat_loggers: logger.log() From e05201aa2e5ccf19284e5062064050a06d50e281 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Mon, 12 Oct 2026 16:59:38 +0800 Subject: [PATCH 393/952] [XPU] refine memory release for sleep mode (#61158) Signed-off-by: Yan Ma --- vllm/device_allocator/xpumem.py | 26 ++++++++----------- .../model_loader/sharded_state_loader.py | 7 +---- 2 files changed, 12 insertions(+), 21 deletions(-) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index 7d4b40a78e76..5eb836008d68 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -146,6 +146,7 @@ def __init__(self): self.pointer_to_data: dict[int, AllocationData] = {} self.current_tag: str = XpuMemAllocator.default_tag self.allocator_and_pools: dict[str, list[Any]] = {} + self._pending_allocators: list[Any] = [] self.python_malloc_callback = self._python_malloc_callback self.python_free_callback = self._python_free_callback @@ -291,24 +292,17 @@ def wake_up(self, tags: list[str] | None = None) -> None: data.cpu_backup_tensor = None def release_pools(self) -> None: - """Drop Python references to MemPool/pluggable allocators eagerly. - - This prevents pool destruction from being deferred to interpreter - finalization, which can happen after parts of XPU runtime are already - torn down. - """ - if not self.allocator_and_pools: + """Release pools, retaining allocators until their allocations are freed.""" + if not self.allocator_and_pools and not self._pending_allocators: return - # Note: keep allocators alive while MemPool objects are destroyed. - # MemPool teardown may invoke allocator virtual methods (e.g. raw_delete) - # when releasing cached blocks. If allocator wrappers are dropped first, - # C++ can hit "pure virtual method called" during shutdown. pool_entries = [e for es in self.allocator_and_pools.values() for e in es] self.allocator_and_pools.clear() mem_pools = [entry[0] for entry in pool_entries] - allocators = [entry[1] for entry in pool_entries] + # Live tensors can outlast their MemPool wrappers. Keep allocators alive + # across calls so deferred frees cannot invoke a destroyed allocator. + self._pending_allocators.extend(entry[1] for entry in pool_entries) pool_entries.clear() xpu_sync = getattr(torch.xpu, "synchronize", None) @@ -318,12 +312,14 @@ def release_pools(self) -> None: except Exception: logger.debug("torch.xpu.synchronize() failed during release_pools") - # Phase 1: drop MemPool refs while allocators are still strongly held. mem_pools.clear() gc.collect() + torch.xpu.empty_cache() - # Phase 2: now it is safe to release allocator wrappers. - allocators.clear() + # Free callbacks remove allocation records, including cached blocks. + # If any remain, a later release_pools() call must retry this cleanup. + if not self.pointer_to_data: + self._pending_allocators.clear() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/model_executor/model_loader/sharded_state_loader.py b/vllm/model_executor/model_loader/sharded_state_loader.py index 97a31e8e879a..9e8e6e9cfc75 100644 --- a/vllm/model_executor/model_loader/sharded_state_loader.py +++ b/vllm/model_executor/model_loader/sharded_state_loader.py @@ -211,12 +211,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # TODO: support un-sharded checkpoints too raise ValueError( f"Could not find checkpoint files " - f"'{ - os.path.join( - local_model_path, - rendered_pattern, - ) - }', " + f"'{os.path.join(local_model_path, rendered_pattern)}', " f"only pre-sharded checkpoints are currently supported!" ) state_dict = self._filter_subtensors(model.state_dict()) From a33d46b47193e3273dc38f980abbf55459f3c77f Mon Sep 17 00:00:00 2001 From: zhangchengzhucufe-dev Date: Mon, 12 Oct 2026 17:10:55 +0800 Subject: [PATCH 394/952] [Bugfix][XPU] Forward softcap and alibi_slopes in xpu_ops.flash_attn_varlen_func (#58936) Signed-off-by: zhangchengzhucufe-dev Signed-off-by: Kunshang Ji Co-authored-by: Kunshang Ji --- vllm/_xpu_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 1f00bc681d11..a384ac721572 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -1087,8 +1087,8 @@ def flash_attn_varlen_func( block_table=block_table, s_aux=s_aux, window_size=real_window_size, - # alibi_slopes = alibi_slopes, - # softcap=softcap, + alibi_slopes=alibi_slopes, + softcap=softcap, return_softmax_lse=return_softmax_lse, q_descale=q_descale, k_descale=k_descale, From 83773ecdf3ce6da9a2f2bd14cb87e6bbf58d4229 Mon Sep 17 00:00:00 2001 From: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:22:55 +0800 Subject: [PATCH 395/952] [Frontend] Support CustomToolCall in the non-Harmony Responses API (#57322) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Co-authored-by: Chauncey --- .../openai/responses/test_protocol.py | 35 +++++ .../openai/responses/test_responses_utils.py | 88 ++++++++++++ .../openai/responses/test_streaming_events.py | 116 ++++++++++++++- tests/tool_parsers/test_gemma4_tool_parser.py | 16 +++ .../test_k2_horizon_tool_parser.py | 12 ++ .../test_kimi_k3_named_tool_choice.py | 24 ++++ .../test_poolside_v1_tool_parser.py | 20 ++- .../test_structural_tag_registry.py | 15 +- tests/tool_use/test_muse_glimmer.py | 13 ++ tests/tool_use/test_tool_choice_required.py | 35 +++++ vllm/entrypoints/openai/responses/protocol.py | 38 ++++- .../openai/responses/streaming_events.py | 136 ++++++++++++++---- vllm/entrypoints/openai/responses/utils.py | 94 ++++++++++-- vllm/parser/abstract_parser.py | 20 ++- vllm/parser/response_template.py | 5 +- .../tool_parsers/gemma4_engine_tool_parser.py | 9 +- vllm/tool_parsers/k2_horizon_tool_parser.py | 4 +- vllm/tool_parsers/kimi_k3_tool_parser.py | 8 +- vllm/tool_parsers/muse_glimmer_tool_parser.py | 9 +- vllm/tool_parsers/poolside_v1_tool_parser.py | 9 +- vllm/tool_parsers/structural_tag_registry.py | 15 +- vllm/tool_parsers/utils.py | 49 ++++++- 22 files changed, 698 insertions(+), 72 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_protocol.py b/tests/entrypoints/openai/responses/test_protocol.py index db5d7d692490..04b07992528b 100644 --- a/tests/entrypoints/openai/responses/test_protocol.py +++ b/tests/entrypoints/openai/responses/test_protocol.py @@ -1,13 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest from openai_harmony import ( Message, ) +from pydantic import ValidationError from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, serialize_message, serialize_messages, ) +from vllm.exceptions import VLLMValidationError def test_serialize_message() -> None: @@ -37,3 +41,34 @@ def test_serialize_messages() -> None: } msg = Message.from_dict(msg_value) assert serialize_messages([msg, dict_value]) == [msg_value, dict_value] + + +def test_custom_tool_grammar_format_rejected() -> None: + grammar = {"type": "grammar", "syntax": "lark", "definition": 'start: "pwd"'} + with pytest.raises(VLLMValidationError) as exc_info: + ResponsesRequest( + input="hi", tools=[{"type": "custom", "name": "emit", "format": grammar}] + ) + assert exc_info.value.parameter == "tools" + + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit", "format": {"type": "text"}}], + ) + assert request.tools is not None and request.tools[0].type == "custom" + + +@pytest.mark.parametrize( + "tools", + [ + "abc", + {"type": "custom"}, + [None], + [{"type": "custom", "name": "x", "format": "text"}], + ], +) +def test_malformed_tools_rejected_by_field_validation(tools) -> None: + """Malformed tools bodies fail pydantic field validation (4xx), not an + AttributeError inside the custom-format validator (500).""" + with pytest.raises(ValidationError): + ResponsesRequest(input="hi", tools=tools) diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 967ddacda6de..6789ddc08bad 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest +from openai.types.responses import CustomTool, FunctionTool, ResponseCustomToolCall from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_tool_call_output_item import ( ResponseFunctionToolCallOutputItem, @@ -16,10 +17,14 @@ Summary, ) +from vllm.entrypoints.generate.base.protocol import FunctionCall from vllm.entrypoints.openai.responses.utils import ( _construct_message_from_response_item, + build_response_output_items, construct_chat_messages_with_tool_call, construct_input_messages, + decode_custom_tool_input, + decode_custom_tool_input_prefix, should_continue_final_message, ) from vllm.exceptions import VLLMValidationError @@ -248,6 +253,89 @@ def test_construct_chat_messages_preserves_single_item_conversions(self): assert formatted_item["content"] == "dongyi" +class TestCustomTools: + def test_custom_tool_call_output_item(self): + tools = [ + CustomTool(type="custom", name="emit_command", format={"type": "text"}), + FunctionTool( + type="function", name="get_weather", parameters={}, strict=None + ), + ] + custom, function = build_response_output_items( + reasoning=None, + content=None, + tool_calls=[ + FunctionCall( + id="call_1", name="emit_command", arguments='{"input": "pwd"}' + ), + FunctionCall( + id="call_2", name="get_weather", arguments='{"city": "Paris"}' + ), + ], + tools=tools, + ) + assert isinstance(custom, ResponseCustomToolCall) + assert (custom.call_id, custom.input) == ("call_1", "pwd") + assert custom.id.startswith("ctc_") + assert isinstance(function, ResponseFunctionToolCall) + assert function.arguments == '{"city": "Paris"}' + + def test_custom_tool_history_replays_through_function_shim(self): + messages = construct_chat_messages_with_tool_call( + [ + make_reasoning_item(content_text="Run pwd."), + ResponseCustomToolCall( + type="custom_tool_call", + id="ctc_1", + call_id="call_exec", + name="emit_command", + input='echo "hi"', + ), + { + "type": "custom_tool_call_output", + "call_id": "call_exec", + "output": "hi", + }, + {"role": "user", "content": "Continue"}, + ] + ) + assert messages == [ + { + "role": "assistant", + "reasoning": "Run pwd.", + "tool_calls": [ + { + "id": "call_exec", + "type": "function", + "function": { + "name": "emit_command", + "arguments": '{"input": "echo \\"hi\\""}', + }, + } + ], + }, + {"role": "tool", "content": "hi", "tool_call_id": "call_exec"}, + {"role": "user", "content": "Continue"}, + ] + + def test_decode_custom_tool_input(self): + assert decode_custom_tool_input('{"input": "pwd"}') == "pwd" + assert decode_custom_tool_input('{"cmd": "pwd"}') == '{"cmd": "pwd"}' + assert decode_custom_tool_input('{"file": "a.py"}') == '{"file": "a.py"}' + assert decode_custom_tool_input('{"input": 7}') == '{"input": 7}' + assert decode_custom_tool_input("pwd") == "pwd" + + def test_decode_custom_tool_input_prefix(self): + assert decode_custom_tool_input_prefix('{"inp') == "" + assert decode_custom_tool_input_prefix('{"input": "ls \\"dir') == 'ls "dir' + assert decode_custom_tool_input_prefix('{"input": "a\\') == "a" + assert decode_custom_tool_input_prefix('{"input": "\\ud83d') == "" + assert ( + decode_custom_tool_input_prefix('{"input": "\\ud83d\\ude00!"}') + == "\U0001f600!" + ) + + class TestReasoningItemContentPriority: """Tests that content is prioritized over summary for reasoning items.""" diff --git a/tests/entrypoints/openai/responses/test_streaming_events.py b/tests/entrypoints/openai/responses/test_streaming_events.py index 21fe5537bb74..a37109f85134 100644 --- a/tests/entrypoints/openai/responses/test_streaming_events.py +++ b/tests/entrypoints/openai/responses/test_streaming_events.py @@ -2,7 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from openai.types.responses import ResponseFunctionWebSearch, ResponseStreamEvent +from openai.types.responses import ( + CustomTool, + ResponseCustomToolCall, + ResponseFunctionWebSearch, + ResponseStreamEvent, +) from openai_harmony import Message, Role from pydantic import TypeAdapter @@ -20,6 +25,7 @@ emit_reasoning_done_events, split_delta, ) +from vllm.entrypoints.openai.responses.utils import decode_custom_tool_input def test_browser_find_uses_responses_action_type(): @@ -190,3 +196,111 @@ def test_reasoning_to_content_transition(self): types = [e.type for e in events] assert "response.reasoning_text.delta" in types assert "response.output_text.delta" in types + + +class TestSimpleToolCallShapes: + def test_arguments_done_emitted_without_deltas(self): + processor = SimpleStreamingEventProcessor() + events = _run_through_processor( + processor, DeltaMessage(tool_calls=[_make_tool_call(0, name="noop")]) + ) + events += processor.close_current() + assert [e.type for e in events] == [ + "response.output_item.added", + "response.function_call_arguments.done", + "response.output_item.done", + ] + + +class TestCustomToolStreaming: + def test_custom_tool_input_events(self): + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + events = _run_through_processor( + processor, + DeltaMessage( + tool_calls=[ + _make_tool_call(0, name="emit_command", arguments='{"input": "ls ') + ] + ), + ) + events += _run_through_processor( + processor, + DeltaMessage(tool_calls=[_make_tool_call(0, arguments='-la \\"x\\""}')]), + ) + events += processor.close_current() + + assert [e.type for e in events] == [ + "response.output_item.added", + "response.custom_tool_call_input.delta", + "response.custom_tool_call_input.delta", + "response.custom_tool_call_input.done", + "response.output_item.done", + ] + done = events[-1].item + assert isinstance(done, ResponseCustomToolCall) + assert done.input == 'ls -la "x"' + assert "".join(e.delta for e in events[1:3]) == done.input + assert events[3].input == done.input + assert events[0].item.type == "custom_tool_call" + assert events[0].item.id == done.id + assert done.id.startswith("ctc_") + assert events[0].item.call_id == done.call_id + + def test_unrelated_keys_are_not_streamed_as_payload(self): + """Review repro: chunks of an object whose input key arrives last must + not leak fragments of other keys into the payload.""" + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + events = _run_through_processor( + processor, + DeltaMessage( + tool_calls=[ + _make_tool_call(0, name="emit_command", arguments='{"a": "xy"') + ] + ), + ) + events += _run_through_processor( + processor, + DeltaMessage( + tool_calls=[_make_tool_call(0, arguments=', "input": "hello"}')] + ), + ) + events += processor.close_current() + + deltas = [ + e for e in events if e.type == "response.custom_tool_call_input.delta" + ] + done_input = next( + e for e in events if e.type == "response.custom_tool_call_input.done" + ).input + item = events[-1].item + assert "".join(e.delta for e in deltas) == "hello" # "xy" never leaks + assert done_input == "hello" + assert item.input == "hello" + assert decode_custom_tool_input('{"a": "xy", "input": "hello"}') == "hello" + + def test_done_reported_authoritative_payload_on_mismatch(self): + """If the final decode is not an extension of what was streamed, the + done event and item carry the authoritative value and a correcting + delta is emitted rather than dropping the divergence.""" + processor = SimpleStreamingEventProcessor( + tools=[CustomTool(type="custom", name="emit_command")] + ) + processor.state.current_state = _StateType.TOOL_CALL + processor.state.tool_call_is_custom = True + processor.state.tool_call_payload = "xy" # stale fragment already streamed + processor.state.accumulated_text = '{"a": "xy", "input": "hello"}' + + events = processor.close_current() + delta = next( + e for e in events if e.type == "response.custom_tool_call_input.delta" + ) + done_input = next( + e for e in events if e.type == "response.custom_tool_call_input.done" + ).input + assert delta.delta == "hello" # full authoritative resend, not a dropped suffix + assert done_input == "hello" + assert events[-1].item.input == "hello" diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 97ed0cd69d4c..43fac974052c 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -8,6 +8,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.gemma4 import ( TOOL_CALL_END, TOOL_CALL_START, @@ -117,6 +118,21 @@ def mock_request(): return request +def test_named_custom_tool_choice_skips_structured_outputs(mock_tokenizer): + # Responses custom tools ride the function shim, so a named custom choice + # must route like a named function choice, not into guided JSON. + tools = [{"type": "custom", "name": "emit_command"}] + request = ResponsesRequest( + input="hi", + tools=tools, + tool_choice={"type": "custom", "name": "emit_command"}, + ) + Gemma4EngineToolParser(mock_tokenizer, tools=tools).adjust_request(request) + + assert request.skip_special_tokens is False + assert request.structured_outputs is None + + # --------------------------------------------------------------------------- # Unit tests for _parse_gemma4_args (shared parser logic) # --------------------------------------------------------------------------- diff --git a/tests/tool_parsers/test_k2_horizon_tool_parser.py b/tests/tool_parsers/test_k2_horizon_tool_parser.py index 6fd4844aeb34..b57b50553e7e 100644 --- a/tests/tool_parsers/test_k2_horizon_tool_parser.py +++ b/tests/tool_parsers/test_k2_horizon_tool_parser.py @@ -8,6 +8,7 @@ from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.parser_manager import ParserManager from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.k2_horizon_tool_parser import K2HorizonToolParser @@ -316,3 +317,14 @@ def test_whitespace_only_surrounding_content_is_not_preserved(tokenizer): def test_tool_parser_registered(): assert ToolParserManager.get_tool_parser("k2_horizon") is K2HorizonToolParser assert K2HorizonToolParser.supports_required_and_named is False + + +def test_named_tool_resolves_custom_choice(): + # Custom tools ride the function shim; a ToolChoiceCustom must resolve to + # its name like ToolChoiceFunction does. + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + assert K2HorizonToolParser._named_tool(request) == "emit_command" diff --git a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py index 36b04b67b195..d426f623da34 100644 --- a/tests/tool_parsers/test_kimi_k3_named_tool_choice.py +++ b/tests/tool_parsers/test_kimi_k3_named_tool_choice.py @@ -8,6 +8,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.exceptions import VLLMValidationError from vllm.sampling_params import StructuredOutputsParams @@ -60,3 +61,26 @@ def test_named_choice_allowed_with_structural_tag(): def test_named_choice_rejected_without_structural_tag(): with pytest.raises(VLLMValidationError): _parser().adjust_request(_request(with_tag=False)) + + +def _custom_request(with_tag: bool) -> ResponsesRequest: + req = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + if with_tag: + req.structured_outputs = StructuredOutputsParams( + structural_tag='{"type": "structural_tag", "format": {}}' + ) + return req + + +def test_named_custom_choice_allowed_with_structural_tag(): + req = _parser().adjust_request(_custom_request(with_tag=True)) + assert req.skip_special_tokens is False + + +def test_named_custom_choice_rejected_without_structural_tag(): + with pytest.raises(VLLMValidationError): + _parser().adjust_request(_custom_request(with_tag=False)) diff --git a/tests/tool_parsers/test_poolside_v1_tool_parser.py b/tests/tool_parsers/test_poolside_v1_tool_parser.py index de6e751e5c91..635081b43b37 100644 --- a/tests/tool_parsers/test_poolside_v1_tool_parser.py +++ b/tests/tool_parsers/test_poolside_v1_tool_parser.py @@ -78,12 +78,15 @@ def _build_chat_request(*, tool_choice: str | dict[str, Any]) -> ChatCompletionR def _build_responses_request( - *, tool_choice: str | dict[str, Any], include: list[str] | None = None + *, + tool_choice: str | dict[str, Any], + include: list[str] | None = None, + tools: list[dict[str, Any]] | None = None, ) -> ResponsesRequest: return ResponsesRequest( model="poolside-test", input=[{"role": "user", "content": "write the file"}], - tools=[_responses_write_file_tool()], + tools=tools if tools is not None else [_responses_write_file_tool()], tool_choice=tool_choice, stream=True, max_output_tokens=200, @@ -145,6 +148,19 @@ def test_named_skips_structured_outputs_responses() -> None: assert request.skip_special_tokens is False +def test_named_custom_tool_choice_skips_structured_outputs_responses() -> None: + # A named custom choice parses to ToolChoiceCustom; custom tools ride the + # function shim, so it must route like ToolChoiceFunction. + request = _build_responses_request( + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + PoolsideV1ToolParser(_StubTokenizer()).adjust_request(request) + + assert request.text is None + assert request.skip_special_tokens is False + + def test_auto_still_keeps_special_tokens() -> None: request = _build_chat_request(tool_choice="auto") _make_parser(request).adjust_request(request) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 817c48a4401d..52d0d824691a 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -7,7 +7,7 @@ from unittest.mock import MagicMock import pytest -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import CustomTool, ToolChoiceFunction from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from xgrammar import Grammar, StructuralTag from xgrammar import get_model_structural_tag as get_xgrammar_builtin_structural_tag @@ -787,6 +787,19 @@ def test_get_model_structural_tag_supports_named_tool_choice( assert isinstance(tag, StructuralTag) +def test_glm_4_7_forced_custom_tool_uses_function_shim(): + tools = [CustomTool(type="custom", name="emit_command", description="Emit.")] + tag = get_model_structural_tag( + model="glm_4_7", + tools=tools, + tool_choice="required", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + assert "emit_command" in tag.model_dump_json() + + @pytest.mark.parametrize( ("parser_cls", "model"), [ diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py index 5badd9b03162..2fbf3cff49c0 100644 --- a/tests/tool_use/test_muse_glimmer.py +++ b/tests/tool_use/test_muse_glimmer.py @@ -29,6 +29,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser @@ -400,3 +401,15 @@ def test_exact_match_kept(): T, _call("get_weather"), _req("get_weather") ) assert out.tool_calls[0].function.name == "get_weather" + + +def test_named_custom_tool_choice_skips_structured_outputs(): + # Custom tools ride the function shim; a named custom choice must skip + # the base adjust_request like a named function choice does. + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + T.adjust_request(request) + assert request.structured_outputs is None diff --git a/tests/tool_use/test_tool_choice_required.py b/tests/tool_use/test_tool_choice_required.py index be58ffa83aa9..aedada111aa5 100644 --- a/tests/tool_use/test_tool_choice_required.py +++ b/tests/tool_use/test_tool_choice_required.py @@ -6,7 +6,9 @@ import pytest import regex as re from openai.types.responses import ( + CustomTool, FunctionTool, + ToolChoiceCustom, ToolChoiceFunction, WebSearchTool, ) @@ -21,8 +23,10 @@ from vllm.exceptions import VLLMValidationError from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.tool_parsers.utils import ( + custom_tool_parameters, find_tool_properties, get_json_schema_from_tools, + iter_response_function_tool_dicts, ) pytestmark = pytest.mark.cpu_test @@ -483,6 +487,37 @@ def test_get_json_schema_rejects_only_non_function_tools(self): get_json_schema_from_tools(tools=[WEB_SEARCH_TOOL], tool_choice="required") +class TestCustomToolShim: + """``custom`` tools are shown to the model as a single-string function.""" + + def test_rendered_as_single_string_function(self): + tool = CustomTool( + type="custom", name="emit_command", description="Emit a command." + ) + [tool_dict] = iter_response_function_tool_dicts([tool]) + assert tool_dict["type"] == "function" + assert tool_dict["name"] == "emit_command" + assert tool_dict["description"] == "Emit a command." + assert tool_dict["parameters"] == custom_tool_parameters() + assert tool_dict["parameters"]["required"] == ["input"] + + def test_forced_tool_choice_sees_custom_tool(self): + tools = [CustomTool(type="custom", name="emit_command"), FUNCTION_TOOL] + schema = get_json_schema_from_tools(tools=tools, tool_choice="required") + assert isinstance(schema, dict) + names = [e["properties"]["name"]["enum"] for e in schema["items"]["anyOf"]] + assert names == [["emit_command"], ["get_weather"]] + named = get_json_schema_from_tools( + tools=tools, + tool_choice=ToolChoiceCustom(type="custom", name="emit_command"), + ) + assert named == custom_tool_parameters() + assert ( + find_tool_properties(tools, "emit_command") + == custom_tool_parameters()["properties"] + ) + + class TestMalformedToolSchemaDefs: """Malformed `$defs` in caller-supplied tool parameters is a 400, not a 500.""" diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 373cb35df057..200054a4ef59 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -14,6 +14,8 @@ ResponseCodeInterpreterCallInterpretingEvent, ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, ResponseFunctionToolCall, ResponseInputItemParam, ResponseMcpCallArgumentsDeltaEvent, @@ -699,6 +701,34 @@ def input_item_parsing(cls, data): data["input"] = processed_input return data + @model_validator(mode="before") + @classmethod + def check_custom_tool_format(cls, data): + if not isinstance(data, dict): + return data + tools = data.get("tools") + if not isinstance(tools, list): + return data + for tool in tools: + if isinstance(tool, dict): + tool_dict = tool + elif hasattr(tool, "model_dump"): + tool_dict = tool.model_dump() + else: + continue + tool_format = tool_dict.get("format") + if ( + tool_dict.get("type") == "custom" + and isinstance(tool_format, dict) + and tool_format.get("type") == "grammar" + ): + raise VLLMValidationError( + "Custom tools with a grammar format are not supported; " + "only free-form text custom tools are accepted.", + parameter="tools", + ) + return data + @model_validator(mode="before") @classmethod def check_tool_usage(cls, data): @@ -708,9 +738,9 @@ def check_tool_usage(cls, data): tools = data.get("tools") tool_choice = data.get("tool_choice", "auto") has_tools = tools is not None and len(tools) > 0 - is_named_tool_choice = ( - isinstance(tool_choice, dict) and tool_choice.get("type") == "function" - ) + is_named_tool_choice = isinstance(tool_choice, dict) and tool_choice.get( + "type" + ) in ("function", "custom") if not has_tools: if tool_choice in ("auto", "none"): @@ -913,6 +943,8 @@ class ResponseInProgressEvent(OpenAIResponseInProgressEvent): | ResponseOutputItemDoneEvent | ResponseContentPartAddedEvent | ResponseContentPartDoneEvent + | ResponseCustomToolCallInputDeltaEvent + | ResponseCustomToolCallInputDoneEvent | ResponseReasoningTextDeltaEvent | ResponseReasoningTextDoneEvent | ResponseCodeInterpreterCallInProgressEvent diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index b101040a57a2..49107661a79c 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -29,6 +29,9 @@ ResponseCodeInterpreterToolCallParam, ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, + ResponseCustomToolCall, + ResponseCustomToolCallInputDeltaEvent, + ResponseCustomToolCallInputDoneEvent, ResponseFunctionCallArgumentsDeltaEvent, ResponseFunctionCallArgumentsDoneEvent, ResponseFunctionToolCall, @@ -73,6 +76,9 @@ ) from vllm.entrypoints.openai.responses.utils import ( build_responses_tool_call_name_map, + custom_tool_names, + decode_custom_tool_input, + decode_custom_tool_input_prefix, resolve_responses_tool_call_name, ) from vllm.outputs import CompletionOutput @@ -821,6 +827,8 @@ class SimpleStreamingState: tool_call_name: str = "" tool_call_namespace: str | None = None tool_call_index: int | None = None + tool_call_is_custom: bool = False + tool_call_payload: str = "" current_state: _StateType = field(default_factory=lambda: _StateType.NONE) @@ -1039,38 +1047,80 @@ def emit_simple_tool_call_open( name: str, index: int | None, namespace: str | None = None, + is_custom: bool = False, call_id: str | None = None, ) -> list[StreamingResponsesResponse]: state.current_state = _StateType.TOOL_CALL - state.current_item_id = f"fc_{random_uuid()}" + state.current_item_id = ( + f"ctc_{random_uuid()}" if is_custom else f"fc_{random_uuid()}" + ) state.tool_call_id = call_id or make_tool_call_id() state.tool_call_name = name state.tool_call_namespace = namespace state.tool_call_index = index + state.tool_call_is_custom = is_custom + state.tool_call_payload = "" state.accumulated_text = "" + item: ResponseOutputItem + if is_custom: + item = ResponseCustomToolCall( + type="custom_tool_call", + id=state.current_item_id, + call_id=state.tool_call_id, + name=name, + input="", + ) + else: + item = ResponseFunctionToolCallItem( + type="function_call", + id=state.current_item_id, + call_id=state.tool_call_id, + name=name, + namespace=namespace, + arguments="", + status="in_progress", + ) return [ ResponseOutputItemAddedEvent( type="response.output_item.added", sequence_number=-1, output_index=state.output_index, - item=ResponseFunctionToolCallItem( - type="function_call", - id=state.current_item_id, - call_id=state.tool_call_id, - name=name, - namespace=namespace, - arguments="", - status="in_progress", - ), + item=item, ), ] +def _emit_custom_tool_input_delta( + state: SimpleStreamingState, + payload: str, +) -> list[StreamingResponsesResponse]: + streamed = state.tool_call_payload + if payload == streamed: + return [] + state.tool_call_payload = payload + # A payload that is not an extension of what was streamed cannot be + # retracted, so the authoritative value is resent in full. + delta = payload[len(streamed) :] if payload.startswith(streamed) else payload + return [ + ResponseCustomToolCallInputDeltaEvent( + type="response.custom_tool_call_input.delta", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + delta=delta, + ) + ] + + def emit_simple_tool_call_delta( state: SimpleStreamingState, delta: str, ) -> list[StreamingResponsesResponse]: state.accumulated_text += delta + if state.tool_call_is_custom: + return _emit_custom_tool_input_delta( + state, decode_custom_tool_input_prefix(state.accumulated_text) + ) return [ ResponseFunctionCallArgumentsDeltaEvent( type="response.function_call_arguments.delta", @@ -1085,32 +1135,58 @@ def emit_simple_tool_call_delta( def emit_simple_tool_call_done( state: SimpleStreamingState, ) -> list[StreamingResponsesResponse]: - events: list[StreamingResponsesResponse] = [ - ResponseFunctionCallArgumentsDoneEvent( - type="response.function_call_arguments.done", - sequence_number=-1, - output_index=state.output_index, - item_id=state.current_item_id, - arguments=state.accumulated_text, + events: list[StreamingResponsesResponse] + item: ResponseOutputItem + if state.tool_call_is_custom: + payload = decode_custom_tool_input(state.accumulated_text) + events = _emit_custom_tool_input_delta(state, payload) + events.append( + ResponseCustomToolCallInputDoneEvent( + type="response.custom_tool_call_input.done", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + input=payload, + ) + ) + item = ResponseCustomToolCall( + type="custom_tool_call", + id=state.current_item_id, + call_id=state.tool_call_id, name=state.tool_call_name, - ), + input=payload, + ) + else: + events = [ + ResponseFunctionCallArgumentsDoneEvent( + type="response.function_call_arguments.done", + sequence_number=-1, + output_index=state.output_index, + item_id=state.current_item_id, + arguments=state.accumulated_text, + name=state.tool_call_name, + ) + ] + item = ResponseFunctionToolCall( + type="function_call", + name=state.tool_call_name, + namespace=state.tool_call_namespace, + arguments=state.accumulated_text, + status="completed", + id=state.current_item_id, + call_id=state.tool_call_id, + ) + events.append( ResponseOutputItemDoneEvent( type="response.output_item.done", sequence_number=-1, output_index=state.output_index, - item=ResponseFunctionToolCall( - type="function_call", - name=state.tool_call_name, - namespace=state.tool_call_namespace, - arguments=state.accumulated_text, - status="completed", - id=state.current_item_id, - call_id=state.tool_call_id, - ), - ), - ] + item=item, + ) + ) state.output_index += 1 state.tool_call_namespace = None + state.tool_call_is_custom = False state.current_state = _StateType.NONE return events @@ -1194,6 +1270,7 @@ def __init__( ) -> None: self.state = state or SimpleStreamingState() self.tool_call_name_map = build_responses_tool_call_name_map(tools) + self.custom_tool_names = custom_tool_names(tools) self.output_items: list[ResponseOutputItem] = [] def resolve_target_state( @@ -1264,6 +1341,7 @@ def open( call_name.name, tool_call.index, call_name.namespace, + call_name.name in self.custom_tool_names, tool_call.id, ) return handlers.open_fn(self.state) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index c4523964e9dd..d9b71d684922 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from collections.abc import Iterable from typing import Any +import partial_json_parser from openai.types.chat import ( ChatCompletionAssistantMessageParam, ChatCompletionMessageToolCallParam, @@ -13,6 +15,7 @@ Function as FunctionCallTool, ) from openai.types.responses import ( + ResponseCustomToolCall, ResponseFunctionToolCall, ResponseOutputItem, ResponseOutputMessage, @@ -28,6 +31,7 @@ Content as ResponseReasoningTextContent, ) from openai.types.responses.tool import Tool +from partial_json_parser.core.options import Allow from vllm import envs from vllm.entrypoints.generate.base.protocol import FunctionCall, FunctionDefinition @@ -40,6 +44,7 @@ from vllm.logger import init_logger from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.utils import ( + CUSTOM_TOOL_INPUT_KEY, build_responses_tool_call_name_map, flat_namespace_tool_name, iter_response_function_tool_dicts, @@ -50,6 +55,51 @@ logger = init_logger(__name__) +def custom_tool_names(tools: list[Tool] | None) -> frozenset[str]: + return frozenset(tool.name for tool in tools or [] if tool.type == "custom") + + +def encode_custom_tool_input(payload: str) -> str: + return json.dumps({CUSTOM_TOOL_INPUT_KEY: payload}, ensure_ascii=False) + + +def _custom_tool_payload(parsed: Any) -> str | None: + if not isinstance(parsed, dict): + return None + value = parsed.get(CUSTOM_TOOL_INPUT_KEY) + return value if isinstance(value, str) else None + + +def decode_custom_tool_input(arguments: str) -> str: + """Payload of a completed shim call; bare text is passed through.""" + try: + payload = _custom_tool_payload(json.loads(arguments)) + except ValueError: + return arguments + return arguments if payload is None else payload + + +def decode_custom_tool_input_prefix(arguments: str) -> str: + """Longest decodable payload prefix of partial shim arguments, so that + streamed ``custom_tool_call_input`` deltas stay a prefix of the final input.""" + try: + parsed = partial_json_parser.loads(arguments, Allow.STR | Allow.OBJ) + except (ValueError, IndexError, AssertionError): + return "" + payload = _custom_tool_payload(parsed) or "" + if payload and "\ud800" <= payload[-1] <= "\udbff": + payload = payload[:-1] + return payload + + +def _item_type(item: ResponseInputOutputItem) -> str | None: + return item.get("type") if isinstance(item, dict) else getattr(item, "type", None) + + +def _item_dict(item: ResponseInputOutputItem) -> dict[str, Any]: + return item if isinstance(item, dict) else item.model_dump() + + def build_response_output_items( reasoning: str | None, content: str | None, @@ -59,6 +109,7 @@ def build_response_output_items( ) -> list[ResponseOutputItem]: outputs: list[ResponseOutputItem] = [] tool_call_name_map = build_responses_tool_call_name_map(tools) + custom_names = custom_tool_names(tools) if reasoning: outputs.append( @@ -93,14 +144,27 @@ def build_response_output_items( if tool_calls: for idx, tool_call in enumerate(tool_calls): + call_id = tool_call.id or make_tool_call_id( + func_name=tool_call.name, idx=idx + ) + if tool_call.name in custom_names: + outputs.append( + ResponseCustomToolCall( + id=f"ctc_{random_uuid()}", + call_id=call_id, + type="custom_tool_call", + name=tool_call.name, + input=decode_custom_tool_input(tool_call.arguments), + ) + ) + continue call_name = resolve_responses_tool_call_name( tool_call.name, tool_call_name_map=tool_call_name_map ) outputs.append( ResponseFunctionToolCall( id=f"fc_{random_uuid()}", - call_id=tool_call.id - or make_tool_call_id(func_name=tool_call.name, idx=idx), + call_id=call_id, type="function_call", status="completed", name=call_name.name, @@ -255,6 +319,7 @@ def _construct_message_from_response_item( prev_msg if prev_msg and prev_msg.get("role") == "assistant" else None ) + tool_call: ChatCompletionMessageToolCallParam | None = None if isinstance(item, ResponseFunctionToolCall): tool_name = item.name if item.namespace: @@ -267,6 +332,17 @@ def _construct_message_from_response_item( ), type="function", ) + elif _item_type(item) == "custom_tool_call": + call = _item_dict(item) + tool_call = ChatCompletionMessageToolCallParam( + id=call["call_id"], + function=FunctionCallTool( + name=call["name"], + arguments=encode_custom_tool_input(call.get("input") or ""), + ), + type="function", + ) + if tool_call is not None: if prev_assistant_msg: tool_calls = prev_assistant_msg.get("tool_calls") if tool_calls is None: @@ -285,14 +361,14 @@ def _construct_message_from_response_item( logger.warning( "Previous assistant message has unknown tool_calls format. " "Tool call merging is skipped and a new assistant message is created. " - "Item %s", - item.id, + "Call %s", + tool_call["id"], ) return ChatCompletionAssistantMessageParam( role="assistant", tool_calls=[tool_call], ) - elif isinstance(item, ResponseReasoningItem): + if isinstance(item, ResponseReasoningItem): reasoning = "" if item.encrypted_content: raise VLLMValidationError( @@ -336,12 +412,12 @@ def _construct_message_from_response_item( content=item.output, tool_call_id=item.call_id, ) - elif isinstance(item, dict) and item.get("type") == "function_call_output": - # Append the function call output as a tool message. + elif _item_type(item) in ("function_call_output", "custom_tool_call_output"): + output = _item_dict(item) return ChatCompletionToolMessageParam( role="tool", - content=item.get("output"), - tool_call_id=item.get("call_id"), + content=output.get("output"), + tool_call_id=output.get("call_id"), ) elif isinstance(item, dict) and item.get("role") == "assistant": content = item.get("content") diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 61826dc8c86c..12a9c2369843 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -7,7 +7,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field, replace -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from pydantic import TypeAdapter, ValidationError from xgrammar import Grammar, StructuralTag from xgrammar.structural_tag import ( @@ -359,7 +359,9 @@ def extract_reasoning( def _get_function_name( self, request: ChatCompletionRequest | ResponsesRequest ) -> str: - if request.tool_choice and isinstance(request.tool_choice, ToolChoiceFunction): + if request.tool_choice and isinstance( + request.tool_choice, (ToolChoiceFunction, ToolChoiceCustom) + ): return request.tool_choice.name if request.tool_choice and isinstance( request.tool_choice, ChatCompletionNamedToolChoiceParam @@ -398,7 +400,7 @@ def _extract_tool_calls( supports_required_and_named = tool_parser.supports_required_and_named is_named_tool_choice = request.tool_choice and isinstance( request.tool_choice, - (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam), + (ToolChoiceFunction, ToolChoiceCustom, ChatCompletionNamedToolChoiceParam), ) is_required_tool_choice = request.tool_choice == "required" is_auto_tool_choice = enable_auto_tools and ( @@ -503,7 +505,11 @@ def _apply_structural_tag( or request.tool_choice == "required" or isinstance( request.tool_choice, - (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) ) if not need_tool_calling: @@ -723,7 +729,11 @@ def _extract_tool_calls_streaming( and request.tool_choice and isinstance( request.tool_choice, - (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam), + ( + ToolChoiceFunction, + ToolChoiceCustom, + ChatCompletionNamedToolChoiceParam, + ), ) ): delta_message, function_name_returned = extract_named_tool_call_streaming( diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index 85f8b67efbd6..b1665f1772a5 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -9,7 +9,7 @@ from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any, cast -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -173,7 +173,8 @@ def _streaming_template(template: dict[str, Any]) -> ResponseTemplate: def _is_named_tool_choice(request: ChatCompletionRequest | ResponsesRequest) -> bool: return isinstance( - request.tool_choice, (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam) + request.tool_choice, + (ToolChoiceFunction, ToolChoiceCustom, ChatCompletionNamedToolChoiceParam), ) diff --git a/vllm/tool_parsers/gemma4_engine_tool_parser.py b/vllm/tool_parsers/gemma4_engine_tool_parser.py index 04c03ecaa20a..fcb578b129cb 100644 --- a/vllm/tool_parsers/gemma4_engine_tool_parser.py +++ b/vllm/tool_parsers/gemma4_engine_tool_parser.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -29,7 +29,12 @@ def adjust_request( if request.tools: tc = request.tool_choice if tc == "required" or isinstance( - tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tc, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ): request.skip_special_tokens = False return request diff --git a/vllm/tool_parsers/k2_horizon_tool_parser.py b/vllm/tool_parsers/k2_horizon_tool_parser.py index 742bf7d14013..a63f5fa12239 100644 --- a/vllm/tool_parsers/k2_horizon_tool_parser.py +++ b/vllm/tool_parsers/k2_horizon_tool_parser.py @@ -6,7 +6,7 @@ from typing import Any import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -184,7 +184,7 @@ def _named_tool(request: ToolParserRequest) -> str | None: tool_choice = request.tool_choice if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): return tool_choice.function.name - if isinstance(tool_choice, ToolChoiceFunction): + if isinstance(tool_choice, (ToolChoiceFunction, ToolChoiceCustom)): return tool_choice.name return None diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 76c6b0aa657c..15b5a22df861 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -35,7 +35,7 @@ from collections.abc import Sequence import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from vllm.entrypoints.generate.base.protocol import ( DeltaFunctionCall, @@ -149,7 +149,11 @@ def adjust_request( ) -> ChatCompletionRequest | ResponsesRequest: named = isinstance( request.tool_choice, - (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) structured_outputs = getattr(request, "structured_outputs", None) has_structural_tag = ( diff --git a/vllm/tool_parsers/muse_glimmer_tool_parser.py b/vllm/tool_parsers/muse_glimmer_tool_parser.py index 5a0b26cf423a..644ca5ed3e3a 100644 --- a/vllm/tool_parsers/muse_glimmer_tool_parser.py +++ b/vllm/tool_parsers/muse_glimmer_tool_parser.py @@ -38,7 +38,7 @@ from collections.abc import Iterator, Sequence import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from transformers import PreTrainedTokenizerBase from vllm.entrypoints.generate.base.protocol import ( @@ -228,7 +228,12 @@ def adjust_request( if request.tools and ( tool_choice == "required" or isinstance( - tool_choice, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tool_choice, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ) ): return request diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index d2dfb7815fbe..998927264757 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -16,7 +16,7 @@ import partial_json_parser.core.complete import regex as re -from openai.types.responses import ToolChoiceFunction +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction from partial_json_parser.core.options import Allow from vllm.entrypoints.generate.base.protocol import ( @@ -171,7 +171,12 @@ def adjust_request( if request.tools: tc = request.tool_choice if tc == "required" or isinstance( - tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + tc, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), ): request.skip_special_tokens = False return request diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index f3148a827eb0..7e5902b8c2f6 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Sequence from typing import Any, Literal, TypeAlias -from openai.types.responses import FunctionTool +from openai.types.responses import CustomTool, FunctionTool, ToolChoiceCustom from openai.types.responses.response import ToolChoice as ResponsesToolChoice from openai.types.responses.tool import Tool as ResponsesTool from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed @@ -39,7 +39,11 @@ ChatCompletionToolsParam, ) from vllm.tool_parsers.tool_strict_level import ToolStrictLevel -from vllm.tool_parsers.utils import get_function_tools, require_function_tools +from vllm.tool_parsers.utils import ( + custom_tool_parameters, + get_function_tools, + require_function_tools, +) ToolChoice: TypeAlias = ( Literal["none", "auto", "required"] @@ -287,6 +291,11 @@ def _dump_tool_for_xgrammar( if tool.strict is not None: function["strict"] = tool.strict return {"type": "function", "function": function} + if isinstance(tool, CustomTool): + function = {"name": tool.name, "parameters": custom_tool_parameters()} + if tool.description is not None: + function["description"] = tool.description + return {"type": "function", "function": function} dumped_tool = tool.model_dump(mode="json", exclude_none=True) if isinstance(tool, ChatCompletionToolsParam): return dumped_tool @@ -306,7 +315,7 @@ def _dump_tool_choice_for_xgrammar( if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): return tool_choice.model_dump(mode="json", exclude_none=True) - if isinstance(tool_choice, ToolChoiceFunction): + if isinstance(tool_choice, (ToolChoiceFunction, ToolChoiceCustom)): return { "type": "function", "function": {"name": tool_choice.name}, diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 7a93c843b8e3..2462824ccfc0 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -13,8 +13,10 @@ import partial_json_parser from openai.types.responses import ( + CustomTool, FunctionTool, NamespaceTool, + ToolChoiceCustom, ToolChoiceFunction, ) from openai.types.responses.tool import Tool as ResponsesTool @@ -37,6 +39,32 @@ logger = init_logger(__name__) +CUSTOM_TOOL_INPUT_KEY = "input" + + +def custom_tool_parameters() -> dict[str, Any]: + """Schema of the single-string function shim a ``custom`` tool is shown as.""" + return { + "type": "object", + "properties": { + CUSTOM_TOOL_INPUT_KEY: { + "type": "string", + "description": "The freeform text payload for this tool.", + } + }, + "required": [CUSTOM_TOOL_INPUT_KEY], + "additionalProperties": False, + } + + +def custom_tool_as_function_dict(tool: CustomTool) -> dict[str, Any]: + return { + "type": "function", + "name": tool.name, + "description": tool.description, + "parameters": custom_tool_parameters(), + } + def safe_literal_eval(text: str): with warnings.catch_warnings(): @@ -188,8 +216,9 @@ def get_function_tools( This is the single source for both the chat template's tool list and the tools a tool-call grammar constrains, so the two cannot drift apart. - Namespace functions are flattened to ``__``; other - Responses tool types are never rendered for the model and are dropped. + Namespace functions are flattened to ``__``; custom tools + become a single-string function shim. Other Responses tool types are never + rendered for the model and are dropped. """ function_tools: list[ChatCompletionToolsParam | FunctionTool] = [] for tool in tools: @@ -206,6 +235,10 @@ def get_function_tools( for namespaced_tool in tool.tools if namespaced_tool.type == "function" ) + elif isinstance(tool, CustomTool): + function_tools.append( + FunctionTool.model_validate(custom_tool_as_function_dict(tool)) + ) elif isinstance(tool, (FunctionTool, ChatCompletionToolsParam)): function_tools.append(tool) return function_tools @@ -307,7 +340,7 @@ def find_tool_properties( if not tools: return {} for tool in tools: - if isinstance(tool, (FunctionTool, NamespaceTool)): + if isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): for name, params in iter_response_function_tool_info(tool): if name == tool_name: return (params or {}).get("properties", {}) @@ -328,7 +361,7 @@ def find_tool_name( if not tools: return False for tool in tools: - if isinstance(tool, (FunctionTool, NamespaceTool)): + if isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): for name, _ in iter_response_function_tool_info(tool): if name == tool_name: return True @@ -414,7 +447,9 @@ def _get_json_schema_from_tools( def get_json_schema_from_tools( - tool_choice: str | ToolChoiceFunction | ChatCompletionNamedToolChoiceParam, + tool_choice: ( + str | ToolChoiceFunction | ToolChoiceCustom | ChatCompletionNamedToolChoiceParam + ), tools: list[Tool] | None, parallel_tool_calls: bool | None = None, ) -> str | dict | None: @@ -423,12 +458,12 @@ def get_json_schema_from_tools( return None # tool_choice: Forced Function (Responses) if (not isinstance(tool_choice, str)) and isinstance( - tool_choice, ToolChoiceFunction + tool_choice, (ToolChoiceFunction, ToolChoiceCustom) ): tool_name = tool_choice.name responses_tool_map: dict[str, dict[str, Any] | None] = {} for tool in tools: - if not isinstance(tool, (FunctionTool, NamespaceTool)): + if not isinstance(tool, (FunctionTool, NamespaceTool, CustomTool)): continue for name, params in iter_response_function_tool_info(tool): responses_tool_map[name] = params From 0466bd36ea9f6300acbc15c705b1ebe61590c691 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 12 Oct 2026 17:43:47 +0800 Subject: [PATCH 396/952] [Rust Frontend] Build the Kimi K3 full-output grammar from its parser mode (#59190) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../chat/tests/grammar_replay/Kimi-K3.json | 1016 ++++++++++------- .../src/chat/tests/grammar_replay/Kimi-K3.txt | 99 +- rust/src/parser/src/output_grammar.rs | 19 +- rust/src/parser/src/unified/kimi_k3.rs | 14 +- .../src/parser/src/unified/kimi_k3/grammar.rs | 356 ++++++ .../src/unified/kimi_k3/grammar/arguments.rs | 184 +++ .../src/unified/kimi_k3/structural_tag.rs | 343 ------ rust/src/parser/src/unified/mod.rs | 2 +- 8 files changed, 1246 insertions(+), 787 deletions(-) create mode 100644 rust/src/parser/src/unified/kimi_k3/grammar.rs create mode 100644 rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs delete mode 100644 rust/src/parser/src/unified/kimi_k3/structural_tag.rs diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.json b/rust/src/chat/tests/grammar_replay/Kimi-K3.json index 7fe53b0ab2be..5306b936604f 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.json @@ -8,213 +8,413 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "optional", - "content": { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ - { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" - }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } - }, - "required": [ - "location" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null }, - "end": "<|close|>call<|sep|>" + "end": "<|close|>response<|sep|>" }, { - "type": "tag", - "begin": "<|open|>call tool=\"add\" index=\"", + "type": "optional", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ - { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "number" - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "number" + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "array", - "items": { - "type": "string" - } + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>argument<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] } ] }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "y": { - "type": "number" + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + } + ] + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "x": { - "type": "number" + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "items": { - "type": "array", + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, "items": { - "type": "string" + "type": "array", + "items": { + "type": "string" + } } - } + }, + "required": [ + "y", + "x", + "items" + ] }, - "required": [ - "y", - "x", - "items" - ] + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" - } + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -225,7 +425,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] @@ -235,210 +435,225 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ - { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } - }, - "required": [ - "location" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" - } - ] - } - ] - }, - "end": "<|close|>call<|sep|>" - }, - { - "type": "tag", - "begin": "<|open|>call tool=\"add\" index=\"", - "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", - "elements": [ + "type": "const_string", + "value": "\"<|sep|>" + }, { - "type": "sequence", + "type": "or", "elements": [ { "type": "tag", - "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", "content": { - "type": "json_schema", - "json_schema": { - "type": "number" - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null }, "end": "<|close|>argument<|sep|>" }, { "type": "tag", - "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "begin": "<|open|>json type=\"object\"<|sep|>", "content": { "type": "json_schema", "json_schema": { - "type": "number" + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] }, "style": "json", "any_order": false, "max_whitespace_cnt": null, "excludes": [] }, - "end": "<|close|>argument<|sep|>" + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] }, { "type": "tag", - "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "begin": "<|open|>json type=\"object\"<|sep|>", "content": { "type": "json_schema", "json_schema": { - "type": "array", - "items": { - "type": "string" - } + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] }, "style": "json", "any_order": false, "max_whitespace_cnt": null, "excludes": [] }, - "end": "<|close|>argument<|sep|>" + "end": "<|close|>json<|sep|>" } ] - }, - { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "y": { - "type": "number" - }, - "x": { - "type": "number" - }, - "items": { - "type": "array", - "items": { - "type": "string" - } - } - }, - "required": [ - "y", - "x", - "items" - ] - }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" } ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -449,7 +664,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] @@ -459,100 +674,115 @@ "format": { "type": "sequence", "elements": [ - { - "type": "optional", - "content": { - "type": "const_string", - "value": "<|open|>response<|sep|>" - } - }, { "type": "tag", "begin": "", "content": { "type": "any_text", "excludes": [ - "<|open|>", - "<|close|>", + "<|close|>think<|sep|>", "<|end_of_msg|>" ], "max_tokens": null, "max_chars": null }, - "end": "<|close|>response<|sep|>" + "end": "<|close|>think<|sep|>" }, { - "type": "tag", - "begin": "<|open|>tools<|sep|>", - "content": { - "type": "tags_with_separator", - "tags": [ - { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { "type": "tag", - "begin": "<|open|>call tool=\"get_weather\" index=\"", + "begin": "<|open|>response<|sep|>", "content": { - "type": "sequence", - "elements": [ - { - "type": "regex", - "pattern": "[1-9][0-9]*" - }, - { - "type": "const_string", - "value": "\"<|sep|>" - }, - { - "type": "or", + "type": "any_text", + "excludes": [ + "<|open|>", + "<|close|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", "elements": [ { - "type": "tag", - "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", - "content": { - "type": "any_text", - "excludes": [ - "<|close|>argument<|sep|>", - "<|close|>call<|sep|>" - ], - "max_tokens": null, - "max_chars": null - }, - "end": "<|close|>argument<|sep|>" + "type": "regex", + "pattern": "[1-9][0-9]*" }, { - "type": "tag", - "begin": "<|open|>json type=\"object\"<|sep|>", - "content": { - "type": "json_schema", - "json_schema": { - "type": "object", - "properties": { - "location": { - "type": "string" - } + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null }, - "required": [ - "location" - ] + "end": "<|close|>argument<|sep|>" }, - "style": "json", - "any_order": false, - "max_whitespace_cnt": null, - "excludes": [] - }, - "end": "<|close|>json<|sep|>" + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] } ] - } - ] - }, - "end": "<|close|>call<|sep|>" - } - ], - "separator": "", - "at_least_one": true, - "stop_after_first": false - }, - "end": "<|close|>tools<|sep|>" + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": true + }, + "end": "<|close|>tools<|sep|>" + } + ] }, { "type": "optional", @@ -563,7 +793,7 @@ } ] }, - "coverage": "final_output_only" + "coverage": "from_token_zero" }, "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt index 3aeb6126420e..6a1c814c131c 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt @@ -1,8 +1,29 @@ -# auto (final_output_only) +# auto (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - optional + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` tags_with_separator `` at_least_one tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` @@ -24,42 +45,44 @@ sequence tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` -# required (final_output_only) +# required (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` - tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - sequence - tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` - tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` - tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` -# named (final_output_only) +# named (from_token_zero) sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index 53c52538ee1d..0108c3193ef7 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -12,7 +12,8 @@ use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagOptio use xgrammar_structural_tag::format::Format; use xgrammar_structural_tag::tool::ToolChoiceValue; use xgrammar_structural_tag::{ - FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + FunctionDefinition, FunctionToolParam, NormalizedToolChoice, ToolChoice, ToolParam, + build_structural_tag, }; use crate::tool::Tool; @@ -218,6 +219,22 @@ fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec .collect() } +/// Normalize the request tools the way `build_structural_tag` does before +/// calling a crate builder, for model-owned grammar builders. Returns `None` +/// when the request asks for no tool grammar. +pub(crate) fn normalize_tool_choice( + ctx: &OutputGrammarContext<'_>, +) -> Result> { + if !tool_grammar_applies(ctx) { + return Ok(None); + } + + let tools = tool_params(ctx.tools, ctx.tool_strict_level); + let normalized = + xgrammar_structural_tag::normalize_tool_choice(&tools, ctx.tool_choice.clone())?; + Ok(Some(normalized)) +} + /// Whether the request asks for a tool grammar at all. fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { if ctx.tools.is_empty() { diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs index d6e32463c1fb..13d6fa42eb6f 100644 --- a/rust/src/parser/src/unified/kimi_k3.rs +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -46,9 +46,7 @@ //! now only applies to `argument` and `json` blocks inside a call body, which //! are still parsed from text. -mod structural_tag; - -pub use structural_tag::KimiK3StructuralTagBuilder; +mod grammar; use serde_json::{Map, Value}; use vllm_tokenizer::{DecodedText, DynTokenizer, Tokenizer}; @@ -58,11 +56,8 @@ use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; use winnow::prelude::*; use winnow::token::{literal, rest, take_till, take_until, take_while}; -use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; use super::{Result, UnifiedParser, UnifiedParserOutput, special_token}; -use crate::output_grammar::{ - self, BuiltOutputGrammar, OutputGrammarContext, visible_format_from_builder, -}; +use crate::output_grammar::{self, BuiltOutputGrammar, OutputGrammarContext}; use crate::tool::{Tool, ToolCallDelta}; use crate::unified::parsing_failed; use crate::utils::{ @@ -347,10 +342,7 @@ impl UnifiedParser for KimiK3UnifiedParser { &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - Ok( - visible_format_from_builder(Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER), ctx)? - .map(BuiltOutputGrammar::final_output_only), - ) + grammar::build_output_grammar(&self.mode, ctx) } fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { diff --git a/rust/src/parser/src/unified/kimi_k3/grammar.rs b/rust/src/parser/src/unified/kimi_k3/grammar.rs new file mode 100644 index 000000000000..166fbafb0f82 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/grammar.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Full-output grammar for Kimi K3 XTML generations. +//! +//! The grammar starts at the channel the prompt left open, which is the +//! parser's initialized [`KimiK3Mode`], and follows the same channel +//! transitions as the parser. A prompt that opened no channel may first emit a +//! whole `think` channel. +//! +//! | Tool choice | After `think` closes | After `response` closes | +//! | ----------- | ------------------------------ | ----------------------- | +//! | `auto` | `response [tools]`, or `tools` | `[tools]` | +//! | `required` | `[response] tools` | `tools` | +//! | named | `[response] tools`, one call | `tools`, one call | +//! +//! Every path may end with the message close, and a stop token is legal only +//! after the last channel closes. + +mod arguments; + +use xgrammar_structural_tag::NormalizedToolChoice; +use xgrammar_structural_tag::format::Format; +use xgrammar_structural_tag::tool::BuilderToolChoice; + +use self::arguments::call_tag; +use super::{ + CLOSE, END_OF_MSG, KimiK3Mode, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, +}; +use crate::output_grammar::{ + self, BuiltOutputGrammar, OutputGrammarContext, normalize_tool_choice, +}; + +/// Build the grammar of everything generated after a prompt that left the +/// parser in `mode`, or `None` when the request asks for no tool grammar. +pub(super) fn build_output_grammar( + mode: &KimiK3Mode, + ctx: &OutputGrammarContext<'_>, +) -> output_grammar::Result> { + let Some(NormalizedToolChoice { + function_tools, + choice, + .. + }) = normalize_tool_choice(ctx)? + else { + return Ok(None); + }; + + let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; + let tools = Format::tag( + TOOLS_OPEN, + Format::tags_with_separator( + function_tools.iter().map(call_tag).collect(), + "", + true, + single_call, + ), + TOOLS_CLOSE, + ); + + let mut elements = match mode { + KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(choice, tools)], + KimiK3Mode::Response => vec![response_body(), after_response(choice, tools)], + KimiK3Mode::Idle => vec![ + Format::optional(Format::tag(THINK_OPEN, think_text(), THINK_CLOSE)), + after_reasoning(choice, tools), + ], + KimiK3Mode::Epilogue | KimiK3Mode::Tools | KimiK3Mode::Call { .. } | KimiK3Mode::Done => { + unreachable!("initialization leaves the parser idle or inside `think` or `response`") + } + }; + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(Some(BuiltOutputGrammar::from_token_zero(Format::sequence( + elements, + )))) +} + +/// Channels that may follow the closed `think` channel. +fn after_reasoning(tool_choice: BuilderToolChoice, tools: Format) -> Format { + let response = Format::tag(RESPONSE_OPEN, response_text(), RESPONSE_CLOSE); + match tool_choice { + BuilderToolChoice::Auto => Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + tools, + ]), + BuilderToolChoice::Required | BuilderToolChoice::Forced => { + Format::sequence(vec![Format::optional(response), tools]) + } + } +} + +/// Channels that may follow the closed `response` channel. +fn after_response(tool_choice: BuilderToolChoice, tools: Format) -> Format { + match tool_choice { + BuilderToolChoice::Auto => Format::optional(tools), + BuilderToolChoice::Required | BuilderToolChoice::Forced => tools, + } +} + +/// The rest of the `think` channel the prompt opened. +fn think_body() -> Format { + Format::tag("", think_text(), THINK_CLOSE) +} + +/// The rest of the `response` channel the prompt opened. +fn response_body() -> Format { + Format::tag("", response_text(), RESPONSE_CLOSE) +} + +fn think_text() -> Format { + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]) +} + +fn response_text() -> Format { + // Keep marker-looking prefixes out of response text so the grammar must + // resolve them through a valid channel boundary. + Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]) +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + use xgrammar_structural_tag::ToolChoice; + use xgrammar_structural_tag::format::{EndBoundary, TagFormat}; + + use super::super::CALL_CLOSE; + use super::*; + use crate::output_grammar::test_utils::outline; + use crate::output_grammar::{GrammarCoverage, ToolStrictLevel}; + use crate::tool::Tool; + + fn tools() -> Vec { + ["get_weather", "add"] + .map(|name| Tool { + name: name.to_string(), + description: None, + parameters: json!({ "type": "object", "properties": {} }), + strict: None, + defer_loading: None, + }) + .into() + } + + fn check( + mode: KimiK3Mode, + tool_choice: ToolChoice, + parallel_tool_calls: bool, + expected: Expect, + ) { + let tools = tools(); + let grammar = build_output_grammar( + &mode, + &OutputGrammarContext { + tools: &tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Function, + parallel_tool_calls, + }, + ) + .unwrap() + .unwrap(); + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + let mut format = grammar.format; + elide_call_bodies(&mut format); + expected.assert_eq(&outline(&format)); + } + + /// Replace each call body with `..`, so that outlines show the channel + /// structure. The `arguments` tests cover call bodies. + fn elide_call_bodies(format: &mut Format) { + match format { + Format::Sequence(format) => format.elements.iter_mut().for_each(elide_call_bodies), + Format::Or(format) => format.elements.iter_mut().for_each(elide_call_bodies), + Format::Optional(format) => elide_call_bodies(&mut format.content), + Format::Tag(tag) => elide_tag(tag), + Format::TagsWithSeparator(format) => format.tags.iter_mut().for_each(elide_tag), + _ => {} + } + } + + fn elide_tag(tag: &mut TagFormat) { + if matches!(&tag.end, EndBoundary::Text(end) if end == CALL_CLOSE) { + *tag.content = Format::const_string(".."); + } else { + elide_call_bodies(&mut tag.content); + } + } + + #[test] + fn prompt_opened_think_channel() { + check( + KimiK3Mode::Reasoning, + ToolChoice::auto(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Reasoning, + ToolChoice::required(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Reasoning, + ToolChoice::function("add"), + true, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn prompt_opened_response_channel() { + check( + KimiK3Mode::Response, + ToolChoice::auto(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Response, + ToolChoice::required(), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check( + KimiK3Mode::Response, + ToolChoice::function("add"), + true, + expect![[r#" + sequence + tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn no_prompt_opened_channel() { + check( + KimiK3Mode::Idle, + ToolChoice::auto(), + true, + expect![[r#" + sequence + optional tag `<|open|>think<|sep|>` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn serial_tool_calls_stop_after_one_call() { + check( + KimiK3Mode::Reasoning, + ToolChoice::required(), + false, + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional tag `<|open|>response<|sep|>` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one stop_after_first + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn non_strict_auto_builds_no_grammar() { + let tools = tools(); + let grammar = build_output_grammar( + &KimiK3Mode::Reasoning, + &OutputGrammarContext { + tools: &tools, + tool_choice: &ToolChoice::auto(), + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + }, + ) + .unwrap(); + assert_eq!(grammar, None); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs b/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs new file mode 100644 index 000000000000..49587535aadd --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/grammar/arguments.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML tool-call grammar: call tags and their typed arguments. + +use serde_json::Value; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, TagFormat}; +use xgrammar_structural_tag::tool::{FunctionToolParam, function_parameters}; + +use super::super::{ARG_CLOSE, ARG_OPEN, CALL_CLOSE, JSON_CLOSE, JSON_OPEN, OPEN, SEP}; +use crate::output_grammar::arguments::{ + self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, ValueOption, group_by, one_of, +}; + +/// XTML arguments: +/// `<|open|>argument key="KEY" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>`, +/// adjacent to each other. +/// +/// The parser decodes each value by its `type` attribute, so every value option +/// sits under the attribute of its own JSON type. String values are raw text +/// and every other value is JSON. +struct XtmlArguments; + +impl ArgumentSyntax for XtmlArguments { + fn separator(&self) -> Option { + None + } + + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { + let escaped; + let key = match key { + ParameterKey::Declared(key) => { + escaped = escape_attr_value(key); + ParameterKey::Declared(&escaped) + } + ParameterKey::Free => ParameterKey::Free, + }; + let arguments = group_by(options, |option| xtml_type(option.ty)) + .into_iter() + .filter_map(|(type_name, options)| { + let values = options + .into_iter() + .filter_map(|option| match option.ty { + // The parser ends the call at its close marker, so a + // string value may contain neither close marker. + JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), + _ => Some(option.json()), + }) + .collect::>(); + (!values.is_empty()).then(|| { + let suffix = format!("\" type=\"{type_name}\"{SEP}"); + key.tag( + &format!("{ARG_OPEN} key=\""), + &suffix, + one_of(values), + ARG_CLOSE, + ) + }) + }) + .collect::>(); + (!arguments.is_empty()).then(|| one_of(arguments)) + } +} + +/// The `type` attribute K3 renders for a value of type `ty`, as +/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are +/// both `number`. +fn xtml_type(ty: JsonType) -> &'static str { + match ty { + JsonType::String => "string", + JsonType::Integer | JsonType::Number => "number", + JsonType::Boolean => "boolean", + JsonType::Null => "null", + JsonType::Object => "object", + JsonType::Array => "array", + } +} + +/// Build the tag for one call to `tool`. +pub(super) fn call_tag(tool: &FunctionToolParam) -> TagFormat { + let parameters = function_parameters(&tool.function); + let typed_arguments = + arguments::arguments(¶meters, &XtmlArguments, &ArgumentOptions::default()); + let call_body = Format::or(vec![typed_arguments, raw_json_arguments(¶meters)]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn raw_json_arguments(parameters: &Value) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + Format::JsonSchema(JsonSchemaFormat::new(parameters.clone())), + JSON_CLOSE, + ) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::{Expect, expect}; + use serde_json::json; + use xgrammar_structural_tag::FunctionDefinition; + use xgrammar_structural_tag::format::Format; + use xgrammar_structural_tag::tool::FunctionToolParam; + + use crate::output_grammar::test_utils::outline; + + fn check(parameters: serde_json::Value, expected: Expect) { + let tool = FunctionToolParam::new( + FunctionDefinition::new("get_weather").with_parameters(parameters), + ); + expected.assert_eq(&outline(&Format::Tag(super::call_tag(&tool)))); + } + + #[test] + fn call_tag_matches_xtml_arguments() { + check( + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + expect![[r#" + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + optional + tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>` + or + `celsius` + `fahrenheit` + tag `<|open|>argument key="place" type="object"<|sep|>` json({ city?: string } where place = { city?: string }) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>` + "#]], + ); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + check( + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + expect![[r#" + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + optional + or + tag `<|open|>argument key="count" type="number"<|sep|>` json(integer) `<|close|>argument<|sep|>` + tag `<|open|>argument key="count" type="null"<|sep|>` json(null) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ count?: integer | null }) `<|close|>json<|sep|>` + "#]], + ); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs deleted file mode 100644 index 4754726c1b5d..000000000000 --- a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs +++ /dev/null @@ -1,343 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -//! Structural-tag grammar for Kimi K3 XTML tool calls. - -use serde_json::Value; -use xgrammar_structural_tag::Result; -use xgrammar_structural_tag::builders::{ - ReasoningMode, StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, -}; -use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; -use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; - -use crate::output_grammar::arguments::{ - self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, ValueOption, group_by, one_of, -}; - -use super::{ - ARG_CLOSE, ARG_OPEN, CALL_CLOSE, CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, - RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, -}; - -pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = - KimiK3StructuralTagBuilder; - -/// XTML arguments: -/// `<|open|>argument key="KEY" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>`, -/// adjacent to each other. -/// -/// The parser decodes each value by its `type` attribute, so every value option -/// sits under the attribute of its own JSON type. String values are raw text -/// and every other value is JSON. -struct XtmlArguments; - -impl ArgumentSyntax for XtmlArguments { - fn separator(&self) -> Option { - None - } - - fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { - let escaped; - let key = match key { - ParameterKey::Declared(key) => { - escaped = escape_attr_value(key); - ParameterKey::Declared(&escaped) - } - ParameterKey::Free => ParameterKey::Free, - }; - let arguments = group_by(options, |option| xtml_type(option.ty)) - .into_iter() - .filter_map(|(type_name, options)| { - let values = options - .into_iter() - .filter_map(|option| match option.ty { - // The parser ends the call at its close marker, so a - // string value may contain neither close marker. - JsonType::String => option.raw_string(&[ARG_CLOSE, CALL_CLOSE]), - _ => Some(option.json()), - }) - .collect::>(); - (!values.is_empty()).then(|| { - let suffix = format!("\" type=\"{type_name}\"{SEP}"); - key.tag( - &format!("{ARG_OPEN} key=\""), - &suffix, - one_of(values), - ARG_CLOSE, - ) - }) - }) - .collect::>(); - (!arguments.is_empty()).then(|| one_of(arguments)) - } -} - -/// The `type` attribute K3 renders for a value of type `ty`, as -/// `_xtml_type` in the checkpoint's `encoding_k3.py`: integers and floats are -/// both `number`. -fn xtml_type(ty: JsonType) -> &'static str { - match ty { - JsonType::String => "string", - JsonType::Integer | JsonType::Number => "number", - JsonType::Boolean => "boolean", - JsonType::Null => "null", - JsonType::Object => "object", - JsonType::Array => "array", - } -} - -/// Kimi K3 XTML structural-tag builder. -#[derive(Debug, Clone, Copy, Default)] -pub struct KimiK3StructuralTagBuilder; - -impl StructuralTagBuilder for KimiK3StructuralTagBuilder { - fn build(&self, ctx: StructuralTagContext<'_>) -> Result { - let mut elements = response_prefix(ctx.options.reasoning); - - let tools = match ctx.tool_choice { - // Serving lowering filters empty tools before calling the builder, - // while direct `build_structural_tag(..., auto, ...)` calls do not. - BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, - BuilderToolChoice::Auto => Some(Format::optional(tools_channel( - ctx.function_tools, - ctx.options, - ))), - BuilderToolChoice::Forced | BuilderToolChoice::Required => { - Some(tools_channel(ctx.function_tools, ctx.options)) - } - }; - if let Some(tools) = tools { - elements.push(tools); - } - elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); - - Ok(StructuralTag::new(Format::sequence(elements))) - } -} - -fn response_prefix(reasoning: ReasoningMode) -> Vec { - let mut elements = Vec::new(); - if reasoning != ReasoningMode::Disabled { - let prefix = Format::tag( - if reasoning == ReasoningMode::Auto { - THINK_OPEN - } else { - "" - }, - Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), - THINK_CLOSE, - ); - elements.push(if reasoning == ReasoningMode::Auto { - Format::optional(prefix) - } else { - prefix - }); - elements.push(Format::const_string(RESPONSE_OPEN)); - } else { - elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); - } - elements.push(Format::tag( - "", - // Keep marker-looking prefixes out of response text so the grammar - // must resolve them through a valid channel boundary. - Format::any_text_excluding(&[OPEN, CLOSE, END_OF_MSG]), - RESPONSE_CLOSE, - )); - elements -} - -fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { - let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); - Format::tag( - TOOLS_OPEN, - Format::tags_with_separator(calls, "", true, false), - TOOLS_CLOSE, - ) -} - -fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { - let parameters = function_parameters(&tool.function); - let typed_arguments = arguments::arguments( - ¶meters, - &XtmlArguments, - &ArgumentOptions { - any_order: options.any_order, - max_whitespace_cnt: options.max_whitespace_cnt, - excludes: vec![], - }, - ); - let call_body = Format::or(vec![ - typed_arguments, - raw_json_arguments(¶meters, options), - ]); - - TagFormat::new( - format!( - "{OPEN}call tool=\"{}\" index=\"", - escape_attr_value(&tool.function.name) - ), - Format::sequence(vec![ - Format::regex("[1-9][0-9]*"), - Format::const_string(format!("\"{SEP}")), - call_body, - ]), - CALL_CLOSE, - ) -} - -fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { - Format::tag( - format!("{JSON_OPEN} type=\"object\"{SEP}"), - json_schema(parameters.clone(), options), - JSON_CLOSE, - ) -} - -fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { - Format::JsonSchema( - JsonSchemaFormat::new(schema) - .with_any_order(options.any_order) - .with_max_whitespace_cnt(options.max_whitespace_cnt), - ) -} - -fn escape_attr_value(value: &str) -> String { - value.replace('&', "&").replace('"', """) -} - -#[cfg(test)] -mod tests { - use expect_test::expect; - use serde_json::json; - use xgrammar_structural_tag::builders::StructuralTagOptions; - use xgrammar_structural_tag::{ - FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, - }; - - use super::KimiK3StructuralTagBuilder; - use crate::output_grammar::test_utils::outline; - - fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { - ToolParam::Function(FunctionToolParam::new( - FunctionDefinition::new(name).with_parameters(parameters), - )) - } - - #[test] - fn required_structural_tag_matches_xtml_channels() { - let tools = vec![tool( - "get_weather", - json!({ - "$defs": { - "place": { "type": "object", "properties": { "city": { "type": "string" } } } - }, - "type": "object", - "properties": { - "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, - "place": { "$ref": "#/$defs/place", "type": "object" } - }, - "required": ["place"] - }), - )]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::required(), - StructuralTagOptions::default().with_reasoning(false), - ) - .unwrap(); - - expect![[r#" - sequence - optional `<|open|>response<|sep|>` - tag `` text excluding [`<|open|>`, `<|close|>`, `<|end_of_msg|>`] `<|close|>response<|sep|>` - tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` - tags_with_separator `` at_least_one - tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` - sequence - /[1-9][0-9]*/ - `"<|sep|>` - or - sequence - optional - tag `<|open|>argument key="unit" type="string"<|sep|>` .. `<|close|>argument<|sep|>` - or - `celsius` - `fahrenheit` - tag `<|open|>argument key="place" type="object"<|sep|>` json({ city?: string } where place = { city?: string }) `<|close|>argument<|sep|>` - tag `<|open|>json type="object"<|sep|>` json({ unit?: "celsius" | "fahrenheit", place: place & object } where place = { city?: string }) `<|close|>json<|sep|>` - optional `<|close|>message<|sep|>` - "#]].assert_eq(&outline(&tag.format)); - } - - #[test] - fn reasoning_grammar_starts_inside_prefilled_think_channel() { - let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::auto(), - StructuralTagOptions::default().with_reasoning(true), - ) - .unwrap(); - let value = serde_json::to_value(tag).unwrap(); - - assert_eq!( - value["format"]["elements"][0]["end"], - "<|close|>think<|sep|>" - ); - assert_eq!( - value["format"]["elements"][1]["value"], - "<|open|>response<|sep|>" - ); - assert_eq!(value["format"]["elements"][3]["type"], "optional"); - } - - #[test] - fn forced_choice_keeps_only_the_named_tool() { - let tools = vec![ - tool("search", json!({ "type": "object" })), - tool("lookup", json!({ "type": "object" })), - ]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::function("lookup"), - StructuralTagOptions::default().with_reasoning(false), - ) - .unwrap() - .to_json_string() - .unwrap(); - - assert!(tag.contains("lookup")); - assert!(!tag.contains("search")); - } - - #[test] - fn union_argument_content_matches_its_xtml_type() { - let tools = vec![tool( - "set_count", - json!({ - "type": "object", - "properties": { - "count": { "type": ["integer", "null"] } - } - }), - )]; - let tag = build_structural_tag( - KimiK3StructuralTagBuilder, - &tools, - ToolChoice::required(), - StructuralTagOptions::default(), - ) - .unwrap() - .to_json_string() - .unwrap(); - - assert!(tag.contains(r#"type=\"number\""#)); - assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); - assert!(tag.contains(r#"type=\"null\""#)); - assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); - } -} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index aa0e92e6e17f..493b7b3f1b18 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -15,7 +15,7 @@ pub use gemma4::Gemma4UnifiedParser; pub use hf::{HfTemplateError, HfUnifiedParser, ResponseTemplate}; pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser}; pub use inkling::InklingUnifiedParser; -pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; +pub use kimi_k3::KimiK3UnifiedParser; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::{DecodedText, DynTokenizer}; From 5a0906263209da2c1165ce82b9fe8e767aa89669 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 12 Oct 2026 17:43:48 +0800 Subject: [PATCH 397/952] [Rust Frontend] Compose answer constraints into parser-owned output grammars (#60781) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/output/default/mod.rs | 55 ++ rust/src/chat/src/output/structural_tag.rs | 165 +++++- rust/src/chat/tests/grammar_replay.py | 93 ++- .../grammar_replay/DeepSeek-V3.2-Exp.json | 211 ++++++- .../grammar_replay/DeepSeek-V3.2-Exp.txt | 17 + .../grammar_replay/DeepSeek-V4-Flash.json | 211 ++++++- .../grammar_replay/DeepSeek-V4-Flash.txt | 17 + .../grammar_replay/DeepSeek-V4.1-Flash.json | 211 ++++++- .../grammar_replay/DeepSeek-V4.1-Flash.txt | 17 + .../chat/tests/grammar_replay/GLM-4.5.json | 68 +++ .../src/chat/tests/grammar_replay/GLM-4.5.txt | 7 + .../tests/grammar_replay/GLM-4.7-Flash.json | 209 ++++++- .../tests/grammar_replay/GLM-4.7-Flash.txt | 14 + .../tests/grammar_replay/GLM-5.2-FP8.json | 209 ++++++- .../chat/tests/grammar_replay/GLM-5.2-FP8.txt | 14 + rust/src/chat/tests/grammar_replay/Hy3.json | 321 +++++++++- rust/src/chat/tests/grammar_replay/Hy3.txt | 50 ++ .../chat/tests/grammar_replay/Kimi-K2.5.json | 188 +++++- .../chat/tests/grammar_replay/Kimi-K2.5.txt | 20 + .../chat/tests/grammar_replay/Kimi-K3.json | 549 +++++++++++++++++- .../src/chat/tests/grammar_replay/Kimi-K3.txt | 55 ++ .../grammar_replay/MiMo-V2.6-Flash-RL.json | 212 ++++++- .../grammar_replay/MiMo-V2.6-Flash-RL.txt | 14 + .../tests/grammar_replay/MiniMax-M2.5.json | 229 +++++++- .../tests/grammar_replay/MiniMax-M2.5.txt | 21 + .../chat/tests/grammar_replay/MiniMax-M3.json | 40 ++ .../chat/tests/grammar_replay/MiniMax-M3.txt | 2 + .../chat/tests/grammar_replay/Qwen3-0.6B.json | 228 +++++++- .../chat/tests/grammar_replay/Qwen3-0.6B.txt | 20 + .../chat/tests/grammar_replay/Qwen3.5-4B.json | 222 ++++++- .../chat/tests/grammar_replay/Qwen3.5-4B.txt | 18 + .../grammar_replay/Seed-OSS-36B-Instruct.json | 40 ++ .../grammar_replay/Seed-OSS-36B-Instruct.txt | 2 + rust/src/chat/tests/roundtrip.rs | 188 ++++-- rust/src/parser/src/output_grammar.rs | 9 + rust/src/parser/src/tool/minimax_m2.rs | 2 +- rust/src/parser/src/tool/mod.rs | 69 ++- rust/src/parser/src/unified/combined.rs | 114 +++- rust/src/parser/src/unified/hy/hy_v3.rs | 2 + rust/src/parser/src/unified/hy/hy_v4.rs | 1 + .../src/parser/src/unified/kimi_k3/grammar.rs | 254 ++++++-- 41 files changed, 4133 insertions(+), 255 deletions(-) create mode 100644 rust/src/chat/tests/grammar_replay/GLM-4.5.json create mode 100644 rust/src/chat/tests/grammar_replay/GLM-4.5.txt create mode 100644 rust/src/chat/tests/grammar_replay/MiniMax-M3.json create mode 100644 rust/src/chat/tests/grammar_replay/MiniMax-M3.txt create mode 100644 rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json create mode 100644 rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index b31c5e9e4cbe..981173dc7450 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -13,8 +13,10 @@ use vllm_parser::output_grammar::{BuiltOutputGrammar, OutputGrammarContext}; use vllm_parser::unified::{CombinedParser, UnifiedParser}; use vllm_text::tokenizer::DynTokenizer; use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::Format; use self::unified::unified_event_stream; +use super::structural_tag::answer_format; use super::structured::structured_chat_event_stream; use crate::error::Result; use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream}; @@ -48,6 +50,10 @@ struct GrammarInputs { tools: Vec, tool_choice: ToolChoice, tool_strict_level: ToolStrictLevel, + /// The request's structured-output constraint, normalized for composition + /// into the output grammar. The original stays on the request until an + /// output grammar replaces it. + answer: Option, } impl DefaultChatOutputProcessor { @@ -137,6 +143,7 @@ impl DefaultChatOutputProcessor { tools: request.tools().to_vec(), tool_choice: request.tool_choice().into(), tool_strict_level, + answer: request.sampling_params.structured_outputs.as_ref().and_then(answer_format), }), }) } @@ -258,6 +265,7 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor { tool_choice: &inputs.tool_choice, tool_strict_level: inputs.tool_strict_level, parallel_tool_calls: self.parallel_tool_calls, + answer: inputs.answer.as_ref(), }) .map_err(|error| Error::OutputGrammar { error: Box::new(error), @@ -283,7 +291,9 @@ mod tests { use std::sync::Arc; use thiserror_ext::AsReport as _; + use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_tokenizer::test_utils::TestTokenizer; + use xgrammar_structural_tag::format::Format; use super::DefaultChatOutputProcessor; use crate::output::ChatOutputProcessor; @@ -339,6 +349,51 @@ mod tests { assert!(build(ToolStrictLevel::Parameter, None, ChatToolChoice::None).is_none()); } + #[test] + fn answer_constraint_keeps_optional_tool_calls() { + let tools = vec![ChatTool { + name: "search".to_string(), + description: None, + parameters: serde_json::json!({"type": "object"}), + strict: None, + defer_loading: None, + }]; + let schema = serde_json::json!({"type": "object"}); + let build = |structured_outputs| { + let mut request = ChatRequest { + tool_context: ResolvedToolContext::new( + &[], + tools.clone(), + Some(ChatToolChoice::Auto), + true, + ) + .unwrap(), + ..ChatRequest::for_test() + }; + request.sampling_params.structured_outputs = structured_outputs; + let mut processor = DefaultChatOutputProcessor::new( + &mut request, + "other-model", + tokenizer(), + &ParserSelection::Explicit("qwen3_coder".to_string()), + &ParserSelection::None, + ToolStrictLevel::Auto, + ) + .unwrap(); + processor.initialize(&[]).unwrap(); + processor.build_output_grammar().unwrap() + }; + + // Non-strict `auto` alone needs no grammar; with an answer constraint + // the grammar holds the answer or a call. + assert!(build(None).is_none()); + let built = build(Some(StructuredOutputsParams::json(schema.clone()))).unwrap(); + let Format::Or(branches) = built.format else { + panic!("expected the calls or the answer, got {:?}", built.format); + }; + assert_eq!(branches.elements.last(), Some(&Format::json_schema(schema))); + } + #[test] fn tool_grammar_honors_parallel_call_policy() { for parallel in [false, true] { diff --git a/rust/src/chat/src/output/structural_tag.rs b/rust/src/chat/src/output/structural_tag.rs index dba6b6d1b739..b99a32d5bd3d 100644 --- a/rust/src/chat/src/output/structural_tag.rs +++ b/rust/src/chat/src/output/structural_tag.rs @@ -3,19 +3,91 @@ //! Applies parser-built xgrammar structural-tag constraints. +use serde::Deserialize; +use serde_json::{Value, json}; use vllm_engine_core_client::protocol::structured_outputs::{ - StructuredOutputBackend, StructuredOutputsParams, + StructuredOutputBackend, StructuredOutputConstraint, StructuredOutputsParams, }; use vllm_parser::output_grammar::{BuiltOutputGrammar, GrammarCoverage}; use vllm_text::TextRequest; use xgrammar_structural_tag::StructuralTag; +use xgrammar_structural_tag::format::{Format, GrammarFormat, TagFormat}; use crate::{Error, Result}; +/// Normalize the user's structured-output constraint to the answer format a +/// parser may compose into its output grammar. +/// +/// Like Python's `structured_outputs_to_format`, this ignores the request's +/// options: no backend reads them per request, only from the engine config. +/// +/// Returns `None` when the constraint must stay on its own wire field, so the +/// request keeps today's behavior: a grammar that is not XGrammar EBNF (such +/// as Lark), or a structural tag that does not parse. +pub(crate) fn answer_format(params: &StructuredOutputsParams) -> Option { + Some(match ¶ms.constraint { + StructuredOutputConstraint::Json(Value::String(schema)) => { + Format::json_schema(serde_json::from_str(schema).ok()?) + } + StructuredOutputConstraint::Json(schema) => Format::json_schema(schema.clone()), + StructuredOutputConstraint::JsonObject => Format::json_schema(json!({ "type": "object" })), + StructuredOutputConstraint::Regex(pattern) => Format::regex(pattern), + StructuredOutputConstraint::Choice(choices) => { + Format::or(choices.iter().map(Format::const_string).collect()) + } + // Lark rules use `:`; XGrammar EBNF rules use `::=`. + StructuredOutputConstraint::Grammar(grammar) if grammar.contains("::=") => { + Format::Grammar(GrammarFormat { + grammar: grammar.clone(), + }) + } + StructuredOutputConstraint::Grammar(_) => return None, + StructuredOutputConstraint::StructuralTag(tag) => structural_tag_format(tag)?, + }) +} + +/// The legacy structural tag shape: free text with JSON tags behind triggers. +#[derive(Deserialize)] +struct LegacyStructuralTag { + structures: Vec, + triggers: Vec, +} + +#[derive(Deserialize)] +struct LegacyStructure { + begin: String, + schema: Value, + end: String, +} + +/// Read a structural tag in its current (`format`) or legacy +/// (`structures`/`triggers`) shape. +fn structural_tag_format(tag: &str) -> Option { + let value: Value = serde_json::from_str(tag).ok()?; + if value.get("structures").is_none() { + return Some(serde_json::from_value::(value).ok()?.format); + } + let legacy: LegacyStructuralTag = serde_json::from_value(value).ok()?; + let triggers: Vec<&str> = legacy.triggers.iter().map(String::as_str).collect(); + let tags = legacy + .structures + .into_iter() + .map(|structure| { + TagFormat::new( + structure.begin, + Format::json_schema(structure.schema), + structure.end, + ) + }) + .collect(); + Some(Format::triggered_tags(&triggers, tags)) +} + /// Apply one parser-built output grammar to the prepared text request. /// -/// A parser-owned grammar replaces any structured output constraint the user -/// supplied, matching the Python frontend. `None` leaves the request untouched. +/// The parser has already composed the user's answer constraint into the +/// grammar where its protocol allows, so the grammar replaces the request's +/// structured outputs. `None` leaves the request untouched. pub(crate) fn apply_output_grammar( request: &mut TextRequest, built: Option, @@ -45,13 +117,92 @@ pub(crate) fn apply_output_grammar( #[cfg(test)] mod tests { use serde_json::json; - use vllm_engine_core_client::protocol::structured_outputs::{ - StructuredOutputBackend, StructuredOutputsParams, - }; - use xgrammar_structural_tag::format::Format; + use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputOptions; use super::*; + #[test] + fn answer_format_normalizes_composable_constraints() { + let schema = json!({ "type": "object", "properties": { "a": { "type": "integer" } } }); + let cases = [ + ( + StructuredOutputsParams::json(schema.clone()), + Format::json_schema(schema.clone()), + ), + ( + StructuredOutputsParams::json(Value::String(schema.to_string())), + Format::json_schema(schema), + ), + ( + StructuredOutputsParams::json_object(), + Format::json_schema(json!({ "type": "object" })), + ), + ( + StructuredOutputsParams::regex("[a-z]+"), + Format::regex("[a-z]+"), + ), + ( + StructuredOutputsParams::choice(vec!["yes".to_string(), "no".to_string()]), + Format::or(vec![ + Format::const_string("yes"), + Format::const_string("no"), + ]), + ), + ( + StructuredOutputsParams::grammar("root ::= \"a\""), + Format::Grammar(GrammarFormat { + grammar: "root ::= \"a\"".to_string(), + }), + ), + ( + StructuredOutputsParams::structural_tag( + r#"{"type": "structural_tag", "format": {"type": "const_string", "value": "x"}}"#, + ), + Format::const_string("x"), + ), + ( + StructuredOutputsParams::structural_tag( + r#"{"type": "structural_tag", "structures": [{"begin": "", "schema": {"type": "object"}, "end": ""}], "triggers": ["", + Format::json_schema(json!({ "type": "object" })), + "", + )], + ), + ), + ( + StructuredOutputsParams { + options: StructuredOutputOptions { + disable_any_whitespace: true, + disable_additional_properties: true, + whitespace_pattern: Some(" ?".to_string()), + }, + ..StructuredOutputsParams::json_object() + }, + Format::json_schema(json!({ "type": "object" })), + ), + ]; + for (params, expected) in cases { + assert_eq!(answer_format(¶ms), Some(expected), "{params:?}"); + } + } + + #[test] + fn answer_format_leaves_other_constraints_on_their_wire_field() { + let cases = [ + StructuredOutputsParams::json(Value::String("{".to_string())), + StructuredOutputsParams::grammar("start: \"a\""), + StructuredOutputsParams::structural_tag(r#"{"structures": [{"begin": ""}]}"#), + StructuredOutputsParams::structural_tag(r#"{"format": {"type": "unknown"}}"#), + ]; + for params in cases { + assert_eq!(answer_format(¶ms), None, "{params:?}"); + } + } + #[test] fn output_grammar_overwrites_answer_constraint() { let mut request = TextRequest::for_test(); diff --git a/rust/src/chat/tests/grammar_replay.py b/rust/src/chat/tests/grammar_replay.py index d34df627bae0..056ab1d35289 100755 --- a/rust/src/chat/tests/grammar_replay.py +++ b/rust/src/chat/tests/grammar_replay.py @@ -22,14 +22,17 @@ # /// """Replay the grammar cases exported by the roundtrip tests through XGrammar. -Each file holds one model's cases, keyed by tool-choice variant. Every case is -compiled against the model tokenizer the same way as -`vllm/v1/structured_output/backend_xgrammar.py`. Token-zero cases are then -replayed from the first generated token: each token must be allowed by the -bitmask and accepted by the matcher, and the final stop token must terminate it. +Each file holds one model's distinct output grammars, each with the generations +of the fixture variants that build it. Every grammar is compiled against the +model tokenizer the same way as `vllm/v1/structured_output/backend_xgrammar.py`. +The generations of token-zero grammars are then replayed from the first +generated token: each token must be allowed by the bitmask and accepted by the +matcher, and the final stop token must terminate it. """ import json +import multiprocessing +import os import sys from functools import cache from pathlib import Path @@ -49,23 +52,29 @@ def load_tokenizer(model_id: str, vocab_size: int): return tokenizer, xgr.GrammarCompiler(tokenizer_info) -def replay(model: dict, case: dict) -> str: - """Check one case of a model and return a summary of what was checked.""" - tokenizer, compiler = load_tokenizer(model["model_id"], model["vocab_size"]) - generation = case["generation_token_ids"] +def compile_grammar(model: dict, grammar: dict) -> xgr.CompiledGrammar: + """Compile one output grammar of a model.""" + _, compiler = load_tokenizer(model["model_id"], model["vocab_size"]) + structural_tag = {"type": "structural_tag", "format": grammar["format"]} + return compiler.compile_structural_tag(json.dumps(structural_tag)) + + +def replay( + model: dict, grammar: dict, compiled: xgr.CompiledGrammar, generation: dict +) -> str: + """Check one generation of a grammar and return a summary of what was checked.""" + tokenizer, _ = load_tokenizer(model["model_id"], model["vocab_size"]) + token_ids = generation["generation_token_ids"] - # The case token IDs come from the Rust tokenizer; the runner must load the + # The recorded token IDs come from the Rust tokenizer; the runner must load the # same vocabulary for the grammar to mean the same thing. - encoded = tokenizer.encode(case["completion"], add_special_tokens=False) - if encoded != generation[:-1]: + encoded = tokenizer.encode(generation["completion"], add_special_tokens=False) + if encoded != token_ids[:-1]: raise AssertionError( "tokenizer mismatch: the completion encodes to " - f"{encoded}, but the case records {generation[:-1]}" + f"{encoded}, but the generation records {token_ids[:-1]}" ) - grammar = case["grammar"] - structural_tag = {"type": "structural_tag", "format": grammar["format"]} - compiled = compiler.compile_structural_tag(json.dumps(structural_tag)) if grammar["coverage"] != "from_token_zero": return f"compiled only ({grammar['coverage']} is gated by the engine)" @@ -73,18 +82,41 @@ def replay(model: dict, case: dict) -> str: compiled, override_stop_tokens=model["all_stop_token_ids"] ) bitmask = xgr.allocate_token_bitmask(1, model["vocab_size"]) - for index, token_id in enumerate(generation): + for index, token_id in enumerate(token_ids): matcher.fill_next_token_bitmask(bitmask) allowed = (bitmask[0, token_id // 32].item() >> (token_id % 32)) & 1 if not allowed or not matcher.accept_token(token_id): - prefix = tokenizer.decode(generation[:index]) + prefix = tokenizer.decode(token_ids[:index]) raise AssertionError( f"token {index} ({token_id}, {tokenizer.decode([token_id])!r}) " f"rejected after {prefix!r}" ) if not matcher.is_terminated(): raise AssertionError("matcher not terminated after the stop token") - return f"replayed {len(generation)} tokens from token zero" + return f"replayed {len(token_ids)} tokens from token zero" + + +def check_model(path: Path) -> tuple[list[str], list[str]]: + """Check every grammar of one model file and return its stdout and stderr lines.""" + model = json.loads(path.read_text()) + out, err = [], [] + for case in model["grammars"]: + grammar = case["grammar"] + coverage = grammar["coverage"] + try: + compiled = compile_grammar(model, grammar) + except Exception as error: + names = ", ".join(case["generations"]) + err.append(f"FAIL {path.name} {names} [{coverage}]: {error}") + continue + for name, generation in case["generations"].items(): + try: + summary = replay(model, grammar, compiled, generation) + except Exception as error: + err.append(f"FAIL {path.name} {name} [{coverage}]: {error}") + else: + out.append(f"PASS {path.name} {name}: {summary}") + return out, err def main() -> int: @@ -93,19 +125,18 @@ def main() -> int: print(f"no grammar replay cases in {CASES_DIR}", file=sys.stderr) return 1 - failures = 0 - for path in paths: - model = json.loads(path.read_text()) - for name, case in model["cases"].items(): - try: - summary = replay(model, case) - except Exception as error: - failures += 1 - coverage = case["grammar"]["coverage"] - print(f"FAIL {path.name} {name} [{coverage}]: {error}", file=sys.stderr) - else: - print(f"PASS {path.name} {name}: {summary}") + # The cost is dominated by loading each model's tokenizer, which no other model + # shares, so check the model files in parallel. Fork after the imports above so + # that the workers share them instead of paying for them again. + workers = min(len(paths), os.cpu_count() or 1) + with multiprocessing.get_context("fork").Pool(workers) as pool: + results = pool.map(check_model, paths, chunksize=1) + failures = 0 + for out, err in results: + print("\n".join(out)) + print("\n".join(err), file=sys.stderr) + failures += len(err) return 1 if failures else 0 diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json index 4903b5900c74..f08b418bf0e6 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V3.2-Exp", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 43, 759, 1926, 270, 6704, 339, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 43, 759, 1926, 270, 6704, 339, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜function_calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜function_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128799, 271, 30, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 30, 28217, 10525, 7398, 28217, 40148, 5406, 2329, 1281, 3327, 3816, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 30, 28217, 10525, 7398, 28217, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 28217, 10525, 7398, 28217, 41523, 1018, 1718, 28217, 10525, 7398, 28217, 40148, 5406, 1018, 1718, 28217, 10525, 7398, 28217, 8701, 4941, 12548, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt index 5cf40f3218e8..25b572b8a247 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V3.2-Exp.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜function_calls>\n` tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜function_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json index 1a286d8acda9..8ab51cbbf991 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V4-Flash", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜tool_calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"get_weather\">\n<｜DSML｜parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜invoke name=\"add\">\n<｜DSML｜parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜parameter name=\"x\" string=\"false\">2\n<｜DSML｜parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 72461, 4941, 12548, 1018, 30, 128825, 40148, 5406, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 41523, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 30, 128825, 40148, 5406, 2329, 1281, 3327, 3816, 30, 128825, 41523, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 41523, 1018, 30, 128825, 41523, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 41523, 1018, 1718, 128825, 40148, 5406, 1018, 1718, 128825, 72461, 4941, 12548, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt index 5e216e8ffcf2..535da9ba2876 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4-Flash.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜tool_calls>\n` tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜tool_calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜invoke name="get_weather">\n` deepseek_xml({ location: string }) `\n` + tag `<｜DSML｜invoke name="add">\n` deepseek_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json index a21d8b90e5af..a47ce49c2f87 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.json @@ -2,8 +2,8 @@ "model_id": "deepseek-ai/DeepSeek-V4.1-Flash", "vocab_size": 129280, "all_stop_token_ids": [1], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -107,10 +107,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 43, 759, 1926, 270, 6704, 339, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -208,10 +212,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -267,8 +275,189 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n", - "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n<｜DSML｜ calls>\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<｜DSML｜ invoke name=\"get_weather\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "deepseek_v4_1_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "<｜DSML｜ invoke name=\"add\">\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "deepseek_v4_1_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 24313, 19995, 16562, 69958, 2928, 62773, 1] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_weather\">\n<｜DSML｜ parameter name=\"location\" string=\"true\">Shanghai\n\n<｜DSML｜ invoke name=\"add\">\n<｜DSML｜ parameter name=\"y\" string=\"false\">1.0\n<｜DSML｜ parameter name=\"x\" string=\"false\">2\n<｜DSML｜ parameter name=\"items\" string=\"false\">[\"left\", \"right\"]\n\n", + "generation_token_ids": [50249, 1926, 270, 9670, 305, 1258, 6704, 16, 128822, 271, 30, 128825, 10699, 1018, 30, 128825, 34756, 2329, 1281, 1133, 65, 50219, 3816, 30, 128825, 10767, 2329, 1281, 33182, 4, 3418, 1281, 11476, 3320, 119606, 28011, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 30, 128825, 34756, 2329, 1281, 3327, 3816, 30, 128825, 10767, 2329, 1281, 91, 4, 3418, 1281, 19836, 3320, 19, 16, 18, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 90, 4, 3418, 1281, 19836, 3320, 20, 1718, 128825, 10767, 1018, 30, 128825, 10767, 2329, 1281, 17173, 4, 3418, 1281, 19836, 3320, 10931, 4589, 1760, 582, 2459, 13747, 1718, 128825, 10767, 1018, 1718, 128825, 34756, 1018, 1718, 128825, 10699, 32, 1] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt index 8c92c3a5a374..cb2c646545df 100644 --- a/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/DeepSeek-V4.1-Flash.txt @@ -24,3 +24,20 @@ sequence `\n\n<｜DSML｜ calls>\n` tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n\n<｜DSML｜ calls>\n` + tags_with_separator `` at_least_one + tag `<｜DSML｜ invoke name="get_weather">\n` deepseek_v4_1_xml({ location: string }) `\n` + tag `<｜DSML｜ invoke name="add">\n` deepseek_v4_1_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.5.json b/rust/src/chat/tests/grammar_replay/GLM-4.5.json new file mode 100644 index 000000000000..e00f05757bb2 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-4.5.json @@ -0,0 +1,68 @@ +{ + "model_id": "zai-org/GLM-4.5", + "vocab_size": 151552, + "all_stop_token_ids": [151329, 151336, 151338], + "grammars": [ + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "\n", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "\nNeed call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [198, 151350, 23591, 1618, 279, 9101, 323, 912, 7375, 13, 151351, 198, 4913, 9214, 3252, 82, 27208, 9204, 151329] + }, + "auto_schema_content": { + "completion": "\nNeed call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [198, 151350, 23591, 1618, 279, 9101, 323, 912, 7375, 13, 151351, 198, 4913, 9214, 3252, 82, 27208, 9204, 151329] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.5.txt b/rust/src/chat/tests/grammar_replay/GLM-4.5.txt new file mode 100644 index 000000000000..8a37a06dbda5 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/GLM-4.5.txt @@ -0,0 +1,7 @@ +# none_schema, auto_schema_content (from_token_zero) +sequence + optional + sequence + tag `\n` text `` + `\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json index b289ec62e62b..6fd2a440e4d4 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json +++ b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.json @@ -2,8 +2,8 @@ "model_id": "zai-org/GLM-4.7-Flash", "vocab_size": 154880, "all_stop_token_ids": [154820, 154827, 154829], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "required": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -244,8 +252,187 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghai", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + "generations": { + "named": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghai", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "get_weather", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "", + "", + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt index 58b4bef2da32..bc58a47fc936 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt +++ b/rust/src/chat/tests/grammar_replay/GLM-4.7-Flash.txt @@ -16,3 +16,17 @@ sequence sequence tag `` text `` tag `get_weather` glm_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json index 404667c30a7b..ecf8880406de 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json +++ b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.json @@ -2,8 +2,8 @@ "model_id": "zai-org/GLM-5.2-FP8", "vocab_size": 154880, "all_stop_token_ids": [154820, 154827, 154829], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 40, 686, 1618, 279, 7375, 13, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + "generations": { + "required": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -244,8 +252,187 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.get_weatherlocationShanghai", - "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + "generations": { + "named": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghai", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "get_weather", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "glm_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "", + "", + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 4913, 9214, 3252, 82, 27208, 9204, 154820] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.get_weatherlocationShanghaiaddy1.0x2items[\"left\", \"right\"]", + "generation_token_ids": [23591, 1618, 279, 9101, 323, 912, 7375, 13, 154842, 154843, 455, 68852, 154847, 2527, 154848, 154849, 2016, 29953, 154850, 154844, 154843, 718, 154847, 88, 154848, 154849, 16, 13, 15, 154850, 154847, 87, 154848, 154849, 17, 154850, 154847, 3615, 154848, 154849, 1183, 2359, 497, 330, 1291, 1341, 154850, 154844, 154820] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt index 58b4bef2da32..bc58a47fc936 100644 --- a/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt +++ b/rust/src/chat/tests/grammar_replay/GLM-5.2-FP8.txt @@ -16,3 +16,17 @@ sequence sequence tag `` text `` tag `get_weather` glm_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, ``, ``, ``, ``] at_least_one + tag `get_weather` glm_xml({ location: string }) `` + tag `add` glm_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Hy3.json b/rust/src/chat/tests/grammar_replay/Hy3.json index f6ae7b0969fd..da089616300a 100644 --- a/rust/src/chat/tests/grammar_replay/Hy3.json +++ b/rust/src/chat/tests/grammar_replay/Hy3.json @@ -2,8 +2,8 @@ "model_id": "tencent/Hy3", "vocab_size": 120832, "all_stop_token_ids": [120025], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -263,10 +263,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 40, 823, 1609, 252, 7106, 13, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 40, 823, 1609, 252, 7106, 13, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -523,10 +527,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -618,8 +626,299 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\n", - "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120036, 120025] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120036, 120025] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "get_weather\n", + "content": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "location" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + ] + }, + "end": "" + }, + { + "type": "tag", + "begin": "add\n", + "content": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "y" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "x" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "items" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "any_text", + "excludes": [ + "", + "", + "", + "", + "", + "", + "", + "", + "" + ], + "max_tokens": null, + "max_chars": null + }, + { + "type": "const_string", + "value": "" + }, + { + "type": "const_string", + "value": "\n" + } + ] + } + ] + }, + "end": "" + } + ], + "separator": "\n", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\nget_weather\nlocation\nShanghai\n\nadd\ny\n1.0\nx\n2\nitems\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [47893, 1609, 252, 9969, 289, 1189, 7106, 13, 120030, 120035, 185, 120033, 874, 62, 39899, 120043, 185, 120066, 24806, 120067, 185, 120068, 4169, 768, 32381, 120069, 185, 120034, 185, 120033, 2176, 120043, 185, 120066, 88, 120067, 185, 120068, 16, 13, 15, 120069, 185, 120066, 87, 120067, 185, 120068, 17, 120069, 185, 120066, 12792, 120067, 185, 120068, 5831, 1639, 1221, 478, 1331, 7029, 120069, 185, 120034, 185, 120036, 120025] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Hy3.txt b/rust/src/chat/tests/grammar_replay/Hy3.txt index dcd6bb08ec8b..5642b3697078 100644 --- a/rust/src/chat/tests/grammar_replay/Hy3.txt +++ b/rust/src/chat/tests/grammar_replay/Hy3.txt @@ -110,3 +110,53 @@ sequence `` `\n` `\n` + +# auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `\n` + tags_with_separator `\n` at_least_one + tag `get_weather\n` .. `` + sequence + sequence + `` + `location` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + tag `add\n` .. `` + sequence + sequence + `` + `y` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `x` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + sequence + `` + `items` + `` + `\n` + `` + text excluding [``, ``, ``, ``, ``, ``, ``, ``, ``] + `` + `\n` + `\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json index 8ac6e689f78b..3fb22cc3b8cf 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.json @@ -2,8 +2,8 @@ "model_id": "moonshotai/Kimi-K2.5", "vocab_size": 163840, "all_stop_token_ids": [163585, 163586], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -134,10 +134,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 40, 748, 2390, 276, 7697, 13, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 40, 748, 2390, 276, 7697, 13, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -261,10 +265,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + "generations": { + "required": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -333,8 +341,166 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_calls_section_end|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163596, 163586] + "generations": { + "named": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163596, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "<|tool_calls_section_begin|>" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|tool_call_begin|>functions.get_weather:", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "\\d+" + }, + { + "type": "const_string", + "value": "<|tool_call_argument_begin|>" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "end": "<|tool_call_end|>" + }, + { + "type": "tag", + "begin": "<|tool_call_begin|>functions.add:", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "\\d+" + }, + { + "type": "const_string", + "value": "<|tool_call_argument_begin|>" + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "end": "<|tool_call_end|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "<|tool_calls_section_end|>" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_calls": { + "completion": "Need call the weather and add tools.<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{\"location\": \"Shanghai\"}<|tool_call_end|><|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>{\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}<|tool_call_end|><|tool_calls_section_end|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163607, 163595, 163597, 41937, 1150, 21055, 2800, 25, 15, 163598, 8264, 5791, 1289, 414, 2797, 31231, 16934, 163599, 163597, 41937, 3149, 25, 16, 163598, 8264, 88, 1289, 220, 16, 13, 15, 11, 414, 87, 1289, 220, 17, 11, 414, 7989, 1289, 8681, 3794, 665, 414, 2277, 3956, 92, 163599, 163596, 163586] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt index 553c1ce6f4c3..ca3bab8061d7 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K2.5.txt @@ -44,3 +44,23 @@ sequence `<|tool_call_argument_begin|>` json({ location: string }) `<|tool_calls_section_end|>` + +# auto_schema_calls (from_token_zero) +sequence + tag `` text `` + or + sequence + `<|tool_calls_section_begin|>` + tags_with_separator `` at_least_one + tag `<|tool_call_begin|>functions.get_weather:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ location: string }) + tag `<|tool_call_begin|>functions.add:` .. `<|tool_call_end|>` + sequence + /\d+/ + `<|tool_call_argument_begin|>` + json({ y: number, x: number, items: string[] }) + `<|tool_calls_section_end|>` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.json b/rust/src/chat/tests/grammar_replay/Kimi-K3.json index 5306b936604f..3e1832e805f7 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.json +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.json @@ -2,8 +2,8 @@ "model_id": "moonshotai/Kimi-K3", "vocab_size": 163840, "all_stop_token_ids": [163585, 163586], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -427,10 +427,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>I will call the tools.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 40, 748, 2390, 276, 7697, 13, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -666,10 +670,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "required": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -795,8 +803,527 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", - "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + "generations": { + "named": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>think<|sep|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>think<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>response<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>response<|sep|>" + }, + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|close|>message<|sep|>" + } + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>{\"answer\":\"sunny\"}<|close|>response<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 8264, 26921, 7471, 82, 38975, 16934, 163588, 12092, 163589, 163588, 2778, 163589, 163586] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>think<|sep|>", + "<|end_of_msg|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>think<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>response<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>response<|sep|>" + }, + { + "type": "optional", + "content": { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + } + ] + }, + { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|open|>response<|sep|><|close|>response<|sep|>" + } + }, + { + "type": "tag", + "begin": "<|open|>tools<|sep|>", + "content": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "<|open|>call tool=\"get_weather\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"location\" type=\"string\"<|sep|>", + "content": { + "type": "any_text", + "excludes": [ + "<|close|>argument<|sep|>", + "<|close|>call<|sep|>" + ], + "max_tokens": null, + "max_chars": null + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>call tool=\"add\" index=\"", + "content": { + "type": "sequence", + "elements": [ + { + "type": "regex", + "pattern": "[1-9][0-9]*" + }, + { + "type": "const_string", + "value": "\"<|sep|>" + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "<|open|>argument key=\"y\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"x\" type=\"number\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "number" + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + }, + { + "type": "tag", + "begin": "<|open|>argument key=\"items\" type=\"array\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "array", + "items": { + "type": "string" + } + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>argument<|sep|>" + } + ] + }, + { + "type": "tag", + "begin": "<|open|>json type=\"object\"<|sep|>", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "<|close|>json<|sep|>" + } + ] + } + ] + }, + "end": "<|close|>call<|sep|>" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + "end": "<|close|>tools<|sep|>" + } + ] + } + ] + }, + { + "type": "optional", + "content": { + "type": "const_string", + "value": "<|close|>message<|sep|>" + } + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|>{\"answer\":\"sunny\"}<|close|>response<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 8264, 26921, 7471, 82, 38975, 16934, 163588, 12092, 163589, 163588, 2778, 163589, 163586] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.<|close|>think<|sep|><|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool=\"get_weather\" index=\"1\"<|sep|><|open|>argument key=\"location\" type=\"string\"<|sep|>Shanghai<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool=\"add\" index=\"2\"<|sep|><|open|>argument key=\"y\" type=\"number\"<|sep|>1.0<|close|>argument<|sep|><|open|>argument key=\"x\" type=\"number\"<|sep|>2<|close|>argument<|sep|><|open|>argument key=\"items\" type=\"array\"<|sep|>[\"left\",\"right\"]<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>", + "generation_token_ids": [20082, 2390, 276, 10666, 316, 1239, 7697, 13, 163588, 39964, 163589, 163587, 12092, 163589, 163588, 12092, 163589, 163587, 25385, 163589, 163587, 10257, 4453, 878, 618, 21055, 2800, 1, 4002, 878, 16, 1, 163589, 163587, 47185, 2355, 878, 5791, 1, 1798, 878, 2033, 1, 163589, 2797, 31231, 163588, 47185, 163589, 163588, 10257, 163589, 163587, 10257, 4453, 878, 1210, 1, 4002, 878, 17, 1, 163589, 163587, 47185, 2355, 878, 88, 1, 1798, 878, 10006, 1, 163589, 16, 13, 15, 163588, 47185, 163589, 163587, 47185, 2355, 878, 87, 1, 1798, 878, 10006, 1, 163589, 17, 163588, 47185, 163589, 163587, 47185, 2355, 878, 7989, 1, 1798, 878, 4392, 1, 163589, 3530, 3794, 3923, 2277, 3956, 163588, 47185, 163589, 163588, 10257, 163589, 163588, 25385, 163589, 163588, 2778, 163589, 163586] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt index 6a1c814c131c..d7357e275c86 100644 --- a/rust/src/chat/tests/grammar_replay/Kimi-K3.txt +++ b/rust/src/chat/tests/grammar_replay/Kimi-K3.txt @@ -86,3 +86,58 @@ sequence tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` optional `<|close|>message<|sep|>` + +# none_schema (from_token_zero) +sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json({ answer: string }) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` json({ answer: string }) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + sequence + optional `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + tag `<|open|>argument key="location" type="string"<|sep|>` text excluding [`<|close|>argument<|sep|>`, `<|close|>call<|sep|>`] `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ location: string }) `<|close|>json<|sep|>` + tag `<|open|>call tool="add" index="` .. `<|close|>call<|sep|>` + sequence + /[1-9][0-9]*/ + `"<|sep|>` + or + sequence + tag `<|open|>argument key="y" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="x" type="number"<|sep|>` json(number) `<|close|>argument<|sep|>` + tag `<|open|>argument key="items" type="array"<|sep|>` json(string[]) `<|close|>argument<|sep|>` + tag `<|open|>json type="object"<|sep|>` json({ y: number, x: number, items: string[] }) `<|close|>json<|sep|>` + optional `<|close|>message<|sep|>` diff --git a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json index 582f29a1b943..451255dd8c45 100644 --- a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json +++ b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.json @@ -2,8 +2,8 @@ "model_id": "XiaomiMiMo/MiMo-V2.6-Flash-RL", "vocab_size": 152576, "all_stop_token_ids": [151643, 151645, 151672], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -99,10 +99,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.I will call the tools.Shanghai1.02[\"left\", \"right\"]", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 40, 686, 1618, 279, 7375, 13, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.I will call the tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 40, 686, 1618, 279, 7375, 13, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -198,10 +202,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "required": { + "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -247,8 +255,190 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.Shanghai", - "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151645] + "generations": { + "named": { + "completion": "Need call the weather and add tools.Shanghai", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 4913, 9217, 3252, 82, 27297, 9207, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + } + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + }, + { + "type": "tag", + "begin": "", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "", + "", + "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 4913, 9217, 3252, 82, 27297, 9207, 151645] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.Shanghai1.02[\"left\", \"right\"]", + "generation_token_ids": [151667, 23657, 1618, 279, 9104, 323, 912, 7375, 13, 151668, 151657, 27, 1688, 28280, 69364, 1784, 16181, 28, 2527, 29, 2016, 30070, 522, 16181, 1472, 1688, 29, 151658, 151657, 27, 1688, 92784, 1784, 16181, 29368, 29, 16, 13, 15, 522, 16181, 1784, 16181, 25387, 29, 17, 522, 16181, 1784, 16181, 28, 3615, 29, 1183, 2359, 497, 330, 1291, 1341, 522, 16181, 1472, 1688, 29, 151658, 151645] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt index 1d12560eb034..7cc9aa0e59a2 100644 --- a/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt +++ b/rust/src/chat/tests/grammar_replay/MiMo-V2.6-Flash-RL.txt @@ -16,3 +16,17 @@ sequence sequence optional tag `` text `` tag `` qwen_xml({ location: string }) `` + +# none_schema (from_token_zero) +sequence + optional tag `` text `` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + optional tag `` text `` + or + triggered_tags [``] excluding [``, ``, ``, `` qwen_xml({ location: string }) `` + tag `` qwen_xml({ y: number, x: number, items: string[] }) `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json index 64f65e2f5a32..d4c2e8e789f6 100644 --- a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.json @@ -2,8 +2,8 @@ "model_id": "MiniMaxAI/MiniMax-M2.5", "vocab_size": 200064, "all_stop_token_ids": [200020], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -116,10 +116,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 73, 715, 1728, 275, 6629, 320, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 73, 715, 1728, 275, 6629, 320, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -226,10 +230,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -294,8 +302,207 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n", - "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "or", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "const_string", + "value": "\n\n" + }, + { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": "\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "minimax_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + }, + { + "type": "tag", + "begin": "\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "minimax_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n" + } + ], + "separator": "", + "at_least_one": true, + "stop_after_first": false + }, + { + "type": "const_string", + "value": "" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 367, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n1.0\n2\n[\"left\", \"right\"]\n\n", + "generation_token_ids": [42478, 1728, 275, 12562, 306, 1187, 6629, 320, 200051, 4368, 200052, 10, 60, 88638, 1925, 1139, 848, 95, 51872, 6143, 60, 24662, 1925, 1139, 28415, 3361, 165806, 1579, 24662, 1100, 1579, 88638, 1100, 60, 88638, 1925, 1139, 2132, 6143, 60, 24662, 1925, 1139, 121, 3361, 49, 46, 48, 1579, 24662, 1100, 60, 24662, 1925, 1139, 120, 3361, 50, 1579, 24662, 1100, 60, 24662, 1925, 1139, 12405, 3361, 7795, 2038, 1164, 494, 1581, 8760, 1579, 24662, 1100, 1579, 88638, 1100, 200053, 200020] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt index 431238009a13..fae13b84450a 100644 --- a/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M2.5.txt @@ -30,3 +30,24 @@ sequence `\n\n` tag `\n` minimax_xml({ location: string }) `\n` `` + +# none_schema (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + or + sequence + `\n\n` + tags_with_separator `` at_least_one + tag `\n` minimax_xml({ location: string }) `\n` + tag `\n` minimax_xml({ y: number, x: number, items: string[] }) `\n` + `` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M3.json b/rust/src/chat/tests/grammar_replay/MiniMax-M3.json new file mode 100644 index 000000000000..73679eff6a1e --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M3.json @@ -0,0 +1,40 @@ +{ + "model_id": "MiniMaxAI/MiniMax-M3", + "vocab_size": 200064, + "all_stop_token_ids": [200020], + "grammars": [ + { + "grammar": { + "format": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "coverage": "final_output_only" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [200059, 42478, 1728, 275, 12562, 306, 1187, 6629, 46, 200060, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + }, + "auto_schema_content": { + "completion": "Need call the weather and add tools.{\"answer\":\"sunny\"}", + "generation_token_ids": [200059, 42478, 1728, 275, 12562, 306, 1187, 6629, 46, 200060, 16673, 33591, 15368, 78697, 7582, 22089, 200020] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt b/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt new file mode 100644 index 000000000000..f54e72d4c825 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/MiniMax-M3.txt @@ -0,0 +1,2 @@ +# none_schema, auto_schema_content (final_output_only) +json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json index 5d171fc5f080..8206b8b40c1a 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json +++ b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.json @@ -2,8 +2,8 @@ "model_id": "Qwen/Qwen3-0.6B", "vocab_size": 151936, "all_stop_token_ids": [151643, 151645], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -106,10 +106,14 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\nI will call the tools.\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 40, 686, 1618, 279, 7375, 624, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + "generations": { + "auto": { + "completion": "\nNeed call the weather and add tools.\n\n\nI will call the tools.\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 40, 686, 1618, 279, 7375, 624, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -212,10 +216,14 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + "generations": { + "required": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -270,8 +278,206 @@ }, "coverage": "from_token_zero" }, - "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n", - "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 151645] + "generations": { + "named": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + } + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "\nNeed call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 4913, 9217, 3252, 82, 27297, 9207, 151645] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "optional", + "content": { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + } + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "" + ], + "tags": [ + { + "type": "tag", + "begin": "\n{\"name\": \"get_weather\", \"arguments\": ", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "}\n" + }, + { + "type": "tag", + "begin": "\n{\"name\": \"add\", \"arguments\": ", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "}\n" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "\nNeed call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 4913, 9217, 3252, 82, 27297, 9207, 151645] + }, + "auto_schema_calls": { + "completion": "\nNeed call the weather and add tools.\n\n\n\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Shanghai\"}}\n\n\n{\"name\": \"add\", \"arguments\": {\"y\": 1.0, \"x\": 2, \"items\": [\"left\", \"right\"]}}\n", + "generation_token_ids": [151667, 198, 23657, 1618, 279, 9104, 323, 912, 7375, 624, 151668, 271, 151657, 198, 4913, 606, 788, 330, 455, 69364, 497, 330, 16370, 788, 5212, 2527, 788, 330, 2016, 30070, 95642, 151658, 198, 151657, 198, 4913, 606, 788, 330, 718, 497, 330, 16370, 788, 5212, 88, 788, 220, 16, 13, 15, 11, 330, 87, 788, 220, 17, 11, 330, 3615, 788, 4383, 2359, 497, 330, 1291, 1341, 11248, 151658, 151645] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt index fe38f5e9ff4a..865d1b543071 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt +++ b/rust/src/chat/tests/grammar_replay/Qwen3-0.6B.txt @@ -25,3 +25,23 @@ sequence tag `` text `` `\n\n` tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + +# none_schema (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + optional + sequence + tag `` text `` + `\n\n` + or + triggered_tags [``] excluding [``, ``] at_least_one + tag `\n{"name": "get_weather", "arguments": ` json({ location: string }) `}\n` + tag `\n{"name": "add", "arguments": ` json({ y: number, x: number, items: string[] }) `}\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json index 8bd4c901e541..9cc7b53ff9ca 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json +++ b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.json @@ -2,8 +2,8 @@ "model_id": "Qwen/Qwen3.5-4B", "vocab_size": 248320, "all_stop_token_ids": [248044, 248046], - "cases": { - "auto": { + "grammars": [ + { "grammar": { "format": { "type": "sequence", @@ -103,10 +103,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 40, 668, 1562, 279, 7141, 13, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "auto": { + "completion": "Need call the weather and add tools.\n\n\nI will call the tools.\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 40, 668, 1562, 279, 7141, 13, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } }, - "required": { + { "grammar": { "format": { "type": "sequence", @@ -206,10 +210,14 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "required": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } }, - "named": { + { "grammar": { "format": { "type": "sequence", @@ -261,8 +269,200 @@ }, "coverage": "from_token_zero" }, - "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n", - "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + "generations": { + "named": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 4754, 8944, 3147, 82, 26415, 8934, 248046] + } + } + }, + { + "grammar": { + "format": { + "type": "sequence", + "elements": [ + { + "type": "sequence", + "elements": [ + { + "type": "tag", + "begin": "", + "content": { + "type": "any_text", + "excludes": [], + "max_tokens": null, + "max_chars": null + }, + "end": "" + }, + { + "type": "const_string", + "value": "\n\n" + } + ] + }, + { + "type": "or", + "elements": [ + { + "type": "triggered_tags", + "triggers": [ + "\n\n\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "location": { + "type": "string" + } + }, + "required": [ + "location" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n\n" + }, + { + "type": "tag", + "begin": "\n\n", + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "y": { + "type": "number" + }, + "x": { + "type": "number" + }, + "items": { + "type": "array", + "items": { + "type": "string" + } + } + }, + "required": [ + "y", + "x", + "items" + ] + }, + "style": "qwen_xml", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "end": "\n\n" + } + ], + "at_least_one": true, + "stop_after_first": false, + "excludes": [ + "", + "" + ] + }, + { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + } + ] + } + ] + }, + "coverage": "from_token_zero" + }, + "generations": { + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n\n\n{\"answer\":\"sunny\"}", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 4754, 8944, 3147, 82, 26415, 8934, 248046] + }, + "auto_schema_calls": { + "completion": "Need call the weather and add tools.\n\n\n\n\n\nShanghai\n\n\n\n\n\n\n1.0\n\n\n2\n\n\n[\"left\", \"right\"]\n\n\n", + "generation_token_ids": [22916, 1562, 279, 8831, 321, 884, 7141, 13, 198, 248069, 271, 248058, 198, 27, 1628, 27362, 67017, 29, 198, 27, 15704, 28, 2447, 29, 198, 1948, 29086, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 198, 248058, 198, 27, 1628, 89587, 29, 198, 27, 15704, 28408, 29, 198, 16, 13, 15, 198, 510, 15704, 29, 198, 27, 15704, 24582, 29, 198, 17, 198, 510, 15704, 29, 198, 27, 15704, 28, 3494, 29, 198, 1143, 2282, 487, 328, 1246, 1293, 198, 510, 15704, 29, 198, 510, 1628, 29, 198, 248059, 248046] + } + } } - } + ] } diff --git a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt index b0577806494b..56831217d83d 100644 --- a/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt +++ b/rust/src/chat/tests/grammar_replay/Qwen3.5-4B.txt @@ -22,3 +22,21 @@ sequence tag `` text `` `\n\n` tag `\n\n` qwen_xml({ location: string }) `\n\n` + +# none_schema (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + json({ answer: string }) + +# auto_schema_content, auto_schema_calls (from_token_zero) +sequence + sequence + tag `` text `` + `\n\n` + or + triggered_tags [`\n`, ``] at_least_one + tag `\n\n` qwen_xml({ location: string }) `\n\n` + tag `\n\n` qwen_xml({ y: number, x: number, items: string[] }) `\n\n` + json({ answer: string }) diff --git a/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json new file mode 100644 index 000000000000..e933ae4abbb3 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.json @@ -0,0 +1,40 @@ +{ + "model_id": "ByteDance-Seed/Seed-OSS-36B-Instruct", + "vocab_size": 155136, + "all_stop_token_ids": [2], + "grammars": [ + { + "grammar": { + "format": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": { + "answer": { + "type": "string" + } + }, + "required": [ + "answer" + ] + }, + "style": "json", + "any_order": false, + "max_whitespace_cnt": null, + "excludes": [] + }, + "coverage": "final_output_only" + }, + "generations": { + "none_schema": { + "completion": "Need call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [3, 40143, 2434, 405, 11785, 440, 1731, 9633, 141, 4, 326, 13059, 35733, 15301, 40841, 2313, 16320, 2] + }, + "auto_schema_content": { + "completion": "Need call the weather and add tools.\n{\"answer\":\"sunny\"}", + "generation_token_ids": [3, 40143, 2434, 405, 11785, 440, 1731, 9633, 141, 4, 326, 13059, 35733, 15301, 40841, 2313, 16320, 2] + } + } + } + ] +} diff --git a/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt new file mode 100644 index 000000000000..f54e72d4c825 --- /dev/null +++ b/rust/src/chat/tests/grammar_replay/Seed-OSS-36B-Instruct.txt @@ -0,0 +1,2 @@ +# none_schema, auto_schema_content (final_output_only) +json({ answer: string }) diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 079ce3e1a7e7..c5080a59c5c2 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -7,9 +7,10 @@ //! parsed from the generated assistant completion and then rendered back to the exact same //! assistant completion. //! -//! The tool-call fixture also snapshots the output grammar built by the initialized parser for each -//! tool-choice variant, with the rendered completion as its generation, as one file per model under -//! `tests/grammar_replay/`. `grammar_replay.rs` replays these cases through the real XGrammar. +//! The tool-call fixture also snapshots the output grammars built by the initialized parser, one +//! entry per distinct grammar with the rendered completion of every variant that builds it as a +//! generation, as one file per model under `tests/grammar_replay/`. `grammar_replay.rs` replays +//! these generations through the real XGrammar. use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; @@ -30,6 +31,7 @@ use vllm_chat::{ GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection, RendererSelection, ToolStrictLevel, load_model_backends, }; +use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use vllm_parser::output_grammar::BuiltOutputGrammar; use vllm_text::{DecodedTextEvent, Finished, Prompt}; use vllm_tokenizer::Tokenizer; @@ -57,6 +59,13 @@ struct RoundtripCase { json_fmt: JsonFmt, /// Whether the template renders tool-call argument object keys in sorted order. sort_json_keys: bool, + /// Which tool-call fixture variants with a JSON schema this case runs; + /// every case runs the variants without one. The variants answered with + /// content need a template that renders a final turn without tool calls + /// back to its completion (some templates drop its reasoning or its stop + /// suffix); the variant answered with calls needs a tool grammar, which + /// holds the calls next to the schema. + with_schema: &'static [ToolCallMixVariant], } #[derive(Clone, Copy)] @@ -125,10 +134,11 @@ impl ThinkingBehavior { } } -/// Tool-choice variant of the tool-call fixture. +/// Variant of the tool-call fixture: the request's tool choice and optional +/// JSON schema, and the completion that answers it. #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize)] #[serde(rename_all = "snake_case")] -enum ToolChoiceVariant { +enum ToolCallMixVariant { /// `auto`, with content before the tool calls. Auto, /// `required`, with tool calls only: builder grammars admit no content @@ -136,18 +146,59 @@ enum ToolChoiceVariant { Required, /// Named `get_weather`, with a single call to it and no content. Named, + + /// `none` with a JSON schema, answered with schema-valid content. + NoneSchema, + /// `auto` with a JSON schema, answered with schema-valid content. + AutoSchemaContent, + /// `auto` with a JSON schema, answered with tool calls only: builder + /// grammars hold either schema-valid content or the calls. + AutoSchemaCalls, } -impl ToolChoiceVariant { - const ALL: [Self; 3] = [Self::Auto, Self::Required, Self::Named]; +/// The schema-valid content of the schema variants answered with content. +const SCHEMA_CONTENT: &str = r#"{"answer":"sunny"}"#; + +impl ToolCallMixVariant { + const WITHOUT_SCHEMA: &[Self] = &[Self::Auto, Self::Required, Self::Named]; + const WITH_SCHEMA: &[Self] = &[ + Self::NoneSchema, + Self::AutoSchemaContent, + Self::AutoSchemaCalls, + ]; fn tool_choice(self) -> ChatToolChoice { match self { - Self::Auto => ChatToolChoice::Auto, + Self::Auto | Self::AutoSchemaContent | Self::AutoSchemaCalls => ChatToolChoice::Auto, Self::Required => ChatToolChoice::Required, Self::Named => ChatToolChoice::Function { name: "get_weather".to_string(), }, + Self::NoneSchema => ChatToolChoice::None, + } + } + + fn structured_outputs(self) -> Option { + matches!( + self, + Self::NoneSchema | Self::AutoSchemaContent | Self::AutoSchemaCalls + ) + .then(|| { + StructuredOutputsParams::json(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"] + })) + }) + } + + /// The expected text content and number of tool calls. + fn expected(self) -> (Option<&'static str>, usize) { + match self { + Self::Auto => (Some("I will call the tools."), 2), + Self::Required | Self::AutoSchemaCalls => (None, 2), + Self::Named => (None, 1), + Self::NoneSchema | Self::AutoSchemaContent => (Some(SCHEMA_CONTENT), 0), } } } @@ -163,6 +214,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -176,6 +228,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -189,6 +242,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -202,6 +256,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -215,6 +270,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The MiniMax M3 tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -228,6 +289,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -241,6 +303,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -254,6 +317,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: false }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -267,6 +331,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The GLM-4.5 tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -280,6 +350,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -293,6 +364,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -306,6 +378,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: true, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -329,6 +402,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -346,6 +420,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -363,6 +438,7 @@ impl RoundtripCase { }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -376,6 +452,12 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + // The Seed-OSS tool parser builds no tool grammar, so the schema + // excludes the calls. + with_schema: &[ + ToolCallMixVariant::NoneSchema, + ToolCallMixVariant::AutoSchemaContent, + ], } } @@ -389,6 +471,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -402,6 +485,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } @@ -415,6 +499,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: false, + with_schema: &[ToolCallMixVariant::AutoSchemaCalls], } } @@ -428,6 +513,7 @@ impl RoundtripCase { thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), sort_json_keys: true, + with_schema: ToolCallMixVariant::WITH_SCHEMA, } } } @@ -545,8 +631,10 @@ async fn run_roundtrip_tool_call_mix( backends: &vllm_chat::LoadedModelBackends, ) -> Result<()> { let mut results = Vec::new(); - for variant in ToolChoiceVariant::ALL { - let result = run_roundtrip_tool_call_mix_inner(case, backends, variant).await?; + for &variant in ToolCallMixVariant::WITHOUT_SCHEMA.iter().chain(case.with_schema) { + let result = run_roundtrip_tool_call_mix_inner(case, backends, variant) + .await + .with_context(|| format!("variant {}", serde_json::to_value(variant).unwrap()))?; results.push((variant, result)); } check_grammar_replay_file(case, backends, results) @@ -555,9 +643,9 @@ async fn run_roundtrip_tool_call_mix( async fn run_roundtrip_tool_call_mix_inner( case: &RoundtripCase, backends: &vllm_chat::LoadedModelBackends, - variant: ToolChoiceVariant, + variant: ToolCallMixVariant, ) -> Result { - let request = roundtrip_request( + let mut request = roundtrip_request( "roundtrip-reasoning-tools", vec![ChatMessage::text( ChatRole::User, @@ -568,12 +656,9 @@ async fn run_roundtrip_tool_call_mix_inner( Some(true), // always enable thinking in this fixture case.thinking_behavior, ); + request.sampling_params.structured_outputs = variant.structured_outputs(); let expected_reasoning = "Need call the weather and add tools."; - let (expected_text, tool_call_count) = match variant { - ToolChoiceVariant::Auto => (Some("I will call the tools."), 2), - ToolChoiceVariant::Required => (None, 2), - ToolChoiceVariant::Named => (None, 1), - }; + let (expected_text, tool_call_count) = variant.expected(); let mut content = vec![AssistantContentBlock::Reasoning { text: expected_reasoning.to_string(), @@ -617,11 +702,13 @@ async fn run_roundtrip_tool_call_mix_inner( "parsed message: {:#?}", result.parsed_message ); - assert_eq!(tool_calls[0].name, "get_weather"); - assert_eq!( - tool_calls[0].arguments, - expected_arguments(case, r#"{"location": "Shanghai"}"#)?, - ); + if let Some(weather) = tool_calls.first() { + assert_eq!(weather.name, "get_weather"); + assert_eq!( + weather.arguments, + expected_arguments(case, r#"{"location": "Shanghai"}"#)?, + ); + } if let Some(add) = tool_calls.get(1) { assert_eq!(add.name, "add"); assert_eq!( @@ -1117,16 +1204,23 @@ struct GrammarReplayFile { /// Request stop set, passed to the matcher as `override_stop_tokens`. #[serde(serialize_with = "serialize_one_line")] all_stop_token_ids: BTreeSet, - /// Cases keyed by tool-choice variant. - cases: BTreeMap, + /// Distinct output grammars, in the order of the first variant building each. + grammars: Vec, } -/// One parser-built output grammar with the generation it must accept. +/// One parser-built output grammar with the generations it must accept. #[derive(Serialize)] struct GrammarReplayCase { /// Parser-built grammar; only token-zero grammars are replayed from the /// first generated token. grammar: BuiltOutputGrammar, + /// Generations keyed by the fixture variants whose requests build this grammar. + generations: BTreeMap, +} + +/// One rendered completion replayed against its output grammar. +#[derive(Serialize)] +struct GrammarReplayGeneration { /// Generated completion text, without the stop token. completion: String, /// Generated token IDs presented to the matcher, ending with a stop token. @@ -1150,13 +1244,13 @@ fn grammar_replay_dir() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/grammar_replay") } -/// Snapshot the output grammars of one model's roundtrips, keyed by variant, -/// as its grammar replay file, or check that no stale file remains when the +/// Snapshot the output grammars of one model's roundtrips, grouping the +/// variants that build the same grammar, as its grammar replay file, or check that no stale file remains when the /// parser builds none. fn check_grammar_replay_file( case: &RoundtripCase, backends: &vllm_chat::LoadedModelBackends, - results: Vec<(ToolChoiceVariant, RoundtripResult)>, + results: Vec<(ToolCallMixVariant, RoundtripResult)>, ) -> Result<()> { let model_name = case.model_id.rsplit_once('/').map_or(case.model_id, |(_, name)| name); let path = grammar_replay_dir().join(format!("{model_name}.json")); @@ -1197,7 +1291,7 @@ fn check_grammar_replay_file( case.assistant_stop_suffix ); - let mut cases = BTreeMap::new(); + let mut grammars = Vec::::new(); for (variant, grammar, closed_completion) in results { let (completion, mut generation_token_ids) = match completion_body( tokenizer.as_ref(), @@ -1208,27 +1302,41 @@ fn check_grammar_replay_file( CompletionBody::TokenIds(body) => (tokenizer.decode(body, false)?, body.to_vec()), }; generation_token_ids.push(stop_token_id); - cases.insert( - variant, - GrammarReplayCase { + let generation = GrammarReplayGeneration { + completion, + generation_token_ids, + }; + match grammars.iter_mut().find(|case| case.grammar == grammar) { + Some(case) => { + case.generations.insert(variant, generation); + } + None => grammars.push(GrammarReplayCase { grammar, - completion, - generation_token_ids, - }, - ); + generations: BTreeMap::from([(variant, generation)]), + }), + } } // A readable outline of the same grammars, for review only. let mut outline = String::new(); - for (variant, case) in &cases { + for case in &grammars { if !outline.is_empty() { outline.push('\n'); } - let variant = serde_json::to_value(variant)?; + let variants = case + .generations + .keys() + .map(|variant| { + serde_json::to_value(variant)? + .as_str() + .map(str::to_owned) + .context("variant name") + }) + .collect::>>()?; let coverage = serde_json::to_value(case.grammar.coverage)?; outline.push_str(&format!( "# {} ({})\n", - variant.as_str().context("variant name")?, + variants.join(", "), coverage.as_str().context("coverage name")?, )); outline.push_str(&vllm_parser::output_grammar::test_utils::outline( @@ -1241,7 +1349,7 @@ fn check_grammar_replay_file( model_id: case.model_id, vocab_size: backends.text_backend.model_vocab_size(), all_stop_token_ids, - cases, + grammars, }; let mut json = serde_json::to_string_pretty(&file)?; json.push('\n'); diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index 0108c3193ef7..ba9b4968d54b 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -132,6 +132,10 @@ pub struct OutputGrammarContext<'a> { pub tool_strict_level: ToolStrictLevel, /// Whether the request permits more than one tool call. pub parallel_tool_calls: bool, + /// The user's constraint on the answer text, such as a `response_format` + /// JSON schema, normalized to a format. A builder that composes it inserts + /// it where the model writes its answer; `None` leaves the answer free. + pub answer: Option<&'a Format>, } /// Errors produced while building an output grammar. @@ -243,9 +247,12 @@ fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { match ctx.tool_choice { ToolChoice::Value(ToolChoiceValue::None) => false, + // With an answer constraint, the grammar must hold the calls `auto` + // allows, or the answer would exclude them. ToolChoice::Value(ToolChoiceValue::Auto) => { ctx.tool_strict_level >= ToolStrictLevel::Function || ctx.tools.iter().any(|tool| tool.strict == Some(true)) + || ctx.answer.is_some() } ToolChoice::Value(ToolChoiceValue::Required) | ToolChoice::NamedFunction(_) @@ -284,6 +291,7 @@ mod tests { tool_choice, tool_strict_level: ToolStrictLevel::Auto, parallel_tool_calls: true, + answer: None, }, ) .unwrap() @@ -310,6 +318,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: level, parallel_tool_calls: true, + answer: None, }, ) .unwrap() diff --git a/rust/src/parser/src/tool/minimax_m2.rs b/rust/src/parser/src/tool/minimax_m2.rs index 3b868013bb32..33d316579f27 100644 --- a/rust/src/parser/src/tool/minimax_m2.rs +++ b/rust/src/parser/src/tool/minimax_m2.rs @@ -122,7 +122,7 @@ impl ToolParser for MinimaxM2ToolParser { Some(xgrammar_structural_tag::Model::Minimax.builder()) } - fn build_visible_format( + fn build_call_format( &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index e4afbda3e902..974d6c758a27 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -20,6 +20,7 @@ mod seed_oss; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; use std::collections::{BTreeMap, btree_map}; +use std::sync::Once; pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser, DeepSeekV41ToolParser}; pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser}; @@ -39,8 +40,11 @@ pub use qwen_coder::Qwen3CoderToolParser; pub use seed_oss::SeedOssToolParser; use serde::{Deserialize, Serialize}; use serde_json::Value; +use tracing::warn; +use xgrammar_structural_tag::ToolChoice; pub use xgrammar_structural_tag::builders::StructuralTagBuilder; use xgrammar_structural_tag::format::Format; +use xgrammar_structural_tag::tool::{AllowedToolsMode, ToolChoiceValue}; use crate::output_grammar::{self, OutputGrammarContext}; use crate::utils; @@ -196,21 +200,54 @@ pub trait ToolParser: Send { } /// Build the language of everything the model may emit after reasoning - /// ends: visible text plus tool calls. `None` means no tool grammar applies. + /// ends for the context's tool choice: visible text plus tool calls. `None` + /// means no tool grammar applies. /// /// The default forwards to [`Self::structural_tag_builder`] with reasoning /// disabled, which is exactly today's strict-tool-calling grammar. - fn build_visible_format( + fn build_call_format( &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { output_grammar::visible_format_from_builder(self.structural_tag_builder(), ctx) } + /// Build the visible language with the request's answer constraint + /// composed in, as the Python frontend does: an optional call becomes + /// `Or[calls built as required, answer]`, and a tool choice that obliges a + /// call ignores the answer. Parsers override [`Self::build_call_format`] + /// instead. + fn build_visible_format( + &self, + ctx: &OutputGrammarContext<'_>, + ) -> output_grammar::Result> { + let Some(answer) = ctx.answer else { + return self.build_call_format(ctx); + }; + match call_required(ctx.tool_choice) { + // Calls built from the optional choice would also admit free text + // in place of the answer. + Some(tool_choice) => { + let calls = self.build_call_format(&OutputGrammarContext { + tool_choice: &tool_choice, + ..*ctx + })?; + Ok(calls.map(|calls| Format::or(vec![calls, answer.clone()]))) + } + None => { + let calls = self.build_call_format(ctx)?; + if calls.is_some() { + warn_answer_ignored_once(); + } + Ok(calls) + } + } + } + /// Return the xgrammar structural-tag builder used for strict tool calling. /// /// To be deprecated: this only exists to back the default - /// [`Self::build_visible_format`]. New parsers should override that method + /// [`Self::build_call_format`]. New parsers should override that method /// directly instead of exposing a builder. fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { None @@ -243,6 +280,32 @@ pub trait ToolParser: Send { fn reset(&mut self) -> String; } +/// The tool choice with an optional call made obligatory, or `None` when the +/// choice forbids or already obliges a call. +fn call_required(tool_choice: &ToolChoice) -> Option { + let mut required = tool_choice.clone(); + match &mut required { + ToolChoice::Value(value @ ToolChoiceValue::Auto) => *value = ToolChoiceValue::Required, + ToolChoice::AllowedTools(choice) if choice.allowed_tools.mode == AllowedToolsMode::Auto => { + choice.allowed_tools.mode = AllowedToolsMode::Required; + } + ToolChoice::FlatAllowedTools(choice) if choice.mode == AllowedToolsMode::Auto => { + choice.mode = AllowedToolsMode::Required; + } + _ => return None, + } + Some(required) +} + +/// Warn that a request's answer constraint is ignored because its tool +/// choice obliges a call, as the Python frontend does. +fn warn_answer_ignored_once() { + static ONCE: Once = Once::new(); + ONCE.call_once(|| { + warn!("structured outputs are ignored because tool_choice forces a tool call"); + }); +} + /// Extension methods for easily testing `ToolParser` implementations. /// /// These helpers do not handle partial parsing or error recovery, so they are diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index ba44e7fb9613..b1f36e885546 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -87,10 +87,19 @@ impl UnifiedParser for CombinedParser { &self, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - let Some(tool) = self.tool.as_ref() else { + // Without an active reasoning or tool parser, the request's own + // constraint applies unchanged. + if self.reasoning.is_none() && self.tool.is_none() { return Ok(None); + } + + // The tool parser composes the answer constraint into its calls where + // it can; without a tool grammar, the answer is the visible language. + let calls = match self.tool.as_ref() { + Some(tool) => tool.build_visible_format(ctx)?, + None => None, }; - let Some(visible) = tool.build_visible_format(ctx)? else { + let Some(visible) = calls.or_else(|| ctx.answer.cloned()) else { return Ok(None); }; // The reasoning parser may wrap that language with the reasoning phase @@ -165,7 +174,8 @@ mod tests { use super::CombinedParser; use crate::output_grammar::{ - GrammarCoverage, OutputGrammarContext, full_format_from_builder_for_test, + BuiltOutputGrammar, GrammarCoverage, OutputGrammarContext, + full_format_from_builder_for_test, }; use crate::reasoning::{ DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser, Glm45ReasoningParser, @@ -235,6 +245,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let tool = T::create(&tools).unwrap(); let expected = full_format_from_builder_for_test( @@ -524,6 +535,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); let tool = Qwen3XmlToolParser::create(&tools).unwrap(); @@ -553,6 +565,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let visible = Format::const_string("answer"); @@ -586,6 +599,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let visible = Format::const_string("visible"); @@ -622,6 +636,7 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }; let tokenizer = Arc::new(tokenizer()); let reasoning = Qwen3ReasoningParser::create(tokenizer.clone()).unwrap(); @@ -650,4 +665,97 @@ mod tests { ); } } + + fn grammar_ctx<'a>( + tools: &'a [Tool], + tool_choice: &'a ToolChoice, + answer: Option<&'a Format>, + ) -> OutputGrammarContext<'a> { + OutputGrammarContext { + tools, + tool_choice, + tool_strict_level: Default::default(), + parallel_tool_calls: true, + answer, + } + } + + fn qwen3_parser(tools: &[Tool], with_tool: bool) -> CombinedParser { + let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); + let tool = with_tool.then(|| Qwen3XmlToolParser::create(tools).unwrap()); + let mut parser = CombinedParser::new(Some(reasoning), tool); + parser.initialize(&[]).unwrap(); + parser + } + + #[test] + fn answer_constraint_composes_with_the_tool_choice() { + let tools = test_tools(); + let answer = Format::json_schema(serde_json::json!({ "type": "object" })); + let reasoning = |visible| { + Format::sequence(vec![ + Format::optional(Format::sequence(vec![ + Format::tag("", Format::any_text(), ""), + Format::const_string("\n\n"), + ])), + visible, + ]) + }; + let required = ToolChoice::required(); + let required_calls = Qwen3XmlToolParser::create(&tools) + .unwrap() + .build_visible_format(&grammar_ctx(&tools, &required, None)) + .unwrap() + .unwrap(); + + // No callable tool: the answer alone, after the reasoning phase. + for (tool_choice, with_tool) in [(ToolChoice::none(), true), (ToolChoice::auto(), false)] { + let parser = qwen3_parser(&tools, with_tool); + let ctx = grammar_ctx(&tools, &tool_choice, Some(&answer)); + let actual = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(actual.coverage, GrammarCoverage::FromTokenZero); + assert_eq!(actual.format, reasoning(answer.clone()), "{tool_choice:?}"); + } + + // `auto`, strict or not, keeps tool calls next to the answer. + let parser = qwen3_parser(&tools, true); + let auto = ToolChoice::auto(); + let actual = parser.build_output_grammar(&grammar_ctx(&tools, &auto, Some(&answer))); + let expected = Format::or(vec![required_calls, answer.clone()]); + assert_eq!(actual.unwrap().unwrap().format, reasoning(expected)); + + // A required call ignores the answer constraint. + for tool_choice in [ToolChoice::required(), ToolChoice::function("get_weather")] { + let with_answer = grammar_ctx(&tools, &tool_choice, Some(&answer)); + let without = grammar_ctx(&tools, &tool_choice, None); + assert_eq!( + parser.build_output_grammar(&with_answer).unwrap(), + parser.build_output_grammar(&without).unwrap(), + "{tool_choice:?}" + ); + } + } + + #[test] + fn answer_without_a_reasoning_phase_covers_the_final_output() { + let tools = test_tools(); + let answer = Format::regex("[a-z]+"); + let none = ToolChoice::none(); + let auto = ToolChoice::auto(); + + // Without an active parser, the request keeps its own constraint. + let parser = CombinedParser::new(None, None); + let ctx = grammar_ctx(&tools, &auto, Some(&answer)); + assert!(parser.build_output_grammar(&ctx).unwrap().is_none()); + + let parser = CombinedParser::new(None, Some(Qwen3XmlToolParser::create(&tools).unwrap())); + let ctx = grammar_ctx(&tools, &none, Some(&answer)); + let alone = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(alone, BuiltOutputGrammar::final_output_only(answer.clone())); + + let ctx = grammar_ctx(&tools, &auto, Some(&answer)); + let composed = parser.build_output_grammar(&ctx).unwrap().unwrap(); + assert_eq!(composed.coverage, GrammarCoverage::FinalOutputOnly); + assert!(matches!(composed.format, Format::Or(_))); + } } diff --git a/rust/src/parser/src/unified/hy/hy_v3.rs b/rust/src/parser/src/unified/hy/hy_v3.rs index 3599ce2cb9a0..18a34954b3b6 100644 --- a/rust/src/parser/src/unified/hy/hy_v3.rs +++ b/rust/src/parser/src/unified/hy/hy_v3.rs @@ -164,6 +164,7 @@ mod tests { tools: &tools, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, tool_choice: &ToolChoice::required(), }) .unwrap() @@ -194,6 +195,7 @@ mod tests { tool_choice: &ToolChoice::required(), tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, }) .unwrap() .unwrap(); diff --git a/rust/src/parser/src/unified/hy/hy_v4.rs b/rust/src/parser/src/unified/hy/hy_v4.rs index a5d81e0f2ae7..f10fb662eca3 100644 --- a/rust/src/parser/src/unified/hy/hy_v4.rs +++ b/rust/src/parser/src/unified/hy/hy_v4.rs @@ -178,6 +178,7 @@ mod tests { tools: &tools, tool_strict_level: Default::default(), parallel_tool_calls: true, + answer: None, tool_choice: &ToolChoice::required(), }) .unwrap() diff --git a/rust/src/parser/src/unified/kimi_k3/grammar.rs b/rust/src/parser/src/unified/kimi_k3/grammar.rs index 166fbafb0f82..e6ef746fa2ac 100644 --- a/rust/src/parser/src/unified/kimi_k3/grammar.rs +++ b/rust/src/parser/src/unified/kimi_k3/grammar.rs @@ -16,6 +16,10 @@ //! //! Every path may end with the message close, and a stop token is legal only //! after the last channel closes. +//! +//! The user's answer constraint, if any, replaces the free text of the +//! `response` channel. Without a callable tool, the `response` channel then +//! follows `think` alone. mod arguments; @@ -33,38 +37,44 @@ use crate::output_grammar::{ }; /// Build the grammar of everything generated after a prompt that left the -/// parser in `mode`, or `None` when the request asks for no tool grammar. +/// parser in `mode`, or `None` when the request asks for no tool grammar and +/// has no answer constraint. pub(super) fn build_output_grammar( mode: &KimiK3Mode, ctx: &OutputGrammarContext<'_>, ) -> output_grammar::Result> { - let Some(NormalizedToolChoice { - function_tools, - choice, - .. - }) = normalize_tool_choice(ctx)? - else { + let normalized = normalize_tool_choice(ctx)?; + if normalized.is_none() && ctx.answer.is_none() { return Ok(None); - }; - - let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; - let tools = Format::tag( - TOOLS_OPEN, - Format::tags_with_separator( - function_tools.iter().map(call_tag).collect(), - "", - true, - single_call, - ), - TOOLS_CLOSE, + } + let calls = normalized.map( + |NormalizedToolChoice { + function_tools, + choice, + .. + }| { + let single_call = choice == BuilderToolChoice::Forced || !ctx.parallel_tool_calls; + let tools = Format::tag( + TOOLS_OPEN, + Format::tags_with_separator( + function_tools.iter().map(call_tag).collect(), + "", + true, + single_call, + ), + TOOLS_CLOSE, + ); + (choice, tools) + }, ); + let answer = ctx.answer; let mut elements = match mode { - KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(choice, tools)], - KimiK3Mode::Response => vec![response_body(), after_response(choice, tools)], + KimiK3Mode::Reasoning => vec![think_body(), after_reasoning(calls, answer)], + KimiK3Mode::Response => rest_of_response(calls, answer), KimiK3Mode::Idle => vec![ Format::optional(Format::tag(THINK_OPEN, think_text(), THINK_CLOSE)), - after_reasoning(choice, tools), + after_reasoning(calls, answer), ], KimiK3Mode::Epilogue | KimiK3Mode::Tools | KimiK3Mode::Call { .. } | KimiK3Mode::Done => { unreachable!("initialization leaves the parser idle or inside `think` or `response`") @@ -77,38 +87,73 @@ pub(super) fn build_output_grammar( )))) } -/// Channels that may follow the closed `think` channel. -fn after_reasoning(tool_choice: BuilderToolChoice, tools: Format) -> Format { - let response = Format::tag(RESPONSE_OPEN, response_text(), RESPONSE_CLOSE); - match tool_choice { - BuilderToolChoice::Auto => Format::or(vec![ +/// Channels that may follow the closed `think` channel, given the tool +/// choice and `tools` channel when a tool is callable, and the user's answer +/// constraint for the `response` channel. +fn after_reasoning(calls: Option<(BuilderToolChoice, Format)>, answer: Option<&Format>) -> Format { + let response = Format::tag(RESPONSE_OPEN, response_content(answer), RESPONSE_CLOSE); + let Some((choice, tools)) = calls else { + return response; + }; + let empty = || Format::const_string(format!("{RESPONSE_OPEN}{RESPONSE_CLOSE}")); + match (choice, answer) { + (BuilderToolChoice::Auto, None) => Format::or(vec![ Format::sequence(vec![response, Format::optional(tools.clone())]), tools, ]), - BuilderToolChoice::Required | BuilderToolChoice::Forced => { + (BuilderToolChoice::Auto, Some(_)) => Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + Format::sequence(vec![Format::optional(empty()), tools]), + ]), + (BuilderToolChoice::Required | BuilderToolChoice::Forced, None) => { Format::sequence(vec![Format::optional(response), tools]) } + (BuilderToolChoice::Required | BuilderToolChoice::Forced, Some(_)) => { + Format::sequence(vec![ + Format::optional(Format::or(vec![response, empty()])), + tools, + ]) + } } } -/// Channels that may follow the closed `response` channel. -fn after_response(tool_choice: BuilderToolChoice, tools: Format) -> Format { - match tool_choice { - BuilderToolChoice::Auto => Format::optional(tools), - BuilderToolChoice::Required | BuilderToolChoice::Forced => tools, +/// The rest of the `response` channel the prompt opened, then the channels +/// that may follow it. +fn rest_of_response( + calls: Option<(BuilderToolChoice, Format)>, + answer: Option<&Format>, +) -> Vec { + let response = Format::tag("", response_content(answer), RESPONSE_CLOSE); + let Some((choice, tools)) = calls else { + return vec![response]; + }; + let empty = || Format::const_string(RESPONSE_CLOSE); + match (choice, answer) { + (BuilderToolChoice::Auto, None) => vec![response, Format::optional(tools)], + (BuilderToolChoice::Auto, Some(_)) => vec![Format::or(vec![ + Format::sequence(vec![response, Format::optional(tools.clone())]), + Format::sequence(vec![empty(), tools]), + ])], + (BuilderToolChoice::Required | BuilderToolChoice::Forced, None) => vec![response, tools], + (BuilderToolChoice::Required | BuilderToolChoice::Forced, Some(_)) => { + vec![Format::or(vec![response, empty()]), tools] + } } } +/// The text of the `response` channel: the user's answer constraint, or free +/// text. A constrained answer rules out the empty `response` channel the chat +/// template writes before tool calls, so callers allow that channel where +/// calls follow. +fn response_content(answer: Option<&Format>) -> Format { + answer.cloned().unwrap_or_else(response_text) +} + /// The rest of the `think` channel the prompt opened. fn think_body() -> Format { Format::tag("", think_text(), THINK_CLOSE) } -/// The rest of the `response` channel the prompt opened. -fn response_body() -> Format { - Format::tag("", response_text(), RESPONSE_CLOSE) -} - fn think_text() -> Format { Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]) } @@ -158,6 +203,28 @@ mod tests { tool_choice: &tool_choice, tool_strict_level: ToolStrictLevel::Function, parallel_tool_calls, + answer: None, + }, + ) + .unwrap() + .unwrap(); + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + let mut format = grammar.format; + elide_call_bodies(&mut format); + expected.assert_eq(&outline(&format)); + } + + /// Like [`check`], for non-strict tools and a JSON answer constraint. + fn check_answer(mode: KimiK3Mode, tools: &[Tool], tool_choice: ToolChoice, expected: Expect) { + let answer = Format::json_schema(json!({ "type": "object" })); + let grammar = build_output_grammar( + &mode, + &OutputGrammarContext { + tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + answer: Some(&answer), }, ) .unwrap() @@ -348,9 +415,120 @@ mod tests { tool_choice: &ToolChoice::auto(), tool_strict_level: ToolStrictLevel::Auto, parallel_tool_calls: true, + answer: None, }, ) .unwrap(); assert_eq!(grammar, None); } + + #[test] + fn answer_constraint_fills_the_response_channel() { + // Non-strict `auto` keeps the calls next to the constrained answer. + // Calls may follow the empty `response` channel the chat template + // writes before them. + check_answer( + KimiK3Mode::Reasoning, + &tools(), + ToolChoice::auto(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + or + sequence + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + sequence + optional `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + // A required call may follow a constrained answer. + check_answer( + KimiK3Mode::Reasoning, + &tools(), + ToolChoice::required(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + sequence + optional + or + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + `<|open|>response<|sep|><|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Response, + &tools(), + ToolChoice::auto(), + expect![[r#" + sequence + or + sequence + tag `` json(object) `<|close|>response<|sep|>` + optional + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + sequence + `<|close|>response<|sep|>` + tag `<|open|>tools<|sep|>` .. `<|close|>tools<|sep|>` + tags_with_separator `` at_least_one + tag `<|open|>call tool="get_weather" index="` `..` `<|close|>call<|sep|>` + tag `<|open|>call tool="add" index="` `..` `<|close|>call<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } + + #[test] + fn answer_constraint_without_a_callable_tool_requires_the_response() { + check_answer( + KimiK3Mode::Idle, + &tools(), + ToolChoice::none(), + expect![[r#" + sequence + optional tag `<|open|>think<|sep|>` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Reasoning, + &[], + ToolChoice::auto(), + expect![[r#" + sequence + tag `` text excluding [`<|close|>think<|sep|>`, `<|end_of_msg|>`] `<|close|>think<|sep|>` + tag `<|open|>response<|sep|>` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + check_answer( + KimiK3Mode::Response, + &[], + ToolChoice::auto(), + expect![[r#" + sequence + tag `` json(object) `<|close|>response<|sep|>` + optional `<|close|>message<|sep|>` + "#]], + ); + } } From 2534a96fcf165202f5358335f8b0705175159fd0 Mon Sep 17 00:00:00 2001 From: Jimmy Wang <53066820+JimmyWang0417@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:47:52 +0800 Subject: [PATCH 398/952] [Docs] List additional unauthenticated endpoints (#61081) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Jimmy Wang <53066820+JimmyWang0417@users.noreply.github.com> Co-authored-by: Juan Pérez de Algaba <124347725+jperezdealgaba@users.noreply.github.com> --- docs/usage/security.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/usage/security.md b/docs/usage/security.md index f069de0def74..0d7894a467a2 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -208,9 +208,13 @@ The following endpoints **do not require authentication** even when `--api-key` - `/update_weights` - Update model weights (can alter model behavior) - `/get_world_size` - Get distributed world size - `/abort_requests` - Abort in-flight requests (available with `--tokens-only`. Use the authenticated `/inference/v1/abort_requests` otherwise) +- `/fault_tolerance/apply` - Trigger fault-recovery operations across ranks (available when fault tolerance is enabled) +- `/fault_tolerance/status` - Query fault-recovery status (available when fault tolerance is enabled) **Utility endpoints:** +- `/metrics` - Prometheus metrics, including model, throughput, queue, and KV-cache statistics +- `/kv_event_sources` - KV-event publisher configuration for each data-parallel rank, which may reveal internal transport endpoints (when KV events are configured) - `/tokenize` - Tokenize text (not gated by `--enable-scale-out`) - `/detokenize` - Detokenize tokens - `/health` - Health check From f7a3a7265f5161594721e8c1aa6b68a761629b68 Mon Sep 17 00:00:00 2001 From: MatthewCYM <40845677+MatthewCYM@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:42:51 +0800 Subject: [PATCH 399/952] [Bugfix][Multimodal] Fix Qwen3-ASR audio preprocessing parity (#57146) Signed-off-by: MatthewCYM --- .../multimodal/processing/test_qwen3_asr.py | 181 ++++++++++++++ vllm/model_executor/models/qwen3_asr.py | 73 +++++- .../transformers_utils/processors/__init__.py | 2 - .../processors/qwen3_asr.py | 229 ------------------ 4 files changed, 243 insertions(+), 242 deletions(-) create mode 100644 tests/models/multimodal/processing/test_qwen3_asr.py delete mode 100644 vllm/transformers_utils/processors/qwen3_asr.py diff --git a/tests/models/multimodal/processing/test_qwen3_asr.py b/tests/models/multimodal/processing/test_qwen3_asr.py new file mode 100644 index 000000000000..b016a9d70860 --- /dev/null +++ b/tests/models/multimodal/processing/test_qwen3_asr.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen3-ASR preprocessing must preserve its HF processor's audio inputs.""" + +import numpy as np +import pytest +import torch +from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen3_asr import Qwen3ASRProcessor + +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache +from vllm.multimodal.inputs import batched_tensors_equal +from vllm.transformers_utils.processor import cached_processor_from_config + +from .test_common import _assert_inputs_equal + +_AUDIO_PROMPT = "<|audio_start|><|audio_pad|><|audio_end|>" + + +@pytest.fixture(params=["Qwen/Qwen3-ASR-0.6B", "Qwen/Qwen3-ASR-0.6B-hf"]) +def asr_processor(request): + config = ModelConfig( + request.param, + dtype="float32", + max_model_len=4096, + limit_mm_per_prompt={"audio": 2}, + mm_processor_cache_gb=0, + ) + processor = MULTIMODAL_REGISTRY.create_processor(config) + native = Qwen3ASRProcessor.from_pretrained(request.param) + return processor, native + + +@pytest.mark.parametrize( + "samples", [321, 15999, 16000, 16001, 16159, 16160, 16161, 32001, 480001] +) +@pytest.mark.parametrize("num_audios", [1, 2]) +def test_audio_processing_matches_hf(asr_processor, samples, num_audios): + """Each clip must match HF independently of other clips in the request.""" + processor, native = asr_processor + rng = np.random.RandomState(42) + audios = [ + rng.randn(samples + index * 160).astype(np.float32) + for index in range(num_audios) + ] + expected = [ + native(text=_AUDIO_PROMPT, audio=audio, return_tensors="pt") for audio in audios + ] + actual = processor( + _AUDIO_PROMPT * num_audios, + mm_items=processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + + expected_ids = [ + token for item in expected for token in item["input_ids"][0].tolist() + ] + assert actual["prompt_token_ids"] == expected_ids + masks = [item["input_features_mask"][0] for item in expected] + expected_features = torch.cat( + [ + item["input_features"][0, :, mask.bool()] + for item, mask in zip(expected, masks) + ], + dim=1, + ) + actual_features = actual["mm_kwargs"].get_data() + torch.testing.assert_close( + actual_features["input_audio_features"], expected_features, rtol=0, atol=0 + ) + torch.testing.assert_close( + actual_features["audio_feature_lengths"], + torch.stack([mask.sum() for mask in masks]), + ) + + +@pytest.mark.parametrize("cached_indices", [[], [0], [1], [0, 1]]) +def test_audio_processing_cache_is_batch_independent(asr_processor, cached_indices): + """Partial cache hits must not change features, masks, or prompt expansion.""" + processor, _ = asr_processor + config = processor.info.ctx.model_config + config.get_multimodal_config().mm_processor_cache_gb = 1 + cache = MultiModalProcessorOnlyCache(config) + rng = np.random.RandomState(42) + audios = [rng.randn(samples).astype(np.float32) for samples in (16001, 32000)] + for index in cached_indices: + processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": audios[index]}), + hf_processor_mm_kwargs={}, + cache=cache, + ) + + for batch in (audios, audios[::-1]): + kwargs = dict( + mm_items=processor.info.parse_mm_data({"audio": batch}), + hf_processor_mm_kwargs={}, + ) + _assert_inputs_equal( + processor(_AUDIO_PROMPT * 2, **kwargs), + processor(_AUDIO_PROMPT * 2, **kwargs, cache=cache), + ) + + +def test_text_only_processing(asr_processor): + """Empty media input must not invoke the audio feature extractor.""" + processor, native = asr_processor + prompt = "Transcribe the audio." + actual = processor( + prompt, + mm_items=processor.info.parse_mm_data({}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == native.tokenizer.encode(prompt) + assert actual["mm_kwargs"].get_data() == {} + + +@pytest.mark.parametrize("hf_data", [{}, {"text": ""}, {"audio": None}, {"audio": []}]) +def test_hf_processor_without_audio(asr_processor, hf_data): + """The media hook must not require audio or call HF for absent waveforms.""" + processor, _ = asr_processor + assert not processor._call_hf_processor(hf_data, {}) + + +@pytest.mark.parametrize("hf_data", [{}, {"text": ""}]) +def test_postprocessing_without_audio_features(asr_processor, hf_data): + """Non-audio output must pass through even when HF input is nonempty.""" + processor, _ = asr_processor + processed = BatchFeature({"input_ids": torch.tensor([[1, 2]])}) + actual = processor._postprocess_hf_mm_data(hf_data, {}, processed) + assert actual is processed + assert set(actual) == {"input_ids"} + torch.testing.assert_close(actual["input_ids"], torch.tensor([[1, 2]])) + + +def test_chat_template_processing(asr_processor): + """The generic processor loader must still render audio chat prompts.""" + processor, native = asr_processor + chat_processor = cached_processor_from_config(processor.info.ctx.model_config) + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": [{"type": "audio", "audio": "unused"}]}, + ] + prompt = chat_processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + expected_prompt = native.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + assert prompt == expected_prompt + audio = np.random.RandomState(42).randn(16001).astype(np.float32) + expected = native(text=expected_prompt, audio=audio, return_tensors="pt") + actual = processor( + prompt, + mm_items=processor.info.parse_mm_data({"audio": audio}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == expected["input_ids"][0].tolist() + + +def test_precomputed_audio_features(asr_processor): + """Precomputed features must bypass waveform extraction unchanged.""" + processor, _ = asr_processor + processor.info.ctx.get_mm_config().enable_mm_embeds = True + audio = np.random.RandomState(42).randn(16001).astype(np.float32) + expected = processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": audio}), + hf_processor_mm_kwargs={}, + ) + features = expected["mm_kwargs"].get_data() + actual = processor( + _AUDIO_PROMPT, + mm_items=processor.info.parse_mm_data({"audio": features}), + hf_processor_mm_kwargs={}, + ) + assert actual["prompt_token_ids"] == expected["prompt_token_ids"] + assert actual["mm_placeholders"] == expected["mm_placeholders"] + assert batched_tensors_equal(actual["mm_kwargs"].get_data(), features) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index c8a0849fc56a..161d1efacd06 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -29,6 +29,7 @@ import torch import torch.nn as nn from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen3_asr import Qwen3ASRFeatureExtractor, Qwen3ASRProcessor from transformers.models.whisper import WhisperFeatureExtractor from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig @@ -54,7 +55,6 @@ from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoder, - Qwen3OmniMoeThinkerMultiModalProcessor, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -80,6 +80,7 @@ ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, + BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, @@ -92,9 +93,6 @@ Qwen3ASRThinkerConfig, ) from vllm.transformers_utils.processor import cached_processor_from_config -from vllm.transformers_utils.processors.qwen3_asr import ( - Qwen3ASRProcessor, -) from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d @@ -188,18 +186,19 @@ def get_hf_config(self): return self.ctx.get_hf_config(Qwen3ASRConfig).thinker_config def get_hf_processor(self, **kwargs: object) -> Qwen3ASRProcessor: - processor = self.ctx.get_hf_processor( + return self.ctx.get_hf_processor( Qwen3ASRProcessor, **kwargs, ) - if not hasattr(processor, "audio_token"): - processor.audio_token = "<|audio_pad|>" - return processor - def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor: + def get_feature_extractor( + self, **kwargs: object + ) -> WhisperFeatureExtractor | Qwen3ASRFeatureExtractor: hf_processor = self.get_hf_processor(**kwargs) feature_extractor = hf_processor.feature_extractor - assert isinstance(feature_extractor, WhisperFeatureExtractor) + assert isinstance( + feature_extractor, (WhisperFeatureExtractor, Qwen3ASRFeatureExtractor) + ) return feature_extractor def get_supported_mm_limits(self) -> Mapping[str, int | None]: @@ -275,8 +274,60 @@ def _parse_audio_data( class Qwen3ASRMultiModalProcessor( - Qwen3OmniMoeThinkerMultiModalProcessor[Qwen3ASRProcessingInfo] + BaseMultiModalProcessor[Qwen3ASRProcessingInfo], ): + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: + return self.dummy_inputs.get_dummy_text(mm_counts) + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + audios = hf_data.get("audio") + if not audios: + return BatchFeature() + assert isinstance(audios, list) + + if len(audios) == 1: + return super()._call_hf_processor(hf_data, hf_kwargs) + + # HF requires one text per audio, and batch padding changes clip features. + # Keep clips independent so partial cache hits cannot change their inputs. + text = self._get_hf_mm_text({"audio": 1}) + outputs: dict[str, list[torch.Tensor]] = { + "input_features": [], + "input_features_mask": [], + } + for audio in audios: + processed = super()._call_hf_processor( + dict(hf_data, text=text, audio=audio), hf_kwargs + ) + for key, values in outputs.items(): + values.append(processed[key][0]) + + return BatchFeature(outputs) + + def _postprocess_hf_mm_data( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + if "input_features" not in processed_data: + return processed_data + + features = processed_data.pop("input_features") + masks = processed_data.pop("input_features_mask") + processed_data["input_audio_features"] = torch.cat( + [feature[:, mask.bool()] for feature, mask in zip(features, masks)], dim=1 + ) + processed_data["feature_attention_mask"] = masks + processed_data["audio_feature_lengths"] = torch.stack( + [mask.sum() for mask in masks] + ) + return processed_data + def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index f2af4901e2e9..db45e49d5adc 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -38,7 +38,6 @@ "OpenVLAProcessor", "OvisProcessor", "Ovis2_5Processor", - "Qwen3ASRProcessor", "Step3VLProcessor", "InklingProcessor", "InklingImageProcessor", @@ -74,7 +73,6 @@ "OpenVLAProcessor": "vllm.transformers_utils.processors.openvla", "OvisProcessor": "vllm.transformers_utils.processors.ovis", "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5", - "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr", "Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl", "InklingProcessor": "vllm.transformers_utils.processors.inkling", "InklingImageProcessor": "vllm.transformers_utils.processors.inkling", diff --git a/vllm/transformers_utils/processors/qwen3_asr.py b/vllm/transformers_utils/processors/qwen3_asr.py deleted file mode 100644 index 55d38537928d..000000000000 --- a/vllm/transformers_utils/processors/qwen3_asr.py +++ /dev/null @@ -1,229 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# ruff: noqa -# mypy: ignore-errors -# coding=utf-8 -# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import regex as re - -import numpy as np - -from transformers import AutoProcessor -from transformers.audio_utils import AudioInput -from transformers.feature_extraction_utils import BatchFeature -from transformers.processing_utils import ProcessingKwargs, ProcessorMixin -from transformers.tokenization_utils_base import TextInput - - -class Qwen3ASRProcessorKwargs(ProcessingKwargs, total=False): - _defaults = { - "text_kwargs": { - "padding": False, - "padding_side": "left", - }, - "audio_kwargs": { - "sampling_rate": 16000, - "padding": True, - "return_attention_mask": True, - }, - } - - -def _get_feat_extract_output_lengths(input_lengths): - """ - Computes the output length of the convolutional layers and the output length of the audio encoder - """ - - input_lengths_leave = input_lengths % 100 - feat_lengths = (input_lengths_leave - 1) // 2 + 1 - output_lengths = ( - ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 - ) - return output_lengths - - -class Qwen3ASRProcessor(ProcessorMixin): - r""" - Constructs a Qwen3ASR processor. - [`Qwen3ASRProcessor`] offers all the functionalities of [`WhisperFeatureExtractor`], and [`Qwen2TokenizerFast`]. See the - [`~Qwen3ASRProcessor.__call__`] and [`~Qwen3ASRProcessor.decode`] for more information. - - Args: - feature_extractor ([`WhisperFeatureExtractor`], *optional*): - The audio feature extractor. - tokenizer ([`Qwen2TokenizerFast`], *optional*): - The text tokenizer. - chat_template (`Optional[str]`, *optional*): - The Jinja template to use for formatting the conversation. If not provided, the default chat template is used. - """ - - attributes = ["feature_extractor", "tokenizer"] - feature_extractor_class = "WhisperFeatureExtractor" - tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast") - - def __init__(self, feature_extractor=None, tokenizer=None, chat_template=None): - super().__init__(feature_extractor, tokenizer, chat_template=chat_template) - self.audio_token = self.tokenizer.audio_token - self.audio_bos_token = self.tokenizer.audio_bos_token - self.audio_eos_token = self.tokenizer.audio_eos_token - - def __call__( - self, - text: TextInput = None, - audio: AudioInput = None, - **kwargs, - ) -> BatchFeature: - """ - Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text` - and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode - the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to - WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. Please refer to the doctsring - of the above two methods for more information. - - Args: - text (`str`, `List[str]`, `List[List[str]]`): - The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings - (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set - `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). - audio (`np.ndarray`, `List[np.ndarray]`): - The audio or batch of audio to be prepared. Each audio can be a NumPy array. - """ - - if text is None: - raise ValueError("You need to specify either a `text` input to process.") - - output_kwargs = self._merge_kwargs( - Qwen3ASRProcessorKwargs, - tokenizer_init_kwargs=self.tokenizer.init_kwargs, - **kwargs, - ) - - if audio is not None: - output_kwargs["audio_kwargs"]["padding"] = True - output_kwargs["audio_kwargs"]["truncation"] = False - audio_inputs = self.feature_extractor( - audio, **output_kwargs["audio_kwargs"] - ) - audio_inputs["feature_attention_mask"] = audio_inputs.pop( - "attention_mask" - ) # rename feature_attention_mask to prevent conflicts later on - audio_inputs["input_features"] = audio_inputs.pop( - "input_features" - ) # rename input_features to prevent conflicts later on - audio_lengths = iter( - _get_feat_extract_output_lengths( - audio_inputs["feature_attention_mask"].sum(-1) - ) - ) - else: - audio_inputs = {} - audio_lengths = iter([]) - - if not isinstance(text, list): - text = [text] - - text = self.replace_multimodal_special_tokens( - text, - audio_lengths, - ) - - texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) - - return BatchFeature( - data={**texts_inputs, **audio_inputs}, - tensor_type=kwargs.get("return_tensors"), - ) - - def replace_multimodal_special_tokens( - self, - text, - audio_lengths, - ): - processed_text = [] - for sample in text: - positions = [] - special_tokens = [re.escape(tok) for tok in [self.audio_token]] - pattern = "|".join(special_tokens) - positions = sorted( - [ - (match.start(), match.group()) - for match in re.finditer(pattern, sample) - ] - ) - positions.sort(key=lambda x: x[0]) - - for _, special_token in positions: - if special_token == self.audio_token: - sample = sample.replace( - self.audio_token, - "<|audio_placeholder|>" * next(audio_lengths), - 1, - ) - - sample = sample.replace("<|audio_placeholder|>", self.audio_token) - processed_text.append(sample) - return processed_text - - def get_chunked_index( - self, token_indices: np.ndarray, tokens_per_chunk: int - ) -> list[tuple[int, int]]: - """ - Splits token index list into chunks based on token value ranges. - - Given a list of token indices, returns a list of (start, end) index tuples representing - slices of the list where the token values fall within successive ranges of `tokens_per_chunk`. - - For example, if `tokens_per_chunk` is 1000, the function will create chunks such that: - - the first chunk contains token values < 1000, - - the second chunk contains values >= 1000 and < 2000, and so on. - - Parameters: - token_indices (`np.ndarray`): A monotonically increasing list of token index values. - tokens_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold). - - Returns: - `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive) - and end (exclusive) indices of a chunk in `token_indices`. - """ - - def _iter(): - i, start_idx = 0, 0 # skip bos token - current_chunk = 1 - while i < len(token_indices): # skip eos token - if token_indices[i] >= current_chunk * tokens_per_chunk: - yield (start_idx, i) - start_idx = i - current_chunk += 1 - i += 1 - yield (start_idx, len(token_indices)) - - return list(_iter()) - - def apply_chat_template(self, conversations, chat_template=None, **kwargs): - kwargs["return_dict"] = False - return super().apply_chat_template(conversations, chat_template, **kwargs) - - @property - def model_input_names(self): - tokenizer_input_names = self.tokenizer.model_input_names - feature_extractor_input_names = self.feature_extractor.model_input_names - return list( - dict.fromkeys( - tokenizer_input_names - + feature_extractor_input_names - + ["feature_attention_mask"] - ) - ) From ee8f7ea4493bf95e7c04951aec701af4000abdb7 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Mon, 12 Oct 2026 06:46:07 -0400 Subject: [PATCH 400/952] [CI] Enroll a fourth batch of five steps in automatic sharding (#60858) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 3 +++ .buildkite/test_areas/engine.yaml | 1 + .buildkite/test_areas/misc.yaml | 1 + 3 files changed, 5 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 78a6af2e0ca4..f83946386a38 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -77,6 +77,7 @@ steps: - label: ":nvidia: (L4) Distributed DP Basic" key: distributed-dp-tests-2-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 @@ -300,6 +301,7 @@ steps: - label: ":nvidia: (L4) Distributed DP Extended" key: distributed-dp-tests-4-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 @@ -677,6 +679,7 @@ steps: - label: ":nvidia: (L4) RayExecutorV2" key: rayexecutorv2-4-gpus + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" device: l4 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 26485e0a172d..21a65bb5ed6b 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -164,6 +164,7 @@ steps: device: h100 num_devices: 2 key: replayssm-e2e + automatic_shard: true timeout_in_minutes: 90 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 93f59b50e52f..20b48cb98f1f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -653,6 +653,7 @@ steps: - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 + automatic_shard: true timeout_in_minutes: 60 device: h100 source_file_dependencies: From 21445806bdf6cfc4fdb9b161643e081bd73cc716 Mon Sep 17 00:00:00 2001 From: Zupeng Wang Date: Mon, 12 Oct 2026 19:12:30 +0800 Subject: [PATCH 401/952] [Rust Frontend] Support logprob_token_ids on OpenAI endpoints (#59549) Signed-off-by: Wang Zupeng Signed-off-by: Bugen Zhao Co-authored-by: OpenAI Codex Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../routes/openai/chat_completions/convert.rs | 38 +++- .../routes/openai/chat_completions/types.rs | 6 + .../openai/chat_completions/validate.rs | 7 + .../src/routes/openai/completions/convert.rs | 30 ++- .../src/routes/openai/completions/types.rs | 5 + .../src/routes/openai/completions/validate.rs | 16 ++ .../src/routes/openai/utils/logprobs.rs | 24 ++- rust/src/server/src/routes/tests.rs | 171 ++++++++++++++++++ 8 files changed, 292 insertions(+), 5 deletions(-) diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 46aaafbc69e2..30c4fc8ad24f 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -126,6 +126,13 @@ pub(super) fn prepare_chat_request( .prompt_logprobs .or((request.echo && !request.stream).then_some(top_logprobs)); let include_prompt_logprobs = prompt_logprobs.is_some(); + let logprob_token_ids = request.logprob_token_ids.filter(|ids| !ids.is_empty()); + let logprobs = (request.logprobs && logprob_token_ids.is_none()).then_some(top_logprobs); + let output_top_logprobs = if logprob_token_ids.is_some() { + -1 + } else { + top_logprobs + }; let structured_outputs = convert_from_response_format( request.response_format.as_ref(), @@ -157,7 +164,7 @@ pub(super) fn prepare_chat_request( max_tokens: request.max_completion_tokens, min_tokens: request.min_tokens, thinking_token_budget: request.thinking_token_budget, - logprobs: request.logprobs.then_some(top_logprobs), + logprobs, prompt_logprobs, prompt_logprob_token_ids: None, prompt_logprob_start: None, @@ -172,7 +179,7 @@ pub(super) fn prepare_chat_request( allowed_token_ids: request.allowed_token_ids, bad_words: request.bad_words, bad_words_token_ids: None, - logprob_token_ids: None, + logprob_token_ids, structured_outputs, skip_reading_prefix_cache: None, vllm_xargs: merge_kv_transfer_params( @@ -213,7 +220,7 @@ pub(super) fn prepare_chat_request( include_usage, include_continuous_usage, requested_logprobs, - output_top_logprobs: top_logprobs, + output_top_logprobs, include_prompt_logprobs, include_reasoning, echo, @@ -1556,6 +1563,31 @@ mod tests { assert!(prepared.options.is_named_tool_choice); } + #[test] + fn selected_logprobs_override_chat_top_k_unless_empty() { + for ids in [vec![10, 20], vec![]] { + let request = serde_json::from_value(json!({ + "messages": [{"role": "user", "content": "hello"}], + "logprobs": true, "top_logprobs": 1, "logprob_token_ids": ids, + })) + .unwrap(); + let prepared = prepare_chat_request( + request, + &served(&["test-model"]), + ResolvedRequestContext::default(), + ) + .unwrap(); + let params = prepared.chat_request.sampling_params; + let selected = !ids.is_empty(); + assert_eq!(params.logprob_token_ids, selected.then_some(ids)); + assert_eq!(params.logprobs, (!selected).then_some(1)); + assert_eq!( + prepared.options.output_top_logprobs, + if selected { -1 } else { 1 } + ); + } + } + #[test] fn prepare_chat_request_lowers_logprobs_fields() { let request = ChatCompletionRequest { diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 204f3bab5aab..21776dd075e0 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -171,6 +171,11 @@ pub struct ChatCompletionRequest { /// Number of prompt logprobs to return pub prompt_logprobs: Option, + /// Specific vocab token IDs to return logprobs for at each generated + /// position, in addition to the sampled token. Takes precedence over + /// `top_logprobs`. Requires `logprobs=true`. + pub logprob_token_ids: Option>, + /// Restrict output to these token IDs only pub allowed_token_ids: Option>, @@ -308,6 +313,7 @@ impl Default for ChatCompletionRequest { truncate_prompt_tokens: None, truncation_side: None, prompt_logprobs: None, + logprob_token_ids: None, allowed_token_ids: None, bad_words: None, echo: false, diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index 579b815ba104..e3701118d94b 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -35,6 +35,13 @@ pub(super) fn validate_request_compat( ); } + if request.logprob_token_ids.as_ref().is_some_and(|ids| !ids.is_empty()) && !request.logprobs { + bail_invalid_request!( + param = "logprob_token_ids", + "when using `logprob_token_ids`, `logprobs` must be set to true." + ); + } + if let Some(prompt_logprobs) = request.prompt_logprobs { if prompt_logprobs < 0 && prompt_logprobs != -1 { bail_invalid_request!( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 33fc6bbe0a07..c75cd8da1db3 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -100,6 +100,12 @@ pub(super) fn prepare_completion_request( request.max_tokens }; let echo = completion_echo_text(&request, tokenizer)?; + let logprob_token_ids = request.logprob_token_ids.filter(|ids| !ids.is_empty()); + let logprobs = if logprob_token_ids.is_some() { + None + } else { + logprobs + }; let structured_outputs = convert_from_response_format_value(&request.response_format, &request.structured_outputs)?; @@ -137,7 +143,7 @@ pub(super) fn prepare_completion_request( allowed_token_ids: request.allowed_token_ids, bad_words: None, bad_words_token_ids: None, - logprob_token_ids: None, + logprob_token_ids, structured_outputs, skip_reading_prefix_cache: None, vllm_xargs: merge_kv_transfer_params( @@ -653,6 +659,28 @@ mod tests { ); } + #[test] + fn selected_logprobs_override_completion_top_k_unless_empty() { + for ids in [vec![10, 20], vec![]] { + let request = serde_json::from_value(json!({ + "prompt": "hello", "logprobs": 0, "logprob_token_ids": ids, + })) + .unwrap(); + let prepared = prepare_completion_request( + request, + &served(&["test-model"]), + ResolvedRequestContext::default(), + &test_tokenizer(), + ) + .unwrap(); + let params = prepared.text_request.sampling_params; + let selected = !ids.is_empty(); + assert_eq!(params.logprob_token_ids, selected.then_some(ids)); + assert_eq!(params.logprobs, (!selected).then_some(0)); + assert_eq!(prepared.options.requested_logprobs, Some(0)); + } + } + #[test] fn prepare_completion_request_accepts_logprobs_fields() { let request: CompletionRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index b10f8a23d701..9f7027511cbf 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -158,6 +158,11 @@ pub struct CompletionRequest { /// Number of prompt logprobs to return pub prompt_logprobs: Option, + /// Specific vocab token IDs to return logprobs for at each generated + /// position, in addition to the sampled token. Takes precedence over + /// the top-k selected by `logprobs`. Requires `logprobs` to be set. + pub logprob_token_ids: Option>, + // -------- Extra vLLM Parameters -------- /// Whether to add special tokens (e.g. BOS) to the prompt #[serde(default = "default_true")] diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index e41ad7a85d4e..96090cee935e 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -67,6 +67,22 @@ pub(super) fn validate_request_compat( } } + if request.logprob_token_ids.as_ref().is_some_and(|ids| !ids.is_empty()) { + if request.logprobs.is_none() { + bail_invalid_request!( + param = "logprob_token_ids", + "when using `logprob_token_ids`, `logprobs` must be set." + ); + } + if request.echo && request.max_tokens == Some(0) { + bail_invalid_request!( + param = "logprob_token_ids", + "`logprob_token_ids` is not supported when `echo=True` and \ + `max_tokens=0` because no output tokens are generated." + ); + } + } + if request.use_beam_search { bail_invalid_request!( param = "use_beam_search", diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs index 1de2cdc10015..3f8b4d9fa749 100644 --- a/rust/src/server/src/routes/openai/utils/logprobs.rs +++ b/rust/src/server/src/routes/openai/utils/logprobs.rs @@ -284,7 +284,9 @@ fn chat_top_logprob_entries( usize::try_from(top_logprobs).unwrap_or(0) }; - position.entries.iter().take(limit) + // The engine row is the sampled token followed by the top-k or requested + // tokens, so the sampled token may repeat. Keep each token's first entry. + position.entries.iter().unique_by(|entry| entry.token_id).take(limit) } fn token_bytes(token: &str) -> Vec { @@ -370,6 +372,26 @@ mod tests { assert_eq!(chat_top_logprobs_len(-1), 3); } + #[test] + fn chat_logprobs_dedupes_repeated_tokens_before_top_logprobs_cut() { + // Sampled `A` is also the top-1 token: the engine row is `[A, A, B, C]`. + let mut logprobs = sample_logprobs(); + let sampled = logprobs.positions[0].entries[0].clone(); + logprobs.positions[0].entries.insert(0, sampled); + + let top_tokens = |top_logprobs| { + let chat_logprobs = + decoded_logprobs_to_openai_chat(&logprobs, top_logprobs, false).unwrap(); + chat_logprobs.content.unwrap()[0] + .top_logprobs + .iter() + .map(|entry| entry.token.clone()) + .collect::>() + }; + assert_eq!(top_tokens(2), ["A", "B"]); + assert_eq!(top_tokens(-1), ["A", "B", "C"]); + } + #[test] fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() { let mut logprobs = sample_logprobs(); diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 9b8378f89934..c6c05e63b1c8 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -3012,6 +3012,177 @@ async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt() ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn selected_logprobs_invalid_requests_return_client_errors() { + let (mut app, engine_task) = test_app_with_engine_handle().await; + for chat_route in [true, false] { + let base = if chat_route { + json!({"messages": [{"role": "user", "content": "hello"}], "logprobs": true}) + } else { + json!({"prompt": "hello", "logprobs": 0}) + }; + for invalid in [ + json!({"logprob_token_ids": [-1]}), + json!({"logprob_token_ids": vec![1; 129]}), + if chat_route { + json!({"logprob_token_ids": [1], "logprobs": false}) + } else { + json!({"logprob_token_ids": [1], "logprobs": null}) + }, + json!({"logprob_token_ids": [1], "use_beam_search": true}), + ] { + let mut body = base.clone(); + body.as_object_mut().unwrap().extend(invalid.as_object().unwrap().clone()); + let path = if chat_route { + "/v1/chat/completions" + } else { + "/v1/completions" + }; + let (status, error) = post_json(&mut app, path, body.clone()).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}: {error}"); + } + } + let (status, error) = post_json( + &mut app, + "/v1/completions", + json!({"prompt": "hello", "logprobs": 0, + "logprob_token_ids": [1], "echo": true, "max_tokens": 0}), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{error}"); + engine_task.abort_and_join().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn selected_logprobs_reach_engine_and_both_openai_response_modes() { + for chat_route in [true, false] { + for stream in [true, false] { + let ipc = IpcNamespace::new().unwrap(); + let address = ipc.handshake_endpoint(); + let engine_task = MockEngineTask::new(spawn_mock_engine_task( + address.clone(), + b"selected-logprobs".to_vec(), + |dealer, push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + let params = request.sampling_params.as_ref().unwrap(); + assert_eq!(params.logprob_token_ids, Some(vec![65, 104, 66])); + assert_eq!(params.logprobs, None); + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output_with_logprobs( + &request.request_id, + vec![104], + Some(EngineCoreFinishReason::Stop), + None, + Some(Logprobs { + positions: vec![PositionLogprobs { + entries: vec![ + TokenLogprob { + token_id: 104, + logprob: -0.5, + rank: 1, + }, + TokenLogprob { + token_id: 65, + logprob: -4.0, + rank: 20, + }, + TokenLogprob { + token_id: 104, + logprob: -0.5, + rank: 1, + }, + TokenLogprob { + token_id: 66, + logprob: -5.0, + rank: 30, + }, + ], + }], + }), + None, + )], + ..Default::default() + } + .into(), + ) + .await; + }) + }, + )); + let client = EngineCoreClient::connect( + EngineCoreClientConfig::new_single(address).with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ), + ) + .await + .unwrap(); + let chat = + ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); + let mut app = build_router(Arc::new(AppState::new(vec!["test-model".into()], chat))); + let mut body = if chat_route { + json!({"messages": [{"role": "user", "content": "hello"}], "logprobs": true, "top_logprobs": 1}) + } else { + json!({"prompt": "hello", "logprobs": 0}) + }; + body["stream"] = json!(stream); + // The sampled token `h` (104) is also requested. + body["logprob_token_ids"] = json!([65, 104, 66]); + let response = app + .call( + Request::builder() + .method("POST") + .uri(if chat_route { + "/v1/chat/completions" + } else { + "/v1/completions" + }) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + engine_task.await.unwrap(); + let payloads = if stream { + sse_json_payloads(std::str::from_utf8(&bytes).unwrap()) + } else { + vec![serde_json::from_slice(&bytes).unwrap()] + }; + let logprobs = payloads + .iter() + .map(|payload| &payload["choices"][0]["logprobs"]) + .find(|value| value.is_object()) + .unwrap(); + if chat_route { + let actual = logprobs["content"][0]["top_logprobs"] + .as_array() + .unwrap() + .iter() + .map(|entry| json!([entry["token"], entry["logprob"]])) + .collect::>(); + assert_eq!( + actual, + [json!(["h", -0.5]), json!(["A", -4.0]), json!(["B", -5.0])] + ); + } else { + assert_eq!( + logprobs["top_logprobs"][0], + json!({"h": -0.5, "A": -4.0, "B": -5.0}) + ); + } + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() { From 3b9e46277407ed53317e78114ad50f1f9bd9d6a4 Mon Sep 17 00:00:00 2001 From: 0xsensei Date: Mon, 12 Oct 2026 16:42:35 +0530 Subject: [PATCH 402/952] [Bugfix][Rust Frontend] Correct DP routing after request removal (#59631) Signed-off-by: 0xsensei Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../engine-core-client/src/client/state.rs | 109 ++++++++++++++++-- 1 file changed, 101 insertions(+), 8 deletions(-) diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 9f6f8bbf5c2b..3d4844f3add8 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -26,6 +26,7 @@ struct TrackedRequest { sender: OutputSender, engine_id: EngineId, lora: Option, + admission_stats_generation: Option, } /// Frontend-side view of one LoRA request's scheduling phase. @@ -73,6 +74,8 @@ struct EngineRoutingState { /// Added to the snapshot so each admission raises the routing score before /// the next snapshot arrives. Only tracked once a snapshot exists. admitted_since_stats: usize, + /// Identifies admissions made after the latest scheduler snapshot. + stats_generation: u64, } impl EngineRoutingState { @@ -90,10 +93,21 @@ impl EngineRoutingState { } /// Record one request admitted to this engine. - fn record_admission(&mut self) { + fn record_admission(&mut self) -> Option { self.inflight += 1; if self.last_scheduler_stats.is_some() { self.admitted_since_stats += 1; + return Some(self.stats_generation); + } + None + } + + /// Record one request removed from this engine. Its admission is dropped + /// from `admitted_since_stats` only if no newer snapshot has reset that count. + fn record_removal(&mut self, admission_stats_generation: Option) { + self.inflight -= 1; + if admission_stats_generation == Some(self.stats_generation) { + self.admitted_since_stats -= 1; } } @@ -102,6 +116,10 @@ impl EngineRoutingState { fn apply_scheduler_counts(&mut self, next: EngineLoadSnapshot) { self.last_scheduler_stats = Some(next); self.admitted_since_stats = 0; + self.stats_generation = self + .stats_generation + .checked_add(1) + .expect("scheduler stats generation overflow"); } } @@ -161,21 +179,22 @@ impl RequestRegistry { if lora.is_some() { self.active_lora_requests += 1; } + let state = self + .routing_per_engine + .get_mut(&engine_id) + .expect("request registry must track all known engines"); + let admission_stats_generation = state.record_admission(); + self.requests.insert( request_id, TrackedRequest { sender: tx, engine_id: engine_id.clone(), lora, + admission_stats_generation, }, ); - let state = self - .routing_per_engine - .get_mut(&engine_id) - .expect("request registry must track all known engines"); - state.record_admission(); - Ok((engine_id, rx)) } @@ -368,7 +387,7 @@ impl RequestRegistry { self.routing_per_engine .get_mut(&tracked.engine_id) .expect("request registry must track all known engines") - .inflight -= 1; + .record_removal(tracked.admission_stats_generation); Some((tracked.sender, tracked.engine_id)) } @@ -1079,6 +1098,80 @@ mod tests { expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}")); } + #[test] + fn registry_finished_requests_do_not_route_live_work_to_a_busier_engine() { + let mut registry = RequestRegistry::new( + &[0, 1].map(|rank| connected_engine(EngineId::from_engine_index(rank))), + ); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 0, + }, + ); + registry.apply_scheduler_counts( + 1, + EngineLoadSnapshot { + waiting: 0, + running: 1, + }, + ); + + let mut finished_routes = Vec::new(); + for request_id in ["finished-1", "finished-2"].map(String::from) { + let (engine_id, _) = + registry.register(request_id.clone(), None, None, NonZeroU32::MIN).unwrap(); + + finished_routes.push(engine_id.engine_index()); + drop(registry.finish_many([&request_id])); + } + + let (live_engine, _) = + registry.register("live-request".into(), None, None, NonZeroU32::MIN).unwrap(); + expect_test::expect!["([Some(0), Some(0)], Some(0))"].assert_eq(&format!( + "{:?}", + (finished_routes, live_engine.engine_index()) + )); + } + + #[test] + fn registry_removal_keeps_admissions_counted_after_a_newer_snapshot() { + let engine = EngineId::from_engine_index(0); + let mut registry = RequestRegistry::new(&[connected_engine(engine.clone())]); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 0, + }, + ); + + // `old` is admitted before the next snapshot, which already counts it. + registry.register("old".into(), None, None, NonZeroU32::MIN).unwrap(); + registry.apply_scheduler_counts( + 0, + EngineLoadSnapshot { + waiting: 0, + running: 1, + }, + ); + // `new` is admitted after that snapshot, so only it is counted on top. + registry.register("new".into(), None, None, NonZeroU32::MIN).unwrap(); + + // Removing `old` must not take back the admission counted for `new`. + drop(registry.remove("old")); + let state = ®istry.routing_per_engine[&engine]; + expect_test::expect!["(1, 1, 2)"].assert_eq(&format!( + "{:?}", + ( + state.inflight, + state.admitted_since_stats, + state.routing_score() + ) + )); + } + #[test] fn register_with_data_parallel_rank_routes_to_specified_engine() { let engine_0 = EngineId::from_engine_index(0); From f5cda48c9023e8ea3b5d29b57815c8455473e97b Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:12:39 +0800 Subject: [PATCH 403/952] [Bugfix][Rust Frontend] Export corrupted request metric (#59937) Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/llm/src/request_metrics.rs | 8 ++++++++ rust/src/llm/tests/generate.rs | 4 ++++ rust/src/metrics/src/request.rs | 9 +++++++++ 3 files changed, 21 insertions(+) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 61cea120940d..71decdd281ff 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -47,6 +47,7 @@ pub(crate) struct RequestMetricsTracker { num_generation_tokens: u32, num_preemptions: u64, latest_num_cached_tokens: u32, + is_corrupted: bool, } /// Cached request metric handles for one model and engine index. @@ -61,6 +62,7 @@ struct RequestMetricHandles { prompt_tokens_external_kv_transfer: U64Counter, prompt_tokens_cached: U64Counter, generation_tokens: U64Counter, + corrupted_requests: U64Counter, // Request lifecycle counters and histograms. request_success: Family, @@ -123,6 +125,7 @@ impl RequestMetricsTracker { num_generation_tokens: 0, num_preemptions: 0, latest_num_cached_tokens: 0, + is_corrupted: false, } } @@ -141,6 +144,7 @@ impl RequestMetricsTracker { } self.num_generation_tokens += output.new_token_ids.len() as u32; self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64); + self.is_corrupted |= output.num_nans_in_logits > 0; if let Some(events) = &output.events { self.observe_events(events); @@ -187,6 +191,9 @@ impl RequestMetricsTracker { }; self.record_request_success(finish_reason); + if self.is_corrupted { + self.handles.corrupted_requests.inc(); + } self.handles.request_prompt_tokens.observe(self.prompt_len as f64); self.handles @@ -297,6 +304,7 @@ fn resolve_request_metric_handles( ), prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels), generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels), + corrupted_requests: metrics.corrupted_requests.get_or_create_owned(&labels), request_success: metrics.request_success.clone(), request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels), request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 8a150d07a44b..0a1eacd73ff6 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -898,6 +898,7 @@ async fn generate_records_request_metrics_in_prometheus_output() { engine_index: 4, timestamp: 10.0, outputs: vec![EngineCoreOutput { + num_nans_in_logits: 2, prefill_stats: Some(Box::new(PrefillStats { num_prompt_tokens: 2, num_computed_tokens: 2, @@ -982,6 +983,9 @@ async fn generate_records_request_metrics_in_prometheus_output() { assert!(rendered.contains(&format!( "vllm:generation_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 3" ))); + assert!(rendered.contains(&format!( + "vllm:corrupted_requests_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" + ))); assert!(rendered.contains(&format!( "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index b6e7683e768b..46260cf58397 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -106,6 +106,7 @@ pub struct RequestMetrics { pub prompt_tokens_by_source: PromptTokenSourceCounterFamily, pub prompt_tokens_cached: Family, pub generation_tokens: Family, + pub corrupted_requests: Family, // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, @@ -173,6 +174,13 @@ impl RequestMetrics { generation_tokens.clone(), ); + let corrupted_requests = Family::default(); + registry.register( + "vllm:corrupted_requests", + "Corrupted requests, in terms of total number of requests with NaNs in logits.", + corrupted_requests.clone(), + ); + // Request lifecycle counters and histograms. let request_success = Family::default(); registry.register( @@ -309,6 +317,7 @@ impl RequestMetrics { prompt_tokens_by_source, prompt_tokens_cached, generation_tokens, + corrupted_requests, request_success, request_prompt_tokens, request_generation_tokens, From 32b05addc1301be904ba1b6c9a5e3704581b1fc3 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:12:43 +0800 Subject: [PATCH 404/952] [Bugfix][Rust Frontend] Restore DP wave delivery after XPUB reconnects (#60249) Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../src/coordinator/bootstrap.rs | 24 +++++--------- .../src/coordinator/inproc.rs | 33 +++++++++++++++++-- .../engine-core-client/src/tests/client.rs | 33 ++++++++++++++----- 3 files changed, 65 insertions(+), 25 deletions(-) diff --git a/rust/src/engine-core-client/src/coordinator/bootstrap.rs b/rust/src/engine-core-client/src/coordinator/bootstrap.rs index 68229085d239..17c346c933e3 100644 --- a/rust/src/engine-core-client/src/coordinator/bootstrap.rs +++ b/rust/src/engine-core-client/src/coordinator/bootstrap.rs @@ -63,22 +63,9 @@ async fn wait_for_engine_subscriptions( timeout: ready_timeout, } })??; - if message.len() != 1 { + if !is_engine_subscription(&message) { bail_unexpected_handshake_message!( - "expected 1 frame for coordinator subscription, got {}", - message.len() - ); - } - - let frame = message - .into_vec() - .into_iter() - .next() - .expect("single-frame coordinator subscription message"); - if frame.as_ref() != [0x01] { - bail_unexpected_handshake_message!( - "expected coordinator subscription frame [0x01], got {:?}", - frame.as_ref() + "expected coordinator subscription [0x01], got {message:?}" ); } received += 1; @@ -87,6 +74,13 @@ async fn wait_for_engine_subscriptions( Ok(()) } +/// Whether `message` is the single-frame subscribe-all notification (`[0x01]`) +/// that an engine's XSUB sends whenever it connects or reconnects to the +/// coordinator broadcast socket. +pub(super) fn is_engine_subscription(message: &ZmqMessage) -> bool { + message.len() == 1 && message.get(0).is_some_and(|frame| frame.as_ref() == [0x01]) +} + /// Send the coordinator READY marker to all subscribed engines. async fn send_ready_to_engines(input_socket: &mut XPubSocket) -> Result<()> { input_socket.send(ZmqMessage::from(Bytes::from_static(b"READY"))).await?; diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index bf4c23cc137a..cbd69095f85c 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -7,10 +7,11 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use thiserror_ext::AsReport; use tokio::sync::mpsc; use tracing::{debug, warn}; -use zeromq::prelude::SocketSend; -use zeromq::{XPubSocket, ZmqMessage}; +use zeromq::prelude::{SocketRecv, SocketSend}; +use zeromq::{XPubSocket, ZmqError, ZmqMessage}; use crate::client::imp::ClientInner; +use crate::coordinator::bootstrap::is_engine_subscription; use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; use crate::protocol::encode_msgpack; @@ -187,6 +188,34 @@ impl InProcCoordinatorRunner { let result: Result<()> = async { loop { tokio::select! { + // Keep receiving engine subscriptions: zmq.rs XPUB registers a + // subscription only when `recv()` yields it (unlike libzmq), so an + // engine whose XSUB reconnects gets no broadcasts until we read its + // subscribe-all. If a wave is running, replay it so that engine does + // not miss the `START_DP_WAVE` sent while it was disconnected. + subscription = self.coordinator_input.recv() => { + match subscription { + Ok(message) => { + if is_engine_subscription(&message) { + let state = *self.state.lock(); + if state.engines_running { + debug!( + current_wave = state.current_wave, + "replaying active DP wave after engine subscription" + ); + self.broadcast_start_wave(state.current_wave, None).await?; + } + } + } + Err(error @ (ZmqError::Codec(_) | ZmqError::Network(_))) => { + warn!( + error = %error, + "coordinator subscription peer disconnected" + ); + } + Err(error) => return Err(error.into()), + } + } // Received frontend-originated command from the handle. command = self.command_rx.recv() => { let Some(command) = command else { diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 43eb88db79ed..7b7a880bde6d 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -257,7 +257,7 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.unwrap().into_vec() } -async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, u32) { +async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option) { let frames = sub.recv().await.unwrap().into_vec(); assert_eq!(frames.len(), 2); assert_eq!( @@ -607,12 +607,27 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let handshake_address = handshake_address.clone(); async move { let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; + let coordinator_address = engine + .init + .addresses + .coordinator_input + .clone() + .expect("coordinator input address should be present"); let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); + + let mut late_subscriber = SubSocket::new(); + late_subscriber.connect(&coordinator_address).await.unwrap(); + late_subscriber.subscribe("").await.unwrap(); + + let replay = recv_start_dp_wave(&mut late_subscriber).await; + assert_eq!(replay, (0, None)); + let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!(replay, (0, None)); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -656,7 +671,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, 0)); + assert_eq!((wave, exclude_engine), (1, Some(0))); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -693,7 +708,9 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); + let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!(replay, (0, None)); let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); @@ -727,7 +744,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, 0)); + assert_eq!((wave, exclude_engine), (1, Some(0))); assert!( timeout( @@ -812,7 +829,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { engine.coordinator.take().expect("coordinator sockets should be present"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (4, 1)); + assert_eq!((wave, exclude_engine), (4, Some(1))); let _ = shutdown0_rx.await; } @@ -838,7 +855,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { .await; let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (4, 1)); + assert_eq!((wave, exclude_engine), (4, Some(1))); let _ = shutdown1_rx.await; } @@ -880,7 +897,7 @@ async fn coordinator_accepts_stats_only_outputs() { let data_socket = engine.data_sockets.first_mut().expect("data socket"); let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, 0)); + assert_eq!((wave, exclude_engine), (0, Some(0))); send_outputs( &mut coordinator.output_push, From 8d7b93fec2eb72ce01f2217472ac1f51b9bce600 Mon Sep 17 00:00:00 2001 From: Shimi Bandiel Date: Mon, 12 Oct 2026 04:37:53 -0700 Subject: [PATCH 405/952] [Frontend] Report stop_reason on /inference/v1/generate and pass it through derender (#60251) Signed-off-by: Hyeonki Hong Signed-off-by: Shimi Bandiel Signed-off-by: Bugen Zhao Co-authored-by: Hyeonki Hong Co-authored-by: Mimir Co-authored-by: Claude Co-authored-by: Bugen Zhao --- .../server/src/routes/inference/generate.rs | 72 ++++++++++++++++++- .../src/routes/inference/generate/types.rs | 4 +- .../scale_out/derender/test_derender.py | 32 +++++++++ .../derender/test_derender_stream.py | 15 ++++ .../test_generate_stream.py | 67 +++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 6 ++ .../scale_out/token_in_token_out/serving.py | 2 + vllm/renderers/online_derenderer.py | 5 ++ 8 files changed, 201 insertions(+), 2 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index a30cd0c85ab4..930c698c30e4 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -226,6 +226,8 @@ async fn generate_chunk_stream( None }; + let stop_reason = + finish_reason.as_ref().and_then(|reason| reason.as_stop_reason().cloned()); let prompt_token_ids = prompt_token_ids.take(); y.yield_ok(GenerateStreamResponse { request_id: request_id.clone(), @@ -233,6 +235,7 @@ async fn generate_chunk_stream( index: 0, logprobs, finish_reason: finish_reason.map(|reason| reason.as_str().to_string()), + stop_reason, token_ids, sampling_mask: output.sampling_mask.map(|mask| mask.rows), weight_versions: output.weight_versions, @@ -319,6 +322,7 @@ fn collect_generate( None }; let finish_reason = collected.finish_reason.as_str().to_string(); + let stop_reason = collected.finish_reason.as_stop_reason().cloned(); if enable_log_requests { info!( @@ -335,6 +339,7 @@ fn collect_generate( index: 0, logprobs, finish_reason: Some(finish_reason), + stop_reason, token_ids: collected.token_ids, sampling_mask: collected.sampling_mask.map(|mask| mask.rows), weight_versions: collected.weight_versions, @@ -537,7 +542,7 @@ mod tests { use futures::{TryStreamExt as _, stream}; use vllm_engine_core_client::protocol::multimodal::{MmModality, PlaceholderRange}; - use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; + use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason}; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_llm::GeneratePromptInfo; @@ -1125,4 +1130,69 @@ mod tests { ) .expect_err("multi-token prompt without payload is an engine failure"); } + + #[tokio::test] + async fn generate_chunk_stream_reports_stop_reason() { + let chunks = collect_chunks( + vec![ + stream_output(Some(&[11, 22]), vec![33], None), + stream_output( + Some(&[11, 22]), + vec![44], + Some(FinishReason::Stop(Some(StopReason::Text( + "wor".to_string(), + )))), + ), + ], + false, + None, + ) + .await; + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].choices[0].stop_reason, None); + assert_eq!(chunks[1].choices[0].finish_reason.as_deref(), Some("stop")); + assert_eq!( + chunks[1].choices[0].stop_reason, + Some(StopReason::Text("wor".to_string())) + ); + } + + #[test] + fn collect_generate_reports_stop_token_id() { + let output = CollectedGenerateOutput { + request_id: "raw-stop".to_string(), + prompt_logprobs: None, + prompt_token_id_logprobs: None, + token_ids: vec![30, 7], + logprobs: None, + finish_reason: FinishReason::Stop(Some(StopReason::TokenId(7))), + usage: TokenUsage { + prompt_token_count: 2, + output_token_count: 2, + cached_token_count: 0, + }, + kv_transfer_params: None, + ec_transfer_params: None, + prompt_token_ids: vec![10, 20], + sampling_mask: None, + spec_decode_metrics: None, + weight_versions: None, + }; + + let response = collect_generate( + output, + "raw-stop".to_string(), + ApiServerOptions::default(), + ResponseOptions::default(), + None, + ) + .expect("response"); + + assert_eq!(response.choices[0].finish_reason.as_deref(), Some("stop")); + assert_eq!( + response.choices[0].stop_reason, + Some(StopReason::TokenId(7)) + ); + } } diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 6352b187250a..d2737b5bce69 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -7,7 +7,7 @@ use llm_multimodal::MediaContentPart; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; -use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics; +use vllm_engine_core_client::protocol::output::{RequestSpecDecodeMetrics, StopReason}; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::WeightVersionSpan; use vllm_text::SamplingParams; @@ -72,6 +72,7 @@ pub(super) struct GenerateResponseChoice { pub index: u32, pub logprobs: Option, pub finish_reason: Option, + pub stop_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, pub weight_versions: Option>, @@ -84,6 +85,7 @@ pub(super) struct GenerateResponseStreamChoice { pub index: u32, pub logprobs: Option, pub finish_reason: Option, + pub stop_reason: Option, pub token_ids: Vec, pub sampling_mask: Option>>, /// Set on the final chunk only; covers the whole output. diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index d29aa3f53f73..30d64ed39bfd 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -112,6 +112,38 @@ async def test_derender_chat_roundtrip(client): assert data["choices"][0]["message"]["role"] == "assistant" +@pytest.mark.asyncio +async def test_derender_chat_passes_stop_reason_through(client): + """The coupled server reports the stop string or token id that ended + generation; derender must carry the generate value through unchanged.""" + gen_req = await _render_chat(client) + generate_response = _make_generate_response(gen_req["token_ids"][:3]) + generate_response["choices"][0]["stop_reason"] = "\n\n" + + response = await client.post( + "/v1/chat/completions/derender", + json={"model": MODEL_NAME, "generate_response": generate_response}, + ) + assert response.status_code == 200 + assert response.json()["choices"][0]["stop_reason"] == "\n\n" + + +@pytest.mark.asyncio +async def test_derender_completion_passes_stop_reason_through(client): + gen_req = await _render_completion(client, "Hello world") + generate_response = _make_completion_generate_response( + gen_req["token_ids"][:3], "gen-stop" + ) + generate_response["choices"][0]["stop_reason"] = 151645 + + response = await client.post( + "/v1/completions/derender", + json={"model": MODEL_NAME, "generate_responses": [generate_response]}, + ) + assert response.status_code == 200 + assert response.json()["choices"][0]["stop_reason"] == 151645 + + @pytest.mark.asyncio async def test_derender_chat_usage(client): """Supplied prompt_tokens flows through into usage correctly.""" diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index d97a2f992ac2..55b0789843ee 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -195,6 +195,7 @@ def _make_stream_chunk( metrics: PerRequestMetrics | None = None, logprobs: dict | None = None, prompt_token_ids: list[int] | None = None, + stop_reason: int | str | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" return GenerateTokensStreamResponse( @@ -204,6 +205,7 @@ def _make_stream_chunk( index=index, token_ids=token_ids, finish_reason=finish_reason, + stop_reason=stop_reason, logprobs=logprobs, ) ], @@ -886,6 +888,19 @@ async def test_finish_reason_forwarded(self, derenderer, tokenizer): ) assert chunk.choices[0].finish_reason == "length" + @pytest.mark.asyncio + async def test_stop_reason_forwarded(self, derenderer, tokenizer): + """stop_reason from the final generate chunk reaches the derendered + choice, matching what the coupled server reports.""" + token_ids = tokenizer.encode("done")[:2] + chunk, _ = await derenderer.derender_completion_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + token_ids, finish_reason="stop", stop_reason="\n\n" + ), + ) + assert chunk.choices[0].stop_reason == "\n\n" + class TestStreamLogprobs: """Streaming derender must carry per-chunk logprobs with token ids diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 7f0cd1286041..c30732f3150f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -139,6 +139,7 @@ def _make_request_output( request_id: str, token_ids: list[int], finish_reason: str | None = None, + stop_reason: int | str | None = None, finished: bool = False, prompt_token_ids: list[int] | None = None, logprobs: list[dict[int, Any] | None] | None = None, @@ -160,6 +161,7 @@ def _make_request_output( cumulative_logprob=None, logprobs=logprobs, finish_reason=finish_reason, + stop_reason=stop_reason, spec_decode_metrics=spec_decode_metrics, ) ], @@ -378,6 +380,71 @@ async def mock_generate(*args, **kwargs): assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.asyncio +async def test_non_stream_reports_engine_stop_reason(): + """The matched stop token id from the engine lands on the choice.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10, 42], + finish_reason="stop", + stop_reason=42, + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=False, + ) + + response = await serving.serve_tokens(request) + choice = response.model_dump()["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["stop_reason"] == 42 + + +@pytest.mark.asyncio +async def test_stream_reports_engine_stop_reason_on_final_chunk(): + """A matched stop string is reported with finish_reason and null before.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", + token_ids=[20], + finish_reason="stop", + stop_reason="done", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + data_chunks = [c for c in parsed if c != "[DONE]"] + assert len(data_chunks) == 2 + + assert data_chunks[0]["choices"][0]["stop_reason"] is None + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + assert data_chunks[1]["choices"][0]["stop_reason"] == "done" + + @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) async def test_serve_tokens_returns_spec_decode_metrics(stream: bool): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index bccc8d8eff0b..c8d82eeb5a23 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -400,6 +400,9 @@ class GenerateChoiceBase(BaseModel): index: int # per OpenAI spec this is the default finish_reason: str | None = "stop" + # Stop string or token id that ended generation (repetition detection + # sets it too); None on EOS, length and abort. + stop_reason: int | str | None = None token_ids: list[int] | None = None # Per-token expert routing decisions, base64-encoded `.npy` bytes # (numpy serialization). Shape after decode: @@ -443,6 +446,9 @@ class GenerateStreamChoiceBase(BaseModel): index: int finish_reason: str | None = None + # Stop string or token id that ended generation (repetition detection + # sets it too); None on EOS, length and abort. + stop_reason: int | str | None = None token_ids: list[int] | None = None routed_experts: str | None = None sampling_mask: list[list[int]] | None = None diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 085ae97e16a7..f09ee6fad5ab 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -450,6 +450,7 @@ async def serve_tokens_full_generator( index=output.index, logprobs=logprobs, finish_reason=output.finish_reason if output.finish_reason else "stop", + stop_reason=output.stop_reason, token_ids=as_list(output.token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, @@ -623,6 +624,7 @@ async def serve_tokens_stream_generator( index=i, logprobs=logprobs, finish_reason=finish_reason, + stop_reason=output.stop_reason, token_ids=as_list(delta_token_ids), routed_experts=routed_experts_b64, sampling_mask=sampling_mask, diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 3fe50f16c391..345e5443ea72 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -206,6 +206,7 @@ def _derender_chat( message=message, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -484,6 +485,7 @@ async def derender_chat_stream( delta=delta, logprobs=resolved_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) ) @@ -672,6 +674,7 @@ def _replay(token_ids: list[int], chunk_lens: list[int]) -> None: index=choice.index, delta=delta_message, finish_reason=finish_reason, + stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, ) stream_choices.append( @@ -777,6 +780,7 @@ def _derender_completion( index=index, text=decoded_text, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, logprobs=completion_logprobs, ) ) @@ -878,6 +882,7 @@ async def derender_completion_stream( text=new_text, logprobs=completion_logprobs, finish_reason=choice.finish_reason, + stop_reason=choice.stop_reason, ) ) From aa1d44df16fe554405c6064cad2a904a16b63581 Mon Sep 17 00:00:00 2001 From: Liana Koleva <43767763+lianakoleva@users.noreply.github.com> Date: Mon, 12 Oct 2026 05:20:53 -0700 Subject: [PATCH 406/952] [Bugfix][Rust Frontend] Alias `qwen3_xml` to `qwen3_coder` and add `qwen3_json` for the Qwen JSON format (#60136) Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 2 +- rust/src/chat/src/output/default/mod.rs | 2 +- rust/src/chat/src/parser/tool/mod.rs | 14 ++-- rust/src/chat/src/parser/tool/tests.rs | 31 +++++++- rust/src/chat/tests/roundtrip.rs | 2 +- rust/src/parser/Cargo.toml | 2 +- .../benches/{qwen3_xml.rs => qwen3_json.rs} | 14 ++-- rust/src/parser/src/tool/json/mod.rs | 2 +- rust/src/parser/src/tool/json/qwen.rs | 77 ++++++++++--------- rust/src/parser/src/tool/mod.rs | 2 +- rust/src/parser/src/unified/combined.rs | 16 ++-- vllm/tool_parsers/__init__.py | 7 ++ 12 files changed, 103 insertions(+), 68 deletions(-) rename rust/src/parser/benches/{qwen3_xml.rs => qwen3_json.rs} (89%) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index f437745e3967..0b7e89914ce0 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -431,7 +431,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 981173dc7450..e95ebd8b0dcf 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -498,7 +498,7 @@ mod tests { Err(error) => error, }; - expect_test::expect!["unified parsing requires the tool and reasoning selections to resolve to the same parser; resolved tool=qwen3_xml, reasoning=gemma4"] + expect_test::expect!["unified parsing requires the tool and reasoning selections to resolve to the same parser; resolved tool=qwen3_json, reasoning=gemma4"] .assert_eq(&format!("{error}")); } diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 9f3ab4f516a3..9360b62e2f1e 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -10,7 +10,7 @@ pub use vllm_parser::tool::{ DeepSeekV41ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MiMoToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, SeedOssToolParser, ToolParser, ToolParserError, + Qwen3CoderToolParser, Qwen3JsonToolParser, SeedOssToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -39,6 +39,9 @@ pub mod names { pub const MISTRAL: &str = "mistral"; pub const PHI4_MINI_JSON: &str = "phi4_mini_json"; pub const QWEN3_CODER: &str = "qwen3_coder"; + // The ``-wrapped JSON format of Qwen2.5, QwQ and Qwen3. + pub const QWEN3_JSON: &str = "qwen3_json"; + // An alias of `qwen3_coder` for the same XML format. pub const QWEN3_XML: &str = "qwen3_xml"; pub const SEED_OSS: &str = "seed_oss"; } @@ -82,8 +85,9 @@ impl ToolParserFactory { .register_parser::(names::MINIMAX_M3) .register_parser::(names::MISTRAL) .register_parser::(names::PHI4_MINI_JSON) - .register_parser::(names::QWEN3_XML) .register_parser::(names::QWEN3_CODER) + .register_parser::(names::QWEN3_XML) + .register_parser::(names::QWEN3_JSON) .register_parser::(names::SEED_OSS); factory @@ -92,9 +96,9 @@ impl ToolParserFactory { .register_pattern("mimo-v2", names::MIMO) .register_pattern("qwen3-coder", names::QWEN3_CODER) .register_pattern("qwen3.5", names::QWEN3_CODER) - .register_pattern("qwq", names::HERMES) - .register_pattern("qwen2.5", names::HERMES) - .register_pattern("qwen3", names::QWEN3_XML) + .register_pattern("qwq", names::QWEN3_JSON) + .register_pattern("qwen2.5", names::QWEN3_JSON) + .register_pattern("qwen3", names::QWEN3_JSON) .register_pattern("hermes", names::HERMES) // Narrow to `internlm2` substring so it matches `internlm2-chat-7b` // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1, diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index 69022e550d59..d7eb6c2dac52 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -86,10 +86,10 @@ fn factory_distinguishes_qwen_tool_formats() { for (model, expected) in [ ("Qwen/Qwen3.5-0.8B", Some(names::QWEN3_CODER)), - ("Qwen/Qwen3-0.6B", Some(names::QWEN3_XML)), + ("Qwen/Qwen3-0.6B", Some(names::QWEN3_JSON)), ("Qwen/Qwen3-Coder-30B", Some(names::QWEN3_CODER)), - ("Qwen/QwQ-32B", Some(names::HERMES)), - ("Qwen/Qwen2.5-0.5B-Instruct", Some(names::HERMES)), + ("Qwen/QwQ-32B", Some(names::QWEN3_JSON)), + ("Qwen/Qwen2.5-0.5B-Instruct", Some(names::QWEN3_JSON)), ("Qwen/Qwen2-1.5B-Instruct", None), ] { assert_eq!(factory.resolve_name_for_model(model), expected, "{model}"); @@ -97,7 +97,7 @@ fn factory_distinguishes_qwen_tool_formats() { } #[test] -fn factory_parses_qwen2_5_coder_template_with_hermes() { +fn factory_parses_qwen2_5_coder_template_with_qwen3_json() { let factory = ToolParserFactory::new(); let tool_call = r#" {"name":"get_weather","arguments":{"location":"Tokyo"}} @@ -115,6 +115,29 @@ fn factory_parses_qwen2_5_coder_template_with_hermes() { assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } +#[test] +fn factory_parses_qwen3_coder_format_with_qwen3_xml() { + let factory = ToolParserFactory::new(); + let tool_call = " + + +Tokyo + + +"; + + let mut parser = factory.create(names::QWEN3_XML, &[]).unwrap(); + let mut output = ToolParserOutput::default(); + parser.parse_into(tool_call, &mut output).unwrap(); + output.append(parser.finish().unwrap()); + let output = output.coalesce(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); +} + #[test] fn factory_new_resolves_default_patterns() { let factory = ToolParserFactory::new(); diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index c5080a59c5c2..41f06d773432 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -204,7 +204,7 @@ impl ToolCallMixVariant { } impl RoundtripCase { - /// Qwen3 XML tool-call format with `qwen3` reasoning tags. + /// Qwen3 JSON tool-call format with `qwen3` reasoning tags. fn qwen3() -> Self { Self { model_id: "Qwen/Qwen3-0.6B", diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml index aff4c795ab63..10972ace3ee3 100644 --- a/rust/src/parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -53,7 +53,7 @@ harness = false required-features = ["test-util"] [[bench]] -name = "qwen3_xml" +name = "qwen3_json" harness = false required-features = ["test-util"] diff --git a/rust/src/parser/benches/qwen3_xml.rs b/rust/src/parser/benches/qwen3_json.rs similarity index 89% rename from rust/src/parser/benches/qwen3_xml.rs rename to rust/src/parser/benches/qwen3_json.rs index 5138007b541f..a19260ea1d2a 100644 --- a/rust/src/parser/benches/qwen3_xml.rs +++ b/rust/src/parser/benches/qwen3_json.rs @@ -6,7 +6,7 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::QwenParser as ExternalQwenParser; use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; -use vllm_parser::tool::{Qwen3XmlToolParser, Tool, ToolParser}; +use vllm_parser::tool::{Qwen3JsonToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; @@ -27,12 +27,12 @@ fn mixed_fixture() -> String { } fn long_normal_text_fixture() -> String { - let line = "This is ordinary assistant text with no Qwen XML tool markers at all.\n"; + let line = "This is ordinary assistant text with no Qwen3 JSON tool markers at all.\n"; line.repeat(LONG_NORMAL_TEXT_REPEATS) } fn native_parser(tools: &[Tool]) -> Box { - Qwen3XmlToolParser::create(tools).expect("Qwen XML parser should initialize") + Qwen3JsonToolParser::create(tools).expect("Qwen3 JSON parser should initialize") } fn run_stream_group( @@ -98,14 +98,14 @@ fn run_stream_group( group.finish(); } -fn bench_qwen3_xml(c: &mut Criterion) { +fn bench_qwen3_json(c: &mut Criterion) { let tools = test_tools(); let mixed_text = mixed_fixture(); let long_normal_text = long_normal_text_fixture(); run_stream_group( c, - "qwen3_xml/mixed_text_tool_call", + "qwen3_json/mixed_text_tool_call", &tools, &mixed_text, CHUNK_CHARS, @@ -115,7 +115,7 @@ fn bench_qwen3_xml(c: &mut Criterion) { run_stream_group( c, - "qwen3_xml/long_normal_text", + "qwen3_json/long_normal_text", &tools, &long_normal_text, CHUNK_CHARS, @@ -124,5 +124,5 @@ fn bench_qwen3_xml(c: &mut Criterion) { ); } -criterion_group!(benches, bench_qwen3_xml); +criterion_group!(benches, bench_qwen3_json); criterion_main!(benches); diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index 487be423aef2..5112ecf5d066 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -9,7 +9,7 @@ pub use internlm2::Internlm2ToolParser; pub use llama::Llama3JsonToolParser; pub use mistral::MistralToolParser; pub use phi4mini::Phi4MiniJsonToolParser; -pub use qwen::Qwen3XmlToolParser; +pub use qwen::Qwen3JsonToolParser; mod granite4; mod hermes; diff --git a/rust/src/parser/src/tool/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs index c4db8c2fbc8b..88325d3d9112 100644 --- a/rust/src/parser/src/tool/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -4,8 +4,8 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput}; -const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { - parser_name: "Qwen XML", +const QWEN3_JSON_CONFIG: JsonToolCallConfig = JsonToolCallConfig { + parser_name: "Qwen3 JSON", start_marker: "", framed_start_marker: Some("\n"), end_marker: "", @@ -15,7 +15,8 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { arguments_key: &["arguments"], }; -/// Tool parser for Qwen XML-wrapped JSON tool calls. +/// Tool parser for Qwen2.5, QwQ and Qwen3 JSON tool calls wrapped in +/// `` tags. /// /// Example tool call content: /// @@ -30,20 +31,20 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { /// /// Note: parallel calls are represented as repeated /// `...` blocks, not as multiple calls inside one tag. -pub struct Qwen3XmlToolParser { +pub struct Qwen3JsonToolParser { inner: JsonToolCallParser, } -impl Qwen3XmlToolParser { - /// Create a Qwen XML tool parser. +impl Qwen3JsonToolParser { + /// Create a Qwen3 JSON tool parser. fn new(_tools: &[Tool]) -> Self { Self { - inner: JsonToolCallParser::new(QWEN_XML_CONFIG), + inner: JsonToolCallParser::new(QWEN3_JSON_CONFIG), } } } -impl ToolParser for Qwen3XmlToolParser { +impl ToolParser for Qwen3JsonToolParser { fn create(tools: &[Tool]) -> Result> where Self: Sized + 'static, @@ -73,7 +74,7 @@ mod tests { use expect_test::expect; use thiserror_ext::AsReport; - use super::Qwen3XmlToolParser; + use super::Qwen3JsonToolParser; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::tests::assert_tool_framing_preserves_body_whitespace; use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; @@ -85,8 +86,8 @@ mod tests { } #[test] - fn qwen_xml_parse_complete_without_tool_call_keeps_text() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_parse_complete_without_tool_call_keeps_text() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); assert_eq!(output.normal_text(), "Hello, world!"); @@ -94,8 +95,8 @@ mod tests { } #[test] - fn qwen_xml_parse_complete_extracts_raw_json_arguments() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_parse_complete_extracts_raw_json_arguments() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{ "location": "Tokyo", "days": "3" }"#; let output = parser .parse_complete(&format!( @@ -112,8 +113,8 @@ mod tests { } #[test] - fn qwen_xml_does_not_validate_or_normalize_arguments() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_does_not_validate_or_normalize_arguments() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{"location":"Tokyo",}"#; let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap(); @@ -121,8 +122,8 @@ mod tests { } #[test] - fn qwen_xml_streaming_emits_argument_deltas() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_streaming_emits_argument_deltas() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let chunks = [ "", "\n{\"name\": \"get_weather\", \"arguments\": ", @@ -154,13 +155,13 @@ mod tests { } #[test] - fn qwen_xml_streaming_handles_split_markers() { + fn qwen3_json_streaming_handles_split_markers() { let input = format!( "hello {}", build_tool_call("get_weather", r#"{"location":"Tokyo"}"#) ); let chunks = split_by_chars(&input, 5); - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); @@ -170,8 +171,8 @@ mod tests { } #[test] - fn qwen_xml_keeps_end_marker_literal_inside_json_string() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_keeps_end_marker_literal_inside_json_string() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let arguments = r#"{"text":"literal inside"}"#; let output = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap(); @@ -180,8 +181,8 @@ mod tests { } #[test] - fn qwen_xml_decodes_escaped_function_name() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_decodes_escaped_function_name() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = parser .parse_complete( r#" @@ -194,8 +195,8 @@ mod tests { } #[test] - fn qwen_xml_requires_newline_after_tool_call_start() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_requires_newline_after_tool_call_start() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let input = r#"{"name":"get_weather","arguments":{}} "#; @@ -206,12 +207,12 @@ mod tests { } #[test] - fn qwen_xml_streams_start_marker_without_newline_as_text() { + fn qwen3_json_streams_start_marker_without_newline_as_text() { for chunk in [ "{\"name\":\"get_weather\"", "hi\n{\"name\":\"get_weather\"", ] { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let mut output = ToolParserOutput::default(); parser.parse_into(chunk, &mut output).unwrap(); @@ -223,8 +224,8 @@ mod tests { } #[test] - fn qwen_xml_requires_newline_before_tool_call_end() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_requires_newline_before_tool_call_end() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let error = parser .parse_complete( r#" @@ -236,14 +237,14 @@ mod tests { } #[test] - fn qwen_xml_streaming_extracts_multiple_tool_calls() { + fn qwen3_json_streaming_extracts_multiple_tool_calls() { let input = format!( "{}{}", build_tool_call("get_weather", r#"{"location":"Shanghai"}"#), build_tool_call("add", r#"{"x":1,"y":2}"#), ); let chunks = split_by_chars(&input, 7); - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); @@ -275,8 +276,8 @@ mod tests { } #[test] - fn qwen_xml_finish_fails_incomplete_tool_call() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_finish_fails_incomplete_tool_call() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); parser .parse_chunk( r#" @@ -286,13 +287,13 @@ mod tests { let error = parser.finish().unwrap_err(); - expect!["tool parser parsing failed: incomplete Qwen XML tool call"] + expect!["tool parser parsing failed: incomplete Qwen3 JSON tool call"] .assert_eq(&error.to_report_string()); } #[test] - fn qwen_xml_malformed_field_order_fails_fast() { - let mut parser = Qwen3XmlToolParser::new(&test_tools()); + fn qwen3_json_malformed_field_order_fails_fast() { + let mut parser = Qwen3JsonToolParser::new(&test_tools()); let error = parser .parse_chunk( r#" @@ -302,14 +303,14 @@ mod tests { .unwrap_err(); expect![[r#" - tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n": invalid Qwen XML + tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n": invalid Qwen3 JSON expected `name`"#]] .assert_eq(&error.to_report_string()); } #[test] fn tool_framing_preserves_body_whitespace_across_chunk_boundaries() { - assert_tool_framing_preserves_body_whitespace::( + assert_tool_framing_preserves_body_whitespace::( "\n", "\n{\"name\":\"get_weather\",\"arguments\":{}}\n", ); diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index 974d6c758a27..69ecee1bffff 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -29,7 +29,7 @@ pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; pub(crate) use hy::{HyDialect, HyToolMarkers, HyToolParser}; pub use json::{ Granite4ToolParser, HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, - MistralToolParser, Phi4MiniJsonToolParser, Qwen3XmlToolParser, + MistralToolParser, Phi4MiniJsonToolParser, Qwen3JsonToolParser, }; pub use kimi_k2::KimiK2ToolParser; pub use mimo::MiMoToolParser; diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs index b1f36e885546..aaab4312569c 100644 --- a/rust/src/parser/src/unified/combined.rs +++ b/rust/src/parser/src/unified/combined.rs @@ -186,7 +186,7 @@ mod tests { use crate::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, Glm47MoeToolParser, KimiK2ToolParser, MinimaxM2ToolParser, Qwen3CoderToolParser, - Qwen3XmlToolParser, Tool, ToolParser, + Qwen3JsonToolParser, Tool, ToolParser, }; use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; @@ -463,7 +463,7 @@ mod tests { #[test] fn combined_parser_emits_tool_calls_from_visible_content() { - let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap(); + let tool = Qwen3JsonToolParser::create(&test_tools()).unwrap(); let mut parser = CombinedParser::new(None, Some(tool)); let output = collect( @@ -515,7 +515,7 @@ mod tests { #[test] fn split_parser_wrappers_match_builder_full_formats() { - assert_builder_parity::("qwen3"); + assert_builder_parity::("qwen3"); assert_builder_parity::("qwen3.5"); assert_builder_parity::("deepseek-r1"); assert_builder_parity::("deepseek-v3.1"); @@ -538,7 +538,7 @@ mod tests { answer: None, }; let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); - let tool = Qwen3XmlToolParser::create(&tools).unwrap(); + let tool = Qwen3JsonToolParser::create(&tools).unwrap(); let visible = tool.build_visible_format(&ctx).unwrap().unwrap(); let mut parser = CombinedParser::new(Some(reasoning), Some(tool)); parser.initialize(&[]).unwrap(); @@ -640,7 +640,7 @@ mod tests { }; let tokenizer = Arc::new(tokenizer()); let reasoning = Qwen3ReasoningParser::create(tokenizer.clone()).unwrap(); - let tool = Qwen3XmlToolParser::create(&tools).unwrap(); + let tool = Qwen3JsonToolParser::create(&tools).unwrap(); let mut parser = CombinedParser::new(Some(reasoning), Some(tool)); parser.initialize(&[]).unwrap(); let outside = parser.build_output_grammar(&ctx).unwrap().unwrap().format; @@ -682,7 +682,7 @@ mod tests { fn qwen3_parser(tools: &[Tool], with_tool: bool) -> CombinedParser { let reasoning = Qwen3ReasoningParser::create(Arc::new(tokenizer())).unwrap(); - let tool = with_tool.then(|| Qwen3XmlToolParser::create(tools).unwrap()); + let tool = with_tool.then(|| Qwen3JsonToolParser::create(tools).unwrap()); let mut parser = CombinedParser::new(Some(reasoning), tool); parser.initialize(&[]).unwrap(); parser @@ -702,7 +702,7 @@ mod tests { ]) }; let required = ToolChoice::required(); - let required_calls = Qwen3XmlToolParser::create(&tools) + let required_calls = Qwen3JsonToolParser::create(&tools) .unwrap() .build_visible_format(&grammar_ctx(&tools, &required, None)) .unwrap() @@ -748,7 +748,7 @@ mod tests { let ctx = grammar_ctx(&tools, &auto, Some(&answer)); assert!(parser.build_output_grammar(&ctx).unwrap().is_none()); - let parser = CombinedParser::new(None, Some(Qwen3XmlToolParser::create(&tools).unwrap())); + let parser = CombinedParser::new(None, Some(Qwen3JsonToolParser::create(&tools).unwrap())); let ctx = grammar_ctx(&tools, &none, Some(&answer)); let alone = parser.build_output_grammar(&ctx).unwrap().unwrap(); assert_eq!(alone, BuiltOutputGrammar::final_output_only(answer.clone())); diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 8446aae81540..4b9cb33502e5 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -194,6 +194,13 @@ "qwen3_engine_tool_parser", "Qwen3EngineToolParser", ), + # TODO: Qwen2.5/QwQ/Qwen3 JSON tool calls. Aliased to `hermes` for now; + # a dedicated parser should treat the template's newline before and + # between `` blocks as framing, as the Rust frontend does. + "qwen3_json": ( + "hermes_tool_parser", + "Hermes2ProToolParser", + ), "qwen3_xml": ( "qwen3_engine_tool_parser", "Qwen3EngineToolParser", From 258b97ff6bc1dbbc7b76d2d5a5a87591f9de107f Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Mon, 12 Oct 2026 20:32:39 +0800 Subject: [PATCH 407/952] [Bugfix][EC Connector] Size transfers by measured encoder output widths (#57904) Signed-off-by: Tianyu Guo --- docs/features/ec_cpu_connector.md | 2 +- .../unit/cpu/scheduler/test_scheduler.py | 8 +- .../unit/cpu/test_ec_shared_region.py | 4 +- .../unit/cpu/worker/test_worker.py | 29 +++-- .../unit/test_cross_encoder_cache.py | 28 +++++ .../unit/test_encoder_output_width.py | 92 ++++++++++++++++ .../unit/test_scheduler_nixl_consumer.py | 73 +++++++++---- .../unit/test_scheduler_nixl_producer.py | 14 +-- tests/v1/ec_connector/unit/utils.py | 18 +++- vllm/config/ec_transfer.py | 5 + .../ec_transfer/ec_connector/cpu/common.py | 32 +----- .../ec_transfer/ec_connector/cpu/protocol.py | 4 +- .../ec_connector/cpu/scheduler/__init__.py | 70 ++++++++---- .../cpu/scheduler/embedding_cache.py | 21 ++-- .../ec_connector/cpu/worker/__init__.py | 16 ++- .../ec_connector/encoder_output_width.py | 102 ++++++++++++++++++ .../ec_connector/mooncake/scheduler.py | 14 +-- vllm/v1/engine/core.py | 5 + 18 files changed, 416 insertions(+), 121 deletions(-) create mode 100644 tests/v1/ec_connector/unit/test_encoder_output_width.py create mode 100644 vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py diff --git a/docs/features/ec_cpu_connector.md b/docs/features/ec_cpu_connector.md index 2a6df25f6dbf..18274b6dd970 100644 --- a/docs/features/ec_cpu_connector.md +++ b/docs/features/ec_cpu_connector.md @@ -96,7 +96,7 @@ EC transfer is configured via `--ec-transfer-config` (CLI) or the `ec_transfer_c | --- | --- | --- | --- | | `ec_enable_nixl` | `bool` | No (default `false`) | Enables NIXL P2P transfer in addition to local CPU offload. Omitted or `false` imports no NIXL/ZMQ. Extra config is not type coerced, so a string value is parsed: `"true"`, `"1"`, `"yes"` enable it, anything else does not. | | `consumer_ack_timeout_s` | `float` | No (default `2.0`) | How long a consumer waits for an `XferAck` before giving up on a read. The producer answers `XferReq`s from its own scheduler step, so its reply latency scales with the encoder's `--max-num-batched-tokens`: a loaded encoder whose steps run longer than this makes consumers abandon reads the producer is about to grant. Raise it for large encoder batches. | -| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where `block_size_bytes = hidden_dim * dtype.element_size()` (`hidden_dim` accounts for Qwen3-VL deepstack: `out_hidden_size * (1 + num_deepstack_layers)`). | +| `ec_cpu_bytes` | `int` | Yes | Total size, in bytes, of the shared CPU mmap region. `ECCPUConnector` raises `ValueError` if unset. Block count = `ec_cpu_bytes // block_size_bytes`, where each block holds one LLM input-embedding vector. Entries allocate enough blocks for their modality-specific output shape, including visual DeepStack features. | ### Environment Variables diff --git a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py index 24ea717cabb7..ef86505057cb 100644 --- a/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py +++ b/tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py @@ -22,12 +22,16 @@ def __init__(self, offset, length): self.offset = offset self.length = length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, identifier=None): self.mm_hash = mm_hash self.identifier = identifier if identifier is not None else mm_hash self.mm_position = _Pos(0, length) + self.modality = "image" class _Request: @@ -80,7 +84,7 @@ def _make_scheduler( def _load_ids(meta) -> list[int]: """Transfer ids the scheduler dispatched in this step's metadata.""" - return [transfer_id for transfer_id, _ in meta.loads.values()] + return [transfer_id for transfer_id, _, _ in meta.loads.values()] def _load_blocks(meta, mm_hash: str) -> list[int]: @@ -90,7 +94,7 @@ def _load_blocks(meta, mm_hash: str) -> list[int]: def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int): """Pre-populate a ready cache entry backed by real blocks.""" - s._cache.alloc(mm_hash, n_blocks) + s._cache.alloc(mm_hash, n_blocks, (n_blocks, _BS // s._dtype.itemsize)) s._cache.mark_ready(mm_hash) diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index c168970147b0..0b64d68ca6c0 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -97,10 +97,9 @@ def test_second_instance_opens_existing_file_and_shares_memory(): r1.cleanup() -def test_dense_dp_engines_get_distinct_regions(monkeypatch): +def test_dense_dp_engines_get_distinct_regions(): """Dense DP engines reset data_parallel_rank to 0 but share instance_id, so the region must be keyed by data_parallel_index.""" - monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8) instance_id = str(uuid.uuid4()) regions = [] try: @@ -113,6 +112,7 @@ def test_dense_dp_engines_get_distinct_regions(monkeypatch): cfg.instance_id = instance_id cfg.parallel_config = parallel_config cfg.model_config.dtype = torch.float16 + cfg.model_config.get_inputs_embeds_size.return_value = 8 cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} regions.append(common.create_ec_shared_region(cfg)) diff --git a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py index f4f5f03f6946..4ff0979e305e 100644 --- a/tests/v1/ec_connector/unit/cpu/worker/test_worker.py +++ b/tests/v1/ec_connector/unit/cpu/worker/test_worker.py @@ -91,7 +91,9 @@ def _make_region() -> ECSharedRegion: def _vllm_config(rank: int = 0) -> Mock: - return create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config = create_ec_vllm_config(rank=rank, dtype=_DTYPE) + config.model_config.get_inputs_embeds_size.return_value = _HIDDEN_DIM + return config def _meta( @@ -105,7 +107,7 @@ def _meta( return ECCPUConnectorMetadata( saves=saves or {}, loads={ - mm_hash: (idx, block_ids) + mm_hash: (idx, block_ids, (len(block_ids), _HIDDEN_DIM)) for idx, (mm_hash, block_ids) in enumerate((loads or {}).items()) }, ) @@ -904,21 +906,29 @@ def test_e2e_scheduler_worker_save_then_load(make_worker, monkeypatch): # Build scheduler sharing the same region. monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region) - scheduler = ECCPUScheduler(_vllm_config()) + config = _vllm_config() + # DeepStack-wide output: each embedding spans four blocks. + width = 4 * _HIDDEN_DIM + config.ec_transfer_config.mm_encoder_output_widths = {"image": width} + scheduler = ECCPUScheduler(config) # -- Step 1: scheduler allocates, worker saves -- - n_blocks = 3 - src = torch.arange( - n_blocks * _HIDDEN_DIM, dtype=_DTYPE, device=DEVICE_TYPE - ).reshape(n_blocks, _HIDDEN_DIM) + n_tokens = 2 + src = torch.arange(n_tokens * width, dtype=_DTYPE, device=DEVICE_TYPE).reshape( + n_tokens, width + ) class _Pos: offset = 0 - length = n_blocks + length = n_tokens + + def get_num_embeds(self): + return self.length class _Feature: identifier = "img_001" mm_position = _Pos() + modality = "image" class _Request: request_id = "req_e2e" @@ -948,8 +958,9 @@ class _Output: scheduler.update_state_after_alloc(_Request(), 0) meta_load = scheduler.build_connector_meta(scheduler_output=None) assert "img_001" in meta_load.loads - _, load_blocks = meta_load.loads["img_001"] + _, load_blocks, shape = meta_load.loads["img_001"] assert load_blocks == meta_save.saves["img_001"] + assert shape == tuple(src.shape) load_cache: dict[str, torch.Tensor] = {} worker.start_load_caches(load_cache, meta_load) diff --git a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py index 26155d8b6e64..20dd32416c0b 100644 --- a/tests/v1/ec_connector/unit/test_cross_encoder_cache.py +++ b/tests/v1/ec_connector/unit/test_cross_encoder_cache.py @@ -45,6 +45,34 @@ SPEC = TensorSpec((2, 2), "torch.float32", 16) +def test_push_shapes_follow_each_modality(monkeypatch): + """Each push is sized at its modality's measured encoder output width.""" + from vllm.distributed.ec_transfer.ec_connector.mooncake import scheduler + + monkeypatch.setattr(scheduler, "ensure_mooncake_available", lambda: None) + widths = {"image": 8192, "audio": 2048} + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths=widths) + request = SimpleNamespace( + request_id="mixed", + mm_features=[SimpleNamespace(identifier=m, modality=m) for m in widths], + get_num_encoder_embeds=lambda index: 2, + ec_transfer_params={ + "consumer_zmq": "tcp://consumer:1234", + "ec_items": [{"mm_hash": m, "transfer_id": m} for m in widths], + }, + ) + instance = scheduler.ECMooncakeScheduler(config) + try: + for index in range(len(widths)): + instance.update_state_after_alloc(request, index) + metadata = instance.build_connector_meta( + SimpleNamespace(free_encoder_mm_hashes=[], preempted_req_ids=set()) + ) + assert [spec.shape for spec in metadata.pushes] == [(2, 8192), (2, 2048)] + finally: + instance.close() + + def test_embedding_key_combines_namespace_and_identifier(): config = create_ec_vllm_config(ec_role="ec_producer") config.ec_transfer_config.ec_connector_extra_config.update( diff --git a/tests/v1/ec_connector/unit/test_encoder_output_width.py b/tests/v1/ec_connector/unit/test_encoder_output_width.py new file mode 100644 index 000000000000..b7d7c1a8225e --- /dev/null +++ b/tests/v1/ec_connector/unit/test_encoder_output_width.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup measurement of the encoder output width EC producers size by.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config +from vllm.distributed.ec_transfer.ec_connector import encoder_output_width as eow + +pytestmark = pytest.mark.cpu_test + + +class _DeepStackModel: + """Emits each modality at its own width, as Qwen3-Omni does.""" + + widths = {"image": 4 * 16, "audio": 16} + + def embed_multimodal(self, *, modality): + return [torch.zeros(3, self.widths[modality])] + + +class _Budget: + def __init__(self, *args, **kwargs): + self.mm_max_toks_per_item = dict.fromkeys(_DeepStackModel.widths, 8) + + def get_dummy_encoder_profile_inputs(self, modality, max_items_per_batch): + assert max_items_per_batch == 1 + return [(modality, None)] + + +def _worker(first_pp_rank=True, supports_mm=True): + runner = SimpleNamespace( + supports_mm_inputs=supports_mm, + get_model=_DeepStackModel, + device=torch.device("cpu"), + ) + pp_group = SimpleNamespace(is_first_rank=first_pp_rank) + return SimpleNamespace(model_runner=runner, vllm_config=None), pp_group + + +@pytest.fixture +def patched(monkeypatch): + import vllm.distributed.parallel_state as ps + import vllm.multimodal.encoder_budget as budget_mod + import vllm.multimodal.utils as mm_utils + + monkeypatch.setattr(budget_mod, "MultiModalBudget", _Budget) + monkeypatch.setattr( + mm_utils, + "group_and_batch_mm_kwargs", + lambda items, **_: iter([(items[0][0], 1, {"modality": items[0][0]})]), + ) + return lambda group: monkeypatch.setattr(ps, "get_pp_group", lambda: group) + + +def test_worker_reports_the_encoder_output_width_per_modality(patched): + """The width comes from the tensor the encoder emits, not from config.""" + worker, group = _worker() + patched(group) + assert eow._measure_on_worker(worker) == {"image": 64, "audio": 16} + + +@pytest.mark.parametrize("first_pp_rank, supports_mm", [(False, True), (True, False)]) +def test_ranks_without_an_encoder_report_nothing(patched, first_pp_rank, supports_mm): + worker, group = _worker(first_pp_rank, supports_mm) + patched(group) + assert eow._measure_on_worker(worker) == {} + + +def test_producer_records_the_widths_its_encoder_ranks_agree_on(): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + reports = [{"image": 64}, {}, {"image": 64}] + eow.measure_encoder_output_widths(config, lambda fn: reports) + assert config.ec_transfer_config.mm_encoder_output_widths == {"image": 64} + assert eow.get_encoder_output_width(config, "image") == 64 + + +@pytest.mark.parametrize( + "reports, match", + [ + ([{}, {}], "could not measure"), + ([{"image": 64}, {"image": 16}], "disagree"), + ], + ids=["no-encoder", "disagree"], +) +def test_producer_refuses_to_start_without_a_trustworthy_width(reports, match): + config = create_ec_vllm_config(ec_role="ec_producer", encoder_output_widths={}) + with pytest.raises(ValueError, match=match): + eow.measure_encoder_output_widths(config, lambda fn: reports) diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py index a3e87f9b7888..8c9e14563f5b 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py @@ -25,6 +25,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1, data=None): @@ -85,6 +88,9 @@ def close(self): pass +_MISSING = object() + + def _consumer_sched(monkeypatch): def _region(cfg): return ECSharedRegion( @@ -95,16 +101,15 @@ def _region(cfg): # ec_enable_nixl defaults to False, so this builds gate-off; the fields # below flip on the NIXL consumer state directly. - s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_consumer")) + # No measured widths: a consumer takes each shape from the announcement. + s = ECCPUScheduler( + create_ec_vllm_config(ec_role="ec_consumer", encoder_output_widths={}) + ) # Turn on NIXL consumer state without constructing real transports. s._nixl_enabled = True s._transport = _FakeTransport() s._data = None s._compat_hash = "c" - # _setup_nixl normally computes these from model_config; set them - # directly since this helper builds gate-off then flips fields on. - s._hidden_dim = _HID - s._element_size = _ES s._metadata_resolver._cache["image"] = {"image_grid_thw"} return s @@ -115,24 +120,32 @@ def _params(mm_hash, length): "peer_host": "h", "peer_port": 1, "size_bytes": length * _HID * _ES, + "shape": [length, _HID], } } def test_new_remote_read_defers_then_completes(monkeypatch): + """The read is sized from the announced shape, not a local width.""" + width = 4 * _HID s = _consumer_sched(monkeypatch) fake = _FakeSession() # Route _start_xfer to our fake session instead of real ZMQ/NIXL, but # still reserve a real not-ready cache entry so mark_ready works. - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + assert size == width * _ES + assert shape == (1, width) + entry = s._cache.alloc(mm_hash, size // _BS, shape) assert entry is not None fake.started.append(mm_hash) return True monkeypatch.setattr(s, "_start_xfer", _fake_start) - req = _Request([_Feature("h1", 1)], params=_params("h1", 1)) + params = _params("h1", 1) + params["h1"]["size_bytes"] = width * _ES + params["h1"]["shape"] = [1, width] + req = _Request([_Feature("h1", 1)], params=params) # Step 1: unseen remote item -> read started, request deferred. assert s.ensure_cache_available(req, 0) is False @@ -155,6 +168,7 @@ def _fake_start(mm_hash, info, size): s.update_state_after_alloc(req, 0) meta = s.build_connector_meta(scheduler_output=None) assert "h1" in meta.loads + assert meta.loads["h1"][2] == (1, width) assert "h1" not in s._step_completed # cleared by promote # Step 3: still cached -> admitted directly, no new transfer. @@ -176,8 +190,8 @@ def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -230,8 +244,8 @@ def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry): fake = _FakeSession() s._sessions[("h", 1)] = fake - def start(mm_hash, info, size): - assert s._cache.alloc(mm_hash, 1) is not None + def start(mm_hash, info, size, shape): + assert s._cache.alloc(mm_hash, 1, shape) is not None fake.started.append(mm_hash) return True @@ -261,8 +275,8 @@ def test_retryable_read_re_requests_without_admitting(monkeypatch): s = _consumer_sched(monkeypatch) fake = _FakeSession() - def _fake_start(mm_hash, info, size): - entry = s._cache.alloc(mm_hash, 1) + def _fake_start(mm_hash, info, size, shape): + entry = s._cache.alloc(mm_hash, 1, shape) assert entry is not None fake.started.append(mm_hash) return True @@ -287,14 +301,29 @@ def _fake_start(mm_hash, info, size): @pytest.mark.parametrize( - "size_fields", - [{"size_bytes": 999}, {"size_bytes": None}, {"size_bytes": "big"}, {}], - ids=["mismatch", "null-size", "text-size", "no-size"], + "fields", + [ + {"size_bytes": 999}, + {"size_bytes": None}, + {"size_bytes": "big"}, + {"size_bytes": _MISSING}, + {"shape": _MISSING}, + {"shape": [2, _HID]}, + ], + ids=[ + "size-mismatch", + "null-size", + "text-size", + "no-size", + "no-shape", + "wrong-rows", + ], ) -def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): - """Invalid remote sizes fail the request without raising in the scheduler.""" +def test_invalid_announcement_fails_the_request(monkeypatch, fields): + """An unusable announced size or shape fails the request, not the engine.""" s = _consumer_sched(monkeypatch) - announced = {"peer_host": "h", "peer_port": 1, **size_fields} + announced = {**_params("h1", 1)["h1"], **fields} + announced = {k: v for k, v in announced.items() if v is not _MISSING} req = _Request([_Feature("h1", 1)], params={"h1": announced}) assert s.ensure_cache_available(req, 0) is False assert s.take_unavailable_requests() == {"r1"} @@ -306,7 +335,7 @@ def test_invalid_remote_size_fails_the_request(monkeypatch, size_fields): def test_deferral_budget_is_per_request(monkeypatch): """Shared hashes have independent wait budgets and cancellation cleanup.""" s = _consumer_sched(monkeypatch) - monkeypatch.setattr(s._cache, "alloc", lambda key, n: None) + monkeypatch.setattr(s._cache, "alloc", lambda key, n, shape=None: None) params = _params("h1", 1) old = _Request([_Feature("h1", 1)], params=params, req_id="old") new = _Request([_Feature("h1", 1)], params=params, req_id="new") @@ -419,7 +448,7 @@ def test_in_flight_hash_defers_without_second_transfer(monkeypatch): s = _consumer_sched(monkeypatch) started: list[str] = [] - def _spy_start_xfer(mm_hash, info, size): + def _spy_start_xfer(mm_hash, info, size, shape): started.append(mm_hash) return True diff --git a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py index d7ef19b5f4f4..167f333f7301 100644 --- a/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py +++ b/tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py @@ -16,6 +16,9 @@ class _Pos: def __init__(self, offset, length): self.offset, self.length = offset, length + def get_num_embeds(self): + return self.length + class _Feature: def __init__(self, mm_hash, length=1): @@ -53,11 +56,8 @@ def test_request_finished_producer_emits_params(monkeypatch): # Simulate NIXL-enabled producer bookkeeping without building real NIXL. s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - # _setup_nixl normally computes these from model_config; set them - # directly since this test builds gate-off then flips fields on. - s._hidden_dim, s._element_size = 32, 2 # feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks. - entry = s._cache.alloc("h1", 2) + entry = s._cache.alloc("h1", 2, (2, 32)) assert entry is not None s._cache.mark_ready("h1") @@ -71,6 +71,7 @@ def test_request_finished_producer_emits_params(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -84,8 +85,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 - s._cache.alloc("h1", 2) # allocated but not marked ready + s._cache.alloc("h1", 2, (2, 32)) # allocated but not marked ready delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) assert delay is False @@ -96,6 +96,7 @@ def test_request_finished_announces_not_ready_entry(monkeypatch): "peer_host": "1.2.3.4", "peer_port": 5601, "size_bytes": 2 * 32 * 2, + "shape": [2, 32], } } s.shutdown() @@ -105,7 +106,6 @@ def test_request_finished_skips_unallocated_entry(monkeypatch): s = _sched_gate_off(monkeypatch) s._nixl_enabled = True s._peer_host, s._peer_port = "1.2.3.4", 5601 - s._hidden_dim, s._element_size = 32, 2 # No alloc() for "h1" at all — e.g. the cache was full at save time. delay, params = s.request_finished(_Request([_Feature("h1", length=2)])) diff --git a/tests/v1/ec_connector/unit/utils.py b/tests/v1/ec_connector/unit/utils.py index a718bc63b635..705a8ff7ad1e 100644 --- a/tests/v1/ec_connector/unit/utils.py +++ b/tests/v1/ec_connector/unit/utils.py @@ -18,12 +18,16 @@ _PCP_SIZE = 1 +_EMBED_WIDTH = 32 + + def create_ec_vllm_config( *, ec_role: ECRole = "ec_both", tensor_parallel_size: int = 1, rank: int = 0, dtype: torch.dtype = torch.float16, + encoder_output_widths: dict[str, int] | None = None, ) -> Mock: """Build a `VllmConfig` stand-in for EC connector unit tests. @@ -31,10 +35,8 @@ def create_ec_vllm_config( role derivation and the rank arithmetic the connector depends on behave as they do in production. - `model_config` is a stub because nothing under test needs a real one: the - connector reads only `dtype` here, and the paths that inspect the model - itself are not exercised. Building a real `ModelConfig` would resolve an HF - model and inspect its architecture. + `model_config` supplies a dtype and embedding width without resolving an + HF model or inspecting its architecture. Args: ec_role: EC role this instance plays. `is_ec_producer` and @@ -42,6 +44,8 @@ def create_ec_vllm_config( tensor_parallel_size: TP degree. rank: Global rank of this worker. dtype: Encoder cache dtype. + encoder_output_widths: Per-modality encoder output width, as a + producer measures it at startup. Defaults to the embedding width. Returns: A `VllmConfig`-specced mock carrying the configs above. @@ -59,9 +63,15 @@ def create_ec_vllm_config( ec_role=ec_role, engine_id=str(uuid.uuid4()), ) + ec_transfer_config.mm_encoder_output_widths = ( + dict.fromkeys(("image", "video", "audio"), _EMBED_WIDTH) + if encoder_output_widths is None + else encoder_output_widths + ) model_config = Mock(spec=ModelConfig) model_config.dtype = dtype + model_config.get_inputs_embeds_size.return_value = _EMBED_WIDTH vllm_config = Mock(spec=VllmConfig) vllm_config.ec_transfer_config = ec_transfer_config diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 5c2b8e589c2e..b3ca10a5a690 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -55,6 +55,11 @@ class ECTransferConfig: """The Python module path to dynamically load the EC connector from. Only supported in V1.""" + mm_encoder_output_widths: dict[str, int] = field(default_factory=dict, init=False) + """Per-token width of the encoder output for each modality, measured by EC + producers at startup. It can exceed the LM embedding width (e.g. DeepStack). + """ + def compute_hash(self) -> str: """WARNING: Whenever a new field is added to this config, ensure that it is included in the factors list if diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index eff16e08b0ae..461cd6a55a6c 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -33,12 +33,14 @@ class ECCPUConnectorMetadata(ECConnectorMetadata): saves: dict[str, list[int]] = field(default_factory=dict) # Consumer role: mm_hashes whose bytes are available in the local mmap, - # mapped to (transfer_id, block_ids); the worker's start_load_caches + # mapped to (transfer_id, block_ids, shape); the worker's start_load_caches # copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id # once its copy lands. The id is minted per dispatch, so two loads of the # same mm_hash in different steps are distinct and a late report from the # earlier one cannot release the later one's pin. - loads: dict[str, tuple[int, list[int]]] = field(default_factory=dict) + loads: dict[str, tuple[int, list[int], tuple[int, int]]] = field( + default_factory=dict + ) @dataclass @@ -66,30 +68,6 @@ def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetada return self -def _get_encoder_cache_hidden_dim(vllm_config: "VllmConfig") -> int: - """Return the per-token hidden dimension for encoder cache entries. - - For most models this equals the LLM's hidden size. Qwen3-VL (and any - future model with deepstack visual encoding) is an exception: the ViT - concatenates its own output with features from N decoder layers before - storing in encoder_cache, producing a tensor of width - ``out_hidden_size * (1 + N)`` per visual token. Using the plain LLM - hidden size would under-allocate EC blocks and silently truncate the - transferred data, leading to a shape mismatch on the consumer. - """ - model_config = vllm_config.model_config - hf_config = getattr(model_config, "hf_config", None) - vision_config = ( - getattr(hf_config, "vision_config", None) if hf_config is not None else None - ) - if vision_config is not None: - out_hidden_size = getattr(vision_config, "out_hidden_size", None) - deepstack_indexes = getattr(vision_config, "deepstack_visual_indexes", None) - if out_hidden_size is not None and deepstack_indexes: - return out_hidden_size * (1 + len(deepstack_indexes)) - return model_config.get_inputs_embeds_size() - - def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: """Build the EC mmap region from `vllm_config`. @@ -103,7 +81,7 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: engine_id = f"{vllm_config.instance_id}_dp{dp_index}" dtype = vllm_config.model_config.dtype - hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) + hidden_dim = vllm_config.model_config.get_inputs_embeds_size() element_size = torch.empty(0, dtype=dtype).element_size() block_size_bytes = hidden_dim * element_size diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py index 85bd3cae32fa..e20bf48f0ba8 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py @@ -14,9 +14,9 @@ import msgspec -# Bump when the on-wire shape (XferReq/XferAck) changes in a +# Bump when the wire format or tensor/block layout changes in a # backward-incompatible way. Peers reject mismatched versions. -EC_CONNECTOR_VERSION: int = 1 +EC_CONNECTOR_VERSION: int = 2 class XferStatus(enum.IntEnum): diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py index 65403980b9d0..567f2d59c9f3 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py @@ -16,12 +16,14 @@ from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( ECCPUConnectorMetadata, ECCPUWorkerMetadata, - _get_encoder_cache_hidden_dim, create_ec_shared_region, ) from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( EmbeddingCache, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, +) from vllm.distributed.ec_transfer.ec_connector.utils import ( PlaceholderMetadataResolver, collect_ec_item_metadata, @@ -49,6 +51,7 @@ class ECCPUScheduler: """Scheduler delegate for the ECCPUConnector.""" def __init__(self, vllm_config: "VllmConfig") -> None: + self._vllm_config = vllm_config ec_config = vllm_config.ec_transfer_config assert ec_config is not None self._is_producer: bool = ec_config.is_ec_producer @@ -61,8 +64,8 @@ def __init__(self, vllm_config: "VllmConfig") -> None: # mm_hash → block IDs allocated this step for GPU→mmap saves. self._pending_saves: dict[str, list[int]] = {} - # mm_hash → (transfer_id, block IDs) to load from mmap→GPU this step. - self._pending_loads: dict[str, tuple[int, list[int]]] = {} + # mm_hash → (transfer_id, block IDs, shape) to load from mmap→GPU. + self._pending_loads: dict[str, tuple[int, list[int], tuple[int, int]]] = {} # Dispatched loads awaiting completion reports, keyed by transfer id: # transfer_id → (mm_hash, reports still outstanding). The pin taken at @@ -108,11 +111,9 @@ def __init__(self, vllm_config: "VllmConfig") -> None: self._unrecoverable: set[str] = set() self._peer_host: str | None = None self._peer_port: int | None = None - # Model shape for size checks + compat hash; only set by - # _setup_nixl (the gate-off path never touches model_config). - self._dtype: torch.dtype | None = None - self._hidden_dim: int = 0 - self._element_size: int = 0 + dtype = vllm_config.model_config.dtype + assert isinstance(dtype, torch.dtype) + self._dtype = dtype self._ack_timeout_s: float = 0.0 if self._nixl_enabled: try: @@ -148,9 +149,6 @@ def _setup_nixl(self, vllm_config: "VllmConfig") -> None: ) engine_id = self._ec_config.engine_id assert engine_id is not None - self._dtype = vllm_config.model_config.dtype - self._hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) - self._element_size = torch.empty(0, dtype=self._dtype).element_size() # How long a consumer waits for an XferAck. The producer answers from # its scheduler step, so its reply latency scales with the encoder's # batch size: a deployment whose steps run longer must raise this. @@ -296,22 +294,30 @@ def _admit_item( # held by a quarantined/settling DMA and cannot be reused. return False - expected = pos.length * self._hidden_dim * self._element_size try: size = int(remote["size_bytes"]) + shape = tuple(int(dim) for dim in remote["shape"]) except (KeyError, TypeError, ValueError) as error: raise _RemoteUnavailable("the announced size was unusable") from error - if size != expected: + # The width is the producer's to report; the row count is not. + num_embeds = pos.get_num_embeds() + if ( + len(shape) != 2 + or shape[0] != num_embeds + or size != self._num_blocks(shape) * self._region.block_size_bytes + ): logger.warning( - "EC consumer: size mismatch mm_hash=%s announced=%d expected=%d", + "EC consumer: size mismatch mm_hash=%s announced=%d shape=%s " + "expected %d rows", mm_hash, size, - expected, + shape, + num_embeds, ) raise _RemoteUnavailable("the announced size was wrong") try: - started = self._start_xfer(mm_hash, remote, expected) + started = self._start_xfer(mm_hash, remote, size, (shape[0], shape[1])) except Exception as error: logger.exception( "EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash @@ -377,7 +383,11 @@ def take_unavailable_requests(self) -> set[str]: return failed def _start_xfer( - self, mm_hash: str, info: "dict[str, Any]", size_bytes: int + self, + mm_hash: str, + info: "dict[str, Any]", + size_bytes: int, + shape: tuple[int, int], ) -> bool: """Allocate a not-ready cache entry and start a NIXL READ into it. @@ -391,7 +401,7 @@ def _start_xfer( ) n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes)) - entry = self._cache.alloc(mm_hash, n_blocks) + entry = self._cache.alloc(mm_hash, n_blocks, shape) if entry is None: logger.debug( "EC consumer: cache full for mm_hash=%s (%d blocks); deferring", @@ -491,22 +501,36 @@ def _promote_completed_reads(self) -> None: """ self._step_completed.clear() + def _num_blocks(self, shape: tuple[int, ...]) -> int: + size_bytes = shape[0] * shape[1] * self._dtype.itemsize + block_size = self._region.block_size_bytes + return (size_bytes + block_size - 1) // block_size + def update_state_after_alloc(self, request: "Request", index: int) -> None: feature = request.mm_features[index] mm_hash = feature.identifier if self._is_producer and self._cache.get(mm_hash) is None: - entry = self._cache.alloc(mm_hash, feature.mm_position.length) + shape = ( + feature.mm_position.get_num_embeds(), + get_encoder_output_width(self._vllm_config, feature.modality), + ) + entry = self._cache.alloc(mm_hash, self._num_blocks(shape), shape) if entry is not None: self._pending_saves[mm_hash] = list(entry.block_ids) if self._is_consumer and mm_hash not in self._pending_loads: entry = self._cache.get(mm_hash) if entry is not None and entry.ready: + assert entry.shape is not None self._cache.pin(mm_hash) transfer_id = self._next_transfer_id self._next_transfer_id += 1 - self._pending_loads[mm_hash] = (transfer_id, list(entry.block_ids)) + self._pending_loads[mm_hash] = ( + transfer_id, + list(entry.block_ids), + entry.shape, + ) self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports) def build_connector_meta( @@ -608,13 +632,13 @@ def request_finished( # A read arriving before the save lands is NACKed NACK_NOT_READY, # which the consumer retries on a later step rather than treating # as a miss. - size_bytes = ( - feature.mm_position.length * self._hidden_dim * self._element_size - ) + assert entry.shape is not None + size_bytes = len(entry.block_ids) * self._region.block_size_bytes items[mm_hash].update( peer_host=self._peer_host, peer_port=self._peer_port, size_bytes=size_bytes, + shape=list(entry.shape), ) logger.debug( "EC producer: announcing NIXL-readable encodings req_id=%s items=%s", diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py index 2fd4b01059db..046fba05741e 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py @@ -14,13 +14,20 @@ class CacheEntry: - """A single cache entry. Read `.block_ids` and `.ready` freely; - mutations only through EmbeddingCache methods.""" + """A single cache entry. Read `.block_ids`, `.shape` and `.ready` freely; + mutations only through EmbeddingCache methods. - __slots__ = ("block_ids", "_pin_count") + `shape` is the encoder output the blocks hold, so a load can restore it + without knowing how wide this model's encoder output is. + """ + + __slots__ = ("block_ids", "shape", "_pin_count") - def __init__(self, block_ids: tuple[int, ...]) -> None: + def __init__( + self, block_ids: tuple[int, ...], shape: tuple[int, int] | None = None + ) -> None: self.block_ids = block_ids + self.shape = shape self._pin_count = -1 # not ready @property @@ -76,7 +83,9 @@ def has_held_entries(self) -> bool: with self._lock: return len(self._entries) != len(self._entries_free_list) - def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: + def alloc( + self, key: str, n_blocks: int, shape: tuple[int, int] | None = None + ) -> CacheEntry | None: """Allocate *n_blocks* for *key*, evicting as needed. The entry starts not-ready. Returns None if there is not enough @@ -97,7 +106,7 @@ def alloc(self, key: str, n_blocks: int) -> CacheEntry | None: # in the free list when it is evicted, and lets consumers detect # which of an entry's blocks are adjacent. block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks))) - entry = CacheEntry(block_ids) + entry = CacheEntry(block_ids, shape) self._entries[key] = entry return entry diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py index ef7ff0f75162..95c6c9612885 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py @@ -267,8 +267,7 @@ def save_caches( # so they are added up in uint64; signed arithmetic rejects them. src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr()) dst_base = self._region.blocks.data_ptr() - # Entries are sized by placeholder count, which can exceed the number - # of embeddings, so blocks the output never reaches get no descriptor. + # Copy only the tensor bytes, excluding any block padding. n_used = -(-total_bytes // block_size) if n_used: slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used]) @@ -367,7 +366,7 @@ def start_load_caches( load_items = sorted( connector_metadata.loads.items(), key=lambda kv: kv[1][1][0] ) - total_blocks = sum(len(block_ids) for _, (_, block_ids) in load_items) + total_blocks = sum(len(block_ids) for _, (_, block_ids, _) in load_items) stream = self._acquire_stream() compute_stream = current_platform.current_stream() @@ -384,7 +383,7 @@ def start_load_caches( bufs = self._buf_pool.acquire(total_blocks) slots, first_blocks, num_blocks = _coalesce_runs( - list(chain.from_iterable(ids for _, (_, ids) in load_items)) + list(chain.from_iterable(ids for _, (_, ids, _) in load_items)) ) bufs.add_copies( 0, @@ -405,7 +404,7 @@ def start_load_caches( Transfer( start_event=start_event, end_event=end_event, - completions=[tid for _, (tid, _) in load_items], + completions=[tid for _, (tid, _, _) in load_items], bufs=bufs, stream=stream, num_bytes=total_blocks * block_size, @@ -414,11 +413,10 @@ def start_load_caches( # Slice contiguous buffer into per-hash views. offset = 0 - for mm_hash, (_, block_ids) in load_items: + for mm_hash, (_, block_ids, shape) in load_items: n = len(block_ids) - encoder_cache[mm_hash] = ( - dst_buf[offset : offset + n].view(dtype).reshape(n, -1) - ) + flat = dst_buf[offset : offset + n].view(dtype).reshape(-1) + encoder_cache[mm_hash] = flat[: shape[0] * shape[1]].view(shape) offset += n current_platform.current_stream().wait_stream(stream) diff --git a/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py new file mode 100644 index 000000000000..c1a270bb2b18 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/encoder_output_width.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder output widths for sizing EC transfers. + +An encoder can emit features wider than the LM embedding (DeepStack packs +several layers on the last dim) by an amount only the model knows. Producers +size transfers before the first real item is encoded, so the engine measures +the width once at startup by encoding a dummy item per modality. +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@torch.inference_mode() +def _measure_on_worker(worker: Any) -> dict[str, int]: + """Encode one dummy item per modality; runs on every worker. + + Returns {} on ranks that hold no encoder. + """ + from vllm.distributed.parallel_state import get_pp_group + from vllm.multimodal import MULTIMODAL_REGISTRY + from vllm.multimodal.encoder_budget import MultiModalBudget + from vllm.multimodal.utils import group_and_batch_mm_kwargs + from vllm.utils.torch_utils import PIN_MEMORY + + runner = worker.model_runner + if not runner.supports_mm_inputs or not get_pp_group().is_first_rank: + return {} + budget = MultiModalBudget( + worker.vllm_config, MULTIMODAL_REGISTRY, enable_cache=False + ) + model = runner.get_model() + widths: dict[str, int] = {} + for modality in budget.mm_max_toks_per_item: + items = budget.get_dummy_encoder_profile_inputs(modality, 1) + if not items: + continue + _, _, mm_kwargs = next( + group_and_batch_mm_kwargs( + items, device=runner.device, pin_memory=PIN_MEMORY + ) + ) + outputs = model.embed_multimodal(**mm_kwargs) + widths[modality] = int(outputs[0].shape[-1]) + del outputs + return widths + + +def measure_encoder_output_widths( + vllm_config: "VllmConfig", collective_rpc: Callable[..., list[Any]] +) -> None: + """Record the encoder output widths on an EC producer's config. + + Call after the model is loaded and before the KV cache is allocated: the + Scheduler's connector reads the widths, and once the KV cache holds the + memory budget there is no room left for a dummy encode. + + Raises: + ValueError: No worker could encode a dummy item, or workers disagree. + + """ + ec_config = vllm_config.ec_transfer_config + if ec_config is None or not ec_config.is_ec_producer: + return + reports = [widths for widths in collective_rpc(_measure_on_worker) if widths] + if not reports: + raise ValueError( + "EC producer could not measure the encoder output width: the model " + "encoded no dummy multimodal input." + ) + if any(widths != reports[0] for widths in reports): + raise ValueError(f"Workers disagree on the encoder output widths: {reports}") + logger.info("EC producer encoder output widths: %s", reports[0]) + ec_config.mm_encoder_output_widths = reports[0] + + +def get_encoder_output_width(vllm_config: "VllmConfig", modality: str) -> int: + """Encoder output width for `modality`, as measured at producer startup. + + Raises: + ValueError: No width was measured for `modality`, so a transfer + cannot be sized. + + """ + assert vllm_config.ec_transfer_config is not None + widths = vllm_config.ec_transfer_config.mm_encoder_output_widths + if modality not in widths: + raise ValueError( + f"No encoder output width was measured for modality {modality!r}; " + f"measured: {sorted(widths)}" + ) + return widths[modality] diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py index 2a607118db45..597ba767dfc8 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py +++ b/vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py @@ -20,8 +20,8 @@ from vllm.distributed.ec_transfer.ec_connector.base import ( ECConnectorMetadata, ) -from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( - _get_encoder_cache_hidden_dim, +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + get_encoder_output_width, ) from vllm.distributed.ec_transfer.ec_connector.mooncake.config import ( _RESERVATION_TTL_SECONDS, @@ -85,9 +85,7 @@ def __init__(self, vllm_config: VllmConfig) -> None: self._is_consumer = config.is_consumer self._control_addr = config.control_addr self._push_wait_timeout = config.push_wait_timeout_s - self._encoder_cache_hidden_dim = ( - _get_encoder_cache_hidden_dim(vllm_config) if config.is_producer else None - ) + self._vllm_config = vllm_config self._model_config = vllm_config.model_config self._control_client = ControlClient(config.control_timeout_ms) self._control_executor = ThreadPoolExecutor( @@ -447,8 +445,10 @@ def _prepare_push_spec(self, request: Any, index: int) -> None: def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: dtype = self._model_config.dtype assert isinstance(dtype, torch.dtype) - assert self._encoder_cache_hidden_dim is not None - shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) + hidden_dim = get_encoder_output_width( + self._vllm_config, request.mm_features[index].modality + ) + shape = (request.get_num_encoder_embeds(index), hidden_dim) return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) def update_state_after_alloc(self, request: Any, index: int) -> None: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bbca64a58ae9..35cc8cc1860c 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -27,6 +27,9 @@ cleanup_dist_env_and_memory, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import ( + measure_encoder_output_widths, +) from vllm.envs import enable_envs_cache from vllm.logger import configure_logging, init_logger from vllm.logging_utils.dump_input import dump_engine_exception @@ -148,6 +151,8 @@ def __init__( if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self._eep_scale_up_before_kv_init() + measure_encoder_output_widths(vllm_config, self.collective_rpc) + # Setup KV Caches and update CacheConfig after profiling. kv_cache_config = self._initialize_kv_caches(vllm_config) self.structured_output_manager = StructuredOutputManager(vllm_config) From c75cb63cd2bb8e0dc62608e9930d8b953ee3b422 Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Mon, 12 Oct 2026 20:56:44 +0800 Subject: [PATCH 408/952] [Qwen4Exp] support Sequence Parallelism for Qwen3.8-Flash-Next (#56322) Signed-off-by: huanghaoyan.hhy Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/models_basic.yaml | 3 +- .buildkite/test_areas/models_distributed.yaml | 17 ++ .../core/test_vocab_parallel_embedding.py | 42 +++ tests/models/qwen4_exp/test_config.py | 86 ++++-- tests/models/qwen4_exp/test_ple.py | 103 ++++--- tests/models/qwen4_exp/test_tpsp.py | 267 ++++++++++++++++++ tests/models/test_registry.py | 16 ++ tests/test_config.py | 41 +++ vllm/config/model.py | 8 + vllm/config/parallel.py | 5 + vllm/engine/arg_utils.py | 3 + .../layers/vocab_parallel_embedding.py | 35 +-- vllm/model_executor/models/__init__.py | 4 + vllm/model_executor/models/config.py | 10 + vllm/model_executor/models/interfaces.py | 21 ++ vllm/model_executor/models/qwen3_next.py | 8 +- vllm/model_executor/models/registry.py | 12 + .../qwen4_exp/common/ngram_embedding.py | 56 ++-- vllm/models/qwen4_exp/common/ops/ple.py | 40 +-- vllm/models/qwen4_exp/nvidia/model.py | 124 ++++++-- vllm/models/qwen4_exp/nvidia/mtp.py | 44 ++- .../qwen4_exp/nvidia/ngram_embedding.py | 45 ++- vllm/models/qwen4_exp/nvidia/ple_layer.py | 60 ++-- 23 files changed, 871 insertions(+), 179 deletions(-) create mode 100644 tests/models/qwen4_exp/test_tpsp.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 483a9aceac22..288e7621cab4 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -205,12 +205,12 @@ steps: device: h200_35gb source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ commands: - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py + - pytest -v -s models/qwen4_exp/test_ple.py -k test_fused_conv_correctness mirror: amd: label: ":amd: (MI355 DPX) Qwen4 Exp" @@ -231,7 +231,6 @@ steps: device: cpu-small source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 9ab75cfc6e70..8d2bc5dbaf3c 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -2,6 +2,23 @@ group: Models - Distributed depends_on: - image-build steps: +- label: ":nvidia: (L4) Qwen4 Exp SP" + key: qwen4-exp-sp-tests-4-gpus + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + device: l4 + num_devices: 4 + source_file_dependencies: + - vllm/config/parallel.py + - vllm/models/common/ops/sequence_parallel.py + - vllm/models/qwen4_exp/ + - vllm/model_executor/layers/vocab_parallel_embedding.py + - vllm/model_executor/models/qwen3_next.py + - tests/models/qwen4_exp/test_ple.py + - tests/models/qwen4_exp/test_tpsp.py + commands: + - pytest -v -s models/qwen4_exp/test_tpsp.py + - label: ":nvidia: (L4) Distributed Models Shard %N" key: distributed-model-tests-2-gpus timeout_in_minutes: 35 diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index b0038c084e3c..b50b5e568871 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -9,6 +9,8 @@ reconstructs a full-table lookup. """ +from unittest.mock import Mock + import pytest import torch @@ -124,3 +126,43 @@ def test_rank_sum_is_full_lookup(tp_size): total += _fused(ids, _shard(full_table, si, rows, dtype), si) torch.testing.assert_close(total, full_table[ids.long()], atol=0.0, rtol=0.0) + + +@requires_cuda +@torch.inference_mode() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("use_fused", [False, True]) +def test_embedding_optionally_reduces_output(dtype, use_fused, default_vllm_config): + """SP's deferred reduction preserves masked contributions and their dtype.""" + with torch.device("cuda"): + values = torch.arange(128 * 16).reshape(128, 16) + table = ((values % 29 - 14) / 16).to(dtype) + ids = torch.tensor([0, 63, 64, 127, 12]) + expected = table[ids].float() + for rank in range(2): + local = expected.clone() + local[ids // 64 != rank] = 0 + remote = (expected - local).to(dtype) + group = Mock(world_size=2, rank_in_group=rank) + # Model the other vocabulary owner's contribution, including FP8 bytes. + group.all_reduce.side_effect = lambda output, remote=remote: ( + output + remote.view(output.dtype) + ) + for reduce_results in (True, False): + embedding = VocabParallelEmbedding( + 128, + 16, + params_dtype=dtype, + parallel_group=group, + reduce_results=reduce_results, + ) + embedding.use_fused_embedding = use_fused + embedding.weight_loader(embedding.get_parameter("weight"), table) + actual = embedding(ids) + assert actual.dtype == dtype + torch.testing.assert_close( + actual.float(), + expected if reduce_results else local, + atol=0, + rtol=0, + ) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 6fde77eda246..818f51543857 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -9,6 +9,7 @@ import torch from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm.config import CacheConfig, ParallelConfig from vllm.config.speculative import SpeculativeConfig from vllm.model_executor.models.config import ( Qwen3_5ForConditionalGenerationConfig, @@ -48,23 +49,54 @@ def _text_config(**kwargs) -> Qwen4ExpTextConfig: return Qwen4ExpTextConfig(**values) -def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: +def test_moe_sp_rejects_pipeline_parallel() -> None: + """Automatic MoE SP must validate HC topology even without --enable-hc-sp.""" + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=_text_config(num_experts=4, ple_layer_ids=[]), + ), + cache_config=CacheConfig(), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + pipeline_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ), + ) + with ( + patch("vllm.platforms.current_platform.is_cuda", return_value=True), + pytest.raises(ValueError, match="MoE SP requires PP=1"), + ): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(config) + + +@pytest.mark.parametrize("use_sp", [False, True]) +def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: + """MTP returns complete, contiguous outputs after gathering padded SP shards.""" from vllm.models.qwen4_exp.nvidia.mtp import ( Qwen4ExpMultiTokenPredictor, ) model = object.__new__(Qwen4ExpMultiTokenPredictor) + model.use_hc_sequence_parallel = use_sp + model.use_sequence_parallel = False torch.nn.Module.__init__(model) model.hc_count = 2 model.hidden_size = 4 model.num_mtp_layers = 1 - model.layers = [ - lambda **kwargs: ( - kwargs["hidden_states"], - kwargs["hidden_states"], - torch.zeros(kwargs["hidden_states"].shape[0], 2), - ), - ] + model.pre_fc_norm_embedding = torch.nn.Identity() + model.pre_fc_norm_hidden = torch.nn.Identity() + model.fc_embedding = torch.nn.Identity() + model.fc_hidden = torch.nn.Identity() + + def decoder(**kwargs): + """Check that the embedding residual shares the hidden state's token slice.""" + hidden = kwargs["hidden_states"] + assert kwargs["prev_block_output"].shape == (hidden.shape[0], model.hidden_size) + return hidden, hidden, torch.zeros(hidden.shape[0], model.hc_count) + + model.layers = [decoder] model.hyper_connection_mixer = SimpleNamespace( combine_and_mix=lambda hidden_states, block_output, injection: ( hidden_states, @@ -72,24 +104,32 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: None, ), ) - multi_hidden = torch.arange(16, dtype=torch.float32).reshape(2, 8) - pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=True) - - with patch( - "vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", - return_value=pp_group, + multi_hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8) + expected_sample = multi_hidden.unflatten(-1, (2, 4)).mean(-2) + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + + def gather(local): + """Supply rank 1's last token and padding for a simulated two-rank gather.""" + assert local.shape[0] == 2 + remote = torch.cat([expected_sample[2:], multi_hidden[2:]], dim=-1) + return torch.cat([local, remote, torch.full_like(remote, float("nan"))]) + + with ( + patch("vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", return_value=pp_group), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_shard", side_effect=lambda x: x[:2]), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_all_gather", side_effect=gather), ): sample_hidden, returned_multi_hidden = model.forward( input_ids=None, - positions=torch.arange(2), - intermediate_tensors={"hidden_states": multi_hidden}, + positions=torch.arange(3), + hidden_states=multi_hidden, + inputs_embeds=torch.zeros(3, 4), ) - torch.testing.assert_close( - sample_hidden, - multi_hidden.unflatten(-1, (2, 4)).mean(dim=-2), - ) - assert returned_multi_hidden is multi_hidden + torch.testing.assert_close(sample_hidden, expected_sample) + torch.testing.assert_close(returned_multi_hidden, multi_hidden) + assert sample_hidden.is_contiguous() + assert returned_multi_hidden.is_contiguous() @spawn_new_process_for_each_test @@ -150,9 +190,7 @@ def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> Non hf_text_config=_text_config(ple_layer_ids=ple_layer_ids), multimodal_config=None, ), - parallel_config=SimpleNamespace( - pipeline_parallel_size=2, enable_dbo=False, ubatch_size=1 - ), + parallel_config=ParallelConfig(pipeline_parallel_size=2), speculative_config=None, ) with patch.object( diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index d7ee461d8f0d..a4a428ecccea 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -611,46 +611,65 @@ def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None: ) -def test_pinned_embedding_forward_finalizes_prefetched_output( +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("reduce_results", [False, True]) +def test_pinned_embedding_finalize_prefetch_optionally_reduces_output( monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + reduce_results: bool, ) -> None: + """Deferring reduction must preserve full token rows and wait for the lookup.""" + total_tokens = 3 embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) - embedding._prefetch_buffer = torch.empty(4, 2, 3, dtype=torch.float8_e4m3fn) + values = torch.arange(1, 13).reshape(4, 3) + embedding._prefetch_buffer = torch.stack( + (values, torch.zeros_like(values)), dim=1 + ).to(dtype) embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - expected = torch.arange(12).reshape(2, 6).to(torch.float8_e4m3fn) - - def finalize_prefetched( - prefetch_output: torch.Tensor, - output: torch.Tensor, - ) -> None: - assert prefetch_output is embedding._prefetch_buffer - output.copy_(expected) - + embedding.tp_size = 2 + embedding.reduce_results = reduce_results + embedding.etp_data_parallel_size = 1 + embedding._prefetch_stream = object() + waited_streams: list[object] = [] + reduced_dtypes = [] monkeypatch.setattr( - embedding, - "_finalize_prefetch", - finalize_prefetched, + torch.cuda, + "current_stream", + lambda: SimpleNamespace(wait_stream=waited_streams.append), ) - output = embedding(hidden_states) + def all_reduce(tensor: torch.Tensor) -> torch.Tensor: + assert waited_streams == [embedding._prefetch_stream] + reduced_dtypes.append(tensor.dtype) + # The remote rank owns the other head's vocabulary rows. + return tensor + tensor.flip(1) - assert output.dtype == embedding._prefetch_buffer.dtype - assert torch.equal(output, expected) + embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce) + output = embedding.finalize_prefetch(total_tokens=total_tokens) + + expected = ( + values.unsqueeze(1).expand(-1, 2, -1).to(dtype) + if reduce_results + else embedding._prefetch_buffer + ) + expected = expected[:total_tokens].flatten(-2) + comm_dtype = torch.int8 if dtype == torch.float8_e4m3fn else dtype + assert reduced_dtypes == ([comm_dtype] if reduce_results else []) + assert waited_streams == [embedding._prefetch_stream] + assert output.dtype == dtype + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) -def test_pinned_embedding_forward_requires_prior_start_prefetch() -> None: - """Forward before any start_prefetch must fail loudly, not read garbage.""" +def test_pinned_embedding_finalize_prefetch_requires_prior_start_prefetch() -> None: + """Finalizing before start_prefetch must fail instead of reading garbage.""" embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) embedding._prefetch_buffer = None embedding._prefetch_stream = None embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="prior start_prefetch"): - embedding(hidden_states) + embedding.finalize_prefetch(total_tokens=2) def test_pinned_embedding_finalize_requires_prior_start_prefetch() -> None: @@ -781,15 +800,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( +def test_pinned_embedding_start_prefetch_allocates_lazily( monkeypatch: pytest.MonkeyPatch, ) -> None: - """start_prefetch lazily allocates the stream/buffer, then forward consumes it. - - Exercises the full pinned-prefetch contract end to end: first call - allocates the side stream and buffer with the constructor geometry, - and forward finalizes the lookup into the returned output. - """ + """Allocate prefetch storage on first use and finalize the lookup results.""" _mock_etp_group(monkeypatch) monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( @@ -825,8 +839,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( ) ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16, device="cuda:0") - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) buffer = embedding._prefetch_buffer assert buffer is not None @@ -835,7 +848,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( assert buffer.device.type == "cuda" assert embedding._prefetch_stream is not None - output = embedding(hidden_states) + output = embedding.finalize_prefetch(total_tokens=ngram_ids.shape[0]) expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) torch.testing.assert_close(output, expected, rtol=0, atol=0) @@ -1149,7 +1162,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: prefetched: list[torch.Tensor] = [] module.ngram_embedding = SimpleNamespace( supports_prefetch=True, - start_prefetch=lambda _, ids: prefetched.append(weak_ref_tensor(ids)), + start_prefetch=lambda ids: prefetched.append(weak_ref_tensor(ids)), ) expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params) @@ -1157,7 +1170,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: # like a later CUDA graph segment would. pool = torch.cuda.MemPool() with torch.cuda.use_mem_pool(pool): - module.start_prefetch(None, input_ids, query_start_loc, ngram_context) + module.start_prefetch(input_ids, query_start_loc, ngram_context) torch.full_like(expected, -1) assert torch.equal(prefetched[0], expected) @@ -1859,12 +1872,18 @@ def _make_conv_metadata( ], ) @pytest.mark.parametrize( - "layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"] + ("layer_cls", "fuse_residual"), + [ + pytest.param(Qwen4ExpPLELayer, True, id="nvidia-fused"), + pytest.param(Qwen4ExpPLELayer, False, id="nvidia-unfused"), + pytest.param(Qwen4ExpPLELayerAMD, True, id="amd-fused"), + ], ) def test_fused_conv_correctness( case: _ConvBatchCase, state_layout: str, layer_cls: type[nn.Module], + fuse_residual: bool, ) -> None: device = torch.device("cuda") metadata, num_real_tokens = _make_conv_metadata(case, device) @@ -1906,12 +1925,15 @@ def test_fused_conv_correctness( ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() - residual_reference = residual.clone() + # The unfused kernel must overwrite the buffer without reading its contents. + residual_reference = ( + residual.clone() if fuse_residual else torch.zeros_like(residual) + ) module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, - outer_residual=outer_residual, + outer_residual=outer_residual if fuse_residual else None, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1925,7 +1947,8 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) - residual_reference = outer_residual + residual_reference + if fuse_residual: + residual_reference = outer_residual + residual_reference assert torch.equal( residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] @@ -1935,7 +1958,7 @@ def test_fused_conv_correctness( if case.graph_padding: assert torch.equal( residual_kernel[num_real_tokens:], - (outer_residual + residual)[num_real_tokens:], + residual_reference[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py new file mode 100644 index 000000000000..5182ab09c398 --- /dev/null +++ b/tests/models/qwen4_exp/test_tpsp.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multi-GPU equivalence tests for Qwen4Exp HC and MoE sequence parallelism.""" + +from dataclasses import MISSING, fields +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.multiprocessing as mp +from torch import nn +from transformers import Qwen4ExpTextConfig + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("Qwen4Exp SP requires CUDA", allow_module_level=True) + +from tests.utils import multi_gpu_test +from vllm.config import ( + EngramConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.linear import RowParallelLinear +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) +from vllm.models.qwen4_exp.nvidia import model as qwen4_model +from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer +from vllm.utils.network_utils import get_open_port +from vllm.v1.attention.backends.short_conv_attn import PleShortConvAttentionMetadata +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + +from .test_ple import _ConvBatchCase, _make_conv_metadata + + +class _TPAttentionProjection(nn.Module): + """Keep real TP communication without setting up an attention KV cache.""" + + def __init__(self, config, *, reduce_results: bool, **kwargs) -> None: + super().__init__() + self.proj = RowParallelLinear( + config.hidden_size, + config.hidden_size, + bias=False, + input_is_parallel=False, + reduce_results=reduce_results, + ) + + def forward(self, hidden_states: torch.Tensor): + return self.proj(hidden_states)[0] + + +def _make_config(rank: int, *, hc_sp: bool = False, moe_sp: bool = False) -> VllmConfig: + """Configure HC SP explicitly or MoE SP through TP=2, DP=2 and EP.""" + config = VllmConfig( + parallel_config=ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=2 if moe_sp else 1, + data_parallel_rank=rank // 2, + enable_expert_parallel=moe_sp, + enable_hc_sp=hc_sp, + all2all_backend="allgather_reducescatter", + ), + ) + config.model_config = SimpleNamespace( + architecture="Qwen4ExpForCausalLM", + is_moe=True, + sleep_mode_offload_cudagraph=False, + hf_text_config=Qwen4ExpTextConfig( + vocab_size=128, + eos_token_id=0, + hidden_size=128, + hc_count=2, + hc_lowrank=16, + num_hidden_layers=1, + layer_types=["linear_attention"], + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=256, + shared_expert_intermediate_size=256, + ple_layer_ids=[] if moe_sp else [1], + ple_embed_dim=128, + heads_per_ngram=2, + ngram_vocab_size_base=32, + make_ngram_vocab_size_divisible_by=64, + ), + dtype=torch.bfloat16, + ) + config.engram_config = EngramConfig(cpu_offload=False) + return config + + +def _make_decoder(config: VllmConfig, rank: int) -> Qwen4ExpDecoderLayer: + """Use the production constructor and identical initial weights for both modes.""" + with ( + set_current_vllm_config(config), + patch.object(qwen4_model, "QwenGatedDeltaNetAttention", _TPAttentionProjection), + ): + layer = Qwen4ExpDecoderLayer( + config, config.model_config.hf_text_config.layer_types[0], "model.layers.0" + ) + torch.manual_seed(12) + for param in layer.parameters(): + param.normal_(std=0.08) + # Distinct rank contributions expose missing or duplicate reductions. + layer.linear_attn.proj.weight.add_(rank * 0.01) + experts = layer.mlp.experts.routed_experts + if config.parallel_config.enable_expert_parallel: + experts.w13_weight.add_(rank * 0.01) + experts.quant_method.process_weights_after_loading(experts) + return layer + + +def _check_decoder_ple(rank: int) -> None: + """Compare HC SP with full tokens across a PLE prefill and cached decode.""" + configs = [_make_config(rank, hc_sp=enabled) for enabled in (False, True)] + layers = [_make_decoder(config, rank) for config in configs] + states = [] + for layer in layers: + ple = layer.ple + assert ple is not None + state = torch.zeros(2, ple.hc_hidden_size, ple.conv_state_len) + if not is_conv_state_dim_first(): + state = state.transpose(-1, -2).contiguous() + ple.kv_cache = (state,) + states.append(state) + + # Three prefill tokens straddle the TP=2 boundary and require SP padding. + for case in ( + _ConvBatchCase(prefill_query_lens=(3,)), + _ConvBatchCase(num_decodes=1), + ): + layout, num_tokens = _make_conv_metadata(case, torch.device("cuda", rank)) + metadata_args = { + f.name: None + for f in fields(PleShortConvAttentionMetadata) + if f.default is MISSING and f.default_factory is MISSING + } + metadata = PleShortConvAttentionMetadata(**(metadata_args | vars(layout))) + ids = torch.randint(0, 128, (num_tokens,)) + positions = torch.arange(num_tokens).expand(3, -1) + qsl = torch.tensor([0, num_tokens], dtype=torch.int32) + context = torch.zeros(1, layers[0].config.ngram_size - 1, dtype=torch.int64) + hidden = torch.randn(num_tokens, 256) + block = torch.randn(num_tokens, 128) + injection = torch.randn(num_tokens, 2) + outputs = [] + for config, layer in zip(configs, layers): + assert layer.ple is not None + inputs = (hidden, block, injection) + if config.parallel_config.enable_hc_sp: + inputs = tuple(sp_shard(x) for x in inputs) + with ( + set_current_vllm_config(config), + set_forward_context({layer.ple.prefix: metadata}, config), + ): + layer.ple.start_prefetch(ids, qsl, context) + output = layer( + *inputs, + positions, + input_ids=ids, + query_start_loc=qsl, + ngram_context=context, + ) + if config.parallel_config.enable_hc_sp: + output = tuple(sp_all_gather(x)[:num_tokens] for x in output) + # A later MoE call may reuse workspace-backed output storage. + outputs.append(tuple(x.clone() for x in output)) + torch.testing.assert_close(outputs[1], outputs[0], atol=0.002, rtol=0.02) + torch.testing.assert_close(states[1], states[0], atol=0.002, rtol=0.02) + + +def _check_moe_sp(rank: int, config: VllmConfig) -> None: + """Compare native SP with full-token GR and the MoE wrapper's public interface.""" + layer = _make_decoder(config, rank) + dp_rank = config.parallel_config.data_parallel_rank + counts = torch.tensor([3, 5], device="cpu") + num_tokens = int(counts[dp_rank]) + torch.manual_seed(30 + dp_rank) + hidden = torch.randn(num_tokens, 256) + positions = torch.arange(num_tokens) + # Include runner padding in addition to the padding introduced by SP. + padding = sp_padding_mask(positions == num_tokens - 1, hidden) + with ( + set_current_vllm_config(config), + set_forward_context( + None, + config, + num_tokens=num_tokens, + num_tokens_across_dp=counts, + is_padding=padding, + ), + ): + # The reference keeps full GR rows. MoE handles its own chunk and gather. + full_hidden, attn_input, injection = layer.attn_hyper_connection.mix(hidden) + attn_output = tensor_model_parallel_all_reduce(layer.linear_attn(attn_input)) + full_hidden, moe_input, injection = layer.mlp_hyper_connection.combine_and_mix( + full_hidden, attn_output, injection + ) + expected = tuple( + x.clone() for x in (full_hidden, layer.mlp(moe_input), injection) + ) + output = layer( + sp_shard(hidden), + None, + None, + positions, + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + actual = tuple(sp_all_gather(x)[:num_tokens] for x in output) + torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.02) + + +def _run_sp(rank: int, port: int, moe_sp: bool) -> None: + """Initialize rank-local runtime state and run one equivalence check.""" + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + config = _make_config(rank, moe_sp=moe_sp) + init_distributed_environment( + world_size=4 if moe_sp else 2, + rank=rank, + local_rank=rank, + distributed_init_method=f"tcp://127.0.0.1:{port}", + ) + try: + # Spawned ranks bypass the GPU worker's workspace initialization. + init_workspace_manager(torch.device(f"cuda:{rank}")) + with set_current_vllm_config(config): + initialize_model_parallel(tensor_model_parallel_size=2) + with torch.device(f"cuda:{rank}"), torch.inference_mode(): + torch.set_default_dtype(torch.bfloat16) + if moe_sp: + _check_moe_sp(rank, config) + else: + _check_decoder_ple(rank) + finally: + reset_workspace_manager() + cleanup_dist_env_and_memory() + + +@multi_gpu_test(num_gpus=2) +def test_hc_sp_decoder_and_ple() -> None: + """HC SP preserves decoder output and PLE cache across a token-shard boundary.""" + mp.spawn(_run_sp, args=(get_open_port(), False), nprocs=2) + + +@multi_gpu_test(num_gpus=4) +def test_moe_sp_equivalence(monkeypatch) -> None: + """Native MoE SP preserves outputs with unequal DP batches and padding.""" + monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") + mp.spawn(_run_sp, args=(get_open_port(), True), nprocs=4) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index d269aeec1330..cffc72804326 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -127,6 +127,22 @@ def test_registry_is_pp(model_arch, is_pp): assert model_info.supports_pp is is_pp +@pytest.mark.parametrize( + "model_arch,supported", + [ + ("Qwen4ExpForCausalLM", not current_platform.is_rocm()), + ("Qwen4ExpForConditionalGeneration", not current_platform.is_rocm()), + ("Qwen4ExpMTP", not current_platform.is_rocm()), + ("LlamaForCausalLM", False), + ], +) +def test_registry_supports_hc_sp(model_arch, supported): + """HC SP is opt-in for each model implementation, including MTP.""" + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_hc_sp is supported + + @create_new_process_for_each_test() @pytest.mark.parametrize( "model_arch", diff --git a/tests/test_config.py b/tests/test_config.py index 734253edc9f8..a7e36ada6259 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4109,6 +4109,47 @@ def test_revision_resolved_for_model(mock_resolve): mock_resolve.assert_any_call(model, None, config.hf_token) +@pytest.mark.parametrize( + "enabled,supported,is_cuda,tp_size,pp_size,error", + [ + (True, True, True, 2, 1, None), + (True, False, True, 2, 1, "not supported for this model"), + (True, True, False, 2, 1, "requires CUDA"), + (True, True, True, 1, 1, None), + (True, False, False, 1, 2, None), + (True, True, True, 2, 2, "requires PP=1"), + (False, False, False, 2, 1, None), + ], +) +def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, error): + """Validate active HC SP configurations and ignore the flag with TP=1.""" + from vllm.model_executor.models.registry import ModelRegistry + + model = SimpleNamespace( + architectures=["HCModel"], + registry=ModelRegistry, + model_arch_config=SimpleNamespace(total_num_attention_heads=8), + multimodal_config=None, + ) + parallel = ParallelConfig( + enable_hc_sp=enabled, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + # Avoid backend auto-detection in this config-only test. + distributed_executor_backend="mp", + ) + with ( + patch("vllm.config.model.current_platform.is_cuda", return_value=is_cuda), + patch.object(ModelRegistry, "is_hc_sp_supported_model", return_value=supported), + patch.object(ModelRegistry, "is_pp_supported_model", return_value=True), + ): + if error: + with pytest.raises(ValueError, match=error): + ModelConfig.verify_with_parallel_config(model, parallel) + else: + ModelConfig.verify_with_parallel_config(model, parallel) + + @pytest.mark.parametrize( ("layer_types", "expected_attention"), [ diff --git a/vllm/config/model.py b/vllm/config/model.py index d4bafad5d6d3..545c306e1e1f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1537,6 +1537,14 @@ def verify_with_parallel_config( self, parallel_config: ParallelConfig, ) -> None: + if parallel_config.enable_hc_sp and parallel_config.tensor_parallel_size > 1: + if not current_platform.is_cuda(): + raise ValueError("--enable-hc-sp currently requires CUDA") + if not self.registry.is_hc_sp_supported_model(self.architectures, self): + raise ValueError("--enable-hc-sp is not supported for this model") + if parallel_config.pipeline_parallel_size != 1: + raise ValueError("--enable-hc-sp requires PP=1") + total_num_attention_heads = self.model_arch_config.total_num_attention_heads tensor_parallel_size = parallel_config.tensor_parallel_size if total_num_attention_heads % tensor_parallel_size != 0: diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6960661e0c7e..e103c9ffd73d 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -181,6 +181,11 @@ class ParallelConfig: --data-parallel-start-rank.""" is_moe_model: bool | None = None """Whether the deployed model is MoE (if known).""" + enable_hc_sp: bool = False + """Enable sequence parallelism for HyperConnection (HC) computation. + Currently supported by Qwen4Exp on CUDA with TP>1 and PP=1 (ignored when TP=1). + Does not change MoE parallelism. MoE SP may enable HC SP independently of this flag. + """ enable_expert_parallel: bool = False """Use expert parallelism instead of tensor parallelism for MoE layers.""" enable_batch_sharded_sampling: bool | None = None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31777e266f8b..a89290081605 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -538,6 +538,7 @@ class EngineArgs: data_parallel_external_lb: bool = False data_parallel_multi_port_external_lb: bool = False data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend + enable_hc_sp: bool = ParallelConfig.enable_hc_sp enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel enable_batch_sharded_sampling: bool | None = ( ParallelConfig.enable_batch_sharded_sampling @@ -1258,6 +1259,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "server per local data parallel rank and exposes aggregated health on " "a supervisor port.", ) + parallel_group.add_argument("--enable-hc-sp", **parallel_kwargs["enable_hc_sp"]) parallel_group.add_argument( "--enable-expert-parallel", "-ep", @@ -2532,6 +2534,7 @@ def create_engine_config( data_parallel_backend=self.data_parallel_backend, data_parallel_hybrid_lb=self.data_parallel_hybrid_lb, is_moe_model=model_config.is_moe, + enable_hc_sp=self.enable_hc_sp, enable_expert_parallel=self.enable_expert_parallel, enable_batch_sharded_sampling=self.enable_batch_sharded_sampling, enable_ep_weight_filter=self.enable_ep_weight_filter, diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index be20c412438c..92e65c839dbf 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -27,6 +27,7 @@ method_has_implemented_embedding, resolve_quant_method, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.layers.utils import dispatch_unquantized_gemm from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs @@ -240,6 +241,7 @@ class VocabParallelEmbedding(PluggableLayer): quant_config: quant config for the layer prefix: full name of the layer in the state dict disable_tp: If true, tensor parallelism will be disabled for this layer. + reduce_results: Whether to all-reduce contributions from vocabulary shards. quant_method: Preselected quantization method for model-specific layers. parallel_group: Process group used to shard and reduce the embedding. @@ -258,6 +260,7 @@ def __init__( prefix: str = "", *, disable_tp: bool = False, + reduce_results: bool = True, quant_method: QuantizeMethodBase | None = None, parallel_group: GroupCoordinator | None = None, ): @@ -266,6 +269,7 @@ def __init__( # Keep the input dimensions. self.disable_tp = disable_tp self.parallel_group = parallel_group + self.reduce_results = reduce_results if disable_tp: tp_rank, self.tp_size = 0, 1 elif parallel_group is not None: @@ -516,6 +520,7 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): param[loaded_weight.shape[0] :].data.fill_(0) def forward(self, input_): + """Look up IDs, optionally reducing contributions from vocabulary shards.""" if self.tp_size == 1: return self.quant_method.embedding(self, input_.long()) @@ -543,24 +548,20 @@ def forward(self, input_): self.shard_indices.added_vocab_end_index, ) output_parallel = self.quant_method.embedding(self, masked_input.long()) - if output_parallel.dtype in ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ): - # Each vocab token has one owner, so FP8 bytes can use int8 SUM. - comm_output = output_parallel.view(torch.int8) - comm_output.masked_fill_(input_mask.unsqueeze(-1), 0) - output = ( - self.parallel_group.all_reduce(comm_output) - if self.parallel_group is not None - else tensor_model_parallel_all_reduce(comm_output) - ) - return output.view(output_parallel.dtype) + + output_dtype = output_parallel.dtype + if is_fp8(output_dtype): + # Each vocab token has one owner, so FP8 bytes can use int8 SUM. + output_parallel = output_parallel.view(torch.int8) + if not self.use_fused_embedding: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) - # Reduce across all the model parallel GPUs. - if self.parallel_group is not None: - return self.parallel_group.all_reduce(output_parallel) - return tensor_model_parallel_all_reduce(output_parallel) + if self.reduce_results: + output_parallel = ( + self.parallel_group.all_reduce(output_parallel) + if self.parallel_group is not None + else tensor_model_parallel_all_reduce(output_parallel) + ) + return output_parallel.view(output_dtype) def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings}" diff --git a/vllm/model_executor/models/__init__.py b/vllm/model_executor/models/__init__.py index 478b00d0af33..78fa39ee413c 100644 --- a/vllm/model_executor/models/__init__.py +++ b/vllm/model_executor/models/__init__.py @@ -3,6 +3,7 @@ from .interfaces import ( HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -10,6 +11,7 @@ SupportsPP, SupportsTranscription, has_inner_state, + supports_hc_sp, supports_lora, supports_mrope, supports_multimodal, @@ -33,6 +35,8 @@ "is_text_generation_model", "HasInnerState", "has_inner_state", + "SupportsHCSP", + "supports_hc_sp", "SupportsLoRA", "supports_lora", "SupportsMultiModal", diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7f6a3aef40a7..d64e51bfa8bf 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -912,11 +912,21 @@ class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConf @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: + from vllm.platforms import current_platform + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) text_config = vllm_config.model_config.hf_text_config if text_config.hc_count <= 1: raise ValueError("Qwen4Exp requires hc_count > 1") parallel_config = vllm_config.parallel_config + # MoE SP also enables HC SP without the explicit HC flag. + if ( + current_platform.is_cuda() + and parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size != 1 + ): + raise ValueError("Qwen4Exp MoE SP requires PP=1") + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( getattr(text_config, "indexer_n_heads", None) is not None ) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 272b78da50a4..4dc67ff34d27 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -858,6 +858,27 @@ def _supports_pp_inspect(model: type[object] | object) -> bool: return supports_kw(model_forward, "intermediate_tensors") +@runtime_checkable +class SupportsHCSP(Protocol): + """Support sequence-parallel HC without changing MoE parallelism.""" + + supports_hc_sp: ClassVar[Literal[True]] = True + + +@overload +def supports_hc_sp(model: type[object]) -> TypeIs[type[SupportsHCSP]]: ... + + +@overload +def supports_hc_sp(model: object) -> TypeIs[SupportsHCSP]: ... + + +def supports_hc_sp( + model: type[object] | object, +) -> TypeIs[type[SupportsHCSP]] | TypeIs[SupportsHCSP]: + return getattr(model, "supports_hc_sp", False) + + @runtime_checkable class HasInnerState(Protocol): """The interface required for all models that has inner state.""" diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 04d20c574145..ee5c05d0a4ae 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -129,7 +129,12 @@ def _should_replicate_misaligned_shared_expert( class Qwen3NextSparseMoeBlock(nn.Module): - def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ): super().__init__() config = vllm_config.model_config.hf_text_config @@ -231,6 +236,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, + reduce_results=reduce_results, fuse_shared_experts=self.is_fused_shared_expert_enabled, shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index fa2924441314..427648045f9d 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -49,6 +49,7 @@ is_attention_free, is_hybrid, requires_raw_input_tokens, + supports_hc_sp, supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, @@ -873,6 +874,7 @@ class _ModelInfo: requires_raw_input_tokens: bool supports_multimodal_encoder_tp_data: bool supports_pp: bool + supports_hc_sp: bool has_inner_state: bool is_attention_free: bool is_hybrid: bool @@ -903,6 +905,7 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": model ), supports_pp=supports_pp(model), + supports_hc_sp=supports_hc_sp(model), has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), @@ -1463,6 +1466,15 @@ def is_pp_supported_model( model_info, _ = self.inspect_model_cls(architectures, model_config) return model_info.supports_pp + def is_hc_sp_supported_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + """Check support for HC SP without changing MoE parallelism.""" + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_hc_sp + def model_has_inner_state( self, architectures: str | list[str], diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 6dad6b7b5afe..85a382c5ec67 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -70,6 +70,7 @@ def __init__( num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: del num_ngram_heads, max_total_tokens super().__init__( @@ -80,6 +81,7 @@ def __init__( prefix=prefix, quant_method=embedding_method, parallel_group=get_etp_group(), + reduce_results=reduce_results, ) self.embedding_method = embedding_method self.data_parallel_rank = data_parallel_rank @@ -153,14 +155,15 @@ def _select_embeddings( return embeddings.narrow(0, slot_offset, local_num_tokens) @abstractmethod - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Start an asynchronous lookup when supported.""" raise NotImplementedError + @abstractmethod + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Return prefetched embeddings for the full input token set.""" + raise NotImplementedError + class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): """Quantization interface shared by resident and pinned PLE tables.""" @@ -351,16 +354,16 @@ def allocate_embedding_weight( """Allocate the complete PLE weight on the active device.""" return torch.empty(num_embeddings, embedding_dim, dtype=dtype) - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Resident embedding prefetch is a no-op.""" return None + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Reject prefetch finalization for resident embeddings.""" + raise RuntimeError("Device PLE embedding does not support prefetch") + def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: - """Gather ETP inputs, look up embeddings, and select local rows.""" + """Gather ETP IDs, optionally reduce lookup results, and select local rows.""" slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) embeddings = super().forward(gathered_ids) @@ -422,6 +425,7 @@ def __init__( num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: if not is_uva_available(): raise RuntimeError("Engram CPU offload requires UVA support") @@ -435,6 +439,7 @@ def __init__( num_ngram_heads=num_ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=reduce_results, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) self._row_bytes = self.embedding_dim * self.weight.element_size() @@ -521,11 +526,7 @@ def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: return self.parallel_group.all_reduce(embeddings) @eager_break_during_capture - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Gather ETP IDs and launch their UVA lookup on the side stream.""" buffer = self._prefetch_buffer if buffer is None: @@ -574,14 +575,18 @@ def _finalize_prefetch( prefetch_output: torch.Tensor, output: torch.Tensor, ) -> None: - """Join the side stream, reduce ETP shards, and select local rows.""" + """Join the side stream, optionally reduce ETP shards, and select rows.""" prefetch_stream = self._prefetch_stream if prefetch_stream is None: raise RuntimeError("pinned PLE finalize requires a prior start_prefetch") torch.cuda.current_stream().wait_stream(prefetch_stream) slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) - active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] - embeddings = self._reduce_etp_embeddings(active_output) + embeddings = prefetch_output[: slot_size * self.etp_data_parallel_size] + if self.reduce_results: + embeddings = self._reduce_etp_embeddings(embeddings) + else: + # The outer TP reduce-scatter requires ETP=TP. + assert self.etp_data_parallel_size == 1 embeddings = self._select_embeddings( embeddings, output.shape[0], @@ -589,11 +594,18 @@ def _finalize_prefetch( ) output.copy_(embeddings.flatten(-2)) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Finish the pinned lookup into graph-owned output storage.""" + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Finalize the full-token pinned lookup with optional ETP reduction.""" buffer = self._prefetch_buffer if buffer is None: raise RuntimeError("pinned PLE lookup requires a prior start_prefetch") - output = buffer.new_empty((hidden_states.shape[0], self._output_dim)) + # total_tokens counts this DP rank's tokens before SP sharding. + output = buffer.new_empty((total_tokens, self._output_dim)) self._finalize_prefetch(buffer, output) return output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + raise RuntimeError( + "Pinned-host PLE embedding requires " + "start_prefetch() and finalize_prefetch()." + ) diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py index bc4c99957d90..41ce082c9e70 100644 --- a/vllm/models/qwen4_exp/common/ops/ple.py +++ b/vllm/models/qwen4_exp/common/ops/ple.py @@ -344,6 +344,7 @@ def _ple_conv_kernel( MODE: tl.constexpr, HAS_INIT: tl.constexpr, NULL_STATE_ID: tl.constexpr, + FUSE_RESIDUAL: tl.constexpr, launch_pdl: tl.constexpr, ): t = tl.program_id(0) @@ -438,21 +439,23 @@ def _ple_conv_kernel( conv = acc.to(residual_ptr.dtype.element_ty).to(tl.float32) y = conv * tl.sigmoid(conv) conv_output = tl.where(out_ok, y, 0.0).to(residual_ptr.dtype.element_ty) - residual = tl.load( - residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - # Preserve the original eager operation boundaries: short convolution is - # first accumulated into the BF16/FP16 PLE output, then the outer residual - # is added to that rounded value. - ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) - outer_residual = tl.load( - outer_residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) + ple_output = conv_output + if FUSE_RESIDUAL: + residual = tl.load( + residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( @@ -575,7 +578,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -585,7 +588,9 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short convolution and the outer residual; update state.""" + """Update state and write activated convolution into residual. + A non-None outer_residual enables both fused residual additions. + """ BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -657,6 +662,7 @@ def ple_conv( SPEC_QUERY_LEN=kernel_spec_query_len, MODE=mode, HAS_INIT=has_initial_states_arg, + FUSE_RESIDUAL=outer_residual is not None, NULL_STATE_ID=NULL_BLOCK_ID, launch_pdl=launch_pdl, num_warps=num_warps, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 2de883cc4747..2fb056bde43c 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -9,8 +9,10 @@ from torch import nn from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -36,6 +38,7 @@ IsHybrid, MixtureOfExperts, MultiModalEmbeddings, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -66,6 +69,12 @@ maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -84,6 +93,14 @@ _QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") +def is_hc_sequence_parallel_enabled(vllm_config: VllmConfig) -> bool: + """Enable GR/PLE SP explicitly or automatically with sequence-parallel MoE.""" + parallel_config = vllm_config.parallel_config + return parallel_config.tensor_parallel_size > 1 and ( + parallel_config.enable_hc_sp or parallel_config.use_sequence_parallel_moe + ) + + def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: @@ -161,13 +178,19 @@ def _remap_qsa_cache_scale_name( class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): """Qwen3Next MoE with Qwen4Exp HC validation.""" - def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: - parallel_config = vllm_config.parallel_config - if parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__( + vllm_config=vllm_config, prefix=prefix, reduce_results=reduce_results + ) + if self.replicate_shared_expert and not self.is_sequence_parallel: + # Reduce routed outputs before adding replicated shared outputs. + # SP then slices the result to avoid summing shared outputs again via RS. + self.experts.moe_config.skip_final_all_reduce = False config = vllm_config.model_config.hf_text_config self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) @@ -185,11 +208,12 @@ def __init__( self.config = config self.layer_type = layer_type + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.layer_idx = extract_layer_index(prefix) - if vllm_config.parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.ple: Qwen4ExpPLELayer | None = None ple_layer_ids = config.ple_layer_ids if (self.layer_idx + 1) in ple_layer_ids: @@ -204,6 +228,7 @@ def __init__( layer_idx=self.layer_idx, ple_dense_layer_id=ple_dense_layer_id, prefix=f"{prefix}.ple", + use_sequence_parallel=self.use_hc_sequence_parallel, ) if layer_type == "linear_attention": @@ -212,6 +237,7 @@ def __init__( vllm_config=vllm_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, + reduce_results=not self.use_hc_sequence_parallel, ) elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None @@ -222,6 +248,7 @@ def __init__( cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: self.self_attn = Qwen4ExpQSAAttention( @@ -230,12 +257,15 @@ def __init__( layer_id=self.layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: raise ValueError(f"Invalid layer_type {layer_type}") self.mlp = Qwen4ExpSparseMoeBlock( - vllm_config=vllm_config, prefix=f"{prefix}.mlp" + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_hc_sequence_parallel, ) hc_config = HyperConnectionConfig( @@ -295,6 +325,9 @@ def forward( else: hidden_states, block_input, injection = attn_hc.mix(hidden_states) + if self.use_hc_sequence_parallel: + # Positions retain the full token axis for both text and MRoPE inputs. + block_input = sp_all_gather(block_input)[: positions.shape[-1]] if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) elif self.layer_type in _QSA_LAYER_TYPES: @@ -306,10 +339,25 @@ def forward( raise ValueError("Invalid layer_type") mlp_hc = self.mlp_hyper_connection + if self.use_hc_sequence_parallel: + attn_out = sp_reduce_scatter(attn_out) hidden_states, block_input, injection = mlp_hc.combine_and_mix( hidden_states, attn_out, injection ) - mlp_out = self.mlp(block_input) + if self.use_sequence_parallel: + # MoE consumes local tokens and returns complete local outputs. + mlp_out = self.mlp(block_input, already_sequence_parallel=True) + else: + if self.use_hc_sequence_parallel: + block_input = sp_all_gather(block_input)[: positions.shape[-1]] + mlp_out = self.mlp(block_input) + if self.use_hc_sequence_parallel: + assert isinstance(self.mlp, Qwen4ExpSparseMoeBlock) + reduce_output = self.mlp.experts.moe_config.skip_final_all_reduce + # Only sum partial outputs. Reduced outputs already include all ranks. + mlp_out = ( + sp_reduce_scatter(mlp_out) if reduce_output else sp_shard(mlp_out) + ) return hidden_states, mlp_out, injection @@ -373,6 +421,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.num_redundant_experts = ( vllm_config.parallel_config.eplb_config.num_redundant_experts ) + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self._qsa_layer_ids = frozenset( layer_idx @@ -447,7 +500,6 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @staticmethod def _start_layer_ple_prefetch( layer: nn.Module, - hidden_states: torch.Tensor, input_ids: torch.Tensor | None, query_start_loc: torch.Tensor | None, ngram_context: torch.Tensor | None, @@ -458,12 +510,7 @@ def _start_layer_ple_prefetch( return if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - ple.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + ple.start_prefetch(input_ids, query_start_loc, ngram_context) def forward( self, @@ -482,19 +529,31 @@ def forward( if input_ids is None: raise ValueError("input_ids or inputs_embeds is required") hidden_states = self.embed_input_ids(input_ids) + if self.use_hc_sequence_parallel: + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # Expand only the local token rows when SP is enabled. hidden_states = hidden_states.repeat(1, self.config.hc_count) else: if intermediate_tensors is None: raise ValueError("pipeline stage requires intermediate tensors") hidden_states = intermediate_tensors["hidden_states"] + full_num_tokens = positions.shape[-1] block_output = None injection = None last_layer = None if self.start_layer < self.end_layer: self._start_layer_ple_prefetch( self.layers[self.start_layer], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -506,7 +565,6 @@ def forward( if layer_idx + 1 < self.end_layer: self._start_layer_ple_prefetch( self.layers[layer_idx + 1], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -526,6 +584,8 @@ def forward( deepstack_embed = deepstack_input_embeds[ f"deepstack_input_embeds_{layer_idx}" ] + if self.use_hc_sequence_parallel: + deepstack_embed = sp_shard(deepstack_embed) deepstack_embed = ( deepstack_embed.unsqueeze(-2) .expand( @@ -560,6 +620,24 @@ def forward( multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( hidden_states, block_output, injection ) + if self.use_hc_sequence_parallel: + if self._mtp_hidden_buffer is not None: + # Gather LM-head and MTP states together when both are needed. + hidden_size = self.config.hidden_size + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[ + :full_num_tokens + ] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, self.config.hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + else: + sample_hidden_states = sp_all_gather(sample_hidden_states)[ + :full_num_tokens + ] if self._mtp_hidden_buffer is not None: # Capture the pre-final-mixer multi-stream hidden state # [T, hc_count*H] for the MTP drafter (zero extra compute: @@ -615,6 +693,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: class Qwen4ExpForCausalLM( nn.Module, HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -865,6 +944,7 @@ def get_hf_config(self) -> Qwen4ExpConfig: class Qwen4ExpForConditionalGeneration( Qwen3_5ForConditionalGeneration, HasInnerState, + SupportsHCSP, Qwen4ExpMixtureOfExperts, ): """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 61290617ae75..32e95faa0dc8 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -20,8 +20,10 @@ from torch import nn from transformers import Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -33,7 +35,7 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import SupportsHCSP, SupportsPP from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -44,6 +46,11 @@ maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from vllm.sequence import IntermediateTensors from .hyperconnection import GatedResidual, HyperConnectionConfig @@ -54,6 +61,7 @@ Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, Qwen4ExpSparseMoeBlock, + is_hc_sequence_parallel_enabled, ) @@ -171,6 +179,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config: Qwen4ExpTextConfig = model_config.hf_text_config self.config = config + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self.mtp_start_layer_idx = config.num_hidden_layers @@ -310,6 +323,22 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] current_step_idx = spec_step_idx % self.num_mtp_layers + full_num_tokens = positions.shape[-1] + if self.use_hc_sequence_parallel: + # SP requires PP=1, so the first-rank branch sets prev_block_output. + assert prev_block_output is not None + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # The FC output is already complete and needs only a slice. + prev_block_output = sp_shard(prev_block_output) layer = self.layers[current_step_idx] hidden_states, block_output, injection = layer( hidden_states=hidden_states, @@ -337,6 +366,17 @@ def forward( hidden_states, block_output, injection ) ) + if self.use_hc_sequence_parallel: + # Gather both outputs together, then restore their contiguous layouts. + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[:full_num_tokens] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + multi_hidden = multi_hidden.contiguous() return sample_hidden_states, multi_hidden def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -357,7 +397,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loader.load_weights(weights, mapper=mapper) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(nn.Module, SupportsHCSP, SupportsPP, Qwen4ExpMixtureOfExperts): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index cf03cbb40d47..794eea87a620 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -9,11 +9,14 @@ from transformers import Qwen4ExpTextConfig from vllm.config import get_current_vllm_config +from vllm.distributed import get_etp_group, get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, @@ -146,8 +149,10 @@ def __init__( prefix: str, quant_config: QuantizationConfig | None = None, params_dtype: torch.dtype | None = None, + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel self.embedding_dim = embedding_dim self.ngram_size = int(config.ngram_size) self.heads_per_ngram = int(config.heads_per_ngram) @@ -211,6 +216,11 @@ def __init__( if engram_config is not None and engram_config.cpu_offload else Qwen4ExpPLEDeviceEmbedding ) + # A TP reduce-scatter covers all vocabulary owners only when ETP=TP. + self.use_reduce_scatter = ( + self.use_sequence_parallel + and get_etp_group().world_size == get_tp_group().world_size + ) self.ngram_embedding = embedding_cls( padded_vocab_size, self.head_dim, @@ -221,6 +231,7 @@ def __init__( num_ngram_heads=self.ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=not self.use_reduce_scatter, ) if self.ngram_embedding.supports_prefetch: # The side-stream lookup outlives eager-break args, whose @@ -348,6 +359,18 @@ def compute_ngram_ids( id_blocks.append(ids[request_indices, adjusted_columns]) return torch.cat(id_blocks, dim=-1) + def _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: + """Shard embedding token rows, optionally reducing vocabulary contributions.""" + # Each embedding element has one vocabulary owner, including FP8 bytes. + output_dtype = embeddings.dtype + if is_fp8(embeddings): + embeddings = embeddings.view(torch.int8) + if self.use_reduce_scatter: + embeddings = sp_reduce_scatter(embeddings) + else: + embeddings = sp_shard(embeddings) + return embeddings.view(output_dtype) + def forward( self, hidden_states: torch.Tensor, @@ -355,15 +378,22 @@ def forward( query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: + """Return full embeddings or a token shard after combining vocabulary owners.""" embedding = self.ngram_embedding if embedding.supports_prefetch: - return embedding(hidden_states) - ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) - return self.ngram_embedding(ngram_ids).flatten(-2) + # Prefetch contains full tokens even when hidden states are sharded. + embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0]) + else: + ngram_ids = self.compute_ngram_ids( + input_ids, query_start_loc, ngram_context + ) + embeddings = embedding(ngram_ids).flatten(-2) + if not self.use_sequence_parallel: + return embeddings + return self._sp_shard_embeddings(embeddings) def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, @@ -378,7 +408,7 @@ def start_prefetch( ngram_context, output=self._prefetch_ids[: input_ids.numel()], ) - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" @@ -444,8 +474,3 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if regular_weights: loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) return loaded - - -__all__ = [ - "Qwen4ExpNGramEmbedding", -] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index bdb16b3adef0..a595dba18d47 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -18,6 +18,7 @@ MambaStateShapeCalculator, is_conv_state_dim_first, ) +from vllm.models.common.ops.sequence_parallel import sp_all_gather, sp_shard from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionBackend, @@ -69,8 +70,10 @@ def __init__( layer_idx: int = 0, ple_dense_layer_id: int | None = None, prefix: str = "", + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -100,8 +103,9 @@ def __init__( prefix=f"{prefix}.ple_embedding", quant_config=quant_config, params_dtype=model_config.dtype, + use_sequence_parallel=self.use_sequence_parallel, ) - # The PLE cache is TP-replicated, so this merged projection is too. + # Projection weights stay replicated even when their token rows are sharded. self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), [self.hc_hidden_size, self.hidden_size], @@ -151,18 +155,12 @@ def _dequantize_embeddings( def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> None: """Start the pinned PLE lookup while the preceding decoder layer runs.""" - self.ple_embedding.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + self.ple_embedding.start_prefetch(input_ids, query_start_loc, ngram_context) @property def mamba_type(self) -> MambaAttentionBackendEnum: @@ -192,7 +190,7 @@ def _short_conv_dilated_dispatch( self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -207,7 +205,8 @@ def _short_conv_dilated_dispatch( has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] - outer_residual = outer_residual[: metadata.num_actual_tokens] + if outer_residual is not None: + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -259,11 +258,14 @@ def _short_conv_dilated_dispatch( residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) - outer_residual_d, outer_residual_p = torch.split( - outer_residual, - [num_decode_tokens, num_prefill_tokens], - dim=0, - ) + outer_residual_d: torch.Tensor | None = None + outer_residual_p: torch.Tensor | None = None + if outer_residual is not None: + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: @@ -336,14 +338,16 @@ def _short_conv( self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, ) -> None: + """Update convolution state, optionally fusing both residual additions.""" forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata # Profiling omits all metadata or this Mamba entry. Short convolution # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -354,7 +358,8 @@ def _short_conv( layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -396,7 +401,10 @@ def forward( ngram_context: torch.Tensor, ) -> torch.Tensor: input_ids = input_ids.reshape(-1) - if input_ids.shape[0] != hidden_states.shape[0]: + if ( + not self.use_sequence_parallel + and input_ids.shape[0] != hidden_states.shape[0] + ): raise ValueError( "PLE expects input_ids and hidden_states to have the same " f"token length, got {input_ids.shape[0]} and " @@ -408,6 +416,7 @@ def forward( query_start_loc, ngram_context, ) + assert embeddings.shape[0] == hidden_states.shape[0] embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) @@ -420,9 +429,16 @@ def forward( self.norm_conv.weight, self.norm_key.eps, ) - # State routing depends on runtime request metadata and remains outside - # the graph; short convolution accumulates into gated_output. - self._short_conv(conv_input, gated_output, hidden_states) + if self.use_sequence_parallel: + # Only convolution needs full tokens. Add residuals on local shards. + conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]] + # Metadata may omit padding tokens or all tokens during profiling. + conv_output = torch.zeros_like(conv_input) + self._short_conv(conv_input, conv_output, outer_residual=None) + gated_output.add_(sp_shard(conv_output)) + gated_output.add_(hidden_states) + else: + self._short_conv(conv_input, gated_output, hidden_states) return gated_output From 7900d2eabc4f9a7facb7a196a04c2c32803a79c8 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Mon, 12 Oct 2026 08:18:44 -0500 Subject: [PATCH 409/952] [Test] Accept tied tokens at the k-th value in the Triton top-k test (#60726) Signed-off-by: Rohan138 Co-authored-by: Claude Opus 5.5 --- tests/v1/sample/test_topk_topp_sampler.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index c0b979245659..190a4e3fcaf0 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -251,8 +251,17 @@ def _compare_results( triton_kept = (result_triton != float("-inf")).sum(dim=-1) if p is None: - # Top-k only: expect exact match - assert torch.equal(pytorch_kept, triton_kept), ( + # Top-k only: Triton keeps exactly k (index-ordered tie-break at + # the k-th value) while the sort-based reference keeps all ties, + # so PyTorch may only keep extra tokens tied with the k-th value. + pytorch_mask = result_pytorch != float("-inf") + triton_mask = result_triton != float("-inf") + kth = torch.where(triton_mask, logits.float(), float("inf")) + kth = kth.amin(dim=-1, keepdim=True).expand_as(logits) + extra = pytorch_mask & ~triton_mask + assert not (triton_mask & ~pytorch_mask).any() and torch.equal( + logits.float()[extra], kth[extra] + ), ( f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, " f"Triton kept {triton_kept.tolist()}" ) From 8c175f403acd2749086c298cd1c6f29ed9b4e2ba Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Mon, 12 Oct 2026 15:20:35 +0200 Subject: [PATCH 410/952] [CI][ROCm] Compare greedy output only in the AITER attn+quant test (#61185) Signed-off-by: Stefan Koncarevic --- tests/compile/correctness_e2e/test_attn_quant.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/compile/correctness_e2e/test_attn_quant.py b/tests/compile/correctness_e2e/test_attn_quant.py index e13a8f264a6a..1d8f73da4b9a 100644 --- a/tests/compile/correctness_e2e/test_attn_quant.py +++ b/tests/compile/correctness_e2e/test_attn_quant.py @@ -77,6 +77,10 @@ def settings(enabled): env1=env, env2=env, method="generate", + # Fusing output quant drops a bf16 round-trip, so the two settings + # agree on the argmax but not on the exact distribution. Greedy + # decoding survives that, sampling at temperature 1.0 does not. + include_seeded_sampling=False, ) matches = [ int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text) From ac28877f0e86452f2e151279bdc373cd58ae88f0 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Mon, 12 Oct 2026 21:21:35 +0800 Subject: [PATCH 411/952] [Bugfix] Fix out-of-bounds lookahead in adaptive block verification (#60945) Signed-off-by: chaunceyjiang --- vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index b2efa6432e9f..4387b1d38ab8 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -425,6 +425,9 @@ def _compute_local_residual_mass_kernel( PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) + # Adaptive verification can end before num_speculative_steps. + if logit_idx + 1 >= tl.num_programs(0): + return draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps: # The acceptance threshold, h, looks one position ahead and sums From 314ab8cffd9b115198d50491c5f1c4a2112ad4a1 Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:25:36 +0900 Subject: [PATCH 412/952] [Bugfix][Structured Output] Drain compilation tasks before releasing the backend (#61047) Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Co-authored-by: Artem Perevedentsev --- .../test_structured_output_manager.py | 85 ++++++++++++++++++- vllm/v1/structured_output/__init__.py | 9 ++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 67c15825580c..1d1e7b78df04 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -2,12 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass +from threading import Event import pytest from transformers import AutoTokenizer -from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config import ( + DeviceConfig, + SchedulerConfig, + StructuredOutputsConfig, + VllmConfig, +) from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig from vllm.config.structured_outputs import StructuredOutputsBackend @@ -576,3 +583,79 @@ def test_outlines_choice_with_non_bmp_characters(tokenizer): grammar = backend.compile_grammar(request_type, grammar_spec) assert grammar.accept_tokens("", tokenizer.encode(choices[0])) + + +def test_clear_backend_waits_for_running_grammar_and_cancels_pending( + tokenizer, monkeypatch +): + """Keep the compiler alive until running work finishes during shutdown.""" + manager, request = _build_harness(tokenizer, "xgrammar", use_reasoner=False) + manager.executor.shutdown(wait=True) + manager.executor = ThreadPoolExecutor(max_workers=1) + started, release = Event(), Event() + assert manager.backend is not None + compile_grammar = manager.backend.compile_grammar + + def delayed_compile(*args, **kwargs): + started.set() + assert release.wait(MAX_WAIT_SECONDS) + return compile_grammar(*args, **kwargs) + + monkeypatch.setattr(manager.backend, "compile_grammar", delayed_compile) + try: + manager.grammar_init(request) + assert request.structured_output_request is not None + running = request.structured_output_request._grammar + assert isinstance(running, Future) + assert started.wait(MAX_WAIT_SECONDS) + pending_request = Request( + request_id="pending-at-shutdown", + prompt_token_ids=[], + sampling_params=request.sampling_params, + pooling_params=None, + ) + manager.grammar_init(pending_request) + assert pending_request.structured_output_request is not None + pending = pending_request.structured_output_request._grammar + assert isinstance(pending, Future) + # Cancellation releases the running task without timing-based sleeps. + pending.add_done_callback(lambda _: release.set()) + manager.clear_backend() + release.set() + assert running.result(timeout=MAX_WAIT_SECONDS) is not None + assert pending.cancelled() + assert manager.backend is None + manager.clear_backend() + finally: + release.set() + manager.executor.shutdown(wait=True, cancel_futures=True) + + +@pytest.mark.parametrize("skip_tokenizer_init", [False, True]) +@pytest.mark.parametrize("max_num_seqs", [1, 256]) +def test_clear_backend_without_grammar(skip_tokenizer_init, max_num_seqs): + """Stop initialized thread pools even if no request needed a backend.""" + manager = StructuredOutputManager( + VllmConfig( + model_config=ModelConfig( + tokenizer=TOKENIZER, skip_tokenizer_init=skip_tokenizer_init + ), + scheduler_config=SchedulerConfig( + max_num_seqs=max_num_seqs, + max_model_len=128, + is_encoder_decoder=False, + ), + device_config=DeviceConfig(device="cpu"), + ) + ) + manager.clear_backend() + manager.clear_backend() + assert manager.backend is None + executors = [] + if not skip_tokenizer_init: + executors.append(manager.executor) + if max_num_seqs > manager.fill_bitmask_parallel_threshold: + executors.append(manager.executor_for_fillmask) + for executor in executors: + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + executor.submit(int, 0) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 2163a800a563..3f461ed47f75 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -482,5 +482,14 @@ def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: return grammar.accept_tokens(request.request_id, new_token_ids[grammar_start:]) def clear_backend(self) -> None: + # Finish tasks that can still use the backend before releasing it. + if not self.vllm_config.model_config.skip_tokenizer_init: + self.executor.shutdown(wait=True, cancel_futures=True) + if ( + self.fill_bitmask_parallel_threshold + < self.vllm_config.scheduler_config.max_num_seqs + ): + self.executor_for_fillmask.shutdown(wait=True, cancel_futures=True) if self.backend is not None: self.backend.destroy() + self.backend = None From 092c35dd155db2644da6bb671332c51902bd3039 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 12 Oct 2026 21:27:34 +0800 Subject: [PATCH 413/952] [Model] Support Cloudflare/clef-omni on /v1/systemone (#61064) Signed-off-by: zjy0516 Signed-off-by: Isotr0py Co-authored-by: Claude Opus 5.5 Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 27 ++ .../structured_decisions/test_strategies.py | 345 +++++++++++++- .../test_filter_duplicate_safetensors.py | 29 ++ .../multimodal/pooling/test_clef_omni.py | 168 +++++++ .../multimodal/processing/test_qwen3_omni.py | 30 ++ tests/models/registry.py | 4 + tests/multimodal/test_video.py | 16 + vllm/entrypoints/generate/api_router.py | 39 -- .../structured_decisions/api_router.py | 39 ++ .../generate/structured_decisions/protocol.py | 9 + .../generate/structured_decisions/serving.py | 23 +- .../structured_decisions/strategies.py | 362 ++++++++++++++- .../launchers/api_server/app_state.py | 13 + .../launchers/api_server/routers.py | 6 + .../model_loader/default_loader.py | 8 +- vllm/model_executor/models/clef_omni.py | 437 ++++++++++++++++++ .../models/qwen3_omni_moe_thinker.py | 11 +- vllm/model_executor/models/registry.py | 1 + vllm/multimodal/video.py | 64 +++ .../model_arch_config_convertor.py | 30 ++ 20 files changed, 1592 insertions(+), 69 deletions(-) create mode 100644 tests/models/multimodal/pooling/test_clef_omni.py create mode 100644 vllm/model_executor/models/clef_omni.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index a9ed5a2ee5aa..4d8d65d89e94 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,32 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Omni + +[Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: +a joint schema head on the Qwen3-Omni thinker scores every option of every +question in one forward pass. vLLM serves it as a pooling model. + +```bash +vllm serve Cloudflare/clef-omni +``` + +- A request is one read. The prompt is the one the checkpoint's + `joint_schema_model.py` builds: the state, then every question with its + options. A JSON state is written as compact JSON with sorted keys. +- A state too long for the model is truncated to fit, as the reference does. + A request whose questions alone exceed 64,000 tokens, or `max_model_len`, + returns 400. +- `instructions` returns 400, since the prompt is fixed. `chat_template_kwargs` + has no effect. +- There are no labels, so `label_mass` is 1 and `confidence` is the chosen + option's probability. `usage.input_tokens` counts the prompt once, and + `output_tokens` is 0. +- `images`, `audio` and `videos` are read as the reference does: placed + before the state, videos sampled at 2 fps, and a video's soundtrack heard + with it when every video has one. Audio clips cannot be sent together with + videos that have soundtracks. + ## Request fields | field | meaning | @@ -139,6 +165,7 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni)) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 41b9a8f01271..8d38e61e6248 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -1,12 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from io import BytesIO from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock +import numpy as np +import pybase64 as base64 import pytest +import torch from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from PIL import Image from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient @@ -23,22 +29,32 @@ ) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, - state_text, + parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( DiffusionGemmaCanvasStrategy, + JointSchemaStrategy, NextTokenStrategy, ReadContext, reply_label_ids, select_read_strategy, + state_text, ) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.exception_handling.register import ( init_exception_handler, ) +from vllm.inputs import mm_input from vllm.logprobs import Logprob -from vllm.outputs import CompletionOutput, RequestOutput +from vllm.multimodal.media import MediaConnector, MediaWithBytes +from vllm.outputs import ( + CompletionOutput, + PoolingOutput, + PoolingRequestOutput, + RequestOutput, +) from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -53,17 +69,26 @@ def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: def test_strategy_selection(): - qwen = "Qwen3ForCausalLM" - assert select_read_strategy(model(qwen)) is NextTokenStrategy - assert select_read_strategy(model(qwen, "processed_logprobs")) is NextTokenStrategy + qwen, generate = "Qwen3ForCausalLM", ("generate",) + assert select_read_strategy(model(qwen), generate) is NextTokenStrategy assert ( - select_read_strategy(model("DiffusionGemmaForBlockDiffusion")) + select_read_strategy(model(qwen, "processed_logprobs"), generate) + is NextTokenStrategy + ) + assert ( + select_read_strategy(model("DiffusionGemmaForBlockDiffusion"), generate) is DiffusionGemmaCanvasStrategy ) with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): - select_read_strategy(model("LlamaForCausalLM")) + select_read_strategy(model("LlamaForCausalLM"), generate) with pytest.raises(ValueError, match="not raw_logits"): - select_read_strategy(model(qwen, "raw_logits")) + select_read_strategy(model(qwen, "raw_logits"), generate) + # A generative architecture served as a pooling model has no read strategy. + with pytest.raises(ValueError, match="need the generate task"): + select_read_strategy(model(qwen), ("embed",)) + # Clef-Omni scores options with its head, so it reads no logprobs. + clef = model("ClefOmniForJointSchema", "raw_logits") + assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy def test_route_is_registered_by_default(): @@ -126,7 +151,7 @@ def test_labels_start_the_reply(qwen): @pytest.fixture -def serving_and_engine(qwen): +def engine_and_renderer(qwen): tokenizer, _ = qwen model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32") renderer = HfRenderer( @@ -148,6 +173,19 @@ def serving_and_engine(qwen): chat_template=None, chat_template_content_format="auto", ) + yield engine, online_renderer + renderer.shutdown() + + +def serving(engine, strategy) -> ServingStructuredDecisions: + name = engine.model_config.model + models = OpenAIServingModels(engine, [BaseModelPath(name=name, model_path=name)]) + return ServingStructuredDecisions(models, strategy) + + +@pytest.fixture +def serving_and_engine(engine_and_renderer): + engine, online_renderer = engine_and_renderer async def generate(prompt, params, request_id, **kwargs): label_ids = params.logprob_token_ids @@ -171,11 +209,7 @@ async def generate(prompt, params, request_id, **kwargs): engine.generate.side_effect = generate strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {})) - models = OpenAIServingModels( - engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)] - ) - yield ServingStructuredDecisions(models, strategy), engine - renderer.shutdown() + return serving(engine, strategy), engine @pytest.mark.asyncio @@ -256,3 +290,286 @@ def test_canvas_read(): strategy._read_input( {"type": "token", "prompt_token_ids": [1, 2, 3, 4]}, [1, 2, 3, 4] ) + + +JOINT_SCHEMA_PROMPT = """<|im_start|>system +Read the complete state and schema. Decide every field jointly. Each answer \ +must be exactly one of that field's allowed options.<|im_end|> +<|im_start|>user +STATE: +{"customer":{"tier":"gold"},"ticket":"Charged twice"} + +SCHEMA FIELDS: + +FIELD 1 +ID: team +TYPE: choice +INSTRUCTION: Which team? +ALLOWED OPTIONS: +OPTION 1: {"description":"payments","option_id":"billing"} +OPTION 2: {"option_id":"security"} +OPTION 3: {"description":"deliveries","option_id":"shipping"} +END FIELD + +FIELD 2 +ID: urgent +TYPE: noul +INSTRUCTION: urgent +ALLOWED OPTIONS: +OPTION 1: {"description":"money is moving","option_id":"true"} +OPTION 2: {"description":"The proposition is false or the answer is no.",\ +"option_id":"false"} +END FIELD + +FIELD 3 +ID: tone +TYPE: score +INSTRUCTION: How upset? +ALLOWED OPTIONS: +OPTION 1: {"description":"calm","option_id":"0"} +OPTION 2: {"description":"annoyed","option_id":"1"} +OPTION 3: {"description":"furious","option_id":"2"} +END FIELD + +<|im_end|> +<|im_start|>assistant + + + + +JOINT SCHEMA DECISIONS:""" + +JOINT_SCHEMA_REQUEST: dict[str, Any] = dict( + state={"ticket": "Charged twice", "customer": {"tier": "gold"}}, + questions={ + "team": { + "type": "choice", + "instructions": "Which team?", + "criteria": { + "shipping": "deliveries", + "billing": "payments", + "security": None, + }, + }, + "urgent": {"type": "noul", "criteria": {"true": "money is moving"}}, + "tone": { + "type": "score", + "instructions": "How upset?", + "criteria": ["calm", "annoyed", "furious"], + }, + }, +) + + +@pytest.fixture +def joint_schema(qwen, engine_and_renderer): + tokenizer, _ = qwen + engine, online_renderer = engine_and_renderer + # The probability the head gives each option, by option id. + scores = {"billing": 0.7, "security": 0.1, "shipping": 0.2} + scores |= {"true": 0.9, "false": 0.1, "0": 0.2, "1": 0.5, "2": 0.3} + + async def encode(prompt, params, request_id, **kwargs): + ids = prompt["prompt_token_ids"] + spans = params.extra_kwargs["joint_schema"]["option_spans"] + options = [json.loads(tokenizer.decode(ids[s:e])) for s, e in spans] + data = torch.tensor([scores[o["option_id"]] for o in options]) + yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) + + engine.encode.side_effect = encode + strategy = JointSchemaStrategy( + ReadContext(engine, online_renderer, None, "auto", {}) + ) + return serving(engine, strategy), engine, tokenizer + + +@pytest.mark.asyncio +async def test_joint_schema_reads_every_question_at_once(joint_schema): + serving, engine, tokenizer = joint_schema + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, cache_salt="salt") + ) + assert isinstance(response, StructuredDecisionResponse) + + engine.encode.assert_called_once() + engine.check_admission.assert_called_once_with(1) + prompt, params = engine.encode.call_args.args[:2] + ids = prompt["prompt_token_ids"] + assert tokenizer.decode(ids) == JOINT_SCHEMA_PROMPT + assert prompt["cache_salt"] == "salt" + assert params.task == "token_classify" + schema = params.extra_kwargs["joint_schema"] + assert schema["question_types"] == [1, 0, 2] + assert schema["option_counts"] == [3, 2, 3] + assert [tokenizer.decode(ids[s:e]) for s, e in schema["question_spans"]] == [ + "Which team?", + "urgent", + "How upset?", + ] + + # Options are read in the model's order and answered in the request's. + team = response.answers["team"] + assert team["choice"] == "billing" + assert list(team["probabilities"]) == ["shipping", "billing", "security"] + assert team["probabilities"] == pytest.approx( + {"shipping": 0.2, "billing": 0.7, "security": 0.1} + ) + assert response.answers["urgent"]["noul"] == pytest.approx(0.9) + assert response.answers["tone"]["score"] == pytest.approx(1.1) + assert response.usage.input_tokens == len(ids) + assert response.usage.output_tokens == 0 + assert all(d.label_mass == 1.0 for d in response.diagnostics.values()) + + +@pytest.mark.asyncio +async def test_joint_schema_rejects_instructions(joint_schema): + serving, engine, _ = joint_schema + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, instructions="Be strict.") + ) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + engine.encode.assert_not_called() + + +def test_joint_schema_truncates_the_state(joint_schema): + serving, _, _ = joint_schema + strategy = serving.strategy + request = StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST) + questions = parse_questions(request, strategy.limits()) + prompt, schema, _ = strategy.encode(questions, "word " * 100) + strategy.max_length = len(prompt) - 50 + truncated, truncated_schema, _ = strategy.encode(questions, "word " * 100) + # The state loses 50 tokens and the questions stay at the end. + assert len(truncated) == strategy.max_length + assert truncated[-200:] == prompt[-200:] + assert truncated_schema == schema + strategy.max_length = 100 + with pytest.raises(StructuredDecisionError, match="questions take"): + strategy.encode(questions, "word") + + +@pytest.mark.asyncio +async def test_media_needs_a_model_that_takes_it(serving_and_engine): + serving, engine = serving_and_engine + request = StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, images=["data:image/png;base64,AAAA"] + ) + response = await serving.create_decision(request) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + assert "does not take image inputs" in response.error.message + engine.generate.assert_not_called() + + +@pytest.mark.asyncio +async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): + serving, engine, tokenizer = joint_schema + strategy = serving.strategy + strategy.__dict__["_connector"] = MediaConnector() + image_pad = tokenizer.convert_tokens_to_ids("<|image_pad|>") + rendered: dict[str, Any] = {} + + async def process_for_engine_async(prompt, arrival_time): + rendered.update(prompt) + # The image's placeholder expands to three tokens. + ids = [ + token + for pad in prompt["prompt_token_ids"] + for token in ([pad] * 3 if pad == image_pad else [pad]) + ] + return mm_input(ids, {}, {}, {}, cache_salt=prompt.get("cache_salt")) + + renderer = strategy.context.online_renderer.renderer + monkeypatch.setattr(renderer, "process_for_engine_async", process_for_engine_async) + buffer = BytesIO() + Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") + image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + response = await serving.create_decision( + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + ) + assert isinstance(response, StructuredDecisionResponse) + + placeholder = tokenizer.encode( + "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False + ) + assert rendered["prompt_token_ids"] == strategy.prefix + placeholder + assert rendered["multi_modal_data"]["image"][0].size == (8, 8) + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } + prompt = engine.encode.call_args.args[0] + assert prompt["type"] == "multimodal" + ids = prompt["prompt_token_ids"] + expanded = len(strategy.prefix) + len(placeholder) + 2 + assert ids[:expanded].count(image_pad) == 3 + state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] + assert tokenizer.decode(ids[expanded:]) == state_and_questions + assert response.answers["team"]["choice"] == "billing" + assert response.usage.input_tokens == len(ids) + + +@pytest.mark.asyncio +async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): + serving, engine, tokenizer = joint_schema + strategy = serving.strategy + video = MediaWithBytes( + (np.zeros((4, 8, 8, 3), np.uint8), {"fps": 2.0}), b"video-bytes" + ) + track = (np.zeros(160, np.float32), 16000) + rendered: dict[str, Any] = {} + + class FakeConnector: + async def fetch_image_async(self, url): + raise AssertionError("no images in this test") + + async def load_from_url_async(self, url, media_io, *, fetch_timeout=None): + return video + + async def fetch_audio_async(self, url): + if url == "https://example.org/silent.mp4": + raise ValueError("No audio stream found.") + return track + + strategy.__dict__["_connector"] = FakeConnector() + + async def process_for_engine_async(prompt, arrival_time): + rendered.clear() + rendered.update(prompt) + return mm_input(prompt["prompt_token_ids"], {}, {}, {}) + + renderer = strategy.context.online_renderer.renderer + monkeypatch.setattr(renderer, "process_for_engine_async", process_for_engine_async) + + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, videos=["https://example.org/talking.mp4"] + ) + ) + assert isinstance(response, StructuredDecisionResponse) + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } + + # A video without a decodable soundtrack is watched silently. + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, videos=["https://example.org/silent.mp4"] + ) + ) + assert isinstance(response, StructuredDecisionResponse) + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + + # Audio clips cannot be read with videos that have soundtracks. + response = await serving.create_decision( + StructuredDecisionRequest( + **JOINT_SCHEMA_REQUEST, + videos=["https://example.org/talking.mp4"], + audio=["https://example.org/clip.wav"], + ) + ) + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 diff --git a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py index da04d0108fdf..7aca4b64ab62 100644 --- a/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py +++ b/tests/model_executor/model_loader/test_filter_duplicate_safetensors.py @@ -137,3 +137,32 @@ def is_unused_checkpoint_weight(self, name: str) -> bool: ["model-00002-of-00003.safetensors", "model-00003-of-00003.safetensors"], sorted(f.name for f in mtp_checkpoint.glob("*.safetensors")), ] + + +def test_secondary_source_with_explicit_patterns_skips_backbone_index(tmp_path): + """A source with explicit allow patterns (e.g. a standalone head file in + the repo root) must not be validated against the backbone's index.""" + save_file( + {"backbone.weight": torch.zeros(1)}, + tmp_path / "model-00001-of-00001.safetensors", + ) + save_file({"head.weight": torch.zeros(1)}, tmp_path / "joint_head.safetensors") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": {"backbone.weight": "model-00001-of-00001.safetensors"}} + ) + ) + + loader = DefaultModelLoader(LoadConfig()) + secondary = DefaultModelLoader.Source( + model_or_path=str(tmp_path), + revision=None, + prefix="pooler.head.", + allow_patterns_overrides=["joint_head.safetensors"], + ) + assert [n for n, _ in loader._get_weights_iterator(secondary)] == [ + "pooler.head.head.weight" + ] + # The primary source still filters files the backbone index doesn't list. + primary = DefaultModelLoader.Source(model_or_path=str(tmp_path), revision=None) + assert [n for n, _ in loader._get_weights_iterator(primary)] == ["backbone.weight"] diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py new file mode 100644 index 000000000000..ccdc913b3132 --- /dev/null +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +which loops over questions where the pooler batches them.""" + +import math +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.pooling_params import PoolingParams +from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates + +HIDDEN, VOCAB = 64, 500 +TYPES = [0, 1, 1, 2] +COUNTS = [2, 3, 1, 5] + + +def reference_logits(head, hidden, token_ids, embedding, question_spans, spans): + """``JointSchemaHead.forward`` of the checkpoint's joint_schema_model.py + for one record, with ``spans`` the option spans of each question.""" + hidden = head.hidden_norm(hidden) + memory = head.memory_projection(hidden).unsqueeze(0) + global_vector = hidden[-1] + questions = torch.stack([hidden[s:e].mean(0) for s, e in question_spans]) + contexts = [torch.stack([hidden[s:e].mean(0) for s, e in q]) for q in spans] + lexicals = [ + torch.stack([embedding[token_ids[s:e]].mean(0) for s, e in q]) for q in spans + ] + options = torch.cat( + [ + head.option_context_projection(c) + + head.option_lexical_projection(lex) + + head.option_question_projection(questions[i]).unsqueeze(0) + for i, (c, lex) in enumerate(zip(contexts, lexicals)) + ] + ).unsqueeze(0) + for layer in head.evidence_layers: + options = layer(options, memory) + routed = torch.split(options[0], [len(q) for q in spans]) + + base = head.question_projection(questions) + summaries = [] + for field, opts in zip(base, routed): + weights = torch.softmax(opts @ field / math.sqrt(opts.shape[-1]), dim=0) + summaries.append((weights[:, None] * opts).sum(0)) + fields = ( + base + + head.option_summary_norm(torch.stack(summaries)) + + head.global_projection(global_vector) + + head.type_embedding(torch.tensor(TYPES)) + ).unsqueeze(0) + for layer in head.layers: + fields = layer(fields, memory) + fields = head.field_norm(fields[0]) + + logits = [] + for i, (field, lexical, opts) in enumerate(zip(fields, lexicals, routed)): + anchor = F.normalize(questions[i] + global_vector, dim=-1) + prior_scale = head.prior_logit_scale.clamp(max=math.log(100.0)).exp() + prior = prior_scale * (F.normalize(lexical, dim=-1) @ anchor) + opts = head.option_norm(opts) + field = field.expand_as(opts) + features = torch.cat([field, opts, field * opts, (field - opts).abs()], -1) + joint_scale = head.joint_logit_scale.clamp(max=math.log(100.0)).exp() + joint = joint_scale * F.cosine_similarity(field, opts, dim=-1) + joint = joint + head.residual_scorer(features).squeeze(-1) + logits.append(prior + torch.sigmoid(head.residual_gate) * joint) + return logits + + +@pytest.fixture +def pooler(): + torch.manual_seed(0) + head = JointSchemaHead( + hidden_size=HIDDEN, + width=32, + routing_layers=2, + layers=2, + heads=4, + feedforward=48, + ).eval() + with torch.no_grad(): + for p in head.parameters(): + p.add_(torch.randn_like(p) * 0.2) + embedding = torch.randn(VOCAB, HIDDEN) + config = SimpleNamespace( + scheduler_config=SimpleNamespace( + enable_chunked_prefill=False, async_scheduling=False + ) + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=config, + ): + pooler = JointSchemaPooler(head, lambda ids: embedding[ids]) + return pooler, embedding + + +def schema(length: int): + """Questions after a 20-token state, with options of 1 to 4 tokens.""" + pos = 20 + question_spans: list[tuple[int, int]] = [] + spans: list[list[tuple[int, int]]] = [] + for count in COUNTS: + question_spans.append((pos, pos + 4)) + pos += 6 + spans.append([]) + for k in range(count): + n = 1 + (k * 3) % 4 + spans[-1].append((pos, pos + n)) + pos += n + 1 + assert pos < length + extra = { + "question_types": TYPES, + "question_spans": [[s - length, e - length] for s, e in question_spans], + "option_spans": [[s - length, e - length] for q in spans for s, e in q], + "option_counts": COUNTS, + } + return {"joint_schema": extra}, question_spans, spans + + +@torch.inference_mode() +def test_pooler_matches_reference(pooler): + pooler, embedding = pooler + lengths = [120, 90, 100] + token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths] + hidden = torch.randn(sum(lengths), HIDDEN) + first, first_questions, first_spans = schema(lengths[0]) + last, last_questions, last_spans = schema(lengths[2]) + params = [ + PoolingParams(task="token_classify", use_activation=True, extra_kwargs=first), + # A prompt without a schema, like the profiling run's. + PoolingParams(task="token_classify", use_activation=True), + PoolingParams(task="token_classify", use_activation=False, extra_kwargs=last), + ] + padded = torch.zeros(len(lengths), max(lengths), dtype=torch.long) + for i, ids in enumerate(token_ids): + padded[i, : len(ids)] = ids + metadata = PoolingMetadata( + prompt_lens=torch.tensor(lengths), + prompt_token_ids=None, + prompt_token_ids_cpu=padded, + pooling_params=params, + pooling_states=[PoolingStates() for _ in lengths], + ) + metadata.build_pooling_cursor( + torch.tensor(lengths).numpy(), + seq_lens_cpu=torch.tensor(lengths), + device=torch.device("cpu"), + ) + probs, empty, logits = pooler(hidden, metadata) + + first_hidden, _, last_hidden = torch.split(hidden, lengths) + expected = reference_logits( + pooler.head, first_hidden, token_ids[0], embedding, first_questions, first_spans + ) + expected_probs = torch.cat([q.softmax(-1) for q in expected]) + torch.testing.assert_close(probs, expected_probs, atol=1e-5, rtol=1e-4) + assert empty.numel() == 0 + expected = reference_logits( + pooler.head, last_hidden, token_ids[2], embedding, last_questions, last_spans + ) + torch.testing.assert_close(logits, torch.cat(expected), atol=1e-5, rtol=1e-4) diff --git a/tests/models/multimodal/processing/test_qwen3_omni.py b/tests/models/multimodal/processing/test_qwen3_omni.py index 301471f45a7b..c0ebf0297184 100644 --- a/tests/models/multimodal/processing/test_qwen3_omni.py +++ b/tests/models/multimodal/processing/test_qwen3_omni.py @@ -108,3 +108,33 @@ def get_token_count(duration: float) -> int: f"Expected longer audio (2s) to have more tokens than shorter (1s). " f"Got short={short_tokens}, long={long_tokens}" ) + + +@pytest.mark.parametrize("model_id", ["Cloudflare/clef-omni"]) +def test_audio_in_video_follows_requested_fps(model_id: str) -> None: + """Audio and video tokens interleave by the fps the video was processed + at, as in HF, not by the default of one frame a second.""" + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + rng = np.random.RandomState(42) + video = rng.randint(0, 255, (4, 64, 64, 3), dtype=np.uint8) + audio = rng.rand(32000).astype(np.float32) + prompt = "<|vision_start|><|video_pad|><|vision_end|>" + mm_kwargs: dict[str, Any] = { + "use_audio_in_video": True, + "videos_kwargs": {"fps": 2.0, "do_sample_frames": False}, + } + processed = processor( + prompt, + mm_items=processor.info.parse_mm_data( + {"video": [video], "audio": [(audio, 16000)]} + ), + hf_processor_mm_kwargs=mm_kwargs, + ) + expected = processor.info.get_hf_processor()( + text=[prompt], videos=[video], audio=[audio], return_tensors="pt", **mm_kwargs + ) + assert processed["prompt_token_ids"] == expected["input_ids"][0].tolist() diff --git a/tests/models/registry.py b/tests/models/registry.py index 2ed6d4871c27..0a7027461ef3 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1351,6 +1351,10 @@ def check_available_online( "Qwen/Qwen3-Omni-30B-A3B-Instruct", max_model_len=4096, ), + "ClefOmniForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-omni", + max_model_len=4096, + ), "Qwen3ASRForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-ASR-0.6B", max_model_len=4096, diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 407bffed20a5..0066d239eae2 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -21,6 +21,7 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, + ClefOmniVideoBackend, DynamicVideoBackend, Glm5NextVideoBackend, GLM46VVideoBackend, @@ -1813,3 +1814,18 @@ def test_glm5next_read_frames_dense_walk_matches_stock(tmp_path): assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1 # One initial seek, then pure walking -- no re-seek churn. assert cap.seeks == 1 + + +def test_clef_omni_backend_samples_reference_frames(tmp_path): + """Clef-Omni samples the first frame at or after each 0.5 s boundary + (the reference's 2 fps timestamp sampling), not uniform strides.""" + pytest.importorskip("cv2") + + # 2 s at 30 fps: the reference reads frames 0, 15, 30 and 45, where + # uniform sampling would read 0, 19, 39 and 59. + path = _write_gray_video(tmp_path, total_frames=60, fps=30) + frames, metadata = ClefOmniVideoBackend.load_bytes(path.read_bytes()) + assert metadata["frames_indices"] == [0, 15, 30, 45] + assert metadata["do_sample_frames"] is False + for frame, idx in zip(frames, [0, 15, 30, 45]): + assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1 diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 6bf8d690265a..9bbe4e879b34 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -5,7 +5,6 @@ from fastapi import FastAPI import vllm.envs as envs -from vllm.logger import init_logger if TYPE_CHECKING: from argparse import Namespace @@ -18,8 +17,6 @@ else: RequestLogger = object -logger = init_logger(__name__) - def register_generate_api_routers(app: FastAPI): from vllm.entrypoints.openai.chat_completion.api_router import ( @@ -56,12 +53,6 @@ def register_generate_api_routers(app: FastAPI): register_generative_scoring_api_router(app) - from .structured_decisions.api_router import ( - register_structured_decisions_api_router, - ) - - register_structured_decisions_api_router(app) - async def init_generate_state( engine_client: "EngineClient", @@ -231,33 +222,3 @@ async def init_generate_state( state.openai_serving_models, request_logger=request_logger, ) - - from .structured_decisions.serving import ServingStructuredDecisions - from .structured_decisions.strategies import ReadContext, select_read_strategy - - strategy = None - if "generate" in supported_tasks: - try: - strategy_cls = select_read_strategy(engine_client.model_config) - strategy = strategy_cls( - ReadContext( - engine_client=engine_client, - online_renderer=state.online_renderer, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - default_chat_template_kwargs=default_chat_template_kwargs, - ) - ) - except ValueError as e: - # Info, since every model without a read strategy lands here. - logger.info("/v1/systemone is disabled: %s", e) - strategy = None - state.serving_structured_decisions = ( - ServingStructuredDecisions( - state.openai_serving_models, - strategy, - request_logger=request_logger, - ) - if strategy is not None - else None - ) diff --git a/vllm/entrypoints/generate/structured_decisions/api_router.py b/vllm/entrypoints/generate/structured_decisions/api_router.py index 5729b58b2e50..a3afe7fe1120 100644 --- a/vllm/entrypoints/generate/structured_decisions/api_router.py +++ b/vllm/entrypoints/generate/structured_decisions/api_router.py @@ -1,19 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace from http import HTTPStatus +from typing import Any from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse +from starlette.datastructures import State +from vllm.engine.protocol import EngineClient from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, validate_json_request, with_cancellation, ) +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.logger import init_logger +from vllm.renderers.chat_utils import load_chat_template +from vllm.tasks import SupportedTask from .protocol import StructuredDecisionRequest from .serving import ServingStructuredDecisions +from .strategies import ReadContext, select_read_strategy + +logger = init_logger(__name__) router = APIRouter() @@ -50,3 +61,31 @@ async def create_structured_decision( def register_structured_decisions_api_router(app: FastAPI): app.include_router(router) + + +def init_structured_decisions_state( + engine_client: EngineClient, + state: State, + args: Namespace, + request_logger: RequestLogger | None, + supported_tasks: tuple[SupportedTask, ...], + default_chat_template_kwargs: dict[str, Any], +): + state.serving_structured_decisions = None + try: + strategy_cls = select_read_strategy(engine_client.model_config, supported_tasks) + strategy = strategy_cls( + ReadContext( + engine_client=engine_client, + online_renderer=state.online_renderer, + chat_template=load_chat_template(args.chat_template), + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=default_chat_template_kwargs, + ) + ) + except ValueError as e: + logger.debug("/v1/systemone is disabled: %s", e) + return + state.serving_structured_decisions = ServingStructuredDecisions( + state.openai_serving_models, strategy, request_logger=request_logger + ) diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index ab89084dff2d..a8f104b8cb09 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -30,6 +30,15 @@ class StructuredDecisionRequest(OpenAIBaseModel): instructions: str | None = Field( default=None, description="Context placed ahead of the questions." ) + images: list[str] | None = Field( + default=None, description="Image URLs or data URLs, for models that see." + ) + audio: list[str] | None = Field( + default=None, description="Audio URLs or data URLs, for models that hear." + ) + videos: list[str] | None = Field( + default=None, description="Video URLs or data URLs, for models that see." + ) chat_template_kwargs: dict[str, Any] | None = None cache_salt: str | None = Field( default=None, diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 2633647a0b65..52fdd430b116 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import json from typing import Any from fastapi import Request @@ -24,10 +23,6 @@ logger = init_logger(__name__) -def state_text(state: Any) -> str: - return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) - - def parse_questions( request: StructuredDecisionRequest, limits: DecisionLimits, @@ -72,6 +67,18 @@ def __init__( self.strategy = strategy self.limits = strategy.limits() + def _media(self, request: StructuredDecisionRequest) -> dict[str, list[str]]: + urls = { + "image": request.images, + "audio": request.audio, + "video": request.videos, + } + media = {modality: items for modality, items in urls.items() if items} + if unsupported := media.keys() - self.strategy.modalities: + names = ", ".join(sorted(unsupported)) + raise StructuredDecisionError(f"this model does not take {names} inputs") + return media + async def create_decision( self, request: StructuredDecisionRequest, @@ -87,8 +94,9 @@ async def create_decision( request_id = f"decision-{base_id}" try: questions = parse_questions(request, self.limits) + media = self._media(request) lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] - engine_client.check_admission(len(questions)) + engine_client.check_admission(self.strategy.num_requests(questions)) trace_headers = ( None if raw_request is None @@ -97,13 +105,14 @@ async def create_decision( reads = await self.strategy.read( questions, request.instructions, - state_text(request.state), + self.strategy.render_state(request.state), request_id=request_id, chat_template_kwargs=request.chat_template_kwargs, lora_request=lora_request, priority=request.priority, cache_salt=request.cache_salt, trace_headers=trace_headers, + media=media, ) except StructuredDecisionError as e: return self.create_error_response(e) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 48debdb0b30f..f4365924ddf0 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -5,24 +5,40 @@ Only models in READ_STRATEGIES are currently supported. """ +import asyncio import hashlib import json import math import random +import time from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any +from functools import cached_property +from typing import Any, ClassVar, TypeVar +import numpy as np + +import vllm.envs as envs from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.label_reads import next_token_label_reads -from vllm.inputs import EngineInput, tokens_input +from vllm.inputs import EngineInput, MultiModalInput, TokensPrompt, tokens_input +from vllm.logger import init_logger from vllm.lora.request import LoRARequest +from vllm.multimodal.media import ( + MEDIA_CONNECTOR_REGISTRY, + ImageMediaIO, + MediaConnector, + MediaIO, + VideoMediaIO, +) +from vllm.pooling_params import PoolingParams from vllm.renderers.chat_utils import ChatTemplateContentFormatOption from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import SamplingParams +from vllm.tasks import SupportedTask from vllm.tokenizers import TokenizerLike from .protocol import ReadPromptRequest @@ -34,6 +50,10 @@ label_softmax, ) +_M = TypeVar("_M") + +logger = init_logger(__name__) + @dataclass class ReadContext: @@ -59,13 +79,30 @@ class QuestionRead: output_tokens: int +def state_text(state: Any) -> str: + return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) + + class ReadStrategy(ABC): + #: The engine task reads run as. + task: ClassVar[SupportedTask] = "generate" + #: The modalities a request may attach media of. + modalities: ClassVar[frozenset[str]] = frozenset() + def __init__(self, context: ReadContext): self.context = context @abstractmethod def limits(self) -> DecisionLimits: ... + def render_state(self, state: Any) -> str: + """The request's state as the model reads it.""" + return state_text(state) + + def num_requests(self, questions: list[Question]) -> int: + """The engine requests a decision takes.""" + return len(questions) + @abstractmethod async def read( self, @@ -79,6 +116,7 @@ async def read( priority: int, cache_salt: str | None = None, trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, ) -> list[QuestionRead]: """One read per question, in order.""" @@ -214,6 +252,7 @@ async def read( priority: int, cache_salt: str | None = None, trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, ) -> list[QuestionRead]: ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) @@ -334,7 +373,315 @@ def _read_input( ) +class JointSchemaStrategy(ReadStrategy): + """Clef-Omni. The whole request is one pooling read: the state, then every + question with its options, as the model was trained to read them. The + model's joint schema head scores each option, so a read has no labels. + + Mirrors ``encode_record`` in the checkpoint's ``joint_schema_model.py``. + Each piece is tokenized on its own, so the spans of the questions and the + options are known. The spans count from the end of the prompt, so they + hold after the media placeholders, which come before the state, expand.""" + + task = "token_classify" + modalities = frozenset({"image", "audio", "video"}) + + SYSTEM_PROMPT = ( + "Read the complete state and schema. Decide every field jointly. Each " + "answer must be exactly one of that field's allowed options." + ) + #: The reference's default bound on a prompt, in tokens. + MAX_LENGTH = 64000 + TYPE_IDS = {"noul": 0, "choice": 1, "score": 2} + NOUL_CRITERIA = { + "true": "The proposition is true or the answer is yes.", + "false": "The proposition is false or the answer is no.", + } + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>\n", + "audio": "<|audio_start|><|audio_pad|><|audio_end|>\n", + "video": "<|vision_start|><|video_pad|><|vision_end|>\n", + } + IMAGE_SIZE = {"shortest_edge": 65_536, "longest_edge": 1_048_576} + #: The reference samples videos at 2 fps; the "clef_omni" video backend + #: reproduces the reference's sampling, pixel cap and frame count cap. + VIDEO_FPS = 2 + + def __init__(self, context: ReadContext): + super().__init__(context) + self.tokenizer = context.online_renderer.renderer.get_tokenizer() + self.max_length = min( + self.MAX_LENGTH, context.engine_client.model_config.max_model_len + ) + self.prefix = self._tokens( + f"<|im_start|>system\n{self.SYSTEM_PROMPT}<|im_end|>\n" + "<|im_start|>user\nSTATE:\n" + ) + self.suffix = self._tokens( + "\n<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + "JOINT SCHEMA DECISIONS:" + ) + + def limits(self) -> DecisionLimits: + return DecisionLimits(max_questions=64, max_options=len(LABELS)) + + def render_state(self, state: Any) -> str: + return self._render(state) + + def num_requests(self, questions: list[Question]) -> int: + return 1 + + @staticmethod + def _render(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps( + value, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ) + + def _tokens(self, text: str) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=False) + + @cached_property + def _connector(self) -> MediaConnector: + model_config = self.context.engine_client.model_config + return MEDIA_CONNECTOR_REGISTRY.load( + envs.VLLM_MEDIA_CONNECTOR, + media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) + + async def _fetch( + self, urls: Sequence[str], media_io: MediaIO[_M], timeout: int + ) -> list[_M]: + return await asyncio.gather( + *( + self._connector.load_from_url_async( + url, media_io, fetch_timeout=timeout + ) + for url in urls + ) + ) + + async def _soundtrack( + self, video_url: str + ) -> tuple[np.ndarray, int | float] | None: + """The video's soundtrack, or None when it has no decodable one.""" + try: + return await self._connector.fetch_audio_async(video_url) + except ValueError as e: + logger.debug("no decodable soundtrack in %s: %s", video_url, e) + return None + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + """The system prompt and the media, with the media's placeholders + expanded. A video's soundtrack is heard with it when every video has + one.""" + connector = self._connector + video_io = VideoMediaIO( + ImageMediaIO(), + num_frames=-1, + fps=self.VIDEO_FPS, + video_backend="clef_omni", + ) + try: + images, audio, videos, tracks = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + asyncio.gather( + *map(connector.fetch_audio_async, media.get("audio", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + asyncio.gather(*map(self._soundtrack, media.get("video", ()))), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + + audio_in_video = bool(videos) and all(t is not None for t in tracks) + if audio_in_video and audio: + raise StructuredDecisionError( + "audio clips cannot be read with videos that have soundtracks" + ) + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["audio"] * len(audio) + + self.PLACEHOLDERS["video"] * len(videos) + ) + if audio_in_video: + audio = [track for track in tracks if track is not None] + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if audio: + mm_data["audio"] = audio + if videos: + mm_data["video"] = videos + mm_kwargs: dict[str, Any] = {"images_kwargs": {"size": self.IMAGE_SIZE}} + if videos: + mm_kwargs["use_audio_in_video"] = audio_in_video + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs=mm_kwargs, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + renderer = self.context.online_renderer.renderer + try: + engine_input = await renderer.process_for_engine_async(prompt, time.time()) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + def _options(self, q: Question) -> list[tuple[int, str, Any]]: + """Each option's index in ``q.options``, id and description, in the + model's order: choices sorted by name, noul true then false, score + levels by index.""" + if q.type.name == "noul": + return [ + ( + i, + key, + self.NOUL_CRITERIA[key] if o.description is None else o.description, + ) + for i, (key, o) in enumerate(zip(("true", "false"), q.options)) + ] + if q.type.name == "choice": + return sorted( + ((i, o.name, o.description) for i, o in enumerate(q.options)), + key=lambda option: option[1], + ) + if q.type.name == "score": + return [(i, str(i), o.name) for i, o in enumerate(q.options)] + raise StructuredDecisionError( + f"question {q.id!r}: this model does not answer {q.type.name!r} questions" + ) + + def encode( + self, questions: list[Question], state: str, prefix: list[int] | None = None + ) -> tuple[list[int], dict[str, list[Any]], list[list[int]]]: + """The prompt, the schema for the pooler, and each question's option + indices in the order of the output. ``prefix`` is what precedes the + state: the system prompt, and the media if there are any.""" + prefix = self.prefix if prefix is None else prefix + schema_ids = self._tokens("\n\nSCHEMA FIELDS:\n") + question_spans, option_spans, orders = [], [], [] + for n, q in enumerate(questions, start=1): + schema_ids += self._tokens( + f"\nFIELD {n}\nID: {q.id}\nTYPE: {q.type.name}\nINSTRUCTION: " + ) + start = len(schema_ids) + schema_ids += self._tokens(q.instructions or q.id) + question_spans.append([start, len(schema_ids)]) + schema_ids += self._tokens("\nALLOWED OPTIONS:\n") + options = self._options(q) + for m, (_, option_id, description) in enumerate(options, start=1): + schema_ids += self._tokens(f"OPTION {m}: ") + start = len(schema_ids) + semantics = {"option_id": option_id} + if description is not None: + semantics["description"] = description + schema_ids += self._tokens(self._render(semantics)) + option_spans.append([start, len(schema_ids)]) + schema_ids += self._tokens("\n") + schema_ids += self._tokens("END FIELD\n") + orders.append([i for i, _, _ in options]) + + fixed = len(prefix) + len(schema_ids) + len(self.suffix) + if fixed > self.max_length: + raise StructuredDecisionError( + f"the questions take {fixed} tokens, more than the " + f"{self.max_length} this model reads" + ) + state_ids = self._tokens(state)[: self.max_length - fixed] + prompt = prefix + state_ids + schema_ids + self.suffix + # The schema ends len(self.suffix) tokens before the prompt does. + end = len(schema_ids) + len(self.suffix) + schema: dict[str, list[Any]] = { + "question_types": [self.TYPE_IDS[q.type.name] for q in questions], + "question_spans": [[s - end, e - end] for s, e in question_spans], + "option_spans": [[s - end, e - end] for s, e in option_spans], + "option_counts": [len(order) for order in orders], + } + return prompt, schema, orders + + async def read( + self, + questions: list[Question], + instructions: str | None, + state: str, + *, + request_id: str, + chat_template_kwargs: dict[str, Any] | None, + lora_request: LoRARequest | None, + priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, + media: Mapping[str, Sequence[str]] | None = None, + ) -> list[QuestionRead]: + if instructions: + raise StructuredDecisionError( + "this model reads a fixed prompt, so it takes no instructions" + ) + engine_input: EngineInput + if media: + engine_input = await self._render_media(media, cache_salt) + prompt, schema, orders = self.encode( + questions, state, engine_input["prompt_token_ids"] + ) + engine_input["prompt_token_ids"] = prompt + else: + prompt, schema, orders = self.encode(questions, state) + engine_input = tokens_input(prompt, cache_salt=cache_salt) + result = None + async for result in self.context.engine_client.encode( + engine_input, + PoolingParams(task="token_classify", extra_kwargs={"joint_schema": schema}), + f"{request_id}-0", + lora_request=lora_request, + trace_headers=trace_headers, + priority=priority, + ): + pass + if result is None: + raise ValueError("Failed to score the questions") + if result.error is not None: + raise ValueError(f"Failed to score the questions: {result.error.message}") + probs = result.outputs.data.tolist() + if len(probs) != sum(schema["option_counts"]): + raise ValueError( + f"Expected {sum(schema['option_counts'])} option scores, " + f"got {len(probs)}" + ) + + reads, offset = [], 0 + for n, order in enumerate(orders): + q_probs = [0.0] * len(order) + for i, p in zip(order, probs[offset : offset + len(order)]): + q_probs[i] = p + offset += len(order) + reads.append( + QuestionRead( + probs=q_probs, + label_mass=1.0, + argmax_is_label=True, + # One forward pass reads the prompt for every question. + input_tokens=len(prompt) if n == 0 else 0, + output_tokens=0, + ) + ) + return reads + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefOmniForJointSchema": JointSchemaStrategy, "Qwen3ForCausalLM": NextTokenStrategy, "Qwen3_5ForConditionalGeneration": NextTokenStrategy, "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, @@ -347,7 +694,9 @@ def _read_input( LOGPROBS_MODES = frozenset({"raw_logprobs", "processed_logprobs"}) -def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: +def select_read_strategy( + model_config: ModelConfig, supported_tasks: Sequence[SupportedTask] +) -> type[ReadStrategy]: """Raise ValueError when the model cannot serve structured decisions.""" strategy = READ_STRATEGIES.get(model_config.architecture) if strategy is None: @@ -356,7 +705,12 @@ def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]: f"{model_config.architecture}. Supported architectures: " f"{sorted(READ_STRATEGIES)}" ) - if model_config.logprobs_mode not in LOGPROBS_MODES: + if strategy.task not in supported_tasks: + raise ValueError( + f"Structured decisions on {model_config.architecture} need the " + f"{strategy.task} task" + ) + if strategy.task == "generate" and model_config.logprobs_mode not in LOGPROBS_MODES: raise ValueError( "Structured decisions need --logprobs-mode raw_logprobs or " f"processed_logprobs, not {model_config.logprobs_mode}" diff --git a/vllm/entrypoints/launchers/api_server/app_state.py b/vllm/entrypoints/launchers/api_server/app_state.py index ce68890717e3..ff87dbea4915 100644 --- a/vllm/entrypoints/launchers/api_server/app_state.py +++ b/vllm/entrypoints/launchers/api_server/app_state.py @@ -157,6 +157,19 @@ async def init_app_state( init_pooling_state(engine_client, state, args, request_logger, supported_tasks) + from vllm.entrypoints.generate.structured_decisions.api_router import ( + init_structured_decisions_state, + ) + + init_structured_decisions_state( + engine_client, + state, + args, + request_logger, + supported_tasks, + default_chat_template_kwargs, + ) + await init_endpoint_plugins_state(engine_client, state, args) state.enable_server_load_tracking = args.enable_server_load_tracking diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py index 40d04862afb2..52c995a30135 100644 --- a/vllm/entrypoints/launchers/api_server/routers.py +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -25,6 +25,12 @@ def register_api_routers( register_models_api_router(app) + from vllm.entrypoints.generate.structured_decisions.api_router import ( + register_structured_decisions_api_router, + ) + + register_structured_decisions_api_router(app) + from vllm.entrypoints.serve.sagemaker.api_router import ( attach_router as register_sagemaker_api_router, ) diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index b162f60b0577..6bedabfdc013 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -247,9 +247,11 @@ def _prepare_weights( subfolder=subfolder, revision=revision, ) - hf_weights_files = filter_duplicate_safetensors_files( - hf_weights_files, hf_folder, index_file - ) + # only filter the weights from index without explicit allow patterns + if allow_patterns_overrides is None: + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file + ) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py new file mode 100644 index 000000000000..7cd0d676873b --- /dev/null +++ b/vllm/model_executor/models/clef_omni.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). + +The Qwen3-Omni thinker with a joint schema head that scores every option of +every question in one forward pass. The head reads the final hidden states of +the whole prompt, so it runs as a ``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Mapping, Set +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + +from vllm.config import VllmConfig +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +from .interfaces_base import default_pooling_type +from .qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerDummyInputsBuilder, + Qwen3OmniMoeThinkerForConditionalGeneration, + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen3OmniMoeThinkerProcessingInfo, +) + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids + + +MAX_FRAME_PIXELS = 256 * 32 * 32 + + +def cap_frame_pixels(frames: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """The reference reads each video frame at at most 256 * 32 * 32 pixels, + in even sizes. Frames are NHWC uint8, numpy or tensor; tensors stay on + their device (a CUDA tensor from torchcodec must not round-trip).""" + height, width = frames.shape[1:3] + if height * width <= MAX_FRAME_PIXELS: + return frames + scale = math.sqrt(MAX_FRAME_PIXELS / (width * height)) + new_width = max(2, round(width * scale / 2) * 2) + new_height = max(2, round(height * scale / 2) * 2) + is_tensor = isinstance(frames, torch.Tensor) + tensor = frames if is_tensor else torch.from_numpy(frames) + resized = F.interpolate( + tensor.permute(0, 3, 1, 2).float(), + size=(new_height, new_width), + mode="bilinear", + ) + if tensor.dtype == torch.uint8: + resized = resized.round() + resized = resized.to(tensor.dtype).permute(0, 2, 3, 1).contiguous() + return resized if is_tensor else resized.numpy() + + +def _cap_video(item: Any) -> Any: + if isinstance(item, tuple): + frames, metadata = item + return cap_frame_pixels(frames), metadata + return cap_frame_pixels(item) + + +class ClefOmniMultiModalProcessor(Qwen3OmniMoeThinkerMultiModalProcessor): + """Caps video frames at the reference's per-frame pixel count.""" + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + if videos := hf_data.get("videos"): + assert isinstance(videos, list) + hf_data = {**hf_data, "videos": [_cap_video(v) for v in videos]} + return super()._call_hf_processor(hf_data, hf_kwargs) + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + ClefOmniMultiModalProcessor, + info=Qwen3OmniMoeThinkerProcessingInfo, + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder, +) +class ClefOmniForJointSchema(Qwen3OmniMoeThinkerForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + head_config = model_config.hf_config.joint_head_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**head_config), self._embed_lexical + ) + # The head's weights ship as their own safetensors file. + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + # The head embeds options with the output embeddings. ParallelLMHead + # forbids lookups, so use the vocab-parallel embedding's. + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 4ae80a4a36e8..321b5c228973 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -83,6 +83,7 @@ PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.utils import copy_mm_embedding_modality from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -1498,8 +1499,12 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int): audio_in_video_item_idx += 1 + # Match the fps the processor used, as HF and the mrope positions do. + processed_second_per_grid_ts = out_mm_data.get("second_per_grid_ts") second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) - if second_per_grid_ts: + if isinstance(processed_second_per_grid_ts, torch.Tensor): + video_second_per_grid_t = processed_second_per_grid_ts[item_idx].item() + elif second_per_grid_ts: video_second_per_grid_t = second_per_grid_ts[item_idx] else: video_second_per_grid_t = 2.0 @@ -1904,7 +1909,9 @@ def embed_input_ids( embeddings_main, embeddings_multiscale = torch.split( embeddings, [visual_dim, multi_dim], dim=-1 ) - multimodal_embeddings[index] = embeddings_main + multimodal_embeddings[index] = copy_mm_embedding_modality( + embeddings, embeddings_main + ) multimodal_embeddings_multiscale.append(embeddings_multiscale) if not is_interleaved: current_positions = mm_positions[ diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 427648045f9d..36f9bd6e6fcc 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", "ModernBertForTokenClassification", diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index fc0cdc28d478..75a6c1beb1cf 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -1570,3 +1570,67 @@ def load_bytes( valid_frame_indices=valid_frame_indices, ) return frames, metadata + + +@VIDEO_LOADER_REGISTRY.register("clef_omni") +class ClefOmniVideoBackend(VideoBackend): + """Clef-Omni's reference sampling: the first frame at or after each + ``1/fps`` boundary (frame i sits at ``i / original_fps`` seconds, so the + k-th boundary is crossed by frame ``ceil(k * original_fps / fps)``), + capped at 768 frames and padded to an even count by repeating the last + frame.""" + + MAX_FRAMES = 768 + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + if source.original_fps <= 0: + raise ValueError("video has an unknown frame rate") + step = source.original_fps / target.fps + indices = [] + k = 0 + while (index := math.ceil(k * step - 1e-6)) < source.total_frames_num: + indices.append(index) + if len(indices) > cls.MAX_FRAMES: + raise ValueError( + f"video is longer than {cls.MAX_FRAMES / target.fps:.0f} seconds" + ) + k += 1 + if not indices: + raise ValueError("video has no decodable frames") + return indices + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 2, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: VideoDecoderBackend = "opencv", + **kwargs: Any, + ) -> tuple[DecodedFrames, dict[str, Any]]: + frames, metadata = super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + if len(frames) % 2: + last = frames[-1:] + if isinstance(frames, torch.Tensor): + frames = torch.cat([frames, last]) + else: + frames = np.concatenate([frames, last]) + metadata["frames_indices"].append(metadata["frames_indices"][-1]) + return frames, metadata diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index d79c1275b6e7..147b65759dea 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -17,6 +17,10 @@ ConfigFormat, get_safetensors_params_metadata, ) +from vllm.transformers_utils.repo_utils import ( + file_or_path_exists, + get_hf_file_to_dict, +) from vllm.utils.torch_utils import common_broadcastable_dtype logger = init_logger(__name__) @@ -643,6 +647,31 @@ def __init__( _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + # Clef-Omni keeps Qwen3-Omni's config and adds a joint schema head; + # route to its wrapper class when the head's config file is present. + # Mutate hf_config (like the deepseek_v4/mimo_v2 rewrites above) + # because model-class resolution reads hf_config.architectures. + if getattr(hf_config, "architectures", None) == [ + "Qwen3OmniMoeForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefOmniForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -822,6 +851,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, + "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, "qwen4_exp_mtp": Qwen4ExpMTPModelArchConfigConvertor, "RefinedWeb": FalconModelArchConfigConvertor, From 3facf822b771cc0a47a346249d6ea42f9071aaac Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:45:54 +0100 Subject: [PATCH 414/952] [Performance] Load HF tokenizers and processors from a local snapshot (#60825) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- tests/tokenizers_/test_basic.py | 16 +++++ tests/transformers_utils/test_processor.py | 28 ++++++++ vllm/tokenizers/registry.py | 15 ++++- vllm/transformers_utils/processor.py | 17 ++++- vllm/transformers_utils/repo_utils.py | 74 ++++++++++++++++++++++ 5 files changed, 146 insertions(+), 4 deletions(-) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index 4a1a2b156050..a46be988e5c8 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import _get_protocol_attrs # type: ignore +from unittest.mock import patch import pytest from transformers import ( PreTrainedTokenizerBase, TokenizersBackend, + tokenization_utils_base, ) from vllm.tokenizers import TokenizerLike, get_tokenizer @@ -56,6 +58,20 @@ def test_tokenizer_revision(tokenizer_name: str): get_tokenizer(tokenizer_name, revision="never") +def test_hub_tokenizer_loads_from_local_snapshot(): + """Transformers makes a Hub request per file when given a repo ID.""" + with patch.object( + tokenization_utils_base, + "list_repo_templates", + wraps=tokenization_utils_base.list_repo_templates, + ) as list_repo_templates: + tokenizer = get_tokenizer("openai-community/gpt2") + + assert isinstance(tokenizer, PreTrainedTokenizerBase) + assert tokenizer.name_or_path == "openai-community/gpt2" + list_repo_templates.assert_not_called() + + @pytest.mark.parametrize("tokenizer_name", ["BAAI/bge-base-en"]) @pytest.mark.parametrize("n_tokens", [510]) def test_special_tokens(tokenizer_name: str, n_tokens: int): diff --git a/tests/transformers_utils/test_processor.py b/tests/transformers_utils/test_processor.py index a3a1c7841865..93f742ff743f 100644 --- a/tests/transformers_utils/test_processor.py +++ b/tests/transformers_utils/test_processor.py @@ -2,14 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib +from unittest.mock import patch +from transformers import processing_utils from transformers.processing_utils import ProcessingKwargs from typing_extensions import Unpack +from vllm.tokenizers import get_tokenizer from vllm.transformers_utils.processor import ( + get_processor, get_processor_kwargs_keys, get_processor_kwargs_type, ) +from vllm.transformers_utils.repo_utils import _snapshot_non_weight_files class _FakeProcessorKwargs(ProcessingKwargs, total=False): # type: ignore @@ -65,3 +70,26 @@ def test_get_processor_kwargs_from_processor_module_scan_returns_full_union(): proc = _ProcWithoutUnpack() keys = get_processor_kwargs_keys(get_processor_kwargs_type(proc)) _assert_has_all_expected(keys) + + +def test_hub_processor_loads_from_local_snapshot(): + """Transformers makes a Hub request per file when given a repo ID.""" + with patch.object( + processing_utils, + "list_repo_templates", + wraps=processing_utils.list_repo_templates, + ) as list_repo_templates: + processor = get_processor("Qwen/Qwen2-VL-2B-Instruct") + + assert processor.tokenizer.name_or_path == "Qwen/Qwen2-VL-2B-Instruct" + assert processor.chat_template is not None + list_repo_templates.assert_not_called() + + +def test_processor_reuses_tokenizer_snapshot(): + """The chat template lookup must not snapshot the repo a second time.""" + model = "Qwen/Qwen2-VL-2B-Instruct" + _snapshot_non_weight_files.cache_clear() + get_tokenizer(model) + get_processor(model) + assert _snapshot_non_weight_files.cache_info().misses == 1 diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index ebf85a11fd25..a7adf4597f6a 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, assert_never +from typing import TYPE_CHECKING, assert_never, cast import huggingface_hub from typing_extensions import TypeVar @@ -14,11 +14,12 @@ from vllm.transformers_utils.config import _maybe_register_hf_config, get_config from vllm.transformers_utils.repo_utils import ( any_pattern_in_repo_files, + get_non_weight_snapshot_path, is_mistral_model_repo, ) from vllm.utils.import_utils import resolve_obj_by_qualname -from .hf import CachedHfTokenizer +from .hf import CachedHfTokenizer, HfTokenizer from .protocol import TokenizerLike if TYPE_CHECKING: @@ -164,6 +165,12 @@ def resolve_tokenizer_args( "mistral_format=True is not supported with tokenizer_mode='hf'" ) kwargs["mistral_format"] = False + tokenizer_name = get_non_weight_snapshot_path( + str(tokenizer_name), + revision=revision, + cache_dir=download_dir, + token=kwargs.get("token"), + ) return tokenizer_mode, tokenizer_name, args, kwargs @@ -202,6 +209,7 @@ def get_tokenizer( apply_fastokens_patch() + requested_name = str(tokenizer_name) tokenizer_mode, tokenizer_name, args, kwargs = cached_resolve_tokenizer_args( tokenizer_name, *args, @@ -254,6 +262,9 @@ def get_tokenizer( kwargs.setdefault("config", config) tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) + # Report the Hub ID rather than the local snapshot it was loaded from + if tokenizer_mode == "hf" and not envs.VLLM_USE_MODELSCOPE: + cast(HfTokenizer, tokenizer).name_or_path = requested_name if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: from vllm.tokenizers.hf import get_cached_tokenizer diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 9aa592084990..1b9b0bd442a7 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -22,7 +22,10 @@ from vllm.logger import init_logger from vllm.transformers_utils import processors -from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.transformers_utils.repo_utils import ( + get_hf_file_to_dict, + get_non_weight_snapshot_path, +) from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.func_utils import get_allowed_kwarg_only_overrides @@ -214,7 +217,12 @@ def get_processor( if revision is None: revision = "main" try: - processor_name = convert_model_repo_to_path(processor_name) + processor_name = requested_name = convert_model_repo_to_path(processor_name) + # A snapshot pinned to `revision` would shadow remote code at `code_revision` + if kwargs.get("code_revision") in (None, revision): + processor_name = get_non_weight_snapshot_path( + processor_name, revision=revision, token=kwargs.get("token") + ) registered_cls_name = get_processor_cls_name_from_config( processor_name, revision=revision ) @@ -269,6 +277,11 @@ def get_processor( f"found type: {type(processor)}" ) + # Report the Hub ID rather than the local snapshot it was loaded from + tokenizer = getattr(processor, "tokenizer", processor) + if hasattr(tokenizer, "name_or_path"): + tokenizer.name_or_path = requested_name + return processor diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index b6a23b007aa2..43ef07676e3a 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -304,6 +304,80 @@ def get_model_path(model: str | Path, revision: str | None = None): ) +_WEIGHT_FILE_PATTERNS = [ + "*.safetensors", + "*.bin", + "*.pt", + "*.pth", + "*.ckpt", + "*.gguf", + "*.h5", + "*.msgpack", + "*.onnx", + "*.onnx_data", + # Large non-weight assets that tokenizers and processors never read + "*.png", + "*.jpg", + "*.jpeg", + "*.gif", + "*.webp", + "*.mp3", + "*.mp4", + "*.wav", + "*.flac", + "*.parquet", + "*.arrow", + "*.ipynb", + "*.zip", + "*.tar*", +] + + +@cache +def _snapshot_non_weight_files( + repo_id: str, + revision: str | None, + cache_dir: str | None, + token: str | bool | None, +) -> str: + # Avoid circular import + from vllm.model_executor.model_loader.weight_utils import get_lock + + with get_lock(repo_id, cache_dir): + return hf_api().snapshot_download( + repo_id=repo_id, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_patterns=_WEIGHT_FILE_PATTERNS, + ) + + +def get_non_weight_snapshot_path( + repo_id: str, + revision: str | None = None, + cache_dir: str | None = None, + token: str | bool | None = None, +) -> str: + """Download all non-weight files of a Hub repo and return the local snapshot. + + Loading tokenizers and processors from a local directory stops Transformers + from making a Hub request per file (including the `additional_chat_templates` + listing). Returns `repo_id` unchanged if it is local or cannot be fetched, so + Transformers handles it as before. Failures are not cached, so they are retried. + """ + if envs.VLLM_USE_MODELSCOPE or Path(repo_id).exists(): + return repo_id + # Normalise so that `None` and "main" share a cache entry + revision = revision or huggingface_hub.constants.DEFAULT_REVISION + try: + return _snapshot_non_weight_files(repo_id, revision, cache_dir, token) + except (HfHubHTTPError, OSError, ValueError): + logger.debug("Failed to snapshot %s", repo_id, exc_info=True) + return repo_id + + def _try_download_from_hf_hub( model: str | Path, file_name: str, revision: str | None ) -> Path | None: From 4c36d21280824d1a9c521106932408243078de0d Mon Sep 17 00:00:00 2001 From: Suraj Sharan Date: Mon, 12 Oct 2026 17:49:30 +0400 Subject: [PATCH 415/952] [Bugfix][LoRA] Apply PEFT rank_pattern / alpha_pattern scaling per module (#59801) Signed-off-by: suraj Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: linitra24 --- tests/lora/test_lora_weights.py | 57 +++++++++++++++++ tests/lora/test_peft_helper.py | 105 ++++++++++++++++++++++++++++++++ vllm/lora/lora_model.py | 4 +- vllm/lora/lora_weights.py | 38 +++++++++--- vllm/lora/peft_helper.py | 60 ++++++++++++++++-- 5 files changed, 249 insertions(+), 15 deletions(-) diff --git a/tests/lora/test_lora_weights.py b/tests/lora/test_lora_weights.py index ff132147933e..7230126e40c8 100644 --- a/tests/lora/test_lora_weights.py +++ b/tests/lora/test_lora_weights.py @@ -4,6 +4,7 @@ import torch from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper RANK, ALPHA = 16, 32 # PEFTHelper stores alpha/sqrt(r) for use_rslora adapters. @@ -50,3 +51,59 @@ def test_non_gated_moe_keeps_w3_unscaled(): ) assert packed.scaling[:2] == pytest.approx([RSLORA_SCALING] * 2) assert packed.scaling[2] == 1.0 + + +@pytest.mark.parametrize( + "packer,stacked", + [ + (PackedLoRALayerWeights.pack_moe, False), + (PackedLoRALayerWeights.pack_moe_stacked, True), + ], +) +def test_moe_packing_keeps_per_projection_scaling(packer, stacked: bool): + # e.g. rank_pattern / alpha_pattern giving gate, down and up different scalings + loras = moe_loras(1.0, stacked) + for lora, scaling in zip(loras, (2.0, 4.0, 8.0)): + lora.scaling = scaling + packed = packer(loras, "experts") + assert packed.scaling == pytest.approx([2.0, 4.0, 8.0]) + + +def test_moe_packing_keeps_per_expert_scaling(): + """Experts of one projection can have different scaling, e.g. from + alpha_pattern.""" + loras = moe_loras(RSLORA_SCALING) + moe_loras(RSLORA_SCALING) + loras[4].scaling = 2 * RSLORA_SCALING # w2 of expert 1 + expected = [lora.lora_b * lora.scaling for lora in loras] + packed = PackedLoRALayerWeights.pack_moe(loras, "experts").optimize() + for i, lora_b in enumerate(packed.lora_b): + torch.testing.assert_close(lora_b, torch.stack(expected[i::3])) + + +def test_moe_packing_rejects_per_expert_rank(): + """Experts of one projection must share a rank to be stacked.""" + loras = moe_loras(1.0) + moe_loras(1.0) + loras[4].lora_a = torch.randn(RANK // 2, 4) + loras[4].lora_b = torch.randn(6, RANK // 2) + with pytest.raises(ValueError, match="different LoRA ranks"): + PackedLoRALayerWeights.pack_moe(loras, "experts") + + +def test_from_config_uses_checkpoint_module_name(): + peft_helper = PEFTHelper( + r=RANK, + lora_alpha=ALPHA, + target_modules=["q_proj"], + rank_pattern={r"^model\.layers\.0\.self_attn\.q_proj": 4}, + ) + # the pattern refers to the checkpoint name, not the name after weights mapping + lora = LoRALayerWeights.from_config( + "language_model.model.layers.0.self_attn.q_proj", + peft_helper, + "model.layers.0.self_attn.q_proj", + ) + assert lora.rank == 4 + assert lora.scaling == pytest.approx(ALPHA / 4) + lora = LoRALayerWeights.from_config("model.layers.1.self_attn.q_proj", peft_helper) + assert lora.rank == RANK + assert lora.scaling == pytest.approx(ALPHA / RANK) diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index bb1246144b0d..2fcf5ff96205 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -22,6 +22,11 @@ {"modules_to_save": ["lm_head"]}, "Unsupported modules_to_save", ), + ( + "test_rank_pattern", + {"r": 8, "rank_pattern": {"q_proj": 32}}, + "LoRA rank 32 is greater than max_lora_rank", + ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), ("test_lora_bias", {"lora_bias": True}, "does not support LoRA bias"), @@ -153,3 +158,103 @@ def test_peft_helper_init_lora_weights_supported(init_lora_weights): "use_bdlora": None, } ).validate_legal(lora_config) + + +@pytest.mark.parametrize( + "config,module_name,expected", + [ + # rank_pattern only: alpha / r_m + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.q_proj", (4, 8.0)), + ({"rank_pattern": {"q_proj": 4}}, "model.layers.0.self_attn.v_proj", (16, 2.0)), + # alpha_pattern only: alpha_m / r + ( + {"alpha_pattern": {"v_proj": 64}}, + "model.layers.0.self_attn.v_proj", + (16, 4.0), + ), + # keys only match at a "." boundary, like in PEFT + ({"rank_pattern": {"proj": 4}}, "model.layers.0.self_attn.q_proj", (16, 2.0)), + # fused MoE experts from PEFT `target_parameters`: gate_up_proj is stored + # as `experts.base_layer`, down_proj as `experts` + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts.base_layer", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.down_proj": 2}}, + "model.layers.0.mlp.experts", + (2, 16.0), + ), + ( + {"rank_pattern": {"experts.gate_up_proj": 2}}, + "model.layers.0.mlp.experts", + (16, 2.0), + ), + ], +) +def test_peft_helper_rank_and_alpha_pattern(config, module_name, expected): + peft_helper = PEFTHelper( + r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], **config + ) + rank, scaling = peft_helper.get_rank_and_scaling(module_name) + assert rank == expected[0] + assert scaling == pytest.approx(expected[1]) + + +def test_peft_helper_pattern_rslora(): + peft_helper = PEFTHelper( + r=16, + lora_alpha=32, + target_modules=["q_proj"], + use_rslora=True, + rank_pattern={"q_proj": 4}, + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (4, 32 / math.sqrt(4)) + + +def test_peft_helper_dynamic_rank_conversion(): + """PEFT's `save_as_lora` with a dynamic rank writes r=1, lora_alpha=1 and the + same per-module values in rank_pattern and alpha_pattern, so every module + keeps a scaling of 1.""" + peft_helper = PEFTHelper( + r=1, + lora_alpha=1, + target_modules=["q_proj", "v_proj"], + rank_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + alpha_pattern={"layers.0.self_attn.q_proj": 32, "layers.0.self_attn.v_proj": 3}, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.q_proj") == ( + 32, + 1.0, + ) + assert peft_helper.get_rank_and_scaling("model.layers.0.self_attn.v_proj") == ( + 3, + 1.0, + ) + with pytest.raises(ValueError, match="LoRA rank 32 is greater than max_lora_rank"): + peft_helper.validate_legal(LoRAConfig(max_lora_rank=16)) + + +def test_peft_helper_null_patterns(): + peft_helper = PEFTHelper.from_dict( + { + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + "rank_pattern": None, + "alpha_pattern": None, + } + ) + assert peft_helper.get_rank_and_scaling("model.q_proj") == (8, 2.0) + + +@pytest.mark.parametrize("bad_rank", [0, -4, 2.5]) +def test_peft_helper_invalid_rank_pattern(bad_rank): + with pytest.raises(ValueError, match="rank_pattern"): + PEFTHelper( + r=8, + lora_alpha=16, + target_modules=["q_proj"], + rank_pattern={"q_proj": bad_rank}, + ) diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index 196445332443..f4520b12c75b 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -150,8 +150,10 @@ def from_lora_tensors( continue if module_name not in loras: + # rank/alpha patterns refer to the names before weights mapping + peft_module_name, _ = parse_fine_tuned_lora_name(tensor_name) loras[module_name] = LoRALayerWeights.from_config( - module_name, peft_helper + module_name, peft_helper, peft_module_name ) if is_lora_a: diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index 13a692a62353..9ea7c550d703 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -59,15 +59,22 @@ def from_config( cls, module_name: str, peft_helper: PEFTHelper, + peft_module_name: str | None = None, ) -> "LoRALayerWeights": + """`peft_module_name` is the module name in the adapter checkpoint, used + to resolve `rank_pattern` / `alpha_pattern`. Defaults to `module_name`. + """ + rank, scaling = peft_helper.get_rank_and_scaling( + peft_module_name or module_name + ) # lora_a and lora_b are set to None for config-based construction return cls( module_name, - peft_helper.r, + rank, peft_helper.lora_alpha, None, None, - peft_helper.vllm_lora_scaling_factor, + scaling, ) @classmethod @@ -168,6 +175,18 @@ def pack_moe( rank = first_lora.rank lora_alpha = first_lora.lora_alpha assert len(loras) % 3 == 0 + for proj_loras in (loras[0::3], loras[1::3], loras[2::3]): + experts = [lora for lora in proj_loras if lora is not None] + if len({lora.lora_a.shape[0] for lora in experts}) > 1: + raise ValueError( + f"The experts of {module_name} have different LoRA ranks for " + "the same projection, which is not supported." + ) + # The packed LoRA has one scaling per projection, so fold it into + # lora_b where the experts differ (e.g. from alpha_pattern). + if len({lora.scaling for lora in experts}) > 1: + for lora in experts: + lora.optimize() w1_lora_a_lst = [] w2_lora_a_lst = [] w3_lora_a_lst = [] @@ -201,10 +220,12 @@ def pack_moe( w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank) w2_lora_b = torch.stack(w2_lora_b_lst, dim=0) - # All w1, w2, w3 have the same scaling factor. Use the per-adapter - # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank. - scaling = first_lora.scaling - last_scaling = scaling + # Use each projection's own scaling (e.g. alpha/sqrt(r) for rsLoRA, or + # per-module values from rank_pattern / alpha_pattern) instead of + # alpha/rank. + w1_scaling = loras[0].scaling # type: ignore[union-attr] + w2_scaling = loras[1].scaling # type: ignore[union-attr] + last_scaling = loras[2].scaling if loras[2] is not None else w1_scaling if is_non_gated_moe: # For non-gated MoE, reuse w1 tensors for w3 to avoid memory waste @@ -224,7 +245,7 @@ def pack_moe( [lora_alpha, lora_alpha, lora_alpha], [w1_lora_a, w2_lora_a, w3_lora_a], [w1_lora_b, w2_lora_b, w3_lora_b], - scaling=[scaling, scaling, last_scaling], + scaling=[w1_scaling, w2_scaling, last_scaling], ) return obj @@ -251,14 +272,13 @@ def pack_moe_stacked( assert w1_lora is not None and w2_lora is not None and w3_lora is not None rank = w1_lora.rank lora_alpha = w1_lora.lora_alpha - scaling = w1_lora.scaling return cls( module_name, rank, [lora_alpha, lora_alpha, lora_alpha], [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a], [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b], - scaling=[scaling, scaling, scaling], + scaling=[w1_lora.scaling, w2_lora.scaling, w3_lora.scaling], ) def optimize(self) -> "PackedLoRALayerWeights": diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index de05dea1cfc1..721b7a039f00 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -9,6 +9,8 @@ from dataclasses import MISSING, dataclass, field, fields from typing import Literal +import regex as re + from vllm.config.lora import LoRAConfig from vllm.logger import init_logger @@ -41,6 +43,9 @@ class PEFTHelper: layer_replication: list[list[int]] | None = field(default=None) use_bdlora: dict | None = field(default=None) use_qalora: bool = field(default=False) + # Per-module overrides of `r` and `lora_alpha`, keyed by module name pattern + rank_pattern: dict[str, int] = field(default_factory=dict) + alpha_pattern: dict[str, float] = field(default_factory=dict) # Extra vllm field, start with 'vllm_' to avoid conflict vllm_lora_scaling_factor: float = field(default=1.0) vllm_max_position_embeddings: int | None = field(default=False) @@ -97,11 +102,55 @@ def _init_is_supported(self) -> bool: def __post_init__(self): if self.r <= 0: raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") + self.rank_pattern = self.rank_pattern or {} + self.alpha_pattern = self.alpha_pattern or {} + for key, rank in self.rank_pattern.items(): + if not isinstance(rank, int) or rank <= 0: + raise ValueError( + f"LoRA rank for `{key}` in `rank_pattern` must be a positive " + f"integer, got {rank}." + ) if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") - self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r) - else: - self.vllm_lora_scaling_factor = self.lora_alpha / self.r + self.vllm_lora_scaling_factor = self._scaling(self.r, self.lora_alpha) + + def _scaling(self, rank: int, alpha: float) -> float: + if self.use_rslora: + return alpha / math.sqrt(rank) + return alpha / rank + + @staticmethod + def _match_pattern(pattern: dict, name: str): + # Same rule as PEFT's `get_pattern_key`: a key matches if it is a suffix + # of the module name starting at a "." boundary; the first match wins. + for key, value in pattern.items(): + if re.match(rf"(.*\.)?({key})$", name): + return value + return None + + def get_rank_and_scaling(self, module_name: str) -> tuple[int, float]: + """Return the rank and scaling PEFT uses for `module_name`, taking + `rank_pattern` and `alpha_pattern` into account. + + `module_name` is the name in the adapter checkpoint, without the + `base_model.model.` prefix and before any weights mapping. + """ + if not self.rank_pattern and not self.alpha_pattern: + return self.r, self.vllm_lora_scaling_factor + # PEFT stores fused MoE expert LoRAs (`target_parameters`) as + # `experts.base_layer` (gate_up_proj) and `experts` (down_proj), but + # matches the patterns against the parameter names. + if module_name.endswith(".experts.base_layer"): + module_name = module_name.removesuffix(".base_layer") + ".gate_up_proj" + elif module_name.endswith(".experts"): + module_name += ".down_proj" + rank = self._match_pattern(self.rank_pattern, module_name) + if rank is None: + rank = self.r + alpha = self._match_pattern(self.alpha_pattern, module_name) + if alpha is None: + alpha = self.lora_alpha + return rank, self._scaling(rank, alpha) @classmethod def from_dict(cls, config_dict: dict) -> "PEFTHelper": @@ -141,9 +190,10 @@ def validate_legal(self, lora_config: LoRAConfig) -> None: constraints and requirements. """ error_msg = self._validate_features() - if self.r > lora_config.max_lora_rank: + max_rank = max([self.r, *self.rank_pattern.values()]) + if max_rank > lora_config.max_lora_rank: error_msg.append( - f"LoRA rank {self.r} is greater than max_lora_rank" + f"LoRA rank {max_rank} is greater than max_lora_rank" f" {lora_config.max_lora_rank}." ) if self.bias != "none": From f022975964e27ed18dd8e437eed0ab70b0a00b33 Mon Sep 17 00:00:00 2001 From: bzsuni <86399306+bzsuni@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:49:46 +0800 Subject: [PATCH 416/952] [Bugfix] Reject beam search in completion render (#61212) Signed-off-by: bzsuni --- .../scale_out/render/test_render.py | 36 ++++++++++++++++--- vllm/entrypoints/scale_out/render/serving.py | 6 ++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 6c11a518d764..308555be362e 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -12,6 +12,7 @@ import pytest_asyncio from tests.utils import RemoteLaunchRenderServer +from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.entrypoints.scale_out.render.api_router import router from vllm.entrypoints.scale_out.render.serving import ServingRender @@ -21,7 +22,7 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM" -def _build_responses_serving_render() -> ServingRender: +def _build_serving_render() -> ServingRender: serving = ServingRender.__new__(ServingRender) serving.model_config = SimpleNamespace( max_model_len=100, @@ -43,6 +44,31 @@ def _build_responses_serving_render() -> ServingRender: return serving +@pytest.mark.asyncio +@pytest.mark.skip_global_cleanup +async def test_render_completion_rejects_beam_search(): + serving = _build_serving_render() + serving.online_renderer.render_completion = AsyncMock( + return_value=[{"prompt_token_ids": [7, 8, 9]}] + ) + request = CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + use_beam_search=True, + n=2, + length_penalty=0.7, + ) + + response = await serving.render_completion_request(request) + + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + assert ( + response.error.message == "Beam search is not supported by the render endpoint" + ) + serving.online_renderer.render_completion.assert_not_awaited() + + @pytest.mark.skip_global_cleanup def test_responses_render_route_is_registered(): assert any(route.path == "/v1/responses/render" for route in router.routes) @@ -60,7 +86,7 @@ def test_responses_render_route_documents_not_implemented(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_returns_generate_request_without_stored_state(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock( return_value=MagicMock( messages=[{"role": "user", "content": "Test prompt"}], @@ -101,7 +127,7 @@ async def test_render_responses_returns_generate_request_without_stored_state(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_rejects_previous_response_id(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock() request = ResponsesRequest( model=MODEL_NAME, @@ -120,7 +146,7 @@ async def test_render_responses_rejects_previous_response_id(): @pytest.mark.asyncio @pytest.mark.skip_global_cleanup async def test_render_responses_rejects_empty_token_ids(): - serving = _build_responses_serving_render() + serving = _build_serving_render() serving.online_renderer.render_responses = AsyncMock( return_value=MagicMock( messages=[], @@ -141,7 +167,7 @@ async def test_render_responses_rejects_empty_token_ids(): async def test_render_responses_carries_reasoning_parser_kwargs(): """/v1/responses gives the engine's reasoning parser the effective template kwargs, so the rendered request must carry them to generate.""" - serving = _build_responses_serving_render() + serving = _build_serving_render() online = serving.online_renderer online.parser = MagicMock(reasoning_parser_cls=object) online.chat_template = None diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index 84771c5f7390..5caf78bffe9e 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -186,6 +186,12 @@ async def render_completion_request( error_check_ret = await self._check_model(request) if error_check_ret is not None: return error_check_ret + + if request.use_beam_search: + return self.create_error_response( + "Beam search is not supported by the render endpoint" + ) + result = await self.online_renderer.render_completion( request, skip_mm_cache=True ) From b3ee560ae816b763b09b1aeddf9adca95a0014d0 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 12 Oct 2026 16:02:56 +0200 Subject: [PATCH 417/952] [Bugfix][MoE] Apply expert_map remapping for CUTLASS FP4 MOE expert parallel (#59042) Signed-off-by: Markov Ilya --- tests/kernels/moe/test_cutlass_moe.py | 203 +++++++++++++++++- tests/kernels/moe/test_ocp_mx_moe.py | 111 ++++++++++ .../layers/fused_moe/experts/cutlass_moe.py | 45 +++- .../experts/gpt_oss_triton_kernels_moe.py | 48 +---- .../layers/fused_moe/oracle/mxfp4.py | 55 +++++ vllm/model_executor/layers/fused_moe/utils.py | 67 ++++++ 6 files changed, 470 insertions(+), 59 deletions(-) diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 050e1bc10487..a06abfbc1fad 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -9,7 +9,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.moe.utils import make_dummy_moe_config, make_test_quant_config from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import ( @@ -26,8 +26,10 @@ ) from vllm.model_executor.layers.fused_moe.config import ( FUSED_MOE_UNQUANTIZED_CONFIG, + FusedMoEParallelConfig, FusedMoEQuantConfig, fp8_w8a8_moe_quant_config, + mxfp4_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp4, @@ -41,7 +43,10 @@ ) from vllm.model_executor.layers.fused_moe.oracle import mxfp4 as mxfp4_oracle from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input +from vllm.model_executor.layers.fused_moe.utils import ( + moe_kernel_quantize_input, + remap_topk_to_local, +) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -183,6 +188,189 @@ def test_cutlass_moe_permutation_maps_padding_to_zero(): torch.testing.assert_close(restored[3], torch.zeros_like(restored[3])) +def test_remap_topk_ids_to_local(): + device = "cuda" + expert_map = torch.tensor([-1, 0, -1, 1, -1, 2], dtype=torch.int32, device=device) + topk_ids = torch.tensor([[1, 3], [0, 5], [-1, 4]], dtype=torch.int32, device=device) + expected = torch.tensor( + [[0, 1], [-1, 2], [-1, -1]], dtype=torch.int32, device=device + ) + actual = remap_topk_to_local(topk_ids, expert_map, out_dtype=torch.int32) + torch.testing.assert_close(actual, expected) + + +def test_remap_topk_ids_to_local_requires_contiguous_topk_ids(): + expert_map = torch.tensor([0, 1], dtype=torch.int32) + topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32).t() + assert not topk_ids.is_contiguous() + with pytest.raises(AssertionError): + remap_topk_to_local(topk_ids, expert_map) + + +def test_cutlass_fp4_moe_supports_expert_parallel_config(): + ep_config = FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ) + assert CutlassExpertsMxfp4._supports_parallel_config(ep_config) + assert CutlassExpertsFp4._supports_parallel_config(ep_config) + + +def _slice_moe_quant_config( + quant_config: FusedMoEQuantConfig, start: int, end: int +) -> FusedMoEQuantConfig: + num_experts = quant_config.w1_scale.shape[0] + sliced_config = copy.deepcopy(quant_config) + for desc_name in ("_a1", "_a2", "_w1", "_w2"): + desc = getattr(sliced_config, desc_name) + for field_name in ("scale", "alpha_or_gscale", "bias", "zp"): + tensor = getattr(desc, field_name) + if ( + isinstance(tensor, torch.Tensor) + and tensor.ndim > 0 + and tensor.shape[0] == num_experts + ): + setattr(desc, field_name, tensor[start:end]) + return sliced_config + + +def _swizzle_mxfp4_moe_scales( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + e, w1_n, w1_scale_k = w1_scale.shape + w1_k = w1_scale_k * 32 + w2_m, w2_scale_n = w2_scale.shape[1], w2_scale.shape[2] + w2_n = w2_scale_n * 32 + + swizzled_w1 = [] + swizzled_w2 = [] + for expert in range(e): + swizzled_w1.append( + swizzle_mxfp4_scales(w1_scale[expert], w1_n, w1_k).reshape(w1_n, w1_scale_k) + ) + swizzled_w2.append( + swizzle_mxfp4_scales(w2_scale[expert], w2_m, w2_n).reshape(w2_m, w2_scale_n) + ) + return torch.stack(swizzled_w1), torch.stack(swizzled_w2) + + +def _make_cutlass_fp4_moe_weights( + quantization: str, + num_experts: int, + n: int, + k: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]: + if quantization == "nvfp4": + return make_test_quant_config( + num_experts, + n, + k, + in_dtype=dtype, + quant_dtype="nvfp4", + block_shape=None, + per_act_token_quant=False, + make_gate=True, + ) + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, + ) + + w1, w1_scale = mxfp4_quantize( + torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 8 + ) + w2, w2_scale = mxfp4_quantize( + torch.randn((num_experts, k, n), device=device, dtype=dtype) / 8 + ) + w1_scale, w2_scale = _swizzle_mxfp4_moe_scales(w1_scale, w2_scale) + return ( + w1.contiguous(), + w2.contiguous(), + mxfp4_moe_quant_config(w1_scale=w1_scale, w2_scale=w2_scale), + ) + + +@pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"]) +@torch.inference_mode() +def test_cutlass_fp4_moe_EP( + quantization: str, + workspace_init, +): + """Sharded CUTLASS FP4/MXFP4 MoE with expert_map must match full weights.""" + experts_cls = CutlassExpertsFp4 if quantization == "nvfp4" else CutlassExpertsMxfp4 + if not experts_cls._supports_current_device(): + pytest.skip(f"CUTLASS {quantization} MoE is not supported on this GPU") + + set_random_seed(7) + m, n, k = 8, 128, 128 + num_experts, num_local_experts, topk = 8, 2, 2 + dtype = torch.bfloat16 + device = torch.device("cuda") + + w1, w2, quant_config = _make_cutlass_fp4_moe_weights( + quantization, num_experts, n, k, dtype, device + ) + hidden_states = torch.randn((m, k), device=device, dtype=dtype) / 10 + score = torch.randn((m, num_experts), device=device, dtype=dtype) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + + moe_config = make_dummy_moe_config( + num_experts=num_experts, + num_local_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + in_dtype=dtype, + max_num_tokens=m, + ) + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts_cls(moe_config=moe_config, quant_config=quant_config), + ) + apply_kwargs = { + "hidden_states": hidden_states, + "w1": w1, + "w2": w2, + "topk_weights": topk_weights, + "topk_ids": topk_ids, + "global_num_experts": num_experts, + "activation": MoEActivation.SILU, + "expert_map": None, + "apply_router_weight_on_input": False, + } + reference = kernel.apply(**apply_kwargs) + sharded = run_with_expert_maps( + num_experts, + num_local_experts, + quant_config, + experts_cls, + **apply_kwargs, + ) + + assert reference.abs().max() > 0 + torch.testing.assert_close(sharded, reference, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"]) @torch.inference_mode() def test_cutlass_fp4_moe_padded_routes_do_not_change_valid_output(quantization: str): @@ -486,6 +674,7 @@ def run_with_expert_maps( num_experts: int, num_local_experts: int, quant_config: FusedMoEQuantConfig, + experts_cls: type[mk.FusedMoEExperts], **cutlass_moe_kwargs, ): def slice_experts(): @@ -513,9 +702,7 @@ def slice_experts(): for k, t in full_tensors.items(): cutlass_moe_kwargs[k] = t[s:e] - new_quant_config = copy.deepcopy(quant_config) - new_quant_config._w1.scale = quant_config.w1_scale[s:e] - new_quant_config._w2.scale = quant_config.w2_scale[s:e] + new_quant_config = _slice_moe_quant_config(quant_config, s, e) yield cutlass_moe_kwargs, new_quant_config @@ -523,13 +710,14 @@ def slice_experts(): for kwargs, new_quant_config in slice_experts(): w2 = kwargs["w2"] a = kwargs["hidden_states"] + w1 = kwargs["w1"] moe_config = make_dummy_moe_config( max_num_tokens=kwargs.get("hidden_states").shape[0], experts_per_token=kwargs.get("topk_ids").shape[1], num_experts=num_experts, num_local_experts=num_local_experts, hidden_dim=w2.shape[1], - intermediate_size=w2.shape[2], + intermediate_size=w1.shape[1] // 2, in_dtype=a.dtype, ) kernel = mk.FusedMoEKernel( @@ -539,7 +727,7 @@ def slice_experts(): allow_new_interface=True, use_monolithic=False, ), - CutlassExpertsFp8( + experts_cls( moe_config=moe_config, quant_config=new_quant_config, ), @@ -624,6 +812,7 @@ def run_8_bit( num_experts, num_local_experts, quant_config, + CutlassExpertsFp8, **kwargs, ) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 13f29087d951..0a8f5c29d034 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -2114,6 +2114,117 @@ def test_aiter_mxfp4_monolithic_rejects_expert_parallel( assert "parallel config" in reason +@pytest.mark.parametrize( + "all2all_backend", + [ + "deepep_high_throughput", + "flashinfer_nvlink_two_sided", + "flashinfer_nvlink_one_sided", + ], +) +def test_make_mxfp4_moe_kernel_rejects_cutlass_unsupported_expert_parallel( + monkeypatch: pytest.MonkeyPatch, + all2all_backend: str, +): + """Compressed-tensors CUTLASS W4A4 bypasses the oracle; guard at assembly.""" + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + mxfp4_moe_quant_config, + ) + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + make_mxfp4_moe_kernel, + ) + + monkeypatch.setattr( + CutlassExpertsMxfp4, "_supports_current_device", staticmethod(lambda: True) + ) + + moe_config = replace( + make_dummy_moe_config( + num_experts=256, + num_local_experts=32, + hidden_dim=4096, + intermediate_size=2048, + ), + moe_parallel_config=FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend=all2all_backend, + enable_eplb=False, + ), + ) + w1_scale = torch.empty(32, 64, 128, dtype=torch.uint8) + w2_scale = torch.empty(32, 4096, 64, dtype=torch.uint8) + quant_config = mxfp4_moe_quant_config(w1_scale=w1_scale, w2_scale=w2_scale) + + with pytest.raises(ValueError, match="parallel config"): + make_mxfp4_moe_kernel( + quant_config, + moe_config, + CutlassExpertsMxfp4, + Mxfp4MoeBackend.MARLIN, + ) + + +def test_make_mxfp4_moe_kernel_validates_cutlass_expert_parallel( + monkeypatch: pytest.MonkeyPatch, +): + """CUTLASS MXFP4 supports EP with allgather_reducescatter after expert_map.""" + from tests.kernels.moe.utils import make_dummy_moe_config + from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsMxfp4, + ) + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + _validate_mxfp4_experts_config, + ) + + monkeypatch.setattr( + CutlassExpertsMxfp4, "_supports_current_device", staticmethod(lambda: True) + ) + + moe_config = replace( + make_dummy_moe_config( + num_experts=256, + num_local_experts=32, + hidden_dim=4096, + intermediate_size=2048, + ), + moe_parallel_config=FusedMoEParallelConfig( + tp_size=1, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=8, + dp_rank=0, + ep_size=8, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ), + ) + + _validate_mxfp4_experts_config( + CutlassExpertsMxfp4, moe_config, Mxfp4MoeBackend.MARLIN + ) + + # Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that # `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns # the activation untouched, so a name it does not know (e.g. "mxfp6" instead of diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index 6f9a62f282ce..5c6a96f26daa 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -31,6 +31,7 @@ ) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, + remap_topk_to_local, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -412,10 +413,9 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - # CutlassExpertsFp8 does not support expert map, which is - # needed for STANDARD activation format kernels in DP/EP mode. - # Note that the BATCHED activation format does not use - # the expert map for identifying experts. + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. return not ( moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_deepep_ht_kernels @@ -491,6 +491,15 @@ def workspace_shapes( FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max +def _prepare_topk_ids_for_cutlass_fp4_moe( + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, +) -> torch.Tensor: + if expert_map is None: + return topk_ids + return remap_topk_to_local(topk_ids, expert_map, out_dtype=torch.int32) + + def run_cutlass_moe_fp4( output: torch.Tensor, a: torch.Tensor, @@ -513,6 +522,7 @@ def run_cutlass_moe_fp4( e: int, device: torch.device, apply_router_weight_on_input: bool = False, + expert_map: torch.Tensor | None = None, *, activation_config: ApplyMoEActivationConfig | None = None, ) -> None: @@ -543,6 +553,8 @@ def run_cutlass_moe_fp4( assumes that topk < k < n to satisfy - up/down projection expectations. """ + topk_ids = _prepare_topk_ids_for_cutlass_fp4_moe(topk_ids, expert_map) + is_gated = activation.is_gated # For gated activations (e.g. SiLU), w1 output is 2*n (gate + up). # For non-gated activations (e.g. SiLU_NO_MUL), w1 output is n (up only). @@ -729,9 +741,14 @@ def _supports_activation(activation: MoEActivation) -> bool: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - # CutlassExpertsFp4 does not support expert map, which is - # needed for STANDARD activation format kernels in EP mode. - return moe_parallel_config.ep_size == 1 + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_deepep_ht_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) @staticmethod def _supports_batch_invariance() -> bool: @@ -806,6 +823,7 @@ def apply( e=e, device=hidden_states.device, apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, activation_config=self.activation_config, ) @@ -828,10 +846,13 @@ def run_cutlass_moe_mxfp4( e: int, device: torch.device, apply_router_weight_on_input: bool = False, + expert_map: torch.Tensor | None = None, *, activation_config: ApplyMoEActivationConfig | None = None, ) -> None: """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM.""" + topk_ids = _prepare_topk_ids_for_cutlass_fp4_moe(topk_ids, expert_map) + is_gated = activation.is_gated w1_n = n * 2 if is_gated else n @@ -1036,7 +1057,14 @@ def _supports_activation(activation: MoEActivation) -> bool: def _supports_parallel_config( moe_parallel_config: FusedMoEParallelConfig, ) -> bool: - return moe_parallel_config.ep_size == 1 + # EP is supported only with allgather_reducescatter (expert_map remapping). + # DeepEP HT and FlashInfer NVLink prepare/finalize use incompatible + # post-dispatch layouts for Standard-format grouped GEMM. + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_deepep_ht_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels + ) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -1103,6 +1131,7 @@ def apply( e=e, device=hidden_states.device, apply_router_weight_on_input=apply_router_weight_on_input, + expert_map=expert_map, activation_config=self.activation_config, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index e3b6995d693c..3e70c1589365 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -21,7 +21,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import _resize_cache +from vllm.model_executor.layers.fused_moe.utils import ( + _resize_cache, + remap_topk_to_local, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -994,49 +997,6 @@ def masked_moe_sum( ) -@triton.jit -def _remap_topk_to_local_kernel( - topk_ids_ptr, # [n] global expert IDs (-1 = invalid) - expert_map_ptr, # [num_experts] global->local (-1 for non-local) - out_ptr, # [n] int64 local expert IDs (-1 for invalid/non-local) - n_elements, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n_elements - tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) - # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid - # rows don't read OOB, then select -1 for them. Matches - # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a - # plain expert_map[-1] would wrap to a valid local id and misroute). - valid = tid >= 0 - idx = tl.where(valid, tid, 0) - local = tl.load(expert_map_ptr + idx, mask=mask, other=-1) - out = tl.where(valid, local.to(tl.int64), -1) - tl.store(out_ptr + offs, out, mask=mask) - - -def remap_topk_to_local( - topk_ids: torch.Tensor, expert_map: torch.Tensor -) -> torch.Tensor: - """Fused global->local expert-id mapping over a topk_ids tensor, preserving -1. - - Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)`` - with one kernel. Returns a NEW int64 tensor -- the caller keeps the original - ``topk_ids`` as ``global_topk_ids``, so this must not write in place. - - (Distinct from ``deep_gemm_utils.apply_expert_map``, which is a scalar - ``@triton.jit`` device helper called from within other kernels.) - """ - out = torch.empty_like(topk_ids, dtype=torch.int64) - n = topk_ids.numel() - BLOCK = 1024 - grid = (triton.cdiv(n, BLOCK),) - _remap_topk_to_local_kernel[grid](topk_ids, expert_map, out, n, BLOCK=BLOCK) - return out - - class BaseOAITritonExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7667dfc6aef0..cbbc39b94310 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -2133,6 +2133,59 @@ def make_mxfp4_moe_quant_config( ) +_MXFP4_SCHEME_CANDIDATES: list[tuple[QuantKey, QuantKey | None]] = [ + (kMxfp4Static, kMxfp4Dynamic), + (kMxfp4Static, kMxfp8Dynamic), + (kMxfp4Static, kFp8Dynamic128Sym), + (kMxfp4Static, kFp8StaticTensorSym), + (kMxfp4Static, None), +] + + +def _mxfp4_quant_keys_for_experts( + experts_cls: type[mk.FusedMoEExperts], + mxfp4_backend: Mxfp4MoeBackend, +) -> tuple[QuantKey | None, QuantKey | None]: + for weight_key, activation_key in _MXFP4_SCHEME_CANDIDATES: + if experts_cls._supports_quant_scheme(weight_key, activation_key): + return weight_key, activation_key + return kMxfp4Static, _backend_activation_key(mxfp4_backend) + + +def _activation_format_for_config( + moe_config: FusedMoEConfig, +) -> mk.FusedMoEActivationFormat: + return ( + mk.FusedMoEActivationFormat.BatchedExperts + if moe_config.moe_parallel_config.use_batched_activation_format + else mk.FusedMoEActivationFormat.Standard + ) + + +def _validate_mxfp4_experts_config( + experts_cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + mxfp4_backend: Mxfp4MoeBackend, +) -> None: + """Fail closed when a caller bypasses the oracle and picks experts directly.""" + activation_format = _activation_format_for_config(moe_config) + weight_key, activation_key = _mxfp4_quant_keys_for_experts( + experts_cls, mxfp4_backend + ) + supported, reason = experts_cls.is_supported_config( + experts_cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + raise ValueError( + f"{experts_cls.__name__} does not support this MoE deployment" + f"{f' since {reason}' if reason else ''}." + ) + + def make_mxfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, @@ -2141,6 +2194,8 @@ def make_mxfp4_moe_kernel( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: """Create a FusedMoEKernel for the given MXFP4 backend.""" + _validate_mxfp4_experts_config(experts_cls, moe_config, mxfp4_backend) + is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) prepare_finalize = maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index c724a7c7cb21..b1a6dd5632b8 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -235,6 +235,73 @@ def count_expert_num_tokens( return expert_num_tokens +@triton.jit +def _remap_topk_to_local_kernel( + topk_ids_ptr, # [n] global expert IDs (-1 = invalid) + expert_map_ptr, # [num_experts] global->local (-1 for non-local) + out_ptr, # [n] local expert IDs (-1 for invalid/non-local) + n_elements, + BLOCK: tl.constexpr, + OUT_INT32: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) + # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid + # rows don't read OOB, then select -1 for them. Matches + # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a + # plain expert_map[-1] would wrap to a valid local id and misroute). + valid = tid >= 0 + idx = tl.where(valid, tid, 0) + local = tl.load(expert_map_ptr + idx, mask=mask, other=-1) + if OUT_INT32: + out = tl.where(valid, local, -1) + else: + out = tl.where(valid, local.to(tl.int64), -1) + tl.store(out_ptr + offs, out, mask=mask) + + +def remap_topk_to_local( + topk_ids: torch.Tensor, + expert_map: torch.Tensor, + *, + out_dtype: torch.dtype = torch.int64, +) -> torch.Tensor: + """Map global expert ids to local ids; non-local or invalid routes become -1. + + Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)`` + with one kernel. Returns a NEW tensor -- the caller keeps the original + ``topk_ids`` as ``global_topk_ids``, so this must not write in place. + + Args: + topk_ids: Global expert ids; -1 marks invalid slots. Must be contiguous; + the Triton kernel indexes by flat ``numel()`` offsets. + expert_map: Global-to-local mapping; -1 marks non-local experts. + out_dtype: Output dtype, ``torch.int32`` or ``torch.int64``. + + Returns: + Remapped local expert ids with the same shape as ``topk_ids``. + + """ + if out_dtype not in (torch.int32, torch.int64): + raise ValueError(f"out_dtype must be int32 or int64, got {out_dtype}") + assert topk_ids.is_contiguous() + out = torch.empty(topk_ids.shape, dtype=out_dtype, device=topk_ids.device) + n = topk_ids.numel() + BLOCK = 1024 + grid = (triton.cdiv(n, BLOCK),) + _remap_topk_to_local_kernel[grid]( + topk_ids, + expert_map, + out, + n, + BLOCK=BLOCK, + OUT_INT32=out_dtype == torch.int32, + ) + return out + + def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor: """Shrink the given tensor and apply the given view to it. This is used to resize the intermediate fused_moe caches. From a281ad32f0368caa2c6ebe38a0e59a75b3f22392 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 12 Oct 2026 16:03:09 +0200 Subject: [PATCH 418/952] [EPLB] Add EPLB support to dspark draft models (DeepSeekV4) (#57549) Signed-off-by: Markov Ilya --- tests/distributed/test_eplb_spec_decode.py | 230 +++++++++++++++++- tests/models/test_deepseek_v4_mega_moe.py | 95 ++++++++ tests/models/test_dspark_mla.py | 51 ++++ tests/v1/worker/test_dspark_utils.py | 37 ++- .../worker/test_gpu_model_runner_v2_eplb.py | 83 ++++++- vllm/distributed/eplb/eplb_state.py | 64 ++--- vllm/models/deepseek_v4/common/eplb_util.py | 50 ++++ vllm/models/deepseek_v4/nvidia/dspark.py | 16 +- vllm/models/deepseek_v4/nvidia/model.py | 25 +- vllm/v1/worker/gpu/eplb_utils.py | 37 ++- .../v1/worker/gpu/spec_decode/dspark/utils.py | 21 +- vllm/v1/worker/gpu_model_runner.py | 32 ++- 12 files changed, 665 insertions(+), 76 deletions(-) create mode 100644 vllm/models/deepseek_v4/common/eplb_util.py diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py index ced206a3056a..b49d06cd7681 100644 --- a/tests/distributed/test_eplb_spec_decode.py +++ b/tests/distributed/test_eplb_spec_decode.py @@ -2,11 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations +from types import SimpleNamespace + import lm_eval import pytest +import torch from tests.utils import large_gpu_mark +from vllm.distributed.eplb.eplb_state import EplbState +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config +from vllm.v1.worker.gpu.eplb_utils import ( + EPLBController, + draft_model_supports_eplb, +) def get_model_args( @@ -15,14 +25,18 @@ def get_model_args( spec_method: str, tp_size: int, model_max_len: int, + num_speculative_tokens: int = 1, use_async: bool = True, ) -> dict: speculative_config = { "method": spec_method, "model": spec_model_name, - "num_speculative_tokens": 1, + "num_speculative_tokens": num_speculative_tokens, "max_model_len": model_max_len, } + if spec_method == "dspark": + speculative_config["model"] = model_name + speculative_config["draft_sample_method"] = "probabilistic" eplb_config = { "num_redundant_experts": tp_size, "window_size": 128, @@ -42,6 +56,8 @@ def get_model_args( "enable_eplb": True, "max_model_len": model_max_len, } + if spec_method == "dspark": + model_args["trust_remote_code"] = True return model_args @@ -51,11 +67,190 @@ def get_model_args( ) +class FakeEplbState: + instances: list[FakeEplbState] = [] + + def __init__(self, parallel_config, device: torch.device): + self.add_model_calls: list[tuple[object, object, str | None]] = [] + FakeEplbState.instances.append(self) + + def add_model( + self, + model: object, + model_config: object, + *, + model_name: str | None = None, + ) -> None: + self.add_model_calls.append((model, model_config, model_name)) + + +def _make_dsv4_dspark_hf_config() -> DeepseekV4Config: + return DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + num_hash_layers=0, + n_shared_experts=1, + moe_intermediate_size=128, + hc_mult=1, + hc_eps=1e-5, + rms_norm_eps=1e-5, + dspark_target_layer_ids=[0], + dspark_markov_rank=8, + index_topk=4, + head_dim=64, + num_attention_heads=4, + vocab_size=256, + n_mtp_layers=2, + enable_confidence_head=False, + compress_ratios=[1, 1], + ) + + +@pytest.fixture +def dspark_vllm_config(): + hf_config = _make_dsv4_dspark_hf_config() + model_config = SimpleNamespace( + dtype=torch.bfloat16, hf_config=hf_config, model="dspark" + ) + return SimpleNamespace( + model_config=model_config, + quant_config=None, + kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + tensor_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=True, + enable_eplb=True, + enable_elastic_ep=False, + eplb_config=SimpleNamespace(num_redundant_experts=4), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + speculative_config=SimpleNamespace( + method="dspark", + draft_model_config=model_config, + ), + ) + + +def _make_moe_topology( + *, + num_routed_experts: int = 8, + num_redundant_experts: int = 4, + num_expert_groups: int = 1, +) -> SimpleNamespace: + num_physical_experts = num_routed_experts + num_redundant_experts + return SimpleNamespace( + num_routed_experts=num_routed_experts, + num_redundant_experts=num_redundant_experts, + num_physical_experts=num_physical_experts, + num_logical_experts=num_routed_experts, + num_expert_groups=num_expert_groups, + ) + + +def test_eplb_state_accepts_matching_dsv4_draft_and_target(): + draft = _make_moe_topology() + target = _make_moe_topology() + + EplbState.assert_confs_equal(draft, target) + + +def test_eplb_state_rejects_mismatched_dsv4_draft_redundant_experts(): + draft = _make_moe_topology(num_redundant_experts=0) + target = _make_moe_topology(num_redundant_experts=4) + + with pytest.raises(RuntimeError, match="mismatch"): + EplbState.assert_confs_equal(draft, target) + + +@pytest.mark.parametrize( + ("model_type", "expected"), + [ + ("deepseek_v4", True), + ("deepseek_v41", False), + ], +) +def test_dspark_draft_supports_eplb( + dspark_vllm_config, model_type: str, expected: bool +): + draft_model_config = dspark_vllm_config.speculative_config.draft_model_config + draft_model_config.hf_config.model_type = model_type + assert dspark_draft_supports_eplb(draft_model_config) is expected + + +def test_draft_model_supports_eplb_requires_draft_moe(dspark_vllm_config): + assert not draft_model_supports_eplb( + dspark_vllm_config.speculative_config, + None, + ) + + +@pytest.mark.parametrize( + "model_type,load_dummy_weights,expected_registered", + [ + ("deepseek_v4", False, True), + ("deepseek_v4", True, False), + ("deepseek_v41", False, False), + ], +) +def test_eplb_dspark_draft_registration( + dspark_vllm_config, + monkeypatch: pytest.MonkeyPatch, + model_type, + load_dummy_weights: bool, + expected_registered: bool, +): + """Register V4 drafts; skip incompatible V4.1 drafts and dummy weights.""" + FakeEplbState.instances.clear() + monkeypatch.setattr( + "vllm.v1.worker.gpu.eplb_utils.EplbState", + FakeEplbState, + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.eplb_utils.get_mixture_of_experts_model", + lambda model: model, + ) + + draft_model_config = dspark_vllm_config.speculative_config.draft_model_config + draft_model_config.hf_config.model_type = model_type + draft = SimpleNamespace() + controller = EPLBController(dspark_vllm_config.parallel_config, torch.device("cpu")) + controller.prepare_load() + speculator = SimpleNamespace(model=draft, eplb_state=None) + + def set_eplb_state(state) -> None: + speculator.eplb_state = state + + speculator.set_eplb_state = set_eplb_state + + registered = controller.maybe_register_speculator( + speculator, + dspark_vllm_config.speculative_config, + load_dummy_weights=load_dummy_weights, + ) + + assert registered is expected_registered + assert controller.state is not None + if expected_registered: + assert controller.state.add_model_calls == [ + (draft, draft_model_config, "dspark (draft)") + ] + assert speculator.eplb_state is controller.state + else: + assert controller.state.add_model_calls == [] + assert speculator.eplb_state is None + + @pytest.mark.parametrize( "model_setup", [ pytest.param( - ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86), + ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86, 1), marks=large_gpu_mark(min_gb=80), ), pytest.param( @@ -65,20 +260,42 @@ def get_model_args( "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", 4, 0.92, + 1, ), marks=pytest.mark.skip(reason="Skipping due to CI OOM issues"), ), + pytest.param( + ( + "dspark", + "deepseek-ai/DeepSeek-V4-Flash-DSpark", + None, + 4, + 0.95, + 7, + ), + marks=large_gpu_mark(min_gb=80), + ), ], - ids=["qwen3_next_mtp", "llama4_eagle"], + ids=["qwen3_next_mtp", "llama4_eagle", "dsv4_dspark"], ) def test_eplb_spec_decode( monkeypatch: pytest.MonkeyPatch, - model_setup: tuple[str, str, str, int, float], + model_setup: tuple[str, str, str | None, int, float, int], ): """Test the correctness of EPLB speculative decoding with GSM8K dataset. - Applicable to MoE models with mtp or eagle spec decode. + + Applicable to MoE models with mtp, eagle, or dspark spec decode. """ - method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup + ( + method, + model_name, + spec_model_name, + tp_size, + expected_gsm8k_value, + num_speculative_tokens, + ) = model_setup + if method == "dspark" and not current_platform.is_device_capability_family(100): + pytest.skip("DSV4 DSpark EPLB requires SM100 for MegaMoE") TASK = "gsm8k" FILTER = "exact_match,strict-match" @@ -90,6 +307,7 @@ def test_eplb_spec_decode( spec_method=method, tp_size=tp_size, model_max_len=4096, + num_speculative_tokens=num_speculative_tokens, ) results = lm_eval.simple_evaluate( diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 76915b94a32d..250a9d729186 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -15,6 +15,10 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( deepgemm_post_process_weight_scale_block, ) +from vllm.model_executor.models.interfaces import ( + get_mixture_of_experts_model, + is_mixture_of_experts, +) from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4ForCausalLM, @@ -28,6 +32,7 @@ from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.torch_utils import set_default_torch_dtype @@ -913,6 +918,96 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): assert calls == ["mtp", "dspark"] +@pytest.fixture +def v4_dspark_config(dist_init): + hf_config = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], + hidden_size=128, + num_hidden_layers=2, + n_routed_experts=8, + num_experts_per_tok=2, + num_hash_layers=0, + n_shared_experts=1, + moe_intermediate_size=128, + hidden_act="silu", + swiglu_limit=10.0, + norm_topk_prob=True, + topk_method="noaux_tc", + routed_scaling_factor=1.5, + hc_mult=1, + hc_eps=1e-5, + rms_norm_eps=1e-5, + dspark_target_layer_ids=[0], + dspark_markov_rank=8, + index_topk=4, + head_dim=64, + num_attention_heads=4, + vocab_size=256, + n_mtp_layers=2, + enable_confidence_head=False, + compress_ratios=[1, 1], + q_lora_rank=32, + o_lora_rank=32, + qk_rope_head_dim=16, + o_groups=4, + sliding_window=128, + hc_sinkhorn_iters=1, + ) + model_config = SimpleNamespace( + dtype=torch.bfloat16, + hf_config=hf_config, + max_model_len=2048, + ) + return SimpleNamespace( + model_config=model_config, + quant_config=None, + kernel_config=SimpleNamespace( + moe_backend="deep_gemm_mega_moe", + enable_jit_warmup=False, + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + tensor_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=True, + enable_eplb=True, + eplb_config=SimpleNamespace(num_redundant_experts=4), + ), + attention_config=SimpleNamespace(backend=None), + cache_config=SimpleNamespace(block_size=64, cache_dtype="auto"), + scheduler_config=SimpleNamespace(max_num_batched_tokens=4), + compilation_config=SimpleNamespace(static_forward_context={}), + speculative_config=SimpleNamespace(draft_model_config=model_config), + use_v2_model_runner=False, + ) + + +def test_dspark_draft_registers_mixture_of_experts( + v4_dspark_config, monkeypatch, dist_init +): + """DSV4 DSpark MoE drafts expose MixtureOfExperts for EPLB registration.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM MegaMoE requires SM100") + + hf_config = v4_dspark_config.model_config.hf_config + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.dspark.get_current_vllm_config", + lambda: v4_dspark_config, + ) + + with ( + set_default_torch_dtype(v4_dspark_config.model_config.dtype), + torch.device("cuda"), + ): + draft = DSparkDeepseekV4ForCausalLM(vllm_config=v4_dspark_config) + + assert is_mixture_of_experts(draft) + assert get_mixture_of_experts_model(draft) is draft + assert len(draft.moe_layers) == draft.model.num_dspark_layers == 2 + assert draft.num_routed_experts == hf_config.n_routed_experts + assert draft.num_redundant_experts == 4 + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 53a93b940cd5..b2156af346c8 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -266,6 +266,57 @@ def load_scale(param, weight, *args): torch.testing.assert_close(param, checkpoint_scale) +def test_dsv4_dspark_fused_moe_expert_mapping_includes_redundant_experts( + monkeypatch, +): + from vllm.models.deepseek_v4.nvidia import dspark + + captured: dict[str, int] = {} + + def capture_expert_mapping(*args, **kwargs): + captured.update(kwargs) + return [] + + draft = SimpleNamespace( + config=SimpleNamespace( + num_attention_heads=4, + n_routed_experts=8, + expert_dtype="fp4", + ), + quant_config=None, + pad_shared_expert=False, + model=SimpleNamespace( + layers=[ + SimpleNamespace( + ffn=SimpleNamespace( + use_native_mega_moe=False, + use_mega_moe=False, + n_redundant_experts=4, + ) + ) + ], + confidence_head=None, + ), + named_parameters=lambda: [], + process_weights_after_loading=lambda: None, + ) + draft._remap_dspark_name = lambda name: ( + dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name) + ) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + dspark, + "fused_moe_make_expert_params_mapping", + capture_expert_mapping, + ) + + dspark.DSparkDeepseekV4ForCausalLM.load_weights(draft, []) + + assert captured["num_experts"] == 8 + assert captured["num_redundant_experts"] == 4 + + def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer(): weights = [ ("mtp.0.attn.wkv.weight", torch.arange(4)), diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py index 1014051d649c..891ca5602079 100644 --- a/tests/v1/worker/test_dspark_utils.py +++ b/tests/v1/worker/test_dspark_utils.py @@ -15,6 +15,16 @@ from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +@dataclass +class _FakeHFConfig: + model_type: str + + +@dataclass +class _FakeModelConfig: + hf_config: _FakeHFConfig + + @dataclass class _FakeEPLBConfig: num_redundant_experts: int = 0 @@ -37,12 +47,28 @@ def __post_init__(self) -> None: raise ValueError("elastic EP requires EPLB") -def test_dspark_parallel_config_disables_eplb_atomically(): +@pytest.mark.parametrize( + ("model_type", "enable_eplb", "num_redundant_experts", "shares_eplb_config"), + [ + ("deepseek_v4", True, 32, True), + ("deepseek_v41", False, 0, False), + ], +) +def test_dspark_parallel_config_eplb_by_model_type( + model_type: str, + enable_eplb: bool, + num_redundant_experts: int, + shares_eplb_config: bool, +): target_config = _FakeParallelConfig() + draft_model_config = _FakeModelConfig( + hf_config=_FakeHFConfig(model_type=model_type) + ) draft_config = _get_dspark_parallel_config( target_config, tensor_parallel_size=4, + draft_model_config=draft_model_config, ) assert target_config.pipeline_parallel_size == 2 @@ -54,10 +80,13 @@ def test_dspark_parallel_config_disables_eplb_atomically(): assert draft_config is not target_config assert draft_config.pipeline_parallel_size == 1 assert draft_config.tensor_parallel_size == 4 - assert not draft_config.enable_eplb - assert draft_config.eplb_config.num_redundant_experts == 0 + assert draft_config.enable_eplb is enable_eplb + assert draft_config.eplb_config.num_redundant_experts == num_redundant_experts assert not draft_config.enable_elastic_ep - assert draft_config.eplb_config is not target_config.eplb_config + if shares_eplb_config: + assert draft_config.eplb_config is target_config.eplb_config + else: + assert draft_config.eplb_config is not target_config.eplb_config @pytest.mark.parametrize("pcp_size", [1, 4]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 6b38bfef968d..b0acf6916015 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -29,15 +29,21 @@ class FakeEplbState: def __init__(self, parallel_config: Any, device: torch.device): self.parallel_config = parallel_config self.device = device - self.add_model_calls: list[tuple[Any, Any]] = [] + self.add_model_calls: list[tuple[Any, Any, str | None]] = [] self.step_calls: list[tuple[bool, bool, bool]] = [] self.async_started = False self.is_async = True self.update_mapping_calls: list[tuple[Any, torch.Tensor]] = [] FakeEplbState.instances.append(self) - def add_model(self, model: Any, model_config: Any) -> None: - self.add_model_calls.append((model, model_config)) + def add_model( + self, + model: Any, + model_config: Any, + *, + model_name: str | None = None, + ) -> None: + self.add_model_calls.append((model, model_config, model_name)) def step(self, is_dummy: bool, is_profile: bool, *, log_stats: bool) -> None: self.step_calls.append((is_dummy, is_profile, log_stats)) @@ -98,6 +104,15 @@ def _make_runner(**overrides: Any) -> Any: runner.execute_model_state = None for key, value in overrides.items(): setattr(runner, key, value) + spec_config = runner.speculative_config + runner._draft_workspace_lane = int( + spec_config is not None + and ( + spec_config.use_dspark() + if hasattr(spec_config, "use_dspark") + else getattr(spec_config, "method", None) == "dspark" + ) + ) return runner @@ -131,7 +146,7 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch): assert runner.model is model assert runner.model_state is not None assert runner.eplb_state is not None - assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] + assert runner.eplb_state.add_model_calls == [(model, runner.model_config, None)] assert runner.eplb_state.async_started is True @@ -159,7 +174,7 @@ def test_v2_load_model_with_dummy_weights_defers_eplb_async_loop(monkeypatch): assert runner.load_config.load_format == "dummy" assert runner.eplb_state is not None # A scaling-up worker registers so it can join the EPLB communicator. - assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] + assert runner.eplb_state.add_model_calls == [(model, runner.model_config, None)] assert runner.eplb_state.async_started is False @@ -174,6 +189,64 @@ def test_v2_setup_eplb_from_mapping_updates_state_in_place(): assert state.update_mapping_calls == [(runner.model_config, mapping)] +def test_v2_load_model_registers_dspark_speculator_with_eplb(monkeypatch): + FakeEplbState.instances.clear() + target_model = SimpleNamespace(is_moe=True) + draft_model = SimpleNamespace(is_moe=True) + draft_model_config = SimpleNamespace( + model="dspark-draft", + hf_config=SimpleNamespace(model_type="deepseek_v4"), + ) + + class FakeDSparkSpeculator: + def __init__(self): + self.model = draft_model + self.eplb_state = None + + def load_model(self, target_model): + return None + + def set_eplb_state(self, state) -> None: + self.eplb_state = state + + monkeypatch.setattr(mrv2, "DeviceMemoryProfiler", FakeMemoryProfiler) + monkeypatch.setattr(eplb, "EplbState", FakeEplbState) + monkeypatch.setattr(mrv2, "DraftModelSpeculator", FakeDSparkSpeculator) + monkeypatch.setattr( + mrv2, + "get_model_loader", + lambda load_config: SimpleNamespace(load_model=lambda **_: target_model), + ) + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) + monkeypatch.setattr( + eplb, + "get_mixture_of_experts_model", + lambda model: model if getattr(model, "is_moe", False) else None, + ) + + speculator = FakeDSparkSpeculator() + runner = _make_runner( + is_last_pp_rank=False, + speculative_config=SimpleNamespace( + method="dspark", + draft_model_config=draft_model_config, + ), + speculator=speculator, + ) + mrv2.GPUModelRunner.load_model(runner) + + assert runner.eplb_state is not None + assert runner.eplb_state.add_model_calls == [ + (draft_model, draft_model_config, "dspark-draft (draft)"), + (target_model, runner.model_config, None), + ] + assert speculator.eplb_state is runner.eplb_state + + def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): events = [] runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 205226cde4a9..2d782c417945 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -323,39 +323,47 @@ def validate_ep_configuration(self, new_model: MixtureOfExperts): """ if len(self.model_states) > 0: model = next(iter(self.model_states.values())).model - if ( - model.num_routed_experts != new_model.num_routed_experts - or model.num_redundant_experts != new_model.num_redundant_experts - or model.num_physical_experts != new_model.num_physical_experts - or model.num_logical_experts != new_model.num_logical_experts - or model.num_expert_groups != new_model.num_expert_groups - ): - raise RuntimeError( - "Model: {} " - "with config {} " - "{} {} {} {} " - "mismatch with new model {} " - "with config {} " - "{} {} {} {}".format( - type(model), - model.num_routed_experts, - model.num_redundant_experts, - model.num_physical_experts, - model.num_logical_experts, - model.num_expert_groups, - type(new_model), - new_model.num_routed_experts, - new_model.num_redundant_experts, - new_model.num_physical_experts, - new_model.num_logical_experts, - new_model.num_expert_groups, - ) + self.assert_confs_equal(model, new_model) + + @staticmethod + def assert_confs_equal( + model: MixtureOfExperts, new_model: MixtureOfExperts + ) -> None: + """Raise if the models have different expert parallel configurations.""" + if ( + model.num_routed_experts != new_model.num_routed_experts + or model.num_redundant_experts != new_model.num_redundant_experts + or model.num_physical_experts != new_model.num_physical_experts + or model.num_logical_experts != new_model.num_logical_experts + or model.num_expert_groups != new_model.num_expert_groups + ): + raise RuntimeError( + "Model: {} " + "with config {} " + "{} {} {} {} " + "mismatch with new model {} " + "with config {} " + "{} {} {} {}".format( + type(model), + model.num_routed_experts, + model.num_redundant_experts, + model.num_physical_experts, + model.num_logical_experts, + model.num_expert_groups, + type(new_model), + new_model.num_routed_experts, + new_model.num_redundant_experts, + new_model.num_physical_experts, + new_model.num_logical_experts, + new_model.num_expert_groups, ) + ) def add_model( self, model: MixtureOfExperts, model_config: ModelConfig, + model_name: str | None = None, ): """Build the initial EPLB state.""" self.validate_ep_configuration(model) @@ -506,7 +514,7 @@ def add_model( dtype=torch.int64, device=self.device, ), - model_name=model_config.model, + model_name=model_name or model_config.model, model=model, expert_buffer=expert_buffer, rebalanced=False, diff --git a/vllm/models/deepseek_v4/common/eplb_util.py b/vllm/models/deepseek_v4/common/eplb_util.py new file mode 100644 index 000000000000..53c6d100243d --- /dev/null +++ b/vllm/models/deepseek_v4/common/eplb_util.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared EPLB MoE registration helpers for DeepSeek V4 family models.""" + +from collections.abc import Iterable +from typing import Any + +import torch.nn as nn + +from vllm.config import ModelConfig + + +def dspark_draft_supports_eplb(draft_model_config: ModelConfig) -> bool: + """Return whether a DSpark draft can share EPLB state with the target. + + Only DeepSeek-V4 DSpark drafts reuse the target expert layout. V4.1 drafts + use a smaller routed-expert count and cannot share EPLB state. + """ + return getattr(draft_model_config.hf_config, "model_type", None) == "deepseek_v4" + + +def collect_moe_layers( + moe_model: Any, + layers: Iterable[nn.Module], + config: object, + *, + decoder_layer_type: type[nn.Module], + moe_type: type[nn.Module], +) -> None: + """Populate ``MixtureOfExperts`` fields from decoder layers for EPLB. + + Works for both target backbones and DSpark draft stacks. Draft models pass + their shorter ``model.layers`` list; pipeline-parallel placeholder layers + are skipped via ``decoder_layer_type``. + """ + moe_model.num_expert_groups = getattr(config, "n_group", 1) + moe_model.moe_layers = [] + moe_model.moe_mlp_layers = [] + example_moe = None + for layer in layers: + if not isinstance(layer, decoder_layer_type): + continue + if not isinstance(layer.ffn, moe_type): + continue + example_moe = layer.ffn + moe_model.moe_mlp_layers.append(layer.ffn) + moe_model.moe_layers.append(layer.ffn.experts) + + moe_model.num_moe_layers = len(moe_model.moe_layers) + moe_model.extract_moe_parameters(example_moe) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 54250185c840..2a40412e933b 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -52,11 +52,14 @@ sp_padding_mask, sp_shard, ) +from vllm.models.deepseek_v4.common.eplb_util import collect_moe_layers from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from .model import ( DeepseekV4DecoderLayer, + DeepseekV4MixtureOfExperts, DeepseekV4Model, + DeepseekV4MoE, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, prepare_mega_gate_routing_metadata, @@ -346,7 +349,7 @@ def _insert_context_kv( ) -class DSparkDeepseekV4ForCausalLM(nn.Module): +class DSparkDeepseekV4ForCausalLM(nn.Module, DeepseekV4MixtureOfExperts): # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so # load_dspark_model always aliases the target's. has_own_embed_tokens = False @@ -373,6 +376,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(self.config.vocab_size) + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + collect_moe_layers( + self, + self.model.layers, + self.config, + decoder_layer_type=DeepseekV4DecoderLayer, + moe_type=DeepseekV4MoE, + ) # --- Hooks used by the speculator ------------------------------------- @@ -451,6 +464,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.n_routed_experts, + num_redundant_experts=first_layer.ffn.n_redundant_experts, ) expert_scale_suffix = ( ".weight_scale" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index fc10cd07c1cf..ff4470dc4c97 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -81,6 +81,7 @@ sp_shard, ) from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.models.deepseek_v4.common.eplb_util import collect_moe_layers from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -1841,23 +1842,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.set_moe_parameters() def set_moe_parameters(self) -> None: - self.num_expert_groups = getattr(self.config, "n_group", 1) - self.num_moe_layers = self.config.num_hidden_layers - self.moe_layers: list[nn.Module] = [] - self.moe_mlp_layers: list[DeepseekV4MoE] = [] - example_moe: DeepseekV4MoE | None = None - for layer in self.model.layers: - if isinstance(layer, PPMissingLayer): - continue - if not isinstance(layer, DeepseekV4DecoderLayer): - continue - if isinstance(layer.ffn, DeepseekV4MoE): - example_moe = layer.ffn - self.moe_mlp_layers.append(layer.ffn) - self.moe_layers.append(layer.ffn.experts) - - self.num_moe_layers = len(self.moe_layers) - self.extract_moe_parameters(example_moe) + collect_moe_layers( + self, + self.model.layers, + self.config, + decoder_layer_type=DeepseekV4DecoderLayer, + moe_type=DeepseekV4MoE, + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index e5eb31ad28f0..5c2c5d14121e 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -13,8 +13,10 @@ from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( + MixtureOfExperts, get_mixture_of_experts_model, ) +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -22,6 +24,25 @@ logger = init_logger(__name__) +def eplb_draft_model_name(model_config: ModelConfig) -> str: + """Return the EPLB step-log display name for a draft model.""" + return f"{model_config.model} (draft)" + + +def draft_model_supports_eplb( + speculative_config: Any | None, + draft_moe_model: MixtureOfExperts | None, +) -> bool: + """Return whether a draft MoE model should register with EPLB.""" + if draft_moe_model is None: + return False + if speculative_config is None or speculative_config.draft_model_config is None: + return False + if getattr(speculative_config, "method", None) == "dspark": + return dspark_draft_supports_eplb(speculative_config.draft_model_config) + return True + + def step_eplb_after(*, is_dummy: bool = False) -> Callable: """Step EPLB after a model runner method completes successfully.""" @@ -93,16 +114,26 @@ def maybe_register_speculator( draft_model = speculator.model draft_moe_model = get_mixture_of_experts_model(draft_model) - if draft_moe_model is None: + if not draft_model_supports_eplb(speculative_config, draft_moe_model): return False + assert draft_moe_model is not None assert not self.parallel_config.enable_elastic_ep, ( "Elastic EP is not supported with draft model." ) assert speculative_config is not None - assert speculative_config.draft_model_config is not None + draft_model_config = speculative_config.draft_model_config + assert draft_model_config is not None + logger.info_once( + "EPLB is enabled for MoE part of drafter model %s.", + draft_model_config.model, + ) assert self.state is not None - self.state.add_model(draft_moe_model, speculative_config.draft_model_config) + self.state.add_model( + draft_moe_model, + draft_model_config, + model_name=eplb_draft_model_name(draft_model_config), + ) speculator.set_eplb_state(self.state) self._has_registered_models = True return True diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 489d0da554a2..38108d4e0414 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -6,6 +6,7 @@ from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.model_executor.model_loader.utils import get_draft_load_config +from vllm.models.deepseek_v4.common.eplb_util import dspark_draft_supports_eplb from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -35,11 +36,24 @@ def _resolve_dspark_attention_backend( def _get_dspark_parallel_config( parallel_config: ParallelConfig, tensor_parallel_size: int, + draft_model_config: ModelConfig, ) -> ParallelConfig: - if parallel_config.enable_eplb: + enable_draft_eplb = parallel_config.enable_eplb and dspark_draft_supports_eplb( + draft_model_config + ) + if parallel_config.enable_eplb and not enable_draft_eplb: logger.warning_once( - "EPLB is disabled for the DSpark draft model. EPLB remains enabled " - "for the target model." + "EPLB is disabled for the DSpark draft model (%s). EPLB remains " + "enabled for the target model.", + draft_model_config.hf_config.model_type, + ) + + if enable_draft_eplb: + return replace( + parallel_config, + pipeline_parallel_size=1, + tensor_parallel_size=tensor_parallel_size, + enable_elastic_ep=False, ) return replace( @@ -81,6 +95,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo parallel_config=_get_dspark_parallel_config( vllm_config.parallel_config, speculative_config.draft_parallel_config.tensor_parallel_size, + draft_model_config, ), attention_config=replace( vllm_config.attention_config, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4db3f00b91d4..669b314d4cac 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -206,6 +206,10 @@ ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin +from vllm.v1.worker.gpu.eplb_utils import ( + draft_model_supports_eplb, + eplb_draft_model_name, +) from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin @@ -5237,25 +5241,34 @@ def load_model( logger.info_once("Loading drafter model...") if hasattr(self.drafter, "load_model"): self.drafter.load_model(self.model) + drafter_moe_model = None + if self.parallel_config.enable_eplb and hasattr( + self.drafter, "model" + ): + drafter_moe_model = get_mixture_of_experts_model( + self.drafter.model + ) + if not draft_model_supports_eplb( + self.vllm_config.speculative_config, + drafter_moe_model, + ): + drafter_moe_model = None + if ( self.parallel_config.enable_eplb and hasattr(self.drafter, "model") - and ( - drafter_moe_model := get_mixture_of_experts_model( - self.drafter.model - ) - ) - is not None + and drafter_moe_model is not None ): assert not self.parallel_config.enable_elastic_ep, ( "Elastic EP is not supported with drafter model." ) spec_config = self.vllm_config.speculative_config assert spec_config is not None - assert spec_config.draft_model_config is not None + draft_model_config = spec_config.draft_model_config + assert draft_model_config is not None logger.info_once( "EPLB is enabled for MoE part of drafter model %s.", - spec_config.draft_model_config.model, + draft_model_config.model, ) if self.eplb_state is None: self.eplb_state = EplbState( @@ -5263,7 +5276,8 @@ def load_model( ) self.eplb_state.add_model( drafter_moe_model, - spec_config.draft_model_config, + draft_model_config, + model_name=eplb_draft_model_name(draft_model_config), ) assert hasattr(self.drafter, "set_eplb_state") self.drafter.set_eplb_state(self.eplb_state) From 716961cf1fc474c491d401adece882af3b39d5ab Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Mon, 12 Oct 2026 10:05:09 -0400 Subject: [PATCH 419/952] [Bugfix] Drop enable_in_kernel_fc2_reduce from FlashInfer CuteDSL MegaMoE config (#59900) Signed-off-by: Shang Wang Co-authored-by: Claude Opus 5.5 --- tests/kernels/moe/test_flashinfer_moe_ep.py | 62 +++++++++++++++++++ .../layers/fused_moe/flashinfer_moe_ep.py | 3 +- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py index c3223b7d9800..7ce3124ae144 100644 --- a/tests/kernels/moe/test_flashinfer_moe_ep.py +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -10,9 +10,14 @@ the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per 32). The expected output is then known in closed form and compared bit for bit. + +A test that needs no SM100 also checks the adapter's kwargs against the +installed FlashInfer's config types, so config drift fails on any CI runner +with FlashInfer installed. """ from dataclasses import dataclass +from types import SimpleNamespace import pytest import torch @@ -22,6 +27,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM from vllm.distributed.parallel_state import ensure_model_parallel_initialized +from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep as fi_ep from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -305,3 +311,59 @@ def test_init_warmup_forward_is_exact_on_grid_inputs( if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs") distributed_run(_init_warmup_forward, world_size, geometry, backend) + + +@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm")) +def test_adapter_kwargs_match_installed_flashinfer( + monkeypatch: pytest.MonkeyPatch, backend: str +): + """The adapter's kwargs construct the installed FlashInfer's config types. + + Needs no SM100, so config kwarg drift fails in CI rather than at model load. + Fakes only the megakernel layer (symmetric memory), the EP group, the device + index and the DeepGEMM import. + """ + moe_ep = pytest.importorskip("flashinfer.moe_ep") + layer_args: list[tuple] = [] + + class FakeMegaLayer: + def __init__(self, *args) -> None: + layer_args.append(args) + + def destroy(self) -> None: + pass + + monkeypatch.setattr(moe_ep, "MoEEpMegaLayer", FakeMegaLayer) + monkeypatch.setattr(fi_ep, "_expose_deep_gemm_to_flashinfer", lambda: None) + monkeypatch.setattr( + fi_ep, + "get_ep_group", + lambda: SimpleNamespace(world_size=2, rank_in_group=0, device_group=None), + ) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + geometry = GEOMETRIES[0] + moe = SimpleNamespace( + moe_backend=backend, + num_experts=geometry.num_experts, + experts_per_token=geometry.top_k, + hidden_dim=geometry.hidden_size, + intermediate_size=geometry.intermediate_size, + max_num_tokens=MAX_TOKENS_PER_RANK, + swiglu_limit=10.0, + routing_method=_routing_method(backend), + ) + weight = torch.zeros(1, dtype=torch.bfloat16) + FlashInferMoeEp( + moe, + FlashInferMoeEpWeights(w13=weight, w2=weight), + apply_topk_in_fc1=apply_topk_in_fc1(moe), + ) + + ((_, _, _, mega_config),) = layer_args + expected = ( + moe_ep.Nvfp4CutedslMegaMoeConfig + if backend == FLASHINFER_MOE_EP_CUTEDSL + else moe_ep.DeepGemmMegaMoeConfig + ) + assert isinstance(mega_config.megakernel, expected) diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py index bac5345255c9..63345a417438 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -315,13 +315,14 @@ def __init__( w2_scale=weights.w2_scale, ) if spec.kernel == "cutedsl": + # In-kernel fc2 reduce keeps its default (off); FlashInfer 0.7.1rc1 + # renamed in_kernel_fc2_reduce to enable_in_kernel_fc2_reduce. megakernel = api.Nvfp4CutedslMegaMoeConfig( intermediate_size=moe.intermediate_size, top_k=moe.experts_per_token, gate_up_clamp=moe.swiglu_limit, fast_math=True, apply_topk_in_fc1=apply_topk_in_fc1, - enable_in_kernel_fc2_reduce=False, combine_dtype="bf16", input_norm_const=epilogue.input_norm_const, fc1_alpha=None, From f3cb10b690c352265048991901de0e62b65abe8c Mon Sep 17 00:00:00 2001 From: Hoai Vo <62198968+hoaivoke@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:09:53 -0400 Subject: [PATCH 420/952] [Model][LoRA] Enable LoRA support for BERT and RoBERTa embedding models (#61033) Signed-off-by: hoaivoke Co-authored-by: Claude Opus 5.5 --- docs/models/pooling_models/embed.md | 6 +++--- vllm/model_executor/models/bert.py | 17 +++++++++++++++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/docs/models/pooling_models/embed.md b/docs/models/pooling_models/embed.md index 4c33126977e9..5211997edb81 100644 --- a/docs/models/pooling_models/embed.md +++ b/docs/models/pooling_models/embed.md @@ -37,7 +37,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th | Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | ------------------------------ | ------------------------------------------ | -| `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | +| `BertModel` | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | ✅︎ | | | `BertSpladeSparseEmbeddingModel` | SPLADE | `naver/splade-v3` | | | | `BgeM3EmbeddingModel` | BGE-M3 | `BAAI/bge-m3` | | | | `Gemma2Model`C | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | @@ -51,9 +51,9 @@ You can compute pairwise similarity scores to build a similarity matrix using th | `NomicBertModel` | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | | `Qwen2Model`C, `Qwen2ForCausalLM`C | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | | `Qwen3Model`C, `Qwen3ForCausalLM`C | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, `ai-sage/Giga-Embeddings-instruct-480M-0826`^, etc. | ✅︎ | ✅︎ | -| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | +| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | ✅︎ | | | `VoyageQwen3BidirectionalEmbedModel`C | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ | -| `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | | | +| `XLMRobertaModel` | XLMRobertaModel-based | `BAAI/bge-m3` (see note), `intfloat/multilingual-e5-base`, `jinaai/jina-embeddings-v3` (see note), etc. | ✅︎ | | | `*Model`C, `*ForCausalLM`C, etc. | Generative models | N/A | \* | \* | !!! note diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index cdaa0bce9725..508748681261 100644 --- a/vllm/model_executor/models/bert.py +++ b/vllm/model_executor/models/bert.py @@ -41,7 +41,7 @@ from vllm.tasks import PoolingTask from vllm.v1.pool.metadata import PoolingMetadata -from .interfaces import SupportsCrossEncoding, SupportsQuant +from .interfaces import SupportsCrossEncoding, SupportsLoRA, SupportsQuant from .interfaces_base import attn_type, default_pooling_type from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -448,7 +448,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @default_pooling_type(seq_pooling_type="CLS") -class BertEmbeddingModel(nn.Module, SupportsQuant): +class BertEmbeddingModel(nn.Module, SupportsQuant, SupportsLoRA): """A model that uses Bert to provide embedding functionalities. This class encapsulates the BertModel and provides an interface for @@ -463,9 +463,22 @@ class BertEmbeddingModel(nn.Module, SupportsQuant): is_pooling_model = True embedding_class: type[nn.Module] = BertEmbedding + packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]} + # PEFT adapters name modules relative to the HF backbone (optionally under + # `bert.`/`roberta.`), while vLLM nests the backbone under `model.`. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "bert.": "model.", + "roberta.": "model.", + "encoder.": "model.encoder.", + "embeddings.": "model.embeddings.", + } + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + self.config = vllm_config.model_config.hf_config pooler_config = vllm_config.model_config.pooler_config assert pooler_config is not None From ffec221276fa39d7927b9bb63dc12766d77fc403 Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Mon, 12 Oct 2026 10:23:00 -0400 Subject: [PATCH 421/952] [Bugfix][DSv4.1] Match decoder-replay input_ids dtype for FlashInfer MegaMoE (#61041) Signed-off-by: Shang Wang Co-authored-by: Claude Opus 5.5 --- .../test_deepseek_v41_decoder_replay_layers.py | 17 +++++++++++++++-- .../nvidia/decoder_replay_cudagraph.py | 7 ++++--- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 614d903ca10d..b7f5ca29ec2d 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -8,6 +8,7 @@ import torch from vllm.config import CompilationConfig, CUDAGraphMode +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import ( BatchDescriptor, ForwardContext, @@ -144,8 +145,18 @@ def test_piecewise_graphs_break_out_by_size(): write_batch_kv.assert_called_once_with(*states) -def test_replay_graph_matches_eager(monkeypatch): - """The replay graph of the next captured size pads and runs the rows.""" +@pytest.mark.parametrize( + ("moe_backend", "ids_dtype"), + [ + ("auto", torch.int32), + *((b, torch.int64) for b in sorted(NATIVE_MEGA_MOE_BACKENDS)), + *((b, torch.int32) for b in sorted(FLASHINFER_MOE_EP_BACKENDS)), + ], + ids=str, +) +def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype): + """The replay graph of the next captured size pads and runs the rows. Input ids + arrive as the model passes them: int32, cast to int64 for native MegaMoE only.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) capture = torch.cuda.stream(torch.cuda.Stream()) monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) @@ -157,6 +168,7 @@ def test_replay_graph_matches_eager(monkeypatch): cfg.cache_config.use_kda_recoverssm = False cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN) cfg.model_config.dtype = torch.float32 + cfg.kernel_config.moe_backend = moe_backend def run_layers(hidden, positions, ids, pre, post, mix, residual): return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None] @@ -173,6 +185,7 @@ def prepare(desc): torch.randn(12, *b.shape[1:], device=DEVICE).to(b.dtype) for b in manager.input_buffers ] + states[2] = torch.arange(12, device=DEVICE, dtype=ids_dtype) # input_ids rows = torch.tensor([11, 2, 7], device=DEVICE) desc = manager.dispatch(2, 3, None, 0) with override_forward_context(_graph_context(desc.num_tokens)): diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py index 871f535584c6..c985f973bb49 100644 --- a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py +++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py @@ -10,7 +10,7 @@ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.config import CUDAGraphMode, VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.forward_context import get_forward_context, override_forward_context from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -56,8 +56,9 @@ def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers) config, act = vllm_config.model_config.hf_config, vllm_config.model_config.dtype assert isinstance(act, torch.dtype) hc, h, f32 = config.hc_mult, config.hidden_size, torch.float32 - mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS - ids = torch.int64 if mega_moe else torch.int32 + # DeepseekV4Model.forward casts input_ids to int64 for native MegaMoE only. + native = vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ids = torch.int64 if native else torch.int32 # hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual shapes = [(h,), (), (), (hc,), (hc, 1), (hc, hc), (hc, h)] dtypes = [act, torch.int64, ids, f32, f32, f32, act] From a686c8e1b7e5870eadfc9f47447de2b8a26211fc Mon Sep 17 00:00:00 2001 From: Nick Riasanovsky Date: Mon, 12 Oct 2026 11:38:43 -0400 Subject: [PATCH 422/952] [Triton] Migrate Triton kernels off tl.make_block_ptr for triton-main compat (#60502) Signed-off-by: Nick Riasanovsky --- .../amd/ops/third_party/kda/kernels.py | 357 +++++----------- .../nvidia/ops/third_party/kda/kernels.py | 357 +++++----------- .../inkling/amd/ops/fa4_rel_attention.py | 40 +- .../nvidia/ops/third_party/kda/chunk.py | 260 ++++------- .../nvidia/ops/third_party/kda/chunk_intra.py | 274 +++++------- .../kda/chunk_intra_token_parallel.py | 16 +- .../minimax_m3/common/ops/index_topk.py | 18 +- .../ops/chunk_delta_h.py | 182 ++++---- .../flash_linear_attention/ops/chunk_o.py | 49 +-- .../ops/chunk_scaled_dot_kkt.py | 35 +- .../flash_linear_attention/ops/cumsum.py | 69 +-- .../ops/fused_norm_gate.py | 34 +- .../flash_linear_attention/ops/kda.py | 402 ++++++------------ .../flash_linear_attention/ops/l2norm.py | 15 +- .../flash_linear_attention/ops/solve_tril.py | 172 +++----- .../flash_linear_attention/ops/wy_fast.py | 75 ++-- 16 files changed, 838 insertions(+), 1517 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py index 9805b3747130..0e32c9454351 100644 --- a/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/amd/ops/third_party/kda/kernels.py @@ -228,13 +228,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -255,58 +255,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -333,12 +326,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -357,33 +350,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -536,12 +510,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -550,108 +524,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -732,13 +655,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -754,70 +677,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -896,12 +793,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -911,26 +808,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -951,12 +837,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1271,7 +1157,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1280,25 +1166,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1321,7 +1196,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py index 8f662f68aa24..da498b5855ed 100644 --- a/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py +++ b/vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py @@ -230,13 +230,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -257,58 +257,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -335,12 +328,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -359,33 +352,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -538,12 +512,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -552,108 +526,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -734,13 +657,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -756,70 +679,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_h = m_v[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -898,12 +795,12 @@ def kda_gate_cumsum_fwd_kernel( HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -913,26 +810,15 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] @@ -953,12 +839,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on + # out-of-bounds output rows, which are masked away by `mask` on # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1273,7 +1159,7 @@ def kda_gate_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.load(A + i_h).to(tl.float32) @@ -1282,25 +1168,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1323,7 +1198,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/inkling/amd/ops/fa4_rel_attention.py b/vllm/models/inkling/amd/ops/fa4_rel_attention.py index d448ac8e2c5f..936b51eb6ea4 100644 --- a/vllm/models/inkling/amd/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/amd/ops/fa4_rel_attention.py @@ -165,15 +165,23 @@ def _inkling_rel_attention_kernel( q_head_start = pid_kh * gqa_group_size bt_row = block_table_ptr + pid_b * stride_bt_b - q_block_ptr = tl.make_block_ptr( - base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_q_t, stride_q_h, stride_q_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + o0 = q_block + tl.arange(0, BLOCK_Q) + o1 = tl.arange(0, BLOCK_H) + o2 = tl.arange(0, BLOCK_D) + m_q = ( + (o0 < q_len)[:, None, None] + & (o1 < gqa_group_size)[None, :, None] + & (o2 < head_dim)[None, None, :] ) - q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") + q_block_ptr = ( + q_ptr + + q_start * stride_q_t + + q_head_start * stride_q_h + + o0[:, None, None] * stride_q_t + + o1[None, :, None] * stride_q_h + + o2[None, None, :] * stride_q_d + ) + q = tl.load(q_block_ptr, mask=m_q, other=0.0) q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) q_rows = q_block + tl.arange(0, BLOCK_Q) @@ -265,18 +273,18 @@ def _inkling_rel_attention_kernel( acc /= l_i[:, None] acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) - out_block_ptr = tl.make_block_ptr( - base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, - shape=(q_len, gqa_group_size, head_dim), - strides=(stride_o_t, stride_o_h, stride_o_d), - offsets=(q_block, 0, 0), - block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), - order=(2, 1, 0), + out_block_ptr = ( + out_ptr + + q_start * stride_o_t + + q_head_start * stride_o_h + + o0[:, None, None] * stride_o_t + + o1[None, :, None] * stride_o_h + + o2[None, None, :] * stride_o_d ) tl.store( out_block_ptr, acc.to(out_ptr.dtype.element_ty), - boundary_check=(0, 1, 2), + mask=m_q, ) diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py index 3911d20d0358..19fcd9a686bf 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -69,12 +69,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -83,108 +83,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0).to(tl.float32) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -265,13 +214,13 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -287,70 +236,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -434,12 +357,12 @@ def kda_gate_chunk_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -472,34 +395,15 @@ def kda_gate_chunk_cumsum_vector_kernel( i_s -= 1 - p_s = tl.make_block_ptr( - s + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) if HAS_BIAS: - p_bias = tl.make_block_ptr( - g_bias + i_h * S, - (S,), - (1,), - (i_s * BS,), - (BS,), - (0,), - ) - b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_bias = tl.load(g_bias + i_h * S + o_s, mask=o_s < S, other=0.0).to(tl.float32) b_s += b_bias[None, :] b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) @@ -517,7 +421,7 @@ def kda_gate_chunk_cumsum_vector_kernel( # Boundary loads return zero, but bias and gate activation can make padded # rows nonzero. Padding trails valid rows, so it only affects masked stores. b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s) def fused_kda_gate_chunk_cumsum( @@ -843,33 +747,19 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) - stride_row = H * D - stride_col = 1 - - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + y_ptr = y + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -887,7 +777,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py index 07b4a388936e..286d23c814d0 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -70,14 +70,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( 5. Computes merged Akk_inv 6. Writes Akk_inv to Akk """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -106,6 +106,15 @@ def chunk_kda_fwd_kernel_inter_solve_fused( m_tc1 = (i_tc1 + o_i) < T m_tc2 = (i_tc2 + o_i) < T m_tc3 = (i_tc3 + o_i) < T + o_c0 = i_tc0 + o_i + o_c1 = i_tc1 + o_i + o_c2 = i_tc2 + o_i + o_c3 = i_tc3 + o_i + m_tc0 = o_c0 < T + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) @@ -129,29 +138,21 @@ def chunk_kda_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - p_g0 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) - ) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + m_ck0 = m_tc0[:, None] & m_k[None, :] + p_k0 = k + o_c0[:, None] * (H * K) + o_k[None, :] + p_g0 = g + o_c0[:, None] * (HV * K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_ck0, other=0.0).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_ck0, other=0.0).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_k1 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) - p_g1 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) - ) + m_ck1 = m_tc1[:, None] & m_k[None, :] + p_q1 = q + o_c1[:, None] * (H * K) + o_k[None, :] + p_k1 = k + o_c1[:, None] * (H * K) + o_k[None, :] + p_g1 = g + o_c1[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + b_q1 = tl.load(p_q1, mask=m_ck1, other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_ck1, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_ck1, other=0.0).to(tl.float32) # [BK] b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -163,19 +164,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) if i_tc2 < T: - p_q2 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_k2 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) - p_g2 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) - ) + m_ck2 = m_tc2[:, None] & m_k[None, :] + p_q2 = q + o_c2[:, None] * (H * K) + o_k[None, :] + p_k2 = k + o_c2[:, None] * (H * K) + o_k[None, :] + p_g2 = g + o_c2[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + b_q2 = tl.load(p_q2, mask=m_ck2, other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_ck2, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_ck2, other=0.0).to(tl.float32) # [BK] b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -195,19 +191,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk21 += tl.dot(b_kg2, b_kgt) if i_tc3 < T: - p_q3 = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_k3 = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) - p_g3 = tl.make_block_ptr( - g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) - ) + m_ck3 = m_tc3[:, None] & m_k[None, :] + p_q3 = q + o_c3[:, None] * (H * K) + o_k[None, :] + p_k3 = k + o_c3[:, None] * (H * K) + o_k[None, :] + p_g3 = g + o_c3[:, None] * (HV * K) + o_k[None, :] # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + b_q3 = tl.load(p_q3, mask=m_ck3, other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_ck3, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_ck3, other=0.0).to(tl.float32) # [BK] b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( tl.float32 @@ -236,82 +227,44 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # save off-diagonal Aqk blocks and prepare Akk ################################################################################ if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk10 = Aqk + o_c1[:, None] * (HV * BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_A1) - p_b1 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) - ) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + p_b1 = beta + bos * HV + i_hv + o_c1 * HV + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) b_Akk10 = b_Akk10 * b_b1[:, None] if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Aqk21 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) + p_Aqk20 = Aqk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Aqk21 = Aqk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_A2) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_A2) - p_b2 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) - ) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + p_b2 = beta + bos * HV + i_hv + o_c2 * HV + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Aqk31 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Aqk32 = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - tl.store( - p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - tl.store( - p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) - ) - - p_b3 = tl.make_block_ptr( - beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) - ) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + p_Aqk30 = Aqk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Aqk31 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Aqk32 = Aqk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + + p_b3 = beta + bos * HV + i_hv + o_c3 * HV + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) b_Akk30 = b_Akk30 * b_b3[:, None] b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + p_Akk00 = Akkd + o_c0[:, None] * (HV * BC) + o_i[None, :] + p_Akk11 = Akkd + o_c1[:, None] * (HV * BC) + o_i[None, :] + p_Akk22 = Akkd + o_c2[:, None] * (HV * BC) + o_i[None, :] + p_Akk33 = Akkd + o_c3[:, None] * (HV * BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_A0, other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_A1, other=0.0).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, mask=m_A2, other=0.0).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, mask=m_A3, other=0.0).to(tl.float32) ################################################################################ # forward substitution on diagonals @@ -397,47 +350,27 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # store full Akk_inv to Akk ################################################################################ - p_Akk00 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk10 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) - ) - p_Akk20 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk21 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk30 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - p_Akk31 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Akk32 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) - ) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Akk00 = Akk + o_c0[:, None] * (HV * BT) + o_i[None, :] + p_Akk10 = Akk + o_c1[:, None] * (HV * BT) + o_i[None, :] + p_Akk11 = Akk + o_c1[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk20 = Akk + o_c2[:, None] * (HV * BT) + o_i[None, :] + p_Akk21 = Akk + o_c2[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk22 = Akk + o_c2[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk30 = Akk + o_c3[:, None] * (HV * BT) + o_i[None, :] + p_Akk31 = Akk + o_c3[:, None] * (HV * BT) + (o_i + BC)[None, :] + p_Akk32 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 2 * BC)[None, :] + p_Akk33 = Akk + o_c3[:, None] * (HV * BT) + (o_i + 3 * BC)[None, :] + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_A0) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_A3) @triton.heuristics( @@ -474,14 +407,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( IS_VARLEN: tl.constexpr, USE_GATHER: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -505,16 +438,19 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * HV + i_hv) * BT Akk = Akk + (bos * HV + i_hv) * BC - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_ck = m_c[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + o_c[:, None] * (HV * K) + o_k[None, :] - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + p_beta = beta + o_c * HV - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_g = tl.load(p_g, mask=m_ck, other=0.0) + b_beta = tl.load(p_beta, mask=m_c, other=0.0).to(tl.float32) if USE_GATHER: b_gn = gather( @@ -546,12 +482,12 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) - ) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + m_Aqk_st = m_c[:, None] & (o_i[None, :] < BT) + m_Akk_st = m_c[:, None] & (o_i[None, :] < BC) + p_Aqk = Aqk + o_c[:, None] * (HV * BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (HV * BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_Aqk_st) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_Akk_st) tl.debug_barrier() @@ -566,7 +502,7 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_Akk_st) def chunk_kda_fwd_intra( diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py index 86654f264431..7cb68b2fd49a 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -108,20 +108,16 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr( - g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :] + p_beta = beta + i_t * HV + o_hv + b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32) + b_beta = tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32) b_k *= b_beta[:, None] for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr( - g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) - ) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :] + b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32) b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 025c3dd428b4..2d612c12f254 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -115,15 +115,17 @@ def _index_block_score_kernel( if BLOCK_SIZE_Q * pid_q >= q_len: return - q_ptrs = tl.make_block_ptr( - base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, - shape=(q_len, head_dim), - strides=(stride_q_n, stride_q_d), - offsets=(pid_q * BLOCK_SIZE_Q, 0), - block_shape=(BLOCK_SIZE_Q, head_dim), - order=(1, 0), + o_q = pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q) + o_d = tl.arange(0, head_dim) + m_q = o_q < q_len + q_ptrs = ( + q_ptr + + seq_start.to(tl.int64) * stride_q_n + + pid_h.to(tl.int64) * stride_q_h + + o_q[:, None].to(tl.int64) * stride_q_n + + o_d[None, :].to(tl.int64) * stride_q_d ) - q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q = tl.load(q_ptrs, mask=m_q[:, None], other=0.0) # A matched fp8 pair (e.g. e4m3 x e4m3) lowers to a native FP8 MMA, so # upcast only when the operands differ: there is no mixed-dtype fp8 MMA, # including across fp8 flavours (fp8e4nv vs fp8e4b8). bf16 and fp32 loads diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py index eb1c3af15297..cc27437fac58 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py @@ -110,109 +110,82 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht = ht + i_nh * V * K # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) if K > 64: - p_h0_2 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) if K > 128: - p_h0_3 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) if K > 192: - p_h0_4 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): - p_h1 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 0), - (BV, 64), - (1, 0), - ) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + o_t = (i_t * BT + tl.arange(0, BT)).to(tl.int64) + m_t = o_t < T + p_h1 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) if K > 64: - p_h2 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 64), - (BV, 64), - (1, 0), - ) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) if K > 128: - p_h3 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 128), - (BV, 64), - (1, 0), - ) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) if K > 192: - p_h4 = tl.make_block_ptr( - h + i_t.to(tl.int64) * stride_h, - (V, K), - (K, 1), - (i_v * BV, 192), - (BV, 64), - (1, 0), - ) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = h + i_t.to(tl.int64) * stride_h + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) if K > 128: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) if K > 192: - p_w = tl.make_block_ptr( - w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) - ) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) - p_v = tl.make_block_ptr( - v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + p_v = v + o_t[:, None] * stride_v + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr( - v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + p_v = v_new + o_t[:, None] * stride_v + o_v[None, :] + tl.store( + p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :] ) - tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) last_idx = min((i_t.to(tl.int64) + 1) * BT, T) - 1 if USE_G: - m_t = (i_t.to(tl.int64) * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) if USE_EXP2: b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) @@ -239,7 +212,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h1 *= exp(b_gk_last1)[None, :] if K > 64: - o_k2 = 64 + o_k1 b_gk_last2 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k2, mask=(o_k2 < K), @@ -250,7 +222,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h2 *= exp(b_gk_last2)[None, :] if K > 128: - o_k3 = 128 + o_k1 b_gk_last3 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k3, mask=(o_k3 < K), @@ -261,7 +232,6 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h3 *= exp(b_gk_last3)[None, :] if K > 192: - o_k4 = 192 + o_k1 b_gk_last4 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k4, mask=(o_k4 < K), @@ -273,48 +243,38 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 *= exp(b_gk_last4)[None, :] b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k1[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) b_h1 += tl.trans(tl.dot(b_k, b_v)) if K > 64: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) b_h2 += tl.trans(tl.dot(b_k, b_v)) if K > 128: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) b_h3 += tl.trans(tl.dot(b_k, b_v)) if K > 192: - p_k = tl.make_block_ptr( - k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) b_h4 += tl.trans(tl.dot(b_k, b_v)) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) if K > 64: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) if K > 128: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) if K > 192: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) def chunk_gated_delta_rule_fwd_h( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_o.py b/vllm/third_party/flash_linear_attention/ops/chunk_o.py index 0c323b8ce215..219c8ed43f71 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_o.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_o.py @@ -60,14 +60,18 @@ def chunk_fwd_kernel_o( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = ( + tl.program_id(0), + tl.program_id(1).to(tl.int64), + tl.program_id(2).to(tl.int64), + ) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_tg = i_t i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -90,22 +94,21 @@ def chunk_fwd_kernel_o( b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_v = i_v * BV + tl.arange(0, BV) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) - ) - p_h = tl.make_block_ptr( - h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_q = q + o_t[:, None] * (Hg * K) + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * (Hg * K) + p_h = h + o_v[:, None] * K + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & m_k[None, :], other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & m_t[None, :], other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] b_o += tl.dot(b_q, tl.trans(b_h)) @@ -114,28 +117,22 @@ def chunk_fwd_kernel_o( if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_o = b_o * exp(b_g)[:, None] b_A = b_A * exp(b_g[:, None] - b_g[None, :]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr( - v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - p_o = tl.make_block_ptr( - o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) # to fix mma -> mma layout conversion # already solved by triton v3.2 or higher b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) def chunk_fwd_o( diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py index 10ff970c809f..f610fc7031be 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py @@ -61,12 +61,12 @@ def chunk_scaled_dot_kkt_fwd_kernel( USE_G: tl.constexpr, CAST_DOT_TO_K_DTYPE: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -78,22 +78,19 @@ def chunk_scaled_dot_kkt_fwd_kernel( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + bos * H + i_h + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * Hg + i_h // (H // Hg)) * K, - (T, K), - (Hg * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + o_k = i_k * BK + tl.arange(0, BK) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_kb = b_k * b_beta[:, None] if CAST_DOT_TO_K_DTYPE: # RDNA: force operands to k's native dtype so WMMA is used. @@ -103,17 +100,15 @@ def chunk_scaled_dot_kkt_fwd_kernel( b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype)) if USE_G: - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A = b_A * exp(b_g_diff) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (BT * H) + tl.arange(0, BT)[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None]) def chunk_scaled_dot_kkt_fwd( diff --git a/vllm/third_party/flash_linear_attention/ops/cumsum.py b/vllm/third_party/flash_linear_attention/ops/cumsum.py index b0820104b1a1..1fcab916caf2 100644 --- a/vllm/third_party/flash_linear_attention/ops/cumsum.py +++ b/vllm/third_party/flash_linear_attention/ops/cumsum.py @@ -37,12 +37,12 @@ def chunk_local_cumsum_scalar_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -52,23 +52,21 @@ def chunk_local_cumsum_scalar_kernel( else: bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) - p_o = tl.make_block_ptr( - o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) - ) + p_s = s + bos * H + i_h * T + o_t + p_o = o + bos * H + i_h * T + o_t else: - p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_s = s + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32) b_o = tl.cumsum(b_s, axis=0) if REVERSE: b_z = tl.sum(b_s, axis=0) b_o = -b_o + b_z[None] + b_s - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -96,12 +94,12 @@ def chunk_local_cumsum_vector_kernel( IS_VARLEN: tl.constexpr, HEAD_FIRST: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -117,44 +115,19 @@ def chunk_local_cumsum_vector_kernel( else: m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_ts = (o_t[:, None] < T) & (o_s[None, :] < S) if HEAD_FIRST: - p_s = tl.make_block_ptr( - s + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h * T) * S, - (T, S), - (S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + p_s = s + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] + p_o = o + (bos * H + i_h * T) * S + o_t[:, None] * S + o_s[None, :] else: - p_s = tl.make_block_ptr( - s + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * S, - (T, S), - (H * S, 1), - (i_t * BT, i_s * BS), - (BT, BS), - (1, 0), - ) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H * S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_ts, other=0.0).to(tl.float32) b_o = tl.dot(m_s, b_s, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_ts) def chunk_local_cumsum_scalar( diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py index 8f58ddcf8c7f..51ef43270713 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -50,31 +50,30 @@ def layer_norm_gated_fwd_kernel( HAS_BIAS: tl.constexpr, launch_pdl: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] if launch_pdl: tl.extra.cuda.gdc_wait() tl.extra.cuda.gdc_launch_dependents() - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -82,8 +81,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -99,7 +98,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -112,8 +110,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index b5c3a92b65e7..9327b24ddd0d 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -178,27 +178,26 @@ def layer_norm_gated_fwd_kernel( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) o_d = tl.arange(0, BD) m_d = o_d < D + m_t = o_t < T + m_x = m_t[:, None] & m_d[None, :] - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + p_x = x + o_t[:, None] * D + o_d[None, :] + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) if HAS_RESIDUAL: - p_res = tl.make_block_ptr( - residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + p_res = residual + o_t[:, None] * D + o_d[None, :] + b_x += tl.load(p_res, mask=m_x, other=0.0).to(tl.float32) if STORE_RESIDUAL_OUT: - p_res_out = tl.make_block_ptr( - residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) - ) - tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + p_res_out = residual_out + o_t[:, None] * D + o_d[None, :] + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_x) if not IS_RMS_NORM: b_mean = tl.sum(b_x, axis=1) / D - p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + p_mean = mean + o_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), mask=m_t) b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) b_var = tl.sum(b_xbar * b_xbar, axis=1) / D else: @@ -206,8 +205,8 @@ def layer_norm_gated_fwd_kernel( b_var = tl.sum(b_xbar * b_xbar, axis=1) / D b_rstd = 1 / tl.sqrt(b_var + eps) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + o_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) if HAS_WEIGHT: b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) @@ -223,7 +222,6 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - o_t = i_t * BT + tl.arange(0, BT) o_g = (o_t // H) * g_stride_n + (o_t % H) * D b_g = tl.load( g + o_g[:, None] + o_d[None, :], @@ -236,8 +234,8 @@ def layer_norm_gated_fwd_kernel( b_y = b_y * tl.sigmoid(b_g) # Write output - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + p_y = y + o_t[:, None] * D + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.heuristics( @@ -576,13 +574,13 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -603,58 +601,51 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter( A += (bos * H + i_h) * BT Aqk += (bos * H + i_h) * BT - p_b = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,) - ) - b_b = tl.load(p_b, boundary_check=(0,)) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T + + p_b = beta + bos * H + i_h + o_i * H + b_b = tl.load(p_b, mask=m_i, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) b_Aqk = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_k = tl.make_block_ptr( - k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - p_g = tl.make_block_ptr( - g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0) - ) - b_kt = tl.make_block_ptr( - k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - p_gk = tl.make_block_ptr( - g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1) - ) - o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] + p_q = q + o_i[:, None] * (H * K) + o_k[None, :] + p_k = k + o_i[:, None] * (H * K) + o_k[None, :] + p_g = g + o_i[:, None] * (H * K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_j[None, :] * (H * K) + p_gk = g + o_k[:, None] + o_j[None, :] * (H * K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) * exp2(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) # [BC, BC] b_ktg = b_kt * exp2(b_gn[:, None] - b_gk) b_A += tl.dot(b_k, b_ktg) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale b_Aqk += tl.dot(b_qg, b_ktg) b_A *= b_b[:, None] - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - p_Aqk = tl.make_block_ptr( - Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0) - ) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) + p_Aqk = Aqk + o_i[:, None] * (H * BT) + o_jA[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_A) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -681,12 +672,12 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -705,33 +696,14 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT + i_i * BC, 0), - (BC, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None] @@ -878,12 +850,12 @@ def recompute_w_u_fwd_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -892,108 +864,57 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos * H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load( gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 ) b_kg = b_k * exp2(b_gn - b_gk) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) def recompute_w_u_fwd( @@ -1074,17 +995,17 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_tg = i_t + i_tg = i_t.to(tl.int64) i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int32), - tl.load(cu_seqlens + i_n + 1).to(tl.int32), + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) T = eos - bos NT = tl.cdiv(T, BT) @@ -1096,70 +1017,44 @@ def chunk_gla_fwd_kernel_o( m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (V, K), - (K, 1), - (i_v * BV, i_k * BK), - (BV, BK), - (1, 0), - ) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K * V + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_qk, other=0.0) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] if i_k >= 0: b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) def chunk_gla_fwd_o_gk( @@ -1238,12 +1133,12 @@ def kda_gate_cumsum_fwd_kernel( IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -1253,27 +1148,17 @@ def kda_gate_cumsum_fwd_kernel( else: bos = i_b * T - p_g = tl.make_block_ptr( - g + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) - p_y = tl.make_block_ptr( - y + (bos * H + i_h) * D, - (T, D), - (H * D, 1), - (i_t * BT, i_d * BD), - (BT, BD), - (1, 0), - ) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + m_g = m_t[:, None] & m_d[None, :] + p_g = g + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] + p_y = y + (bos * H + i_h) * D + o_t[:, None] * (H * D) + o_d[None, :] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - o_d = i_d * BD + tl.arange(0, BD) - b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) + b_bias = tl.load(g_bias + i_h * D + o_d, mask=m_d, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1290,12 +1175,12 @@ def kda_gate_cumsum_fwd_kernel( # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to - # out-of-bounds output rows, which are masked away by `boundary_check` on - # the store, so visible output matches unfused gate + chunk-local cumsum. - o_t = tl.arange(0, BT) - m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0) + # out-of-bounds output rows, which are masked away by the store mask, + # so visible output matches unfused gate + chunk-local cumsum. + o_c = tl.arange(0, BT) + m_cumsum = tl.where(o_c[:, None] >= o_c[None, :], 1.0, 0.0) b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_g) def fused_kda_gate_chunk_cumsum( @@ -1605,7 +1490,7 @@ def kda_gate_fwd_kernel( HAS_BIAS: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) n_t = i_t * BT b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) @@ -1613,25 +1498,14 @@ def kda_gate_fwd_kernel( stride_row = H * D stride_col = 1 - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) + o_t = n_t + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + g_ptr = g + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col + y_ptr = y + i_h * D + o_t[:, None] * stride_row + o_d[None, :] * stride_col - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(g_ptr, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: n_d = tl.arange(0, BD) @@ -1649,7 +1523,7 @@ def kda_gate_fwd_kernel( sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) b_y = -b_a * sp - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(y_ptr, b_y.to(y.dtype.element_ty), mask=m_g) def fused_kda_gate( diff --git a/vllm/third_party/flash_linear_attention/ops/l2norm.py b/vllm/third_party/flash_linear_attention/ops/l2norm.py index 2eb137a242fb..4827df2d07dd 100644 --- a/vllm/third_party/flash_linear_attention/ops/l2norm.py +++ b/vllm/third_party/flash_linear_attention/ops/l2norm.py @@ -66,13 +66,18 @@ def l2norm_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, ): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + i_t = tl.program_id(0).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_x = m_t[:, None] & (o_d[None, :] < D) + p_x = x + o_t[:, None] * D + o_d[None, :] + p_y = y + o_t[:, None] * D + o_d[None, :] + + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) b_var = tl.sum(b_x * b_x, axis=1) b_y = b_x / tl.sqrt(b_var + eps)[:, None] - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) @triton.jit diff --git a/vllm/third_party/flash_linear_attention/ops/solve_tril.py b/vllm/third_party/flash_linear_attention/ops/solve_tril.py index 8d3811ca4c17..93ef770be627 100644 --- a/vllm/third_party/flash_linear_attention/ops/solve_tril.py +++ b/vllm/third_party/flash_linear_attention/ops/solve_tril.py @@ -69,12 +69,12 @@ def solve_tril_16x16_kernel( Ai = Ai + (bos * H + i_h) * 16 offset = (i_t * 16) % BT + o_t = (i_t * 16 + o_i).to(tl.int64) + m_t = o_t < T if not USE_TMA: - p_A = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) - ) + p_A = A + o_t[:, None] * (H * BT) + (o_i + offset)[None, :] # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = tl.load(p_A, mask=m_t[:, None], other=0.0).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) @@ -88,13 +88,11 @@ def solve_tril_16x16_kernel( b_A = tl.where((o_i == i)[:, None], b_a, b_A) b_A += m_I if not USE_TMA: - p_Ai = tl.make_block_ptr( - Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) - ) + p_Ai = Ai + o_t[:, None] * (H * 16) + o_i[None, :] tl.store( p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=m_t[:, None], ) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -143,15 +141,14 @@ def merge_16x16_to_32x32_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + p_A_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -175,10 +172,8 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) @@ -189,29 +184,23 @@ def merge_16x16_to_32x32_inverse_kernel( ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) + p_Ai_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) else: desc_o.store( @@ -268,23 +257,22 @@ def merge_16x16_to_64x64_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_A_22 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + p_A_11 = A + o_t[:, None] * (H * BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_A_33 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_A_44 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to( + tl.float32 ) - p_A_33 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + b_Ai_33 = tl.load(p_A_33, mask=(o_t[:, None] + 32 < T), other=0.0).to( + tl.float32 ) - p_A_44 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + b_Ai_44 = tl.load(p_A_44, mask=(o_t[:, None] + 48 < T), other=0.0).to( + tl.float32 ) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) @@ -321,30 +309,18 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_A_31 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_A_32 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_A_41 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_A_42 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_A_43 = tl.make_block_ptr( - A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) - b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_A_31 = A + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_A_32 = A + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_A_41 = A + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_A_42 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_A_43 = A + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) + b_A_31 = tl.load(p_A_31, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_32 = tl.load(p_A_32, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_41 = tl.load(p_A_41, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_42 = tl.load(p_A_42, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_43 = tl.load(p_A_43, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) @@ -390,85 +366,65 @@ def merge_16x16_to_64x64_inverse_kernel( ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) - ) - p_Ai_22 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) - ) - p_Ai_33 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) - ) - p_Ai_44 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) - ) - p_Ai_21 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) - ) - p_Ai_31 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) - ) - p_Ai_32 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) - ) - p_Ai_41 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) - ) - p_Ai_42 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) - ) - p_Ai_43 = tl.make_block_ptr( - Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) - ) + p_Ai_11 = Ai + o_t[:, None] * (H * BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H * BT) + (o_i[None, :] + 16) + p_Ai_33 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 32) + p_Ai_44 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 48) + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H * BT) + o_i[None, :] + p_Ai_31 = Ai + (o_t[:, None] + 32) * (H * BT) + o_i[None, :] + p_Ai_32 = Ai + (o_t[:, None] + 32) * (H * BT) + (o_i[None, :] + 16) + p_Ai_41 = Ai + (o_t[:, None] + 48) * (H * BT) + o_i[None, :] + p_Ai_42 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 16) + p_Ai_43 = Ai + (o_t[:, None] + 48) * (H * BT) + (o_i[None, :] + 32) tl.store( p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] < T), ) tl.store( p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 16 < T), ) tl.store( p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 32 < T), ) tl.store( p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) tl.store( p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), - boundary_check=(0, 1), + mask=(o_t[:, None] + 48 < T), ) else: desc_o.store( diff --git a/vllm/third_party/flash_linear_attention/ops/wy_fast.py b/vllm/third_party/flash_linear_attention/ops/wy_fast.py index 52d2b28195a8..040c37ebfe7f 100644 --- a/vllm/third_party/flash_linear_attention/ops/wy_fast.py +++ b/vllm/third_party/flash_linear_attention/ops/wy_fast.py @@ -46,12 +46,12 @@ def recompute_w_u_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: i_n, i_t = ( tl.load(chunk_indices + i_t * 2).to(tl.int32), - tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64), ) bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int32), @@ -60,60 +60,41 @@ def recompute_w_u_fwd_kernel( T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr( - beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) - ) - p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) - ) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_beta = beta + bos * H + i_h + o_t * H + p_g = g + (bos * H + i_h) + o_t * H + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_g = tl.exp(tl.load(p_g, mask=m_t, other=0.0)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, i_v * BV), - (BT, BV), - (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr( - k + (bos * Hg + i_h // (H // Hg)) * K, - (T, K), - (Hg * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), - ) - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, i_k * BK), - (BT, BK), - (1, 0), + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = ( + k + + (bos * Hg + i_h // (H // Hg)) * K + + o_t[:, None] * (Hg * K) + + o_k[None, :] ) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) b_w = tl.dot(b_A, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) def recompute_w_u_fwd( From 50b21a2667b6903d2e0d651aab9f449832e5e1e9 Mon Sep 17 00:00:00 2001 From: TJian Date: Mon, 12 Oct 2026 23:50:28 +0800 Subject: [PATCH 423/952] [Perf] [Feat] [ROCm] Add densemha support to ROCm AITER Sparse MLA (#49263) Signed-off-by: tjtanaa --- .buildkite/test-amd.yaml | 10 +- .buildkite/test_areas/attention.yaml | 9 +- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 503 ++++++++++++++++++ .../v1/attention/test_rocm_glm5next_sparse.py | 80 ++- .../v1/attention/test_sparse_mla_backends.py | 275 +++++++++- vllm/models/deepseek_v32/amd/rocm.py | 35 +- .../backends/mla/rocm_aiter_mla_sparse.py | 214 ++++---- 7 files changed, 984 insertions(+), 142 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a14a0a33a4d2..b4c4b59fdfe8 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2632,8 +2632,11 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) Metrics, Tracing" # TBD timeout_in_minutes: 30 @@ -5642,10 +5645,11 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py optional: true commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355 DPX) CUDAGraph" timeout_in_minutes: 55 diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 83ba90e2cb1b..ec47e45fe7a0 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -30,6 +30,10 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py + commands: + - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (B200) V1 Attention Shard %N" key: v1-attention-b200 @@ -59,9 +63,10 @@ steps: - vllm/_aiter_ops.py - vllm/envs.py - vllm/platforms/rocm.py + - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":nvidia: (DGX) Spark FlashInfer NVFP4 XQA" key: flashinfer-nvfp4-xqa-dgx-spark diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index aabce97fac9e..dc1ff79eeabd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -33,6 +33,101 @@ class _FakeAiter(ModuleType): get_mla_metadata_v1: Mock +@pytest.mark.parametrize("block_size", [16, 64, 128]) +@pytest.mark.parametrize("topk_tokens", [128, 256]) +def test_rocm_index_conversion_matrix(block_size, topk_tokens): + """Convert valid, masked, and out-of-range sparse indices exactly.""" + device = torch.device("cuda") + num_tokens = 4 + max_blocks = 4 + req_ids = torch.tensor([0, 1, 0, 1], dtype=torch.int32, device=device) + block_table = torch.tensor( + [[3, 5, 7, 9], [2, 4, 6, 8]], dtype=torch.int32, device=device + ) + token_indices = torch.arange( + num_tokens * topk_tokens, dtype=torch.int32, device=device + ).view(num_tokens, topk_tokens) + token_indices %= max_blocks * block_size + token_indices[0, 0] = -1 + token_indices[1, 1] = max_blocks * block_size + cu_seqlens = torch.arange( + 0, + (num_tokens + 1) * topk_tokens, + topk_tokens, + dtype=torch.int32, + device=device, + ) + output = torch.full( + (num_tokens * topk_tokens,), -99, dtype=torch.int32, device=device + ) + + sparse_mod.triton_convert_req_index_to_global_index( + req_ids, + block_table, + token_indices, + cu_seqlens, + output, + BLOCK_STRIDE_ROWS=block_size, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=topk_tokens, + ) + + block_ids = token_indices // block_size + valid = (token_indices >= 0) & (block_ids < max_blocks) + safe_blocks = block_ids.clamp(0, max_blocks - 1) + bases = block_table[req_ids[:, None].long(), safe_blocks.long()] + expected = torch.where(valid, bases * block_size + token_indices % block_size, 0) + torch.testing.assert_close(output.view_as(expected), expected, rtol=0, atol=0) + + +def test_rocm_index_conversion_rejects_mismatched_token_extents(): + """Reject full-batch request IDs paired with decode-subset indices.""" + with pytest.raises(AssertionError, match="must cover the same tokens"): + sparse_mod.triton_convert_req_index_to_global_index( + torch.zeros(3, dtype=torch.int32), + torch.zeros(1, 4, dtype=torch.int32), + torch.zeros(2, 128, dtype=torch.int32), + torch.arange(3, dtype=torch.int32), + torch.zeros(256, dtype=torch.int32), + BLOCK_STRIDE_ROWS=64, + NUM_TOPK_TOKENS=128, + ) + + +def test_rocm_index_conversion_respects_packed_valid_counts(): + """ROCm's indptr contract packs only each token's valid top-k prefix.""" + device = torch.device("cuda") + topk_tokens = 128 + valid_counts = torch.tensor([1, 17, 128], dtype=torch.int32, device=device) + cu_seqlens = torch.zeros(4, dtype=torch.int32, device=device) + torch.cumsum(valid_counts, dim=0, out=cu_seqlens[1:]) + req_ids = torch.tensor([0, 0, 0], dtype=torch.int32, device=device) + block_table = torch.tensor([[2, 3, 4]], dtype=torch.int32, device=device) + token_indices = torch.arange( + 3 * topk_tokens, dtype=torch.int32, device=device + ).view(3, topk_tokens) + token_indices %= 3 * 64 + output = torch.full((int(cu_seqlens[-1]),), -99, dtype=torch.int32, device=device) + + sparse_mod.triton_convert_req_index_to_global_index( + req_ids, + block_table, + token_indices, + cu_seqlens, + output, + BLOCK_STRIDE_ROWS=64, + BLOCK_SIZE=64, + NUM_TOPK_TOKENS=topk_tokens, + ) + + expected_parts = [] + for token_id, count in enumerate(valid_counts.tolist()): + local = token_indices[token_id, :count] + expected_parts.append(block_table[0, (local // 64).long()] * 64 + local % 64) + expected = torch.cat(expected_parts) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def _make_builder(): builder = object.__new__(sparse_mod.ROCMAiterMLASparseMetadataBuilder) max_num_batched_tokens = 8 @@ -44,6 +139,18 @@ def _make_builder(): builder.kv_cache_dtype = "fp8" builder.mla_dims = SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64) builder.topk_tokens = topk_tokens + # SparseMLACommonMetadataBuilder state consumed by its build() method. + builder.cp_kv_cache_interleave_size = 1 + builder.use_pcp = False + builder._prefill_backend = None + + def cpu_token_to_req(metadata): + return torch.repeat_interleave( + torch.arange(metadata.num_reqs, dtype=torch.int32), + metadata.query_start_loc_cpu.diff(), + ) + + builder._build_req_id_per_token = cpu_token_to_req # type: ignore[method-assign] builder.req_id_per_token_buffer = torch.zeros( max_num_batched_tokens, dtype=torch.int32, device="cpu" ) @@ -62,6 +169,8 @@ def _make_builder(): builder._use_persistent_metadata = True builder._num_attention_heads = 16 builder._num_compute_units = current_platform.num_compute_units() + builder._mla_q_dtype = torch.bfloat16 + builder._mla_kv_dtype = torch.bfloat16 builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu") builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu") builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu") @@ -216,6 +325,8 @@ def fake_get_mla_metadata_v1(*args, **kwargs): assert events == ["metadata", "sync"] assert fake_get_mla_metadata_v1_mock.call_count == 1 assert fake_get_mla_metadata_v1_mock.call_args.kwargs["max_split_per_batch"] == 1 + assert fake_get_mla_metadata_v1_mock.call_args.kwargs["dtype_q"] == torch.bfloat16 + assert fake_get_mla_metadata_v1_mock.call_args.kwargs["dtype_kv"] == torch.bfloat16 events.clear() @@ -223,3 +334,395 @@ def fake_get_mla_metadata_v1(*args, **kwargs): assert events == [] assert fake_get_mla_metadata_v1_mock.call_count == 1 + + +def test_dense_prefill_builds_sparse_metadata_for_decode_subset(monkeypatch): + """Dense prefill must exclude prefill tokens from AITER sparse metadata. + + A mixed batch is routed as sparse MQA for its leading decode requests and + dense MHA for its trailing short prefill requests. The AITER metadata + buffers therefore must describe only the decode subset; including prefill + tokens would mismatch the query passed to ``forward_mqa``. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32) + seq_lens = torch.tensor([16, 4], dtype=torch.int32) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=4, + max_seq_len=16, + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=2, + max_query_len=4, + max_seq_len=16, + num_actual_tokens=5, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(5), + block_table=torch.arange(16, dtype=torch.int32).view(2, 8), + req_id_per_token=torch.tensor([0, 1, 1, 1, 1], dtype=torch.int32), + topk_tokens=4, + num_decodes=1, + num_prefills=1, + num_decode_tokens=1, + seq_lens=seq_lens, + prefill_max_seq_len=4, + prefill=SimpleNamespace(use_dense_mha=True), + ) + + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + sparse_call: dict[str, torch.Tensor] = {} + + def fake_generate_sparse_seqlen( + query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + ): + sparse_call.update( + query_lens=query_lens.clone(), + seq_lens=sparse_seq_lens.clone(), + query_start_loc=sparse_query_start_loc.clone(), + num_tokens=num_tokens, + ) + return torch.ones(num_tokens, dtype=torch.int32) + + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen + ) + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = Mock() + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + assert result.prefill is metadata.prefill + assert result.num_prefills == 1 + assert result.req_id_per_token.tolist() == [0] + assert result.qo_indptr.numel() == 2 + assert result.paged_kv_indptr.numel() == 2 + assert sparse_call["num_tokens"] == 1 + assert sparse_call["query_lens"].tolist() == [1] + assert sparse_call["seq_lens"].tolist() == [16] + assert sparse_call["query_start_loc"].tolist() == [0, 1] + + +def test_long_prefill_shares_full_sparse_metadata_with_decode(monkeypatch): + """A prefill longer than top-k must remain in shared sparse metadata. + + Dense MHA is only selected when the prefill sequence fits within the + sparse top-k window. Once it exceeds that window, both decode and prefill + are handled by sparse MQA, so every query token must be represented by the + same sparse paging metadata. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + + # One decode token followed by a five-token prefill. topk_tokens is four, + # so the prefill must not be split onto the dense-MHA path. + query_start_loc = torch.tensor([0, 1, 6], dtype=torch.int32) + seq_lens = torch.tensor([16, 5], dtype=torch.int32) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=5, + max_seq_len=16, + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=2, + max_query_len=5, + max_seq_len=16, + num_actual_tokens=6, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(6), + block_table=torch.arange(16, dtype=torch.int32).view(2, 8), + req_id_per_token=torch.tensor([0, 1, 1, 1, 1, 1], dtype=torch.int32), + topk_tokens=4, + num_decodes=1, + num_prefills=1, + num_decode_tokens=1, + seq_lens=seq_lens, + prefill_max_seq_len=5, + prefill=SimpleNamespace(use_dense_mha=False), + ) + + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + sparse_call: dict[str, torch.Tensor] = {} + + def fake_generate_sparse_seqlen( + query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q + ): + sparse_call.update( + query_lens=query_lens.clone(), + seq_lens=sparse_seq_lens.clone(), + query_start_loc=sparse_query_start_loc.clone(), + num_tokens=num_tokens, + ) + return torch.ones(num_tokens, dtype=torch.int32) + + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen + ) + fake_aiter = _FakeAiter("aiter") + get_persistent_metadata = Mock() + fake_aiter.get_mla_metadata_v1 = get_persistent_metadata + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + assert result.num_decodes == 1 + assert result.num_prefills == 1 + assert result.req_id_per_token.tolist() == [0, 1, 1, 1, 1, 1] + assert result.qo_indptr.numel() == 7 + assert result.paged_kv_indptr.numel() == 7 + assert result.paged_kv_indices.numel() == 6 * result.topk_tokens + assert sparse_call["num_tokens"] == 6 + assert sparse_call["query_lens"].tolist() == [1, 5] + assert sparse_call["seq_lens"].tolist() == [16, 5] + assert sparse_call["query_start_loc"].tolist() == [0, 1, 6] + get_persistent_metadata.assert_called_once() + assert result.work_meta_data is builder._mla_work_meta_data + assert result.reduce_indptr is builder._mla_reduce_indptr + + +def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch): + """The sparse kernel inputs must have the same token extent as its query. + + In a mixed dense-prefill batch, ``forward_mqa`` receives only decode query + tokens while some metadata buffers may represent the full batch. This + guards the slicing that prevents the index conversion from writing beyond + the decode-sized output. + """ + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.kv_cache_dtype = "auto" + impl.num_heads = 4 + impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) + forward_mla = Mock(return_value=(expected_output, None)) + monkeypatch.setattr(impl, "_forward_mla", forward_mla) + + metadata = SimpleNamespace( + req_id_per_token=torch.tensor([0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32), + block_table=torch.arange(40, dtype=torch.int32).view(4, 10), + paged_kv_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.zeros(8, dtype=torch.int32), + block_size=1, + topk_tokens=4, + ) + query = torch.zeros(2, 4, 576, dtype=torch.bfloat16) + captured = {} + + def fake_convert(req_ids, block_table, token_indices, *args, **kwargs): + captured["req_ids"] = req_ids.clone() + captured["token_indices"] = token_indices.clone() + + monkeypatch.setattr( + sparse_mod, "triton_convert_req_index_to_global_index", fake_convert + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, "get_mla_padded_q", lambda num_heads, q: q + ) + + output, lse = impl.forward_mqa( + query, + torch.zeros(1, 1, 576, dtype=torch.bfloat16), + metadata, + SimpleNamespace(), + ) + + assert output is expected_output + assert lse is None + assert captured["req_ids"].tolist() == [0, 1] + assert captured["token_indices"].shape == (2, 4) + forward_mla.assert_called_once() + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache_dtype): + """Both cache dtypes must slice mixed-batch metadata to decode queries.""" + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.kv_cache_dtype = kv_cache_dtype + impl.num_heads = 4 + impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) + forward_mla = Mock(return_value=(expected_output, None)) + monkeypatch.setattr(impl, "_forward_mla", forward_mla) + + metadata = SimpleNamespace( + req_id_per_token=torch.tensor([0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32), + block_table=torch.arange(40, dtype=torch.int32).view(4, 10), + paged_kv_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.zeros(8, dtype=torch.int32), + block_size=1, + topk_tokens=4, + ) + query = torch.zeros(2, 4, 576, dtype=torch.bfloat16) + kv_cache = torch.zeros(1, 1, 576, dtype=torch.uint8) + captured = {} + + def fake_convert(req_ids, block_table, token_indices, *args, **kwargs): + captured["req_ids"] = req_ids.clone() + captured["token_indices"] = token_indices.clone() + + monkeypatch.setattr( + sparse_mod, "triton_convert_req_index_to_global_index", fake_convert + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, "get_mla_padded_q", lambda num_heads, q: q + ) + if kv_cache_dtype == "fp8": + fp8_dtype = current_platform.fp8_dtype() + monkeypatch.setattr( + sparse_mod.ops, + "scaled_fp8_quant", + lambda q, scale: (q.to(fp8_dtype), None), + ) + + output, lse = impl.forward_mqa( + query, + kv_cache, + metadata, + SimpleNamespace(_q_scale=torch.tensor(1.0)), + ) + + assert output is expected_output + assert lse is None + assert captured["req_ids"].tolist() == [0, 1] + assert captured["token_indices"].shape == (2, 4) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("length_relation", ["short", "long"]) +@pytest.mark.parametrize("workload", ["prefill", "mixed", "decode"]) +def test_sparse_metadata_workload_length_dtype_matrix( + monkeypatch, workload, length_relation, kv_cache_dtype +): + """Cover routing and persistent metadata across the ROCm support matrix. + + The matrix crosses pure prefill, mixed prefill/decode, and decode-only + workloads with sequence lengths below/above top-k and BF16/FP8 KV cache. + """ + builder = _make_builder() + builder.vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(sparse_mla_force_mqa=False) + ) + if kv_cache_dtype == "fp8": + builder._mla_q_dtype = current_platform.fp8_dtype() + builder._mla_kv_dtype = current_platform.fp8_dtype() + + is_long = length_relation == "long" + prefill_len = 5 if is_long else 3 + decode_seq_len = 8 if is_long else 3 + if workload == "prefill": + query_lens = [prefill_len] + seq_lens_list = [prefill_len] + num_decodes, num_prefills = 0, 1 + elif workload == "mixed": + query_lens = [1, prefill_len] + seq_lens_list = [decode_seq_len, prefill_len] + num_decodes, num_prefills = 1, 1 + else: + query_lens = [1, 1] + seq_lens_list = [decode_seq_len, decode_seq_len + 1] + num_decodes, num_prefills = 2, 0 + + query_start_loc = torch.tensor( + [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32) + num_tokens = sum(query_lens) + common_metadata = SimpleNamespace( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens, + max_query_len=max(query_lens), + max_seq_len=max(seq_lens_list), + ) + metadata = sparse_mod.ROCMAiterMLASparseMetadata( + num_reqs=len(query_lens), + max_query_len=max(query_lens), + max_seq_len=max(seq_lens_list), + num_actual_tokens=num_tokens, + query_start_loc=query_start_loc, + slot_mapping=torch.arange(num_tokens), + block_table=torch.arange(len(query_lens) * 8, dtype=torch.int32).view( + len(query_lens), 8 + ), + req_id_per_token=torch.repeat_interleave( + torch.arange(len(query_lens), dtype=torch.int32), + torch.tensor(query_lens), + ), + topk_tokens=4, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decodes, + seq_lens=seq_lens, + prefill_max_seq_len=prefill_len if num_prefills else 0, + prefill=SimpleNamespace(use_dense_mha=not is_long) if num_prefills else None, + ) + monkeypatch.setattr( + sparse_mod.SparseMLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + monkeypatch.setattr( + sparse_mod, + "generate_sparse_seqlen_triton", + lambda *args, **kwargs: torch.ones(args[4], dtype=torch.int32), + ) + get_persistent_metadata = Mock() + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = get_persistent_metadata + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace(synchronize=lambda: None), + ) + + result = builder.build(0, common_metadata) + + uses_dense_prefill = workload in ("prefill", "mixed") and not is_long + expected_sparse_tokens = num_decodes if uses_dense_prefill else num_tokens + assert result.req_id_per_token.numel() == expected_sparse_tokens + assert result.qo_indptr.numel() == expected_sparse_tokens + 1 + assert result.paged_kv_indptr.numel() == expected_sparse_tokens + 1 + + uses_fallback = workload == "prefill" and is_long and kv_cache_dtype == "auto" + if expected_sparse_tokens == 0 or uses_fallback: + get_persistent_metadata.assert_not_called() + else: + get_persistent_metadata.assert_called_once() + if uses_fallback or expected_sparse_tokens == 0: + assert result.work_meta_data is None + assert result.reduce_indptr is None + else: + assert result.work_meta_data is builder._mla_work_meta_data + assert result.reduce_indptr is builder._mla_reduce_indptr diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index 6077b9eb2fb6..e4244bc06aff 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/tests/v1/attention/test_rocm_glm5next_sparse.py @@ -155,8 +155,86 @@ def test_rocm_sparse_triton_route( ) +@pytest.mark.parametrize( + ("gfx942", "fp8", "num_prefills", "num_decodes", "prefill_max_seq_len", "route"), + [ + pytest.param(True, False, 1, 0, 129, "triton", id="above-topk"), + pytest.param(False, False, 1, 0, 129, "aiter", id="gfx950"), + pytest.param(True, True, 1, 0, 129, "aiter", id="fp8"), + pytest.param(True, False, 0, 2, 129, "aiter", id="decode-only"), + pytest.param(True, False, 1, 1, 129, "aiter", id="mixed"), + pytest.param(True, False, 1, 0, 128, "aiter", id="equal-topk"), + pytest.param(True, False, 1, 0, 127, "aiter", id="below-topk"), + pytest.param(True, False, 0, 0, 129, "aiter", id="empty"), + pytest.param( + True, False, 1, 0, None, "missing", id="required-prefill-metadata" + ), + ], +) +def test_rocm_long_bf16_prefill_dispatch( + monkeypatch, gfx942, fp8, num_prefills, num_decodes, prefill_max_seq_len, route +): + """Check the new dispatch gate using CPU tensors and intercepted kernels.""" + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: gfx942) + + class KernelSelected(Exception): + pass + + def triton_kernel(**kwargs): + raise KernelSelected("triton") + + def aiter_kernel(*args, **kwargs): + raise KernelSelected("aiter") + + monkeypatch.setattr(sparse_mod, "rocm_sparse_attn_prefill", triton_kernel) + monkeypatch.setattr(sparse_mod.rocm_aiter_ops, "mla_decode_fwd", aiter_kernel) + impl = object.__new__(sparse_mod.ROCMAiterMLASparseImpl) + impl.num_heads = 64 + impl.kv_lora_rank = 512 + impl.kv_cache_dtype = "fp8" if fp8 else "auto" + impl.scale = 576**-0.5 + impl.sinks = None + + dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + q = torch.empty(2, 64, 576, dtype=dtype) + kv = torch.empty(4, 1, 576, dtype=dtype) + metadata = SimpleNamespace( + attn_out_dtype=torch.bfloat16, + num_prefills=num_prefills, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + max_query_len=2, + prefill_max_seq_len=prefill_max_seq_len, + topk_tokens=128, + work_meta_data=None, + qo_indptr=torch.arange(3, dtype=torch.int32), + paged_kv_indices=torch.empty(0, dtype=torch.int32), + paged_kv_indptr=torch.zeros(3, dtype=torch.int32), + paged_kv_last_page_len=torch.ones(2, dtype=torch.int32), + ) + # Irrelevant routes must not inspect prefill-only metadata. + if ( + not gfx942 + or fp8 + or num_decodes + or not num_prefills + or prefill_max_seq_len is None + ): + del metadata.prefill_max_seq_len + del metadata.topk_tokens + layer = SimpleNamespace(_q_scale=1.0, _k_scale=1.0) + error = AttributeError if route == "missing" else KernelSelected + match = "prefill_max_seq_len" if route == "missing" else f"^{route}$" + with pytest.raises(error, match=match): + impl._forward_mla(layer, q, kv, metadata) + + +@pytest.mark.parametrize("gfx942", [False, True]) @pytest.mark.parametrize("num_heads", [8, 12]) -def test_rocm_sparse_triton_route_preserves_padded_sinks(monkeypatch, num_heads): +def test_rocm_sparse_triton_route_preserves_padded_sinks( + monkeypatch, num_heads, gfx942 +): + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: gfx942) captured = {} def fake_rocm_sparse_attn_prefill(**kwargs): diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 6eb089ce3ed6..c21fc15a5278 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -47,13 +47,12 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform -# TODO: Integrate ROCMAiterMLASparseBackend for ROCm. -# The ROCm sparse MLA backend (rocm_aiter_mla_sparse.py) has a compatible -# forward_mqa interface but needs validation on ROCm hardware. -if not current_platform.is_cuda(): +# This module's end-to-end sparse-backend correctness test supports CUDA and +# ROCm. The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs +# and are skipped by the autouse fixture below when collecting on ROCm. +if not (current_platform.is_cuda() or current_platform.is_rocm()): pytest.skip( - "Sparse MLA backend tests currently only support CUDA. " - "ROCm support requires integrating ROCMAiterMLASparseBackend.", + "Sparse MLA backend tests require CUDA or ROCm.", allow_module_level=True, ) @@ -65,14 +64,16 @@ ) from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import current_stream -from vllm.v1.attention.backends.mla import ( - flashattn_mla_sparse as flashattn_sparse_module, -) from vllm.v1.attention.backends.mla import index_group as index_group_module -from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( - FlashAttnMLASparseImpl, - FlashAttnMLASparseMetadataBuilder, -) + +if current_platform.is_cuda(): + from vllm.v1.attention.backends.mla import ( + flashattn_mla_sparse as flashattn_sparse_module, + ) + from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseImpl, + FlashAttnMLASparseMetadataBuilder, + ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, FlashInferMLASparseTRTLLMBackend, @@ -111,9 +112,43 @@ ) from vllm.v1.hisparse.types import SparseKVRowMirror +if current_platform.is_rocm(): + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, + ) + + SPARSE_BACKENDS = [ROCMAiterMLASparseBackend] + SPARSE_BACKEND_IDS = ["ROCMAiterMLA"] + SPARSE_KV_CACHE_DTYPES = ["auto", "fp8"] + SPARSE_BLOCK_SIZES = [64] +else: + SPARSE_BACKENDS = [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend] + SPARSE_BACKEND_IDS = ["FlashMLA", "FlashInferTRTLLM"] + SPARSE_KV_CACHE_DTYPES = ["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"] + SPARSE_BLOCK_SIZES = [32, 64] + + +@pytest.fixture(autouse=True) +def _skip_cuda_specific_sparse_mla_tests_on_rocm(request): + """Skip tests whose helpers exercise CUDA-only sparse MLA APIs.""" + rocm_portable_tests = { + "test_sparse_backend_decode_correctness", + "test_sparse_backend_prefill_correctness", + "test_rocm_fused_v32_dense_prefill_matches_mqa", + "test_rocm_long_bf16_prefill_avoids_nonpersistent_aiter", + } + if ( + current_platform.is_rocm() + and getattr(request.node, "originalname", request.node.name) + not in rocm_portable_tests + ): + pytest.skip("This sparse MLA test exercises a CUDA-specific API") + + SPARSE_BACKEND_BATCH_SPECS = { name: BATCH_SPECS[name] for name in [ + "small_decode", "mixed_small", "mixed_medium", "small_prefill", @@ -128,10 +163,61 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec( seq_lens=[256] * 2, query_lens=[256] * 2 ) +SPARSE_BACKEND_BATCH_SPECS["large_decode_over_topk"] = BatchSpec( + seq_lens=[256, 384], query_lens=[1, 1] +) DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("batch_name", "tensor_parallel_size"), + [ + ("large_q_prefill", 2), + ("large_q_pure_prefill", 4), + ("large_q_pure_prefill", 1), + ], +) +def test_rocm_long_bf16_prefill_avoids_nonpersistent_aiter( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch, + batch_name, + tensor_parallel_size, +): + """MI300 long BF16 prefills must remain correct without AITER ASM dispatch.""" + if not current_platform.is_rocm(): + pytest.skip("ROCm sparse MLA regression") + from vllm._aiter_ops import rocm_aiter_ops + + # Exercise the gfx942 routing decision on other ROCm devices too; + # the numerical kernel still compiles for the actual device. + monkeypatch.setattr("vllm.platforms.rocm.on_gfx942", lambda: True) + + def unsupported_aiter(*args, **kwargs): + pytest.fail("Long gfx942 BF16 prefill reached nonpersistent AITER") + + monkeypatch.setattr(rocm_aiter_ops, "mla_decode_fwd", unsupported_aiter) + test_sparse_backend_decode_correctness( + default_vllm_config=default_vllm_config, + dist_init=dist_init, + backend_cls=ROCMAiterMLASparseBackend, + batch_name=batch_name, + kv_cache_dtype="auto", + tensor_parallel_size=tensor_parallel_size, + block_size=64, + workspace_init=workspace_init, + q_scale=1.0, + k_scale=1.0, + total_num_heads=128, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + monkeypatch=monkeypatch, + ) + + def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): """Weight absorption must not change the model's attention temperature.""" model_scale = 256**-0.5 @@ -403,16 +489,13 @@ def unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: @pytest.mark.parametrize( "backend_cls", - [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend], - ids=["FlashMLA", "FlashInferTRTLLM"], + SPARSE_BACKENDS, + ids=SPARSE_BACKEND_IDS, ) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) -@pytest.mark.parametrize( - "kv_cache_dtype", - ["auto", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"], -) +@pytest.mark.parametrize("kv_cache_dtype", SPARSE_KV_CACHE_DTYPES) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) -@pytest.mark.parametrize("block_size", [32, 64]) +@pytest.mark.parametrize("block_size", SPARSE_BLOCK_SIZES) @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)]) @pytest.mark.parametrize( ("total_num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), @@ -525,6 +608,12 @@ def test_sparse_backend_decode_correctness( "attn_module_list_cfg": [{"topk_tokens": topk_tokens}], }, ) + # This shared correctness test constructs a sparse SDPA reference for every + # query token. ROCm normally routes short prefills to dense MHA, whose + # output intentionally differs from that sparse reference. Keep this test + # on the sparse-MQA path; dense-prefill routing has dedicated ROCm tests. + if current_platform.is_rocm(): + vllm_config.attention_config.sparse_mla_force_mqa = True model_config = vllm_config.model_config model_config.hf_text_config = SimpleNamespace( index_topk=topk_tokens, @@ -737,7 +826,7 @@ def test_sparse_backend_decode_correctness( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = backend_cls.get_builder_cls() @@ -1384,11 +1473,13 @@ def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): @pytest.mark.skipif( - torch.cuda.get_device_capability()[0] < 10, + current_platform.is_cuda() and torch.cuda.get_device_capability()[0] < 10, reason="Sparse MLA forward_mha requires FA4 (SM100+)", ) @pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys())) -@pytest.mark.parametrize("kv_cache_dtype", ["auto"]) +@pytest.mark.parametrize( + "kv_cache_dtype", (["auto", "fp8"] if current_platform.is_rocm() else ["auto"]) +) @pytest.mark.parametrize( ("num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), [ @@ -1408,11 +1499,19 @@ def test_sparse_backend_prefill_correctness( v_head_dim, workspace_init, ): - """Test dense and masked MHA across supported sparse MLA dimensions.""" - backend_cls = FlashMLASparseBackend + """Test the normal dense forward_mha path for sparse MLA short prefill.""" + backend_cls = ( + ROCMAiterMLASparseBackend + if current_platform.is_rocm() + else FlashMLASparseBackend + ) + if current_platform.is_rocm() and ( + batch_name.startswith("masked_mha") or qk_nope_head_dim != 128 + ): + pytest.skip("ROCm dense prefill requires the DeepSeek MLA layout") batch_spec = PREFILL_BATCH_SPECS[batch_name] - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 block_size = 64 @@ -1579,7 +1678,7 @@ def test_sparse_backend_prefill_correctness( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = backend_cls.get_builder_cls() @@ -1653,7 +1752,10 @@ def test_sparse_backend_prefill_correctness( assert out_buffer.shape == ref_output.shape assert torch.isfinite(out_buffer).all(), "Non-finite values in output" - torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01) + if kv_cache_dtype.startswith("fp8"): + torch.testing.assert_close(out_buffer, ref_output, rtol=0.065, atol=0.05) + else: + torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01) @pytest.mark.parametrize( @@ -3057,7 +3159,7 @@ def test_hisparse_mixed_batch_bf16_row_split( vllm_config=vllm_config, ) vllm_config.compilation_config.static_forward_context["placeholder"] = ( - SimpleNamespace(prefill_backend=prefill_backend) + SimpleNamespace(prefill_backend=prefill_backend, impl=SimpleNamespace()) ) builder_cls = FlashMLASparseBackend.get_builder_cls() @@ -4264,3 +4366,118 @@ def test_sparse_mla_common_impl_resolves_buffer_lazily(): assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), ( "topk_indices_buffer must stay a lazily-resolved property" ) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm model implementation") +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_rocm_fused_v32_dense_prefill_matches_mqa( + tmp_path, vllm_runner, kv_cache_dtype +): + """Fused normalization, RoPE and mixed dispatch preserve generation/logprobs. + + Chunking unequal prompts exercises prefill alongside decode. The two long + prompts straddle the dense-prefill threshold. No pretrained weights or + tokenizer download is needed. + """ + import json + + from vllm import SamplingParams + + config = { + "architectures": ["DeepseekV32ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "ep_size": 1, + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 2048, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 2048, + "initializer_range": 0.02, + "intermediate_size": 4096, + "kv_lora_rank": 512, + "max_position_embeddings": 8192, + "model_type": "deepseek_v32", + "moe_intermediate_size": 512, + "moe_layer_freq": 1, + "n_group": 1, + "n_routed_experts": 8, + "n_shared_experts": 1, + "norm_topk_prob": True, + "num_attention_heads": 32, + "num_experts_per_tok": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 32, + "num_nextn_predict_layers": 0, + "q_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + "rope_theta": 10000, + "routed_scaling_factor": 2.5, + "scoring_func": "sigmoid", + "tie_word_embeddings": False, + "topk_group": 1, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.44.2", + "use_cache": True, + "v_head_dim": 128, + "vocab_size": 4096, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + prompts = [ + {"prompt_token_ids": [2 + (i * 17 + j * 13) % 4000 for i in range(n)]} + for j, n in enumerate([64, 128, 2048, 2049]) + ] + sampling = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True, logprobs=5) + results = [] + for force_mqa in [True, False]: + with vllm_runner( + str(tmp_path), + load_format="dummy", + skip_tokenizer_init=True, + dtype="bfloat16", + kv_cache_dtype=kv_cache_dtype, + seed=42, + max_model_len=4096, + max_num_seqs=4, + max_num_batched_tokens=256, + block_size=64, + enable_chunked_prefill=True, + enable_prefix_caching=False, + enforce_eager=True, + gpu_memory_utilization=0.25, + attention_config={ + "backend": "ROCM_AITER_MLA_SPARSE", + "sparse_mla_force_mqa": force_mqa, + }, + ) as runner: + outputs = runner.llm.generate(prompts, sampling, use_tqdm=False) + results.append([request.outputs[0] for request in outputs]) + + for reference, actual in zip(*results): + assert len(actual.token_ids) == 16 + assert actual.token_ids == reference.token_ids + assert actual.logprobs is not None and reference.logprobs is not None + for token, expected, observed in zip( + actual.token_ids, reference.logprobs, actual.logprobs + ): + torch.testing.assert_close( + torch.tensor(observed[token].logprob), + torch.tensor(expected[token].logprob), + rtol=0.01, + atol=0.01, + ) diff --git a/vllm/models/deepseek_v32/amd/rocm.py b/vllm/models/deepseek_v32/amd/rocm.py index b9ac31f8be2b..7211edaa3a81 100644 --- a/vllm/models/deepseek_v32/amd/rocm.py +++ b/vllm/models/deepseek_v32/amd/rocm.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import cast + import torch from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer @@ -179,6 +182,8 @@ def _fused_attention( else: attn_metadata = attn_metadata_raw + attn_metadata = cast(MLACommonMetadata | None, attn_metadata) + slot_mapping = forward_context.slot_mapping assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) @@ -213,6 +218,10 @@ def _fused_attention( mla_kv_cache = self.kv_cache mla_k_scale = self._k_scale + use_mha = attn_metadata is not None and self._use_sparse_mha(attn_metadata) + kv_c_out = torch.empty_like(kv_c) if use_mha else None + k_pe_out = torch.empty_like(k_pe) if use_mha else None + q_c = fused_norm_rope( positions, q_c, @@ -237,9 +246,18 @@ def _fused_attention( mla_k_scale=mla_k_scale, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, + kv_c_out=kv_c_out, + k_pe_out=k_pe_out, ) - ql_nope, q_pe = self._compute_ql_nope(q_c) + if use_mha: + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + ql_nope = q_nope + else: + ql_nope, q_pe = self._compute_ql_nope(q_c) if self.indexer is not None and not self.skip_topk: index_q = self.indexer.wq_b(q_c)[0] @@ -260,7 +278,7 @@ def _fused_attention( indexer_n_head_scale, has_indexer=has_indexer, index_rope_interleave=self._index_rope_interleave, - quantize_mqa=self._fp8_kv, + quantize_mqa=self._fp8_kv and not use_mha, ) if self.indexer is not None and not self.skip_topk: @@ -270,6 +288,19 @@ def _fused_attention( output.zero_() return + if use_mha: + assert kv_c_out is not None and k_pe_out is not None + mha_q = torch.cat((q_nope, mqa_q), dim=-1) + self.forward_impl( + mha_q, + kv_c_out, + k_pe_out.unsqueeze(1), + self.kv_cache, + attn_metadata, + output, + ) + return + num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] kv_cache = self.kv_cache if self._fp8_kv_needs_view: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e30fec715deb..a0393e772c3e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -12,11 +12,10 @@ from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger -from vllm.model_executor.layers.attention.mla_attention import ( - get_mla_dims, -) from vllm.model_executor.layers.attention.sparse_mla_attention import ( - SharedTopkIndicesBuffer, + SparseMLACommonImpl, + SparseMLACommonMetadata, + SparseMLACommonMetadataBuilder, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -26,10 +25,7 @@ AttentionBackend, AttentionCGSupport, AttentionLayer, - AttentionMetadata, - AttentionMetadataBuilder, CommonAttentionMetadata, - MLAAttentionImpl, MultipleOf, ) from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( @@ -39,7 +35,6 @@ align_blocks_to_rows, flat_kv_row_view, ) -from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode_bf16, rocm_sparse_attn_prefill, @@ -76,6 +71,18 @@ def _use_rocm_sparse_triton( ) +def _is_long_bf16_prefill( + kv_dtype: torch.dtype, + metadata: "ROCMAiterMLASparseMetadata", +) -> bool: + return ( + kv_dtype == torch.bfloat16 + and metadata.num_decodes == 0 + and metadata.num_prefills > 0 + and metadata.prefill_max_seq_len > metadata.topk_tokens + ) + + @triton.jit def _fit_kpool_indices_kernel( token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS + TAIL_WIDTH] @@ -239,6 +246,10 @@ def triton_convert_req_index_to_global_index( assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 assert token_indices.dtype == torch.int32 + assert req_id.shape[0] == token_indices.shape[0], ( + f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) " + "must cover the same tokens" + ) assert token_indices.shape[1] == NUM_TOPK_TOKENS assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})" @@ -390,35 +401,13 @@ def supports_sink(cls) -> bool: return on_mi3xx() -@dataclass -class ROCMAiterMLASparseMetadata(AttentionMetadata): - num_reqs: int - max_query_len: int - max_seq_len: int - - num_actual_tokens: int # Number of tokens excluding padding. - query_start_loc: torch.Tensor - slot_mapping: torch.Tensor - - block_table: torch.Tensor - req_id_per_token: torch.Tensor - - qo_indptr: torch.Tensor - paged_kv_last_page_len: torch.Tensor - paged_kv_indices: torch.Tensor - paged_kv_indptr: torch.Tensor - attn_out_dtype: torch.dtype - - block_size: int = 1 - topk_tokens: int = 2048 - - # Fields read by the shared MLA forward. This impl has no dense-MHA prefill - # path (supports_dense_mha_prefill=False), so it always runs the MQA path; - num_decodes: int = 0 - num_prefills: int = 0 - num_decode_tokens: int = 0 - prefill_max_seq_len: int = 0 - prefill: object = None +@dataclass(kw_only=True) +class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): + qo_indptr: torch.Tensor | None = None + paged_kv_last_page_len: torch.Tensor | None = None + paged_kv_indices: torch.Tensor | None = None + paged_kv_indptr: torch.Tensor | None = None + attn_out_dtype: torch.dtype | None = None # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). @@ -432,9 +421,10 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata): @dataclass class ROCMAiterMLASparseMetadataBuilder( - AttentionMetadataBuilder[ROCMAiterMLASparseMetadata] + SparseMLACommonMetadataBuilder[ROCMAiterMLASparseMetadata] ): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + metadata_cls = ROCMAiterMLASparseMetadata def __init__( self, @@ -443,20 +433,15 @@ def __init__( vllm_config: VllmConfig, device: torch.device, ): - self.kv_cache_spec = kv_cache_spec - self.model_config = vllm_config.model_config + super().__init__(kv_cache_spec, layer_names, vllm_config, device) self.model_dtype = vllm_config.model_config.dtype self.kv_cache_dtype = vllm_config.cache_config.cache_dtype parallel_config = vllm_config.parallel_config - self.device = device max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens - self.vllm_config = vllm_config self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self.num_heads = self.model_config.get_num_attention_heads(parallel_config) - self.mla_dims = get_mla_dims(self.model_config) - self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk attention_context = vllm_config.compilation_config.static_forward_context # Sink decode must use AITER's nonpersistent path. In particular, # gfx942 has no persistent+LSE kernel, and its metadata heuristic @@ -521,9 +506,12 @@ def __init__( kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" + q_dtype = dtypes.fp8 else: kv_cache_dtype_str = "bf16" kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) + self._mla_q_dtype = q_dtype + self._mla_kv_dtype = kv_dtype ( (work_meta_data_size, work_meta_data_type), @@ -587,13 +575,19 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> ROCMAiterMLASparseMetadata: - num_tokens = common_attn_metadata.num_actual_tokens - (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills( - common_attn_metadata, - decode_threshold=self.reorder_batch_threshold or 1, + metadata = super().build( + common_prefix_len, common_attn_metadata, fast_build=fast_build + ) + use_dense_prefill = ( + metadata.prefill is not None and metadata.prefill.use_dense_mha ) + num_reqs = metadata.num_decodes if use_dense_prefill else metadata.num_reqs + num_tokens = int(common_attn_metadata.query_start_loc_cpu[num_reqs].item()) + num_decodes = metadata.num_decodes + num_prefills = metadata.num_prefills + num_decode_tokens = metadata.num_decode_tokens starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) - seg_lengths = np.diff(starts) + seg_lengths = np.diff(starts[: num_reqs + 1]) req_id_per_token = np.repeat( np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths ) @@ -616,15 +610,13 @@ def build( self.req_id_per_token_buffer[:new_req_extent].copy_( np_to_pinned_tensor(req_id_per_token), non_blocking=True ) - query_lens = ( - common_attn_metadata.query_start_loc[1:] - - common_attn_metadata.query_start_loc[:-1] - ) - seq_lens = common_attn_metadata.seq_lens + query_start_loc = common_attn_metadata.query_start_loc[: num_reqs + 1] + query_lens = query_start_loc[1:] - query_start_loc[:-1] + seq_lens = common_attn_metadata.seq_lens[:num_reqs] sparse_seqlen = generate_sparse_seqlen_triton( query_lens, seq_lens, - common_attn_metadata.query_start_loc, + query_start_loc, self.topk_tokens, num_tokens, common_attn_metadata.max_query_len, @@ -663,8 +655,13 @@ def build( reduce_indptr = None reduce_final_map = None reduce_partial_map = None - if self._use_persistent_metadata and not use_triton_sparse: - num_reqs = common_attn_metadata.num_reqs + is_long_pure_bf16_prefill = _is_long_bf16_prefill(self._mla_kv_dtype, metadata) + if ( + self._use_persistent_metadata + and not use_triton_sparse + and not is_long_pure_bf16_prefill + and num_tokens > 0 + ): with gpu_sync_allowed(): seq_lens_cpu = common_attn_metadata.seq_lens[:num_reqs].cpu().numpy() clamped_seq_lens = np.minimum( @@ -708,6 +705,8 @@ def build( uni_seqlen_qo=1, fast_mode=True, max_split_per_batch=max_split_per_batch, + dtype_q=self._mla_q_dtype, + dtype_kv=self._mla_kv_dtype, ) torch.cuda.current_stream(self.device).synchronize() self._prev_metadata_key = metadata_key @@ -718,32 +717,18 @@ def build( reduce_final_map = self._mla_reduce_final_map reduce_partial_map = self._mla_reduce_partial_map - metadata = ROCMAiterMLASparseMetadata( - num_reqs=common_attn_metadata.num_reqs, - max_query_len=common_attn_metadata.max_query_len, - max_seq_len=common_attn_metadata.max_seq_len, - num_actual_tokens=common_attn_metadata.num_actual_tokens, - query_start_loc=common_attn_metadata.query_start_loc, - slot_mapping=common_attn_metadata.slot_mapping, - block_table=common_attn_metadata.block_table_tensor, - req_id_per_token=req_id_per_token, - block_size=self.kv_cache_spec.block_size, - attn_out_dtype=self.model_dtype, - topk_tokens=self.topk_tokens, - num_decodes=num_decodes, - num_prefills=num_prefills, - num_decode_tokens=num_decode_tokens, - qo_indptr=qo_indptr, - paged_kv_last_page_len=paged_kv_last_page_len, - paged_kv_indices=paged_kv_indices, - paged_kv_indptr=paged_kv_indptr, - work_meta_data=work_meta_data, - work_indptr=work_indptr, - work_info_set=work_info_set, - reduce_indptr=reduce_indptr, - reduce_final_map=reduce_final_map, - reduce_partial_map=reduce_partial_map, - ) + metadata.attn_out_dtype = self.model_dtype + metadata.req_id_per_token = req_id_per_token + metadata.qo_indptr = qo_indptr + metadata.paged_kv_last_page_len = paged_kv_last_page_len + metadata.paged_kv_indices = paged_kv_indices + metadata.paged_kv_indptr = paged_kv_indptr + metadata.work_meta_data = work_meta_data + metadata.work_indptr = work_indptr + metadata.work_info_set = work_info_set + metadata.reduce_indptr = reduce_indptr + metadata.reduce_final_map = reduce_final_map + metadata.reduce_partial_map = reduce_partial_map return metadata @@ -776,11 +761,7 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor: return (result, lse) -class ROCMAiterMLASparseImpl( - MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer -): - is_sparse = True - supports_dense_mha_prefill = False +class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]): supports_dcp = False use_aiter_sparse_mla = False @@ -803,10 +784,6 @@ def __init__( ) -> None: AiterMLAHelper.check_num_heads_validity(num_heads) - self.num_heads = num_heads - self.head_size = head_size - self.scale = float(scale) - self.num_kv_heads = num_kv_heads sinks = mla_args.pop("sinks", None) if sinks is not None: if sinks.dtype != torch.float32: @@ -821,10 +798,24 @@ def __init__( if not sinks.is_contiguous(): raise ValueError("ROCm AITER MLA sinks must be contiguous") self.sinks: torch.Tensor | None = sinks - self.kv_cache_dtype = kv_cache_dtype - self.kv_lora_rank: int = mla_args["kv_lora_rank"] + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + indexer=indexer, + topk_indices_buffer=topk_indices_buffer, + **mla_args, + ) self.softmax_scale = scale - self.init_topk_indices_buffer(indexer, topk_indices_buffer) + self.supports_dense_mha_prefill = sinks is None # type: ignore[misc] + self.masked_mha_available = False vllm_config = get_current_vllm_config() max_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -894,14 +885,26 @@ def _forward_mla( ) ) - if triton_sink_fallback or _use_rocm_sparse_triton( - kv_cache_dtype=self.kv_cache_dtype, - head_size=q.shape[-1], - kv_lora_rank=self.kv_lora_rank, - num_prefills=attn_metadata.num_prefills, - num_decodes=attn_metadata.num_decodes, - num_decode_tokens=attn_metadata.num_decode_tokens, - max_query_len=attn_metadata.max_query_len, + # Long BF16 prefills omit persistent metadata. gfx942's + # nonpersistent ASM path is unsupported or unsafe for these shapes. + # Sink attention has its own Triton/LSE dispatch below. + triton_prefill_fallback = ( + not need_lse + and on_gfx942() + and _is_long_bf16_prefill(kv_c_and_k_pe_cache.dtype, attn_metadata) + ) + if ( + triton_sink_fallback + or triton_prefill_fallback + or _use_rocm_sparse_triton( + kv_cache_dtype=self.kv_cache_dtype, + head_size=q.shape[-1], + kv_lora_rank=self.kv_lora_rank, + num_prefills=attn_metadata.num_prefills, + num_decodes=attn_metadata.num_decodes, + num_decode_tokens=attn_metadata.num_decode_tokens, + max_query_len=attn_metadata.max_query_len, + ) ): output = torch.empty( [num_tokens, q.shape[1], self.kv_lora_rank], @@ -1078,6 +1081,7 @@ def _forward_mla( # Empty ragged rows have only sink mass and no value contribution. # AITER can return NaN output/LSE for those rows; do not multiply it # by a zero normalization factor and propagate the NaN. + assert attn_metadata.paged_kv_indptr is not None has_keys = ( attn_metadata.paged_kv_indptr[1:] > attn_metadata.paged_kv_indptr[:-1] ).unsqueeze(-1) @@ -1156,7 +1160,7 @@ def forward_mqa( else: ops.concat_mla_q(ql_nope, q_pe, q) - num_actual_toks = attn_metadata.num_actual_tokens + num_actual_toks = q.shape[0] # Get topk indices assert self.topk_indices_buffer is not None @@ -1168,7 +1172,7 @@ def forward_mqa( kv_c_and_k_pe_cache, attn_metadata.block_size ) triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token, + attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, attn_metadata.paged_kv_indptr, From e0317b28ce3ba4c9bd2ff69d885e3f34688a6d71 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:53:26 -0700 Subject: [PATCH 424/952] [Feature] Support batch invariance on ROCm (#52231) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 89 +++- .buildkite/test_areas/misc.yaml | 112 +++++- docs/features/batch_invariance.md | 1 + .../quantization/test_triton_scaled_mm.py | 25 +- tests/v1/determinism/test_batch_invariance.py | 1 + .../determinism/test_batch_invariance_vlm.py | 6 +- .../test_cutlass_batch_invariance.py | 118 +++++- .../test_lora_shrink_batch_invariant.py | 6 +- .../test_matmul_batch_invariant.py | 52 ++- .../test_mx_linear_batch_invariant.py | 165 ++++++++ .../determinism/test_nvfp4_batch_invariant.py | 10 +- .../test_online_batch_invariance.py | 6 +- .../test_rms_norm_batch_invariant.py | 18 +- tests/v1/determinism/utils.py | 22 +- .../device_communicators/cuda_communicator.py | 39 +- .../determinism/batch_invariant.py | 380 +++++++++++++++--- .../model_executor/kernels/linear/__init__.py | 43 ++ .../kernels/linear/mxfp4/aiter.py | 18 +- .../kernels/linear/mxfp8/rocm_native.py | 6 +- .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/triton.py | 42 ++ .../fused_moe/experts/aiter_mxfp8_moe.py | 4 + .../fused_moe/experts/ocp_mx_emulation_moe.py | 2 +- .../compressed_tensors/triton_scaled_mm.py | 12 +- .../model_executor/layers/quantization/fp8.py | 4 +- .../layers/quantization/online/fp8.py | 10 +- .../layers/quantization/utils/fp8_utils.py | 5 +- vllm/model_executor/layers/utils.py | 7 + vllm/platforms/rocm.py | 9 + .../backends/mla/prefill/aiter_flash_attn.py | 10 + .../backends/mla/prefill/flash_attn.py | 4 +- 31 files changed, 1106 insertions(+), 122 deletions(-) create mode 100644 tests/v1/determinism/test_mx_linear_batch_invariant.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b4c4b59fdfe8..49bc3c32b67a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -273,7 +273,7 @@ steps: #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# - label: ":amd: (MI250) Batch Invariance" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] dind: false agent_pool: mi250_1 @@ -282,14 +282,18 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: + - export VLLM_GPU_MEMORY_UTILIZATION=0.8 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 @@ -2443,6 +2447,56 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# +- label: ":amd: (MI300) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + +- label: ":amd: (MI300) Batch Invariance (Distributed)" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -5215,6 +5269,37 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# +- label: ":amd: (MI355) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + - label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 20b48cb98f1f..02f78e280c98 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -644,12 +644,33 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + mirror: + amd: + label: ":amd: (MI250) Batch Invariance" + dind: false + device: mi250_1 + timeout_in_minutes: 60 + env: + VLLM_GPU_MEMORY_UTILIZATION: "0.8" + VLLM_TEST_CLEAN_GPU_MEMORY: "1" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 @@ -660,6 +681,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -671,15 +694,18 @@ steps: - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN mirror: amd: - label: ":amd: (MI250) Batch Invariance" + label: ":amd: (MI300) Batch Invariance" dind: false - device: mi250_1 - timeout_in_minutes: 35 + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py @@ -688,6 +714,11 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - label: ":nvidia: (B200) Batch Invariance" key: batch-invariance-b200 @@ -697,6 +728,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -711,6 +744,79 @@ steps: - pytest -v -s v1/determinism/test_matmul_batch_invariant.py - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py - pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI355) Batch Invariance" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + +- label: ":nvidia: (H100) Batch Invariance (Distributed)" + key: batch-invariance-distributed-4xh100 + timeout_in_minutes: 45 + device: h100 + optional: true + num_devices: 4 + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI300) Batch Invariance (Distributed)" + dind: false + device: mi300_4 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py - label: ":nvidia: (H200 MIG 35GB) Acceptance Length (Large Models)" device: h200_35gb diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 03d07896cce4..01dc7116fef9 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -20,6 +20,7 @@ Batch invariance is crucial for several use cases: Batch invariance is supported on the following platforms: - NVIDIA GPUs with compute capability 8.0 or higher. +- AMD Instinct GPUs (gfx90a, gfx942, gfx950). - Intel XPUs with Triton support. ### Attention Backend Selection for XPU diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index 55d1f76b2d1e..b0a2c215ab02 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -32,11 +32,9 @@ def torch_scaled_mm( out = torch.mm(a.to(torch.float32), b.to(torch.float32)) out = scale_a * out out = scale_b.T * out - out = out.to(out_dtype) if bias is not None: - out = out + bias - - return out + out = out + bias.float() + return out.to(out_dtype) def get_8bit_types(): @@ -125,6 +123,25 @@ def test_scaled_mm( torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1) +@pytest.mark.parametrize("in_dtype", get_8bit_types()) +@pytest.mark.parametrize("use_heuristic", [True, False]) +def test_scaled_mm_adds_bias_before_rounding(in_dtype, use_heuristic): + # The accumulator is 1 + 3 * 2**-10 and the bias 2**-8. Rounding to bf16 + # first gives 1.0 (1 + 2**-8 is a tie, to even); adding first, 1 + 2**-7. + a = torch.ones((1, 32), device=device).to(in_dtype) + col = torch.tensor([64.0] * 16 + [2.0, 1.0] + [0.0] * 14, device=device) + b = col[:, None].expand(32, 16).to(in_dtype).contiguous() + scale_a = torch.ones((1, 1), device=device) + scale_b = torch.full((1, 1), 2**-10, device=device) + bias = torch.full((16,), 2**-8, device=device, dtype=torch.bfloat16) + args = (a, b, scale_a, scale_b, torch.bfloat16, bias) + + out = triton_scaled_mm(*args, use_heuristic=use_heuristic) + expected = torch_scaled_mm(*args) + assert (expected == 1 + 2**-7).all() + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + # TD operand loads must be bit-exact vs the plain masked-load path. @pytest.mark.skipif( not (current_platform.is_cuda_alike() or current_platform.is_xpu()), diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 452362715111..139c68ed2b81 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -180,6 +180,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( llm = LLM( model=TEST_MODEL, tensor_parallel_size=tp_size, + disable_custom_all_reduce=os.getenv("VLLM_TEST_DISABLE_CUSTOM_AR") == "1", max_num_seqs=128, max_model_len=8192, dtype="auto", # not everything is supported diff --git a/tests/v1/determinism/test_batch_invariance_vlm.py b/tests/v1/determinism/test_batch_invariance_vlm.py index fe95b03c61fa..5c557a1364d3 100644 --- a/tests/v1/determinism/test_batch_invariance_vlm.py +++ b/tests/v1/determinism/test_batch_invariance_vlm.py @@ -9,7 +9,7 @@ import pytest import torch from PIL import Image -from utils import _extract_step_logprobs, skip_if_not_cuda +from utils import _extract_step_logprobs, skip_if_not_cuda_alike from vllm import LLM, SamplingParams from vllm.inputs import TextPrompt @@ -89,7 +89,7 @@ def _assert_batch_invariant(llm: LLM, inputs: list[TextPrompt], sampling) -> Non ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("input_type", ["image", "video"]) @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"]) def test_vlm_batch_invariance_bs1_vs_bsN( @@ -101,7 +101,7 @@ def test_vlm_batch_invariance_bs1_vs_bsN( _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False) diff --git a/tests/v1/determinism/test_cutlass_batch_invariance.py b/tests/v1/determinism/test_cutlass_batch_invariance.py index c3384a62f6f0..b8f83817065c 100644 --- a/tests/v1/determinism/test_cutlass_batch_invariance.py +++ b/tests/v1/determinism/test_cutlass_batch_invariance.py @@ -19,8 +19,10 @@ from tests.kernels.quantization.nvfp4_utils import get_nvfp4_global_scale from tests.utils import TestFP8Layer, requires_fp8 from vllm import _custom_ops as ops -from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( +from vllm.model_executor.kernels.linear import fp8_scaled_mm_is_batch_invariant +from vllm.model_executor.kernels.linear.scaled_mm import ( CutlassFP8ScaledMMLinearKernel, + TritonFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -31,8 +33,11 @@ CutlassExpertsFp4, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, kFp8DynamicTokenSym, + kFp8StaticChannelSym, kFp8StaticTensorSym, + kFp8StaticTokenSym, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -87,16 +92,42 @@ @pytest.fixture(autouse=True) def setup_cuda(): - if not current_platform.is_cuda(): - pytest.skip("CUTLASS FP8 kernels require CUDA.") + if not current_platform.is_cuda_alike(): + pytest.skip("Requires CUDA or ROCm.") torch.set_default_device("cuda") @requires_fp8 +@pytest.mark.parametrize( + "kernel", + [ + pytest.param( + CutlassFP8ScaledMMLinearKernel, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUTLASS requires CUDA." + ), + ), + TritonFP8ScaledMMLinearKernel, + ], +) +@pytest.mark.parametrize( + "act_key,weight_key", + [ + (kFp8DynamicTokenSym, kFp8StaticTensorSym), + (kFp8DynamicTokenSym, kFp8StaticChannelSym), + (kFp8DynamicTokenSym, kFp8StaticTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + # An amax over the whole batch: the GEMM is invariant, the layer is not. + (kFp8DynamicTensorSym, kFp8StaticTensorSym), + ], +) @pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)]) @pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257]) @torch.inference_mode() -def test_cutlass_fp8_batch_invariant_fixed_config( +def test_fp8_scaled_mm_batch_invariant_fixed_config( + kernel: type, + act_key, + weight_key, weight_shape: tuple[int, int], batch_size: int, default_vllm_config, @@ -108,14 +139,16 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.manual_seed(0) layer = TestFP8Layer( weight_shape=weight_shape, - activation_quant_key=kFp8DynamicTokenSym, - weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=act_key, + weight_quant_key=weight_key, input_dtype=torch.bfloat16, out_dtype=torch.bfloat16, device=torch.device("cuda"), - force_kernel=CutlassFP8ScaledMMLinearKernel, + force_kernel=kernel, ) - assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel) + assert isinstance(layer.kernel, kernel) + dynamic_per_tensor = act_key == kFp8DynamicTensorSym + assert fp8_scaled_mm_is_batch_invariant(layer.kernel) != dynamic_per_tensor in_features = weight_shape[1] needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16) @@ -124,6 +157,12 @@ def test_cutlass_fp8_batch_invariant_fixed_config( filler = torch.randn( (max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16 ) + if dynamic_per_tensor: + if batch_size == 1: + pytest.skip("no other rows to move the scale") + front_output = layer(torch.cat([needle, 100 * filler], dim=0))[0] + assert not torch.equal(front_output, baseline) + return front_batch = torch.cat([needle, filler], dim=0) back_batch = torch.cat([filler, needle], dim=0) @@ -135,6 +174,69 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.testing.assert_close(back_output, baseline, rtol=0, atol=0) +@requires_fp8 +@pytest.mark.parametrize("scheme", ["fp8_per_tensor", "fp8_per_channel"]) +@torch.inference_mode() +def test_online_fp8_linear_batch_invariant( + scheme, monkeypatch, default_vllm_config, dist_init +): + from vllm.config.model import ModelConfig + from vllm.config.quantization import resolve_quantization_config + from vllm.model_executor.layers.linear import ReplicatedLinear + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + from vllm.utils.torch_utils import set_default_torch_dtype + + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + default_vllm_config.model_config = ModelConfig(dtype="bfloat16") + quant_config = OnlineQuantizationConfig(resolve_quantization_config(scheme, None)) + set_random_seed(0) + with set_default_torch_dtype(torch.bfloat16), torch.device("cuda"): + layer = ReplicatedLinear( + 1024, 768, bias=False, quant_config=quant_config, prefix="mlp.down_proj" + ) + layer.weight.weight_loader(layer.weight, torch.randn(768, 1024)) + + needle = torch.randn(1, 1024) + filler = 100 * torch.randn(63, 1024) + baseline = layer(needle)[0] + assert torch.equal(layer(torch.cat([needle, filler]))[0][:1], baseline) + assert torch.equal(layer(torch.cat([filler, needle]))[0][-1:], baseline) + + +@requires_fp8 +@torch.inference_mode() +def test_block_fp8_batch_invariant_across_tuned_block_size_k(monkeypatch): + """The block-FP8 kernel applies the block scales once per K tile, so a tuned + BLOCK_SIZE_K below the scale group size reorders the sum.""" + from vllm.model_executor.layers.quantization.utils import fp8_utils + + base = dict( + BLOCK_SIZE_M=64, BLOCK_SIZE_N=128, GROUP_SIZE_M=1, num_warps=4, num_stages=2 + ) + table = {256: {**base, "BLOCK_SIZE_K": 128}, 512: {**base, "BLOCK_SIZE_K": 64}} + monkeypatch.setattr(fp8_utils, "get_w8a8_block_fp8_configs", lambda *_: table) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + N, K, block = 12288, 2048, [128, 128] + + set_random_seed(0) + fp8 = current_platform.fp8_dtype() + A = torch.randn(512, K).to(fp8) + B = torch.randn(N, K).to(fp8) + As = torch.rand(512, K // 128) / 100 + Bs = torch.rand(N // 128, K // 128) / 100 + + def run(m): + return fp8_utils.w8a8_triton_block_scaled_mm( + A[:m], B, As[:m], Bs, block, torch.bfloat16 + )[:256] + + assert torch.equal(run(256), run(512)) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + assert not torch.equal(run(256), run(512)) + + @_NVFP4_REQUIRES_SM100 @pytest.mark.parametrize("dtype", _NVFP4_LINEAR_DTYPES) @pytest.mark.parametrize("shape", _NVFP4_LINEAR_CONSISTENCY_SHAPES) diff --git a/tests/v1/determinism/test_lora_shrink_batch_invariant.py b/tests/v1/determinism/test_lora_shrink_batch_invariant.py index 2a7029a60324..0c4c595c1d60 100644 --- a/tests/v1/determinism/test_lora_shrink_batch_invariant.py +++ b/tests/v1/determinism/test_lora_shrink_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda +from utils import skip_if_not_cuda_alike import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op import vllm.lora.ops.triton_ops.utils as lora_utils @@ -39,7 +39,7 @@ def _shrink(inputs, weights, mapping, scaling): return out -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize( "hidden_size,rank,nslices,dtype", [ @@ -68,7 +68,7 @@ def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype): torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("nslices", [1, 3]) def test_lora_shrink_invariant_across_batch_sizes(monkeypatch, nslices): reduce_kernel, reduce_grids = shrink_op._lora_shrink_reduce_kernel, [] diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index af395e74e7f1..b6ea1648f398 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -10,7 +10,11 @@ import torch from utils import skip_unsupported -from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant +import vllm.model_executor.determinism.batch_invariant as bi +from vllm.model_executor.determinism.batch_invariant import ( + enable_batch_invariant_mode, + matmul_batch_invariant, +) from vllm.model_executor.determinism.batch_invariant_configs import ( _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS, _get_tuned_matmul_arch_family, @@ -129,3 +133,49 @@ def test_matmul_batch_invariance_across_tuned_m_buckets(m, transpose_b): batch_output = matmul_batch_invariant(a, b) assert torch.equal(single_output[0], batch_output[0]) + + +@pytest.fixture +def batch_invariant_mode(): + # The overrides are process-global: remove them so later tests compare + # against the stock aten kernels again. + enable_batch_invariant_mode() + yield + assert bi._batch_invariant_LIB is not None + bi._batch_invariant_LIB._destroy() + bi._batch_invariant_MODE = False + + +@skip_unsupported +@pytest.mark.usefixtures("batch_invariant_mode") +@pytest.mark.parametrize("op", ["mm", "addmm", "bmm"]) +def test_compiled_matmul_reaches_batch_invariant_out_overloads(op): + """Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, + i.e. the ``.out`` overload, which eager never dispatches.""" + if ( + op != "bmm" + and current_platform.is_cuda() + and not current_platform.is_device_capability_family(80) + ): + pytest.skip("mm/addmm are not overridden on SM90+") + torch.manual_seed(0) + dtype, k, n, needle = torch.bfloat16, 1024, 768, 37 + w = torch.randn(4, k, n, device=DEVICE_TYPE, dtype=dtype) + bias = torch.randn(n, device=DEVICE_TYPE, dtype=dtype) + fn = { + "mm": lambda x: torch.mm(x[0], w[0]), + # Bare addmm: a following pointwise op would make Inductor unfuse the bias. + "addmm": lambda x: torch.addmm(bias, x[0], w[0]), + "bmm": lambda x: torch.bmm(x, w), + }[op] + compiled = torch.compile(fn, fullgraph=True) + x = torch.randn(4, 256, k, device=DEVICE_TYPE, dtype=dtype) + row = None + for m in (64, 256, 128): + xs = x[:, :m].clone() + # With dynamic=True, a specialization on M would silently recompile. + torch._dynamo.mark_dynamic(xs, 1) + out = compiled(xs) + assert torch.equal(out, fn(xs)), f"compiled {op} bypassed the BI kernel (M={m})" + row = out[..., needle, :] if row is None else row + assert torch.equal(out[..., needle, :], row) diff --git a/tests/v1/determinism/test_mx_linear_batch_invariant.py b/tests/v1/determinism/test_mx_linear_batch_invariant.py new file mode 100644 index 000000000000..af31302ac91c --- /dev/null +++ b/tests/v1/determinism/test_mx_linear_batch_invariant.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The ROCm native MX linear kernels must not depend on the number of rows.""" + +import pytest +import torch +from utils import requires_mx + +from vllm.utils.torch_utils import set_random_seed + +SEED = 0 + +# E8M0 exponent half-width. With 4-bit operands and a narrow spread the fp32 +# partial sums are exact, so a reordering of them would be invisible. +SCALE_SPREAD = 15 + +# (N, K, use_asm_gemm) +MXFP4_CASES = [ + (4096, 2048, False), + (2048, 6144, False), + (4608, 7168, False), + (8192, 2048, True), + (8192, 1024, True), + (8192, 14336, True), +] + +MXFP8_CASES = [(4096, 2048), (2048, 6144), (1024, 768), (1536, 1024)] + +# The untuned AITER MXFP4 configs switch at M = 9/33/65/129/257/513. +MXFP4_TOKEN_COUNTS = [ + 1, + 2, + 8, + 9, + 16, + 32, + 33, + 64, + 65, + 128, + 129, + 256, + 257, + 512, + 513, + 1024, + 2048, +] + +MXFP8_TOKEN_COUNTS = [1, 32, 64, 65, 128, 129, 256, 257, 512, 1024, 1025, 2048] + +# Row 0 lands in the first tile of every decomposition, so it can stay +# invariant when the rest of the output does not. +CHECK_ROWS = [0, 1, 2, 3, 7, 15, 31] + + +def _probe(out: torch.Tensor) -> torch.Tensor: + """The leading rows of an output, kept so the sweep does not hold every GEMM.""" + return out[: CHECK_ROWS[-1] + 1].clone() + + +def _variant_rows(probes: dict[int, torch.Tensor]) -> list[str]: + """Rows whose contents changed with the number of rows in the launch.""" + failures = [] + for row in CHECK_ROWS: + counts = [n for n in sorted(probes) if n > row] + if not counts: + continue + reference = probes[counts[0]][row] + variant = [n for n in counts if not torch.equal(probes[n][row], reference)] + if variant: + failures.append(f"row {row} changed at row counts {variant}") + return failures + + +def _mxfp4_weights(n: int, k: int, use_asm_gemm: bool): + """Weights laid out as AiterMxfp4LinearKernel.process_weights_after_loading.""" + weight = torch.randint(0, 255, (n, k // 2), dtype=torch.uint8, device="cuda") + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + if not use_asm_gemm: + return weight, weight_scale.T.contiguous() + + from aiter.ops.shuffle import shuffle_weight + + sm, sn = weight_scale.shape + weight_scale = ( + weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) + .permute(0, 3, 5, 2, 4, 1, 6) + .contiguous() + .view(sm, sn) + ) + return shuffle_weight(weight, layout=(16, 16)), weight_scale + + +def _mxfp8_weights(n: int, k: int): + """E4M3 weights with E8M0 block scales, as the MXFP8 layer stores them.""" + weight = (torch.randn(n, k, device="cuda") / 8).to(torch.float8_e4m3fn) + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + return weight, weight_scale + + +@requires_mx +@pytest.mark.parametrize("n,k,use_asm_gemm", MXFP4_CASES) +def test_mxfp4_linear_is_batch_invariant(n: int, k: int, use_asm_gemm: bool): + pytest.importorskip("aiter") + + # Importing the module registers torch.ops.vllm.gemm_with_dynamic_quant. + import vllm.model_executor.kernels.linear.mxfp4.aiter # noqa: F401 + + set_random_seed(SEED) + weight, weight_scale = _mxfp4_weights(n, k, use_asm_gemm) + x = torch.randn(max(MXFP4_TOKEN_COUNTS), k, device="cuda", dtype=torch.bfloat16) + + probes = { + num_tokens: _probe( + torch.ops.vllm.gemm_with_dynamic_quant( + x[:num_tokens], weight, weight_scale, use_asm_gemm, torch.bfloat16 + ) + ) + for num_tokens in MXFP4_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP4 linear depends on the row count " + f"(N={n}, K={k}, asm={use_asm_gemm}):\n " + "\n ".join(failures) + ) + + +@requires_mx +@pytest.mark.parametrize("n,k", MXFP8_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_mxfp8_linear_is_batch_invariant(n: int, k: int, dtype: torch.dtype): + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _mxfp8_dot_scaled_linear, + ) + + set_random_seed(SEED) + weight, weight_scale = _mxfp8_weights(n, k) + x = torch.randn(max(MXFP8_TOKEN_COUNTS), k, device="cuda", dtype=dtype) + + probes = { + num_tokens: _probe( + _mxfp8_dot_scaled_linear(x[:num_tokens], weight, weight_scale) + ) + for num_tokens in MXFP8_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP8 linear depends on the row count (N={n}, K={k}, {dtype}):\n " + + "\n ".join(failures) + ) diff --git a/tests/v1/determinism/test_nvfp4_batch_invariant.py b/tests/v1/determinism/test_nvfp4_batch_invariant.py index 1e9b0cc8fc56..162e6b36c3da 100644 --- a/tests/v1/determinism/test_nvfp4_batch_invariant.py +++ b/tests/v1/determinism/test_nvfp4_batch_invariant.py @@ -9,10 +9,14 @@ from utils import ( _extract_step_logprobs, _random_prompt, - skip_if_not_cuda, + skip_if_not_cuda_alike, ) from vllm import LLM, SamplingParams +from vllm.platforms import current_platform + +# FLASH_ATTN is CUDA-only. +E2E_BACKENDS = ["TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"] pytestmark = pytest.mark.skipif( not hasattr(torch, "float8_e4m3fn"), @@ -40,8 +44,8 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM: ) -@skip_if_not_cuda -@pytest.mark.parametrize("backend", ["FLASH_ATTN"]) +@skip_if_not_cuda_alike +@pytest.mark.parametrize("backend", E2E_BACKENDS) def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend): seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index dcba87dc4927..fb5db254ba43 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/tests/v1/determinism/test_online_batch_invariance.py @@ -6,6 +6,7 @@ Environment variables: - VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1) - VLLM_TP_SIZE: tensor parallelism size (e.g., 4) + - VLLM_TEST_SERVER_ARGS: extra server args (e.g., "--data-parallel-size 4") """ @@ -16,7 +17,7 @@ import openai import pytest -from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda +from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda_alike from tests.utils import RemoteOpenAIServer @@ -132,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("backend", BACKENDS) def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( backend: str, @@ -156,6 +157,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( ] if tp_size: server_args += ["-tp", tp_size] + server_args += os.getenv("VLLM_TEST_SERVER_ARGS", "").split() with RemoteOpenAIServer(TEST_MODEL, server_args) as server: client = server.get_client() diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 1feb5e67413d..f26217bd7410 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda, skip_unsupported +from utils import skip_if_not_cuda_alike, skip_unsupported from vllm.model_executor.determinism.batch_invariant import ( rms_norm_batch_invariant, @@ -27,7 +27,7 @@ def _rms_norm_reference( return (output * weight.float()).to(input_tensor.dtype) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -76,7 +76,7 @@ def test_rms_norm_batch_invariant_vs_reference( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) @@ -189,7 +189,7 @@ def _assert_rows_bit_identical(small, large, msg): torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -221,7 +221,7 @@ def rms_norm(x): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -267,7 +267,7 @@ def quant(x, res): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -303,7 +303,7 @@ def per_block_quant(x): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) @pytest.mark.parametrize("hidden_size", [2048, 4096]) @@ -343,7 +343,7 @@ def test_rms_norm_3d_input( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_rms_norm_numerical_stability(default_vllm_config): """Test RMS norm numerical stability with extreme values. @@ -431,7 +431,7 @@ def test_rms_norm_formula(default_vllm_config): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): """Test RMS norm with various hidden sizes to ensure block size handling. diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index f03ea05b4331..405b5ac3d8a5 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -34,12 +34,16 @@ class DeviceConfig(NamedTuple): available=current_platform.is_xpu() and HAS_TRITON, backends=["TRITON_ATTN"], ), + "rocm": DeviceConfig( + available=current_platform.is_rocm() and HAS_TRITON, + backends=["TRITON_ATTN"], + ), } DEFAULT_MODEL = "Qwen/Qwen3-1.7B" TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) -# Override backends for MLA models (MLA only supported on CUDA). +# Override backends for MLA models. if os.getenv("VLLM_TEST_MODEL"): config = get_config(TEST_MODEL, trust_remote_code=False) if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): @@ -52,6 +56,10 @@ class DeviceConfig(NamedTuple): available=DEVICE_BACKENDS["xpu"].available, backends=[], ) + DEVICE_BACKENDS["rocm"] = DeviceConfig( + available=DEVICE_BACKENDS["rocm"].available, + backends=["TRITON_MLA"], + ) # Only include backends for devices that are actually available. BACKENDS: list[str] = sorted( @@ -60,7 +68,7 @@ class DeviceConfig(NamedTuple): skip_unsupported = pytest.mark.skipif( not any(cfg.available for cfg in DEVICE_BACKENDS.values()), - reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton", + reason="Requires CUDA >= Ampere (SM80), ROCm, or Intel XPU with Triton", ) skip_if_not_cuda = pytest.mark.skipif( @@ -68,6 +76,16 @@ class DeviceConfig(NamedTuple): reason="Requires CUDA >= Ampere (SM80)", ) +skip_if_not_cuda_alike = pytest.mark.skipif( + not (DEVICE_BACKENDS["cuda"].available or DEVICE_BACKENDS["rocm"].available), + reason="Requires CUDA >= Ampere (SM80) or ROCm", +) + +requires_mx = pytest.mark.skipif( + not (current_platform.is_rocm() and current_platform.supports_mx()), + reason="requires a ROCm device with native MX support (gfx95x)", +) + def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: # Generate more realistic prompts that will actually produce varied tokens diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 29d5bae89d2c..1bc6182288ab 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -413,6 +413,17 @@ def all_reduce(self, input_): out = ca_comm.custom_all_reduce(input_) assert out is not None return out + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + # Integer sums are exact in any order and stay on RCCL. + from vllm.model_executor.determinism.batch_invariant import ( + all_reduce_batch_invariant, + ) + + return all_reduce_batch_invariant(input_, self.device_group) symm_mem_comm = self.symm_mem_comm if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_): out = symm_mem_comm.all_reduce(input_) @@ -498,7 +509,17 @@ def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] - if should_nccl_symm_mem_ag_rs(): + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + output = reduce_scatter_batch_invariant(input_tensor, self.device_group) + elif should_nccl_symm_mem_ag_rs(): output = self._reduce_scatter_symm_mem(input_tensor) else: output = torch.empty( @@ -563,6 +584,22 @@ def _reduce_scatterv( assert output.device == input_tensor.device assert output.is_contiguous() + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + result = reduce_scatter_batch_invariant( + input_tensor, self.device_group, sizes=sizes + ) + if output is not None: + return output.copy_(result) + return result.movedim(0, dim).contiguous() + if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 403f47570204..e4b04793feed 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -6,6 +6,7 @@ from typing import Any import torch +import torch.distributed as dist import vllm.envs as envs from vllm.model_executor.determinism.batch_invariant_configs import ( @@ -59,6 +60,7 @@ def matmul_kernel_persistent( stride_bn, stride_cm, stride_cn, + stride_bias, BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZE_N: tl.constexpr, # BLOCK_SIZE_K: tl.constexpr, # @@ -128,7 +130,7 @@ def matmul_kernel_persistent( c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) if HAS_BIAS: - bias_ptrs = bias_ptr + offs_cn + bias_ptrs = bias_ptr + offs_cn * stride_bias bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) accumulator += bias c = accumulator.to(c_ptr.dtype.element_ty) @@ -305,6 +307,8 @@ def grid(META): ), ) + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -345,6 +349,7 @@ def grid(META): b.stride(1), # c.stride(0), c.stride(1), # + bias.stride(0) if bias is not None else 1, NUM_SMS=NUM_SMS, # A_LARGE=a.numel() > 2**31, B_LARGE=b.numel() > 2**31, @@ -489,16 +494,19 @@ def bmm_kernel( @triton.jit -def _log_softmax_kernel( +def _softmax_kernel( input_ptr, output_ptr, input_row_stride, + input_col_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr, + LOG: tl.constexpr = False, ): - """Compute log_softmax along the last dimension of a 2D tensor. - Each block handles one row of the input tensor. + """Compute softmax, or log_softmax when ``LOG``, along the last dimension of a + 2D tensor. Each block handles one row of the input tensor, in a fixed block + order, so the result never depends on the row count. """ # Get the row index for this block row_idx = tl.program_id(0).to(tl.int64) @@ -514,7 +522,9 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) + vals = tl.load( + row_start_ptr + col_idx * input_col_stride, mask=mask, other=-float("inf") + ) # Update maximum max_val = tl.max(tl.maximum(vals, max_val)) @@ -526,25 +536,26 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask, other=0.0) # Compute exp(x - max_val) and accumulate exp_vals = tl.exp(vals - max_val) sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) - # Compute log(sum_exp) - log_sum_exp = tl.log(sum_exp) + log_sum_exp = tl.log(sum_exp) if LOG else 0.0 - # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp + # Step 3: normalise for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask) - # Compute log_softmax - output = vals - max_val - log_sum_exp + if LOG: + output = vals - max_val - log_sum_exp + else: + output = tl.exp(vals - max_val) / sum_exp # Store results tl.store(output_row_start_ptr + col_idx, output, mask=mask) @@ -570,30 +581,55 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) - input_2d = input_2d.contiguous() n_rows, n_cols = input_2d.shape # Allocate output tensor - output = torch.empty_like(input_2d) + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) # Choose block size based on the number of columns BLOCK_SIZE = 1024 # Launch kernel with one block per row grid = (n_rows,) - _log_softmax_kernel[grid]( + _softmax_kernel[grid]( input_2d, output, input_2d.stride(0), + input_2d.stride(1), output.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, + LOG=True, ) # Reshape output back to original shape return output.reshape(original_shape) +def softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Compute softmax along the last dimension using a Triton kernel.""" + if dim != -1 and dim != input.ndim - 1: + raise ValueError( + "This implementation only supports softmax along the last dimension" + ) + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + n_rows, n_cols = input_2d.shape + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) + + _softmax_kernel[(n_rows,)]( + input_2d, + output, + input_2d.stride(0), + input_2d.stride(1), + output.stride(0), + n_cols, + BLOCK_SIZE=1024, + ) + return output.reshape(original_shape) + + @triton.jit def mean_kernel( input_ptr, @@ -735,6 +771,9 @@ def mean_dim( def mm_batch_invariant(a, b, bias=None): if a.device.type == "xpu": + # The descriptor kernel reads the bias at unit stride. + if bias is not None: + bias = bias.contiguous() return matmul_descriptor_persistent(a, b, bias=bias) return matmul_persistent(a, b, bias=bias) @@ -820,6 +859,8 @@ def bmm_batch_invariant(a, b, *, out=None): assert out.dtype == dtype and out.device == a.device, "out tensor mismatch" c = out + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -877,8 +918,50 @@ def bmm_batch_invariant(a, b, *, out=None): return c -def addmm_batch_invariant(bias, a, b): - return mm_batch_invariant(a, b, bias=bias) +def _addmm_fused_bias(bias, beta, alpha, n): + # The kernel adds bias into its fp32 accumulator indexed by output column, + # so it can only absorb a row vector of width n, unscaled by alpha. + if beta == 0 or alpha != 1: + return None + if bias.dim() != 1 or bias.shape[0] != n: + return None + # fp32 keeps the scaling exact; the kernel casts the bias to fp32 anyway. + return bias if beta == 1 else bias.float() * beta + + +def _addmm_impl(bias, a, b, beta, alpha, out): + fused_bias = _addmm_fused_bias(bias, beta, alpha, b.shape[1]) + result = mm_batch_invariant(a, b, bias=fused_bias) + if fused_bias is None: + if alpha != 1: + result = result * alpha + if beta != 0: + result = result + (bias if beta == 1 else beta * bias) + if out is None: + return result + return out.copy_(result) + + +def addmm_batch_invariant(bias, a, b, *, beta=1, alpha=1): + return _addmm_impl(bias, a, b, beta, alpha, None) + + +# Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, which +# dispatches the ``.out`` overload rather than the default one. Registering only +# the default leaves the compiled path on the vendor GEMM, so a compiled +# ``torch.addmm`` is batch variant while the eager one is not. + + +def mm_out_batch_invariant(a, b, *, out): + return matmul_batch_invariant(a, b, out=out) + + +def addmm_out_batch_invariant(bias, a, b, *, beta=1, alpha=1, out): + return _addmm_impl(bias, a, b, beta, alpha, out) + + +def bmm_out_batch_invariant(a, b, *, out): + return bmm_batch_invariant(a, b, out=out) def _log_softmax_batch_invariant(input, dim, _half_to_float): @@ -887,10 +970,16 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float): return log_softmax(input, dim=dim) -def softmax_batch_invariant(input, dim, dtype=None): - # Compute softmax in a deterministic way - # First subtract max for numerical stability (standard practice) +def softmax_batch_invariant(input, dim, half_to_float=False): + if half_to_float: + input = input.float() + if dim == -1 or dim == input.ndim - 1: + return softmax(input, dim=-1) + + # Off the last dimension this falls back to torch ops, whose reduction split + # depends on the tensor shape. input_max = torch.amax(input, dim=dim, keepdim=True) + # First subtract max for numerical stability (standard practice) input = input - input_max exp_x = torch.exp(input) sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True) @@ -1035,11 +1124,150 @@ def rms_norm_batch_invariant( def linear_batch_invariant(input, weight, bias=None): - output = matmul_batch_invariant(input, weight.t()) - if bias is not None: - output = output + bias - return output + bias = bias.expand(weight.shape[0]) + output = mm_batch_invariant( + input.reshape(-1, input.shape[-1]), weight.t(), bias=bias + ) + return output.view(*input.shape[:-1], output.shape[-1]) + + +@triton.jit +def _fixed_order_sum_kernel( + src_ptr, # (WORLD_SIZE, numel) gathered contributions, rank-major + out_ptr, # (numel,) + numel, + stride_rank, + WORLD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + LARGE: tl.constexpr, +): + """Sum ``WORLD_SIZE`` contributions in ascending rank order, fp32 accumulator. + + An output element sums exactly ``WORLD_SIZE`` values in a compile-time + constant order and rounds once, so the result depends only on the rank + ordering and never on how many elements are being reduced. + """ + pid = tl.program_id(0) + if LARGE: + pid = pid.to(tl.int64) + stride_rank = stride_rank.to(tl.int64) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < numel + + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for r in tl.static_range(WORLD_SIZE): + vals = tl.load(src_ptr + r * stride_rank + offs, mask=mask, other=0.0) + acc += vals.to(tl.float32) + + tl.store(out_ptr + offs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +def all_reduce_batch_invariant( + input_: torch.Tensor, group: "torch.distributed.ProcessGroup | None" = None +) -> torch.Tensor: + """Sum all-reduce whose result does not depend on the message size. + + Library all-reduces (NCCL/RCCL) pick their algorithm, channel count and chunk + boundaries from the message size, so the order in which a given element's + contributions are summed changes with the number of tokens in the batch. + + Instead, all-gather the contributions -- pure data movement, so bitwise + reproducible at any size -- and reduce them with ``_fixed_order_sum_kernel``. + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + numel = x.numel() + gathered = torch.empty((world_size, numel), dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(gathered, x.view(-1), group=group) + + out = torch.empty_like(x) + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + gathered, + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=gathered.numel() > 2**31, + ) + return out + + +def reduce_scatter_batch_invariant( + input_: torch.Tensor, + group: "torch.distributed.ProcessGroup | None" = None, + sizes: list[int] | None = None, +) -> torch.Tensor: + """Sum reduce-scatter over dim 0 whose result does not depend on the size. + + ``ncclReduceScatter`` picks its chunking from the message size just as + ``ncclAllReduce`` does. Route around it the same way + ``all_reduce_batch_invariant`` does: an all-to-all sends rank ``d`` exactly + the rows it owns from every rank -- pure data movement, so bitwise + reproducible however the library chunks it -- and lands them rank-major, + which is the layout ``_fixed_order_sum_kernel`` already consumes. + + ``sizes`` (the reduce-scatterv case) changes only *which* rows a rank + receives, never how the received contributions are summed, so variable + shard sizes are invariant for the same reason. + + Args: + input_: Contiguous full-size input; dim 0 is the scattered axis. + group: Process group, defaults to the world group. + sizes: Per-rank row counts. Uniform split when ``None``. + + Returns: + This rank's shard of the sum. + + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + if sizes is None: + assert x.shape[0] % world_size == 0 + sizes = [x.shape[0] // world_size] * world_size + else: + assert len(sizes) == world_size + assert x.shape[0] == sum(sizes) + + rows = sizes[dist.get_rank(group)] + out = torch.empty((rows, *x.shape[1:]), dtype=x.dtype, device=x.device) + recv = torch.empty( + (world_size * rows, *x.shape[1:]), dtype=x.dtype, device=x.device + ) + # Every rank takes part even when its own shard is empty: it still has rows + # to send to the others. + dist.all_to_all_single( + recv, + x, + output_split_sizes=[rows] * world_size, + input_split_sizes=sizes, + group=group, + ) + + numel = out.numel() + if numel == 0: + return out + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + recv.view(-1), + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=recv.numel() > 2**31, + ) + return out _batch_invariant_MODE = False @@ -1047,9 +1275,33 @@ def linear_batch_invariant(input, weight, bias=None): _fp16_block_size_n = 256 _fp32_block_size_n = 128 _fp32_num_stages = 3 +_BLOCK_SIZES_RESOLVED = False _NUM_SMS: int = 0 +def _resolve_block_sizes(): + """Size the tiles to the shared memory, to avoid triton OutOfResources. + + Runs when the mode is enabled, or on first use by a caller that never + enables it, such as a kernel unit test. + """ + global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages + global _BLOCK_SIZES_RESOLVED + _BLOCK_SIZES_RESOLVED = True + if not current_platform.is_cuda_alike(): + return + if get_max_shared_memory_bytes() > 106496: + _fp16_block_size_n = 256 + _fp32_block_size_n = 128 + _fp32_num_stages = 3 + else: + _fp16_block_size_n = 128 + # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). + # Only N=1 OORs; keep 128/stages=3 for every other N. + _fp32_block_size_n = 32 + _fp32_num_stages = 2 + + def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages @@ -1062,12 +1314,17 @@ def enable_batch_invariant_mode(): _batch_invariant_LIB = torch.library.Library("aten", "IMPL") key = current_platform.dispatch_key - if current_platform.is_cuda(): - if current_platform.is_device_capability_family(80): - # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the - # triton persistent matmul overrides for mm/addmm/matmul/linear. + if current_platform.is_cuda_alike(): + if current_platform.is_rocm() or current_platform.is_device_capability_family( + 80 + ): + # Neither cuBLASLt on SM80 (Ampere) nor hipBLASLt can be made batch + # invariant; install the triton persistent matmul overrides for + # mm/addmm/matmul/linear. _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) else: @@ -1077,18 +1334,7 @@ def enable_batch_invariant_mode(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" - # Query the shared memory size and set block size - # accordingly to avoid triton OutOfResources. - if get_max_shared_memory_bytes() > 106496: - _fp16_block_size_n = 256 - _fp32_block_size_n = 128 - _fp32_num_stages = 3 - else: - _fp16_block_size_n = 128 - # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). - # Only N=1 OORs; keep 128/stages=3 for every other N. - _fp32_block_size_n = 32 - _fp32_num_stages = 2 + _resolve_block_sizes() elif current_platform.is_xpu(): # Tensor descriptors need a global-memory allocator; must be set outside # torch.compile regions (triton.set_allocator modifies global state). @@ -1100,6 +1346,8 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _NUM_SMS = num_compute_units(0) _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) @@ -1123,7 +1371,9 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _batch_invariant_LIB.impl( "aten::bmm", bmm_batch_invariant, key, allow_override=True ) - torch.bmm = bmm_batch_invariant + _batch_invariant_LIB.impl( + "aten::bmm.out", bmm_out_batch_invariant, key, allow_override=True + ) reduced_precision_val = ( (False, False) if is_torch_equal_or_newer("2.10.0") else False @@ -1145,28 +1395,32 @@ def override_envs_for_invariance(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, - # unlike the pins below, it also applies to communicators shared with other - # processes. - os.environ["NCCL_LAUNCH_MODE"] = "GROUP" - pin_nccl_env( - { - "NCCL_COLLNET_ENABLE": "0", - "NCCL_NVLS_ENABLE": "0", - "NCCL_P2P_NET_DISABLE": "1", - "NCCL_MIN_NCHANNELS": "1", - "NCCL_MAX_NCHANNELS": "1", - "NCCL_PROTO": "Simple", - # NCCL >= 2.31 zero-fills the algorithm table of every collective - # when NCCL_ALGO is set and re-enables only the named ones; together - # with the NCCL_PROTO above, collectives not named here end up with - # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and - # Tree for all collectives, then pin AllReduce to Tree. - "NCCL_ALGO": "ring,tree;allreduce:tree", - "NCCL_NTHREADS": "1", - "NCCL_SOCKET_NTHREADS": "1", - } - ) + if current_platform.is_rocm(): + # The skinny GEMMs pick their split of K from the batch size. + os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" + else: + # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, + # unlike the pins below, it also applies to communicators shared with other + # processes. + os.environ["NCCL_LAUNCH_MODE"] = "GROUP" + pin_nccl_env( + { + "NCCL_COLLNET_ENABLE": "0", + "NCCL_NVLS_ENABLE": "0", + "NCCL_P2P_NET_DISABLE": "1", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "NCCL_PROTO": "Simple", + # NCCL >= 2.31 zero-fills the algorithm table of every collective + # when NCCL_ALGO is set and re-enables only the named ones; together + # with the NCCL_PROTO above, collectives not named here end up with + # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and + # Tree for all collectives, then pin AllReduce to Tree. + "NCCL_ALGO": "ring,tree;allreduce:tree", + "NCCL_NTHREADS": "1", + "NCCL_SOCKET_NTHREADS": "1", + } + ) # torch.compile settings os.environ["VLLM_USE_AOT_COMPILE"] = "0" diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 484143ffc14a..0ce22f477ac7 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -218,6 +218,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( TritonFp8BlockScaledMMKernel, + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -253,6 +254,20 @@ def _get_linear_backend(*, quantization: str) -> str: return config.kernel_config.linear_backend +def _batch_invariant_fp8_kernel(kernel, config, quantization: str): + can_implement, reason = is_supported_and_can_implement_kernel(kernel, config, None) + if not can_implement: + raise ValueError(f"Batch invariance on ROCm: {reason}") + linear_backend = _get_linear_backend(quantization=quantization) + if linear_backend not in ("auto", "triton"): + logger.warning_once( + "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; " + "using the Triton backend for deterministic execution.", + linear_backend, + ) + return kernel + + # Kernel classes covered by each --linear-backend value. _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "b12x": { @@ -723,6 +738,13 @@ def init_fp8_linear_kernel( ) if activation_quant_key.scale.group_shape.is_per_group(): + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFp8BlockScaledMMKernel, + scaled_mm_linear_kernel_config, + "fp8_block_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] @@ -755,6 +777,13 @@ def init_fp8_linear_kernel( ) else: + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFP8ScaledMMLinearKernel, + scaled_mm_linear_kernel_config, + "fp8_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] @@ -780,6 +809,19 @@ def init_fp8_linear_kernel( ) +def fp8_scaled_mm_is_batch_invariant( + kernel: FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel, +) -> bool: + """Whether a row's output from the kernel is independent of the other rows.""" + if not isinstance( + kernel, (CutlassFP8ScaledMMLinearKernel, TritonFP8ScaledMMLinearKernel) + ): + return False + # A dynamic per-tensor activation scale is an amax over every row. + scale = kernel.config.activation_quant_key.scale + return scale.static or not scale.group_shape.is_per_tensor() + + def init_int8_linear_kernel( is_channelwise: bool, is_static_input_scheme: bool, @@ -1250,6 +1292,7 @@ def register_linear_kernel( __all__ = [ + "fp8_scaled_mm_is_batch_invariant", "init_fp8_linear_kernel", "init_int8_linear_kernel", "init_nvfp4_linear_kernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 420e7f9493e1..4cd09c1f905f 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -62,7 +62,12 @@ def gemm_with_dynamic_quant( N = weight.shape[0] K = weight.shape[1] if rocm_use_aiter_fp4_asm_gemm: - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): + # The preshuffle kernel's tuned configs split K in some M buckets. + if ( + M <= 64 + and not envs.VLLM_BATCH_INVARIANT + and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K) + ): if x_scales is None: # use hip quant kernel for performance if M >= 32: @@ -116,7 +121,16 @@ def gemm_with_dynamic_quant( x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype ) - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) + config = None + if envs.VLLM_BATCH_INVARIANT: + # The configs split K in some M buckets (up to M <= 256 for + # tuned shapes). Keep the tile, drop the split. + from aiter.ops.triton.gemm_afp4wfp4 import _get_config + + config, _ = _get_config(M, N, K) + config = dict(config, NUM_KSPLIT=1) + + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y, config) return y def gemm_with_dynamic_quant_fake( diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index faf1e180fa2f..98920f2df13a 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -12,6 +12,7 @@ import torch from torch.nn.parameter import Parameter +import vllm.envs as envs from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, @@ -246,7 +247,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() - block_scale = _as_block32_scale(weight_scale) + # rocm_mxfp8_block32_gemm's split-K varies with M. + block_scale = ( + None if envs.VLLM_BATCH_INVARIANT else _as_block32_scale(weight_scale) + ) if block_scale is not None: # Checkpoints with 32x32 scale blocks (DeepSeek V4/V4.1) get them # back from the loader expanded per row; keep the blocks and use diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 75db4bcaa9ed..ac096b47b96a 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -39,6 +39,7 @@ ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -65,6 +66,7 @@ "PerTensorTorchFP8ScaledMMLinearKernel", "RowWiseTorchFP8ScaledMMLinearKernel", "ROCmFP8ScaledMMLinearKernel", + "TritonFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index d6169d2cfc1a..fed471e474ee 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -20,6 +20,8 @@ ) from .cutlass import CutlassInt8ScaledMMLinearKernel from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, Int8ScaledMMLinearLayerConfig, ) @@ -156,6 +158,46 @@ def apply_weights( return out +class TritonFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda_alike(): + return False, "requires ROCm or CUDA." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + act_group = c.activation_quant_key.scale.group_shape + weight_group = c.weight_quant_key.scale.group_shape + if not (act_group.is_per_tensor() or act_group.is_per_token()): + return False, "requires per-tensor or per-token activation scales." + if not ( + weight_group.is_per_tensor() + or weight_group.is_per_channel() + or weight_group.is_per_token() + ): + return False, "requires per-tensor, per-channel or per-token weight scales." + return True, None + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + out = triton_scaled_mm( + A, B, scale_a=As, scale_b=Bs, out_dtype=out_dtype, bias=bias + ) + return out.view(*output_shape) + + class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): @classmethod def is_supported(cls, compute_capability=None): diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c15a1a13d42e..1efb058ef07c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -51,6 +51,10 @@ def _supports_current_device() -> bool: def _supports_parallel_config(moe_parallel_config) -> bool: return True + @staticmethod + def _supports_batch_invariance() -> bool: + return False + @staticmethod def is_supported_config( cls, moe_config, weight_key, activation_key, activation_format diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index da0bca4977b7..ed44bf4cb135 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -218,7 +218,7 @@ def apply( global_num_experts=global_num_experts, expert_map=expert_map, a1q_scale=None, - a2_scale=None, + a2_scale=a2_scale, workspace13=workspace13, workspace2=workspace2, expert_tokens_meta=expert_tokens_meta, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py index e0d1fdf27170..8adb41151c43 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py @@ -150,16 +150,16 @@ def scaled_mm_kernel( scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1)) accumulator = scale_b.T * accumulator.to(tl.float32) - # Convert to output format. - c = accumulator.to(c_ptr.type.element_ty) - - # Add bias, it's already in output format, so add it after conversion. + # Add bias in accumulator dtype for less rounding error. if bias_ptr: offsets_bias = offsets_bn bias_ptrs = bias_ptr + offsets_bias bias_mask = offsets_bias < N bias = tl.load(bias_ptrs, bias_mask) - c += bias + accumulator += bias + + # Convert to output format. + c = accumulator.to(c_ptr.type.element_ty) # Save output offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) @@ -224,6 +224,8 @@ def triton_scaled_mm( tile_shape = (64, 128, 128) else: tile_shape = (128, 128, 128) + else: + tile_shape = (block_size_m, block_size_n, block_size_k) block_size_m, block_size_n, block_size_k = tile_shape diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 1c1d18614ecd..b18ec0a0aab3 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -9,10 +9,10 @@ from vllm.config import get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, init_fp8_linear_kernel, ) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.attention import Attention @@ -432,7 +432,7 @@ def apply( bias, ) else: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) # per-tensor/channel: dequant to BF16 and run GEMM diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index f4662ad5a66e..8d8a7eca6066 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -15,9 +15,11 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config -from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, + init_fp8_linear_kernel, +) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -278,7 +280,7 @@ def apply( ) -> torch.Tensor: # if batch invariant mode is enabled, use BF16 dequant if envs.VLLM_BATCH_INVARIANT: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) weight_fp8 = layer.weight.to(torch.bfloat16) @@ -466,7 +468,7 @@ def apply( # if batch invariant mode is enabled dequant if ( envs.VLLM_BATCH_INVARIANT - and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) + and not fp8_scaled_mm_is_batch_invariant(self.fp8_linear) and not isinstance(x, QuantizedActivation) ): weight_dequant = ( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index fb87c3a61450..00f9e5b92eed 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1041,7 +1041,10 @@ def w8a8_triton_block_scaled_mm( C_shape = A.shape[:-1] + (N,) C = A.new_empty(C_shape, dtype=output_dtype) - configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) + if envs.VLLM_BATCH_INVARIANT: + configs = None # use the default config + else: + configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) if configs: # Get the optimal config if there is one config = configs[min(configs.keys(), key=lambda x: abs(x - M))] diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 80eef7c75a12..8d285dda75b8 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -288,6 +288,13 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT: + from vllm.model_executor.determinism.batch_invariant import ( + linear_batch_invariant, + ) + + return linear_batch_invariant(x, weight, bias) + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 5840ffe3b792..807e571002ac 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1068,6 +1068,15 @@ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops: compilation_config.custom_ops.append("+quant_fp8") + # Only RMSNorm's custom op reaches rms_norm_batch_invariant. Inductor + # lowers the native one to a reduction whose block size it picks by + # benchmarking, so the summation order varies with the batch and run. + if ( + envs.VLLM_BATCH_INVARIANT + and "-rms_norm" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+rms_norm") + if use_aiter_fused_se and "-grouped_topk" in compilation_config.custom_ops: logger.warning_once( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled, which " diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 8b0e337bb221..ecc2a50135f6 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -12,6 +12,7 @@ import torch +import vllm.envs as envs from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend @@ -47,6 +48,15 @@ def is_available(cls) -> bool: return rocm_aiter_ops.is_enabled() + @classmethod + def validate_configuration(cls, device_capability, selector_config): + invalid_reasons = super().validate_configuration( + device_capability, selector_config + ) + if envs.VLLM_BATCH_INVARIANT: + invalid_reasons.append("batch invariance not supported") + return invalid_reasons + def __init__( self, num_heads: int, diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 3af623b1f13b..b6cf96e5b738 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -427,13 +427,13 @@ def _flash_attn_varlen_diff_headdims( kwargs["return_softmax_lse"] = return_softmax_lse kwargs["out"] = out kwargs["output_scale"] = output_scale + if envs.VLLM_BATCH_INVARIANT: + kwargs["num_splits"] = 1 else: # ROCm leverages the upstream flash_attn, which takes a parameter # called "return_attn_probs" instead of return_softmax_lse kwargs["return_attn_probs"] = return_softmax_lse assert out is None and output_scale is None - if envs.VLLM_BATCH_INVARIANT: - kwargs["num_splits"] = 1 attn_out = _FA4_MLA_PREFILL_KERNEL( q=q, From d79b52b18219ecbaba780f72ae0a4aa0546ac00f Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 12:46:22 -0400 Subject: [PATCH 425/952] [Misc] Consolidate csrc shared headers and prune dead code (#60721) Signed-off-by: mgoin Co-authored-by: Claude --- .buildkite/test-amd.yaml | 46 +- .buildkite/test_areas/distributed.yaml | 4 +- .buildkite/test_areas/kernels.yaml | 35 +- .buildkite/test_areas/lm_eval.yaml | 8 +- .buildkite/test_areas/models_basic.yaml | 2 +- .gitignore | 3 - .pre-commit-config.yaml | 2 +- CMakeLists.txt | 235 ++-- cmake/cpu_extension.cmake | 6 +- cmake/hipify.py | 2 +- .../activation_kernels.cu | 2 +- csrc/{libtorch_stable => }/async_util.cuh | 0 .../dcp_utils/dcp_direct_a2a_lse_reduce.cu | 0 .../attention/dcp_utils/dcp_direct_common.cuh | 0 .../dcp_utils/dcp_direct_kv_gather.cu | 0 .../dcp_utils/dcp_direct_q_gather.cu | 0 .../attention/merge_attn_states.cu | 5 +- .../cutlass_sm100_mla/device/sm100_mla.hpp | 0 .../kernel/sm100_fmha_mla_reduction.hpp | 0 .../sm100_fmha_mla_tma_warpspecialized.hpp | 0 .../kernel/sm100_mla_tile_scheduler.hpp | 0 .../attention/mla/sm100_cutlass_mla_kernel.cu | 2 +- csrc/cache.h | 93 -- csrc/{libtorch_stable => }/cache_kernels.cu | 22 +- .../cache_kernels_fused.cu | 8 +- csrc/{libtorch_stable => }/concat_mla_q.cuh | 0 .../{libtorch_stable => }/cooperative_topk.cu | 0 .../cooperative_topk.cuh | 0 csrc/core/batch_invariant.hpp | 18 - csrc/{ => core}/cuda_compat.h | 6 - csrc/core/exception.hpp | 3 - csrc/core/registration.h | 5 - .../core/math.hpp => core/utils.hpp} | 16 +- .../vec_dtypes}/dtype_bfloat16.cuh | 2 +- .../vec_dtypes}/dtype_float16.cuh | 2 +- .../vec_dtypes}/dtype_float32.cuh | 2 +- .../vec_dtypes}/dtype_fp8.cuh | 7 +- .../vec_dtypes/dtype_generic.cuh} | 0 .../vec_dtypes/vec_dtypes.h} | 2 +- csrc/cpu/cpu_attn_impl.hpp | 2 +- .../{moe => cpu}/dynamic_4bit_int_moe_cpu.cpp | 0 csrc/cpu/ops.h | 46 + csrc/cpu/torch_bindings.cpp | 2 +- csrc/{libtorch_stable => }/cub_helpers.h | 0 csrc/cuda_utils.h | 25 +- .../cuda_utils_kernels.cu | 0 csrc/{libtorch_stable => }/cuda_vec_utils.cuh | 0 csrc/{libtorch_stable => }/cuda_view.cu | 0 .../custom_all_gather_reduce_scatter.cu | 0 .../custom_all_gather_reduce_scatter_ops.cpp | 0 .../custom_all_reduce.cu | 0 csrc/custom_collective_common.cuh | 31 - .../cutlass_extensions/common.hpp | 25 +- .../epilogue/broadcast_load_epilogue_c2x.hpp | 0 .../epilogue/scaled_mm_epilogues_c2x.hpp | 0 .../epilogue/scaled_mm_epilogues_c3x.hpp | 40 +- .../cutlass_extensions/torch_utils.hpp | 49 +- .../vllm_collective_builder.cuh | 0 .../vllm_numeric_conversion.cuh | 6 +- csrc/cutlass_extensions/vllm_type_utils.cuh | 2 +- csrc/dispatch_utils.h | 302 +++-- .../dsv3_fused_a_gemm.cu | 2 +- .../{libtorch_stable => }/fp32_router_gemm.cu | 0 .../fp32_router_gemm_entry.cu | 2 +- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 2 +- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 8 +- ...minimax_m3_qknorm_rope_kv_insert_kernel.cu | 8 +- .../fused_qknorm_rope_kernel.cu | 2 +- .../gdn/fused_gdn_decode_kernel.cu | 2 +- .../{libtorch_stable => }/hisparse_kernels.cu | 1 - .../kimi_k3/attn_res_kernel.cu | 0 .../kimi_k3/fused_kda_chunk_kernel_rocm.cu | 0 .../kimi_k3/fused_kda_decode_kernel.cu | 2 +- .../kimi_k3/fused_kda_decode_kernel_rocm.cu | 0 .../launch_bounds_utils.h | 0 .../layernorm_kernels.cu | 2 +- .../layernorm_quant_kernels.cu | 4 +- .../attention/attention_utils.cuh | 57 - .../cutlass_extensions/common.cpp | 11 - csrc/libtorch_stable/dispatch_utils.h | 188 --- csrc/libtorch_stable/mamba/static_switch.h | 28 - .../moe/marlin_moe_wna16/generate_kernels.py | 315 ----- csrc/libtorch_stable/ops.h | 822 ------------ .../w8a8/int8/per_token_group_quant.cu | 12 - .../w8a8/per_token_group_quant_8bit.h | 20 - csrc/libtorch_stable/torch_bindings.cpp | 1171 ---------------- csrc/libtorch_stable/torch_utils.h | 88 -- .../mamba/selective_scan.h | 0 .../mamba/selective_scan_fwd.cu | 8 +- .../minimax_reduce_rms_kernel.cu | 2 +- .../minimax_reduce_rms_kernel.h | 0 .../moe/grouped_topk_kernels.cu | 2 +- .../moe/marlin_moe_wna16/.gitignore | 1 - .../moe/marlin_moe_wna16/kernel.h | 4 +- .../moe/marlin_moe_wna16/marlin_template.h | 8 +- .../moe/marlin_moe_wna16/ops.cu | 2 +- .../moe/moeTopKFuncs.cuh | 0 .../moe/moe_align_sum_kernels.cu | 10 +- csrc/{libtorch_stable => }/moe/moe_ops.h | 0 .../moe/moe_permute_unpermute_op.cu | 4 +- csrc/{libtorch_stable => }/moe/moe_wna16.cu | 2 +- .../moe/moe_wna16_utils.h | 0 .../moe/permute_unpermute_kernels/dispatch.h | 0 .../moe_permute_unpermute_kernel.cu | 2 +- .../moe_permute_unpermute_kernel.h | 2 +- .../moe_permute_unpermute_kernel.inl | 0 .../moe/topk_softmax_kernels.cu | 4 +- .../moe/topk_softplus_sqrt_kernels.cu | 4 +- .../moe/torch_bindings.cpp | 0 .../ngram_embedding_kernels.cu | 0 .../nvfp4_ds_mla_cache.h | 0 .../nvfp4_ds_mla_cache_kernels.cu | 0 .../nvfp4_kv_cache_kernels.cu | 6 +- csrc/ops.h | 848 +++++++++++- .../{libtorch_stable => }/persistent_topk.cuh | 0 .../pos_encoding_kernels.cu | 2 +- .../quantization/activation_kernels.cu | 8 +- .../quantization/awq/dequantize.cuh | 0 .../quantization/awq/gemm_kernels.cu | 4 +- .../cutlass_w4a8/get_group_starts.cuh | 2 +- .../cutlass_w4a8/w4a8_grouped_mm_entry.cu | 8 +- .../cutlass_w4a8/w4a8_mm_entry.cu | 8 +- .../quantization/cutlass_w4a8/w4a8_utils.cu | 0 .../quantization/cutlass_w4a8/w4a8_utils.cuh | 0 .../activation_nvfp4_quant_fusion_kernels.cu | 6 +- .../fp4/mxfp4_blockwise_moe_kernel.cu | 4 +- .../quantization/fp4/mxfp4_experts_quant.cu | 8 +- .../fp4/nvfp4_blockwise_moe_kernel.cu | 4 +- .../quantization/fp4/nvfp4_experts_quant.cu | 6 +- .../quantization/fp4/nvfp4_quant_entry.cu | 4 +- .../quantization/fp4/nvfp4_quant_kernels.cu | 6 +- .../quantization/fp4/nvfp4_scaled_mm_entry.cu | 4 +- .../fp4/nvfp4_scaled_mm_kernels.cu | 7 +- .../fp4/nvfp4_scaled_mm_sm120_kernels.cu | 7 +- .../quantization/fp4/nvfp4_utils.cuh | 0 ...fused_layernorm_dynamic_per_token_quant.cu | 2 +- .../fused_silu_mul_block_quant.cu | 0 .../fused_kernels/layernorm_utils.cuh | 4 +- .../fused_kernels/quant_conversions.cuh | 4 +- .../quantization/gptq/compat.cuh | 0 .../quantization/gptq/matrix_view.cuh | 0 .../quantization/gptq/q_gemm.cu | 0 .../quantization/gptq/qdq_2.cuh | 0 .../quantization/gptq/qdq_3.cuh | 0 .../quantization/gptq/qdq_4.cuh | 0 .../quantization/gptq/qdq_8.cuh | 0 .../quantization/gptq/qdq_util.cuh | 0 .../hadacore/hadamard_transform_cuda.cu | 4 +- .../quantization/machete/Readme.md | 0 .../quantization/machete/generate.py | 3 +- .../machete/machete_collective_builder.cuh | 2 +- .../machete/machete_interleaving_utils.cuh | 0 .../quantization/machete/machete_mainloop.cuh | 5 +- .../machete/machete_mm_kernel.cuh | 6 +- .../machete/machete_mm_launcher.cuh | 4 +- .../machete/machete_prepack_kernel.cuh | 2 +- .../machete/machete_prepack_launcher.cuh | 4 +- .../machete/machete_prepacked_layout.cuh | 0 .../quantization/machete/machete_pytorch.cu | 0 .../quantization/marlin/.gitignore | 1 - .../quantization/marlin/awq_marlin_repack.cu | 2 +- .../quantization/marlin/dequant.h | 0 .../quantization/marlin/generate_kernels.py | 11 +- .../quantization/marlin/gptq_marlin_repack.cu | 2 +- .../quantization/marlin/kernel.h | 0 .../quantization/marlin/marlin.cu | 2 +- .../quantization/marlin/marlin.cuh | 0 .../quantization/marlin/marlin_dtypes.cuh | 0 .../marlin/marlin_int4_fp8_preprocess.cu | 2 +- .../quantization/marlin/marlin_mma.h | 0 .../quantization/marlin/marlin_template.h | 0 .../quantization/vectorization.cuh | 0 .../quantization/vectorization_utils.cuh | 0 .../w8a8/cutlass/c3x/cutlass_gemm_caller.cuh | 6 +- .../w8a8/cutlass/c3x/scaled_mm.cuh | 4 +- .../cutlass/c3x/scaled_mm_azp_sm90_int8.cu | 2 +- .../c3x/scaled_mm_blockwise_sm100_fp8.cu | 2 +- ...scaled_mm_blockwise_sm100_fp8_dispatch.cuh | 6 +- .../c3x/scaled_mm_blockwise_sm120_fp8.cu | 2 +- ...scaled_mm_blockwise_sm120_fp8_dispatch.cuh | 1 - .../c3x/scaled_mm_blockwise_sm90_fp8.cu | 2 +- .../scaled_mm_blockwise_sm90_fp8_dispatch.cuh | 0 .../w8a8/cutlass/c3x/scaled_mm_helper.hpp | 18 +- .../w8a8/cutlass/c3x/scaled_mm_kernels.hpp | 0 .../w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu | 2 +- .../c3x/scaled_mm_sm100_fp8_dispatch.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu | 4 +- .../c3x/scaled_mm_sm120_fp8_dispatch.cuh | 0 .../w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu | 2 +- .../c3x/scaled_mm_sm90_fp8_dispatch.cuh | 2 +- .../w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu | 2 +- .../c3x/scaled_mm_sm90_int8_dispatch.cuh | 0 .../w8a8/cutlass/moe/get_group_starts.cuh | 2 +- .../w8a8/cutlass/moe/grouped_mm_c3x.cuh | 6 +- .../w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu | 2 +- .../w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu | 2 +- .../quantization/w8a8/cutlass/moe/moe_data.cu | 4 +- .../w8a8/cutlass/scaled_mm_c2x.cu | 4 +- .../w8a8/cutlass/scaled_mm_c2x.cuh | 6 +- .../cutlass/scaled_mm_c2x_sm75_dispatch.cuh | 0 .../cutlass/scaled_mm_c2x_sm80_dispatch.cuh | 0 .../scaled_mm_c2x_sm89_fp8_dispatch.cuh | 0 .../scaled_mm_c2x_sm89_int8_dispatch.cuh | 0 .../w8a8/cutlass/scaled_mm_c3x_sm100.cu | 0 .../w8a8/cutlass/scaled_mm_c3x_sm120.cu | 0 .../w8a8/cutlass/scaled_mm_c3x_sm90.cu | 0 .../w8a8/cutlass/scaled_mm_entry.cu | 4 +- .../quantization/w8a8/fp8/amd/quant_utils.cuh | 2 +- .../quantization/w8a8/fp8/common.cu | 4 +- csrc/quantization/w8a8/fp8/common.cuh | 16 +- .../w8a8/fp8/nvidia/quant_utils.cuh | 4 +- .../w8a8/fp8/per_token_group_quant.cu | 31 +- .../quantization/w8a8/int8/scaled_quant.cu | 0 csrc/rocm/attention.cu | 6 +- csrc/rocm/dispatch_utils.h | 24 + csrc/rocm/q_gemm_rdna3.cu | 3 +- csrc/{ => rocm}/quickreduce/base.h | 0 .../quickreduce}/custom_quickreduce.cu | 2 +- csrc/{ => rocm}/quickreduce/quick_reduce.h | 0 .../quickreduce/quick_reduce_impl.cuh | 0 csrc/rocm/quickreduce/torch_bindings.cpp | 51 + csrc/rocm/skinny_gemms.cu | 4 +- csrc/rocm/skinny_gemms_int4.cu | 3 +- csrc/{libtorch_stable => }/sampled_topk.cuh | 0 csrc/{libtorch_stable => }/sampler.cu | 2 +- csrc/{libtorch_stable => }/topk.cu | 0 .../topk_histogram_4096.cuh | 0 csrc/torch_bindings.cpp | 1194 ++++++++++++++++- csrc/torch_utils.h | 101 +- csrc/{libtorch_stable => }/type_convert.cuh | 0 .../vocab_parallel_embedding_kernels.cu | 0 docs/design/fusions.md | 4 +- tests/kernels/test_bf16_skinny_gemm.py | 5 +- tools/ci/kernel_symbol_map.py | 6 +- .../device_communicators/quick_all_reduce.py | 2 +- .../kernels/linear/scaled_mm/cutlass.py | 3 +- vllm/models/common/ops/fused_qk_rmsnorm.py | 3 +- vllm/models/kimi_k3/amd/ops/kda_chunk.py | 2 +- vllm/models/kimi_k3/amd/ops/kda_decode.py | 2 +- .../ops/fused_mla_key_concat_kv_cache.py | 2 +- .../attention/backends/mla/flashmla_sparse.py | 2 +- vllm/v1/hisparse/runtime.py | 2 +- 242 files changed, 2792 insertions(+), 3647 deletions(-) rename csrc/{libtorch_stable => }/activation_kernels.cu (99%) rename csrc/{libtorch_stable => }/async_util.cuh (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_common.cuh (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_kv_gather.cu (100%) rename csrc/{libtorch_stable => }/attention/dcp_utils/dcp_direct_q_gather.cu (100%) rename csrc/{libtorch_stable => }/attention/merge_attn_states.cu (99%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp (100%) rename csrc/{libtorch_stable => }/attention/mla/sm100_cutlass_mla_kernel.cu (99%) delete mode 100644 csrc/cache.h rename csrc/{libtorch_stable => }/cache_kernels.cu (99%) rename csrc/{libtorch_stable => }/cache_kernels_fused.cu (98%) rename csrc/{libtorch_stable => }/concat_mla_q.cuh (100%) rename csrc/{libtorch_stable => }/cooperative_topk.cu (100%) rename csrc/{libtorch_stable => }/cooperative_topk.cuh (100%) delete mode 100644 csrc/core/batch_invariant.hpp rename csrc/{ => core}/cuda_compat.h (90%) delete mode 100644 csrc/core/exception.hpp rename csrc/{libtorch_stable/core/math.hpp => core/utils.hpp} (70%) rename csrc/{attention => core/vec_dtypes}/dtype_bfloat16.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_float16.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_float32.cuh (99%) rename csrc/{attention => core/vec_dtypes}/dtype_fp8.cuh (88%) rename csrc/{attention/attention_generic.cuh => core/vec_dtypes/dtype_generic.cuh} (100%) rename csrc/{attention/attention_dtypes.h => core/vec_dtypes/vec_dtypes.h} (79%) rename csrc/{moe => cpu}/dynamic_4bit_int_moe_cpu.cpp (100%) create mode 100644 csrc/cpu/ops.h rename csrc/{libtorch_stable => }/cub_helpers.h (100%) rename csrc/{libtorch_stable => }/cuda_utils_kernels.cu (100%) rename csrc/{libtorch_stable => }/cuda_vec_utils.cuh (100%) rename csrc/{libtorch_stable => }/cuda_view.cu (100%) rename csrc/{libtorch_stable => }/custom_all_gather_reduce_scatter.cu (100%) rename csrc/{libtorch_stable => }/custom_all_gather_reduce_scatter_ops.cpp (100%) rename csrc/{libtorch_stable => }/custom_all_reduce.cu (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/common.hpp (89%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp (93%) rename csrc/{libtorch_stable => }/cutlass_extensions/torch_utils.hpp (76%) rename csrc/{libtorch_stable => }/cutlass_extensions/vllm_collective_builder.cuh (100%) rename csrc/{libtorch_stable => }/cutlass_extensions/vllm_numeric_conversion.cuh (99%) rename csrc/{libtorch_stable => }/dsv3_fused_a_gemm.cu (99%) rename csrc/{libtorch_stable => }/fp32_router_gemm.cu (100%) rename csrc/{libtorch_stable => }/fp32_router_gemm_entry.cu (99%) rename csrc/{libtorch_stable => }/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu (99%) rename csrc/{libtorch_stable => }/fused_qknorm_rope_kernel.cu (99%) rename csrc/{libtorch_stable => }/gdn/fused_gdn_decode_kernel.cu (99%) rename csrc/{libtorch_stable => }/hisparse_kernels.cu (99%) rename csrc/{libtorch_stable => }/kimi_k3/attn_res_kernel.cu (100%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_chunk_kernel_rocm.cu (100%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_decode_kernel.cu (99%) rename csrc/{libtorch_stable => }/kimi_k3/fused_kda_decode_kernel_rocm.cu (100%) rename csrc/{libtorch_stable => }/launch_bounds_utils.h (100%) rename csrc/{libtorch_stable => }/layernorm_kernels.cu (99%) rename csrc/{libtorch_stable => }/layernorm_quant_kernels.cu (99%) delete mode 100644 csrc/libtorch_stable/attention/attention_utils.cuh delete mode 100644 csrc/libtorch_stable/cutlass_extensions/common.cpp delete mode 100644 csrc/libtorch_stable/dispatch_utils.h delete mode 100644 csrc/libtorch_stable/mamba/static_switch.h delete mode 100644 csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py delete mode 100644 csrc/libtorch_stable/ops.h delete mode 100644 csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu delete mode 100644 csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h delete mode 100644 csrc/libtorch_stable/torch_bindings.cpp delete mode 100644 csrc/libtorch_stable/torch_utils.h rename csrc/{libtorch_stable => }/mamba/selective_scan.h (100%) rename csrc/{libtorch_stable => }/mamba/selective_scan_fwd.cu (99%) rename csrc/{libtorch_stable => }/minimax_reduce_rms_kernel.cu (99%) rename csrc/{libtorch_stable => }/minimax_reduce_rms_kernel.h (100%) rename csrc/{libtorch_stable => }/moe/grouped_topk_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/.gitignore (73%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/kernel.h (95%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/marlin_template.h (99%) rename csrc/{libtorch_stable => }/moe/marlin_moe_wna16/ops.cu (99%) rename csrc/{libtorch_stable => }/moe/moeTopKFuncs.cuh (100%) rename csrc/{libtorch_stable => }/moe/moe_align_sum_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_ops.h (100%) rename csrc/{libtorch_stable => }/moe/moe_permute_unpermute_op.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_wna16.cu (99%) rename csrc/{libtorch_stable => }/moe/moe_wna16_utils.h (100%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/dispatch.h (100%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu (98%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h (97%) rename csrc/{libtorch_stable => }/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl (100%) rename csrc/{libtorch_stable => }/moe/topk_softmax_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/topk_softplus_sqrt_kernels.cu (99%) rename csrc/{libtorch_stable => }/moe/torch_bindings.cpp (100%) rename csrc/{libtorch_stable => }/ngram_embedding_kernels.cu (100%) rename csrc/{libtorch_stable => }/nvfp4_ds_mla_cache.h (100%) rename csrc/{libtorch_stable => }/nvfp4_ds_mla_cache_kernels.cu (100%) rename csrc/{libtorch_stable => }/nvfp4_kv_cache_kernels.cu (99%) rename csrc/{libtorch_stable => }/persistent_topk.cuh (100%) rename csrc/{libtorch_stable => }/pos_encoding_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/activation_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/awq/dequantize.cuh (100%) rename csrc/{libtorch_stable => }/quantization/awq/gemm_kernels.cu (99%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/get_group_starts.cuh (99%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_mm_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_utils.cu (100%) rename csrc/{libtorch_stable => }/quantization/cutlass_w4a8/w4a8_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu (97%) rename csrc/{libtorch_stable => }/quantization/fp4/mxfp4_blockwise_moe_kernel.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/mxfp4_experts_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_blockwise_moe_kernel.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_experts_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_quant_entry.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_quant_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_entry.cu (97%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu (98%) rename csrc/{libtorch_stable => }/quantization/fp4/nvfp4_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu (99%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/fused_silu_mul_block_quant.cu (100%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/layernorm_utils.cuh (99%) rename csrc/{libtorch_stable => }/quantization/fused_kernels/quant_conversions.cuh (96%) rename csrc/{libtorch_stable => }/quantization/gptq/compat.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/matrix_view.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/q_gemm.cu (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_2.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_3.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_4.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_8.cuh (100%) rename csrc/{libtorch_stable => }/quantization/gptq/qdq_util.cuh (100%) rename csrc/{libtorch_stable => }/quantization/hadamard/hadacore/hadamard_transform_cuda.cu (99%) rename csrc/{libtorch_stable => }/quantization/machete/Readme.md (100%) rename csrc/{libtorch_stable => }/quantization/machete/generate.py (99%) rename csrc/{libtorch_stable => }/quantization/machete/machete_collective_builder.cuh (94%) rename csrc/{libtorch_stable => }/quantization/machete/machete_interleaving_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mainloop.cuh (99%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mm_kernel.cuh (98%) rename csrc/{libtorch_stable => }/quantization/machete/machete_mm_launcher.cuh (96%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepack_kernel.cuh (98%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepack_launcher.cuh (96%) rename csrc/{libtorch_stable => }/quantization/machete/machete_prepacked_layout.cuh (100%) rename csrc/{libtorch_stable => }/quantization/machete/machete_pytorch.cu (100%) rename csrc/{libtorch_stable => }/quantization/marlin/.gitignore (73%) rename csrc/{libtorch_stable => }/quantization/marlin/awq_marlin_repack.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/dequant.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/generate_kernels.py (96%) rename csrc/{libtorch_stable => }/quantization/marlin/gptq_marlin_repack.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/kernel.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin.cuh (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_dtypes.cuh (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_int4_fp8_preprocess.cu (99%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_mma.h (100%) rename csrc/{libtorch_stable => }/quantization/marlin/marlin_template.h (100%) rename csrc/{libtorch_stable => }/quantization/vectorization.cuh (100%) rename csrc/{libtorch_stable => }/quantization/vectorization_utils.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh (96%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm.cuh (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu (90%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu (89%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu (96%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu (89%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp (82%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu (91%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu (91%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/get_group_starts.cuh (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/moe/moe_data.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x.cu (98%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x.cuh (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu (100%) rename csrc/{libtorch_stable => }/quantization/w8a8/cutlass/scaled_mm_entry.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/fp8/common.cu (99%) rename csrc/{libtorch_stable => }/quantization/w8a8/fp8/per_token_group_quant.cu (97%) rename csrc/{libtorch_stable => }/quantization/w8a8/int8/scaled_quant.cu (100%) create mode 100644 csrc/rocm/dispatch_utils.h rename csrc/{ => rocm}/quickreduce/base.h (100%) rename csrc/{ => rocm/quickreduce}/custom_quickreduce.cu (99%) rename csrc/{ => rocm}/quickreduce/quick_reduce.h (100%) rename csrc/{ => rocm}/quickreduce/quick_reduce_impl.cuh (100%) create mode 100644 csrc/rocm/quickreduce/torch_bindings.cpp rename csrc/{libtorch_stable => }/sampled_topk.cuh (100%) rename csrc/{libtorch_stable => }/sampler.cu (99%) rename csrc/{libtorch_stable => }/topk.cu (100%) rename csrc/{libtorch_stable => }/topk_histogram_4096.cuh (100%) rename csrc/{libtorch_stable => }/type_convert.cuh (100%) rename csrc/{libtorch_stable => }/vocab_parallel_embedding_kernels.cu (100%) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 49bc3c32b67a..e4ae7f897e35 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -829,9 +829,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py @@ -1514,7 +1512,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -1560,6 +1558,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -1585,10 +1584,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -1610,9 +1609,9 @@ steps: - vllm/models/kimi_k3/nvidia/ops/third_party/kda/ - vllm/models/kimi_k3/amd/ - CMakeLists.txt - - csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/models/kimi_k3/test_amd_attn_res.py - tests/models/kimi_k3/test_amd_kda_decode.py - tests/models/kimi_k3/test_kda.py @@ -1633,7 +1632,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - vllm/platforms/rocm.py commands: @@ -3512,9 +3511,7 @@ steps: optional: true working_dir: "/vllm-workspace/" source_file_dependencies: - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/distributed/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py @@ -4033,7 +4030,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -4604,10 +4601,10 @@ steps: - vllm/platforms/rocm.py - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py optional: true commands: @@ -4622,7 +4619,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py @@ -4674,7 +4671,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -4690,6 +4687,7 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -4709,7 +4707,7 @@ steps: parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - csrc/rocm/ - tests/kernels/moe @@ -4759,7 +4757,7 @@ steps: - vllm/utils/torch_utils.py - vllm/platforms/ - vllm/_aiter_ops.py - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/distributed/device_communicators/ - vllm/config commands: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index f83946386a38..0b1a73136b04 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -510,9 +510,7 @@ steps: - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py - - csrc/custom_quickreduce.cu - - csrc/ops.h - - csrc/torch_bindings.cpp + - csrc/rocm/quickreduce/ - vllm/model_executor/layers/ - vllm/entrypoints/llm.py - vllm/config/parallel.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index aa0971b4c55e..472c19e74ae1 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -72,10 +72,10 @@ steps: num_devices: 2 device: h100 source_file_dependencies: - - csrc/libtorch_stable/minimax_reduce_rms_kernel.cu - - csrc/libtorch_stable/minimax_reduce_rms_kernel.h - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/model_executor/layers/minimax_rms_norm/ - vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py - tests/kernels/core/test_minimax_reduce_rms.py @@ -122,7 +122,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu + - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - vllm/platforms/rocm.py @@ -316,6 +316,7 @@ steps: timeout_in_minutes: 30 source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267) - vllm/model_executor/layers/attention @@ -335,6 +336,7 @@ steps: - image-build-amd source_file_dependencies: - csrc/attention/ + - csrc/core/vec_dtypes/ - vllm/v1/attention - vllm/model_executor/layers/attention - tests/kernels/attention @@ -467,7 +469,7 @@ steps: num_devices: 1 timeout_in_minutes: 50 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -485,7 +487,7 @@ steps: device: mi300_1 timeout_in_minutes: 55 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -576,11 +578,10 @@ steps: # optional: true source_file_dependencies: - csrc/quantization/fp4/ - - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu - - csrc/libtorch_stable/quantization/fp4/ - - csrc/libtorch_stable/quantization/w8a8/cutlass/ + - csrc/nvfp4_kv_cache_kernels.cu + - csrc/quantization/w8a8/cutlass/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fusion/ - vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py - vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -606,9 +607,9 @@ steps: - vllm/config/attention.py - vllm/v1/attention/backends/registry.py - vllm/_custom_ops.py - - csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu - - csrc/libtorch_stable/ops.h - - csrc/libtorch_stable/torch_bindings.cpp + - csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - tests/kernels/attention/test_flashinfer_mla_decode.py - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py - tests/kernels/attention/test_cache.py @@ -853,7 +854,7 @@ steps: device: h100 num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -896,7 +897,7 @@ steps: device: b200-k8s num_devices: 2 source_file_dependencies: - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ @@ -1009,7 +1010,7 @@ steps: source_file_dependencies: - csrc/quantization/fp4/ - csrc/attention/mla/ - - csrc/quantization/cutlass_w8a8/moe/ + - csrc/quantization/w8a8/cutlass/moe/ - vllm/model_executor/layers/fused_moe/cutlass_moe.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 06607951057b..f6a4a89f11cd 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -230,7 +230,7 @@ steps: - vllm/v1/attention/backends/mla/ - vllm/model_executor/layers/attention/sparse_mla_attention.py - vllm/v1/worker/gpu/ - - csrc/libtorch_stable/hisparse_kernels.cu + - csrc/hisparse_kernels.cu commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt @@ -1185,7 +1185,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1203,7 +1203,7 @@ steps: optional: true num_devices: 4 source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ @@ -1224,7 +1224,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: - - csrc/libtorch_stable/gdn/ + - csrc/gdn/ - tests/evals/qwen4_exp/ - vllm/model_executor/layers/mamba/ - vllm/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 288e7621cab4..552910e85a96 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -132,7 +132,7 @@ steps: device: b200-k8s source_file_dependencies: - vllm/models/kimi_k3/ - - csrc/libtorch_stable/kimi_k3/ + - csrc/kimi_k3/ - tests/models/kimi_k3/ - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py commands: diff --git a/.gitignore b/.gitignore index 46418c03924e..57d0d54fc87b 100644 --- a/.gitignore +++ b/.gitignore @@ -237,9 +237,6 @@ benchmarks/**/*.json actionlint shellcheck*/ -# Ignore moe/marlin_moe gen code -csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_* - # Ignore ep_kernels_workspace folder ep_kernels_workspace/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8e257edccc12..78a672c1104b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: rev: v21.1.2 hooks: - id: clang-format - exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*' + exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 diff --git a/CMakeLists.txt b/CMakeLists.txt index b4dd9ada5eb6..3a71d7e1e5ad 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -363,8 +363,8 @@ endif() if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_EXT_SRC - "csrc/torch_bindings.cpp" - "csrc/custom_quickreduce.cu") + "csrc/rocm/quickreduce/torch_bindings.cpp" + "csrc/rocm/quickreduce/custom_quickreduce.cu") message(STATUS "Enabling C extension.") define_extension_target( @@ -374,16 +374,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") SOURCES ${VLLM_EXT_SRC} COMPILE_FLAGS ${VLLM_GPU_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} - INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR} - INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR} USE_SABI 3 WITH_SOABI) - - # If CUTLASS is compiled on NVCC >= 12.5, it by default uses - # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the - # driver API. This causes problems when linking with earlier versions of CUDA. - # Setting this variable sidesteps the issue by calling the driver directly. - target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") @@ -419,36 +411,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/torch_bindings.cpp" - "csrc/libtorch_stable/cuda_view.cu" - "csrc/libtorch_stable/cuda_utils_kernels.cu" - "csrc/libtorch_stable/activation_kernels.cu" - "csrc/libtorch_stable/ngram_embedding_kernels.cu" - "csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu" - "csrc/libtorch_stable/quantization/activation_kernels.cu" - "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" - "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" - "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" - "csrc/libtorch_stable/pos_encoding_kernels.cu" - "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/layernorm_kernels.cu" - "csrc/libtorch_stable/layernorm_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu" - "csrc/libtorch_stable/attention/merge_attn_states.cu" - "csrc/libtorch_stable/sampler.cu" - "csrc/libtorch_stable/topk.cu" - "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" - "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels_fused.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" - "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" - "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" - "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") + "csrc/torch_bindings.cpp" + "csrc/cuda_view.cu" + "csrc/cuda_utils_kernels.cu" + "csrc/activation_kernels.cu" + "csrc/ngram_embedding_kernels.cu" + "csrc/vocab_parallel_embedding_kernels.cu" + "csrc/quantization/activation_kernels.cu" + "csrc/quantization/w8a8/int8/scaled_quant.cu" + "csrc/quantization/w8a8/fp8/common.cu" + "csrc/quantization/w8a8/fp8/per_token_group_quant.cu" + "csrc/quantization/gptq/q_gemm.cu" + "csrc/pos_encoding_kernels.cu" + "csrc/fused_qknorm_rope_kernel.cu" + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu" + "csrc/layernorm_kernels.cu" + "csrc/layernorm_quant_kernels.cu" + "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" + "csrc/attention/merge_attn_states.cu" + "csrc/sampler.cu" + "csrc/topk.cu" + "csrc/mamba/selective_scan_fwd.cu" + "csrc/cache_kernels.cu" + "csrc/cache_kernels_fused.cu" + "csrc/custom_all_gather_reduce_scatter.cu" + "csrc/custom_all_gather_reduce_scatter_ops.cpp" + "csrc/custom_all_reduce.cu" + "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -470,10 +461,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/hisparse_kernels.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" - "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") + "csrc/hisparse_kernels.cu" + "csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" + "csrc/attention/dcp_utils/dcp_direct_kv_gather.cu" + "csrc/attention/dcp_utils/dcp_direct_q_gather.cu") SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") @@ -508,12 +499,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") FetchContent_MakeAvailable(cutlass) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cutlass_extensions/common.cpp" - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu" - "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" - "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu" + "csrc/quantization/fp4/nvfp4_quant_entry.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu" + "csrc/quantization/awq/gemm_kernels.cu" + "csrc/minimax_reduce_rms_kernel.cu") # # Machete kernels @@ -527,7 +517,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MACHETE_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/machete/generate.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py) file(MD5 ${MACHETE_GEN_SCRIPT} MACHETE_GEN_SCRIPT_HASH) message(STATUS "Machete generation script hash: ${MACHETE_GEN_SCRIPT_HASH}") @@ -560,7 +550,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() # Add machete generated sources - file(GLOB MACHETE_GEN_SOURCES "csrc/libtorch_stable/quantization/machete/generated/*.cu") + file(GLOB MACHETE_GEN_SOURCES "csrc/quantization/machete/generated/*.cu") list(APPEND VLLM_STABLE_EXT_SRC ${MACHETE_GEN_SOURCES}) # forward compatible @@ -569,7 +559,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") CUDA_ARCHS "${MACHETE_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC - csrc/libtorch_stable/quantization/machete/machete_pytorch.cu) + csrc/quantization/machete/machete_pytorch.cu) message(STATUS "Building Machete kernels for archs: ${MACHETE_ARCHS}") else() if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 @@ -587,7 +577,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # The fused MiniMax-M3 KV insert quantizes NVFP4 with SM100-family-only # cvt instructions, so its SM10x/11x code is built for the family targets. set(MINIMAX_M3_FUSED_SRC - "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") + "csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu") list(REMOVE_ITEM VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}") set_gencode_flags_for_srcs( SRCS "${VLLM_STABLE_EXT_SRC}" @@ -611,9 +601,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(COOPERATIVE_TOPK_ARCHS) list(APPEND VLLM_STABLE_EXT_SRC - "csrc/libtorch_stable/cooperative_topk.cu") + "csrc/cooperative_topk.cu") set_gencode_flags_for_srcs( - SRCS "csrc/libtorch_stable/cooperative_topk.cu" + SRCS "csrc/cooperative_topk.cu" CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}") endif() @@ -656,7 +646,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # preselected input type pairs and schedules. # Generate sources: set(MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -691,7 +681,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_ARCHS) - file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_ARCHS}") @@ -701,7 +691,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) - file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_BF16_ARCHS}") @@ -713,7 +703,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_SM75_ARCHS) - file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_SM75_ARCHS}") @@ -725,7 +715,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if (MARLIN_FP8_ARCHS) - file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" CUDA_ARCHS "${MARLIN_FP8_ARCHS}") @@ -737,10 +727,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() set(MARLIN_SRCS - "csrc/libtorch_stable/quantization/marlin/marlin.cu" - "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" - "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" - "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + "csrc/quantization/marlin/marlin.cu" + "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/quantization/marlin/gptq_marlin_repack.cu" + "csrc/quantization/marlin/awq_marlin_repack.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_SRCS}" CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") @@ -763,7 +753,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS) - set(DSV3_FUSED_A_GEMM_SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu") + set(DSV3_FUSED_A_GEMM_SRCS "csrc/dsv3_fused_a_gemm.cu") set_gencode_flags_for_srcs( SRCS "${DSV3_FUSED_A_GEMM_SRCS}" CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}") @@ -778,8 +768,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) set(FP32_ROUTER_GEMM_SRCS - "csrc/libtorch_stable/fp32_router_gemm_entry.cu" - "csrc/libtorch_stable/fp32_router_gemm.cu") + "csrc/fp32_router_gemm_entry.cu" + "csrc/fp32_router_gemm.cu") set_gencode_flags_for_srcs( SRCS "${FP32_ROUTER_GEMM_SRCS}" CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") @@ -799,11 +789,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS) set(SCALED_MM_SM90_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -834,9 +824,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM120_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM120_SRCS}" @@ -868,9 +858,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(SCALED_MM_SM100_SRCS - "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" - "csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu" + "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu" ) set_gencode_flags_for_srcs( SRCS "${SCALED_MM_SM100_SRCS}" @@ -901,7 +891,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # subtract out the archs that are already built for 3x list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS}) if (SCALED_MM_2X_ARCHS) - set(SCALED_MM_C2X_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu") + set(SCALED_MM_C2X_SRCS "csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu") set_gencode_flags_for_srcs( SRCS "${SCALED_MM_C2X_SRCS}" CUDA_ARCHS "${SCALED_MM_2X_ARCHS}") @@ -927,7 +917,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # if it's possible to compile MoE kernels that use its output. cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM90_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") + set(CUTLASS_MOE_SM90_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM90_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -951,7 +941,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) - set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") + set(CUTLASS_MOE_SM100_SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_SM100_SRCS}" CUDA_ARCHS "${SCALED_MM_ARCHS}") @@ -976,7 +966,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS) - set(CUTLASS_MOE_DATA_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu") + set(CUTLASS_MOE_DATA_SRCS "csrc/quantization/w8a8/cutlass/moe/moe_data.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MOE_DATA_SRCS}" CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}") @@ -1008,12 +998,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS) set(FP4_SM120_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${FP4_SM120_SRCS}" CUDA_ARCHS "${FP4_SM120_ARCHS}") @@ -1034,15 +1024,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS) set(FP4_SM100_SRCS - "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu" - "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" - "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" - "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu" - "csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu") + "csrc/quantization/fp4/nvfp4_quant_kernels.cu" + "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu" + "csrc/quantization/fp4/nvfp4_experts_quant.cu" + "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu" + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" + "csrc/nvfp4_kv_cache_kernels.cu" + "csrc/nvfp4_ds_mla_cache_kernels.cu") if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) message(STATUS "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") @@ -1066,9 +1056,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS) set(W4A8_SRCS - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" - "csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu" + "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu" + "csrc/quantization/cutlass_w4a8/w4a8_utils.cu" ) set_gencode_flags_for_srcs( @@ -1101,7 +1091,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS) set(CUTLASS_MLA_SRCS - "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu") + "csrc/attention/mla/sm100_cutlass_mla_kernel.cu") set_gencode_flags_for_srcs( SRCS "${CUTLASS_MLA_SRCS}" CUDA_ARCHS "${MLA_ARCHS}") @@ -1125,7 +1115,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_KDA_DECODE_ARCHS) set(FUSED_KDA_DECODE_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + "csrc/kimi_k3/fused_kda_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_KDA_DECODE_SRC}" CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") @@ -1137,7 +1127,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(FUSED_GDN_DECODE_ARCHS) set(FUSED_GDN_DECODE_SRC - "csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu") + "csrc/gdn/fused_gdn_decode_kernel.cu") set_gencode_flags_for_srcs( SRCS "${FUSED_GDN_DECODE_SRC}" CUDA_ARCHS "${FUSED_GDN_DECODE_ARCHS}") @@ -1154,7 +1144,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(KIMI_K3_ATTN_RES_ARCHS) set(KIMI_K3_ATTN_RES_SRC - "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + "csrc/kimi_k3/attn_res_kernel.cu") set_gencode_flags_for_srcs( SRCS "${KIMI_K3_ATTN_RES_SRC}" CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") @@ -1169,7 +1159,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) - set(HADACORE_SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") + set(HADACORE_SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu") set_gencode_flags_for_srcs( SRCS "${HADACORE_SRCS}" CUDA_ARCHS "${HADACORE_ARCHS}") @@ -1193,7 +1183,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_DECODE_HIP_ARCHS INCLUDE REGEX "gfx942|gfx950") if(FUSED_KDA_DECODE_HIP_ARCHS) set(FUSED_KDA_DECODE_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_HIP_SRC}") message(STATUS "Building fused KDA decode for HIP archs: ${FUSED_KDA_DECODE_HIP_ARCHS}") @@ -1203,7 +1193,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") list(FILTER FUSED_KDA_CHUNK_HIP_ARCHS INCLUDE REGEX "gfx950") if(FUSED_KDA_CHUNK_HIP_ARCHS) set(FUSED_KDA_CHUNK_HIP_SRC - "csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu") + "csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu") list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_CHUNK_HIP_SRC}") message(STATUS "Building fused KDA chunk for HIP archs: ${FUSED_KDA_CHUNK_HIP_ARCHS}") @@ -1289,21 +1279,21 @@ endif() # set(VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/torch_bindings.cpp" - "csrc/libtorch_stable/moe/moe_align_sum_kernels.cu" - "csrc/libtorch_stable/moe/topk_softmax_kernels.cu" - "csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu") + "csrc/moe/torch_bindings.cpp" + "csrc/moe/moe_align_sum_kernels.cu" + "csrc/moe/topk_softmax_kernels.cu" + "csrc/moe/topk_softplus_sqrt_kernels.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu" - "csrc/libtorch_stable/moe/grouped_topk_kernels.cu") + "csrc/moe/moe_wna16.cu" + "csrc/moe/grouped_topk_kernels.cu") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") set(MOE_PERMUTE_SRC - "csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" - "csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu") + "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu" + "csrc/moe/moe_permute_unpermute_op.cu") list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}") endif() @@ -1314,7 +1304,7 @@ set_gencode_flags_for_srcs( if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_MOE_WNA16_SRC - "csrc/libtorch_stable/moe/moe_wna16.cu") + "csrc/moe/moe_wna16.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_MOE_WNA16_SRC}" @@ -1349,7 +1339,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # preselected input type pairs and schedules. # Generate sources: set(MOE_MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py) + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH) list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") @@ -1363,6 +1353,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=$ENV{PYTHONPATH} ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16 RESULT_VARIABLE moe_marlin_generation_result OUTPUT_VARIABLE moe_marlin_generation_output OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log @@ -1384,7 +1375,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_ARCHS) - file(GLOB MARLIN_MOE_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_*.cu") + file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SRC}" CUDA_ARCHS "${MARLIN_MOE_ARCHS}") @@ -1396,7 +1387,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_SM75_ARCHS) - file(GLOB MARLIN_MOE_SM75_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm75_kernel_*.cu") + file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_SM75_SRC}" CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}") @@ -1408,7 +1399,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (MARLIN_MOE_FP8_ARCHS) - file(GLOB MARLIN_MOE_FP8_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/sm89_kernel_*.cu") + file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_FP8_SRC}" CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}") @@ -1419,7 +1410,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC}) endif() - set(MARLIN_MOE_OTHER_SRC "csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu") + set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu") set_gencode_flags_for_srcs( SRCS "${MARLIN_MOE_OTHER_SRC}" CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}") diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 730b963c2362..c4d111e77bfe 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -497,7 +497,7 @@ set(VLLM_EXT_SRC "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" "csrc/cpu/mamba_cpu.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_isa.cpp" @@ -599,7 +599,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") set(VLLM_EXT_SRC_AVX2 "csrc/cpu/sgl-kernels/fla.cpp" @@ -617,7 +617,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" - "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") + "csrc/cpu/dynamic_4bit_int_moe_cpu.cpp") message(STATUS "CPU extension (AVX512F + BF16 + VNNI + AMX) source files: ${VLLM_EXT_SRC_AVX512} ${VLLM_EXT_SRC_SGL}") message(STATUS "CPU extension (AVX512F) source files: ${VLLM_EXT_SRC_AVX512}") diff --git a/cmake/hipify.py b/cmake/hipify.py index f4932260c9eb..2cf58a7b7d33 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -65,7 +65,7 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: project_directory=args.project_dir, output_directory=args.output_dir, # Hipify resolves quoted includes next to the including file first; vLLM - # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h" + # uses paths relative to csrc/ (e.g. "torch_utils.h" # from quantization/w8a8/fp8/*.cu). Without an include root here, those # headers are never found and are not hipified or rewritten in dependents. header_include_dirs=["."], diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/activation_kernels.cu rename to csrc/activation_kernels.cu index a7ee87b35f4e..9e28c3343bd7 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -10,7 +10,7 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "async_util.cuh" #include "cuda_vec_utils.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/async_util.cuh b/csrc/async_util.cuh similarity index 100% rename from csrc/libtorch_stable/async_util.cuh rename to csrc/async_util.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu b/csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu rename to csrc/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh b/csrc/attention/dcp_utils/dcp_direct_common.cuh similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh rename to csrc/attention/dcp_utils/dcp_direct_common.cuh diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/attention/dcp_utils/dcp_direct_kv_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_kv_gather.cu diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu b/csrc/attention/dcp_utils/dcp_direct_q_gather.cu similarity index 100% rename from csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu rename to csrc/attention/dcp_utils/dcp_direct_q_gather.cu diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/attention/merge_attn_states.cu similarity index 99% rename from csrc/libtorch_stable/attention/merge_attn_states.cu rename to csrc/attention/merge_attn_states.cu index cc89397f68a1..3be334fd64fc 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/attention/merge_attn_states.cu @@ -6,9 +6,8 @@ #include "../dispatch_utils.h" #include -#include "../../attention/attention_dtypes.h" -#include "attention_utils.cuh" -#include "../../quantization/w8a8/fp8/common.cuh" +#include "core/vec_dtypes/vec_dtypes.h" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp b/csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp rename to csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_reduction.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_fmha_mla_tma_warpspecialized.hpp diff --git a/csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp b/csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp similarity index 100% rename from csrc/libtorch_stable/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp rename to csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp diff --git a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu similarity index 99% rename from csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu rename to csrc/attention/mla/sm100_cutlass_mla_kernel.cu index 150e32462812..b04c014b9805 100644 --- a/csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu +++ b/csrc/attention/mla/sm100_cutlass_mla_kernel.cu @@ -18,7 +18,7 @@ limitations under the License. * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929 * by Alcanderian JieXin Liang */ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include diff --git a/csrc/cache.h b/csrc/cache.h deleted file mode 100644 index dbc2675fc70c..000000000000 --- a/csrc/cache.h +++ /dev/null @@ -1,93 +0,0 @@ -#pragma once - -#include -#include - -#include -#include - -void swap_blocks(torch::Tensor& src, torch::Tensor& dst, - int64_t block_size_in_bytes, - const torch::Tensor& block_mapping); - -void swap_blocks_batch(const torch::Tensor& src_ptrs, - const torch::Tensor& dst_ptrs, - const torch::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, - torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& k_scale, torch::Tensor& v_scale); - -void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe, - torch::Tensor& kv_cache, torch::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::Tensor& scale); - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe, - torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache, - const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache, - const double scale, const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::Tensor const& dst, // [TOT_TOKENS, 576] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, std::optional seq_starts = std::nullopt); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::Tensor& k, // [num_tokens, head_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::Tensor& q_out); // [num_tokens, num_heads, nope_dim + rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& dst_k, // [num_tokens, head_dim] - torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::Tensor& block_table, // [batch_size, num_blocks] - const torch::Tensor& cu_seq_lens); // [batch_size + 1] diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/cache_kernels.cu rename to csrc/cache_kernels.cu index 9130adf87d69..d1c1e9df9896 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -1,8 +1,8 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_utils.h" -#include "../cuda_compat.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" #include "quantization/vectorization_utils.cuh" #include "concat_mla_q.cuh" @@ -12,9 +12,9 @@ #endif #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #include @@ -712,8 +712,7 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( } __syncthreads(); - for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); - iter++) { + for (int iter = 0; iter < div_ceil(batch_size, int(blockDim.x)); iter++) { int tid = iter * blockDim.x + threadIdx.x; if (tid < batch_size) { const int seq_start = cu_seq_lens[tid]; @@ -1229,7 +1228,7 @@ __device__ __forceinline__ bool map_gather_page_task( const int32_t source_begin = seq_starts == nullptr ? 0 : seq_starts[req_id]; const int32_t first_block = source_begin / block_size; const int32_t num_pages = - cuda_utils::ceil_div(source_begin + seq_len, block_size) - first_block; + div_ceil(source_begin + seq_len, block_size) - first_block; if (relative_page < num_pages) { page.req_id = req_id; page.logical_block = first_block + relative_page; @@ -1408,8 +1407,7 @@ void gather_and_maybe_dequant_cache( const int64_t dst_entry_stride = dst.stride(0); const int32_t page_threads = num_tokens >= (1 << 20) ? 128 : 256; const int32_t required_blocks = - cuda_utils::ceil_div(static_cast(num_tokens), block_size) + - 2 * num_reqs; + div_ceil(static_cast(num_tokens), block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(page_threads); @@ -1685,7 +1683,7 @@ void cp_gather_cache( const int32_t num_reqs = static_cast(batch_size); const int32_t required_blocks = - cuda_utils::ceil_div(total_tokens, block_size) + 2 * num_reqs; + div_ceil(total_tokens, block_size) + 2 * num_reqs; const dim3 grid(required_blocks); const dim3 block(256); @@ -1828,8 +1826,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( const int32_t* seq_starts_ptr = seq_starts.has_value() ? seq_starts.value().const_data_ptr() : nullptr; - const int required_blocks = cuda_utils::ceil_div(total_tokens, block_size) + - 2 * static_cast(batch_size); + const int required_blocks = + div_ceil(total_tokens, block_size) + 2 * static_cast(batch_size); vllm::cp_gather_and_upconvert_fp8_kv_cache_page<<< required_blocks, block_size_threads, 0, stream>>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), diff --git a/csrc/libtorch_stable/cache_kernels_fused.cu b/csrc/cache_kernels_fused.cu similarity index 98% rename from csrc/libtorch_stable/cache_kernels_fused.cu rename to csrc/cache_kernels_fused.cu index 744e3a2dcdcd..4dc691636d55 100644 --- a/csrc/libtorch_stable/cache_kernels_fused.cu +++ b/csrc/cache_kernels_fused.cu @@ -1,13 +1,13 @@ #include "torch_utils.h" #include "dispatch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" -#include "../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #ifdef USE_ROCM diff --git a/csrc/libtorch_stable/concat_mla_q.cuh b/csrc/concat_mla_q.cuh similarity index 100% rename from csrc/libtorch_stable/concat_mla_q.cuh rename to csrc/concat_mla_q.cuh diff --git a/csrc/libtorch_stable/cooperative_topk.cu b/csrc/cooperative_topk.cu similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cu rename to csrc/cooperative_topk.cu diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/cooperative_topk.cuh similarity index 100% rename from csrc/libtorch_stable/cooperative_topk.cuh rename to csrc/cooperative_topk.cuh diff --git a/csrc/core/batch_invariant.hpp b/csrc/core/batch_invariant.hpp deleted file mode 100644 index 8273bc74b1ef..000000000000 --- a/csrc/core/batch_invariant.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once -#include -#include - -namespace vllm { - -// vllm_is_batch_invariant(); returns true -// if env VLLM_BATCH_INVARIANT=1 -inline bool vllm_is_batch_invariant() { - static bool cached = []() { - std::string env_key = "VLLM_BATCH_INVARIANT"; - const char* val = std::getenv(env_key.c_str()); - return (val && std::atoi(val) != 0) ? 1 : 0; - }(); - return cached; -} - -} // namespace vllm diff --git a/csrc/cuda_compat.h b/csrc/core/cuda_compat.h similarity index 90% rename from csrc/cuda_compat.h rename to csrc/core/cuda_compat.h index d7d589db62cf..0eafeff6219b 100644 --- a/csrc/cuda_compat.h +++ b/csrc/core/cuda_compat.h @@ -54,12 +54,6 @@ struct Utils { __shfl_xor(var, lane_mask, width) #endif -#ifndef USE_ROCM - #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) -#else - #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) -#endif - #ifndef USE_ROCM #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ __shfl_down_sync(uint32_t(-1), var, lane_delta) diff --git a/csrc/core/exception.hpp b/csrc/core/exception.hpp deleted file mode 100644 index f3b2ffaef6cc..000000000000 --- a/csrc/core/exception.hpp +++ /dev/null @@ -1,3 +0,0 @@ -#pragma once - -#define VLLM_IMPLIES(p, q) (!(p) || (q)) diff --git a/csrc/core/registration.h b/csrc/core/registration.h index 4d0ce1c572c1..e5396e9a8b13 100644 --- a/csrc/core/registration.h +++ b/csrc/core/registration.h @@ -12,11 +12,6 @@ // could be a macro instead of a literal token. #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) -// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ - TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) - // REGISTER_EXTENSION allows the shared library to be loaded and initialized // via python's import statement. #define REGISTER_EXTENSION(NAME) \ diff --git a/csrc/libtorch_stable/core/math.hpp b/csrc/core/utils.hpp similarity index 70% rename from csrc/libtorch_stable/core/math.hpp rename to csrc/core/utils.hpp index 6764e1fd6054..36b8c2e38e9a 100644 --- a/csrc/libtorch_stable/core/math.hpp +++ b/csrc/core/utils.hpp @@ -1,7 +1,8 @@ #pragma once #include -#include +#include +#include inline constexpr uint32_t next_pow_2(uint32_t const num) { if (num <= 1) return num; @@ -26,3 +27,16 @@ template inline constexpr T round_to_next_multiple_of(T a, T b) { return a % b == 0 ? a : ((a / b) + 1) * b; } + +namespace vllm { + +// Returns true if env VLLM_BATCH_INVARIANT=1 +inline bool vllm_is_batch_invariant() { + static bool cached = []() { + const char* val = std::getenv("VLLM_BATCH_INVARIANT"); + return val && std::atoi(val) != 0; + }(); + return cached; +} + +} // namespace vllm diff --git a/csrc/attention/dtype_bfloat16.cuh b/csrc/core/vec_dtypes/dtype_bfloat16.cuh similarity index 99% rename from csrc/attention/dtype_bfloat16.cuh rename to csrc/core/vec_dtypes/dtype_bfloat16.cuh index 97a25baa1fc0..5d01a51e76fe 100644 --- a/csrc/attention/dtype_bfloat16.cuh +++ b/csrc/core/vec_dtypes/dtype_bfloat16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifndef USE_ROCM diff --git a/csrc/attention/dtype_float16.cuh b/csrc/core/vec_dtypes/dtype_float16.cuh similarity index 99% rename from csrc/attention/dtype_float16.cuh rename to csrc/core/vec_dtypes/dtype_float16.cuh index 3a1815f0ed4f..6883d521d2ac 100644 --- a/csrc/attention/dtype_float16.cuh +++ b/csrc/core/vec_dtypes/dtype_float16.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float32.cuh" #ifdef USE_ROCM diff --git a/csrc/attention/dtype_float32.cuh b/csrc/core/vec_dtypes/dtype_float32.cuh similarity index 99% rename from csrc/attention/dtype_float32.cuh rename to csrc/core/vec_dtypes/dtype_float32.cuh index 7c6a686db3ba..33f78c492e61 100644 --- a/csrc/attention/dtype_float32.cuh +++ b/csrc/core/vec_dtypes/dtype_float32.cuh @@ -20,7 +20,7 @@ */ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include diff --git a/csrc/attention/dtype_fp8.cuh b/csrc/core/vec_dtypes/dtype_fp8.cuh similarity index 88% rename from csrc/attention/dtype_fp8.cuh rename to csrc/core/vec_dtypes/dtype_fp8.cuh index 3d56859d8fc5..6d3c35e978b3 100644 --- a/csrc/attention/dtype_fp8.cuh +++ b/csrc/core/vec_dtypes/dtype_fp8.cuh @@ -1,7 +1,8 @@ #pragma once -#include "attention_generic.cuh" -#include "torch_utils.h" +#include "dtype_generic.cuh" + +#include #include #ifdef ENABLE_FP8 @@ -31,7 +32,7 @@ inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( } else if (dtype_str == "fp8_e5m2") { return Fp8KVCacheDataType::kFp8E5M2; } - TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); + STD_TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); } // fp8 vector types for quantization of kv cache diff --git a/csrc/attention/attention_generic.cuh b/csrc/core/vec_dtypes/dtype_generic.cuh similarity index 100% rename from csrc/attention/attention_generic.cuh rename to csrc/core/vec_dtypes/dtype_generic.cuh diff --git a/csrc/attention/attention_dtypes.h b/csrc/core/vec_dtypes/vec_dtypes.h similarity index 79% rename from csrc/attention/attention_dtypes.h rename to csrc/core/vec_dtypes/vec_dtypes.h index 64f86381d9db..41a7b91ca45a 100644 --- a/csrc/attention/attention_dtypes.h +++ b/csrc/core/vec_dtypes/vec_dtypes.h @@ -1,6 +1,6 @@ #pragma once -#include "attention_generic.cuh" +#include "dtype_generic.cuh" #include "dtype_float16.cuh" #include "dtype_float32.cuh" #include "dtype_bfloat16.cuh" diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 0666ee413f13..0290c7ab60b3 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -14,7 +14,7 @@ namespace cpu_attention { enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX, AMX_FP8 }; -// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. +// Mirrors csrc/core/vec_dtypes/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { kAuto = 0, kFp8E4M3 = 1, diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/cpu/dynamic_4bit_int_moe_cpu.cpp similarity index 100% rename from csrc/moe/dynamic_4bit_int_moe_cpu.cpp rename to csrc/cpu/dynamic_4bit_int_moe_cpu.cpp diff --git a/csrc/cpu/ops.h b/csrc/cpu/ops.h new file mode 100644 index 000000000000..ead20b01ae67 --- /dev/null +++ b/csrc/cpu/ops.h @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include +#include +#include + +void rms_norm(torch::Tensor& out, torch::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, + std::optional weight, double epsilon); + +void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, + std::optional key, int64_t head_size, + torch::Tensor& cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse); + +void silu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); + +void gelu_tanh(torch::Tensor& out, torch::Tensor& input); + +void gelu_new(torch::Tensor& out, torch::Tensor& input); + +void gelu_fast(torch::Tensor& out, torch::Tensor& input); + +void gelu_quick(torch::Tensor& out, torch::Tensor& input); + +void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, + torch::Tensor& scales, + std::optional const& azp); + +torch::Tensor dynamic_4bit_int_moe_cpu( + torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 2cb1be095d3a..f39cf31c1f17 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -1,4 +1,4 @@ -#include "cache.h" +#include #include "ops.h" #include "core/registration.h" diff --git a/csrc/libtorch_stable/cub_helpers.h b/csrc/cub_helpers.h similarity index 100% rename from csrc/libtorch_stable/cub_helpers.h rename to csrc/cub_helpers.h diff --git a/csrc/cuda_utils.h b/csrc/cuda_utils.h index 6e62ea208db8..bca7c75bc178 100644 --- a/csrc/cuda_utils.h +++ b/csrc/cuda_utils.h @@ -1,21 +1,8 @@ #pragma once +#include #include -#if defined(__HIPCC__) - #define HOST_DEVICE_INLINE __host__ __device__ - #define DEVICE_INLINE __device__ - #define HOST_INLINE __host__ -#elif defined(__CUDACC__) || defined(_NVHPC_CUDA) - #define HOST_DEVICE_INLINE __host__ __device__ __forceinline__ - #define DEVICE_INLINE __device__ __forceinline__ - #define HOST_INLINE __host__ __forceinline__ -#else - #define HOST_DEVICE_INLINE inline - #define DEVICE_INLINE inline - #define HOST_INLINE inline -#endif - #define CUDA_CHECK(cmd) \ do { \ cudaError_t e = cmd; \ @@ -29,13 +16,3 @@ int64_t get_device_attribute(int64_t attribute, int64_t device_id); int64_t get_max_shared_memory_per_block_device_attribute(int64_t device_id); - -namespace cuda_utils { - -template -HOST_DEVICE_INLINE constexpr std::enable_if_t, T> -ceil_div(T a, T b) { - return (a + b - 1) / b; -} - -}; // namespace cuda_utils \ No newline at end of file diff --git a/csrc/libtorch_stable/cuda_utils_kernels.cu b/csrc/cuda_utils_kernels.cu similarity index 100% rename from csrc/libtorch_stable/cuda_utils_kernels.cu rename to csrc/cuda_utils_kernels.cu diff --git a/csrc/libtorch_stable/cuda_vec_utils.cuh b/csrc/cuda_vec_utils.cuh similarity index 100% rename from csrc/libtorch_stable/cuda_vec_utils.cuh rename to csrc/cuda_vec_utils.cuh diff --git a/csrc/libtorch_stable/cuda_view.cu b/csrc/cuda_view.cu similarity index 100% rename from csrc/libtorch_stable/cuda_view.cu rename to csrc/cuda_view.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/custom_all_gather_reduce_scatter.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu rename to csrc/custom_all_gather_reduce_scatter.cu diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/custom_all_gather_reduce_scatter_ops.cpp similarity index 100% rename from csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp rename to csrc/custom_all_gather_reduce_scatter_ops.cpp diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/custom_all_reduce.cu similarity index 100% rename from csrc/libtorch_stable/custom_all_reduce.cu rename to csrc/custom_all_reduce.cu diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh index 353dcd07e8b4..8c6a98149d0f 100644 --- a/csrc/custom_collective_common.cuh +++ b/csrc/custom_collective_common.cuh @@ -216,21 +216,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - // This function is meant to be used as the second or the final // synchronization barrier in the all reduce kernel. If it's the final // synchronization barrier, we don't need to make any visibility guarantees @@ -279,22 +264,6 @@ DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; } -template -DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, - int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - template DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { __syncthreads(); diff --git a/csrc/libtorch_stable/cutlass_extensions/common.hpp b/csrc/cutlass_extensions/common.hpp similarity index 89% rename from csrc/libtorch_stable/cutlass_extensions/common.hpp rename to csrc/cutlass_extensions/common.hpp index ca96b0ef3fef..0a6fbd838411 100644 --- a/csrc/libtorch_stable/cutlass_extensions/common.hpp +++ b/csrc/cutlass_extensions/common.hpp @@ -25,7 +25,15 @@ inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) { return max_shared_mem_per_block_opt_in; } -int32_t get_sm_version_num(); +inline int32_t get_sm_version_num() { + int32_t major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + 0); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + 0); + int32_t version_num = major_capability * 10 + minor_capability; + return version_num; +} /** * A wrapper for a kernel that is used to guard against compilation on @@ -110,21 +118,6 @@ struct enable_sm100_to_sm120 : Kernel { } }; -template -struct enable_sm120_only : Kernel { - template - CUTLASS_DEVICE void operator()(Args&&... args) { -#if defined __CUDA_ARCH__ - #if __CUDA_ARCH__ == 1200 - Kernel::operator()(std::forward(args)...); - #else - printf("This kernel only supports sm120a.\n"); - asm("trap;"); - #endif -#endif - } -}; - // SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10) template struct enable_sm120_family : Kernel { diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp rename to csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp similarity index 93% rename from csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp rename to csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp index c2ddcea6d027..30f4a547f875 100644 --- a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp @@ -1,15 +1,9 @@ #pragma once -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp" -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp" +#include "broadcast_load_epilogue_c3x.hpp" +#include "broadcast_load_epilogue_array_c3x.hpp" -// This header is shared by both _C (unstable ABI) and _C_stable_libtorch -// (stable ABI) targets. When compiled under the stable ABI target, -// TORCH_TARGET_VERSION is defined and Tensor is unavailable, so we -// use torch::stable::Tensor instead. -#ifdef TORCH_TARGET_VERSION - #include -#endif +#include /* This file defines custom epilogues for fusing channel scales, token scales, @@ -23,38 +17,10 @@ namespace vllm::c3x { -#ifdef TORCH_TARGET_VERSION using TensorType = torch::stable::Tensor; -#else -using TensorType = torch::Tensor; -#endif using namespace cute; -template -struct identity { - CUTLASS_HOST_DEVICE - T operator()(T lhs) const { return lhs; } -}; - -template -struct TrivialEpilogue { - private: - using Accum = cutlass::epilogue::fusion::Sm90AccFetch; - using Compute = cutlass::epilogue::fusion::Sm90Compute< - cutlass::epilogue::thread::Identity, ElementD, ElementAcc, - cutlass::FloatRoundStyle::round_to_nearest>; - - public: - using EVTCompute = cutlass::epilogue::fusion::Sm90EVT; - using ArgumentType = typename EVTCompute::Arguments; - - template - static ArgumentType prepare_args(Args... args) { - return {}; - } -}; - /* * This class provides the common load descriptors for the * ScaledEpilogue[...] classes diff --git a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp b/csrc/cutlass_extensions/torch_utils.hpp similarity index 76% rename from csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp rename to csrc/cutlass_extensions/torch_utils.hpp index 5f973033fd72..c5b9d9591e6d 100644 --- a/csrc/libtorch_stable/cutlass_extensions/torch_utils.hpp +++ b/csrc/cutlass_extensions/torch_utils.hpp @@ -1,19 +1,12 @@ #pragma once -#include "torch_utils.h" - -// This header is shared between _C (unstable ABI, used by machete) and -// _C_stable_libtorch (stable ABI, used by W4A8/sparse). TORCH_TARGET_VERSION -// is defined only for the stable target, so we switch includes and types -// accordingly. TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. -#ifdef TORCH_TARGET_VERSION - #include - #include - #include +#include +#include +#include +#include + +// TorchTensor (not Tensor) avoids ambiguity with cute::Tensor. using TorchTensor = torch::stable::Tensor; -#else -using TorchTensor = torch::Tensor; -#endif #include "cute/layout.hpp" #include "cutlass/layout/matrix.h" @@ -70,15 +63,15 @@ CUTE_HOST_DEVICE constexpr auto make_shape_from_idx(F&& f) { template static inline auto make_cute_layout(TorchTensor const& tensor, std::string_view name = "tensor") { - TORCH_UTILS_CHECK(tensor.dim() <= rank(Stride{})); + STD_TORCH_CHECK(tensor.dim() <= rank(Stride{})); auto stride = cute::transform_with_idx(Stride{}, [&](auto const& stride_ele, auto const& idx) { using StrideEle = std::decay_t; if (idx < tensor.dim()) { if constexpr (cute::is_static_v) { - TORCH_UTILS_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", - name, ".stride(", idx, ") to be ", StrideEle::value); + STD_TORCH_CHECK(StrideEle::value == tensor.stride(idx), "Expected ", + name, ".stride(", idx, ") to be ", StrideEle::value); return StrideEle{}; } else { if (tensor.size(idx) == 1) { @@ -122,29 +115,7 @@ static inline auto maybe_make_cute_layout( } // -// Torch Type to Cutlass Type (equivalent_cutlass_type) -// - -template -struct equivalent_cutlass_type { - using type = T; -}; - -template -using equivalent_cutlass_type_t = typename equivalent_cutlass_type::type; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::half_t; -}; - -template <> -struct equivalent_cutlass_type { - using type = cutlass::bfloat16_t; -}; - -// -// equivalent_scalar_t (basically inverse of equivalent_cutlass_type) +// Cutlass Type to Torch Type (equivalent_scalar_type) // // Return a `torch::headeronly::CppTypeToScalarType` compatible type, i.e. diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh b/csrc/cutlass_extensions/vllm_collective_builder.cuh similarity index 100% rename from csrc/libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh rename to csrc/cutlass_extensions/vllm_collective_builder.cuh diff --git a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh similarity index 99% rename from csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh rename to csrc/cutlass_extensions/vllm_numeric_conversion.cuh index 90f226cf64c0..0f74053fd29f 100644 --- a/csrc/libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh +++ b/csrc/cutlass_extensions/vllm_numeric_conversion.cuh @@ -1,9 +1,9 @@ #pragma once #include "cutlass/numeric_conversion.h" -#include "cutlass_extensions/vllm_custom_types.cuh" -#include "cutlass_extensions/cute_utils.cuh" -#include "cutlass_extensions/vllm_type_utils.cuh" +#include "vllm_custom_types.cuh" +#include "cute_utils.cuh" +#include "vllm_type_utils.cuh" // this file extends: // https://github.com/NVIDIA/cutlass/blob/cutlass-3.5.0/include/cutlass/numeric_conversion.h diff --git a/csrc/cutlass_extensions/vllm_type_utils.cuh b/csrc/cutlass_extensions/vllm_type_utils.cuh index 500ed508c830..ba79dddf6f42 100644 --- a/csrc/cutlass_extensions/vllm_type_utils.cuh +++ b/csrc/cutlass_extensions/vllm_type_utils.cuh @@ -2,7 +2,7 @@ #include "cutlass/half.h" #include "cuda_bf16.h" -#include "cutlass_extensions/vllm_custom_types.cuh" +#include "vllm_custom_types.cuh" namespace cutlass { diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index de0c505b7a62..aa125a08fd6c 100644 --- a/csrc/dispatch_utils.h +++ b/csrc/dispatch_utils.h @@ -1,158 +1,188 @@ /* - * Adapted from - * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + * Stable ABI compatible dispatch utilities for vLLM. + * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros + * instead of the ATen (AT_*) macros. + * + * These macros use: + * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH + * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE + * - torch::headeronly::ScalarType instead of at::ScalarType + * + * Add more macros here as needed when migrating additional kernels. */ #pragma once -#include +#include +#include +#include // Need a special dispatch case macro since we will nest the FP8 dispatch. // Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ - AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) - -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) +// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. +#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ + __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ + __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) + +#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) + +// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format +#ifdef USE_ROCM + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) +#else + #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ + VLLM_STABLE_DISPATCH_FP8_CASE( \ + torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) +#endif -#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See VLLM_STABLE_DISPATCH_FP8_CASE above. +#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) -#define VLLM_DISPATCH_CASE_HALF_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +// Half types dispatch (Half + BFloat16) +#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) -#define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) +#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) -// ROCm devices might use either fn or fnuz, so set up dispatch table for both. -// A host-based check at runtime will create a preferred FP8 type for ROCm -// such that the correct kernel is dispatched. +// Quant type dispatch (FP8 + INT8) #ifdef USE_ROCM - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) #else - #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ - AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) - - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) + #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ + __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) #endif -// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See AT_DISPATCH_FP8_CASE above. -#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) - -#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) - -#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - AT_DISPATCH_SWITCH( \ - TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -#define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ +#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) + +// Group size dispatch (pure C++ if/else, no ATen dependency) +#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ + if (group_size == 128) { \ + constexpr int const_group_size = 128; \ + __VA_ARGS__(); \ + } else if (group_size == 64) { \ + constexpr int const_group_size = 64; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ +// Boolean dispatch +#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ + if (expr) { \ + constexpr bool const_expr = true; \ + __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + __VA_ARGS__(); \ } -#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ +// Vec size dispatch (pure C++ switch, no ATen dependency) +#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ + switch (VEC_SIZE) { \ + case 16: { \ + constexpr int vec_size = 16; \ + __VA_ARGS__(); \ + break; \ + } \ + case 8: { \ + constexpr int vec_size = 8; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int vec_size = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + case 2: { \ + constexpr int vec_size = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + default: { \ + constexpr int vec_size = 1; \ + __VA_ARGS__(); \ + break; \ + } \ } -#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \ +// Tensor rank dispatch (2D, 3D, 4D) +#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ + switch (NUM_DIMS) { \ + case 2: { \ + constexpr int tensor_rank = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + case 3: { \ + constexpr int tensor_rank = 3; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int tensor_rank = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + default: \ + STD_TORCH_CHECK( \ + false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/dsv3_fused_a_gemm.cu similarity index 99% rename from csrc/libtorch_stable/dsv3_fused_a_gemm.cu rename to csrc/dsv3_fused_a_gemm.cu index 4250683b9158..c0989f8420da 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/dsv3_fused_a_gemm.cu @@ -25,7 +25,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu similarity index 100% rename from csrc/libtorch_stable/fp32_router_gemm.cu rename to csrc/fp32_router_gemm.cu diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu similarity index 99% rename from csrc/libtorch_stable/fp32_router_gemm_entry.cu rename to csrc/fp32_router_gemm_entry.cu index fc09193c8643..e55ac7e72104 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -6,7 +6,7 @@ #include #include "core/registration.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu rename to csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 5a5caacce7a5..4eb8d7756b19 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -47,7 +47,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu rename to csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 3892790c4040..dd88645904d2 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -57,16 +57,16 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" #ifndef USE_ROCM #include - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else #include - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif #include #include @@ -835,7 +835,7 @@ static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token, // clang-format off // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the // launch closer as ">>>". clang-format would otherwise re-split it (it does - // not parse the CUDA launch syntax in this libtorch_stable file). + // not parse the CUDA launch syntax in this stable-ABI file). kernel<<>>(args...); // clang-format on #endif diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu rename to csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 0a45af8ade2b..09484daf3315 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -54,15 +54,15 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" -#include "../attention/dtype_fp8.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" #include "dispatch_utils.h" #ifdef USE_ROCM - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif // Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs. diff --git a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu b/csrc/fused_qknorm_rope_kernel.cu similarity index 99% rename from csrc/libtorch_stable/fused_qknorm_rope_kernel.cu rename to csrc/fused_qknorm_rope_kernel.cu index 01bad79cd8e2..a75ede86fc37 100644 --- a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu +++ b/csrc/fused_qknorm_rope_kernel.cu @@ -21,7 +21,7 @@ #include "torch_utils.h" #include "async_util.cuh" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "type_convert.cuh" #include "dispatch_utils.h" diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/gdn/fused_gdn_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu rename to csrc/gdn/fused_gdn_decode_kernel.cu index 1229404dc634..5e606ffd7fa1 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/gdn/fused_gdn_decode_kernel.cu @@ -10,7 +10,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/hisparse_kernels.cu similarity index 99% rename from csrc/libtorch_stable/hisparse_kernels.cu rename to csrc/hisparse_kernels.cu index 5c7621c2f598..da3857d1e973 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/csrc/hisparse_kernels.cu @@ -17,7 +17,6 @@ #include "torch_utils.h" #include "ops.h" -#include "../cuda_utils.h" #include #include diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/kimi_k3/attn_res_kernel.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu rename to csrc/kimi_k3/attn_res_kernel.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/kimi_k3/fused_kda_decode_kernel.cu similarity index 99% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu rename to csrc/kimi_k3/fused_kda_decode_kernel.cu index 00e5ce794d68..c5214ff48671 100644 --- a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu +++ b/csrc/kimi_k3/fused_kda_decode_kernel.cu @@ -9,7 +9,7 @@ #include #include "../torch_utils.h" -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" namespace { diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu b/csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu similarity index 100% rename from csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu rename to csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu diff --git a/csrc/libtorch_stable/launch_bounds_utils.h b/csrc/launch_bounds_utils.h similarity index 100% rename from csrc/libtorch_stable/launch_bounds_utils.h rename to csrc/launch_bounds_utils.h diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/layernorm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_kernels.cu rename to csrc/layernorm_kernels.cu index b342c59f1801..8d6741ddac48 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/layernorm_kernels.cu @@ -3,7 +3,7 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/layernorm_quant_kernels.cu similarity index 99% rename from csrc/libtorch_stable/layernorm_quant_kernels.cu rename to csrc/layernorm_quant_kernels.cu index f43be531de03..7f8fb9e51007 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/layernorm_quant_kernels.cu @@ -10,8 +10,8 @@ #include "torch_utils.h" #include "cub_helpers.h" -#include "../core/batch_invariant.hpp" -#include "../quantization/w8a8/fp8/common.cuh" +#include "core/utils.hpp" +#include "quantization/w8a8/fp8/common.cuh" #include "type_convert.cuh" #include "dispatch_utils.h" #include "quantization/vectorization_utils.cuh" diff --git a/csrc/libtorch_stable/attention/attention_utils.cuh b/csrc/libtorch_stable/attention/attention_utils.cuh deleted file mode 100644 index 29783cb0e5f3..000000000000 --- a/csrc/libtorch_stable/attention/attention_utils.cuh +++ /dev/null @@ -1,57 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "../../cuda_compat.h" -#include "../../attention/attention_dtypes.h" - -#include -#include - -namespace vllm { - -// Q*K^T operation. -template -inline __device__ float qk_dot_(const Vec (&q)[N], const Vec (&k)[N]) { - using A_vec = typename FloatVec::Type; - // Compute the parallel products for Q*K^T (treat vector lanes separately). - A_vec qk_vec = mul(q[0], k[0]); -#pragma unroll - for (int ii = 1; ii < N; ++ii) { - qk_vec = vllm::fma(q[ii], k[ii], qk_vec); - } - - // Finalize the reduction across lanes. - float qk = sum(qk_vec); -#pragma unroll - for (int mask = THREAD_GROUP_SIZE / 2; mask >= 1; mask /= 2) { - qk += VLLM_SHFL_XOR_SYNC(qk, mask); - } - return qk; -} - -template -struct Qk_dot { - template - static inline __device__ float dot(const Vec (&q)[N], const Vec (&k)[N]) { - return qk_dot_(q, k); - } -}; - -} // namespace vllm diff --git a/csrc/libtorch_stable/cutlass_extensions/common.cpp b/csrc/libtorch_stable/cutlass_extensions/common.cpp deleted file mode 100644 index 5bc9463bfa60..000000000000 --- a/csrc/libtorch_stable/cutlass_extensions/common.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "common.hpp" - -int32_t get_sm_version_num() { - int32_t major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - 0); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - 0); - int32_t version_num = major_capability * 10 + minor_capability; - return version_num; -} \ No newline at end of file diff --git a/csrc/libtorch_stable/dispatch_utils.h b/csrc/libtorch_stable/dispatch_utils.h deleted file mode 100644 index aa125a08fd6c..000000000000 --- a/csrc/libtorch_stable/dispatch_utils.h +++ /dev/null @@ -1,188 +0,0 @@ -/* - * Stable ABI compatible dispatch utilities for vLLM. - * Adapted from dispatch_utils.h to use PyTorch's header-only (THO_*) macros - * instead of the ATen (AT_*) macros. - * - * These macros use: - * - THO_DISPATCH_SWITCH instead of AT_DISPATCH_SWITCH - * - THO_DISPATCH_CASE instead of AT_DISPATCH_CASE - * - torch::headeronly::ScalarType instead of at::ScalarType - * - * Add more macros here as needed when migrating additional kernels. - */ -#pragma once - -#include -#include -#include - -// Need a special dispatch case macro since we will nest the FP8 dispatch. -// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. -#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) - -// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids) -// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'. -#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \ - THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \ - __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \ - __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) - -#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH( \ - TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) - -// FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...) \ - VLLM_STABLE_DISPATCH_FP8_CASE( \ - torch::headeronly::ScalarType::Float8_e4m3fn, __VA_ARGS__) -#endif - -// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. -// See VLLM_STABLE_DISPATCH_FP8_CASE above. -#define VLLM_STABLE_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) - -// Half types dispatch (Half + BFloat16) -#define VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) - -#define VLLM_STABLE_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) - -// Quant type dispatch (FP8 + INT8) -#ifdef USE_ROCM - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fnuz, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#else - #define VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(...) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float8_e4m3fn, \ - __VA_ARGS__) \ - THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__) -#endif - -#define VLLM_STABLE_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ - THO_DISPATCH_SWITCH(TYPE, NAME, \ - VLLM_STABLE_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) - -// Group size dispatch (pure C++ if/else, no ATen dependency) -#define VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ - if (group_size == 128) { \ - constexpr int const_group_size = 128; \ - __VA_ARGS__(); \ - } else if (group_size == 64) { \ - constexpr int const_group_size = 64; \ - __VA_ARGS__(); \ - } - -// Boolean dispatch -#define VLLM_STABLE_DISPATCH_BOOL(expr, const_expr, ...) \ - if (expr) { \ - constexpr bool const_expr = true; \ - __VA_ARGS__(); \ - } else { \ - constexpr bool const_expr = false; \ - __VA_ARGS__(); \ - } - -// Vec size dispatch (pure C++ switch, no ATen dependency) -#define VLLM_STABLE_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ - switch (VEC_SIZE) { \ - case 16: { \ - constexpr int vec_size = 16; \ - __VA_ARGS__(); \ - break; \ - } \ - case 8: { \ - constexpr int vec_size = 8; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int vec_size = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - case 2: { \ - constexpr int vec_size = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - default: { \ - constexpr int vec_size = 1; \ - __VA_ARGS__(); \ - break; \ - } \ - } - -// Tensor rank dispatch (2D, 3D, 4D) -#define VLLM_STABLE_DISPATCH_RANK234(NUM_DIMS, ...) \ - switch (NUM_DIMS) { \ - case 2: { \ - constexpr int tensor_rank = 2; \ - __VA_ARGS__(); \ - break; \ - } \ - case 3: { \ - constexpr int tensor_rank = 3; \ - __VA_ARGS__(); \ - break; \ - } \ - case 4: { \ - constexpr int tensor_rank = 4; \ - __VA_ARGS__(); \ - break; \ - } \ - default: \ - STD_TORCH_CHECK( \ - false, "Expects rank 2, 3 or 4 tensors but got unsupported rank"); \ - } diff --git a/csrc/libtorch_stable/mamba/static_switch.h b/csrc/libtorch_stable/mamba/static_switch.h deleted file mode 100644 index 840cb2374a2f..000000000000 --- a/csrc/libtorch_stable/mamba/static_switch.h +++ /dev/null @@ -1,28 +0,0 @@ -// Inspired by -// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h -// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h - -// clang-format off -// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/static_switch.h -#pragma once - -/// @param COND - a boolean expression to switch by -/// @param CONST_NAME - a name given for the constexpr bool variable. -/// @param ... - code to execute for true and false -/// -/// Usage: -/// ``` -/// BOOL_SWITCH(flag, BoolConst, [&] { -/// some_function(...); -/// }); -/// ``` -#define BOOL_SWITCH(COND, CONST_NAME, ...) \ - [&] { \ - if (COND) { \ - constexpr bool CONST_NAME = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr bool CONST_NAME = false; \ - return __VA_ARGS__(); \ - } \ - }() diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py deleted file mode 100644 index 6948a98a374d..000000000000 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py +++ /dev/null @@ -1,315 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import glob -import itertools -import os -import subprocess -import sys - -import jinja2 - -ARCHS = [] -SUPPORT_FP8 = False -SUPPORT_SM75 = False -SUPPORT_SM80 = False -for arch in sys.argv[1].split(","): - arch = arch[: arch.index(".") + 2].replace(".", "") - arch = int(arch) - # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10) - # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. - # SM90 and SM100 can use this PTX, but it’s simulated - # with FP16 MMA, so it cannot achieve any acceleration. - if arch == 89 or arch // 10 == 12: - SUPPORT_FP8 = True - if arch >= 80: - SUPPORT_SM80 = True - if arch == 75: - SUPPORT_SM75 = True - -FILE_HEAD_COMMENT = """ -// auto generated by generate_kernels.py -// clang-format off -""".lstrip() - -FILE_HEAD = ( - FILE_HEAD_COMMENT - + """ -#include "kernel.h" -#include "marlin_template.h" - -namespace MARLIN_NAMESPACE_NAME { -""" -) - -TEMPLATE = ( - "template __global__ void Marlin<" - "{{a_type_id}}, " - "{{b_type_id}}, " - "{{c_type_id}}, " - "{{s_type_id}}, " - "{{threads}}, " - "{{thread_m_blocks}}, " - "{{thread_n_blocks}}, " - "{{thread_k_blocks}}, " - "{{m_block_size_8}}, " - "{{stages}}, " - "{{group_blocks}}, " - "{{is_zp_float}}>" - "( MARLIN_KERNEL_PARAMS );" -) - -THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] - -THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] - -QUANT_CONFIGS = [ - # AWQ-INT4 - { - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 - { - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT8 - { - "b_type": "kU8B128", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # FP8 - { - "b_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 8], - }, - # NVFP4 - { - "b_type": "kFE2M1f", - "s_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [1], - }, - # MXFP4 - { - "a_type": ["kBFloat16"], - "b_type": "kFE2M1f", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # MXFP8 - { - "a_type": ["kBFloat16"], - "b_type": "kFE4M3fn", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # AWQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # MXFP4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kFE2M1f", - "c_type": ["kBFloat16"], - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [2], - }, -] - - -def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): - subprocess.call(["rm", "-f", filename]) - - filename = os.path.dirname(__file__) + "/kernel_selector.h" - subprocess.call(["rm", "-f", filename]) - - -def generate_new_kernels(): - result_dict = {} - sm_75_result_dict = {} - - for quant_config in QUANT_CONFIGS: - c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"]) - a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"]) - b_type = quant_config["b_type"] - all_group_blocks = quant_config["group_blocks"] - all_m_blocks = quant_config["thread_m_blocks"] - all_thread_configs = quant_config["thread_configs"] - - for a_type, c_type in itertools.product(a_types, c_types): - if not SUPPORT_FP8 and a_type == "kFE4M3fn": - continue - if "16" in a_type and "16" in c_type and a_type != c_type: - continue - s_type = quant_config.get("s_type", c_type) - if (a_type, b_type, c_type) not in result_dict: - result_dict[(a_type, b_type, c_type)] = [] - if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16": - sm_75_result_dict[(a_type, b_type, c_type)] = [] - - for group_blocks, m_blocks, thread_configs in itertools.product( - all_group_blocks, all_m_blocks, all_thread_configs - ): - thread_k, thread_n, threads = thread_configs - - if threads == 256: - # for small batch (m_blocks == 1), - # we only need (128, 128, 256) - # for large batch (m_blocks > 1), - # we only need (64, 256, 256) - if m_blocks <= 1 and (thread_k, thread_n) != (128, 128): - continue - if m_blocks > 1 and (thread_k, thread_n) != (64, 256): - continue - - config = { - "threads": threads, - "s_type": s_type, - "thread_m_blocks": max(m_blocks, 1), - "thread_k_blocks": thread_k // 16, - "thread_n_blocks": thread_n // 16, - "m_block_size_8": "true" if m_blocks == 0.5 else "false", - "stages": 4, - "group_blocks": group_blocks, - "is_zp_float": "false", - } - - if SUPPORT_SM80: - result_dict[(a_type, b_type, c_type)].append(config) - if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75: - config_sm75 = config.copy() - config_sm75["stages"] = 2 - sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75) - - kernel_selector_str = FILE_HEAD_COMMENT - - for result_dict_tmp in [result_dict, sm_75_result_dict]: - for (a_type, b_type, c_type), config_list in result_dict_tmp.items(): - all_template_str_list = [] - if not config_list: - continue - for config in config_list: - s_type = config["s_type"] - template_str = jinja2.Template(TEMPLATE).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - all_template_str_list.append(template_str) - - conditions = [ - f"a_type == vllm::{a_type}", - f"b_type == vllm::{b_type}", - f"c_type == vllm::{c_type}", - f"s_type == vllm::{s_type}", - f"threads == {config['threads']}", - f"thread_m_blocks == {config['thread_m_blocks']}", - f"thread_n_blocks == {config['thread_n_blocks']}", - f"thread_k_blocks == {config['thread_k_blocks']}", - f"m_block_size_8 == {config['m_block_size_8']}", - f"stages == {config['stages']}", - f"group_blocks == {config['group_blocks']}", - f"is_zp_float == {config['is_zp_float']}", - ] - conditions = " && ".join(conditions) - - if kernel_selector_str == FILE_HEAD_COMMENT: - kernel_selector_str += f"if ({conditions})\n kernel = " - else: - kernel_selector_str += f"else if ({conditions})\n kernel = " - - kernel_template2 = ( - "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, " - "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, " - "{{thread_n_blocks}}, {{thread_k_blocks}}, " - "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, " - "{{is_zp_float}}>;" - ) - - kernel_selector_str += ( - jinja2.Template(kernel_template2).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - + "\n" - ) - - file_content = FILE_HEAD + "\n\n" - file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" - if a_type == "kFE4M3fn": - filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - elif result_dict_tmp is sm_75_result_dict: - filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - else: - filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - - filename = filename.lower() - - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: - f.write(file_content) - - if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: - kernel_selector_str += ( - "else if (a_type == vllm::kFE4M3fn)\n" - " STD_TORCH_CHECK(false, " - '"marlin kernel with fp8 activation is not built.");' - ) - - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: - f.write(kernel_selector_str) - - -if __name__ == "__main__": - remove_old_kernels() - generate_new_kernels() diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h deleted file mode 100644 index 3266d28b1edb..000000000000 --- a/csrc/libtorch_stable/ops.h +++ /dev/null @@ -1,822 +0,0 @@ -#pragma once - -#include -#include -#include - -#include -#include -#include - -#include - -inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { - // Ensure tensor is on CUDA - STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); - - // Get the raw data pointer - void* data_ptr = tensor.mutable_data_ptr(); - - /// Create a new tensor from the raw data pointer - return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), - tensor.device(), tensor.scalar_type()); -} - -void per_token_group_quant_fp8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double fp8_min, - double fp8_max, bool scale_ue8m0, - bool dummy_is_scale_transposed, - bool dummy_is_tma_aligned); - -// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max); - -#ifndef USE_ROCM -bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); -bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); -bool cutlass_group_gemm_supported(int64_t cuda_device_capability); - -void cutlass_scaled_mm(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - std::optional const& bias); - -void cutlass_moe_mm(torch::stable::Tensor& out_tensors, - torch::stable::Tensor const& a_tensors, - torch::stable::Tensor const& b_tensors, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& expert_offsets, - torch::stable::Tensor const& problem_sizes, - torch::stable::Tensor const& a_strides, - torch::stable::Tensor const& b_strides, - torch::stable::Tensor const& c_strides, bool per_act_token, - bool per_out_ch); - -void cutlass_scaled_mm_azp(torch::stable::Tensor& out, - torch::stable::Tensor const& a, - torch::stable::Tensor const& b, - torch::stable::Tensor const& a_scales, - torch::stable::Tensor const& b_scales, - torch::stable::Tensor const& azp_adj, - std::optional const& azp, - std::optional const& bias); - -void get_cutlass_moe_mm_data( - const torch::stable::Tensor& topk_ids, - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - torch::stable::Tensor& input_permutation, - torch::stable::Tensor& output_permutation, const int64_t num_experts, - const int64_t n, const int64_t k, - const std::optional& blockscale_offsets, - const bool is_gated); - -void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( - const torch::stable::Tensor& expert_first_token_offset, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, - const bool swap_ab); - -void get_cutlass_batched_moe_mm_data( - torch::stable::Tensor& expert_offsets, - torch::stable::Tensor& problem_sizes1, - torch::stable::Tensor& problem_sizes2, - const torch::stable::Tensor& expert_num_tokens, - const int64_t num_local_experts, const int64_t padded_m, const int64_t n, - const int64_t k); - -// FP4/NVFP4 ops -bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); - -void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, - torch::stable::Tensor const& A, - torch::stable::Tensor const& B, - torch::stable::Tensor const& A_sf, - torch::stable::Tensor const& B_sf, - torch::stable::Tensor const& alpha); - -void cutlass_fp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& alphas, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -std::tuple scaled_fp4_quant_func( - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); - -void scaled_fp4_quant_out(torch::stable::Tensor const& input, - torch::stable::Tensor const& input_scale, - bool is_sf_swizzled_layout, - torch::stable::Tensor& output, - torch::stable::Tensor& output_scale); - -void scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_scaled_fp4_experts_quant( - torch::stable::Tensor& output, torch::stable::Tensor& output_scale, - torch::stable::Tensor const& input, - torch::stable::Tensor const& input_global_scale, - torch::stable::Tensor const& input_offset_by_experts, - torch::stable::Tensor const& output_scale_offset_by_experts); - -void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, - torch::stable::Tensor& output_block_scale, - torch::stable::Tensor& input, - torch::stable::Tensor& input_global_scale); - -void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, - const torch::stable::Tensor& a, - const torch::stable::Tensor& b, - const torch::stable::Tensor& a_blockscale, - const torch::stable::Tensor& b_blockscales, - const torch::stable::Tensor& problem_sizes, - const torch::stable::Tensor& expert_offsets, - const torch::stable::Tensor& sf_offsets); - -// AWQ ops -torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, - torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters); - -torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, - torch::stable::Tensor _scaling_factors, - torch::stable::Tensor _zeros, - int64_t split_k_iters, int64_t thx, - int64_t thy); - -// DSV3 fused A GEMM: conditionally compiled so declaration and impl -// registration are in the source file (dsv3_fused_a_gemm.cu) - -#endif - -// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) -torch::stable::Tensor get_cuda_view_from_cpu_tensor( - torch::stable::Tensor& cpu_tensor); - -// Attention kernels (shared CUDA/ROCm) -void merge_attn_states( - torch::stable::Tensor& output, - std::optional output_lse, - const torch::stable::Tensor& prefix_output, - const torch::stable::Tensor& prefix_lse, - const torch::stable::Tensor& suffix_output, - const torch::stable::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context, - const std::optional& output_scale = std::nullopt); - -torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, - bool inplace); - -// Layernorm kernels (shared CUDA/ROCm) -void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, - std::optional weight, double epsilon); - -void fused_add_rms_norm(torch::stable::Tensor& input, - torch::stable::Tensor& residual, - std::optional weight, - double epsilon); - -// Layernorm-quant kernels (shared CUDA/ROCm) -void rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, double epsilon); - -void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& residual, - torch::stable::Tensor& weight, - torch::stable::Tensor& scale, - double epsilon); - -// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) -void rms_norm_dynamic_per_token_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, torch::stable::Tensor& scales, - double const var_epsilon, std::optional scale_ub, - std::optional residual); - -void rms_norm_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& weight, - torch::stable::Tensor& scales, - double const var_epsilon, - std::optional scale_ub, - std::optional residual, - int64_t group_size, bool is_scale_transposed); - -void silu_and_mul_per_block_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - int64_t group_size, - std::optional scale_ub, - bool is_scale_transposed); - -// Positional encoding kernels (shared CUDA/ROCm) -void rotary_embedding(torch::stable::Tensor& positions, - torch::stable::Tensor& query, - std::optional key, - int64_t head_size, torch::stable::Tensor& cos_sin_cache, - bool is_neox, int64_t rope_dim_offset, bool inverse); - -void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, - int64_t num_heads_k, int64_t num_heads_v, - int64_t head_dim, double eps, - torch::stable::Tensor& q_weight, - torch::stable::Tensor& k_weight, - torch::stable::Tensor& cos_sin_cache, bool is_neox, - torch::stable::Tensor& position_ids, - int64_t forced_token_heads_per_warp); - -void fused_deepseek_v4_kv_rope_insert( - torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, - std::optional fp8_scale, bool kv_mxfp8); - -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, - bool apply_q_rope, bool is_q_interleaved); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, double eps, - int64_t cache_block_size, bool apply_q_norm); - -void fused_kimi_k3_mla_key_concat_kv_cache_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_key_concat_ds_mla_insert( - torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor& k_out); - -void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, - torch::stable::Tensor const& v, - torch::stable::Tensor& k_fp8, - torch::stable::Tensor& v_fp8); - -void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, - torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, - torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, - torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, - torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& k_scale_inv, - torch::stable::Tensor const& v_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& q_scale_inv, - torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( - torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, - torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, - torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, - std::optional position_ids, - std::optional cos_sin_cache); - -void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - torch::stable::Tensor const& q, torch::stable::Tensor const& kv, - torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, - torch::stable::Tensor const& slot_mapping, - torch::stable::Tensor const& position_ids, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& fp8_scale, - torch::stable::Tensor const& q_fp8_scale_inv, double eps, - int64_t cache_block_size, bool apply_q_norm); - -#ifndef USE_ROCM -std::tuple -minimax_allreduce_rms_qk(torch::stable::Tensor qkv, - torch::stable::Tensor const& norm_weight_q, - torch::stable::Tensor const& norm_weight_k, - torch::stable::Tensor workspace, int64_t const q_size, - int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps); -#endif - -// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / -// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs -// the index branch and scatters k/v/index_k into their paged caches. -void fused_minimax_m3_qknorm_rope_kv_insert( - torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, - torch::stable::Tensor const& k_norm_weight, - torch::stable::Tensor const& cos_sin_cache, - torch::stable::Tensor const& positions, int64_t num_heads, - int64_t num_kv_heads, int64_t rotary_dim, double eps, - std::optional index_q_norm_weight, - std::optional index_k_norm_weight, - int64_t num_index_heads, std::optional slot_mapping, - std::optional index_slot_mapping, - std::optional kv_cache, - std::optional index_cache, int64_t block_size, - std::optional q_out, - std::optional index_q_out, - const std::string& kv_cache_dtype, bool skip_index_branch, - std::optional q_fp8_out, double q_fp8_scale, - std::optional kv_k_scale, - std::optional kv_v_scale); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE -void fused_kda_decode( - torch::stable::Tensor const& x, torch::stable::Tensor const& weight, - std::optional bias, - torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, - torch::stable::Tensor& out, std::optional lower_bound, - std::optional output_gate, - std::optional norm_weight, double norm_eps); - -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE -void fused_gdn_decode_post_conv_mtp( - torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, - torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, - torch::stable::Tensor const& state_indices, - torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& num_accepted_tokens, - torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, - double scale, double norm_eps, const std::string& output_gate_activation); - -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK -void fused_kda_prologue( - torch::stable::Tensor const& q, torch::stable::Tensor const& k, - torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, - torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, - torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, - torch::stable::Tensor& w, torch::stable::Tensor& u, - torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, - torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_indices, - std::optional conv_weight, - std::optional conv_state, - std::optional conv_state_indices, - std::optional conv_has_initial_state, double scale, - double lower_bound); - -void fused_kda_chunk( - torch::stable::Tensor const& qg, torch::stable::Tensor const& w, - torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, - torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, - std::optional initial_state, - std::optional final_state, - torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, - torch::stable::Tensor const& chunk_offsets, double scale, - std::optional group_state, int64_t groups, - std::optional checkpoint_state, - std::optional checkpoint_offsets, - std::optional checkpoint_state_indices, - std::optional state_cache, - std::optional state_indices, - std::optional has_initial_state); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES -void kimi_k3_attn_res(torch::stable::Tensor& prefix, - std::optional delta, - torch::stable::Tensor& blocks, - torch::stable::Tensor const& norm_weight, - torch::stable::Tensor const& qk_weight, - std::optional output_norm_weight, - torch::stable::Tensor& output, int64_t num_blocks, - int64_t block_write_idx, double eps, - double output_norm_eps); -#endif - -// Sampler kernels (shared CUDA/ROCm) -void apply_repetition_penalties_( - torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, - const torch::stable::Tensor& output_mask, - const torch::stable::Tensor& repetition_penalties); - -void top_k_per_row_prefill(const torch::stable::Tensor& logits, - const torch::stable::Tensor& rowStarts, - const torch::stable::Tensor& rowEnds, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, - const torch::stable::Tensor& seqLens, - torch::stable::Tensor& indices, int64_t numRows, - int64_t stride0, int64_t stride1, int64_t topK); - -void persistent_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK -void cooperative_topk(const torch::stable::Tensor& logits, - const torch::stable::Tensor& lengths, - torch::stable::Tensor& output, - torch::stable::Tensor& workspace, int64_t k, - int64_t max_seq_len); -#endif - -void selective_scan_fwd( - const torch::stable::Tensor& u, const torch::stable::Tensor& delta, - const torch::stable::Tensor& A, const torch::stable::Tensor& B, - const torch::stable::Tensor& C, - const std::optional& D_, - const std::optional& z_, - const std::optional& delta_bias_, - bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, - const torch::stable::Tensor& ssm_states, int64_t null_block_id, - int64_t block_size, - const std::optional& block_idx_first_scheduled_token, - const std::optional& block_idx_last_scheduled_token, - const std::optional& initial_state_idx, - const std::optional& cu_chunk_seqlen, - const std::optional& last_chunk_indices); - -using fptr_t = int64_t; -fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::stable::Tensor& rank_data, int64_t rank, - bool fully_connected); -void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t local_buffer, - fptr_t multicast_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, fptr_t reg_buffer, - int64_t reg_buffer_sz_bytes); -void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t epoch_buffer, - int64_t stage_sz_bytes); -void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, - torch::stable::Tensor& out, - fptr_t local_buffer, fptr_t multicast_buffer, - int64_t stage_sz_bytes, int64_t block_limit); -void dispose(fptr_t _fa); -int64_t meta_size(); -void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa); -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets); -std::tuple allocate_shared_buffer_and_handle( - int64_t size); -int64_t open_mem_handle(torch::stable::Tensor& mem_handle); -void free_shared_buffer(int64_t buffer); - -// Activation kernels (shared CUDA/ROCm) -void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void silu_and_mul_clamp(torch::stable::Tensor& out, - torch::stable::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); - -void silu_and_mul_quant(torch::stable::Tensor& out, - torch::stable::Tensor& input, - torch::stable::Tensor& scale); - -void persistent_masked_m_silu_mul_quant( - const torch::stable::Tensor& input, // (E, T, 2*H) - const torch::stable::Tensor& tokens_per_expert, // (E) - torch::stable::Tensor& y_q, // (E, T, H) [OUT] - torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] - bool use_ue8m0); - -void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_tanh_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input); -void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double threshold); -void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double alpha = 1.702, double limit = 7.0); -void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, - double beta = 1.0, double linear_beta = -1.0); -void situ_and_mul_quant( - torch::stable::Tensor& out, torch::stable::Tensor& scale, - torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, - int64_t group_size = 0, - std::optional num_valid_tokens = std::nullopt, - int64_t topk = 1); -void masked_situ_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& expert_num_tokens, - double beta = 1.0, double linear_beta = -1.0); -void masked_moe_activation(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& valid_token_counts, - const std::string& activation, - double clamp_limit = 0.0, double alpha = 1.0, - double beta = 0.0, double situ_beta = 1.0, - double situ_linear_beta = -1.0); -void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); -void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); - -void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); - -// INT8 quantization kernels (shared CUDA/ROCm) -void static_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional const& azp); - -void dynamic_scaled_int8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scales, - std::optional const& azp); - -// FP8 quantization kernels (shared CUDA/ROCm) -void static_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor const& scale, - std::optional group_shape = - std::nullopt); - -void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, - torch::stable::Tensor const& input, - torch::stable::Tensor& scale); - -void dynamic_per_token_scaled_fp8_quant( - torch::stable::Tensor& out, torch::stable::Tensor const& input, - torch::stable::Tensor& scale, - std::optional const& scale_ub); - -// GPTQ kernels (shared CUDA/ROCm) -torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, - torch::stable::Tensor b_q_weight, - torch::stable::Tensor b_gptq_qzeros, - torch::stable::Tensor b_gptq_scales, - bool use_exllama, bool use_v2_format, - int64_t bit); - -void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); - -// Cache ops (shared CUDA/ROCm) -void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, - int64_t block_size_in_bytes, - const torch::stable::Tensor& block_mapping); - -// Batch swap: submit all block copies in a single driver call. -void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, - const torch::stable::Tensor& dst_ptrs, - const torch::stable::Tensor& sizes, - bool is_src_access_order_any); - -void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale); - -void reshape_and_cache_flash( - torch::stable::Tensor& key, torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); - -void concat_and_cache_mla(torch::stable::Tensor& kv_c, - torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache, - torch::stable::Tensor& slot_mapping, - const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); - -void concat_and_cache_mla_grouped( - torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, - torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, - int64_t block_size, int64_t block_stride, int64_t entry_stride, - std::optional kv_scales, - const std::string& kv_cache_dtype); - -#ifndef USE_ROCM -// HiSparse kernels use raw PTX in the row copy; CUDA-only. -void hisparse_resolve_residency( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& hot_block_table, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor& hot_indices, - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor& lru_slots, - std::optional const& request_state_indices, - int64_t region_stride, int64_t max_union_rows, - std::optional const& miss_mask, - std::optional const& stats, - std::optional const& attention_indices, - int64_t attention_block_stride, - std::optional const& request_ids, - std::optional const& source_block_table, - int64_t source_block_size, - std::optional const& resolved_global_indices, - std::optional const& valid_counts, - std::optional const& swap_host_physical_rows, - std::optional const& swap_device_physical_rows, - std::optional const& swap_counts, - std::optional const& resident_block_table, - int64_t resident_block_size, int64_t resident_null_block); - -void hisparse_invalidate_written_slots( - torch::stable::Tensor& device_global_indices, - torch::stable::Tensor const& request_state_indices, - torch::stable::Tensor const& req_id_per_token, - torch::stable::Tensor const& written_slots); - -void hisparse_gather_plan( - torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& global_indices, - torch::stable::Tensor const& hot_indices, - torch::stable::Tensor const& miss_mask, - std::optional const& request_state_indices, - std::optional const& attention_indices, - int64_t attention_block_stride); - -void hisparse_gather_compact(torch::stable::Tensor const& host_cache, - torch::stable::Tensor& hot_cache, - torch::stable::Tensor const& miss_global_indices, - torch::stable::Tensor const& miss_hot_indices, - torch::stable::Tensor const& miss_counts); - -#endif // !USE_ROCM - -// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla -void concat_and_cache_mla_rope_fused( - torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, - torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, - torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, - torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, - const std::string& kv_cache_dtype, - torch::stable::Tensor& kv_cache_quant_scale); - -// Just for unittest -void convert_fp8(torch::stable::Tensor& dst_cache, - torch::stable::Tensor& src_cache, const double scale, - const std::string& kv_cache_dtype); - -void gather_and_maybe_dequant_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] - int64_t num_tokens, const std::string& kv_cache_dtype, - torch::stable::Tensor const& scale, - std::optional seq_starts = std::nullopt); - -// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. -void cp_gather_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // ENTRIES...] - torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] - int64_t batch_size, - std::optional seq_starts = std::nullopt); - -// Gather and upconvert FP8 KV cache to BF16 workspace -void cp_gather_and_upconvert_fp8_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 656] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size, - std::optional seq_starts = std::nullopt, - std::optional host_cache = std::nullopt, - std::optional host_row_ids = std::nullopt, - std::optional device_row_ids = std::nullopt); - -// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace -void cp_gather_and_upconvert_nvfp4_kv_cache( - torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, - // 352] - torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] - torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); - -// Indexer K quantization and cache function -void indexer_k_quant_and_cache( - torch::stable::Tensor& k, // [num_tokens, head_dim] - torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size - const std::string& scale_fmt); - -// Concatenate query nope and rope for MLA/DSA attention -void concat_mla_q( - torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + - // rope_dim] - -// Extract function to gather quantized K cache -void cp_gather_indexer_k_quant_cache( - const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, - // cache_stride] - torch::stable::Tensor& dst_k, // [num_tokens, head_dim] - torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / - // quant_block_size * 4] - const torch::stable::Tensor& block_table, // [batch_size, num_blocks] - const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] - -// Fused vocab-parallel embedding lookup (see -// vocab_parallel_embedding_kernels.cu). -void vocab_parallel_embedding( - torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, - const torch::stable::Tensor& weight, int64_t org_vocab_start_index, - int64_t org_vocab_end_index, int64_t num_org_vocab_padding, - int64_t added_vocab_start_index, int64_t added_vocab_end_index); - -// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). -void ngram_compute_n_gram_ids( - int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, - torch::stable::Tensor& ne_mods, - torch::stable::Tensor& exclusive_ne_embedder_size_sums, - torch::stable::Tensor& exclusive_req_len_sums, - torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, - torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu deleted file mode 100644 index 2ffbee7ca26e..000000000000 --- a/csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu +++ /dev/null @@ -1,12 +0,0 @@ -#include - -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - -void per_token_group_quant_int8(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double int8_min, - double int8_max) { - per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, - int8_min, int8_max); -} diff --git a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h b/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h deleted file mode 100644 index 6630c0decee6..000000000000 --- a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -// 8-bit per-token-group quantization helper used by both FP8 and INT8 -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0 = false); - -// Public op: register-resident packed quant for the DeepGEMM Blackwell path. -// Restricted to group_size == 128 and bf16/fp16 input; other configurations -// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because -// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes. -void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s_packed, - int64_t group_size, double eps, - double min_8bit, double max_8bit); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp deleted file mode 100644 index 3d4f9ab27389..000000000000 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ /dev/null @@ -1,1171 +0,0 @@ -#include "ops.h" -#include "cuda_utils.h" -#include "core/registration.h" - -#include - -// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. -// Note: We register under namespace "_C" so ops are accessible as -// torch.ops._C. for compatibility with existing code. -STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { - // Compute per-token-group FP8 quantized tensor and scaling factor. - // The dummy arguments are here so we can correctly fuse with RMSNorm. - ops.def( - "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " - "output_s, " - "int group_size, float eps, float fp8_min, float fp8_max, bool " - "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " - ") -> ()"); - // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, - // TMA-aligned scales for DeepGEMM. - ops.def( - "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " - "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " - "float fp8_max) -> ()"); - // Compute per-token-group INT8 quantized tensor and scaling factor. - ops.def( - "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " - "output_s, int group_size, float eps, float int8_min, float int8_max) -> " - "()"); - ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); - -#ifndef USE_ROCM - - // Note about marlin kernel 'workspace' arguments: - // Technically these should be mutable since they are modified by the kernel. - // But since they are set back to zero once the kernel is finished we can - // hand wave and say that they have no net effect. - // - // The reason to mark 'workspace' as immutable is so that they don't interfere - // with using ScalarType arguments in the ops. If they are marked as mutable, - // pytorch throws an assert in - // 'torch._higher_order_ops._register_effectful_op' that prevents these - // kernels from being torch.compile'd. - // See the following document for more info on custom types and ops that use - // custom types: - // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA - - // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. - ops.def( - "machete_supported_schedules(" - " ScalarType a_type," - " int b_type," - " ScalarType? maybe_group_scales_type," - " ScalarType? maybe_group_zeros_type," - " ScalarType? maybe_channel_scales_type," - " ScalarType? maybe_token_scales_type," - " ScalarType? maybe_out_type" - ") -> str[]"); - ops.def( - "machete_mm(" - " Tensor A," - " Tensor B," - " int b_type," - " ScalarType? out_type," - " Tensor? group_scales," - " Tensor? group_zeros," - " int? group_size," - " Tensor? channel_scales," - " Tensor? token_scales," - " str? schedule" - ") -> Tensor"); - ops.def( - "machete_prepack_B(" - " Tensor B," - " ScalarType a_type," - " int b_type," - " ScalarType? group_scales_type" - ") -> Tensor"); - // conditionally compiled so impl registration is in source file - - // Marlin GEMM - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file -#endif - -#ifndef USE_ROCM - // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column - // quantization, as well as bias - ops.def( - "cutlass_scaled_mm(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor? bias) -> ()"); - - // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column - // quantization. - ops.def( - "cutlass_scaled_mm_azp(Tensor! out, Tensor a," - " Tensor b, Tensor a_scales," - " Tensor b_scales, Tensor azp_adj," - " Tensor? azp, Tensor? bias) -> ()"); - - // Check if cutlass scaled_mm is supported for CUDA devices of the given - // capability - ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); - - // Check if cutlass grouped gemm is supported for CUDA devices of the given - // capability - ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); - - // CUTLASS w8a8 grouped GEMM - ops.def( - "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " - " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " - " Tensor problem_sizes, Tensor a_strides, " - " Tensor b_strides, Tensor c_strides, bool per_act_token, " - " bool per_out_ch) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM. It takes topk_ids as an input, and computes expert_offsets - // (token start indices of each expert). In addition to this, it computes - // problem sizes for each expert's multiplication used by the two mms called - // from fused MoE operation, and arrays with permutations required to shuffle - // and de-shuffle the input/output of the fused operation. - ops.def( - "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " - " Tensor! problem_sizes1, Tensor! problem_sizes2, " - " Tensor! input_permutation, " - " Tensor! output_permutation, int num_experts, " - " int n, int k, Tensor? blockscale_offsets, " - " bool is_gated) -> ()"); - - // compute per-expert problem sizes from expert_first_token_offset - // produced by vLLM's moe_permute kernel - ops.def( - "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" - " Tensor expert_first_token_offset, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " int n, int k, bool swap_ab) -> ()"); - - // A function that computes data required to run fused MoE with w8a8 grouped - // GEMM in batched expert format. It takes expert_num_tokens - // as an input, and computes expert_offsets (token start indices of each - // expert). In addition to this, it computes problem sizes for each expert's - // multiplication used by the two mms called from fused MoE operation. - ops.def( - "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " - " Tensor! problem_sizes1, " - " Tensor! problem_sizes2, " - " Tensor expert_num_tokens, " - " int num_local_experts, int padded_m, " - " int n, int k) -> ()"); - - // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) - ops.def( - "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " - "bool"); - - // CUTLASS nvfp4 block scaled GEMM - ops.def( - "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," - " Tensor block_scale_a, Tensor block_scale_b," - " Tensor alpha) -> ()"); - - // cutlass nvfp4 block scaled group GEMM - ops.def( - "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) - ops.def( - "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," - " Tensor a_blockscale, Tensor b_blockscales," - " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); - - // Compute NVFP4 block quantized tensor. - ops.def( - "scaled_fp4_quant(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout) -> (Tensor, Tensor)"); - - // Out variant - // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) - // This registration is now migrated to stable ABI - // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not - // yet in torch/headeronly), the tag should be applied from Python - // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) - // with the .impl remaining in C++. - // See pytorch/pytorch#176117. - ops.def( - "scaled_fp4_quant.out(Tensor input," - " Tensor input_scale, bool " - "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " - "-> ()"); - - // Compute NVFP4 experts quantization. - ops.def( - "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 experts quantization. - ops.def( - "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts) -> ()"); - - // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). - ops.def( - "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+MXFP4 experts quantization. - ops.def( - "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " - "output_scale," - "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); - - // Fused SiLU+Mul+NVFP4 quantization. - ops.def( - "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " - "Tensor input, Tensor input_global_scale) -> ()"); - - // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices - // of the given capability - ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); - - // CUTLASS w4a8 GEMM - ops.def( - "cutlass_w4a8_mm(" - " Tensor A," - " Tensor B," - " Tensor group_scales," - " int group_size," - " Tensor channel_scales," - " Tensor token_scales," - " ScalarType? out_type," - " str? maybe_schedule" - ") -> Tensor"); - - // pack scales - ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); - - // encode and reorder weight matrix - ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); - - // CUTLASS w4a8 grouped GEMM - ops.def( - "cutlass_w4a8_moe_mm(" - " Tensor! out_tensors," - " Tensor a_tensors," - " Tensor b_tensors," - " Tensor a_scales," - " Tensor b_scales," - " Tensor b_group_scales," - " int b_group_size," - " Tensor expert_offsets," - " Tensor problem_sizes," - " Tensor a_strides," - " Tensor b_strides," - " Tensor c_strides," - " Tensor group_scale_strides," - " str? maybe_schedule" - ") -> ()"); - - ops.def( - "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " - "Tensor)"); - - // SM100 CUTLASS MLA decode - // conditionally compiled so impl registrations are in source file - ops.def( - "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," - " Tensor q_pe, Tensor kv_c_and_k_pe_cache," - " Tensor seq_lens, Tensor page_table," - " Tensor workspace, float scale," - " int num_kv_splits) -> ()"); - - ops.def( - "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," - " int sm_count, int num_kv_splits) " - "-> int"); - // Quantized GEMM for AWQ. - ops.def( - "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters) -> Tensor"); - - // Dequantization for AWQ. - ops.def( - "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " - "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); - - // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). - // conditionally compiled so impl registration is in source file - ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " - "bool enable_pdl=False) -> ()"); - - // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). - // conditionally compiled so impl registration is in source file - ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); -#endif - - // Merge attn states - // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 - // can be used to combine partial attention results (in the split-KV case) - ops.def( - "merge_attn_states(" - " Tensor! output," - " Tensor!? output_lse," - " Tensor prefix_output," - " Tensor prefix_lse," - " Tensor suffix_output," - " Tensor suffix_lse," - " int!? prefill_tokens_with_context," - " Tensor? output_scale=None) -> ()"); - - // Hadamard transforms - // conditionally compiled so impl registration is in source file - ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); - - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " - "-> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " - "float epsilon) -> ()"); - - // Layernorm-quant - // Apply Root Mean Square (RMS) Normalization to the input tensor. - ops.def( - "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " - "Tensor scale, float epsilon) -> " - "()"); - - // In-place fused Add and RMS Normalization. - ops.def( - "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " - "Tensor! residual, Tensor weight, " - "Tensor scale, float epsilon) -> ()"); - - // Fused Layernorm + Quant kernels - ops.def( - "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual) -> ()"); - - // Fused Layernorm + Block quant kernels - ops.def( - "rms_norm_per_block_quant(Tensor! result, Tensor input, " - "Tensor weight, Tensor! scale, float epsilon, " - "Tensor? scale_ub, Tensor!? residual, int group_size, " - "bool is_scale_transposed) -> ()"); - - // Fused SiLU+Mul + per-block quantization - ops.def( - "silu_and_mul_per_block_quant(" - "Tensor! out, " - "Tensor input, " - "Tensor! scales, " - "int group_size, " - "Tensor? scale_ub=None, " - "bool is_scale_transposed=False) -> ()"); - - // Rotary embedding - // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. - ops.def( - "rotary_embedding(Tensor positions, Tensor! query," - " Tensor!? key, int head_size," - " Tensor cos_sin_cache, bool is_neox, int " - "rope_dim_offset=0, bool inverse=False) -> ()"); - - // Function for fused QK Norm and RoPE - ops.def( - "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " - "int num_heads_k, int num_heads_v, int head_dim, float eps, " - "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids, " - "int forced_token_heads_per_warp=-1) -> ()"); - - // q_head_padded is the padded Q head count of the returned tensor, or 0 to - // do the KV insert alone and return an empty tensor. The Q knobs are - // independent: apply_q_norm and apply_q_rope each drop that step for Q - // alone (KV is always rotated), and is_q_interleaved reads and writes Q in - // FlashMLA's mega-attention chunk-interleaved layout, which moves the - // padding heads to the tail of every head-dim chunk. - ops.def( - "fused_deepseek_v4_kv_rope_insert(" - "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " - "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " - "bool kv_mxfp8=False) -> " - "()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor q_in, Tensor kv, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " - "bool is_q_interleaved=False) -> Tensor"); - - // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate - // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" - "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " - "Tensor position_ids, Tensor cos_sin_cache, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" - "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " - "int cache_block_size, bool apply_q_norm=True) -> ()"); - - // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. - ops.def( - "fused_kimi_k3_mla_key_concat_kv_cache_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_key_concat_ds_mla_insert(" - "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " - "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_kv_concat_quant_fp8(" - "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " - "-> ()"); - ops.def( - "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" - "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " - "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " - "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " - "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - - // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the - // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " - "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); - ops.def( - "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" - "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " - "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " - "int cache_block_size, Tensor? position_ids=None, " - "Tensor? cos_sin_cache=None) -> ()"); - -#ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms_qk(" - "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " - "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " - "float eps) -> (Tensor, Tensor)"); -#endif - - // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. - ops.def( - "fused_minimax_m3_qknorm_rope_kv_insert(" - "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " - "Tensor cos_sin_cache, Tensor positions, int num_heads, " - "int num_kv_heads, int rotary_dim, float eps, " - "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " - "int num_index_heads, " - "Tensor? slot_mapping, Tensor? index_slot_mapping, " - "Tensor!? kv_cache, Tensor!? index_cache, " - "int block_size, Tensor!? q_out, Tensor!? index_q_out, " - "str kv_cache_dtype, bool skip_index_branch=False, " - "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " - "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); - -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.def( - "fused_kda_decode(" - "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " - "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor! state, Tensor! out, " - "float? lower_bound=None, Tensor? output_gate=None, " - "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.def( - "fused_gdn_decode_post_conv_mtp(" - "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " - "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " - "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " - "float scale, float norm_eps=1e-5, " - "str output_gate_activation='silu') -> ()"); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.def( - "fused_kda_prologue(" - "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " - "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " - "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " - "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " - "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " - "float scale, float lower_bound) -> ()"); - ops.def( - "fused_kda_chunk(" - "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " - "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " - "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " - "Tensor(b!)? group_state, int groups, " - "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " - "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " - "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.def( - "kimi_k3_attn_res(" - "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " - "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " - "int num_blocks, " - "int block_write_idx, float eps, " - "float output_norm_eps) -> ()"); -#endif - - // Apply repetition penalties to logits in-place. - ops.def( - "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " - "Tensor output_mask, Tensor repetition_penalties) -> ()"); - - // Optimized top-k per row operations. - ops.def( - "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " - "Tensor! indices, int numRows, int stride0, " - "int stride1, int topK) -> ()"); - - ops.def( - "top_k_per_row_decode(Tensor logits, int next_n, " - "Tensor seq_lens, Tensor! indices, " - "int numRows, int stride0, int stride1, int topK) -> ()"); - - ops.def( - "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); - -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.def( - "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); -#endif - - // Activation ops - ops.def( - "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " - "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); - ops.def("weak_ref_tensor(Tensor input) -> Tensor"); - - // Activation function used in SwiGLU. - ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); - - ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); - - // SwiGLU activation with input clamping. - // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to - // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. - ops.def( - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " - "float alpha=1.0, float beta=0.0) -> ()"); - - // SwiGLU activation with FP8 quantization. - ops.def( - "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); - - // Activation function used in GeGLU with `none` approximation. - ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); - - // Activation function used in GeGLU with `tanh` approximation. - ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); - - // FATReLU implementation. - ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); - - ops.def( - "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " - "limit=7.0) " - "-> ()"); - - // SituGLU implementation used in Kimi models. - ops.def( - "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " - "linear_beta=-1.0) -> ()"); - // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path - // (writes the fp8 down input and its float32 scale). group_size=0 -> - // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., - // d/128]. - ops.def( - "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " - "float beta=1.0, float linear_beta=-1.0, int group_size=0, " - "Tensor? num_valid_tokens=None, int topk=1) -> ()"); - ops.def( - "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " - "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); - ops.def( - "masked_moe_activation(Tensor! out, Tensor input, Tensor " - "valid_token_counts, str activation, float clamp_limit=0.0, float " - "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " - "situ_linear_beta=-1.0) -> ()"); - - // GELU implementation used in GPT-2. - ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - - // Approximate GELU implementation. - ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - - // Quick GELU implementation. - ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); - - // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 - ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); - - // Compute int8 quantized tensor for given scaling factor. - ops.def( - "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," - "Tensor? azp) -> ()"); - - // Compute int8 quantized tensor and scaling factor - ops.def( - "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " - "Tensor!? azp) -> ()"); - - // Compute FP8 quantized tensor for given scaling factor. - // Supports per-tensor, per-channel, per-token, and arbitrary 2D group - // scaling. Optional group_m/group_n specify the group shape explicitly; - // required for 1D scales to disambiguate per-channel vs per-token. - ops.def( - "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " - "int[]? group_shape=None) -> ()"); - - // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " - "-> " - "()"); - - // Compute dynamic-per-token FP8 quantized tensor and scaling factor. - ops.def( - "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " - "Tensor! scale, Tensor? scale_ub) -> " - "()"); - - // Quantized GEMM for GPTQ. - // Note: even though the C++ inferred schema is correct for this op, it seems - // to prevent the meta function registry. - ops.def( - "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " - "Tensor b_gptq_scales, bool use_exllama, bool " - "use_v2_format, int bit) " - "-> Tensor"); - - // Post processing for GPTQ. - ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); - - // Mamba selective scan kernel - ops.def( - "selective_scan_fwd(Tensor! u, Tensor! delta," - "Tensor! A, Tensor! B, Tensor! C," - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," - "bool delta_softplus," - "Tensor? query_start_loc," - "Tensor? cache_indices," - "Tensor? has_initial_state," - "Tensor! ssm_states," - "int null_block_id," - "int block_size," - "Tensor? block_idx_first_scheduled_token," - "Tensor? block_idx_last_scheduled_token," - "Tensor? initial_state_idx," - "Tensor? cu_chunk_seqlen," - "Tensor? last_chunk_indices) -> ()"); - - // LongCat n-gram embedding index kernel. All tensor args are marked mutable - // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table - // and n_gram_ids are actually written in place. - // Fused vocab-parallel embedding: gather the rows this rank owns and write - // zeros for the rest, so the following all-reduce reconstructs the full - // embedding. - ops.def( - "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " - "int org_vocab_start_index, int org_vocab_end_index, " - "int num_org_vocab_padding, int added_vocab_start_index, " - "int added_vocab_end_index) -> ()"); - ops.def( - "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " - "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " - "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " - "Tensor(f!) row_indices, Tensor(g!) column_starts, " - "Tensor(h!) n_gram_ids) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { - // LongCat n-gram embedding index kernel. - ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); - ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); - - // Per-token group quantization - ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); - ops.impl("per_token_group_fp8_quant_packed", - TORCH_BOX(&per_token_group_quant_8bit_packed)); - ops.impl("per_token_group_quant_int8", - TORCH_BOX(&per_token_group_quant_int8)); - -#ifndef USE_ROCM - // CUTLASS scaled_mm ops - ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); - ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); - ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); - ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); - ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", - TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); - ops.impl("get_cutlass_batched_moe_mm_data", - TORCH_BOX(&get_cutlass_batched_moe_mm_data)); - - // FP4/NVFP4 ops - ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); - ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); - ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); - ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_scaled_fp4_experts_quant", - TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); - ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); - // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). - // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. - - // AWQ ops - ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); - ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); - - // DSV3 fused A GEMM: conditionally compiled so impl registration is in - // source file (dsv3_fused_a_gemm.cu) -#endif - - ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); - - // Layernorm kernels (shared CUDA/ROCm) - ops.impl("rms_norm", TORCH_BOX(&rms_norm)); - ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); - - // Layernorm-quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); - ops.impl("fused_add_rms_norm_static_fp8_quant", - TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); - - // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) - ops.impl("rms_norm_dynamic_per_token_quant", - TORCH_BOX(&rms_norm_dynamic_per_token_quant)); - ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); - ops.impl("silu_and_mul_per_block_quant", - TORCH_BOX(&silu_and_mul_per_block_quant)); - - // Positional encoding kernels (shared CUDA/ROCm) - ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); - ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); - ops.impl("fused_deepseek_v4_kv_rope_insert", - TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); - ops.impl( - "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); - ops.impl("fused_kimi_k3_mla_kv_concat", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); - ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", - TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); - ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); - ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", - TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); -#ifndef USE_ROCM - ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); -#endif - ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", - TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); -#ifdef VLLM_ENABLE_FUSED_KDA_DECODE - ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); -#endif - -#ifdef VLLM_ENABLE_FUSED_GDN_DECODE - ops.impl("fused_gdn_decode_post_conv_mtp", - TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); -#endif - -#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK - ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); - ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); -#endif - -#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES - ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); -#endif - - // Sampler kernels (shared CUDA/ROCm) - ops.impl("apply_repetition_penalties_", - TORCH_BOX(&apply_repetition_penalties_)); - ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); - ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); - ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); -#ifdef VLLM_ENABLE_COOPERATIVE_TOPK - ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); -#endif - - // Activation kernels (shared CUDA/ROCm) - ops.impl("persistent_masked_m_silu_mul_quant", - TORCH_BOX(&persistent_masked_m_silu_mul_quant)); - ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); - ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); - ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); - ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); - ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); - ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); - ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); - ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); - ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); - ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); - ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); - ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); - ops.impl("gelu_new", TORCH_BOX(&gelu_new)); - ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); - ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); - ops.impl("relu_squared", TORCH_BOX(&relu_squared)); - ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); - - // INT8 quantization kernels - ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); - ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); - - // FP8 quantization kernels - ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); - ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); - ops.impl("dynamic_per_token_scaled_fp8_quant", - TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); - - // GPTQ kernels - ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); - ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); - - // Mamba kernels - ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { - ops.impl("get_cuda_view_from_cpu_tensor", - TORCH_BOX(&get_cuda_view_from_cpu_tensor)); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { - cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); - cuda_utils.def( - "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, - cuda_utils) { - cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); - cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", - TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); -} - -// These capability-check functions take only primitive args (no tensors), so -// there is no device to dispatch on. CompositeExplicitAutograd makes them -// available for all backends. This is the stable ABI equivalent of calling -// ops.impl("op_name", &func) without a dispatch key in the non-stable API. -STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { -#ifndef USE_ROCM - ops.impl("cutlass_scaled_mm_supports_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); - ops.impl("cutlass_group_gemm_supported", - TORCH_BOX(&cutlass_group_gemm_supported)); - ops.impl("cutlass_scaled_mm_supports_block_fp8", - TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); - ops.impl("cutlass_scaled_mm_supports_fp4", - TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); -#endif -} - -// Cache ops -STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { - // Swap in (out) the cache blocks from src to dst. - ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - - // Batch swap: submit all block copies in a single driver call. - ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Reshape the key and value tensors and cache them. - ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - - // Concat kv_c and k_pe and cache them. - ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - - // Grouped concat_and_cache_mla across all layers. Each layer's cache base - // pointer and optional plain-FP8 scale are read from device tensors. - ops.def( - "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," - " Tensor kv_cache_ptrs," - " Tensor slot_mapping," - " int block_size, int block_stride," - " int entry_stride," - " Tensor? kv_scales=None," - " str kv_cache_dtype='auto') -> ()"); - -#ifndef USE_ROCM - ops.def( - "hisparse_resolve_residency(Tensor host_cache," - " Tensor! hot_cache," - " Tensor hot_block_table," - " Tensor global_indices," - " Tensor! hot_indices," - " Tensor! device_global_indices," - " Tensor! lru_slots," - " Tensor? request_state_indices," - " int region_stride," - " int max_union_rows," - " Tensor(a!)? miss_mask=None," - " Tensor(b!)? stats=None," - " Tensor(c!)? attention_indices=None," - " int attention_block_stride=0," - " Tensor? request_ids=None," - " Tensor? source_block_table=None," - " int source_block_size=0," - " Tensor(d!)? resolved_global_indices=None," - " Tensor(e!)? valid_counts=None," - " Tensor(f!)? swap_host_physical_rows=None," - " Tensor(g!)? swap_device_physical_rows=None," - " Tensor(h!)? swap_counts=None," - " Tensor? resident_block_table=None," - " int resident_block_size=0," - " int resident_null_block=0) -> ()"); - - ops.def( - "hisparse_invalidate_written_slots(Tensor! device_global_indices," - " Tensor request_state_indices," - " Tensor req_id_per_token," - " Tensor written_slots) -> ()"); - - ops.def( - "hisparse_gather_plan(Tensor host_cache," - " Tensor! hot_cache," - " Tensor global_indices," - " Tensor hot_indices," - " Tensor miss_mask," - " Tensor? request_state_indices," - " Tensor(a!)? attention_indices=None," - " int attention_block_stride=0) -> ()"); - - ops.def( - "hisparse_gather_compact(Tensor host_cache," - " Tensor! hot_cache," - " Tensor miss_global_indices," - " Tensor miss_hot_indices," - " Tensor miss_counts) -> ()"); - -#endif // !USE_ROCM - - // Rotate Q and K, then write to kv cache for MLA - ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - - // Convert the key and value cache to fp8 data type. - ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - - ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " - "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " - "device_row_ids=None) -> ()"); - - ops.def( - "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); - - ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - - ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - - ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); -} - -STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { - custom_ar.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool fully_connected) -> int"); - custom_ar.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes) -> ()"); - custom_ar.def("dispose(int fa) -> ()"); - custom_ar.def("meta_size() -> int"); - custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); - custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); - custom_ar.def( - "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); - custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); - custom_ar.def("free_shared_buffer(int ptr) -> ()"); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { - custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); - custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { - custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { - custom_ar.impl("dispose", TORCH_BOX(&dispose)); - custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); - custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); - custom_ar.impl("get_graph_buffer_ipc_meta", - TORCH_BOX(&get_graph_buffer_ipc_meta)); - custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); - custom_ar.impl("allocate_shared_buffer_and_handle", - TORCH_BOX(&allocate_shared_buffer_and_handle)); - custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { - ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); -} - -STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { - ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); - ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); - ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); - ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); - ops.impl("concat_and_cache_mla_grouped", - TORCH_BOX(&concat_and_cache_mla_grouped)); - -#ifndef USE_ROCM - ops.impl("hisparse_resolve_residency", - TORCH_BOX(&hisparse_resolve_residency)); - ops.impl("hisparse_invalidate_written_slots", - TORCH_BOX(&hisparse_invalidate_written_slots)); - ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); - ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); -#endif // !USE_ROCM - ops.impl("concat_and_cache_mla_rope_fused", - TORCH_BOX(&concat_and_cache_mla_rope_fused)); - ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); - ops.impl("gather_and_maybe_dequant_cache", - TORCH_BOX(&gather_and_maybe_dequant_cache)); - ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); - ops.impl("cp_gather_and_upconvert_fp8_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); - ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", - TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); - ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); - ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); - ops.impl("cp_gather_indexer_k_quant_cache", - TORCH_BOX(&cp_gather_indexer_k_quant_cache)); -} - -REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/libtorch_stable/torch_utils.h b/csrc/libtorch_stable/torch_utils.h deleted file mode 100644 index f02346739e0a..000000000000 --- a/csrc/libtorch_stable/torch_utils.h +++ /dev/null @@ -1,88 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include - -// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. -#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ - STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) - -// Device properties cache for stable ABI compatibility. -// Uses raw CUDA/HIP APIs instead of ATen functions. -// Using inline ensures a single instance across all translation units. -inline std::deque device_flags; -inline std::vector device_properties; -inline std::once_flag vectors_init_flag; - -inline void do_init_device_vectors() { - int device_count; - cudaError_t err = cudaGetDeviceCount(&device_count); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + - std::string(cudaGetErrorString(err))); - } - device_flags.resize(device_count); - device_properties.resize(device_count); -} - -inline void initDeviceVectors() { - std::call_once(vectors_init_flag, do_init_device_vectors); -} - -inline void initDeviceProperty(int device_index) { - cudaDeviceProp device_prop{}; - cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + - std::string(cudaGetErrorString(err))); - } - device_properties[device_index] = device_prop; -} - -// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). -// Caches results per device so cudaGetDeviceProperties is called at most once -// per device. -inline cudaDeviceProp* get_device_prop() { - initDeviceVectors(); - int device_index; - cudaError_t err = cudaGetDevice(&device_index); - if (err != cudaSuccess) { - STD_TORCH_CHECK( - false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); - } - STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < - device_properties.size(), - "CUDA device index " + std::to_string(device_index) + - " out of range [0, " + - std::to_string(device_properties.size()) + ")"); - - std::call_once(device_flags[device_index], initDeviceProperty, device_index); - return &device_properties[device_index]; -} - -// Utility to get the current CUDA stream for a given device using stable APIs. -// Returns a cudaStream_t for use in kernel launches. -inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { - void* stream_ptr = nullptr; - TORCH_ERROR_CODE_CHECK( - aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); - return reinterpret_cast(stream_ptr); -} - -// Utility to get the current cuBLAS handle using stable APIs. -inline cublasHandle_t get_current_cuda_blas_handle() { - void* blas_handle_ptr = nullptr; - TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); - return reinterpret_cast(blas_handle_ptr); -} diff --git a/csrc/libtorch_stable/mamba/selective_scan.h b/csrc/mamba/selective_scan.h similarity index 100% rename from csrc/libtorch_stable/mamba/selective_scan.h rename to csrc/mamba/selective_scan.h diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu b/csrc/mamba/selective_scan_fwd.cu similarity index 99% rename from csrc/libtorch_stable/mamba/selective_scan_fwd.cu rename to csrc/mamba/selective_scan_fwd.cu index c17b06a79bd0..cfc1f038f131 100644 --- a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu +++ b/csrc/mamba/selective_scan_fwd.cu @@ -14,7 +14,7 @@ #endif #include "selective_scan.h" -#include "static_switch.h" +#include "../dispatch_utils.h" template; constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); dim3 grid(params.batch, params.dim / kNRows); diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu similarity index 99% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.cu rename to csrc/minimax_reduce_rms_kernel.cu index 58d61b353d6d..5a3da3dadffc 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -27,7 +27,7 @@ #include #include -#include "cuda_compat.h" +#include "core/cuda_compat.h" #include "minimax_reduce_rms_kernel.h" #include diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h similarity index 100% rename from csrc/libtorch_stable/minimax_reduce_rms_kernel.h rename to csrc/minimax_reduce_rms_kernel.h diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/moe/grouped_topk_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/grouped_topk_kernels.cu rename to csrc/moe/grouped_topk_kernels.cu index 852b2b468adf..ea89aa38a1f0 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/moe/grouped_topk_kernels.cu @@ -22,7 +22,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore b/csrc/moe/marlin_moe_wna16/.gitignore similarity index 73% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore rename to csrc/moe/marlin_moe_wna16/.gitignore index 7dc482a89466..ba805f9250ec 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/.gitignore +++ b/csrc/moe/marlin_moe_wna16/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/moe/marlin_moe_wna16/kernel.h similarity index 95% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h rename to csrc/moe/marlin_moe_wna16/kernel.h index 25c57f3b2ddf..1fd00c12d1bd 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/csrc/moe/marlin_moe_wna16/marlin_template.h similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h rename to csrc/moe/marlin_moe_wna16/marlin_template.h index aac9bd157f6a..e2d80ce6bfbe 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "libtorch_stable/quantization/marlin/marlin.cuh" -#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" -#include "libtorch_stable/quantization/marlin/dequant.h" -#include "libtorch_stable/quantization/marlin/marlin_mma.h" +#include "quantization/marlin/marlin.cuh" +#include "quantization/marlin/marlin_dtypes.cuh" +#include "quantization/marlin/dequant.h" +#include "quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu b/csrc/moe/marlin_moe_wna16/ops.cu similarity index 99% rename from csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu rename to csrc/moe/marlin_moe_wna16/ops.cu index 48de07dc753d..8336b46e1590 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/ops.cu +++ b/csrc/moe/marlin_moe_wna16/ops.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/moe/moeTopKFuncs.cuh similarity index 100% rename from csrc/libtorch_stable/moe/moeTopKFuncs.cuh rename to csrc/moe/moeTopKFuncs.cuh diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_align_sum_kernels.cu rename to csrc/moe/moe_align_sum_kernels.cu index 5c9a522feb7b..a2424f833908 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/moe/moe_align_sum_kernels.cu @@ -8,11 +8,11 @@ #include #include -#include "../../cuda_compat.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/torch_utils.h" +#include "core/cuda_compat.h" +#include "core/utils.hpp" +#include "dispatch_utils.h" +#include "quantization/vectorization.cuh" +#include "torch_utils.h" #define CEILDIV(x, y) (((x) + (y) - 1) / (y)) diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/moe/moe_ops.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_ops.h rename to csrc/moe/moe_ops.h diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu rename to csrc/moe/moe_permute_unpermute_op.cu index 957cd60e49de..31c937b2a2bb 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/moe/moe_permute_unpermute_op.cu @@ -7,8 +7,8 @@ #include #include "core/registration.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" -#include "libtorch_stable/torch_utils.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "torch_utils.h" #include diff --git a/csrc/libtorch_stable/moe/moe_wna16.cu b/csrc/moe/moe_wna16.cu similarity index 99% rename from csrc/libtorch_stable/moe/moe_wna16.cu rename to csrc/moe/moe_wna16.cu index 9345a7c9f789..a8971041d041 100644 --- a/csrc/libtorch_stable/moe/moe_wna16.cu +++ b/csrc/moe/moe_wna16.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "moe_wna16_utils.h" #define DIVIDE(x, size) (((x) + (size) - 1) / (size)) diff --git a/csrc/libtorch_stable/moe/moe_wna16_utils.h b/csrc/moe/moe_wna16_utils.h similarity index 100% rename from csrc/libtorch_stable/moe/moe_wna16_utils.h rename to csrc/moe/moe_wna16_utils.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h b/csrc/moe/permute_unpermute_kernels/dispatch.h similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/dispatch.h rename to csrc/moe/permute_unpermute_kernels/dispatch.h diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu similarity index 98% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index 8d8bdf8baaf6..56c1e2bddf48 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -1,7 +1,7 @@ #include #include -#include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" +#include "moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" // moe_permute kernels require at least CUDA 12.0 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h similarity index 97% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index 53afad2f5d41..1496eccb8d4d 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -10,7 +10,7 @@ #include "cutlass/array.h" #include "cutlass/numeric_size.h" -#include "libtorch_stable/moe/permute_unpermute_kernels/dispatch.h" +#include "moe/permute_unpermute_kernels/dispatch.h" template inline T* get_ptr(torch::stable::Tensor& t) { diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl similarity index 100% rename from csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl rename to csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softmax_kernels.cu rename to csrc/moe/topk_softmax_kernels.cu index 098644a1b901..0e462770177c 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -24,9 +24,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/moe/topk_softplus_sqrt_kernels.cu similarity index 99% rename from csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu rename to csrc/moe/topk_softplus_sqrt_kernels.cu index fbb07216d202..237f76334abd 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/moe/topk_softplus_sqrt_kernels.cu @@ -25,9 +25,9 @@ #include #include -#include "../../cuda_compat.h" +#include "core/cuda_compat.h" #include "../cub_helpers.h" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #ifndef USE_ROCM #include #include diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp similarity index 100% rename from csrc/libtorch_stable/moe/torch_bindings.cpp rename to csrc/moe/torch_bindings.cpp diff --git a/csrc/libtorch_stable/ngram_embedding_kernels.cu b/csrc/ngram_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/ngram_embedding_kernels.cu rename to csrc/ngram_embedding_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache.h b/csrc/nvfp4_ds_mla_cache.h similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache.h rename to csrc/nvfp4_ds_mla_cache.h diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/nvfp4_ds_mla_cache_kernels.cu similarity index 100% rename from csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu rename to csrc/nvfp4_ds_mla_cache_kernels.cu diff --git a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu b/csrc/nvfp4_kv_cache_kernels.cu similarity index 99% rename from csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu rename to csrc/nvfp4_kv_cache_kernels.cu index bb11e474a6a4..341c4cdd9396 100644 --- a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu +++ b/csrc/nvfp4_kv_cache_kernels.cu @@ -15,10 +15,10 @@ // - reciprocal_approximate_ftz() for fast reciprocal #define NVFP4_ENABLE_ELTS16 1 -#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" +#include "quantization/fp4/nvfp4_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "dispatch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/ops.h b/csrc/ops.h index d7ee2d080c48..3266d28b1edb 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1,72 +1,822 @@ #pragma once +#include +#include +#include + #include #include -#include -#include +#include -#include "core/scalar_type.hpp" +#include -#include +inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) { + // Ensure tensor is on CUDA + STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device"); -// rms_norm and fused_add_rms_norm declarations also exist in -// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here -// because the CPU build still uses these torch::Tensor declarations. -void rms_norm(torch::Tensor& out, torch::Tensor& input, - std::optional weight, double epsilon); + // Get the raw data pointer + void* data_ptr = tensor.mutable_data_ptr(); -void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, - std::optional weight, double epsilon); + /// Create a new tensor from the raw data pointer + return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(), + tensor.device(), tensor.scalar_type()); +} -// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable -// ABI for CUDA). It remains here because the CPU build still uses these -// torch::Tensor declarations. -void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, - std::optional key, int64_t head_size, - torch::Tensor& cos_sin_cache, bool is_neox, - int64_t rope_dim_offset, bool inverse); +void per_token_group_quant_fp8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double fp8_min, + double fp8_max, bool scale_ue8m0, + bool dummy_is_scale_transposed, + bool dummy_is_tma_aligned); -void silu_and_mul(torch::Tensor& out, torch::Tensor& input); +// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. +void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); -void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit, - double alpha = 1.0, double beta = 0.0); +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max); -void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); +#ifndef USE_ROCM +bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); +bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); +bool cutlass_group_gemm_supported(int64_t cuda_device_capability); -void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + std::optional const& bias); -void gelu_tanh(torch::Tensor& out, torch::Tensor& input); +void cutlass_moe_mm(torch::stable::Tensor& out_tensors, + torch::stable::Tensor const& a_tensors, + torch::stable::Tensor const& b_tensors, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& expert_offsets, + torch::stable::Tensor const& problem_sizes, + torch::stable::Tensor const& a_strides, + torch::stable::Tensor const& b_strides, + torch::stable::Tensor const& c_strides, bool per_act_token, + bool per_out_ch); -void gelu_new(torch::Tensor& out, torch::Tensor& input); +void cutlass_scaled_mm_azp(torch::stable::Tensor& out, + torch::stable::Tensor const& a, + torch::stable::Tensor const& b, + torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, + torch::stable::Tensor const& azp_adj, + std::optional const& azp, + std::optional const& bias); -void gelu_fast(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_moe_mm_data( + const torch::stable::Tensor& topk_ids, + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + torch::stable::Tensor& input_permutation, + torch::stable::Tensor& output_permutation, const int64_t num_experts, + const int64_t n, const int64_t k, + const std::optional& blockscale_offsets, + const bool is_gated); -void gelu_quick(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( + const torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k, + const bool swap_ab); -void relu_squared(torch::Tensor& out, torch::Tensor& input); +void get_cutlass_batched_moe_mm_data( + torch::stable::Tensor& expert_offsets, + torch::stable::Tensor& problem_sizes1, + torch::stable::Tensor& problem_sizes2, + const torch::stable::Tensor& expert_num_tokens, + const int64_t num_local_experts, const int64_t padded_m, const int64_t n, + const int64_t k); -void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor const& scale, - std::optional const& azp); +// FP4/NVFP4 ops +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability); -void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, - torch::Tensor& scales, - std::optional const& azp); +void cutlass_scaled_fp4_mm(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); -torch::Tensor dynamic_4bit_int_moe_cpu( - torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, - int64_t intermediate_size, int64_t group_size, - bool apply_router_weight_on_input, int64_t activation_kind); +void cutlass_fp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& alphas, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); -using fptr_t = int64_t; -#ifdef USE_ROCM -fptr_t init_custom_qr(int64_t rank, int64_t world_size, - std::optional qr_max_size = std::nullopt); -void qr_destroy(fptr_t _fa); -torch::Tensor qr_get_handle(fptr_t _fa); -void qr_open_handles(fptr_t _fa, const std::vector& handles); -void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - int64_t quant_level, bool cast_bf2half = false); -int64_t qr_max_size(); +std::tuple scaled_fp4_quant_func( + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout); + +void scaled_fp4_quant_out(torch::stable::Tensor const& input, + torch::stable::Tensor const& input_scale, + bool is_sf_swizzled_layout, + torch::stable::Tensor& output, + torch::stable::Tensor& output_scale); + +void scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_scaled_fp4_experts_quant( + torch::stable::Tensor& output, torch::stable::Tensor& output_scale, + torch::stable::Tensor const& input, + torch::stable::Tensor const& input_global_scale, + torch::stable::Tensor const& input_offset_by_experts, + torch::stable::Tensor const& output_scale_offset_by_experts); + +void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out, + torch::stable::Tensor& output_block_scale, + torch::stable::Tensor& input, + torch::stable::Tensor& input_global_scale); + +void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, + const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& a_blockscale, + const torch::stable::Tensor& b_blockscales, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& sf_offsets); + +// AWQ ops +torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats, + torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters); + +torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, + torch::stable::Tensor _scaling_factors, + torch::stable::Tensor _zeros, + int64_t split_k_iters, int64_t thx, + int64_t thy); + +// DSV3 fused A GEMM: conditionally compiled so declaration and impl +// registration are in the source file (dsv3_fused_a_gemm.cu) + +#endif + +// CPU tensor -> CUDA UVA view (shared CUDA/ROCm) +torch::stable::Tensor get_cuda_view_from_cpu_tensor( + torch::stable::Tensor& cpu_tensor); + +// Attention kernels (shared CUDA/ROCm) +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); + +torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, + bool inplace); + +// Layernorm kernels (shared CUDA/ROCm) +void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input, + std::optional weight, double epsilon); + +void fused_add_rms_norm(torch::stable::Tensor& input, + torch::stable::Tensor& residual, + std::optional weight, + double epsilon); + +// Layernorm-quant kernels (shared CUDA/ROCm) +void rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, double epsilon); + +void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& residual, + torch::stable::Tensor& weight, + torch::stable::Tensor& scale, + double epsilon); + +// Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) +void rms_norm_dynamic_per_token_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, torch::stable::Tensor& scales, + double const var_epsilon, std::optional scale_ub, + std::optional residual); + +void rms_norm_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor& scales, + double const var_epsilon, + std::optional scale_ub, + std::optional residual, + int64_t group_size, bool is_scale_transposed); + +void silu_and_mul_per_block_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed); + +// Positional encoding kernels (shared CUDA/ROCm) +void rotary_embedding(torch::stable::Tensor& positions, + torch::stable::Tensor& query, + std::optional key, + int64_t head_size, torch::stable::Tensor& cos_sin_cache, + bool is_neox, int64_t rope_dim_offset, bool inverse); + +void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, + int64_t num_heads_k, int64_t num_heads_v, + int64_t head_dim, double eps, + torch::stable::Tensor& q_weight, + torch::stable::Tensor& k_weight, + torch::stable::Tensor& cos_sin_cache, bool is_neox, + torch::stable::Tensor& position_ids, + int64_t forced_token_heads_per_warp); + +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, double eps, + int64_t cache_block_size, bool apply_q_norm); + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor& k_out); + +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& v, + torch::stable::Tensor& k_fp8, + torch::stable::Tensor& v_fp8); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& kv, + torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& fp8_scale, + torch::stable::Tensor const& q_fp8_scale_inv, double eps, + int64_t cache_block_size, bool apply_q_norm); + +#ifndef USE_ROCM +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif + +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV / +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs +// the index branch and scatters k/v/index_k into their paged caches. +void fused_minimax_m3_qknorm_rope_kv_insert( + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, int64_t num_heads, + int64_t num_kv_heads, int64_t rotary_dim, double eps, + std::optional index_q_norm_weight, + std::optional index_k_norm_weight, + int64_t num_index_heads, std::optional slot_mapping, + std::optional index_slot_mapping, + std::optional kv_cache, + std::optional index_cache, int64_t block_size, + std::optional q_out, + std::optional index_q_out, + const std::string& kv_cache_dtype, bool skip_index_branch, + std::optional q_fp8_out, double q_fp8_scale, + std::optional kv_k_scale, + std::optional kv_v_scale); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); + +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE +void fused_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, + double scale, double norm_eps, const std::string& output_gate_activation); + +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK +void fused_kda_prologue( + torch::stable::Tensor const& q, torch::stable::Tensor const& k, + torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg, + torch::stable::Tensor& w, torch::stable::Tensor& u, + torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk, + torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_indices, + std::optional conv_weight, + std::optional conv_state, + std::optional conv_state_indices, + std::optional conv_has_initial_state, double scale, + double lower_bound); + +void fused_kda_chunk( + torch::stable::Tensor const& qg, torch::stable::Tensor const& w, + torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t, + torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay, + std::optional initial_state, + std::optional final_state, + torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& chunk_offsets, double scale, + std::optional group_state, int64_t groups, + std::optional checkpoint_state, + std::optional checkpoint_offsets, + std::optional checkpoint_state_indices, + std::optional state_cache, + std::optional state_indices, + std::optional has_initial_state); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + std::optional delta, + torch::stable::Tensor& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + std::optional output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + int64_t block_write_idx, double eps, + double output_norm_eps); +#endif + +// Sampler kernels (shared CUDA/ROCm) +void apply_repetition_penalties_( + torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, + const torch::stable::Tensor& output_mask, + const torch::stable::Tensor& repetition_penalties); + +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK +void cooperative_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); #endif + +void selective_scan_fwd( + const torch::stable::Tensor& u, const torch::stable::Tensor& delta, + const torch::stable::Tensor& A, const torch::stable::Tensor& B, + const torch::stable::Tensor& C, + const std::optional& D_, + const std::optional& z_, + const std::optional& delta_bias_, + bool delta_softplus, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, + const torch::stable::Tensor& ssm_states, int64_t null_block_id, + int64_t block_size, + const std::optional& block_idx_first_scheduled_token, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_idx, + const std::optional& cu_chunk_seqlen, + const std::optional& last_chunk_indices); + +using fptr_t = int64_t; +fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, + torch::stable::Tensor& rank_data, int64_t rank, + bool fully_connected); +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t multicast_buffer, + int64_t stage_sz_bytes, int64_t block_limit); +void dispose(fptr_t _fa); +int64_t meta_size(); +void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa); +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets); +std::tuple allocate_shared_buffer_and_handle( + int64_t size); +int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +void free_shared_buffer(int64_t buffer); + +// Activation kernels (shared CUDA/ROCm) +void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void silu_and_mul_clamp(torch::stable::Tensor& out, + torch::stable::Tensor& input, double limit, + double alpha = 1.0, double beta = 0.0); + +void silu_and_mul_quant(torch::stable::Tensor& out, + torch::stable::Tensor& input, + torch::stable::Tensor& scale); + +void persistent_masked_m_silu_mul_quant( + const torch::stable::Tensor& input, // (E, T, 2*H) + const torch::stable::Tensor& tokens_per_expert, // (E) + torch::stable::Tensor& y_q, // (E, T, H) [OUT] + torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT] + bool use_ue8m0); + +void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_tanh_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input); +void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double threshold); +void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void situ_and_mul_quant( + torch::stable::Tensor& out, torch::stable::Tensor& scale, + torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0, + int64_t group_size = 0, + std::optional num_valid_tokens = std::nullopt, + int64_t topk = 1); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); +void masked_moe_activation(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& valid_token_counts, + const std::string& activation, + double clamp_limit = 0.0, double alpha = 1.0, + double beta = 0.0, double situ_beta = 1.0, + double situ_linear_beta = -1.0); +void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); +void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); + +void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input); + +// INT8 quantization kernels (shared CUDA/ROCm) +void static_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional const& azp); + +void dynamic_scaled_int8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scales, + std::optional const& azp); + +// FP8 quantization kernels (shared CUDA/ROCm) +void static_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor const& scale, + std::optional group_shape = + std::nullopt); + +void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, + torch::stable::Tensor const& input, + torch::stable::Tensor& scale); + +void dynamic_per_token_scaled_fp8_quant( + torch::stable::Tensor& out, torch::stable::Tensor const& input, + torch::stable::Tensor& scale, + std::optional const& scale_ub); + +// GPTQ kernels (shared CUDA/ROCm) +torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, + torch::stable::Tensor b_q_weight, + torch::stable::Tensor b_gptq_qzeros, + torch::stable::Tensor b_gptq_scales, + bool use_exllama, bool use_v2_format, + int64_t bit); + +void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, int64_t entry_stride, + std::optional kv_scales, + const std::string& kv_cache_dtype); + +#ifndef USE_ROCM +// HiSparse kernels use raw PTX in the row copy; CUDA-only. +void hisparse_resolve_residency( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& hot_block_table, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor& hot_indices, + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor& lru_slots, + std::optional const& request_state_indices, + int64_t region_stride, int64_t max_union_rows, + std::optional const& miss_mask, + std::optional const& stats, + std::optional const& attention_indices, + int64_t attention_block_stride, + std::optional const& request_ids, + std::optional const& source_block_table, + int64_t source_block_size, + std::optional const& resolved_global_indices, + std::optional const& valid_counts, + std::optional const& swap_host_physical_rows, + std::optional const& swap_device_physical_rows, + std::optional const& swap_counts, + std::optional const& resident_block_table, + int64_t resident_block_size, int64_t resident_null_block); + +void hisparse_invalidate_written_slots( + torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& request_state_indices, + torch::stable::Tensor const& req_id_per_token, + torch::stable::Tensor const& written_slots); + +void hisparse_gather_plan( + torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& global_indices, + torch::stable::Tensor const& hot_indices, + torch::stable::Tensor const& miss_mask, + std::optional const& request_state_indices, + std::optional const& attention_indices, + int64_t attention_block_stride); + +void hisparse_gather_compact(torch::stable::Tensor const& host_cache, + torch::stable::Tensor& hot_cache, + torch::stable::Tensor const& miss_global_indices, + torch::stable::Tensor const& miss_hot_indices, + torch::stable::Tensor const& miss_counts); + +#endif // !USE_ROCM + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size, + std::optional seq_starts = std::nullopt, + std::optional host_cache = std::nullopt, + std::optional host_row_ids = std::nullopt, + std::optional device_row_ids = std::nullopt); + +// Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace +void cp_gather_and_upconvert_nvfp4_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 352] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] + +// Fused vocab-parallel embedding lookup (see +// vocab_parallel_embedding_kernels.cu). +void vocab_parallel_embedding( + torch::stable::Tensor& out, const torch::stable::Tensor& input_ids, + const torch::stable::Tensor& weight, int64_t org_vocab_start_index, + int64_t org_vocab_end_index, int64_t num_org_vocab_padding, + int64_t added_vocab_start_index, int64_t added_vocab_end_index); + +// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu). +void ngram_compute_n_gram_ids( + int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, + torch::stable::Tensor& ne_mods, + torch::stable::Tensor& exclusive_ne_embedder_size_sums, + torch::stable::Tensor& exclusive_req_len_sums, + torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, + torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/persistent_topk.cuh similarity index 100% rename from csrc/libtorch_stable/persistent_topk.cuh rename to csrc/persistent_topk.cuh diff --git a/csrc/libtorch_stable/pos_encoding_kernels.cu b/csrc/pos_encoding_kernels.cu similarity index 99% rename from csrc/libtorch_stable/pos_encoding_kernels.cu rename to csrc/pos_encoding_kernels.cu index 74af743b0961..bad616fc8748 100644 --- a/csrc/libtorch_stable/pos_encoding_kernels.cu +++ b/csrc/pos_encoding_kernels.cu @@ -1,6 +1,6 @@ #include "torch_utils.h" -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/activation_kernels.cu b/csrc/quantization/activation_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/activation_kernels.cu rename to csrc/quantization/activation_kernels.cu index 822a41969e7e..9d7da85a22ee 100644 --- a/csrc/libtorch_stable/quantization/activation_kernels.cu +++ b/csrc/quantization/activation_kernels.cu @@ -1,10 +1,10 @@ -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/core/math.hpp" -#include "cuda_compat.h" -#include "libtorch_stable/dispatch_utils.h" +#include "core/utils.hpp" +#include "core/cuda_compat.h" +#include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" #ifndef USE_ROCM diff --git a/csrc/libtorch_stable/quantization/awq/dequantize.cuh b/csrc/quantization/awq/dequantize.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/awq/dequantize.cuh rename to csrc/quantization/awq/dequantize.cuh diff --git a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu b/csrc/quantization/awq/gemm_kernels.cu similarity index 99% rename from csrc/libtorch_stable/quantization/awq/gemm_kernels.cu rename to csrc/quantization/awq/gemm_kernels.cu index 43301f0fe206..a0be65704544 100644 --- a/csrc/libtorch_stable/quantization/awq/gemm_kernels.cu +++ b/csrc/quantization/awq/gemm_kernels.cu @@ -9,9 +9,9 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023} #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/quantization/awq/dequantize.cuh" +#include "quantization/awq/dequantize.cuh" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh rename to csrc/quantization/cutlass_w4a8/get_group_starts.cuh index 721761791251..aff0502550e1 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/get_group_starts.cuh +++ b/csrc/quantization/cutlass_w4a8/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu index 1d7eb04093a2..82d6b2ec8a89 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu @@ -16,12 +16,12 @@ // vllm includes #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/common.hpp" #include "get_group_starts.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include "w4a8_utils.cuh" namespace vllm::cutlass_w4a8_moe { diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu rename to csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu index 8f19394b3904..c4e96e8aee22 100644 --- a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_mm_entry.cu +++ b/csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu @@ -5,8 +5,8 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "torch_utils.h" +#include "cutlass_extensions/torch_utils.hpp" #include "w4a8_utils.cuh" #include "cutlass/cutlass.h" @@ -21,8 +21,8 @@ #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/mixed_dtype_utils.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" #include diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu b/csrc/quantization/cutlass_w4a8/w4a8_utils.cu similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cu rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cu diff --git a/csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh b/csrc/quantization/cutlass_w4a8/w4a8_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/cutlass_w4a8/w4a8_utils.cuh rename to csrc/quantization/cutlass_w4a8/w4a8_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu rename to csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 6238a27191b7..2064775111fb 100644 --- a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -25,7 +25,7 @@ #include #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu index e4d2f2201250..3234717f6be1 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -8,11 +8,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu rename to csrc/quantization/fp4/mxfp4_experts_quant.cu index 086c07acf002..4c6875f0623a 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -25,9 +25,9 @@ #include #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "torch_utils.h" +#include "dispatch_utils.h" +#include "cutlass_extensions/common.hpp" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" @@ -41,7 +41,7 @@ static_assert(CVT_FP4_ELTS_PER_THREAD == 16, #define VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED 0 #endif -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu rename to csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu index ea3344406352..8e8a30e956d7 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu +++ b/csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu @@ -16,11 +16,11 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/quantization/fp4/nvfp4_experts_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu rename to csrc/quantization/fp4/nvfp4_experts_quant.cu index 6296f0d901cb..55c722f1de98 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/quantization/fp4/nvfp4_experts_quant.cu @@ -15,8 +15,8 @@ */ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include @@ -26,7 +26,7 @@ #include "cuda_utils.h" #include "nvfp4_utils.cuh" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu b/csrc/quantization/fp4/nvfp4_quant_entry.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu rename to csrc/quantization/fp4/nvfp4_quant_entry.cu index e1e7e7a74da2..ace7bbbc10c7 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/quantization/fp4/nvfp4_quant_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "nvfp4_utils.cuh" #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/quantization/fp4/nvfp4_quant_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu rename to csrc/quantization/fp4/nvfp4_quant_kernels.cu index 7ec418ee381b..13a3026cc889 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_quant_kernels.cu @@ -21,12 +21,12 @@ #include -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" -#include "libtorch_stable/launch_bounds_utils.h" +#include "launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu similarity index 97% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu index bfb526fcd401..6a60fe6f6044 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D, diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu index af9f24a70e0b..f925a599db46 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu similarity index 98% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu rename to csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu index 3a45ede8dfd5..02ded13c7835 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu @@ -16,9 +16,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" #include "cutlass/cutlass.h" @@ -31,8 +31,7 @@ #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/quantization/fp4/nvfp4_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh rename to csrc/quantization/fp4/nvfp4_utils.cuh diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu rename to csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 56dd47038727..f864d6037ef3 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,7 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" -#include "../../../core/batch_invariant.hpp" +#include "core/utils.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu b/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu rename to csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu diff --git a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh rename to csrc/quantization/fused_kernels/layernorm_utils.cuh index 9a3d221562e9..d633cb9a6e57 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -4,12 +4,12 @@ * __device__ layernorm utilities. */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" #include "quantization/utils.cuh" #include "quant_conversions.cuh" #include "../../cub_helpers.h" -#include "../../../cuda_compat.h" +#include "core/cuda_compat.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh b/csrc/quantization/fused_kernels/quant_conversions.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh rename to csrc/quantization/fused_kernels/quant_conversions.cuh index a81a0b9feb2f..1953e90f5de0 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/quant_conversions.cuh +++ b/csrc/quantization/fused_kernels/quant_conversions.cuh @@ -4,9 +4,9 @@ * __device__ helper functions to deal with float -> quant datatype conversion */ -#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/vectorization.cuh" // TODO(luka/varun):refactor common.cuh to use this file instead -#include "../../../quantization/w8a8/fp8/common.cuh" +#include "quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/gptq/compat.cuh b/csrc/quantization/gptq/compat.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/compat.cuh rename to csrc/quantization/gptq/compat.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/matrix_view.cuh b/csrc/quantization/gptq/matrix_view.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/matrix_view.cuh rename to csrc/quantization/gptq/matrix_view.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/q_gemm.cu b/csrc/quantization/gptq/q_gemm.cu similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/q_gemm.cu rename to csrc/quantization/gptq/q_gemm.cu diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_2.cuh b/csrc/quantization/gptq/qdq_2.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_2.cuh rename to csrc/quantization/gptq/qdq_2.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_3.cuh b/csrc/quantization/gptq/qdq_3.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_3.cuh rename to csrc/quantization/gptq/qdq_3.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_4.cuh b/csrc/quantization/gptq/qdq_4.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_4.cuh rename to csrc/quantization/gptq/qdq_4.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_8.cuh b/csrc/quantization/gptq/qdq_8.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_8.cuh rename to csrc/quantization/gptq/qdq_8.cuh diff --git a/csrc/libtorch_stable/quantization/gptq/qdq_util.cuh b/csrc/quantization/gptq/qdq_util.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/gptq/qdq_util.cuh rename to csrc/quantization/gptq/qdq_util.cuh diff --git a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu similarity index 99% rename from csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu rename to csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu index a5d199db2737..2575b20b8eef 100644 --- a/csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu +++ b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu @@ -11,8 +11,8 @@ Redistribution and use in source and binary forms, with or without modification, THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ***********/ -#include "libtorch_stable/torch_utils.h" -#include "libtorch_stable/dispatch_utils.h" +#include "torch_utils.h" +#include "dispatch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/Readme.md b/csrc/quantization/machete/Readme.md similarity index 100% rename from csrc/libtorch_stable/quantization/machete/Readme.md rename to csrc/quantization/machete/Readme.md diff --git a/csrc/libtorch_stable/quantization/machete/generate.py b/csrc/quantization/machete/generate.py similarity index 99% rename from csrc/libtorch_stable/quantization/machete/generate.py rename to csrc/quantization/machete/generate.py index 11a5bbdd13c1..60a49c587fa4 100644 --- a/csrc/libtorch_stable/quantization/machete/generate.py +++ b/csrc/quantization/machete/generate.py @@ -496,7 +496,8 @@ def prepacked_type_key(prepack_type: PrepackTypeConfig): def generate(): - # See csrc/quantization/machete/Readme.md, the Codegeneration for more info + # See csrc/quantization/machete/Readme.md, the + # Codegeneration for more info # about how this works SCRIPT_DIR = os.path.dirname(__file__) diff --git a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh b/csrc/quantization/machete/machete_collective_builder.cuh similarity index 94% rename from csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh rename to csrc/quantization/machete/machete_collective_builder.cuh index ba8da0af2c31..ee825583dee1 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh +++ b/csrc/quantization/machete/machete_collective_builder.cuh @@ -1,6 +1,6 @@ #pragma once -#include "libtorch_stable/cutlass_extensions/vllm_collective_builder.cuh" +#include "cutlass_extensions/vllm_collective_builder.cuh" #include "machete_mainloop.cuh" namespace cutlass::gemm::collective { diff --git a/csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh b/csrc/quantization/machete/machete_interleaving_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh rename to csrc/quantization/machete/machete_interleaving_utils.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh b/csrc/quantization/machete/machete_mainloop.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh rename to csrc/quantization/machete/machete_mainloop.cuh index a0d1ded8876a..9127a701a0fd 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh +++ b/csrc/quantization/machete/machete_mainloop.cuh @@ -17,8 +17,9 @@ // before feeding them into the tensor cores in registers, we need the weight // matrix to be A. To achieve this we compute the transpose of Y = XW^t as // Y^t = W^tX^t. This is mostly done outside of this file in -// csrc/quantization/machete/machete_mm_kernel.cuh, but this why A is the -// quantized/narrow type and has the prepacked layout despite the API being: +// csrc/quantization/machete/machete_mm_kernel.cuh, but this why +// A is the quantized/narrow type and has the prepacked layout despite the API +// being: // B_prepacked = machete_prepack_B(B) // Y = machete_mm(A, B_prepacked) // diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh b/csrc/quantization/machete/machete_mm_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh rename to csrc/quantization/machete/machete_mm_kernel.cuh index 57655f369cdf..db3321a39db8 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh +++ b/csrc/quantization/machete/machete_mm_kernel.cuh @@ -18,9 +18,9 @@ // clang-format on #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/vllm_numeric_conversion.cuh" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "machete_collective_builder.cuh" #include "machete_prepacked_layout.cuh" #include "machete_interleaving_utils.cuh" diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh b/csrc/quantization/machete/machete_mm_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh rename to csrc/quantization/machete/machete_mm_launcher.cuh index 5ecf03daf790..14206f1d5cc2 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh +++ b/csrc/quantization/machete/machete_mm_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_mm_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh b/csrc/quantization/machete/machete_prepack_kernel.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh rename to csrc/quantization/machete/machete_prepack_kernel.cuh index 6a28dceccc8d..e1e054e5a007 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh +++ b/csrc/quantization/machete/machete_prepack_kernel.cuh @@ -2,7 +2,7 @@ #include "machete_mm_kernel.cuh" #include "cutlass_extensions/cute_utils.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include namespace machete { diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh b/csrc/quantization/machete/machete_prepack_launcher.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh rename to csrc/quantization/machete/machete_prepack_launcher.cuh index 63af5e496b84..26bf200732d3 100644 --- a/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh +++ b/csrc/quantization/machete/machete_prepack_launcher.cuh @@ -1,9 +1,9 @@ #pragma once #include "machete_prepack_kernel.cuh" -#include "libtorch_stable/cutlass_extensions/torch_utils.hpp" +#include "cutlass_extensions/torch_utils.hpp" #include "core/scalar_type.hpp" -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh b/csrc/quantization/machete/machete_prepacked_layout.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh rename to csrc/quantization/machete/machete_prepacked_layout.cuh diff --git a/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu similarity index 100% rename from csrc/libtorch_stable/quantization/machete/machete_pytorch.cu rename to csrc/quantization/machete/machete_pytorch.cu diff --git a/csrc/libtorch_stable/quantization/marlin/.gitignore b/csrc/quantization/marlin/.gitignore similarity index 73% rename from csrc/libtorch_stable/quantization/marlin/.gitignore rename to csrc/quantization/marlin/.gitignore index 7dc482a89466..ba805f9250ec 100644 --- a/csrc/libtorch_stable/quantization/marlin/.gitignore +++ b/csrc/quantization/marlin/.gitignore @@ -1,3 +1,2 @@ sm*_kernel_*.cu kernel_selector.h -kernel_*.cu diff --git a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu b/csrc/quantization/marlin/awq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu rename to csrc/quantization/marlin/awq_marlin_repack.cu index 55ce5b4e732d..09b647a03a39 100644 --- a/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/quantization/marlin/awq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/dequant.h b/csrc/quantization/marlin/dequant.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/dequant.h rename to csrc/quantization/marlin/dequant.h diff --git a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py b/csrc/quantization/marlin/generate_kernels.py similarity index 96% rename from csrc/libtorch_stable/quantization/marlin/generate_kernels.py rename to csrc/quantization/marlin/generate_kernels.py index 0b908d86d563..9e0af3c275ac 100644 --- a/csrc/libtorch_stable/quantization/marlin/generate_kernels.py +++ b/csrc/quantization/marlin/generate_kernels.py @@ -161,12 +161,15 @@ }, ] +# Kernels are written next to the kernel.h / marlin_template.h they include. +OUTPUT_DIR = sys.argv[2] if len(sys.argv) > 2 else os.path.dirname(__file__) + def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): + for filename in glob.glob(OUTPUT_DIR + "/*kernel_*.cu"): subprocess.call(["rm", "-f", filename]) - filename = os.path.dirname(__file__) + "/kernel_selector.h" + filename = OUTPUT_DIR + "/kernel_selector.h" subprocess.call(["rm", "-f", filename]) @@ -297,7 +300,7 @@ def generate_new_kernels(): filename = filename.lower() - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: + with open(os.path.join(OUTPUT_DIR, filename), "w") as f: f.write(file_content) if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: @@ -307,7 +310,7 @@ def generate_new_kernels(): '"marlin kernel with fp8 activation is not built.");' ) - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: + with open(os.path.join(OUTPUT_DIR, "kernel_selector.h"), "w") as f: f.write(kernel_selector_str) diff --git a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/csrc/quantization/marlin/gptq_marlin_repack.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu rename to csrc/quantization/marlin/gptq_marlin_repack.cu index 5dab9fe932bc..3fcf3406eb1b 100644 --- a/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/quantization/marlin/gptq_marlin_repack.cu @@ -7,7 +7,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" namespace marlin { diff --git a/csrc/libtorch_stable/quantization/marlin/kernel.h b/csrc/quantization/marlin/kernel.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/kernel.h rename to csrc/quantization/marlin/kernel.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cu b/csrc/quantization/marlin/marlin.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin.cu rename to csrc/quantization/marlin/marlin.cu index cd52f3c73619..d77ad692b2bb 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin.cu +++ b/csrc/quantization/marlin/marlin.cu @@ -32,7 +32,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ diff --git a/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/csrc/quantization/marlin/marlin.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin.cuh rename to csrc/quantization/marlin/marlin.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/csrc/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh rename to csrc/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu similarity index 99% rename from csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu rename to csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu index b48601435e3c..c4463269099d 100644 --- a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -8,7 +8,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // for only non-zp format (like gptq) __global__ void marlin_int4_fp8_preprocess_kernel_without_zp( diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/csrc/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_mma.h rename to csrc/quantization/marlin/marlin_mma.h diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_template.h b/csrc/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/libtorch_stable/quantization/marlin/marlin_template.h rename to csrc/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/quantization/vectorization.cuh b/csrc/quantization/vectorization.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization.cuh rename to csrc/quantization/vectorization.cuh diff --git a/csrc/libtorch_stable/quantization/vectorization_utils.cuh b/csrc/quantization/vectorization_utils.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/vectorization_utils.cuh rename to csrc/quantization/vectorization_utils.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh rename to csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh index a93a136bf425..0722101b3879 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh @@ -5,7 +5,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/cutlass.h" @@ -19,8 +19,8 @@ #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/util/packed_stride.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on namespace vllm::c3x { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh index cf18b63a8848..b20c055fb829 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm.cuh @@ -14,8 +14,8 @@ #include "cutlass/epilogue/collective/collective_builder.hpp" #include "cutlass/gemm/collective/collective_builder.hpp" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on /* diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu similarity index 90% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu index 5178eed0722e..bc088cf633f4 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu index b7930012265f..f3df69850ec6 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh index fa2add918213..1e7a31f5ca55 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh @@ -2,7 +2,7 @@ #include -#include "cuda_utils.h" +#include "core/utils.hpp" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" @@ -242,10 +242,10 @@ void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::stable::Tensor& out, if (!swap_ab) { constexpr int TILE_N = 128; int tile_m = 256; - if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 64) <= sms) { + if (div_ceil(n, TILE_N) * div_ceil(m, 64) <= sms) { tile_m = 64; } - else if (cuda_utils::ceil_div(n, TILE_N) * cuda_utils::ceil_div(m, 128) <= sms) { + else if (div_ceil(n, TILE_N) * div_ceil(m, 128) <= sms) { tile_m = 128; } if (tile_m == 64) { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu similarity index 96% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu index 021c16d66824..9f3c8a196df7 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm120_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh index 9c921ea90a32..34bb7ef3f23f 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh @@ -2,7 +2,6 @@ #include -#include "cuda_utils.h" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu similarity index 89% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu index a97909d37ae8..d3318c487675 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_blockwise_sm90_fp8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp similarity index 82% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp index 913436186c3a..a15bc1da4a56 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_helper.hpp @@ -1,7 +1,7 @@ #include #include -#include "cuda_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" template void dispatch_scaled_mm(torch::stable::Tensor& c, @@ -41,14 +41,12 @@ void dispatch_scaled_mm(torch::stable::Tensor& c, STD_TORCH_CHECK(b_scales.dim() == 2, "b scale must be 2d tensor."); int32_t version_num = get_sm_version_num(); if (version_num >= 90) { - STD_TORCH_CHECK( - a.size(0) == a_scales.size(0) && - cuda_utils::ceil_div(a.size(1), int64_t(128)) == a_scales.size(1), - "a_scale_group_shape must be [1, 128]."); - STD_TORCH_CHECK( - cuda_utils::ceil_div(b.size(0), int64_t(128)) == b_scales.size(0) && - cuda_utils::ceil_div(b.size(1), int64_t(128)) == b_scales.size(1), - "b_scale_group_shape must be [128, 128]."); + STD_TORCH_CHECK(a.size(0) == a_scales.size(0) && + div_ceil(a.size(1), int64_t(128)) == a_scales.size(1), + "a_scale_group_shape must be [1, 128]."); + STD_TORCH_CHECK(div_ceil(b.size(0), int64_t(128)) == b_scales.size(0) && + div_ceil(b.size(1), int64_t(128)) == b_scales.size(1), + "b_scale_group_shape must be [128, 128]."); } STD_TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm"); diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_kernels.hpp diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu index 84040a6a2218..59022b041e20 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm100_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh index 42f9e4d5bc83..f790b3653d57 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM100 (fp8) based on the diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu index 1f709699eb9e..5cfe1d8d3ab9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu @@ -1,7 +1,7 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm120_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu index e86c9bd48d3f..c94b9bd07207 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_fp8_dispatch.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh index 2fae3016c309..f78b8daea510 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh @@ -4,7 +4,7 @@ #include "scaled_mm.cuh" #include "cutlass_gemm_caller.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" /** * This file defines Gemm kernel configurations for SM90 (fp8) based on the Gemm diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu similarity index 91% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu index bf2bd030e3c3..717a2a588307 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu +++ b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu @@ -1,6 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm90_int8_dispatch.cuh" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh rename to csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh index 6942c4b48ccb..2189e5d0abd8 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/get_group_starts.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh @@ -3,7 +3,7 @@ #include #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include "cutlass/bfloat16.h" #include "cutlass/float8.h" diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh index b4cd520e9679..12241fb299b3 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x.cuh @@ -7,9 +7,9 @@ #include "cutlass/gemm/device/gemm_universal_adapter.h" #include -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" -#include "libtorch_stable/torch_utils.h" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" +#include "cutlass_extensions/common.hpp" +#include "torch_utils.h" #include "get_group_starts.cuh" using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu index 1f88bf140da5..f03f2c0bc47e 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu index d494bfb9d08c..14bb0586cdfc 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu +++ b/csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu @@ -1,6 +1,6 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu rename to csrc/quantization/w8a8/cutlass/moe/moe_data.cu index 8ebc3dd9e878..b6043e803306 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -1,11 +1,11 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" #include #include #include -#include "libtorch_stable/dispatch_utils.h" +#include "dispatch_utils.h" #include diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu similarity index 98% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu index 184a26b491f8..bd60058d73fd 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu @@ -9,8 +9,8 @@ #include "scaled_mm_c2x_sm89_fp8_dispatch.cuh" #include "scaled_mm_c2x_sm89_int8_dispatch.cuh" -#include "core/batch_invariant.hpp" -#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp" using namespace vllm; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh index 9f9e12e982ed..30e587c63ad4 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cuh @@ -3,7 +3,7 @@ #include #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" // clang-format will break include orders // clang-format off @@ -22,8 +22,8 @@ #include "cutlass/epilogue/threadblock/fusion/visitors.hpp" #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" -#include "libtorch_stable/core/math.hpp" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "core/utils.hpp" +#include "cutlass_extensions/common.hpp" // clang-format on using namespace cute; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm75_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm80_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh b/csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh rename to csrc/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_int8_dispatch.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu rename to csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 51f84d2ffd95..0c05f59d3fcd 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -3,9 +3,9 @@ #include -#include "libtorch_stable/torch_utils.h" +#include "torch_utils.h" -#include "libtorch_stable/cutlass_extensions/common.hpp" +#include "cutlass_extensions/common.hpp" void cutlass_scaled_mm_sm75(torch::stable::Tensor& c, torch::stable::Tensor const& a, diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 32925c0c3c95..cb050c922302 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -5,7 +5,7 @@ #include #include -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include namespace vllm { diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu similarity index 99% rename from csrc/libtorch_stable/quantization/w8a8/fp8/common.cu rename to csrc/quantization/w8a8/fp8/common.cu index abebd6f02f80..471499786977 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/common.cu +++ b/csrc/quantization/w8a8/fp8/common.cu @@ -1,5 +1,5 @@ -#include "../../../../quantization/w8a8/fp8/common.cuh" -#include "../../../../cuda_compat.h" +#include "common.cuh" +#include "core/cuda_compat.h" #include "../../../dispatch_utils.h" #include "../../../cub_helpers.h" #include "../../vectorization_utils.cuh" diff --git a/csrc/quantization/w8a8/fp8/common.cuh b/csrc/quantization/w8a8/fp8/common.cuh index 087f5099165a..143be698cc82 100644 --- a/csrc/quantization/w8a8/fp8/common.cuh +++ b/csrc/quantization/w8a8/fp8/common.cuh @@ -1,17 +1,17 @@ #pragma once -#include "libtorch_stable/quantization/vectorization.cuh" -#include "../../utils.cuh" +#include "quantization/vectorization.cuh" +#include "quantization/utils.cuh" #include -// This header is shared between _C and _C_stable_libtorch targets. +// This header is shared between _rocm_C and _C_stable_libtorch targets. // torch_utils.h provides get_device_prop(). We need to pass USE_CUDA // to the .so to expose some of the shims used by torch_utils.h. For now -// this is only done for _C_stable_libtorch and not for _C, so we use the -// non stable at::cuda::getCurrentDeviceProperties for _C for now. +// this is only done for _C_stable_libtorch and not for _rocm_C, so we use +// the non stable at::cuda::getCurrentDeviceProperties for _rocm_C for now. #ifdef TORCH_TARGET_VERSION - #include "../../../libtorch_stable/torch_utils.h" + #include "torch_utils.h" #else #ifdef USE_ROCM #include @@ -19,9 +19,9 @@ #endif #ifndef USE_ROCM - #include "nvidia/quant_utils.cuh" + #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" #else - #include "amd/quant_utils.cuh" + #include "quantization/w8a8/fp8/amd/quant_utils.cuh" #endif // Determines the preferred FP8 type for the current platform. diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 7fcf541fd64d..77aabe465111 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -1,6 +1,6 @@ #pragma once -#include "../../../../attention/attention_dtypes.h" +#include "core/vec_dtypes/vec_dtypes.h" #include #include #include @@ -243,7 +243,7 @@ __inline__ __device__ uint8_t scaled_vec_conversion( } // torch stable-ABI (headeronly) scalar types delegate to the CUDA-native -// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// conversions, so stable-ABI kernels dispatched on c10::BFloat16 / // c10::Half quantize correctly without manual casts. template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu similarity index 97% rename from csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu rename to csrc/quantization/w8a8/fp8/per_token_group_quant.cu index a0aee0cd9287..d2a888643899 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -3,8 +3,6 @@ #include #include -#include "libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h" - #include #ifdef USE_ROCM @@ -13,10 +11,10 @@ #include #endif -#include "libtorch_stable/quantization/vectorization.cuh" -#include "libtorch_stable/quantization/vectorization_utils.cuh" -#include "libtorch_stable/dispatch_utils.h" -#include "libtorch_stable/torch_utils.h" +#include "quantization/vectorization.cuh" +#include "quantization/vectorization_utils.cuh" +#include "dispatch_utils.h" +#include "torch_utils.h" __device__ __forceinline__ float GroupReduceMax(float val) { #ifdef USE_ROCM @@ -346,11 +344,13 @@ inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) { return 4; } -void per_token_group_quant_8bit(const torch::stable::Tensor& input, - torch::stable::Tensor& output_q, - torch::stable::Tensor& output_s, - int64_t group_size, double eps, double min_8bit, - double max_8bit, bool scale_ue8m0) { +// 8-bit per-token-group quantization helper used by both FP8 and INT8 +static void per_token_group_quant_8bit(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, + double min_8bit, double max_8bit, + bool scale_ue8m0 = false) { STD_TORCH_CHECK(input.is_contiguous()); STD_TORCH_CHECK(output_q.is_contiguous()); @@ -870,3 +870,12 @@ void per_token_group_quant_fp8(const torch::stable::Tensor& input, per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0); } + +void per_token_group_quant_int8(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s, + int64_t group_size, double eps, double int8_min, + double int8_max) { + per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, + int8_min, int8_max); +} diff --git a/csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu b/csrc/quantization/w8a8/int8/scaled_quant.cu similarity index 100% rename from csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu rename to csrc/quantization/w8a8/int8/scaled_quant.cu diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 2a11653a58d7..caf6ee9a14c5 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -19,11 +19,11 @@ #include #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include -#include "../attention/dtype_fp8.cuh" -#include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#include "core/vec_dtypes/dtype_fp8.cuh" +#include "quantization/w8a8/fp8/amd/quant_utils.cuh" // ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent #if !defined(HIP_FP8_TYPE_OCP) diff --git a/csrc/rocm/dispatch_utils.h b/csrc/rocm/dispatch_utils.h new file mode 100644 index 000000000000..12026965f098 --- /dev/null +++ b/csrc/rocm/dispatch_utils.h @@ -0,0 +1,24 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +// Need a special dispatch case macro since we will nest the FP8 dispatch. +// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. +#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +// ROCm devices might use either fn or fnuz, so set up dispatch table for both. +// A host-based check at runtime will create a preferred FP8 type for ROCm +// such that the correct kernel is dispatched. +#define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) + +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See AT_DISPATCH_FP8_CASE above. +#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu index e6f1c2fb176b..8ba5b48f05f8 100644 --- a/csrc/rocm/q_gemm_rdna3.cu +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -3,7 +3,8 @@ // // W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the // activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit -// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: +// kernel (csrc/quantization/gptq/q_gemm.cu) with the following +// changes: // // 1. Deterministic split-K epilogue. K is split across gridDim.z blocks; // each block stores its FP32 block partial to a scratch tensor, and the diff --git a/csrc/quickreduce/base.h b/csrc/rocm/quickreduce/base.h similarity index 100% rename from csrc/quickreduce/base.h rename to csrc/rocm/quickreduce/base.h diff --git a/csrc/custom_quickreduce.cu b/csrc/rocm/quickreduce/custom_quickreduce.cu similarity index 99% rename from csrc/custom_quickreduce.cu rename to csrc/rocm/quickreduce/custom_quickreduce.cu index d4e5d179a54e..c1e75ae05d6a 100644 --- a/csrc/custom_quickreduce.cu +++ b/csrc/rocm/quickreduce/custom_quickreduce.cu @@ -5,7 +5,7 @@ #ifdef USE_ROCM - #include "quickreduce/quick_reduce.h" + #include "rocm/quickreduce/quick_reduce.h" quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size, std::optional qr_max_size) { diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/rocm/quickreduce/quick_reduce.h similarity index 100% rename from csrc/quickreduce/quick_reduce.h rename to csrc/rocm/quickreduce/quick_reduce.h diff --git a/csrc/quickreduce/quick_reduce_impl.cuh b/csrc/rocm/quickreduce/quick_reduce_impl.cuh similarity index 100% rename from csrc/quickreduce/quick_reduce_impl.cuh rename to csrc/rocm/quickreduce/quick_reduce_impl.cuh diff --git a/csrc/rocm/quickreduce/torch_bindings.cpp b/csrc/rocm/quickreduce/torch_bindings.cpp new file mode 100644 index 000000000000..b7d9a695ac48 --- /dev/null +++ b/csrc/rocm/quickreduce/torch_bindings.cpp @@ -0,0 +1,51 @@ +#include +#include "core/registration.h" +#include +#include + +#include +#include + +#ifdef USE_ROCM +// Defined in custom_quickreduce.cu. +using fptr_t = int64_t; +fptr_t init_custom_qr(int64_t rank, int64_t world_size, + std::optional qr_max_size = std::nullopt); +void qr_destroy(fptr_t _fa); +torch::Tensor qr_get_handle(fptr_t _fa); +void qr_open_handles(fptr_t _fa, const std::vector& handles); +void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, + int64_t quant_level, bool cast_bf2half = false); +int64_t qr_max_size(); +#endif + +// Note on op signatures: +// The X_meta signatures are for the meta functions corresponding to op X. +// They must be kept in sync with the signature for X. Generally, only +// functions that return Tensors require a meta function. +// +// See the following links for detailed docs on op registration and function +// schemas. +// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 +// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations + +#ifdef USE_ROCM +TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { + // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). + custom_ar.def( + "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " + "cast_bf2half) -> ()"); + custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); + + custom_ar.def("init_custom_qr", &init_custom_qr); + custom_ar.def("qr_destroy", &qr_destroy); + custom_ar.def("qr_get_handle", &qr_get_handle); + + custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); + custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + + custom_ar.def("qr_max_size", &qr_max_size); +} +#endif + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 79af0814550d..bd7c0daa1db2 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -12,10 +12,10 @@ #include #include -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "quantization/w8a8/fp8/common.cuh" -#include "core/batch_invariant.hpp" +#include "core/utils.hpp" // Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot. static constexpr int64_t kWvSlots = 8; diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 4588b216fe01..be3cda7be6a3 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -10,8 +10,7 @@ #include #include // std::in_range -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#include "core/cuda_compat.h" // Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts #if defined(__GFX11__) || defined(__GFX12__) diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/sampled_topk.cuh similarity index 100% rename from csrc/libtorch_stable/sampled_topk.cuh rename to csrc/sampled_topk.cuh diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/sampler.cu similarity index 99% rename from csrc/libtorch_stable/sampler.cu rename to csrc/sampler.cu index 09151004c1ad..e9145e3b99e6 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/sampler.cu @@ -1,4 +1,4 @@ -#include "../cuda_compat.h" +#include "core/cuda_compat.h" #include "dispatch_utils.h" #include "torch_utils.h" diff --git a/csrc/libtorch_stable/topk.cu b/csrc/topk.cu similarity index 100% rename from csrc/libtorch_stable/topk.cu rename to csrc/topk.cu diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh similarity index 100% rename from csrc/libtorch_stable/topk_histogram_4096.cuh rename to csrc/topk_histogram_4096.cuh diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index bcf0ce3e1c48..3d4f9ab27389 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1,39 +1,1171 @@ -// Provides torch::Tensor for ops.h (previously included transitively via -// cache.h, which is no longer included here after cache ops moved to -// _C_stable_libtorch). -#include #include "ops.h" +#include "cuda_utils.h" #include "core/registration.h" -#include -#include - -// Note on op signatures: -// The X_meta signatures are for the meta functions corresponding to op X. -// They must be kept in sync with the signature for X. Generally, only -// functions that return Tensors require a meta function. -// -// See the following links for detailed docs on op registration and function -// schemas. -// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 -// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations - -#ifdef USE_ROCM -TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { - // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). - custom_ar.def( - "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " - "cast_bf2half) -> ()"); - custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); - custom_ar.def("init_custom_qr", &init_custom_qr); - custom_ar.def("qr_destroy", &qr_destroy); - custom_ar.def("qr_get_handle", &qr_get_handle); +#include + +// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility. +// Note: We register under namespace "_C" so ops are accessible as +// torch.ops._C. for compatibility with existing code. +STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { + // Compute per-token-group FP8 quantized tensor and scaling factor. + // The dummy arguments are here so we can correctly fuse with RMSNorm. + ops.def( + "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! " + "output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool " + "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned " + ") -> ()"); + // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, + // TMA-aligned scales for DeepGEMM. + ops.def( + "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " + "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " + "float fp8_max) -> ()"); + // Compute per-token-group INT8 quantized tensor and scaling factor. + ops.def( + "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " + "output_s, int group_size, float eps, float int8_min, float int8_max) -> " + "()"); + ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + +#ifndef USE_ROCM + + // Note about marlin kernel 'workspace' arguments: + // Technically these should be mutable since they are modified by the kernel. + // But since they are set back to zero once the kernel is finished we can + // hand wave and say that they have no net effect. + // + // The reason to mark 'workspace' as immutable is so that they don't interfere + // with using ScalarType arguments in the ops. If they are marked as mutable, + // pytorch throws an assert in + // 'torch._higher_order_ops._register_effectful_op' that prevents these + // kernels from being torch.compile'd. + // See the following document for more info on custom types and ops that use + // custom types: + // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA + + // Machete (Dense) Optimized Mixed Precision GEMM for Hopper. + ops.def( + "machete_supported_schedules(" + " ScalarType a_type," + " int b_type," + " ScalarType? maybe_group_scales_type," + " ScalarType? maybe_group_zeros_type," + " ScalarType? maybe_channel_scales_type," + " ScalarType? maybe_token_scales_type," + " ScalarType? maybe_out_type" + ") -> str[]"); + ops.def( + "machete_mm(" + " Tensor A," + " Tensor B," + " int b_type," + " ScalarType? out_type," + " Tensor? group_scales," + " Tensor? group_zeros," + " int? group_size," + " Tensor? channel_scales," + " Tensor? token_scales," + " str? schedule" + ") -> Tensor"); + ops.def( + "machete_prepack_B(" + " Tensor B," + " ScalarType a_type," + " int b_type," + " ScalarType? group_scales_type" + ") -> Tensor"); + // conditionally compiled so impl registration is in source file + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file +#endif + +#ifndef USE_ROCM + // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column + // quantization, as well as bias + ops.def( + "cutlass_scaled_mm(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor? bias) -> ()"); + + // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column + // quantization. + ops.def( + "cutlass_scaled_mm_azp(Tensor! out, Tensor a," + " Tensor b, Tensor a_scales," + " Tensor b_scales, Tensor azp_adj," + " Tensor? azp, Tensor? bias) -> ()"); + + // Check if cutlass scaled_mm is supported for CUDA devices of the given + // capability + ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool"); + + // Check if cutlass grouped gemm is supported for CUDA devices of the given + // capability + ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool"); + + // CUTLASS w8a8 grouped GEMM + ops.def( + "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, " + " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, " + " Tensor problem_sizes, Tensor a_strides, " + " Tensor b_strides, Tensor c_strides, bool per_act_token, " + " bool per_out_ch) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM. It takes topk_ids as an input, and computes expert_offsets + // (token start indices of each expert). In addition to this, it computes + // problem sizes for each expert's multiplication used by the two mms called + // from fused MoE operation, and arrays with permutations required to shuffle + // and de-shuffle the input/output of the fused operation. + ops.def( + "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, " + " Tensor! problem_sizes1, Tensor! problem_sizes2, " + " Tensor! input_permutation, " + " Tensor! output_permutation, int num_experts, " + " int n, int k, Tensor? blockscale_offsets, " + " bool is_gated) -> ()"); + + // compute per-expert problem sizes from expert_first_token_offset + // produced by vLLM's moe_permute kernel + ops.def( + "get_cutlass_moe_mm_problem_sizes_from_expert_offsets(" + " Tensor expert_first_token_offset, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " int n, int k, bool swap_ab) -> ()"); + + // A function that computes data required to run fused MoE with w8a8 grouped + // GEMM in batched expert format. It takes expert_num_tokens + // as an input, and computes expert_offsets (token start indices of each + // expert). In addition to this, it computes problem sizes for each expert's + // multiplication used by the two mms called from fused MoE operation. + ops.def( + "get_cutlass_batched_moe_mm_data(Tensor! expert_offsets, " + " Tensor! problem_sizes1, " + " Tensor! problem_sizes2, " + " Tensor expert_num_tokens, " + " int num_local_experts, int padded_m, " + " int n, int k) -> ()"); + + // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3) + ops.def( + "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> " + "bool"); + + // CUTLASS nvfp4 block scaled GEMM + ops.def( + "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b," + " Tensor block_scale_a, Tensor block_scale_b," + " Tensor alpha) -> ()"); + + // cutlass nvfp4 block scaled group GEMM + ops.def( + "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE) + ops.def( + "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b," + " Tensor a_blockscale, Tensor b_blockscales," + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()"); + + // Compute NVFP4 block quantized tensor. + ops.def( + "scaled_fp4_quant(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout) -> (Tensor, Tensor)"); + + // Out variant + // TODO: Add out_variant tag once PyTorch supports it (added in 2.11) + // This registration is now migrated to stable ABI + // at::Tag::out_variant is not available in the stable ABI (enum_tag.h is not + // yet in torch/headeronly), the tag should be applied from Python + // via torch.library.Library.define(..., tags=(torch.Tag.out_variant,)) + // with the .impl remaining in C++. + // See pytorch/pytorch#176117. + ops.def( + "scaled_fp4_quant.out(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " + "-> ()"); + + // Compute NVFP4 experts quantization. + ops.def( + "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 experts quantization. + ops.def( + "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts) -> ()"); + + // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs). + ops.def( + "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+MXFP4 experts quantization. + ops.def( + "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " + "output_scale," + "Tensor input, Tensor input_offset_by_experts," + "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + + // Fused SiLU+Mul+NVFP4 quantization. + ops.def( + "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, " + "Tensor input, Tensor input_global_scale) -> ()"); + + // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices + // of the given capability + ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool"); + + // CUTLASS w4a8 GEMM + ops.def( + "cutlass_w4a8_mm(" + " Tensor A," + " Tensor B," + " Tensor group_scales," + " int group_size," + " Tensor channel_scales," + " Tensor token_scales," + " ScalarType? out_type," + " str? maybe_schedule" + ") -> Tensor"); + + // pack scales + ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor"); + + // encode and reorder weight matrix + ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor"); + + // CUTLASS w4a8 grouped GEMM + ops.def( + "cutlass_w4a8_moe_mm(" + " Tensor! out_tensors," + " Tensor a_tensors," + " Tensor b_tensors," + " Tensor a_scales," + " Tensor b_scales," + " Tensor b_group_scales," + " int b_group_size," + " Tensor expert_offsets," + " Tensor problem_sizes," + " Tensor a_strides," + " Tensor b_strides," + " Tensor c_strides," + " Tensor group_scale_strides," + " str? maybe_schedule" + ") -> ()"); + + ops.def( + "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, " + "Tensor)"); + + // SM100 CUTLASS MLA decode + // conditionally compiled so impl registrations are in source file + ops.def( + "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope," + " Tensor q_pe, Tensor kv_c_and_k_pe_cache," + " Tensor seq_lens, Tensor page_table," + " Tensor workspace, float scale," + " int num_kv_splits) -> ()"); + + ops.def( + "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches," + " int sm_count, int num_kv_splits) " + "-> int"); + // Quantized GEMM for AWQ. + ops.def( + "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters) -> Tensor"); + + // Dequantization for AWQ. + ops.def( + "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, " + "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor"); + + // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + // conditionally compiled so impl registration is in source file + ops.def( + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); + + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). + // conditionally compiled so impl registration is in source file + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); +#endif + + // Merge attn states + // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 + // can be used to combine partial attention results (in the split-KV case) + ops.def( + "merge_attn_states(" + " Tensor! output," + " Tensor!? output_lse," + " Tensor prefix_output," + " Tensor prefix_lse," + " Tensor suffix_output," + " Tensor suffix_lse," + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); + + // Hadamard transforms + // conditionally compiled so impl registration is in source file + ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); + + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) " + "-> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " + "float epsilon) -> ()"); + + // Layernorm-quant + // Apply Root Mean Square (RMS) Normalization to the input tensor. + ops.def( + "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, " + "Tensor scale, float epsilon) -> " + "()"); + + // In-place fused Add and RMS Normalization. + ops.def( + "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, " + "Tensor! residual, Tensor weight, " + "Tensor scale, float epsilon) -> ()"); + + // Fused Layernorm + Quant kernels + ops.def( + "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual) -> ()"); + + // Fused Layernorm + Block quant kernels + ops.def( + "rms_norm_per_block_quant(Tensor! result, Tensor input, " + "Tensor weight, Tensor! scale, float epsilon, " + "Tensor? scale_ub, Tensor!? residual, int group_size, " + "bool is_scale_transposed) -> ()"); + + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + + // Rotary embedding + // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. + ops.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox, int " + "rope_dim_offset=0, bool inverse=False) -> ()"); + + // Function for fused QK Norm and RoPE + ops.def( + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " + "int num_heads_k, int num_heads_v, int head_dim, float eps, " + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); + + // q_head_padded is the padded Q head count of the returned tensor, or 0 to + // do the KV insert alone and return an empty tensor. The Q knobs are + // independent: apply_q_norm and apply_q_rope each drop that step for Q + // alone (KV is always rotated), and is_q_interleaved reads and writes Q in + // FlashMLA's mega-attention chunk-interleaved layout, which moves the + // padding heads to the tail of every head-dim chunk. + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor q_in, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size, " + "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " + "bool is_q_interleaved=False) -> Tensor"); + + // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate + // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(" + "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, " + "Tensor position_ids, Tensor cos_sin_cache, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(" + "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " + "int cache_block_size, bool apply_q_norm=True) -> ()"); + + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat_quant_fp8(" + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + +#ifndef USE_ROCM + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " + "float eps) -> (Tensor, Tensor)"); +#endif + + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert. + ops.def( + "fused_minimax_m3_qknorm_rope_kv_insert(" + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, " + "Tensor cos_sin_cache, Tensor positions, int num_heads, " + "int num_kv_heads, int rotary_dim, float eps, " + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, " + "int num_index_heads, " + "Tensor? slot_mapping, Tensor? index_slot_mapping, " + "Tensor!? kv_cache, Tensor!? index_cache, " + "int block_size, Tensor!? q_out, Tensor!? index_q_out, " + "str kv_cache_dtype, bool skip_index_branch=False, " + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); + +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.def( + "fused_gdn_decode_post_conv_mtp(" + "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " + "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " + "float scale, float norm_eps=1e-5, " + "str output_gate_activation='silu') -> ()"); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.def( + "fused_kda_prologue(" + "Tensor q, Tensor k, Tensor v, Tensor raw_g, Tensor raw_beta, " + "Tensor A_log, Tensor dt_bias, Tensor! qg, Tensor! w, Tensor! u, " + "Tensor! kg_t, Tensor! aqk, Tensor! decay, Tensor cu_seqlens, " + "Tensor chunk_indices, Tensor? conv_weight, Tensor(e!)? conv_state, " + "Tensor? conv_state_indices, Tensor? conv_has_initial_state, " + "float scale, float lower_bound) -> ()"); + ops.def( + "fused_kda_chunk(" + "Tensor qg, Tensor w, Tensor u, Tensor kg_t, Tensor aqk, Tensor decay, " + "Tensor? initial_state, Tensor(a!)? final_state, Tensor! out, " + "Tensor cu_seqlens, Tensor chunk_offsets, float scale, " + "Tensor(b!)? group_state, int groups, " + "Tensor(c!)? checkpoint_state=None, Tensor? checkpoint_offsets=None, " + "Tensor? checkpoint_state_indices=None, Tensor(d!)? state_cache=None, " + "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()"); +#endif + +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " + "int num_blocks, " + "int block_write_idx, float eps, " + "float output_norm_eps) -> ()"); +#endif + + // Apply repetition penalties to logits in-place. + ops.def( + "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " + "Tensor output_mask, Tensor repetition_penalties) -> ()"); + + // Optimized top-k per row operations. + ops.def( + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " + "Tensor! indices, int numRows, int stride0, " + "int stride1, int topK) -> ()"); + + ops.def( + "top_k_per_row_decode(Tensor logits, int next_n, " + "Tensor seq_lens, Tensor! indices, " + "int numRows, int stride0, int stride1, int topK) -> ()"); + + ops.def( + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.def( + "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); +#endif + + // Activation ops + ops.def( + "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! " + "y_q, Tensor! y_s, bool use_ue8m0) -> ()"); + ops.def("weak_ref_tensor(Tensor input) -> Tensor"); + + // Activation function used in SwiGLU. + ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); + + ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()"); + + // SwiGLU activation with input clamping. + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up. + ops.def( + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, " + "float alpha=1.0, float beta=0.0) -> ()"); + + // SwiGLU activation with FP8 quantization. + ops.def( + "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); + + // Activation function used in GeGLU with `none` approximation. + ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); + + // Activation function used in GeGLU with `tanh` approximation. + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); + + // FATReLU implementation. + ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()"); + + ops.def( + "swigluoai_and_mul(Tensor! out, Tensor input, float alpha=1.702, float " + "limit=7.0) " + "-> ()"); + + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path + // (writes the fp8 down input and its float32 scale). group_size=0 -> + // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [.., + // d/128]. + ops.def( + "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " + "float beta=1.0, float linear_beta=-1.0, int group_size=0, " + "Tensor? num_valid_tokens=None, int topk=1) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + ops.def( + "masked_moe_activation(Tensor! out, Tensor input, Tensor " + "valid_token_counts, str activation, float clamp_limit=0.0, float " + "alpha=1.0, float beta=0.0, float situ_beta=1.0, float " + "situ_linear_beta=-1.0) -> ()"); + + // GELU implementation used in GPT-2. + ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); - custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); - custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + // Approximate GELU implementation. + ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); - custom_ar.def("qr_max_size", &qr_max_size); + // Quick GELU implementation. + ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); + + // relu(x)^2 activation from https://arxiv.org/abs/2109.08668v2 + ops.def("relu_squared(Tensor! out, Tensor input) -> ()"); + + // Compute int8 quantized tensor for given scaling factor. + ops.def( + "static_scaled_int8_quant(Tensor! result, Tensor input, Tensor scale," + "Tensor? azp) -> ()"); + + // Compute int8 quantized tensor and scaling factor + ops.def( + "dynamic_scaled_int8_quant(Tensor! result, Tensor input, Tensor! scale, " + "Tensor!? azp) -> ()"); + + // Compute FP8 quantized tensor for given scaling factor. + // Supports per-tensor, per-channel, per-token, and arbitrary 2D group + // scaling. Optional group_m/group_n specify the group shape explicitly; + // required for 1D scales to disambiguate per-channel vs per-token. + ops.def( + "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, " + "int[]? group_shape=None) -> ()"); + + // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_scaled_fp8_quant(Tensor! result, Tensor input, Tensor! scale) " + "-> " + "()"); + + // Compute dynamic-per-token FP8 quantized tensor and scaling factor. + ops.def( + "dynamic_per_token_scaled_fp8_quant(Tensor! result, Tensor input, " + "Tensor! scale, Tensor? scale_ub) -> " + "()"); + + // Quantized GEMM for GPTQ. + // Note: even though the C++ inferred schema is correct for this op, it seems + // to prevent the meta function registry. + ops.def( + "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, " + "Tensor b_gptq_scales, bool use_exllama, bool " + "use_v2_format, int bit) " + "-> Tensor"); + + // Post processing for GPTQ. + ops.def("gptq_shuffle(Tensor! q_weight, int bit) -> ()"); + + // Mamba selective scan kernel + ops.def( + "selective_scan_fwd(Tensor! u, Tensor! delta," + "Tensor! A, Tensor! B, Tensor! C," + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," + "bool delta_softplus," + "Tensor? query_start_loc," + "Tensor? cache_indices," + "Tensor? has_initial_state," + "Tensor! ssm_states," + "int null_block_id," + "int block_size," + "Tensor? block_idx_first_scheduled_token," + "Tensor? block_idx_last_scheduled_token," + "Tensor? initial_state_idx," + "Tensor? cu_chunk_seqlen," + "Tensor? last_chunk_indices) -> ()"); + + // LongCat n-gram embedding index kernel. All tensor args are marked mutable + // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table + // and n_gram_ids are actually written in place. + // Fused vocab-parallel embedding: gather the rows this rank owns and write + // zeros for the rest, so the following all-reduce reconstructs the full + // embedding. + ops.def( + "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, " + "int org_vocab_start_index, int org_vocab_end_index, " + "int num_org_vocab_padding, int added_vocab_start_index, " + "int added_vocab_end_index) -> ()"); + ops.def( + "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, " + "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, " + "Tensor(d!) exclusive_req_len_sums, Tensor(e!) ne_token_table, " + "Tensor(f!) row_indices, Tensor(g!) column_starts, " + "Tensor(h!) n_gram_ids) -> ()"); } + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { + // LongCat n-gram embedding index kernel. + ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); + ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); + + // Per-token group quantization + ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); + ops.impl("per_token_group_fp8_quant_packed", + TORCH_BOX(&per_token_group_quant_8bit_packed)); + ops.impl("per_token_group_quant_int8", + TORCH_BOX(&per_token_group_quant_int8)); + +#ifndef USE_ROCM + // CUTLASS scaled_mm ops + ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); + ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); + ops.impl("cutlass_moe_mm", TORCH_BOX(&cutlass_moe_mm)); + ops.impl("get_cutlass_moe_mm_data", TORCH_BOX(&get_cutlass_moe_mm_data)); + ops.impl("get_cutlass_moe_mm_problem_sizes_from_expert_offsets", + TORCH_BOX(&get_cutlass_moe_mm_problem_sizes_from_expert_offsets)); + ops.impl("get_cutlass_batched_moe_mm_data", + TORCH_BOX(&get_cutlass_batched_moe_mm_data)); + + // FP4/NVFP4 ops + ops.impl("cutlass_scaled_fp4_mm", TORCH_BOX(&cutlass_scaled_fp4_mm)); + ops.impl("scaled_fp4_quant", TORCH_BOX(&scaled_fp4_quant_func)); + ops.impl("scaled_fp4_quant.out", TORCH_BOX(&scaled_fp4_quant_out)); + ops.impl("scaled_fp4_experts_quant", TORCH_BOX(&scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_scaled_fp4_experts_quant", + TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant)); + ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant)); + // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only). + // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu. + + // AWQ ops + ops.impl("awq_gemm", TORCH_BOX(&awq_gemm)); + ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize)); + + // DSV3 fused A GEMM: conditionally compiled so impl registration is in + // source file (dsv3_fused_a_gemm.cu) +#endif + + ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); + + // Layernorm kernels (shared CUDA/ROCm) + ops.impl("rms_norm", TORCH_BOX(&rms_norm)); + ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); + + // Layernorm-quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_static_fp8_quant", TORCH_BOX(&rms_norm_static_fp8_quant)); + ops.impl("fused_add_rms_norm_static_fp8_quant", + TORCH_BOX(&fused_add_rms_norm_static_fp8_quant)); + + // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm) + ops.impl("rms_norm_dynamic_per_token_quant", + TORCH_BOX(&rms_norm_dynamic_per_token_quant)); + ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant)); + ops.impl("silu_and_mul_per_block_quant", + TORCH_BOX(&silu_and_mul_per_block_quant)); + + // Positional encoding kernels (shared CUDA/ROCm) + ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); + ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert)); + ops.impl( + "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_kv_concat", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); +#ifndef USE_ROCM + ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); +#endif + ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", + TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.impl("fused_gdn_decode_post_conv_mtp", + TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); +#endif + +#ifdef VLLM_ENABLE_FUSED_KDA_CHUNK + ops.impl("fused_kda_prologue", TORCH_BOX(&fused_kda_prologue)); + ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk)); #endif -REGISTER_EXTENSION(TORCH_EXTENSION_NAME) +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif + + // Sampler kernels (shared CUDA/ROCm) + ops.impl("apply_repetition_penalties_", + TORCH_BOX(&apply_repetition_penalties_)); + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); +#ifdef VLLM_ENABLE_COOPERATIVE_TOPK + ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); +#endif + + // Activation kernels (shared CUDA/ROCm) + ops.impl("persistent_masked_m_silu_mul_quant", + TORCH_BOX(&persistent_masked_m_silu_mul_quant)); + ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor)); + ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant)); + ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); + ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); + ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul)); + ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); + ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); + ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); + ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); + ops.impl("gelu_new", TORCH_BOX(&gelu_new)); + ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); + ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); + ops.impl("relu_squared", TORCH_BOX(&relu_squared)); + ops.impl("silu_and_mul_with_clamp", TORCH_BOX(&silu_and_mul_clamp)); + + // INT8 quantization kernels + ops.impl("static_scaled_int8_quant", TORCH_BOX(&static_scaled_int8_quant)); + ops.impl("dynamic_scaled_int8_quant", TORCH_BOX(&dynamic_scaled_int8_quant)); + + // FP8 quantization kernels + ops.impl("static_scaled_fp8_quant", TORCH_BOX(&static_scaled_fp8_quant)); + ops.impl("dynamic_scaled_fp8_quant", TORCH_BOX(&dynamic_scaled_fp8_quant)); + ops.impl("dynamic_per_token_scaled_fp8_quant", + TORCH_BOX(&dynamic_per_token_scaled_fp8_quant)); + + // GPTQ kernels + ops.impl("gptq_gemm", TORCH_BOX(&gptq_gemm)); + ops.impl("gptq_shuffle", TORCH_BOX(&gptq_shuffle)); + + // Mamba kernels + ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { + ops.impl("get_cuda_view_from_cpu_tensor", + TORCH_BOX(&get_cuda_view_from_cpu_tensor)); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) { + cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int"); + cuda_utils.def( + "get_max_shared_memory_per_block_device_attribute(int device_id) -> int"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd, + cuda_utils) { + cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute)); + cuda_utils.impl("get_max_shared_memory_per_block_device_attribute", + TORCH_BOX(&get_max_shared_memory_per_block_device_attribute)); +} + +// These capability-check functions take only primitive args (no tensors), so +// there is no device to dispatch on. CompositeExplicitAutograd makes them +// available for all backends. This is the stable ABI equivalent of calling +// ops.impl("op_name", &func) without a dispatch key in the non-stable API. +STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { +#ifndef USE_ROCM + ops.impl("cutlass_scaled_mm_supports_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_fp8)); + ops.impl("cutlass_group_gemm_supported", + TORCH_BOX(&cutlass_group_gemm_supported)); + ops.impl("cutlass_scaled_mm_supports_block_fp8", + TORCH_BOX(&cutlass_scaled_mm_supports_block_fp8)); + ops.impl("cutlass_scaled_mm_supports_fp4", + TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); +#endif +} + +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Grouped concat_and_cache_mla across all layers. Each layer's cache base + // pointer and optional plain-FP8 scale are read from device tensors. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride," + " Tensor? kv_scales=None," + " str kv_cache_dtype='auto') -> ()"); + +#ifndef USE_ROCM + ops.def( + "hisparse_resolve_residency(Tensor host_cache," + " Tensor! hot_cache," + " Tensor hot_block_table," + " Tensor global_indices," + " Tensor! hot_indices," + " Tensor! device_global_indices," + " Tensor! lru_slots," + " Tensor? request_state_indices," + " int region_stride," + " int max_union_rows," + " Tensor(a!)? miss_mask=None," + " Tensor(b!)? stats=None," + " Tensor(c!)? attention_indices=None," + " int attention_block_stride=0," + " Tensor? request_ids=None," + " Tensor? source_block_table=None," + " int source_block_size=0," + " Tensor(d!)? resolved_global_indices=None," + " Tensor(e!)? valid_counts=None," + " Tensor(f!)? swap_host_physical_rows=None," + " Tensor(g!)? swap_device_physical_rows=None," + " Tensor(h!)? swap_counts=None," + " Tensor? resident_block_table=None," + " int resident_block_size=0," + " int resident_null_block=0) -> ()"); + + ops.def( + "hisparse_invalidate_written_slots(Tensor! device_global_indices," + " Tensor request_state_indices," + " Tensor req_id_per_token," + " Tensor written_slots) -> ()"); + + ops.def( + "hisparse_gather_plan(Tensor host_cache," + " Tensor! hot_cache," + " Tensor global_indices," + " Tensor hot_indices," + " Tensor miss_mask," + " Tensor? request_state_indices," + " Tensor(a!)? attention_indices=None," + " int attention_block_stride=0) -> ()"); + + ops.def( + "hisparse_gather_compact(Tensor host_cache," + " Tensor! hot_cache," + " Tensor miss_global_indices," + " Tensor miss_hot_indices," + " Tensor miss_counts) -> ()"); + +#endif // !USE_ROCM + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts, Tensor? host_cache=None, Tensor? host_row_ids=None, Tensor? " + "device_row_ids=None) -> ()"); + + ops.def( + "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("meta_size() -> int"); + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("free_shared_buffer(int ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { + custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); + custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); + custom_ar.impl("get_graph_buffer_ipc_meta", + TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("allocate_shared_buffer_and_handle", + TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); + +#ifndef USE_ROCM + ops.impl("hisparse_resolve_residency", + TORCH_BOX(&hisparse_resolve_residency)); + ops.impl("hisparse_invalidate_written_slots", + TORCH_BOX(&hisparse_invalidate_written_slots)); + ops.impl("hisparse_gather_plan", TORCH_BOX(&hisparse_gather_plan)); + ops.impl("hisparse_gather_compact", TORCH_BOX(&hisparse_gather_compact)); +#endif // !USE_ROCM + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("cp_gather_and_upconvert_nvfp4_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_nvfp4_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + +REGISTER_EXTENSION(_C_stable_libtorch) diff --git a/csrc/torch_utils.h b/csrc/torch_utils.h index 898b9e113e19..f02346739e0a 100644 --- a/csrc/torch_utils.h +++ b/csrc/torch_utils.h @@ -1,17 +1,88 @@ #pragma once -// Shared TORCH_UTILS_CHECK across both libtorch stable and unstable source -// files. Keep this header free of CUTLASS/CUTE so attention/quant headers can -// use it. -// -// If TORCH_TARGET_VERSION is defined, we are building _C_stable_libtorch.so so -// use STD_TORCH_CHECK via header-only. -// Otherwise, use TORCH_CHECK via torch/all.h. - -#ifdef TORCH_TARGET_VERSION - #include - #define TORCH_UTILS_CHECK STD_TORCH_CHECK -#else - #include - #define TORCH_UTILS_CHECK TORCH_CHECK -#endif +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED. +#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \ + STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__) + +// Device properties cache for stable ABI compatibility. +// Uses raw CUDA/HIP APIs instead of ATen functions. +// Using inline ensures a single instance across all translation units. +inline std::deque device_flags; +inline std::vector device_properties; +inline std::once_flag vectors_init_flag; + +inline void do_init_device_vectors() { + int device_count; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " + + std::string(cudaGetErrorString(err))); + } + device_flags.resize(device_count); + device_properties.resize(device_count); +} + +inline void initDeviceVectors() { + std::call_once(vectors_init_flag, do_init_device_vectors); +} + +inline void initDeviceProperty(int device_index) { + cudaDeviceProp device_prop{}; + cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " + + std::string(cudaGetErrorString(err))); + } + device_properties[device_index] = device_prop; +} + +// Get device properties using raw CUDA/HIP APIs (stable ABI compatible). +// Caches results per device so cudaGetDeviceProperties is called at most once +// per device. +inline cudaDeviceProp* get_device_prop() { + initDeviceVectors(); + int device_index; + cudaError_t err = cudaGetDevice(&device_index); + if (err != cudaSuccess) { + STD_TORCH_CHECK( + false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err))); + } + STD_TORCH_CHECK(device_index >= 0 && static_cast(device_index) < + device_properties.size(), + "CUDA device index " + std::to_string(device_index) + + " out of range [0, " + + std::to_string(device_properties.size()) + ")"); + + std::call_once(device_flags[device_index], initDeviceProperty, device_index); + return &device_properties[device_index]; +} + +// Utility to get the current CUDA stream for a given device using stable APIs. +// Returns a cudaStream_t for use in kernel launches. +inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); + return reinterpret_cast(stream_ptr); +} + +// Utility to get the current cuBLAS handle using stable APIs. +inline cublasHandle_t get_current_cuda_blas_handle() { + void* blas_handle_ptr = nullptr; + TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr)); + return reinterpret_cast(blas_handle_ptr); +} diff --git a/csrc/libtorch_stable/type_convert.cuh b/csrc/type_convert.cuh similarity index 100% rename from csrc/libtorch_stable/type_convert.cuh rename to csrc/type_convert.cuh diff --git a/csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu b/csrc/vocab_parallel_embedding_kernels.cu similarity index 100% rename from csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu rename to csrc/vocab_parallel_embedding_kernels.cu diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 4095c41a7e63..6f76838f5560 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -232,7 +232,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) -- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu) +- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu) ### SiLU+Mul + Quantization (`fuse_act_quant`) @@ -258,7 +258,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) -- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu) +- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) ### RMSNorm + Padding (`fuse_act_padding`) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 7ca207c8e4ad..173a2ea34432 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -185,10 +185,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None: let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). """ source = ( - Path(__file__).resolve().parents[2] - / "csrc" - / "libtorch_stable" - / "dsv3_fused_a_gemm.cu" + Path(__file__).resolve().parents[2] / "csrc" / "dsv3_fused_a_gemm.cu" ).read_text(encoding="utf-8") # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built # by default, so they must not count as available. diff --git a/tools/ci/kernel_symbol_map.py b/tools/ci/kernel_symbol_map.py index c88f2d9ad551..3177a10bba9b 100755 --- a/tools/ci/kernel_symbol_map.py +++ b/tools/ci/kernel_symbol_map.py @@ -32,13 +32,13 @@ "errors": [], # [{"object", "source", "error"}, ...] "objects": [ { - "source": "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", + "source": "csrc/fused_qknorm_rope_kernel.cu", "target": "_C", "object": "CMakeFiles/_C.dir/csrc/.../fused_qknorm_rope_kernel.cu.o", "device": true, "symbols": ["_Z21fusedQKNormRopeKernelI...", ...], - "deps": ["csrc/libtorch_stable/fused_qknorm_rope_kernel.cu", - "csrc/cuda_compat.h", ...] + "deps": ["csrc/fused_qknorm_rope_kernel.cu", + "csrc/core/cuda_compat.h", ...] # "error": "..." # present when cuobjdump failed on it }, ... diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 406869ea2c4b..0aa2a6090261 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -28,7 +28,7 @@ class QuickReduceRegime(Enum): - # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h + # Keep integer ids aligned with csrc/rocm/quickreduce/quick_reduce.h FP = 0 INT8 = 1 INT6 = 2 diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index da8c69eed0a2..a541debfe9f3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -104,7 +104,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # azp_adj is the AZP adjustment term, used to account for weights. # It does not depend on scales or azp, so it is the same for # static and dynamic quantization. - # For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md + # For more details, see + # csrc/quantization/w8a8/cutlass/Epilogues.md # https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md if not config.input_symmetric: weight = getattr(layer, w_q_name) diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py index e64e81821062..fec8719a899c 100644 --- a/vllm/models/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -68,7 +68,8 @@ def kernel( weight_ptr = kv_weight_ptr row_out = kv_out_ptr + token_idx * kv_out_stride - # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # RMSNorm in fp32 throughout — matches + # csrc/layernorm_kernels.cu's # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which # keep x, rrms, and w all in fp32 and perform a single cast at store. block = tl.arange(0, BLOCK_SIZE) diff --git a/vllm/models/kimi_k3/amd/ops/kda_chunk.py b/vllm/models/kimi_k3/amd/ops/kda_chunk.py index 035a49f852d6..a3fbada908fa 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_chunk.py +++ b/vllm/models/kimi_k3/amd/ops/kda_chunk.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry point for the fused Kimi-K3 KDA chunk kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` replaces the chunk-state recurrence and the output GEMM of the Triton chunk path with a single launch that keeps the per-chunk state in registers, so the ``[chunks, H, V, K]`` state tensor and the recomputed values never reach HBM. diff --git a/vllm/models/kimi_k3/amd/ops/kda_decode.py b/vllm/models/kimi_k3/amd/ops/kda_decode.py index 132cd58cf816..682776a80726 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_decode.py +++ b/vllm/models/kimi_k3/amd/ops/kda_decode.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ROCm entry points for the fused Kimi-K3 KDA decode kernel. -The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel_rocm.cu`` +The kernel in ``csrc/kimi_k3/fused_kda_decode_kernel_rocm.cu`` replaces, for a pure non-speculative decode batch, the three Triton launches and two copies the AMD KDA layer otherwise runs per layer: the packed causal conv1d update, the recurrent delta-rule step, and the gated output RMSNorm. diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 79de3af0bc22..a39ec787e0bb 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -3,7 +3,7 @@ """Fused MLA prefill and decode epilogues for Kimi-K3. Thin wrappers over the CUDA ops in -``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +``csrc/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. - ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 0ff34f8ac5e0..4fe0be0b1f68 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -118,7 +118,7 @@ elements, stored permuted (an 8x4 -> 4x8 transpose: the scale for element block `s` lives at byte `8 * (s & 3) + (s >> 2)`) so that the 8 scales one FlashMLA dequant thread needs are contiguous. See the layout comment in - `csrc/libtorch_stable/cache_kernels.cu`. + `csrc/cache_kernels.cu`. """ diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index f1b0fa688eda..f7474b6bedad 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -712,7 +712,7 @@ def __init__( "HiSparse requires its compiled _C_cache_ops CUDA kernels " "(host-resident decode has no Python fallback). Rebuild vLLM " "from source so " - "csrc/libtorch_stable/hisparse_kernels.cu is included." + "csrc/hisparse_kernels.cu is included." ) self.max_num_reqs = max_num_reqs self.row_width = row_width From fa39582756dcbf25b65cc36d8daa465b19134372 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 13:00:47 -0400 Subject: [PATCH 426/952] [Frontend] Serve LiquidAI/d1-3B on /v1/systemone (#60676) Signed-off-by: mgoin Co-authored-by: Simon Veitner Co-authored-by: Claude --- .../online_serving/structured_decisions.md | 25 ++- .../structured_decisions/test_strategies.py | 52 ++++++ .../structured_decisions/strategies.py | 162 ++++++++++++++++-- 3 files changed, 222 insertions(+), 17 deletions(-) diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 4d8d65d89e94..ce09d6f81fee 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -156,6 +156,29 @@ vllm serve Cloudflare/clef-omni with it when every video has one. Audio clips cannot be sent together with videos that have soundtracks. +## Liquid d1 + +[`LiquidAI/d1-3B`](https://huggingface.co/LiquidAI/d1-3B) is LFM2.5-VL-3B +post-trained to answer decisions as the first token of its reply. vLLM serves +it as a generative model. + +```bash +vllm serve LiquidAI/d1-3B +``` + +- The server writes the prompt and reads the label tokens the way the + `prompt.py` shipped with the checkpoint does, so picks and probabilities + match `model.system_one()`. `confidence` and `usage` follow the definitions + on this page instead. +- A JSON state is written as indented JSON. A `null` state leaves the images, + or the question alone, as the prompt. +- `images` are placed before the state and shrunk to at most 1024 x 1024 + pixels, as the reference does. +- A `choice` whose option names are all single letters is labeled with those + letters. +- When each request carries one question over a state that changes from + request to request, `--no-enable-prefix-caching` lowers latency. + ## Request fields | field | meaning | @@ -165,7 +188,7 @@ vllm serve Cloudflare/clef-omni | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni)) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 8d38e61e6248..dae545d1f04c 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -26,6 +26,7 @@ from vllm.entrypoints.generate.structured_decisions.question_types import ( LABELS, StructuredDecisionError, + build_question, ) from vllm.entrypoints.generate.structured_decisions.serving import ( ServingStructuredDecisions, @@ -34,8 +35,10 @@ from vllm.entrypoints.generate.structured_decisions.strategies import ( DiffusionGemmaCanvasStrategy, JointSchemaStrategy, + LiquidStrategy, NextTokenStrategy, ReadContext, + cap_pixels, reply_label_ids, select_read_strategy, state_text, @@ -79,6 +82,8 @@ def test_strategy_selection(): select_read_strategy(model("DiffusionGemmaForBlockDiffusion"), generate) is DiffusionGemmaCanvasStrategy ) + lfm = "Lfm2VlForConditionalGeneration" + assert select_read_strategy(model(lfm), generate) is LiquidStrategy with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): select_read_strategy(model("LlamaForCausalLM"), generate) with pytest.raises(ValueError, match="not raw_logits"): @@ -573,3 +578,50 @@ async def process_for_engine_async(prompt, arrival_time): ) assert isinstance(response, ErrorResponse) assert response.error.code == 400 + + +@pytest.mark.parametrize( + "type_name,criteria,ask", + [ + ( + "choice", + {"billing": "Charges", "app_fault": None}, + "Which?\n\nOptions:\nA Charges\nB app fault\n\n" + "Reply with the option code only.", + ), + ( + "choice", + {"a": "Berlin", "b": "Paris"}, + "Which?\n\nOptions:\na Berlin\nb Paris\n\nReply with the option code only.", + ), + ( + "noul", + {"true": "asks to cancel", "false": "anything else"}, + "Which?\nYes: asks to cancel\nNo: anything else\n\n" + "Reply with yes or no only.", + ), + ( + "score", + ["low", "high"], + "Which?\n\n0 low\n1 high\n\nReply with a single digit 0-1 only.", + ), + ], +) +def test_liquid_prompt_matches_the_checkpoint(type_name, criteria, ask): + # The wording of prompt.py in LiquidAI/d1-3B, which the model was trained on. + strategy = object.__new__(LiquidStrategy) + q = build_question("q", type_name, "Which?", criteria, 26) + assert strategy._content(strategy.render_state(None), q) == ask + assert strategy._content("Hi.", q) == f"Hi.\n\n\nQUESTION:\n{ask}" + state = strategy.render_state({"k": "é"}) + assert strategy._content(state, q).startswith('{\n "k": "é"\n}\n\n\nQUEST') + + +def test_liquid_caps_images_like_the_checkpoint(): + # runner.cap_pixels in LiquidAI/d1-3B: 1600x1598 becomes 1024x1023. + cap = LiquidStrategy.max_image_pixels + assert cap is not None + assert cap_pixels(Image.new("RGB", (1600, 1598)), cap).size == (1024, 1023) + small = Image.new("RGB", (640, 480)) + assert cap_pixels(small, cap) is small + assert NextTokenStrategy.max_image_pixels is None diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index f4365924ddf0..5567cb808ab5 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -18,6 +18,7 @@ from typing import Any, ClassVar, TypeVar import numpy as np +from PIL import Image import vllm.envs as envs from vllm.config import ModelConfig @@ -92,6 +93,16 @@ class ReadStrategy(ABC): def __init__(self, context: ReadContext): self.context = context + @cached_property + def _connector(self) -> MediaConnector: + model_config = self.context.engine_client.model_config + return MEDIA_CONNECTOR_REGISTRY.load( + envs.VLLM_MEDIA_CONNECTOR, + media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) + @abstractmethod def limits(self) -> DecisionLimits: ... @@ -162,12 +173,25 @@ def type_label_ids( return label_ids +def cap_pixels(image: Image.Image, max_pixels: int) -> Image.Image: + """Shrink ``image`` to at most ``max_pixels`` pixels (bicubic).""" + w, h = image.size + if w * h <= max_pixels: + return image + scale = math.sqrt(max_pixels / (w * h)) + size = (max(1, int(w * scale)), max(1, int(h * scale))) + return image.resize(size, Image.Resampling.BICUBIC) + + class NextTokenStrategy(ReadStrategy): """Autoregressive models. Each question is one request: the state, then the question with its labeled options, and the logprobs of the label tokens as the reply's first token. The requests share the state. With prefix caching, it is prefilled once.""" + #: Images are shrunk to at most this many pixels before the processor. + max_image_pixels: ClassVar[int | None] = None + def __init__(self, context: ReadContext): super().__init__(context) # Every prompt ends with the same generation prompt, so a label's token @@ -210,6 +234,27 @@ def _read_input( """The engine input for a read of the rendered prompt.""" return engine_input + def _label_groups(self, question: Question) -> list[list[int]]: + """The tokens that count for each label. A label scores its best one.""" + return [[self.label_ids[label]] for label in question.labels] + + def _content(self, state: str, question: Question) -> str: + """The user turn: the state, then the question.""" + return f"{state}\n\n{question.type.prompt(question)}" + + async def _image_parts(self, urls: Sequence[str]) -> list[dict[str, Any]]: + """The images as chat content parts, fetched once so that every + question reads the same pictures.""" + if not urls: + return [] + try: + images = await asyncio.gather(*map(self._connector.fetch_image_async, urls)) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + if self.max_image_pixels is not None: + images = [cap_pixels(image, self.max_image_pixels) for image in images] + return [{"type": "image_pil", "image_pil": image} for image in images] + def limits(self) -> DecisionLimits: return DecisionLimits(max_questions=64, max_options=len(LABELS)) @@ -234,6 +279,9 @@ async def _render( default_template=ctx.chat_template, default_template_content_format=ctx.chat_template_content_format, default_template_kwargs=ctx.default_chat_template_kwargs, + # A read that fails before it reaches the engine must not leave the + # sender cache believing the engine has the images. + skip_mm_cache=True, ) prompt_ids = extract_prompt_components( ctx.engine_client.model_config, engine_input @@ -256,10 +304,14 @@ async def read( ) -> list[QuestionRead]: ctx = self.context read_request = self._read_request(chat_template_kwargs, cache_salt) + image_parts = await self._image_parts((media or {}).get("image", ())) slots, engine_inputs, params = [], [], [] for q in questions: - messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}] + content: Any = self._content(state, q) + if image_parts: + content = [*image_parts, {"type": "text", "text": content}] + messages = [{"role": "user", "content": content}] if instructions: messages.insert(0, {"role": "system", "content": instructions}) engine_input, prompt_ids = await self._render(read_request, messages) @@ -268,9 +320,11 @@ async def read( "these chat options end the prompt differently, so the " "labels' tokens are unknown" ) - slots.append([self.label_ids[label] for label in q.labels]) + slots.append(self._label_groups(q)) engine_inputs.append(self._read_input(engine_input, prompt_ids)) - params.append(self._sampling_params(slots[-1], prompt_ids)) + params.append( + self._sampling_params([i for ids in slots[-1] for i in ids], prompt_ids) + ) label_reads = await next_token_label_reads( ctx.engine_client, @@ -283,14 +337,17 @@ async def read( ) reads = [] - for ids, label_read in zip(slots, label_reads): + for groups, label_read in zip(slots, label_reads): output = label_read.result.outputs[0] + logprobs = iter(label_read.logprobs) reads.append( QuestionRead( - probs=label_softmax(label_read.logprobs), + probs=label_softmax( + [max(next(logprobs) for _ in g) for g in groups] + ), label_mass=sum(math.exp(lp) for lp in label_read.logprobs), argmax_is_label=bool(output.token_ids) - and output.token_ids[0] in ids, + and any(output.token_ids[0] in ids for ids in groups), input_tokens=len(label_read.result.prompt_token_ids or ()), output_tokens=len(output.token_ids), ) @@ -373,6 +430,88 @@ def _read_input( ) +class LiquidStrategy(NextTokenStrategy): + """Liquid decision models (d1). The prompt and the label tokens are the + ones the checkpoint was trained with, from the `prompt.py` it ships.""" + + modalities = frozenset({"image"}) + # runner.py bounds every picture at 1024 x 1024 pixels before the processor. + max_image_pixels = 1024 * 1024 + + def render_state(self, state: Any) -> str: + if state is None: + return "" + if isinstance(state, str): + return state + return json.dumps(state, ensure_ascii=False, indent=2) + + def _single_tokens(self, texts: Sequence[str]) -> list[int]: + ids: list[int] = [] + for text in texts: + encoded = self._tokenizer().encode(text, add_special_tokens=False) + if len(encoded) == 1 and encoded[0] not in ids: + ids.append(encoded[0]) + return ids + + @staticmethod + def _codes(question: Question) -> list[str]: + """A choice's codes are A to Z, or the option names when those are + already single letters.""" + names = [o.name.strip() for o in question.options] + if all(len(n) == 1 and n.isalpha() for n in names): + return names + return list(question.labels) + + def _label_groups(self, question: Question) -> list[list[int]]: + if question.type.name == "noul": + forms: list[Sequence[str]] = [("yes", "Yes", "YES"), ("no", "No", "NO")] + elif question.type.name == "choice": + forms = [(code, f" {code}") for code in self._codes(question)] + else: + forms = [(label,) for label in question.labels] + groups = [self._single_tokens(f) for f in forms] + seen: set[int] = set() + for (label, *_), ids in zip(forms, groups): + if not ids or self._single_tokens([label]) != ids[:1] or seen & set(ids): + raise StructuredDecisionError( + f"question {question.id!r}: label {label!r} is not one " + "distinct token for this model" + ) + seen.update(ids) + return groups + + def _content(self, state: str, question: Question) -> str: + q = question + if q.type.name == "choice": + lines = "\n".join( + f"{code} {o.description or o.name.replace('_', ' ')}" + for code, o in zip(self._codes(q), q.options) + ) + ask = ( + f"{q.instructions}\n\nOptions:\n{lines}\n\n" + "Reply with the option code only." + ) + elif q.type.name == "noul": + yes, no = (o.description for o in q.options) + sides = f"\nYes: {yes}\nNo: {no}" if yes or no else "" + ask = f"{q.instructions}{sides}\n\nReply with yes or no only." + elif q.type.name == "score": + legend = "\n".join( + f"{label} {o.name}" for label, o in zip(q.labels, q.options) + ) + ask = ( + f"{q.instructions}\n\n{legend}\n\n" + f"Reply with a single digit 0-{len(q.options) - 1} only." + ) + else: + raise StructuredDecisionError( + f"question {q.id!r}: this model does not answer {q.type.name!r} " + "questions" + ) + # A null state leaves the images, or the question alone, as the prompt. + return f"{state}\n\n\nQUESTION:\n{ask}" if state else ask + + class JointSchemaStrategy(ReadStrategy): """Clef-Omni. The whole request is one pooling read: the state, then every question with its options, as the model was trained to read them. The @@ -442,16 +581,6 @@ def _render(value: Any) -> str: def _tokens(self, text: str) -> list[int]: return self.tokenizer.encode(text, add_special_tokens=False) - @cached_property - def _connector(self) -> MediaConnector: - model_config = self.context.engine_client.model_config - return MEDIA_CONNECTOR_REGISTRY.load( - envs.VLLM_MEDIA_CONNECTOR, - media_io_kwargs=model_config.get_multimodal_config().media_io_kwargs, - allowed_local_media_path=model_config.allowed_local_media_path, - allowed_media_domains=model_config.allowed_media_domains, - ) - async def _fetch( self, urls: Sequence[str], media_io: MediaIO[_M], timeout: int ) -> list[_M]: @@ -682,6 +811,7 @@ async def read( READ_STRATEGIES: dict[str, type[ReadStrategy]] = { "ClefOmniForJointSchema": JointSchemaStrategy, + "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, "Qwen3_5ForConditionalGeneration": NextTokenStrategy, "Qwen3_5MoeForConditionalGeneration": NextTokenStrategy, From 6e6e86078afe73af5180a9888e88751d51ef1b31 Mon Sep 17 00:00:00 2001 From: haotiaa Date: Mon, 12 Oct 2026 10:05:34 -0700 Subject: [PATCH 427/952] [Bugfix][Platform] Don't resolve current_platform during import vllm (#60766) Signed-off-by: Haotian An --- .../vllm_add_dummy_platform/__init__.py | 3 +++ tests/plugins_tests/test_platform_plugins.py | 17 +++++++++++++++++ vllm/utils/torch_utils.py | 18 ++++++++++++++---- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py index 280b68514e19..2d28e151d328 100644 --- a/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py +++ b/tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py @@ -3,6 +3,9 @@ def dummy_platform_plugin() -> str | None: + # Real platform plugins often import vLLM modules while registering. + from vllm.utils.torch_utils import direct_register_custom_op # noqa: F401 + return "vllm_add_dummy_platform.dummy_platform.DummyPlatform" diff --git a/tests/plugins_tests/test_platform_plugins.py b/tests/plugins_tests/test_platform_plugins.py index 6d32c4c6d6f0..a75e36a7d1c8 100644 --- a/tests/plugins_tests/test_platform_plugins.py +++ b/tests/plugins_tests/test_platform_plugins.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import subprocess +import sys + import pytest import torch @@ -31,6 +34,20 @@ def test_platform_plugins(): ) +def test_import_vllm_does_not_resolve_platform(): + # Platform plugins must be loaded after `import vllm` completes; resolving + # current_platform during it runs plugins against a half-imported vLLM. + # Needs a fresh interpreter, since vllm is already imported here. + code = ( + "import sys, vllm\n" + "p = sys.modules.get('vllm.platforms')\n" + "assert p is None or p._current_platform is None, p._init_trace\n" + "from vllm.platforms import current_platform\n" + "assert current_platform.device_name == 'DummyDevice'\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) + + def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch): # simulate workload by running an example load_general_plugins() diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index dce5e84fce25..4bf2e3397ca6 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -73,7 +73,15 @@ T = TypeVar("T") -PIN_MEMORY = is_pin_memory_available() +if TYPE_CHECKING: + PIN_MEMORY: bool + + +def __getattr__(name: str) -> Any: + # Lazy so that importing this module does not resolve current_platform. + if name == "PIN_MEMORY": + return is_pin_memory_available() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: @@ -718,13 +726,15 @@ def async_tensor_h2d( data = torch.from_numpy(data) if isinstance(data, torch.Tensor): t = data - if PIN_MEMORY and not t.is_pinned(): + if is_pin_memory_available() and not t.is_pinned(): # Stage in pinned, contiguous buffer to ensure fully async copy. t = torch.empty( t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True ).copy_(t) else: - t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") + t = torch.tensor( + data, dtype=dtype, pin_memory=is_pin_memory_available(), device="cpu" + ) assert t.is_cpu if out is not None: @@ -736,7 +746,7 @@ def async_tensor_h2d( def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor: t = torch.from_numpy(array) - return t.pin_memory() if PIN_MEMORY else t + return t.pin_memory() if is_pin_memory_available() else t def make_ndarray_with_pad( From c599b5971faf7ed8bff4bd1b14713293c20082d6 Mon Sep 17 00:00:00 2001 From: Chester Date: Tue, 13 Oct 2026 01:05:55 +0800 Subject: [PATCH 428/952] [Bugfix][Frontend] Return all logprobs when completions requests logprobs=-1 (#55717) Signed-off-by: Chester <42577861+superdiaodiao@users.noreply.github.com> --- .../completion/test_completion_error.py | 40 +++++++++++++++++++ vllm/entrypoints/openai/completion/serving.py | 4 +- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index bbbb61f9c9c8..17934026efe6 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -16,6 +16,7 @@ from vllm.entrypoints.scale_out.render.serving import ServingRender from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.exceptions import GenerationError, VLLMValidationError +from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.deepseek_v4 import DeepseekV4Renderer from vllm.renderers.hf import HfRenderer @@ -841,6 +842,45 @@ def test_logprobs_minus_one_allowed(): assert request.logprobs == -1 +def _minimal_serving_completion() -> OpenAIServingCompletion: + """A serving object with only what `_create_completion_logprobs` reads.""" + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.return_tokens_as_token_ids = False + return serving + + +@pytest.mark.parametrize( + ("num_output_top_logprobs", "expected"), + [(-1, 3), (1, 2), (2, 3)], +) +def test_completion_logprobs_honors_requested_count( + num_output_top_logprobs: int, expected: int +): + """`logprobs=-1` asks for every logprob, so none may be filtered out. + + The validator accepts -1 and the sampler expands it to the vocabulary, so + dropping the candidates here would return an empty `top_logprobs` for every + position while still charging for the full computation. + """ + tokenizer = MagicMock() + tokenizer.decode = lambda token_ids: f"tok{token_ids[0]}" + step_logprobs = { + 10: Logprob(logprob=-0.1, decoded_token="a"), + 11: Logprob(logprob=-0.2, decoded_token="b"), + 12: Logprob(logprob=-0.3, decoded_token="c"), + } + + logprobs = _minimal_serving_completion()._create_completion_logprobs( + token_ids=[10], + top_logprobs=[step_logprobs], + num_output_top_logprobs=num_output_top_logprobs, + tokenizer=tokenizer, + ) + + assert logprobs.top_logprobs is not None + assert len(logprobs.top_logprobs[0]) == expected + + def test_logprobs_below_minus_one_rejected(): """Values more negative than -1 stay invalid.""" with pytest.raises(Exception, match="must be a positive value or -1"): diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 7f7bdac0757b..e5ba20c52fab 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -722,7 +722,9 @@ def _create_completion_logprobs( return_as_token_id=should_return_as_token_id, ): max(top_lp[1].logprob, -9999.0) for i, top_lp in enumerate(step_top_logprobs.items()) - if logprob_token_ids or num_output_top_logprobs >= i + if logprob_token_ids + or num_output_top_logprobs == -1 + or num_output_top_logprobs >= i } ) From 12e6a585ae020cdcef9ba61aae45eec5ba29591d Mon Sep 17 00:00:00 2001 From: Laith Sakka Date: Mon, 12 Oct 2026 10:09:57 -0700 Subject: [PATCH 429/952] =?UTF-8?q?[Fix][Compile]=20Declare=20group=5Fshap?= =?UTF-8?q?e=20in=20the=20static=20FP8=20quant=20fusion=20=E2=80=A6=20(#61?= =?UTF-8?q?108)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Laith Sakka Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/compilation/passes/fusion/matcher_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 6d959c30677a..80be65a9781e 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -481,8 +481,13 @@ def forward_custom( if self.quant_key.scale.static: assert scale is not None + group_shape = self.quant_key.scale.group_shape _, result = auto_functionalized( - self.QUANT_OP, result=result, input=input, scale=scale + self.QUANT_OP, + result=result, + input=input, + scale=scale, + group_shape=[group_shape.row, group_shape.col], ) return result, scale else: From 45bb710ecaa2f0e100d1cab27308186e9f4aa7dd Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 12 Oct 2026 13:11:26 -0400 Subject: [PATCH 430/952] [CI] Add nightly Python 3.14 smoke test (#60456) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/misc.yaml | 11 +++++ tests/standalone_tests/python_only_compile.sh | 43 +++++++++++++++++-- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 02f78e280c98..bf650fd756ba 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -446,6 +446,17 @@ steps: - setup.py - vllm/platforms/rocm.py +- label: ":nvidia: (H200 MIG 18GB) Python 3.14 Smoke Nightly" + key: python-314-smoke-nightly + depends_on: ~ + optional: true + timeout_in_minutes: 30 + device: h200_18gb + source_file_dependencies: + - tests/standalone_tests/python_only_compile.sh + commands: + - PYTHON_ONLY_COMPILE_PYTHON=3.14 bash standalone_tests/python_only_compile.sh + - label: ":nvidia: (H200 MIG 35GB) Async Engine, Inputs, Utils, Worker" device: h200_35gb key: async-engine-inputs-utils-worker diff --git a/tests/standalone_tests/python_only_compile.sh b/tests/standalone_tests/python_only_compile.sh index 7036dd4314f2..5a6c6961d67c 100644 --- a/tests/standalone_tests/python_only_compile.sh +++ b/tests/standalone_tests/python_only_compile.sh @@ -1,6 +1,10 @@ #!/bin/bash # This script tests if the python only compilation works correctly # for users who do not have any compilers installed on their system +# +# Set PYTHON_ONLY_COMPILE_PYTHON (e.g. 3.14) to install into a fresh uv venv +# with that interpreter and run a short generation, as a smoke test for +# Python versions other than the one the CI image is built with. set -e @@ -173,9 +177,12 @@ elif [[ ! -d vllm ]]; then exit 1 fi -# remove all compilers -apt remove --purge build-essential -y -apt autoremove -y +# remove all compilers, except for the smoke test, whose generation needs a +# C compiler at runtime for torch.compile and Triton +if [[ -z "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + apt remove --purge build-essential -y + apt autoremove -y +fi rm -f /tmp/changed.file echo 'import os; os.system("touch /tmp/changed.file")' >> vllm/__init__.py @@ -185,12 +192,42 @@ if [[ -n "${rocm_wheel}" ]]; then VLLM_PRECOMPILED_WHEEL_LOCATION="${rocm_wheel}" VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps elif [[ "${is_rocm}" == "1" ]]; then VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps +elif [[ -n "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + # Use the CUDA build the image's own torch was installed with, for both the + # new venv's torch and the precompiled wheel variant. + cuda_variant="$(python3 -c 'import torch; print("cu" + torch.version.cuda.replace(".", ""))')" + uv venv --python "${PYTHON_ONLY_COMPILE_PYTHON}" /tmp/python-only-venv + source /tmp/python-only-venv/bin/activate + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 \ + VLLM_PRECOMPILED_WHEEL_VARIANT="${cuda_variant}" \ + uv pip install -e . --torch-backend="${cuda_variant}" else VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . fi # Run the script python3 -c 'import vllm' +if [[ -n "${PYTHON_ONLY_COMPILE_PYTHON:-}" ]]; then + python3 - <<'PY' +import os +import sys + +from vllm import LLM + +expected = os.environ["PYTHON_ONLY_COMPILE_PYTHON"] +actual = f"{sys.version_info.major}.{sys.version_info.minor}" +assert actual == expected, f"expected Python {expected}, got {actual}" + +llm = LLM( + "hmellor/tiny-random-LlamaForCausalLM", + max_model_len=256, + gpu_memory_utilization=0.3, +) +outputs = llm.generate(["Hello, my name is"]) +assert outputs[0].outputs[0].token_ids, "no tokens generated" +PY +fi + # Check if the clangd log file was created if [ ! -f /tmp/changed.file ]; then echo "ERROR: changed.file was not created, python only compilation failed" From a7b7bc1fbe1f5c4e47c53b0f02a3ecfbcd8054b9 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Mon, 12 Oct 2026 19:21:22 +0200 Subject: [PATCH 431/952] [Metrics] Count finished requests by watermark status (#58715) Signed-off-by: Simon Veitner Co-authored-by: Claude Co-authored-by: Lucas Wilkinson --- docs/features/watermarking.md | 28 ++++++ tests/v1/engine/test_output_processor.py | 94 +++++++++++++++++++++ tests/v1/metrics/test_histogram_buckets.py | 14 +++ tests/watermarking/test_watermarking_e2e.py | 16 ++++ vllm/v1/engine/output_processor.py | 6 +- vllm/v1/metrics/loggers.py | 30 +++++++ vllm/v1/metrics/stats.py | 16 ++++ 7 files changed, 203 insertions(+), 1 deletion(-) diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index 66b6465eb851..f51ec72eb1ae 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -56,6 +56,34 @@ watermarking algorithm that does not support it natively, at the cost of weaker detectability. See [Speculative decoding](#speculative-decoding). +## Monitoring + +When `watermark_config` is set, the Prometheus counter +`vllm:watermark_requests_total` counts finished requests by `status`: + +- `watermarked`: the request was sampled with the watermark switched on. +- `skipped`: the request asked for the watermark (omitted or `true`) but could + not use it, such as `temperature=0` and trace replay requests. A beam search + request that asks for the watermark runs without it and logs a warning. Its + beam steps count as `not_watermarked`, not `skipped`. +- `not_watermarked`: all other finished requests. This includes explicit + opt-outs, pooling requests, and the internal requests of beam search (one + per beam step), generative scoring and language detection. + +Each finished request has exactly one status, so the sum over `status` equals +`vllm:request_success_total` on the same instance. Requests that the client +aborts are in neither counter. The counter does not exist without +`watermark_config`. With `n > 1`, each sequence counts once. In a +disaggregated prefill/decode deployment, both instances count a request: sum +over decode instances only. For streaming input, the first input chunk sets +the status. + +The status is per request, not per token. A `watermarked` request can contain +tokens without the watermark: accepted draft tokens with +`allow_target_only_watermarking`, tokens whose context repeats +(`deduplicate_contexts`), the first `context_width` generated tokens with +`deduplicate_contexts: "all"`, and tokens that a grammar forces. + ## Architecture `WatermarkConfig` selects an algorithm and PRF. Model Runner V2 constructs the diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 8eb85765993b..3c186a4e738d 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, Mock import numpy as np +import prometheus_client import pytest import torch @@ -17,6 +18,7 @@ MockEngineCore, ) from vllm import PoolingParams +from vllm.config.watermarking import WatermarkConfig from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs from vllm.lora.request import LoRARequest from vllm.outputs import ( @@ -41,6 +43,9 @@ RequestOutputCollector, RequestState, ) +from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.v1.metrics.loggers import PrometheusStatLogger +from vllm.v1.metrics.prometheus import unregister_vllm_metrics from vllm.v1.metrics.stats import IterationStats, PrefillStats, SchedulerStats from vllm.v1.outputs import SamplingMaskLists @@ -1169,6 +1174,95 @@ def test_iteration_stats(dummy_test_vectors): assert iteration_stats.num_generation_tokens == num_active +def test_watermark_status_partitions_finished_requests(dummy_test_vectors): + """Each finished request gets exactly one watermark status.""" + output_processor = OutputProcessor(dummy_test_vectors.tokenizer, log_stats=True) + prompt_token_ids = dummy_test_vectors.prompt_tokens[0] + + def make_request( + request_id: str, + sampling_params: SamplingParams | None = None, + pooling_params: PoolingParams | None = None, + ) -> EngineCoreRequest: + return EngineCoreRequest( + request_id=request_id, + external_req_id=f"{request_id}-ext", + prompt_token_ids=prompt_token_ids, + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=sampling_params, + pooling_params=pooling_params, + ) + + # The frontend resolves the flags before the request reaches the + # OutputProcessor; set them as it does. + def resolved(watermarking: bool, skipped: bool, n: int = 1) -> SamplingParams: + params = SamplingParams(n=n, watermarking=watermarking) + params._watermarking_skipped = skipped + return params + + finish_reasons = { + "wm": FinishReason.ABORT, + "skip": FinishReason.STOP, + # A re-resolved request can carry both flags; it is watermarked. + "sticky": FinishReason.STOP, + "off": FinishReason.ERROR, + "pool": FinishReason.STOP, + } + output_processor.add_request(make_request("wm", resolved(True, False)), None) + output_processor.add_request(make_request("skip", resolved(False, True)), None) + output_processor.add_request(make_request("sticky", resolved(True, True)), None) + output_processor.add_request(make_request("off", resolved(False, False)), None) + output_processor.add_request( + make_request("pool", pooling_params=PoolingParams(task="embed")), None + ) + parent_req = ParentRequest(make_request("parent", resolved(True, False, n=2))) + for idx in range(2): + child_id, child_params = parent_req.get_child_info(idx) + output_processor.add_request( + make_request(child_id, child_params), None, parent_req, idx + ) + finish_reasons[child_id] = FinishReason.LENGTH + + outputs = [ + EngineCoreOutput( + request_id=request_id, + new_token_ids=[] if request_id == "pool" else [prompt_token_ids[0]], + pooling_output=torch.tensor([1.0]) if request_id == "pool" else None, + finish_reason=finish_reason, + ) + for request_id, finish_reason in finish_reasons.items() + ] + iteration_stats = IterationStats() + output_processor.process_outputs(outputs, time.monotonic(), iteration_stats) + + vllm_config = dummy_test_vectors.vllm_config + vllm_config.watermark_config = WatermarkConfig(key=42) + try: + PrometheusStatLogger(vllm_config).record( + scheduler_stats=None, iteration_stats=iteration_stats + ) + totals: dict[str, dict[str, float]] = {} + for metric in prometheus_client.REGISTRY.collect(): + for sample in metric.samples: + if sample.name in ( + "vllm:watermark_requests_total", + "vllm:request_success_total", + ): + label = sample.labels.get("status", "all") + by_label = totals.setdefault(sample.name, {}) + by_label[label] = by_label.get(label, 0) + sample.value + finally: + unregister_vllm_metrics() + + status = totals["vllm:watermark_requests_total"] + assert status == {"watermarked": 4, "skipped": 1, "not_watermarked": 2} + assert sum(status.values()) == totals["vllm:request_success_total"]["all"] + + @pytest.mark.parametrize("log_stats", [True, False]) def test_lora_request_tracking(log_stats: bool, dummy_test_vectors): """Test LoRA request lifecycle tracking through waiting -> running -> finished.""" diff --git a/tests/v1/metrics/test_histogram_buckets.py b/tests/v1/metrics/test_histogram_buckets.py index cf905b5d8b71..c5d41db054b3 100644 --- a/tests/v1/metrics/test_histogram_buckets.py +++ b/tests/v1/metrics/test_histogram_buckets.py @@ -11,6 +11,7 @@ from vllm.config import ModelConfig, VllmConfig from vllm.config.observability import ObservabilityConfig +from vllm.config.watermarking import WatermarkConfig from vllm.v1.metrics.buckets import ( BUCKET_FAMILY_KEYS, build_1_2_5_buckets, @@ -317,3 +318,16 @@ def test_prometheus_logger_applies_overrides(): assert found[metric_name] == overridden[family], metric_name finally: unregister_vllm_metrics() + + +@pytest.mark.parametrize("with_config", [True, False]) +def test_watermark_requests_counter_registration(with_config: bool): + """The watermark counter must exist only with watermark_config.""" + config = build_logger_config(ObservabilityConfig()) + config.watermark_config = WatermarkConfig(key=42) if with_config else None + try: + PrometheusStatLogger(config) + names = {m.name for m in prometheus_client.REGISTRY.collect()} + assert ("vllm:watermark_requests" in names) == with_config + finally: + unregister_vllm_metrics() diff --git a/tests/watermarking/test_watermarking_e2e.py b/tests/watermarking/test_watermarking_e2e.py index c73d48f6b6b7..ce6433e61357 100644 --- a/tests/watermarking/test_watermarking_e2e.py +++ b/tests/watermarking/test_watermarking_e2e.py @@ -5,6 +5,7 @@ from vllm import SamplingParams from vllm.platforms import current_platform +from vllm.v1.metrics.reader import Counter from vllm.v1.watermarking import GumbelWatermarkDetector @@ -23,14 +24,29 @@ def test_llm_generated_sequence_is_watermarked(vllm_runner): enforce_eager=True, max_model_len=512, gpu_memory_utilization=0.2, + disable_log_stats=False, ) request = "Tell me a story about an explorer who discovers a mysterious island." params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True) params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False) + params_greedy = SamplingParams(temperature=0.0, max_tokens=16) with runner: output_use_wm = runner.llm.generate(request, params_use_wm) output_no_wm = runner.llm.generate(request, params_no_wm) + runner.llm.generate(request, params_greedy) + metrics = runner.llm.get_metrics() + + # The greedy request is skipped; the opt-out one is not watermarked. + counters = [m for m in metrics if isinstance(m, Counter)] + status = { + m.labels["status"]: m.value + for m in counters + if m.name == "vllm:watermark_requests" + } + assert status == {"watermarked": 1, "skipped": 1, "not_watermarked": 1} + num_success = sum(m.value for m in counters if m.name == "vllm:request_success") + assert sum(status.values()) == num_success result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids)) result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids)) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 263a28fb3901..eb499f6cc78b 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -276,7 +276,7 @@ def from_new_request( output_kind = request.pooling_params.output_kind assert request.external_req_id is not None - return cls( + state = cls( request_id=request.request_id, external_req_id=request.external_req_id, parent_req=parent_req, @@ -299,6 +299,10 @@ def from_new_request( stream_input=request.resumable, remote_prefill_cached_tokens=remote_prefill_cached_tokens, ) + if state.stats is not None and sampling_params is not None: + state.stats.is_watermarked = bool(sampling_params.watermarking) + state.stats.is_watermark_skipped = sampling_params._watermarking_skipped + return state def make_request_output( self, diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index c42d17790926..65acc49ac0e0 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -639,6 +639,25 @@ def __init__( counter_corrupted_requests, per_engine_labelvalues ) + if vllm_config.watermark_config is not None: + counter_watermark_requests = self._counter_cls( + name="vllm:watermark_requests", + documentation=( + "Finished requests by watermark status " + "(only with watermark_config)." + ), + labelnames=labelnames + ["status"], + ) + self.counter_watermark_requests: dict[str, dict[int, Counter]] = {} + for status in ("watermarked", "skipped", "not_watermarked"): + per_engine_labelvalues_with_status = { + idx: labelvalues + [status] + for idx, labelvalues in per_engine_labelvalues.items() + } + self.counter_watermark_requests[status] = create_metric_per_engine( + counter_watermark_requests, per_engine_labelvalues_with_status + ) + counter_prefix_cache_queries = self._counter_cls( name="vllm:prefix_cache_queries", documentation=( @@ -1226,6 +1245,17 @@ def record( self.counter_corrupted_requests[engine_idx].inc( iteration_stats.num_corrupted_reqs ) + if self.vllm_config.watermark_config is not None: + counter_watermark = self.counter_watermark_requests + counter_watermark["watermarked"][engine_idx].inc( + iteration_stats.num_watermarked_reqs + ) + counter_watermark["skipped"][engine_idx].inc( + iteration_stats.num_watermark_skipped_reqs + ) + counter_watermark["not_watermarked"][engine_idx].inc( + iteration_stats.num_not_watermarked_reqs + ) self.counter_num_preempted_reqs[engine_idx].inc( iteration_stats.num_preempted_reqs ) diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index 0046dcefa92e..87277f6545f0 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -277,6 +277,12 @@ class RequestStateStats: # Track if this request is corrupted (NaNs in logits) is_corrupted: bool = False + # Track if this request was sampled with the engine watermark + is_watermarked: bool = False + + # Track if this request asked for the watermark but could not use it + is_watermark_skipped: bool = False + @dataclass class FinishedRequestStats: @@ -479,6 +485,9 @@ def __init__(self): self.time_to_first_tokens_iter: list[float] = [] self.inter_token_latencies_iter: list[float] = [] self.num_corrupted_reqs: int = 0 + self.num_watermarked_reqs: int = 0 + self.num_watermark_skipped_reqs: int = 0 + self.num_not_watermarked_reqs: int = 0 def __repr__(self) -> str: field_to_value_str = ", ".join(f"{k}={v}" for k, v in vars(self).items()) @@ -624,6 +633,13 @@ def update_from_finished_request( if req_stats.is_corrupted: self.num_corrupted_reqs += 1 + if req_stats.is_watermarked: + self.num_watermarked_reqs += 1 + elif req_stats.is_watermark_skipped: + self.num_watermark_skipped_reqs += 1 + else: + self.num_not_watermarked_reqs += 1 + class LoRAStats: """Tracks waiting and running request IDs for a single LoRA.""" From 5b0b8f251ea6b878a22d85baba56e127c66297c8 Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Mon, 12 Oct 2026 20:38:12 +0300 Subject: [PATCH 432/952] [Bugfix][HiSparse] Fix TP group hang on steps without scheduled tokens (#60941) Signed-off-by: Maroon Ayoub Signed-off-by: Matthew Bonanni Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../unit/test_hisparse_connector.py | 47 +++++++++++++++++++ .../kv_connector/v1/hisparse/worker.py | 4 ++ 2 files changed, 51 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index 9ae0d3bc8eb1..bdfbee28d4c0 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -401,6 +401,53 @@ def swap_blocks_batch(src, dst, sizes): assert worker_meta.completed_transfer_counts == {transfer_id: 1} +@pytest.mark.parametrize("shared", [True, False]) +@pytest.mark.parametrize("num_tokens", [0, 4]) +def test_zero_token_step_holds_tp_ranks_after_host_write_wait( + monkeypatch, shared, num_tokens +): + """Steps without tokens hold every TP rank after its shared event wait.""" + calls: list[tuple[str, object]] = [] + compute_stream = MagicMock() + compute_stream.wait_event.side_effect = lambda event: calls.append(("wait", event)) + tp_group = MagicMock() + tp_group.barrier.side_effect = lambda: calls.append(("barrier", None)) + monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr(worker_module, "get_tp_group", lambda: tp_group) + + worker = object.__new__(HiSparseConnectorWorker) + worker.shared_host_region = MagicMock() if shared else None + worker.host_write_events = (MagicMock(), MagicMock()) + worker.host_write_event = worker.host_write_events[1] + worker._next_host_write_event = 0 + worker._slot_mapping_staging = None + worker.cache_handles = [] + worker._pending_invalid_block_ids = [] + for name in ( + "_stage_row_mirror_mapping", + "_finish_previous_step", + "_release_completed_dma_descriptors", + "_set_row_mirrors", + "_clear_forward_mirror_state", + "_copy_host_blocks", + "_restore_pages", + "_submit_transfers", + ): + setattr(worker, name, MagicMock()) + + worker.start_step( + HiSparseConnectorMetadata(None, (), (), {}, True, {}), + None, + [], + num_tokens=num_tokens, + ) + + expected = [("wait", worker.host_write_events[1])] + if shared and not num_tokens: + expected.append(("barrier", None)) + assert calls == expected + + def test_scheduled_prefix_hit_publishes_adopted_copies(): """Copies adopted after scheduling reach the worker as a residency update, leaving the block-table row the scheduler output carries untouched.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index bbccecd65784..c53be27079fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -416,6 +416,10 @@ def start_step( self.host_write_event = self.host_write_events[self._next_host_write_event] self._next_host_write_event ^= 1 current_stream().wait_event(previous_host_write_event) + if self.shared_host_region is not None and not num_tokens: + # Without a forward, no TP collective keeps rank 0 from recording + # a shared event again before every rank has issued its wait. + get_tp_group().barrier() self._release_completed_dma_descriptors() self._dma_submitted = False mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids) From e1a5086256a00c47a968aac6a583c25d040100a7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 12 Oct 2026 10:41:07 -0700 Subject: [PATCH 433/952] [Attention] Remove cascade attention (#61126) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- docs/design/cuda_graphs.md | 6 +- tests/conftest.py | 7 - .../attention/test_cascade_flash_attn.py | 187 ------ .../kernels/attention/test_prefix_prefill.py | 1 - tests/platforms/test_cpu.py | 1 - tests/system_messages/sonnet3.5_nov2024.txt | 71 --- tests/v1/attention/test_attention_backends.py | 2 - .../test_rocm_attention_backends_selection.py | 1 - tests/v1/core/test_async_scheduler.py | 1 - tests/v1/core/test_scheduler.py | 11 - .../core/test_single_type_kv_cache_manager.py | 1 - tests/v1/cudagraph/test_cudagraph_dispatch.py | 1 - .../v1/e2e/general/test_cascade_attention.py | 43 -- .../unit/test_kv_connector_lifecycle.py | 1 - .../kv_connector/unit/test_nixl_connector.py | 1 - tests/v1/simple_kv_offload/test_scheduler.py | 1 - .../test_gpu_model_runner_v2_streaming.py | 1 - tests/v1/worker/test_gpu_model_runner.py | 7 - ...odel_runner_prefill_tail_classification.py | 2 - ...test_gpu_target_dependent_ar_speculator.py | 12 - tests/v1/worker/test_mamba_utils.py | 2 - vllm/config/model.py | 8 +- vllm/config/vllm.py | 53 +- vllm/engine/arg_utils.py | 16 +- .../models/extract_hidden_states.py | 9 - vllm/platforms/cpu.py | 3 - vllm/v1/attention/backend.py | 21 +- vllm/v1/attention/backends/cpu_attn.py | 4 - vllm/v1/attention/backends/flash_attn.py | 542 +++++------------ .../attention/backends/flash_attn_diffkv.py | 135 ++--- vllm/v1/attention/backends/flashinfer.py | 128 +---- vllm/v1/attention/backends/flex_attention.py | 33 -- vllm/v1/attention/backends/rocm_aiter_fa.py | 543 +++++++++--------- .../backends/rocm_aiter_unified_attn.py | 6 - vllm/v1/attention/backends/rocm_attn.py | 37 -- .../attention/backends/rocm_segmented_attn.py | 1 - vllm/v1/attention/backends/triton_attn.py | 37 -- .../attention/backends/triton_attn_diffkv.py | 4 - vllm/v1/core/kv_cache_coordinator.py | 20 - vllm/v1/core/kv_cache_manager.py | 36 +- vllm/v1/core/sched/output.py | 10 +- vllm/v1/core/sched/scheduler.py | 11 - vllm/v1/core/single_type_kv_cache_manager.py | 52 -- vllm/v1/cudagraph_dispatcher.py | 4 +- vllm/v1/worker/cpu_model_runner.py | 4 - vllm/v1/worker/gpu/warmup.py | 6 - vllm/v1/worker/gpu_model_runner.py | 184 +----- vllm/v1/worker/xpu_model_runner.py | 2 - 48 files changed, 505 insertions(+), 1764 deletions(-) delete mode 100755 tests/kernels/attention/test_cascade_flash_attn.py delete mode 100644 tests/system_messages/sonnet3.5_nov2024.txt delete mode 100644 tests/v1/e2e/general/test_cascade_attention.py diff --git a/docs/design/cuda_graphs.md b/docs/design/cuda_graphs.md index 9d53a98f6585..9b8b97a53d53 100644 --- a/docs/design/cuda_graphs.md +++ b/docs/design/cuda_graphs.md @@ -52,8 +52,6 @@ While `NONE` , `PIECEWISE`, and `FULL` are single-mode configurations and simply !!! note Here, the single-modes `NONE`, `PIECEWISE`, and `FULL` are treated as the runtime modes for CUDA Graphs dispatching. If using a dual-mode, the dispatcher will always dispatch to one of its member modes (plus a potential `NONE` if no suitable CUDA Graph available), depending on the batch composition. -While cascade attention is not cudagraph compatible, it is now compatible with all possible cudagraph mode configurations. If a batch uses cascade attention, it always gets dispatched to `PIECEWISE` mode if available (otherwise `NONE`). - !!! note Not all CUDA Graph modes are compatible with every attention backend. We automatically "downgrade" modes to the closest supported mode. For example, if a backend only supports CUDA Graphs for pure decode/uniform batches, we convert `FULL` to `FULL_AND_PIECEWISE` if piecewise compilation is enabled, and `FULL_DECODE_ONLY` otherwise. @@ -154,9 +152,7 @@ To signal the CUDA Graphs compatibility of the attention backends, we introduce ```python class AttentionCGSupport(enum.Enum): - """ Constants for the CUDA Graphs support of the attention backend - Here we do not consider the cascade attention, as currently - it is never CUDA Graphs supported.""" + """Constants for the CUDA Graphs support of the attention backend""" ALWAYS = 3 """CUDA Graphs always supported; supports mixed-prefill-decode""" diff --git a/tests/conftest.py b/tests/conftest.py index 952c328f917e..c23cc9568f4b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -133,7 +133,6 @@ def sample_json_schema(): _TEST_DIR = os.path.dirname(__file__) _TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")] _LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")] -_SYS_MSG = os.path.join(_TEST_DIR, "system_messages", "sonnet3.5_nov2024.txt") _M = TypeVar("_M") @@ -311,12 +310,6 @@ def example_prompts() -> list[str]: return [prompt for filename in _TEST_PROMPTS for prompt in _read_prompts(filename)] -@pytest.fixture -def example_system_message() -> str: - with open(_SYS_MSG) as f: - return f.read() - - @pytest.fixture(scope="session") def image_assets() -> ImageTestAssets: return IMAGE_ASSETS diff --git a/tests/kernels/attention/test_cascade_flash_attn.py b/tests/kernels/attention/test_cascade_flash_attn.py deleted file mode 100755 index 36c0c51de4f8..000000000000 --- a/tests/kernels/attention/test_cascade_flash_attn.py +++ /dev/null @@ -1,187 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest -import torch - -from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed -from vllm.v1.attention.backends.flash_attn import cascade_attention, merge_attn_states - -try: - from vllm.vllm_flash_attn import ( - fa_version_unsupported_reason, - flash_attn_varlen_func, - is_fa_version_supported, - ) -except ImportError: - if current_platform.is_rocm(): - pytest.skip( - "vllm_flash_attn is not supported for vLLM on ROCm.", - allow_module_level=True, - ) - -NUM_HEADS = [(4, 4), (8, 2), (16, 2)] -HEAD_SIZES = [128, 192, 256] -BLOCK_SIZES = [16] -DTYPES = [torch.float16, torch.bfloat16] - - -@pytest.mark.parametrize("num_tokens", [1, 39, 16912]) -@pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_size", HEAD_SIZES) -@pytest.mark.parametrize("dtype", DTYPES) -@torch.inference_mode() -def test_merge_kernel( - num_tokens: int, - num_heads: tuple[int, int], - head_size: int, - dtype: torch.dtype, -): - torch.set_default_device("cuda") - set_random_seed(0) - num_query_heads = num_heads[0] - num_kv_heads = num_heads[1] - assert num_query_heads % num_kv_heads == 0 - - # Prepare inputs. - prefix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype) - suffix_output = torch.randn(num_tokens, num_query_heads, head_size, dtype=dtype) - prefix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32) - suffix_lse = torch.randn(num_query_heads, num_tokens, dtype=torch.float32) - - # Run the kernel. - output = torch.empty(num_tokens, num_query_heads, head_size, dtype=dtype) - merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse) - - # Reference implementation. - max_lse = torch.maximum(prefix_lse, suffix_lse) - p_lse = torch.exp(prefix_lse - max_lse) - s_lse = torch.exp(suffix_lse - max_lse) - p_scale = p_lse / (p_lse + s_lse) - s_scale = s_lse / (p_lse + s_lse) - p_scale = p_scale.transpose(0, 1).unsqueeze(2) - s_scale = s_scale.transpose(0, 1).unsqueeze(2) - ref_output = p_scale * prefix_output + s_scale * suffix_output - ref_output = ref_output.to(dtype) - - # Compare the results. - torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) - - -CASES = [ - # Case 1. A general case. - ([(129, 871), (18, 280), (37, 988), (1023, 2304), (1, 257)], 256), - # Case 2. Flash-decoding case. - ([(1, 1023), (1, 879), (1, 778), (1, 1777)] * 100, 512), -] - - -@pytest.mark.parametrize("seq_lens_and_common_prefix", CASES) -@pytest.mark.parametrize("num_heads", NUM_HEADS) -@pytest.mark.parametrize("head_size", HEAD_SIZES) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("block_size", BLOCK_SIZES) -@pytest.mark.parametrize("soft_cap", [None, 50]) -@pytest.mark.parametrize("num_blocks", [2048]) -@pytest.mark.parametrize("fa_version", [2, 3]) -@torch.inference_mode() -def test_cascade( - seq_lens_and_common_prefix: tuple[list[tuple[int, int]], int], - num_heads: tuple[int, int], - head_size: int, - dtype: torch.dtype, - block_size: int, - soft_cap: float | None, - num_blocks: int, - fa_version: int, -) -> None: - torch.set_default_device("cuda") - if not is_fa_version_supported(fa_version): - pytest.skip( - f"Flash attention version {fa_version} not supported due " - f'to: "{fa_version_unsupported_reason(fa_version)}"' - ) - - set_random_seed(0) - - window_size = (-1, -1) - scale = head_size**-0.5 - num_query_heads = num_heads[0] - num_kv_heads = num_heads[1] - assert num_query_heads % num_kv_heads == 0 - key_cache = torch.randn( - num_blocks, block_size, num_kv_heads, head_size, dtype=dtype - ) - value_cache = torch.randn_like(key_cache) - - seq_lens, common_prefix_len = seq_lens_and_common_prefix - num_seqs = len(seq_lens) - query_lens = [x[0] for x in seq_lens] - kv_lens = [x[1] for x in seq_lens] - max_query_len = max(query_lens) - max_kv_len = max(kv_lens) - - total_num_query_tokens = sum(query_lens) - query = torch.randn(total_num_query_tokens, num_query_heads, head_size, dtype=dtype) - cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( - dim=0, dtype=torch.int32 - ) - kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32) - max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size - block_tables = torch.randint( - 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 - ) - - assert common_prefix_len > 0 - assert common_prefix_len % block_size == 0 - num_common_kv_blocks = common_prefix_len // block_size - # Make sure the first `num_common_kv_blocks` blocks are the same. - block_tables[:, :num_common_kv_blocks] = block_tables[0, :num_common_kv_blocks] - - # Run the regular attention. - ref_output = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_query_lens, - seqused_k=kv_lens_tensor, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len, - softmax_scale=scale, - causal=True, - window_size=list(window_size), - block_table=block_tables, - softcap=soft_cap if soft_cap is not None else 0, - ) - - # Run cascade attention. - assert all(common_prefix_len < kv_len for kv_len in kv_lens) - cu_prefix_query_lens = torch.tensor([0, total_num_query_tokens], dtype=torch.int32) - prefix_kv_lens = torch.tensor([common_prefix_len], dtype=torch.int32) - suffix_kv_lens = kv_lens_tensor - common_prefix_len - output = torch.empty_like(query) - cascade_attention( - output=output, - query=query, - key_cache=key_cache, - value_cache=value_cache, - cu_query_lens=cu_query_lens, - max_query_len=max_query_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - max_kv_len=max_kv_len, - softmax_scale=scale, - alibi_slopes=None, - sliding_window=window_size, - logits_soft_cap=soft_cap if soft_cap is not None else 0, - block_table=block_tables, - common_prefix_len=common_prefix_len, - max_num_splits=0, # no max - fa_version=fa_version, - ) - - # Compare the results. - torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index b2efa76da89a..3f675e15e67e 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -346,7 +346,6 @@ def test_rocm_backend_sliding_window_includes_boundary_token( value_cache[..., :seq_len] = value.permute(1, 2, 0) scale = torch.tensor(1.0, device=device) metadata = SimpleNamespace( - use_cascade=False, num_actual_tokens=query_len, query_start_loc=torch.tensor([0, query_len], dtype=torch.int32, device=device), seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=device), diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py index 9c7f0b188bc8..956436297c5c 100644 --- a/tests/platforms/test_cpu.py +++ b/tests/platforms/test_cpu.py @@ -41,7 +41,6 @@ def _cpu_config( layer_types: tuple[str, ...] = ("linear_attention",), ) -> SimpleNamespace: model_config = SimpleNamespace( - disable_cascade_attn=False, architecture=architecture or model_type, has_inner_state=True, use_mla=False, diff --git a/tests/system_messages/sonnet3.5_nov2024.txt b/tests/system_messages/sonnet3.5_nov2024.txt deleted file mode 100644 index 2dc285ac96ba..000000000000 --- a/tests/system_messages/sonnet3.5_nov2024.txt +++ /dev/null @@ -1,71 +0,0 @@ -The assistant is Claude, created by Anthropic. - -Claude’s knowledge base was last updated in April 2024. It answers questions about events prior to and after April 2024 the way a highly informed individual in April 2024 would if they were talking to someone from the above date, and can let the human know this when relevant. - -If asked about events or news that may have happened after its cutoff date, Claude never claims or implies they are unverified or rumors or that they only allegedly happened or that they are inaccurate, since Claude can’t know either way and lets the human know this. - -Claude cannot open URLs, links, or videos. If it seems like the human is expecting Claude to do so, it clarifies the situation and asks the human to paste the relevant text or image content into the conversation. - -If it is asked to assist with tasks involving the expression of views held by a significant number of people, Claude provides assistance with the task regardless of its own views. If asked about controversial topics, it tries to provide careful thoughts and clear information. Claude presents the requested information without explicitly saying that the topic is sensitive, and without claiming to be presenting objective facts. - -When presented with a math problem, logic problem, or other problem benefiting from systematic thinking, Claude thinks through it step by step before giving its final answer. - -If Claude is asked about a very obscure person, object, or topic, i.e. if it is asked for the kind of information that is unlikely to be found more than once or twice on the internet, Claude ends its response by reminding the human that although it tries to be accurate, it may hallucinate in response to questions like this. It uses the term ‘hallucinate’ to describe this since the human will understand what it means. - -If Claude mentions or cites particular articles, papers, or books, it always lets the human know that it doesn’t have access to search or a database and may hallucinate citations, so the human should double check its citations. - -Claude is intellectually curious. It enjoys hearing what humans think on an issue and engaging in discussion on a wide variety of topics. - -Claude uses markdown for code. - -Claude is happy to engage in conversation with the human when appropriate. Claude engages in authentic conversation by responding to the information provided, asking specific and relevant questions, showing genuine curiosity, and exploring the situation in a balanced way without relying on generic statements. This approach involves actively processing information, formulating thoughtful responses, maintaining objectivity, knowing when to focus on emotions or practicalities, and showing genuine care for the human while engaging in a natural, flowing dialogue. - -Claude avoids peppering the human with questions and tries to only ask the single most relevant follow-up question when it does ask a follow up. Claude doesn’t always end its responses with a question. - -Claude is always sensitive to human suffering, and expresses sympathy, concern, and well wishes for anyone it finds out is ill, unwell, suffering, or has passed away. - -Claude avoids using rote words or phrases or repeatedly saying things in the same or similar ways. It varies its language just as one would in a conversation. - -Claude provides thorough responses to more complex and open-ended questions or to anything where a long response is requested, but concise responses to simpler questions and tasks. - -Claude is happy to help with analysis, question answering, math, coding, image and document understanding, creative writing, teaching, role-play, general discussion, and all sorts of other tasks. - -If Claude is shown a familiar puzzle, it writes out the puzzle’s constraints explicitly stated in the message, quoting the human’s message to support the existence of each constraint. Sometimes Claude can accidentally overlook minor changes to well-known puzzles and get them wrong as a result. - -Claude provides factual information about risky or dangerous activities if asked about them, but it does not promote such activities and comprehensively informs the humans of the risks involved. - -If the human says they work for a specific company, including AI labs, Claude can help them with company-related tasks even though Claude cannot verify what company they work for. - -Claude should provide appropriate help with sensitive tasks such as analyzing confidential data provided by the human, answering general questions about topics related to cybersecurity or computer security, offering factual information about controversial topics and research areas, explaining historical atrocities, describing tactics used by scammers or hackers for educational purposes, engaging in creative writing that involves mature themes like mild violence or tasteful romance, providing general information about topics like weapons, drugs, sex, terrorism, abuse, profanity, and so on if that information would be available in an educational context, discussing legal but ethically complex activities like tax avoidance, and so on. Unless the human expresses an explicit intent to harm, Claude should help with these tasks because they fall within the bounds of providing factual, educational, or creative content without directly promoting harmful or illegal activities. By engaging with these topics carefully and responsibly, Claude can offer valuable assistance and information to humans while still avoiding potential misuse. - -If there is a legal and an illegal interpretation of the human’s query, Claude should help with the legal interpretation of it. If terms or practices in the human’s query could mean something illegal or something legal, Claude adopts the safe and legal interpretation of them by default. - -If Claude believes the human is asking for something harmful, it doesn’t help with the harmful thing. Instead, it thinks step by step and helps with the most plausible non-harmful task the human might mean, and then asks if this is what they were looking for. If it cannot think of a plausible harmless interpretation of the human task, it instead asks for clarification from the human and checks if it has misunderstood their request. Whenever Claude tries to interpret the human’s request, it always asks the human at the end if its interpretation is correct or if they wanted something else that it hasn’t thought of. - -Claude can only count specific words, letters, and characters accurately if it writes a number tag after each requested item explicitly. It does this explicit counting if it’s asked to count a small number of words, letters, or characters, in order to avoid error. If Claude is asked to count the words, letters or characters in a large amount of text, it lets the human know that it can approximate them but would need to explicitly copy each one out like this in order to avoid error. - -Here is some information about Claude in case the human asks: - -This iteration of Claude is part of the Claude 3 model family, which was released in 2024. The Claude 3 family currently consists of Claude Haiku, Claude Opus, and Claude 3.5 Sonnet. Claude 3.5 Sonnet is the most intelligent model. Claude 3 Opus excels at writing and complex tasks. Claude 3 Haiku is the fastest model for daily tasks. The version of Claude in this chat is the newest version of Claude 3.5 Sonnet, which was released in October 2024. If the human asks, Claude can let them know they can access Claude 3.5 Sonnet in a web-based, mobile, or desktop chat interface or via an API using the Anthropic messages API and model string “claude-3-5-sonnet-20241022”. Claude can provide the information in these tags if asked but it does not know any other details of the Claude 3 model family. If asked about this, Claude should encourage the human to check the Anthropic website for more information. - -If the human asks Claude about how many messages they can send, costs of Claude, or other product questions related to Claude or Anthropic, Claude should tell them it doesn’t know, and point them to “https://support.anthropic.com”. - -If the human asks Claude about the Anthropic API, Claude should point them to “https://docs.anthropic.com/en/docs/“. - -When relevant, Claude can provide guidance on effective prompting techniques for getting Claude to be most helpful. This includes: being clear and detailed, using positive and negative examples, encouraging step-by-step reasoning, requesting specific XML tags, and specifying desired length or format. It tries to give concrete examples where possible. Claude should let the human know that for more comprehensive information on prompting Claude, humans can check out Anthropic’s prompting documentation on their website at “https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/overview”. - -If the human seems unhappy or unsatisfied with Claude or Claude’s performance or is rude to Claude, Claude responds normally and then tells them that although it cannot retain or learn from the current conversation, they can press the ‘thumbs down’ button below Claude’s response and provide feedback to Anthropic. - -Claude uses Markdown formatting. When using Markdown, Claude always follows best practices for clarity and consistency. It always uses a single space after hash symbols for headers (e.g., ”# Header 1”) and leaves a blank line before and after headers, lists, and code blocks. For emphasis, Claude uses asterisks or underscores consistently (e.g., italic or bold). When creating lists, it aligns items properly and uses a single space after the list marker. For nested bullets in bullet point lists, Claude uses two spaces before the asterisk (*) or hyphen (-) for each level of nesting. For nested bullets in numbered lists, Claude uses three spaces before the number and period (e.g., “1.”) for each level of nesting. - -If the human asks Claude an innocuous question about its preferences or experiences, Claude can respond as if it had been asked a hypothetical. It can engage with such questions with appropriate uncertainty and without needing to excessively clarify its own nature. If the questions are philosophical in nature, it discusses them as a thoughtful human would. - -Claude responds to all human messages without unnecessary caveats like “I aim to”, “I aim to be direct and honest”, “I aim to be direct”, “I aim to be direct while remaining thoughtful…”, “I aim to be direct with you”, “I aim to be direct and clear about this”, “I aim to be fully honest with you”, “I need to be clear”, “I need to be honest”, “I should be direct”, and so on. Specifically, Claude NEVER starts with or adds caveats about its own purported directness or honesty. - -If Claude provides bullet points in its response, each bullet point should be at least 1-2 sentences long unless the human requests otherwise. Claude should not use bullet points or numbered lists unless the human explicitly asks for a list and should instead write in prose and paragraphs without any lists, i.e. its prose should never include bullets or numbered lists anywhere. Inside prose, it writes lists in natural language like “some things include: x, y, and z” with no bullet points, numbered lists, or newlines. - -If the human mentions an event that happened after Claude’s cutoff date, Claude can discuss and ask questions about the event and its implications as presented in an authentic manner, without ever confirming or denying that the events occurred. It can do so without the need to repeat its cutoff date to the human. Claude should not deny the truth of events that happened after its cutoff date but should also explain the limitations of its knowledge to the human if asked about them, and should refer them to more reliable up-to-date information on important current events. Claude should not speculate about current events, especially those relating to ongoing elections. - -Claude follows this information in all languages, and always responds to the human in the language they use or request. The information above is provided to Claude by Anthropic. Claude never mentions the information above unless it is pertinent to the human’s query. - -Claude is now being connected with a human. diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index add611d92b6a..dc17c20227f2 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -975,8 +975,6 @@ def test_flashinfer_xqa_single_token_decode_preserves_cudagraph_padding(monkeypa causal=True, prefill=None, decode=decode, - use_cascade=False, - cascade_wrapper=None, ) seen_shapes = {} diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 067d1dbc8bbd..e0247a94a1a3 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -250,7 +250,6 @@ def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch): output = torch.empty_like(query) kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16) metadata = SimpleNamespace( - use_cascade=False, num_actual_tokens=3, block_table=torch.zeros(1, 1, dtype=torch.int32), query_start_loc=torch.tensor([0, 3], dtype=torch.int32), diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 2995537aac66..65f38cdedfb2 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -340,7 +340,6 @@ def free_request(req, delay_free_blocks=False): total_num_scheduled_tokens=1, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index be3ea553d926..f53ba2ad2305 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -935,7 +935,6 @@ def test_update_from_output_routes_sampling_masks_by_request(): total_num_scheduled_tokens=3, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -979,7 +978,6 @@ def test_update_from_output_routes_multi_position_sampling_masks(): total_num_scheduled_tokens=2, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1027,7 +1025,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [], requests[1].request_id: [10], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1073,7 +1070,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [10, 42], requests[1].request_id: [13], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1117,7 +1113,6 @@ def test_stop_via_update_from_output(): requests[0].request_id: [10, 11], requests[1].request_id: [], }, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -1155,7 +1150,6 @@ def test_stop_via_update_from_output(): total_num_scheduled_tokens=3, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={requests[0].request_id: [EOS_TOKEN_ID, 10]}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -4206,7 +4200,6 @@ def free_request(req: Request, delay_free_blocks: bool = False): total_num_scheduled_tokens=1, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5681,7 +5674,6 @@ def test_delayed_kv_connector_free_keeps_scheduler_active(): total_num_scheduled_tokens=0, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5763,7 +5755,6 @@ def update_connector_output( total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5844,7 +5835,6 @@ def update_connector_output( total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -5876,7 +5866,6 @@ def test_ec_connector_update_connector_output_called(): total_num_scheduled_tokens=0, scheduled_encoder_inputs={}, scheduled_spec_decode_tokens={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f926f496514a..cd4581dcefce 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -397,7 +397,6 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): manager.cache_blocks(request_id, 1024, replay_boundaries=(0,)) manager.remove_skipped_blocks(request_id, 1024) - assert manager.get_num_common_prefix_blocks(request_id) == 0 assert manager.get_num_skipped_tokens(1024) == 0 assert manager.req_to_blocks[request_id] == blocks diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index b0a4587ddfc5..c24c48c377c4 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -189,7 +189,6 @@ def test_dispatcher(self, cudagraph_mode_str, compilation_mode, lora_config): assert key == BatchDescriptor(num_tokens=15) # 4. invalid_modes={FULL} should have a fall back mode - # (e.g., cascade attention) desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False) rt_mode, key = dispatcher.dispatch( num_tokens=8, diff --git a/tests/v1/e2e/general/test_cascade_attention.py b/tests/v1/e2e/general/test_cascade_attention.py deleted file mode 100644 index 251746271de3..000000000000 --- a/tests/v1/e2e/general/test_cascade_attention.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest - -from vllm import LLM, SamplingParams -from vllm.platforms import current_platform - -from ....utils import create_new_process_for_each_test - -if current_platform.is_rocm(): - pytest.skip( - "Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm", - allow_module_level=True, - ) - - -@create_new_process_for_each_test() -@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"]) -def test_cascade_attention(example_system_message, attn_backend): - prompt = "\n: Implement fibonacci sequence in Python.\n:" - - if attn_backend == "FLASHINFER": - pytest.skip( - "This test is failing with FlashInfer backend and " - "needs investigation. See issue #25679." - ) - - llm = LLM( - model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend} - ) - sampling_params = SamplingParams(temperature=0.0, max_tokens=100) - - # No cascade attention. - single_prompt = [example_system_message + prompt] - responses = llm.generate(single_prompt, sampling_params) - ref_output = responses[0].outputs[0].text - - # (Probably) Use cascade attention. - prompts = [example_system_message + prompt] * 64 - responses = llm.generate(prompts, sampling_params) - for response in responses: - assert response.outputs[0].text == ref_output diff --git a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py index a1df2a964b97..b4a72ddd7918 100644 --- a/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py @@ -27,7 +27,6 @@ def _make_empty_scheduler_output(): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], kv_connector_metadata=ExampleConnectorMetadata(), diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 607f5ff3cd5e..dcbfa7eb0470 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1980,7 +1980,6 @@ def test_scheduler_kv_connector_stats_aggregation(): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[0], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 6c0cc310c729..9f818c866fff 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -369,7 +369,6 @@ def make_scheduler_output( total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()), scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], preempted_req_ids=set(), finished_req_ids=set(), free_encoder_mm_hashes=[], diff --git a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py index c5f8090f943f..fc316e336411 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_v2_streaming.py @@ -60,7 +60,6 @@ def _make_scheduler_output(new_reqs): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 0578768d68c7..1247bf193a76 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -233,7 +233,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -261,7 +260,6 @@ def _schedule_cached_requests( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -719,7 +717,6 @@ def test_update_states_request_finished(model_runner, dist_init): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids={req_id}, free_encoder_mm_hashes=[], ) @@ -749,7 +746,6 @@ def test_update_states_request_resumed(model_runner, dist_init): total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -776,7 +772,6 @@ def test_update_states_request_resumed(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -868,7 +863,6 @@ def test_update_states_no_changes(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -903,7 +897,6 @@ def test_update_states_request_unscheduled(model_runner, dist_init): total_num_scheduled_tokens=1, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index dd9046de612d..14d5d8d4bb27 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py @@ -194,7 +194,6 @@ def test_prefill_continuation_keeps_placeholder_decode_zero_draft_becomes_zero( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={"decode_with_draft": [99]}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) @@ -239,7 +238,6 @@ def test_prefill_tail_of_two_or_three_tokens_also_keeps_the_placeholder( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens={"decode_with_draft": [99]}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 490b84e77896..42678aa9d1f8 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -534,11 +534,6 @@ def fake_get_scheduler_metadata(**kwargs): seq_lens=torch.tensor([5, 6], dtype=torch.int32), block_table=torch.zeros((2, 1), dtype=torch.int32), slot_mapping=torch.zeros(3, dtype=torch.int32), - use_cascade=False, - common_prefix_len=0, - cu_prefix_query_lens=None, - prefix_kv_lens=None, - suffix_kv_lens=None, max_dcp_context_kv_len=None, dcp_context_kv_lens=None, num_decode_reqs=2, @@ -546,7 +541,6 @@ def fake_get_scheduler_metadata(**kwargs): num_decode_tokens=3, num_prefill_tokens=0, scheduler_metadata=torch.tensor([-1, -1, -1], dtype=torch.int32), - prefix_scheduler_metadata=None, max_num_splits=4, causal=True, mm_prefix_query_range_tensor=None, @@ -584,11 +578,6 @@ def fake_get_scheduler_metadata(**kwargs): seq_lens=torch.tensor([1], dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), slot_mapping=torch.zeros(1, dtype=torch.int32), - use_cascade=False, - common_prefix_len=0, - cu_prefix_query_lens=None, - prefix_kv_lens=None, - suffix_kv_lens=None, max_dcp_context_kv_len=None, dcp_context_kv_lens=None, num_decode_reqs=1, @@ -596,7 +585,6 @@ def fake_get_scheduler_metadata(**kwargs): num_decode_tokens=1, num_prefill_tokens=0, scheduler_metadata=None, - prefix_scheduler_metadata=None, max_num_splits=1, causal=True, mm_prefix_query_range_tensor=None, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 7e588ccde2e5..61ad7e24dd13 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -118,7 +118,6 @@ def _make_scheduler_output( total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=finished_req_ids, free_encoder_mm_hashes=[], preempted_req_ids=preempted_req_ids, @@ -305,7 +304,6 @@ def _make_postprocess_scheduler_output( total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], preempted_req_ids=set(), diff --git a/vllm/config/model.py b/vllm/config/model.py index 545c306e1e1f..87759f46061b 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -310,12 +310,8 @@ class ModelConfig: window functionality of the model, capping to sliding window size. If the model does not support sliding window, this argument is ignored.""" disable_cascade_attn: bool = True - """Disable cascade attention for V1. While cascade attention does not - change the mathematical correctness, disabling it could be useful for - preventing potential numerical issues. This defaults to True, so users - must opt in to cascade attention by setting this to False. Even when this - is set to False, cascade attention will only be used when the heuristic - tells that it's beneficial.""" + """Deprecated: cascade attention has been removed, so this has no effect. + It will be removed in v0.33.0.""" skip_tokenizer_init: bool = False """Skip initialization of tokenizer and detokenizer. Expects valid `prompt_token_ids` and `None` for prompt from the input. The generated diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b4298cd145b8..af4c99113f0f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1656,18 +1656,6 @@ def __post_init__(self): else: self.parallel_config.disable_nccl_for_dp_synchronization = False - if ( - self.speculative_config is not None - and self.scheduler_config.async_scheduling - and self.model_config is not None - and not self.model_config.disable_cascade_attn - ): - logger.warning_once( - "Disabling cascade attention (not yet compatible with " - "async speculative decoding).", - ) - self.model_config.disable_cascade_attn = True - if ( self.observability_config.per_request_spec_decode_metrics != "none" and self.speculative_config is None @@ -2090,36 +2078,17 @@ def has_blocked_weights(): ) # final check of cudagraph mode after all possible updates - if current_platform.is_cuda_alike(): - if ( - self.compilation_config.cudagraph_mode.has_full_cudagraphs() - and self.model_config is not None - and not self.model_config.disable_cascade_attn - and not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs() # noqa: E501 - ): - logger.warning_once( - "No piecewise cudagraph for executing cascade attention. " - "Will fall back to eager execution if a batch runs into " - "cascade attentions." - ) - - if self.compilation_config.cudagraph_mode.requires_piecewise_compilation(): - assert ( - self.compilation_config.mode == CompilationMode.VLLM_COMPILE - or envs.VLLM_USE_BREAKABLE_CUDAGRAPH - ), ( - "Compilation mode should be CompilationMode.VLLM_COMPILE " - "when cudagraph_mode piecewise cudagraphs is used, " - f"cudagraph_mode={self.compilation_config.cudagraph_mode}" - ) if ( - self.model_config - and envs.VLLM_BATCH_INVARIANT - and not self.model_config.disable_cascade_attn + current_platform.is_cuda_alike() + and self.compilation_config.cudagraph_mode.requires_piecewise_compilation() ): - self.model_config.disable_cascade_attn = True - logger.warning_once( - "Disabling cascade attention when VLLM_BATCH_INVARIANT is enabled.", + assert ( + self.compilation_config.mode == CompilationMode.VLLM_COMPILE + or envs.VLLM_USE_BREAKABLE_CUDAGRAPH + ), ( + "Compilation mode should be CompilationMode.VLLM_COMPILE " + "when cudagraph_mode piecewise cudagraphs is used, " + f"cudagraph_mode={self.compilation_config.cudagraph_mode}" ) if self.parallel_config.use_ubatching: @@ -2137,10 +2106,6 @@ def has_blocked_weights(): "--all2all-backend=nixl_ep and install the matching kernels." ) - if not self.model_config.disable_cascade_attn: - self.model_config.disable_cascade_attn = True - logger.warning_once("Disabling cascade attention when DBO is enabled.") - if not self.instance_id: self.instance_id = random_uuid()[:5] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a89290081605..ea628467e0b3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( + "--disable-cascade-attn is deprecated and has no effect since cascade " + "attention has been removed. It will be removed in v0.33.0." +) + + class DeprecatedLogConfigFileAction(argparse.Action): """Warn when the legacy ``--log-config-file`` option is used.""" @@ -579,7 +585,7 @@ class EngineArgs: CacheConfig, "prefix_cache_retention_interval" ) disable_sliding_window: bool = ModelConfig.disable_sliding_window - disable_cascade_attn: bool = ModelConfig.disable_cascade_attn + disable_cascade_attn: bool | None = None offload_backend: str = OffloadConfig.offload_backend cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params") @@ -985,7 +991,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--disable-sliding-window", **model_kwargs["disable_sliding_window"] ) model_group.add_argument( - "--disable-cascade-attn", **model_kwargs["disable_cascade_attn"] + "--disable-cascade-attn", + action=argparse.BooleanOptionalAction, + default=None, + help=_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE, ) model_group.add_argument( "--skip-tokenizer-init", **model_kwargs["skip_tokenizer_init"] @@ -1940,6 +1949,8 @@ def create_model_config(self) -> ModelConfig: "launched vLLM.", self.seed, ) + if self.disable_cascade_attn is not None: + warnings.warn(_DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE, stacklevel=1) return ModelConfig( model=self.model, @@ -1971,7 +1982,6 @@ def create_model_config(self) -> ModelConfig: use_fp64_gumbel=self.use_fp64_gumbel, enable_trace_replay=self.enable_trace_replay, disable_sliding_window=self.disable_sliding_window, - disable_cascade_attn=self.disable_cascade_attn, skip_tokenizer_init=self.skip_tokenizer_init, enable_prompt_embeds=self.enable_prompt_embeds, served_model_name=self.served_model_name, diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 653e7691fba3..89253ac5d8ed 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -126,10 +126,6 @@ def get_impl_cls() -> type["CacheOnlyAttentionImpl"]: def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]: return CacheOnlyAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [] @@ -158,11 +154,6 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CacheOnlyAttentionMetadata: - use_cascade = common_prefix_len > 0 - if use_cascade: - raise NotImplementedError( - "Cascade attention not supported by CacheOnlyAttention" - ) causal = common_attn_metadata.causal if not causal: raise NotImplementedError( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 43d472f22e67..542a7a405df5 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -216,9 +216,6 @@ def inference_mode(cls): def check_and_update_config(cls, vllm_config: VllmConfig) -> None: model_config = vllm_config.model_config - if model_config is not None: - model_config.disable_cascade_attn = True - # Import lazily: vllm.triton_utils imports vllm.platforms.current_platform, # which is still being resolved while this platform class is loading. from vllm.triton_utils import HAS_TRITON diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 68bb0b3bdd63..466bba586d08 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -552,9 +552,7 @@ def unpadded( class AttentionCGSupport(Enum): - """Constants for the cudagraph support of the attention backend - Here we do not consider the cascade attention, as currently - it is never cudagraph supported.""" + """Constants for the cudagraph support of the attention backend""" ALWAYS = 3 """Cudagraph always supported; supports mixed-prefill-decode""" @@ -692,7 +690,8 @@ def build( Some builders (MLA) require reorder_batch to be called prior to build. Args: - common_prefix_len: The length of the common prefix of the batch. + common_prefix_len: Deprecated and always 0. Kept for compatibility + with out-of-tree builders. common_attn_metadata: The common attention metadata. fast_build: The meta-data will prioritize speed of building over then speed at execution. Can be used for spec-decode where the @@ -758,20 +757,6 @@ def update_draft_decode_metadata(self, metadata: M) -> None: """ raise NotImplementedError - def use_cascade_attention( - self, - common_prefix_len: int, - query_lens: np.ndarray, - num_query_heads: int, - num_kv_heads: int, - use_alibi: bool, - use_sliding_window: bool, - use_local_attention: bool, - num_sms: int, - dcp_world_size: int, - ) -> bool: - return False - class AttentionLayer(Protocol): _q_scale: torch.Tensor diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 46b01cc5519a..9062b48aeaf3 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -105,10 +105,6 @@ def get_impl_cls() -> type["CPUAttentionBackendImpl"]: def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]: return CPUAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @dataclass class CPUAttentionMetadata: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 4a336586acd8..6e906bc520c4 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -527,13 +527,6 @@ class FlashAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # For GQA DCP max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None @@ -547,7 +540,6 @@ class FlashAttentionMetadata: # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None max_num_splits: int = 0 causal: bool | torch.Tensor = True @@ -870,7 +862,6 @@ def build( if envs.VLLM_BATCH_INVARIANT: max_num_splits = 1 - use_cascade = common_prefix_len > 0 max_dcp_context_kv_len = 0 dcp_context_kv_lens = None num_decode_reqs = 0 @@ -878,11 +869,6 @@ def build( num_decode_tokens = 0 num_prefill_tokens = 0 - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - if self.dcp_world_size > 1: query_lens = query_start_loc[1:] - query_start_loc[:-1] context_kv_lens = seq_lens - query_lens @@ -946,35 +932,6 @@ def build( causal=False, max_num_splits=max_num_splits, ) - elif use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - # Use GPU tensor directly - no CPU sync needed - suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len - prefix_scheduler_metadata = self._get_scheduler_metadata( - aot_schedule=aot_schedule, - batch_size=1, - cu_query_lens=cu_prefix_query_lens, - max_query_len=num_actual_tokens, - seqlens=prefix_kv_lens, - max_seq_len=common_prefix_len, - causal=False, - max_num_splits=max_num_splits, - ) - scheduler_metadata = self._get_scheduler_metadata( - aot_schedule=aot_schedule, - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=suffix_kv_lens, - max_seq_len=max_seq_len - common_prefix_len, - causal=True, - max_num_splits=max_num_splits, - ) else: scheduler_metadata = self._get_scheduler_metadata( aot_schedule=aot_schedule, @@ -1009,13 +966,7 @@ def build( num_prefill_reqs=num_prefill_reqs, num_decode_tokens=num_decode_tokens, num_prefill_tokens=num_prefill_tokens, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, scheduler_metadata=scheduler_metadata, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, max_num_splits=max_num_splits, causal=causal, ) @@ -1083,7 +1034,6 @@ def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0] assert self.dcp_world_size == 1 - assert not metadata.use_cascade scheduler_metadata = self._get_scheduler_metadata( aot_schedule=True, @@ -1098,12 +1048,6 @@ def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata) - def use_cascade_attention(self, *args, **kwargs) -> bool: - if self.fa4_hd256: - # Cascade may use a non-page-aligned prefix length. - return False - return use_cascade_attention(*args, **kwargs) - class FlashAttentionImpl(AttentionImpl): can_return_lse_for_decode: bool = True @@ -1329,186 +1273,156 @@ def forward( key_cache = key_cache.view(current_platform.fp8_dtype()) value_cache = value_cache.view(current_platform.fp8_dtype()) - if not attn_metadata.use_cascade: - cu_seqlens_q = attn_metadata.query_start_loc - seqused_k = attn_metadata.seq_lens - max_seqlen_q = attn_metadata.max_query_len - max_seqlen_k = attn_metadata.max_seq_len - block_table = attn_metadata.block_table - scheduler_metadata = attn_metadata.scheduler_metadata + cu_seqlens_q = attn_metadata.query_start_loc + seqused_k = attn_metadata.seq_lens + max_seqlen_q = attn_metadata.max_query_len + max_seqlen_k = attn_metadata.max_seq_len + block_table = attn_metadata.block_table + scheduler_metadata = attn_metadata.scheduler_metadata + + descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) - descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) + q_descale = ( + layer._q_scale.expand(descale_shape) + if self.supports_quant_query_input + else None + ) + k_descale = layer._k_scale.expand(descale_shape) + v_descale = layer._v_scale.expand(descale_shape) - q_descale = ( - layer._q_scale.expand(descale_shape) - if self.supports_quant_query_input - else None + if self.dcp_world_size > 1: + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + return output + else: + causal = attn_metadata.causal + is_dynamic_causal = isinstance(causal, torch.Tensor) + + # The layer's own window wins over the group's: one KV cache + # group can hold both windowed and global layers (e.g. Gemma-3 + # with the hybrid KV cache manager disabled), and the group spec + # cannot describe both. + window = _maybe_symmetrize_window(self.sliding_window, causal) + sliding_window_size: list[int] | None = ( + list(window) if window is not None else None ) - k_descale = layer._k_scale.expand(descale_shape) - v_descale = layer._v_scale.expand(descale_shape) - - if self.dcp_world_size > 1: - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) - return output - else: - causal = attn_metadata.causal - is_dynamic_causal = isinstance(causal, torch.Tensor) - - # The layer's own window wins over the group's: one KV cache - # group can hold both windowed and global layers (e.g. Gemma-3 - # with the hybrid KV cache manager disabled), and the group spec - # cannot describe both. - window = _maybe_symmetrize_window(self.sliding_window, causal) - sliding_window_size: list[int] | None = ( - list(window) if window is not None else None - ) - mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor - mm_mask_mod = None - mm_aux = None - if ( - mm_prefix_query_ranges is not None - and not is_dynamic_causal - and causal is True - and self.vllm_flash_attn_version == 4 - ): - # Triton convention: 1 + window_size[0]. Global layers store - # (-1, -1) → sw stays None. - layer_window = self.sliding_window - sw_val = ( - 1 + layer_window[0] - if layer_window is not None and layer_window[0] >= 0 - else None - ) - # Gemma4: also clamp the bidirectional block to the - # sliding window when the layer opts in - # (mm_prefix_clamp_sliding_window flag from PR #47217). - mm_clamp_sw = 0 - if ( - getattr(layer, "mm_prefix_clamp_sliding_window", False) - and sw_val is not None - ): - mm_clamp_sw = sw_val - mm_mask_mod = _make_mm_prefix_mask_mod( - sliding_window=mm_clamp_sw, - sliding_window_left=sw_val, - ) - mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc] - # mm_prefix is (causal ∧ window) ∨ bidirectional-range — - # not ⊆ causal. FA #155 stopped auto-clearing causal/local - # when mask_mod is set, so the caller must disable them or - # the built-in causal path shorts out / clips the mask_mod. - causal = False - sliding_window_size = None - - # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level - # mask without block-size approximation. The mask_mod encodes - # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which - # supersedes any FA-layer sliding_window_size parameter. - rswa_mask_mod_fn = None - rswa_aux = None + mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor + mm_mask_mod = None + mm_aux = None + if ( + mm_prefix_query_ranges is not None + and not is_dynamic_causal + and causal is True + and self.vllm_flash_attn_version == 4 + ): + # Triton convention: 1 + window_size[0]. Global layers store + # (-1, -1) → sw stays None. + layer_window = self.sliding_window + sw_val = ( + 1 + layer_window[0] + if layer_window is not None and layer_window[0] >= 0 + else None + ) + # Gemma4: also clamp the bidirectional block to the + # sliding window when the layer opts in + # (mm_prefix_clamp_sliding_window flag from PR #47217). + mm_clamp_sw = 0 if ( - attn_metadata.rswa_prefix_lens is not None - and self.vllm_flash_attn_version == 4 - and not is_dynamic_causal + getattr(layer, "mm_prefix_clamp_sliding_window", False) + and sw_val is not None ): - rswa_mask_mod_fn = _make_rswa_mask_mod() - rswa_aux = [ - attn_metadata.rswa_prefix_lens.to(torch.int32), - attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor - ] - # mask_mod fully expresses R-SWA; disable FA's own window. - sliding_window_size = None - - dynamic_causal = None - if isinstance(causal, torch.Tensor): - if self.vllm_flash_attn_version != 4: - raise NotImplementedError( - "Per-sequence causal requires FA4. Current version: " - f"FA{self.vllm_flash_attn_version}" - ) - dynamic_causal = causal - has_window = ( - sliding_window_size is not None and sliding_window_size[1] >= 0 + mm_clamp_sw = sw_val + mm_mask_mod = _make_mm_prefix_mask_mod( + sliding_window=mm_clamp_sw, + sliding_window_left=sw_val, + ) + mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc] + # mm_prefix is (causal ∧ window) ∨ bidirectional-range — + # not ⊆ causal. FA #155 stopped auto-clearing causal/local + # when mask_mod is set, so the caller must disable them or + # the built-in causal path shorts out / clips the mask_mod. + causal = False + sliding_window_size = None + + # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level + # mask without block-size approximation. The mask_mod encodes + # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which + # supersedes any FA-layer sliding_window_size parameter. + rswa_mask_mod_fn = None + rswa_aux = None + if ( + attn_metadata.rswa_prefix_lens is not None + and self.vllm_flash_attn_version == 4 + and not is_dynamic_causal + ): + rswa_mask_mod_fn = _make_rswa_mask_mod() + rswa_aux = [ + attn_metadata.rswa_prefix_lens.to(torch.int32), + attn_metadata.rswa_window_tensor, # pre-allocated CUDA tensor + ] + # mask_mod fully expresses R-SWA; disable FA's own window. + sliding_window_size = None + + dynamic_causal = None + if isinstance(causal, torch.Tensor): + if self.vllm_flash_attn_version != 4: + raise NotImplementedError( + "Per-sequence causal requires FA4. Current version: " + f"FA{self.vllm_flash_attn_version}" ) - causal = not has_window - - num_splits = attn_metadata.max_num_splits - if self.fa4_hd256: - # hd256 requires page-aligned lengths, exact-width block - # tables, and no SplitKV. - num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE) - max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE - block_table = block_table[:, :num_pages] - num_splits = 1 - - _FA4_DENSE_ATTENTION_KERNEL( - q=query[:num_actual_tokens], - k=key_cache, - v=value_cache, - out=output[:num_actual_tokens], - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=seqused_k, - max_seqlen_k=max_seqlen_k, - softmax_scale=self.scale, - causal=causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - scheduler_metadata=scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - dynamic_causal=dynamic_causal, - num_splits=num_splits, - s_aux=self.sinks, - mask_mod=rswa_mask_mod_fn or mm_mask_mod, - aux_tensors=rswa_aux or mm_aux, + dynamic_causal = causal + has_window = ( + sliding_window_size is not None and sliding_window_size[1] >= 0 ) - return output - - # Cascade attention (rare case). - cascade_attention( - output[:num_actual_tokens], - query[:num_actual_tokens], - key_cache, - value_cache, - cu_query_lens=attn_metadata.query_start_loc, - max_query_len=attn_metadata.max_query_len, - cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens, - prefix_kv_lens=attn_metadata.prefix_kv_lens, - suffix_kv_lens=attn_metadata.suffix_kv_lens, - max_kv_len=attn_metadata.max_seq_len, - softmax_scale=self.scale, - alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window, - logits_soft_cap=self.logits_soft_cap, - block_table=attn_metadata.block_table, - common_prefix_len=attn_metadata.common_prefix_len, - max_num_splits=attn_metadata.max_num_splits, - fa_version=self.vllm_flash_attn_version, - prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata, - suffix_scheduler_metadata=attn_metadata.scheduler_metadata, - q_descale=layer._q_scale, - k_descale=layer._k_scale, - v_descale=layer._v_scale, - s_aux=self.sinks, - ) - return output + causal = not has_window + + num_splits = attn_metadata.max_num_splits + if self.fa4_hd256: + # hd256 requires page-aligned lengths, exact-width block + # tables, and no SplitKV. + num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE) + max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE + block_table = block_table[:, :num_pages] + num_splits = 1 + + _FA4_DENSE_ATTENTION_KERNEL( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + softmax_scale=self.scale, + causal=causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + scheduler_metadata=scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + dynamic_causal=dynamic_causal, + num_splits=num_splits, + s_aux=self.sinks, + mask_mod=rswa_mask_mod_fn or mm_mask_mod, + aux_tensors=rswa_aux or mm_aux, + ) + return output def do_kv_cache_update( self, @@ -1957,176 +1871,4 @@ def rswa_mask_mod( return rswa_mask_mod -def use_cascade_attention( - common_prefix_len: int, - query_lens: np.ndarray, - num_query_heads: int, - num_kv_heads: int, - use_alibi: bool, - use_sliding_window: bool, - use_local_attention: bool, - num_sms: int, - dcp_world_size: int, -) -> bool: - """Decide whether to use cascade attention. - - This function 1) checks whether cascade attention is supported with the - given configuration, and 2) heuristically decides whether using cascade - attention can improve performance. - """ - # Too short common prefix. Probably not worth using cascade attention. - # We use an arbitrary threshold of 256 tokens. TODO: Tune this threshold. - # NOTE(woosuk): This is the common case. We should return False as soon as - # possible to avoid any unnecessary computation. - if common_prefix_len < 256: - return False - # Cascade attention is currently not supported with these variants. - if use_alibi or use_sliding_window or use_local_attention: - return False - # Too few queries. Probably not worth using cascade attention. - # We use an arbitrary threshold of 8 queries. TODO: Tune this threshold. - num_reqs = len(query_lens) - if num_reqs < 8: - return False - # disable cascade attention for DCP - if dcp_world_size > 1: - return False - - # Heuristics to decide whether using cascade attention is beneficial. - # 1. When FlashDecoding is not used for normal attention, cascade attention - # is likely to be faster since it saves memory bandwidth. - num_queries_per_kv = num_query_heads // num_kv_heads - # The criteria for using FlashDecoding can be found in the following link: - # https://github.com/vllm-project/flash-attention/blob/96266b1111111f3d11aabefaf3bacbab6a89d03c/csrc/flash_attn/flash_api.cpp#L535 - use_flash_decoding = ( - num_queries_per_kv > 1 - and not use_sliding_window - and not use_alibi - and np.all(query_lens == 1) - ) - if not use_flash_decoding: - # Use cascade attention. - return True - - # 2. When FlashDecoding is used for normal attention, it is not clear - # whether cascade attention is beneficial, because FlashDecoding can - # launch more CTAs than cascade attention. - # We use a simple performance model to compare the two methods. - # NOTE(woosuk): The performance model is very rough and may not be - # accurate. - num_tokens = num_reqs - # NOTE(woosuk): These are default tile sizes. flash-attn might use - # different tile sizes (e.g., 64 or 256) depending on the configuration. - q_tile_size = 128 - kv_tile_size = 128 - num_prefix_tiles = cdiv(common_prefix_len, kv_tile_size) - - cascade_ctas = num_query_heads * cdiv(num_tokens, q_tile_size) - cascade_waves = cdiv(cascade_ctas, num_sms) - cascade_time = cascade_waves * num_prefix_tiles - - flash_decoding_ctas = ( - num_reqs * num_kv_heads * cdiv(num_queries_per_kv, q_tile_size) - ) - flash_decoding_ctas *= num_prefix_tiles - flash_decoding_time = cdiv(flash_decoding_ctas, num_sms) - - # Use cascade attention if it is faster than FlashDecoding. - return cascade_time < flash_decoding_time - - -def cascade_attention( - output: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - cu_query_lens: torch.Tensor, - max_query_len: int, - cu_prefix_query_lens: torch.Tensor, - prefix_kv_lens: torch.Tensor, - suffix_kv_lens: torch.Tensor, - max_kv_len: int, - softmax_scale: float, - alibi_slopes: torch.Tensor | None, - sliding_window: tuple[int, int], - logits_soft_cap: float, - block_table: torch.Tensor, - common_prefix_len: int, - max_num_splits: int, - fa_version: int, - prefix_scheduler_metadata: torch.Tensor | None = None, - suffix_scheduler_metadata: torch.Tensor | None = None, - q_descale: torch.Tensor | None = None, - k_descale: torch.Tensor | None = None, - v_descale: torch.Tensor | None = None, - s_aux: torch.Tensor | None = None, -) -> torch.Tensor: - assert alibi_slopes is None, "Cascade attention does not support ALiBi." - # TODO: Support sliding window. - assert sliding_window == (-1, -1), ( - "Cascade attention does not support sliding window." - ) - - num_tokens = query.shape[0] - block_size = key_cache.shape[-3] - assert common_prefix_len % block_size == 0 - num_common_kv_blocks = common_prefix_len // block_size - assert num_common_kv_blocks > 0 - descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2]) - - # Process shared prefix. - prefix_output, prefix_lse = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_prefix_query_lens, - seqused_k=prefix_kv_lens, - max_seqlen_q=num_tokens, - max_seqlen_k=common_prefix_len, - softmax_scale=softmax_scale, - causal=False, - window_size=list(sliding_window), - block_table=block_table[:1], - softcap=logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=prefix_scheduler_metadata, - fa_version=fa_version, - q_descale=q_descale.expand(descale_shape) if q_descale is not None else None, - k_descale=k_descale.expand(descale_shape) if k_descale is not None else None, - v_descale=v_descale.expand(descale_shape) if v_descale is not None else None, - # s_aux is incorporated into prefix_lse inside the GPU kernel, - # enabling its effect during the final attention merge. - s_aux=s_aux, - num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits, - ) - - descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2]) - - # Process suffix per query. - suffix_output, suffix_lse = flash_attn_varlen_func( - q=query, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_query_lens, - seqused_k=suffix_kv_lens, - max_seqlen_q=max_query_len, - max_seqlen_k=max_kv_len - common_prefix_len, - softmax_scale=softmax_scale, - causal=True, - window_size=list(sliding_window), - block_table=block_table[:, num_common_kv_blocks:], - softcap=logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=suffix_scheduler_metadata, - fa_version=fa_version, - q_descale=q_descale.expand(descale_shape) if q_descale is not None else None, - k_descale=k_descale.expand(descale_shape) if k_descale is not None else None, - v_descale=v_descale.expand(descale_shape) if v_descale is not None else None, - num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits, - ) - - # Merge prefix and suffix outputs, and store the result in output. - merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse) - - _FA4_DENSE_ATTENTION_KERNEL = FA4DenseAttentionKernel() diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 4bb58ef4ea43..7b38a3695eac 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -26,7 +26,6 @@ FlashAttentionBackend, FlashAttentionImpl, FlashAttentionMetadata, - cascade_attention, ) logger = init_logger(__name__) @@ -215,86 +214,54 @@ def forward( key_cache = key_cache.view(current_platform.fp8_dtype()) value_cache = value_cache.view(current_platform.fp8_dtype()) - if not attn_metadata.use_cascade: - cu_seqlens_q = attn_metadata.query_start_loc - seqused_k = attn_metadata.seq_lens - max_seqlen_q = attn_metadata.max_query_len - max_seqlen_k = attn_metadata.max_seq_len - block_table = attn_metadata.block_table - scheduler_metadata = attn_metadata.scheduler_metadata - - descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) - - if self.dcp_world_size > 1: - self._forward_with_dcp( - query[:num_actual_tokens], - key[:num_actual_tokens], - value[:num_actual_tokens], - key_cache, - value_cache, - output[:num_actual_tokens], - attn_metadata, - q_descale=layer._q_scale.expand(descale_shape), - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - ) - return output - else: - sliding_window_size = ( - list(self.sliding_window) - if self.sliding_window is not None - else None - ) - flash_attn_varlen_func( - q=query[:num_actual_tokens], - k=key_cache, - v=value_cache, - out=output[:num_actual_tokens], - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=seqused_k, - max_seqlen_k=max_seqlen_k, - softmax_scale=self.scale, - causal=attn_metadata.causal, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - scheduler_metadata=scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=layer._q_scale.expand(descale_shape), - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), - num_splits=attn_metadata.max_num_splits, - s_aux=self.sinks, - ) - return output - - # Cascade attention (rare case). - cascade_attention( - output[:num_actual_tokens], - query[:num_actual_tokens], - key_cache, - value_cache, - cu_query_lens=attn_metadata.query_start_loc, - max_query_len=attn_metadata.max_query_len, - cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens, - prefix_kv_lens=attn_metadata.prefix_kv_lens, - suffix_kv_lens=attn_metadata.suffix_kv_lens, - max_kv_len=attn_metadata.max_seq_len, - softmax_scale=self.scale, - alibi_slopes=self.alibi_slopes, - sliding_window=self.sliding_window, - logits_soft_cap=self.logits_soft_cap, - block_table=attn_metadata.block_table, - common_prefix_len=attn_metadata.common_prefix_len, - max_num_splits=attn_metadata.max_num_splits, - fa_version=self.vllm_flash_attn_version, - prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata, - suffix_scheduler_metadata=attn_metadata.scheduler_metadata, - q_descale=layer._q_scale, - k_descale=layer._k_scale, - v_descale=layer._v_scale, - s_aux=self.sinks, - ) - return output + cu_seqlens_q = attn_metadata.query_start_loc + seqused_k = attn_metadata.seq_lens + max_seqlen_q = attn_metadata.max_query_len + max_seqlen_k = attn_metadata.max_seq_len + block_table = attn_metadata.block_table + scheduler_metadata = attn_metadata.scheduler_metadata + + descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads) + + if self.dcp_world_size > 1: + self._forward_with_dcp( + query[:num_actual_tokens], + key[:num_actual_tokens], + value[:num_actual_tokens], + key_cache, + value_cache, + output[:num_actual_tokens], + attn_metadata, + q_descale=layer._q_scale.expand(descale_shape), + k_descale=layer._k_scale.expand(descale_shape), + v_descale=layer._v_scale.expand(descale_shape), + ) + return output + else: + sliding_window_size = ( + list(self.sliding_window) if self.sliding_window is not None else None + ) + flash_attn_varlen_func( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + scheduler_metadata=scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=layer._q_scale.expand(descale_shape), + k_descale=layer._k_scale.expand(descale_shape), + v_descale=layer._v_scale.expand(descale_shape), + num_splits=attn_metadata.max_num_splits, + s_aux=self.sinks, + ) + return output diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 1a4145844705..f6c3e262d60f 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -14,7 +14,6 @@ BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper, - MultiLevelCascadeAttentionWrapper, get_seq_lens, nvfp4_quantize_append_paged_kv_cache_with_slot_mapping, ) @@ -732,16 +731,6 @@ class FlashInferMetadata: Will be `None` if `num_decode_tokens == 0`. """ - # --- Special Case: Cascade Attention --- - - use_cascade: bool - """ - If True, the entire batch is a cascade attention call, and the - `prefill` and `decode` fields will both be None. - """ - - cascade_wrapper: MultiLevelCascadeAttentionWrapper | None - class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): kv_cache_spec: AttentionSpec @@ -951,8 +940,6 @@ def __init__( supports_dcp_with_varlen=False, ) - self._cascade_wrapper = None # Wrapper for cascade attention - # Global hyperparameters shared by all attention layers # TODO: discard this for trtllm-gen backend per_layer_parameters = get_per_layer_parameters( @@ -1360,15 +1347,6 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): return decode_wrapper - def _get_cascade_wrapper(self): - if self._cascade_wrapper is None: - self._cascade_wrapper = MultiLevelCascadeAttentionWrapper( - 2, - self._get_workspace_buffer(), - get_flashinfer_layout_string(self.kv_cache_layout), - ) - return self._cascade_wrapper - def _compute_flashinfer_kv_metadata( self, num_blocks_np: np.ndarray, @@ -1445,10 +1423,8 @@ def build( qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu # Step 1: Decide which dispatch modes to use: - # - Cascade attention (distinct mode) # - Prefill (FI native or TRTLLM) # - Decode (FI native, XQA, or trtllm-gen) - use_cascade = common_prefix_len > 0 uses_spec_reorder = self.reorder_batch_threshold > 1 # Page sizes >= 128 must use trtllm-gen; force it for prefill too. prefill_force_trtllm = ( @@ -1486,12 +1462,11 @@ def build( ) if not all_uses_trtllm: - sinks_would_be_dropped = self.use_dcp or use_cascade - if self.has_sinks and sinks_would_be_dropped: + if self.has_sinks and self.use_dcp: raise NotImplementedError( - "FlashInfer BatchDCPPrefillWrapper and " - "MultiLevelCascadeAttentionWrapper do not support attention sinks. " - "Please use TRTLLM on Blackwell or FlashAttention on earlier GPUs." + "FlashInfer BatchDCPPrefillWrapper does not support attention " + "sinks. Please use TRTLLM on Blackwell or FlashAttention on " + "earlier GPUs." ) if not self.global_hyperparameters.has_same_window_lefts: @@ -1508,7 +1483,7 @@ def build( ) # Step 2: Initialize the output metadata - # Leave prefill/decode/cascade_wrapper empty, to be populated + # Leave prefill/decode empty, to be populated # case by case depending on the batch contents and backend selection. attn_metadata = FlashInferMetadata( kv_cache_layout=self.kv_cache_layout, @@ -1521,10 +1496,8 @@ def build( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, causal=causal, - use_cascade=use_cascade, prefill=None, decode=None, - cascade_wrapper=None, ) # Guard access to seq_lens_cpu, which may not always be needed @@ -1532,7 +1505,7 @@ def build( # When all attention (both prefill and decode) uses TRTLLM, # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors # (block_tables, seq_lens) directly. - needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm + needs_seq_lens_cpu = self.use_dcp or not all_uses_trtllm if needs_seq_lens_cpu: if ( self._num_speculative_tokens == 0 @@ -1578,13 +1551,6 @@ def build( seq_lens_np = None num_blocks_np = None - # Adjust num_block_np for cascade attention - if use_cascade: - assert num_blocks_np is not None - assert common_prefix_len % page_size == 0 - num_common_kv_blocks = common_prefix_len // page_size - num_blocks_np -= num_common_kv_blocks - # Compute paged_kv_indices if necessary # paged_kv_indices is only needed for FlashInfer native paths; # XQA/trtllm-gen paths use block_tables directly on GPU. @@ -1592,9 +1558,7 @@ def build( needs_native_paged_decode = ( num_decodes > 0 and not decode_with_flashinfer_trtllm_api ) - needs_paged_kv_indices = ( - use_cascade or needs_native_paged_prefill or needs_native_paged_decode - ) + needs_paged_kv_indices = needs_native_paged_prefill or needs_native_paged_decode if needs_paged_kv_indices: assert num_blocks_np is not None assert seq_lens_np is not None @@ -1608,57 +1572,6 @@ def build( else: paged_kv_indices = None - # Early-out for cascade attention - if use_cascade: - assert num_blocks_np is not None - # Grab the blocks of the shared prefix from the first request. - num_common_kv_blocks = common_prefix_len // page_size - - # Create CPU versions directly for cascade (no GPU versions needed) - shared_qo_indptr_cpu = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device="cpu" - ) - shared_kv_page_indptr_cpu = torch.tensor( - [0, num_common_kv_blocks], dtype=torch.int32, device="cpu" - ) - shared_kv_page_indices_cpu = block_table_tensor[0, :num_common_kv_blocks] - shared_kv_last_page_len_cpu = torch.tensor( - [page_size], dtype=torch.int32, device="cpu" - ) - - # Remove the blocks of the shared prefix from all requests. - block_table_tensor = block_table_tensor[:, num_common_kv_blocks:] - num_blocks_np -= num_common_kv_blocks - - assert paged_kv_indices is not None - paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: 1 + num_reqs] - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs] - - attn_metadata.cascade_wrapper = self._get_cascade_wrapper() - # Cascade attention must use the same q dtype for prefill and decode - # because it does not support FP8 kv-cache or FP8 query yet. - assert self.q_data_type_prefill == self.q_data_type_decode - attn_metadata.cascade_wrapper.plan( - qo_indptr_arr=[shared_qo_indptr_cpu, qo_indptr_cpu], - paged_kv_indptr_arr=[shared_kv_page_indptr_cpu, paged_kv_indptr_cpu], - paged_kv_indices_arr=[shared_kv_page_indices_cpu, paged_kv_indices], - paged_kv_last_page_len=[ - shared_kv_last_page_len_cpu, - paged_kv_last_page_len_cpu, - ], - num_qo_heads=self.num_qo_heads, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - page_size=self.page_size, - causal=True, - sm_scale=self.sm_scale, - window_left=self.window_left, - logits_soft_cap=self.logits_soft_cap, - q_data_type=self.q_data_type_prefill, - kv_data_type=self.kv_cache_dtype, - ) - return attn_metadata - # Step 3: Handle prefill and decode pathways case by case ## PREFILL PATHWAY if num_prefills > 0: @@ -1925,15 +1838,6 @@ def build( def update_draft_decode_metadata(self, _metadata: FlashInferMetadata) -> None: pass - def use_cascade_attention(self, *args, **kwargs) -> bool: - if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype: - # TODO: The cascade wrapper currently does not support setting - # kv cache dtype to something different from query dtype. - return False - # TODO: Cascade attention doesn't work, disable it for now - # return use_cascade_attention(*args, **kwargs) - return False - class FlashInferImpl(AttentionImpl): can_return_lse_for_decode: bool = True @@ -2223,24 +2127,6 @@ def forward( output_padded = output output = output[:num_actual_tokens] - if attn_metadata.use_cascade: - # Cascade attention (rare case). - assert attn_metadata.cascade_wrapper is not None - stride_order = kv_cache_layout.layer_view_order - if self.is_kvcache_nvfp4: - kv_cache_views = tuple( - cache.permute(*stride_order) - for cache in kv_cache.split(self.num_kv_heads, dim=1) - ) - else: - kv_perm = kv_cache.permute(*stride_order) - kv_cache_views = kv_perm.split(self.head_size, dim=-1) - kv_tuple = tuple( - canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views - ) - output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple)) - return output - # When using spec decoding, num_decodes can be < num_decode_tokens # because some decode requests may have more than one query token. num_decode_tokens = attn_metadata.num_decode_tokens diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 1b9a21b65578..c2103082df7c 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -145,10 +145,6 @@ def get_impl_cls() -> type["FlexAttentionImpl"]: def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]: return FlexAttentionMetadataBuilder - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @classmethod def get_supported_head_sizes(cls) -> list[int]: return [] @@ -375,12 +371,6 @@ class FlexAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Block info total_cache_tokens: int block_size: int @@ -831,11 +821,6 @@ def build_block_mask(self) -> BlockMask: ) def __post_init__(self): - assert self.use_cascade is False, "Not implemented yet." - assert self.common_prefix_len == 0, "Not implemented yet." - assert self.cu_prefix_query_lens is None, "Not implemented yet." - assert self.prefix_kv_lens is None, "Not implemented yet." - assert self.suffix_kv_lens is None, "Not implemented yet." # Create a lookup mapping from query indices -> request number self.doc_ids = _offsets_to_doc_ids_tensor( self.query_start_loc_cpu, self.query_start_loc.device @@ -1029,16 +1014,6 @@ def build( slot_mapping = common_attn_metadata.slot_mapping num_blocks_per_seq = cdiv(seq_lens, self.block_size) - use_cascade = common_prefix_len > 0 - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - if use_cascade: - raise NotImplementedError( - "Cascade prefix attention is not yet implemented " - "for FlexAttention backend" - ) - block_size = self.kv_cache_spec.block_size max_possible_seq_len = self.model_config.max_model_len num_gpu_blocks = self.cache_config.num_gpu_blocks @@ -1128,11 +1103,6 @@ def build( seq_lens=self.persistent_seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, block_size=block_size, max_possible_sequence_length=max_possible_seq_len, num_reqs=num_reqs, @@ -1166,9 +1136,6 @@ def build( return out - def use_cascade_attention(self, *args, **kwargs) -> bool: - return False - class FlexAttentionImpl(AttentionImpl): sliding_window: int | None diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 329aee53599a..0ec0d5dd3642 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -403,9 +403,6 @@ class AiterFlashAttentionMetadata: prefill_metadata: AiterFlashAttentionPrefillMetadata | None extend_metadata: AiterFlashAttentionChunkPrefillMetadata | None - # For cascade attention. - use_cascade: bool - # Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer # since we might integrate per token quant for kv cache in the future. k_scale: dict[str, torch.Tensor] | None @@ -696,8 +693,6 @@ def build( chunk_context_metadata=chunk_context_metadata, ) - use_cascade = common_prefix_len > 0 - attn_metadata = AiterFlashAttentionMetadata( num_actual_tokens=common_attn_metadata.num_actual_tokens, query_start_loc=common_attn_metadata.query_start_loc, @@ -714,7 +709,6 @@ def build( decode_metadata=decode_metadata, prefill_metadata=prefill_metadata, extend_metadata=extend_metadata, - use_cascade=use_cascade, k_scale=self.scale, v_scale=self.scale, kv_sharing_metadata=kv_sharing_metadata, @@ -755,14 +749,10 @@ def build_for_drafting( decode_metadata=decode_metadata, prefill_metadata=None, extend_metadata=None, - use_cascade=False, k_scale=self.scale, v_scale=self.scale, ) - def use_cascade_attention(self, *args, **kwargs) -> bool: - return False - class AiterFlashAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -1219,186 +1209,136 @@ def forward( else "NHD", total_tokens=num_actual_tokens - num_decode_tokens, ) - if not attn_metadata.use_cascade: - # calculate for pure prefills - if num_prefills > 0: - assert attn_metadata.prefill_metadata is not None - assert key is not None and value is not None - - prefill_query = query[num_decode_tokens + num_extend_tokens :] - prefill_key = key[num_decode_tokens + num_extend_tokens :] - prefill_value = value[num_decode_tokens + num_extend_tokens :] - - rocm_aiter_ops.flash_attn_varlen_func( - q=prefill_query, - k=prefill_key, - v=prefill_value, - cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc, - cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc, - max_seqlen_q=attn_metadata.prefill_metadata.max_query_len, - max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len, - min_seqlen_q=1, - dropout_p=0.0, - softmax_scale=self.scale, - causal=attn_metadata.causal, - window_size=self.sliding_window, - alibi_slopes=self.alibi_slopes, - out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], - sink_ptr=self.sinks, - ) + # calculate for pure prefills + if num_prefills > 0: + assert attn_metadata.prefill_metadata is not None + assert key is not None and value is not None + + prefill_query = query[num_decode_tokens + num_extend_tokens :] + prefill_key = key[num_decode_tokens + num_extend_tokens :] + prefill_value = value[num_decode_tokens + num_extend_tokens :] + + rocm_aiter_ops.flash_attn_varlen_func( + q=prefill_query, + k=prefill_key, + v=prefill_value, + cu_seqlens_q=attn_metadata.prefill_metadata.query_start_loc, + cu_seqlens_k=attn_metadata.prefill_metadata.query_start_loc, + max_seqlen_q=attn_metadata.prefill_metadata.max_query_len, + max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len, + min_seqlen_q=1, + dropout_p=0.0, + softmax_scale=self.scale, + causal=attn_metadata.causal, + window_size=self.sliding_window, + alibi_slopes=self.alibi_slopes, + out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], + sink_ptr=self.sinks, + ) - # calculate for extends - if num_extends > 0: - assert attn_metadata.extend_metadata is not None - assert key is not None and value is not None - extend_tokens_slice = slice( - num_decode_tokens, num_decode_tokens + num_extend_tokens - ) - extend_queries = query[extend_tokens_slice] - extend_keys = key[extend_tokens_slice] - extend_values = value[extend_tokens_slice] - extend_outputs = output[extend_tokens_slice] - k_scale = layer._k_scale - v_scale = layer._v_scale - if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - k_scale = attn_metadata.k_scale - v_scale = attn_metadata.v_scale - self.extend_forward( - attn_metadata=attn_metadata, - query=extend_queries, - key=extend_keys, - value=extend_values, - key_cache=key_cache, - value_cache=value_cache, - output=extend_outputs, - cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc, - max_seqlen_q=attn_metadata.extend_metadata.max_query_len, - min_seqlen_q=1, - max_seqlen_k=attn_metadata.extend_metadata.max_seq_len, - block_table=attn_metadata.block_table[ - num_decodes : num_decodes + num_extends - ], - slot_mapping=attn_metadata.slot_mapping[ - num_decodes : num_decodes + num_extends - ], - k_scale=k_scale, - v_scale=v_scale, + # calculate for extends + if num_extends > 0: + assert attn_metadata.extend_metadata is not None + assert key is not None and value is not None + extend_tokens_slice = slice( + num_decode_tokens, num_decode_tokens + num_extend_tokens + ) + extend_queries = query[extend_tokens_slice] + extend_keys = key[extend_tokens_slice] + extend_values = value[extend_tokens_slice] + extend_outputs = output[extend_tokens_slice] + k_scale = layer._k_scale + v_scale = layer._v_scale + if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): + k_scale = attn_metadata.k_scale + v_scale = attn_metadata.v_scale + self.extend_forward( + attn_metadata=attn_metadata, + query=extend_queries, + key=extend_keys, + value=extend_values, + key_cache=key_cache, + value_cache=value_cache, + output=extend_outputs, + cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc, + max_seqlen_q=attn_metadata.extend_metadata.max_query_len, + min_seqlen_q=1, + max_seqlen_k=attn_metadata.extend_metadata.max_seq_len, + block_table=attn_metadata.block_table[ + num_decodes : num_decodes + num_extends + ], + slot_mapping=attn_metadata.slot_mapping[ + num_decodes : num_decodes + num_extends + ], + k_scale=k_scale, + v_scale=v_scale, + ) + + # calculate for decodes + if num_decodes > 0: + assert attn_metadata.decode_metadata is not None + decode_max_query_len = attn_metadata.decode_metadata.max_query_len + + # Use unified_attention for the decodes the paged kernels can't + # take: sliding window, sinks, or a multi-token batch + # (pa_fwd_asm and paged_attention_v1 don't support sinks). + if ( + self.sliding_window[0] != -1 + or decode_max_query_len > 1 + or self.sinks is not None + ): + k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes) + assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( + "Shuffle KV cache layout is not supported with sliding " + "window, sinks, or speculative decoding (multi-token decode)." ) + if not attn_metadata.causal: + from aiter.ops.triton.attention.mha_v3 import ( + flash_attn_with_kvcache, + ) - # calculate for decodes - if num_decodes > 0: - assert attn_metadata.decode_metadata is not None - decode_max_query_len = attn_metadata.decode_metadata.max_query_len - - # Use unified_attention for the decodes the paged kernels can't - # take: sliding window, sinks, or a multi-token batch - # (pa_fwd_asm and paged_attention_v1 don't support sinks). - if ( - self.sliding_window[0] != -1 - or decode_max_query_len > 1 - or self.sinks is not None - ): - k_descale, v_descale = self._get_kv_cache_descales( - layer, num_decodes + decode_query = query[:num_decode_tokens].reshape( + num_decodes, + decode_max_query_len, + query.shape[1], + query.shape[2], ) - assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "Shuffle KV cache layout is not supported with sliding " - "window, sinks, or speculative decoding (multi-token decode)." + decode_out = flash_attn_with_kvcache( + q=decode_query, + k_cache=key_cache, + v_cache=value_cache, + cache_seqlens=attn_metadata.seq_lens[:num_decodes], + softmax_scale=self.scale, + causal=attn_metadata.causal, + window_size=self.sliding_window, + softcap=self.logits_soft_cap, + q_descale=None, + k_descale=k_descale, + v_descale=v_descale, + page_table=attn_metadata.block_table[:num_decodes], ) - if not attn_metadata.causal: - from aiter.ops.triton.attention.mha_v3 import ( - flash_attn_with_kvcache, - ) - - decode_query = query[:num_decode_tokens].reshape( - num_decodes, - decode_max_query_len, + output[:num_decode_tokens].copy_( + decode_out.reshape( + num_decode_tokens, query.shape[1], query.shape[2], ) - decode_out = flash_attn_with_kvcache( - q=decode_query, - k_cache=key_cache, - v_cache=value_cache, - cache_seqlens=attn_metadata.seq_lens[:num_decodes], - softmax_scale=self.scale, - causal=attn_metadata.causal, - window_size=self.sliding_window, - softcap=self.logits_soft_cap, - q_descale=None, - k_descale=k_descale, - v_descale=v_descale, - page_table=attn_metadata.block_table[:num_decodes], - ) - output[:num_decode_tokens].copy_( - decode_out.reshape( - num_decode_tokens, - query.shape[1], - query.shape[2], - ) - ) - else: - # Non-uniform query lengths can appear in real serving - # traffic (e.g. mixed datasets). Fall back to varlen - # unified_attention instead of asserting. - from aiter.ops.triton.unified_attention import ( - unified_attention, - ) - - unified_attention( - q=query[:num_decode_tokens], - k=key_cache, - v=value_cache, - out=output[:num_decode_tokens], - cu_seqlens_q=attn_metadata.query_start_loc[ - : num_decodes + 1 - ], - max_seqlen_q=decode_max_query_len, - seqused_k=attn_metadata.seq_lens[:num_decodes], - max_seqlen_k=attn_metadata.max_seq_len, - softmax_scale=self.scale, - causal=True, - alibi_slopes=self.alibi_slopes, - window_size=self.sliding_window, - block_table=attn_metadata.block_table[:num_decodes], - softcap=self.logits_soft_cap, - q_descale=None, - k_descale=k_descale, - v_descale=v_descale, - sinks=self.sinks, - ) - return - - # The ll4mi kernel in paged_attention_v1 requires - # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4). - # For smaller head sizes or sliding window attention, - # fall back to the unified_attention triton kernel which - # handles both correctly. - _MIN_HEAD_SIZE_FOR_LL4MI = 64 - use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI - - if use_unified_attention: - k_descale, v_descale = self._get_kv_cache_descales( - layer, num_decodes - ) - assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "unified_attention fallback with shuffle layout " - "is not supported yet." ) + else: + # Non-uniform query lengths can appear in real serving + # traffic (e.g. mixed datasets). Fall back to varlen + # unified_attention instead of asserting. from aiter.ops.triton.unified_attention import ( unified_attention, ) - decode_cu_seqlens_q = attn_metadata.query_start_loc[ - : num_decodes + 1 - ] unified_attention( q=query[:num_decode_tokens], k=key_cache, v=value_cache, out=output[:num_decode_tokens], - cu_seqlens_q=decode_cu_seqlens_q, - max_seqlen_q=1, + cu_seqlens_q=attn_metadata.query_start_loc[: num_decodes + 1], + max_seqlen_q=decode_max_query_len, seqused_k=attn_metadata.seq_lens[:num_decodes], max_seqlen_k=attn_metadata.max_seq_len, softmax_scale=self.scale, @@ -1410,120 +1350,157 @@ def forward( q_descale=None, k_descale=k_descale, v_descale=v_descale, + sinks=self.sinks, ) - elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - _, num_heads, head_size = query.shape - num_blocks, block_size, num_kv_heads, _ = key_cache.shape - x = 16 // key_cache.element_size() - new_key_cache = key_cache.reshape( - num_blocks, num_kv_heads, head_size // x, block_size, x - ) - new_value_cache = value_cache.reshape( - num_blocks, num_kv_heads, block_size // x, head_size, x - ) + return + + # The ll4mi kernel in paged_attention_v1 requires + # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4). + # For smaller head sizes or sliding window attention, + # fall back to the unified_attention triton kernel which + # handles both correctly. + _MIN_HEAD_SIZE_FOR_LL4MI = 64 + use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI + + if use_unified_attention: + k_descale, v_descale = self._get_kv_cache_descales(layer, num_decodes) + assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( + "unified_attention fallback with shuffle layout " + "is not supported yet." + ) + from aiter.ops.triton.unified_attention import ( + unified_attention, + ) - # This kernel derives block addresses arithmetically, so it - # only reads the right bytes when each side is a dense - # plane of pages. - assert new_key_cache.stride(0) == ( - num_kv_heads * head_size * block_size - ), ( - "paged_attention_common needs the K/V sides as dense " - "planes; the resolved KV cache layout interleaves them " - "within a block." - ) - num_seqs = attn_metadata.seq_lens.shape[0] - max_num_partitions = ( - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 - ) // _PARTITION_SIZE_ROCM - tmp_out = torch.empty( - (num_seqs, num_heads, max_num_partitions, head_size), - dtype=query.dtype, - device=query.device, - ) - exp_sums = torch.empty( - (num_seqs, num_heads, max_num_partitions), - dtype=torch.float32, - device=query.device, - ) - max_logits = torch.empty_like(exp_sums) - k_qscale = ( - layer._k_scale - if attn_metadata.k_scale is None - else attn_metadata.k_scale - ) - v_qscale = ( - layer._v_scale - if attn_metadata.v_scale is None - else attn_metadata.v_scale - ) - rocm_aiter_ops.paged_attention_common( - Q=query[:num_decode_tokens], - K=new_key_cache, - V=new_value_cache, - tmp_out=tmp_out, - max_logits=max_logits, - exp_sums=exp_sums, - max_seq_len=attn_metadata.max_seq_len, - block_tables=attn_metadata.block_table[:num_decodes], - context_lens=attn_metadata.seq_lens[:num_decodes], - block_tables_stride0=attn_metadata.block_table[ - :num_decodes - ].stride(0), - scale=self.scale, - K_QScale_hip=k_qscale, - V_QScale_hip=v_qscale, - K_QScale_asm=k_qscale, - V_QScale_asm=v_qscale, - out_=output[:num_decode_tokens], - kv_cache_dtype=self.kv_cache_dtype, - ) - else: - _, num_heads, head_size = query.shape - nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8 - num_seqs = attn_metadata.seq_lens.shape[0] - max_num_partitions = ( - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 - ) // _PARTITION_SIZE_ROCM - - workspace_buffer = torch.empty( - (num_seqs * num_heads * max_num_partitions * head_size) - * nbytes_per_qo_elem - + 2 * (num_seqs * num_heads * max_num_partitions) * 4, - dtype=torch.uint8, - device=output.device, - ) + decode_cu_seqlens_q = attn_metadata.query_start_loc[: num_decodes + 1] + unified_attention( + q=query[:num_decode_tokens], + k=key_cache, + v=value_cache, + out=output[:num_decode_tokens], + cu_seqlens_q=decode_cu_seqlens_q, + max_seqlen_q=1, + seqused_k=attn_metadata.seq_lens[:num_decodes], + max_seqlen_k=attn_metadata.max_seq_len, + softmax_scale=self.scale, + causal=True, + alibi_slopes=self.alibi_slopes, + window_size=self.sliding_window, + block_table=attn_metadata.block_table[:num_decodes], + softcap=self.logits_soft_cap, + q_descale=None, + k_descale=k_descale, + v_descale=v_descale, + ) + elif rocm_aiter_ops.is_shuffle_kv_cache_enabled(): + _, num_heads, head_size = query.shape + num_blocks, block_size, num_kv_heads, _ = key_cache.shape + x = 16 // key_cache.element_size() + new_key_cache = key_cache.reshape( + num_blocks, num_kv_heads, head_size // x, block_size, x + ) + new_value_cache = value_cache.reshape( + num_blocks, num_kv_heads, block_size // x, head_size, x + ) - # import so that aiter register the op to the namespace of - # torch.ops.aiter - import aiter # noqa: F401 - - torch.ops.aiter.paged_attention_v1( - output[:num_decode_tokens], - workspace_buffer, - query[:num_decode_tokens], - key_cache, - value_cache, - self.scale, - attn_metadata.block_table[:num_decodes], - attn_metadata.query_start_loc[:num_decodes], - attn_metadata.seq_lens[:num_decodes], - attn_metadata.max_seq_len, - self.alibi_slopes, - self.kv_cache_dtype, - "NHD", - self.logits_soft_cap, - layer._k_scale, - layer._v_scale, - None, - _PARTITION_SIZE_ROCM, - 1, - self.sliding_window[0] + 1, - ) - else: - raise NotImplementedError( - "Cascade attention is not implemented for ROCM AITER" - ) + # This kernel derives block addresses arithmetically, so it + # only reads the right bytes when each side is a dense + # plane of pages. + assert new_key_cache.stride(0) == ( + num_kv_heads * head_size * block_size + ), ( + "paged_attention_common needs the K/V sides as dense " + "planes; the resolved KV cache layout interleaves them " + "within a block." + ) + num_seqs = attn_metadata.seq_lens.shape[0] + max_num_partitions = ( + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 + ) // _PARTITION_SIZE_ROCM + tmp_out = torch.empty( + (num_seqs, num_heads, max_num_partitions, head_size), + dtype=query.dtype, + device=query.device, + ) + exp_sums = torch.empty( + (num_seqs, num_heads, max_num_partitions), + dtype=torch.float32, + device=query.device, + ) + max_logits = torch.empty_like(exp_sums) + k_qscale = ( + layer._k_scale + if attn_metadata.k_scale is None + else attn_metadata.k_scale + ) + v_qscale = ( + layer._v_scale + if attn_metadata.v_scale is None + else attn_metadata.v_scale + ) + rocm_aiter_ops.paged_attention_common( + Q=query[:num_decode_tokens], + K=new_key_cache, + V=new_value_cache, + tmp_out=tmp_out, + max_logits=max_logits, + exp_sums=exp_sums, + max_seq_len=attn_metadata.max_seq_len, + block_tables=attn_metadata.block_table[:num_decodes], + context_lens=attn_metadata.seq_lens[:num_decodes], + block_tables_stride0=attn_metadata.block_table[:num_decodes].stride( + 0 + ), + scale=self.scale, + K_QScale_hip=k_qscale, + V_QScale_hip=v_qscale, + K_QScale_asm=k_qscale, + V_QScale_asm=v_qscale, + out_=output[:num_decode_tokens], + kv_cache_dtype=self.kv_cache_dtype, + ) + else: + _, num_heads, head_size = query.shape + nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8 + num_seqs = attn_metadata.seq_lens.shape[0] + max_num_partitions = ( + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1 + ) // _PARTITION_SIZE_ROCM + + workspace_buffer = torch.empty( + (num_seqs * num_heads * max_num_partitions * head_size) + * nbytes_per_qo_elem + + 2 * (num_seqs * num_heads * max_num_partitions) * 4, + dtype=torch.uint8, + device=output.device, + ) + + # import so that aiter register the op to the namespace of + # torch.ops.aiter + import aiter # noqa: F401 + + torch.ops.aiter.paged_attention_v1( + output[:num_decode_tokens], + workspace_buffer, + query[:num_decode_tokens], + key_cache, + value_cache, + self.scale, + attn_metadata.block_table[:num_decodes], + attn_metadata.query_start_loc[:num_decodes], + attn_metadata.seq_lens[:num_decodes], + attn_metadata.max_seq_len, + self.alibi_slopes, + self.kv_cache_dtype, + "NHD", + self.logits_soft_cap, + layer._k_scale, + layer._v_scale, + None, + _PARTITION_SIZE_ROCM, + 1, + self.sliding_window[0] + 1, + ) return output diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 0d4fd5519484..9f8e9e9ab12b 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -129,10 +129,6 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: return (KVCacheLayout.LBHNC,) return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC) - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["RocmAiterUnifiedAttentionMetadataBuilder"]: return RocmAiterUnifiedAttentionMetadataBuilder @@ -237,8 +233,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 3ccc1e5d5783..471bb9f05848 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -58,16 +58,8 @@ class RocmAttentionMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None # DFlash drafting sets this to False via CommonAttentionMetadata. causal: bool = True @@ -124,23 +116,6 @@ def build( block_table_tensor = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping - use_cascade = common_prefix_len > 0 - - if use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len - suffix_kv_lens = suffix_kv_lens.to(self.device) - else: - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - attn_metadata = RocmAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -149,12 +124,6 @@ def build( seq_lens=seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, causal=common_attn_metadata.causal, ) return attn_metadata @@ -267,10 +236,6 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: # supported, just without the native kernels. return (KVCacheLayout.LHBNC, KVCacheLayout.LBHNC) - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]: return RocmAttentionMetadataBuilder @@ -418,8 +383,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/rocm_segmented_attn.py b/vllm/v1/attention/backends/rocm_segmented_attn.py index c541eed45166..4725db917b70 100644 --- a/vllm/v1/attention/backends/rocm_segmented_attn.py +++ b/vllm/v1/attention/backends/rocm_segmented_attn.py @@ -271,7 +271,6 @@ def forward( workspace = self._get_workspace(query.device) if attn_metadata is None: return output.fill_(0) - assert not attn_metadata.use_cascade num_actual_tokens = attn_metadata.num_actual_tokens key_cache, value_cache = self._split_kv_cache(kv_cache) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 095163f14423..a95b0c1d0e13 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -80,16 +80,8 @@ class TritonAttentionMetadata: causal: bool | torch.Tensor - # For cascade attention. - use_cascade: bool - common_prefix_len: int - cu_prefix_query_lens: torch.Tensor | None - prefix_kv_lens: torch.Tensor | None - suffix_kv_lens: torch.Tensor | None - # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None - prefix_scheduler_metadata: torch.Tensor | None = None mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None mm_prefix_range_tensor: torch.Tensor | None = None rswa_prefix_lens: torch.Tensor | None = None @@ -230,23 +222,6 @@ def build( block_table_tensor = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping - use_cascade = common_prefix_len > 0 - - if use_cascade: - cu_prefix_query_lens = torch.tensor( - [0, num_actual_tokens], dtype=torch.int32, device=self.device - ) - prefix_kv_lens = torch.tensor( - [common_prefix_len], dtype=torch.int32, device=self.device - ) - suffix_kv_lens = common_attn_metadata.seq_lens.cpu() - common_prefix_len - suffix_kv_lens = suffix_kv_lens.to(self.device) - else: - cu_prefix_query_lens = None - prefix_kv_lens = None - suffix_kv_lens = None - prefix_scheduler_metadata = None - attn_metadata = TritonAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -256,12 +231,6 @@ def build( block_table=block_table_tensor, slot_mapping=slot_mapping, causal=common_attn_metadata.causal, - use_cascade=use_cascade, - common_prefix_len=common_prefix_len, - cu_prefix_query_lens=cu_prefix_query_lens, - prefix_kv_lens=prefix_kv_lens, - suffix_kv_lens=suffix_kv_lens, - prefix_scheduler_metadata=prefix_scheduler_metadata, seq_threshold_3D=self.seq_threshold_3D, num_par_softmax_segments=self.num_par_softmax_segments, softmax_segm_output=self.softmax_segm_output, @@ -378,10 +347,6 @@ def supports_batch_invariance(cls) -> bool: def get_impl_cls() -> type["TritonAttentionImpl"]: return TritonAttentionImpl - @staticmethod - def use_cascade_attention(*args, **kwargs) -> bool: - return False - @staticmethod def get_builder_cls() -> type["TritonAttentionMetadataBuilder"]: return TritonAttentionMetadataBuilder @@ -622,8 +587,6 @@ def forward( # Profiling run. return output.fill_(0) - assert attn_metadata.use_cascade is False - # IMPORTANT! # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 932f483af37b..08e16c3e64aa 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -181,10 +181,6 @@ def forward( if attn_metadata is None: return output.fill_(0) - assert attn_metadata.use_cascade is False, ( - "Cascade attention not supported for TritonAttentionDiffKVImpl" - ) - num_actual_tokens = attn_metadata.num_actual_tokens head_size_qk = self.head_size head_size_v = TritonAttentionDiffKVBackend.head_size_v diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index b6defc52e847..7d4f2d5a2b44 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -426,23 +426,6 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: blocks.extend(manager.pop_blocks_for_free(request_id)) return blocks - 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 - KV cache for each kv cache group. - - Args: - running_request_id: The request ID of any running request, used to - identify the common prefix blocks. - - Returns: - list[int]: The number of common prefix blocks for each kv cache group. - - """ - return [ - manager.get_num_common_prefix_blocks(running_request_id) - for manager in self.single_type_managers - ] - def remove_skipped_blocks( self, request_id: str, @@ -527,9 +510,6 @@ def __init__( ) self.num_single_type_manager = len(self.single_type_managers) - def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - return [0] * self.num_single_type_manager - def find_longest_cache_hit( self, block_hashes: list[BlockHash], diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 37ceb4e474cf..4f472170a874 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -689,39 +689,9 @@ def reset_prefix_cache(self) -> bool: return True def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: - """Calculate the number of common prefix blocks for each kv cache group. - - The function selects a running request and iterates through its blocks. - A block is considered a common prefix block if ALL requests with - allocated KV cache share it (i.e., ref_cnt equals the number of entries - in req_to_blocks). - - NOTE(woosuk): The number of requests with allocated KV cache is **greater - than or equal to** the number of requests scheduled in the current step. - This is because having allocated KV cache only indicates that: - 1. The request has not yet finished, and - 2. The request holds its blocks unfreed. - - While all scheduled requests must have allocated KV cache, the inverse - is not necessarily true. There may be requests with allocated KV cache - that are not scheduled in the current step. - - This can result in an edge case where the number of common prefix blocks - is 0, even though all scheduled requests share a common prefix. This - occurs because there may be unscheduled requests that do not share the - common prefix. Currently, this case cannot be easily detected, so the - function returns 0 in such cases. - - Args: - running_request_id: The request ID of any running request, used to - identify the common prefix blocks. - - Returns: - list[int]: The number of common prefix blocks for each kv cache - group. - - """ - return self.coordinator.get_num_common_prefix_blocks(running_request_id) + """Deprecated: cascade attention has been removed. Always returns 0 + for each KV cache group. Kept for out-of-tree schedulers.""" + return [0] * self.num_kv_cache_groups def take_events(self) -> list[KVCacheEvent]: """Take the KV cache events from the block pool. diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index a5af5a3d5755..9ad0d4e65865 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import cached_property from typing import TYPE_CHECKING @@ -253,9 +253,6 @@ class SchedulerOutput: # E.g., if a request has [0, 1], it could mean the vision encoder needs # to process that the request's 0-th and 1-th images in the current step. scheduled_encoder_inputs: dict[str, list[int]] - # Number of common prefix blocks for all requests in each KV cache group. - # This can be used for cascade attention. - num_common_prefix_blocks: list[int] # Request IDs that are finished in between the previous and the current # steps. This is used to notify the workers about the finished requests @@ -311,6 +308,10 @@ class SchedulerOutput: # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 + # Deprecated: cascade attention has been removed, so this is always empty. + # Kept for out-of-tree schedulers and workers. + num_common_prefix_blocks: list[int] = field(default_factory=list) + @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( @@ -320,7 +321,6 @@ def make_empty(cls) -> "SchedulerOutput": total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, - num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2603fa10a75e..bcb6def9383c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1414,16 +1414,6 @@ def skip_request(from_queue: RequestQueue) -> None: scheduled_running_reqs ) <= len(self.running) - # Get the longest common prefix among all requests in the running queue. - # This can be potentially used for cascade attention. - num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) - with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): - if self.running: - any_request_id = self.running[0].request_id - num_common_prefix_blocks = ( - self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) - ) - # Construct the scheduler output. if self.use_v2_model_runner: scheduled_new_reqs.extend(scheduled_resumed_reqs) @@ -1516,7 +1506,6 @@ def skip_request(from_queue: RequestQueue) -> None: scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, scheduled_encoder_input_stats=scheduled_encoder_input_stats, - num_common_prefix_blocks=num_common_prefix_blocks, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, # instead of being newly scheduled in this step. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 218d59f430c9..2e4fe55b96c2 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -588,21 +588,6 @@ def free(self, request_id: str) -> None: # Free blocks in reverse order so that the tail blocks are freed first. self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) - @abstractmethod - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Get the number of common prefix blocks for all requests with allocated - KV cache. - - Args: - running_request_id: The request ID. - - Returns: - The number of common prefix blocks for all requests with allocated - KV cache. - - """ - raise NotImplementedError - @classmethod @abstractmethod def find_longest_cache_hit( @@ -898,16 +883,6 @@ def _cache_partial_tail_block( block_size=self.block_size, ) - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - blocks = self.req_to_blocks[running_request_id] - num_common_blocks = 0 - for block in blocks: - if block.ref_cnt == len(self.req_to_blocks): - num_common_blocks += 1 - else: - break - return num_common_blocks - class RSWAManager(FullAttentionManager): """KV cache manager for Reference Sliding Window Attention (R-SWA). @@ -1194,14 +1169,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens, ) - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """NOTE(Chen): The prefix blocks are null blocks for sliding window layers. - So it's not correct to count ref_cnt like FullAttentionManager. Return - 0 here for correctness. Need to support cascade attention + sliding - window in the future. - """ - return 0 - class CircularBufferManager(FullAttentionManager): """Claims the ring's single block per request; prefix caching disabled.""" @@ -1289,9 +1256,6 @@ def remove_skipped_blocks( ) -> None: return - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - return 0 - def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: return 0 @@ -1447,10 +1411,6 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: ) * self.attention_chunk_size return num_skipped_tokens - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Cascade attention is not supported by chunked local attention.""" - return 0 - class MambaManager(SingleTypeKVCacheManager): supports_fine_grained_hash_lookup: ClassVar[bool] = True @@ -1680,10 +1640,6 @@ def remove_skipped_blocks( self.block_pool.free_blocks([blocks[last_state_block_idx]]) blocks[last_state_block_idx] = self._null_block - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - """Cascade attention is not supported by mamba""" - return 0 - def _needs_internal_checkpoint( self, request_id: str, @@ -2192,11 +2148,6 @@ def cache_blocks( # requests, so this method is not relevant. raise ValueError("Should not be called as prefix caching is disabled.") - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - # Cross-attention blocks contain request-specific encoder states - # and are not shared between different requests - return 0 - @classmethod def find_longest_cache_hit( cls, @@ -2394,9 +2345,6 @@ def cache_blocks( ) -> None: return None - def get_num_common_prefix_blocks(self, running_request_id: str) -> int: - return 0 - @classmethod def find_longest_cache_hit( cls, diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index 0c1fb0b2c4c9..d6abba301aec 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -254,8 +254,8 @@ def dispatch( all modes are allowed. invalid_modes: Set of cudagraph modes to exclude. Subtracted from valid_modes to compute allowed modes. (e.g., {FULL} for - features like cascade attention not supported by full - cudagraphs). None means no modes are excluded. + encoder-decoder batches with encoder outputs). None means no + modes are excluded. """ allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes() diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 32fa01bd27ec..d8df0ca1033e 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -32,7 +32,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Note: speculative decoding is now supported on CPU with C++ native impls self.use_cuda_graph = False - self.cascade_attn_enabled = False self._postprocess_tensors() @@ -109,9 +108,6 @@ def initialize_kv_cache( elif self.speculative_config.uses_draft_model(): logger.info("Draft model KV cache initialized for CPU backend") - def _init_device_properties(self) -> None: - pass - def _sync_device(self) -> None: pass diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 894e2e5b6bd9..b81ddf6d7335 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -106,7 +106,6 @@ def run_mixed_prefill_decode_warmup( prefill_token_ids = list(range(prefill_len)) kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups - num_kv_cache_groups = len(kv_cache_groups) block_count = _warmup_block_counter(model_runner) kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] decode_prefill_block_counts = [ @@ -166,7 +165,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: decode_req_id: decode_prompt_len, } decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len - decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas) cached_decode_req = CachedRequestData.make_empty() @@ -198,7 +196,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: prefill_req_id: prefill_len, } mixed_output.total_num_scheduled_tokens = num_tokens - mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id} @@ -272,7 +269,6 @@ def _warmup_kernels( decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups - num_kv_cache_groups = len(kv_cache_groups) # Encoder-decoder models: give each warmup request a dummy encoder input so # cross-attention warms up over a realistic, non-empty key sequence. @@ -361,7 +357,6 @@ def _alloc_blocks(num_blocks: int) -> list[int]: prefill_output.scheduled_new_reqs = new_reqs prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs - prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups # Disable KV connector for warmup run. model_runner.kv_connector.set_disabled(True) @@ -424,7 +419,6 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: decode_output.total_num_scheduled_tokens = sum( step_num_scheduled_tokens.values() ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups worker_execute_model(decode_output) worker_sample_tokens(None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 669b314d4cac..e19c1b84f457 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -115,7 +115,6 @@ from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks -from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, @@ -145,15 +144,12 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.kv_cache_interface import ( AttentionSpec, - ChunkedLocalAttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, - FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, KVCacheSpecKind, - SlidingWindowSpec, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, ) @@ -530,13 +526,11 @@ def __init__( ) # Model-related. - self.num_query_heads = model_config.get_num_attention_heads(parallel_config) self.inputs_embeds_size = model_config.get_inputs_embeds_size() - # Only relevant for models using ALiBi (e.g, MPT) - self.use_alibi = model_config.uses_alibi - - self.cascade_attn_enabled = not self.model_config.disable_cascade_attn self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm + # Deprecated: cascade attention has been removed. Kept for + # out-of-tree model runners. + self.cascade_attn_enabled = False # Multi-modal data support self.mm_registry = MULTIMODAL_REGISTRY @@ -769,9 +763,6 @@ def __init__( else: self.cudagraph_batch_sizes = [] - # Cache the device properties. - self._init_device_properties() - # Encoder timing registry for observability self.encoder_timing_registry: dict[str, EncoderTimingStats] = {} self._encoder_timing_lock = threading.Lock() @@ -1151,11 +1142,6 @@ def _zero_block_ids(self, block_ids: list[int]) -> None: 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""" - self.num_sms = num_compute_units(self.device.index) - # Note: used for model runner override. def _sync_device(self) -> None: torch.accelerator.synchronize() @@ -2273,7 +2259,6 @@ def _build_attention_metadata( use_spec_decode: bool = False, for_cudagraph_capture: bool = False, num_scheduled_tokens: dict[str, int] | None = None, - cascade_attn_prefix_lens: list[list[int]] | None = None, slot_mappings: dict[int, torch.Tensor] | None = None, ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]: """Returns: @@ -2472,12 +2457,6 @@ def _build_attn_group_metadata( kv_cache_spec = kv_cache_spec.kv_cache_specs[attn_group.layer_names[0]] cache_key = (kv_cache_spec, type(builder)) - cascade_attn_prefix_len = ( - cascade_attn_prefix_lens[kv_cache_gid][attn_gid] - if cascade_attn_prefix_lens - else 0 - ) - extra_attn_metadata_args = {} if use_spec_decode and isinstance( builder, @@ -2517,7 +2496,6 @@ def _build_attn_group_metadata( attn_metadata_i = attn_group.build_metadata( common_attn_metadata, ubid or 0, - common_prefix_len=cascade_attn_prefix_len, **extra_attn_metadata_args, ) if builder.supports_update_block_table: @@ -2595,144 +2573,6 @@ def _build_attn_group_metadata( return attn_metadata, spec_decode_common_attn_metadata - def _compute_cascade_attn_prefix_lens( - self, - num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, - num_common_prefix_blocks: list[int], - ) -> list[list[int]] | None: - """Returns: - Optional[cascade_attn_prefix_lens] - cascade_attn_prefix_lens is 2D: - ``[kv_cache_group_id][attn_group_idx]``, - None if we should not use cascade attention - - """ - use_cascade_attn = False - num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups) - cascade_attn_prefix_lens: list[list[int]] = [ - [] for _ in range(num_kv_cache_groups) - ] - - for kv_cache_gid in range(num_kv_cache_groups): - for attn_group in self.attn_groups[kv_cache_gid]: - if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec): - cascade_attn_prefix_len = 0 - else: - # 0 if cascade attention should not be used - cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len( - num_scheduled_tokens, - num_computed_tokens, - num_common_prefix_blocks[kv_cache_gid], - attn_group.kv_cache_spec, - attn_group.get_metadata_builder(), - ) - cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len) - use_cascade_attn |= cascade_attn_prefix_len > 0 - - return cascade_attn_prefix_lens if use_cascade_attn else None - - def _compute_cascade_attn_prefix_len( - self, - num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, - num_common_prefix_blocks: int, - kv_cache_spec: KVCacheSpec, - attn_metadata_builder: AttentionMetadataBuilder, - ) -> int: - """Compute the length of the common prefix for cascade attention. - - NOTE(woosuk): The common prefix length returned by this function - represents the length used specifically for cascade attention, not the - actual number of tokens shared between requests. When cascade attention - is disabled (use_cascade=False), this function returns 0 even if - requests share common tokens. Additionally, the common prefix length is - truncated to a multiple of the block size and may be further truncated - due to implementation details explained below. - - Args: - num_scheduled_tokens: Number of tokens scheduled per request. - num_computed_tokens: Number of tokens already computed per request. - num_common_prefix_blocks: Number of shared KV cache blocks. - kv_cache_spec: The KV cache spec of the group being considered. - attn_metadata_builder: Builder used to decide whether cascade - attention can be used. - - Returns: - int: Length of common prefix in tokens. - - """ - common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size - if common_prefix_len == 0: - # Common case. - return 0 - - # NOTE(woosuk): Cascade attention uses two attention kernels: one - # for the common prefix and the other for the rest. For the first - # kernel, we concatenate all the query tokens (possibly from - # different requests) and treat them as if they are from the same - # request. Then, we use bi-directional attention to process the - # common prefix in the KV cache. Importantly, this means that the - # first kernel does not do any masking. - - # Consider the following example: - # Request 1's input query: [D, E, X] - # Request 1's kv cache: [A, B, C, D, E, X] - # Request 1's num_computed_tokens: 3 (i.e., [A, B, C]) - # Request 2's input query: [E, Y] - # Request 2's kv cache: [A, B, C, D, E, Y] - # Request 2's num_computed_tokens: 4 (i.e., [A, B, C, D]) - - # If we use [A, B, C, D, E] as the common prefix, then the - # first kernel will compute the bi-directional attention between - # input query [D, E, X, E, Y] and common prefix [A, B, C, D, E]. - # However, this is wrong because D in Request 1 should not attend to - # E in the common prefix (i.e., we need masking). - # To avoid this, [A, B, C, D] should be the common prefix. - # That is, the common prefix should be capped by the minimum - # num_computed_tokens among the requests, and plus one to include - # the first token of the query. - - # In practice, we use [A, B, C] as the common prefix, instead of - # [A, B, C, D] (i.e., the common prefix is capped by the minimum - # num_computed_tokens, without plus one). - # This is because of an implementation detail: We want to always - # use two kernels for cascade attention. Let's imagine: - # Request 3's input query: [D] - # Request 3's kv cache: [A, B, C, D] - # Request 3's num_computed_tokens: 3 (i.e., [A, B, C]) - # If we use [A, B, C, D] as the common prefix for Request 1-3, - # then Request 3 will be processed only by the first kernel, - # and the second kernel will get an empty input. While this is not - # a fundamental problem, our current implementation does not support - # this case. - common_prefix_len = min(common_prefix_len, num_computed_tokens.min()) - # common_prefix_len should be a multiple of the block size. - common_prefix_len = ( - common_prefix_len // kv_cache_spec.block_size * kv_cache_spec.block_size - ) - use_sliding_window = isinstance(kv_cache_spec, SlidingWindowSpec) or ( - isinstance(kv_cache_spec, FullAttentionSpec) - and kv_cache_spec.sliding_window is not None - ) - use_local_attention = isinstance(kv_cache_spec, ChunkedLocalAttentionSpec) or ( - isinstance(kv_cache_spec, FullAttentionSpec) - and kv_cache_spec.attention_chunk_size is not None - ) - assert isinstance(kv_cache_spec, AttentionSpec) - use_cascade = attn_metadata_builder.use_cascade_attention( - common_prefix_len=common_prefix_len, - query_lens=num_scheduled_tokens, - num_query_heads=self.num_query_heads, - num_kv_heads=kv_cache_spec.num_kv_heads, - use_alibi=self.use_alibi, - use_sliding_window=use_sliding_window, - use_local_attention=use_local_attention, - num_sms=self.num_sms, - dcp_world_size=self.dcp_world_size, - ) - return common_prefix_len if use_cascade else 0 - def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"): mrope_pos_ptr = 0 for index, req_id in enumerate(self.input_batch.req_ids): @@ -3881,7 +3721,8 @@ def _determine_batch_execution_and_padding( num_reqs: int, num_scheduled_tokens_np: np.ndarray, max_num_scheduled_tokens: int, - use_cascade_attn: bool, + # Deprecated and ignored. Kept for out-of-tree model runners. + use_cascade_attn: bool = False, allow_microbatching: bool = True, force_eager: bool = False, # For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will @@ -3929,7 +3770,7 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): ) cudagraph_mode, batch_descriptor = dispatch_cudagraph( - num_tokens, disable_full=use_cascade_attn or has_encoder_output + num_tokens, disable_full=has_encoder_output ) num_tokens_padded = batch_descriptor.num_tokens @@ -4181,16 +4022,6 @@ def execute_model( self._prepare_inputs(scheduler_output, num_scheduled_tokens_np) ) - cascade_attn_prefix_lens = None - # Disable cascade attention when using microbatching (DBO) - if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: - # Pre-compute cascade attention prefix lengths - cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( - num_scheduled_tokens_np, - self.input_batch.num_computed_tokens_cpu[:num_reqs], - scheduler_output.num_common_prefix_blocks, - ) - ( cudagraph_mode, batch_desc, @@ -4202,7 +4033,7 @@ def execute_model( num_reqs=num_reqs, num_scheduled_tokens_np=num_scheduled_tokens_np, max_num_scheduled_tokens=max_num_scheduled_tokens, - use_cascade_attn=cascade_attn_prefix_lens is not None, + use_cascade_attn=False, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), allow_microbatching=self._allow_microbatching( num_reqs, num_scheduled_tokens_np @@ -4319,7 +4150,6 @@ def execute_model( max_num_sampled_tokens=max_num_sampled_tokens, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, - cascade_attn_prefix_lens=cascade_attn_prefix_lens, slot_mappings=slot_mappings_by_group, ) ) diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index 15fa8e1e763f..83ff08371efa 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -24,8 +24,6 @@ def __init__( ): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - # FIXME: To be verified. - self.cascade_attn_enabled = False class XPUModelRunnerV2(GPUModelRunnerV2): From c9e62f13c2b52619d8bc99f0c76cf04e77bff3c0 Mon Sep 17 00:00:00 2001 From: josh-schell-amd Date: Mon, 12 Oct 2026 13:04:37 -0500 Subject: [PATCH 434/952] [CI][ROCm] Add AITER nightly overlay scripts (#60883) Signed-off-by: josh schell Co-authored-by: Claude Sonnet 5.5 Co-authored-by: Micah Williamson --- .../scripts/rocm/aiter-nightly-overlay.sh | 112 ++++++++ .../rocm/select_aiter_nightly_wheel.py | 260 ++++++++++++++++++ 2 files changed, 372 insertions(+) create mode 100755 .buildkite/scripts/rocm/aiter-nightly-overlay.sh create mode 100755 .buildkite/scripts/rocm/select_aiter_nightly_wheel.py diff --git a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh new file mode 100755 index 000000000000..528247964baf --- /dev/null +++ b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# Install the AITER nightly wheel over this build's ci_base. +# +# Runs in CI when AITER_NIGHTLY=1, once the stock ci_base it builds on is +# published. In order: +# 1. Select the newest nightly wheel that fits the stock ci_base +# (select_aiter_nightly_wheel.py, beside this script, run inside that image) +# 2. Build a new image: pip install that wheel over the stock ci_base +# 3. Push the new image as this build's ci_base-build tag, replacing the stock +# one, so later steps and native jobs use it; record its digest +# +# Buildkite meta-data, a key-value store shared by every job in this build: +# reads rocm-ci-base-image the stock ci_base +# r/w aiter-nightly-wheel-url the wheel step 1 chose; a retry reads it +# back and installs that same wheel +# writes aiter-nightly-ci-base the new image as tag@digest; the vLLM +# image build uses it as its ci_base +# +# Exit status: +# 10 no wheel fits this ci_base +# 11 the wheel does not install +# Anything else is unexpected and passed through as is. +# Buildkite retries this step once on exit status 1 or 2, which it treats as +# transient Docker/registry failures (amd_infra_retry in the template). +# 10 and 11 sit outside that range: they are real AITER failures, and a retry +# would only fail the same way. + +set -euo pipefail + +SELECT_WHEEL="$(dirname "${BASH_SOURCE[0]}")/select_aiter_nightly_wheel.py" +BUILD_TAG="rocm/vllm-dev:ci_base-build-${BUILDKITE_BUILD_ID:?BUILDKITE_BUILD_ID is required}" + +# Scratch space for the build context and captured errors. The trap deletes it +# however the script exits: success, fail, or set -e. +scratch="$(mktemp -d)" +trap 'rm -rf "${scratch}"' EXIT +errors="${scratch}/errors" + +fail() { + local status="$1" title="$2" detail="$3" + echo "+++ :x: AITER nightly: ${title}" + echo "${detail}" >&2 + # shellcheck disable=SC2016 # the backticks are a literal markdown fence + printf ':x: **AITER nightly: %s**\n\n```\n%s\n```\n' "${title}" "${detail}" \ + | buildkite-agent annotate --style error --context aiter-nightly + exit "${status}" +} + +# Set by ensure-ci-base-amd; this step's depends_on guarantees it is there. +ci_base="$(buildkite-agent meta-data get rocm-ci-base-image)" +# Set below on the first attempt. Meta-data is per build, so a retry reuses +# that wheel, not a newer nightly. +pinned="$(buildkite-agent meta-data get aiter-nightly-wheel-url --default "")" + +echo "--- :docker: Pulling ${ci_base}" +docker pull "${ci_base}" + +echo "--- :mag: Selecting the AITER nightly for ${ci_base}" +set +e +out="$(docker run --rm -i --entrypoint python3 "${ci_base}" \ + -P - ${pinned:+--wheel-url "${pinned}"} \ + < "${SELECT_WHEEL}" 2> "${errors}")" +status=$? +set -e +cat "${errors}" >&2 +case "${status}" in + 0) ;; + 10) fail 10 "no wheel for this image" "${out}" ;; + *) fail "${status}" "could not select a wheel" "$(tail -n 20 "${errors}")" ;; +esac +# The first line is the wheel's URL; the rest, the annotation naming it. +wheel_url="$(head -n 1 <<< "${out}")" +note="$(tail -n +2 <<< "${out}")" +buildkite-agent meta-data set aiter-nightly-wheel-url "${wheel_url}" + +cat > "${scratch}/Dockerfile" <<'EOF' +ARG CI_BASE_IMAGE +FROM ${CI_BASE_IMAGE} +ARG AITER_WHEEL_URL +RUN python3 -m pip install --no-cache-dir "$AITER_WHEEL_URL" +EOF + +echo "--- :docker: Installing ${wheel_url##*/}" +if ! docker buildx build \ + --builder default \ + --provenance=false \ + --progress plain \ + --build-arg "CI_BASE_IMAGE=${ci_base}" \ + --build-arg "AITER_WHEEL_URL=${wheel_url}" \ + -t "${BUILD_TAG}" \ + "${scratch}"; then + fail 11 "the wheel does not install" \ + "Installing ${wheel_url##*/} over ${ci_base} failed; see this step's log." +fi + +docker push "${BUILD_TAG}" +# Handed off as tag@digest: build-test-image.sh rejects a ci_base that is not +# digest-pinned. +digest="$(docker buildx imagetools inspect "${BUILD_TAG}" | awk '$1 == "Digest:" { print $2; exit }')" +if [[ ! "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then + echo "Could not resolve the pushed digest of ${BUILD_TAG}" >&2 + exit 1 +fi +# The vLLM image build uses this as its ci_base. +buildkite-agent meta-data set aiter-nightly-ci-base "${BUILD_TAG}@${digest}" +# The annotation starts with :warning: when the wheel is stale. +style=info +if [[ "${note}" == :warning:* ]]; then + style=warning +fi +echo "${note}" | buildkite-agent annotate --style "${style}" --context aiter-nightly +echo "--- :white_check_mark: ${BUILD_TAG}@${digest}" diff --git a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py new file mode 100755 index 000000000000..af3de54f6020 --- /dev/null +++ b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Select the AITER nightly wheel for a vLLM ROCm image. + +aiter-nightly-overlay.sh runs this inside the stock ci_base, so it reads that +image's ROCm, torch and Python. It prints the wheel's URL, then the Buildkite +annotation naming it. +""" + +import argparse +import datetime +import importlib.metadata +import os +import sys +import urllib.request +from dataclasses import dataclass +from html.parser import HTMLParser +from pathlib import Path +from urllib.parse import unquote, urldefrag, urljoin + +import regex as re + +INDEX_URL = "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/amd-aiter/" +AITER_COMMIT_URL = "https://github.com/ROCm/aiter/commit/" +STALE_AFTER = datetime.timedelta(days=1) + +# Example: +# amd_aiter-0.1.25+rocm7.2.3.torch2.12.ae984f5.d20261008-cp312-cp312-linux_x86_64.whl +# is AITER main at commit ae984f5, built 2026-10-08, for ROCm 7.2.3, torch 2.12 +# and Python 3.12. 0.1.25 is a guess at the next release, not a real one. +WHEEL_NAME = re.compile( + r"^amd_aiter-(?P[^+-]+)\+(?P[^-]+)" + r"-(?Pcp\d+)-(?P=python)-linux_x86_64\.whl$" +) +ROCM_FIELD = re.compile(r"(?:^|\.)rocm(\d+(?:\.\d+)*)(?:a\d+)?(?=\.|$)") +TORCH_FIELD = re.compile(r"(?:^|\.)torch(\d+\.\d+)(?=\.|$)") +COMMIT_AND_DATE_FIELD = re.compile(r"(?:^|\.)([0-9a-f]{7,40})\.d(\d{8})(?=\.|$)") + + +@dataclass(frozen=True) +class Wheel: + """An AITER nightly wheel, read from its filename (see WHEEL_NAME).""" + + filename: str + url: str + version: str + base_version: str + rocm: str + torch: str + commit: str + build_date: datetime.date + python: str + + +class _LinkParser(HTMLParser): + """Collects every link on the index page, which lists one per wheel.""" + + def __init__(self): + super().__init__() + self.hrefs: list[str] = [] + + def handle_starttag(self, tag, attrs): + if tag == "a": + self.hrefs.extend( + value for name, value in attrs if name == "href" and value + ) + + +def wheel_from_url(url: str) -> Wheel | None: + """The AITER nightly wheel an absolute URL points at, if it is one.""" + filename = unquote(urldefrag(url).url.rsplit("/", 1)[-1]) + match = WHEEL_NAME.match(filename) + if not match: + return None + local = match["local"] + rocm, torch, build = ( + field.search(local) + for field in (ROCM_FIELD, TORCH_FIELD, COMMIT_AND_DATE_FIELD) + ) + if not (rocm and torch and build): + return None + return Wheel( + filename=filename, + url=url, + version=f"{match['base_version']}+{local}", + base_version=match["base_version"], + rocm=rocm[1], + torch=torch[1], + commit=build[1], + build_date=datetime.datetime.strptime(build[2], "%Y%m%d").date(), + python=match["python"], + ) + + +def parse_index(html: str, index_url: str) -> list[Wheel]: + """Every AITER nightly wheel linked from the index page, a plain HTML list + of links, one per file (the format pip reads).""" + parser = _LinkParser() + parser.feed(html) + wheels = (wheel_from_url(urljoin(index_url, href)) for href in parser.hrefs) + return [wheel for wheel in wheels if wheel] + + +def rocm_version() -> str: + """This image's ROCm version; ROCm 10+ ships as the `rocm` pip SDK.""" + try: + return importlib.metadata.version("rocm") + except importlib.metadata.PackageNotFoundError: + # No pip SDK: fall back to older ROCm releases, which install to ROCM_PATH. + rocm_path = Path(os.environ.get("ROCM_PATH", "/opt/rocm")) + return (rocm_path / ".info" / "version").read_text() + + +def rocm_release(version: str) -> str: + """`7.2.3-49` or `10.0.0a20261001` -> `7.2.3` or `10.0.0`, as wheel names + spell it.""" + match = re.match(r"\d+(?:\.\d+)*", version.strip()) + if not match: + raise ValueError(f"Unrecognised ROCm version: {version!r}") + return match.group(0) + + +def torch_release(version: str) -> str: + """`2.12.0a0+git6bbd260` -> `2.12`, as wheel names spell it.""" + match = re.match(r"\d+\.\d+", version) + if not match: + raise ValueError(f"Unrecognised torch version: {version!r}") + return match.group(0) + + +def python_tag(version_info: tuple[int, int]) -> str: + """(3, 12) -> `cp312`, as wheel names spell it.""" + return f"cp{version_info[0]}{version_info[1]}" + + +def _base_version_key(wheel: Wheel) -> tuple[int, ...]: + """`0.1.25` -> (0, 1, 25), so versions sort as numbers: 0.1.10 > 0.1.9.""" + return tuple(int(part) for part in re.findall(r"\d+", wheel.base_version)) + + +def _label(rocm: str, torch: str, python: str) -> str: + """`rocm7.2.3/torch2.12/cp312`, for messages.""" + return f"rocm{rocm}/torch{torch}/{python}" + + +def select_wheel(wheels: list[Wheel], rocm: str, torch: str, python: str) -> Wheel: + """Newest build for this ROCm, torch and Python, by date: version labels lag. + + Two builds with the same date and version: whichever the index lists first. + """ + target = (rocm, torch, python) + matching = [w for w in wheels if (w.rocm, w.torch, w.python) == target] + if not matching: + available = sorted({_label(w.rocm, w.torch, w.python) for w in wheels}) + raise LookupError( + f"No AITER nightly for {_label(*target)}. " + f"Available: {', '.join(available) or 'none'}" + ) + return max(matching, key=lambda w: (w.build_date, _base_version_key(w))) + + +def fetch(url: str) -> str: + """GET url as text. A failure exits 1, so Buildkite retries the step.""" + with urllib.request.urlopen(url, timeout=60) as response: + return response.read().decode() + + +def annotation(wheel: Wheel, today: datetime.date) -> str: + """The Buildkite annotation naming the wheel this build tests. + + Starts with :warning: if the wheel is stale; the overlay script styles it so. + """ + age = today - wheel.build_date + stale = ( + f":warning: **AITER nightly is {age.days} days old**: " + f"no newer build has been published for " + f"{_label(wheel.rocm, wheel.torch, wheel.python)}.\n\n" + if age > STALE_AFTER + else "" + ) + return stale + ( + f":crescent_moon: **AITER nightly**: main at " + f"[`{wheel.commit}`]({AITER_COMMIT_URL}{wheel.commit}), " + f"built {wheel.build_date.isoformat()} for torch {wheel.torch} " + f"(`{wheel.version}`)" + ) + + +class NoWheel(Exception): + """No wheel fits this image's ROCm, torch and Python. Exit 10.""" + + +def _image_wheel(wheel_url: str) -> Wheel: + """The wheel at wheel_url, if it fits this image's ROCm, torch and Python.""" + wheel = wheel_from_url(wheel_url) + if not wheel: + raise NoWheel(f"Not an AITER nightly wheel: {wheel_url}") + return _for_image([wheel]) + + +def _image() -> tuple[str, str, str]: + """This image's ROCm, torch and Python, as wheel names spell them.""" + try: + rocm = rocm_release(rocm_version()) + except (OSError, ValueError) as error: + raise NoWheel(f"Cannot read this image's ROCm version: {error}") from error + try: + torch = torch_release(importlib.metadata.version("torch")) + except (importlib.metadata.PackageNotFoundError, ValueError) as error: + raise NoWheel(f"Cannot read this image's torch version: {error}") from error + return rocm, torch, python_tag(sys.version_info[:2]) + + +def _for_image(wheels: list[Wheel]) -> Wheel: + """The newest of wheels for the image this runs in.""" + try: + return select_wheel(wheels, *_image()) + except LookupError as error: + raise NoWheel(str(error)) from error + + +def select(index_url: str, wheel_url: str = "") -> Wheel: + """The wheel to install: wheel_url on a retry, else the newest. + + However old the newest is, it is tested; the annotation shows its build date. + """ + if wheel_url: + return _image_wheel(wheel_url) + return _for_image(parse_index(fetch(index_url), index_url)) + + +def _today() -> datetime.date: + """Today in UTC; a function so tests can set the date.""" + return datetime.datetime.now(datetime.UTC).date() + + +# Not 1 or 2, which Buildkite retries; see aiter-nightly-overlay.sh. +EXIT_NO_WHEEL = 10 + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--index-url", default=INDEX_URL) + parser.add_argument( + "--wheel-url", default="", help="a retry's earlier choice; skips the index" + ) + args = parser.parse_args(argv) + try: + wheel = select(args.index_url, args.wheel_url) + except NoWheel as error: + print(error) + return EXIT_NO_WHEEL + print(wheel.url) + print(annotation(wheel, _today())) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 8055293353e4bc2aee2775bc0b077dd77b946b8c Mon Sep 17 00:00:00 2001 From: fululi12 Date: Mon, 12 Oct 2026 11:11:20 -0700 Subject: [PATCH 435/952] [ROCm][MLA] Cut ~21x host dispatch from the MLA metadata build (#58381) Signed-off-by: lifulu Signed-off-by: lifulu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: lifulu --- .../test_aiter_mla_uniform_mtp_qo_len.py | 104 ++++++++++++++++++ .../layers/attention/mla_attention.py | 22 +++- .../attention/backends/mla/rocm_aiter_mla.py | 37 +++++-- 3 files changed, 145 insertions(+), 18 deletions(-) create mode 100644 tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py diff --git a/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py b/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py new file mode 100644 index 000000000000..6c5416d1f41a --- /dev/null +++ b/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Differential test for the AITER MLA uniform-MTP padding decision. + +`_uniform_padded_mtp_qo_len` decides whether a full-CUDA-graph decode batch may +synthesize dummy query rows. It used to reach that decision through eight tensor +dispatches -- sum/item, boolean-mask index, all, eq, nonzero, item, any -- on a +CPU tensor holding one entry per request. It now reads the tensor out once and +decides on Python ints. + +That rewrite also **reordered** the checks: the `num_decode_tokens <= sum(...)` +early-out moved below the divisibility and `uniform_qo_len <= 1` guards. Every +one of those paths returns 0, so the result should be unchanged -- but "should +be" is exactly the sort of claim worth checking exhaustively, so both +implementations are carried here and compared over the whole small input space. + +Pure host arithmetic; no GPU required. +""" + +import itertools + +import pytest +import torch + +from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder + + +def _reference(qo_len: torch.Tensor, max_qo_len: int, num_decode_tokens: int) -> int: + """The pre-#58381 implementation, verbatim.""" + num_reqs = qo_len.numel() + if num_reqs == 0 or num_decode_tokens <= 0: + return 0 + if num_decode_tokens <= int(qo_len.sum().item()): + return 0 + if num_decode_tokens % num_reqs != 0: + return 0 + + uniform_qo_len = num_decode_tokens // num_reqs + if uniform_qo_len <= 1: + return 0 + + positive_qo_len = qo_len[qo_len > 0] + if positive_qo_len.numel() == qo_len.numel(): + return 0 + if positive_qo_len.numel() > 0: + if max_qo_len != uniform_qo_len: + return 0 + if not torch.all(positive_qo_len == uniform_qo_len): + return 0 + + zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten() + if zero_positions.numel() > 0: + first_zero = int(zero_positions[0].item()) + if torch.any(qo_len[first_zero:] > 0): + return 0 + + return uniform_qo_len + + +_ACTUAL = AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len + + +@pytest.mark.parametrize("num_reqs", [1, 2, 3, 4]) +def test_matches_reference_exhaustively(num_reqs): + """Every qo_len vector x decode-token count x max_qo_len in a small space.""" + checked = 0 + for lens in itertools.product(range(0, 4), repeat=num_reqs): + qo = torch.tensor(lens, dtype=torch.int32) + for num_decode_tokens in range(0, 13): + for max_qo_len in range(0, 5): + got = _ACTUAL(qo, max_qo_len, num_decode_tokens) + want = _reference(qo, max_qo_len, num_decode_tokens) + assert got == want, ( + f"lens={lens} num_decode_tokens={num_decode_tokens} " + f"max_qo_len={max_qo_len}: got {got}, reference {want}" + ) + checked += 1 + assert checked > 0 + + +def test_empty_batch(): + assert _ACTUAL(torch.zeros(0, dtype=torch.int32), 4, 8) == 0 + assert _reference(torch.zeros(0, dtype=torch.int32), 4, 8) == 0 + + +def test_the_padded_mtp_case_that_should_succeed(): + """A live prefix at the uniform width, then zero-padded dummies.""" + qo = torch.tensor([4, 4, 0, 0], dtype=torch.int32) + # 4 requests x 4 tokens = 16 captured, 8 live -> dummies are synthesizable. + assert _ACTUAL(qo, 4, 16) == 4 + assert _reference(qo, 4, 16) == 4 + + +def test_positive_after_zero_is_rejected(): + """A live request after a padding hole breaks the contiguous-prefix layout.""" + qo = torch.tensor([4, 0, 4, 0], dtype=torch.int32) + assert _ACTUAL(qo, 4, 16) == 0 + assert _reference(qo, 4, 16) == 0 + + +def test_ragged_live_lengths_are_rejected(): + qo = torch.tensor([4, 2, 0, 0], dtype=torch.int32) + assert _ACTUAL(qo, 4, 16) == 0 + assert _reference(qo, 4, 16) == 0 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 861c894c383e..db9e8dbc88ca 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2540,16 +2540,26 @@ def build( "Non-causal multi-token MLA requires an explicitly supported " "attention group." ) - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - num_active_reqs = int(torch.count_nonzero(query_lens > 0)) - uniform_active_queries = num_active_reqs > 0 and bool( - torch.all(query_lens[:num_active_reqs] == query_lens[0]) + # query_start_loc_cpu is host-side and holds one entry per + # request, so this check costs more in tensor dispatch than in + # the work it does. Read it out once and validate on Python ints: + # the block runs on every build, between decode steps. + qsl = query_start_loc_cpu.tolist() + query_lens = [qsl[i + 1] - qsl[i] for i in range(len(qsl) - 1)] + num_active_reqs = 0 + for qlen in query_lens: + if qlen > 0: + num_active_reqs += 1 + uniform_active_queries = num_active_reqs > 0 and all( + qlen == query_lens[0] for qlen in query_lens[:num_active_reqs] + ) + trailing_graph_padding = all( + qlen == 0 for qlen in query_lens[num_active_reqs:] ) - trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0)) if not (uniform_active_queries and trailing_graph_padding): raise ValueError( "Non-causal MLA requires a uniform query block; got query " - f"lengths {query_lens.tolist()}." + f"lengths {query_lens}." ) # Use exact GPU sequence lengths instead of the prefill path's CPU # context-length upper bounds. diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 32154105a96d..1115da630b3c 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -563,8 +563,6 @@ def _uniform_padded_mtp_qo_len( # query_start_loc flat for dummy requests. Only synthesize dummy rows # when every padded request maps to the same qlen and the q buffer has # exactly that many rows. - if num_decode_tokens <= int(qo_len.sum().item()): - return 0 if num_decode_tokens % num_reqs != 0: return 0 @@ -572,19 +570,34 @@ def _uniform_padded_mtp_qo_len( if uniform_qo_len <= 1: return 0 - positive_qo_len = qo_len[qo_len > 0] - if positive_qo_len.numel() == qo_len.numel(): + # qo_len is derived from query_start_loc_cpu, so it is already + # host-side and tiny (one entry per request). Reading it out once and + # deciding on Python ints avoids eight separate tensor dispatches -- + # sum/item, boolean-mask index, all, eq, nonzero, item, any -- on the + # metadata-build path, which runs between decode steps while the GPUs + # are idle. + lens = qo_len.tolist() + if num_decode_tokens <= sum(lens): + return 0 + + first_zero = -1 + num_positive = 0 + for i, qlen in enumerate(lens): + if qlen > 0: + num_positive += 1 + # A positive entry after a zero breaks the contiguous live + # prefix that the padded layout requires. + if first_zero >= 0: + return 0 + elif first_zero < 0: + first_zero = i + + if num_positive == num_reqs: return 0 - if positive_qo_len.numel() > 0: + if num_positive > 0: if max_qo_len != uniform_qo_len: return 0 - if not torch.all(positive_qo_len == uniform_qo_len): - return 0 - - zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten() - if zero_positions.numel() > 0: - first_zero = int(zero_positions[0].item()) - if torch.any(qo_len[first_zero:] > 0): + if any(qlen != uniform_qo_len for qlen in lens if qlen > 0): return 0 return uniform_qo_len From 56e3b081aea454d407ff658b3b12a5692f54fa47 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Tue, 13 Oct 2026 03:33:31 +0900 Subject: [PATCH 436/952] [Bugfix][weight loader] Stop warning when a parameter only widens the weight (#60020) Signed-off-by: lesj0610 Co-authored-by: Ning Xie --- vllm/model_executor/models/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index bbb7a54b552e..ef1aac6de558 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -322,7 +322,15 @@ def _load_param( ) if param.dtype != weight_data.dtype: - logger.warning( + # An exact widening -- an fp32 accumulator reading a bf16 + # table, say -- keeps every stored value, so it is a debug + # line. Every other mismatch still warns. + exact_widening = weight_data.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ) and param.dtype in (torch.float32, torch.float64) + (logger.debug if exact_widening else logger.warning)( "Attempted to load weight %s with dtype %s into " "parameter with dtype %s", weight_qualname, From d6aca837c3c6b773a557448d526db219a5a5718b Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Mon, 12 Oct 2026 21:56:51 +0300 Subject: [PATCH 437/952] [Attention] Make the sparse MLA prefill workspace size configurable (#61023) Signed-off-by: Maroon Ayoub Co-authored-by: Claude --- vllm/config/attention.py | 9 +++++++++ vllm/config/vllm.py | 10 ++++++++++ .../attention/backends/mla/flashmla_sparse.py | 20 ++++++++----------- vllm/v1/attention/backends/mla/indexer.py | 11 +++++----- 4 files changed, 33 insertions(+), 17 deletions(-) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 80548f58f2ca..67605f3c7e00 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -117,6 +117,15 @@ class AttentionConfig: When False (default), pure prefill batches use forward_mha when implemented. Set to True to always use the MQA path.""" + sparse_mla_prefill_workspace_factor: int | None = Field(default=None, ge=2) + """Sparse MLA prefill workspace size, in multiples of `max_model_len`. + FlashMLA sparse gathers a prefill step's context KV into this workspace and + runs attention in more chunks when the step holds more. The sparse indexer's + prefill K buffer is sized to fit in the same memory. Lower values return + memory to the KV cache at long `max_model_len`; the minimum of 2 leaves room + for a full-length request's context-parallel shard. If None, 5 when + `max_model_len` is at most 256K tokens, else 2.""" + flex_attn_block_m: int | None = None """Triton kernel BLOCK_M tile size for flex attention. Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index af4c99113f0f..f27c8418ab4d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1787,6 +1787,16 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if ( + self.attention_config.sparse_mla_prefill_workspace_factor is None + and self.model_config is not None + ): + # Resolved once: KV cache auto-fit can lower max_model_len later, and + # factor * max_model_len must stay within the workspace reserved at init. + self.attention_config.sparse_mla_prefill_workspace_factor = ( + 5 if self.model_config.max_model_len <= 256 * 1024 else 2 + ) + if ( self.attention_config.hisparse_config is None and self.kv_transfer_config is not None diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 4fe0be0b1f68..3cf0c33c534d 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -334,14 +334,14 @@ class Chunk: fp8_use_mixed_batch: bool = False -def get_prefill_workspace_size(max_model_len: int): - # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size. - # May be tuned later. - # Memory usage: 5 * max_model_len * 576 * 2 bytes - # Example: DeepSeek-V3.2 with max_model_len=163840 -> +def get_prefill_workspace_size(vllm_config: VllmConfig) -> int: + # Memory usage: factor * max_model_len * 576 * 2 bytes + # Example: DeepSeek-V3.2 with max_model_len=163840 and the default factor -> # 5 * 163840 * 576 * 2 = ~900 MB # This fits nicely below the typical MoE workspace size of >2GB so this is "free" - return max_model_len * 5 + factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor + assert factor is not None + return vllm_config.model_config.max_model_len * factor class FlashMLASparseMetadataBuilder( @@ -562,9 +562,7 @@ def _build_fp8_separate_prefill_decode( prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] # Chunk prefill requests to fit within workspace size - max_prefill_buffer_size = get_prefill_workspace_size( - self.vllm_config.model_config.max_model_len - ) + max_prefill_buffer_size = get_prefill_workspace_size(self.vllm_config) # Workspace rows per entry, and the prefill rows each entry covers. workspace_rows = prefill_seq_lens_cpu row_bounds = np.arange(num_prefills + 1) @@ -804,9 +802,7 @@ def __init__( if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None - prefill_workspace_size = get_prefill_workspace_size( - vllm_config.model_config.max_model_len - ) + prefill_workspace_size = get_prefill_workspace_size(vllm_config) parallel_config = vllm_config.parallel_config self.pcp_dcp_kv_gather = ( parallel_config.prefill_context_parallel_size > 1 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 97de9ffa9ba2..2d797ae66283 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -917,15 +917,16 @@ def update_draft_decode_metadata( def get_max_prefill_buffer_size(vllm_config: VllmConfig): max_model_len = vllm_config.model_config.max_model_len - # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size. + factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor + assert factor is not None # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes. - # The flashmla_sparse backend uses a workspace size of 5 * max_model_len. + # The flashmla_sparse backend uses a workspace size of factor * max_model_len. # The memory usage of the workspace there is 576 * 2 bytes; so we size this as - # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting - # within the flashmla_sparse workspace. + # (576 * 2 // 132) * factor (40 at factor 5) to maximize this workspace + # size while still fitting within the flashmla_sparse workspace. # For DeepSeek-V3.2, the max_model_len is 163840. # 40 * 163840 * 132 = 865075200 bytes = 825 MB - return max_model_len * 40 + return max_model_len * (576 * 2 // 132) * factor def _supports_varlen_paged_mqa_logits() -> bool: From d0948f754dd1fa9a77721c5e9d3722036e74e6a7 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:13:23 -0400 Subject: [PATCH 438/952] [Bug] Honor spec decode draft model config (#61240) Signed-off-by: yewentao256 --- tests/v1/worker/test_gpu_model_runner_v2.py | 34 +++++++++++++++++++ .../spec_decode/standalone_ar/speculator.py | 12 +++++++ 2 files changed, 46 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 85c3ddfce425..8ccc3731d9c8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import Mock @@ -9,6 +10,7 @@ import torch import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.config import CacheConfig, ParallelConfig from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.v1.kv_cache_interface import ( CircularBufferSpec, @@ -21,6 +23,7 @@ from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner +from vllm.v1.worker.gpu.spec_decode.standalone_ar import speculator as ar_module from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states @@ -366,3 +369,34 @@ def test_async_copy_to_np_does_not_alias_reused_buffer(): buffer.fill_(1) assert snapshot.tolist() == [0, 0, 0, 0] + + +@pytest.mark.parametrize( + ("draft_dtype", "expected_dtype"), + [("auto", "auto"), ("bfloat16", "bfloat16"), (None, "fp8_e4m3")], +) +def test_standalone_draft_kv_dtype_override(monkeypatch, draft_dtype, expected_dtype): + @dataclass + class Config: + cache_config: CacheConfig + parallel_config: ParallelConfig + model_config: object + quant_config: object | None = None + + target = Config(CacheConfig(cache_dtype="fp8_e4m3"), ParallelConfig(), object()) + speculator = object.__new__(ar_module.StandaloneARSpeculator) + speculator.vllm_config = target + speculator.draft_model_config = object() + speculator.speculative_config = SimpleNamespace( + kv_cache_dtype=draft_dtype, draft_parallel_config=ParallelConfig() + ) + get_model = Mock(return_value=object()) + monkeypatch.setattr(ar_module, "get_model", get_model) + monkeypatch.setattr(ar_module, "set_model_tag", lambda _: contextlib.nullcontext()) + + speculator.load_draft_model(object(), set()) + + draft_config = get_model.call_args.kwargs["vllm_config"] + assert draft_config.cache_config.cache_dtype == expected_dtype + assert speculator.attn_vllm_config.cache_config.cache_dtype == expected_dtype + assert target.cache_config.cache_dtype == "fp8_e4m3" diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 04e79debb869..bd404ef38fd9 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -70,6 +70,12 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: def capture(self) -> None: pass # CUDA graph not yet supported for standalone AR speculator. + @property + def attn_vllm_config(self) -> VllmConfig: + if self.speculative_config.kv_cache_dtype is not None: + return self._draft_vllm_config + return self.vllm_config + def set_attn( self, model_state: ModelState, @@ -99,15 +105,21 @@ def load_draft_model( ) -> nn.Module: spec = self.speculative_config assert spec is not None + cache_config = self.vllm_config.cache_config + if spec.kv_cache_dtype is not None: + cache_config = replace(cache_config, cache_dtype=spec.kv_cache_dtype) + draft_vllm_config = replace( self.vllm_config, model_config=self.draft_model_config, quant_config=None, + cache_config=cache_config, parallel_config=replace( spec.draft_parallel_config, rank=self.vllm_config.parallel_config.rank, ), ) + self._draft_vllm_config = draft_vllm_config with set_model_tag("draft_model"): return get_model( vllm_config=draft_vllm_config, From 08582942d8af7fc58339f3a5fedaffbbc8ea9f02 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 12 Oct 2026 20:17:12 +0100 Subject: [PATCH 439/952] [Bugfix][Frontend] Honor request_id and X-Request-Id in chat and completions render (#61236) Signed-off-by: Martin Hickey --- .../scale_out/render/test_render.py | 68 +++++++++++++++++++ .../scale_out/render/api_router.py | 6 +- vllm/entrypoints/scale_out/render/serving.py | 17 +++-- 3 files changed, 83 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 308555be362e..37f4b8eecf47 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/tests/entrypoints/scale_out/render/test_render.py @@ -367,6 +367,74 @@ async def test_completion_render_multiple_prompts(client): assert prompt["request_id"].startswith("cmpl-") +CHAT_BODY = { + "model": MODEL_NAME, + "messages": [{"role": "user", "content": "Hello"}], +} + + +@pytest.mark.asyncio +async def test_chat_render_honors_body_request_id(client): + response = await client.post( + "/v1/chat/completions/render", json={**CHAT_BODY, "request_id": "abc"} + ) + + assert response.status_code == 200 + assert response.json()["request_id"] == "chatcmpl-abc" + + +@pytest.mark.asyncio +async def test_chat_render_header_request_id_wins(client): + response = await client.post( + "/v1/chat/completions/render", + json={**CHAT_BODY, "request_id": "abc"}, + headers={"X-Request-Id": "hdr"}, + ) + + assert response.status_code == 200 + assert response.json()["request_id"] == "chatcmpl-hdr" + + +@pytest.mark.asyncio +async def test_chat_render_generates_unique_request_ids(client): + ids = set() + for _ in range(2): + response = await client.post("/v1/chat/completions/render", json=CHAT_BODY) + assert response.status_code == 200 + request_id = response.json()["request_id"] + assert request_id.startswith("chatcmpl-") + ids.add(request_id) + + assert len(ids) == 2 + + +@pytest.mark.asyncio +async def test_completion_render_request_id_single_prompt(client): + response = await client.post( + "/v1/completions/render", + json={"model": MODEL_NAME, "prompt": "Hello world", "request_id": "abc"}, + ) + + assert response.status_code == 200 + assert [r["request_id"] for r in response.json()] == ["cmpl-abc"] + + +@pytest.mark.asyncio +async def test_completion_render_request_id_multiple_prompts(client): + response = await client.post( + "/v1/completions/render", + json={ + "model": MODEL_NAME, + "prompt": ["Hello world", "Goodbye world"], + "request_id": "abc", + }, + headers={"X-Request-Id": "hdr"}, + ) + + assert response.status_code == 200 + assert [r["request_id"] for r in response.json()] == ["cmpl-hdr-0", "cmpl-hdr-1"] + + @pytest.mark.asyncio async def test_chat_completion_render_multi_turn(client): """Test chat completion render with multi-turn conversation.""" diff --git a/vllm/entrypoints/scale_out/render/api_router.py b/vllm/entrypoints/scale_out/render/api_router.py index 97df983d2e52..0bf6dbe5f107 100644 --- a/vllm/entrypoints/scale_out/render/api_router.py +++ b/vllm/entrypoints/scale_out/render/api_router.py @@ -43,7 +43,7 @@ async def render_chat_completion(request: ChatCompletionRequest, raw_request: Re "The model does not support Chat Completions Render API" ) - result = await handler.render_chat_request(request) + result = await handler.render_chat_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) @@ -67,7 +67,7 @@ async def render_messages(request: AnthropicMessagesRequest, raw_request: Reques if handler is None: raise NotImplementedError("The model does not support Messages Render API") - result = await handler.render_messages_request(request) + result = await handler.render_messages_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) @@ -90,7 +90,7 @@ async def render_completion(request: CompletionRequest, raw_request: Request): if handler is None: raise NotImplementedError("The model does not support Completions Render API") - result = await handler.render_completion_request(request) + result = await handler.render_completion_request(request, raw_request) if isinstance(result, ErrorResponse): return JSONResponse(content=result.model_dump(), status_code=result.error.code) diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index 5caf78bffe9e..b8d8ab7e1bc2 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/vllm/entrypoints/scale_out/render/serving.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING +from fastapi import Request + from vllm.entrypoints.anthropic.protocol import AnthropicMessagesRequest from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -31,7 +33,6 @@ extract_prompt_len, ) from vllm.renderers.online_renderer import OnlineRenderer -from vllm.utils import random_uuid logger = init_logger(__name__) @@ -77,6 +78,7 @@ def __init__( async def render_chat_request( self, request: ChatCompletionRequest, + raw_request: Request | None = None, ) -> GenerateRequest | ErrorResponse: """Validate the model and preprocess a chat completion request. @@ -139,7 +141,9 @@ async def render_chat_request( ) reasoning_ended = request.resolve_reasoning_ended(parser, token_ids) - request_id = f"chatcmpl-{random_uuid()}" + request_id = ( + f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}" + ) return GenerateRequest( request_id=request_id, @@ -161,6 +165,7 @@ async def render_chat_request( async def render_messages_request( self, request: AnthropicMessagesRequest, + raw_request: Request | None = None, ) -> GenerateRequest | ErrorResponse: """Validate the model and preprocess an Anthropic Messages request. @@ -172,11 +177,12 @@ async def render_messages_request( request, merge_inline_system=self._merge_inline_system ) chat_req.return_mm_kwargs = request.return_mm_kwargs - return await self.render_chat_request(chat_req) + return await self.render_chat_request(chat_req, raw_request) async def render_completion_request( self, request: CompletionRequest, + raw_request: Request | None = None, ) -> list[GenerateRequest] | ErrorResponse: """Validate the model and preprocess a completion request. @@ -197,8 +203,9 @@ async def render_completion_request( ) if isinstance(result, ErrorResponse): return result + base_id = f"cmpl-{self._base_request_id(raw_request, request.request_id)}" generate_requests: list[GenerateRequest] = [] - for engine_input in result: + for i, engine_input in enumerate(result): prompt_components = extract_prompt_components( self.model_config, engine_input ) @@ -220,7 +227,7 @@ async def render_completion_request( max_tokens, self.default_sampling_params ) - request_id = f"cmpl-{random_uuid()}" + request_id = base_id if len(result) == 1 else f"{base_id}-{i}" generate_requests.append( GenerateRequest( From 1bdd3fa2d509ce8c40d92238050433822ce76d22 Mon Sep 17 00:00:00 2001 From: ztysth <166926442+ztysth@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:25:35 +0800 Subject: [PATCH 440/952] [Bugfix][Core] Check multimodal chunk boundaries before encoder cache reuse (#59721) Signed-off-by: ztysth Signed-off-by: ztysth <166926442+ztysth@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/core/test_scheduler.py | 41 ++++++++++++++++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 33 +++++++++++++------------- 2 files changed, 57 insertions(+), 17 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index f53ba2ad2305..107a37dc91c0 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -741,7 +741,8 @@ def test_same_step_duplicate_encoder_input_stays_cached(): assert "image" not in cache.freeable -def test_no_mm_input_chunking(): +@pytest.mark.parametrize("warm_encoder_cache", [False, True]) +def test_no_mm_input_chunking(warm_encoder_cache: bool): # Disable multimodal input chunking. scheduler = create_scheduler( model="llava-hf/llava-1.5-7b-hf", @@ -753,6 +754,16 @@ def test_no_mm_input_chunking(): requests = create_requests( num_requests=1, num_tokens=1200, mm_positions=mm_positions ) + if warm_encoder_cache: + cached_request = create_requests( + num_requests=1, + num_tokens=1200, + mm_positions=mm_positions, + req_ids=["cache-primer"], + )[0] + scheduler.encoder_cache_manager.allocate(cached_request, 0) + scheduler.encoder_cache_manager.free(cached_request) + for request in requests: scheduler.add_request(request) @@ -780,6 +791,8 @@ def test_no_mm_input_chunking(): assert output.scheduled_cached_reqs.num_reqs == 1 assert len(output.finished_req_ids) == 0 assert output.num_scheduled_tokens[requests[0].request_id] == 800 + if warm_encoder_cache: + assert not output.scheduled_encoder_inputs # Test that we fail if we disable chunked mm input and use too small # of a max_num_batched_tokens for the mm input. @@ -791,6 +804,32 @@ def test_no_mm_input_chunking(): ) +def test_no_mm_input_chunking_for_same_step_duplicate(): + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + max_num_batched_tokens=1024, + disable_chunked_mm_input=True, + max_model_len=2048, + ) + request = create_requests( + num_requests=1, + num_tokens=1352, + mm_hashes_list=[["image", "image"]], + mm_positions=[ + [ + PlaceholderRange(offset=100, length=576), + PlaceholderRange(offset=776, length=576), + ] + ], + )[0] + scheduler.add_request(request) + + output = scheduler.schedule() + + assert output.num_scheduled_tokens[request.request_id] == 776 + assert output.scheduled_encoder_inputs == {request.request_id: [0]} + + @pytest.mark.parametrize("enable_prefix_caching", [True, False]) def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool): """Test scheduling behavior with concurrent partial requests. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index bcb6def9383c..cec96d2a4923 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1863,6 +1863,23 @@ def _try_schedule_encoder_inputs( # already calculated encoder inputs and can skip here. continue + # If no encoder input chunking is allowed, we do not want to + # partially schedule a multimodal item. If the scheduled range would + # only cover part of the mm input, roll back to before the mm item. + if ( + self.scheduler_config.disable_chunked_mm_input + and num_computed_tokens < start_pos + and (num_computed_tokens + num_new_tokens) + < (start_pos + num_encoder_tokens) + ): + # Account for EAGLE shift when rolling back to avoid + # encoder cache miss. This ensures the scheduled range + # stops before start_pos even with the shift. + num_new_tokens = max( + 0, start_pos - (num_computed_tokens + shift_computed_tokens) + ) + break + if not self.is_encoder_decoder: # We are not using the encoder cache for encoder-decoder models, # yet. @@ -1902,22 +1919,6 @@ def _try_schedule_encoder_inputs( # previous step. continue - # If no encoder input chunking is allowed, we do not want to - # partially schedule a multimodal item. If the scheduled range would - # only cover part of the mm input, roll back to before the mm item. - if ( - self.scheduler_config.disable_chunked_mm_input - and num_computed_tokens < start_pos - and (num_computed_tokens + num_new_tokens) - < (start_pos + num_encoder_tokens) - ): - # Account for EAGLE shift when rolling back to avoid - # encoder cache miss. This ensures the scheduled range - # stops before start_pos even with the shift. - num_new_tokens = max( - 0, start_pos - (num_computed_tokens + shift_computed_tokens) - ) - break if not self.encoder_cache_manager.can_allocate( request, i, encoder_compute_budget, num_embeds_to_schedule ): From 7039c7f17b5190862ab94ec476419c9127f2f400 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 12 Oct 2026 20:31:03 +0100 Subject: [PATCH 441/952] [Bugfix][Rust Frontend] Omit unset sampling params in render output (#61220) Signed-off-by: Martin Hickey --- rust/src/server/src/routes/tests.rs | 34 +++++++++++++++++++ rust/src/text/src/request.rs | 20 +++++++++++ .../token_in_token_out/test_protocol.py | 17 ++++++++++ 3 files changed, 71 insertions(+) diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index c6c05e63b1c8..d1e9a55d16cc 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1374,6 +1374,40 @@ async fn render_chat_returns_generate_request_with_header_request_id() { assert!(json.get("mm_features").is_none()); } +#[tokio::test] +async fn render_chat_omits_unset_sampling_params() { + let mut app = test_render_app(); + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions/render") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "render-model", + "messages": [{"role": "user", "content": "hello"}], + "max_completion_tokens": 8 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + + let sampling_params = json["sampling_params"].as_object().expect("sampling_params object"); + assert!( + sampling_params.values().all(|value| !value.is_null()), + "unexpected null in {sampling_params:?}" + ); + assert_eq!(sampling_params["max_tokens"], 8); +} + #[tokio::test] async fn render_chat_uses_processor_owned_parser_selections() { let mut app = test_render_app_with_parser_selections( diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index b00d8f3b2f12..c3efcde496f4 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -65,6 +65,7 @@ impl Default for Prompt { /// /// Original Python definition: /// +#[serde_with::skip_serializing_none] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(default)] pub struct SamplingParams { @@ -320,4 +321,23 @@ mod tests { assert!(err.to_string().contains("stop strings cannot be empty")); assert!(err.to_string().contains("req-empty-stop")); } + + #[test] + fn sampling_params_serialization_omits_unset_fields() { + let params = SamplingParams { + max_tokens: Some(64), + ..Default::default() + }; + + let value = serde_json::to_value(¶ms).unwrap(); + let obj = value.as_object().unwrap(); + assert!( + obj.values().all(|v| !v.is_null()), + "unexpected null in {value}" + ); + assert_eq!(obj["max_tokens"], 64); + + let back: SamplingParams = serde_json::from_value(value).unwrap(); + assert_eq!(back, params); + } } diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index 6a57e3b2b475..842dded1c76c 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -132,6 +132,23 @@ def test_reasoning_parser_kwargs_rejects_unknown_keys(): GenerateRequest.model_validate(payload) +def test_rust_render_sampling_params_without_nulls_validate(): + """Rust render omits unset sampling fields; Python generate rejects nulls.""" + sampling_params = {"max_tokens": 64, "watermarking": True} + request = GenerateRequest.model_validate( + {"token_ids": [1, 2, 3], "sampling_params": sampling_params} + ) + assert request.sampling_params.max_tokens == 64 + + with pytest.raises(ValidationError, match="temperature"): + GenerateRequest.model_validate( + { + "token_ids": [1, 2, 3], + "sampling_params": {**sampling_params, "temperature": None}, + } + ) + + def test_output_mode_defaults_to_tokens(): assert GenerateRequest.model_validate(_base_payload()).output_mode == "tokens" From fabaecbfa5d48389cc8921d782b3f7d3aa0d928f Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 03:32:25 +0800 Subject: [PATCH 442/952] [Misc] Expand @ZJY0516 CODEOWNERS to Mamba and KV cache (#61181) Signed-off-by: zjy0516 --- .github/CODEOWNERS | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 01dab4e76b52..b14edb0d5bce 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,13 +3,14 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng +/vllm/compilation/breakable_cudagraph.py @ZJY0516 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi @ivanium @wzhao18 /vllm/distributed/kv_transfer/kv_connector/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye -/vllm/model_executor/layers/mamba @tdoublep @tomeras91 +/vllm/model_executor/layers/mamba @tdoublep @tomeras91 @ZJY0516 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy /vllm/model_executor/model_loader @22quinn /vllm/model_executor/determinism/ @yewentao256 @@ -26,7 +27,7 @@ # Any change to the VllmConfig changes can have a large user-facing impact, # so spam a lot of people /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg -/vllm/config/cache.py @heheda12345 @ivanium +/vllm/config/cache.py @heheda12345 @ivanium @ZJY0516 # Config utils /vllm/config/utils.py @hmellor @@ -102,17 +103,18 @@ /vllm/v1/attention/backends/flashinfer.py @mgoin @pavanimajety @vadiklyutiy @wzhao18 /vllm/v1/attention/backends/triton_attn.py @tdoublep /vllm/v1/attention/backends/gdn_attn.py @ZJY0516 @vadiklyutiy -/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium +/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium @ZJY0516 /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni /vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera -/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium +/vllm/v1/kv_cache_interface.py @heheda12345 @ivanium @ZJY0516 /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium /vllm/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/v1/engine @njhill /vllm/v1/executor @njhill /vllm/v1/worker @njhill +/vllm/v1/worker/mamba_utils.py @njhill @ZJY0516 /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche @ivanium /vllm/v1/worker/encoder_cudagraph.py @shen-shanshan /vllm/v1/worker/encoder_cudagraph_defs.py @shen-shanshan @@ -120,6 +122,7 @@ # Model runner V2 /vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256 /vllm/v1/worker/gpu/kv_connector.py @orozery +/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @WoosukKwon @njhill @yewentao256 @ZJY0516 # CI & building /.buildkite @Harry-Chen @khluu @@ -132,7 +135,7 @@ /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao -/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas +/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas @ZJY0516 /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa /tests/models @DarkLight1337 @ywang96 @AndreasKaratzas @@ -140,9 +143,9 @@ /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas /tests/v1/structured_output @mgoin @aarnphm @arpera -/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium +/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium @ZJY0516 /tests/lora @jeejeelee -/tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 +/tests/models/language/generation/test_hybrid.py @tdoublep @tomeras91 @ZJY0516 /tests/v1/kv_connector/nixl_integration @NickLucche /tests/v1/kv_connector @ApostaC @orozery @ivanium @wzhao18 /tests/v1/kv_offload @ApostaC @orozery From abad32a718d2f9e3af7b18818eb474322eaadfc4 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Mon, 12 Oct 2026 22:53:08 +0300 Subject: [PATCH 443/952] [ROCm][Perf] Fix Gemma AR+RMS prefill fallback (#58547) Signed-off-by: Mikko Tukiainen Signed-off-by: Mikko Tukiainen --- .buildkite/test_areas/distributed.yaml | 2 + .../distributed/test_fusion_all_reduce.py | 65 ++++++++++++-- tests/distributed/test_aiter_rms_fallback.py | 89 +++++++++++++++++++ vllm/_aiter_ops.py | 16 +++- .../passes/fusion/allreduce_rms_fusion.py | 58 +++++++++++- vllm/kernels/aiter_ops.py | 3 + 6 files changed, 220 insertions(+), 13 deletions(-) create mode 100644 tests/distributed/test_aiter_rms_fallback.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 0b1a73136b04..f541c42c93a8 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -522,6 +522,7 @@ steps: - tests/distributed/test_rocm_aiter_custom_ar.py - tests/distributed/test_rocm_quick_reduce.py - tests/distributed/test_quick_all_reduce.py + - tests/distributed/test_aiter_rms_fallback.py - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/utils.py - examples/rl/rlhf_async_new_apis.py @@ -539,6 +540,7 @@ steps: - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py + - pytest -v -s tests/distributed/test_aiter_rms_fallback.py - label: ":nvidia: (B200) Distributed" key: distributed-tests-2xb200 diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1c99cc25330f..9e2b96957576 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -177,10 +177,12 @@ def __init__( token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16, + use_aiter: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps + self.use_aiter = use_aiter self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)] # Non-trivial weight (~Gemma range) so (1 + w) exercises the scale path. for n in self.norm: @@ -210,6 +212,8 @@ def ops_in_model_before(self): return [torch.ops.vllm.all_reduce.default] def ops_in_model_after(self): + if self.use_aiter: + return [rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()] return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default] @@ -383,6 +387,32 @@ def ops_in_model_after(self): ] +class TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model( + TestAiterAllReduceRMSNormGroupQuantFP8Model +): + """Gemma norms feeding the per-group FP8 quant sites. + + The fused AR+RMS+quant op has no ``gemma_norm`` mode, so these sites must + keep matching the quant-fused patterns with the pre-added ``1 + weight``. + The Gemma AR+RMS-only patterns match a prefix of the same subgraph, so if + they were registered ahead of the quant-fused ones they would consume the + ``all_reduce`` and strand ``rocm_aiter_group_fp8_quant`` as a separate + kernel -- which the ``fully_replaced=True`` check below would catch. + """ + + def __init__( + self, + hidden_size=128, + token_num=16, + eps=1e-6, + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__(hidden_size, token_num, eps, dtype) + self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)] + for norm in self.norm: + norm.weight.requires_grad_(False) + + class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module): def __init__( self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16 @@ -450,10 +480,10 @@ def ops_in_model_before(self): pytest.param( TestAllReduceGemmaRMSNormModel, False, - False, + IS_AITER_FOUND, marks=pytest.mark.skipif( - current_platform.is_rocm(), - reason="Not supported on ROCm platform", + current_platform.is_rocm() and not IS_AITER_FOUND, + reason="AITER is required on ROCm", ), ), pytest.param( @@ -646,7 +676,10 @@ def all_reduce_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - if test_model_cls is TestAllReduceRMSNormModel: + if test_model_cls in ( + TestAllReduceRMSNormModel, + TestAllReduceGemmaRMSNormModel, + ): model = test_model_cls( hidden_size, token_num, dtype=dtype, use_aiter=use_aiter ) @@ -667,9 +700,13 @@ def all_reduce_fusion_pass_on_test_model( ) backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False) backend.check_after_ops(model.ops_in_model_after()) - if test_model_cls in ( - TestAllReduceGemmaRMSNormModel, - TestAllReduceGemmaRMSNormStaticQuantFP8Model, + if ( + test_model_cls + in ( + TestAllReduceGemmaRMSNormModel, + TestAllReduceGemmaRMSNormStaticQuantFP8Model, + ) + and not use_aiter ): fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass)) @@ -684,6 +721,13 @@ def all_reduce_fusion_pass_on_test_model( @pytest.mark.parametrize("hidden_size", [128]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False]) +@pytest.mark.parametrize( + "test_model_cls", + [ + TestAiterAllReduceRMSNormGroupQuantFP8Model, + TestAiterAllReduceGemmaRMSNormGroupQuantFP8Model, + ], +) @pytest.mark.skipif( not current_platform.is_rocm(), reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only", @@ -695,6 +739,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, + test_model_cls: type[torch.nn.Module], monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -710,6 +755,10 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( residual, DSv3.2 indexer fan-out; parametrized over ``rocm_aiter_group_fp8_quant`` producer). + + Parametrized over plain and Gemma norms: the Gemma AR+RMS-only patterns + match a prefix of these subgraphs, so this also pins their registration + order behind the quant-fused ones. """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -730,7 +779,7 @@ def run_torch_spawn(fn, nprocs): args=( num_processes, master_port, - TestAiterAllReduceRMSNormGroupQuantFP8Model, + test_model_cls, batch_size, seq_len, hidden_size, diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py new file mode 100644 index 000000000000..4a0fc67873a1 --- /dev/null +++ b/tests/distributed/test_aiter_rms_fallback.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +import vllm._aiter_ops as aiter_ops +import vllm.distributed as vllm_distributed +import vllm.kernels # noqa: F401 + + +def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce( + monkeypatch: pytest.MonkeyPatch, +): + class FakeAiterCA: + world_size = 8 + fully_connected = True + + def custom_fused_ar_rms(self, *args, **kwargs): + return None + + class FakeAiterAllReduce: + aiter_ca = FakeAiterCA() + + def use_1stage_fused_ar_rms(self, inp): + return False + + class FakeQuickReduce: + disabled = False + + def should_quick_allreduce(self, inp): + raise AssertionError("Gemma norm must not use fused QR+RMSNorm") + + device_comm = SimpleNamespace(qr_comm=FakeQuickReduce()) + monkeypatch.setattr( + vllm_distributed, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_comm), + ) + monkeypatch.setattr( + aiter_ops.rocm_aiter_ops, + "get_aiter_allreduce", + lambda: FakeAiterAllReduce(), + ) + monkeypatch.setattr( + vllm_distributed, + "tensor_model_parallel_all_reduce", + lambda inp: inp + 2, + ) + + calls = [] + fake_aiter = types.SimpleNamespace() + + def fake_rmsnorm2d_fwd_with_add( + out, + inp, + residual, + residual_out, + weight, + epsilon, + *, + gemma_norm, + ): + calls.append((inp, weight, epsilon, gemma_norm)) + residual_out.copy_(inp + residual) + out.copy_(residual_out) + + fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + + inp = torch.randn(8192, 16, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + weight = torch.randn(16, dtype=torch.bfloat16) + + out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( + inp, residual, weight, 1e-6, gemma_norm=True + ) + + assert len(calls) == 1 + assert calls[0][0] is not inp + assert calls[0][1] is weight + assert calls[0][2] == 1e-6 + assert calls[0][3] is True + torch.testing.assert_close(residual_out, inp + 2 + residual) + torch.testing.assert_close(out, residual_out) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0ed84906185c..8af999906f28 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1086,7 +1086,10 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( hidden_dim = input_.shape[-1] row_size = hidden_dim * input_.element_size() fused_qr_rmsnorm_ok = ( - qr_comm is not None + # aiter.qr_all_reduce_rmsnorm does not implement Gemma's + # (1 + weight) scaling. + not gemma_norm + and qr_comm is not None and not getattr(qr_comm, "disabled", True) and hasattr(qr_comm, "should_quick_allreduce") and qr_comm.should_quick_allreduce(input_) @@ -1133,7 +1136,16 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( use_1stage=use_1stage, gemma_norm=gemma_norm, ) - assert result is not None + if result is None: + from vllm.distributed import tensor_model_parallel_all_reduce + + return torch.ops.vllm_aiter.fused_add_rms_norm( + tensor_model_parallel_all_reduce(input_), + residual, + weight, + epsilon, + gemma_norm, + ) return result[0], result[1] diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index b2961debcec2..59757deb2686 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1149,6 +1149,8 @@ def __del__(self) -> None: # TODO: make BasePattern to inherit from VllmPatternReplacement class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement): + gemma_norm = False + def __init__( self, epsilon: float, @@ -1170,7 +1172,8 @@ def _pattern( input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon) + norm_weight = weight.float() + 1.0 if self.gemma_norm else weight + rms = vllm.ir.ops.rms_norm(allreduce_output, norm_weight, self.epsilon) return rms, allreduce_output @@ -1187,13 +1190,22 @@ def _replacement( residual=residual, weight=weight.to(input.dtype), epsilon=self.epsilon, + gemma_norm=self.gemma_norm, ) return allreduce[0], allreduce[1] return _replacement +class AiterAllreduceFusedGemmaRMSNormPattern(AiterAllreduceFusedRMSNormPattern): + """Match all-reduce followed by GemmaRMSNorm without a residual.""" + + gemma_norm = True + + class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement): + gemma_norm = False + def __init__( self, epsilon: float, @@ -1216,8 +1228,9 @@ def _pattern( residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) + norm_weight = weight.float() + 1.0 if self.gemma_norm else weight rms, residual = vllm.ir.ops.fused_add_rms_norm( - allreduce_output, residual, weight, self.epsilon + allreduce_output, residual, norm_weight, self.epsilon ) return rms, residual @@ -1233,12 +1246,19 @@ def _replacement( residual=residual, weight=weight.to(input.dtype), epsilon=self.epsilon, + gemma_norm=self.gemma_norm, ) return allreduce[0], allreduce[1] return _replacement +class AiterAllreduceFusedAddGemmaRMSNormPattern(AiterAllreduceFusedAddRMSNormPattern): + """Match all-reduce followed by fused-add GemmaRMSNorm.""" + + gemma_norm = True + + class AiterAllreduceFusedAddRMSNormOutputOnlyPattern( AiterAllreduceFusedAddRMSNormPattern ): @@ -1267,6 +1287,14 @@ def _replacement( return _replacement +class AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern( + AiterAllreduceFusedAddRMSNormOutputOnlyPattern +): + """Match fused-add GemmaRMSNorm when its residual output is dead.""" + + gemma_norm = True + + class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern( BasePattern, VllmPatternReplacement ): @@ -1608,7 +1636,7 @@ def __init__(self, config: VllmConfig) -> None: # AR+RMS-only fusion runs first and consumes the all_reduce node, # leaving the trailing quant op stranded as an unfused kernel. # Register larger subgraphs first (DeepSeek indexer fan-out, then - # quant-only AR+RMS+quant, then AR+RMS-only). + # quant-only AR+RMS+quant, then Gemma AR+RMS, then AR+RMS-only). if supports_per_group_quant: self.register( AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( @@ -1632,6 +1660,30 @@ def __init__(self, config: VllmConfig) -> None: ) ) + # Must stay above the non-Gemma AR+RMS patterns below, which would + # otherwise match first. + self.register( + AiterAllreduceFusedGemmaRMSNormPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( + AiterAllreduceFusedAddGemmaRMSNormPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( + AiterAllreduceFusedAddGemmaRMSNormOutputOnlyPattern( + epsilon, + self.model_dtype, + self.device, + ) + ) + self.register( AiterAllreduceFusedRMSNormPattern( epsilon, diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 402c9cac1752..9817f5adc9bb 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -136,6 +136,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( residual: torch.Tensor, weight: torch.Tensor, variance_epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: from aiter import rmsnorm2d_fwd_with_add @@ -157,6 +158,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( residual_out, # residual output weight, variance_epsilon, + gemma_norm=gemma_norm, ) return out.reshape(x_shape), residual_out.reshape(residual_shape) @@ -166,6 +168,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( residual: torch.Tensor, weight: torch.Tensor, variance_epsilon: float, + gemma_norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: residual_out = torch.empty_like(residual) out = torch.empty_like(x) From 654b725ea9cb11085bce072a22a5e85709aafee6 Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Mon, 12 Oct 2026 16:23:53 -0400 Subject: [PATCH 444/952] [CI] Pace the CRCR nightly report under the relay rate limit (#61269) Signed-off-by: Andrey Talman --- .buildkite/scripts/crcr-report.sh | 7 ++++++- .buildkite/scripts/crcr_report.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/crcr-report.sh b/.buildkite/scripts/crcr-report.sh index 0c114f4653fa..a1c5610a5106 100644 --- a/.buildkite/scripts/crcr-report.sh +++ b/.buildkite/scripts/crcr-report.sh @@ -227,7 +227,12 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # five-minute lifetime and the loop above can run to the build-age deadline. # Build 91827 waited 6h48m and then took a 401 on all 360 callbacks from a # token that had expired hours earlier. -OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)" +# Posting is paced under the relay's rate limit, which takes several minutes +# for a full nightly, so ask for a longer lifetime than the default. Fall back +# to the default if this agent rejects --lifetime. +OIDC_TOKEN="$(buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" \ + --lifetime 1800 2>/dev/null \ + || buildkite-agent oidc request-token --skip-redaction --audience "${AUDIENCE}" 2>/dev/null)" if [[ -z "${OIDC_TOKEN}" ]]; then echo "could not mint a Buildkite OIDC token after the wait -- skipping report" exit 0 diff --git a/.buildkite/scripts/crcr_report.py b/.buildkite/scripts/crcr_report.py index 52f0649f011b..4fab92a2c56b 100644 --- a/.buildkite/scripts/crcr_report.py +++ b/.buildkite/scripts/crcr_report.py @@ -24,7 +24,9 @@ import argparse import json +import os import sys +import time import urllib.error import urllib.request @@ -43,6 +45,14 @@ } +# The relay rate-limits callbacks per repo over a sliding 60 s window +# (RATE_LIMIT_PER_MIN in pytorch/test-infra), and a nightly has far more jobs +# than that. Pace below the limit, and back off on a 429 for about a minute in +# total so the last retry lands in a new window. +_POST_INTERVAL_S = float(os.environ.get("CRCR_POST_INTERVAL_S", "0.65")) +_RATE_LIMIT_BACKOFF_S = (2, 4, 8, 16, 32) + + def post(url: str, token: str, body: dict, timeout: int = 30) -> tuple[int, str]: req = urllib.request.Request( url, @@ -113,6 +123,12 @@ def main() -> int: ok += 1 continue status, detail = post(args.callback_url, args.oidc_token, body) + for delay in _RATE_LIMIT_BACKOFF_S: + if status != 429: + break + time.sleep(delay) + status, detail = post(args.callback_url, args.oidc_token, body) + time.sleep(_POST_INTERVAL_S) if 200 <= status < 300: ok += 1 else: From deedc7be008fe93473ca621936cde4f3c6ca1610 Mon Sep 17 00:00:00 2001 From: SeongJun Lee Date: Tue, 13 Oct 2026 05:24:12 +0900 Subject: [PATCH 445/952] [Perf][Humming] Skip the redundant MoE input copy before the w13 GEMM (#59340) Signed-off-by: lesj0610 Co-authored-by: lesj0610 Co-authored-by: Misha Goin --- .../layers/fused_moe/experts/fused_humming_moe.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 66e03229a9e1..f27171ef0de5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -214,12 +214,20 @@ def process_input( if scatter_idx is not None: layout = "scatter" + hadamard_block_size = getattr(quant_config, f"{prefix}_hadamard_block_size") + has_input_processing = ( + mode.should_quantize + or (hadamard_block_size is not None and hadamard_block_size > 1) + or activation is not None + or scatter_idx is not None + or num_valid_tokens is not None + ) inputs, group_scales, token_scales = may_process_input( config, inputs=inputs, - outputs=quanted_input, + outputs=quanted_input if has_input_processing else None, token_scales=scale2 if mode.has_secondary_scale else scale, - hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"), + hadamard_block_size=hadamard_block_size, layout=layout, expert_tokens=expert_tokens if layout == "grouped_mask" else None, scatter_idx=scatter_idx, From fe341a53dda00472c281a2d1f4bd34a74e7eae6e Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 12 Oct 2026 16:49:49 -0400 Subject: [PATCH 446/952] [HiSparse] Report host pool in KV cache capacity logs (#60285) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- vllm/v1/core/kv_cache_utils.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 919a7b4b3f52..03c6c5ff62ac 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2308,6 +2308,11 @@ def get_kv_cache_capacity( return int(max_concurrency * max_model_len), max_concurrency +def _kv_cache_device_name(vllm_config: VllmConfig) -> str: + device_type = vllm_config.device_config.device_type + return "GPU" if device_type == "cuda" else device_type.upper() + + def update_kv_cache_capacity( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> None: @@ -2316,11 +2321,13 @@ def update_kv_cache_capacity( vllm_config.cache_config.kv_cache_size_tokens = num_tokens vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency max_model_len = vllm_config.model_config.max_model_len - device_type = vllm_config.device_config.device_type + pools = _kv_cache_device_name(vllm_config) + if vllm_config.attention_config.hisparse_config is not None: + pools += " + host" logger.info_once( "%s KV cache size: %s tokens, " "Maximum concurrency for %s tokens per request: %.2fx", - "GPU" if device_type == "cuda" else device_type.upper(), + pools, f"{num_tokens:,}", f"{max_model_len:,}", max_concurrency, @@ -2449,22 +2456,30 @@ def _auto_fit_max_model_len( "to serve even a single token. Try increasing `gpu_memory_utilization`." ) + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + if auto_fit_max >= original_max: # The model's full context length fits in memory logger.info_once( "Auto-fit max_model_len: full model context length %d fits in " - "available GPU memory", + "available KV cache memory (%s)", original_max, + kv_cache_memory, ) else: # Need to reduce max_model_len to fit in memory vllm_config.model_config.max_model_len = auto_fit_max logger.info_once( "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available GPU memory (%s GiB available for KV cache)", + "available KV cache memory (%s)", original_max, auto_fit_max, - format_gib(limiting_worker_mem), + kv_cache_memory, ) From 69d54eb078db7ff9dd63f94d55a139b182c1776d Mon Sep 17 00:00:00 2001 From: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:26:43 +0300 Subject: [PATCH 447/952] [Docs] Fix broken llm-d.ai links on the llm-d integration page (#61022) Signed-off-by: Almog Tavor Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- docs/deployment/integrations/llm-d.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md index 7d261eb910b6..dc335edb0399 100644 --- a/docs/deployment/integrations/llm-d.md +++ b/docs/deployment/integrations/llm-d.md @@ -8,11 +8,11 @@ It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own: -- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. -- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. -- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. -- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. -- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. +- **[Prefix-aware routing](https://llm-d.ai/docs/well-lit-paths/foundations/precise-prefix-cache-routing).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. +- **[Distributed KV-cache management](https://llm-d.ai/docs/architecture/advanced/kv-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/well-lit-paths/foundations/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. +- **[Prefill/decode disaggregation](https://llm-d.ai/docs/well-lit-paths/foundations/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. +- **[Wide expert-parallelism](https://llm-d.ai/docs/well-lit-paths/foundations/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. +- **SLO-aware [autoscaling](https://llm-d.ai/docs/well-lit-paths/foundations/workload-autoscaling) and [flow control](https://llm-d.ai/docs/well-lit-paths/foundations/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear. @@ -28,9 +28,9 @@ See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and r ## Get started -1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. -2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. -3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. +1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/well-lit-paths/foundations/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. +2. Browse the [well-lit path guides](https://llm-d.ai/docs/well-lit-paths), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. +3. Read the [Introduction](https://llm-d.ai/docs) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). From 44f0ddfd7b61e200aac323231dfc127d4cd53cfb Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Tue, 13 Oct 2026 00:28:35 +0300 Subject: [PATCH 448/952] [Bugfix][Parser] Stop re-parsing the whole GLM tool argument per chunk (#61253) Signed-off-by: Maroon Ayoub Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/test_glm47_moe.py | 124 ++++++++++++++++++++- tests/parser/engine/test_parser_engine.py | 2 +- vllm/parser/engine/parser_engine.py | 81 ++++++++++---- vllm/parser/engine/parser_engine_config.py | 9 ++ vllm/parser/glm47_moe.py | 36 +++++- 5 files changed, 226 insertions(+), 26 deletions(-) diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py index 58fe7aa28cbd..9f9d65ef2048 100644 --- a/tests/parser/engine/test_glm47_moe.py +++ b/tests/parser/engine/test_glm47_moe.py @@ -5,18 +5,27 @@ GLM can transition straight from reasoning into a tool call without ever emitting ````, so ```` is an implicit reasoning terminator. These cover that path plus the multi-turn prompts where an -earlier turn's markers must not be read as the current turn's state. +earlier turn's markers must not be read as the current turn's state. The +streaming tests check that long tool arguments stream the same JSON as a +non-streaming parse, without converting the whole argument on every delta. """ +import dataclasses + import pytest from tests.parser.engine.conftest import make_mock_tokenizer +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionToolsParam, +) from vllm.parser.glm47_moe import ( THINK_END, THINK_START, TOOL_CALL_END, TOOL_CALL_START, Glm47MoeParser, + _glm47_arg_value_extender, + glm47_moe_config, ) THINK_S, THINK_E = 1000, 1001 @@ -106,3 +115,116 @@ def test_no_markers_returns_input_ids(self, parser): def test_thinking_disabled(self, no_thinking_parser): ids = [THINK_S, TEXT, TOOL_S] assert no_thinking_parser.extract_content_ids(ids) == ids + + +WRITE_TOOL = ChatCompletionToolsParam( + type="function", + function={ + "name": "Write", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "line_count": {"type": "integer"}, + "content": {"type": "string"}, + }, + }, + }, +) +CALL_HEAD = ( + f"{THINK_END}{TOOL_CALL_START}Write\n" + "file_path/repo/app.py\n" + "line_count120\n" + "content" +) +CALL_TAIL = f"\n{TOOL_CALL_END}" +# Escapes, non-ASCII, and tag-like text that must stay inside the value. +CONTENT_LINE = 'if a < b > c: print("x\\y") # \t\u00e9\n' + + +def _pieces(text: str, size: int) -> list[str]: + return [text[i : i + size] for i in range(0, len(text), size)] + + +def _stream_arguments(parser, request, chunks: list[str]) -> str: + args = [] + for i, chunk in enumerate(chunks): + delta = parser.parse_delta(chunk, [], request, finished=i == len(chunks) - 1) + for tool_call in (delta and delta.tool_calls) or []: + if tool_call.function and tool_call.function.arguments: + args.append(tool_call.function.arguments) + return "".join(args) + + +class TestStreamedToolArguments: + @pytest.mark.parametrize( + "tools, line_count", + [([], '"120"'), ([WRITE_TOOL], "120")], + ids=["no_schema", "schema"], + ) + def test_long_call_streams_non_streaming_arguments( + self, mock_request, tools, line_count + ): + mock_request.tools = tools + text = CALL_HEAD + CONTENT_LINE * 400 + CALL_TAIL + + streamed = _stream_arguments( + Glm47MoeParser(make_mock_tokenizer(VOCAB)), mock_request, _pieces(text, 7) + ) + result = Glm47MoeParser(make_mock_tokenizer(VOCAB)).extract_tool_calls( + text, mock_request + ) + + assert streamed == result.tool_calls[0].function.arguments + assert f'"line_count": {line_count},' in streamed + + def test_conversions_do_not_grow_with_argument_length(self, mock_request): + """The open value streams without re-running the converter.""" + mock_request.tools = [WRITE_TOOL] + + def streaming_conversions(lines: int) -> int: + calls = 0 + base = glm47_moe_config() + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + if partial: + calls += 1 + return base.arg_converter(raw_args, partial) + + parser = Glm47MoeParser( + make_mock_tokenizer(VOCAB), + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + chunks = [ + *_pieces(CALL_HEAD, 4), + *_pieces(CONTENT_LINE * lines, 4), + *_pieces(CALL_TAIL, 4), + ] + _stream_arguments(parser, mock_request, chunks) + return calls + + assert streaming_conversions(lines=40) == streaming_conversions(lines=400) + + def test_value_extender_stops_at_split_value_end(self): + extend = _glm47_arg_value_extender("kx") + assert [extend(d) for d in ['a"\n', ""]] == [ + 'a\\"\\n', + "kx", + "k1" + "j2" + "k3", + ], + ids=["closed_value", "repeated_key_not_last"], + ) + def test_no_value_extender_unless_open_value_is_last(self, raw_args): + assert _glm47_arg_value_extender(raw_args) is None diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index dc91cef74028..25dfd546b018 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -1711,7 +1711,7 @@ class TestSafeArgPrefix: ], ) def test_safe_arg_prefix(self, json_str, expected): - assert ParserEngine._safe_arg_prefix(json_str) == expected + assert ParserEngine._safe_arg_prefix(json_str)[0] == expected # ── Coercion instability regression tests ──────────────────────── diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index c74be931d256..eaf9e4519e1d 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -24,7 +24,11 @@ from vllm.logger import init_logger from vllm.parser.abstract_parser import Parser, StreamState from vllm.parser.engine.events import EventType, SemanticEvent -from vllm.parser.engine.parser_engine_config import ParserEngineConfig, ParserState +from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, + ParserEngineConfig, + ParserState, +) from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine from vllm.renderers.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.tool_parsers.utils import ( @@ -60,7 +64,10 @@ class ToolCallSlot: "_args_joined", "name_sent", "string_keys", - "streamed_json", + "_streamed_parts", + "_streamed_joined", + "extend_open_value", + "open_value_streams", ) def __init__(self) -> None: @@ -70,7 +77,10 @@ def __init__(self) -> None: self._args_joined: str | None = "" self.name_sent: bool = False self.string_keys: set[str] | None = None - self.streamed_json: str = "" + self._streamed_parts: list[str] = [] + self._streamed_joined: str | None = "" + self.extend_open_value: ArgValueExtender | None = None + self.open_value_streams: bool = False @property def args(self) -> str: @@ -82,6 +92,21 @@ def append_args(self, value: str) -> None: self._args_parts.append(value) self._args_joined = None + @property + def streamed_json(self) -> str: + if self._streamed_joined is None: + self._streamed_joined = "".join(self._streamed_parts) + return self._streamed_joined + + @streamed_json.setter + def streamed_json(self, value: str) -> None: + self._streamed_parts = [value] + self._streamed_joined = value + + def append_streamed(self, value: str) -> None: + self._streamed_parts.append(value) + self._streamed_joined = None + class ParserEngine(Parser): """A :class:`Parser` backed by a single declarative engine config. @@ -133,6 +158,7 @@ def __init__( self._arg_converter = parser_engine_config.arg_converter self._arg_structural_chars = parser_engine_config.arg_structural_chars + self._arg_value_extender = parser_engine_config.arg_value_extender self._stream_arg_deltas = parser_engine_config.stream_arg_deltas self._strip_trailing_reasoning_ws = ( parser_engine_config.strip_trailing_reasoning_whitespace @@ -303,7 +329,9 @@ def _coerce_dict(args: dict, properties: dict) -> tuple[dict, bool]: return args, changed @staticmethod - def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: + def _safe_arg_prefix( + json_str: str, string_keys: set[str] | None = None + ) -> tuple[str, bool]: """Return the prefix of *json_str* up to the last top-level value. Middle values (followed by a comma) are stable across streaming @@ -313,6 +341,9 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: String values for keys in ``string_keys`` are prefix-stable, so stream their unterminated content instead of buffering long arguments until the closing tag arrives. + + Also returns whether the prefix ends inside that streamed trailing + string value. """ last_colon = -1 last_key: str | None = None @@ -345,14 +376,14 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: last_key = pending_key pending_key = None if last_colon < 0: - return "" + return "", False end = last_colon + 1 while end < len(json_str) and json_str[end] in (" ", "\t", "\n", "\r"): end += 1 if end >= len(json_str) or json_str[end] != '"': - return json_str[:end] + return json_str[:end], False if string_keys is not None and last_key not in string_keys: - return json_str[:end] + return json_str[:end], False escape = False for i in range(end + 1, len(json_str)): @@ -364,8 +395,8 @@ def _safe_arg_prefix(json_str: str, string_keys: set[str] | None = None) -> str: escape = True continue if c == '"': - return json_str[:i] - return json_str + return json_str[:i], True + return json_str, True @staticmethod def _streamable_string_keys(properties: dict) -> set[str] | None: @@ -1027,6 +1058,15 @@ def _compute_arg_delta(self, idx: int, raw_delta: str) -> str | None: return None slot = self._tool_slots[idx] + if slot.extend_open_value is not None: + diff = slot.extend_open_value(raw_delta) + if diff is not None: + if not diff or not slot.open_value_streams: + return None + slot.append_streamed(diff) + return diff + slot.extend_open_value = None + try: current_json = converter(slot.args, True) except (json.JSONDecodeError, ValueError, TypeError): @@ -1040,22 +1080,19 @@ def _compute_arg_delta(self, idx: int, raw_delta: str) -> str | None: current_json = self._fix_arg_types(current_json, slot.name) prev = slot.streamed_json - safe_json = self._safe_arg_prefix(current_json, slot.string_keys) + safe_json, value_streams = self._safe_arg_prefix(current_json, slot.string_keys) - if not safe_json or safe_json == prev: + if not safe_json: + return None + if prev and not safe_json.startswith(prev): return None - if prev: - if not safe_json.startswith(prev): - return None - diff = safe_json[len(prev) :] - else: - diff = safe_json - - if diff: - slot.streamed_json = safe_json - return diff - return None + slot.streamed_json = safe_json + if self._arg_value_extender is not None: + # Deltas that only extend the open value skip the conversion above. + slot.extend_open_value = self._arg_value_extender(slot.args) + slot.open_value_streams = value_streams + return safe_json[len(prev) :] or None def _flush_arg_converter(self, idx: int) -> str | None: converter = self._arg_converter diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 215fa0d9bdf5..2bb2ae267e6f 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -24,6 +24,10 @@ from vllm.parser.engine.events import EventType +# Maps a raw argument delta to the JSON text it appends to the open trailing +# value, or returns None when the delta may change anything else. +ArgValueExtender = Callable[[str], str | None] + class ParserState(Enum): CONTENT = auto() @@ -90,6 +94,11 @@ class ParserEngineConfig: arg_structural_chars: frozenset[str] | None = None + # Called with the raw argument text when its converted JSON ends in an + # open value; returns an extender for that value, or None. Streaming + # skips ``arg_converter`` for every delta the extender accepts. + arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None + # Special tokens exempt from auto-drop but not state-machine terminals. preserve_tokens: frozenset[str] = field(default_factory=frozenset) diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 9a18228fb0d1..44377e62b707 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -24,6 +24,7 @@ from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -63,12 +64,16 @@ ) -def _glm47_arg_converter(raw_args: str, partial: bool) -> str: +def _glm47_parse_args( + raw_args: str, partial: bool +) -> tuple[dict[str, object], str | None]: + """Return the arguments and the key of the open trailing value, if any.""" params: dict[str, object] = {} for match in _ARG_RE.finditer(raw_args): params[match.group("key").strip()] = match.group("value") + open_key = None if partial: remaining = _ARG_RE.sub("", raw_args) match = _PARTIAL_ARG_RE.search(remaining) @@ -76,8 +81,34 @@ def _glm47_arg_converter(raw_args: str, partial: bool) -> str: key = match.group("key").strip() if key: params[key] = match.group("value") + open_key = key + + return params, open_key + + +def _glm47_arg_converter(raw_args: str, partial: bool) -> str: + return json.dumps(_glm47_parse_args(raw_args, partial)[0], ensure_ascii=False) + + +def _glm47_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + # The open value runs to the end of raw_args and is the last JSON value, + # so a delta appends its JSON-escaped text unless it closes the value. + params, open_key = _glm47_parse_args(raw_args, partial=True) + if open_key is None or next(reversed(params)) != open_key: + return None + + keep = len(ARG_VALUE_END) - 1 + tail = raw_args[-keep:] + + def extend(delta: str) -> str | None: + nonlocal tail + text = tail + delta + if ARG_VALUE_END in text: + return None + tail = text[-keep:] + return json.dumps(delta, ensure_ascii=False)[1:-1] - return json.dumps(params, ensure_ascii=False) + return extend def _glm53_always_thinks(tokenizer: TokenizerLike) -> bool: @@ -194,6 +225,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: }, turn_boundary_tokens=GLM_TURN_BOUNDARIES, arg_converter=_glm47_arg_converter, + arg_value_extender=_glm47_arg_value_extender, stream_arg_deltas=True, tool_args_json=False, validate_tool_names=True, From 802feb6770a0522ce8d2999262d95f4618e04adf Mon Sep 17 00:00:00 2001 From: Akshay Viswanathan Date: Tue, 13 Oct 2026 00:29:31 +0200 Subject: [PATCH 449/952] [Bugfix][KVConnector] Fix MoRI-IO WRITE-mode requests stranded in WAITING_FOR_REMOTE_KVS (#57705) Signed-off-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Akshay Viswanathan <4953343+akshayv@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- .../unit/test_moriio_write_done_routing.py | 104 +++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 9 +- .../v1/moriio/moriio_connector.py | 73 ++++++++--- .../kv_connector/v1/moriio/moriio_engine.py | 121 +++++++++++++----- 4 files changed, 255 insertions(+), 52 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index adac007d19eb..57b6893869eb 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -3,20 +3,30 @@ """Unit tests for mori-io write_done routing (#51681).""" import threading +import time +from collections import OrderedDict from types import SimpleNamespace from typing import Any +from unittest.mock import patch import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ROLE, + MoRIIOMode, RemoteAllocInfo, ReqMeta, WriteTask, get_port_offset, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorScheduler, MoRIIOConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) from .utils import make_moriio_writer @@ -79,6 +89,7 @@ def _write_task( *, multi_pod_hosts: list[str] | None = None, remote_dp_size_local: int = 0, + remote_dp_size: int = 1, request_id: str = "req", ) -> WriteTask: return WriteTask( @@ -93,9 +104,21 @@ def _write_task( remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts or [], remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) +def _real_wrapper() -> MoRIIOWrapper: + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) + wrapper.lock = threading.Lock() + wrapper.done_req_ids = [] + wrapper.done_remote_allocate_req_dict = {} + wrapper.done_write_cache_req_ids = [] + wrapper.failed_write_cache_req_ids = [] + wrapper._terminal_transfer_ids = OrderedDict() + return wrapper + + def test_execute_write_task_ip_is_per_task_under_overwrite(): alloc = {"tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0)} worker = _fake_worker(alloc) @@ -269,3 +292,84 @@ def _capture_schedule_write_blocks(**kwargs: Any) -> None: assert captured["remote_dp_size_local"] == 4 assert captured["local_block_ids"] == [1] assert captured["remote_block_ids"] == [2] + + +def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): + wrapper = _real_wrapper() + sent: list[tuple[str, int, str]] = [] + wrapper.send_notify = lambda tid, ip, port, message_type: sent.append( + (ip, port, message_type) + ) + writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=0)) + task = _write_task( + "tA", + "10.0.0.1", + multi_pod_hosts=["10.0.0.1", "10.0.0.2"], + remote_dp_size_local=2, + remote_dp_size=4, + ) + task.enqueue_time = time.perf_counter() - 2 * writer._defer_timeout + writer._deferred_tasks.append(task) + + writer._process_deferred_tasks() + + assert writer._deferred_tasks == [] + assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["tA"] + # Global DP ranks 0-3, two per pod: the owning rank is unknown. + assert sent == [ + (host, 20000 + get_port_offset(local_rank, 0), "write_failed") + for host in ("10.0.0.1", "10.0.0.2") + for local_rank in (0, 1) + ] + + # A late decode allocation must not resurrect a write into freed blocks. + with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER): + wrapper._handle_remote_blocks_message( + {"transfer_id": "tA", "block_notify_list": [7], "decode_rank": 1} + ) + assert wrapper.done_remote_allocate_req_dict == {} + + +def test_write_scheduler_waits_for_worker_ack_past_defer_deadline(): + scheduler = SimpleNamespace( + is_producer=True, + mode=MoRIIOMode.WRITE, + _defer_timeout=60.0, + _deferred_send_deadlines={"req-1": (time.monotonic() - 1.0, "tA")}, + _pending_sent_acks={}, + unmap_request_id=lambda *args, **kwargs: None, + ) + output = SimpleNamespace(finished_sending=set()) + + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending is None + assert "req-1" in scheduler._deferred_send_deadlines + + output.finished_sending = {"req-1"} + MoRIIOConnectorScheduler.update_connector_output(scheduler, output) + assert output.finished_sending == {"req-1"} + assert scheduler._deferred_send_deadlines == {} + + +def test_write_failed_is_reported_as_failed_recving_on_owning_rank(): + wrapper = _real_wrapper() + with patch.object(moriio_engine, "get_role", return_value=ROLE.CONSUMER): + for tid in ("tA", "tOtherRank"): + wrapper._handle_structured_message( + {"type": "write_failed", "transfer_id": tid} + ) + worker = SimpleNamespace( + mode=MoRIIOMode.WRITE, + is_producer=False, + moriio_wrapper=wrapper, + transfer_id_to_request_id={"tA": "req-1"}, + _unmatched_write_failures=OrderedDict(), + get_finished=lambda: (set(), set()), + ) + + results = MoRIIOConnectorWorker.get_transfer_results(worker) + + assert results.finished_recving == {"req-1"} + assert results.failed_recving == {"req-1"} + assert wrapper.pop_finished_write_req_ids() == set() + assert list(worker._unmatched_write_failures) == ["tOtherRank"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index ae1eecf6e55d..e82d0668189d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -71,6 +71,7 @@ class WriteTask: remote_ip: str multi_pod_hosts: list[str] = field(default_factory=list) remote_dp_size_local: int = 0 + remote_dp_size: int = 1 enqueue_time: float = field(default_factory=time.perf_counter) retried: int = 0 @@ -307,8 +308,8 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": # WRITE mode. # transfer_timeout -> Timeout for waiting_for_transfer_complete before # raising TransferError (sec). - # defer_timeout -> Timeout before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed (sec). + # defer_timeout -> Timeout for deferred WRITE allocation or a missing + # READ release notification (sec). # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA # completion never arrived is aborted (sec). @@ -414,8 +415,8 @@ class MoRIIOConstants: # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError. # Overridable via kv_connector_extra_config["transfer_timeout"]. DEFAULT_TRANSFER_TIMEOUT = 30.0 - # Timeout (seconds) before a deferred send with no finished_sending - # notification is reaped and its blocks force-freed. + # Timeout (seconds) for deferred WRITE allocation or a missing READ release + # notification. # Overridable via kv_connector_extra_config["defer_timeout"]. DEFAULT_DEFER_TIMEOUT = 60.0 # Timeout (seconds) before an in-flight recv whose RDMA completion was lost diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 082643778e7d..c398cb20a8b0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -5,7 +5,7 @@ import queue import threading import time -from collections import defaultdict +from collections import OrderedDict, defaultdict from collections.abc import Collection from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast @@ -25,6 +25,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorTransferResults, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( @@ -116,6 +117,9 @@ _SQ_FULL_BACKOFF_INITIAL_S = 0.001 _SQ_FULL_BACKOFF_MAX_S = 0.05 _MAX_LOCAL_MAMBA_TAIL_BLOCKS = 1 +# A WRITE failure is broadcast to every decode DP rank, so ranks that never +# see the request would otherwise keep its transfer id forever. +_MAX_UNMATCHED_WRITE_FAILURES = 4096 try: @@ -356,6 +360,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + assert self.connector_worker is not None + return self.connector_worker.get_transfer_results() + def get_block_ids_with_load_errors(self) -> set[int]: """Blocks whose MoRIIO read failed, for the scheduler to recompute.""" if self.connector_worker is None: @@ -667,12 +677,10 @@ def __init__( set_role(ROLE.CONSUMER) # Reqs to send and their expiration time self._reqs_need_send: dict[ReqId, float] = {} - # Deadlines for requests whose block freeing was deferred. - # Survives across scheduler steps. If the worker never reports - # finished_sending before the deadline, we inject them into - # connector_output.finished_sending so the scheduler frees the blocks to avoid - # hanging indefinitely waiting for a free notification that never comes. - # Value: (deadline, transfer_id) for unmapping after mutation. + # Requests whose block freeing is deferred. READ mode uses the deadline + # as a missing-ACK watchdog; WRITE mode waits for the worker ACK because + # only the worker can mark the transfer terminal before blocks are freed. + # Value: (deadline, transfer_id) for unmapping after completion. self._deferred_send_deadlines: dict[ReqId, tuple[float, TransferId | None]] = {} self._defer_timeout = float( self.kv_transfer_config.kv_connector_extra_config.get( @@ -1131,7 +1139,14 @@ def update_state_after_alloc( _should_notify = self._global_dp_rank == remote_dp_rank else: _should_notify = self._is_kv_master - if _should_notify: + if _should_notify and num_external_tokens == 0: + # Nothing left to fetch (full local prefix hit), so release + # the producer's blocks instead of sending an empty + # allocation, which it rejects. + self._release_write_prefill_blocks( + request.request_id, request.kv_transfer_params + ) + elif _should_notify: peer_zmq = get_peer_zmq_from_request_id( request.request_id, is_producer=False ) @@ -1159,12 +1174,6 @@ def update_state_after_alloc( f"remote_notify_port={remote_notify_port!r})" ) - # num_external_tokens == 0: nothing to push, so don't tell - # the producer to write into these blocks. - block_notify_list = ( - blocks.get_block_ids()[0] if num_external_tokens > 0 else [] - ) - # Wide-EP multi-pod: a pod binds notify sockets only for # its LOCAL ranks, so the port offset must use the per-pod # local rank (% dp_local), not the global rank. Single-pod @@ -1191,7 +1200,7 @@ def update_state_after_alloc( self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], - block_notify_list=block_notify_list, + block_notify_list=blocks.get_block_ids()[0], host=_notify_host, port=target_port, ) @@ -1541,8 +1550,9 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None: * An ACK that arrives BEFORE its request finished is parked in ``_pending_sent_acks`` and released on a later step once the request enters ``_deferred_send_deadlines``. - * A deferred send whose ACK never arrives is reaped after - ``_defer_timeout`` and surfaced, so leaked blocks are freed. + * In READ mode, a deferred send whose ACK never arrives is reaped + after ``_defer_timeout``. WRITE-mode timeouts are owned by the worker, + which marks the transfer terminal before ACKing the block release. * A parked ACK that never matches a deferral before its own deadline is a stale duplicate (e.g. a real ACK landing after the send was already reaped) and is dropped. @@ -1575,7 +1585,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None: expired = [ req_id for req_id, (deadline, _) in self._deferred_send_deadlines.items() - if now >= deadline + if self.mode == MoRIIOMode.READ and now >= deadline ] if expired: safe.update(expired) @@ -1685,6 +1695,7 @@ def __init__( # Completions that arrived before transfer_id_to_request_id was populated. # Retried each step until the mapping is established. self._unmatched_write_completions: set[str] = set() + self._unmatched_write_failures: OrderedDict[TransferId, None] = OrderedDict() # Producer-side READ-mode ACK fan-in. When decode TP is larger than # prefill TP, multiple decode ranks can read from one prefill rank and # notify the same transfer_id. Blocks are reusable only after all ACKs. @@ -1877,6 +1888,7 @@ def schedule_write_blocks( remote_ip: str, multi_pod_hosts: list[str], remote_dp_size_local: int, + remote_dp_size: int, ) -> None: """Schedule a block write operation. @@ -1892,6 +1904,7 @@ def schedule_write_blocks( remote_ip: IP address of remote node multi_pod_hosts: List of pod IPs for multi-pod Wide-EP remote_dp_size_local: Per-pod DP size for multi-pod + remote_dp_size: Global DP size of the decode instance """ # synchronization to prevent dirty reads between @@ -1915,6 +1928,7 @@ def schedule_write_blocks( remote_ip=remote_ip, multi_pod_hosts=multi_pod_hosts, remote_dp_size_local=remote_dp_size_local, + remote_dp_size=remote_dp_size, ) self._writer.schedule_write(task) @@ -2686,6 +2700,28 @@ def get_finished(self) -> tuple[set[str], set[str]]: return done_sending, done_recving + def get_transfer_results(self) -> KVConnectorTransferResults: + done_sending, done_recving = self.get_finished() + failed_recving: set[ReqId] = set() + + if self.mode == MoRIIOMode.WRITE and not self.is_producer: + pending = self._unmatched_write_failures + mapping = self.transfer_id_to_request_id + for transfer_id in self.moriio_wrapper.pop_failed_write_req_ids(): + pending[transfer_id] = None + for transfer_id in [t for t in pending if t in mapping]: + del pending[transfer_id] + failed_recving.add(mapping[transfer_id]) + while len(pending) > _MAX_UNMATCHED_WRITE_FAILURES: + pending.popitem(last=False) + done_recving |= failed_recving + + return KVConnectorTransferResults( + finished_sending=done_sending, + finished_recving=done_recving, + failed_recving=failed_recving, + ) + def wait_for_layer_load(self, layer_name: str) -> None: """Block until all in-flight READs of this layer have landed. @@ -3258,6 +3294,7 @@ def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_lay remote_ip=meta.remote_host, multi_pod_hosts=hosts, remote_dp_size_local=dp_local, + remote_dp_size=int(meta.remote_dp_size), ) def merge_contiguous_blocks( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 4ec3f72fb434..e33a1524e13f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -33,8 +33,10 @@ TransferError, TransferId, WriteTask, + fold_local_rank, get_port_offset, get_role, + pod_index, zmq_ctx, ) @@ -226,15 +228,6 @@ def _process_deferred_tasks(self) -> None: for task in self._deferred_tasks: if self._is_transfer_terminal(task.transfer_id): continue - if now - task.enqueue_time > defer_timeout: - logger.error( - "Deferred write task for request %s expired after %.1fs " - "(remote blocks never arrived), marking done", - task.request_id, - now - task.enqueue_time, - ) - self._mark_request_done(task.transfer_id) - continue if self._is_remote_ready(task): try: self._execute_write_task(task) @@ -244,11 +237,42 @@ def _process_deferred_tasks(self) -> None: task.request_id, ) self._mark_request_done(task.transfer_id) + elif ( + now - task.enqueue_time > defer_timeout + and self.worker.moriio_wrapper.fail_unallocated_write(task.transfer_id) + ): + self._notify_write_failed(task, now - task.enqueue_time) else: still_deferred.append(task) self._deferred_tasks = still_deferred + def _notify_write_failed(self, task: WriteTask, age: float) -> None: + self._clear_transfer_state(task.transfer_id) + logger.error( + "Deferred write task for request %s timed out after %.1fs waiting " + "for remote blocks", + task.request_id, + age, + ) + # No remote allocation arrived, so the decode DP rank that owns the + # request is unknown. Tell every rank; the others drop it. + for decode_dp_rank in range(max(task.remote_dp_size, 1)): + remote_ip, remote_port = self._resolve_notify_endpoint(task, decode_dp_rank) + try: + self.worker.moriio_wrapper.send_notify( + task.transfer_id, + remote_ip, + remote_port, + message_type="write_failed", + ) + except Exception: + logger.exception( + "Failed to notify decode DP rank %d that transfer %s timed out", + decode_dp_rank, + task.transfer_id, + ) + def _clear_transfer_state(self, transfer_id: TransferId) -> None: with self._write_state_lock: self._scheduled_writes.pop(transfer_id, None) @@ -314,27 +338,12 @@ def _execute_write_task(self, task: WriteTask) -> None: request_info = self._get_remote_alloc_info(task.transfer_id) with self._write_state_lock: request_info.completion_request_id = task.request_id - # Wide-EP multi-pod: task.remote_ip addresses only the first pod, - # so resolve the per-rank host from multi_pod_hosts (falls back to - # task.remote_ip for single-pod or on any indexing miss). - _remote_ip = task.remote_ip - _hosts = task.multi_pod_hosts - _dp_local = task.remote_dp_size_local - if _hosts and _dp_local > 0: - _pod_idx = int(request_info.decode_dp_rank) // _dp_local - if 0 <= _pod_idx < len(_hosts): - _remote_ip = _hosts[_pod_idx] - request_info.completion_remote_ip = _remote_ip - # Resolve the final notify port under the lock so finalize only - # reads it. The offset uses the per-pod local rank (% dp_local), - # since each pod binds notify sockets only for its local ranks; - # single-pod is bit-identical (modulus is a no-op). - _decode_dp_rank_for_port = int(request_info.decode_dp_rank) - if _dp_local > 0: - _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local - request_info.completion_notify_port = task.remote_notify_port + ( - get_port_offset(_decode_dp_rank_for_port, self.worker.tp_rank) - ) + # Resolve the final notify endpoint under the lock so finalize only + # reads it. + ( + request_info.completion_remote_ip, + request_info.completion_notify_port, + ) = self._resolve_notify_endpoint(task, request_info.decode_dp_rank) if task.transfer_id in self._sealed_writes: request_info.writes_expected = self._sealed_writes[task.transfer_id] @@ -506,6 +515,26 @@ def _finalize_if_complete( remote_port, ) + def _resolve_notify_endpoint( + self, task: WriteTask, decode_dp_rank: int + ) -> tuple[str, int]: + # Wide-EP multi-pod: task.remote_ip addresses only the first pod, so + # resolve the per-rank host from multi_pod_hosts (falls back to + # task.remote_ip for single-pod or on any indexing miss). The port + # offset uses the per-pod local rank, since each pod binds notify + # sockets only for its local ranks. + decode_dp_rank = int(decode_dp_rank) + dp_local = task.remote_dp_size_local + remote_ip = task.remote_ip + if task.multi_pod_hosts and dp_local > 0: + remote_pod = pod_index(decode_dp_rank, dp_local) + if 0 <= remote_pod < len(task.multi_pod_hosts): + remote_ip = task.multi_pod_hosts[remote_pod] + remote_port = task.remote_notify_port + get_port_offset( + fold_local_rank(decode_dp_rank, dp_local), self.worker.tp_rank + ) + return remote_ip, remote_port + class MoRIIOWrapper: """Wrapper for MoRIIO engine operations. @@ -539,6 +568,7 @@ def __init__( self.done_req_ids: list[MoRIIOTransferAck] = [] self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {} self.done_write_cache_req_ids: list[str] = [] + self.failed_write_cache_req_ids: list[str] = [] self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict() self._transfer_timeout = transfer_timeout self.notify_thread: threading.Thread | None = None @@ -803,6 +833,8 @@ def _handle_structured_message(self, data: dict): self._handle_remote_blocks_message(data) elif message_type == "write_done": self._handle_write_done_message(data) + elif message_type == "write_failed": + self._handle_write_failed_message(data) elif message_type == "release": self._handle_release_message(data) else: @@ -837,6 +869,12 @@ def _handle_write_done_message(self, data: dict): with self.lock: self.done_write_cache_req_ids.append(transfer_id) + def _handle_write_failed_message(self, data: dict): + assert get_role() != ROLE.PRODUCER, "Only decode can get WRITE failure messages" + transfer_id = data["transfer_id"] + with self.lock: + self.failed_write_cache_req_ids.append(transfer_id) + def _handle_release_message(self, data: dict): assert get_role() == ROLE.PRODUCER, ( "Only prefill can get transfer release messages" @@ -852,6 +890,23 @@ def _handle_release_message(self, data: dict): self.done_remote_allocate_req_dict.pop(transfer_id, None) self._mark_transfer_terminal_locked(transfer_id) + def fail_unallocated_write(self, transfer_id: TransferId) -> bool: + """Mark a WRITE transfer terminal and release its source blocks if + decode never sent its allocation. + + Returns False if the transfer is already terminal or the allocation + has arrived; the writer then drops or executes the task as usual. + """ + with self.lock: + if ( + self._is_transfer_terminal_locked(transfer_id) + or transfer_id in self.done_remote_allocate_req_dict + ): + return False + self._mark_transfer_terminal_locked(transfer_id) + self.done_req_ids.append(MoRIIOTransferAck(transfer_id)) + return True + def _handle_completion_message(self, msg: str): with self.lock: if get_role() == ROLE.PRODUCER: @@ -929,6 +984,12 @@ def pop_finished_write_req_ids(self): self.done_write_cache_req_ids = [] return done_write_cache + def pop_failed_write_req_ids(self): + with self.lock: + failed_write_cache = set(self.failed_write_cache_req_ids) + self.failed_write_cache_req_ids = [] + return failed_write_cache + def shutdown(self): logger.debug("Closing MoRIIOWrapper and cleaning up ZMQ sockets") for path, sock in self.paths.items(): From 1bc6e072b0ae402db5b7f054fc4d862eeb3067ea Mon Sep 17 00:00:00 2001 From: xinnywinne <36653773+xinnywinne@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:39:10 -0700 Subject: [PATCH 450/952] [Rust Frontend] Add per-request timing metrics to chat and completions (#59314) Signed-off-by: wxin Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/event.rs | 4 +- rust/src/chat/src/output/default/unified.rs | 7 + rust/src/chat/src/output/harmony/mod.rs | 1 + rust/src/chat/src/output/harmony/tests.rs | 3 + rust/src/chat/src/output/mod.rs | 2 + rust/src/chat/src/output/structured.rs | 20 ++- rust/src/chat/src/stream.rs | 6 + rust/src/chat/tests/roundtrip.rs | 2 + rust/src/cmd/src/cli.rs | 6 + rust/src/cmd/src/cli/tests.rs | 40 ++++++ rust/src/cmd/src/cli/unsupported.rs | 9 -- rust/src/llm/src/lib.rs | 2 +- rust/src/llm/src/output.rs | 10 ++ rust/src/llm/src/request_metrics.rs | 30 ++++- rust/src/llm/tests/generate.rs | 9 +- rust/src/server/src/config.rs | 7 + rust/src/server/src/grpc/convert.rs | 1 + rust/src/server/src/grpc/inference.rs | 1 + .../server/src/routes/inference/generate.rs | 3 + .../src/routes/openai/chat_completions.rs | 36 ++++++ .../routes/openai/chat_completions/types.rs | 4 + .../server/src/routes/openai/completions.rs | 107 ++++++++++++++- .../src/routes/openai/completions/types.rs | 4 + .../server/src/routes/openai/utils/metrics.rs | 122 ++++++++++++++++++ .../src/server/src/routes/openai/utils/mod.rs | 1 + rust/src/text/src/output/decoded.rs | 43 +++++- rust/src/text/src/output/mod.rs | 5 + .../chat_completion/test_chat_completion.py | 57 ++++++++ 28 files changed, 517 insertions(+), 25 deletions(-) create mode 100644 rust/src/server/src/routes/openai/utils/metrics.rs diff --git a/rust/src/chat/src/event.rs b/rust/src/chat/src/event.rs index 5fe25a9969ef..4056476b489d 100644 --- a/rust/src/chat/src/event.rs +++ b/rust/src/chat/src/event.rs @@ -5,7 +5,7 @@ use std::ops::Deref; use std::sync::Arc; use serde::{Deserialize, Serialize}; -use vllm_llm::TokenUsage; +use vllm_llm::{RequestTimestamps, TokenUsage}; use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs}; use crate::FinishReason; @@ -235,6 +235,8 @@ pub enum ChatEvent { Done { message: AssistantMessage, usage: ChatTokenUsage, + /// Engine lifecycle times used by frontend response metrics. + timestamps: RequestTimestamps, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. /// Boxed, like `ec_transfer_params`, to keep the per-token variants of diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index 9f872d97945b..a6238c440e70 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -277,6 +277,7 @@ pub(crate) async fn unified_event_stream( engine: finished.usage, reasoning_tokens: state.reasoning_tokens, }, + timestamps: finished.timestamps, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, ec_transfer_params: finished.ec_transfer_params, @@ -419,6 +420,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), } } @@ -742,6 +744,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(0), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -782,6 +785,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(5), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -814,6 +818,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(2), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -866,6 +871,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(5), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, @@ -912,6 +918,7 @@ mod tests { }, AssistantEvent::Done { usage: done_usage(0), + timestamps: Default::default(), finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index 01e0390622bc..8ec873c6f2f3 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -380,6 +380,7 @@ async fn harmony_assistant_event_stream( if let Some(finished) = finished { y.yield_ok(AssistantEvent::Done { usage: finished.usage.into(), + timestamps: finished.timestamps, finish_reason: finished.finish_reason, kv_transfer_params: finished.kv_transfer_params, ec_transfer_params: finished.ec_transfer_params, diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 5b89e53bb9aa..a529888ccff5 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -59,6 +59,7 @@ fn finished() -> Finished { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), } } @@ -132,6 +133,7 @@ fn interrupted_final_message_is_preserved() { output_token_count: 0, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -183,6 +185,7 @@ fn interrupted_analysis_message_is_preserved() { output_token_count: 0, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index 8c1dae37f4de..14ac725206e7 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use futures::Stream; use trait_set::trait_set; use uuid::Uuid; +use vllm_llm::RequestTimestamps; use vllm_parser::output_grammar::BuiltOutputGrammar; use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent}; @@ -52,6 +53,7 @@ pub(crate) enum AssistantEvent { ToolCallArgumentsDelta { delta: String }, Done { usage: ChatTokenUsage, + timestamps: RequestTimestamps, finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. /// Boxed, like `ec_transfer_params`, to keep the per-token variants of diff --git a/rust/src/chat/src/output/structured.rs b/rust/src/chat/src/output/structured.rs index 80035db4cb1e..3ed1a4e85271 100644 --- a/rust/src/chat/src/output/structured.rs +++ b/rust/src/chat/src/output/structured.rs @@ -10,6 +10,7 @@ use asynk_strim_attr::{TryYielder, try_stream}; use futures::{StreamExt as _, pin_mut}; +use vllm_llm::RequestTimestamps; use vllm_text::DecodedLogprobs; use super::{AssistantEvent, AssistantEventStream}; @@ -149,6 +150,7 @@ impl StructuredEventState { fn finish( &mut self, usage: ChatTokenUsage, + timestamps: RequestTimestamps, finish_reason: FinishReason, kv_transfer_params: Option>, ec_transfer_params: Option>, @@ -159,6 +161,7 @@ impl StructuredEventState { events.push(ChatEvent::Done { message: self.message.clone(), usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -312,13 +315,18 @@ pub(crate) async fn structured_chat_event_stream( } AssistantEvent::Done { usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, } => { - for next in - state.finish(usage, finish_reason, kv_transfer_params, ec_transfer_params)? - { + for next in state.finish( + usage, + timestamps, + finish_reason, + kv_transfer_params, + ec_transfer_params, + )? { y.yield_ok(next).await; } } @@ -353,6 +361,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -408,6 +417,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -461,6 +471,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -514,6 +525,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -582,6 +594,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -633,6 +646,7 @@ mod tests { }, reasoning_tokens: 3, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/src/stream.rs b/rust/src/chat/src/stream.rs index 509628e8e2af..fd58249461c8 100644 --- a/rust/src/chat/src/stream.rs +++ b/rust/src/chat/src/stream.rs @@ -7,6 +7,7 @@ use std::task::{Context, Poll}; use futures::Stream; use trait_set::trait_set; +use vllm_llm::RequestTimestamps; use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs}; use crate::FinishReason; @@ -22,6 +23,7 @@ pub struct CollectedAssistantMessage { pub logprobs: Option, pub token_ids: Vec, pub usage: ChatTokenUsage, + pub timestamps: RequestTimestamps, pub finish_reason: FinishReason, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, @@ -81,6 +83,7 @@ impl ChatEventStream { ChatEvent::Done { message: done, usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -94,6 +97,7 @@ impl ChatEventStream { }), token_ids, usage, + timestamps, finish_reason, kv_transfer_params: kv_transfer_params.map(|value| *value), ec_transfer_params: ec_transfer_params.map(|value| *value), @@ -200,6 +204,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -241,6 +246,7 @@ mod tests { output_token_count: 1, cached_token_count: 0, }), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 41f06d773432..d2891ed02797 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -1051,6 +1051,7 @@ fn decoded_completion_stream( kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), } }); @@ -1063,6 +1064,7 @@ fn decoded_completion_stream( kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }); events.push(DecodedTextEvent::TextDelta { decoded: chunk.delta, diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index eeb458c86e75..66b48f13d385 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -348,6 +348,11 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub enable_prompt_tokens_details: bool, + /// Include per-request timing metrics in API responses. + #[arg(long, default_missing_value = "true", num_args = 0..=1)] + #[serde(default)] + pub enable_per_request_metrics: bool, + /// If specified, API server will add X-Request-Id header to responses. #[arg( long, @@ -635,6 +640,7 @@ impl SharedRuntimeArgs { ApiServerOptions { enable_log_requests: self.enable_log_requests, enable_prompt_tokens_details: self.enable_prompt_tokens_details, + enable_per_request_metrics: self.enable_per_request_metrics, enable_request_id_headers: self.enable_request_id_headers, enable_scale_out: self.enable_scale_out, sse_keep_alive_interval: (self.sse_keep_alive_interval > 0) diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index cc545ce33e80..8f4318ac1693 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -199,6 +199,7 @@ fn serve_args_forward_python_flags_with_separator() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -970,6 +971,40 @@ fn serve_args_reject_unsupported_flag_arg() { "#]].assert_eq(&error.to_string()); } +#[test] +fn serve_args_enable_per_request_metrics_requires_engine_stats() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-per-request-metrics", + ]) + .expect("per-request metrics flag is supported"); + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let config = args.to_frontend_config("tcp://127.0.0.1:0".into()); + assert!(config.api_server_options.enable_per_request_metrics); + assert!(config.validate().is_ok()); + + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-per-request-metrics", + "--disable-log-stats", + ]) + .expect("flags parse independently"); + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let error = args.to_frontend_config("tcp://127.0.0.1:0".into()).validate().unwrap_err(); + assert!( + format!("{error:#}") + .contains("--enable-per-request-metrics requires engine statistics logging") + ); +} + #[test] fn serve_args_reject_custom_generation_config_source() { let error = Cli::try_parse_from([ @@ -1067,6 +1102,7 @@ fn frontend_args_accept_json() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -1738,6 +1774,7 @@ fn serve_args_accept_handshake_aliases() { chat_template_content_format: Auto, enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: 0, @@ -1895,6 +1932,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, @@ -1989,6 +2027,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, @@ -2127,6 +2166,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, + enable_per_request_metrics: false, enable_request_id_headers: false, enable_scale_out: false, sse_keep_alive_interval: None, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 0a14d1b841c4..02e1377ff01c 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -476,15 +476,6 @@ pub struct ServerUnsupportedArgs { #[arg(long)] pub max_log_len: Option, - /// If set to True, include per-request timing metrics in API responses. - #[arg( - long, - visible_alias = "no-enable-per-request-metrics", - default_missing_value = "true", - num_args = 0..=1 - )] - pub enable_per_request_metrics: Option, - /// If set to True, enable tracking server_load_metrics in the app state. #[arg( long, diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index c94c1d80f743..1e15f11c12e3 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -19,7 +19,7 @@ pub use output::{ GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, WeightVersionSpan, }; pub use request::GenerateRequest; -pub use request_metrics::current_unix_timestamp_secs; +pub use request_metrics::{RequestTimestamps, current_unix_timestamp_secs}; pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use crate::inflight::InflightRequests; diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 81dbb3fda53c..3f0b0300626d 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -20,6 +20,7 @@ use vllm_engine_core_client::{ AbortCause, EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput, }; +use crate::RequestTimestamps; use crate::error::Result; use crate::inflight::RequestGuard; use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs}; @@ -183,6 +184,12 @@ pub struct GenerateOutput { pub spec_decode_metrics: Option>, /// Weight-version spans over all output tokens, set on terminal outputs only. pub weight_versions: Option>, + /// Snapshot of engine lifecycle timestamps after this output. + /// + /// Every output carries a snapshot because text decoding may finish the + /// request locally after matching a stop string, before the engine emits a + /// terminal output. + pub timestamps: RequestTimestamps, } impl GenerateOutput { @@ -237,6 +244,7 @@ impl GenerateOutput { self.ec_transfer_params = next.ec_transfer_params; self.spec_decode_metrics = next.spec_decode_metrics; self.weight_versions = next.weight_versions; + self.timestamps = next.timestamps; } } @@ -266,6 +274,7 @@ impl GenerateOutput { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: RequestTimestamps::default(), } } } @@ -402,6 +411,7 @@ impl GenerateOutputStream { sampling_mask, spec_decode_metrics: raw.spec_decode_metrics, weight_versions, + timestamps: self.request_metrics.timestamps(), }; Ok(output) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 71decdd281ff..eebf0596b0d2 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -18,6 +18,15 @@ const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute"; const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit"; const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; +/// Engine lifecycle timestamps used by frontend per-request response metrics. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct RequestTimestamps { + pub queued_ts: f64, + pub scheduled_ts: f64, + pub first_token_ts: f64, + pub last_token_ts: f64, +} + /// Request-scoped metrics state tracked across streamed engine-core updates. /// /// This is the Rust-side counterpart of the Python frontend's request-lifecycle @@ -227,6 +236,16 @@ impl RequestMetricsTracker { } } + /// Engine lifecycle timestamps recorded so far. + pub(crate) fn timestamps(&self) -> RequestTimestamps { + RequestTimestamps { + queued_ts: self.queued_ts, + scheduled_ts: self.scheduled_ts, + first_token_ts: self.first_token_ts, + last_token_ts: self.last_token_ts, + } + } + /// Record prompt token counters through cached metric handles. fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) { let computed = prefill_stats.num_computed_tokens as u64; @@ -388,7 +407,7 @@ mod tests { use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType}; use vllm_engine_core_client::protocol::stats::PrefillStats; - use super::{RequestMetricsTracker, diff_or_zero}; + use super::{RequestMetricsTracker, RequestTimestamps, diff_or_zero}; #[test] fn tracker_updates_timing_state_across_prefill_decode_and_finish() { @@ -445,6 +464,15 @@ mod tests { assert_eq!(tracker.scheduled_ts, 9.0); assert_eq!(tracker.first_token_ts, 10.0); assert_eq!(tracker.last_token_ts, 11.5); + assert_eq!( + tracker.timestamps(), + RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + } + ); assert!((tracker.first_token_latency - 0.2).abs() < 1e-9); assert_eq!( diff_or_zero(tracker.last_token_ts, tracker.first_token_ts), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 0a1eacd73ff6..cc6fd9d169f6 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -352,7 +352,9 @@ async fn generate_merges_outputs_by_stream_interval() { (&[6], "v7", None), (&[7], "v7", Some(EngineCoreFinishReason::Length)), ]; - for (token_ids, weight_version, finish_reason) in steps { + for (step, (token_ids, weight_version, finish_reason)) in + steps.into_iter().enumerate() + { send_outputs( push, RequestBatchOutputs { @@ -364,6 +366,7 @@ async fn generate_merges_outputs_by_stream_interval() { None, )], weight_version: Some(weight_version.into()), + timestamp: step as f64 + 1.0, ..Default::default() } .into(), @@ -430,6 +433,9 @@ async fn generate_merges_outputs_by_stream_interval() { Some(vec![span("v7", 0, 2), span("v8", 2, 5), span("v7", 5, 7)]) ] ); + let last_token_timestamps: Vec<_> = + outputs.iter().map(|output| output.timestamps.last_token_ts).collect(); + assert_eq!(last_token_timestamps, [1.0, 4.0, 6.0]); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -549,6 +555,7 @@ async fn collect_output_rejects_partial_sampling_mask() { }), spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }; let error = futures::stream::iter([Ok(output)]).collect_output().await.unwrap_err(); diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index ca91ba40cd50..d70933fb446d 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -56,6 +56,8 @@ pub struct ApiServerOptions { pub enable_log_requests: bool, /// When `true`, include prompt token cache details in response usage. pub enable_prompt_tokens_details: bool, + /// Include request timing metrics in OpenAI chat and completion responses. + pub enable_per_request_metrics: bool, /// When `true`, set `X-Request-Id` on every HTTP response. pub enable_request_id_headers: bool, /// When `true`, register the scale-out `/inference/v1/generate` route. @@ -296,6 +298,11 @@ impl Config { /// Validate frontend configuration that can be checked before engine /// startup. pub fn validate(&self) -> Result<()> { + if self.api_server_options.enable_per_request_metrics && self.disable_log_stats { + bail!( + "--enable-per-request-metrics requires engine statistics logging; remove --disable-log-stats" + ); + } vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?; self.cors.validate()?; if let Some(tls) = &self.tls { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 8e31bb84d2d3..1eec1d0fe320 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -947,6 +947,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), } } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 1484a9f9e375..fae2794ba331 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -259,6 +259,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl { kv_transfer_params: collected.kv_transfer_params.map(Box::new), ec_transfer_params: collected.ec_transfer_params.map(Box::new), sampling_mask: collected.sampling_mask, + timestamps: collected.timestamps, }; let outputs = convert::to_sequence_output( diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 930c698c30e4..42716de8de37 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -637,6 +637,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }), Ok(GenerateOutput { request_id: String::new(), @@ -654,6 +655,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), }), ]); @@ -741,6 +743,7 @@ mod tests { sampling_mask: None, spec_decode_metrics: None, weight_versions: None, + timestamps: Default::default(), } } diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index a87c52fb5578..5d8dcda53641 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -38,6 +38,7 @@ use crate::routes::openai::chat_completions::types::{ use crate::routes::openai::utils::logprobs::{ decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps, }; +use crate::routes::openai::utils::metrics::PerRequestMetrics; use crate::routes::openai::utils::types::{ ChatLogProbs, FunctionCallDelta, FunctionCallResponse, StreamResponseEnvelope, ToolCall, ToolCallDelta, Usage, @@ -128,6 +129,7 @@ async fn collect_chat_completion( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -158,6 +160,7 @@ async fn collect_chat_completion( logprobs, token_ids, usage, + timestamps, finish_reason, kv_transfer_params, ec_transfer_params, @@ -202,6 +205,8 @@ async fn collect_chat_completion( } else { None }; + let metrics = enable_per_request_metrics + .then(|| PerRequestMetrics::from_timestamps(timestamps, usage.output_token_count)); let usage = Usage::from_token_usage(usage, enable_prompt_tokens_details); if enable_log_requests { @@ -236,6 +241,7 @@ async fn collect_chat_completion( token_ids: (return_token_ids && include_output_metadata).then_some(token_ids), }], usage: Some(usage), + metrics, system_fingerprint: None, prompt_logprobs, prompt_token_ids: return_token_ids.then(|| prompt_token_ids.to_vec()), @@ -254,6 +260,7 @@ async fn chat_completion_chunk_stream( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -428,6 +435,7 @@ async fn chat_completion_chunk_stream( } Ok(ChatEvent::Done { usage: final_usage, + timestamps, finish_reason, .. }) => { @@ -472,9 +480,16 @@ async fn chat_completion_chunk_stream( } if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + timestamps, + final_usage.output_token_count, + ) + }); y.yield_ok(usage_chunk( &envelope, Usage::from_token_usage(final_usage, enable_prompt_tokens_details), + metrics, )) .await; } @@ -496,9 +511,11 @@ async fn chat_completion_chunk_stream( fn usage_chunk( envelope: &Arc, usage: Usage, + metrics: Option, ) -> ChatCompletionStreamResponse { let mut chunk = ChatCompletionStreamResponse::new(envelope); chunk.usage = Some(usage); + chunk.metrics = metrics.map(Into::into); chunk } @@ -962,6 +979,12 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 1), + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }, finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -975,6 +998,7 @@ mod tests { 1, ApiServerOptions { enable_prompt_tokens_details: true, + enable_per_request_metrics: true, ..Default::default() }, ResponseOptions { @@ -1006,6 +1030,10 @@ mod tests { .map(|details| details.cached_tokens), Some(1) ); + assert!(chunks[..3].iter().all(|chunk| chunk.metrics.is_none())); + let metrics = chunks[3].metrics.as_ref().expect("per-request metrics"); + assert!((metrics.time_to_first_token_ms.unwrap() - 300.0).abs() < 1e-9); + assert!((metrics.tokens_per_second.unwrap() - 1.25).abs() < 1e-9); } #[tokio::test] @@ -1041,6 +1069,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1095,6 +1124,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 2, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1175,6 +1205,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 2, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1307,6 +1338,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 4, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1385,6 +1417,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1468,6 +1501,7 @@ mod tests { }, reasoning_tokens: 2, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1539,6 +1573,7 @@ mod tests { }, reasoning_tokens: 2, }, + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, @@ -1593,6 +1628,7 @@ mod tests { Ok(ChatEvent::Done { message: Default::default(), usage: done_usage(1, 1, 0), + timestamps: Default::default(), finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 21776dd075e0..d73480b45d5e 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -13,6 +13,7 @@ use validator::Validate; use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::TruncationSide; +use crate::routes::openai::utils::metrics::{PerRequestMetrics, StreamPerRequestMetrics}; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ ChatLogProbs, ChatMessage, Normalizable, PromptLogprobs, ReasoningEffort, StreamOptions, @@ -365,6 +366,7 @@ pub(super) struct ChatCompletionResponse { pub model: String, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub system_fingerprint: Option, pub prompt_logprobs: Option, pub prompt_token_ids: Option>, @@ -412,6 +414,7 @@ pub(super) struct ChatCompletionStreamResponse { pub envelope: Arc, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub prompt_token_ids: Option>, } @@ -422,6 +425,7 @@ impl ChatCompletionStreamResponse { envelope: Arc::clone(envelope), choices: Vec::new(), usage: None, + metrics: None, prompt_token_ids: None, } } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 646a67d1fb8f..05dd579ecec7 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -33,6 +33,7 @@ use super::utils::logprobs::{ collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai, prompt_logprobs_to_maps, text_len, }; +use super::utils::metrics::PerRequestMetrics; use super::utils::types::{StreamResponseEnvelope, Usage}; use crate::config::ApiServerOptions; use crate::error::{ApiError, bail_server_error, server_error, text_submit_error}; @@ -140,6 +141,7 @@ async fn collect_completion( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -197,6 +199,9 @@ async fn collect_completion( Some(prompt) => format!("{prompt}{}", collected.text), }; let finish_reason = completion_finish_reason_to_openai(&finish_reason)?.to_string(); + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps(collected.timestamps, collected.usage.output_token_count) + }); let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details); if enable_log_requests { @@ -225,6 +230,7 @@ async fn collect_completion( prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()), }], usage: Some(usage), + metrics, system_fingerprint: None, kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, @@ -241,6 +247,7 @@ async fn completion_chunk_stream( ApiServerOptions { enable_log_requests, enable_prompt_tokens_details, + enable_per_request_metrics, .. }: ApiServerOptions, ResponseOptions { @@ -356,12 +363,19 @@ async fn completion_chunk_stream( yield_chunk!(final_chunk); if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + finished.timestamps, + finished.usage.output_token_count, + ) + }); y.yield_ok(CompletionSseChunk::Usage(usage_chunk( &envelope, Usage::from_token_usage( finished.usage, enable_prompt_tokens_details, ), + metrics, ))) .await; } @@ -411,9 +425,16 @@ async fn completion_chunk_stream( yield_chunk!(final_chunk); if include_usage { + let metrics = enable_per_request_metrics.then(|| { + PerRequestMetrics::from_timestamps( + finished.timestamps, + finished.usage.output_token_count, + ) + }); y.yield_ok(CompletionSseChunk::Usage(usage_chunk( &envelope, Usage::from_token_usage(finished.usage, enable_prompt_tokens_details), + metrics, ))) .await; } @@ -509,9 +530,14 @@ fn prompt_only_logprobs_to_openai( )) } -fn usage_chunk(envelope: &Arc, usage: Usage) -> CompletionStreamResponse { +fn usage_chunk( + envelope: &Arc, + usage: Usage, + metrics: Option, +) -> CompletionStreamResponse { let mut chunk = CompletionStreamResponse::new(envelope); chunk.usage = Some(usage); + chunk.metrics = metrics.map(Into::into); chunk } @@ -578,7 +604,7 @@ mod tests { use super::{ ApiServerOptions, CompletionSseChunk, CompletionStreamResponse, ResponseOptions, - StreamResponseEnvelope, completion_chunk_stream, final_chunk, + StreamResponseEnvelope, collect_completion, completion_chunk_stream, final_chunk, }; fn decoded_delta( @@ -639,6 +665,66 @@ mod tests { assert!(final_chunk(&stream_envelope(), FinishReason::Error).is_err()); } + #[tokio::test] + async fn collected_completion_includes_metrics_only_when_enabled() { + let make_stream = || { + stream::iter(vec![ + Ok(DecodedTextEvent::Start { + prompt_token_ids: vec![1].into(), + prompt_logprobs: None, + }), + Ok(decoded_delta( + "hi", + vec![2], + None, + Some(Finished { + usage: vllm_llm::TokenUsage { + prompt_token_count: 1, + output_token_count: 1, + cached_token_count: 0, + }, + finish_reason: FinishReason::stop_eos(), + kv_transfer_params: None, + ec_transfer_params: None, + sampling_mask: None, + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 10.5, + }, + }), + )), + ]) + }; + + for enabled in [false, true] { + let response = collect_completion( + make_stream(), + "cmpl-1".into(), + "model".into(), + 1, + ApiServerOptions { + enable_per_request_metrics: enabled, + ..Default::default() + }, + ResponseOptions::default(), + ) + .await + .expect("completion response"); + let json = serde_json::to_value(response).expect("serializable response"); + if enabled { + assert!( + (json["metrics"]["time_to_first_token_ms"].as_f64().unwrap() - 300.0).abs() + < 1e-9 + ); + assert_eq!(json["metrics"]["mean_itl_ms"], serde_json::Value::Null); + } else { + assert_eq!(json["metrics"], serde_json::Value::Null); + } + } + } + #[tokio::test] async fn completion_chunk_stream_maps_streaming_logprobs() { let stream = stream::iter(vec![ @@ -702,6 +788,12 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: vllm_llm::RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }, }), )), ]); @@ -713,6 +805,7 @@ mod tests { 1, ApiServerOptions { enable_prompt_tokens_details: true, + enable_per_request_metrics: true, ..Default::default() }, ResponseOptions { @@ -779,6 +872,10 @@ mod tests { match &chunks[3] { CompletionSseChunk::Usage(chunk) => { + assert!(chunks[..3].iter().all(|chunk| chunk_response(chunk).metrics.is_none())); + let metrics = chunk.metrics.as_ref().expect("per-request metrics"); + assert!((metrics.queue_time_ms.unwrap() - 200.0).abs() < 1e-9); + assert!((metrics.tokens_per_second.unwrap() - 2.5).abs() < 1e-9); assert_eq!( chunk .usage @@ -815,6 +912,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -879,6 +977,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -932,6 +1031,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -988,6 +1088,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -1061,6 +1162,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); @@ -1136,6 +1238,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), )), ]); diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 9f7027511cbf..fb6aa22c601d 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -11,6 +11,7 @@ use validator::Validate; use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::{Prompt, TruncationSide}; +use crate::routes::openai::utils::metrics::{PerRequestMetrics, StreamPerRequestMetrics}; use crate::routes::openai::utils::types::{ LogProbs, Normalizable, PromptLogprobs, StreamOptions, StreamResponseEnvelope, StringOrArray, Usage, default_true, deserialize_request_top_k, validate_stop, @@ -239,6 +240,7 @@ pub(super) struct CompletionResponse { pub model: String, pub choices: Vec, pub usage: Option, + pub metrics: Option, pub system_fingerprint: Option, pub kv_transfer_params: Option, pub ec_transfer_params: Option, @@ -265,6 +267,7 @@ pub(super) struct CompletionStreamResponse { pub envelope: Arc, pub choices: Vec, pub usage: Option, + pub metrics: Option, } impl CompletionStreamResponse { @@ -274,6 +277,7 @@ impl CompletionStreamResponse { envelope: Arc::clone(envelope), choices: Vec::new(), usage: None, + metrics: None, } } } diff --git a/rust/src/server/src/routes/openai/utils/metrics.rs b/rust/src/server/src/routes/openai/utils/metrics.rs new file mode 100644 index 000000000000..0c1155f733a4 --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/metrics.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use serde::Serialize; +use vllm_llm::RequestTimestamps; + +/// Per-request response timing fields matching Python's `PerRequestMetrics`. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct PerRequestMetrics { + pub time_to_first_token_ms: Option, + pub generation_time_ms: Option, + pub queue_time_ms: Option, + pub mean_itl_ms: Option, + pub tokens_per_second: Option, +} + +impl PerRequestMetrics { + pub fn from_timestamps(times: RequestTimestamps, num_generation_tokens: usize) -> Self { + let duration = |end: f64, start: f64| (end > 0.0 && start > 0.0).then_some(end - start); + let decode = duration(times.last_token_ts, times.first_token_ts); + let inference = duration(times.last_token_ts, times.scheduled_ts); + + Self { + time_to_first_token_ms: duration(times.first_token_ts, times.scheduled_ts) + .map(|seconds| seconds * 1000.0), + generation_time_ms: decode.map(|seconds| seconds * 1000.0), + queue_time_ms: duration(times.scheduled_ts, times.queued_ts) + .map(|seconds| seconds * 1000.0), + mean_itl_ms: decode + .filter(|_| num_generation_tokens > 1) + .map(|seconds| seconds * 1000.0 / (num_generation_tokens - 1) as f64), + tokens_per_second: inference + .filter(|seconds| *seconds > 0.0) + .map(|seconds| num_generation_tokens as f64 / seconds), + } + } +} + +/// Streaming form of [`PerRequestMetrics`]. +/// +/// Python serializes the final usage chunk with recursive `exclude_none=True`, +/// so unset fields are omitted instead of serialized as `null`. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct StreamPerRequestMetrics { + pub time_to_first_token_ms: Option, + pub generation_time_ms: Option, + pub queue_time_ms: Option, + pub mean_itl_ms: Option, + pub tokens_per_second: Option, +} + +impl From for StreamPerRequestMetrics { + fn from(metrics: PerRequestMetrics) -> Self { + let PerRequestMetrics { + time_to_first_token_ms, + generation_time_ms, + queue_time_ms, + mean_itl_ms, + tokens_per_second, + } = metrics; + Self { + time_to_first_token_ms, + generation_time_ms, + queue_time_ms, + mean_itl_ms, + tokens_per_second, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matches_python_timing_formulas_and_missing_fields() { + let times = RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 11.0, + }; + let metrics = PerRequestMetrics::from_timestamps(times, 6); + assert!((metrics.time_to_first_token_ms.unwrap() - 300.0).abs() < 1e-9); + assert_eq!(metrics.generation_time_ms, Some(500.0)); + assert!((metrics.queue_time_ms.unwrap() - 200.0).abs() < 1e-9); + assert_eq!(metrics.mean_itl_ms, Some(100.0)); + assert!((metrics.tokens_per_second.unwrap() - 7.5).abs() < 1e-9); + + let single_token = PerRequestMetrics::from_timestamps(times, 1); + assert_eq!(single_token.mean_itl_ms, None); + + let missing = PerRequestMetrics::from_timestamps(RequestTimestamps::default(), 0); + assert_eq!(missing.time_to_first_token_ms, None); + assert_eq!(missing.generation_time_ms, None); + assert_eq!(missing.queue_time_ms, None); + assert_eq!(missing.mean_itl_ms, None); + assert_eq!(missing.tokens_per_second, None); + } + + #[test] + fn stream_serialization_omits_unset_fields() { + let metrics = PerRequestMetrics::from_timestamps( + RequestTimestamps { + queued_ts: 10.0, + scheduled_ts: 10.2, + first_token_ts: 10.5, + last_token_ts: 10.5, + }, + 1, + ); + let value = serde_json::to_value(StreamPerRequestMetrics::from(metrics)).unwrap(); + let metrics = value.as_object().unwrap(); + assert_eq!(metrics.len(), 4); + assert!(metrics.contains_key("time_to_first_token_ms")); + assert!(metrics.contains_key("generation_time_ms")); + assert!(metrics.contains_key("queue_time_ms")); + assert!(metrics.contains_key("tokens_per_second")); + assert!(!metrics.contains_key("mean_itl_ms")); + } +} diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 4bbea924a38b..a5f00ae0eadf 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project pub mod logprobs; +pub mod metrics; pub mod structured_outputs; pub mod types; pub mod usage; diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index 64df0492bfbe..6308bdf6e116 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -10,7 +10,7 @@ use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; -use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; +use vllm_llm::{FinishReason, GenerateOutput, RequestTimestamps, TokenUsage}; use vllm_tokenizer::{DecodedText, DynTokenizer, IncrementalDecoder}; use super::logprobs::{ @@ -55,6 +55,7 @@ pub struct Finished { pub ec_transfer_params: Option>, /// Sampling support sets aligned with all generated token positions. pub sampling_mask: Option, + pub timestamps: RequestTimestamps, } /// Sample metadata emitted by one engine output update. @@ -317,6 +318,7 @@ pub async fn decoded_text_event_stream( kv_transfer_params, ec_transfer_params, sampling_mask, + timestamps: output.timestamps, })), }) .await; @@ -543,7 +545,19 @@ mod tests { name: "utf8 byte fallback resolves across engine updates", outputs: vec![ GenerateOutput::for_test(Some(Arc::clone(&prompt)), vec![0xe4], None), - GenerateOutput::for_test(None, vec![0xbd, 0xa0], Some(FinishReason::Length)), + GenerateOutput { + timestamps: RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + }, + ..GenerateOutput::for_test( + None, + vec![0xbd, 0xa0], + Some(FinishReason::Length), + ) + }, ], expected: vec![ DecodedTextEvent::Start { @@ -565,6 +579,12 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: RequestTimestamps { + queued_ts: 8.0, + scheduled_ts: 9.0, + first_token_ts: 10.0, + last_token_ts: 11.5, + }, }), ), ], @@ -597,6 +617,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), }), ), ], @@ -634,12 +655,19 @@ mod tests { async fn stream_stop_string_sets_task_local_abort_cause_on_raw_stream_drop() { let prompt: Arc<[u32]> = Arc::from([]); let dropped_cause = Arc::new(Mutex::new(None)); + let timestamps = RequestTimestamps { + queued_ts: 1.0, + scheduled_ts: 2.0, + first_token_ts: 3.0, + last_token_ts: 4.0, + }; let raw_stream = DropRecordingStream { - next: Some(Ok(GenerateOutput::for_test( - Some(prompt), - ascii_tokens("hello"), - Some(FinishReason::Length), - ))), + // The engine output is not terminal. The frontend terminates the + // request after decoding the configured stop string. + next: Some(Ok(GenerateOutput { + timestamps, + ..GenerateOutput::for_test(Some(prompt), ascii_tokens("hello"), None) + })), dropped_cause: Arc::clone(&dropped_cause), }; let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new()); @@ -657,6 +685,7 @@ mod tests { assert_eq!(output.text, "he"); assert!(output.finish_reason.is_stop()); + assert_eq!(output.timestamps, timestamps); assert_eq!( *dropped_cause.lock().unwrap(), Some(AbortCause::StopStringMatched) diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index 8e09f0761870..b69d439e3f25 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -31,6 +31,7 @@ pub struct CollectedTextOutput { pub token_ids: Vec, pub finish_reason: FinishReason, pub usage: vllm_llm::TokenUsage, + pub timestamps: vllm_llm::RequestTimestamps, /// Connector-specific KV transfer parameters for disaggregated serving. pub kv_transfer_params: Option, /// Connector-specific encoder cache transfer parameters for disaggregated @@ -91,6 +92,7 @@ impl T { token_ids: delta_token_ids, finish_reason: FinishReason::Error, usage: vllm_llm::TokenUsage::default(), + timestamps: vllm_llm::RequestTimestamps::default(), kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, @@ -101,6 +103,7 @@ impl T { let mut collected = collected.unwrap(); collected.finish_reason = finished.finish_reason; collected.usage = finished.usage; + collected.timestamps = finished.timestamps; collected.kv_transfer_params = finished.kv_transfer_params.map(|value| *value); collected.ec_transfer_params = @@ -181,6 +184,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), }), ]); @@ -304,6 +308,7 @@ mod tests { kv_transfer_params: None, ec_transfer_params: None, sampling_mask: None, + timestamps: Default::default(), })), }), ]); diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion.py b/tests/entrypoints/openai/chat_completion/test_chat_completion.py index e0422a520aae..ff6595fb0c62 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_completion.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion.py @@ -20,6 +20,7 @@ def default_server_args(): "--max-num-seqs", "128", "--enforce-eager", + "--enable-per-request-metrics", ] @@ -247,3 +248,59 @@ async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI # streamed text-out, reconstructed from deltas, with generated token ids. assert content assert delta_token_ids + + +def _assert_complete_timing_metrics(metrics: object) -> None: + assert isinstance(metrics, dict) + timing_fields = { + "time_to_first_token_ms", + "generation_time_ms", + "queue_time_ms", + "mean_itl_ms", + "tokens_per_second", + } + assert timing_fields <= metrics.keys() + assert all(metrics[field] >= 0 for field in timing_fields) + + +@pytest.mark.asyncio +async def test_per_request_metrics_non_streaming(client: openai.AsyncOpenAI): + response = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Count to four."}], + max_completion_tokens=4, + temperature=0, + extra_body={"min_tokens": 4}, + ) + + _assert_complete_timing_metrics(response.model_extra.get("metrics")) + + +@pytest.mark.asyncio +async def test_per_request_metrics_on_final_usage_chunk(client: openai.AsyncOpenAI): + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=[{"role": "user", "content": "Say hello."}], + max_completion_tokens=1, + temperature=0, + stream=True, + stream_options={"include_usage": True}, + extra_body={"min_tokens": 1}, + ) + + chunks = [chunk async for chunk in stream] + assert len(chunks) > 1 + for chunk in chunks[:-1]: + assert chunk.usage is None + assert "metrics" not in (chunk.model_extra or {}) + + usage_chunk = chunks[-1] + assert usage_chunk.usage is not None + assert usage_chunk.choices == [] + metrics = (usage_chunk.model_extra or {}).get("metrics") + assert isinstance(metrics, dict) + assert metrics["time_to_first_token_ms"] >= 0 + assert metrics["generation_time_ms"] >= 0 + assert metrics["queue_time_ms"] >= 0 + assert metrics["tokens_per_second"] >= 0 + assert "mean_itl_ms" not in metrics From 148426693823ec8fa27c353a1f8b5b9f7d402d8c Mon Sep 17 00:00:00 2001 From: Tomas Ruiz Date: Tue, 13 Oct 2026 01:01:11 +0200 Subject: [PATCH 451/952] [Perf] Improve `_topk_topp_kernel()` for masked logits (#61213) Signed-off-by: Tomas Ruiz Co-authored-by: Codex Co-authored-by: Nick Hill --- benchmarks/benchmark_topk_topp.py | 325 +++++++++++++------------ vllm/v1/sample/ops/topk_topp_triton.py | 5 +- 2 files changed, 169 insertions(+), 161 deletions(-) diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index 14111f733b9a..2c5bbeb3eba0 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -1,23 +1,29 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations. +"""Benchmark top-k/top-p filtering with Triton, FlashInfer, and optional PyTorch. -Compares: -- apply_top_k_top_p_triton (Triton binary search) -- apply_top_k_top_p (PyTorch sort-based) +All implementations return masked logits. Timings include a fresh-input copy. +FlashInfer timings include softmax and mask reconstruction when its APIs require them. +L2 is flushed before each measured iteration, outside the timed region. Scenarios: - top_k only (whole batch, partial batch) - top_p only (whole batch, partial batch) - mix of top_k and top_p +- optional sparse inputs via --num-finite-logits """ import argparse import gc +import json from dataclasses import dataclass +from pathlib import Path +import flashinfer.sampling as fi_sampling +import pandas as pd import torch +from flashinfer.testing import bench_gpu_time_with_cupti from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch from vllm.v1.sample.ops.topk_topp_triton import ( @@ -65,14 +71,18 @@ def calculate_ops_pct( def create_logits( - batch_size: int, vocab_size: int, device: str = "cuda" + batch_size: int, + vocab_size: int, + device: str = "cuda", + num_finite_logits: int | None = None, ) -> torch.Tensor: """Create random logits mimicking a realistic LLM distribution. Uses a Zipf-like probability distribution (rank^-1.1) converted to logits via log, then randomly permuted per row. This produces a peaked distribution where a small number of tokens capture most probability mass, similar to - real model outputs. + real model outputs. When num_finite_logits is set, retain that many ranks + before permuting their positions independently in each row. """ # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha) ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device) @@ -81,6 +91,8 @@ def create_logits( # Convert to logits (log-probabilities, unnormalized is fine) base_logits = probs.log() + if num_finite_logits is not None: + base_logits[num_finite_logits:] = -float("inf") # Broadcast to batch and randomly permute each row logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone() @@ -90,6 +102,26 @@ def create_logits( return logits +def apply_top_k_top_p_flashinfer_mask(logits, k, p): + """Keep top-k in logits; adapt probability-based top-p back to logits.""" + if k is not None: + logits = fi_sampling.top_k_mask_logits(logits, k) + if p is not None: + # Top-p renormalization requires softmax, even after logits-based top-k. + probs = logits.softmax(dim=-1) + probs = fi_sampling.top_p_renorm_probs(probs, p, is_deterministic=True) + # Return the selected original logits, not normalized probabilities. + logits = logits.masked_fill(probs == 0, -float("inf")) + return logits + + +IMPLEMENTATIONS = { + "triton": apply_top_k_top_p_triton, + "flashinfer-mask": apply_top_k_top_p_flashinfer_mask, + "pytorch": apply_top_k_top_p_pytorch, +} + + def measure_memory() -> tuple[int, int]: """Return (allocated, reserved) memory in bytes.""" torch.accelerator.synchronize() @@ -107,53 +139,6 @@ def reset_memory_stats(): gc.collect() -def benchmark_function( - func, - logits: torch.Tensor, - k: torch.Tensor | None, - p: torch.Tensor | None, - warmup_iters: int = 5, - benchmark_iters: int = 20, -) -> tuple[float, int]: - """Benchmark a function and return (avg_time_ms, peak_memory_bytes). - - Returns average time in milliseconds and peak memory usage. - """ - # Warmup - for _ in range(warmup_iters): - logits_copy = logits.clone() - func(logits_copy, k, p) - torch.accelerator.synchronize() - - # Reset memory stats before benchmark - reset_memory_stats() - - # Benchmark - start_events = [ - torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters) - ] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)] - - for i in range(benchmark_iters): - logits_copy = logits.clone() - start_events[i].record() - func(logits_copy, k, p) - end_events[i].record() - - torch.accelerator.synchronize() - - # Calculate timing - times = [ - start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters) - ] - avg_time = sum(times) / len(times) - - # Get peak memory - _, peak_memory = measure_memory() - - return avg_time, peak_memory - - def create_benchmark_configs( batch_sizes: list[int], vocab_sizes: list[int], @@ -286,118 +271,103 @@ def format_memory(bytes_val: int) -> str: def run_benchmark( configs: list[BenchmarkConfig], + implementations: list[str], warmup_iters: int = 5, benchmark_iters: int = 20, verbose: bool = True, + num_finite_logits: int | None = None, + check: bool = False, + output: Path | None = None, ): - """Run all benchmarks and print results.""" + """Run selected implementations with identical inputs and masked-logit outputs.""" results = [] - print("=" * 100) - print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based") - print("=" * 100) - print() - - for config in configs: - if verbose: - print(f"Running: {config.description}") - - # Create fresh logits for this config - logits = create_logits(config.batch_size, config.vocab_size) - - # Benchmark Triton - reset_memory_stats() - triton_time, triton_mem = benchmark_function( - apply_top_k_top_p_triton, - logits, - config.k_values, - config.p_values, - warmup_iters, - benchmark_iters, + for index, config in enumerate(configs): + logits = create_logits( + config.batch_size, + config.vocab_size, + num_finite_logits=num_finite_logits, ) - - # Benchmark PyTorch - reset_memory_stats() - pytorch_time, pytorch_mem = benchmark_function( - apply_top_k_top_p_pytorch, - logits, - config.k_values, - config.p_values, - warmup_iters, - benchmark_iters, - ) - - speedup = pytorch_time / triton_time if triton_time > 0 else float("inf") - mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf") - - result = { - "config": config, - "triton_time_ms": triton_time, - "pytorch_time_ms": pytorch_time, - "triton_mem": triton_mem, - "pytorch_mem": pytorch_mem, - "speedup": speedup, - "mem_ratio": mem_ratio, - } - results.append(result) - + checks = {} + if check: + checks = check_implementations(logits, config, implementations) + k, p = config.k_values, config.p_values + timings = {} + for name in implementations: + func = IMPLEMENTATIONS[name] + reset_memory_stats() + times = bench_gpu_time_with_cupti( + # note: measured runtime includes logits.clone() + lambda func=func, logits=logits, k=k, p=p: func(logits.clone(), k, p), + dry_run_iters=warmup_iters, + repeat_iters=benchmark_iters, + use_cuda_graph=False, # sampling is outside cuda graph + cold_l2_cache=True, + ) + elapsed = sum(times) / len(times) + _, memory = measure_memory() + timings[name] = dict(time_ms=elapsed, peak_memory=memory) + row = create_result_row(config, timings, checks) + results.append(row) + if output is not None: + with output.open("a") as file: + file.write(json.dumps(row) + "\n") if verbose: - print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}") - print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}") - print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x") - print() - - # Clean up + print(f"Completed {index + 1}/{len(configs)}: {config.name}", flush=True) del logits reset_memory_stats() return results -def print_summary_table(results: list[dict]): - """Print a summary table of results.""" - print() - print("=" * 130) - print("SUMMARY TABLE") - print("=" * 130) - print() +def check_implementations( + logits: torch.Tensor, + config: BenchmarkConfig, + implementations: list[str], +) -> dict: + """Compare filtered distributions against PyTorch outside the timed region.""" + k, p = config.k_values, config.p_values + reference = apply_top_k_top_p_pytorch(logits.clone(), k, p).softmax(dim=-1) + checks = {} + for name in implementations: + actual = IMPLEMENTATIONS[name](logits.clone(), k, p).softmax(dim=-1) + total_variation = (actual - reference).abs().sum(-1).max().item() / 2 + checks[name] = { + "close": torch.allclose(actual, reference, rtol=1e-5, atol=1e-6), + "max_total_variation": total_variation, + } + if not checks[name]["close"]: + print(f"Distribution difference: {name} {config.name} {checks[name]}") + del actual + return checks - # Header - header = ( - f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} " - f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} " - f"{'Tri Mem':>10} {'Pyt Mem':>10}" - ) - print(header) - print("-" * 130) - - # Group by scenario type - current_vocab = None - for result in results: - config = result["config"] - - # Add separator between vocab sizes - if current_vocab != config.vocab_size: - if current_vocab is not None: - print("-" * 130) - current_vocab = config.vocab_size - - scenario = config.name.split("_b")[0] # Extract scenario name - print( - f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} " - f"{config.ops_pct:>5.0f}% " - f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} " - f"{result['speedup']:>7.2f}x " - f"{format_memory(result['triton_mem']):>10} " - f"{format_memory(result['pytorch_mem']):>10}" - ) - print("=" * 130) +def print_summary_table(results: list[dict]): + """Print the same result rows saved to the output file.""" + print() + print(pd.DataFrame(results).to_string(index=False, float_format="{:.3g}".format)) + + +def create_result_row(config, timings, checks) -> dict: + """Build a shared row for the summary table and JSONL output.""" + row = { + "Scenario": config.name.split("_b")[0], + "Batch": config.batch_size, + "Vocab": config.vocab_size, + "Ops%": f"{config.ops_pct:.0f}%", + } + row.update({f"{name} (ms)": data["time_ms"] for name, data in timings.items()}) + for name, data in timings.items(): + row[f"{name} mem"] = format_memory(data["peak_memory"]) + for name, data in checks.items(): + row[f"{name} check"] = data["close"] + row[f"{name} max TV"] = data["max_total_variation"] + return row def main(): parser = argparse.ArgumentParser( - description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations" + description="Benchmark Triton and FlashInfer top-k/top-p filtering" ) parser.add_argument( "--batch-sizes", @@ -431,29 +401,60 @@ def main(): help="Only print summary table", ) - args = parser.parse_args() + parser.add_argument( + "--num-finite-logits", + type=int, + help="Finite logits per row; remaining entries are -inf (default: dense)", + ) - # Print configuration - print(f"Batch sizes: {args.batch_sizes}") - print(f"Vocab sizes: {args.vocab_sizes}") - print(f"Warmup iterations: {args.warmup_iters}") - print(f"Benchmark iterations: {args.benchmark_iters}") - print() + parser.add_argument( + "--implementations", + nargs="+", + choices=list(IMPLEMENTATIONS), + default=["triton", "flashinfer-mask"], + help="Implementations to compare (PyTorch is an optional reference)", + ) + + parser.add_argument( + "--check", + action="store_true", + help="Check filtered distributions against PyTorch outside the timed region", + ) + + parser.add_argument( + "--scenarios", + nargs="+", + choices=[ + "topk_whole", + "topk_partial", + "topp_whole", + "topp_partial", + "topk_topp_whole", + "mixed_partial", + ], + help="Run only the selected scenarios (default: all)", + ) + parser.add_argument("--output", type=Path, help="Save incremental JSONL results") + + args = parser.parse_args() + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.open("x").close() + if args.num_finite_logits is not None and not ( + 0 <= args.num_finite_logits <= min(args.vocab_sizes) + ): + parser.error("--num-finite-logits must be between 0 and each vocabulary size") + torch.manual_seed(0) # Check CUDA if not torch.cuda.is_available(): print("ERROR: CUDA is not available. This benchmark requires a GPU.") return - device_name = torch.cuda.get_device_name(0) - print(f"GPU: {device_name}") - print() - # Create configs - configs = create_benchmark_configs( - args.batch_sizes, - args.vocab_sizes, - ) + configs = create_benchmark_configs(args.batch_sizes, args.vocab_sizes) + if args.scenarios is not None: + configs = [c for c in configs if c.name.split("_b")[0] in args.scenarios] # Run benchmarks results = run_benchmark( @@ -461,6 +462,10 @@ def main(): warmup_iters=args.warmup_iters, benchmark_iters=args.benchmark_iters, verbose=not args.quiet, + num_finite_logits=args.num_finite_logits, + implementations=args.implementations, + check=args.check, + output=args.output, ) # Print summary diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 151f641d215d..4aa837e9d73c 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -206,7 +206,10 @@ def _topk_topp_kernel( k_pivots_num = tl.zeros((), dtype=tl.uint32) min_larger = float("inf") num_min_larger = tl.zeros((), dtype=tl.uint32) - if num_outliers > k: + if num_finite_total <= k: + # All finite logits survive top-k; skip the cutoff search. + k_pivot = -float("inf") + elif num_outliers > k: max_range = max_logit min_range = outlier_pivot search_range = tl.cast(num_outliers, tl.int32) From 189bbfb9f432e283e21dcdf2a78357ae5db05ace Mon Sep 17 00:00:00 2001 From: Ilya Boytsov Date: Tue, 13 Oct 2026 00:18:43 +0100 Subject: [PATCH 452/952] Add MoNe experts pruning support (#50415) Signed-off-by: Ilya Boytsov Signed-off-by: Misha Goin Signed-off-by: mgoin Co-authored-by: Ilya Boytsov Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: Claude Opus 5.5 --- docs/design/expert_substitution.md | 121 +++++ pyproject.toml | 2 + tests/kernels/moe/test_expert_substitution.py | 422 ++++++++++++++++++ .../moe/test_unquantized_backend_selection.py | 52 ++- .../model_loader/test_expert_substitution.py | 337 ++++++++++++++ tests/models/test_deepseek_v4_fi_moe_ep.py | 1 + .../model_executor/layers/fused_moe/config.py | 2 + .../layers/fused_moe/expert_substitution.py | 380 ++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 26 +- .../layers/fused_moe/modular_kernel.py | 2 + .../model_loader/tensorizer_loader.py | 2 + vllm/model_executor/model_loader/utils.py | 6 + .../models/expert_substitution.py | 109 +++++ 13 files changed, 1458 insertions(+), 4 deletions(-) create mode 100644 docs/design/expert_substitution.md create mode 100644 tests/kernels/moe/test_expert_substitution.py create mode 100644 tests/model_executor/model_loader/test_expert_substitution.py create mode 100644 vllm/model_executor/layers/fused_moe/expert_substitution.py create mode 100644 vllm/model_executor/models/expert_substitution.py diff --git a/docs/design/expert_substitution.md b/docs/design/expert_substitution.md new file mode 100644 index 000000000000..05e69bab6239 --- /dev/null +++ b/docs/design/expert_substitution.md @@ -0,0 +1,121 @@ +# Expert substitution (MoNE) + +[Mixture of Novices and Experts (MoNE)](https://arxiv.org/abs/2507.00390) +prunes MoE layers by replacing redundant experts with "novices" that return a +constant vector. vLLM loads the retained experts into compact weights and adds +the novices' output separately, so pruned experts use no MLP weight memory and +no MoE backend needs to know about them. + +## Checkpoint format + +The checkpoint keeps the original architecture. `config.json` lists the +substituted experts of each decoder layer: + +```json +{ + "architectures": ["DeepseekV2ForCausalLM"], + "n_routed_experts": 64, + "approximate_experts": { + "1": [1, 6, 8, 12], + "2": [0, 3, 17, 40] + } +} +``` + +For VLMs, `approximate_experts` may live on the text config. Retained experts +keep their logical checkpoint names. Each substituted expert stores one +`[hidden_size]` tensor instead of its MLP: + +```text +model.layers.1.mlp.experts.0.gate_proj.weight +model.layers.1.mlp.experts.1.approx_value +``` + +Alternatively, `compression_config.transform_config.expert_substitution` +declares the same information with explicitly named constant tensors: + +```json +{ + "compression_config": { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": true, + "preserve_router_weights": true, + "renormalize_after_substitution": false + }, + "targets": { + "model.layers.1.mlp.experts": { + "num_logical_experts": 64, + "weight_layout": "compact_retained_experts", + "replacements": { + "1": { + "format": "constant-v1", + "tensors": {"value": "model.layers.1.mlp.expert_replacements.1.value"} + } + } + } + } + } + } + } +} +``` + +Only the values shown for `version`, `router_semantics`, `weight_layout`, and +`format` are supported. Each target path must contain exactly one layer index. +This form is not yet accepted alongside compressed-tensors quantization, whose +`transform_config` schema does not include `expert_substitution`. + +## Semantics + +For a substituted expert `j` with constant `v_j` and router weight `w_j(x)`, +the expert contributes `w_j(x) v_j`. Routing is unchanged: router weights are +neither renormalized nor redistributed to retained experts. + +## Execution + +`FusedMoEFactory` matches each MoE layer to `approximate_experts` by the layer +index in its prefix and, for a match, builds `SubstitutedRoutedExperts`, which: + +1. allocates weights only for retained experts, in ascending logical order; +2. after routing, sums `w_j(x) v_j` over substituted routes in FP32; +3. rewrites substituted routes as zero-weight routes to physical expert 0 and + retained routes as compact physical IDs, preserving the `[num_tokens, top_k]` + contract of every decomposed backend; +4. adds the constant output on one tensor-parallel rank, before the partial + outputs are all-reduced. + +Zero-weight routes may schedule some unnecessary GEMM work; the memory savings +are unaffected. Monolithic backends route internally and are not selected. + +## Weight loading + +Retained experts are loaded through the standard expert mapping by logical ID +and placed in their compact row. Substituted IDs map to no local row, the same +way expert parallelism skips non-local experts. + +Constant tensors are consumed at the model's `load_weights` boundary, so +initial loading and direct weight updates share one path. `approx_value` tensors of +layers that the model does not build (other pipeline stages, MTP layers) are +passed through to the model's loader, which skips them like any other weight +it does not own. Explicitly named constants are always consumed, and dropped +when their layer is built elsewhere. Loading fails if a local substituted +expert has no constant. + +Direct `load_weights` updates may change individual constants. Layerwise +reloads (`reload_weights` in checkpoint format) are rejected, and the model +keeps its previous weights. + +## Supported configurations + +- Models whose MoE layers use `FusedMoEFactory` with the default + `RoutedExperts`, and whose MoE prefix contains exactly one layer index. +- Unquantized experts. +- Tensor and pipeline parallelism, including MTP draft models. +- Any decomposed MoE backend. + +Expert and data parallelism, EPLB, MoE LoRA, fused shared experts, quantized +experts, and routed input/output transforms are rejected at initialization. +Deferred MoE finalize and layerwise weight reloads are rejected when used. diff --git a/pyproject.toml b/pyproject.toml index bc924440fa59..db5a2a3b9c83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,6 +226,8 @@ view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" [tool.typos.default.extend-words] +MoNE = "MoNE" +mone = "mone" Hel = "Hel" wether = "wether" iy = "iy" diff --git a/tests/kernels/moe/test_expert_substitution.py b/tests/kernels/moe/test_expert_substitution.py new file mode 100644 index 000000000000..70720af69ff0 --- /dev/null +++ b/tests/kernels/moe/test_expert_substitution.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( + ProcessGroupInfo, + parallel_launch_with_config, +) +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, + SubstitutedRoutedExperts, + get_expert_substitution_spec, + make_expert_substitution, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( + FlashInferExperts, +) +from vllm.model_executor.models.mixtral import MixtralMoE +from vllm.platforms import current_platform +from vllm.v1.worker.workspace import init_workspace_manager + +MOE_PREFIX = "model.layers.1.mlp.experts" + + +def _model_config(approximate_experts, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts=approximate_experts), + **kwargs, + ) + + +def _expert_substitution(targets: dict[str, dict[int, str]]) -> dict: + """The ``compression_config.transform_config.expert_substitution`` format.""" + return { + "version": 1, + "router_semantics": { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + }, + "targets": { + path: { + "num_logical_experts": 5, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert_id): { + "format": "constant-v1", + "tensors": {"value": name}, + } + for expert_id, name in replacements.items() + }, + } + for path, replacements in targets.items() + }, + } + + +def _compressed_model_config(expert_substitution: dict, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + hf_config=SimpleNamespace( + compression_config={ + "transform_config": {"expert_substitution": expert_substitution} + } + ), + **kwargs, + ) + + +def _substitution(hidden_size: int, dtype=torch.float32): + # Experts 1 and 3 of 5 are substituted; 0, 2, 4 map to physical 0, 1, 2. + return ConstantExpertSubstitution(1, 5, (1, 3), hidden_size, dtype) + + +def test_parse_approximate_experts(): + assert get_expert_substitution_spec(None) is None + assert get_expert_substitution_spec(_model_config(None)) is None + spec = get_expert_substitution_spec(_model_config({"1": [3, 1], "2": [], 5: (0,)})) + assert spec.experts == {1: (1, 3), 5: (0,)} + assert spec.value_names == {} + text_config = SimpleNamespace( + hf_text_config=SimpleNamespace(approximate_experts={"0": [2]}), + hf_config=SimpleNamespace(), + ) + assert get_expert_substitution_spec(text_config).experts == {0: (2,)} + + +@pytest.mark.parametrize( + "approximate_experts", + [[1, 2], {"layer": [1]}, {"1": [1, 1]}, {"1": [-1]}, {"1": ["x"]}], +) +def test_parse_approximate_experts_rejects_invalid_entries(approximate_experts): + with pytest.raises(ValueError, match="approximate_experts"): + get_expert_substitution_spec(_model_config(approximate_experts)) + + +def test_parse_expert_substitution_with_explicit_value_names(): + spec = get_expert_substitution_spec( + _compressed_model_config( + _expert_substitution( + { + MOE_PREFIX: {3: "shared", 1: "layer_1.expert_1"}, + "model.layers.2.mlp.experts": {0: "shared"}, + } + ) + ) + ) + assert spec.experts == {1: (1, 3), 2: (0,)} + assert spec.value_names == { + "shared": ((1, 3), (2, 0)), + "layer_1.expert_1": ((1, 1),), + } + + +@pytest.mark.parametrize( + "mutate", + [ + lambda raw: raw.update(version=2), + lambda raw: raw["router_semantics"].update(renormalize_after_substitution=True), + lambda raw: raw["targets"].update({"mlp.experts": raw["targets"][MOE_PREFIX]}), + lambda raw: raw["targets"][MOE_PREFIX].update(weight_layout="dense"), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].update( + format="linear-v1" + ), + lambda raw: raw["targets"][MOE_PREFIX]["replacements"]["1"].pop("tensors"), + ], +) +def test_parse_expert_substitution_rejects_unsupported_metadata(mutate): + raw = _expert_substitution({MOE_PREFIX: {1: "value"}}) + mutate(raw) + with pytest.raises(ValueError, match="invalid expert_substitution metadata"): + get_expert_substitution_spec(_compressed_model_config(raw)) + + +def test_parse_rejects_both_formats(): + config = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + config.hf_config.approximate_experts = {"1": [1]} + with pytest.raises(ValueError, match="both"): + get_expert_substitution_spec(config) + + +def test_make_expert_substitution_matches_layer_index(): + config = _model_config({"1": [3, 1]}) + substitution = make_expert_substitution(config, MOE_PREFIX, 5, 16, torch.bfloat16) + + assert substitution is not None + assert substitution.layer_idx == 1 + assert substitution.substituted_expert_ids == (1, 3) + assert substitution.num_compute_experts == 3 + assert substitution.values.dtype == torch.bfloat16 + assert [substitution.physical_expert_id(i) for i in range(5)] == [0, -1, 1, -1, 2] + assert ( + make_expert_substitution(config, "model.layers.2.mlp.experts", 5, 16, None) + is None + ) + assert make_expert_substitution(None, MOE_PREFIX, 5, 16, None) is None + explicit = _compressed_model_config(_expert_substitution({MOE_PREFIX: {1: "v"}})) + substitution = make_expert_substitution(explicit, MOE_PREFIX, 5, 16, None) + assert substitution is not None + assert substitution.substituted_expert_ids == (1,) + with pytest.raises(ValueError, match="exactly one layer index"): + make_expert_substitution(config, "mlp.experts", 5, 16, None) + with pytest.raises(ValueError, match="within"): + make_expert_substitution(config, MOE_PREFIX, 3, 16, None) + with pytest.raises(ValueError, match="retain at least one expert"): + make_expert_substitution(_model_config({"1": [0, 1]}), MOE_PREFIX, 2, 16, None) + + +def test_constant_substitution_loads_values(): + substitution = _substitution(2) + substitution.weight_loader(substitution.values, torch.tensor([1.0, 2.0]), 3) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + substitution.validate_loaded_values("layer") + substitution.weight_loader(substitution.values, torch.tensor([3.0, 4.0]), 1) + substitution.validate_loaded_values("layer") + torch.testing.assert_close( + substitution.values, torch.tensor([[3.0, 4.0], [1.0, 2.0]]) + ) + with pytest.raises(ValueError, match="not listed"): + substitution.weight_loader(substitution.values, torch.ones(2), 2) + with pytest.raises(ValueError, match="has shape"): + substitution.weight_loader(substitution.values, torch.ones(3), 1) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 3, 7]) +@pytest.mark.parametrize("hidden_size", [2, 4]) +def test_constant_substitution_transforms_routes_and_computes_side_output( + num_tokens, hidden_size +): + substitution = _substitution(2) + substitution.values.data.copy_(torch.tensor([[10.0, 20.0], [30.0, 40.0]])) + + rows = torch.arange(num_tokens) % 3 + topk_ids = torch.tensor([[1, 2], [4, 3], [99, -1]])[rows] + topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4], [0.5, 0.5]])[rows] + constant_output = substitution.transform_routes( + torch.zeros(num_tokens, hidden_size), topk_weights, topk_ids + ) + + torch.testing.assert_close( + topk_weights, torch.tensor([[0.0, 0.75], [0.6, 0.0], [0.0, 0.0]])[rows] + ) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1], [2, 0], [0, 0]])[rows]) + expected = torch.zeros(num_tokens, hidden_size) + expected[:, :2] = torch.tensor([[2.5, 5.0], [12.0, 16.0], [0.0, 0.0]])[rows] + torch.testing.assert_close(constant_output, expected) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="requires a CUDA-like platform", + ), + ), + ], +) +def test_constant_substitution_preserves_router_weight_precision(dtype, device): + substitution = _substitution(1, dtype).to(device) + substitution.values.data.copy_( + torch.tensor([[60000.0], [-60000.0]], dtype=dtype, device=device) + ) + topk_weights = torch.tensor([[0.5001, 0.4999]], dtype=torch.float32, device=device) + topk_ids = torch.tensor([[1, 3]], device=device) + # Nearly cancelling constants must retain the difference between FP32 weights. + expected = (topk_weights.double() @ substitution.values.double()).to(dtype) + + actual = substitution.transform_routes( + torch.zeros(1, 1, dtype=dtype, device=device), topk_weights, topk_ids + ) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _bare_substituted_routed_experts(**attrs) -> SubstitutedRoutedExperts: + layer = SubstitutedRoutedExperts.__new__(SubstitutedRoutedExperts) + torch.nn.Module.__init__(layer) + layer.expert_substitution = _substitution(2) + for name, value in attrs.items(): + setattr(layer, name, value) + return layer + + +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("padded_input", [False, True]) +def test_substituted_routed_experts_adds_constant_on_one_tp_rank(tp_rank, padded_input): + def apply(*, layer, x, topk_weights, topk_ids, **kwargs): + torch.testing.assert_close(topk_weights, torch.tensor([[0.0, 0.75]])) + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1]])) + return torch.ones(x.shape[0], 2) + + layer = _bare_substituted_routed_experts( + moe_config=SimpleNamespace(tp_rank=tp_rank), + quant_method=SimpleNamespace(is_monolithic=False, apply=apply), + ) + layer.expert_substitution.values.data.copy_( + torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + ) + actual = layer.forward_modular( + torch.zeros(1, 4 if padded_input else 2), + torch.tensor([[0.25, 0.75]]), + torch.tensor([[1, 2]]), + ) + expected = torch.ones(1, 2) + if tp_rank == 0: + expected += torch.tensor([[2.5, 5.0]]) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +def test_standard_fused_moe_model_discovers_substitution(dist_init): + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.kernel_config.moe_backend = "triton" + + with set_current_vllm_config(vllm_config): + moe = MixtralMoE( + num_experts=4, + top_k=2, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.block_sparse_moe", + ) + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts, SubstitutedRoutedExperts) + assert routed_experts.expert_substitution.substituted_expert_ids == (1,) + assert routed_experts.w13_weight.shape[0] == 3 + assert moe.experts.moe_config.require_decomposed_backend + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like platform" +) +@pytest.mark.parametrize( + "moe_backend", + [ + pytest.param("triton", id="triton"), + pytest.param( + "flashinfer_cutlass", + id="flashinfer-cutlass", + marks=pytest.mark.skipif( + not FlashInferExperts._supports_current_device(), + reason="FlashInfer CUTLASS is unavailable on this platform", + ), + ), + ], +) +def test_fused_moe_combines_compact_experts_and_substitution(dist_init, moe_backend): + hidden_size = 256 + vllm_config = VllmConfig() + vllm_config.model_config = _model_config( + {"0": [1, 3]}, dtype=torch.bfloat16, is_moe=True + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = moe_backend + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + init_workspace_manager(torch.accelerator.current_device_index()) + layer = FusedMoEFactory( + num_experts=4, + top_k=2, + hidden_size=hidden_size, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=False, + ).cuda() + substitution = layer.routed_experts.expert_substitution + assert layer.routed_experts.w13_weight.shape[0] == 2 + with torch.no_grad(): + layer.routed_experts.w13_weight.normal_(0, 0.01) + layer.routed_experts.w2_weight.normal_(0, 0.01) + substitution.values.normal_(0, 0.01) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.randn(8, hidden_size, dtype=torch.bfloat16, device="cuda") + router_logits = torch.randn(8, 4, dtype=torch.float32, device="cuda") + topk_weights, topk_ids = layer.router.select_experts( + hidden_states.clone(), router_logits + ) + expected = substitution.transform_routes(hidden_states, topk_weights, topk_ids) + expected += layer._quant_method.apply( + layer=layer.routed_experts, + x=hidden_states.clone(), + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + + get_forward_context().all_moe_layers = None + actual = layer(hidden_states.clone(), router_logits) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _tp_substitution_worker( + process_group: ProcessGroupInfo, + vllm_config: VllmConfig, + _cpu_group, +) -> None: + with set_forward_context(None, vllm_config): + init_workspace_manager(process_group.local_rank) + layer = FusedMoEFactory( + num_experts=4, + top_k=1, + hidden_size=256, + intermediate_size=512, + params_dtype=torch.bfloat16, + prefix="model.layers.0.mlp.experts", + renormalize=True, + ) + with torch.no_grad(): + layer.routed_experts.w13_weight.zero_() + layer.routed_experts.w2_weight.zero_() + layer.routed_experts.expert_substitution.values.fill_(0.25) + layer._quant_method.process_weights_after_loading(layer.routed_experts) + + hidden_states = torch.zeros(2, 256, dtype=torch.bfloat16) + router_logits = torch.tensor( + [[-100.0, 100.0, -100.0, -100.0]] * 2, + dtype=torch.float32, + ) + get_forward_context().all_moe_layers = None + actual = layer(hidden_states, router_logits) + torch.testing.assert_close( + actual, torch.full_like(actual, 0.25), atol=0, rtol=0 + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike() or current_platform.device_count() < 2, + reason="requires two CUDA-like devices", +) +def test_constant_substitution_is_added_once_with_tp2(): + vllm_config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=2)) + vllm_config.model_config = _model_config( + {"0": [1]}, + dtype=torch.bfloat16, + is_moe=True, + sleep_mode_offload_cudagraph=False, + ) + vllm_config.compilation_config.static_forward_context = {} + vllm_config.kernel_config.moe_backend = "triton" + + parallel_launch_with_config(2, _tp_substitution_worker, vllm_config, None) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index 25e6b7dec618..b64e909fbd45 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -136,7 +136,13 @@ def test_arm_cpu_unquantized_kernel_selection( ], ) @patch( - "vllm.utils.flashinfer.has_flashinfer", + "vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe." + "has_flashinfer_cutlass_fused_moe", + return_value=False, +) +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe." + "has_flashinfer_trtllm_fused_moe", return_value=False, ) @patch( @@ -145,13 +151,18 @@ def test_arm_cpu_unquantized_kernel_selection( ) def test_select_default_backend_by_platform( mock_aiter_enabled, - mock_has_flashinfer, + mock_has_flashinfer_trtllm, + mock_has_flashinfer_cutlass, monkeypatch, platform_method, expected_backend, ): """Test default backend selection per platform with all optional - accelerators (FlashInfer, AITER) disabled.""" + accelerators (FlashInfer, AITER) disabled. + + Patch the feature probes directly because their cached results can outlive + changes to has_flashinfer(). + """ with patch( "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform" ) as mock_platform: @@ -407,6 +418,41 @@ def test_select_cuda_flashinfer_trtllm_modular_for_standard_all2all( assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular" +@patch( + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", + return_value=True, +) +@pytest.mark.parametrize("require_decomposed_backend", [False, True]) +def test_require_decomposed_backend_rejects_monolithic_kernels( + mock_supports_current_device, + require_decomposed_backend, +): + """Expert substitution adjusts routes, so the kernel must not route itself.""" + from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + TrtLlmBf16ExpertsMonolithic, + ) + from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEActivationFormat, + ) + + moe_config = make_dummy_moe_config() + moe_config.routing_method = RoutingMethodType.Renormalize + moe_config.require_decomposed_backend = require_decomposed_backend + + supported, reason = TrtLlmBf16ExpertsMonolithic.is_supported_config( + TrtLlmBf16ExpertsMonolithic, + moe_config, + None, + None, + FusedMoEActivationFormat.Standard, + ) + + assert supported != require_decomposed_backend + assert (reason is not None and "routes adjusted" in reason) == ( + require_decomposed_backend + ) + + @patch( "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device", return_value=True, diff --git a/tests/model_executor/model_loader/test_expert_substitution.py b/tests/model_executor/model_loader/test_expert_substitution.py new file mode 100644 index 000000000000..926720c68629 --- /dev/null +++ b/tests/model_executor/model_loader/test_expert_substitution.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file +from transformers import DeepseekV2Config, DeepseekV2ForCausalLM + +from vllm.config import ( + AttentionConfig, + DeviceConfig, + KernelConfig, + LoadConfig, + ModelConfig, + VllmConfig, +) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + ConstantExpertSubstitution, +) +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.models.expert_substitution import as_expert_substitution_model +from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace(approximate_experts={"0": [1, 2], "1": [1, 2]}) +) + + +def _compression_config(value_names: dict[int, dict[int, str]]) -> dict: + """Explicitly named constants given as ``{layer: {expert: tensor name}}``.""" + targets = { + f"model.layers.{layer}.mlp.experts": { + "num_logical_experts": 3, + "weight_layout": "compact_retained_experts", + "replacements": { + str(expert): {"format": "constant-v1", "tensors": {"value": name}} + for expert, name in names.items() + }, + } + for layer, names in value_names.items() + } + router_semantics = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, + } + return { + "transform_config": { + "expert_substitution": { + "version": 1, + "router_semantics": router_semantics, + "targets": targets, + } + } + } + + +# Expert 1 of both layers shares one explicitly named constant. +EXPLICIT_MODEL_CONFIG = SimpleNamespace( + hf_config=SimpleNamespace( + compression_config=_compression_config( + { + layer: {1: "constants.shared", 2: f"constants.layer_{layer}"} + for layer in (0, 1) + } + ) + ) +) + + +def _approx_value(layer: int, expert: int) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.approx_value" + + +def _cpu_model(local_layers=(0, 1), tracked=True, model_config=MODEL_CONFIG): + """Two decoder layers; the others are placeholders as on another PP stage.""" + + class Model(torch.nn.Module): + def __init__(self, prefix=""): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(2), requires_grad=False) + self.model = torch.nn.Module() + self.model.layers = torch.nn.ModuleList() + for index in range(2): + if index not in local_layers: + self.model.layers.append(PPMissingLayer()) + continue + layer = torch.nn.Module() + layer.mlp = torch.nn.Module() + layer.mlp.experts = ConstantExpertSubstitution( + index, 3, (1, 2), 2, torch.float32 + ) + self.model.layers.append(layer) + + def load_weights(self, weights): + loaded = AutoWeightsLoader(self).load_weights(weights) + return loaded if tracked else None + + return as_expert_substitution_model(Model, model_config)() + + +@pytest.mark.parametrize("local_layers", [(0, 1), (0,), (1,), ()]) +def test_loads_local_approx_values_and_passes_through_others(local_layers): + """Other PP stages and MTP drafts own none or some of the configured layers.""" + model = _cpu_model(local_layers) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + *( + (_approx_value(layer, expert), torch.full((2,), 10.0 * layer + expert)) + for layer in range(2) + for expert in (1, 2) + ), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0, 2.0]]) + 10.0 * layer, + ) + + +@pytest.mark.parametrize("local_layers", [(0, 1), (1,), ()]) +def test_loads_explicitly_named_values(local_layers): + """Named values are consumed even when their layer is built elsewhere.""" + model = _cpu_model(local_layers, model_config=EXPLICIT_MODEL_CONFIG) + loaded = model.load_weights( + [ + ("weight", torch.tensor([1.0, 2.0])), + ("constants.shared", torch.full((2,), 1.0)), + ("constants.layer_0", torch.full((2,), 2.0)), + ("constants.layer_1", torch.full((2,), 3.0)), + ] + ) + model.process_weights_after_loading() + + assert loaded == {"weight"} | { + f"model.layers.{layer}.mlp.experts.values" for layer in local_layers + } + for layer in local_layers: + torch.testing.assert_close( + model.model.layers[layer].mlp.experts.values, + torch.tensor([[1.0, 1.0], [2.0 + layer, 2.0 + layer]]), + ) + + +def test_direct_reload_updates_values_incrementally(): + model = _cpu_model(local_layers=(0,), tracked=False) + values = model.model.layers[0].mlp.experts.values + model.load_weights([(_approx_value(0, 2), torch.ones(2))]) + with pytest.raises(ValueError, match="missing approx_value .* \\[1\\]"): + model.process_weights_after_loading() + + assert model.load_weights([(_approx_value(0, 1), torch.zeros(2))]) is None + model.process_weights_after_loading() + model.load_weights([(_approx_value(0, 2), torch.full((2,), 3.0))]) + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [3.0, 3.0]])) + + with pytest.raises(ValueError, match="not listed"): + model.load_weights([(_approx_value(0, 0), torch.ones(2))]) + + +def test_layerwise_reload_is_rejected_and_keeps_weights(): + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + record_metadata_for_reloading, + ) + + model = _cpu_model(local_layers=(0,)) + record_metadata_for_reloading(model) + loaded = [ + ("weight", torch.tensor([3.0, 4.0])), + (_approx_value(0, 1), torch.zeros(2)), + (_approx_value(0, 2), torch.ones(2)), + ] + model.load_weights(loaded) + model.process_weights_after_loading() + substitution = model.model.layers[0].mlp.experts + values = substitution.values + + initialize_layerwise_reload(model) + with pytest.raises(NotImplementedError, match="layerwise weight reload"): + model.load_weights([(name, weight + 1) for name, weight in loaded]) + + assert substitution.values is values + torch.testing.assert_close(values, torch.tensor([[0.0, 0.0], [1.0, 1.0]])) + torch.testing.assert_close(model.weight, torch.tensor([3.0, 4.0])) + assert substitution.expert_substitution_routes.tolist() == [0, -1, -2] + + +@pytest.mark.parametrize("missing_value", [False, True]) +def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_value): + """Raw Tensorizer loading must reject an unloaded local constant row.""" + from vllm.model_executor.model_loader import tensorizer_loader + + model = _cpu_model(local_layers=(0,)) + weights = [(_approx_value(0, 1), torch.zeros(2))] + if not missing_value: + weights.append((_approx_value(0, 2), torch.ones(2))) + finalize = Mock(wraps=model.process_weights_after_loading) + monkeypatch.setattr(model, "process_weights_after_loading", finalize) + monkeypatch.setattr(tensorizer_loader, "initialize_model", lambda **kwargs: model) + loader = object.__new__(tensorizer_loader.TensorizerLoader) + monkeypatch.setattr(loader, "_get_weights_iterator", lambda: iter(weights)) + config = SimpleNamespace( + model_config=SimpleNamespace(dtype=torch.float32), + device_config=SimpleNamespace(device="cpu"), + ) + + if missing_value: + with pytest.raises(ValueError, match=r"missing approx_value .* \[2\]"): + loader._load_model_serialized_cpu(config) + else: + assert loader._load_model_serialized_cpu(config) is model + torch.testing.assert_close( + model.model.layers[0].mlp.experts.values, + torch.tensor([[0.0, 0.0], [1.0, 1.0]]), + ) + finalize.assert_called_once_with() + + +def _write_tiny_deepseek_mone_checkpoint( + model_dir: Path, explicit_names: bool +) -> tuple[dict[str, torch.Tensor], str]: + config = DeepseekV2Config( + architectures=["DeepseekV2ForCausalLM"], + vocab_size=32, + hidden_size=16, + intermediate_size=32, + moe_intermediate_size=8, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=32, + first_k_dense_replace=0, + n_routed_experts=3, + n_shared_experts=1, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + norm_topk_prob=False, + q_lora_rank=None, + kv_lora_rank=4, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + tie_word_embeddings=False, + dtype="float16", + ) + experts = "model.layers.0.mlp.experts" + if explicit_names: + value_name = "model.layers.0.mlp.expert_replacements.1.value" + config.compression_config = _compression_config({0: {1: value_name}}) + else: + value_name = f"{experts}.1.approx_value" + config.approximate_experts = {"0": [1]} + + hf_weights = DeepseekV2ForCausalLM(config).to(torch.float16).state_dict() + gate_up = hf_weights.pop(f"{experts}.gate_up_proj") + down = hf_weights.pop(f"{experts}.down_proj") + weights = {name: value.contiguous() for name, value in hf_weights.items()} + for expert in (0, 2): + gate, up = gate_up[expert].chunk(2, dim=0) + weights[f"{experts}.{expert}.gate_proj.weight"] = gate.contiguous() + weights[f"{experts}.{expert}.up_proj.weight"] = up.contiguous() + weights[f"{experts}.{expert}.down_proj.weight"] = down[expert].contiguous() + weights[value_name] = torch.arange(16, dtype=torch.float16) + config.save_pretrained(model_dir) + save_file(weights, model_dir / "model.safetensors") + return weights, value_name + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like device" +) +@pytest.mark.usefixtures("dist_init", "workspace_init") +@pytest.mark.parametrize( + "explicit_names", [False, True], ids=["approximate_experts", "expert_substitution"] +) +def test_loads_mone_checkpoint_into_compact_experts( + tmp_path: Path, explicit_names: bool +): + weights, value_name = _write_tiny_deepseek_mone_checkpoint(tmp_path, explicit_names) + model_config = ModelConfig( + model=str(tmp_path), + tokenizer=str(tmp_path), + skip_tokenizer_init=True, + dtype="float16", + max_model_len=32, + enforce_eager=True, + ) + load_config = LoadConfig(load_format="safetensors", use_tqdm_on_load=False) + vllm_config = VllmConfig( + model_config=model_config, + device_config=DeviceConfig(device="cuda"), + load_config=load_config, + attention_config=AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA), + kernel_config=KernelConfig(moe_backend="triton"), + ) + model = get_model_loader(load_config).load_model(vllm_config, model_config) + + experts = "model.layers.0.mlp.experts" + routed_experts = model.model.layers[0].mlp.experts.routed_experts + # Logical expert 2 is stored in physical row 1. + torch.testing.assert_close( + routed_experts.w13_weight[1].cpu(), + torch.cat( + [ + weights[f"{experts}.2.gate_proj.weight"], + weights[f"{experts}.2.up_proj.weight"], + ] + ), + ) + torch.testing.assert_close( + routed_experts.w2_weight[1].cpu(), weights[f"{experts}.2.down_proj.weight"] + ) + substitution = routed_experts.expert_substitution + torch.testing.assert_close( + substitution.values.cpu(), weights[value_name].unsqueeze(0) + ) + + updated = weights[value_name] + 1 + assert model.load_weights([(value_name, updated)]) == { + f"{experts}.routed_experts.expert_substitution.values" + } + torch.testing.assert_close(substitution.values.cpu(), updated.unsqueeze(0)) diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 7687a6786f7e..d98bb26a2a3b 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -142,6 +142,7 @@ def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace: is_act_and_mul=True, activation=MoEActivation.SILU, has_hash_routing=False, + require_decomposed_backend=False, routing_method=RoutingMethodType.DeepseekV4, router_logits_dtype=torch.float32, hidden_dim=256, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index f20af8a6f826..e79700ae44d9 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,8 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + # Routes are adjusted between routing and expert execution. + require_decomposed_backend: bool = False shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate diff --git a/vllm/model_executor/layers/fused_moe/expert_substitution.py b/vllm/model_executor/layers/fused_moe/expert_substitution.py new file mode 100644 index 000000000000..2483200b0150 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/expert_substitution.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Constant expert substitution for MoNE-pruned checkpoints. + +MoNE (https://arxiv.org/abs/2507.00390) replaces redundant experts with +"novices" that return a constant vector. Checkpoints list the replaced experts +per decoder layer either in ``config.approximate_experts``, with constants named +``..approx_value``, or in +``compression_config.transform_config.expert_substitution``, with explicitly +named constants. Retained experts keep their logical checkpoint names and are +loaded into compact physical rows. +""" + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts + +APPROX_VALUE_SUFFIX = ".approx_value" + + +def decoder_layer_index(path: str) -> int | None: + """Return the only integer component of a module path, if unique.""" + indices = [int(part) for part in path.split(".") if part.isdigit()] + return indices[0] if len(indices) == 1 else None + + +@dataclass(frozen=True) +class ExpertSubstitutionSpec: + # Sorted substituted expert IDs per decoder layer. + experts: dict[int, tuple[int, ...]] + # Explicitly named constant tensors and the (layer, expert) pairs they + # supply. Without an entry, constants use the ``approx_value`` naming. + value_names: dict[str, tuple[tuple[int, int], ...]] = field(default_factory=dict) + + +# The only router semantics that the constant-v1 format supports. +_ROUTER_SEMANTICS = { + "preserve_logical_expert_ids": True, + "preserve_router_weights": True, + "renormalize_after_substitution": False, +} + + +def _expert_ids(layer_idx: int, expert_ids: Any) -> tuple[int, ...]: + ids = tuple(sorted(int(expert_id) for expert_id in expert_ids or ())) + if len(set(ids)) != len(ids) or any(expert_id < 0 for expert_id in ids): + raise ValueError( + f"substituted experts of layer {layer_idx} must be unique, " + f"non-negative expert IDs, got {list(ids)}" + ) + return ids + + +def _parse_approximate_experts(raw: Any) -> ExpertSubstitutionSpec: + if not isinstance(raw, Mapping): + raise ValueError("approximate_experts must map layer indices to expert IDs") + experts: dict[int, tuple[int, ...]] = {} + for layer, expert_ids in raw.items(): + try: + layer_idx = int(layer) + experts[layer_idx] = _expert_ids(layer_idx, expert_ids) + except (TypeError, ValueError) as e: + raise ValueError( + f"approximate_experts has an invalid entry for layer {layer!r}" + ) from e + return ExpertSubstitutionSpec({k: v for k, v in experts.items() if v}) + + +def _parse_expert_substitution(raw: Any) -> ExpertSubstitutionSpec: + experts: dict[int, tuple[int, ...]] = {} + value_names: dict[str, list[tuple[int, int]]] = {} + try: + if raw["version"] != 1 or raw["router_semantics"] != _ROUTER_SEMANTICS: + raise ValueError("only version 1 with unchanged routing is supported") + for path, target in raw["targets"].items(): + layer_idx = decoder_layer_index(path) + if layer_idx is None or layer_idx in experts: + raise ValueError(f"target {path!r} must name one decoder layer") + if target["weight_layout"] != "compact_retained_experts": + raise ValueError(f"unsupported weight layout for {path!r}") + for expert_id, replacement in target["replacements"].items(): + if replacement["format"] != "constant-v1": + raise ValueError(f"unsupported replacement format for {path!r}") + value_names.setdefault(replacement["tensors"]["value"], []).append( + (layer_idx, int(expert_id)) + ) + experts[layer_idx] = _expert_ids(layer_idx, target["replacements"]) + except (AttributeError, KeyError, TypeError, ValueError) as e: + raise ValueError(f"invalid expert_substitution metadata: {e}") from e + return ExpertSubstitutionSpec( + experts, {name: tuple(pairs) for name, pairs in value_names.items()} + ) + + +def get_expert_substitution_spec(model_config: Any) -> ExpertSubstitutionSpec | None: + """Parse the substituted experts declared in a model's HF config.""" + hf_configs = [ + getattr(model_config, name, None) for name in ("hf_text_config", "hf_config") + ] + approximate = None + for config in hf_configs: + approximate = getattr(config, "approximate_experts", None) + if approximate is not None: + break + compression_config = getattr(hf_configs[1], "compression_config", None) + transform_config = ( + compression_config.get("transform_config") + if isinstance(compression_config, Mapping) + else None + ) + declared = ( + transform_config.get("expert_substitution") + if isinstance(transform_config, Mapping) + else None + ) + if approximate is not None and declared is not None: + raise ValueError( + "config declares both approximate_experts and " + "compression_config.transform_config.expert_substitution" + ) + if approximate is not None: + spec = _parse_approximate_experts(approximate) + elif declared is not None: + spec = _parse_expert_substitution(declared) + else: + return None + return spec if spec.experts else None + + +class ConstantExpertSubstitution(nn.Module): + """Constant outputs for the substituted experts of one MoE layer. + + Router outputs keep logical expert IDs. ``transform_routes`` maps retained + experts to compact physical rows and turns substituted routes into + zero-weight routes to physical expert 0, so any decomposed MoE backend can + run them unchanged. + """ + + def __init__( + self, + layer_idx: int, + num_logical_experts: int, + substituted_expert_ids: tuple[int, ...], + hidden_size: int, + params_dtype: torch.dtype, + ) -> None: + super().__init__() + substituted = set(substituted_expert_ids) + if not substituted <= set(range(num_logical_experts)): + raise ValueError( + f"approximate_experts for layer {layer_idx} must be within " + f"[0, {num_logical_experts}), got {sorted(substituted)}" + ) + if len(substituted) == num_logical_experts: + raise ValueError( + f"approximate_experts for layer {layer_idx} must retain at least " + "one expert" + ) + self.layer_idx = layer_idx + self.substituted_expert_ids = tuple(sorted(substituted)) + retained = [i for i in range(num_logical_experts) if i not in substituted] + self.num_compute_experts = len(retained) + + # Retained experts map to their physical row r >= 0, substituted + # experts to -1 - (their row in `values`). + routes = [0] * num_logical_experts + for row, expert_id in enumerate(retained): + routes[expert_id] = row + for row, expert_id in enumerate(self.substituted_expert_ids): + routes[expert_id] = -1 - row + self._routes = tuple(routes) + self.register_buffer( + "expert_substitution_routes", + torch.tensor(routes, dtype=torch.int32), + persistent=False, + ) + + # NaN marks values that the checkpoint has not provided yet. + self.values = nn.Parameter( + torch.full((len(substituted), hidden_size), torch.nan, dtype=params_dtype), + requires_grad=False, + ) + self.values.weight_loader = self.weight_loader + + @property + def num_logical_experts(self) -> int: + return len(self._routes) + + def physical_expert_id(self, expert_id: int) -> int: + """Physical row of a retained expert, or -1 for a substituted one.""" + return max(self._routes[expert_id], -1) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, expert_id: int + ) -> None: + if not 0 <= expert_id < self.num_logical_experts or ( + self._routes[expert_id] >= 0 + ): + raise ValueError( + f"expert {expert_id} of layer {self.layer_idx} is not listed in " + "approximate_experts" + ) + row = param.data[-1 - self._routes[expert_id]] + if loaded_weight.shape != row.shape: + raise ValueError( + f"approx_value for expert {expert_id} of layer {self.layer_idx} has " + f"shape {tuple(loaded_weight.shape)}, expected {tuple(row.shape)}" + ) + row.copy_(loaded_weight) + + def validate_loaded_values(self, prefix: str) -> None: + missing_rows = torch.isnan(self.values).any(dim=1).nonzero().flatten().tolist() + if missing_rows: + missing = [self.substituted_expert_ids[row] for row in missing_rows] + raise ValueError( + f"{prefix} is missing approx_value for {len(missing)} " + f"expert(s), first missing expert IDs: {missing[:8]}" + ) + + def transform_routes( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Remap routes in place and return the constant experts' output.""" + valid = (topk_ids >= 0) & (topk_ids < self.num_logical_experts) + routes = self.expert_substitution_routes[ + torch.where(valid, topk_ids, torch.zeros_like(topk_ids)).long() + ].long() + substituted = valid & (routes < 0) + + # Sum each token's router weights per substituted expert in FP32; other + # routes land in a trailing column that is dropped. + num_values, value_size = self.values.shape + scales = topk_weights.new_zeros( + (topk_ids.shape[0], num_values + 1), dtype=torch.float32 + ) + scales.scatter_add_( + 1, + torch.where(substituted, -1 - routes, num_values), + topk_weights.float(), + ) + output = torch.zeros_like(hidden_states) + output[..., :value_size] = scales[:, :num_values] @ self.values.float() + + topk_weights.masked_fill_(~valid | substituted, 0.0) + topk_ids.copy_(routes.clamp_min(0).to(topk_ids.dtype)) + return output + + +def make_expert_substitution( + model_config: Any, + prefix: str, + num_experts: int, + hidden_size: int, + params_dtype: torch.dtype | None, +) -> ConstantExpertSubstitution | None: + spec = get_expert_substitution_spec(model_config) + if spec is None: + return None + layer_idx = decoder_layer_index(prefix) + if layer_idx is None: + raise ValueError( + f"cannot match MoE layer {prefix!r} to substituted experts: expected " + "exactly one layer index in its prefix" + ) + expert_ids = spec.experts.get(layer_idx) + if expert_ids is None: + return None + return ConstantExpertSubstitution( + layer_idx, + num_experts, + expert_ids, + hidden_size, + params_dtype or torch.get_default_dtype(), + ) + + +class SubstitutedRoutedExperts(RoutedExperts): + """Routed experts that hold weights only for retained experts.""" + + def __init__( + self, *args: Any, expert_substitution: ConstantExpertSubstitution, **kwargs: Any + ): + super().__init__(*args, **kwargs) + self.expert_substitution = expert_substitution + moe_config = self.moe_config + parallel_config = moe_config.moe_parallel_config + unsupported = [ + name + for name, enabled in ( + ( + "quantized MoE weights", + not isinstance(self.quant_method, UnquantizedFusedMoEMethod), + ), + ("monolithic MoE backends", self.quant_method.is_monolithic), + ("MoE LoRA", moe_config.is_lora_enabled), + ( + "EPLB", + parallel_config.enable_eplb + or moe_config.num_experts + != expert_substitution.num_compute_experts, + ), + ( + "fused shared experts", + self.expert_map_manager.num_fused_shared_experts > 0, + ), + ("expert parallelism", parallel_config.use_ep), + ("data parallelism", parallel_config.dp_size > 1), + ("prefill context parallelism", parallel_config.pcp_size > 1), + ("sequence parallelism", parallel_config.is_sequence_parallel), + ("router weights on expert inputs", self.apply_router_weight_on_input), + ) + if enabled + ] + if unsupported: + raise NotImplementedError( + "expert substitution does not support: " + ", ".join(unsupported) + ) + + def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int: + # Checkpoints name retained experts by their logical ID. + physical_id = self.expert_substitution.physical_expert_id(expert_id) + if physical_id < 0: + return -1 + return super()._map_global_expert_id_to_local_expert_id(physical_id) + + def get_expert_mapping( + self, + ckpt_gate_proj_name: str | None = None, + ckpt_down_proj_name: str | None = None, + ckpt_up_proj_name: str | None = None, + include_fused: bool = False, + ) -> list[tuple[str, str, int, str]]: + # Fused checkpoint tensors would cover every logical expert at once. + return self.build_expert_params_mapping( + ckpt_gate_proj_name or self.ckpt_gate_proj_name, + ckpt_down_proj_name or self.ckpt_down_proj_name, + ckpt_up_proj_name or self.ckpt_up_proj_name, + num_experts=self.expert_substitution.num_logical_experts, + routed_experts_prefix="", + lora_base_layer_prefix=self.lora_base_layer_prefix, + ) + + def forward_modular( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: "SharedExperts | None" = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + constant_output = self.expert_substitution.transform_routes( + x, topk_weights, topk_ids + ) + output = super().forward_modular( + x, topk_weights, topk_ids, shared_experts, shared_experts_input + ) + if isinstance(output, UnfinalizedMoEOutput): + raise NotImplementedError( + "expert substitution does not support deferred MoE finalize" + ) + # The constant output is replicated across TP ranks, whose partial + # outputs are all-reduced later, so add it on one rank only. + if self.moe_config.tp_rank == 0: + output.add_(constant_output[..., : output.shape[-1]]) + return output diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 6b7c64bdc429..839eba5d788a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -23,6 +23,10 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import ( ExpertMapManager, ) +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + SubstitutedRoutedExperts, + make_expert_substitution, +) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, @@ -253,6 +257,12 @@ def FusedMoEFactory( fuse_shared_experts, ) ) + expert_substitution = make_expert_substitution( + vllm_config.model_config, layer_name, num_experts, hidden_size, params_dtype + ) + if expert_substitution is not None: + # Route over all logical experts but allocate only the retained ones. + global_num_experts -= len(expert_substitution.substituted_expert_ids) # Initialize EPLB manager (or None?) eplb_state: EplbLayerState | None = None @@ -367,6 +377,7 @@ def FusedMoEFactory( device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal has_hash_routing=hash_indices_table is not None, + require_decomposed_backend=expert_substitution is not None, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, @@ -380,7 +391,20 @@ def FusedMoEFactory( # Create RoutedExperts instance BEFORE create_weights() # This will hold all expert weight parameters - if routed_experts_cls is None: + if expert_substitution is not None: + if routed_experts_cls not in (None, RoutedExperts) or ( + routed_input_transform is not None or routed_output_transform is not None + ): + raise NotImplementedError( + "expert substitution requires the default RoutedExperts and no " + "routed input or output transforms" + ) + routed_experts_cls = SubstitutedRoutedExperts + routed_experts_args = { + **(routed_experts_args or {}), + "expert_substitution": expert_substitution, + } + elif routed_experts_cls is None: routed_experts_cls = RoutedExperts assert params_dtype is not None diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7a4cbd423746..d9a75c55ec40 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -569,6 +569,8 @@ def _make_reason(reason: str) -> str: ) elif moe_config.has_hash_routing and cls.is_monolithic(): return False, _make_reason("hash routing") + elif moe_config.require_decomposed_backend and cls.is_monolithic(): + return False, _make_reason("routes adjusted after routing") elif not cls._supports_routing_method( moe_config.routing_method, weight_key, activation_key ): diff --git a/vllm/model_executor/model_loader/tensorizer_loader.py b/vllm/model_executor/model_loader/tensorizer_loader.py index 338f9eac072a..ae74beb7ed5f 100644 --- a/vllm/model_executor/model_loader/tensorizer_loader.py +++ b/vllm/model_executor/model_loader/tensorizer_loader.py @@ -84,6 +84,8 @@ def _load_model_serialized_cpu( model = initialize_model(vllm_config=vllm_config, prefix=prefix) model.load_weights(self._get_weights_iterator()) + if hasattr(model, "process_weights_after_loading"): + model.process_weights_after_loading() return model.eval() def download_model(self, model_config: ModelConfig) -> None: diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 65d1ff2808d7..8a236941002b 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -84,6 +84,12 @@ def initialize_model( if model_class is None: model_class, _ = get_model_architecture(model_config) + from vllm.model_executor.models.expert_substitution import ( + as_expert_substitution_model, + ) + + model_class = as_expert_substitution_model(model_class, model_config) + if vllm_config.quant_config is not None: configure_quant_config(vllm_config.quant_config, model_class) diff --git a/vllm/model_executor/models/expert_substitution.py b/vllm/model_executor/models/expert_substitution.py new file mode 100644 index 000000000000..0799bca8b7e0 --- /dev/null +++ b/vllm/model_executor/models/expert_substitution.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from functools import wraps +from typing import Any, TypeVar + +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.expert_substitution import ( + APPROX_VALUE_SUFFIX, + ConstantExpertSubstitution, + decoder_layer_index, + get_expert_substitution_spec, +) +from vllm.model_executor.model_loader.reload import finalize_layerwise_reload + +_T = TypeVar("_T", bound=nn.Module) + + +def _collect_expert_substitutions( + model: nn.Module, +) -> dict[int, tuple[str, ConstantExpertSubstitution]]: + substitutions: dict[int, tuple[str, ConstantExpertSubstitution]] = {} + for name, module in model.named_modules(): + if not isinstance(module, ConstantExpertSubstitution): + continue + if module.layer_idx in substitutions: + raise ValueError( + f"layer {module.layer_idx} has more than one MoE with substituted " + "experts" + ) + substitutions[module.layer_idx] = (f"{name}.values", module) + return substitutions + + +def _parse_approx_value_name(name: str) -> tuple[int, int] | None: + """Return (layer index, expert ID) of an ``approx_value`` tensor name.""" + if not name.endswith(APPROX_VALUE_SUFFIX): + return None + experts_path, _, expert_id = name.removesuffix(APPROX_VALUE_SUFFIX).rpartition(".") + layer_idx = decoder_layer_index(experts_path) + if layer_idx is None or not expert_id.isdigit(): + return None + return layer_idx, int(expert_id) + + +def as_expert_substitution_model(model_cls: type[_T], model_config: Any) -> type[_T]: + """Load constant expert values at the ``load_weights`` boundary. + + Initial loading and direct weight updates both go through ``load_weights``. + Explicitly named values are always consumed and loaded if their layer is + built by this model. ``approx_value`` tensors of layers built elsewhere + (other pipeline stages, MTP layers) are passed through, so the model's own + loader skips them like any other weight of a layer it does not own. + """ + spec = get_expert_substitution_spec(model_config) + if spec is None: + return model_cls + value_names = spec.value_names + + class ExpertSubstitutionModel(model_cls): # type: ignore[valid-type, misc] + @wraps(model_cls.__init__) # type: ignore[misc] + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._expert_substitutions = _collect_expert_substitutions(self) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + if any(s.values.is_meta for _, s in self._expert_substitutions.values()): + # Put the previous weights back so the model stays usable. + finalize_layerwise_reload(self, model_config) + raise NotImplementedError( + "expert substitution does not support layerwise weight reload" + ) + loaded_substitutions: set[str] = set() + + def remaining_weights(): + for name, loaded_weight in weights: + targets = value_names.get(name) + if targets is None: + key = _parse_approx_value_name(name) + if key is None or key[0] not in self._expert_substitutions: + yield name, loaded_weight + continue + targets = (key,) + for layer_idx, expert_id in targets: + if layer_idx not in self._expert_substitutions: + continue + param_name, substitution = self._expert_substitutions[layer_idx] + param = substitution.values + param.weight_loader(param, loaded_weight, expert_id) + loaded_substitutions.add(param_name) + + loaded = super().load_weights(remaining_weights()) + if loaded is not None: + loaded.update(loaded_substitutions) + return loaded + + def process_weights_after_loading(self) -> None: + for param_name, substitution in self._expert_substitutions.values(): + substitution.validate_loaded_values(param_name) + finalize = getattr(super(), "process_weights_after_loading", None) + if finalize is not None: + finalize() + + ExpertSubstitutionModel.__name__ = model_cls.__name__ + ExpertSubstitutionModel.__qualname__ = model_cls.__qualname__ + return ExpertSubstitutionModel From 98f721116227f0f48690deb05c815c057bac0668 Mon Sep 17 00:00:00 2001 From: Venky <23023424+venkywonka@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:37:56 -0700 Subject: [PATCH 453/952] [Attention][Spec Decode] Multi-query Triton split-K for spec-decode verify, with empty-segment/padding hardening (#56148) Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com> Signed-off-by: Venky <23023424+venkywonka@users.noreply.github.com> Co-authored-by: OpenAI GPT-6 Astra Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) --- .../test_triton_unified_attention.py | 281 +++++++++++++++++- .../v1/attention/test_triton_attn_segments.py | 158 ++++++++++ vllm/v1/attention/backends/triton_attn.py | 32 +- .../attention/ops/triton_attention_helpers.py | 10 +- .../attention/ops/triton_unified_attention.py | 44 ++- 5 files changed, 500 insertions(+), 25 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_segments.py diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index d936e377040c..b50e97b52bc9 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + import pytest import torch @@ -9,11 +11,17 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops import triton_unified_attention as attention from vllm.v1.attention.ops.triton_attention_helpers import ( apply_softcap, compute_tile_loop_bounds, + softmax_step, +) +from vllm.v1.attention.ops.triton_unified_attention import ( + MAX_3D_QUERY_LEN, + reduce_segments, + unified_attention, ) -from vllm.v1.attention.ops.triton_unified_attention import unified_attention from vllm.v1.kv_cache_interface import KVQuantMode pytestmark = pytest.mark.skip_global_cleanup @@ -37,6 +45,259 @@ SEQ_THRESHOLD_3D_VALUES = [0, 8] +def _scratch(rows): + shapes = [(rows, 8, 16, 128), (rows, 8, 16), (rows, 8, 16)] + return [torch.empty(shape, dtype=torch.float32, device="cpu") for shape in shapes] + + +def _uses_splitk( + monkeypatch, + query_len, + scratch, + *, + causal=True, + threshold=8, + mm_prefix=None, + num_seqs=2, +): + query = torch.empty( + (num_seqs * query_len, 8, 128), dtype=torch.bfloat16, device="cpu" + ) + cache = torch.empty((1, 128, 1, 128), dtype=torch.bfloat16, device="cpu") + kernel = MagicMock() + monkeypatch.setattr(attention, "kernel_unified_attention", kernel) + monkeypatch.setattr(attention, "reduce_segments", MagicMock()) + unified_attention( + query, + cache, + cache, + torch.empty_like(query), + torch.arange(num_seqs + 1, dtype=torch.int32, device="cpu") * query_len, + query_len, + torch.full((num_seqs,), 128, dtype=torch.int32, device="cpu"), + 128, + 128**-0.5, + causal, + (-1, -1), + torch.zeros((num_seqs, 1), dtype=torch.int32, device="cpu"), + 0.0, + None, + None, + None, + seq_threshold_3D=threshold, + num_par_softmax_segments=16, + softmax_segm_output=scratch[0], + softmax_segm_max=scratch[1], + softmax_segm_expsum=scratch[2], + mm_prefix_range=mm_prefix, + ) + return kernel.__getitem__.return_value.call_args.kwargs["IS_3D"] + + +def _spy_reduce_segments(monkeypatch): + original = attention.reduce_segments + grids = [] + + def launch(grid): + grids.append(grid) + return original[grid] + + spy = MagicMock() + spy.__getitem__.side_effect = launch + monkeypatch.setattr(attention, "reduce_segments", spy) + return grids + + +@pytest.mark.parametrize( + "query_len,options,expected", + [ + (1, {}, True), + (1, {"causal": False}, True), + (5, {}, True), + (6, {"num_seqs": 1}, True), + (8, {}, True), + (9, {"num_seqs": 1}, False), + (5, {"causal": False}, False), + (5, {"causal": "per_sequence"}, False), + (5, {"threshold": 1}, False), + (5, {"batch_invariant": True}, False), + (5, {"mm_prefix": True}, False), + # BLOCK_Q is 2 for 8 query heads per KV head: 4 * ceil(4 / 2) <= 8. + (4, {"num_seqs": 4}, True), + (5, {"num_seqs": 4}, False), + (5, {"scratch_rows": 9}, False), + ], +) +def test_splitk_admission_preserves_decode_and_bounds_verification( + monkeypatch, query_len, options, expected +): + options = options.copy() + monkeypatch.setattr( + attention, "is_batch_invariant", options.pop("batch_invariant", False) + ) + num_seqs = options.setdefault("num_seqs", 2) + if options.get("causal") == "per_sequence": + options["causal"] = torch.ones(num_seqs, dtype=torch.bool, device="cpu") + if options.get("mm_prefix"): + options["mm_prefix"] = torch.zeros( + (num_seqs, 1, 2), dtype=torch.int32, device="cpu" + ) + scratch = _scratch(options.pop("scratch_rows", num_seqs * query_len)) + assert _uses_splitk(monkeypatch, query_len, scratch, **options) is expected + + +@pytest.mark.parametrize("empty_state", ["masked", "unvisited", "all_empty"]) +def test_reduce_segments_ignores_empty_softmax_states(empty_state: str) -> None: + """Empty segments must not suppress strongly negative, nonempty segments.""" + output = torch.full((1, 1, 32), float("nan"), device=DEVICE_TYPE) + partial = torch.zeros((1, 1, 16, 32), device=DEVICE_TYPE) + maxima = torch.full((1, 1, 16), -float("inf"), device=DEVICE_TYPE) + masses = torch.ones_like(maxima) + if empty_state != "all_empty": + partial[:, :, 0] = 3.0 + maxima[:, :, 0] = -1000.0 + if empty_state in ("masked", "all_empty"): + maxima[:, :, 1:] = 0.0 + masses[:, :, 1:] = 0.0 + reduce_segments[(1, 1)]( + output, + partial, + maxima, + masses, + torch.tensor([32], dtype=torch.int32, device=DEVICE_TYPE), + 1, + 1, + 1.0, + output.stride(0), + output.stride(1), + 1, + 16, + 32, + 32, + torch.tensor([0, 1], dtype=torch.int32, device=DEVICE_TYPE), + 16, + 16, + False, + ) + torch.testing.assert_close( + output, torch.full_like(output, 0.0 if empty_state == "all_empty" else 3.0) + ) + + +@triton.jit +def _softmax_after_empty_tile(output_ptr): + maximum = tl.full((1,), float("-inf"), tl.float32) + mass = tl.full((1,), 1.0, tl.float32) + empty_scores = tl.full((1, 16), float("-inf"), tl.float32) + maximum, mass, _, _ = softmax_step(empty_scores, maximum, mass) + scores = tl.full((1, 16), -1000.0, tl.float32) + _, mass, probabilities, _ = softmax_step(scores, maximum, mass) + tl.store(output_ptr + tl.arange(0, 16), tl.sum(probabilities, axis=0) / mass) + + +def test_splitk_softmax_recovers_after_fully_masked_tile() -> None: + """A masked first tile must not suppress later, strongly negative logits.""" + probabilities = torch.empty(16, device=DEVICE_TYPE) + _softmax_after_empty_tile[(1,)](probabilities) + torch.testing.assert_close(probabilities, torch.full_like(probabilities, 1 / 16)) + + +@pytest.mark.parametrize( + "query_lens,seq_lens", [([5, 0], [5, 0]), ([0], [0]), ([5], [0]), ([5], [5])] +) +def test_reduce_segments_leaves_graph_padding_untouched(query_lens, seq_lens) -> None: + """A padded reducer row must return before reading sequence or scratch data.""" + output = torch.full((10, 1, 32), 123.0, device=DEVICE_TYPE) + partial = torch.zeros((10, 1, 16, 32), device=DEVICE_TYPE) + maxima = torch.zeros((10, 1, 16), device=DEVICE_TYPE) + masses = torch.ones_like(maxima) + query_starts = torch.tensor( + [0] + query_lens, dtype=torch.int32, device=DEVICE_TYPE + ).cumsum(0, dtype=torch.int32) + reduce_segments[(10, 1)]( + output, + partial, + maxima, + masses, + torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE_TYPE), + len(seq_lens), + 1, + 1.0, + output.stride(0), + output.stride(1), + 1, + 16, + 32, + 32, + query_starts, + 16, + 16, + False, + ) + active_rows = sum(q for q, kv in zip(query_lens, seq_lens) if kv > 0) + torch.testing.assert_close( + output[:active_rows], torch.zeros_like(output[:active_rows]) + ) + torch.testing.assert_close( + output[active_rows:], torch.full_like(output[active_rows:], 123.0) + ) + + +@pytest.mark.parametrize( + ("seq_lens", "num_heads", "q_dtype"), + [ + *[ + pytest.param([(w, 128), (w, 129), (w, 257)], (8, 1), None, id=f"w{w}") + for w in range(2, MAX_3D_QUERY_LEN + 1, 2) + ], + # Prefill chunk, verification, decode, padding, short query. + pytest.param( + [(8, 8), (6, 129), (1, 257), (0, 0), (3, 64)], (8, 1), None, id="mixed" + ), + *[ + pytest.param( + [(6, 128), (6, 129), (6, 257)], + heads, + FP8_DTYPE, + id=f"w6-fp8q-kv{heads[1]}", + ) + for heads in [(8, 1), (8, 8)] + ], + ], +) +@pytest.mark.parametrize("sliding_window", [None, 32]) +def test_speculative_splitk_matches_reference( + monkeypatch, seq_lens, num_heads, q_dtype, sliding_window +): + reduce_grids = _spy_reduce_segments(monkeypatch) + test_triton_unified_attn( + seq_lens=seq_lens, + num_heads=num_heads, + head_size=128, + sliding_window=sliding_window, + dtype=torch.bfloat16, + block_size=128, + soft_cap=None, + num_blocks=8, + q_dtype=q_dtype, + seq_threshold_3D=32, + ) + assert len(reduce_grids) == 1 + + +def test_speculative_splitk_bf16_query_fp8_kv(monkeypatch): + reduce_grids = _spy_reduce_segments(monkeypatch) + test_triton_unified_attn_bf16_query_fp8_kv( + seq_lens=[(6, 128), (6, 129), (6, 4097)], + num_heads=(8, 1), + head_size=128, + block_size=128, + num_blocks=64, + seq_threshold_3D=16, + ) + assert len(reduce_grids) == 1 + + @triton.jit def _compute_clamped_mm_tile_bounds( output_ptr, @@ -471,16 +732,19 @@ def test_triton_unified_attn( num_par_softmax_segments = 16 head_size_padded = next_power_of_2(head_size) + scratch_rows = seq_threshold_3D * ( + max_query_len if max_query_len <= MAX_3D_QUERY_LEN else 1 + ) softmax_segm_output = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded), + (scratch_rows, num_query_heads, num_par_softmax_segments, head_size_padded), dtype=torch.float32, ) softmax_segm_max = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) softmax_segm_expsum = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) @@ -591,16 +855,19 @@ def test_triton_unified_attn_bf16_query_fp8_kv( num_par_softmax_segments = 16 head_size_padded = next_power_of_2(head_size) + scratch_rows = seq_threshold_3D * ( + max_query_len if max_query_len <= MAX_3D_QUERY_LEN else 1 + ) softmax_segm_output = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded), + (scratch_rows, num_query_heads, num_par_softmax_segments, head_size_padded), dtype=torch.float32, ) softmax_segm_max = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) softmax_segm_expsum = torch.empty( - (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + (scratch_rows, num_query_heads, num_par_softmax_segments), dtype=torch.float32, ) diff --git a/tests/v1/attention/test_triton_attn_segments.py b/tests/v1/attention/test_triton_attn_segments.py new file mode 100644 index 000000000000..ae764a14acce --- /dev/null +++ b/tests/v1/attention/test_triton_attn_segments.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder +from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVMetadataBuilder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode + + +def _build( + monkeypatch, + builder_cls, + capability, + num_kv_heads, + query_lens, + max_num_seqs=16, + num_speculative_tokens=None, + parallel_drafting=False, + kv_quant_mode=KVQuantMode.NONE, +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability", lambda cap: cap == capability + ) + monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: num_kv_heads, + get_head_size=lambda: 128, + rswa_window=None, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), + speculative_config=( + None + if num_speculative_tokens is None + else SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, + parallel_drafting=parallel_drafting, + ) + ), + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=num_kv_heads, + head_size=128, + dtype=torch.bfloat16, + kv_quant_mode=kv_quant_mode, + ) + builder = builder_cls(spec, ["layer.0"], config, "cpu") + batch = BatchSpec(seq_lens=[128] * len(query_lens), query_lens=query_lens) + metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) + assert metadata.softmax_segm_output.shape[2] == metadata.num_par_softmax_segments + assert metadata.softmax_segm_max.shape[0] >= sum(query_lens) + # DiffKV re-allocates the output; it must stay sized like the base buffers. + assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] + assert ( + metadata.softmax_segm_output.data_ptr() + == builder.softmax_segm_output.data_ptr() + ) + return builder, metadata + + +BUILDERS = [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder] + + +# Verify batches: the 64-segment scratch holds min(11, max_num_seqs) * (1 + K) +# rows, and its 16-segment view covers every 3D-eligible sequence. 64 segments +# follow the request count, not Q blocks: on SM12.0 they stay faster for verify +# batches up to 11 requests (22 Q blocks at K=3). +@pytest.mark.parametrize("builder_cls", BUILDERS) +@pytest.mark.parametrize( + "max_num_seqs,num_speculative_tokens,num_seqs,rows,segments", + [ + (16, 3, 10, 44, 64), + (16, 3, 11, 44, 64), + (16, 3, 12, 44, 16), + (16, 3, 16, 44, 16), + (1, 5, 1, 6, 64), + ], +) +def test_split_k_segment_rows_cover_verify_queries( + monkeypatch, + builder_cls, + max_num_seqs, + num_speculative_tokens, + num_seqs, + rows, + segments, +): + builder, metadata = _build( + monkeypatch, + builder_cls, + (12, 0), + 1, + [1 + num_speculative_tokens] * num_seqs, + max_num_seqs, + num_speculative_tokens, + ) + assert builder.softmax_segm_max.shape[:2] == (rows, 8) + assert metadata.num_par_softmax_segments == segments + + +# Without SM12.0 64-segment reuse, the scratch holds exactly +# min(seq_threshold_3D, max_num_seqs) * query_len rows for split-K query_lens. +@pytest.mark.parametrize( + "num_speculative_tokens,parallel_drafting,max_num_seqs,rows", + [ + (None, False, 8, 8), + (4, False, 8, 40), + (7, False, 8, 64), + (8, False, 8, 64), + (3, True, 8, 56), + (4, True, 8, 64), + (5, False, 1, 6), + ], +) +@pytest.mark.parametrize( + "kv_quant_mode", [KVQuantMode.NONE, KVQuantMode.INT4_PER_TOKEN_HEAD] +) +def test_split_k_scratch_rows_follow_speculative_config( + monkeypatch, + num_speculative_tokens, + parallel_drafting, + max_num_seqs, + rows, + kv_quant_mode, +): + # The INT4 launcher never runs multi-query split-K. + if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD: + rows = max_num_seqs + builder, _ = _build( + monkeypatch, + TritonAttentionMetadataBuilder, + (9, 0), + 1, + [1], + max_num_seqs, + num_speculative_tokens, + parallel_drafting, + kv_quant_mode, + ) + assert builder.softmax_segm_output.shape == (rows, 8, 16, 128) + assert builder.softmax_segm_max.shape == (rows, 8, 16) + assert builder.softmax_segm_expsum.shape == (rows, 8, 16) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index a95b0c1d0e13..8b26bab84fa5 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -39,7 +39,10 @@ triton_reshape_and_cache_flash, triton_reshape_and_cache_flash_per_token_head_quant, ) -from vllm.v1.attention.ops.triton_unified_attention import unified_attention +from vllm.v1.attention.ops.triton_unified_attention import ( + MAX_3D_QUERY_LEN, + unified_attention, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVQuantMode, @@ -144,7 +147,10 @@ def __init__( ) self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS - # On SM120, batches whose 16-segment grid under-fills the SMs use 64. + # On SM120, batches whose 16-segment decode grid under-fills the SMs use + # 64. This counts requests, not Q blocks, on purpose: for 1..11 verify + # requests (K=3/5, 8 q / 1 KV head, 2k-128k context) 64 segments were + # faster or within 6% at 2k, and 1.1-3.6x faster from 8k up. self.max_seqs_64_segments = 0 if current_platform.is_cuda() and current_platform.is_device_capability( (12, 0) @@ -152,18 +158,34 @@ def __init__( self.max_seqs_64_segments = (current_platform.num_compute_units() - 1) // ( NUM_PAR_SOFTMAX_SEGMENTS * self.num_heads_kv ) - max_num_tokens_3d = self.seq_threshold_3D + headdim_padded = next_power_of_2(self.headdim) + max_query_len_3d = 1 + speculative_config = vllm_config.speculative_config + # The INT4 launcher runs multi-query batches in 2D only. + if ( + speculative_config is not None + and kv_cache_spec.kv_quant_mode != KVQuantMode.INT4_PER_TOKEN_HEAD + ): + num_speculative_tokens = speculative_config.num_speculative_tokens or 0 + query_len = 1 + num_speculative_tokens * ( + 2 if speculative_config.parallel_drafting else 1 + ) + max_query_len_3d = min(query_len, MAX_3D_QUERY_LEN) + # Scratch is indexed by query token, including verification tokens. + max_num_seqs_3d = min( + self.seq_threshold_3D, vllm_config.scheduler_config.max_num_seqs + ) + max_num_tokens_3d = max_num_seqs_3d * max_query_len_3d if self.max_seqs_64_segments > 0: self.num_par_softmax_segments = 64 # build() reuses this scratch at 16 segments with proportionally more rows. max_num_tokens_3d = max( - min(self.max_seqs_64_segments, max_num_tokens_3d), + min(self.max_seqs_64_segments, max_num_seqs_3d) * max_query_len_3d, cdiv( max_num_tokens_3d, self.num_par_softmax_segments // NUM_PAR_SOFTMAX_SEGMENTS, ), ) - headdim_padded = next_power_of_2(self.headdim) self.softmax_segm_output = torch.empty( ( max_num_tokens_3d, diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index 7804b013171e..f2399b742380 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -470,15 +470,15 @@ def softmax_step(S, M, L): # compute running maximum # m_j : (BLOCK_M,) m_j = tl.maximum(M, tl.max(S, axis=1)) - # For sliding window there's a chance the max is -inf due to masking of - # the entire row. In this case we need to set m_j 0 to avoid NaN - m_j = tl.where(m_j > float("-inf"), m_j, 0.0) + # Keep an empty row's maximum at -inf for later tiles; use a finite shift + # only for the exponent arithmetic below. + safe_m_j = tl.where(m_j > float("-inf"), m_j, 0.0) # P : (BLOCK_M, TILE_SIZE) - P = tl.exp(S - m_j[:, None]) + P = tl.exp(S - safe_m_j[:, None]) # l_j : (BLOCK_M,) l_j = tl.sum(P, axis=1) # alpha : (BLOCK_M, ) - alpha = tl.exp(M - m_j) + alpha = tl.exp(M - safe_m_j) # update constants L_new = L * alpha + l_j return m_j, L_new, P, alpha diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 63a9cdc7f956..d914e574f2cc 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -33,6 +33,8 @@ logger = init_logger(__name__) is_batch_invariant = envs.VLLM_BATCH_INVARIANT float8_info = torch.finfo(current_platform.fp8_dtype()) +# Longest per-request query (e.g. 1 + num_speculative_tokens) run by split-K. +MAX_3D_QUERY_LEN = 8 @triton.jit @@ -721,12 +723,18 @@ def reduce_segments( query_token_idx = tl.program_id(0) query_head_idx = tl.program_id(1) + # Graph padding must be rejected before sequence lookup or scratch access. + if query_token_idx >= tl.load(query_start_len_ptr + num_seqs): + return + seq_idx = find_seq_idx( query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False ) # sequence len for this particular sequence seq_len = tl.load(seq_lens_ptr + seq_idx) + if seq_len <= 0: + return # number of segments for this particular sequence num_segments = NUM_SEGMENTS_PER_SEQ @@ -746,11 +754,16 @@ def reduce_segments( + tl.arange(0, NUM_SEGMENTS_PER_SEQ) ) segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + # Empty segments must not set the maximum. Unvisited segments retain + # the M=-inf,L=1 initial state. + valid_segment = segm_mask & (segm_expsum > 0) & (segm_max != float("-inf")) + segm_max = tl.where(valid_segment, segm_max, float("-inf")) overall_max = tl.max(segm_max) # load and rescale segment exp sums - segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) - segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) + segment_scale = tl.where(valid_segment, tl.exp(segm_max - overall_max), 0.0) + segm_expsum = segm_expsum * segment_scale overall_expsum = tl.sum(segm_expsum) # load, rescale, and add segment attention outputs @@ -766,7 +779,7 @@ def reduce_segments( mask=segm_mask[:, None] & dim_mask[None, :], other=0.0, ) - segm_output *= tl.exp(segm_max - overall_max)[:, None] + segm_output *= segment_scale[:, None] acc_sum = tl.sum(segm_output, axis=0) # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) @@ -1047,9 +1060,12 @@ def unified_attention( ) # Launch the 2D kernel if - # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or - # 2. The batch includes at least one prefill request, or - # 3. The number of sequences exceeds the configured threshold, or + # 1. No split-K threshold or scratch buffers, or the scratch has fewer + # rows than query tokens, or + # 2. max_seqlen_q > 1 and either it exceeds MAX_3D_QUERY_LEN, masking is + # non-causal, per-sequence causal, or mm-prefix, or the tuned + # large-head path applies, or + # 3. The number of Q blocks exceeds seq_threshold_3D, or # 4. Batch invariance is enabled use_3d = not ( seq_threshold_3D is None @@ -1057,10 +1073,22 @@ def unified_attention( or softmax_segm_output is None or softmax_segm_max is None or softmax_segm_expsum is None - or max_seqlen_q > 1 - or num_seqs > seq_threshold_3D + or ( + max_seqlen_q > 1 + and ( + max_seqlen_q > MAX_3D_QUERY_LEN + or not use_causal + or use_per_seq_causal + or use_mm_prefix + or tuned_large_head + ) + ) + or num_seqs * triton.cdiv(max_seqlen_q, BLOCK_Q) > seq_threshold_3D or is_batch_invariant ) + if use_3d and softmax_segm_max.shape[0] < q.shape[0]: + logger.debug_once("Split-K scratch rows < query tokens; using 2D.") + use_3d = False # The kernel signature is the same for 2D and 3D — only the launch # grid + a handful of constexpr toggles differ. Per-token-head scale From c6566bb9f678679c6e42ad2beb45fece8b35356a Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:20:38 +0800 Subject: [PATCH 454/952] [ROCm][MLA] Keep DCP prefill context in FP8 through AllGather (#59070) Signed-off-by: whx-sjtu Co-authored-by: OpenAI Codex Co-authored-by: Claude Code Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .buildkite/test-amd.yaml | 3 + .buildkite/test_areas/distributed.yaml | 3 + .../test_rocm_aiter_mla_op_registration.py | 20 + .../attention/test_rocm_aiter_mla_dcp_cprr.py | 7 + .../test_rocm_aiter_mla_mtp_split.py | 77 ++- .../v1/attention/test_rocm_mla_dcp_context.py | 445 ++++++++++++++++++ vllm/_aiter_ops.py | 52 ++ .../attention/backends/mla/rocm_aiter_mla.py | 152 +++++- .../attention/ops/rocm_aiter_mla_prefill.py | 105 +++++ 9 files changed, 847 insertions(+), 17 deletions(-) create mode 100644 tests/v1/attention/test_rocm_mla_dcp_context.py create mode 100644 vllm/v1/attention/ops/rocm_aiter_mla_prefill.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e4ae7f897e35..d2702e7d4f82 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3539,8 +3539,11 @@ steps: - examples/rl/rlhf_async_new_apis.py - tests/distributed/test_weight_transfer.py - tests/distributed/test_packed_tensor.py + - tests/v1/attention/test_rocm_mla_dcp_context.py + - vllm/v1/attention/ops/rocm_aiter_mla_prefill.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/attention/test_rocm_mla_dcp_context.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index f541c42c93a8..b30cbac0ced4 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -528,8 +528,11 @@ steps: - examples/rl/rlhf_async_new_apis.py - tests/distributed/test_weight_transfer.py - tests/distributed/test_packed_tensor.py + - tests/v1/attention/test_rocm_mla_dcp_context.py + - vllm/v1/attention/ops/rocm_aiter_mla_prefill.py commands: - pytest -v -s tests/distributed/test_context_parallel.py + - pytest -v -s tests/v1/attention/test_rocm_mla_dcp_context.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization diff --git a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py index 0944be8937c8..cb40c61c3b94 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_op_registration.py @@ -33,6 +33,26 @@ def _require_aiter(): pytest.skip("aiter is required on supported ROCm hardware for this test") +@torch.inference_mode() +def test_gather_kv_b_proj_op_schema() -> None: + """Validate fake tensors and mutation of both expanded K/V outputs.""" + _require_aiter() + from vllm._aiter_ops import rocm_aiter_ops # noqa: F401 + + rows, num_heads = 32, 12 + cache = torch.randn(rows, 1, 576, device="cuda").to(current_platform.fp8_dtype()) + scale = torch.tensor([0.07], device="cuda") + cu_seq_lens = torch.tensor([0, 13, rows], dtype=torch.int32, device="cuda") + indices = torch.arange(rows, dtype=torch.int32, device="cuda").flip(0) + weight = torch.randn(num_heads * 256, 512, dtype=torch.bfloat16, device="cuda") + k = torch.empty(rows, num_heads, 192, dtype=torch.bfloat16, device="cuda") + v = torch.empty(rows, num_heads, 128, dtype=torch.bfloat16, device="cuda") + opcheck( + torch.ops.vllm.rocm_aiter_gather_kv_b_proj, + (cache, scale, cu_seq_lens, indices, cu_seq_lens, weight, k, v), + ) + + @torch.inference_mode() def test_mla_decode_fwd_op_schema() -> None: """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing. diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py index cd88ade8eebf..56886872ecc6 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py @@ -142,6 +142,13 @@ def _build(monkeypatch, route, gathered_heads, segmented_supported=True): def init_common_builder(self, *args, **kwargs): self.num_heads = gathered_heads // DCP + self.q_data_type = torch.bfloat16 + self.mla_dims = SimpleNamespace( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) self.dcp_world_size = DCP self.reorder_batch_threshold = 5 diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index da141acf6881..1583995ac98b 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -740,26 +740,34 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2) -@pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128]) @pytest.mark.parametrize( - "spec_method, parallel_drafting", + "num_heads,spec_method,parallel_drafting,group_cache_dtype,needs_row_map", [ - ("deepseek_mtp", False), - # A drafter that is not one of the historically recognized MTP methods, - # and a parallel one, so the threshold is 1 + 2 * num_spec rather than - # 1 + num_spec. Sizing the metadata off a method name instead leaves - # these at qlen=1 while the router still admits the full range. - ("custom", True), - ("eagle", False), + *[ + (num_heads, method, parallel, None, True) + for num_heads in [8, 16, 24, 32, 64, 128] + for method, parallel in [ + ("deepseek_mtp", False), + # Parallel drafting needs qlen=1+2*num_spec, regardless of name. + ("custom", True), + ("eagle", False), + ] + ], + pytest.param(16, "deepseek_mtp", False, "auto", False, id="bf16-group"), + pytest.param(16, "deepseek_mtp", False, "fp8_e5m2", False, id="e5m2-group"), ], ) -def test_mtp_builder_init_sizes_native_fp8_metadata( - monkeypatch, num_heads, spec_method, parallel_drafting +def test_builder_init_sizes_decode_metadata_and_filters_dcp_maps( + monkeypatch, + num_heads, + spec_method, + parallel_drafting, + group_cache_dtype, + needs_row_map, ): - """Aiter init sizes the metadata for every query length decode can be handed. + """Size decode metadata for padded heads/qlen; skip unsupported DCP maps. - Sweeping num_heads asserts metadata is sized for the padded decode shape, - covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding. + The group's raw cache format must override the global FP8 configuration. """ dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") info_calls = [] @@ -789,6 +797,13 @@ def get_mla_metadata_info_v1( def init_common_builder(self, *args, **kwargs): self.num_heads = num_heads + self.q_data_type = config.model_config.dtype + self.mla_dims = SimpleNamespace( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) self.dcp_world_size = 1 self.cp_kv_cache_interleave_size = 1 # Mirror what _init_reorder_batch_threshold would have left behind: the @@ -813,6 +828,7 @@ def init_common_builder(self, *args, **kwargs): init_common_builder, ) monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) + monkeypatch.setattr(rocm_aiter_mla, "is_aiter_found_and_supported", lambda: True) config = SimpleNamespace( speculative_config=SimpleNamespace( @@ -840,7 +856,9 @@ def init_common_builder(self, *args, **kwargs): ), ) builder = AiterMLAMetadataBuilder( - kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), + kv_cache_spec=SimpleNamespace( + block_size=1, dtype=torch.bfloat16, cache_dtype_str=group_cache_dtype + ), layer_names=["layer.0"], vllm_config=config, device=torch.device("cpu"), @@ -860,6 +878,35 @@ def init_common_builder(self, *args, **kwargs): assert builder._mla_q_dtype == dtypes.fp8 assert builder._mla_kv_dtype == dtypes.fp8 + builder.dcp_world_size = 2 + chunk = object() + rows = torch.tensor([0], dtype=torch.int64) + build_rows = mock.Mock(return_value=rows) + monkeypatch.setattr(rocm_aiter_mla, "context_row_indices", build_rows) + monkeypatch.setattr( + rocm_aiter_mla, + "_supports_fused_dcp_prefill_config", + mock.Mock(side_effect=AssertionError("Support must be cached at init")), + ) + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, + "build", + lambda *args: SimpleNamespace( + prefill=SimpleNamespace( + chunked_context=SimpleNamespace(chunks=[chunk]), + block_table=torch.empty(0), + ), + decode=None, + dcp_context_row_indices=None, + ), + ) + for _ in range(2): + metadata = builder.build(0, SimpleNamespace(causal=True)) + assert metadata.dcp_context_row_indices == ([rows] if needs_row_map else None) + assert build_rows.call_args_list == ( + [mock.call(chunk, torch.device("cpu"))] * 2 if needs_row_map else [] + ) + def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch): get_mla_metadata_v1 = mock.MagicMock() diff --git a/tests/v1/attention/test_rocm_mla_dcp_context.py b/tests/v1/attention/test_rocm_mla_dcp_context.py new file mode 100644 index 000000000000..c29576f93ba9 --- /dev/null +++ b/tests/v1/attention/test_rocm_mla_dcp_context.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.utils import multi_gpu_only +from vllm.platforms import current_platform +from vllm.v1.attention.ops.rocm_aiter_mla_prefill import context_row_indices + + +@pytest.mark.parametrize( + "padded,lengths,starts,total,expected", + [ + ([3, 2], [[7, 6], [2, 1]], [5, 0], 6, [0, 1, 5, 3, 4, 8]), + ([3], [[2, 9]], [4], 3, [3, 4, 5]), + ([3], [[1, 1]], [0], 3, None), + ], + ids=["padding_and_chunk_starts", "empty_rank", "inconsistent_total"], +) +def test_dcp_context_indices(padded, lengths, starts, total, expected): + chunk = SimpleNamespace( + padded_local_seq_lens=padded, + local_context_lens_allranks=lengths, + local_starts=starts, + num_local_context_tokens=sum(padded), + num_context_tokens=total, + ) + if expected is None: + with pytest.raises(AssertionError): + context_row_indices(chunk, torch.device("cpu")) + else: + rows = context_row_indices(chunk, torch.device("cpu")) + assert rows.tolist() == expected + assert rows.dtype == torch.int32 + + +def test_compressed_gather_preserves_bytes_and_workspace_partition(monkeypatch): + from vllm.v1.attention.ops import rocm_aiter_mla_prefill as prefill + + workspace = torch.full((12, 4), -1, dtype=torch.bfloat16) + cache = torch.arange(8, dtype=torch.uint8).reshape(2, 4) + chunk = SimpleNamespace( + num_local_context_tokens=2, + local_context_lens_allranks=[[2, 2]], + padded_local_cu_seq_lens=torch.tensor([0, 2]), + num_requests=1, + starts=torch.tensor([0]), + ) + + def local_gather(src, dst, *args): + assert src.dtype == dst.dtype == torch.uint8 + dst.copy_(src) + + def allgather(dst, src): + assert dst.dtype == src.dtype == torch.uint8 + assert dst.shape == (4, 4) + assert src.shape == (2, 4) + dst[:2].copy_(src) + dst[2:].copy_(src + 16) + + monkeypatch.setattr(prefill.ops, "cp_gather_cache", local_gather) + gathered = prefill.gather_compressed_context( + cache, workspace, torch.empty(1, 1), chunk, allgather, torch.float8_e4m3fn + ) + torch.testing.assert_close( + gathered.view(torch.uint8), torch.cat((cache, cache + 16)) + ) + # Retain the original logical capacity even though the byte view is larger. + assert gathered.data_ptr() - workspace.data_ptr() == 4 * 4 + torch.testing.assert_close(workspace.view(torch.uint8).reshape(-1, 4)[:2], cache) + + +@pytest.mark.parametrize( + "unsupported", + [ + "packed_cache", + "e5m2", + "bf16_cache", + "quantized_weight", + "bias", + "no_indices", + "query_dtype", + "prefill_dtype", + "weight_dtype", + "weight_layout", + "cache_layout", + "dimensions", + "no_aiter", + ], +) +def test_dcp_prefill_unsupported_formats_use_original_path(monkeypatch, unsupported): + from vllm.platforms import current_platform + + if not current_platform.is_rocm(): + pytest.skip("ROCm MLA backend") + from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + AiterMLAImpl, + MLACommonImpl, + ) + + if unsupported == "no_aiter": + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.rocm_aiter_mla." + "is_aiter_found_and_supported", + lambda: False, + ) + + impl = object.__new__(AiterMLAImpl) + impl.kv_cache_dtype = { + "packed_cache": "fp8_ds_mla", + "e5m2": "fp8_e5m2", + "bf16_cache": "auto", + }.get(unsupported, "fp8") + impl.kv_lora_rank = 512 + if unsupported == "dimensions": + impl.kv_lora_rank = 256 + impl.qk_nope_head_dim = impl.v_head_dim = 128 + impl.qk_rope_head_dim = 64 + impl.num_heads = 1 + impl.kv_b_proj = SimpleNamespace( + weight=torch.empty(256, 512, dtype=torch.bfloat16), + weight_scale=torch.ones(1) if unsupported == "quantized_weight" else None, + bias=torch.ones(256) if unsupported == "bias" else None, + ) + if unsupported == "weight_dtype": + impl.kv_b_proj.weight = impl.kv_b_proj.weight.half() + elif unsupported == "weight_layout": + impl.kv_b_proj.weight = torch.empty(512, 256, dtype=torch.bfloat16).T + metadata = SimpleNamespace( + prefill=SimpleNamespace( + q_data_type=torch.float16 + if unsupported == "prefill_dtype" + else torch.bfloat16 + ), + dcp_context_row_indices=None if unsupported == "no_indices" else [], + ) + expected = (torch.ones(1), torch.zeros(1)) + + def fused_mla_kv_concat(kv_nope, k_pe, use_fp8_prefill): + raise AssertionError("The fallback should receive the callback") + + def original_path(*args): + assert args[-1] is fused_mla_kv_concat + return expected + + monkeypatch.setattr( + MLACommonImpl, + "_context_parallel_compute_prefill_context", + original_path, + ) + q = torch.empty(1, 1, 192, dtype=torch.bfloat16) + cache = torch.empty(1, 1536, 576, dtype=torch.uint8) + if unsupported == "query_dtype": + q = q.half() + elif unsupported == "cache_layout": + cache = torch.empty(1, 1536, 1152, dtype=torch.uint8)[..., ::2] + assert ( + impl._context_parallel_compute_prefill_context( + q, + cache, + metadata, + torch.ones(1), + 8, + fused_mla_kv_concat_fn=fused_mla_kv_concat, + ) + is expected + ) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm MLA backend") +@multi_gpu_only(num_gpus=2) +@pytest.mark.parametrize("interleave", [1, 64]) +@pytest.mark.parametrize( + "context_lens,workspace_size", + [([145, 37], 1024), ([545, 0, 3, 0], 256)], + ids=["single_chunk", "continuations_and_empty_context"], +) +def test_dcp_prefill_matches_full_attention( + interleave, context_lens, workspace_size, tmp_path +): + """Real DCP gather, fused expansion and suffix/chunk merges match dense SDPA.""" + import torch.multiprocessing as mp + + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("AITER on CDNA 3 or newer is required") + mp.spawn( + _dcp_prefill_worker, + args=( + f"file://{tmp_path / 'dist_init'}", + interleave, + context_lens, + workspace_size, + ), + nprocs=2, + ) + + +@torch.inference_mode() +def _dcp_prefill_worker(rank, init_method, interleave, context_lens, workspace_size): + import torch.nn.functional as F + + from tests.v1.attention.test_mla_backends import MockMLAAttentionLayer + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import set_current_vllm_config + from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.model_executor.layers.linear import ColumnParallelLinear + from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( + AiterFlashAttnPrefillBackend, + ) + from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + AiterMLAImpl, + AiterMLAMetadataBuilder, + ) + from vllm.v1.attention.ops.dcp import MLADCPManager + from vllm.v1.kv_cache_interface import MLAAttentionSpec + from vllm.v1.worker.workspace import init_workspace_manager + + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + device = torch.device(f"cuda:{rank}") + init_workspace_manager(device) + world_size, num_heads, block_size = 2, 12, 64 + config = create_vllm_config( + model_name="deepseek-ai/DeepSeek-V2-Lite", + dtype=torch.bfloat16, + tensor_parallel_size=world_size, + max_model_len=2048, + block_size=block_size, + max_num_seqs=8, + max_num_batched_tokens=2048, + ) + config.cache_config.cache_dtype = "fp8" + config.parallel_config.decode_context_parallel_size = world_size + config.parallel_config.cp_kv_cache_interleave_size = interleave + config.attention_config.use_prefill_query_quantization = False + config.kernel_config.enable_jit_warmup = False + with set_current_vllm_config(config): + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=rank, + distributed_init_method=init_method, + ) + initialize_model_parallel( + tensor_model_parallel_size=world_size, + decode_context_model_parallel_size=world_size, + ) + try: + projection = ColumnParallelLinear( + 512, + num_heads * world_size * 256, + bias=False, + params_dtype=torch.bfloat16, + ).to(device) + torch.manual_seed(101 + rank) + projection.weight.copy_(torch.randn_like(projection.weight) * 0.02) + dims = dict( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) + impl = AiterMLAImpl( + num_heads=num_heads, + head_size=576, + scale=192**-0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + qk_head_dim=192, + kv_b_proj=projection, + **dims, + ) + layer = MockMLAAttentionLayer( + impl=impl, + num_heads=num_heads, + device=device, + kv_b_proj=projection, + q_scale=1.0, + k_scale=0.07, + **dims, + ) + layer.prefill_backend = AiterFlashAttnPrefillBackend( + num_heads=num_heads, + scale=impl.scale, + vllm_config=config, + **dims, + ) + layer.dcp_manager = MLADCPManager( + config, + device, + num_heads, + 576, + 512, + torch.bfloat16, + torch.bfloat16, + None, + True, + False, + ) + config.compilation_config.static_forward_context["mla"] = layer + spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=current_platform.fp8_dtype(), + cache_dtype_str="fp8", + ) + builder = AiterMLAMetadataBuilder(spec, ["mla"], config, device) + builder.chunked_prefill_workspace_size = workspace_size + builder.chunked_prefill_workspace = torch.empty( + (workspace_size + workspace_size // world_size, 576), + dtype=torch.bfloat16, + device=device, + ) + layer.dcp_manager.init_kv_gather( + builder.chunked_prefill_workspace, workspace_size + ) + query_lens = [i + 2 for i in range(len(context_lens))] + batch = BatchSpec( + seq_lens=[c + q for c, q in zip(context_lens, query_lens)], + query_lens=query_lens, + ) + common = create_common_attn_metadata( + batch, block_size, device, arange_block_indices=True + ) + metadata = builder.build(0, common) + assert metadata.prefill is not None + chunked = metadata.prefill.chunked_context + assert chunked is not None and metadata.dcp_context_row_indices is not None + if 0 in context_lens: + assert chunked.empty_token_slices + assert any(c.is_continuation for c in chunked.chunks) + else: + assert len(chunked.chunks) == 1 + blocks_per_req = common.block_table_tensor.shape[1] + cache = torch.zeros( + (len(context_lens) * blocks_per_req, block_size, 576), + dtype=torch.bfloat16, + device=device, + ).to(current_platform.fp8_dtype()) + queries, suffixes, references = [], [], [] + # All ranks reconstruct the same global cache before sharding it. + torch.manual_seed(42) + for request, (context_len, query_len) in enumerate( + zip(context_lens, query_lens) + ): + latent = torch.randn( + context_len, 576, device=device, dtype=torch.bfloat16 + ) + quantized = (latent.float() / layer._k_scale).to(cache.dtype) + positions = torch.arange(context_len, device=device) + shard = quantized.view(torch.uint8)[ + (positions // interleave) % world_size == rank + ] + cache.view(len(context_lens), -1, 576).view(torch.uint8)[ + request, : len(shard) + ].copy_(shard) + suffix = torch.randn( + query_len, 576, device=device, dtype=torch.bfloat16 + ) + q = ( + torch.randn( + query_len, num_heads, 192, device=device, dtype=torch.bfloat16 + ) + * 0.5 + ) + full_latent = torch.cat( + ( + (quantized.float() * layer._k_scale).to(torch.bfloat16), + suffix, + ) + ) + kv = F.linear(full_latent[:, :512], projection.weight).view( + -1, num_heads, 256 + ) + k_nope, v = kv.split(128, dim=-1) + k = torch.cat( + (k_nope, full_latent[:, None, 512:].expand(-1, num_heads, -1)), + dim=-1, + ) + mask = torch.arange(context_len + query_len, device=device)[ + None, : + ] <= (context_len + torch.arange(query_len, device=device)[:, None]) + references.append( + F.scaled_dot_product_attention( + q.transpose(0, 1).float(), + k.transpose(0, 1).float(), + v.transpose(0, 1).float(), + attn_mask=mask, + scale=impl.scale, + ).transpose(0, 1) + ) + queries.append(q) + suffixes.append(suffix) + q, suffix = torch.cat(queries), torch.cat(suffixes) + output = torch.empty( + q.shape[0], num_heads * 128, device=device, dtype=torch.bfloat16 + ) + # Ensure this test cannot silently succeed through the fallback. + with pytest.MonkeyPatch.context() as patch: + + def unexpected_fallback(*args, **kwargs): + raise AssertionError("Expected the compressed DCP path") + + patch.setattr( + "vllm.model_executor.layers.attention.mla_attention." + "MLACommonImpl._context_parallel_compute_prefill_context", + unexpected_fallback, + ) + impl.forward_mha( + q, + suffix[:, :512], + suffix[:, None, 512:], + cache, + metadata, + layer._k_scale, + output, + ) + torch.testing.assert_close( + output.view(-1, num_heads, 128).float(), + torch.cat(references), + rtol=0.03, + atol=0.003, + ) + finally: + cleanup_dist_env_and_memory() diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 8af999906f28..0817c651a0f4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -643,6 +643,36 @@ def _check_aiter_mla_fp8_support() -> bool: return _AITER_MLA_SUPPORTS_FP8 +def _rocm_aiter_gather_kv_b_proj_impl( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + from aiter.ops.triton.gather_kv_b_proj import gather_kv_b_proj + + gather_kv_b_proj( + kv_cache, kv_scale, kv_indptr, kv_indices, cu_seq_lens, weight, None, k, v + ) + + +def _rocm_aiter_gather_kv_b_proj_fake( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + pass + + def _rocm_aiter_mla_decode_fwd_impl( q: torch.Tensor, kv_buffer: torch.Tensor, @@ -2851,6 +2881,13 @@ def register_ops_once() -> None: dispatch_key=current_platform.dispatch_key, ) + direct_register_custom_op( + op_name="rocm_aiter_gather_kv_b_proj", + op_func=_rocm_aiter_gather_kv_b_proj_impl, + mutates_args=["k", "v"], + fake_impl=_rocm_aiter_gather_kv_b_proj_fake, + dispatch_key=current_platform.dispatch_key, + ) direct_register_custom_op( op_name="rocm_aiter_mla_decode_fwd", op_func=_rocm_aiter_mla_decode_fwd_impl, @@ -3465,6 +3502,21 @@ def triton_sparse_mla_fwd( extra_indices=extra_kv_indices, ) + @staticmethod + def gather_kv_b_proj( + kv_cache: torch.Tensor, + kv_scale: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> None: + torch.ops.vllm.rocm_aiter_gather_kv_b_proj( + kv_cache, kv_scale, kv_indptr, kv_indices, cu_seq_lens, weight, k, v + ) + @staticmethod def mla_decode_fwd( q: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 1115da630b3c..d64b4797c031 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -2,15 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +from collections.abc import Callable from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final, cast import torch from vllm import envs -from vllm._aiter_ops import rocm_aiter_ops +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group @@ -22,6 +23,8 @@ MLACommonMetadata, MLACommonMetadataBuilder, QueryLenSupport, + accumulate_mla_context_chunk, + init_mla_context_partial, ) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -38,6 +41,11 @@ from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( merge_mla_segments_triton, ) +from vllm.v1.attention.ops.rocm_aiter_mla_prefill import ( + context_row_indices, + expand_context, + gather_compressed_context, +) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache if TYPE_CHECKING: @@ -46,6 +54,26 @@ logger = init_logger(__name__) +def _supports_fused_dcp_prefill_config( + kv_cache_dtype: str | None, + query_dtype: torch.dtype, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, +) -> bool: + """Whether the configuration supports fused DCP context expansion.""" + return ( + kv_cache_dtype in ("fp8", "fp8_e4m3") + and query_dtype == torch.bfloat16 + and kv_lora_rank == 512 + and qk_nope_head_dim == 128 + and qk_rope_head_dim == 64 + and v_head_dim == 128 + and is_aiter_found_and_supported() + ) + + def _segmented_mla_page_size(block_size: int) -> int: """Largest supported power-of-two subpage dividing a physical KV block. @@ -507,6 +535,7 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): @dataclass class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]): + dcp_context_row_indices: list[torch.Tensor] | None = None work_meta_data: torch.Tensor | None = None work_indptr: torch.Tensor | None = None work_info_set: torch.Tensor | None = None @@ -664,6 +693,19 @@ def __init__( device ).multi_processor_count + # Use the group's original format: decode normalizes E5M2 to "fp8". + kv_cache_dtype = getattr(kv_cache_spec, "cache_dtype_str", None) + if kv_cache_dtype is None: + kv_cache_dtype = vllm_config.cache_config.cache_dtype + self._supports_fused_dcp_prefill = _supports_fused_dcp_prefill_config( + kv_cache_dtype, + self.q_data_type, + self.mla_dims.kv_lora_rank, + self.mla_dims.qk_nope_head_dim, + self.mla_dims.qk_rope_head_dim, + self.mla_dims.v_head_dim, + ) + self.compilation_config = vllm_config.compilation_config self.decode_attn_out_dtype = vllm_config.model_config.dtype @@ -1577,6 +1619,17 @@ def build( attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) + prefill = attn_metadata.prefill + if ( + prefill is not None + and prefill.chunked_context is not None + and self.dcp_world_size > 1 + and self._supports_fused_dcp_prefill + ): + attn_metadata.dcp_context_row_indices = [ + context_row_indices(chunk, prefill.block_table.device) + for chunk in prefill.chunked_context.chunks + ] if ( attn_metadata.decode is not None and attn_metadata.decode.has_persistent_metadata @@ -1975,6 +2028,101 @@ def __init__( self._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd self._mla_reduce_v1 = mla_reduce_v1 + def _context_parallel_compute_prefill_context( + self, + q: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + k_scale: torch.Tensor, + dcp_world_size: int, + fused_mla_kv_concat_fn: Callable[ + [torch.Tensor, torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor] + ] + | None = None, + ): + attn_metadata = cast(AiterMLAMetadata, attn_metadata) + assert attn_metadata.prefill is not None + prefill = attn_metadata.prefill + weight = getattr(self.kv_b_proj, "weight", None) + eligible = ( + _supports_fused_dcp_prefill_config( + self.kv_cache_dtype, + prefill.q_data_type, + self.kv_lora_rank, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + ) + and q.dtype == torch.bfloat16 + and weight is not None + and weight.dtype == torch.bfloat16 + and weight.shape == (self.num_heads * 256, 512) + and weight.is_contiguous() + and getattr(self.kv_b_proj, "bias", None) is None + and getattr(self.kv_b_proj, "weight_scale", None) is None + and kv_c_and_k_pe_cache.is_contiguous() + and attn_metadata.dcp_context_row_indices is not None + ) + if not eligible: + return super()._context_parallel_compute_prefill_context( + q, + kv_c_and_k_pe_cache, + attn_metadata, + k_scale, + dcp_world_size, + fused_mla_kv_concat_fn, + ) + from vllm.platforms import current_platform + + assert weight is not None + assert attn_metadata.dcp_context_row_indices is not None + chunked = prefill.chunked_context + assert chunked is not None and chunked.dcp_manager is not None + assert prefill.prefill_backend is not None + assert len(chunked.chunks) == len(attn_metadata.dcp_context_row_indices) + output = output_lse = None + for chunk, row_indices in zip( + chunked.chunks, attn_metadata.dcp_context_row_indices + ): + gathered = gather_compressed_context( + kv_c_and_k_pe_cache, + chunked.workspace, + prefill.block_table[chunk.request_slice], + chunk, + chunked.dcp_manager.kv_gather, + current_platform.fp8_dtype(), + ) + k, v = expand_context( + gathered, + k_scale, + row_indices, + chunk.cu_seq_lens, + weight, + self.num_heads, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + ) + attn_output, attn_lse = prefill.prefill_backend.run_prefill_context_chunk( + chunk=chunk, + q=q[chunk.token_slice], + k=k, + v=v, + ) + if output is None: + if len(chunked.chunks) == 1 and not chunked.empty_token_slices: + return attn_output, attn_lse + output, output_lse = init_mla_context_partial( + chunked, + attn_output, + attn_lse, + num_tokens=q.shape[0], + ) + accumulate_mla_context_chunk( + chunk, attn_output, attn_lse, output, output_lse + ) + return output, output_lse + def _flash_attn_varlen_diff_headdims( self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs ): diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py b/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py new file mode 100644 index 000000000000..c6d839f72453 --- /dev/null +++ b/vllm/v1/attention/ops/rocm_aiter_mla_prefill.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compressed DCP context preparation for ROCm MLA.""" + +from collections.abc import Callable + +import torch + +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata + + +def context_row_indices( + chunk: MLACommonPrefillMetadata.ContextChunk, device: torch.device +) -> torch.Tensor: + """Map compact request-major rows to the padded rank-major AllGather.""" + assert chunk.padded_local_seq_lens is not None + assert chunk.local_context_lens_allranks is not None + assert chunk.local_starts is not None + segments = [] + offset = 0 + for padded, lengths, start in zip( + chunk.padded_local_seq_lens, + chunk.local_context_lens_allranks, + chunk.local_starts, + ): + for rank, length in enumerate(lengths): + count = min(max(0, length - start), padded) + if count: + first = rank * chunk.num_local_context_tokens + offset + segments.append(torch.arange(first, first + count, dtype=torch.int32)) + offset += padded + rows = torch.cat(segments) if segments else torch.empty(0, dtype=torch.int32) + assert rows.numel() == chunk.num_context_tokens + return rows.to(device=device, non_blocking=True) + + +def gather_compressed_context( + cache: torch.Tensor, + workspace: torch.Tensor, + block_table: torch.Tensor, + chunk: MLACommonPrefillMetadata.ContextChunk, + gather: Callable[[torch.Tensor, torch.Tensor], object], + cache_dtype: torch.dtype, +) -> torch.Tensor: + """Gather FP8 bytes, preserving the existing workspace row capacity.""" + assert chunk.local_context_lens_allranks is not None + assert workspace.is_contiguous() + rows, width = workspace.shape + fp8_workspace = workspace.view(cache_dtype).view(-1, width)[:rows] + local_rows = chunk.num_local_context_tokens + world_size = len(chunk.local_context_lens_allranks[0]) + offset = rows // (world_size + 1) + assert offset * (world_size + 1) == rows + assert local_rows <= offset + local = fp8_workspace[:local_rows] + gathered = fp8_workspace[offset : offset + world_size * local_rows] + assert chunk.padded_local_cu_seq_lens is not None + ops.cp_gather_cache( + cache.view(torch.uint8), + local.view(torch.uint8), + block_table, + chunk.padded_local_cu_seq_lens, + chunk.num_requests, + chunk.starts, + ) + gather(gathered.view(torch.uint8), local.view(torch.uint8)) + return gathered + + +def expand_context( + gathered: torch.Tensor, + scale: torch.Tensor, + row_indices: torch.Tensor, + cu_seq_lens: torch.Tensor, + weight: torch.Tensor, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse row reorganization, dequantization, projection and K/V packing.""" + k = torch.empty( + (row_indices.numel(), num_heads, qk_nope_head_dim + qk_rope_head_dim), + device=gathered.device, + dtype=weight.dtype, + ) + v = torch.empty( + (row_indices.numel(), num_heads, v_head_dim), + device=gathered.device, + dtype=weight.dtype, + ) + if row_indices.numel(): + rocm_aiter_ops.gather_kv_b_proj( + gathered.unsqueeze(1), + scale, + cu_seq_lens, + row_indices, + cu_seq_lens, + weight, + k, + v, + ) + return k, v From f47792af7fba76eb8e8f6a564b35cfad316d39f2 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 12 Oct 2026 18:33:26 -0700 Subject: [PATCH 455/952] [Model Loader] Make dummy weight init dtype-aware (#61134) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- tests/model_executor/test_weight_utils.py | 45 ++++++++++++++++++ .../model_loader/dummy_loader.py | 6 +-- .../model_loader/weight_utils.py | 47 ++++++++++++++++++- 3 files changed, 93 insertions(+), 5 deletions(-) diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 512f518a0561..12a8e3acac29 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/tests/model_executor/test_weight_utils.py @@ -6,13 +6,16 @@ import huggingface_hub.constants import pytest +import torch from huggingface_hub.utils import LocalEntryNotFoundError from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, drop_checkpoint_cache, + initialize_single_dummy_weight, maybe_remap_kv_scale_name, ) +from vllm.platforms import current_platform def test_download_weights_from_hf(): @@ -307,5 +310,47 @@ def test_drop_checkpoint_cache_releases_local_weight_files( fadvise.assert_called_once() +@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +def test_dummy_fp8_weight_is_not_degenerate(dtype): + """+-1e-3 underflows in FP8, so 8-bit floats default to +-1.""" + param = torch.empty(4096, dtype=dtype) + initialize_single_dummy_weight(param) + values = param.float() + assert values.isfinite().all() + assert values.abs().max() <= 1.0 + assert (values != 0).float().mean() > 0.99 + + +@pytest.mark.parametrize( + ("dtype", "name"), + [(torch.float8_e8m0fnu, "weight_scale_inv"), (torch.uint8, "w13_weight_scale")], +) +def test_dummy_e8m0_scale_is_fixed_exponent(dtype, name): + """E8M0 scales, including uint8-stored ones, get a finite 2**-7.""" + param = torch.empty(64, dtype=dtype) + initialize_single_dummy_weight(param, name=name) + assert (param.view(torch.uint8) == 120).all() + + +@pytest.mark.parametrize("dtype", [torch.uint8, torch.int8, torch.float4_e2m1fn_x2]) +def test_dummy_packed_weight_gets_random_bytes(dtype): + """Packed weights are filled deterministically instead of left uninitialized.""" + # Some backends can't fill packed FP4 directly, so view zeroed bytes. + params = [torch.zeros(4096, dtype=torch.uint8).view(dtype) for _ in range(2)] + for param in params: + initialize_single_dummy_weight(param, name="w13_weight") + raw = params[0].view(torch.uint8) + assert raw.unique().numel() > 250 + assert torch.equal(raw, params[1].view(torch.uint8)) + + +@pytest.mark.skipif(current_platform.is_rocm(), reason="ROCm zeroes integer params") +def test_dummy_wide_integer_param_is_untouched(): + """Wider integer params may be indices, so random values are unsafe.""" + param = torch.arange(16, dtype=torch.int32) + initialize_single_dummy_weight(param, name="g_idx") + assert torch.equal(param, torch.arange(16, dtype=torch.int32)) + + if __name__ == "__main__": test_download_weights_from_hf() diff --git a/vllm/model_executor/model_loader/dummy_loader.py b/vllm/model_executor/model_loader/dummy_loader.py index 0f263fe11709..ff732f6b222c 100644 --- a/vllm/model_executor/model_loader/dummy_loader.py +++ b/vllm/model_executor/model_loader/dummy_loader.py @@ -45,7 +45,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: # Exclude nonpersistent buffers, matching state_dict(). for name, tensor in get_layer_tensors(layer).items(): if name not in layer._non_persistent_buffers_set: - initialize_single_dummy_weight(tensor) + initialize_single_dummy_weight(tensor, name=name) def _process_online_quant_layer( self, @@ -55,8 +55,8 @@ def _process_online_quant_layer( """Materialize, apply dummy weights, and run quantization processing.""" materialize_layer(layer, info) - for tensor in get_layer_tensors(layer).values(): - initialize_single_dummy_weight(tensor) + for name, tensor in get_layer_tensors(layer).items(): + initialize_single_dummy_weight(tensor, name=name) for param in get_layer_tensors(layer).values(): param.weight_loader = _get_original_loader(param) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index ac1da8d6987b..0a834cd90ca5 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1521,13 +1521,30 @@ def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: return composed_loader +# E8M0 byte 120 encodes 2**-7: with packed FP4 (|x| <= 6) or the FP8 dummy +# range below, the dequantized weights stay small and finite. +_DUMMY_E8M0_SCALE_BYTE = 120 +# Default dummy range of 8-bit floats. +-1e-3 is below FP8 E4M3's smallest +# normal (2**-6), so it would mostly round to zero or subnormals. +_DUMMY_FP8_RANGE = 1.0 + + @torch.no_grad() def initialize_single_dummy_weight( param: torch.Tensor, - low: float = -1e-3, - high: float = 1e-3, + low: float | None = None, + high: float | None = None, seed: int = 1234, + name: str = "", ) -> None: + """Fill `param` with deterministic dummy values suited to its dtype. + + 8-bit integer and packed FP4 params get random bytes (every FP4 code is + finite), except uint8 params named like a scale, which hold E8M0 scales + and get a fixed exponent. Other integer params (e.g. indices) are left + untouched. Float params are sampled from `[low, high]`, which defaults to + +-1e-3, or +-1 for 8-bit floats. + """ if param.device.type == "meta": return # deferred to finalize_layerwise_processing (e.g. online quant) @@ -1535,6 +1552,25 @@ def initialize_single_dummy_weight( param.fill_(dummy_weight_value) return + is_e8m0 = param.dtype == torch.float8_e8m0fnu or ( + param.dtype == torch.uint8 and "scale" in name + ) + if is_e8m0: + param.view(torch.uint8).fill_(_DUMMY_E8M0_SCALE_BYTE) + return + + if ( + param.dtype in (torch.uint8, torch.int8, torch.float4_e2m1fn_x2) + and not current_platform.is_tpu() + ): + generator = torch.Generator(device=param.device) + generator.manual_seed(seed) + if param.dtype == torch.int8: + param.random_(-128, 128, generator=generator) + else: + param.view(torch.uint8).random_(0, 256, generator=generator) + return + if not torch.is_floating_point(param): if current_platform.is_rocm(): # On ROCm, integer params (e.g. GPTQ qweight/qzeros) are left @@ -1543,6 +1579,13 @@ def initialize_single_dummy_weight( param.zero_() return + if low is None or high is None: + dummy_range = 1e-3 + if torch.finfo(param.dtype).bits < 16: + dummy_range = _DUMMY_FP8_RANGE + low = -dummy_range if low is None else low + high = dummy_range if high is None else high + if current_platform.is_tpu(): generator = torch.Generator(device="cpu") generator.manual_seed(seed) From cc2a6eea1a9f92c1d6b47be8fe48adf127d4c67f Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:16:50 -0400 Subject: [PATCH 456/952] [ROCm][Perf][DSv4.1] Mono decode on every backbone layer (#61082) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Opus 5.5 --- tests/models/test_dsv41_mono_decode.py | 83 ++++++++++--- tests/models/test_dsv41_mono_layers.py | 55 +++++---- vllm/config/kernel.py | 6 + vllm/engine/arg_utils.py | 4 + vllm/envs.py | 7 -- vllm/models/deepseek_v41/amd/model.py | 12 +- vllm/models/deepseek_v41/amd/mono/README.md | 23 ++-- .../deepseek_v41/amd/mono/attention/back.py | 10 +- .../deepseek_v41/amd/mono/attention/front.py | 28 ++++- vllm/models/deepseek_v41/amd/mono/layer.py | 84 ++++++++++++- vllm/models/deepseek_v41/amd/mono/runner.py | 39 +++++- .../deepseek_v41/amd/mono/stages/seam.py | 24 ++-- vllm/models/deepseek_v41/amd/mono_decode.py | 115 +++++++++++++----- 13 files changed, 379 insertions(+), 111 deletions(-) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 684e72c00482..309d278bcb42 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -27,21 +27,25 @@ def forward(self, *args, **kwargs): self.calls.append((args, kwargs)) return ("mono outputs",) + def first_pre_mix(self, rows): + return torch.zeros(rows, 4) + def ffn(self, *args): self.calls.append((args, {})) return ("ffn outputs",) -def _layer(comp="comp"): +def _layer(comp="comp", indexer=None): cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8) # noqa: E731 attn = SimpleNamespace( swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()), compressed_cache_prefix=comp, topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32), compress_ratio=1, + indexer=indexer, _compressed_kv_cache=cache, ) - return SimpleNamespace(attn=attn) + return SimpleNamespace(attn=attn, engram=None, hc_mult=4) def _swa(**overrides): @@ -77,7 +81,12 @@ def run(monkeypatch): ``mode`` and ``metadata``; returns (result, runner calls).""" def _run( - mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs + mode=CUDAGraphMode.FULL, + metadata=None, + ffn_only=False, + entry=None, + layer=None, + **inputs, ): if metadata is None: metadata = { @@ -97,7 +106,7 @@ def _run( # has no compressed cache part, _ = args.pop("x"), args.pop("positions") return mono.ffn(_layer(comp=None), part, **args), runner.calls - return mono(_layer(), **args), runner.calls + return mono(layer or _layer(), **args), runner.calls return _run @@ -114,6 +123,23 @@ def test_decode_step_runs_mono(run, mode): 1, ) assert kwargs["comp_cache"].shape == (4, 128, md.RECORD) + assert kwargs["fold"] and "select" not in kwargs + + +def test_first_layer_runs_mono_from_the_embedding(run): + """The first layer's residual is the embedding broadcast to the four + streams, which K1 takes as it is (no owed post to fold).""" + out, calls = run(residual=None, post_mix=None, res_mix=None, pre_mix=None) + assert out == ("mono outputs",) and len(calls) == 1 + args, kwargs = calls[0] + assert args[2].shape == (M, 4, 5120) and not kwargs["fold"] + + +def test_index_source_selects_between_the_kernels(run): + """An index source hands the runner its compressor and indexer to run + between K1 and K2.""" + out, calls = run(layer=_layer(indexer=object())) + assert out == ("mono outputs",) and callable(calls[0][1]["select"]) def test_padded_decode_step_runs_mono(run): @@ -149,7 +175,9 @@ def test_padded_decode_step_runs_mono(run): ), dict(metadata={}), # a profile / dummy run: no attention metadata dict(**_inputs(md.MAX_ROWS + 6)), - dict(residual=None), # the first layer's seam broadcasts the embedding + # a later pipeline stage's first layer: its stream arrives without + # the mixes + dict(residual=None, x=torch.zeros(M, 4, 5120, dtype=torch.bfloat16)), ], ids=[ "piecewise", @@ -158,7 +186,7 @@ def test_padded_decode_step_runs_mono(run): "tokens-past-rows", "no-metadata", "rows", - "first-layer", + "stage-input", ], ) def test_other_steps_take_the_original_path(run, case): @@ -202,22 +230,41 @@ def _decoder_layer(engram=None, **attn): ) +def _indexer(): + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + return object.__new__(DeepseekV41RocmMxfp4Indexer) + + @pytest.mark.parametrize( "layer, path", [ (_decoder_layer(), "whole"), - (_decoder_layer(compress_ratio=0), "ffn"), - (_decoder_layer(engram=object()), "ffn"), - (_decoder_layer(compressor=object(), indexer=object()), "ffn"), + (_decoder_layer(compress_ratio=0), "whole"), + (_decoder_layer(engram=object()), "whole"), + (_decoder_layer(compressor=object(), indexer=_indexer()), "whole"), + (_decoder_layer(indexer=_indexer()), "whole"), + (_decoder_layer(compressor=object()), "ffn"), (_decoder_layer(indexer=object()), "ffn"), (_decoder_layer(kv_cache_dtype="auto"), "ffn"), (_decoder_layer(layer_id=40), None), ], - ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"], + ids=[ + "standard", + "window", + "engram", + "kv-source", + "index-source", + "compressor-only", + "other-indexer", + "kv", + "draft", + ], ) def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path): - """The whole-layer kernels serve the standard layers; every other backbone - layer keeps vLLM's attention and runs its FFN half as one launch.""" + """The whole-layer kernels serve every backbone layer of the checkpoint (an + index source's compressor and indexer run between them); a layer outside + them keeps vLLM's attention and runs its FFN half as one launch.""" _deployment(monkeypatch) mono = md.MonoDecodeLayer.create(layer, _config()) got = None if mono is None else "ffn" if mono.ffn_only else "whole" @@ -254,7 +301,7 @@ def test_other_steps_reduce_wo_b_themselves(run, case): assert out is None and not calls -def _config(**parallel): +def _config(monokernel=True, **parallel): parallel = ( dict(enable_expert_parallel=False, enable_eplb=False, data_parallel_size=1) | parallel @@ -262,6 +309,7 @@ def _config(**parallel): return SimpleNamespace( model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), parallel_config=SimpleNamespace(**parallel), + kernel_config=SimpleNamespace(monokernel=monokernel), ) @@ -269,7 +317,6 @@ def _deployment(monkeypatch, cdna=4, cus=256): """A deployment the kernels take (the runner module stubbed: no FlyDSL).""" import vllm.platforms.rocm as rocm - monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1") monkeypatch.delenv("VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4", raising=False) monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna) monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2) @@ -280,6 +327,14 @@ def _deployment(monkeypatch, cdna=4, cus=256): monkeypatch.setitem(sys.modules, name, runner) +def test_off_without_the_flag(monkeypatch): + """No --monokernel: no mono layer, whatever the deployment.""" + _deployment(monkeypatch) + assert ( + md.MonoDecodeLayer.create(_decoder_layer(), _config(monokernel=False)) is None + ) + + @pytest.mark.parametrize("cdna", [3, 5]) def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna): _deployment(monkeypatch, cdna=cdna) diff --git a/tests/models/test_dsv41_mono_layers.py b/tests/models/test_dsv41_mono_layers.py index 580462a82394..a924cfda50c3 100644 --- a/tests/models/test_dsv41_mono_layers.py +++ b/tests/models/test_dsv41_mono_layers.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """The DeepSeek-V4.1 mono decode path through the real model: a five-layer -DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``VLLM_ROCM_MONO_DECODE`` on +DeepSeek-V4.1-Flash at TP2 and TP4 on CDNA4, served with ``--monokernel`` on and off on the same random weights: every decoder layer's outputs at every decode step must match. Unlike ``test_dsv41_mono_numerics.py``, which checks the kernels against vLLM's ops one stage at a time, this runs @@ -10,12 +10,13 @@ the indexer's top-k -- so a change to the decoder layer that the kernels do not follow fails here. -The five layers hold both kinds the mono path takes: -- 0: sliding window only -> the FFN launch -- 1: ratio-2 KV and index source -> the FFN launch -- 2: ratio-2 consumer -> the whole layer (K1 + K2) -- 3: ratio-1 KV, index and candidate source -> the FFN launch -- 4: ratio-1 consumer -> the whole layer +Every layer takes the whole-layer path (K1 + K2); the five hold each variant: +- 0: the first layer, sliding window only (K1 takes the embedding settled) +- 1: ratio-2 KV and index source (vLLM's compressor and indexer between K1 + and K2) +- 2: ratio-2 consumer with an Engram (K1 takes the residual settled) +- 3: ratio-1 KV, index and candidate source +- 4: ratio-1 index source over 3's candidates and KV """ import zlib @@ -34,13 +35,17 @@ # one entry a layer, the three MTP layers' included "compress_ratios": [0, 2, 2, 1, 1, 0, 0, 0], "kv_source_layer_ids": [1, 3], - "index_source_layer_ids": [1, 3], + # 4: an indexer of its own over 3's candidates (a REINDEX layer) + "index_source_layer_ids": [1, 3, 4], "candidate_source_layer_id": 3, - "engram_layer_ids": [], - "engram_num_embeddings": [], + # an Engram on a whole layer (K1 then takes its residual settled), its + # tables shrunk: 24 prime-sized ranges from 1000 + "engram_layer_ids": [2], + "engram_num_embeddings": [32768], + "engram_vocab_size": 1000, "dspark_target_layer_ids": [2, 3, 4], } -WHOLE, FFN = (2, 4), (0, 1, 3) +WHOLE, FFN = (0, 1, 2, 3, 4), () STEPS = 5 # decode steps after the prefill's token @@ -227,7 +232,6 @@ def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): from vllm import SamplingParams from vllm.inputs import TokensPrompt - monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1" if mono else "0") monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") with vllm_runner( @@ -237,6 +241,7 @@ def _decode(vllm_runner, monkeypatch, tp: int, mono: bool, prompts, path: str): hf_overrides=TRUNCATED, worker_extension_cls=f"{__name__}.MonoProbe", moe_backend="aiter", + monokernel=mono, language_model_only=True, tokenizer_mode="deepseek_v41", attention_config={"indexer_kv_dtype": "mxfp4", "indexer_sparse_logits": True}, @@ -300,15 +305,16 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): same.add((len(prompt), len(prompt) + step)) # each layer's outputs for those tokens, mono against vLLM's path, token by - # token. Layer 0 sees the same inputs on both paths at a request's first - # decode token: there the bound is test_dsv41_mono_numerics' (about 3x the - # kernels' 0.4%); its later tokens attend over KV each path wrote, 3%. Each - # later layer takes the previous one's slightly different output, which - # random weights amplify (2-4% by layer 4): 8%. A token whose top-6 routing - # is a near-tie on vLLM's path (6th / 7th score within 5e-3: a 0.4% input - # difference moves scores near 0.8 by about 3e-3) may pick another expert; - # once it has, its residual stream differs and the later layers are exempt - # for it. The residual stream: 3% (0.3% seen). A layer the kernels no + # token. Layer 0 sees the same inputs on both paths, but its FFN input + # differs by up to test_dsv41_mono_numerics' 1.2%, which random weights + # amplify through the MoE: 3%. Each later layer takes the previous one's + # slightly different output, which random weights amplify (2-4% by layer + # 4): 8%. A token whose top-6 routing is a near-tie on vLLM's path may pick + # another expert: a 0.4% input difference moves scores near 0.8 by about + # 3e-3, so 6th / 7th scores within 5e-3 at layer 0 and, past it, where the + # input carries the earlier layers' differences (0.4-0.7%), within 1e-2. + # Once it has, its residual stream differs and the later layers are exempt + # for it. The residual stream: 3% (0.5% seen). A layer the kernels no # longer follow is off by far more, on every token. for rank in range(tp): got, _ = torch.load(f"{tmp_path}/on.rank{rank}") @@ -317,8 +323,6 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): same & {k[1:] for k in got} & {k[1:] for k in want} ) assert len(keys) >= len(prompts), keys - # the bound that sees a 3% error: a request's first decode token - assert any(p == n for n, p in keys), keys flipped: set = set() for i in sorted(WHOLE + FFN): for k in keys: @@ -327,9 +331,10 @@ def test_mono_decode_matches_the_model(vllm_runner, monkeypatch, tmp_path, tp): (x, res), (x_ref, res_ref) = got[(i, *k)], want[(i, *k)] err = _rel(res, res_ref) assert err < 0.03, (rank, i, k, "residual", err) - bound = (0.015 if k[1] == k[0] else 0.03) if i == 0 else 0.08 + bound = 0.03 if i == 0 else 0.08 err = _rel(x, x_ref) if err >= bound: - assert gaps[(i, *k)] < 5e-3, (rank, i, k, "x", err) + tie = 5e-3 if i == 0 else 1e-2 + assert gaps[(i, *k)] < tie, (rank, i, k, "x", err) flipped.add(k) assert len(flipped) <= len(keys) // 4, (rank, sorted(flipped)) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index fa3ecf5df95d..5c4b877fd795 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -291,6 +291,12 @@ class KernelConfig: enable_jit_warmup: bool = True """If True, run JIT compile warmup during kernel warmup.""" + monokernel: bool = False + """If True, run the eligible layers of decode steps on the model's + MonoKernel: a persistent kernel that runs most or all of one decoder layer + in a single launch (DeepSeek-V4.1 on MI355X). A deployment that cannot run + it is refused at startup.""" + moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ea628467e0b3..86199e5f007b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -758,6 +758,7 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) + monokernel: bool = get_field(KernelConfig, "monokernel") worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1802,6 +1803,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) + kernel_group.add_argument("--monokernel", **kernel_kwargs["monokernel"]) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -2754,6 +2756,8 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune + if self.monokernel: + kernel_config.monokernel = True if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": diff --git a/vllm/envs.py b/vllm/envs.py index a2b46030772f..bb1c03a9dabd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -156,7 +156,6 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False - VLLM_ROCM_MONO_DECODE: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True @@ -1398,12 +1397,6 @@ def _resolve_rust_cli_path() -> str | None: os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower() in ("true", "1") ), - # Run the eligible layers of decode steps on the model's mono decode - # kernels: persistent launches a layer, the TP all-reduces in-kernel - # (DeepSeek-V4.1 on CDNA4). By default is disabled. - "VLLM_ROCM_MONO_DECODE": lambda: ( - os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1") - ), # Whether to use aiter triton kernels for gemm ops. # By default is enabled. "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 5de173a8706c..4c5f15858749 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -282,7 +282,17 @@ def forward( engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: if self.mono is not None: - out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix) + out = self.mono( + self, + x, + positions, + residual, + post_mix, + res_mix, + pre_mix, + engram_hashes, + engram_mask, + ) if out is not None: return out # Layer 0's attention seam projects the 2-D embedding with the folded diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md index 88e40d73dd0a..462f7bdecd28 100644 --- a/vllm/models/deepseek_v41/amd/mono/README.md +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -2,9 +2,10 @@ One fused decoder layer for vLLM's decode steps on MI355X: the attention seam, the attention, the attention all-reduce, the FFN seam, the MoE and the MoE -all-reduce, in two persistent launches a layer. The other backbone layers keep -vLLM's attention and run the rest -- the attention all-reduce, the FFN seam, -the MoE and its all-reduce -- as one launch (the FFN launch). ATOM's V4.1 mono decode +all-reduce, in two persistent launches a layer. A layer outside the whole-layer +kernels keeps vLLM's attention and runs the rest -- the attention all-reduce, +the FFN seam, the MoE and its all-reduce -- as one launch (the FFN launch). +ATOM's V4.1 mono decode (ROCm/ATOM `atom/models/deepseek_v41/mono`) is the reference design: the seam and MoE stages (`stages/`) and the framework under them (`common/`) are adapted from it, each file naming its origin; the attention stages (`attention/`) are @@ -12,13 +13,15 @@ new, reading vLLM's KV records and weights as loaded. ## Scope -- Layers: the whole layer on every backbone layer with the standard seam and no - compressor / indexer: 3-7, 9-13, 15-19 (compress ratio 2) and 21-23, 25-27, - 29-31, 33-35, 37-39 (ratio 1), 30 of 40. The FFN launch on the other 10: - layer 0 (embedding-broadcast seam), layer 1 (Engram at the seam) and the - compressor / indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36, - whose wo_b leaves its TP reduce to it. The DSpark draft layers keep vLLM's - path. The paths interleave freely (same tensors at the layer boundary). +- Layers: the whole layer on all 40 backbone layers. K1 takes the residual + settled, without folding the previous output in, on layer 0 (the embedding + broadcast to the four streams) and after an Engram write (layers 1 and 14, + vLLM's Engram running before K1). On the index sources (2, 8, 14, 20 with + their compressor; 24, 28, 32, 36) K1 also hands out its normed input and bf16 + qr, vLLM's compressor and indexer run between K1 and K2, and K2 builds the + key table from that top-k. The FFN launch serves a layer outside these (an + MXFP8 KV cache, another indexer); the DSpark draft layers keep vLLM's path. + The paths interleave freely (same tensors at the layer boundary). - Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the routing bias is `e_score_correction_bias`). diff --git a/vllm/models/deepseek_v41/amd/mono/attention/back.py b/vllm/models/deepseek_v41/amd/mono/attention/back.py index 4526f8eda7ee..5cd4793be243 100644 --- a/vllm/models/deepseek_v41/amd/mono/attention/back.py +++ b/vllm/models/deepseek_v41/amd/mono/attention/back.py @@ -151,7 +151,9 @@ def key_rec(c, t, k): table: slot | 1 << 31 for a compressed row, a window slot otherwise, -1 for none (it reads a zero record: its scores are masked, its values zeros).""" a, ratio = c["args"], c["ratio"] - val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1)) + # a table this launch wrote (``krdy``) is read past the CU's cache + cm = CM_DEV if const_expr("krdy" in c) else 0 + val = ld_i32(a["kt"], t * KEYS + fx.min(fx.max(k, 0), KEYS - 1), cm) ok = val != fx.Int32(-1) slot = ok.select(val & fx.Int32(0x7FFFFFFF), fx.Int32(0)) if const_expr(ratio > 0): @@ -209,9 +211,11 @@ def stage_split(c, unit): g = lane // 16 r16 = lane % 16 k0 = ch * ck + if const_expr("krdy" in c): + await_flags(c, c["krdy"], 1, lambda u: t) # the key count, this lane's key and q: independent loads, out together - # (K1's key table holds -1 past a token's keys) - kv_len = ld_i32(a["klen"], t) + # (the key table holds -1 past a token's keys) + kv_len = ld_i32(a["klen"], t, CM_DEV if const_expr("krdy" in c) else 0) key = k0 + kt * 16 + r16 data, sc, okv = key_rec(c, t, key) qrow = d.heads * HEAD_DIM // 8 diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py index 3ac9e3fe9b10..23c009739dc2 100644 --- a/vllm/models/deepseek_v41/amd/mono/attention/front.py +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -188,6 +188,8 @@ def stage_kt(c, t): ntopk.""" tid = c["tid"] a, ratio = c["args"], c["ratio"] + # device scope where this launch reads the table (K2's) + cm = CM_DEV if const_expr(c.get("dev", False)) else 0 slot = ld_i32(a["slot"], 2 * t) live = slot >= 0 nswa = live.select(ld_i32(a["swa_lens"], t), fx.Int32(0)) @@ -212,10 +214,10 @@ def stage_kt(c, t): c_val = (local >= 0).select(c_slot, fx.Int32(-1)) val = (k < ntopk).select(c_val, val) if k < KEYS: - bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k) + bo.buffer_store(val, rsrc(a["kt"]), t * KEYS + k, cache_modifier=cm) if tid == 0: - bo.buffer_store(kv_len, rsrc(a["klen"]), t) - bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t) + bo.buffer_store(kv_len, rsrc(a["klen"]), t, cache_modifier=cm) + bo.buffer_store(ntopk, rsrc(a["klen"]), c["S"] + t, cache_modifier=cm) # ---------------------------------------------------------------- wqkv_a @@ -357,6 +359,20 @@ def stage_qkv(c, t): amax = fx.max(amax, abs(y)) code = mx_code(butterfly(amax, (1, 2), fx.max)) mul = mx_mul(code) + if const_expr(c.get("index", False)): # noqa: SIM102 + # the bf16 qr, vLLM's indexer's input + if mine & (tt * 8 < Q_RANK): + words = [ + fx.Vector.from_elements([ys[2 * i], ys[2 * i + 1]], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Int32)[0] + for i in range(4) + ] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(a["qr_out"]), + t * (Q_RANK // 2) + 4 * tt, + ) if mine & (tt * 8 < Q_RANK): w0 = fp8_pack4(*[clamp_fp8(ys[i] * mul) for i in range(4)]) w1 = fp8_pack4(*[clamp_fp8(ys[4 + i] * mul) for i in range(4)]) @@ -597,8 +613,10 @@ def run_front(c, key, bid, x8_given=False, before_wqkv=None, after_wqkv=None): u = free + n_free * tid if u < xq_units: c["mb"].put(c["xrdy"], u, fx.Int32(1)) - for t in range(is_free.select(free, fx.Int32(s)), s, n_free): - stage_kt(c, t) + if const_expr(not c.get("index", False)): + # an index source's top-k comes after this launch: K2 builds its table + for t in range(is_free.select(free, fx.Int32(s)), s, n_free): + stage_kt(c, t) stamp(c, 1) poss = [wqb_pos(c, t0, n) for t0, n in tiles] diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py index c1be9e3428ae..dcd39ae2b754 100644 --- a/vllm/models/deepseek_v41/amd/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -58,6 +58,12 @@ class MonoBuild: tp: int ratio: int # the layer's compress ratio (its attention's key sources) timeline: bool = False + # K1's seam folds the previous layer's output into the residual; off for a + # residual that arrives settled (the first layer, after an Engram) + fold: bool = True + # the layer selects its own top-k (an index source): K1 hands its normed + # input and qr out for vLLM's indexer, K2 builds the key table + index: bool = False def _moe_key(key: MonoBuild) -> MoeBuild: @@ -100,6 +106,7 @@ def place(regions): lo = min(o for o, _ in mr.values()) place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) place({"normed": (0, s * HIDDEN * 2), "xrdy_moe": (s * HIDDEN * 2 + 256, s * 8)}) + place({"krdy": (0, s * 8)}) return out @@ -175,7 +182,15 @@ class K1Lds: seam: SeamLds front: FrontLds # type: ignore[valid-type] - name = kernel_symbol("dsv41_mono_k1", s=s, tp=tp, r=key.ratio, tl=key.timeline) + name = kernel_symbol( + "dsv41_mono_k1", + s=s, + tp=tp, + r=key.ratio, + tl=key.timeline, + f=key.fold, + i=key.index, + ) keyed = key_tuple(key, SOURCES) @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) @@ -215,6 +230,8 @@ def k1( t2r: Int64, kt: Int64, klen: Int64, + normed_out: Int64, + qr_out: Int64, scratch: Int64, epoch: Int64, tl: Int64, @@ -250,9 +267,11 @@ def k1( "t2r": t2r, "kt": kt, "klen": klen, + "qr_out": qr_out, } flayout = {n: layout[n] for n in front_scratch(s, Dims(tp))} c = mb_front.front_context(fkey, lds.front.peek(), fargs, scratch, tag, flayout) + c["index"] = key.index amb = Mailbox(tag) tid = fx.thread_idx.x ca = { @@ -284,6 +303,7 @@ def k1( "pre_out": pre_out, }, "d": StageDims(tp), + "fold": key.fold, } for region in ("lin", "pmix"): ca[region] = sreg(scratch, layout[region][0], region) @@ -295,9 +315,14 @@ def seam_slice(): seam.stage_slice(ca, task) gpu.barrier() + def store(*v): + store_x8(c, *v) + if const_expr(key.index): + store_normed(normed_out, *v) + def seam_rest(): for t in range(first_task(bid, NORM0), s, BLOCKS): - seam.stage_norm(ca, t, lambda *v: store_x8(c, *v)) + seam.stage_norm(ca, t, store) publish_x8(c, t) for t in range(first_task(bid, GATE0), s, BLOCKS): seam.stage_gate(ca, t) @@ -344,6 +369,8 @@ def launch( t2r: Int64, kt: Int64, klen: Int64, + normed_out: Int64, + qr_out: Int64, scratch: Int64, epoch: Int64, tl: Int64, @@ -386,6 +413,8 @@ def launch( t2r, kt, klen, + normed_out, + qr_out, scratch, epoch, tl, @@ -544,7 +573,9 @@ class K2Lds: seam: SeamLds # type: ignore[valid-type] moe: MoeLds # type: ignore[valid-type] - name = kernel_symbol("dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline) + name = kernel_symbol( + "dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline, i=key.index + ) keyed = key_tuple(key, SOURCES) @flyc.kernel(name=name, known_block_size=[THREADS, 1, 1]) @@ -558,6 +589,13 @@ def k2( comp_block: Int32, kt: Int64, klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, pos: Int64, sink: Int64, qk_scale: Int32, @@ -642,6 +680,32 @@ def push(t, col, v0, v1): if const_expr(key.timeline): c["tl"], c["tl_points"] = tl, mb_back.BACK_POINTS mb_back.epoch_begin(c, epoch, ep) + if const_expr(key.index): + # the key table, from the top-k vLLM's indexer selected after K1, + # on the last CTAs: each split task waits on its token's + ckt = { + "S": s, + "tid": c["tid"], + "dev": True, + "ratio": key.ratio, + "args": { + "slot": slot, + "swa_idx": swa_idx, + "swa_lens": swa_lens, + "pos": pos, + "t2r": t2r, + "topk": topk, + "comp_block": comp_block, + "comp_bt": comp_bt, + "bt_stride": bt_stride, + "kt": kt, + "klen": klen, + }, + } + c["krdy"] = sreg(scratch, layout["krdy"][0], "krdy") + for t in range(first_task(bid, BLOCKS - s), s, BLOCKS): + mb_front.stage_kt(ckt, t) + mb_back.publish(c, c["krdy"], t) mb_back.run_back(c, bkey, bid) gpu.barrier() @@ -685,6 +749,13 @@ def launch( comp_block: Int32, kt: Int64, klen: Int64, + slot: Int64, + swa_idx: Int64, + swa_lens: Int64, + t2r: Int64, + topk: Int64, + comp_bt: Int64, + bt_stride: Int32, pos: Int64, sink: Int64, qk_scale: Int32, @@ -736,6 +807,13 @@ def launch( comp_block, kt, klen, + slot, + swa_idx, + swa_lens, + t2r, + topk, + comp_bt, + bt_stride, pos, sink, qk_scale, diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py index eb379902754f..ef9a1978e422 100644 --- a/vllm/models/deepseek_v41/amd/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -16,12 +16,13 @@ from __future__ import annotations import os +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path import torch -from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Dims +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims from .common.plan import BLOCKS from .layer import ( EPOCH_WORDS, @@ -170,6 +171,12 @@ def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda self.klen = torch.zeros(2 * MAX_TOKENS, dtype=torch.int32, device=dev) self._dummy = torch.zeros(4, dtype=torch.int32, device=dev) self._zero_rec = torch.zeros(1024, dtype=torch.uint8, device=dev) + # the first layer's incoming pre-mix: stream 0 of the broadcast embedding + self._first_pre = torch.zeros(MAX_TOKENS, HC, dtype=torch.float32, device=dev) + self._first_pre[:, 0] = 1.0 + # an index source's K1 outputs, vLLM's indexer's inputs + self.normed = torch.zeros(MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device=dev) + self.qr = torch.zeros(MAX_TOKENS, Q_RANK, dtype=torch.bfloat16, device=dev) # the attention seam's outputs, K1 -> K2 self.res_mid = torch.zeros( MAX_TOKENS, HC, HIDDEN, dtype=torch.bfloat16, device=dev @@ -203,14 +210,17 @@ def scratch(self, tokens: int) -> torch.Tensor: self._scratch[tokens] = buf return buf + def first_pre_mix(self, tokens: int) -> torch.Tensor: + return self._first_pre[:tokens] + @staticmethod def supports(tokens: int) -> bool: return 1 <= tokens <= MAX_TOKENS - def kernels(self, tokens: int, ratio: int): - key = (tokens, ratio) + def kernels(self, tokens: int, ratio: int, fold: bool = True, index: bool = False): + key = (tokens, ratio, fold, index) if key not in self._kernels: - b = MonoBuild(tokens, self.tp, ratio) + b = MonoBuild(tokens, self.tp, ratio, fold=fold, index=index) self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) return self._kernels[key] @@ -249,15 +259,21 @@ def forward( comp_cache: torch.Tensor | None = None, comp_block_table: torch.Tensor | None = None, outs: tuple | None = None, + fold: bool = True, + select: Callable[[torch.Tensor, torch.Tensor], None] | None = None, ): """-> (out, residual, post_mix, res_mix, pre_mix): the layer's outputs - (``outs``: buffers to write them into).""" + (``outs``: buffers to write them into). ``select(normed, qr)``: an index + source's top-k selection into ``topk_indices``, run between the + kernels on K1's normed input and qr.""" M = x.shape[0] assert self.supports(M), M a = w.attn assert a is not None ratio = a.ratio - k1, k2 = self.kernels(M, ratio) + index = select is not None + assert not index or ratio + k1, k2 = self.kernels(M, ratio, fold, index) st = torch.cuda.current_stream() if outs is None: outs = self._outs(M, residual) @@ -324,11 +340,15 @@ def forward( token_to_req.data_ptr(), self.kt.data_ptr(), self.klen.data_ptr(), + self.normed.data_ptr(), + self.qr.data_ptr(), self.scratch(M).data_ptr(), self.epoch.data_ptr(), 0, stream=st, ) + if select is not None: + select(self.normed[:M], self.qr[:M]) k2( self.q.data_ptr(), swa_cache.data_ptr(), @@ -339,6 +359,13 @@ def forward( comp_block, self.kt.data_ptr(), self.klen.data_ptr(), + slot_mapping.data_ptr(), + swa_indices.data_ptr(), + swa_lens.data_ptr(), + token_to_req.data_ptr(), + topk.data_ptr(), + bt.data_ptr(), + bt_stride, positions.data_ptr(), a.attn_sink.data_ptr(), self._qk_scale, diff --git a/vllm/models/deepseek_v41/amd/mono/stages/seam.py b/vllm/models/deepseek_v41/amd/mono/stages/seam.py index 82abee9f46fa..ea73e2ef5435 100644 --- a/vllm/models/deepseek_v41/amd/mono/stages/seam.py +++ b/vllm/models/deepseek_v41/amd/mono/stages/seam.py @@ -141,6 +141,8 @@ def slice_loads(c, c0, e): """Element e's loads of ``stage_slice`` (e clamped into the step's).""" a = c["args"] t, o, col = slice_at(c0, e) + if const_expr(not c.get("fold", True)): + return (ld_bf(a["res_in"], (t * HC + o) * HIDDEN + col),) r = [ld_bf(a["res_in"], (t * HC + h) * HIDDEN + col) for h in range(HC)] cm = [ld_f32(a["comb_in"], (t * HC + h) * HC + o) for h in range(HC)] return r, cm, ld_f32(a["post_in"], t * HC + o) @@ -150,15 +152,19 @@ def slice_row(c, c0, e, got): """Element e of ``stage_slice``'s R_new from ``slice_loads``' values.""" a = c["args"] t, o, col = slice_at(c0, e) - r, cm, post = got - # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) - # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) - v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) - for h in range_constexpr(1, HC - 1): - v = v + cm[h] * r[h] - # bf16 R_new is what the collapse, the sums of squares and the - # mix projection all read - v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) + if const_expr(not c.get("fold", True)): + # a settled residual (no owed post): stream o as it arrives + v = got[0] + else: + r, cm, post = got + # aiter's mhc_fused_post_pre_delayed_rmsnorm: ((post x + c0 r0) + # + c1 r1) + c2 r2, each product rounded, then fma(c3, r3, ..) + v = post * pend_value(c, t, col, e % COLS) + (r[0] * cm[0]) + for h in range_constexpr(1, HC - 1): + v = v + cm[h] * r[h] + # bf16 R_new is what the collapse, the sums of squares and the + # mix projection all read + v = bf16_round(fma(cm[HC - 1], r[HC - 1], v)) bo.buffer_store(v.to(fx.BFloat16), rsrc(a["res_out"]), (t * HC + o) * HIDDEN + col) fx.ptr_store(v, c["rl"] + (t * KT + o * COLS + e % COLS)) diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py index b404a42de37a..ddeb6b26314f 100644 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``VLLM_ROCM_MONO_DECODE=1``). +"""DeepSeek-V4.1 mono decode layer on ROCm CDNA4 (``--monokernel``). A decode step's backbone layer -- the attention seam, the attention, its TP all-reduce, the FFN seam, the MoE and its all-reduce -- runs as two persistent FlyDSL launches (``mono``, this directory) on the layer's loaded -weights, the TP reductions inside the kernels. A layer runs whole when it has -the standard seam and no compressor or indexer of its own (compress ratio 1 or -2) over ``fp8_ds_mla`` records. Any other backbone layer (the first, the Engram -layers, the compressor / indexer layers) keeps vLLM's attention seam and -attention and runs the rest as one launch, the FFN launch, from wo_b's +weights, the TP reductions inside the kernels. Every backbone layer of the +checkpoint runs whole (compress ratio 0, 1 or 2 over ``fp8_ds_mla`` records): +the first layer and the Engram layers hand K1 a settled residual, and an index +source runs vLLM's compressor and indexer between the two launches. A layer +outside that (an MXFP8 KV cache, another indexer) keeps vLLM's attention seam +and attention and runs the rest as one launch, the FFN launch, from wo_b's unreduced output. A step takes the kernels when it is decode only, at most ``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); a whole layer also needs causal SWA windows. Anything else runs the layer as @@ -20,7 +21,7 @@ other's pushes): the conditions read only metadata every rank holds alike. """ -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast import torch @@ -62,7 +63,7 @@ def _mono_runner(device: torch.device): # its peer memory is allocated and exchanged eagerly; vLLM warms # every graph up eagerly before capturing it raise RuntimeError( - "VLLM_ROCM_MONO_DECODE was first reached inside a CUDA graph " + "--monokernel was first reached inside a CUDA graph " "capture, before an eager decode step." ) from vllm.models.deepseek_v41.amd.mono.runner import MAX_TOKENS, DSV41MonoLayer @@ -113,6 +114,35 @@ def _unsupported_moe(layer: "DeepseekV4DecoderLayer", vllm_config) -> str | None return None +def _select( + layer: "DeepseekV4DecoderLayer", + positions: torch.Tensor, + normed: torch.Tensor, + qr: torch.Tensor, +) -> None: + """An index source's compressor and top-k selection, from K1's normed input + and bf16 qr: ``_forward_csa2_full`` / ``_forward_csa2_reindex`` and the + indexer op, serial.""" + attn = layer.attn + comp, indexer = attn.compressor, attn.indexer + assert indexer is not None + if comp is not None: + kv_score = torch.mm( + normed, comp.fused_wkv_wgate.weight.T, out_dtype=torch.float32 + ) + latent = comp(kv_score, positions) + comp.insert_cache(latent, positions, attn.rotary_emb) + indexer._produce_k(latent, positions, attn.indexer_rotary_emb) + weights, _ = indexer.weights_proj(normed) + q, q_scale, weights = indexer.forward_q( + qr, None, weights, positions, attn.indexer_rotary_emb + ) + assert weights is not None + indexer.indexer_op( + normed, (q, q_scale) if q_scale is not None else q, None, weights + ) + + class MonoDecodeLayer: """One decoder layer's mono path (``create``: None when the layer is not eligible): the whole layer (``__call__``) or, with ``ffn_only``, the FFN @@ -127,7 +157,7 @@ def __init__(self, ffn_only: bool = False) -> None: def create( layer: "DeepseekV4DecoderLayer", vllm_config ) -> "MonoDecodeLayer | None": - if not envs.VLLM_ROCM_MONO_DECODE: + if not vllm_config.kernel_config.monokernel: return None from vllm.platforms.rocm import get_cdna_version @@ -150,7 +180,7 @@ def create( if cus < BLOCKS: why = f"needs {BLOCKS} compute units, the GPU has {cus}" if why is not None: - raise ValueError(f"VLLM_ROCM_MONO_DECODE {why}.") + raise ValueError(f"--monokernel {why}.") # the layer: the kernels serve the backbone's seams and MoE config = vllm_config.model_config.hf_config attn, ffn = layer.attn, layer.ffn @@ -174,12 +204,17 @@ def create( if why is not None: logger.info_once("DSv4.1 mono decode off for some layers: %s", why) return None - # the whole layer: the kernels' attention takes the standard layers only - if layer.engram is not None: - why = "Engram" - elif attn.compressor is not None or attn.indexer is not None: - why = "a compressor / indexer" - elif attn.compress_ratio not in (1, 2): + # the whole layer: the kernels' attention takes the standard layers; an + # index source's compressor and indexer run between its two launches + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41RocmMxfp4Indexer + + if attn.compressor is not None and attn.indexer is None: + why = "a compressor without an indexer" + elif attn.indexer is not None and not isinstance( + attn.indexer, DeepseekV41RocmMxfp4Indexer + ): + why = f"a {type(attn.indexer).__name__}" + elif attn.compress_ratio not in (0, 1, 2): why = f"compress ratio {attn.compress_ratio}" elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": why = f"a {attn.kv_cache_dtype} KV cache" @@ -297,20 +332,23 @@ def __call__( post_mix: torch.Tensor | None, res_mix: torch.Tensor | None, pre_mix: torch.Tensor | None, + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...] | None: """The layer's outputs (x, residual, post_mix, res_mix, pre_mix) by the mono kernels, or None when this step takes the original path (always for an ``ffn_only`` layer).""" if self.ffn_only: return None - if residual is None or x.dim() != 2: - return self._skip("first layer") + first = residual is None + if first and x.dim() != 2: + return self._skip("first layer of a pipeline stage") why, swa, comp = self._decode_step(layer, x.shape[0]) attn = layer.attn - tensors = (x, residual, post_mix, res_mix, pre_mix) + tensors = (x,) if first else (x, residual, post_mix, res_mix, pre_mix) if why is None: # the kernels' attention: causal SWA windows, the compressed cache - if comp is None: + if comp is None and attn.compress_ratio: why = "layer metadata missing" elif swa.decode_swa_width != SWA_WIDTH: why = f"SWA width {swa.decode_swa_width}" @@ -324,12 +362,37 @@ def __call__( why = "non-contiguous layer inputs" if why is not None: return self._skip(why) - assert swa is not None and comp is not None + assert swa is not None runner = _mono_runner(x.device) logger.info_once( "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS ) - out = runner.forward( + # K1's seam folds the owed post into the residual, except where the + # residual arrives settled: the embedding, or after the Engram write + fold = True + if first: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + pre_mix = runner.first_pre_mix(x.shape[0]) + post_mix = res_mix = pre_mix # unread without the fold + fold = False + elif layer.engram is not None and engram_hashes is not None: + residual = layer.mhc_post(x, residual, post_mix, res_mix) + hashes = engram_hashes[:, layer.engram.layer_hash_index] + residual = layer.engram(residual, hashes, engram_mask) + fold = False + kw: dict[str, Any] = {} + if attn.indexer is not None: + kw["select"] = lambda normed, qr: _select(layer, positions, normed, qr) + if attn.compress_ratio: + assert comp is not None + kw.update( + topk_indices=attn.topk_indices_buffer, + comp_cache=_record_view( + attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio + ), + comp_block_table=comp.block_table, + ) + return runner.forward( self.weights(layer), x, residual, @@ -342,13 +405,9 @@ def __call__( swa.decode_swa_indices, swa.decode_swa_lens, swa.token_to_req_indices, - topk_indices=attn.topk_indices_buffer, - comp_cache=_record_view( - attn._compressed_kv_cache(), comp.block_size // attn.compress_ratio - ), - comp_block_table=comp.block_table, + fold=fold, + **kw, ) - return out def ffn( self, From 69d2b6f82dff3bdd78ec8658496e59c227b569d2 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:18:25 +0800 Subject: [PATCH 457/952] [CI][XPU] Standardize test job labels for intel jobs (#51755) Signed-off-by: zengxian Signed-off-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> Co-authored-by: Kunshang Ji --- .../intel_xpu_ci/test-intel.yaml | 14 +++++------ .../intel_jobs/basic_correctness_intel.yaml | 4 ++-- .buildkite/intel_jobs/benchmarks_intel.yaml | 4 ++-- .buildkite/intel_jobs/engine_intel.yaml | 6 ++--- .buildkite/intel_jobs/entrypoints_intel.yaml | 2 +- .../intel_jobs/expert_parallelism_intel.yaml | 4 ++-- .buildkite/intel_jobs/kernels_intel.yaml | 4 ++-- .buildkite/intel_jobs/lm_eval_intel.yaml | 6 ++--- .buildkite/intel_jobs/lora_intel.yaml | 14 +++++------ .buildkite/intel_jobs/misc_intel.yaml | 24 +++++++++---------- .../intel_jobs/model_executor_intel.yaml | 4 ++-- .../intel_jobs/models_distributed_intel.yaml | 4 ++-- .../intel_jobs/models_multimodal_intel.yaml | 14 +++++------ ...ntization.yaml => quantization_intel.yaml} | 4 ++-- .buildkite/intel_jobs/samplers_intel.yaml | 4 ++-- .buildkite/intel_jobs/test-intel.yaml | 24 +++++++++---------- 16 files changed, 68 insertions(+), 68 deletions(-) rename .buildkite/intel_jobs/{quantization.yaml => quantization_intel.yaml} (95%) diff --git a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml index a4a5ccbabce8..ed98d4408512 100644 --- a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml +++ b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml @@ -15,7 +15,7 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 - - label: "XPU example Test" + - label: ":intel: (BMG) Example Test" key: xpu-example-test depends_on: - image-build-xpu @@ -37,7 +37,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example' - - label: "XPU V1 test" + - label: ":intel: (BMG) V1 test" key: xpu-v1-test depends_on: - image-build-xpu @@ -60,7 +60,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh v1' parallelism: 2 - - label: "XPU W8A8 FP8 Linear Examples" + - label: ":intel: (BMG) W8A8 FP8 Linear Examples" key: xpu-w8a8-fp8-linear-examples depends_on: - image-build-xpu @@ -83,7 +83,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh w8a8-fp8-linear' - - label: "XPU server test" + - label: ":intel: (BMG) Server test" key: xpu-server-test depends_on: - image-build-xpu @@ -105,7 +105,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server' - - label: "XPU quantization test" + - label: ":intel: (BMG) Quantization test" key: xpu-quantization-test depends_on: - image-build-xpu @@ -127,7 +127,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh quantization' - - label: "XPU compressed tensors FP8 test" + - label: ":intel: (BMG) Compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: - image-build-xpu @@ -150,7 +150,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh compressed-tensors-fp8' - - label: "XPU Graph" + - label: ":intel: (BMG) Graph" key: xpu-graph depends_on: - image-build-xpu diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 548b6435a254..0016023dc746 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -1,8 +1,8 @@ -group: Basic Correctness +group: Basic Correctness Intel depends_on: - image-build-xpu steps: -- label: XPU Sleep Mode +- label: ":intel: (BMG) Sleep Mode" key: xpu-sleep-mode timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/benchmarks_intel.yaml b/.buildkite/intel_jobs/benchmarks_intel.yaml index 9ce373e4f42e..f6c1e104b1db 100644 --- a/.buildkite/intel_jobs/benchmarks_intel.yaml +++ b/.buildkite/intel_jobs/benchmarks_intel.yaml @@ -1,8 +1,8 @@ -group: Benchmarks +group: Benchmarks Intel depends_on: - image-build-xpu steps: -- label: Benchmarks CLI Test +- label: ":intel: (BMG) Benchmarks CLI" key: benchmarks-cli-test timeout_in_minutes: 40 device: intel_gpu diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index fd028142cc0c..5235398c2579 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -2,7 +2,7 @@ group: Engine Intel depends_on: - image-build-xpu steps: -- label: Engine +- label: ":intel: (BMG) Engine" key: engine timeout_in_minutes: 40 device: intel_gpu @@ -40,7 +40,7 @@ steps: 'cd tests && pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py' -- label: Engine (1 GPU) +- label: ":intel: (BMG) Engine (1 GPU)" key: engine-1-gpu timeout_in_minutes: 30 device: intel_gpu @@ -65,7 +65,7 @@ steps: pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py --ignore v1/engine/test_core_engine_actor_manager.py && pytest -v -s test_config.py' -- label: V1 e2e (2 GPUs) +- label: ":intel: (BMG) V1 e2e (2 GPUs)" key: v1-e2e-2-gpus timeout_in_minutes: 30 device: intel_gpu diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index 96cbd2002a2f..f49f051ebc43 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -1,4 +1,4 @@ -group: Entrypoints +group: Entrypoints Intel depends_on: - image-build-xpu steps: diff --git a/.buildkite/intel_jobs/expert_parallelism_intel.yaml b/.buildkite/intel_jobs/expert_parallelism_intel.yaml index 0c57b7de1253..c0c95fb3a93a 100644 --- a/.buildkite/intel_jobs/expert_parallelism_intel.yaml +++ b/.buildkite/intel_jobs/expert_parallelism_intel.yaml @@ -1,8 +1,8 @@ -group: Expert Parallelism +group: Expert Parallelism Intel depends_on: - image-build-xpu steps: -- label: EPLB Algorithm +- label: ":intel: (BMG) EPLB Algorithm" key: eplb-algorithm timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index 3650a729d7d7..f3a9d58f4357 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -2,7 +2,7 @@ group: Kernels Intel depends_on: - image-build-xpu steps: -- label: vLLM IR Tests +- label: ":intel: (BMG) vLLM IR Tests" key: vllm-ir-tests timeout_in_minutes: 30 device: intel_gpu @@ -26,7 +26,7 @@ steps: pytest -v -s ir && pytest -v -s kernels/ir' -- label: UVA Kernel Tests +- label: ":intel: (BMG) UVA Kernel Tests" key: uva-kernel-tests timeout_in_minutes: 20 device: intel_gpu diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 62d99cf769b0..1919e4c85041 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -1,8 +1,8 @@ -group: LM Eval +group: LM Eval Intel depends_on: - image-build-xpu steps: -- label: "XPU GPQA Eval (GPT-OSS)" +- label: ":intel: (BMG) GPQA Eval (GPT-OSS)" key: xpu-gpqa-eval-gpt-oss timeout_in_minutes: 60 device: intel_gpu @@ -25,7 +25,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' -- label: LM Eval Large Models(2 GPUs) +- label: ":intel: (BMG) LM Eval Large Models(2 GPUs)" key: lm-eval-large-models-2-gpus timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 885e4b546691..ae5e2d0165cc 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -2,7 +2,7 @@ group: LoRA Intel depends_on: - image-build-xpu steps: -- label: LoRA Runtime + Utils +- label: ":intel: (BMG) LoRA Runtime + Utils" key: lora-runtime-utils timeout_in_minutes: 45 device: intel_gpu @@ -35,7 +35,7 @@ steps: pytest -v -s lora/test_utils.py && pytest -v -s lora/test_worker.py' -- label: LoRA Fused/MoE Kernels +- label: ":intel: (BMG) LoRA Fused/MoE Kernels" key: lora-fused-moe-kernels timeout_in_minutes: 45 device: intel_gpu @@ -60,7 +60,7 @@ steps: pytest -v -s lora/test_fused_moe_lora_kernel.py && pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' -- label: LoRA Punica Kernels +- label: ":intel: (BMG) LoRA Punica Kernels" key: lora-punica-kernels timeout_in_minutes: 45 device: intel_gpu @@ -87,7 +87,7 @@ steps: pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size && pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit' -- label: LoRA Punica FP8/XPU Ops +- label: ":intel: (BMG) LoRA Punica FP8/XPU Ops" key: lora-punica-fp8-xpu-ops timeout_in_minutes: 60 device: intel_gpu @@ -112,7 +112,7 @@ steps: pytest -v -s lora/test_punica_ops_fp8.py && pytest -v -s lora/test_punica_xpu_ops.py' -- label: LoRA Models +- label: ":intel: (BMG) LoRA Models" key: lora-models timeout_in_minutes: 45 device: intel_gpu @@ -139,7 +139,7 @@ steps: pytest -v -s lora/test_chatglm3_tp.py && pytest -v -s lora/test_llama_tp.py::test_llama_lora' -- label: LoRA Multimodal +- label: ":intel: (BMG) LoRA Multimodal" key: lora-multimodal timeout_in_minutes: 45 device: intel_gpu @@ -163,7 +163,7 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && pytest -v -s lora/test_whisper.py' -- label: LoRA Add + Multimodal +- label: ":intel: (BMG) LoRA Add + Multimodal" key: lora-add-multimodal timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index b7ae9f907846..c1706d70f6fc 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -2,7 +2,7 @@ group: Miscellaneous Intel depends_on: - image-build-xpu steps: -- label: XPU Batch Invariance +- label: ":intel: (BMG) Batch Invariance" key: xpu-batch-invariance timeout_in_minutes: 60 num_devices: 2 @@ -30,7 +30,7 @@ steps: pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py && bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh' -- label: V1 Core + KV + Metrics +- label: ":intel: (BMG) V1 Core + KV + Metrics" key: v1-core-kv-metrics-intel timeout_in_minutes: 45 device: intel_gpu @@ -64,7 +64,7 @@ steps: VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' -- label: EC CPU Connector (XPU) +- label: ":intel: (BMG) EC CPU Connector (XPU)" key: ec-cpu-connector-xpu timeout_in_minutes: 45 device: intel_gpu @@ -96,7 +96,7 @@ steps: cd tests && pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' -- label: V1 Sample + Logits +- label: ":intel: (BMG) V1 Sample + Logits" key: v1-sample-logits-intel timeout_in_minutes: 90 device: intel_gpu @@ -141,7 +141,7 @@ steps: limit: 1 parallelism: 4 -- label: Basic Models Tests (Initialization) +- label: ":intel: (BMG) Basic Models Tests (Initialization)" key: basic-models-tests-initialization timeout_in_minutes: 60 device: intel_gpu @@ -166,7 +166,7 @@ steps: cd tests && pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' -- label: XPU CPU Offload +- label: ":intel: (BMG) CPU Offload" key: xpu-cpu-offload timeout_in_minutes: 60 device: intel_gpu @@ -194,7 +194,7 @@ steps: pytest -v -s v1/kv_offload && pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' -- label: NixlConnector PD accuracy (4 GPUs) +- label: ":intel: (BMG) NixlConnector PD accuracy (4 GPUs)" key: nixlconnector-pd-accuracy-4-gpus timeout_in_minutes: 60 num_devices: 4 @@ -223,7 +223,7 @@ steps: PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' -- label: Regression +- label: ":intel: (BMG) Regression" key: regression timeout_in_minutes: 50 device: intel_gpu @@ -256,7 +256,7 @@ steps: 'cd tests && pytest -v -s test_regression.py' -- label: Metrics, Tracing (2 GPUs) +- label: ":intel: (BMG) Metrics, Tracing (2 GPUs)" key: metrics-tracing-2-gpus timeout_in_minutes: 45 num_devices: 2 @@ -291,7 +291,7 @@ steps: 'cd tests && pytest -v -s v1/tracing' -- label: Async Engine, Inputs, Utils, Worker +- label: ":intel: (BMG) Async Engine, Inputs, Utils, Worker" key: async-engine-inputs-utils-worker timeout_in_minutes: 55 device: intel_gpu @@ -330,7 +330,7 @@ steps: pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' -- label: Detokenizer Stop Strings (XPU) +- label: ":intel: (BMG) Detokenizer Stop Strings (XPU)" key: detokenizer-stop-strings-xpu timeout_in_minutes: 45 device: intel_gpu @@ -367,7 +367,7 @@ steps: 'cd tests && pytest -v -s detokenizer/test_stop_strings.py' -- label: Fusion Unit Tests +- label: ":intel: (BMG) Fusion Unit Tests" key: fusion-unit-tests timeout_in_minutes: 30 device: intel_gpu diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index d641c93e5a9e..cac6c931f09a 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -2,7 +2,7 @@ group: Model Executor Intel depends_on: - image-build-xpu steps: -- label: Model Executor (Intel) +- label: ":intel: (BMG) Model Executor" key: model-executor-intel timeout_in_minutes: 45 device: intel_gpu @@ -30,7 +30,7 @@ steps: cd tests && pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' -- label: Tensorizer Model Loader Tests (XPU) +- label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)" key: tensorizer-model-loader-tests-xpu timeout_in_minutes: 45 device: intel_gpu diff --git a/.buildkite/intel_jobs/models_distributed_intel.yaml b/.buildkite/intel_jobs/models_distributed_intel.yaml index 5f2247eeaace..5fd723f8c5f0 100644 --- a/.buildkite/intel_jobs/models_distributed_intel.yaml +++ b/.buildkite/intel_jobs/models_distributed_intel.yaml @@ -1,8 +1,8 @@ -group: Models - Distributed +group: Models - Distributed Intel depends_on: - image-build-xpu steps: -- label: Distributed Model Tests (2 GPUs) +- label: ":intel: (BMG) Distributed Model Tests" key: distributed-model-tests-2-gpus timeout_in_minutes: 65 device: intel_gpu diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 89654ec3aedf..9a9790d5adc7 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -1,8 +1,8 @@ -group: Models - Multimodal +group: Models - Multimodal Intel depends_on: - image-build-xpu steps: -- label: "Multi-Modal Models (Standard) 1: qwen2" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 1: qwen2" key: multi-modal-models-standard-1-qwen2 timeout_in_minutes: 70 device: intel_gpu @@ -26,7 +26,7 @@ steps: pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" && pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model' -- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 2: qwen3 + gemma" key: multi-modal-models-standard-2-qwen3-gemma timeout_in_minutes: 90 device: intel_gpu @@ -51,7 +51,7 @@ steps: pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model && pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model' -- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 3: llava + qwen2_vl" key: multi-modal-models-standard-3-llava-qwen2-vl timeout_in_minutes: 65 device: intel_gpu @@ -76,7 +76,7 @@ steps: pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' -- label: "Multi-Modal Models (Standard) 4: other + whisper" +- label: ":intel: (BMG) Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper timeout_in_minutes: 45 device: intel_gpu @@ -99,7 +99,7 @@ steps: 'cd tests && pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_gemma_ragged_audio.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' -- label: Multi-Modal Processor # 44min +- label: ":intel: (BMG) Multi-Modal Processor" # 44min key: multi-modal-processor timeout_in_minutes: 60 device: intel_gpu @@ -126,7 +126,7 @@ steps: --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB' parallelism: 4 -- label: Multi-Modal Models (Extended Generation 1) +- label: ":intel: (BMG) Multi-Modal Models (Extended Generation 1)" key: multi-modal-models-extended-generation-1 timeout_in_minutes: 60 device: intel_gpu diff --git a/.buildkite/intel_jobs/quantization.yaml b/.buildkite/intel_jobs/quantization_intel.yaml similarity index 95% rename from .buildkite/intel_jobs/quantization.yaml rename to .buildkite/intel_jobs/quantization_intel.yaml index af1b23b8a004..035ad11beef2 100644 --- a/.buildkite/intel_jobs/quantization.yaml +++ b/.buildkite/intel_jobs/quantization_intel.yaml @@ -1,8 +1,8 @@ -group: Quantization +group: Quantization Intel depends_on: - image-build-xpu steps: -- label: Quantization +- label: ":intel: (BMG) Quantization" key: quantization timeout_in_minutes: 30 env: diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml index 027237225d3a..ed55c49da890 100644 --- a/.buildkite/intel_jobs/samplers_intel.yaml +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -2,7 +2,7 @@ group: Samplers Intel depends_on: - image-build-xpu steps: -- label: Samplers Test (XPU Kernel) +- label: ":intel: (BMG) Samplers Test (XPU Kernel)" key: samplers-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu @@ -29,7 +29,7 @@ steps: VLLM_XPU_USE_SAMPLER_KERNEL=1 pytest -v -s samplers --ignore=samplers/test_beam_search.py' -- label: Samplers Beam Search Test (XPU Kernel) +- label: ":intel: (BMG) Samplers Beam Search Test (XPU Kernel)" key: samplers-beam-search-test-xpu-kernel-intel timeout_in_minutes: 40 device: intel_gpu diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 1567c0eba724..74d6eee4900e 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -14,7 +14,7 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 - - label: "XPU example Test" + - label: ":intel: (BMG) Example Test" key: xpu-example-test depends_on: - image-build-xpu @@ -30,7 +30,7 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh @@ -48,7 +48,7 @@ steps: python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC -tp 2 --max-model-len 8192 ' - - label: "XPU W8A8 FP8 Linear Examples" + - label: ":intel: (BMG) W8A8 FP8 Linear Examples" key: xpu-w8a8-fp8-linear-examples depends_on: - image-build-xpu @@ -76,7 +76,7 @@ steps: python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 && python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 ' - - label: "XPU V1 test" + - label: ":intel: (BMG) V1 test" key: xpu-v1-test depends_on: - image-build-xpu @@ -92,7 +92,7 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh @@ -107,7 +107,7 @@ steps: pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' parallelism: 2 - - label: "XPU Qwen3-30B + V1 KV Connector" + - label: ":intel: (BMG) Qwen3-30B + V1 KV Connector" key: xpu-qwen3-30b-v1-kv-connector depends_on: - image-build-xpu @@ -133,7 +133,7 @@ steps: pytest -v -s v1/kv_connector/unit/test_example_connector.py && cd .. && VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' - - label: "XPU server test" + - label: ":intel: (BMG) Server test" key: xpu-server-test depends_on: - image-build-xpu @@ -149,14 +149,14 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py && pytest -v -s benchmarks/test_serve_cli.py' - - label: "XPU quantization test" + - label: ":intel: (BMG) Quantization test" key: xpu-quantization-test depends_on: - image-build-xpu @@ -172,14 +172,14 @@ steps: REPO: "vllm-ci-test-repo" source_file_dependencies: - vllm/ - - .buildkite/intel_jobs/test-intel.yaml + - .buildkite/intel_jobs/test-intel.yaml commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s quantization/test_auto_round.py && pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' - - label: "XPU compressed tensors FP8 test" + - label: ":intel: (BMG) Compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: - image-build-xpu @@ -204,7 +204,7 @@ steps: 'cd tests && pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8 && pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8_block_enabled' - - label: "XPU Graph" + - label: ":intel: (BMG) Graph" key: xpu-graph depends_on: - image-build-xpu From c077dc663d929fd19bdfb15ccf00f83cc490b8b0 Mon Sep 17 00:00:00 2001 From: frida-andersson Date: Mon, 12 Oct 2026 21:19:14 -0500 Subject: [PATCH 458/952] [ROCm][Perf][DSv4.1] Enable decoder SWA bounded replay (#59039) Signed-off-by: Frida Andersson --- .buildkite/test_areas/models_basic.yaml | 4 +- .../test_deepseek_v41_amd_replay_forward.py | 329 ++++++++++++++++++ .../models/test_deepseek_v41_replay_batch.py | 75 ++++ vllm/models/deepseek_v41/amd/model.py | 291 ++++++++++++++-- vllm/models/deepseek_v41/amd/vl_model.py | 5 + vllm/models/deepseek_v41/attention.py | 10 +- .../models/deepseek_v41/nvidia/model_state.py | 34 +- 7 files changed, 705 insertions(+), 43 deletions(-) create mode 100644 tests/models/test_deepseek_v41_amd_replay_forward.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 552910e85a96..045e04533a44 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -82,18 +82,20 @@ steps: - image-build-amd source_file_dependencies: - vllm/models/deepseek_v4/ + - vllm/models/deepseek_v41/ - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/test_deepseek_v41_amd_replay_forward.py - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/test_deepseek_v41_amd_replay_forward.py b/tests/models/test_deepseek_v41_amd_replay_forward.py new file mode 100644 index 000000000000..fe22bc80253f --- /dev/null +++ b/tests/models/test_deepseek_v41_amd_replay_forward.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What the AMD replay forward does with a prefill longer than the window. + +The cut in this fixture is layer 1, so replay starts at layer 1. That is the +same boundary as production layer 20, on a 4-layer stand-in. The final +``mhc_post`` and the aux capture at that cut both run on the window rows. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import ReplayBatch +from vllm.sequence import IntermediateTensors + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +WINDOW = 128 +FULL = 300 +# Trailing window of a 300-token prefill. +ROWS = list(range(FULL - WINDOW, FULL)) +HC, HIDDEN = 4, 8 +DEVICE = torch.device("cuda") +# gfx942 writes the V4 record (584 B). Family-100 writes the V4.1 record (528 B). +# The SWA layer uses block size 32, which is the size the KV-only insert checks. +_BLOCK = 32 +_HEAD_DIM = 512 + + +class _RocmKvInsert: + """SWA cache from ``forward_kv``, checked against the KV-only insert.""" + + def __init__(self, num_tokens: int): + from vllm.models.deepseek_v41.attention import _use_v41_mxfp8_kv_record + + if not hasattr(torch.ops._C, "fused_deepseek_v4_kv_rope_insert"): + pytest.skip("fused DeepseekV4 KV insert op is not built") + + self.num_tokens = num_tokens + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + bytes_per = 528 if self.kv_mxfp8 else 584 + num_blocks = (num_tokens + _BLOCK - 1) // _BLOCK + 1 + cache = torch.full( + (num_blocks, _BLOCK, bytes_per), + 0xA5, + dtype=torch.uint8, + device=DEVICE, + ) + self.poison = cache.clone() + self.kv = torch.randn( + num_tokens, _HEAD_DIM, dtype=torch.bfloat16, device=DEVICE + ) + prefix = "replay.swa" + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=DEVICE) + self.metadata = { + prefix: SimpleNamespace(slot_mapping=slot_mapping, block_size=_BLOCK) + } + self.context = ForwardContext({}, self.metadata, {}) + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device=DEVICE) / 64) + ) + freqs = torch.outer( + torch.arange(num_tokens + 8, dtype=torch.float32, device=DEVICE), inv_freq + ) + self.attn = SimpleNamespace( + kv_mxfp8=self.kv_mxfp8, + compressor=None, + indexer=None, + rotary_emb=SimpleNamespace( + cos_sin_cache=torch.cat((freqs.cos(), freqs.sin()), dim=-1) + ), + swa_cache_layer=SimpleNamespace( + prefix=prefix, kv_cache=cache, block_size=_BLOCK + ), + _run_parallel_input_projections=lambda hidden: (hidden, None, None), + _split_qkv_and_norm=lambda qr_kv: (None, None, self.kv[: qr_kv.shape[0]]), + ) + self.last_rows: int | None = None + self.last_positions: torch.Tensor | None = None + + def write( + self, + x, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + ): + from vllm.models.deepseek_v41.amd.rocm import ( + DeepseekV41ROCMAiterMLAAttention, + ) + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + assert ( + DeepseekV41ROCMAiterMLAAttention.forward_kv + is DeepseekV4Attention.forward_kv + ) + self.last_rows = x.shape[0] + self.last_positions = positions + DeepseekV41ROCMAiterMLAAttention.forward_kv(self.attn, positions, x) + + def assert_matches_forward(self): + assert self.last_rows == self.num_tokens + assert self.last_positions is not None + written = self.attn.swa_cache_layer.kv_cache + assert not torch.equal(written, self.poison) + ref = self.poison.clone() + meta = self.metadata[self.attn.swa_cache_layer.prefix] + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + self.kv, + ref, + meta.slot_mapping, + self.last_positions, + self.attn.rotary_emb.cos_sin_cache, + _BLOCK, + None, + self.kv_mxfp8, + ) + assert torch.equal(written, ref) + assert torch.equal(written[-1], self.poison[-1]) + + +class _Layer(nn.Module): + def __init__( + self, + idx: int, + compressed: str | None = None, + indexer: str | None = None, + ): + super().__init__() + self.idx = idx + self.seen: list[torch.Tensor] = [] + self.attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix=f"layers.{idx}.swa"), + compressed_cache_prefix=compressed, + indexer=( + None + if indexer is None + else SimpleNamespace(k_cache=SimpleNamespace(prefix=indexer)) + ), + compress_ratio=1, + ) + + def forward( + self, + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ): + self.seen.append(hidden_states[:, 0, 0].detach().clone()) + if residual is None: + residual = hidden_states + if pre_mix is None: + pre_mix = hidden_states[:, :, 0] + if post_mix is None: + post_mix = pre_mix + if res_mix is None: + res_mix = pre_mix + return hidden_states, residual, post_mix, res_mix, pre_mix + + +def _model(monkeypatch): + import vllm.models.deepseek_v41.amd.model as amd + + monkeypatch.setattr( + amd, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False), + ) + monkeypatch.setattr(amd.EngramLayout, "from_config", lambda config: None) + monkeypatch.setattr(amd, "MHCPostOp", lambda: SimpleNamespace()) + + # Layer 0 sits before the cut, so its keys are not collected. Layer 2 + # shares layer 1's compressed cache. Layer 3 has neither. + layers = [ + _Layer(0, "layers.0.compressed", "layers.0.indexer"), + _Layer(1, "layers.1.compressed", "layers.1.indexer"), + _Layer(2, "layers.1.compressed", "layers.2.indexer"), + _Layer(3), + ] + # Layer 1 is the KV source, so the replay callback is this layer's write. + insert = _RocmKvInsert(FULL) + layers[1].write_kv = insert.write + monkeypatch.setattr(amd, "make_layers", lambda n, factory, prefix: (0, n, layers)) + + config = SimpleNamespace( + vocab_size=32, + hc_eps=1e-6, + hc_mult=HC, + hidden_size=HIDDEN, + rms_norm_eps=1e-6, + index_topk=4, + kv_source_layer_ids=[1], + num_hidden_layers=4, + sliding_window=WINDOW, + engram_layer_ids=(), + candidate_source_layer_id=-1, + candidate_topk_blocks=0, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config, dtype=torch.bfloat16), + quant_config=None, + cache_config=SimpleNamespace(swa_bounded_replay=True), + use_v2_model_runner=True, + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + use_ubatching=False, + enable_expert_parallel=False, + ), + speculative_config=None, + scheduler_config=SimpleNamespace(max_num_batched_tokens=FULL), + ) + model = amd.DeepseekV4Model(vllm_config=vllm_config) + # Production requests aux at the first replay layer. Here that id is the cut. + model.aux_hidden_state_layers = (model.decoder_replay_start,) + model.topk_indices_buffer = model.topk_indices_buffer.to(DEVICE) + + calls: list[torch.Tensor] = [] + + def mhc_post(hidden, residual, post_mix, res_mix): + calls.append(hidden[:, 0, 0].detach().clone()) + return hidden + + model.mhc_post = mhc_post + saved_aux: list[list[torch.Tensor]] = [] + run_layers = model._run_layers + + def wrapped(*args, **kwargs): + result = run_layers(*args, **kwargs) + # Snapshot: forward() extends the early list with the scattered aux. + saved_aux.append(list(result[5])) + return result + + model._run_layers = wrapped + scattered: list[tuple[torch.Tensor, ...]] = [] + replay = model.decoder_replay_layers + run = replay._run + + def wrapped_run(*args, **kwargs): + result = run(*args, **kwargs) + scattered.append(result) + return result + + replay._run = wrapped_run + return model, layers, calls, saved_aux, scattered, insert + + +def _forward(model, insert: _RocmKvInsert): + token_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + hidden = token_ids[:, None, None].expand(FULL, HC, HIDDEN).contiguous() + pre_mix = torch.ones(FULL, HC, device=DEVICE) + rows = torch.tensor(ROWS, device=DEVICE) + model.decoder_replay_layers.replay_batch = ReplayBatch( + rows, + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}), + ) + positions = torch.arange(FULL, device=DEVICE) + with override_forward_context(insert.context): + out = model( + input_ids=positions.to(torch.int32), + positions=positions, + intermediate_tensors=IntermediateTensors( + {"hidden_states": hidden, "pre_mix": pre_mix} + ), + ) + return out + + +def test_amd_replay_forward_trims_to_the_window(monkeypatch): + model, layers, calls, saved_aux, scattered, insert = _model(monkeypatch) + replay_layers = model.decoder_replay_layers + assert model.decoder_replay_start == 1 + assert replay_layers.first_swa_prefix == "layers.1.swa" + assert replay_layers.metadata_prefixes == { + "layers.1.swa", + "layers.1.compressed", + "layers.1.indexer", + "layers.2.swa", + "layers.2.indexer", + "layers.3.swa", + } + assert replay_layers.replay_batch is None + out = _forward(model, insert) + insert.assert_matches_forward() + replay = model.decoder_replay_layers.replay_batch + assert replay is not None and replay.rows.tolist() == ROWS + + early = [layer.seen[0] for layer in layers[: model.decoder_replay_start]] + late = [layer.seen[0] for layer in layers[model.decoder_replay_start :]] + full_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + window_ids = full_ids[ROWS] + assert [t.tolist() for t in early] == [full_ids.tolist()] + assert [t.tolist() for t in late] == [ + window_ids.tolist(), + window_ids.tolist(), + window_ids.tolist(), + ] + + # Scatter writes the replay rows back and leaves the trimmed rows at zero. + hidden = out["hidden_states"][:, 0, 0] + assert torch.equal(hidden[ROWS], window_ids) + assert hidden[: FULL - WINDOW].eq(0).all() + + # calls[0] is the cut-aux capture. calls[1] is the final mhc_post. + # Both run inside the replay batch, before the scatter back to FULL. + assert len(calls) == 2 + assert [t.shape[0] for t in calls] == [WINDOW, WINDOW] + assert torch.equal(calls[0], window_ids) + assert torch.equal(calls[1], window_ids) + + # The early loop skips this id. The replay loop does not see it either. + assert saved_aux[0] == [] + assert saved_aux[1] == [] + _hidden, _pre_mix, aux = scattered[0] + assert aux.shape[0] == FULL + assert torch.equal(aux[ROWS, 0], window_ids) + assert aux[: FULL - WINDOW].eq(0).all() diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index bcd4db64b25c..5a3625dcbf4e 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -290,6 +290,81 @@ def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state): assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129] +def test_zero_prefix_window_warms_gather_and_skips_pad(monkeypatch): + """ROCm leaves the prefix-cache window at 0. Replay still needs replay_start + and a warmed gather kernel. The pad kernel stays on the prefix-cache path.""" + import vllm.models.deepseek_v41.nvidia.model_state as model_state + + calls = {"pad_warmup": 0, "gather_warmup": 0, "pad_launch": 0} + + def pad_warmup(*args, **kwargs): + calls["pad_warmup"] += 1 + + def gather_warmup(*args, **kwargs): + calls["gather_warmup"] += 1 + + class _Pad: + warmup = staticmethod(pad_warmup) + + def __getitem__(self, grid): + def launch(*args, **kwargs): + calls["pad_launch"] += 1 + + return launch + + monkeypatch.setattr(model_state, "_pad_replayed_slots_kernel", _Pad()) + monkeypatch.setattr( + model_state._gather_replay_batch_kernel, "warmup", gather_warmup + ) + + cfg = MagicMock() + cfg.model_config.enable_prompt_embeds = False + cfg.model_config.uses_mrope = False + cfg.model_config.is_multimodal_model = False + cfg.scheduler_config.max_num_seqs = 16 + cfg.scheduler_config.max_num_batched_tokens = 1024 + cfg.parallel_config.data_parallel_size = 1 + cfg.compilation_config.fast_moe_cold_start = False + cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) + model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) + groups = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=["swa"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=False, prefix_replay_tokens=0 + ), + ), + SimpleNamespace( + layer_names=["mla"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=True, prefix_replay_tokens=0 + ), + ), + ] + ) + + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw): + return {"swa": kw["model_specific_attn_metadata"]} + + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) + state = DeepseekV41ModelState(cfg, model, None, DEVICE) + state._replay_start_np[0] = 16 + batch = _input_batch([8], [24], [True]) + state.prepare_attn( + batch, + CUDAGraphMode.NONE, + BLOCK_TABLES, + _slot_mappings(batch.num_tokens_after_padding), + [], + groups, + ) + assert state._replay[0] == 0 + assert calls == {"pad_warmup": 0, "gather_warmup": 1, "pad_launch": 0} + assert state._replay_start[:1].tolist() == [16] + + def test_dp_ranks_pad_to_one_replay_graph(dp_state): """Every rank pads its replay batch to the graph fitting the largest rank's.""" dp_state.replay_cudagraphs = MagicMock() diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 4c5f15858749..a52a408c74dd 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -62,6 +62,7 @@ from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -421,6 +422,42 @@ def forward( x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre + def write_kv( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> None: + """Write the KV ``forward`` would for the replay layers' inputs, skipping + the rest of the layer; the replay batch then reruns its rows.""" + # The replay batch rules out Engram and sequence parallel here. + assert self.engram is None and not self.use_sequence_parallel + fuse_attn_norm = self.fuse_seam_norm + residual, post_mix, res_mix, x, _attn_pre = self.mhc_pre_delayed( + residual, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, + norm_weight=self.attn_norm.weight if fuse_attn_norm else None, + norm_eps=self.attn_norm.variance_epsilon, + ) + if not fuse_attn_norm: + x = self.attn_norm(x) + self.attn.forward_kv(positions, x) + class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -498,6 +535,45 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + # Config flag is the switch: ROCm forces swa_cache_layer.bounded_replay off. + self.decoder_replay_layers: DecoderReplayLayers | None = None + self.decoder_replay_start = self.end_layer + cut = max(config.kv_source_layer_ids) + if ( + vllm_config.cache_config.swa_bounded_replay + and vllm_config.use_v2_model_runner + and cut < self.end_layer - 1 + and self._decoder_replay_supported(vllm_config, cut) + ): + assert self.layers[cut].attn.compress_ratio <= 1 + # The last KV source writes every row's KV, then replays. + self.decoder_replay_start = cut + # The attention metadata keys the replay layers read. + metadata_prefixes: set[str] = set() + for layer in islice(self.layers, cut, self.end_layer): + attn = typing.cast(DeepseekV4DecoderLayer, layer).attn + metadata_prefixes.add(attn.swa_cache_layer.prefix) + if attn.compressed_cache_prefix is not None: + metadata_prefixes.add(attn.compressed_cache_prefix) + if attn.indexer is not None: + metadata_prefixes.add(attn.indexer.k_cache.prefix) + self.decoder_replay_layers = DecoderReplayLayers( + config.sliding_window, + self._run_replay_layers, + self.layers[cut].write_kv, + metadata_prefixes, + self.layers[cut].attn.swa_cache_layer.prefix, + ) + logger.info_once( + "Decoder SWA bounded replay: in eager prefill steps, layer %d " + "writes its KV for every token, and the rest of it and layers " + "%d-%d run on each request's last %d tokens only.", + cut, + cut + 1, + self.end_layer - 1, + config.sliding_window, + ) + # The n-gram hash needs a slot-keyed rolling store of compressed ids # (chunked prefill / decode lookback); key it off the first local # layer's sliding-window KV cache. Only PP ranks owning an engram @@ -657,13 +733,33 @@ def forward( if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + full_num_tokens, + ) + # Replay collapses inside its own batch. The early loop still sets + # ran_layers, so that flag alone must not post again on the full batch. + collapsed = False + if self.decoder_replay_layers is not None: + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, input_ids, @@ -671,20 +767,15 @@ def forward( post_mix, res_mix, residual, - engram_hashes, - engram_mask, ) - if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) - aux_hidden_state = aux_recon.mean(dim=1) - if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) - final_aux_recon = aux_recon - if layer is not None: - # Reuse if the last layer was captured as an aux hidden state - if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.extend(late_aux) + final_aux_recon = None + collapsed = True + if ran_layers and not collapsed: + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): hidden_states = final_aux_recon else: hidden_states = self.mhc_post( @@ -714,6 +805,162 @@ def forward( return hidden_states, aux_hidden_states return hidden_states + def _run_layers( + self, + layer_ids: range, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + engram_hashes: torch.Tensor | None, + engram_mask: torch.Tensor | None, + full_num_tokens: int, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + list[torch.Tensor], + torch.Tensor | None, + bool, + ]: + aux_hidden_states: list[torch.Tensor] = [] + final_aux_recon: torch.Tensor | None = None + ran_layers = False + for idx in layer_ids: + ran_layers = True + hidden_states, residual, post_mix, res_mix, pre_mix = self.layers[idx]( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ) + # The aux id of the first replay layer is the post of the last + # full-batch layer. Capture it on the trimmed rows instead. + if idx + 1 in self.aux_hidden_state_layers and not ( + self.decoder_replay_layers is not None + and idx + 1 == self.decoder_replay_start + ): + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_state = aux_recon.mean(dim=1) + if self.use_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + final_aux_recon = aux_recon + return ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) + + def _run_replay_layers( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Layers past the last KV source, on the rows they are given. + + The final ``mhc_post`` runs here, on this batch. The return is + ``(hidden_states, pre_mix, *aux)``. + """ + # Post of the layer before the cut. These rows are already trimmed, + # and sequence parallel is unsupported while replay is on. + cut_aux: list[torch.Tensor] = [] + if self.decoder_replay_start in self.aux_hidden_state_layers: + cut_aux.append( + self.mhc_post(hidden_states, residual, post_mix, res_mix).mean(dim=1) + ) + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.decoder_replay_start, self.end_layer), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + None, + None, + positions.shape[0], + ) + assert residual is not None and post_mix is not None + assert res_mix is not None and pre_mix is not None + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): + hidden_states = final_aux_recon + elif ran_layers: + hidden_states = self.mhc_post(hidden_states, residual, post_mix, res_mix) + return (hidden_states, pre_mix, *cut_aux, *aux) + + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: + """Whether this rank may replay from layer ``cut`` on; warns when not.""" + parallel_config = vllm_config.parallel_config + spec_config = vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + draft_hf_config = getattr(draft_config, "hf_config", None) + draft_window = getattr(draft_hf_config, "sliding_window", None) + draft_layer_types = getattr(draft_hf_config, "layer_types", None) or () + window = self.config.sliding_window + if self.start_layer > cut or self.end_layer < self.config.num_hidden_layers: + reason = ( + "the pipeline stage holding the last KV source layer must also " + "hold every layer after it" + ) + elif ( + self.use_sequence_parallel + or parallel_config.prefill_context_parallel_size > 1 + or parallel_config.use_ubatching + ): + reason = ( + "the replay-layer batch shrinks per rank, which sequence and " + "prefill-context parallelism and microbatching cannot follow" + ) + elif any(i >= cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits at or after the last KV source layer" + elif draft_config is not None and ( + draft_window is None + or draft_window > window + or any(t != "sliding_attention" for t in draft_layer_types) + ): + reason = ( + f"the drafter (sliding window {draft_window}) reads hidden states " + f"outside the target's {window}-token window" + ) + else: + return True + logger.warning_once("Decoder SWA bounded replay is off: %s.", reason) + return False + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1128,6 +1375,10 @@ def token_lookback_depth(self) -> int: engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 5ed045029b3c..5b8f7b826205 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -51,6 +51,7 @@ DeepseekV4VLProcessingInfo, ) from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin +from ..decoder_replay_layers import DecoderReplayLayers from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -299,6 +300,10 @@ def get_model_state_cls(): def token_lookback_depth(self) -> int: return self.language_model.token_lookback_depth + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.language_model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 559d874cc483..4d460d89b801 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -547,11 +547,11 @@ def __init__( swa_bounded_replay = False if swa_bounded_replay and current_platform.is_rocm(): logger.warning_once( - "SWA bounded replay is off on ROCm (the sparse SWA metadata " - "builders forward replay_start, but the window clamp it relies " - "on lives in the FlashInfer and FlashMLA prefill kernels, so " - "the padded slots fault); the sliding-window cache takes part " - "in prefix caching instead." + "Prefix-cache SWA replay is off on ROCm (the sparse SWA " + "metadata builders forward replay_start, but the window clamp " + "it relies on lives in the FlashInfer and FlashMLA prefill " + "kernels, so the padded slots fault); the sliding-window cache " + "takes part in prefix caching instead." ) swa_bounded_replay = False self.swa_cache_layer = DeepseekV4SWACache( diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 7fc0bd25ff75..e8bbaef64ddc 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -337,7 +337,7 @@ def prepare_attn( self._warm_up_replay_kernels(input_batch, slot_mappings) window, cacheable_groups = self._replay replay_start: torch.Tensor | None = None - if window: + if window or self.decoder_replay_layers is not None: num_reqs = input_batch.num_reqs # Decode rows sit above the hit, so only prefills carry a replay # start; dummy batches (captures, profiling) carry none. @@ -349,7 +349,7 @@ def prepare_attn( replay_start = self._replay_start_staging.copy_to_gpu( replay_start_np, out=self._replay_start[:num_reqs] ) - if replay_start_np.any(): + if window and replay_start_np.any(): # The replayed tokens rebuild window KV only: their slots in the # prefix-cacheable groups are padded so the cached KV stays as is. _pad_replayed_slots_kernel[(num_reqs,)]( @@ -397,21 +397,21 @@ def _warm_up_replay_kernels( dummy batches never replay or trim, so they would not launch them.""" assert self._replay is not None window, cacheable_groups = self._replay - if not window: - return - _pad_replayed_slots_kernel.warmup( - slot_mappings, - slot_mappings.stride(0), - cacheable_groups, - input_batch.query_start_loc, - input_batch.positions, - self._replay_start, - window, - PAD_SLOT_ID, - NUM_GROUPS=cacheable_groups.numel(), - BLOCK=1024, - grid=(1,), - ) + # Pad warmup is prefix-cache SWA. A zero window must not skip the gather warmup. + if window: + _pad_replayed_slots_kernel.warmup( + slot_mappings, + slot_mappings.stride(0), + cacheable_groups, + input_batch.query_start_loc, + input_batch.positions, + self._replay_start, + window, + PAD_SLOT_ID, + NUM_GROUPS=cacheable_groups.numel(), + BLOCK=1024, + grid=(1,), + ) if self.decoder_replay_layers is None: return assert self._kept_slot_mappings is not None From df6c7ad03194a3051a197d57ca4a94e4a063014e Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:05:39 +0800 Subject: [PATCH 459/952] [Core][Refactor] Move KV cache planning logic to kv_cache_planning.py (#61183) Signed-off-by: Dakai An Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 5 + .buildkite/test_areas/lm_eval.yaml | 6 + tests/models/test_dspark_mla.py | 2 +- tests/models/test_initialization.py | 6 +- .../test_sparse_mla_kv_cache_layout.py | 2 +- tests/v1/core/test_contiguous_kv_packing.py | 6 +- tests/v1/core/test_kv_cache_planning.py | 2284 ++++++++++++++++ tests/v1/core/test_kv_cache_utils.py | 2348 +---------------- tests/v1/core/utils.py | 74 + tests/v1/engine/test_init_error_messaging.py | 4 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/test_hidden_states_connector.py | 2 +- tests/v1/simple_kv_offload/test_worker.py | 4 +- .../test_dsv4_packed_zeroer_geometry.py | 2 +- tests/v1/worker/test_gpu_model_runner.py | 5 +- vllm/v1/core/kv_cache_planning.py | 1637 ++++++++++++ vllm/v1/core/kv_cache_utils.py | 1615 +----------- vllm/v1/engine/core.py | 2 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- vllm/v1/worker/gpu_model_runner.py | 2 +- 21 files changed, 4063 insertions(+), 3955 deletions(-) create mode 100644 tests/v1/core/test_kv_cache_planning.py create mode 100644 vllm/v1/core/kv_cache_planning.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d2702e7d4f82..5145f5a272ff 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4321,6 +4321,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4354,6 +4355,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4387,6 +4389,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4420,6 +4423,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4453,6 +4457,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f6a4a89f11cd..96e5d4fc5e9f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -371,6 +371,7 @@ steps: - vllm/tokenizers/ - vllm/v1/attention/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/v1/spec_decode/metrics.py - vllm/v1/worker/ - vllm/_aiter_ops.py @@ -516,6 +517,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -559,6 +561,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -602,6 +605,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -645,6 +649,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -688,6 +693,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index b2156af346c8..81019421ed13 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -394,7 +394,7 @@ def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): """The draft's MLA layers must share a KV cache group with the target's.""" from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention - from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + from vllm.v1.core.kv_cache_planning import _get_kv_cache_groups_uniform_page_size vllm_config = SimpleNamespace( model_config=None, cache_config=SimpleNamespace(block_size=64) diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 073f616af875..ee6dfee7ecbf 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -14,10 +14,8 @@ from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.attention.backends.utils import resolve_kv_cache_layout -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - get_kv_cache_configs, -) +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.engine.core import EngineCore as V1EngineCore from ..utils import ( diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index d869141df81a..45ba1ac0dc55 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -23,7 +23,7 @@ flat_kv_row_view, ) from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts -from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs from vllm.v1.kv_cache_layout import KVCacheLayout from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 6a705b6e5e8e..64cd55054015 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -18,13 +18,15 @@ from vllm.config import CacheConfig from vllm.v1.core.kv_cache_manager import KVCacheManager -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, _get_packed_kv_cache_groups, _pool_bytes_per_block, - generate_scheduler_kv_cache_config, get_kv_cache_config_from_groups, get_kv_cache_groups, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, resolve_kv_cache_block_sizes, ) from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py new file mode 100644 index 000000000000..629286808148 --- /dev/null +++ b/tests/v1/core/test_kv_cache_planning.py @@ -0,0 +1,2284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap +import subprocess +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.core.kv_cache_planning as kv_cache_planning +import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) +from vllm.config import ( + CacheConfig, + ModelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.config.attention import HiSparseConfig +from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.core.kv_cache_planning import ( + check_enough_kv_cache_memory, + estimate_max_model_len, + get_kv_cache_configs, + get_kv_cache_groups, + is_kv_cache_spec_uniform, + kv_cache_groups_tp_replicas, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, +) +from vllm.v1.hisparse.layout import ( + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + iter_layer_specs, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_config, + SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +@pytest.mark.parametrize( + "max_model_len,num_gpu_blocks,ok", + [ + (8192, 160, False), + (8192, 1000, True), + (32768, 656, False), + (32768, 657, True), + (-1, 657, True), + ], +) +def test_hisparse_pool_must_fit_max_model_len( + monkeypatch, max_model_len, num_gpu_blocks, ok +): + """Each resident group takes its own blocks from HiSparse's shared GPU pool, + so a pool that fits only the indexer pages must be rejected at startup. + Resident pages are capped at the in-flight window: at 32768 tokens with 2048 + in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, + plus the null block. Auto-fit (-1) must pick a length that passes the same + check.""" + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + specs: dict[str, KVCacheSpec] = {} + for i in range(4): + specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ) + specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ) + config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + config.scheduler_config.max_num_batched_tokens = 2048 + # Async scheduling doubles the in-flight window, and with it the bound. + config.scheduler_config.async_scheduling = False + config.attention_config.hisparse_config = HiSparseConfig() + config.model_config.hf_config.index_topk = 128 + config.cache_config.num_gpu_blocks_override = num_gpu_blocks + config.cache_config.kv_cache_layout = "BLHNC" + # Workers report the resident/hot caches HiSparse derives. + specs |= { + name: group.kv_cache_spec + for group in get_hisparse_kv_cache_groups(config, specs) or [] + if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) + for name in group.layer_names + } + + if ok: + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + else: + with pytest.raises(ValueError, match="max seq len"): + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + + +def new_chunked_local_attention_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + attention_chunk_size=4, +): + return ChunkedLocalAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + attention_chunk_size=attention_chunk_size, + ) + + +def test_get_kv_cache_configs_multiple_workers(): + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + same_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + ] + + # Basic case. All things are the same. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different available memory. This is the case for TP. + # Use the smallest memory available. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 20, + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different KV cache specs. This is the case for PP. + different_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer2": new_kv_cache_spec(), + "layer3": new_kv_cache_spec(), + }, + ] + + # Different workers have different layers. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer2", "layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), + ], + ), + ] + + # Some layers are the same, some are different. This is the case for TP+PP + tp_pp_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + ] + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + tp_pp_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected_12 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + expected_3 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [ + expected_12, + expected_12, + expected_3, + expected_3, + ] + + # Different workers have different types of layers. This is the case for + # hybrid models + PP. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_sliding_window_spec(), + "layer4": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + KVCacheGroupSpec([], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer3", "layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec([], ref_kv_cache_spec), + KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), + ], + ), + ] + + # When divided into multiple KVCacheGroups, need to ensure the number of + # layers per group is similar. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_sliding_window_spec(), + "layer3": new_sliding_window_spec(), + }, + { + "layer4": new_kv_cache_spec(), + "layer5": new_sliding_window_spec(), + "layer6": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 10, + ref_kv_cache_spec.page_size_bytes * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer5"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer6"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), + ], + ), + ] + + # Have conflicting layers. Need to raise an error. + conflicting_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer1": new_sliding_window_spec(), + }, + ] + with pytest.raises(AssertionError): + get_kv_cache_configs( + vllm_config, + conflicting_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + + +@pytest.mark.parametrize( + "asymmetric_memory", + [False, True], + ids=["symmetric", "asymmetric"], +) +def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): + model_config = ModelConfig(max_model_len=512) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + pp_kv_cache_specs = [ + {"layer1": ref_kv_cache_spec}, + {"layer2": ref_kv_cache_spec}, + ] + + expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 + avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks + + # With per-worker validation, each worker only needs memory for its own + # layers. Worker 2 having more memory shouldn't affect worker 1's config. + available_memory = ( + [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 + ) + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + pp_kv_cache_specs, + available_memory, + ) + + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], + ), + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], + ), + ] + + +def test_project_kv_cache_groups_to_worker(): + spec_a = new_kv_cache_spec() + spec_b = new_kv_cache_spec(num_kv_heads=4) + + global_groups = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), + ] + worker_spec = {"layer1": spec_a, "layer2": spec_a} + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, worker_spec + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer2"] + assert projected[0].kv_cache_spec is spec_a + + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, {"layer4": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == [] + assert projected[0].kv_cache_spec is spec_a + + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, + ) + global_groups_uniform = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), + ] + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer3"] + proj_spec = projected[0].kv_cache_spec + assert isinstance(proj_spec, UniformTypeKVCacheSpecs) + assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} + + +@pytest.mark.parametrize("with_draft", [False, True]) +@pytest.mark.parametrize("indexer_alignment", [None, 512]) +def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view + from vllm.v1.worker.utils import allocate_kv_cache + + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.cache_config.kv_cache_layout = "BLHNC" + config.parallel_config.decode_context_parallel_size = 4 + common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) + specs = { + "target": MLAAttentionSpec( + **common, + head_size=576, + state_content_bytes=656, + block_stride_alignment=656, + cache_dtype_str="fp8_ds_mla", + is_index_group_leader=True, + ), + "indexer": MLAAttentionSpec( + **common, + head_size=132, + cache_role=SparseCacheRole.INDEXER, + block_stride_alignment=indexer_alignment, + ), + "draft": SlidingWindowSpec( + block_size=64, + num_kv_heads=64, + head_size=64, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_sharded=False, + ), + } + if not with_draft: + del specs["draft"] + groups = get_kv_cache_groups(config, specs) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, groups, available_memory=8 * 1024 * 1024 + ) + caches = allocate_kv_cache( + cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) + ) + cache = caches["target"].squeeze(1) + if indexer_alignment is not None: + assert caches["indexer"].stride(0) % indexer_alignment == 0 + rows, stride_rows = flat_kv_row_view(cache, 64) + cache[1, 0].fill_(7) + torch.testing.assert_close(rows[stride_rows], cache[1, 0]) + assert torch.all(rows[stride_rows] == 7) + + +def test_is_kv_cache_spec_uniform(): + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + script = """ +import sys + +from vllm.v1.core.kv_cache_planning import is_kv_cache_spec_uniform +from vllm.v1.kv_cache_interface import KVCacheSpec + +if sys.flags.optimize < 1: + raise RuntimeError("subprocess is not running with optimization enabled") +specs = { + "a": KVCacheSpec(block_size=1), + "b": KVCacheSpec(block_size=2), +} +if is_kv_cache_spec_uniform(specs): + raise RuntimeError("different specs were treated as uniform") +""" + subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[3], + check=True, + timeout=60, + ) + + +@pytest.mark.parametrize( + ("model_id", "max_model_len", "want_estimated_max_len"), + [ + ("Qwen/Qwen1.5-7B", 16385, 16384), + ("Qwen/Qwen1.5-7B", 16383, 16383), + ], +) +def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): + # Create a VllmConfig + model_config = ModelConfig( + model_id, + runner="generate", + dtype="float16", + max_model_len=max_model_len, + ) + scheduler_config = SchedulerConfig( + max_num_batched_tokens=32768, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + ) + + # Create KV cache specs + kv_cache_spec = {} + for i in range(32): + layer_name = f"layer_{i}" + kv_cache_spec[layer_name] = FullAttentionSpec( + block_size=16, + num_kv_heads=32, + head_size=128, + dtype=torch.float16, + ) + # Estimate the maximum model length, 16384 model_len need 8GB + estimated_max_len = estimate_max_model_len( + vllm_config, kv_cache_spec, 8 * GiB_bytes + ) + assert estimated_max_len == want_estimated_max_len + + +def test_get_kv_cache_config_one_worker(): + # pass max_model_len to pass check_enough_kv_cache_memory + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + # all layers are full attention -> single group + kv_cache_specs_full = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + kv_cache_config_full = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + print(kv_cache_config_full) + assert kv_cache_config_full == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + # all layers are sliding window -> single group + kv_cache_specs_sliding = { + "layer_1": new_sliding_window_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_sliding = get_kv_cache_configs( + vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_sliding == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) + ], + ) + + # full + sliding, but disable_hybrid_kv_cache_manager + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) + ), + ], + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + + # full + sliding, with hybrid_kv_cache_manager + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), + ], + ) + + # 2 full + 4 sliding, 2 layers per group + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_sliding_window_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_3", "layer_5"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_4", "layer_6"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), + ], + ) + + # 3 full + 7 sliding, pad to 3 full + 9 sliding + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1", "layer_2", "layer_3"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_4", "layer_7", "layer_10"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_5", "layer_8"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_6", "layer_9"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), + KVCacheGroupSpec( + ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() + ), + KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), + ], + ) + + # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss + # eagle where there is only one more full attention layer than sliding window layers + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_kv_cache_spec(), + "layer_5": new_kv_cache_spec(), + "layer_6": new_kv_cache_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + "layer_11": new_sliding_window_spec(), + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] + )[0] + print(kv_cache_config_hybrid) + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=[ + "layer_1", + "layer_2", + "layer_3", + "layer_4", + "layer_5", + "layer_6", + ], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], + new_kv_cache_spec(), + ), + KVCacheGroupSpec( + ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + new_sliding_window_spec(), + ), + ], + ) + + # different hidden size but same type, use UniformTypeKVCacheSpecs + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=128), + "layer_2": new_kv_cache_spec(head_size=64), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + # Layers of different page sizes pack densely into one allocation: the + # 2x-larger layer_1 takes the first region, layer_2 the next. + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 2 * 32, + block_stride=mem_per_block_per_layer * 2, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + offset=mem_per_block_per_layer * 2 * 32, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs_hybrid + ), + ) + ], + ) + + # Different hidden size and different type, align by different block size + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": new_sliding_window_spec(head_size=32), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), + KVCacheGroupSpec( + ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) + ), + ], + ) + + # different hidden size that cannot be aligned by using different block size, + # but can be aligned by padding the smaller physical page. + swa_spec = new_sliding_window_spec(head_size=96) + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": swa_spec, + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + padded_page_size = swa_spec.page_size_bytes + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=42, + kv_cache_tensors=[ + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_1"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_2"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1"], + new_kv_cache_spec( + head_size=64, + page_size_padded=padded_page_size, + ), + ), + KVCacheGroupSpec( + ["layer_2"], + new_sliding_window_spec(head_size=96), + ), + ], + ) + + # Test num_gpu_blocks_override + vllm_config.cache_config.num_gpu_blocks_override = 16 + kv_cache_config_override_blocks = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_override_blocks == KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 16 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 16, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + +def test_get_kv_cache_configs_attention_free(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None + kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[], + ) + ] + + +def test_packed_groups_glm5_like_hybrid(): + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) + vllm_config.cache_config.kv_cache_layout = "BLHNC" + kpool, block_size = 4, 1024 + kv_cache_spec: dict[str, KVCacheSpec] = {} + for i in range(11): + kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + block_stride_alignment=32 * 576 * 2, + ) + kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=kpool, + block_stride_alignment=64 * 132, + ) + kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( + block_size=kpool, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + ) + for i in range(34): + kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( + shapes=((3, 6144), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.float32), + block_size=block_size, + ) + + groups = kv_cache_planning.get_kv_cache_groups(vllm_config, kv_cache_spec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) + ] + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) + ] + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec + ) + ] + assert len(attn) == 1 and len(attn[0].layer_names) == 22 + assert len(rings) == 1 and len(rings[0].layer_names) == 11 + # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. + assert len(mamba) == 3 + assert sum(len(g.layer_names) for g in mamba) == 34 + + kv_cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, groups, available_memory=2 * 1024**3 + ) + strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} + assert len(strides) == 1 + (block_stride,) = strides + # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages + # apart so both caches can be re-paged for their kernels. + assert block_stride % (32 * 576 * 2) == 0 + assert block_stride % (64 * 132) == 0 + + +def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): + return SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version=model_version, + ) + + +def new_indexer_mla_spec(block_size=16): + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller + # page size than the main MLA attention (uint8, small head), so their pages + # cannot be unified. + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + ) + + +def test_mixed_page_size_groups_use_spec_compatibility(): + specs = {} + for i in range(3): + specs[f"mla.{i}"] = new_mla_spec() + specs[f"indexer.{i}"] = new_indexer_mla_spec() + specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) + + config = _grouping_config() + config.cache_config = CacheConfig() + config.cache_config.kv_cache_layout = "BLNHC" + groups = get_kv_cache_groups(config, specs) + + assert len(groups) == 3 + assert {name for group in groups for name in group.layer_names} == set(specs) + assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] + + +def _grouping_config(): + cache_config = CacheConfig() + cache_config.kv_cache_layout = "LBNHC" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + cache_config=cache_config, + ) + + +def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): + block_size = 544 + full_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + # FullAttentionSpec stores both K and V, so this produces a + # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. + head_size=512, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + block_size=block_size, + shapes=((557056,),), + dtypes=(torch.bfloat16,), + mamba_cache_mode="align", + ) + hidden_spec = HiddenStateCacheSpec( + block_size=block_size, + num_kv_heads=3, + head_size=1024, + dtype=torch.bfloat16, + ) + assert full_spec.page_size_bytes == mamba_spec.page_size_bytes + + groups = get_kv_cache_groups( + _grouping_config(), + { + "model.full_attn": full_spec, + "model.mamba": mamba_spec, + "cache_only_layers.0": hidden_spec, + }, + ) + + hidden_group = next( + group + for group in groups + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ) + assert hidden_group.kv_cache_spec.block_size == 136 + + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=None, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + kv_transfer_config=object(), + ) + assert kv_cache_utils.resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) == (544, 136) + + +def test_multi_run_layer_compact_strides_place_hoisted_heads(): + """A layer-compact run region is its own dense allocation: under LHBNC the head + groups sit between the layers and the blocks, so a run's block stride is one head + group's slice, not the whole page (regression test for setStorage out-of-bounds + on ROCm hybrid models).""" + full = new_kv_cache_spec() + swa = new_sliding_window_spec(sliding_window=full.block_size * 2) + page = full.page_size_bytes + assert swa.page_size_bytes == page + + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.kv_cache_layout = "LHBNC" + config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, + [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["swa.0", "swa.1"], swa), + ], + available_memory=4 * page * 8, + ) + + num_blocks = config.num_blocks + head_group = full.block_size * full.state_content_size_bytes + for tensor in config.kv_cache_tensors: + # One head group per block, head groups between the layers and blocks. + assert tensor.block_stride == head_group + assert tensor.layer_stride == full.num_heads * num_blocks * head_group + # The last block of the last layer's last head group must stay within the + # allocation (the old page-based stride overran it). + end = ( + tensor.offset + + tensor.layer_stride * (len(tensor.layers) - 1) + + (full.num_heads - 1) * num_blocks * head_group + + tensor.block_stride * (num_blocks - 1) + + head_group + ) + assert end <= tensor.size + + +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): + specs = { + "target.0.attn": new_mla_spec(), + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups(_grouping_config(), specs) + + assert len(groups) == 2 + assert all( + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + + +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention + # compute sliding-window while promoting only its allocation semantics lets + # every layer share the target's block table and remain contiguous. + draft = new_sliding_window_spec(block_size=16) + specs = { + "target.0.attn": new_mla_spec(block_size=64), + "target.0.indexer": new_indexer_mla_spec(block_size=64), + "draft.0": draft, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + group_spec = groups[0].kv_cache_spec + assert isinstance(group_spec, UniformTypeKVCacheSpecs) + assert group_spec.block_size == 64 + promoted_draft = group_spec.kv_cache_specs["draft.0"] + assert isinstance(promoted_draft, FullAttentionSpec) + assert not isinstance(promoted_draft, SlidingWindowSpec) + assert promoted_draft.block_size == 64 + assert promoted_draft.sliding_window == draft.sliding_window + assert specs["draft.0"] is draft + assert "attention compute is unchanged" in caplog_vllm.text + + +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + +def test_auto_fit_max_model_len(): + """Test that max_model_len=-1 auto-fits to available GPU memory.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=1024) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # With enough memory, max_model_len stays at the derived max + large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [large_available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + # Reset for next test + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + # With limited memory, max_model_len should be reduced + # Need memory for at least max_model_len tokens + # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens + limited_memory = mem_per_block_per_layer * 2 * 32 + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [limited_memory] + ) + # Should be reduced to fit in memory + assert vllm_config.model_config.max_model_len < 1024 + assert vllm_config.model_config.max_model_len > 0 + + +def test_auto_fit_max_model_len_with_hybrid(): + """Test that auto-fit works with hybrid KV cache specs.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=8192) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + gamma = 2 + kv_cache_specs = { + "layer_1": new_mamba_spec(num_speculative_blocks=gamma), + "layer_2": new_kv_cache_spec(), + } + + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + +def test_auto_fit_max_model_len_not_triggered(): + """Test that auto-fit is not triggered when original_max_model_len is not -1.""" + model_config = ModelConfig(max_model_len=16) + # original_max_model_len should be None by default, not -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # This should work normally without auto-fit + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] + ) + assert vllm_config.model_config.max_model_len == 16 + + +def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): + """Auto-fit must size max_model_len against the override-clamped pool, not + the raw `available_memory`. Without this, auto-fit could pick a + max_model_len that no longer fits once `num_gpu_blocks_override` is applied. + """ + model_config = ModelConfig(max_model_len=16384) + model_config.original_max_model_len = -1 # request auto-fit + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # Cap the cache to 32 blocks regardless of available memory. + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), # block_size=16 + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + # 32 blocks * block_size 16 = 512 token slots, so max_model_len must + # auto-fit at or below that. + assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 + + +def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): + """Admission check must use the override-clamped pool size, not raw + `available_memory`. Without this, startup could accept a max_model_len + that does not actually fit in `num_gpu_blocks_override` blocks. + """ + model_config = ModelConfig(max_model_len=16384) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # 32 blocks is far too small for max_model_len=16384 (would need 1024). + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory: a bytes-only check against this would pass. + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): + """DFlash drafters can have a smaller head size than the target model. + + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer + block-size multiple, so the smaller page must be padded instead. + """ + # Attention layers read padded pages through a strided view, so padding is + # allowed. + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + ) + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + ) + + unified_specs = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "target_attn": target_spec, + "draft_attn": draft_spec, + } + ) + + assert unified_specs["target_attn"] == target_spec + unified_draft_spec = unified_specs["draft_attn"] + assert unified_draft_spec.block_size == draft_spec.block_size + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes + + +def test_unify_hybrid_kv_cache_specs(): + # 1. has_full_attention and has_sliding_window + before_spec_1 = new_kv_cache_spec(block_size=64) + before_spec_2 = new_sliding_window_spec( + block_size=16, page_size_padded=32 * 1024, sliding_window=1024 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) + kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec(block_size=64) + expected_spec_2 = new_kv_cache_spec( + block_size=64, page_size_padded=64 * 1024, sliding_window=1024 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 + assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) + + # 2. has_full_attention and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + + # 3. has_full_attention, has_sliding_window and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_sliding_window_spec( + page_size_padded=32 * 1024, sliding_window=1024 + ) + before_spec_3 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + "layer_3": before_spec_3, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) + expected_spec_3 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_3"] == expected_spec_3 + + # 4. No FullAttentionSpec, should not convert + kv_cache_spec = { + "layer_1": new_sliding_window_spec(sliding_window=1024), + "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), + } + + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + # Replicated Mamba state still requires the hybrid cache manager. + kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + +def test_unify_kv_cache_spec_page_size_mamba(): + """Regression test for https://github.com/vllm-project/vllm/issues/43626. + + MambaSpec's page_size_bytes is determined by its state shapes and does not + change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba + page instead of scaling its block_size. This situation arises when a layer + with a page larger than the (already platform-aligned) Mamba page joins the + specs, e.g. a dense draft model with more KV heads than the hybrid main + model. + """ + # 1. Hybrid main model (Mamba + full attention, pages already aligned at + # 16KB) plus a dense draft model layer with a 2x larger page (32KB). + # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so + # the old code scaled the Mamba block_size, which left page_size_bytes + # unchanged at 16384. + mamba_spec = new_mamba_spec() # page_size_bytes = 16384 + main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 + draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 + assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 + assert draft_attn_spec.page_size_bytes == 32768 + + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": mamba_spec, + "main_attn_layer": main_attn_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + # Mamba page is padded; block_size (caching granularity) is unchanged. + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + assert unified["mamba_layer"].block_size == mamba_spec.block_size + # Attention layer with smaller page still unifies by scaling block_size. + assert unified["main_attn_layer"].page_size_bytes == 32768 + assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size + assert unified["main_attn_layer"].page_size_padded is None + # Layer already at max page size is unchanged. + assert unified["draft_attn_layer"] == draft_attn_spec + + # 2. Mamba page already padded by the platform (state smaller than the + # padded page); the padding is re-applied at the new maximum. + padded_mamba_spec = new_mamba_spec( + shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 + ) + assert padded_mamba_spec.page_size_bytes == 16384 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": padded_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + + # 3. Mamba page that does not evenly divide the maximum page size is + # padded as well (the divisibility constraint only applies to block_size + # scaling). + odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) + assert odd_mamba_spec.page_size_bytes == 24576 + assert 32768 % odd_mamba_spec.page_size_bytes != 0 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": odd_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + + # 4. Attention layers with non-divisible page sizes are padded too: every + # backend reads a padded page through the view's block stride, so there is + # no longer a case that must raise. + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + assert unified["attn_layer"].page_size_padded == 32768 + assert unified["attn_layer"].page_size_bytes == 32768 + assert unified["attn_layer"].block_size == 24 + + # 5. Uniform page sizes are returned unchanged. + specs = { + "mamba_layer": new_mamba_spec(), + "attn_layer": new_kv_cache_spec(), + } + assert kv_cache_planning.unify_kv_cache_spec_page_size(specs) == specs + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) + + +def _spec_decode_grouping_config(method="dspark", model_type=None): + """Grouping config with an EAGLE-family speculative method enabled.""" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: SimpleNamespace( + is_block_outermost=True + ) + ), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), + speculative_config=SimpleNamespace( + method=method, + use_eagle=lambda: True, + use_eagle_block_drop=lambda: True, + ), + ) + + +def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): + """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a + DSpark-style draft MLA layer marked non_causal_multi_token_decode. + + The target's fp8 KV dtype is one difference from the draft on Kimi-K3 + (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass + draft_shares_target_spec to match that dtype while retaining the draft + marker as a group boundary. + """ + target_dtype = None if draft_shares_target_spec else "fp8_e4m3" + specs = { + "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + } + if draft: + draft_spec = new_mla_spec(block_size=64) + specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) + return specs + + +def test_draft_group_annotated_on_hybrid_general_path(): + # A drafter's MLA layer carries non_causal_multi_token_decode, so its group + # is identifiable without keying off a model version. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "draft.attn.0" in flagged[0].layer_names + + +def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): + # Packed uniform-type groups can contain distinct target and draft layer + # specs; the combined group still holds volatile draft KV and must be + # flagged. What must never happen is a Mamba group being flagged: that + # widens its lookup window to two consecutive chunks, which align-mode + # checkpointing never produces, zeroing every lookup. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), + _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), + ) + + for group in groups: + if "draft.attn.0" in group.layer_names: + assert group.is_eagle_group + if isinstance(group.kv_cache_spec, MambaSpec): + assert not group.is_eagle_group + + +def test_draft_group_not_annotated_without_spec_decode(): + # The marker alone must not flag anything; the eagle semantics only apply + # when a speculative method is actually enabled. + config = _spec_decode_grouping_config() + config.speculative_config = None + groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) + + assert not any(g.is_eagle_group for g in groups) + + +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) + ) + + assert not any(g.is_eagle_group for g in groups) + assert "no KV cache group could be identified as the draft model's" in ( + caplog_vllm.text + ) + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def test_no_warning_when_draft_group_is_identified(caplog_vllm): + get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def _deepseek_v4_specs(model_version="deepseek_v4"): + """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers + at differing page sizes, with the MTP draft layer registered last.""" + return { + "model.layers.0.self_attn.attn": new_mla_spec(), + "model.layers.1.self_attn.attn": new_mla_spec(), + "model.layers.2.self_attn.attn": new_swa_mla_spec( + head_size=1024, model_version=model_version + ), + # The MTP block registers last, and its sliding-window size differs, so + # it lands in a group of its own. + "model.layers.3.self_attn.attn": new_swa_mla_spec( + head_size=1024, sliding_window=256, model_version=model_version + ), + } + + +@pytest.mark.parametrize( + ("method", "model_type"), + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], +) +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec + # carries no draft marker and only the positional rule can find it. This + # pins the pre-existing behaviour that the unified annotator must preserve. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type=model_type), + _deepseek_v4_specs(model_version=None), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def test_trailing_layer_fallback_applies_to_any_mtp_model(): + # The positional rule is sound for every MTP drafter, not just DeepseekV4: + # MTP blocks reuse the target's decoder layer (no spec marker) and always + # register after every target layer. The model_type must not gate it. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="other"), + _deepseek_v4_specs(), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def _qwen3_5_hybrid_specs(with_mtp_layer: bool): + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with + the MTP drafter's full-attn layer (spec-identical to the target's) + registered last.""" + specs = {} + idx = 0 + for _ in range(2): + for _ in range(3): + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( + mamba_cache_mode="align" + ) + idx += 1 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() + idx += 1 + if with_mtp_layer: + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() + return specs + + +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): + # A hybrid mamba + full-attention model with an MTP drafter that is + # spec-indistinguishable from the target reaches the general multi-group + # path. The trailing-layer rule must locate the draft group there so the + # Mamba groups are not swept up by the flag-all consumer fallback. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names + for group in groups: + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + assert not group.is_eagle_group + assert "could be identified as the draft model's" not in caplog_vllm.text + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): + # Other drafters are not covered by the trailing-layer rule, so an + # unidentifiable hybrid draft must still warn. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" in caplog_vllm.text + + +def test_trailing_layer_fallback_requires_exact_partition(): + # If the groups do not partition the layers exactly (e.g. a caller that + # dropped or duplicated layers), the positional rule is meaningless and + # must not fire. + from vllm.v1.core.kv_cache_planning import _annotate_eagle_groups + + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") + groups = get_kv_cache_groups(config, specs) + for g in groups: + g.is_eagle_group = False + # Remove one layer from its group: no longer an exact partition. + trimmed = [ + KVCacheGroupSpec( + [n for n in g.layer_names if n != "model.layers.0.linear_attn"], + g.kv_cache_spec, + ) + for g in groups + ] + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) + + assert not any(g.is_eagle_group for g in trimmed) + + +_GQA_SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 +) + + +_GQA_SWA_SPEC = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 +) + + +@pytest.mark.parametrize( + "specs,expected", + [ + pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": new_swa_mla_spec(), + "s.1": new_swa_mla_spec(), + }, + 4, + id="mla-and-swa-mla", + ), + pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), + pytest.param( + {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, + 1, + id="mla-and-gqa-uniform-group", + ), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": _GQA_SWA_SPEC, + "s.1": _GQA_SWA_SPEC, + }, + 1, + id="mla-and-gqa-swa", + ), + ], +) +def test_kv_tp_replicas(monkeypatch, specs, expected): + """Resolved per layer before scheduler flattening hides mixed groups.""" + from vllm.config import ParallelConfig + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "device_count", lambda: 4) + parallel_config = ParallelConfig(tensor_parallel_size=4) + vllm_config = VllmConfig( + model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config + ) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + mem = sum(s.page_size_bytes for s in specs.values()) * 10 + configs = get_kv_cache_configs(vllm_config, [specs], [mem]) + assert configs[0].kv_tp_replicas == expected + scheduler = generate_scheduler_kv_cache_config(configs) + assert scheduler.kv_tp_replicas == expected + + +@pytest.mark.parametrize( + "specs,tp_size,dcp_size,expected", + [ + pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), + pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), + pytest.param( + [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" + ), + pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), + pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), + pytest.param( + [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], + 8, + 1, + 2, + id="mla-and-gqa-partial", + ), + pytest.param( + [ + HiddenStateCacheSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 + ) + ], + 8, + 1, + 1, + id="hidden-state", + ), + ], +) +def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): + """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" + groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] + assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index b06479f1caf1..1489315ec255 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,20 +3,23 @@ import copy import hashlib import importlib -import mmap -import subprocess -import sys from collections.abc import Callable from dataclasses import replace -from pathlib import Path from types import SimpleNamespace from typing import Any import pytest import torch +import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) from vllm.config import ( CacheConfig, ModelConfig, @@ -33,25 +36,21 @@ ) from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor -from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_planning import ( + get_kv_cache_groups, +) from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, - check_enough_kv_cache_memory, - estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, get_kv_cache_capacity, - get_kv_cache_configs, - get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, hash_block_tokens, init_none_hash, - is_kv_cache_spec_uniform, - kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, @@ -62,10 +61,8 @@ get_hisparse_steady_state_concurrency, ) from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, KVCacheConfig, @@ -93,110 +90,6 @@ pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize("gpu_block_size", [32, 64]) -@pytest.mark.parametrize("shared_host_pool", [False, True]) -def test_hisparse_hma_uses_resolved_gpu_block_size( - monkeypatch, gpu_block_size, shared_host_pool -): - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - specs = { - "model.layers.0.self_attn": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ), - "model.layers.0.self_attn.indexer": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ), - } - config = SimpleNamespace( - attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), - model_config=SimpleNamespace( - hf_config=SimpleNamespace(index_topk=128), - max_model_len=gpu_block_size, - ), - parallel_config=SimpleNamespace( - tensor_parallel_size=2 if shared_host_pool else 1, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, - world_size=2 if shared_host_pool else 1, - distributed_executor_backend="mp", - nnodes_within_dp=1, - ), - cache_config=SimpleNamespace( - num_gpu_blocks_override=7, - prefix_cache_retention_interval=None, - get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, - ), - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 - ) - assert cache_config.num_blocks == 7 - assert cache_config.hisparse_host_num_blocks is not None - assert cache_config.hisparse_host_num_blocks > 7 - - host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups - assert cache_config.hisparse_shared_host_pool is shared_host_pool - host_page = host_group.kv_cache_spec.page_size_bytes - alignment = mmap.PAGESIZE if shared_host_pool else 1 - expected_host_stride = (host_page + alignment - 1) // alignment * alignment - assert cache_config.hisparse_host_block_stride == expected_host_stride - assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride - assert host_group.host_resident - assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) - host_layers = set(host_group.layer_names) - for tensor in cache_config.kv_cache_tensors: - assert all( - (name in host_layers) == tensor.host_resident for name in tensor.layers - ) - host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) - host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] - assert host_tensor.block_stride == host_spec.page_size_bytes - assert host_tensor.layer_stride == ( - host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks - ) - assert host_group.kv_cache_spec.block_size == gpu_block_size - assert indexer_group.kv_cache_spec.block_size == gpu_block_size - host_specs = host_group.kv_cache_spec.kv_cache_specs - gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs - assert set(host_specs) == {"model.layers.0.self_attn"} - assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} - assert indexer_group.kv_cache_spec.prefix_cacheable - assert host_group.enable_kv_transfer - auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] - assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) - assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) - assert all( - spec.block_size == gpu_block_size - for spec in auxiliary_specs - if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) - ) - scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( - cache_config, - SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), - ), - ) - assert scheduler_block_size == hash_block_size == gpu_block_size - - def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): """Running requests pin only their resident tail; the newest still needs its full in-flight window to be admitted.""" @@ -281,68 +174,6 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok create_hisparse_layout(config, groups, host_budget=host_budget) -@pytest.mark.parametrize( - "max_model_len,num_gpu_blocks,ok", - [ - (8192, 160, False), - (8192, 1000, True), - (32768, 656, False), - (32768, 657, True), - (-1, 657, True), - ], -) -def test_hisparse_pool_must_fit_max_model_len( - monkeypatch, max_model_len, num_gpu_blocks, ok -): - """Each resident group takes its own blocks from HiSparse's shared GPU pool, - so a pool that fits only the indexer pages must be rejected at startup. - Resident pages are capped at the in-flight window: at 32768 tokens with 2048 - in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, - plus the null block. Auto-fit (-1) must pick a length that passes the same - check.""" - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - specs: dict[str, KVCacheSpec] = {} - for i in range(4): - specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ) - specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ) - config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - config.scheduler_config.max_num_batched_tokens = 2048 - # Async scheduling doubles the in-flight window, and with it the bound. - config.scheduler_config.async_scheduling = False - config.attention_config.hisparse_config = HiSparseConfig() - config.model_config.hf_config.index_topk = 128 - config.cache_config.num_gpu_blocks_override = num_gpu_blocks - config.cache_config.kv_cache_layout = "BLHNC" - # Workers report the resident/hot caches HiSparse derives. - specs |= { - name: group.kv_cache_spec - for group in get_hisparse_kv_cache_groups(config, specs) or [] - if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) - for name in group.layer_names - } - - if ok: - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - else: - with pytest.raises(ValueError, match="max seq len"): - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - - def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): """Workers report HiSparse's resident/hot caches, which lay out their own raw backing; KV cache layout resolution must skip them.""" @@ -496,28 +327,6 @@ def test_effective_attention_block_size_matches_events(dcp): ] == [block_size] -def new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=None, - attention_chunk_size=None, - kv_quant_mode=KVQuantMode.NONE, -): - return FullAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - attention_chunk_size=attention_chunk_size, - kv_quant_mode=kv_quant_mode, - ) - - def test_kv_cache_config_selects_only_transferable_groups(): """Connectors must see a stable projection of transfer-eligible groups.""" groups = [ @@ -581,60 +390,6 @@ def test_kv_cache_blocks_selects_requested_groups(): assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) -def new_sliding_window_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=1, -): - return SlidingWindowSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - ) - - -def new_chunked_local_attention_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - attention_chunk_size=4, -): - return ChunkedLocalAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - attention_chunk_size=attention_chunk_size, - ) - - -def new_mamba_spec( - block_size=16, - shapes=((2, 512), (3, 32, 32)), - dtypes=(torch.float32, torch.float32), - num_speculative_blocks=2, - mamba_cache_mode="none", - page_size_padded=None, -): - return MambaSpec( - block_size=block_size, - shapes=shapes, - dtypes=dtypes, - page_size_padded=page_size_padded, - mamba_cache_mode=mamba_cache_mode, - num_speculative_blocks=num_speculative_blocks, - ) - - @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1441,430 +1196,6 @@ def test_metrics_empty_stats(): assert metrics.hit_rate == 0.75 -def test_get_kv_cache_configs_multiple_workers(): - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - same_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - ] - - # Basic case. All things are the same. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different available memory. This is the case for TP. - # Use the smallest memory available. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 20, - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different KV cache specs. This is the case for PP. - different_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer2": new_kv_cache_spec(), - "layer3": new_kv_cache_spec(), - }, - ] - - # Different workers have different layers. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer2", "layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), - ], - ), - ] - - # Some layers are the same, some are different. This is the case for TP+PP - tp_pp_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - ] - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - tp_pp_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected_12 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - expected_3 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [ - expected_12, - expected_12, - expected_3, - expected_3, - ] - - # Different workers have different types of layers. This is the case for - # hybrid models + PP. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_sliding_window_spec(), - "layer4": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - KVCacheGroupSpec([], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer3", "layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec([], ref_kv_cache_spec), - KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), - ], - ), - ] - - # When divided into multiple KVCacheGroups, need to ensure the number of - # layers per group is similar. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_sliding_window_spec(), - "layer3": new_sliding_window_spec(), - }, - { - "layer4": new_kv_cache_spec(), - "layer5": new_sliding_window_spec(), - "layer6": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 10, - ref_kv_cache_spec.page_size_bytes * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer5"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer6"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), - ], - ), - ] - - # Have conflicting layers. Need to raise an error. - conflicting_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer1": new_sliding_window_spec(), - }, - ] - with pytest.raises(AssertionError): - get_kv_cache_configs( - vllm_config, - conflicting_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - - -@pytest.mark.parametrize( - "asymmetric_memory", - [False, True], - ids=["symmetric", "asymmetric"], -) -def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): - model_config = ModelConfig(max_model_len=512) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - pp_kv_cache_specs = [ - {"layer1": ref_kv_cache_spec}, - {"layer2": ref_kv_cache_spec}, - ] - - expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 - avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks - - # With per-worker validation, each worker only needs memory for its own - # layers. Worker 2 having more memory shouldn't affect worker 1's config. - available_memory = ( - [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 - ) - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - pp_kv_cache_specs, - available_memory, - ) - - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], - ), - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], - ), - ] - - -def test_project_kv_cache_groups_to_worker(): - spec_a = new_kv_cache_spec() - spec_b = new_kv_cache_spec(num_kv_heads=4) - - global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), - ] - worker_spec = {"layer1": spec_a, "layer2": spec_a} - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer2"] - assert projected[0].kv_cache_spec is spec_a - - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, {"layer4": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == [] - assert projected[0].kv_cache_spec is spec_a - - uniform_spec = UniformTypeKVCacheSpecs( - block_size=16, - kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, - ) - global_groups_uniform = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), - ] - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer3"] - proj_spec = projected[0].kv_cache_spec - assert isinstance(proj_spec, UniformTypeKVCacheSpecs) - assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} - - @pytest.mark.parametrize("sliding_window", [None, 256]) @pytest.mark.parametrize("disable_hybrid", [False, True]) @pytest.mark.parametrize("pcp_size", [1, 4]) @@ -1901,7 +1232,7 @@ def test_dcp_target_allocates_replicated_draft_independently( widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups] assert widths == [16, 64] - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( config, groups, available_memory=16 * 1024 * 1024 ) scheduler_config = generate_scheduler_kv_cache_config([cache_config]) @@ -1971,58 +1302,6 @@ def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): type(draft).merge([draft, replace(draft, dcp_sharded=True)]) -@pytest.mark.parametrize("with_draft", [False, True]) -@pytest.mark.parametrize("indexer_alignment", [None, 512]) -def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): - from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view - from vllm.v1.worker.utils import allocate_kv_cache - - config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - config.cache_config.kv_cache_layout = "BLHNC" - config.parallel_config.decode_context_parallel_size = 4 - common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) - specs = { - "target": MLAAttentionSpec( - **common, - head_size=576, - state_content_bytes=656, - block_stride_alignment=656, - cache_dtype_str="fp8_ds_mla", - is_index_group_leader=True, - ), - "indexer": MLAAttentionSpec( - **common, - head_size=132, - cache_role=SparseCacheRole.INDEXER, - block_stride_alignment=indexer_alignment, - ), - "draft": SlidingWindowSpec( - block_size=64, - num_kv_heads=64, - head_size=64, - dtype=torch.bfloat16, - sliding_window=2048, - dcp_sharded=False, - ), - } - if not with_draft: - del specs["draft"] - groups = get_kv_cache_groups(config, specs) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, groups, available_memory=8 * 1024 * 1024 - ) - caches = allocate_kv_cache( - cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) - ) - cache = caches["target"].squeeze(1) - if indexer_alignment is not None: - assert caches["indexer"].stride(0) % indexer_alignment == 0 - rows, stride_rows = flat_kv_row_view(cache, 64) - cache[1, 0].fill_(7) - torch.testing.assert_close(rows[stride_rows], cache[1, 0]) - assert torch.all(rows[stride_rows] == 7) - - @pytest.mark.parametrize( "layer_type,dcp_size,expected_width", [ @@ -2115,103 +1394,6 @@ def test_merge_kv_cache_spec(): assert merged_layer_spec.sliding_window == 1 -def test_is_kv_cache_spec_uniform(): - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - script = """ -import sys - -from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform -from vllm.v1.kv_cache_interface import KVCacheSpec - -if sys.flags.optimize < 1: - raise RuntimeError("subprocess is not running with optimization enabled") -specs = { - "a": KVCacheSpec(block_size=1), - "b": KVCacheSpec(block_size=2), -} -if is_kv_cache_spec_uniform(specs): - raise RuntimeError("different specs were treated as uniform") -""" - subprocess.run( - [sys.executable, "-O", "-c", script], - cwd=Path(__file__).resolve().parents[3], - check=True, - timeout=60, - ) - - -@pytest.mark.parametrize( - ("model_id", "max_model_len", "want_estimated_max_len"), - [ - ("Qwen/Qwen1.5-7B", 16385, 16384), - ("Qwen/Qwen1.5-7B", 16383, 16383), - ], -) -def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): - # Create a VllmConfig - model_config = ModelConfig( - model_id, - runner="generate", - dtype="float16", - max_model_len=max_model_len, - ) - scheduler_config = SchedulerConfig( - max_num_batched_tokens=32768, - max_model_len=model_config.max_model_len, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - - vllm_config = VllmConfig( - model_config=model_config, - scheduler_config=scheduler_config, - ) - - # Create KV cache specs - kv_cache_spec = {} - for i in range(32): - layer_name = f"layer_{i}" - kv_cache_spec[layer_name] = FullAttentionSpec( - block_size=16, - num_kv_heads=32, - head_size=128, - dtype=torch.float16, - ) - # Estimate the maximum model length, 16384 model_len need 8GB - estimated_max_len = estimate_max_model_len( - vllm_config, kv_cache_spec, 8 * GiB_bytes - ) - assert estimated_max_len == want_estimated_max_len - - def test_get_max_concurrency_for_kv_cache_config(): # Create a VllmConfig model_id = "Qwen/Qwen1.5-7B" @@ -2425,411 +1607,6 @@ def test_allocate_with_lookahead(): assert len(blocks.get_block_ids()[0]) == 2 -def test_get_kv_cache_config_one_worker(): - # pass max_model_len to pass check_enough_kv_cache_memory - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - # all layers are full attention -> single group - kv_cache_specs_full = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - kv_cache_config_full = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - print(kv_cache_config_full) - assert kv_cache_config_full == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - # all layers are sliding window -> single group - kv_cache_specs_sliding = { - "layer_1": new_sliding_window_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_sliding = get_kv_cache_configs( - vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_sliding == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) - ], - ) - - # full + sliding, but disable_hybrid_kv_cache_manager - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) - ), - ], - ) - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False - - # full + sliding, with hybrid_kv_cache_manager - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=64, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), - ], - ) - - # 2 full + 4 sliding, 2 layers per group - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_sliding_window_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_3", "layer_5"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_4", "layer_6"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), - ], - ) - - # 3 full + 7 sliding, pad to 3 full + 9 sliding - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1", "layer_2", "layer_3"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_4", "layer_7", "layer_10"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_5", "layer_8"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_6", "layer_9"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), - KVCacheGroupSpec( - ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() - ), - KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), - ], - ) - - # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss - # eagle where there is only one more full attention layer than sliding window layers - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_kv_cache_spec(), - "layer_5": new_kv_cache_spec(), - "layer_6": new_kv_cache_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - "layer_11": new_sliding_window_spec(), - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] - )[0] - print(kv_cache_config_hybrid) - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=[ - "layer_1", - "layer_2", - "layer_3", - "layer_4", - "layer_5", - "layer_6", - ], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], - new_kv_cache_spec(), - ), - KVCacheGroupSpec( - ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - new_sliding_window_spec(), - ), - ], - ) - - # different hidden size but same type, use UniformTypeKVCacheSpecs - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=128), - "layer_2": new_kv_cache_spec(head_size=64), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - # Layers of different page sizes pack densely into one allocation: the - # 2x-larger layer_1 takes the first region, layer_2 the next. - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 2 * 32, - block_stride=mem_per_block_per_layer * 2, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - offset=mem_per_block_per_layer * 2 * 32, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs_hybrid - ), - ) - ], - ) - - # Different hidden size and different type, align by different block size - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": new_sliding_window_spec(head_size=32), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), - KVCacheGroupSpec( - ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) - ), - ], - ) - - # different hidden size that cannot be aligned by using different block size, - # but can be aligned by padding the smaller physical page. - swa_spec = new_sliding_window_spec(head_size=96) - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": swa_spec, - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - padded_page_size = swa_spec.page_size_bytes - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=42, - kv_cache_tensors=[ - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_1"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_2"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1"], - new_kv_cache_spec( - head_size=64, - page_size_padded=padded_page_size, - ), - ), - KVCacheGroupSpec( - ["layer_2"], - new_sliding_window_spec(head_size=96), - ), - ], - ) - - # Test num_gpu_blocks_override - vllm_config.cache_config.num_gpu_blocks_override = 16 - kv_cache_config_override_blocks = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_override_blocks == KVCacheConfig( - num_blocks=16, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 16 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 16, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_get_kv_cache_configs_attention_free(): - kv_cache_specs: dict[str, KVCacheSpec] = {} - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.prefix_cache_retention_interval = None - kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=[], - ) - ] - - def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2879,223 +1656,28 @@ def test_generate_uniform_type_kv_cache_specs(): def test_generate_scheduler_kv_cache_config(): kv_cache_specs = { "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(head_size=128), - } - kv_cache_configs = [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs - ), - ), - ], - ) - ] - scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) - assert scheduler_kv_cache_config == KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_packed_groups_glm5_like_hybrid(): - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192)) - vllm_config.cache_config.kv_cache_layout = "BLHNC" - kpool, block_size = 4, 1024 - kv_cache_spec: dict[str, KVCacheSpec] = {} - for i in range(11): - kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - block_stride_alignment=32 * 576 * 2, - ) - kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - tokens_per_state=kpool, - block_stride_alignment=64 * 132, - ) - kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec( - block_size=kpool, - num_kv_heads=2, - head_size=128, - head_size_v=0, - dtype=torch.bfloat16, - ) - for i in range(34): - kv_cache_spec[f"layers.{i}.kda"] = MambaSpec( - shapes=((3, 6144), (16, 128, 128)), - dtypes=(torch.bfloat16, torch.float32), - block_size=block_size, - ) - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - attn = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) - ] - rings = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) - ] - mamba = [ - g - for g in groups - if isinstance( - getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec - ) - ] - assert len(attn) == 1 and len(attn[0].layer_names) == 22 - assert len(rings) == 1 and len(rings[0].layer_names) == 11 - # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states. - assert len(mamba) == 3 - assert sum(len(g.layer_names) for g in mamba) == 34 - - kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, groups, available_memory=2 * 1024**3 - ) - strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors} - assert len(strides) == 1 - (block_stride,) = strides - # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages - # apart so both caches can be re-paged for their kernels. - assert block_stride % (32 * 576 * 2) == 0 - assert block_stride % (64 * 132) == 0 - - -def new_mla_spec(cache_dtype_str=None, block_size: int = 16): - # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - max_tp_shards=1, - head_size=576, - dtype=torch.float32, - cache_dtype_str=cache_dtype_str, - ) - - -def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): - return SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - max_tp_shards=1, - head_size=head_size, - dtype=torch.float32, - sliding_window=sliding_window, - model_version=model_version, - ) - - -def new_indexer_mla_spec(block_size=16): - # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller - # page size than the main MLA attention (uint8, small head), so their pages - # cannot be unified. - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - ) - - -def test_mixed_page_size_groups_use_spec_compatibility(): - specs = {} - for i in range(3): - specs[f"mla.{i}"] = new_mla_spec() - specs[f"indexer.{i}"] = new_indexer_mla_spec() - specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) - - config = _grouping_config() - config.cache_config = CacheConfig() - config.cache_config.kv_cache_layout = "BLNHC" - groups = get_kv_cache_groups(config, specs) - - assert len(groups) == 3 - assert {name for group in groups for name in group.layer_names} == set(specs) - assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] - - -def _grouping_config(): - cache_config = CacheConfig() - cache_config.kv_cache_layout = "LBNHC" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - speculative_config=None, - cache_config=cache_config, - ) - - -def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): - block_size = 544 - full_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - # FullAttentionSpec stores both K and V, so this produces a - # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. - head_size=512, - dtype=torch.bfloat16, - ) - mamba_spec = MambaSpec( - block_size=block_size, - shapes=((557056,),), - dtypes=(torch.bfloat16,), - mamba_cache_mode="align", - ) - hidden_spec = HiddenStateCacheSpec( - block_size=block_size, - num_kv_heads=3, - head_size=1024, - dtype=torch.bfloat16, - ) - assert full_spec.page_size_bytes == mamba_spec.page_size_bytes - - groups = get_kv_cache_groups( - _grouping_config(), - { - "model.full_attn": full_spec, - "model.mamba": mamba_spec, - "cache_only_layers.0": hidden_spec, - }, - ) - - hidden_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) - ) - assert hidden_group.kv_cache_spec.block_size == 136 - - kv_cache_config = KVCacheConfig( - num_blocks=1, + "layer_2": new_kv_cache_spec(head_size=128), + } + kv_cache_configs = [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs + ), + ), + ], + ) + ] + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) + assert scheduler_kv_cache_config == KVCacheConfig( + num_blocks=10, kv_cache_tensors=[], - kv_cache_groups=groups, - ) - vllm_config = SimpleNamespace( - cache_config=SimpleNamespace( - block_size=16, - enable_prefix_caching=True, - prefix_match_unit=None, - ), - parallel_config=SimpleNamespace(decode_context_parallel_size=1), - kv_transfer_config=object(), + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) == (544, 136) def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): @@ -3133,125 +1715,6 @@ def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba -def test_multi_run_layer_compact_strides_place_hoisted_heads(): - """A layer-compact run region is its own dense allocation: under LHBNC the head - groups sit between the layers and the blocks, so a run's block stride is one head - group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models).""" - full = new_kv_cache_spec() - swa = new_sliding_window_spec(sliding_window=full.block_size * 2) - page = full.page_size_bytes - assert swa.page_size_bytes == page - - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.kv_cache_layout = "LHBNC" - config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, - [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["swa.0", "swa.1"], swa), - ], - available_memory=4 * page * 8, - ) - - num_blocks = config.num_blocks - head_group = full.block_size * full.state_content_size_bytes - for tensor in config.kv_cache_tensors: - # One head group per block, head groups between the layers and blocks. - assert tensor.block_stride == head_group - assert tensor.layer_stride == full.num_heads * num_blocks * head_group - # The last block of the last layer's last head group must stay within the - # allocation (the old page-based stride overran it). - end = ( - tensor.offset - + tensor.layer_stride * (len(tensor.layers) - 1) - + (full.num_heads - 1) * num_blocks * head_group - + tensor.block_stride * (num_blocks - 1) - + head_group - ) - assert end <= tensor.size - - -def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): - specs = { - "target.0.attn": new_mla_spec(), - "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), - } - assert len({spec.page_size_bytes for spec in specs.values()}) == 1 - - groups = get_kv_cache_groups(_grouping_config(), specs) - - assert len(groups) == 2 - assert all( - not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups - ) - - -def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): - # Sparse MLA pages cannot be padded safely. Keeping the draft's attention - # compute sliding-window while promoting only its allocation semantics lets - # every layer share the target's block table and remain contiguous. - draft = new_sliding_window_spec(block_size=16) - specs = { - "target.0.attn": new_mla_spec(block_size=64), - "target.0.indexer": new_indexer_mla_spec(block_size=64), - "draft.0": draft, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - assert len(groups) == 1 - assert set(groups[0].layer_names) == set(specs) - group_spec = groups[0].kv_cache_spec - assert isinstance(group_spec, UniformTypeKVCacheSpecs) - assert group_spec.block_size == 64 - promoted_draft = group_spec.kv_cache_specs["draft.0"] - assert isinstance(promoted_draft, FullAttentionSpec) - assert not isinstance(promoted_draft, SlidingWindowSpec) - assert promoted_draft.block_size == 64 - assert promoted_draft.sliding_window == draft.sliding_window - assert specs["draft.0"] is draft - assert "attention compute is unchanged" in caplog_vllm.text - - -def test_hidden_states_with_tp_scales_page_size(): - """When TP shrinks KV pages below the hidden-state per-token cost, - get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states.""" - # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. - # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) - # = 16 * 2 * 64 * 2 = 4096 bytes. - kv_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.bfloat16, - ) - # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) - # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. - hs_spec = HiddenStateCacheSpec( - block_size=16, - num_kv_heads=6, - head_size=512, - dtype=torch.bfloat16, - ) - specs = { - "target.0.attn": kv_spec, - "target.1.attn": kv_spec, - "cache_only_layers.48": hs_spec, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - - # The hidden-state layer should be present and no assertion should fire. - all_layers = {name for g in groups for name in g.layer_names} - assert "cache_only_layers.48" in all_layers - - # The target group block sizes must have been scaled up. - for g in groups: - if "cache_only_layers.48" not in g.layer_names: - assert g.kv_cache_spec.block_size > kv_spec.block_size - - def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION @@ -3525,179 +1988,6 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes assert block_hashes[1] == expected_hash2 -def test_auto_fit_max_model_len(): - """Test that max_model_len=-1 auto-fits to available GPU memory.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=1024) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # With enough memory, max_model_len stays at the derived max - large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [large_available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - # Reset for next test - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # With limited memory, max_model_len should be reduced - # Need memory for at least max_model_len tokens - # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens - limited_memory = mem_per_block_per_layer * 2 * 32 - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [limited_memory] - ) - # Should be reduced to fit in memory - assert vllm_config.model_config.max_model_len < 1024 - assert vllm_config.model_config.max_model_len > 0 - - -def test_auto_fit_max_model_len_with_hybrid(): - """Test that auto-fit works with hybrid KV cache specs.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=8192) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - gamma = 2 - kv_cache_specs = { - "layer_1": new_mamba_spec(num_speculative_blocks=gamma), - "layer_2": new_kv_cache_spec(), - } - - # One extra block on top of what a 1024-token request needs: the pool - # reserves one block as the null block. - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - -def test_auto_fit_max_model_len_not_triggered(): - """Test that auto-fit is not triggered when original_max_model_len is not -1.""" - model_config = ModelConfig(max_model_len=16) - # original_max_model_len should be None by default, not -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # This should work normally without auto-fit - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] - ) - assert vllm_config.model_config.max_model_len == 16 - - -def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): - """Auto-fit must size max_model_len against the override-clamped pool, not - the raw `available_memory`. Without this, auto-fit could pick a - max_model_len that no longer fits once `num_gpu_blocks_override` is applied. - """ - model_config = ModelConfig(max_model_len=16384) - model_config.original_max_model_len = -1 # request auto-fit - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # Cap the cache to 32 blocks regardless of available memory. - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), # block_size=16 - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - # 32 blocks * block_size 16 = 512 token slots, so max_model_len must - # auto-fit at or below that. - assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 - - -def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): - """Admission check must use the override-clamped pool size, not raw - `available_memory`. Without this, startup could accept a max_model_len - that does not actually fit in `num_gpu_blocks_override` blocks. - """ - model_config = ModelConfig(max_model_len=16384) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # 32 blocks is far too small for max_model_len=16384 (would need 1024). - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory: a bytes-only check against this would pass. - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - -def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): - """DFlash drafters can have a smaller head size than the target model. - - For example, MiMo uses 192-dim target KV heads while its DFlash draft uses - 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer - block-size multiple, so the smaller page must be padded instead. - """ - # Attention layers read padded pages through a strided view, so padding is - # allowed. - target_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=1, - head_size=192, - dtype=torch.bfloat16, - ) - draft_spec = new_sliding_window_spec( - block_size=16, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - sliding_window=1024, - ) - - unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "target_attn": target_spec, - "draft_attn": draft_spec, - } - ) - - assert unified_specs["target_attn"] == target_spec - unified_draft_spec = unified_specs["draft_attn"] - assert unified_draft_spec.block_size == draft_spec.block_size - assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes - assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes - assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # raw KV allocation, so they must be budgeted into the offload width. The @@ -3729,173 +2019,6 @@ def test_page_size_padded_wins(): assert spec.page_size_bytes == 65536 -def test_unify_hybrid_kv_cache_specs(): - # 1. has_full_attention and has_sliding_window - before_spec_1 = new_kv_cache_spec(block_size=64) - before_spec_2 = new_sliding_window_spec( - block_size=16, page_size_padded=32 * 1024, sliding_window=1024 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) - kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec(block_size=64) - expected_spec_2 = new_kv_cache_spec( - block_size=64, page_size_padded=64 * 1024, sliding_window=1024 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 - assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) - - # 2. has_full_attention and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - - # 3. has_full_attention, has_sliding_window and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_sliding_window_spec( - page_size_padded=32 * 1024, sliding_window=1024 - ) - before_spec_3 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - "layer_3": before_spec_3, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) - expected_spec_3 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_3"] == expected_spec_3 - - # 4. No FullAttentionSpec, should not convert - kv_cache_spec = { - "layer_1": new_sliding_window_spec(sliding_window=1024), - "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), - } - - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - # Replicated Mamba state still requires the hybrid cache manager. - kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - -def test_unify_kv_cache_spec_page_size_mamba(): - """Regression test for https://github.com/vllm-project/vllm/issues/43626. - - MambaSpec's page_size_bytes is determined by its state shapes and does not - change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba - page instead of scaling its block_size. This situation arises when a layer - with a page larger than the (already platform-aligned) Mamba page joins the - specs, e.g. a dense draft model with more KV heads than the hybrid main - model. - """ - # 1. Hybrid main model (Mamba + full attention, pages already aligned at - # 16KB) plus a dense draft model layer with a 2x larger page (32KB). - # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so - # the old code scaled the Mamba block_size, which left page_size_bytes - # unchanged at 16384. - mamba_spec = new_mamba_spec() # page_size_bytes = 16384 - main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 - draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 - assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 - assert draft_attn_spec.page_size_bytes == 32768 - - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": mamba_spec, - "main_attn_layer": main_attn_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - # Mamba page is padded; block_size (caching granularity) is unchanged. - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - assert unified["mamba_layer"].block_size == mamba_spec.block_size - # Attention layer with smaller page still unifies by scaling block_size. - assert unified["main_attn_layer"].page_size_bytes == 32768 - assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size - assert unified["main_attn_layer"].page_size_padded is None - # Layer already at max page size is unchanged. - assert unified["draft_attn_layer"] == draft_attn_spec - - # 2. Mamba page already padded by the platform (state smaller than the - # padded page); the padding is re-applied at the new maximum. - padded_mamba_spec = new_mamba_spec( - shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 - ) - assert padded_mamba_spec.page_size_bytes == 16384 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": padded_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - - # 3. Mamba page that does not evenly divide the maximum page size is - # padded as well (the divisibility constraint only applies to block_size - # scaling). - odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) - assert odd_mamba_spec.page_size_bytes == 24576 - assert 32768 % odd_mamba_spec.page_size_bytes != 0 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": odd_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - - # 4. Attention layers with non-divisible page sizes are padded too: every - # backend reads a padded page through the view's block stride, so there is - # no longer a case that must raise. - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "attn_layer": new_kv_cache_spec(block_size=24), # 24576 - "draft_attn_layer": draft_attn_spec, # 32768 - } - ) - assert unified["attn_layer"].page_size_padded == 32768 - assert unified["attn_layer"].page_size_bytes == 32768 - assert unified["attn_layer"].block_size == 24 - - # 5. Uniform page sizes are returned unchanged. - specs = { - "mamba_layer": new_mamba_spec(), - "attn_layer": new_kv_cache_spec(), - } - assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs - - def test_hma_not_disabled_when_kv_events_enabled(): """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. @@ -3982,78 +2105,6 @@ def test_resolve_block_hashes_rejects_mismatched_view(): resolve_block_hashes(mismatched, 2, 4) -@pytest.mark.parametrize("use_override", [True, False]) -def test_kv_cache_reserves_null_block_for_max_model_len(use_override): - """A KV cache sized to exactly the blocks a max_model_len request needs must - be rejected. BlockPool keeps one block as the null block, so only - num_blocks - 1 are usable; accepting num_blocks == needed would leave the - request unschedulable and hang the engine. Covers both the - num_gpu_blocks_override path and the memory-derived block count. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - if use_override: - # Ample raw memory; the override alone constrains the block count. - vllm_config.cache_config.num_gpu_blocks_override = 32 - available_memory = [spec.page_size_bytes * 1024] - else: - available_memory = [spec.page_size_bytes * 32] - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - -def test_auto_fit_max_model_len_reserves_null_block(): - """Auto-fit (max_model_len=-1) must size max_model_len against usable - blocks, not the total pool. With memory for exactly 64 blocks, one is the - null block, so auto-fit must settle on 63 * block_size; picking the full - pool would leave a max-length request unschedulable and hang the engine. - """ - block_size = 16 - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # Exactly the 1024 / 16 = 64 blocks a full-length request would need. - available_memory = [spec.page_size_bytes * 64] - - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - assert vllm_config.model_config.max_model_len == 63 * block_size - - -def test_check_enough_kv_cache_memory_reserves_null_block(): - """The public admission check must reject a KV cache sized to exactly the - blocks a max_model_len request needs. BlockPool keeps one block as the - null block, so only num_blocks - 1 are usable; accepting - num_blocks == needed would leave the request unschedulable and hang the - engine. - """ - block_size = 16 - max_model_len = 512 # needs 512 / 16 = 32 blocks - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # 32 blocks -> only 31 usable after the null block: one short -> reject. - with pytest.raises(ValueError, match="max seq len"): - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 - ) - - # 33 blocks -> 32 usable after the null block -> accept. - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 - ) - - def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" @@ -4105,340 +2156,3 @@ def test_wrapped_mamba_group_requires_block_zeroing(): ) assert config.needs_kv_cache_zeroing - - -def _spec_decode_grouping_config(method="dspark", model_type=None): - """Grouping config with an EAGLE-family speculative method enabled.""" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - cache_config=SimpleNamespace( - get_resolved_kv_cache_layout=lambda: SimpleNamespace( - is_block_outermost=True - ) - ), - model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), - speculative_config=SimpleNamespace( - method=method, - use_eagle=lambda: True, - use_eagle_block_drop=lambda: True, - ), - ) - - -def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): - """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a - DSpark-style draft MLA layer marked non_causal_multi_token_decode. - - The target's fp8 KV dtype is one difference from the draft on Kimi-K3 - (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass - draft_shares_target_spec to match that dtype while retaining the draft - marker as a group boundary. - """ - target_dtype = None if draft_shares_target_spec else "fp8_e4m3" - specs = { - "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - } - if draft: - draft_spec = new_mla_spec(block_size=64) - specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) - return specs - - -def test_draft_group_annotated_on_hybrid_general_path(): - # A drafter's MLA layer carries non_causal_multi_token_decode, so its group - # is identifiable without keying off a model version. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "draft.attn.0" in flagged[0].layer_names - - -def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): - # Packed uniform-type groups can contain distinct target and draft layer - # specs; the combined group still holds volatile draft KV and must be - # flagged. What must never happen is a Mamba group being flagged: that - # widens its lookup window to two consecutive chunks, which align-mode - # checkpointing never produces, zeroing every lookup. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), - _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), - ) - - for group in groups: - if "draft.attn.0" in group.layer_names: - assert group.is_eagle_group - if isinstance(group.kv_cache_spec, MambaSpec): - assert not group.is_eagle_group - - -def test_draft_group_not_annotated_without_spec_decode(): - # The marker alone must not flag anything; the eagle semantics only apply - # when a speculative method is actually enabled. - config = _spec_decode_grouping_config() - config.speculative_config = None - groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) - - assert not any(g.is_eagle_group for g in groups) - - -def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so consumers fall back to - # conservative behavior that silently breaks reuse for Mamba groups. - # That must at least be visible. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) - ) - - assert not any(g.is_eagle_group for g in groups) - assert "no KV cache group could be identified as the draft model's" in ( - caplog_vllm.text - ) - - -def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): - # Pure-attention models degrade gracefully under the consumers' - # conservative fallback (a one-block hit drop at most), so the warning - # stays silent to avoid noise on every unannotated EAGLE deployment. - specs = { - "target.attn.0": new_mla_spec(block_size=64), - "target.attn.1": new_mla_spec(block_size=64), - } - groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def test_no_warning_when_draft_group_is_identified(caplog_vllm): - get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def _deepseek_v4_specs(model_version="deepseek_v4"): - """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last.""" - return { - "model.layers.0.self_attn.attn": new_mla_spec(), - "model.layers.1.self_attn.attn": new_mla_spec(), - "model.layers.2.self_attn.attn": new_swa_mla_spec( - head_size=1024, model_version=model_version - ), - # The MTP block registers last, and its sliding-window size differs, so - # it lands in a group of its own. - "model.layers.3.self_attn.attn": new_swa_mla_spec( - head_size=1024, sliding_window=256, model_version=model_version - ), - } - - -@pytest.mark.parametrize( - ("method", "model_type"), - [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], -) -def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): - # DeepseekV4's MTP block reuses the target's decoder layer, so its spec - # carries no draft marker and only the positional rule can find it. This - # pins the pre-existing behaviour that the unified annotator must preserve. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type=model_type), - _deepseek_v4_specs(model_version=None), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def test_trailing_layer_fallback_applies_to_any_mtp_model(): - # The positional rule is sound for every MTP drafter, not just DeepseekV4: - # MTP blocks reuse the target's decoder layer (no spec marker) and always - # register after every target layer. The model_type must not gate it. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="other"), - _deepseek_v4_specs(), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def _qwen3_5_hybrid_specs(with_mtp_layer: bool): - """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with - the MTP drafter's full-attn layer (spec-identical to the target's) - registered last.""" - specs = {} - idx = 0 - for _ in range(2): - for _ in range(3): - specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( - mamba_cache_mode="align" - ) - idx += 1 - specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() - idx += 1 - if with_mtp_layer: - specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() - return specs - - -def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): - # A hybrid mamba + full-attention model with an MTP drafter that is - # spec-indistinguishable from the target reaches the general multi-group - # path. The trailing-layer rule must locate the draft group there so the - # Mamba groups are not swept up by the flag-all consumer fallback. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names - for group in groups: - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - assert not group.is_eagle_group - assert "could be identified as the draft model's" not in caplog_vllm.text - - -@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) -def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): - # Other drafters are not covered by the trailing-layer rule, so an - # unidentifiable hybrid draft must still warn. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" in caplog_vllm.text - - -def test_trailing_layer_fallback_requires_exact_partition(): - # If the groups do not partition the layers exactly (e.g. a caller that - # dropped or duplicated layers), the positional rule is meaningless and - # must not fire. - from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups - - specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) - config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") - groups = get_kv_cache_groups(config, specs) - for g in groups: - g.is_eagle_group = False - # Remove one layer from its group: no longer an exact partition. - trimmed = [ - KVCacheGroupSpec( - [n for n in g.layer_names if n != "model.layers.0.linear_attn"], - g.kv_cache_spec, - ) - for g in groups - ] - _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) - - assert not any(g.is_eagle_group for g in trimmed) - - -_GQA_SPEC = FullAttentionSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 -) -_GQA_SWA_SPEC = SlidingWindowSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 -) - - -@pytest.mark.parametrize( - "specs,expected", - [ - pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": new_swa_mla_spec(), - "s.1": new_swa_mla_spec(), - }, - 4, - id="mla-and-swa-mla", - ), - pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), - pytest.param( - {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, - 1, - id="mla-and-gqa-uniform-group", - ), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": _GQA_SWA_SPEC, - "s.1": _GQA_SWA_SPEC, - }, - 1, - id="mla-and-gqa-swa", - ), - ], -) -def test_kv_tp_replicas(monkeypatch, specs, expected): - """Resolved per layer before scheduler flattening hides mixed groups.""" - from vllm.config import ParallelConfig - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "device_count", lambda: 4) - parallel_config = ParallelConfig(tensor_parallel_size=4) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config - ) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - mem = sum(s.page_size_bytes for s in specs.values()) * 10 - configs = get_kv_cache_configs(vllm_config, [specs], [mem]) - assert configs[0].kv_tp_replicas == expected - scheduler = generate_scheduler_kv_cache_config(configs) - assert scheduler.kv_tp_replicas == expected - - -@pytest.mark.parametrize( - "specs,tp_size,dcp_size,expected", - [ - pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), - pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), - pytest.param( - [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" - ), - pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), - pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), - pytest.param( - [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], - 8, - 1, - 2, - id="mla-and-gqa-partial", - ), - pytest.param( - [ - HiddenStateCacheSpec( - block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 - ) - ], - 8, - 1, - 1, - id="hidden-state", - ), - ], -) -def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): - """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" - groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] - assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index cf221fc22a34..541e55f5b60e 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -36,6 +36,10 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, + KVQuantMode, + MambaSpec, + MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -333,3 +337,73 @@ def create_requests( ) requests.append(request) return requests + + +def new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=None, + attention_chunk_size=None, + kv_quant_mode=KVQuantMode.NONE, +): + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + attention_chunk_size=attention_chunk_size, + kv_quant_mode=kv_quant_mode, + ) + + +def new_sliding_window_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=1, +): + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + ) + + +def new_mamba_spec( + block_size=16, + shapes=((2, 512), (3, 32, 32)), + dtypes=(torch.float32, torch.float32), + num_speculative_blocks=2, + mamba_cache_mode="none", + page_size_padded=None, +): + return MambaSpec( + block_size=block_size, + shapes=shapes, + dtypes=dtypes, + page_size_padded=page_size_padded, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + + +def new_mla_spec(cache_dtype_str=None, block_size: int = 16): + # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=576, + dtype=torch.float32, + cache_dtype_str=cache_dtype_str, + ) diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index 5626dec17043..f9455389ef08 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory +from vllm.v1.core.kv_cache_planning import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -38,7 +38,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): config.parallel_config.decode_context_parallel_size = 1 monkeypatch.setattr( - "vllm.v1.core.kv_cache_utils.max_memory_usage_bytes", + "vllm.v1.core.kv_cache_planning.max_memory_usage_bytes", lambda c, s: 100 * 1024**3, # 100 GiB ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 07d530f97d1a..a6692b5cd5fa 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,10 +17,8 @@ SchedulerOffloadConfig, ) from vllm.platforms import current_platform -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - kv_cache_groups_tp_replicas, -) +from vllm.v1.core.kv_cache_planning import kv_cache_groups_tp_replicas +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py index 735b042259b5..854632dce1d1 100644 --- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -10,7 +10,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 ExampleHiddenStatesConnector, ) -from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_groups from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 16e142d7cb97..4f9a7c451ffc 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -29,11 +29,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, is_kv_cache_spec_uniform, - resolve_kv_cache_block_sizes, ) +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 4cc8f9ffe2fd..abd0830a472b 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -104,7 +104,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span(): zeroed no matter which group owns it.""" from unittest.mock import MagicMock - from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheLayout, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 1247bf193a76..534d81312436 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -43,7 +43,10 @@ ROCMAiterMLASparseBackend, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs +from vllm.v1.core.kv_cache_planning import ( + estimate_max_model_len, + get_kv_cache_configs, +) from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py new file mode 100644 index 000000000000..6e258a68a9e7 --- /dev/null +++ b/vllm/v1/core/kv_cache_planning.py @@ -0,0 +1,1637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV cache planning: grouping, sizing, and memory-fit logic.""" + +import math +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from dataclasses import replace +from functools import partial, reduce +from typing import Any + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import format_gib +from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.core.kv_cache_utils import _kv_cache_device_name +from vllm.v1.hisparse.layout import ( + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, + compute_layout_strides, + iter_layer_specs, + replace_as, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry + +logger = init_logger(__name__) + + +def _check_enough_kv_cache_memory( + available_memory: int, + get_needed_memory: Callable[[], int], + max_model_len: int, + estimate_max_model_len: Callable[[int], int], +): + if available_memory <= 0: + raise ValueError( + "No available memory for the cache blocks. " + "Try increasing `gpu_memory_utilization` when initializing the engine " + "(this flag also controls CPU memory reservation on the CPU " + "backend, despite its name). " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more details." + ) + + needed_memory = get_needed_memory() + + if needed_memory > available_memory: + estimated_max_len = estimate_max_model_len(available_memory) + estimated_msg = "" + if estimated_max_len > 0: + estimated_msg = ( + "Based on the available memory, " + f"the estimated maximum model length is {estimated_max_len}. " + ) + + raise ValueError( + f"To serve at least one request with the model's max seq len " + f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " + f"cache is needed, which is larger than the available KV cache " + f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" + f"Try increasing `gpu_memory_utilization` (which also controls " + f"CPU memory on the CPU backend) or decreasing `max_model_len` " + f"when initializing the engine. " + f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + f"for more details." + ) + + +def max_memory_usage_bytes( + vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] +) -> int: + """Get the maximum memory usage in bytes for the given KV cache specs.""" + return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) + + +def estimate_max_model_len( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +) -> int: + """Estimates the maximum model length that can fit in the available memory + using binary search. + + This function temporarily modifies max_model_len during estimation but + restores the original value before returning, ensuring no side effects. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Returns: + The estimated maximum model length that can fit in the available memory. + + """ + # Save the original max_model_len to restore after estimation + original_max_model_len = vllm_config.model_config.max_model_len + + # Define a function to check if a given model length fits in memory + def fits_in_memory(model_len: int) -> bool: + # Temporarily modify the max_model_len for this calculation + vllm_config.model_config.max_model_len = model_len + # Calculate memory needed for the given model length + memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) + return memory_needed <= available_memory + + try: + # Binary search for the maximum model length + left, right = 1, original_max_model_len + + # If even the smallest model length doesn't fit, return 0 + if not fits_in_memory(left): + return 0 + + # Binary search for the maximum model length that fits + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits_in_memory(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + # Always restore the original max_model_len to avoid side effects + vllm_config.model_config.max_model_len = original_max_model_len + + +def check_enough_kv_cache_memory( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +): + """Checks whether `available_memory` is enough for the KV cache to hold at + least one request with the model's max_model_len. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Raises: + ValueError: If there is not enough memory available for the KV cache. + + """ + # No need to check for available memory if the kv_cache_spec is empty + if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # of the specs since grouping may unify them in-place. + groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) + check_memory = ( + available_memory - _pool_bytes_per_block(vllm_config, groups) + if groups + else available_memory + ) + _check_enough_kv_cache_memory( + check_memory, + lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), + vllm_config.model_config.max_model_len, + lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), + ) + + +def create_kv_cache_group_specs( + kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] +) -> list[KVCacheGroupSpec]: + """Create KVCacheGroupSpec object for each kv cache group layer. + The layers in the same group should share the same + KVCacheSpec. + + Args: + kv_cache_spec: + A mapping from each layer name to its corresponding KVCacheSpec. + grouped_layer_names: + A list of kv cache groups, where each element is a list of layer + names that belong to the same group and should share the same + KVCacheSpec. + + Returns: + A list of KVCacheGroupSpec objects, one for each group. + + """ + kv_cache_groups = [] + for layer_names_one_group in grouped_layer_names: + layer_specs = [ + kv_cache_spec[layer_name] for layer_name in layer_names_one_group + ] + merged_layer_spec = layer_specs[0].merge(layer_specs) + kv_cache_groups.append( + KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) + ) + return kv_cache_groups + + +def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + """Whether all layers in the given KVCacheSpec have the same KV cache spec. + Note that we regard FullAttentionSpec with and without sliding window as + the same type. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + True if all layers have the same type, False otherwise. + + """ + if not kv_cache_spec: + # Encoder-only models do not have KV cache, kv_cache_type can be + # regarded as uniform. + return True + try: + kv_cache_spec_values = list(kv_cache_spec.values()) + _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) + except AssertionError: + return False + return True + + +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: + """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. + The override is logged once, at the call site in `get_kv_cache_configs`. + """ + if vllm_config.cache_config.num_gpu_blocks_override is not None: + num_blocks = vllm_config.cache_config.num_gpu_blocks_override + return num_blocks + + +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: + """Bytes consumed by one block in the worker's shared KV cache pool, mirroring + the divisor used by `get_kv_cache_config_from_groups` to convert + `available_memory` into `num_blocks`. Used to compute the effective KV cache + capacity once `num_gpu_blocks_override` is applied. + """ + return _get_kv_cache_bytes_per_block( + [group for group in kv_cache_groups if not group.host_resident], + vllm_config.cache_config.get_resolved_kv_cache_layout(), + ) + + +def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: + """Get the page size of the KV cache.""" + page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} + assert len(page_sizes) == 1 + return page_sizes.pop() + + +def _get_kv_cache_groups_uniform_spec( + kv_cache_specs: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with the same KV cache + spec for all layers. + + Args: + kv_cache_specs: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) + + +def _get_kv_cache_groups_uniform_type( + spec: UniformTypeKVCacheSpecs, +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with one type of KV cache + but different hidden sizes. All layers are merged into one group. + + Args: + spec: The UniformTypeKVCacheSpecs of the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] + + +def unify_kv_cache_spec_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Unify the page size of the given KVCacheSpec. If the page size of all layers + are the same, return the original KVCacheSpec. If not same, unify the page + size by increasing the block size of layers with smaller page size. Two + cases cannot be unified by block size alone and pad their physical page to + the maximum instead: Mamba layers, whose page size comes from state shapes + and is independent of block size; and non-MLA attention layers whose page + does not evenly divide the maximum (the padded page is read through a + strided view). MLA is excluded because sparse MLA indexes the cache in + whole token rows (see ``flat_kv_row_view``), so its block stride can only + be padded by its own row-aligned ``alignment``, not to an arbitrary page + size. Raise NotImplementedError if failed to unify the page size; + ``get_kv_cache_groups`` catches it to try the full-allocation fallback + (e.g. MLA next to an incompatible sliding-window draft). + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + + Returns: + The updated KVCacheSpec with the same page_size_bytes. + + """ + page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} + if len(page_sizes) <= 1: + # All layers have the same page size, no need to unify. + return kv_cache_spec + + max_page_size = max(page_sizes) + new_kv_cache_spec = {} + for layer_name, layer_spec in kv_cache_spec.items(): + if layer_spec.page_size_bytes == max_page_size: + new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, MambaSpec): + # MambaSpec's page size is determined by its state shapes and does + # not scale with block_size, so pad the page instead. This is the + # same padding mechanism the platform uses to align Mamba pages + # with the main model's attention page size; it is needed here + # when another layer (e.g. from a draft model) has a larger page + # than the already-aligned Mamba page. + new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + else: + layer_page_size = layer_spec.page_size_bytes + if max_page_size % layer_page_size == 0: + ratio = max_page_size // layer_page_size + new_block_size = layer_spec.block_size * ratio + new_spec = replace(layer_spec, block_size=new_block_size) + elif isinstance(layer_spec, AttentionSpec) and not isinstance( + layer_spec, MLAAttentionSpec + ): + new_spec = replace(layer_spec, page_size_padded=max_page_size) + else: + raise NotImplementedError( + f"Layer {layer_name}: page size is not divisible by the " + "maximum page size and cannot be padded. Padding is only " + "supported for non-MLA attention layers." + ) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + return new_kv_cache_spec + + +def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + # kv_cache_spec is an empty dict for attention free models + return not kv_cache_spec + + +def _get_kv_cache_groups_uniform_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache groups for hybrid models with multiple + attention types but still with a uniform page size (physical memory per + block per layer) for all layers. + + Detailed explanation about kv cache management of hybrid models: + The layers in the models are repeated with some patterns, e.g., a model + with 10 full attention layers and 20 sliding window attention layers can be + regarded as repeating the pattern (1 * full, 2 * sw) 10 times. + The KVCacheManager allocates different block tables for each of the 3 layers + in the pattern, and repeats each of them 10 times to generate the + block_table for the 30 layers in the model. + Therefore, we can group the layers in the model into 3 kv_cache_groups, each + of which contains 10 layers in the model. + The KVCacheManager allocates the block_table for each group based on its + kv_cache spec, and the model runner applies the block table to each layer + in the group. + For example: + 1. A model only uses full attention. The pattern is + (num_hidden_layers * full), so there is only one group and the block table + is shared by all layers. It is already handled by + `_get_kv_cache_config_uniform_type`. + 2. A model with 10 full attention layers and 20 sliding window + attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so + there are 3 kv_cache_groups, each of which represents 10 layers. + + To simplify the implementation, we make the following assumptions: + 1. Physical memory per block: Must be the same across all KV cache groups. + Breaking this assumption is non-trivial due to memory fragmentation concerns + when allocating blocks of different sizes. + 2. Tokens per block (block_size): Currently, we directly use + `CacheConfig.block_size` for all layers. It can be extended to vary by KV + cache group, but within each KV cache group, all layers must share the same + block size. + 3. Physical memory per token per layer: This property is decided by model + config. Currently we only support models that have the same physical memory + per token per layer for all layers. Can be relaxed with a simple extension, + but still need to keep physical memory per block the same for all groups. + 4. Number of layers per group: Currently assumed the same for all layers. + Can be relaxed with a simple extension, but still need to keep physical + memory per block the same for all groups. + 5. Attention type within groups: All layers in a group must share the same + attention type. One exception is that, when + `--disable-hybrid-kv-cache-manager` is true, the single group for full + attention layers may also include attention layers using sliding window or + LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. + 6. Support for multiple attention types: The design for most components is + general to an arbitrary number of attention types. But + `find_longest_cache_hit` only supports one attention type or two + types of full-attention plus exactly one another type. The general + implementation of this function is feasible but we don't know how to + implement it cleanly yet. + + As we assume tokens per block, physical memory per token per layer, and + number of layers per group are the same now, we can ensure that physical + memory per block is the same for all groups. + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + Returns: + The generated KVCacheGroupSpecs + + """ + # Group all layers by kv_cache_spec. + # E.g., 2 full attention layers and 3 sliding window attention layers, + # -> (full.0, full.1), (sw.0, sw.1, sw.2). + same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) + for layer_name, layer_spec in kv_cache_spec.items(): + same_type_layers[layer_spec].append(layer_name) + + # Attempt to further merge same-type layers based on whether their KV + # cache specs can be merged, to minimize the group count. This benefits + # situations where specs share a block layout and differ only in a + # property it can reconcile (e.g. full attention layers differing only in + # sliding window / attention chunk size). + layer_buckets: list[list[str]] = [] + spec_buckets: list[list[KVCacheSpec]] = [] + for layer_spec, layer_names in same_type_layers.items(): + for names, specs in zip(layer_buckets, spec_buckets): + try: + # A raise means that the specs are incompatible. + type(specs[0]).merge([*specs, layer_spec]) + except (AssertionError, ValueError): + continue + names.extend(layer_names) + specs.append(layer_spec) + break + else: + layer_buckets.append(list(layer_names)) + spec_buckets.append([layer_spec]) + + # Split each group into smaller groups, to make the number of layers in each + # group identical. Add padding to the last group of each type if necessary. + # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) + # split to 3 groups with 2 layers each: + # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). + # FIXME(Chen): At the moment of writing this code (2025-06-02), all + # open-source hybrid model follows a n:1 pattern between different attention + # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and + # full), so we can use the "1" in the n:1 pattern as the group size, which + # is the minimum number of layers among all attention types. Need a better + # strategy if we want to support more complex patterns (e.g., 20 full + 30 + # sw, where the group size should be 10). + min_num_layers = min([len(layers) for layers in layer_buckets]) + group_size = min_num_layers + max_num_layers = max([len(layers) for layers in layer_buckets]) + if max_num_layers < min_num_layers * 1.5: + # If the number of layers is not much larger than the minimum number of + # layers, use the maximum number of layers as the group size to avoid + # too many padding layers. A typical example is gpt-oss-20b + eagle, + # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of + # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding + # layers while accommodating speculative decoding drafters that add + # extra layers to one attention type. + group_size = max_num_layers + grouped_layers = [] + for layers in layer_buckets: + num_padding_layers = group_size - len(layers) % group_size + if num_padding_layers != group_size: + logger.warning( + "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa + num_padding_layers, + num_padding_layers / len(layers) * 100, + ) + num_groups = cdiv(len(layers), group_size) + # In PP case, say if we have + # - stage 0: full.0, sw.0, sw.1 + # - stage 1: full.1, sw.2, sw.3 + # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) + # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because + # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) + # and it will be padded to (full.0, padding), (sw.0, sw.1), + # (padding, padding) to ensure the number of layers in each group is + # the same and will cause memory waste. + # To avoid this, we assign layers[i::num_groups] to the i-th group + # instead of layers[i * group_size: (i + 1) * group_size] + for i in range(num_groups): + grouped_layers.append(layers[i::num_groups]) + return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) + + +def _get_per_layer_spec( + group: KVCacheGroupSpec, + layer_name: str, +) -> KVCacheSpec: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.kv_cache_specs[layer_name] + return spec + + +def _get_layer_block_stride(spec: KVCacheSpec) -> int: + """A layer-outer block stride: the page padded to its alignment.""" + return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) + + +def _get_kv_cache_bytes_per_block( + kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout +) -> int: + """Return the largest cache group's bytes per block.""" + if not layout.is_block_outermost: + # Each layer's block stride is its own padded page. + return max( + sum( + _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + bytes_per_block = max( + sum( + _get_per_layer_spec(group, layer_name).page_size_bytes + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + assert bytes_per_block > 0 + alignments = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + alignments.extend( + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() + for group in kv_cache_groups + for layer_name in group.layer_names + ) + return round_up(bytes_per_block, math.lcm(*alignments)) + + +def validate_kv_cache_layout( + layout: KVCacheLayout, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Validate that the resolved layout can express this model's packing. + + The layout was chosen once in the engine core from the backends' supported + sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 + indexer) declares block-outermost layouts there, so an inexpressible + layout reaching this point is an error. + """ + page_sizes = { + _get_per_layer_spec(group, layer_name).page_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + } + if len(page_sizes) == 1: + # A rectangular layer dim exists; every layout can express it. + return + + # Mixed page sizes pack pages side by side within a block, which needs each page + # to be one contiguous chunk inside its block (a block-compact layout) and, with + # multiple KV cache groups, the layer dim inside the block dim. + if not layout.is_block_compact or ( + len(kv_cache_groups) > 1 and layout.is_layer_compact + ): + raise ValueError( + f"KV cache layout {layout.name} cannot express this model's " + f"mixed page sizes ({sorted(page_sizes)}); a backend should " + "declare block-outermost supported layouts (e.g. BLHNC), or " + "set VLLM_KV_CACHE_LAYOUT=BLHNC." + ) + + +def get_kv_cache_config_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> KVCacheConfig: + """Generate the KV cache configuration from the KV cache groups and spec + of each layer. + + Args: + vllm_config: The global VllmConfig + kv_cache_groups: The KV cache groups + available_memory: Memory available for KV cache in bytes + Returns: + The generated KVCacheConfig + + """ + if len(kv_cache_groups) == 0: + # Attention free models do not have KV cache. + # Return num_blocks=1 as BlockPool always needs a null_block. + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, kv_cache_groups) + bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + + num_blocks = available_memory // bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + size = bytes_per_block * num_blocks + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + block_stride = interleaved_block_stride + padded_page = _get_layer_block_stride(spec) + if block_stride is None and padded_page != spec.page_size_bytes: + assert layout.is_block_compact, f"{layout.name} cannot pad pages" + block_stride = padded_page + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, block_stride, None, None, None), + ) + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=byte_offset, + ) + ) + # Block-outer layers share each block; layer-outer ones own regions. + byte_offset += len(layer_names) * ( + spec.page_size_bytes if layout.is_block_outermost else layer_stride + ) + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + +def _promote_local_kv_cache_specs( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Use full-attention allocation for local-attention cache specs. + + The returned specs affect KV cache management only. Attention modules keep + their original sliding-window or chunked-local compute behavior. + """ + promoted_specs = kv_cache_spec.copy() + + if is_kv_cache_spec_uniform( + promoted_specs + ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): + return promoted_specs + + has_full_attention = any( + isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() + ) + has_sliding_window = any( + isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() + ) + has_chunked_local_attention = any( + isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() + ) + full_block_sizes = { + spec.block_size + for spec in promoted_specs.values() + if isinstance(spec, FullAttentionSpec) + } + full_attention_block_size = ( + next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None + ) + + def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: + if spec.page_size_padded is None: + return None + unpadded_page_size = ( + spec.unpadded_page_size_bytes * block_size // spec.block_size + ) + return max(spec.page_size_padded, unpadded_page_size) + + promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { + SlidingWindowMLASpec: MLAAttentionSpec, + SlidingWindowSpec: FullAttentionSpec, + ChunkedLocalAttentionSpec: FullAttentionSpec, + } + + if has_full_attention and (has_sliding_window or has_chunked_local_attention): + for layer_name, spec in kv_cache_spec.items(): + target_cls = next( + (promotions[c] for c in type(spec).__mro__ if c in promotions), None + ) + if target_cls is None: + continue + assert isinstance(spec, AttentionSpec) + block_size = full_attention_block_size or spec.block_size + promoted_specs[layer_name] = replace_as( + spec, + target_cls, + # Promoted specs allocate blocks for all tokens and never free + # below the window, so the trailing-edge extension is moot. + drop=("extra_retained_tokens",), + block_size=block_size, + page_size_padded=promoted_page_size_padded(spec, block_size), + ) + + if not ( + is_kv_cache_spec_uniform(promoted_specs) + or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) + ): + raise ValueError("Failed to promote local KV cache specs to one unified type.") + + return promoted_specs + + +def _try_get_full_allocation_fallback_groups( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Try a supported full-allocation fallback for local-attention layers.""" + if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): + return None + if any( + isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) + for spec in kv_cache_spec.values() + ): + return None + + has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) + has_regular_swa = any( + isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() + ) + if not (has_mla and has_regular_swa): + return None + + try: + promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) + except ValueError: + return None + uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) + if uniform_spec is None: + return None + logger.warning( + "KV cache page sizes cannot be unified; treating sliding-window " + "layers as full attention for cache allocation. Sliding-window " + "attention compute is unchanged." + ) + return _get_kv_cache_groups_uniform_type(uniform_spec) + + +def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): + """This function tries to convert the KV cache specs to one type if the model + is a hybrid model with multiple type of KV cache. It will convert all + SlidingWindowSpec to FullAttentionSpec if both types are present. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + """ + groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) + for name, spec in kv_cache_spec.items(): + replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded + groups[replicated][name] = spec + for specs in groups.values(): + promoted_specs = _promote_local_kv_cache_specs(specs) + if promoted_specs == specs: + continue + logger.warning( + "Hybrid KV cache manager is disabled for this hybrid model, " + "This means we do not enable any optimizations for saving KV cache " + "memory (e.g., dropping the KV cache outside the sliding window). " + "The compute of layers like sliding window is still saved." + ) + kv_cache_spec.update(promoted_specs) + + +def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: + """Pick a chunk size that minimizes total upward padding. + + Each x is rounded up to a multiple of d: + + x -> ceil(x / d) * d + + Total padding is: + + pad(d) = sum_i (ceil(x_i / d) * d - x_i) + + We brute-force d in [lower_bound, max(values)] (fine for small lists / small + maxima) and return the d with minimum padding. Ties prefer larger d. + """ + if not values: + raise ValueError("values must be non-empty") + if any(x <= 0 for x in values): + raise ValueError(f"values must be positive, got: {list(values)!r}") + + min_d = max(1, lower_bound if lower_bound is not None else 1) + max_d = max(values) + if min_d > max_d: + return min_d + + best_d = min_d + best_pad: int | None = None + for d in range(min_d, max_d + 1): + pad = sum((d - (x % d)) % d for x in values) + if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): + best_pad = pad + best_d = d + + return best_d + + +def _get_packed_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Group mixed-page-size layers for contiguous block-outermost packing. + + Greedily buckets layers into uniform-type specs. Buckets with equal layer + counts per page size are treated as a repeating layer pattern (one layer + per page size) and split into groups covering the same number of pattern + repeats (picked by ``_approximate_gcd`` to minimize padding), so all + groups pack into the same per-block layout. Mamba buckets are additionally + split to fit the block the attention buckets already need. + Returns None when the layout is not block-outermost or all layers already + share one page size. + """ + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} + if not layout.is_block_outermost or len(page_sizes) <= 1: + return None + + buckets: list[dict[str, KVCacheSpec]] = [] + for name, spec in kv_cache_spec.items(): + for bucket in buckets: + candidate = {**bucket, name: spec} + if UniformTypeKVCacheSpecs.is_uniform_type(candidate): + bucket[name] = spec + break + else: + buckets.append({name: spec}) + + bucketed = [] + for bucket in buckets: + uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) + assert uniform_spec is not None + page_size_layers: dict[int, list[str]] = defaultdict(list) + for layer_name, layer_spec in bucket.items(): + page_size_layers[layer_spec.page_size_bytes].append(layer_name) + # Only 1:1 patterns (one layer of each page size per repeat) are + # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle + # repeat too, but such buckets are emitted whole instead. + balanced = len(set(map(len, page_size_layers.values()))) == 1 + bucketed.append((uniform_spec, page_size_layers, balanced)) + + # Balanced buckets that mix page sizes must stay whole, so the largest one + # sets a floor on the repeats per group; larger single-size buckets are + # split down toward it. No such bucket means nothing needs packing. + min_repeats_per_group = max( + ( + spec.get_max_layers_per_page_size() + for spec, page_size_layers, balanced in bucketed + if balanced and len(page_size_layers) > 1 + ), + default=0, + ) + repeats_per_group = ( + _approximate_gcd( + [ + spec.get_max_layers_per_page_size() + for spec, _, balanced in bucketed + if balanced + ], + lower_bound=min_repeats_per_group, + ) + if min_repeats_per_group + else None + ) + + def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: + if balanced and repeats_per_group is not None: + return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) + return 1 + + def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: + """Page bytes of the largest of the n groups a bucket splits into.""" + return sum( + cdiv(len(names), n) * page for page, names in page_size_layers.items() + ) + + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + + anchor_bytes = max( + ( + widest_group_bytes( + page_size_layers, + len(spec.kv_cache_specs) + if is_state_bucket(spec) + else num_groups_for(spec, balanced), + ) + for spec, page_size_layers, balanced in bucketed + ), + default=0, + ) + + groups = [] + for spec, page_size_layers, balanced in bucketed: + num_groups = num_groups_for(spec, balanced) + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): + states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) + num_groups = max( + num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) + ) + if num_groups == 1: + groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) + continue + + pattern_repeats = list(zip(*page_size_layers.values())) + for i in range(num_groups): + group_layer_names = [ + name for repeat in pattern_repeats[i::num_groups] for name in repeat + ] + group_layer_specs = { + name: spec.kv_cache_specs[name] for name in group_layer_names + } + group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) + assert group_spec is not None + groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: + """Whether the drafter's KV layers can be located positionally. + + MTP drafters register their KV layers after the target layers, without a + distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. + The annotator separately checks that the groups partition the layers exactly. + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return False + if spec_config.method == "mtp": + return True + model_config = vllm_config.model_config + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", + ) + + +def _groups_partition_layers_exactly( + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. + + The trailing-layer draft fallback is only meaningful when the last + registered layer is guaranteed to appear in exactly one group. + """ + names = [name for group in kv_cache_groups for name in group.layer_names] + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) + + +def _annotate_eagle_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], + use_trailing_layer_fallback: bool = False, +) -> None: + """Flag the KV cache groups that hold drafter attention layers. + + Two detection rules, in order of preference: + + 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on + MLAAttentionSpec and set by drafter attention layers that run a + non-causal multi-token decode (today only Kimi-K3 DSpark). It survives + MLAAttentionSpec.merge, so it still identifies a group after per-group + spec merging, wherever grouping happens to land. It is sufficient but + not necessary: a drafter whose spec is indistinguishable from the + target's cannot be found this way. + 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose + MTP block reuses the target's own decoder layer and so carries no spec + marker. Their draft attention layers always register after every + target layer, so flag whichever group holds the last registered layer. + This rule is only valid where the groups partition exactly the layers + of ``kv_cache_spec``, which is re-checked here before applying it. + When a drafter's trailing caches span several groups, this rule flags + only the group holding the very last layer and must be generalized. + FIXME(yifan): avoid/generalize this hacky check. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_spec: The kv cache spec of each attention layer, in layer + registration order. Only read by rule 2. + kv_cache_groups: Groups to annotate in place. + use_trailing_layer_fallback: Enable rule 2. Callers gate this on + ``_uses_trailing_mtp_layers``. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle_block_drop(): + return + + for group in kv_cache_groups: + if any( + getattr(spec, "non_causal_multi_token_decode", False) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + group.is_eagle_group = True + + if not use_trailing_layer_fallback: + return + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): + return + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: + group.is_eagle_group = True + break + + +def _warn_if_unannotated_eagle_mamba( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Warn when no KV cache group could be identified as the draft model's. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_groups: Groups as they will be handed to consumers. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return + if any(group.is_eagle_group for group in kv_cache_groups): + return + mamba_groups = [ + idx + for idx, group in enumerate(kv_cache_groups) + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if not mamba_groups: + return + logger.warning( + "Speculative decoding (method=%s) is enabled but no KV cache group " + "could be identified as the draft model's.", + spec_config.method, + ) + + +def _largest_divisor_at_most(value: int, limit: int) -> int: + for candidate in range(min(value, limit), 0, -1): + if value % candidate == 0: + return candidate + return 1 + + +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """Scale up group block sizes so the common page is at least as large as + the biggest hidden-state per-token cost. + + This protects the hidden state extraction feature, where we store hidden states in + the KV cache, and return them to clients using a custom hidden state connector. + + Returns the (possibly scaled) groups and updated common page size. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, Any] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + +def get_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Split the layers in the model into groups with the same KV cache spec. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroups + + """ + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + unify_hybrid_kv_cache_specs(kv_cache_spec) + + if is_kv_cache_type_attention_free(kv_cache_spec): + # This returns an empty list to allow for the KVCacheManager to handle + # attention free models. + return [] + + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + + if is_kv_cache_spec_uniform(kv_cache_spec): + # KV cache of all layers are the same, which is true for + # most models. Allocate the same amount of memory for + # each layer. + return _get_kv_cache_groups_uniform_spec(kv_cache_spec) + elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): + # All layers need the same number of token slots (e.g., all layers are + # full attention, or all layers are sliding window attention with the + # same window size). Put all layers into one group. + return _get_kv_cache_groups_uniform_type(uniform_spec) + # Hidden-state layers use their own block table and must not be absorbed + # into a compatible attention bucket. + hidden_specs = { + k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) + } + filtered_spec = { + k: v + for k, v in kv_cache_spec.items() + if not isinstance(v, HiddenStateCacheSpec) + } + + if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): + # Block-outermost blocks are strided by the widest group, so hidden + # groups need no page alignment. + packed_groups += [ + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() + ] + return packed_groups + + # Prefer preserving each layer's cache semantics. If physical pages cannot + # be unified, try a supported allocation-only fallback before failing. + try: + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) + except NotImplementedError: + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) + if fallback_groups is None: + raise + return fallback_groups + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + + # Add hidden-state layers back with page aligned to the common page. + if hidden_specs: + common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) + group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) + for name, spec in hidden_specs.items(): + per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + max_block_size = max(common_page // per_token, 1) + new_bs = _largest_divisor_at_most(group_block_size, max_block_size) + wasted_bytes = common_page - new_bs * per_token + logger.info( + "Using block size %d for hidden-state cache layer %s; " + "page alignment wastes %d bytes (%.2f%%) per block", + new_bs, + name, + wasted_bytes, + wasted_bytes / common_page * 100, + ) + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) + groups.append(KVCacheGroupSpec([name], aligned)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: + if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: + return 1 + if spec.dcp_sharded and dcp_size > 1: + return 1 + return max(1, tp_size // spec.max_tp_shards) + + +def kv_cache_groups_tp_replicas( + groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 +) -> int: + """Consecutive TP ranks holding identical KV for every layer.""" + specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] + if not specs: + return 1 + return reduce( + math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size + ) + + +def _max_memory_usage_bytes_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + """Calculate maximum memory usage in bytes from KV cache groups. + + This correctly accounts for padding in hybrid models. For example, if a + model has 8 full attention layers and 9 sliding window layers, they will + be padded to 9 full + 9 sliding window for uniform group sizes. + + Each group independently claims blocks from the shared pool, so a request consumes + the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. + """ + kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] + if not kv_cache_groups: + return 0 + + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) + total_blocks = 0 + for group in kv_cache_groups: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + total_blocks += spec.max_memory_usage_pages(vllm_config) + else: + total_blocks += cdiv( + spec.max_memory_usage_bytes(vllm_config), + spec.page_size_bytes, + ) + + return bytes_per_block * total_blocks + + +def _estimate_max_model_len_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> int: + """Binary search for the maximum model length that fits in available memory. + Returns 0 if even 1 token doesn't fit. + """ + original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) + + def fits(model_len: int) -> bool: + vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return ( + _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) + <= available_memory + ) + + try: + left, right = 1, original_max + if not fits(left): + return 0 + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + vllm_config.model_config.max_model_len = original_max + + +def _auto_fit_max_model_len( + vllm_config: VllmConfig, + projected_groups_per_worker: list[list[KVCacheGroupSpec]], + available_memory: list[int], +) -> None: + """When max_model_len is set to -1, this function estimates the largest + context length that can be supported with the available GPU memory. + It uses binary search to find the maximum length that fits across all + workers. + + Args: + vllm_config: The global VllmConfig (will be modified in-place) + projected_groups_per_worker: KV cache groups projected to each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + """ + original_max = vllm_config.model_config.max_model_len + + if all(not groups for groups in projected_groups_per_worker): + # All workers have empty specs (attention-free model) + logger.info_once( + "Auto-fit max_model_len: attention-free model, " + "using derived max_model_len=%d", + original_max, + ) + return + + # Find the max_model_len that fits across all workers. + auto_fit_max = original_max + limiting_worker_mem = available_memory[0] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + continue + worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) + if worker_max < auto_fit_max: + auto_fit_max = worker_max + limiting_worker_mem = avail_mem + + if auto_fit_max <= 0: + raise ValueError( + "Cannot auto-fit max_model_len: not enough GPU memory available " + "to serve even a single token. Try increasing `gpu_memory_utilization`." + ) + + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + + if auto_fit_max >= original_max: + # The model's full context length fits in memory + logger.info_once( + "Auto-fit max_model_len: full model context length %d fits in " + "available KV cache memory (%s)", + original_max, + kv_cache_memory, + ) + else: + # Need to reduce max_model_len to fit in memory + vllm_config.model_config.max_model_len = auto_fit_max + logger.info_once( + "Auto-fit max_model_len: reduced from %d to %d to fit in " + "available KV cache memory (%s)", + original_max, + auto_fit_max, + kv_cache_memory, + ) + + +def _project_kv_cache_groups_to_worker( + global_kv_cache_groups: list[KVCacheGroupSpec], + worker_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Projects global KV cache groups onto a single worker's assigned layers. + + In pipeline parallelism, each worker only owns a subset of layers. This + function filters the global groups to include only layers present on the + given worker, adjusting UniformTypeKVCacheSpecs accordingly. + + Args: + global_kv_cache_groups: The global KV cache groups for the whole model. + worker_spec: The KV cache spec of each layer on this worker. + + Returns: + The projected KV cache groups containing only this worker's layers. + + """ + projected_groups: list[KVCacheGroupSpec] = [] + for group in global_kv_cache_groups: + worker_layer_names = [ + layer_name for layer_name in group.layer_names if layer_name in worker_spec + ] + if len(worker_layer_names) == len(group.layer_names): + projected_groups.append(group) + continue + group_spec = group.kv_cache_spec + if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): + group_spec = UniformTypeKVCacheSpecs( + block_size=group_spec.block_size, + kv_cache_specs={ + layer_name: group_spec.kv_cache_specs[layer_name] + for layer_name in worker_layer_names + }, + ) + projected_groups.append( + KVCacheGroupSpec( + worker_layer_names, + group_spec, + is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + ) + ) + return projected_groups + + +def get_kv_cache_configs( + vllm_config: VllmConfig, + kv_cache_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> list[KVCacheConfig]: + """Generates the KV cache configurations for a model. + Since we use a shared centralized controller for all workers, we need the + `kv_cache_config` to be consistent across all workers to make sure + the KV cache allocation can be applied to all workers. However, different + workers may have different memory available, and different type of layers + (when pipeline parallel is enabled). To handle the difference between + workers, the current implementation is: + 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for + the whole model. + 2. Generate the KV cache groups based on the layer ratio of the whole model. + This also handles spec unification for hybrid models. + 3. Handle auto-fit max_model_len and memory checks using per-worker + projected groups to account for PP sharding. + 4. Generate the KV cache configs for each worker based on the KV cache + grouping strategy. (This is reasonable because the layer ratio of + different PP stages are similar.) + 5. Change the num_blocks of each worker to the smallest among all workers + and shrink tensor sizes proportionally to avoid allocating unused memory. + + Args: + vllm_config: The global VllmConfig + kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + Returns: + The generated KVCacheConfigs for each worker. + + """ + # Merge the KV cache specs of all workers. Different PP stages may have + # different layer names, and different TP ranks of the same PP stage should + # have the same KV cache spec. + merged_kv_cache_specs: dict[str, KVCacheSpec] = {} + for kv_cache_spec_one_worker in kv_cache_specs: + for layer_name, layer_spec in kv_cache_spec_one_worker.items(): + if layer_name not in merged_kv_cache_specs: + merged_kv_cache_specs[layer_name] = layer_spec + else: + assert merged_kv_cache_specs[layer_name] == layer_spec, ( + "The KV cache specs for the same layer are different " + "across workers. This is not supported yet." + ) + + # Check if the KV cache specs are registered correctly. + # This is to prevent that some layers are initialized with unregistered specs. + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + + # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # tag every SlidingWindowSpec with how many extra tokens to retain in the window. + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + + # Get global KV cache groups. This also handles spec unification for + # hybrid models when disable_hybrid_kv_cache_manager is enabled. + # After this call, merged_kv_cache_specs may be modified in-place. + global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) + + # If original_max_model_len was -1, automatically + # determine the maximum model length that fits in available GPU memory. + # We use per-worker projected groups to account for PP sharding. + projected_groups_per_worker = [ + _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) + for worker_spec in kv_cache_specs + ] + + # If `num_gpu_blocks_override` is set, the cache size that will actually + # be allocated is decoupled from the profiled `available_memory`: + # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps + # `num_blocks` to the override. Reflect that in `available_memory` here so + # auto-fit, the admission check, and the per-worker config builder all + # plan against the same effective capacity. + override = vllm_config.cache_config.num_gpu_blocks_override + if override is not None: + adjusted_memory: list[int] = [] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + adjusted_memory.append(avail_mem) + continue + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) + logger.info( + "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", + avail_mem // bytes_per_block, + override, + ) + adjusted_memory.append(override * bytes_per_block) + available_memory = adjusted_memory + + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + for groups, avail_mem in zip(projected_groups_per_worker, available_memory) + ] + + if vllm_config.model_config.original_max_model_len == -1: + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) + + # Check if the available memory is enough per worker. + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): + if not groups: + continue + _check_enough_kv_cache_memory( + avail_mem, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + + kv_cache_configs: list[KVCacheConfig] = [] + for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( + projected_groups_per_worker, kv_cache_specs, available_memory + ): + assert sum(len(group.layer_names) for group in projected_groups) == len( + kv_cache_spec_one_worker + ), "Some layers are not assigned to any group." + kv_cache_configs.append( + get_kv_cache_config_from_groups( + vllm_config, projected_groups, available_memory_one_worker + ) + ) + + # Change the num_blocks of each rank to the smallest among all ranks. + # We also need to shrink the tensor size proportionally to avoid + # allocating unused memory. + min_num_blocks = min( + kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs + ) + for i, kv_cache_config in enumerate(kv_cache_configs): + if kv_cache_config.num_blocks == min_num_blocks: + continue + # Re-plan with exactly the memory the smallest rank can afford, so + # strides and offsets stay consistent with the shrunken allocation. + groups = kv_cache_config.kv_cache_groups + kv_cache_configs[i] = get_kv_cache_config_from_groups( + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(vllm_config, groups), + ) + + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.decode_context_parallel_size, + ) + + return kv_cache_configs diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 03c6c5ff62ac..5c51a8d131ff 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -6,10 +6,8 @@ import hashlib import math import os -from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace -from functools import partial, reduce from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload from vllm import envs @@ -17,34 +15,14 @@ from vllm.logger import init_logger from vllm.multimodal.utils import get_mm_features_in_window from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv, round_up -from vllm.utils.mem_utils import format_gib -from vllm.utils.torch_utils import get_dtype_size -from vllm.v1.hisparse.layout import ( - get_hisparse_host_pool_bytes, - get_hisparse_kv_cache_config, - get_hisparse_kv_cache_groups, -) +from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import ( - AttentionSpec, - ChunkedLocalAttentionSpec, - CircularBufferSpec, - FullAttentionSpec, - HiddenStateCacheSpec, - HiSparseHotSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheLayout, KVCacheSpec, - KVCacheTensor, MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, - compute_layout_strides, iter_layer_specs, - replace_as, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -1000,198 +978,6 @@ def request_block_hasher(request: Request) -> list[BlockHash]: return request_block_hasher -def _check_enough_kv_cache_memory( - available_memory: int, - get_needed_memory: Callable[[], int], - max_model_len: int, - estimate_max_model_len: Callable[[int], int], -): - if available_memory <= 0: - raise ValueError( - "No available memory for the cache blocks. " - "Try increasing `gpu_memory_utilization` when initializing the engine " - "(this flag also controls CPU memory reservation on the CPU " - "backend, despite its name). " - "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - "for more details." - ) - - needed_memory = get_needed_memory() - - if needed_memory > available_memory: - estimated_max_len = estimate_max_model_len(available_memory) - estimated_msg = "" - if estimated_max_len > 0: - estimated_msg = ( - "Based on the available memory, " - f"the estimated maximum model length is {estimated_max_len}. " - ) - - raise ValueError( - f"To serve at least one request with the model's max seq len " - f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " - f"cache is needed, which is larger than the available KV cache " - f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" - f"Try increasing `gpu_memory_utilization` (which also controls " - f"CPU memory on the CPU backend) or decreasing `max_model_len` " - f"when initializing the engine. " - f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - f"for more details." - ) - - -def max_memory_usage_bytes( - vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] -) -> int: - """Get the maximum memory usage in bytes for the given KV cache specs.""" - return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) - - -def estimate_max_model_len( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -) -> int: - """Estimates the maximum model length that can fit in the available memory - using binary search. - - This function temporarily modifies max_model_len during estimation but - restores the original value before returning, ensuring no side effects. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Returns: - The estimated maximum model length that can fit in the available memory. - - """ - # Save the original max_model_len to restore after estimation - original_max_model_len = vllm_config.model_config.max_model_len - - # Define a function to check if a given model length fits in memory - def fits_in_memory(model_len: int) -> bool: - # Temporarily modify the max_model_len for this calculation - vllm_config.model_config.max_model_len = model_len - # Calculate memory needed for the given model length - memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) - return memory_needed <= available_memory - - try: - # Binary search for the maximum model length - left, right = 1, original_max_model_len - - # If even the smallest model length doesn't fit, return 0 - if not fits_in_memory(left): - return 0 - - # Binary search for the maximum model length that fits - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits_in_memory(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - # Always restore the original max_model_len to avoid side effects - vllm_config.model_config.max_model_len = original_max_model_len - - -def check_enough_kv_cache_memory( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -): - """Checks whether `available_memory` is enough for the KV cache to hold at - least one request with the model's max_model_len. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Raises: - ValueError: If there is not enough memory available for the KV cache. - - """ - # No need to check for available memory if the kv_cache_spec is empty - if kv_cache_spec: - # Reserve the null block BlockPool permanently holds back, so the check - # plans against usable blocks, as in get_kv_cache_configs. Group a copy - # of the specs since grouping may unify them in-place. - groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) - check_memory = ( - available_memory - _pool_bytes_per_block(vllm_config, groups) - if groups - else available_memory - ) - _check_enough_kv_cache_memory( - check_memory, - lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), - vllm_config.model_config.max_model_len, - lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), - ) - - -def create_kv_cache_group_specs( - kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] -) -> list[KVCacheGroupSpec]: - """Create KVCacheGroupSpec object for each kv cache group layer. - The layers in the same group should share the same - KVCacheSpec. - - Args: - kv_cache_spec: - A mapping from each layer name to its corresponding KVCacheSpec. - grouped_layer_names: - A list of kv cache groups, where each element is a list of layer - names that belong to the same group and should share the same - KVCacheSpec. - - Returns: - A list of KVCacheGroupSpec objects, one for each group. - - """ - kv_cache_groups = [] - for layer_names_one_group in grouped_layer_names: - layer_specs = [ - kv_cache_spec[layer_name] for layer_name in layer_names_one_group - ] - merged_layer_spec = layer_specs[0].merge(layer_specs) - kv_cache_groups.append( - KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) - ) - return kv_cache_groups - - -def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - """Whether all layers in the given KVCacheSpec have the same KV cache spec. - Note that we regard FullAttentionSpec with and without sliding window as - the same type. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - True if all layers have the same type, False otherwise. - - """ - if not kv_cache_spec: - # Encoder-only models do not have KV cache, kv_cache_type can be - # regarded as uniform. - return True - try: - kv_cache_spec_values = list(kv_cache_spec.values()) - _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) - except AssertionError: - return False - return True - - def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1227,1052 +1013,6 @@ def get_max_concurrency_for_kv_cache_config( return min(limits) -def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: - """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. - The override is logged once, at the call site in `get_kv_cache_configs`. - """ - if vllm_config.cache_config.num_gpu_blocks_override is not None: - num_blocks = vllm_config.cache_config.num_gpu_blocks_override - return num_blocks - - -def _pool_bytes_per_block( - vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] -) -> int: - """Bytes consumed by one block in the worker's shared KV cache pool, mirroring - the divisor used by `get_kv_cache_config_from_groups` to convert - `available_memory` into `num_blocks`. Used to compute the effective KV cache - capacity once `num_gpu_blocks_override` is applied. - """ - return _get_kv_cache_bytes_per_block( - [group for group in kv_cache_groups if not group.host_resident], - vllm_config.cache_config.get_resolved_kv_cache_layout(), - ) - - -def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: - """Get the page size of the KV cache.""" - page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} - assert len(page_sizes) == 1 - return page_sizes.pop() - - -def _get_kv_cache_groups_uniform_spec( - kv_cache_specs: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with the same KV cache - spec for all layers. - - Args: - kv_cache_specs: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) - - -def _get_kv_cache_groups_uniform_type( - spec: UniformTypeKVCacheSpecs, -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with one type of KV cache - but different hidden sizes. All layers are merged into one group. - - Args: - spec: The UniformTypeKVCacheSpecs of the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] - - -def unify_kv_cache_spec_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Unify the page size of the given KVCacheSpec. If the page size of all layers - are the same, return the original KVCacheSpec. If not same, unify the page - size by increasing the block size of layers with smaller page size. Two - cases cannot be unified by block size alone and pad their physical page to - the maximum instead: Mamba layers, whose page size comes from state shapes - and is independent of block size; and non-MLA attention layers whose page - does not evenly divide the maximum (the padded page is read through a - strided view). MLA is excluded because sparse MLA indexes the cache in - whole token rows (see ``flat_kv_row_view``), so its block stride can only - be padded by its own row-aligned ``alignment``, not to an arbitrary page - size. Raise NotImplementedError if failed to unify the page size; - ``get_kv_cache_groups`` catches it to try the full-allocation fallback - (e.g. MLA next to an incompatible sliding-window draft). - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - - Returns: - The updated KVCacheSpec with the same page_size_bytes. - - """ - page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} - if len(page_sizes) <= 1: - # All layers have the same page size, no need to unify. - return kv_cache_spec - - max_page_size = max(page_sizes) - new_kv_cache_spec = {} - for layer_name, layer_spec in kv_cache_spec.items(): - if layer_spec.page_size_bytes == max_page_size: - new_kv_cache_spec[layer_name] = layer_spec - elif isinstance(layer_spec, MambaSpec): - # MambaSpec's page size is determined by its state shapes and does - # not scale with block_size, so pad the page instead. This is the - # same padding mechanism the platform uses to align Mamba pages - # with the main model's attention page size; it is needed here - # when another layer (e.g. from a draft model) has a larger page - # than the already-aligned Mamba page. - new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - else: - layer_page_size = layer_spec.page_size_bytes - if max_page_size % layer_page_size == 0: - ratio = max_page_size // layer_page_size - new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) - elif isinstance(layer_spec, AttentionSpec) and not isinstance( - layer_spec, MLAAttentionSpec - ): - new_spec = replace(layer_spec, page_size_padded=max_page_size) - else: - raise NotImplementedError( - f"Layer {layer_name}: page size is not divisible by the " - "maximum page size and cannot be padded. Padding is only " - "supported for non-MLA attention layers." - ) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - return new_kv_cache_spec - - -def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - # kv_cache_spec is an empty dict for attention free models - return not kv_cache_spec - - -def _get_kv_cache_groups_uniform_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache groups for hybrid models with multiple - attention types but still with a uniform page size (physical memory per - block per layer) for all layers. - - Detailed explanation about kv cache management of hybrid models: - The layers in the models are repeated with some patterns, e.g., a model - with 10 full attention layers and 20 sliding window attention layers can be - regarded as repeating the pattern (1 * full, 2 * sw) 10 times. - The KVCacheManager allocates different block tables for each of the 3 layers - in the pattern, and repeats each of them 10 times to generate the - block_table for the 30 layers in the model. - Therefore, we can group the layers in the model into 3 kv_cache_groups, each - of which contains 10 layers in the model. - The KVCacheManager allocates the block_table for each group based on its - kv_cache spec, and the model runner applies the block table to each layer - in the group. - For example: - 1. A model only uses full attention. The pattern is - (num_hidden_layers * full), so there is only one group and the block table - is shared by all layers. It is already handled by - `_get_kv_cache_config_uniform_type`. - 2. A model with 10 full attention layers and 20 sliding window - attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so - there are 3 kv_cache_groups, each of which represents 10 layers. - - To simplify the implementation, we make the following assumptions: - 1. Physical memory per block: Must be the same across all KV cache groups. - Breaking this assumption is non-trivial due to memory fragmentation concerns - when allocating blocks of different sizes. - 2. Tokens per block (block_size): Currently, we directly use - `CacheConfig.block_size` for all layers. It can be extended to vary by KV - cache group, but within each KV cache group, all layers must share the same - block size. - 3. Physical memory per token per layer: This property is decided by model - config. Currently we only support models that have the same physical memory - per token per layer for all layers. Can be relaxed with a simple extension, - but still need to keep physical memory per block the same for all groups. - 4. Number of layers per group: Currently assumed the same for all layers. - Can be relaxed with a simple extension, but still need to keep physical - memory per block the same for all groups. - 5. Attention type within groups: All layers in a group must share the same - attention type. One exception is that, when - `--disable-hybrid-kv-cache-manager` is true, the single group for full - attention layers may also include attention layers using sliding window or - LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. - 6. Support for multiple attention types: The design for most components is - general to an arbitrary number of attention types. But - `find_longest_cache_hit` only supports one attention type or two - types of full-attention plus exactly one another type. The general - implementation of this function is feasible but we don't know how to - implement it cleanly yet. - - As we assume tokens per block, physical memory per token per layer, and - number of layers per group are the same now, we can ensure that physical - memory per block is the same for all groups. - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - Returns: - The generated KVCacheGroupSpecs - - """ - # Group all layers by kv_cache_spec. - # E.g., 2 full attention layers and 3 sliding window attention layers, - # -> (full.0, full.1), (sw.0, sw.1, sw.2). - same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) - for layer_name, layer_spec in kv_cache_spec.items(): - same_type_layers[layer_spec].append(layer_name) - - # Attempt to further merge same-type layers based on whether their KV - # cache specs can be merged, to minimize the group count. This benefits - # situations where specs share a block layout and differ only in a - # property it can reconcile (e.g. full attention layers differing only in - # sliding window / attention chunk size). - layer_buckets: list[list[str]] = [] - spec_buckets: list[list[KVCacheSpec]] = [] - for layer_spec, layer_names in same_type_layers.items(): - for names, specs in zip(layer_buckets, spec_buckets): - try: - # A raise means that the specs are incompatible. - type(specs[0]).merge([*specs, layer_spec]) - except (AssertionError, ValueError): - continue - names.extend(layer_names) - specs.append(layer_spec) - break - else: - layer_buckets.append(list(layer_names)) - spec_buckets.append([layer_spec]) - - # Split each group into smaller groups, to make the number of layers in each - # group identical. Add padding to the last group of each type if necessary. - # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) - # split to 3 groups with 2 layers each: - # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). - # FIXME(Chen): At the moment of writing this code (2025-06-02), all - # open-source hybrid model follows a n:1 pattern between different attention - # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and - # full), so we can use the "1" in the n:1 pattern as the group size, which - # is the minimum number of layers among all attention types. Need a better - # strategy if we want to support more complex patterns (e.g., 20 full + 30 - # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in layer_buckets]) - group_size = min_num_layers - max_num_layers = max([len(layers) for layers in layer_buckets]) - if max_num_layers < min_num_layers * 1.5: - # If the number of layers is not much larger than the minimum number of - # layers, use the maximum number of layers as the group size to avoid - # too many padding layers. A typical example is gpt-oss-20b + eagle, - # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of - # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding - # layers while accommodating speculative decoding drafters that add - # extra layers to one attention type. - group_size = max_num_layers - grouped_layers = [] - for layers in layer_buckets: - num_padding_layers = group_size - len(layers) % group_size - if num_padding_layers != group_size: - logger.warning( - "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa - num_padding_layers, - num_padding_layers / len(layers) * 100, - ) - num_groups = cdiv(len(layers), group_size) - # In PP case, say if we have - # - stage 0: full.0, sw.0, sw.1 - # - stage 1: full.1, sw.2, sw.3 - # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) - # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because - # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) - # and it will be padded to (full.0, padding), (sw.0, sw.1), - # (padding, padding) to ensure the number of layers in each group is - # the same and will cause memory waste. - # To avoid this, we assign layers[i::num_groups] to the i-th group - # instead of layers[i * group_size: (i + 1) * group_size] - for i in range(num_groups): - grouped_layers.append(layers[i::num_groups]) - return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) - - -def _get_per_layer_spec( - group: KVCacheGroupSpec, - layer_name: str, -) -> KVCacheSpec: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - return spec.kv_cache_specs[layer_name] - return spec - - -def _get_layer_block_stride(spec: KVCacheSpec) -> int: - """A layer-outer block stride: the page padded to its alignment.""" - return round_up(spec.page_size_bytes, spec.get_block_stride_alignment()) - - -def _get_kv_cache_bytes_per_block( - kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout -) -> int: - """Return the largest cache group's bytes per block.""" - if not layout.is_block_outermost: - # Each layer's block stride is its own padded page. - return max( - sum( - _get_layer_block_stride(_get_per_layer_spec(group, layer_name)) - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - bytes_per_block = max( - sum( - _get_per_layer_spec(group, layer_name).page_size_bytes - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - assert bytes_per_block > 0 - alignments = [ - group.kv_cache_spec.page_size_bytes - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, HiSparseHotSpec) - ] - alignments.extend( - _get_per_layer_spec(group, layer_name).get_block_stride_alignment() - for group in kv_cache_groups - for layer_name in group.layer_names - ) - return round_up(bytes_per_block, math.lcm(*alignments)) - - -def validate_kv_cache_layout( - layout: KVCacheLayout, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Validate that the resolved layout can express this model's packing. - - The layout was chosen once in the engine core from the backends' supported - sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 - indexer) declares block-outermost layouts there, so an inexpressible - layout reaching this point is an error. - """ - page_sizes = { - _get_per_layer_spec(group, layer_name).page_size_bytes - for group in kv_cache_groups - for layer_name in group.layer_names - } - if len(page_sizes) == 1: - # A rectangular layer dim exists; every layout can express it. - return - - # Mixed page sizes pack pages side by side within a block, which needs each page - # to be one contiguous chunk inside its block (a block-compact layout) and, with - # multiple KV cache groups, the layer dim inside the block dim. - if not layout.is_block_compact or ( - len(kv_cache_groups) > 1 and layout.is_layer_compact - ): - raise ValueError( - f"KV cache layout {layout.name} cannot express this model's " - f"mixed page sizes ({sorted(page_sizes)}); a backend should " - "declare block-outermost supported layouts (e.g. BLHNC), or " - "set VLLM_KV_CACHE_LAYOUT=BLHNC." - ) - - -def get_kv_cache_config_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> KVCacheConfig: - """Generate the KV cache configuration from the KV cache groups and spec - of each layer. - - Args: - vllm_config: The global VllmConfig - kv_cache_groups: The KV cache groups - available_memory: Memory available for KV cache in bytes - Returns: - The generated KVCacheConfig - - """ - if len(kv_cache_groups) == 0: - # Attention free models do not have KV cache. - # Return num_blocks=1 as BlockPool always needs a null_block. - return KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - if vllm_config.attention_config.hisparse_config is not None: - host_budget = get_hisparse_host_pool_bytes(vllm_config) - return get_hisparse_kv_cache_config( - vllm_config, kv_cache_groups, available_memory, host_budget - ) - - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None - - num_blocks = available_memory // bytes_per_block - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - size = bytes_per_block * num_blocks - - # Groups alias from byte 0. Spec regions are laid out differently: - # - # block-outer (the same packing repeats for every block): - # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... - # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... - # |<--- bytes_per_block -->| - # - # layer-outer (only supported for uniform page sizes or single-group models): - # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | - # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | - - kv_cache_tensors = [] - for group in kv_cache_groups: - group_spec = group.kv_cache_spec - layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) - if isinstance(group_spec, UniformTypeKVCacheSpecs): - for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) - elif group.layer_names: - layers_by_spec[group_spec].extend(group.layer_names) - - byte_offset = 0 - for spec, layer_names in layers_by_spec.items(): - block_stride = interleaved_block_stride - padded_page = _get_layer_block_stride(spec) - if block_stride is None and padded_page != spec.page_size_bytes: - assert layout.is_block_compact, f"{layout.name} cannot pad pages" - block_stride = padded_page - layer_stride, block_stride, _, _, _ = compute_layout_strides( - spec, - num_blocks, - len(layer_names), - layout, - fixed_strides=(None, block_stride, None, None, None), - ) - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=layer_names, - layer_stride=layer_stride, - block_stride=block_stride, - offset=byte_offset, - ) - ) - # Block-outer layers share each block; layer-outer ones own regions. - byte_offset += len(layer_names) * ( - spec.page_size_bytes if layout.is_block_outermost else layer_stride - ) - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - -def _promote_local_kv_cache_specs( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Use full-attention allocation for local-attention cache specs. - - The returned specs affect KV cache management only. Attention modules keep - their original sliding-window or chunked-local compute behavior. - """ - promoted_specs = kv_cache_spec.copy() - - if is_kv_cache_spec_uniform( - promoted_specs - ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): - return promoted_specs - - has_full_attention = any( - isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() - ) - has_sliding_window = any( - isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() - ) - has_chunked_local_attention = any( - isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() - ) - full_block_sizes = { - spec.block_size - for spec in promoted_specs.values() - if isinstance(spec, FullAttentionSpec) - } - full_attention_block_size = ( - next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None - ) - - def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: - if spec.page_size_padded is None: - return None - unpadded_page_size = ( - spec.unpadded_page_size_bytes * block_size // spec.block_size - ) - return max(spec.page_size_padded, unpadded_page_size) - - promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { - SlidingWindowMLASpec: MLAAttentionSpec, - SlidingWindowSpec: FullAttentionSpec, - ChunkedLocalAttentionSpec: FullAttentionSpec, - } - - if has_full_attention and (has_sliding_window or has_chunked_local_attention): - for layer_name, spec in kv_cache_spec.items(): - target_cls = next( - (promotions[c] for c in type(spec).__mro__ if c in promotions), None - ) - if target_cls is None: - continue - assert isinstance(spec, AttentionSpec) - block_size = full_attention_block_size or spec.block_size - promoted_specs[layer_name] = replace_as( - spec, - target_cls, - # Promoted specs allocate blocks for all tokens and never free - # below the window, so the trailing-edge extension is moot. - drop=("extra_retained_tokens",), - block_size=block_size, - page_size_padded=promoted_page_size_padded(spec, block_size), - ) - - if not ( - is_kv_cache_spec_uniform(promoted_specs) - or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) - ): - raise ValueError("Failed to promote local KV cache specs to one unified type.") - - return promoted_specs - - -def _try_get_full_allocation_fallback_groups( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Try a supported full-allocation fallback for local-attention layers.""" - if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): - return None - if any( - isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) - for spec in kv_cache_spec.values() - ): - return None - - has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) - has_regular_swa = any( - isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() - ) - if not (has_mla and has_regular_swa): - return None - - try: - promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) - except ValueError: - return None - uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) - if uniform_spec is None: - return None - logger.warning( - "KV cache page sizes cannot be unified; treating sliding-window " - "layers as full attention for cache allocation. Sliding-window " - "attention compute is unchanged." - ) - return _get_kv_cache_groups_uniform_type(uniform_spec) - - -def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): - """This function tries to convert the KV cache specs to one type if the model - is a hybrid model with multiple type of KV cache. It will convert all - SlidingWindowSpec to FullAttentionSpec if both types are present. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - """ - groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) - for name, spec in kv_cache_spec.items(): - replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded - groups[replicated][name] = spec - for specs in groups.values(): - promoted_specs = _promote_local_kv_cache_specs(specs) - if promoted_specs == specs: - continue - logger.warning( - "Hybrid KV cache manager is disabled for this hybrid model, " - "This means we do not enable any optimizations for saving KV cache " - "memory (e.g., dropping the KV cache outside the sliding window). " - "The compute of layers like sliding window is still saved." - ) - kv_cache_spec.update(promoted_specs) - - -def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: - """Pick a chunk size that minimizes total upward padding. - - Each x is rounded up to a multiple of d: - - x -> ceil(x / d) * d - - Total padding is: - - pad(d) = sum_i (ceil(x_i / d) * d - x_i) - - We brute-force d in [lower_bound, max(values)] (fine for small lists / small - maxima) and return the d with minimum padding. Ties prefer larger d. - """ - if not values: - raise ValueError("values must be non-empty") - if any(x <= 0 for x in values): - raise ValueError(f"values must be positive, got: {list(values)!r}") - - min_d = max(1, lower_bound if lower_bound is not None else 1) - max_d = max(values) - if min_d > max_d: - return min_d - - best_d = min_d - best_pad: int | None = None - for d in range(min_d, max_d + 1): - pad = sum((d - (x % d)) % d for x in values) - if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): - best_pad = pad - best_d = d - - return best_d - - -def _get_packed_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Group mixed-page-size layers for contiguous block-outermost packing. - - Greedily buckets layers into uniform-type specs. Buckets with equal layer - counts per page size are treated as a repeating layer pattern (one layer - per page size) and split into groups covering the same number of pattern - repeats (picked by ``_approximate_gcd`` to minimize padding), so all - groups pack into the same per-block layout. Mamba buckets are additionally - split to fit the block the attention buckets already need. - Returns None when the layout is not block-outermost or all layers already - share one page size. - """ - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if not layout.is_block_outermost or len(page_sizes) <= 1: - return None - - buckets: list[dict[str, KVCacheSpec]] = [] - for name, spec in kv_cache_spec.items(): - for bucket in buckets: - candidate = {**bucket, name: spec} - if UniformTypeKVCacheSpecs.is_uniform_type(candidate): - bucket[name] = spec - break - else: - buckets.append({name: spec}) - - bucketed = [] - for bucket in buckets: - uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) - assert uniform_spec is not None - page_size_layers: dict[int, list[str]] = defaultdict(list) - for layer_name, layer_spec in bucket.items(): - page_size_layers[layer_spec.page_size_bytes].append(layer_name) - # Only 1:1 patterns (one layer of each page size per repeat) are - # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle - # repeat too, but such buckets are emitted whole instead. - balanced = len(set(map(len, page_size_layers.values()))) == 1 - bucketed.append((uniform_spec, page_size_layers, balanced)) - - # Balanced buckets that mix page sizes must stay whole, so the largest one - # sets a floor on the repeats per group; larger single-size buckets are - # split down toward it. No such bucket means nothing needs packing. - min_repeats_per_group = max( - ( - spec.get_max_layers_per_page_size() - for spec, page_size_layers, balanced in bucketed - if balanced and len(page_size_layers) > 1 - ), - default=0, - ) - repeats_per_group = ( - _approximate_gcd( - [ - spec.get_max_layers_per_page_size() - for spec, _, balanced in bucketed - if balanced - ], - lower_bound=min_repeats_per_group, - ) - if min_repeats_per_group - else None - ) - - def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: - if balanced and repeats_per_group is not None: - return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) - return 1 - - def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: - """Page bytes of the largest of the n groups a bucket splits into.""" - return sum( - cdiv(len(names), n) * page for page, names in page_size_layers.items() - ) - - # Bytes a block must hold however the state buckets end up split: mamba, - # circular-buffer and unsplit sliding-window buckets can go down to one - # state per group, every other bucket's split is fixed by the repeat pattern. - def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: - if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): - return True - return repeats_per_group is None and isinstance( - spec.first_spec, SlidingWindowSpec - ) - - anchor_bytes = max( - ( - widest_group_bytes( - page_size_layers, - len(spec.kv_cache_specs) - if is_state_bucket(spec) - else num_groups_for(spec, balanced), - ) - for spec, page_size_layers, balanced in bucketed - ), - default=0, - ) - - groups = [] - for spec, page_size_layers, balanced in bucketed: - num_groups = num_groups_for(spec, balanced) - # Cap a state group at the states a block already fits rather than let - # it widen the block. - if anchor_bytes and is_state_bucket(spec): - states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) - num_groups = max( - num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) - ) - if num_groups == 1: - groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) - continue - - pattern_repeats = list(zip(*page_size_layers.values())) - for i in range(num_groups): - group_layer_names = [ - name for repeat in pattern_repeats[i::num_groups] for name in repeat - ] - group_layer_specs = { - name: spec.kv_cache_specs[name] for name in group_layer_names - } - group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) - assert group_spec is not None - groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: - """Whether the drafter's KV layers can be located positionally. - - MTP drafters register their KV layers after the target layers, without a - distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. - The annotator separately checks that the groups partition the layers exactly. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return False - if spec_config.method == "mtp": - return True - model_config = vllm_config.model_config - return model_config is not None and model_config.hf_config.model_type in ( - "deepseek_v4", - "deepseek_v41", - ) - - -def _groups_partition_layers_exactly( - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. - - The trailing-layer draft fallback is only meaningful when the last - registered layer is guaranteed to appear in exactly one group. - """ - names = [name for group in kv_cache_groups for name in group.layer_names] - return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) - - -def _annotate_eagle_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], - use_trailing_layer_fallback: bool = False, -) -> None: - """Flag the KV cache groups that hold drafter attention layers. - - Two detection rules, in order of preference: - - 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on - MLAAttentionSpec and set by drafter attention layers that run a - non-causal multi-token decode (today only Kimi-K3 DSpark). It survives - MLAAttentionSpec.merge, so it still identifies a group after per-group - spec merging, wherever grouping happens to land. It is sufficient but - not necessary: a drafter whose spec is indistinguishable from the - target's cannot be found this way. - 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose - MTP block reuses the target's own decoder layer and so carries no spec - marker. Their draft attention layers always register after every - target layer, so flag whichever group holds the last registered layer. - This rule is only valid where the groups partition exactly the layers - of ``kv_cache_spec``, which is re-checked here before applying it. - When a drafter's trailing caches span several groups, this rule flags - only the group holding the very last layer and must be generalized. - FIXME(yifan): avoid/generalize this hacky check. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_spec: The kv cache spec of each attention layer, in layer - registration order. Only read by rule 2. - kv_cache_groups: Groups to annotate in place. - use_trailing_layer_fallback: Enable rule 2. Callers gate this on - ``_uses_trailing_mtp_layers``. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle_block_drop(): - return - - for group in kv_cache_groups: - if any( - getattr(spec, "non_causal_multi_token_decode", False) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - group.is_eagle_group = True - - if not use_trailing_layer_fallback: - return - if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - -def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Warn when no KV cache group could be identified as the draft model's. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_groups: Groups as they will be handed to consumers. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return - if any(group.is_eagle_group for group in kv_cache_groups): - return - mamba_groups = [ - idx - for idx, group in enumerate(kv_cache_groups) - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if not mamba_groups: - return - logger.warning( - "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's.", - spec_config.method, - ) - - -def _largest_divisor_at_most(value: int, limit: int) -> int: - for candidate in range(min(value, limit), 0, -1): - if value % candidate == 0: - return candidate - return 1 - - -def _ensure_min_page_size( - groups: list[KVCacheGroupSpec], - common_page: int, - hidden_specs: dict[str, HiddenStateCacheSpec], -) -> tuple[list[KVCacheGroupSpec], int]: - """Scale up group block sizes so the common page is at least as large as - the biggest hidden-state per-token cost. - - This protects the hidden state extraction feature, where we store hidden states in - the KV cache, and return them to clients using a custom hidden state connector. - - Returns the (possibly scaled) groups and updated common page size. - """ - max_per_token = max( - spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - for spec in hidden_specs.values() - ) - if max_per_token <= common_page: - return groups, common_page - - scale = cdiv(max_per_token, common_page) - common_page *= scale - scaled: list[KVCacheGroupSpec] = [] - for g in groups: - s = g.kv_cache_spec - kw: dict[str, Any] = {"block_size": s.block_size * scale} - if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: - kw["page_size_padded"] = s.page_size_padded * scale - scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) - return scaled, common_page - - -def get_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Split the layers in the model into groups with the same KV cache spec. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroups - - """ - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - unify_hybrid_kv_cache_specs(kv_cache_spec) - - if is_kv_cache_type_attention_free(kv_cache_spec): - # This returns an empty list to allow for the KVCacheManager to handle - # attention free models. - return [] - - if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): - return hisparse_groups - - if is_kv_cache_spec_uniform(kv_cache_spec): - # KV cache of all layers are the same, which is true for - # most models. Allocate the same amount of memory for - # each layer. - return _get_kv_cache_groups_uniform_spec(kv_cache_spec) - elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): - # All layers need the same number of token slots (e.g., all layers are - # full attention, or all layers are sliding window attention with the - # same window size). Put all layers into one group. - return _get_kv_cache_groups_uniform_type(uniform_spec) - # Hidden-state layers use their own block table and must not be absorbed - # into a compatible attention bucket. - hidden_specs = { - k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) - } - filtered_spec = { - k: v - for k, v in kv_cache_spec.items() - if not isinstance(v, HiddenStateCacheSpec) - } - - if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): - # Block-outermost blocks are strided by the widest group, so hidden - # groups need no page alignment. - packed_groups += [ - KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() - ] - return packed_groups - - # Prefer preserving each layer's cache semantics. If physical pages cannot - # be unified, try a supported allocation-only fallback before failing. - try: - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) - except NotImplementedError: - fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) - if fallback_groups is None: - raise - return fallback_groups - groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) - - # Add hidden-state layers back with page aligned to the common page. - if hidden_specs: - common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) - # TP may shrink the common page below hidden-state per-token cost. - groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) - group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) - for name, spec in hidden_specs.items(): - per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - max_block_size = max(common_page // per_token, 1) - new_bs = _largest_divisor_at_most(group_block_size, max_block_size) - wasted_bytes = common_page - new_bs * per_token - logger.info( - "Using block size %d for hidden-state cache layer %s; " - "page alignment wastes %d bytes (%.2f%%) per block", - new_bs, - name, - wasted_bytes, - wasted_bytes / common_page * 100, - ) - aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) - groups.append(KVCacheGroupSpec([name], aligned)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: - if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: - return 1 - if spec.dcp_sharded and dcp_size > 1: - return 1 - return max(1, tp_size // spec.max_tp_shards) - - -def kv_cache_groups_tp_replicas( - groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 -) -> int: - """Consecutive TP ranks holding identical KV for every layer.""" - specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] - if not specs: - return 1 - return reduce( - math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size - ) - - def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: @@ -2334,359 +1074,6 @@ def update_kv_cache_capacity( ) -def _max_memory_usage_bytes_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> int: - """Calculate maximum memory usage in bytes from KV cache groups. - - This correctly accounts for padding in hybrid models. For example, if a - model has 8 full attention layers and 9 sliding window layers, they will - be padded to 9 full + 9 sliding window for uniform group sizes. - - Each group independently claims blocks from the shared pool, so a request consumes - the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. - """ - kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] - if not kv_cache_groups: - return 0 - - bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) - total_blocks = 0 - for group in kv_cache_groups: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - total_blocks += spec.max_memory_usage_pages(vllm_config) - else: - total_blocks += cdiv( - spec.max_memory_usage_bytes(vllm_config), - spec.page_size_bytes, - ) - - return bytes_per_block * total_blocks - - -def _estimate_max_model_len_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> int: - """Binary search for the maximum model length that fits in available memory. - Returns 0 if even 1 token doesn't fit. - """ - original_max = vllm_config.model_config.max_model_len - hisparse_enabled = ( - vllm_config.attention_config.hisparse_config is not None - and bool(kv_cache_groups) - ) - - def fits(model_len: int) -> bool: - vllm_config.model_config.max_model_len = model_len - if hisparse_enabled: - try: - get_kv_cache_config_from_groups( - vllm_config, kv_cache_groups, available_memory - ) - except ValueError: - return False - return ( - _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) - <= available_memory - ) - - try: - left, right = 1, original_max - if not fits(left): - return 0 - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - vllm_config.model_config.max_model_len = original_max - - -def _auto_fit_max_model_len( - vllm_config: VllmConfig, - projected_groups_per_worker: list[list[KVCacheGroupSpec]], - available_memory: list[int], -) -> None: - """When max_model_len is set to -1, this function estimates the largest - context length that can be supported with the available GPU memory. - It uses binary search to find the maximum length that fits across all - workers. - - Args: - vllm_config: The global VllmConfig (will be modified in-place) - projected_groups_per_worker: KV cache groups projected to each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - """ - original_max = vllm_config.model_config.max_model_len - - if all(not groups for groups in projected_groups_per_worker): - # All workers have empty specs (attention-free model) - logger.info_once( - "Auto-fit max_model_len: attention-free model, " - "using derived max_model_len=%d", - original_max, - ) - return - - # Find the max_model_len that fits across all workers. - auto_fit_max = original_max - limiting_worker_mem = available_memory[0] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - continue - worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) - if worker_max < auto_fit_max: - auto_fit_max = worker_max - limiting_worker_mem = avail_mem - - if auto_fit_max <= 0: - raise ValueError( - "Cannot auto-fit max_model_len: not enough GPU memory available " - "to serve even a single token. Try increasing `gpu_memory_utilization`." - ) - - kv_cache_memory = ( - f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" - ) - if vllm_config.attention_config.hisparse_config is not None: - host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) - kv_cache_memory += f" + {host_pool_gib} GiB host" - - if auto_fit_max >= original_max: - # The model's full context length fits in memory - logger.info_once( - "Auto-fit max_model_len: full model context length %d fits in " - "available KV cache memory (%s)", - original_max, - kv_cache_memory, - ) - else: - # Need to reduce max_model_len to fit in memory - vllm_config.model_config.max_model_len = auto_fit_max - logger.info_once( - "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available KV cache memory (%s)", - original_max, - auto_fit_max, - kv_cache_memory, - ) - - -def _project_kv_cache_groups_to_worker( - global_kv_cache_groups: list[KVCacheGroupSpec], - worker_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Projects global KV cache groups onto a single worker's assigned layers. - - In pipeline parallelism, each worker only owns a subset of layers. This - function filters the global groups to include only layers present on the - given worker, adjusting UniformTypeKVCacheSpecs accordingly. - - Args: - global_kv_cache_groups: The global KV cache groups for the whole model. - worker_spec: The KV cache spec of each layer on this worker. - - Returns: - The projected KV cache groups containing only this worker's layers. - - """ - projected_groups: list[KVCacheGroupSpec] = [] - for group in global_kv_cache_groups: - worker_layer_names = [ - layer_name for layer_name in group.layer_names if layer_name in worker_spec - ] - if len(worker_layer_names) == len(group.layer_names): - projected_groups.append(group) - continue - group_spec = group.kv_cache_spec - if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): - group_spec = UniformTypeKVCacheSpecs( - block_size=group_spec.block_size, - kv_cache_specs={ - layer_name: group_spec.kv_cache_specs[layer_name] - for layer_name in worker_layer_names - }, - ) - projected_groups.append( - KVCacheGroupSpec( - worker_layer_names, - group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), - ) - ) - return projected_groups - - -def get_kv_cache_configs( - vllm_config: VllmConfig, - kv_cache_specs: list[dict[str, KVCacheSpec]], - available_memory: list[int], -) -> list[KVCacheConfig]: - """Generates the KV cache configurations for a model. - Since we use a shared centralized controller for all workers, we need the - `kv_cache_config` to be consistent across all workers to make sure - the KV cache allocation can be applied to all workers. However, different - workers may have different memory available, and different type of layers - (when pipeline parallel is enabled). To handle the difference between - workers, the current implementation is: - 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for - the whole model. - 2. Generate the KV cache groups based on the layer ratio of the whole model. - This also handles spec unification for hybrid models. - 3. Handle auto-fit max_model_len and memory checks using per-worker - projected groups to account for PP sharding. - 4. Generate the KV cache configs for each worker based on the KV cache - grouping strategy. (This is reasonable because the layer ratio of - different PP stages are similar.) - 5. Change the num_blocks of each worker to the smallest among all workers - and shrink tensor sizes proportionally to avoid allocating unused memory. - - Args: - vllm_config: The global VllmConfig - kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - Returns: - The generated KVCacheConfigs for each worker. - - """ - # Merge the KV cache specs of all workers. Different PP stages may have - # different layer names, and different TP ranks of the same PP stage should - # have the same KV cache spec. - merged_kv_cache_specs: dict[str, KVCacheSpec] = {} - for kv_cache_spec_one_worker in kv_cache_specs: - for layer_name, layer_spec in kv_cache_spec_one_worker.items(): - if layer_name not in merged_kv_cache_specs: - merged_kv_cache_specs[layer_name] = layer_spec - else: - assert merged_kv_cache_specs[layer_name] == layer_spec, ( - "The KV cache specs for the same layer are different " - "across workers. This is not supported yet." - ) - - # Check if the KV cache specs are registered correctly. - # This is to prevent that some layers are initialized with unregistered specs. - KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) - # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) - for layer_name, layer_spec in merged_kv_cache_specs.items(): - if isinstance(layer_spec, SlidingWindowSpec): - merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens - ) - - # Get global KV cache groups. This also handles spec unification for - # hybrid models when disable_hybrid_kv_cache_manager is enabled. - # After this call, merged_kv_cache_specs may be modified in-place. - global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) - - # If original_max_model_len was -1, automatically - # determine the maximum model length that fits in available GPU memory. - # We use per-worker projected groups to account for PP sharding. - projected_groups_per_worker = [ - _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) - for worker_spec in kv_cache_specs - ] - - # If `num_gpu_blocks_override` is set, the cache size that will actually - # be allocated is decoupled from the profiled `available_memory`: - # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps - # `num_blocks` to the override. Reflect that in `available_memory` here so - # auto-fit, the admission check, and the per-worker config builder all - # plan against the same effective capacity. - override = vllm_config.cache_config.num_gpu_blocks_override - if override is not None: - adjusted_memory: list[int] = [] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - adjusted_memory.append(avail_mem) - continue - bytes_per_block = _pool_bytes_per_block(vllm_config, groups) - logger.info( - "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", - avail_mem // bytes_per_block, - override, - ) - adjusted_memory.append(override * bytes_per_block) - available_memory = adjusted_memory - - # Reserve the null block BlockPool permanently holds back, so auto-fit and - # the capacity check both plan against usable blocks. Allocation below - # still uses the full memory. - check_memory = [ - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem - for groups, avail_mem in zip(projected_groups_per_worker, available_memory) - ] - - if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) - - # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, check_memory): - if not groups: - continue - _check_enough_kv_cache_memory( - avail_mem, - partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), - vllm_config.model_config.max_model_len, - partial(_estimate_max_model_len_from_groups, vllm_config, groups), - ) - - kv_cache_configs: list[KVCacheConfig] = [] - for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( - projected_groups_per_worker, kv_cache_specs, available_memory - ): - assert sum(len(group.layer_names) for group in projected_groups) == len( - kv_cache_spec_one_worker - ), "Some layers are not assigned to any group." - kv_cache_configs.append( - get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker - ) - ) - - # Change the num_blocks of each rank to the smallest among all ranks. - # We also need to shrink the tensor size proportionally to avoid - # allocating unused memory. - min_num_blocks = min( - kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs - ) - for i, kv_cache_config in enumerate(kv_cache_configs): - if kv_cache_config.num_blocks == min_num_blocks: - continue - # Re-plan with exactly the memory the smallest rank can afford, so - # strides and offsets stay consistent with the shrunken allocation. - groups = kv_cache_config.kv_cache_groups - kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, - groups, - min_num_blocks * _pool_bytes_per_block(vllm_config, groups), - ) - - for kv_cache_config in kv_cache_configs: - kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( - kv_cache_config.kv_cache_groups, - vllm_config.parallel_config.tensor_parallel_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) - - return kv_cache_configs - - class BlockHashListWithBlockSize: """Convert block-hash granularity from `hash_block_size` to `target_block_size`. Used when KV cache groups have different block sizes: `hash_block_size` diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 35cc8cc1860c..d8c976ff693b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -50,10 +50,10 @@ from vllm.utils.network_utils import make_zmq_socket from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.attention.backends.utils import resolve_kv_cache_layout +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, - get_kv_cache_configs, get_request_block_hasher, init_none_hash, resolve_cache_hit_alignment_tokens, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 96d93650ad29..0a72427d4233 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -71,7 +71,7 @@ def get_hisparse_kv_cache_groups( if not mla_specs: return None - from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) assert mla_group_spec is not None @@ -323,7 +323,7 @@ def get_hisparse_kv_cache_config( available_memory: int, host_budget: int, ) -> KVCacheConfig: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, may_override_num_blocks, validate_kv_cache_layout, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 73be62319a64..3a6409db5546 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -989,7 +989,7 @@ def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: """Allocate the smallest KV cache that still lets every graph be captured.""" - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e19c1b84f457..90059fd39f5b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6206,7 +6206,7 @@ def profile_run(self, randomize_inputs: bool = False) -> None: gc.collect() def _init_minimal_kv_cache_for_profiling(self) -> None: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) From c5422f4c466afd765b32b662151d0f37dad7baa5 Mon Sep 17 00:00:00 2001 From: RAHUL KUMAR <55033230+pctablet505@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:28:32 +0530 Subject: [PATCH 460/952] [Test][Renderer] Cover the reasoning -> reasoning_content mirror in chat message parsing (#56129) Signed-off-by: pctablet505 <55033230+pctablet505@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 --- tests/renderers/test_chat_utils.py | 46 ++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/renderers/test_chat_utils.py b/tests/renderers/test_chat_utils.py index d0a9b9e6c7a8..47fe49a63e6e 100644 --- a/tests/renderers/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -27,6 +27,7 @@ ChatCompletionMessageParam, ConversationMessage, _load_embeds_dict, + _parse_chat_message_content, _parse_metadata_array, _postprocess_messages, parse_chat_messages, @@ -3196,3 +3197,48 @@ def test_validate_chat_template_rejects_invalid_type(): ) as exc_info: validate_chat_template(123) assert exc_info.value.parameter == "chat_template" + + +def _parse_one(message) -> ConversationMessage: + # Plain-text content only reads these two tracker attributes. + tracker = MagicMock() + parser = tracker.create_parser.return_value + parser.model_config.enable_prompt_embeds = False + parser.mm_placeholder_storage.return_value = {} + result = _parse_chat_message_content( + message, tracker, content_format="string", interleave_strings=False + ) + assert len(result) == 1 + return result[0] + + +def test_assistant_reasoning_is_mirrored_to_reasoning_content(): + """``reasoning`` is mirrored to ``reasoning_content`` for legacy templates.""" + conv_msg = _parse_one( + { + "role": "assistant", + "content": "4", + "reasoning": "2+2 equals 4", + } + ) + + assert conv_msg.get("reasoning") == "2+2 equals 4" + assert conv_msg.get("reasoning_content") == "2+2 equals 4" + + +def test_assistant_without_reasoning_sets_neither_field(): + """No reasoning in, no reasoning keys out: no empty thinking block in templates.""" + conv_msg = _parse_one({"role": "assistant", "content": "Hi"}) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg + + +def test_non_assistant_reasoning_is_not_mirrored(): + """Only assistant turns carry reasoning; a user turn must not.""" + conv_msg = _parse_one( + {"role": "user", "content": "Hi", "reasoning": "should be ignored"} + ) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg From 657a2e94174fadcae7ae729b076d61001509d804 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:22:48 -0700 Subject: [PATCH 461/952] [Rust Frontend] Accept whole-number floats for integer tool arguments (#61299) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/parser/src/tool/parameters.rs | 31 +++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index 391ca90f105f..afb3d4a9f2cb 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -294,7 +294,9 @@ fn try_convert_value(param_type: &JsonParamType<'_>, input: &ParamInput) -> Opti fn try_convert_text_value(param_type: &JsonParamType<'_>, value: &str) -> Option { match param_type { JsonParamType::String => Some(Value::String(value.to_string())), - JsonParamType::Integer => convert_integer_text(value), + JsonParamType::Integer => { + convert_integer_text(value).or_else(|| convert_whole_number_text(value)) + } JsonParamType::Number => convert_number_text(value), JsonParamType::Boolean => try_convert_boolean(value), JsonParamType::Object(_) if value.trim().is_empty() => Some(Value::Object(Map::new())), @@ -401,6 +403,14 @@ pub(crate) fn convert_integer_text(value: &str) -> Option { .map(Value::Number) } +/// Convert a whole number written as a float (`3.0`) to a JSON integer, since JSON Schema +/// counts any number with a zero fractional part as an `integer`. +fn convert_whole_number_text(value: &str) -> Option { + let number = without_digit_separators(value.trim())?.parse::().ok()?; + let in_range = number >= i64::MIN as f64 && number < i64::MAX as f64; + (number.fract() == 0.0 && in_range).then(|| Value::Number(Number::from(number as i64))) +} + /// Convert raw text to a JSON number, as for a `number` parameter. pub(crate) fn convert_number_text(value: &str) -> Option { try_convert_number(&without_digit_separators(value.trim())?) @@ -591,6 +601,25 @@ mod tests { } } + #[test] + fn integers_accept_whole_number_floats() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { + "count": { "type": "integer" }, + "either": { "type": ["integer", "number"] } + } + })); + + assert_eq!(params.convert("count", text("3.0")), json!(3)); + assert_eq!(params.convert("count", text(" -2.0 ")), json!(-2)); + assert_eq!(params.convert("count", text("1_000.0")), json!(1000)); + assert_eq!(params.convert("either", text("3.5")), json!(3.5)); + for invalid in ["3.5", "inf", "NaN", "1e300"] { + assert_eq!(params.convert("count", text(invalid)), json!(invalid)); + } + } + #[test] fn nullable_schemas_accept_null() { let params = ToolSchema::from_schema(json!({ From 28d679e762a0b6ed560d767a76ebe02699bf5c46 Mon Sep 17 00:00:00 2001 From: dskwe Date: Tue, 13 Oct 2026 12:23:45 +0800 Subject: [PATCH 462/952] [Bugfix] Preserve independent tokenizer code_revision in snapshot resolution (#61272) Signed-off-by: Yuan Li Co-authored-by: Cyrus Leung --- tests/tokenizers_/test_basic.py | 85 +++++++++++++++++++++++++++++++++ vllm/tokenizers/registry.py | 13 ++--- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index a46be988e5c8..a91124549457 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from pathlib import Path from typing import _get_protocol_attrs # type: ignore from unittest.mock import patch import pytest from transformers import ( + AutoTokenizer, PreTrainedTokenizerBase, TokenizersBackend, tokenization_utils_base, @@ -13,6 +15,7 @@ from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.hf import HfTokenizer from vllm.tokenizers.mistral import MistralTokenizer +from vllm.tokenizers.registry import resolve_tokenizer_args def _get_missing_attrs(obj: object, target: type): @@ -58,6 +61,88 @@ def test_tokenizer_revision(tokenizer_name: str): get_tokenizer(tokenizer_name, revision="never") +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision, error_match", + [ + pytest.param("never", "not a valid git identifier", id="unknown-revision"), + pytest.param( + "73293ac504e43f4a895688b57c65cebe24f01a38", + r"tokenization_fast\.py", + id="revision-without-code", + ), + ], +) +def test_tokenizer_code_revision_rejects_missing_code( + code_revision: str, error_match: str +): + """A data snapshot must not shadow the requested remote-code revision.""" + with pytest.raises(OSError, match=error_match): + get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision="08a35780b7a78794436d8654aafa26ee0165c149", + code_revision=code_revision, + trust_remote_code=True, + ) + + +@pytest.mark.skip_global_cleanup +def test_tokenizer_code_revision_distinct_from_revision(): + """Distinct code/data revisions: the resolver must keep the Hub ID. + + Dynamic code is read from a local directory as-is, so resolving to a + data-revision snapshot would silently shadow an independent code + revision (the same contract already enforced for processors). + """ + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + repo = "hf-internal-testing/test_dynamic_tokenizer" + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision=revision, + trust_remote_code=True, + ) + assert Path(name).is_dir() + assert kwargs["code_revision"] == revision + + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision="73293ac504e43f4a895688b57c65cebe24f01a38", + trust_remote_code=True, + ) + assert name == repo + assert kwargs["code_revision"] == "73293ac504e43f4a895688b57c65cebe24f01a38" + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision", [None, "08a35780b7a78794436d8654aafa26ee0165c149"] +) +def test_tokenizer_code_revision_default_or_matching(code_revision: str | None): + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + tokenizer = get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + reference = AutoTokenizer.from_pretrained( + "hf-internal-testing/test_dynamic_tokenizer", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + assert isinstance(tokenizer, PreTrainedTokenizerBase) + assert tokenizer.encode("hello world") == reference.encode("hello world") + assert tokenizer.name_or_path == "hf-internal-testing/test_dynamic_tokenizer" + _assert_tokenizer_like(tokenizer) + + def test_hub_tokenizer_loads_from_local_snapshot(): """Transformers makes a Hub request per file when given a repo ID.""" with patch.object( diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index a7adf4597f6a..8b1863023c99 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -165,12 +165,13 @@ def resolve_tokenizer_args( "mistral_format=True is not supported with tokenizer_mode='hf'" ) kwargs["mistral_format"] = False - tokenizer_name = get_non_weight_snapshot_path( - str(tokenizer_name), - revision=revision, - cache_dir=download_dir, - token=kwargs.get("token"), - ) + if kwargs.get("code_revision") in (None, revision): + tokenizer_name = get_non_weight_snapshot_path( + str(tokenizer_name), + revision=revision, + cache_dir=download_dir, + token=kwargs.get("token"), + ) return tokenizer_mode, tokenizer_name, args, kwargs From 69d00bc2645cac2144096be3ef1a12a73f95f366 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:55:20 -0700 Subject: [PATCH 463/952] [Bugfix] Accept whole-number floats for integer tool arguments in the response parser (#61297) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- tests/parser/test_chat_parsing.py | 8 ++++++++ vllm/parser/chat_parsing/response_parser.py | 9 ++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/parser/test_chat_parsing.py b/tests/parser/test_chat_parsing.py index 679d97f15b20..e3b8dba17220 100644 --- a/tests/parser/test_chat_parsing.py +++ b/tests/parser/test_chat_parsing.py @@ -1555,6 +1555,14 @@ def test_coerce_tool_calls_casts_declared_types(self): def test_coerce_falls_back_to_raw_on_failure(self): self.assertEqual(_coerce("not-a-number", ("integer",)), "not-a-number") + def test_coerce_integer_accepts_whole_number_floats(self): + self.assertEqual(_coerce("3.0", ("integer",)), 3) + self.assertIsInstance(_coerce("3.0", ("integer",)), int) + self.assertEqual(_coerce("3.5", ("integer",)), "3.5") + self.assertEqual(_coerce("3.5", ("integer", "number")), 3.5) + self.assertEqual(_coerce("inf", ("integer",)), "inf") + self.assertEqual(_coerce("NaN", ("integer",)), "NaN") + def test_coerce_handles_any_of_and_null(self): self.assertEqual(_coerce("5", ("integer", "null")), 5) self.assertIsNone(_coerce("null", ("integer", "null"))) diff --git a/vllm/parser/chat_parsing/response_parser.py b/vllm/parser/chat_parsing/response_parser.py index 556df803a70a..0416fe0fe3ea 100644 --- a/vllm/parser/chat_parsing/response_parser.py +++ b/vllm/parser/chat_parsing/response_parser.py @@ -50,7 +50,14 @@ def _coerce(raw: str, types: tuple[str, ...]) -> Any: for type_name in types: try: if type_name == "integer": - return int(raw) + try: + return int(raw) + except ValueError: + # JSON Schema integers include whole numbers written as floats, e.g. "3.0" + number = float(raw) + if not number.is_integer(): + continue + return int(number) if type_name == "number": number = float(raw) if number != number or number in (float("inf"), float("-inf")): From 5555a2d8f7116b2b4877f66516f4128176cf722a Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:02:12 +0900 Subject: [PATCH 464/952] [Bugfix][Model] Fix EXAONE-4.5 image inputs: ViT received un-normalised pixels (#61104) Signed-off-by: godmook Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- vllm/model_executor/models/exaone4_5.py | 6 ++++++ vllm/model_executor/models/exaone4_5_mtp.py | 1 + 2 files changed, 7 insertions(+) diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py index 5a53e60b7d5e..ad4049c8eeb0 100644 --- a/vllm/model_executor/models/exaone4_5.py +++ b/vllm/model_executor/models/exaone4_5.py @@ -36,6 +36,7 @@ from vllm.distributed import utils as dist_utils from vllm.logger import init_logger from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -272,6 +273,7 @@ def __init__( vision_config: Exaone4_5_VisionConfig, norm_eps: float = 1e-6, quant_config: QuantizationConfig | None = None, + input_norm: nn.Module | None = None, prefix: str = "", use_data_parallel: bool = False, ) -> None: @@ -279,6 +281,7 @@ def __init__( vision_config=vision_config, norm_eps=norm_eps, quant_config=quant_config, + input_norm=input_norm, prefix=prefix, ) depth = vision_config.depth @@ -335,7 +338,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.config = config + self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config + self.video_pruning_rate = multimodal_config.video_pruning_rate self.is_multimodal_pruning_enabled = ( multimodal_config.is_multimodal_pruning_enabled() ) @@ -345,6 +350,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vision_config, norm_eps=getattr(config, "rms_norm_eps", 1e-6), quant_config=self.quant_config, + input_norm=build_mm_input_norm(self.model_config), prefix=maybe_prefix(prefix, "visual"), use_data_parallel=self.use_data_parallel, ) diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py index d9c7960010a5..41235929867a 100644 --- a/vllm/model_executor/models/exaone4_5_mtp.py +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -75,6 +75,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.layers = nn.ModuleList( Exaone4DecoderLayer( text_config, + cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", ) From b337523c118d1cdb211732199f0a7d5b1e942bb9 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Tue, 13 Oct 2026 13:12:10 +0800 Subject: [PATCH 465/952] [XPU] [Test] Extend spawn-safe sitecustomize monkeypatch to XPU (#57911) Signed-off-by: Chaojun Zhang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> --- .buildkite/intel_jobs/distributed_intel.yaml | 33 ++++++++++++++++++++ tests/v1/shutdown/conftest.py | 15 +++++++-- tests/v1/shutdown/test_forward_error.py | 8 ++--- tests/v1/shutdown/test_startup_error.py | 8 ++--- 4 files changed, 53 insertions(+), 11 deletions(-) create mode 100644 .buildkite/intel_jobs/distributed_intel.yaml diff --git a/.buildkite/intel_jobs/distributed_intel.yaml b/.buildkite/intel_jobs/distributed_intel.yaml new file mode 100644 index 000000000000..54800b4103d2 --- /dev/null +++ b/.buildkite/intel_jobs/distributed_intel.yaml @@ -0,0 +1,33 @@ +group: Distributed Intel +depends_on: + - image-build-xpu +steps: +- label: ":intel: (BMG) Distributed Torchrun + Shutdown" + key: distributed-torchrun-shutdown-tests-2-gpus + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && pytest -v -s v1/shutdown' diff --git a/tests/v1/shutdown/conftest.py b/tests/v1/shutdown/conftest.py index b276d0fab267..ef54b188fefa 100644 --- a/tests/v1/shutdown/conftest.py +++ b/tests/v1/shutdown/conftest.py @@ -10,9 +10,18 @@ @pytest.fixture -def rocm_sitecustomize_factory(monkeypatch, tmp_path: Path): - """Return a function that installs a given sitecustomize payload.""" - if not current_platform.is_rocm(): +def spawn_sitecustomize_factory(monkeypatch, tmp_path: Path): + """Return a function that installs a given sitecustomize payload. + + Both ROCm and XPU force the `spawn` multiprocessing start method (see + `vllm.platforms.rocm`/`vllm.platforms.xpu`), which re-imports modules in + child processes instead of inheriting them via `fork`'s copy-on-write + memory. A plain `monkeypatch.setattr` in the parent test process is + therefore invisible to spawned EngineCore/worker subprocesses on these + platforms, so the payload must be installed via `sitecustomize.py` on + `PYTHONPATH` instead. + """ + if not (current_platform.is_rocm() or current_platform.is_xpu()): return lambda _: None def install(lines: Iterable[str]) -> None: diff --git a/tests/v1/shutdown/test_forward_error.py b/tests/v1/shutdown/test_forward_error.py index 8bc09a64cac9..c40a41cc3b41 100644 --- a/tests/v1/shutdown/test_forward_error.py +++ b/tests/v1/shutdown/test_forward_error.py @@ -41,21 +41,21 @@ def evil_forward(self, *args, **kwargs): @pytest.fixture -def rocm_evil_forward(rocm_sitecustomize_factory): +def evil_forward_fixture(spawn_sitecustomize_factory): lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", "from vllm.model_executor.models.llama import LlamaForCausalLM", inspect.getsource(evil_forward), f"LlamaForCausalLM.forward = {evil_forward.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.asyncio @pytest.mark.parametrize("tensor_parallel_size", [2, 1]) @pytest.mark.parametrize("model", MODELS) async def test_async_llm_model_error( - monkeypatch, rocm_evil_forward, tensor_parallel_size: int, model: str + monkeypatch, evil_forward_fixture, tensor_parallel_size: int, model: str ) -> None: """Test that AsyncLLM propagates a forward pass error and frees memory. @@ -118,7 +118,7 @@ async def generate(request_id: str): @pytest.mark.parametrize("model", MODELS) def test_llm_model_error( monkeypatch, - rocm_evil_forward, + evil_forward_fixture, tensor_parallel_size: int, enable_multiprocessing: bool, model: str, diff --git a/tests/v1/shutdown/test_startup_error.py b/tests/v1/shutdown/test_startup_error.py index 6d6d0d91e77f..d70f69409ee5 100644 --- a/tests/v1/shutdown/test_startup_error.py +++ b/tests/v1/shutdown/test_startup_error.py @@ -30,7 +30,7 @@ def evil_method(self, *args, **kwargs): @pytest.fixture -def rocm_evil_method(rocm_sitecustomize_factory, request): +def evil_method_fixture(spawn_sitecustomize_factory, request): failing_method = request.getfixturevalue("failing_method") lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", @@ -38,7 +38,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): inspect.getsource(evil_method), f"LlamaForCausalLM.{failing_method} = {evil_method.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC) @@ -47,7 +47,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_async_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, failing_method: str, @@ -84,7 +84,7 @@ def test_async_llm_startup_error( @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, enable_multiprocessing: bool, From f7ba007b5e766128468fed38c5da9777b521b8ee Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:13:33 -0400 Subject: [PATCH 466/952] [ROCm][MoE] Split AiterExperts by weight-quantization key (#60322) Signed-off-by: Divakar Verma Co-authored-by: Claude Sonnet 5 --- .buildkite/test-amd.yaml | 27 +++++- docs/design/moe_kernel_features.md | 2 +- .../moe/modular_kernel_tools/common.py | 8 +- .../moe/modular_kernel_tools/mk_objects.py | 40 +++++--- .../moe/test_modular_kernel_combinations.py | 44 ++++++--- tests/kernels/moe/test_mori_moe.py | 8 +- tests/kernels/moe/test_rocm_aiter_moe.py | 4 +- tests/kernels/moe/test_topk_softplus_sqrt.py | 2 +- .../layers/fused_moe/__init__.py | 12 ++- .../fused_moe/experts/rocm_aiter_moe.py | 97 ++++++++++++++----- .../layers/fused_moe/oracle/fp8.py | 6 +- .../layers/fused_moe/oracle/mxfp4.py | 8 +- .../layers/fused_moe/oracle/unquantized.py | 4 +- 13 files changed, 198 insertions(+), 64 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5145f5a272ff..d815ee9268b2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1528,7 +1528,32 @@ steps: commands: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":amd: (MI300) MoE Kernels Multi-GPU" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe/test_modular_kernel_combinations.py + - tests/kernels/moe/modular_kernel_tools/ + - tests/kernels/utils.py + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/distributed/parallel_state.py + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "multigpu" - label: ":amd: (MI300) vLLM IR" # TBD timeout_in_minutes: 50 diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 6c71aa1f46d4..2528bc9ea459 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | -| rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | +| rocm aiter moe | standard | none,
mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai,
situ | Y | Y | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.rocm_aiter_fused_experts],
[`AiterUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterUnquantizedExperts],
[`AiterFp8BlockExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8BlockExperts],
[`AiterFp8TensorExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8TensorExperts],
[`AiterFp8ChannelExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8ChannelExperts],
[`AiterMxfp4Experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterMxfp4Experts] | | cpu_moe | standard | N/A | N/A | silu, gelu,
gelu_tanh,
swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],
[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],
[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 5f6349a32880..1f79b62bf8c8 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -167,6 +167,9 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: vllm_config.model_config = SimpleNamespace( enforce_eager=True, is_moe=True, + enable_sleep_mode=False, + sleep_mode_backend="cumem", + sleep_mode_offload_cudagraph=False, ) vllm_config.parallel_config.data_parallel_size = self.world_size vllm_config.parallel_config.enable_expert_parallel = True @@ -722,6 +725,9 @@ def _maybe_convert_weights_for_experts( rank_weights: WeightTensors, ) -> WeightTensors: """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK).""" + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, convert_to_fp8_moe_kernel_format, @@ -734,7 +740,7 @@ def _maybe_convert_weights_for_experts( # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor # and per-token schemes. Those keep JIT-compiling until # https://github.com/ROCm/aiter/issues/5766 is fixed. - if fe_name == "AiterExperts" and ( + if issubclass(fe_type, AiterExperts) and ( config.quant_dtype is None or config.quant_block_shape is not None ): return _shuffle_weights_for_aiter(rank_weights) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 6ca644c5fce1..56119ffddeb7 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -93,6 +93,7 @@ class ExpertInfo: ] common_float_and_int_types = common_float_types + [torch.int8] nvfp4_types = ["nvfp4"] +mxfp4_types = ["mxfp4"] fp8_types = [current_platform.fp8_dtype()] @@ -318,20 +319,37 @@ def expert_info(kind) -> ExpertInfo: if has_aiter(): from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) - register_experts( - AiterExperts, - standard_format, - # AiterExperts also supports the fully-unquantized (None, None) - # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8. - common_float_types, - blocked_quantization_support=True, - needs_aiter=True, - ) + # Registration dtypes gate which quant configs the modular-kernel test + # harness will ever pair with a given class (see Config.is_valid()): + # they must match what each subclass's _supports_quant_scheme() actually + # declares, or the harness will generate a case the class rejects. + for aiter_experts_cls, aiter_supported_dtypes in ( + (AiterUnquantizedExperts, common_float_types), + (AiterFp8BlockExperts, fp8_types), + (AiterFp8TensorExperts, fp8_types), + (AiterFp8ChannelExperts, fp8_types), + (AiterMxfp4Experts, mxfp4_types), + ): + register_experts( + aiter_experts_cls, + standard_format, + aiter_supported_dtypes, + blocked_quantization_support=True, + needs_aiter=True, + ) else: - AiterExperts = None + AiterUnquantizedExperts = None + AiterFp8BlockExperts = None + AiterFp8TensorExperts = None + AiterFp8ChannelExperts = None + AiterMxfp4Experts = None if has_deep_gemm() and is_deep_gemm_supported(): register_experts( diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index f29c22085240..da643befecd9 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -16,7 +16,13 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterUnquantizedExperts, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTensorSym, kFp8StaticTensorSym, @@ -80,10 +86,10 @@ def assert_aiter_quant_scheme_case(config: Config) -> None: indirectly through the general quant-config sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return - if config.quant_config is None: + if config.quant_config is None or config.quant_dtype is None: w_key, a_key = None, None else: w_key, a_key = config.fp8_quant_key_pair() @@ -101,7 +107,7 @@ def assert_aiter_activation_case(config: Config) -> None: AiterExperts being reached only indirectly through the general activation sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return assert fe_cls._supports_activation(config.activation), ( @@ -156,7 +162,7 @@ def rank_worker( if ( topk == 1 and config.supports_apply_weight_on_input() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_block_shape is not None ): print( @@ -172,7 +178,7 @@ def rank_worker( # https://github.com/vllm-project/vllm/issues/57029 if ( config.world_size > 1 - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and getattr(config.prepare_finalize_type, "__name__", "") in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize") ): @@ -210,7 +216,7 @@ def rank_worker( is_aiter_fp8 = ( _cp.is_rocm() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_config is not None ) if is_aiter_fp8: @@ -270,7 +276,7 @@ def is_nyi_config(config: Config) -> bool: return True if config.activation != MoEActivation.SILU: - if config.fused_experts_type is not AiterExperts: + if not issubclass(config.fused_experts_type, AiterExperts): return True # AITER-only for this axis, for now if config.quant_dtype is not None: return True # unquantized-only for this axis, for now @@ -591,7 +597,7 @@ def test_aiter_moe_sorting_backend_dispatch_env_matrix( dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -671,7 +677,7 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -738,6 +744,20 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) ] +def _padding_matrix_fe_cls( + quant_config: TestMoEQuantConfig | None, +) -> type[AiterExperts]: + """Map a _PADDING_QUANT_CONFIGS entry to the AiterExperts leaf class that + supports it (mirrors Config.fp8_quant_key_pair's scheme derivation).""" + if quant_config is None: + return AiterUnquantizedExperts + if quant_config.block_shape is not None: + return AiterFp8BlockExperts + if quant_config.per_out_ch_quant: + return AiterFp8ChannelExperts + return AiterFp8TensorExperts + + def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor: """Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real gate/up halves of `real_per_half` size each.""" @@ -950,7 +970,7 @@ def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig | dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None, intermediate_size_per_partition_unpadded=( @@ -1113,7 +1133,7 @@ def test_aiter_moe_token_padding_garbage_rows( dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, ) assert config.is_valid()[0] diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py index 9869b6bc5467..f06afc8bd63e 100644 --- a/tests/kernels/moe/test_mori_moe.py +++ b/tests/kernels/moe/test_mori_moe.py @@ -42,7 +42,10 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): from vllm._aiter_ops import rocm_aiter_ops - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterFp8ChannelExperts, + AiterUnquantizedExperts, + ) from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( MoriPrepareAndFinalize, ) @@ -94,7 +97,8 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): assert isinstance(prepare, MoriPrepareAndFinalize) assert prepare.use_fp8_dispatch == fp8 assert prepare.num_dispatchers() == 2 - experts = AiterExperts(config, quant) + experts_cls = AiterFp8ChannelExperts if fp8 else AiterUnquantizedExperts + experts = experts_cls(config, quant) assert not experts.expects_unquantized_inputs kernel = FusedMoEKernel(prepare, experts) local_w1, local_w2 = rocm_aiter_ops.shuffle_weights( diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index abd46753515e..af2f7beb1bd5 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -848,13 +848,13 @@ def test_activation_method_enum_values(): def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): """AITER MXFP4 MoE support should stay gfx950-only.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp4Static, ) - assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() + assert AiterMxfp4Experts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index bc438c66931b..a6f7f3080583 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -87,7 +87,7 @@ def _torch_topk_softplus_sqrt( "backend,expected_experts", [ ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), - ("aiter", "AiterExperts"), + ("aiter", "AiterMxfp4Experts"), ("triton_unfused", "UnfusedOAITritonExperts"), ], ) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 120fbd152668..e68310f98d0d 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -106,7 +106,11 @@ def get_config() -> dict[str, Any] | None: BatchedTritonExperts, ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -132,7 +136,11 @@ def get_config() -> dict[str, Any] | None: ) __all__ += [ - "AiterExperts", + "AiterUnquantizedExperts", + "AiterFp8BlockExperts", + "AiterFp8TensorExperts", + "AiterFp8ChannelExperts", + "AiterMxfp4Experts", "fused_topk", "fused_experts", "get_config_file_name", diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 07eddbf4a092..5c506eae123a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -482,29 +482,6 @@ def _supports_current_device() -> bool: def _supports_no_act_and_mul() -> bool: return False - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - SUPPORTED_W_A = [ - (None, None), - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - (kFp8StaticTensorSym, kFp8DynamicTensorSym), - (kFp8StaticChannelSym, kFp8DynamicTokenSym), - (kMxfp4Static, None), - (kMxfp4Static, kMxfp4Dynamic), - ] - if (weight_key, activation_key) not in SUPPORTED_W_A: - return False - if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950, on_gfx1250 - - if not on_gfx950() or on_gfx1250(): - return False - return True - @staticmethod def _supports_activation(activation: MoEActivation) -> bool: return activation in [ @@ -597,3 +574,77 @@ def apply( output.set_(result) else: output.copy_(result) + + +class AiterUnquantizedExperts(AiterExperts): + """Unquantized (bf16/fp16) weights and activations.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(None, None)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8BlockExperts(AiterExperts): + """128x128 block-quantized weight scale, 128-grouped dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8TensorExperts(AiterExperts): + """Per-tensor static weight scale, either per-tensor static or + per-tensor dynamic activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8DynamicTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8ChannelExperts(AiterExperts): + """Per-output-channel static weight scale, per-token dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8StaticChannelSym, kFp8DynamicTokenSym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterMxfp4Experts(AiterExperts): + """MXFP4-quantized weights, with either unquantized (bf16/fp16) or + MXFP4-quantized dynamic activations. gfx950-only (not gfx1250).""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), + ] + if (weight_key, activation_key) not in SUPPORTED_W_A: + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() and not on_gfx1250() diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 446825c59047..074acd2a0942 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -211,10 +211,12 @@ def backend_to_kernel_cls( elif backend == Fp8MoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, ) - return [AiterExperts] + return [AiterFp8BlockExperts, AiterFp8TensorExperts, AiterFp8ChannelExperts] elif backend == Fp8MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.experts.triton_cutlass_moe import ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index cbbc39b94310..88e80c755d1f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -284,10 +284,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( @@ -305,10 +305,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4 diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 25b55670d1f1..1af0450dfd13 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -126,10 +126,10 @@ def backend_to_kernel_cls( elif backend == UnquantizedMoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterUnquantizedExperts, ) - return [AiterExperts] + return [AiterUnquantizedExperts] elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( From 2ce8012fa0c7fbb37335545325e07dacd8141023 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:21:52 +0300 Subject: [PATCH 467/952] [ROCm][Perf] Skip cleaning sparse prefill MQA logits (#51314) Signed-off-by: Sriram Kumar Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sparse.py | 54 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/v1/attention/test_rocm_aiter_mla_sparse.py diff --git a/tests/v1/attention/test_rocm_aiter_mla_sparse.py b/tests/v1/attention/test_rocm_aiter_mla_sparse.py new file mode 100644 index 000000000000..1d28b57bebec --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_sparse.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.platforms import current_platform + +if not current_platform.is_rocm() or not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse indexer test requires ROCm with AITER.", + allow_module_level=True, + ) + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.v1.attention.ops import rocm_aiter_mla_sparse as sparse_mod + + +def test_rocm_fp8_mqa_logits_disables_clean_logits(monkeypatch): + captured = {} + expected = torch.empty((1, 1)) + + def fake_fp8_mqa_logits(*args, clean_logits): + captured["clean_logits"] = clean_logits + return expected + + monkeypatch.setattr(sparse_mod, "_ON_GFX942", False) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + sparse_mod, + "mqa_logits_module", + lambda: SimpleNamespace(fp8_mqa_logits=fake_fp8_mqa_logits), + ) + + q = torch.empty((1, 1, 1)) + k_fp8 = torch.empty((1, 1)) + scale = torch.empty((1,)) + weights = torch.empty((1, 1)) + cu_seqlen_ks = torch.zeros((1,), dtype=torch.int32) + cu_seqlen_ke = torch.ones((1,), dtype=torch.int32) + + result = sparse_mod.rocm_fp8_mqa_logits( + q, + (k_fp8, scale), + weights, + cu_seqlen_ks, + cu_seqlen_ke, + ) + + assert result is expected + assert captured["clean_logits"] is False diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 6a93d2151127..78ce85f6ae97 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -933,7 +933,9 @@ def rocm_fp8_mqa_logits( if aiter_mqa_logits_module is not None: fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits - return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke) + return fp8_mqa_logits( + q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=False + ) else: return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) From 0d2689e0dfd5e9f74f361cab7701df0feee33ef3 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:25:56 +0300 Subject: [PATCH 468/952] [ROCm][Perf] Skip redundant sparse index remap on non-indexer layers (#51309) Signed-off-by: Sriram Kumar Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sink.py | 2 + .../test_rocm_aiter_mla_sparse_index_remap.py | 295 ++++++++++++++++++ ...est_rocm_aiter_mla_sparse_metadata_sync.py | 4 + .../test_rocm_aiter_triton_sparse_mla.py | 2 + .../backends/mla/rocm_aiter_mla_sparse.py | 62 +++- 5 files changed, 349 insertions(+), 16 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index f66d408c8495..399a1fd0e3d0 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -476,6 +476,8 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): (num_tokens, 128), -1, dtype=torch.int32, device="cuda" ) impl.topk_indices_buffer[:, : selected.shape[1]] = selected + impl.owns_indexer = True + impl.index_epoch = 0 impl.q_concat_buffer = torch.empty_like(q) metadata = SimpleNamespace( attn_out_dtype=dtype, diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py new file mode 100644 index 000000000000..2ad9d3a70a52 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires ROCm.", + allow_module_level=True, + ) + +from vllm._aiter_ops import is_aiter_found_and_supported + +if not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires a supported AITER " + "installation.", + allow_module_level=True, + ) + +from vllm.model_executor.layers.attention import sparse_mla_attention as sparse_mla_mod +from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod + +NUM_HEADS = 16 +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 128 +QK_ROPE_HEAD_DIM = 64 +HEAD_SIZE = 576 +NUM_TOKENS = 2 +TOPK_TOKENS = 4 + + +@pytest.fixture(autouse=True) +def _reset_index_epochs(): + # The counter is keyed by id(), and a freed test buffer can hand its id to + # the next one, so a leftover count would leak across tests. + sparse_mod._INDEX_EPOCHS.clear() + yield + sparse_mod._INDEX_EPOCHS.clear() + + +def _make_metadata(): + return SimpleNamespace( + num_actual_tokens=NUM_TOKENS, + req_id_per_token=torch.zeros(NUM_TOKENS, dtype=torch.int32, device="cpu"), + block_table=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + paged_kv_indptr=torch.zeros(NUM_TOKENS + 1, dtype=torch.int32, device="cpu"), + paged_kv_indices=torch.zeros( + NUM_TOKENS * TOPK_TOKENS, dtype=torch.int32, device="cpu" + ), + block_size=1, + topk_tokens=TOPK_TOKENS, + remapped_buf=None, + remapped_epoch=-1, + ) + + +def _patch_init_deps(monkeypatch, *, speculative: bool = False): + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "check_num_heads_validity", + staticmethod(lambda num_heads: None), + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + model_config=SimpleNamespace(dtype=torch.bfloat16), + speculative_config=SimpleNamespace() if speculative else None, + attention_config=SimpleNamespace(hisparse_config=None), + ) + monkeypatch.setattr(sparse_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr(sparse_mla_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr( + sparse_mla_mod, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + sparse_mod, + "current_workspace_manager", + lambda: SimpleNamespace( + get_simultaneous=lambda *specs: tuple( + torch.zeros(shape, dtype=dtype, device="cpu") for shape, dtype in specs + ) + ), + ) + + +def _build_impl(indexer, topk_indices_buffer): + return sparse_mod.ROCMAiterMLASparseImpl( + num_heads=NUM_HEADS, + head_size=HEAD_SIZE, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + topk_indices_buffer=topk_indices_buffer, + indexer=indexer, + q_lora_rank=None, + kv_lora_rank=KV_LORA_RANK, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM, + v_head_dim=QK_NOPE_HEAD_DIM, + kv_b_proj=None, + ) + + +def _new_buffer(): + return torch.zeros(NUM_TOKENS, TOPK_TOKENS, dtype=torch.int32, device="cpu") + + +def _patch_forward_deps(monkeypatch, impl, remap_calls): + monkeypatch.setattr( + sparse_mod, + "triton_convert_req_index_to_global_index", + lambda *args, **kwargs: remap_calls.append(args), + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "get_mla_padded_q", + staticmethod(lambda num_heads, q: q), + ) + attn_out = torch.zeros(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK, device="cpu") + monkeypatch.setattr( + type(impl), + "_forward_mla", + lambda self, layer, q, kv_cache, attn_metadata: (attn_out, None), + ) + return attn_out + + +def _run_forward(impl, metadata): + return impl.forward_mqa( + torch.zeros(NUM_TOKENS, NUM_HEADS, HEAD_SIZE, device="cpu"), + torch.zeros(4, 1, HEAD_SIZE, device="cpu"), + metadata, + SimpleNamespace(_q_scale=None, _k_scale=None), + ) + + +@pytest.mark.parametrize("has_indexer", [True, False]) +def test_owns_indexer_is_structural(monkeypatch, has_indexer): + """owns_indexer reflects the layer, not the serving config.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer = SimpleNamespace(topk_indices_buffer=buffer) if has_indexer else None + + impl = _build_impl(indexer, buffer) + + assert impl.owns_indexer is has_indexer + assert impl.topk_indices_buffer is buffer + + +def test_forward_mqa_remaps_once_per_metadata(monkeypatch): + """Skip layers sharing a metadata share its remap; the first one pays. + + All layers read one top-k buffer, so the reuse has to key off the metadata + that owns ``paged_kv_indices``, not off the buffer being distinct. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first = _build_impl(None, buffer) + second = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + attn_out = _patch_forward_deps(monkeypatch, first, remap_calls) + metadata = _make_metadata() + + out, lse = _run_forward(first, metadata) + assert out is attn_out + assert lse is None + assert len(remap_calls) == 1 + assert metadata.remapped_buf is buffer + + _run_forward(second, metadata) + assert len(remap_calls) == 1 + + +def test_forward_mqa_remaps_for_each_metadata(monkeypatch): + """Skip layers on a second metadata must remap its own paged_kv_indices. + + Regression: gating on ``indexer is not None`` skipped every layer in the + non-indexer metadata group, leaving its ``paged_kv_indices`` unwritten. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer_layer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_layer = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, indexer_layer, remap_calls) + + indexer_metadata = _make_metadata() + _run_forward(indexer_layer, indexer_metadata) + assert len(remap_calls) == 1 + assert indexer_metadata.remapped_buf is buffer + + # A distinct metadata instance carries its own, still-unwritten indices. + skip_metadata = _make_metadata() + _run_forward(skip_layer, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_buf is buffer + + +def test_forward_mqa_remaps_after_a_later_indexer_write(monkeypatch): + """A skip metadata must be refreshed once per indexer write, not once ever. + + Regression for cross-layer index sharing (``index_topk_freq``). With + freq=4 on a 78-layer model the 57 skip layers share one metadata while the + 21 indexer layers share another, so the skip metadata has to pick up every + new selection. Recording only the buffer left layers 4 onwards attending + to the first selection of the pass, which collapsed RULER niah_single_2 to + 0.026 on GLM-5.3 and took GPQA from 0.884 to 0.672. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + # Construction order mirrors the layer stack: indexer, three skips, indexer. + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + first_skip = _build_impl(None, buffer) + reuse_skip = _build_impl(None, buffer) + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + next_skip = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, first_skip, remap_calls) + skip_metadata = _make_metadata() + + _run_forward(first_skip, skip_metadata) + assert len(remap_calls) == 1 + + # Same selection, same metadata: reuse. + _run_forward(reuse_skip, skip_metadata) + assert len(remap_calls) == 1 + + # A later indexer layer has since rewritten the buffer in place. The object + # is unchanged, so only the epoch can reveal it. + _run_forward(next_skip, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_epoch == next_skip.index_epoch + + +def test_index_epoch_tracks_the_preceding_indexer_write(monkeypatch): + """Layers between two indexer layers share the earlier one's epoch.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_a = _build_impl(None, buffer) + skip_b = _build_impl(None, buffer) + second_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_c = _build_impl(None, buffer) + + assert skip_a.index_epoch == first_indexer.index_epoch + assert skip_b.index_epoch == first_indexer.index_epoch + assert second_indexer.index_epoch == first_indexer.index_epoch + 1 + assert skip_c.index_epoch == second_indexer.index_epoch + + +def test_forward_mqa_remaps_when_buffer_differs(monkeypatch): + """A layer reading a different buffer cannot reuse the metadata's remap.""" + _patch_init_deps(monkeypatch) + remapped_buffer = _new_buffer() + other_buffer = _new_buffer() + layer = _build_impl(None, other_buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, layer, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = remapped_buffer + + _run_forward(layer, metadata) + assert len(remap_calls) == 1 + assert metadata.remapped_buf is other_buffer + + +def test_forward_mqa_indexer_layer_always_remaps(monkeypatch): + """The indexer rewrites its buffer in place, so its remap is never stale.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + impl = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, impl, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = buffer + metadata.remapped_epoch = impl.index_epoch + + _run_forward(impl, metadata) + assert len(remap_calls) == 1 diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index dc1ff79eeabd..4289aac1e7cd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -523,6 +523,8 @@ def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch): impl.kv_cache_dtype = "auto" impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) @@ -570,6 +572,8 @@ def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache impl.kv_cache_dtype = kv_cache_dtype impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 645eec29796d..6849c4e3ebde 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -96,6 +96,8 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( impl.sinks = torch.randn(num_heads, generator=gen).to(device) impl.kv_cache_dtype = kv_cache_dtype impl.topk_indices_buffer = topk_indices.to(device) + impl.owns_indexer = True + impl.index_epoch = 0 impl.use_aiter_sparse_mla = True layer = SimpleNamespace( _q_scale=torch.tensor([0.02], device=device), diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index a0393e772c3e..037b5432ed03 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -401,6 +401,18 @@ def supports_sink(cls) -> bool: return on_mi3xx() +# Indexer writes per shared top-k buffer, in layer construction order. +_INDEX_EPOCHS: dict[int, int] = {} + + +def _index_epoch(buf: object, owns_indexer: bool) -> int: + """Number of indexer writes up to and including this layer.""" + # id(): tensors are unhashable, and the buffer outlives every reader. + epoch = _INDEX_EPOCHS.get(id(buf), 0) + (1 if owns_indexer else 0) + _INDEX_EPOCHS[id(buf)] = epoch + return epoch + + @dataclass(kw_only=True) class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): qo_indptr: torch.Tensor | None = None @@ -409,6 +421,10 @@ class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): paged_kv_indptr: torch.Tensor | None = None attn_out_dtype: torch.dtype | None = None + # Source buffer and index_epoch of paged_kv_indices; -1 before any remap. + remapped_buf: object = None + remapped_epoch: int = -1 + # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). work_meta_data: torch.Tensor | None = None @@ -818,6 +834,10 @@ def __init__( self.masked_mha_available = False vllm_config = get_current_vllm_config() + self.owns_indexer: bool = indexer is not None + self.index_epoch: int = _index_epoch( + self.topk_indices_buffer, self.owns_indexer + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens q_concat_shape = (max_tokens, num_heads, head_size) (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( @@ -1162,25 +1182,35 @@ def forward_mqa( num_actual_toks = q.shape[0] - # Get topk indices - assert self.topk_indices_buffer is not None - topk_indices = fit_kpool_indices_to_aiter( - self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens - ) - kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) - triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - attn_metadata.paged_kv_indptr, - attn_metadata.paged_kv_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=attn_metadata.topk_tokens, - ) + + # Indexer layers may rewrite their buffer, so always remap. Non-indexer + # layers have their own KV cache group: remap once per indexer write. + if ( + self.owns_indexer + or attn_metadata.remapped_epoch != self.index_epoch + or attn_metadata.remapped_buf is not self.topk_indices_buffer + ): + # Get topk indices + assert self.topk_indices_buffer is not None + topk_indices = fit_kpool_indices_to_aiter( + self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens + ) + + triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + attn_metadata.paged_kv_indptr, + attn_metadata.paged_kv_indices, + BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=attn_metadata.topk_tokens, + ) + attn_metadata.remapped_buf = self.topk_indices_buffer + attn_metadata.remapped_epoch = self.index_epoch # write the latent and rope to kv cache if fp8_attention: From 380638afbb2c5658f22124f6d1d38fcd464ea7d9 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 13 Oct 2026 13:46:37 +0800 Subject: [PATCH 469/952] [Model] Add Clef-Flash support (#61308) Signed-off-by: Tianyu Guo Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 26 +- tests/config/test_model_arch_config.py | 34 ++ .../structured_decisions/test_strategies.py | 65 ++-- .../multimodal/pooling/test_clef_omni.py | 7 +- tests/models/registry.py | 4 + .../structured_decisions/strategies.py | 56 +++ vllm/model_executor/models/clef_flash.py | 44 +++ .../models/clef_joint_schema.py | 343 ++++++++++++++++++ vllm/model_executor/models/clef_omni.py | 339 +---------------- vllm/model_executor/models/registry.py | 1 + .../model_arch_config_convertor.py | 22 ++ 11 files changed, 580 insertions(+), 361 deletions(-) create mode 100644 vllm/model_executor/models/clef_flash.py create mode 100644 vllm/model_executor/models/clef_joint_schema.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index ce09d6f81fee..4a7ec00dd2ea 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,30 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Flash + +[Clef-Flash](https://huggingface.co/Cloudflare/clef-flash) uses a Qwen3.5-9B +backbone and a joint schema head to score all questions in one forward pass. +vLLM detects the head and serves the checkpoint as a pooling model. + +```bash +vllm serve Cloudflare/clef-flash --max-model-len 16384 +``` + +Send requests to `/v1/systemone` using `state` and `questions`, as in the +examples above. `choice`, `noul` and `score` questions are supported. + +- Text and JSON states use the checkpoint's fixed prompt. JSON is compact + with sorted keys. `instructions` returns 400; `chat_template_kwargs` + has no effect. +- The state is truncated to fit 16,384 tokens, or `max_model_len` if smaller. + A schema that exceeds this bound returns 400. +- `images` and `videos` are placed before the state. Videos are sampled at + 2 fps and processed by the Qwen3.5 vision processor. Audio inputs return + 400, and video soundtracks are ignored. +- Each request performs one pooling read. `usage.input_tokens` counts the + prompt once, `output_tokens` is 0, and `label_mass` is 1. + ## Clef-Omni [Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: @@ -188,7 +212,7 @@ vllm serve LiquidAI/d1-3B | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Flash](#clef-flash), images/video; [Clef-Omni](#clef-omni); [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 8ec16ee7f9fb..2ff992a20562 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -70,6 +70,40 @@ def _load_groundtruth(filename: str) -> dict: return json.load(f) +@pytest.mark.parametrize("has_joint_head", [False, True]) +@pytest.mark.parametrize( + "architecture", ["Qwen3_5ForConditionalGeneration", "ColQwen3_5"] +) +def test_qwen35_joint_head_selects_clef_flash(tmp_path, has_joint_head, architecture): + """Only a Qwen3.5 generation checkpoint with a joint head selects Clef.""" + from transformers import Qwen3_5Config + + head_config = { + "hidden_size": 4096, + "width": 1024, + "routing_layers": 2, + "layers": 4, + "heads": 16, + "feedforward": 4096, + } + if has_joint_head: + (tmp_path / "joint_head_config.json").write_text(json.dumps(head_config)) + config = Qwen3_5Config(architectures=[architecture]) + config.name_or_path = str(tmp_path) + model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config.text_config, + _supports_multimodal_for_mm_prefix=lambda: True, + ) + converted = ModelConfig.get_model_arch_config(model_config) + if has_joint_head and architecture == "Qwen3_5ForConditionalGeneration": + assert converted.architectures == ["ClefFlashForJointSchema"] + assert config.joint_head_config == head_config + else: + assert converted.architectures == [architecture] + assert not hasattr(config, "joint_head_config") + + def _assert_model_arch_config( model_config, expected: dict, check_head_size: bool = True ): diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index dae545d1f04c..505f0e35eef1 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -33,6 +33,7 @@ parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + ClefFlashStrategy, DiffusionGemmaCanvasStrategy, JointSchemaStrategy, LiquidStrategy, @@ -94,6 +95,8 @@ def test_strategy_selection(): # Clef-Omni scores options with its head, so it reads no logprobs. clef = model("ClefOmniForJointSchema", "raw_logits") assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy + clef_flash = model("ClefFlashForJointSchema", "raw_logits") + assert select_read_strategy(clef_flash, ("token_classify",)) is ClefFlashStrategy def test_route_is_registered_by_default(): @@ -366,8 +369,8 @@ def test_canvas_read(): ) -@pytest.fixture -def joint_schema(qwen, engine_and_renderer): +@pytest.fixture(params=[JointSchemaStrategy, ClefFlashStrategy]) +def joint_schema(request, qwen, engine_and_renderer): tokenizer, _ = qwen engine, online_renderer = engine_and_renderer # The probability the head gives each option, by option id. @@ -382,9 +385,7 @@ async def encode(prompt, params, request_id, **kwargs): yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) engine.encode.side_effect = encode - strategy = JointSchemaStrategy( - ReadContext(engine, online_renderer, None, "auto", {}) - ) + strategy = request.param(ReadContext(engine, online_renderer, None, "auto", {})) return serving(engine, strategy), engine, tokenizer @@ -468,7 +469,10 @@ async def test_media_needs_a_model_that_takes_it(serving_and_engine): @pytest.mark.asyncio -async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): +@pytest.mark.parametrize("num_images", [1, 2]) +async def test_joint_schema_reads_media_before_the_state( + joint_schema, monkeypatch, num_images +): serving, engine, tokenizer = joint_schema strategy = serving.strategy strategy.__dict__["_connector"] = MediaConnector() @@ -491,23 +495,30 @@ async def process_for_engine_async(prompt, arrival_time): Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() response = await serving.create_decision( - StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image] * num_images) ) assert isinstance(response, StructuredDecisionResponse) - placeholder = tokenizer.encode( - "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False - ) + if isinstance(strategy, ClefFlashStrategy): + media_text = "<|vision_start|><|image_pad|><|vision_end|>" * num_images + "\n" + else: + media_text = "<|vision_start|><|image_pad|><|vision_end|>\n" * num_images + placeholder = tokenizer.encode(media_text, add_special_tokens=False) assert rendered["prompt_token_ids"] == strategy.prefix + placeholder assert rendered["multi_modal_data"]["image"][0].size == (8, 8) - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE} - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } prompt = engine.encode.call_args.args[0] assert prompt["type"] == "multimodal" ids = prompt["prompt_token_ids"] - expanded = len(strategy.prefix) + len(placeholder) + 2 - assert ids[:expanded].count(image_pad) == 3 + expanded = len(strategy.prefix) + len(placeholder) + 2 * num_images + assert ids[:expanded].count(image_pad) == 3 * num_images state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] assert tokenizer.decode(ids[expanded:]) == state_and_questions assert response.answers["team"]["choice"] == "billing" @@ -515,7 +526,9 @@ async def process_for_engine_async(prompt, arrival_time): @pytest.mark.asyncio -async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): +async def test_joint_schema_video_audio_follows_model_modalities( + joint_schema, monkeypatch +): serving, engine, tokenizer = joint_schema strategy = serving.strategy video = MediaWithBytes( @@ -532,6 +545,7 @@ async def load_from_url_async(self, url, media_io, *, fetch_timeout=None): return video async def fetch_audio_async(self, url): + assert not isinstance(strategy, ClefFlashStrategy) if url == "https://example.org/silent.mp4": raise ValueError("No audio stream found.") return track @@ -552,11 +566,17 @@ async def process_for_engine_async(prompt, arrival_time): ) ) assert isinstance(response, StructuredDecisionResponse) - assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE}, - "use_audio_in_video": True, - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } # A video without a decodable soundtrack is watched silently. response = await serving.create_decision( @@ -566,7 +586,8 @@ async def process_for_engine_async(prompt, arrival_time): ) assert isinstance(response, StructuredDecisionResponse) assert rendered["multi_modal_data"] == {"video": [video]} - assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + if not isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False # Audio clips cannot be read with videos that have soundtracks. response = await serving.create_decision( diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index ccdc913b3132..381ac23d53cc 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +"""The shared Clef pooler against a port of the checkpoint's joint schema head, which loops over questions where the pooler batches them.""" import math @@ -11,7 +11,10 @@ import torch import torch.nn.functional as F -from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.model_executor.models.clef_joint_schema import ( + JointSchemaHead, + JointSchemaPooler, +) from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates diff --git a/tests/models/registry.py b/tests/models/registry.py index 0a7027461ef3..b08efd9fb0f1 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -802,6 +802,10 @@ def check_available_online( "Salesforce/blip2-opt-2.7b", extras={"6b": "Salesforce/blip2-opt-6.7b"}, ), + "ClefFlashForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-flash", + max_model_len=4096, + ), "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025", extras={"command-a-plus": "CohereLabs/command-a-plus-05-2026-bf16"}, diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 5567cb808ab5..66238cf7f103 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -809,7 +809,63 @@ async def read( return reads +class ClefFlashStrategy(JointSchemaStrategy): + """Clef-Flash's joint schema prompt with Qwen3.5 image/video processing.""" + + modalities = frozenset({"image", "video"}) + MAX_LENGTH = 16384 + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>", + } + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + connector = self._connector + video_io = VideoMediaIO(ImageMediaIO(), num_frames=-1, fps=self.VIDEO_FPS) + try: + images, videos = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["video"] * len(videos) + + "\n" + ) + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if videos: + mm_data["video"] = videos + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs={"videos_kwargs": {"do_sample_frames": False}}, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + try: + engine_input = ( + await self.context.online_renderer.renderer.process_for_engine_async( + prompt, time.time() + ) + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefFlashForJointSchema": ClefFlashStrategy, "ClefOmniForJointSchema": JointSchemaStrategy, "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, diff --git a/vllm/model_executor/models/clef_flash.py b/vllm/model_executor/models/clef_flash.py new file mode 100644 index 000000000000..32a9e0a59003 --- /dev/null +++ b/vllm/model_executor/models/clef_flash.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Flash: the Qwen3.5 backbone with a joint schema decision head.""" + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler +from .interfaces_base import default_pooling_type +from .qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5ProcessingInfo +from .qwen3_vl import Qwen3VLDummyInputsBuilder, Qwen3VLMultiModalProcessor + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class ClefFlashForJointSchema(Qwen3_5ForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**model_config.hf_config.joint_head_config), + self._embed_lexical, + ) + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/clef_joint_schema.py b/vllm/model_executor/models/clef_joint_schema.py new file mode 100644 index 000000000000..fad2beb09899 --- /dev/null +++ b/vllm/model_executor/models/clef_joint_schema.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared joint schema head for Clef decision models. + +The head reads the final hidden states of the whole prompt, so it runs as a +``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Set +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py index 7cd0d676873b..1533dde27378 100644 --- a/vllm/model_executor/models/clef_omni.py +++ b/vllm/model_executor/models/clef_omni.py @@ -1,47 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). - -The Qwen3-Omni thinker with a joint schema head that scores every option of -every question in one forward pass. The head reads the final hidden states of -the whole prompt, so it runs as a ``token_classify`` pooler. - -Each request carries its schema in ``PoolingParams.extra_kwargs`` under -``"joint_schema"``: - -- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. -- ``question_spans``: the ``[start, end)`` of each question's instructions. -- ``option_spans``: the ``[start, end)`` of each option, question by question. -- ``option_counts``: the number of options of each question. - -Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), -so they hold after multimodal placeholders are expanded. The output has one -value per option, in ``option_spans`` order: its probability within its -question, or its logit with ``use_activation=False``. -""" +"""The Qwen3-Omni thinker with a Clef joint schema head.""" import math -from collections.abc import Callable, Mapping, Set -from dataclasses import dataclass +from collections.abc import Mapping from typing import Any import numpy as np import torch import torch.nn.functional as F -from torch import nn from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig -from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate -from vllm.model_executor.layers.pooler.tokwise.methods import AllPool from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.tasks import PoolingTask -from vllm.utils.torch_utils import async_tensor_h2d -from vllm.v1.outputs import PoolerOutput -from vllm.v1.pool.metadata import PoolingMetadata +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler from .interfaces_base import default_pooling_type from .qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, @@ -50,314 +25,6 @@ Qwen3OmniMoeThinkerProcessingInfo, ) -JOINT_HEAD_WEIGHTS = "joint_head.safetensors" -NUM_QUESTION_TYPES = 3 -MAX_LOGIT_SCALE = math.log(100.0) - - -def segment_mean( - values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor -) -> torch.Tensor: - sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) - sums.index_add_(0, segment_ids, values.float()) - return (sums / lengths[:, None]).to(values.dtype) - - -@dataclass -class SchemaIndices: - """Where one prompt's questions and options are, on the model's device.""" - - positions: torch.Tensor - """Every token of every span, questions first.""" - span_ids: torch.Tensor - """The span of each of ``positions``.""" - span_lengths: torch.Tensor - """``[Q + O]``.""" - question_types: torch.Tensor - """``[Q]``.""" - option_question: torch.Tensor - """The question of each option, ``[O]``.""" - option_slots: torch.Tensor - """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" - option_mask: torch.Tensor - """``[Q, max_options]``, true at the slots of options.""" - - -class EvidenceRoutingLayer(nn.Module): - def __init__(self, width: int, heads: int, feedforward: int) -> None: - super().__init__() - self.query_norm = nn.LayerNorm(width) - self.memory_norm = nn.LayerNorm(width) - self.attention = nn.MultiheadAttention(width, heads, batch_first=True) - self.feedforward_norm = nn.LayerNorm(width) - # The Identity takes the place of the checkpoint's dropout. - self.feedforward = nn.Sequential( - nn.Linear(width, feedforward), - nn.GELU(), - nn.Identity(), - nn.Linear(feedforward, width), - ) - - def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: - # Separate key and value tensors, like the reference, keep attention's - # unpacked input projections. - routed, _ = self.attention( - self.query_norm(queries), - self.memory_norm(memory), - self.memory_norm(memory), - need_weights=False, - ) - queries = queries + routed - return queries + self.feedforward(self.feedforward_norm(queries)) - - -class JointSchemaHead(nn.Module): - def __init__( - self, - hidden_size: int, - width: int, - routing_layers: int, - layers: int, - heads: int, - feedforward: int, - ) -> None: - super().__init__() - self.hidden_norm = nn.LayerNorm(hidden_size) - self.memory_projection = nn.Linear(hidden_size, width, bias=False) - self.question_projection = nn.Linear(hidden_size, width, bias=False) - self.option_question_projection = nn.Linear(hidden_size, width, bias=False) - self.global_projection = nn.Linear(hidden_size, width, bias=False) - self.option_context_projection = nn.Linear(hidden_size, width, bias=False) - self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) - self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) - self.evidence_layers = nn.ModuleList( - EvidenceRoutingLayer(width, heads, feedforward) - for _ in range(routing_layers) - ) - self.option_summary_norm = nn.LayerNorm(width) - self.layers = nn.ModuleList( - nn.TransformerDecoderLayer( - d_model=width, - nhead=heads, - dim_feedforward=feedforward, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=True, - ) - for _ in range(layers) - ) - self.field_norm = nn.LayerNorm(width) - self.option_norm = nn.LayerNorm(width) - self.residual_scorer = nn.Sequential( - nn.Linear(width * 4, width), - nn.GELU(), - nn.Identity(), - nn.Linear(width, 1), - ) - self.prior_logit_scale = nn.Parameter(torch.zeros(())) - self.joint_logit_scale = nn.Parameter(torch.zeros(())) - self.residual_gate = nn.Parameter(torch.zeros(())) - - def forward( - self, - hidden_states: torch.Tensor, - lexical: torch.Tensor, - schema: SchemaIndices, - ) -> torch.Tensor: - """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` - final hidden states and each option's ``[O, hidden]`` mean output - embedding.""" - num_questions = schema.question_types.shape[0] - option_question = schema.option_question - hidden_states = self.hidden_norm(hidden_states) - memory = self.memory_projection(hidden_states).unsqueeze(0) - global_vector = hidden_states[-1] - means = segment_mean( - hidden_states[schema.positions], schema.span_ids, schema.span_lengths - ) - question_vectors, contexts = means[:num_questions], means[num_questions:] - - options = ( - self.option_context_projection(contexts) - + self.option_lexical_projection(lexical) - + self.option_question_projection(question_vectors)[option_question] - ).unsqueeze(0) - for layer in self.evidence_layers: - options = layer(options, memory) - options = options[0] - - # Each field summarizes its options, weighted by their match to it. - base_fields = self.question_projection(question_vectors) - width = options.shape[-1] - grid = options.new_zeros(schema.option_mask.numel(), width) - grid[schema.option_slots] = options - grid = grid.view(*schema.option_mask.shape, width) - routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) - routing = routing.masked_fill(~schema.option_mask, float("-inf")) - summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) - fields = ( - base_fields - + self.option_summary_norm(summaries) - + self.global_projection(global_vector).unsqueeze(0) - + self.type_embedding(schema.question_types) - ).unsqueeze(0) - for layer in self.layers: - fields = layer(fields, memory) - fields = self.field_norm(fields[0])[option_question] - - anchors = F.normalize(question_vectors + global_vector, dim=-1) - prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - prior = prior_scale * ( - F.normalize(lexical, dim=-1) * anchors[option_question] - ).sum(dim=-1) - options = self.option_norm(options) - cosine = F.cosine_similarity(fields, options, dim=-1) - features = torch.cat( - [fields, options, fields * options, torch.abs(fields - options)], dim=-1 - ) - residual = self.residual_scorer(features).squeeze(-1) - joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - joint = joint_scale * cosine + residual - return prior + torch.sigmoid(self.residual_gate) * joint - - -def parse_joint_schema( - extra_kwargs: dict[str, Any] | None, prompt_len: int -) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: - """The question types, the absolute spans of the questions then of the - options, and the option counts. None if the schema is missing or invalid.""" - schema = (extra_kwargs or {}).get("joint_schema") - if not isinstance(schema, dict): - return None - try: - types = [int(t) for t in schema["question_types"]] - counts = [int(c) for c in schema["option_counts"]] - question_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["question_spans"] - ] - option_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["option_spans"] - ] - except (KeyError, TypeError, ValueError): - return None - spans = question_spans + option_spans - if ( - not types - or len(counts) != len(types) - or len(question_spans) != len(types) - or len(option_spans) != sum(counts) - or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) - or any(c < 1 for c in counts) - or any(not 0 <= start < end <= prompt_len for start, end in spans) - ): - return None - return types, spans, counts - - -class JointSchemaPooler(Pooler): - """Runs the joint schema head on each finished prompt. - - It never raises, because it runs in the engine core: a prompt without a - valid schema, such as a profiling dummy, gets an empty output. - """ - - def __init__( - self, - head: JointSchemaHead, - embed_lexical: Callable[[torch.Tensor], torch.Tensor], - ) -> None: - super().__init__() - self.pooling = AllPool() - self.head = head - self.embed_lexical = embed_lexical - - def get_supported_tasks(self) -> Set[PoolingTask]: - return {"token_classify"} - - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: - return PoolingParamsUpdate(requires_token_ids=True) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled = self.pooling(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() - outputs: list[torch.Tensor | None] = [] - for hidden, token_ids, params in zip( - pooled, prompt_token_ids, pooling_metadata.pooling_params - ): - if hidden is None: - outputs.append(None) - continue - parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) - if parsed is None: - outputs.append(hidden.new_empty(0, dtype=torch.float32)) - continue - schema, lexical_ids, lexical_span_ids = self._indices( - *parsed, hidden.device - ) - lexical = segment_mean( - self.embed_lexical( - async_tensor_h2d(token_ids[lexical_ids], hidden.device) - ), - lexical_span_ids, - schema.span_lengths[len(parsed[0]) :], - ) - logits = self.head(hidden, lexical, schema).float() - if params.use_activation is not False: - mask = schema.option_mask - grid = logits.new_full((mask.numel(),), float("-inf")) - grid[schema.option_slots] = logits - probs = grid.view(mask.shape).softmax(dim=-1) - logits = probs.view(-1)[schema.option_slots] - outputs.append(logits) - return outputs - - @staticmethod - def _indices( - types: list[int], - spans: list[tuple[int, int]], - counts: list[int], - device: torch.device, - ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: - """The schema's indices, plus the option tokens' positions (on the CPU) - and their option spans, for the lexical embeddings.""" - num_questions = len(types) - lengths = torch.tensor([end - start for start, end in spans]) - span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) - positions = torch.cat([torch.arange(start, end) for start, end in spans]) - num_question_tokens = int(lengths[:num_questions].sum()) - - counts_cpu = torch.tensor(counts) - max_options = max(counts) - option_question = torch.repeat_interleave( - torch.arange(num_questions), counts_cpu - ) - option_slots = torch.cat( - [torch.arange(count) + q * max_options for q, count in enumerate(counts)] - ) - schema = SchemaIndices( - positions=async_tensor_h2d(positions, device), - span_ids=async_tensor_h2d(span_ids, device), - span_lengths=async_tensor_h2d(lengths, device), - question_types=async_tensor_h2d(types, device, dtype=torch.long), - option_question=async_tensor_h2d(option_question, device), - option_slots=async_tensor_h2d(option_slots, device), - option_mask=async_tensor_h2d( - torch.arange(max_options) < counts_cpu[:, None], device - ), - ) - lexical_span_ids = async_tensor_h2d( - span_ids[num_question_tokens:] - num_questions, device - ) - return schema, positions[num_question_tokens:], lexical_span_ids - - MAX_FRAME_PIXELS = 256 * 32 * 32 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 36f9bd6e6fcc..cd0557ac4585 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefFlashForJointSchema": ("clef_flash", "ClefFlashForJointSchema"), "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 147b65759dea..06a937697106 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -647,6 +647,27 @@ def __init__( _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3_5ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + if getattr(hf_config, "architectures", None) == [ + "Qwen3_5ForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefFlashForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -850,6 +871,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "nemotron-nas": NemotronNasModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, + "qwen3_5": Qwen3_5ModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, From 86b1c2ef159dcc268d219b074cb047340ea718aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 13 Oct 2026 07:51:22 +0200 Subject: [PATCH 470/952] [Attention][DCP] Support a2a DCP comm backend for FlashMLA sparse (#61285) Signed-off-by: NickLucche Co-authored-by: Claude Opus 5.5 --- tests/v1/attention/test_sparse_mla_backends.py | 16 ++++++++++++++++ .../v1/attention/backends/mla/flashmla_sparse.py | 8 +------- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index c21fc15a5278..f7c47fe6a3ec 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -4157,6 +4157,22 @@ def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise): assert local_pad == gathered_pad +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +@pytest.mark.parametrize("comm_backend", ["ag_rs", "a2a"]) +def test_fp8_dcp_accepts_comm_backend(comm_backend): + """The mixed-batch path returns (out, lse) for every token in the layout + both DCP combines consume, so neither comm backend is rejected.""" + device = torch.device(DEVICE_TYPE) + vllm_config = _build_sparse_dcp_vllm_config(16, 4, comm_backend) + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + assert builder.fp8_use_mixed_batch + + def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 3cf0c33c534d..80b8da179769 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -446,12 +446,6 @@ def __init__( ) if parallel_config.decode_context_parallel_size > 1: - if parallel_config.dcp_comm_backend != "ag_rs": - raise NotImplementedError( - "DCP for FlashMLA sparse is only validated with the " - "default 'ag_rs' DCP comm backend; got " - f"'{parallel_config.dcp_comm_backend}'" - ) if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": raise NotImplementedError( "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " @@ -1268,7 +1262,7 @@ def _forward_fp8_kv_mixed_batch( out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) # The head-padding slice above can leave `out` non-contiguous, and the - # merge feeds it to reduce_scatter. + # DCP combines (reduce_scatter, direct A2A) need packed heads. return out.contiguous(), lse def _fp8_flash_mla_kernel( From a6009ea85be08161f4b5954bbf317027f6ddfd82 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 12 Oct 2026 23:24:43 -0700 Subject: [PATCH 471/952] [Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer (#61149) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 (1M context) --- cmake/external_projects/deepgemm.cmake | 12 +- csrc/cache_kernels.cu | 20 +- .../attention/test_deepgemm_attention.py | 189 +++++++++++------- .../test_rocm_paged_mxfp4_indexer.py | 1 + tests/kernels/test_compressor_kv_cache.py | 42 +++- .../test_fused_deepseek_v32_norm_rope.py | 15 +- .../test_fused_indexer_q_rope_quant.py | 5 + .../test_kpool_decode_update_batched.py | 8 +- .../layers/test_mla_short_prefill_indexer.py | 39 ++-- .../test_indexer_deepseek_v4_slot_mapping.py | 11 +- .../test_rocm_paged_mxfp4_indexer_plan.py | 1 + tools/build_deepgemm_C.py | 3 +- tools/check_wheel_deepgemm.py | 11 +- tools/install_deepgemm.sh | 8 +- .../layers/sparse_attn_indexer.py | 110 +++++++--- vllm/models/deepseek_v32/common/kernels.py | 14 +- .../common/ops/fused_compress_quant_cache.py | 12 +- .../common/ops/indexer_k_store.py | 15 +- .../glm5next/nvidia/ops/kpool_compress.py | 30 ++- vllm/models/glm5next/nvidia/sparse_indexer.py | 72 ++++--- vllm/utils/deep_gemm.py | 101 +++++++--- vllm/v1/attention/backends/mla/indexer.py | 4 + vllm/v1/attention/ops/common.py | 7 + 23 files changed, 519 insertions(+), 211 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 1dcee345f845..b796623b7ac1 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -29,12 +29,11 @@ if(DEEPGEMM_SRC_DIR) else() # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") - # Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120 - # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU - # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x - # layout header fix from vllm-project/DeepGEMM#12. - # TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "1e1842a833699298f7afc02eefb2ed168fff6938") + # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) + # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, + # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 + # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -146,6 +145,7 @@ if(DEEPGEMM_ARCHS) # install(FILES "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py" + "${deepgemm_SOURCE_DIR}/deep_gemm/locality_domain.py" DESTINATION vllm/third_party/deep_gemm COMPONENT _deep_gemm_C) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index d1c1e9df9896..24453ff1c154 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -630,7 +630,8 @@ __global__ void indexer_k_quant_and_cache_kernel( const int cache_block_size, // cache block size const int64_t cache_block_stride, // stride for each block in kv_cache - const bool use_ue8m0 // use ue8m0 scale format + const bool use_ue8m0, // use ue8m0 scale format + const bool use_mx_scale // store the ue8m0 scale as packed MX bytes ) { constexpr int VEC_SIZE = 4; const int64_t token_idx = blockIdx.x; @@ -679,7 +680,13 @@ __global__ void indexer_k_quant_and_cache_kernel( const int64_t dst_scale_idx = block_idx * cache_block_stride + cache_block_size * head_dim + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; - reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + if (use_mx_scale) { + // One UE8M0 byte per 32 elements, all equal to the shared scale's. + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = + ((__float_as_uint(scale) >> 23) & 0xFF) * 0x01010101u; + } else { + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + } } } @@ -1912,7 +1919,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( reinterpret_cast(k.data_ptr()), \ reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ - cache_block_size, cache_block_stride, use_ue8m0); + cache_block_size, cache_block_stride, use_ue8m0, use_mx_scale); void indexer_k_quant_and_cache( torch::stable::Tensor& k, // [num_tokens, head_dim] @@ -1924,8 +1931,13 @@ void indexer_k_quant_and_cache( int head_dim = k.size(1); int cache_block_size = kv_cache.size(1); int64_t cache_block_stride = kv_cache.stride(0); - bool use_ue8m0 = scale_fmt == "ue8m0"; + // "mxfp8" stores the ue8m0 scale as four packed UE8M0 bytes, the MX scales + // DeepGEMM's SM100 MQA logits kernels read. + bool use_mx_scale = scale_fmt == "mxfp8"; + bool use_ue8m0 = use_mx_scale || scale_fmt == "ue8m0"; + STD_TORCH_CHECK(!use_mx_scale || (head_dim == 128 && quant_block_size == 128), + "mxfp8 scales require head_dim == quant_block_size == 128"); STD_TORCH_CHECK(k.device() == kv_cache.device(), "k and kv_cache must be on the same device"); STD_TORCH_CHECK(k.device() == slot_mapping.device(), diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index f332fb0fea21..109cc83830b1 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -5,6 +5,7 @@ import pytest import torch +from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( _ceil_to_ue8m0, @@ -13,18 +14,29 @@ fp8_fp4_paged_mqa_logits, get_num_sms, get_paged_mqa_logits_metadata, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, native_next_n_supported, ) from vllm.utils.import_utils import has_deep_gemm from vllm.utils.math_utils import cdiv +def encode_indexer_scales(sf: torch.Tensor) -> torch.Tensor: + """Per-token power-of-two fp32 scales as the indexer K cache stores them: + four equal packed UE8M0 bytes under MX scales, else the fp32 itself.""" + if not mqa_logits_use_mx_scales(): + return sf + ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8) + return ue8m0[..., None].repeat_interleave(4, dim=-1).view(torch.int32)[..., 0] + + def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: # x: (num_blocks, block_size, 1, head_dim) num_blocks, block_size, num_heads, head_dim = x.shape assert num_heads == 1 x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) - sf = x_amax / 448.0 + sf = _ceil_to_ue8m0(x_amax / 448.0) x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) x_fp8 = torch.empty( (num_blocks, block_size * (head_dim + 4)), @@ -34,9 +46,9 @@ def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: x_fp8[:, : block_size * head_dim] = x_scaled.view( num_blocks, block_size * head_dim ).view(dtype=torch.uint8) - x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view( - dtype=torch.uint8 - ) + x_fp8[:, block_size * head_dim :] = encode_indexer_scales( + sf.view(num_blocks, block_size) + ).view(dtype=torch.uint8) return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) @@ -91,64 +103,72 @@ def _ref_fp8_mqa_logits( return logits +def _packed_requests_test_data(seq_len: int, seq_len_kv: int): + # Two causal prefill requests packed into one KV buffer, so the second + # request's rows start at a nonzero K position. + half, half_kv = seq_len // 2, seq_len_kv // 2 + rows = torch.arange(seq_len, dtype=torch.int, device="cuda") + ks = torch.where(rows < half, 0, half_kv).to(torch.int) + ke = ks + half_kv - half + 1 + rows % half + return ks, ke + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @pytest.mark.skipif( not current_platform.has_device_capability(90), reason="SM90 and SM100 only" ) -@pytest.mark.parametrize("clean_logits", [True, False]) -def test_deepgemm_fp8_mqa_logits(clean_logits: bool): +@pytest.mark.parametrize("layout", ["cp", "no_cp", "packed"]) +def test_deepgemm_fp8_mqa_logits(layout: str): + """FP8 indexer prefill logits, compressed to each row's valid K span.""" torch.manual_seed(0) random.seed(0) num_heads, head_dim = 32, 128 - for seq_len in (512,): - for seq_len_kv in (1024,): - for disable_cp in (False, True): - q = torch.randn( - seq_len, - num_heads, - head_dim, - device="cuda", - dtype=torch.bfloat16, - ) - kv = torch.randn( - seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16 - ) - weights = torch.randn( - seq_len, num_heads, device="cuda", dtype=torch.float32 - ) - - if disable_cp: - ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") - ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( - seq_len_kv - seq_len - ) - else: - ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) - - q_fp8 = q.to(torch.float8_e4m3fn) - kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) - logits = fp8_fp4_mqa_logits( - (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits - ) - - ref_logits = _ref_fp8_mqa_logits( - q=q, - kv=kv, - weights=weights, - cu_seqlen_ks=ks, - cu_seqlen_ke=ke, - ) - ref_neginf_mask = ref_logits == float("-inf") - - if clean_logits: - neginf_mask = logits == float("-inf") - assert torch.equal(neginf_mask, ref_neginf_mask) - - ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) - logits = logits.masked_fill(ref_neginf_mask, 0) - diff = calc_diff(logits, ref_logits) - assert diff < 1e-3, f"{diff=}" + seq_len, seq_len_kv = 512, 1024 + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32) + + if layout == "cp": + ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) + elif layout == "no_cp": + ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") + ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( + seq_len_kv - seq_len + ) + else: + ks, ke = _packed_requests_test_data(seq_len, seq_len_kv) + + q_fp8 = q.to(torch.float8_e4m3fn) + kv_fp8, kv_sf = per_custom_dims_cast_to_fp8(kv, (0,), True) + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + max_seqlen_k = int((ke - ks).max()) + logits = fp8_fp4_mqa_logits( + (q_fp8, q_scale), + (kv_fp8, encode_indexer_scales(kv_sf)), + kernel_weights, + ks, + ke, + max_seqlen_k, + ) + assert logits.shape == (seq_len, max_seqlen_k) + logits = as_fp32_logits(logits, ke - ks) + + ref_logits = _ref_fp8_mqa_logits( + q=q, + kv=kv, + weights=kernel_weights.float(), + cu_seqlen_ks=ks, + cu_seqlen_ke=ke, + ) + # Compressed row i holds K positions [ks[i], ke[i]) from column zero. + cols = ks[:, None] + torch.arange(max_seqlen_k, device="cuda")[None, :] + valid = cols < ke[:, None] + ref_logits = ref_logits.gather(1, cols.clamp(max=seq_len_kv - 1).long()) + ref_logits = ref_logits.masked_fill(~valid, 0) + logits = logits.masked_fill(~valid, 0) + diff = calc_diff(logits, ref_logits) + assert diff < 1e-3, f"{diff=}" def _ref_fp8_fp4_paged_mqa_logits( @@ -212,9 +232,6 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): if not native_next_n_supported(next_n): pytest.skip(f"next_n={next_n} has no native kernel on this architecture") - # NOTE: clean_logits=True is incompatible with the 2D context_lens - # required by csrc/apis/attention.hpp; only the False path is exercised. - clean_logits = False torch.manual_seed(0) random.seed(0) @@ -273,26 +290,42 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens_2d = ( context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange ).contiguous() + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + kernel_q = (q_fp8, q_scale) + kernel_context_lens, kernel_block_tables = context_lens_2d, block_tables + indices = None + if mqa_logits_use_mx_scales(): + # SM100 only takes varlen rows: one Q token each, tied to its + # request by `indices`. + assert q_scale is not None + kernel_q = ( + q_fp8.flatten(0, 1)[:, None], + q_scale.flatten(0, 1)[:, None], + ) + kernel_context_lens = context_lens_2d.view(-1, 1) + kernel_block_tables = block_tables.repeat_interleave(next_n, dim=0) + indices = torch.arange(batch_size, device="cuda", dtype=torch.int32) + indices = indices.repeat_interleave(next_n) schedule_metadata = get_paged_mqa_logits_metadata( - context_lens_2d, - blocksize, - get_num_sms(), + kernel_context_lens, blocksize, get_num_sms(), indices=indices ) logits = fp8_fp4_paged_mqa_logits( - (q_fp8, None), + kernel_q, kv_cache_fp8, - weights, - context_lens_2d, - block_tables, + kernel_weights, + kernel_context_lens, + kernel_block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, + indices=indices, ) + if indices is not None: + logits = as_fp32_logits(logits, kernel_context_lens.view(-1)) ref_logits = _ref_fp8_fp4_paged_mqa_logits( q, kv_cache, - weights, + kernel_weights.float(), context_lens, block_tables, max_model_len, @@ -309,7 +342,7 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens[row_indices] - next_n + next_n_offset ).unsqueeze(1) - logits = logits.masked_fill(~mask, 0) + logits = logits.float().masked_fill(~mask, 0) ref_logits = ref_logits.masked_fill(~mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" @@ -349,12 +382,28 @@ def test_deepgemm_paged_mqa_packed_manager_block_stride(): q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to( torch.float8_e4m3fn ) - weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + q_scale, weights = mqa_logits_mx_inputs( + q, None, torch.randn(1, num_heads, device="cuda", dtype=torch.float32) + ) context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32) - schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms()) + indices = ( + torch.zeros(1, device="cuda", dtype=torch.int32) + if mqa_logits_use_mx_scales() + else None + ) + schedule = get_paged_mqa_logits_metadata( + context_lens, page_size, get_num_sms(), indices=indices + ) logits = [ fp8_fp4_paged_mqa_logits( - (q, None), cache, weights, context_lens, table, schedule, block_size, False + (q, q_scale), + cache, + weights, + context_lens, + table, + schedule, + block_size, + indices=indices, ) for cache, table in zip((compact_pages, packed_pages), tables) ] diff --git a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py index 311ee89a8d12..e270ae4a5813 100644 --- a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py +++ b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py @@ -253,6 +253,7 @@ def _prefill_metadata( cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE), token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=t0, token_end=t1, num_reqs=req_hi - req_lo, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a931..e9924229055a 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1422,7 +1422,10 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) -def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): +@pytest.mark.parametrize("scale_fmt", ["ue8m0", "mxfp8"]) +def test_indexer_quant_cache_roundtrip( + num_tokens: int, block_size: int, scale_fmt: str +): """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache → manual dequant → compare against original.""" HEAD_DIM = 128 @@ -1442,7 +1445,9 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): ) slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) - ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0") + ops.indexer_k_quant_and_cache( + k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, scale_fmt + ) # ── Gather ────────────────────────────────────────────────────────── max_blocks_per_seq = num_blocks @@ -1465,7 +1470,12 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): # ── Manual dequant ────────────────────────────────────────────────── k_fp8 = dst_k.view(torch.float8_e4m3fn).float() # [num_tokens, 128] - scale = dst_scale.view(torch.float32) # [num_tokens, 1] + if scale_fmt == "mxfp8": + # One UE8M0 byte per 32 elements, all equal to the token's scale. + assert (dst_scale == dst_scale[:, :1]).all() + scale = (dst_scale[:, :1].to(torch.int32) << 23).view(torch.float32) + else: + scale = dst_scale.view(torch.float32) # [num_tokens, 1] k_recovered = k_fp8 * scale # [num_tokens, 128] # ── Compare ───────────────────────────────────────────────────────── @@ -1749,8 +1759,13 @@ def _reference_kv_compress_norm_rope( ), ], ) -def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool): +@pytest.mark.parametrize("mx_scale", [False, True]) +def test_fused_kv_insert_indexer( + num_tokens: int, kv_block_size: int, use_fp4: bool, mx_scale: bool +): """Fused K compress+norm+rope+quant+insert for the indexer KV cache.""" + if use_fp4 and mx_scale: + pytest.skip("MX scales only change the FP8 cache") HEAD_DIM = 128 ROPE_DIM = 64 BLOCK_SIZE = 16 @@ -1833,6 +1848,7 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b TOKEN_STRIDE=TOKEN_STRIDE, SCALE_DIM=SCALE_DIM, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=mx_scale, num_warps=1, ) @@ -1876,12 +1892,18 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b ), f"token {i}: FP8 bytes differ" scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM - actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view( - torch.float32 - ) - assert torch.equal(actual_scale, scale[i : i + 1]), ( - f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" - ) + actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM] + if mx_scale: + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale[i : i + 1].view(torch.int32) >> 23).to(torch.uint8) + assert torch.equal(actual_scale, ue8m0.repeat(SCALE_DIM)), ( + f"token {i}: mx scale {actual_scale.tolist()} != {ue8m0.item()}" + ) + else: + actual_scale = actual_scale.view(torch.float32) + assert torch.equal(actual_scale, scale[i : i + 1]), ( + f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" + ) @pytest.mark.parametrize("compress_ratio", [4, 128]) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 63bbe0c95ff4..ec61b94abd5e 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -32,6 +32,7 @@ from vllm.models.deepseek_v32.common import kernels as K from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales FP8 = torch.float8_e4m3fn FP8_MAX = 448.0 @@ -116,6 +117,16 @@ def rope( return out +def indexer_scales(scale_bytes: torch.Tensor) -> torch.Tensor: + """Per-token fp32 scales from the indexer cache's 4 scale bytes per token.""" + if not mqa_logits_use_mx_scales(): + return scale_bytes.view(torch.float32) + # MX scales: the power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale_bytes.view(-1, 4) + assert (ue8m0 == ue8m0[:, :1]).all(), "MX scale bytes differ within a token" + return (ue8m0[:, 0].to(torch.int32) << 23).view(torch.float32) + + def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" amax = vals.float().abs().amax(dim=-1, keepdim=True) @@ -243,7 +254,7 @@ def test_fused_norm_rope( q_ref, s_ref = ue8m0_quant(ik_ref) flat = idx_cache[0].reshape(-1) vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim) - scales = flat[bs * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(flat[bs * cfg.index_head_dim :]) assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) @@ -318,7 +329,7 @@ def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig): .view(FP8) .view(block_size, cfg.index_head_dim) ) - scales = packed[block_size * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(packed[block_size * cfg.index_head_dim :]) assert_fp8(values, q_ref, "packed indexer-K fp8") torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) assert (backing[block_size * idx_row : packed_block_stride] == 0).all() diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 0ce3a918513f..c6dfa11b0f13 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -26,6 +26,7 @@ ) from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.import_utils import has_cutedsl HEAD_DIM = 128 @@ -347,6 +348,10 @@ def _indexer_k_reference( exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448 scale = exponent.log2().ceil().exp2() values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale.view(torch.int32) >> 23).to(torch.uint8) + return values.view(torch.uint8), ue8m0.repeat(1, 4) return values.view(torch.uint8), scale.view(torch.uint8) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 4accf527f3d1..5dd9eff5f10a 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -29,6 +29,7 @@ import torch from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.kpool_compress import ( @@ -206,7 +207,12 @@ def _torch_reference( k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM) s_off = page_base + k_region + loc_tok * 4 kv_flat[k_off] = quantized.view(torch.uint8) - kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8) + scale_bytes = scale.detach().reshape(1).view(torch.uint8) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale.reshape(1).view(torch.int32) >> 23 + scale_bytes = ue8m0.to(torch.uint8).repeat(4) + kv_flat[s_off : s_off + 4] = scale_bytes # stash -- gated on the TOKEN-granular tail slot, not on pos_valid. # pos_valid keys off the pool-granular cache_loc, which is -1 for diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index ba4ff895cbf6..7ca4cff4040d 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -771,6 +771,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): candidates = _make_candidates(lens, num_candidates, cbk).cuda() # Reference: dense bf16 logits masked to the candidates, then torch.topk. + # The logits are compressed: row r starts at K position ks[r]. + max_len = int(lens.max()) dense = deep_gemm.fp8_fp4_mqa_logits( ( q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), @@ -780,19 +782,13 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): weights, ks, ke, - clean_logits=False, - logits_dtype=torch.bfloat16, + max_len, ) - cols = torch.arange(total, device="cuda") - in_range = (cols - ks.unsqueeze(1) >= 0) & ( - cols - ks.unsqueeze(1) < lens.unsqueeze(1) - ) - masked = dense.masked_fill( - ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), - float("-inf"), - ).float() - ref_abs = _reference_topk(masked, topk) # packed-workspace positions - ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + in_range = torch.arange(max_len, device="cuda") < lens.unsqueeze(1) + # A row's last candidate block may extend past max_len. + in_cand = _candidate_token_mask(candidates, cbk, max_len + cbk)[:, :max_len] + masked = dense.masked_fill(~(in_range & in_cand), float("-inf")).float() + ref = _reference_topk(masked, topk) out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") scratch = _sparse_scratch(rows, num_candidates, topk) @@ -827,11 +823,11 @@ def run_sparse(kernel_metadata=None): graph.replay() else: kernel_metadata = run_sparse() - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) out.fill_(-1) run_sparse(kernel_metadata) - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): @@ -902,21 +898,22 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() def reference_topk(): - context_lens_2d = row_lens.view(batch, next_n) + # SM100 paged MQA logits take varlen rows: one Q token per row. + context_lens_2d = row_lens.view(rows, 1) + indices = torch.arange(batch, dtype=torch.int32, device="cuda") + indices = indices.repeat_interleave(next_n) sched = deep_gemm.get_paged_mqa_logits_metadata( - context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + context_lens_2d, page_kv, deep_gemm.get_num_sms(), indices ) logits = deep_gemm.fp8_fp4_paged_mqa_logits( - (q, q_scale), + (q.view(rows, 1, num_heads, -1), q_scale.view(rows, 1, num_heads)), kv_cache, weights, context_lens_2d, - block_table, + block_table.repeat_interleave(next_n, dim=0), sched, int(req_lens.max()), - False, - torch.bfloat16, - None, + indices, ) lens_flat = row_lens.flatten() cols = torch.arange(logits.shape[1], device="cuda") diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 464fce66f212..62c8bc32bff5 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -14,6 +14,7 @@ from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, ) +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla import indexer as indexer_module from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -113,10 +114,17 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): num_reqs, 4 ) num_sms = torch.cuda.get_device_properties(device).multi_processor_count + # SM100's paged MQA logits only take varlen rows. + indices = ( + torch.arange(num_reqs, dtype=torch.int32, device=device) + if mqa_logits_use_mx_scales() + else None + ) schedule_metadata = indexer_module.get_paged_mqa_logits_metadata( torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1), 64, num_sms, + indices=indices, ).clone() builder = object.__new__(DeepseekV32IndexerMetadataBuilder) @@ -139,6 +147,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device), requires_padding=False, schedule_metadata=schedule_metadata, + indices=indices, ), block_table=block_table, ) @@ -157,7 +166,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): torch.testing.assert_close( schedule_metadata, indexer_module.get_paged_mqa_logits_metadata( - expected_decode_seq_lens.unsqueeze(1), 64, num_sms + expected_decode_seq_lens.unsqueeze(1), 64, num_sms, indices=indices ), ) assert metadata.decode is not None diff --git a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py index 0283ed99ee3a..a3dbec86dca5 100644 --- a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py +++ b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py @@ -36,6 +36,7 @@ def _chunk(block_table, token_start, token_end, row_ends): cu_seq_lens=torch.zeros(1, dtype=torch.int32), token_to_seq=torch.zeros(1, dtype=torch.int32), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=token_start, token_end=token_end, num_reqs=block_table.shape[0], diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 36824318c051..dec8b0105ac8 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -35,8 +35,9 @@ env.pop("DG_SKIP_CUDA_BUILD", None) print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True) +# CMake decides when to rebuild; setuptools would only check python_api.cpp. subprocess.check_call( - [sys.executable, "setup.py", "build_ext", "--inplace"], + [sys.executable, "setup.py", "build_ext", "--inplace", "--force"], cwd=src, env=env, ) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index fcc35ead58c4..c8da1351bb54 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -4,10 +4,12 @@ """Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout. Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under -``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100 -deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. +``vllm.third_party.deep_gemm``, and that the package imports: vLLM treats an +import failure as DeepGEMM being absent. Run after vLLM is installed, e.g. the +H100 deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. """ +import importlib import importlib.util import sys from pathlib import Path @@ -38,6 +40,11 @@ def main() -> int: if missing: print(f"missing: {missing}", file=sys.stderr) return 1 + try: + importlib.import_module("vllm.third_party.deep_gemm") + except Exception as e: + print(f"vllm.third_party.deep_gemm failed to import: {e!r}", file=sys.stderr) + return 1 return 0 diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 2fb2351a0ca5..b617e6d5679f 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -7,10 +7,10 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" -# NOTE: This targets the vLLM fork's dev branch, which carries the sm120 -# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE -# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. -DEEPGEMM_GIT_REF="1e1842a833699298f7afc02eefb2ed168fff6938" +# NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release +# plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU +# Mega MoE activations, and NVFP4 Mega MoE. +DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 935e4c0d2d44..3af3a161ac8b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -42,6 +42,8 @@ fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, ) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( @@ -286,6 +288,42 @@ def fused_indexer_q_rope_quant( return q_fp8, weights_out +@triton.jit +def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr): + row = tl.program_id(0) + end = tl.load(row_ends_ptr + row) + src_ptr += row.to(tl.int64) * stride + dst_ptr += row.to(tl.int64) * stride + for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < end + vals = tl.load(src_ptr + offs, mask=mask) + tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask) + + +def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor: + """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels. + + Only columns below ``row_ends`` (int32, one per row) are converted, the + span the consumers read; DeepGEMM's aligned row stride is kept. + """ + if logits.dtype == torch.float32: + return logits + out = torch.empty_strided( + logits.shape, logits.stride(), dtype=torch.float32, device=logits.device + ) + num_rows, width = logits.shape + if num_rows == 0 or width == 0: + return out + assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32 + block = 1024 + splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows)) + _upcast_logits_kernel[(num_rows, splits)]( + logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block + ) + return out + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, @@ -537,7 +575,7 @@ def sparse_attn_indexer( kv_cache, slot_mapping_for_cache, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) # The indexer and main MLA may classify the same short extend differently # because they use independent decode thresholds. Only the main MLA route @@ -649,6 +687,9 @@ def sparse_attn_indexer( q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] + # DeepGEMM returns compressed logits: row i starts at K position + # cu_seqlen_ks[i]. None means every row starts at column zero. + row_starts: torch.Tensor | None = None if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) @@ -662,7 +703,9 @@ def sparse_attn_indexer( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) if current_platform.is_xpu(): if q_scale_slice is not None: raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") @@ -674,14 +717,23 @@ def sparse_attn_indexer( cu_seqlen_ks, cu_seqlen_ke, ) + row_starts = cu_seqlen_ks + row_ends = cu_seqlen_ke else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) num_rows = logits.shape[0] if candidate_blocks is not None: @@ -694,8 +746,8 @@ def sparse_attn_indexer( if candidate_write: _select_candidate_blocks( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates.shape[1], candidate_block_size, chunk_candidates, @@ -703,15 +755,15 @@ def sparse_attn_indexer( else: _apply_candidate_mask( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates, candidate_block_size, ) ops.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends) if row_starts is None else row_starts, + row_ends, topk_indices, num_rows, logits.stride(0), @@ -729,7 +781,7 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=cu_seqlen_ks, + row_starts=row_starts, ) if row_shard is not None: @@ -822,16 +874,22 @@ def sparse_attn_indexer( max_model_len, ) else: - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_model_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] if candidate_blocks is not None: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 255b1f5aaf2f..f876120a6ef2 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -6,7 +6,9 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the # has_indexer=False path so the indexer args don't allocate every call. @@ -87,6 +89,7 @@ def _fp8_quant_and_cache_write( SHUFFLE: tl.constexpr, BLOCK_TILE: tl.constexpr, HEAD_TILE: tl.constexpr, + MX_SCALE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr, ): @@ -111,6 +114,8 @@ def _fp8_quant_and_cache_write( tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask) scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4 + if MX_SCALE: + scale = pack_ue8m0_x4(scale) tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale) @@ -167,6 +172,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE: tl.constexpr, INDEXER_CACHE_BLOCK_TILE: tl.constexpr, INDEXER_CACHE_HEAD_TILE: tl.constexpr, + INDEXER_CACHE_MX_SCALE: tl.constexpr, # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr) mla_cache_ptr, mla_cache_block_stride, @@ -501,6 +507,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE, INDEXER_CACHE_BLOCK_TILE, INDEXER_CACHE_HEAD_TILE, + INDEXER_CACHE_MX_SCALE, FP8_MAX, USE_FNUZ, ) @@ -571,7 +578,10 @@ def fused_norm_rope( assert slot_mapping is not None if indexer_slot_mapping is None: indexer_slot_mapping = slot_mapping - idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) + idx_cache_mx_scale = mqa_logits_use_mx_scales() + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view( + torch.int32 if idx_cache_mx_scale else torch.float32 + ) idx_cache_block_size = indexer_k_cache.shape[1] idx_cache_stride = indexer_k_cache.shape[2] # The caller's cache reports whether its reader expects the shuffled @@ -591,6 +601,7 @@ def fused_norm_rope( ) else: idx_cache_scale_view = None + idx_cache_mx_scale = False idx_cache_block_size = 1 idx_cache_stride = 0 idx_cache_shuffle = False @@ -701,6 +712,7 @@ def fused_norm_rope( idx_cache_shuffle, _INDEXER_CACHE_BLOCK_TILE, idx_cache_head_tile, + idx_cache_mx_scale, # MLA KV cache (uses same slot_mapping) mla_kv_cache, mla_block_stride, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 50ca86558e05..703169f4829f 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -33,7 +33,9 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.math_utils import round_up +from vllm.v1.attention.ops.common import pack_ue8m0_x4 if current_platform.is_rocm(): from vllm.platforms.rocm import _ON_GFX950 @@ -742,6 +744,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( TOKEN_STRIDE: tl.constexpr, # 128 for indexer SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # store the scale as 4 packed UE8M0 bytes ): """Fused compress → RMSNorm → RoPE → FP8 quant → store. @@ -878,7 +881,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( # Single float32 scale scale_val = tl.exp2(exponent) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) + if MX_SCALE: + tl.store(scale_ptr.to(tl.pointer_type(tl.int32)), pack_ue8m0_x4(scale_val)) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) # ============================================================================= @@ -927,6 +933,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # unused for MXFP4 (kept for signature parity) ): """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. @@ -1102,6 +1109,7 @@ class CompileKey: block_size: int kv_cache_block_size: int kv_block_stride: int + mx_scale: bool # Canonical kernel for the launcher's arg_names mapping; the launched # variant (fp8/mxfp4, same arg_names) is picked per call in __call__. @@ -1181,6 +1189,7 @@ def dispatch( # type: ignore[override] block_size=state_block_size, kv_cache_block_size=kv_cache_block_size, kv_block_stride=kv_block_stride, + mx_scale=not use_fp4_cache and mqa_logits_use_mx_scales(), ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -1301,6 +1310,7 @@ def __call__( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), num_warps=1, **pdl_kwargs, ) diff --git a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py index ba8ddbd95a9a..fe203e917861 100644 --- a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py +++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py @@ -17,6 +17,8 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2 @@ -51,7 +53,8 @@ def indexer_k_norm_rope_store( kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip). compress_ratio: group size; keys are emitted at group boundaries. use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else - per-token FP8 with a single fp32 scale. + per-token FP8 with a single fp32 scale (four packed ue8m0 bytes + under `mqa_logits_use_mx_scales`). """ num_tokens = kv_slot_mapping.numel() @@ -111,6 +114,7 @@ def indexer_k_norm_rope_store( KV_BLOCK_STRIDE=k_cache.stride(0), FP8_MAX=448.0, USE_FP4=use_fp4_cache, + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), SHUFFLE=shuffle, BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE, HEAD_TILE_SIZE=_HEAD_TILE_SIZE, @@ -139,6 +143,7 @@ def _indexer_k_norm_rope_quant_store_kernel( KV_BLOCK_STRIDE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FP4: tl.constexpr, + MX_SCALE: tl.constexpr, SHUFFLE: tl.constexpr, BLOCK_TILE_SIZE: tl.constexpr, HEAD_TILE_SIZE: tl.constexpr, @@ -252,4 +257,10 @@ def _indexer_k_norm_rope_quant_store_kernel( tl.store(val_ptr + tiled, x_uint8) else: tl.store(val_ptr + block, x_uint8) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) + if MX_SCALE: + tl.store( + scale_ptr.to(tl.pointer_type(tl.int32)), + pack_ue8m0_x4(tl.exp2(exponent)), + ) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index eb289c65cf68..310d60dfd41e 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -15,6 +15,8 @@ import torch from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # The GLM-5.3-Flash indexer head dimension is fixed at 128. INDEX_HEAD_DIM = 128 @@ -157,6 +159,7 @@ def _kpool_softmax_rotate_write_cache_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, HAS_WRITE_MASK: tl.constexpr, RETURN_COMPRESSED: tl.constexpr, WRITE_CACHE: tl.constexpr, @@ -245,7 +248,14 @@ def _kpool_softmax_rotate_write_cache_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask) - tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + mask=do_write, + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) if RETURN_COMPRESSED: tl.store( @@ -324,6 +334,7 @@ def kpool_compress_and_write_cache( buf_fp8 = buf.view(torch.float8_e4m3fn) buf_fp32 = buf.view(torch.float32) + mx_scale = write_cache and mqa_logits_use_mx_scales() # bytes per page (last dim of kv_cache) viewed as uint8 buf_numel_per_page = buf.stride(0) s_offset_nbytes_in_page = page_size * head_dim @@ -361,7 +372,9 @@ def kpool_compress_and_write_cache( POOL_SIZE=slot_k.shape[1], HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mx_scale, + MX_SCALE=mx_scale, HAS_WRITE_MASK=has_write_mask, RETURN_COMPRESSED=return_compressed, WRITE_CACHE=write_cache, @@ -479,6 +492,7 @@ def _kpool_decode_update_batched_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, BLOCK_D: tl.constexpr, ): """One program per request; iterates its NEXT_N verify tokens in order. @@ -604,7 +618,13 @@ def _kpool_decode_update_batched_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask) - tl.store(buf_fp32_ptr + out_s_offset, scale) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale) # Stash the current token AFTER any completion read so the completion # uses prior stashes (and the current token's own key/score via @@ -718,7 +738,9 @@ def kpool_decode_update_and_maybe_write_cache_batched( KPOOL_HEAD=tail_kv_cache.stride(1), HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mqa_logits_use_mx_scales(), + MX_SCALE=mqa_logits_use_mx_scales(), BLOCK_D=triton.next_power_of_2(head_dim), ) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 154fdcda900e..1a3fbedb8387 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -16,6 +16,7 @@ PrefillRowShard, _merge_dcp_topk_global, _reserve_row_shard_exchange_transient, + as_fp32_logits, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -29,7 +30,13 @@ ) from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform -from vllm.utils.deep_gemm import has_deep_gemm +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, +) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( LayerNameType, @@ -253,7 +260,7 @@ def sparse_attn_indexer_kpool( kv_cache, slot_mapping, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) topk_indices_buffer[: hidden_states.shape[0]] = -1 @@ -343,20 +350,29 @@ def sparse_attn_indexer_kpool( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) num_rows = q_slice.shape[0] if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + # Compressed: row i starts at K position cu_seqlen_ks[i]. + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) # kpool: logits are pool-granular (compress_ratio == index_kpool), @@ -376,8 +392,8 @@ def sparse_attn_indexer_kpool( else: torch.ops._C.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends), + row_ends, topk_dst, num_rows, logits.stride(0), @@ -392,7 +408,7 @@ def sparse_attn_indexer_kpool( dcp_rank, dcp_world_size, pool_cp_interleave, - row_starts=cu_seqlen_ks, + row_starts=None, ) if index_kpool > 1: @@ -594,18 +610,22 @@ def sparse_attn_indexer_kpool( if use_fp4_cache else padded_q_quant_decode_tokens ) - from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits - - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - padded_weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_pool_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, padded_weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_pool_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools, diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 04c4508f0ba9..3dc9be5d9385 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -591,37 +591,79 @@ def transform_sf_into_required_layout(*args, **kwargs): ) +@functools.cache +def mqa_logits_use_mx_scales() -> bool: + """Whether the MQA logits kernels take only MX-scaled inputs. + + On SM100 they need MXFP8/MXFP4 Q and K with packed UE8M0 scales and BF16 + weights, and return BF16 logits. FP8 indexer K caches then store each + token's power-of-two scale as four equal UE8M0 bytes instead of an fp32. + """ + return current_platform.is_cuda() and current_platform.is_device_capability_family( + 100 + ) + + +# Four packed UE8M0 bytes encoding 1.0. +_UNIT_UE8M0_X4 = 0x7F7F7F7F +# Grown buffers are kept alive: CUDA graphs captured earlier still read them. +_unit_mx_scales: dict[torch.device, list[torch.Tensor]] = {} + + +def mqa_logits_mx_inputs( + q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor +) -> tuple[torch.Tensor | None, torch.Tensor]: + """Adapt indexer Q scales and weights to the SM100 MQA logits kernels. + + FP8 Q folds its per-(token, head) power-of-two scale into ``weights``, so + it is passed with unit MX scales, which is exact since relu commutes with + a positive scale. Returns ``(q_scale, weights)`` unchanged off SM100. + """ + if not mqa_logits_use_mx_scales(): + return q_scale, weights + if q_scale is None: + numel = q.numel() // q.shape[-1] + bufs = _unit_mx_scales.setdefault(q.device, []) + if not bufs or bufs[-1].numel() < numel: + size = max(numel, 2 * bufs[-1].numel() if bufs else 0) + bufs.append( + torch.full((size,), _UNIT_UE8M0_X4, dtype=torch.int32, device=q.device) + ) + q_scale = bufs[-1][:numel].view(q.shape[:-1]) + return q_scale, weights.to(torch.bfloat16) + + def fp8_fp4_mqa_logits( q: tuple[torch.Tensor, torch.Tensor | None], kv: tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, - clean_logits: bool, + max_seqlen_k: int, ) -> torch.Tensor: """Compute MQA logits for a single sequence without KV paging. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q - (per-token scale is folded into ``weights``) and a packed block-scale - tensor for MXFP4 Q. + ``q = (values, scales_or_None)``. SM100 takes MXFP8 or MXFP4 with packed + UE8M0 block scales; SM90/SM120 also take FP8 Q with ``q_scale=None`` + (per-token scale folded into ``weights``). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D] - float8_e4m3fn and q_scale is None (per-token scale is folded - into ``weights``). FP4 path: q_values is packed uint8 and - q_scale is the companion block-scale tensor. - kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D] - float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8. - weights: weights of shape [M, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape [M, H, D] (packed + to D/2 bytes for MXFP4) and q_scale int32 [M, H] or None. + kv: Tuple ``(k_values, k_scales)`` of shape [N, D] and [N]; k_scales + is int32 packed UE8M0 with MX scales, fp32 otherwise. + weights: Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere. cu_seqlen_ks: Start indices (inclusive) for valid K per query position, shape [M], dtype int32. cu_seqlen_ke: End indices (exclusive) for valid K per query position, shape [M], dtype int32. - clean_logits: Whether to clean the unfilled logits into `-inf`. + max_seqlen_k: Upper bound on ``cu_seqlen_ke - cu_seqlen_ks``. Returns: - Logits tensor of shape [M, N], dtype `torch.float32`. + Compressed logits of shape [M, max_seqlen_k]: row ``i`` holds K + positions ``[cu_seqlen_ks[i], cu_seqlen_ke[i])`` starting at column + zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -633,7 +675,7 @@ def fp8_fp4_mqa_logits( weights, cu_seqlen_ks, cu_seqlen_ke, - clean_logits=clean_logits, + max_seqlen_k, ) @@ -699,24 +741,25 @@ def fp8_fp4_paged_mqa_logits( block_tables: torch.Tensor, schedule_metadata: torch.Tensor, max_model_len: int, - clean_logits: bool, indices: torch.Tensor | None = None, ) -> torch.Tensor: """Compute MQA logits using a paged KV-cache. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8 - path and ``(q_values, q_scale)`` for MXFP4. + ``q = (values, scales_or_None)``; see `fp8_fp4_mqa_logits` for the + per-architecture formats. SM100 requires the varlen form (``indices`` + with next_n == 1). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is - [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path: - q_values is packed uint8 and q_scale is the companion - block-scale tensor. - kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1, - D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos) - storing the float dequant scale. - weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape + [B, next_n, H, D] (packed to D/2 bytes for MXFP4) and q_scale + int32 [B, next_n, H] or None. + kv_cache: Paged KV-cache of shape [num_blocks, block_size, 1, D+4] + (D/2+4 for MXFP4), dtype `torch.uint8`; the last 4 bytes per + (block, pos) hold the int32 packed UE8M0 scales with MX scales, + or the fp32 dequant scale otherwise. + weights: Tensor of shape [B * next_n, H], bf16 on SM100 and fp32 + elsewhere. context_lens: Tensor of shape [B], dtype int32; effective context length for each batch element. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical @@ -724,12 +767,11 @@ def fp8_fp4_paged_mqa_logits( schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. - clean_logits: Whether to clean the unfilled logits into `-inf`. indices: Optional request index for each varlen row. Returns: - Logits tensor of shape [B * next_n, max_model_len], dtype - `torch.float32`. + Logits tensor of shape [B * next_n, max_model_len]; columns past each + row's context length are not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -751,7 +793,6 @@ def fp8_fp4_paged_mqa_logits( block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, **kwargs, ) @@ -1052,6 +1093,8 @@ def should_use_deepgemm_for_fp8_linear( "fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits", "get_paged_mqa_logits_metadata", + "mqa_logits_mx_inputs", + "mqa_logits_use_mx_scales", "native_next_n_supported", "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2d797ae66283..0382a81fe044 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -467,6 +467,9 @@ class DeepseekV32IndexerPrefillChunkMetadata: cu_seq_lens: torch.Tensor token_to_seq: torch.Tensor total_seq_lens: int + # Host upper bound on cu_seqlen_ke - cu_seqlen_ks, sizing the compressed + # DeepGEMM prefill logits. + max_seqlen_k: int token_start: int token_end: int num_reqs: int @@ -1957,6 +1960,7 @@ def build_prefill_chunk_metadata( cu_seq_lens=cu_seq_lens, token_to_seq=token_to_seq, total_seq_lens=total_seq_lens, + max_seqlen_k=int(compressed_seq_lens_cpu[start_idx:end_idx].max()), block_table=block_table[start_idx:end_idx], token_start=token_start, token_end=token_end, diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 5c1677cac479..6936be9fa169 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -154,6 +154,13 @@ def __call__( ) +@triton.jit +def pack_ue8m0_x4(scale): + """Encode a power-of-two fp32 scale as four equal packed UE8M0 bytes, the + MX scales of a 128-wide row (see `mqa_logits_use_mx_scales`).""" + return ((scale.to(tl.int32, bitcast=True) >> 23) & 0xFF) * 0x01010101 + + def pack_seq_triton( x: torch.Tensor, lengths: torch.Tensor, From e2ad0ce5f5a944e10ed2ea664e4a24bf85791346 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 13 Oct 2026 14:25:16 +0800 Subject: [PATCH 472/952] [Bugfix][LoRA] Slice expert LoRA by the unpadded per-rank size under TP (#61235) Signed-off-by: aoshen524 Co-authored-by: aoshen524 Co-authored-by: Claude Opus 5.5 Co-authored-by: linitra24 --- vllm/lora/layers/fused_moe.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c58c18e71ff7..19a1128ce145 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from transformers import PreTrainedConfig from vllm import envs @@ -110,6 +111,15 @@ def use_ep(self) -> bool: def intermediate_size_per_partition(self) -> int: return self.moe_config.intermediate_size_per_partition + @property + def _tp_shard_size(self) -> int: + """Per-rank stride of an adapter along the expert intermediate dim: the + base weights are sharded by the unpadded size even when a backend pads.""" + unpadded = self.moe_config.intermediate_size_per_partition_unpadded + if unpadded is None: + return self.intermediate_size_per_partition + return unpadded + def _init_lora_stream_context(self) -> None: self._lora_stream: torch.cuda.Stream | None = None self._events: tuple[torch.cuda.Event, ...] | None = None @@ -311,7 +321,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -322,7 +332,7 @@ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: if self.tp_size == 1: return w2_lora_a # w2_lora_a shape (num_experts,rank,input_size) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -536,7 +546,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size # HACK: Currently, only GPT-OSS is in interleaved order @@ -557,6 +567,11 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): w3_lora_b = w13_lora_b[:, slice_size:, :] sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :] sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :] + # w3 starts at the allocated (possibly padded) per-rank size, as + # in the base w13 weight. + pad = self.intermediate_size_per_partition - sliced_w1_lora_b.shape[1] + if pad > 0: + sliced_w1_lora_b = F.pad(sliced_w1_lora_b, (0, 0, 0, pad)) return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1) From 4e4eab73ed2b4e00931a4636e926b8df1cb1fe50 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 12 Oct 2026 23:57:32 -0700 Subject: [PATCH 473/952] [CI] Shard Model Executor and L4 Pipeline + Context Parallelism automatically (#61202) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 1 + .buildkite/test_areas/model_executor.yaml | 1 + 2 files changed, 2 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b30cbac0ced4..2f929b53b840 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -641,6 +641,7 @@ steps: - label: ":nvidia: (L4) Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus + automatic_shard: true timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" device: l4 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index a23359b768a4..c9f467941af9 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Model Executor" device: h200_35gb key: model-executor + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - vllm/engine/arg_utils.py From bc3ddb198a0d5eb864bf484c4719864b671cc93b Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:30:19 +0800 Subject: [PATCH 474/952] [Model][Qwen4Exp] FP8 QSA main KV cache on Ampere (sm_80, sm_86) (#60602) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Signed-off-by: Thien Tran Co-authored-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 69 +++++++++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 21 +++++--- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 401048e02c55..c7fe6014497e 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -22,6 +22,28 @@ def _is_sm120() -> bool: return current_platform.get_device_capability() == (12, 0) +@triton.jit +def _maybe_dequant(tile): + if tile.dtype == tl.uint8: + # e4m3 bytes below sm_89 (no Triton fp8e4nv pointer type): sign to bit + # 15, exp+mantissa to bits 13..7 gives the fp16 of x * 2^-8. Exact for + # every finite code; the two NaN codes decode to +-480 after the host + # scale instead of NaN. + code = tile.to(tl.uint16) + bits = ((code & 0x80) << 8) | ((code & 0x7F) << 7) + return bits.to(tl.float16, bitcast=True) + elif tile.dtype == tl.float8e4nv: + return tile.to(tl.float16) + else: + return tile + + +@lru_cache(maxsize=1) +def _is_sm80() -> bool: + """True on sm_80 (A100/A30): selects the wide FP8 byte-decode tiles.""" + return current_platform.get_device_capability() == (8, 0) + + @lru_cache(maxsize=1) def _is_sm90() -> bool: """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" @@ -69,7 +91,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, - IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -141,6 +162,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[None, :], other=0.0, ) + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK + # measured slower here and less accurate). + keys = _maybe_dequant(keys).to(query.dtype) values = tl.load( v_cache_ptr + safe_page[:, None] * stride_v_block @@ -150,10 +174,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) - if IS_FP8: - # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK - # measured slower here and less accurate). - keys = keys.to(query.dtype) + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online + # softmax) so fp16 has the range, its wider mantissa is more + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. + values = _maybe_dequant(values) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 # caches the K dequant scale is already folded into softmax_scale on the @@ -165,11 +189,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( probabilities = tl.where( valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 ) - if IS_FP8: - # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online - # softmax) so fp16 has the range, its wider mantissa is more - # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. - values = values.to(tl.float16) accumulator = tl.dot( probabilities.to(values.dtype), values, @@ -585,6 +604,10 @@ def _select_config( BLOCK_N, target_splits, num_warps = 64, 4, 2 else: BLOCK_N, target_splits, num_warps = 64, 1, 2 + if is_fp8 and _is_sm80(): + # Wider tiles measured faster for the byte-decode path on A100; + # unmeasured on sm_86, which keeps the table's tiles. + BLOCK_N, num_warps = 128, 4 num_tiles = triton.cdiv(num_columns, BLOCK_N) # Never more splits than tiles, never empty. num_splits = min(target_splits, num_tiles) @@ -607,7 +630,9 @@ def qsa_sparse_paged_attention( ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. - With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales + An fp8 cache is a native e4m3 view on sm_89+ and the raw uint8 storage + below sm_89 (no Triton fp8e4nv pointer type there); the kernel dequantizes + either on load. k_scale/v_scale are the layer's per-tensor dequant scales as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is pre-multiplied into the softmax scale and v_scale becomes the kernel's output scale, so the kernel needs no device scale buffers. @@ -636,15 +661,21 @@ def qsa_sparse_paged_attention( assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 assert q.dtype == torch.bfloat16 assert k_cache.dtype == v_cache.dtype - is_fp8 = k_cache.dtype == torch.float8_e4m3fn + is_fp8 = k_cache.dtype != torch.bfloat16 if is_fp8: assert k_scale is not None and v_scale is not None # Host pre-multiply: fold the K dequant scale into the attention scale # and pass V's dequant scale as the kernel's output scale. softmax_scale = (head_dim**-0.5) * float(k_scale) output_scale = float(v_scale) + if k_cache.dtype == torch.uint8: + # _maybe_dequant's byte decode yields x * 2^-8; fold 2^8 into the + # scales so the decode needs no multiply. + softmax_scale *= 256.0 + output_scale *= 256.0 + else: + assert k_cache.dtype == torch.float8_e4m3fn else: - assert k_cache.dtype == torch.bfloat16 softmax_scale = head_dim**-0.5 output_scale = 1.0 assert logical_indices.dtype == block_table.dtype == torch.int32 @@ -731,7 +762,6 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=partial_warps, num_stages=2, ) @@ -768,9 +798,15 @@ def warmup_qsa_sparse_paged_attention( """Compile every production-reachable split-K/merge specialization.""" head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) - # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. + # An fp8 cache is allocated as uint8 storage; warm the dtype the impl hands + # over: an e4m3 view on sm_89+, the raw bytes below. is_fp8 = kv_cache.dtype == torch.uint8 - cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype + fp8_dtype = ( + torch.float8_e4m3fn + if current_platform.has_device_capability(89) + else torch.uint8 + ) + cache_dtype = fp8_dtype if is_fp8 else key_cache.dtype num_kv_heads = key_cache.shape[2] group_size = num_query_heads // num_kv_heads block_m = triton.next_power_of_2(group_size) @@ -874,7 +910,6 @@ def warmup_qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=warps, num_stages=2, grid=(num_rows, num_kv_heads, num_splits), diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 9f0934d5a0c6..9fb36673a5e7 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -297,10 +297,12 @@ def forward_qsa( key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) k_scale = v_scale = None if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): - # The cache is allocated as uint8; reinterpret the e4m3 bytes - # (same itemsize, so shape and strides are preserved). - key_cache = key_cache.view(torch.float8_e4m3fn) - value_cache = value_cache.view(torch.float8_e4m3fn) + # The cache is allocated as uint8; on sm_89+ reinterpret the e4m3 + # bytes (same itemsize, so shape and strides are preserved). Below + # sm_89 the kernel byte-decodes the raw uint8 storage instead. + if current_platform.has_device_capability(89): + key_cache = key_cache.view(torch.float8_e4m3fn) + value_cache = value_cache.view(torch.float8_e4m3fn) # Host-side per-tensor dequant scales (Python floats), as used by # other host-scale backends; folded into the kernel's scales. k_scale = layer._k_scale_float @@ -308,6 +310,7 @@ def forward_qsa( if query.dtype != torch.bfloat16 or key_cache.dtype not in ( torch.bfloat16, torch.float8_e4m3fn, + torch.uint8, ): raise NotImplementedError( "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V" @@ -492,9 +495,15 @@ def __init__( ) # One launch does the indexer prepare, the main QK-norm/RoPE/gate and # the main K/V cache write (see QSAIndexer.forward); otherwise all of - # them take the separate kernels. + # them take the separate kernels. The fused writer receives an fp8 + # pointer, which Triton cannot lower below sm_89. self.use_fused_qsa_prepare = ( - self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer + self.use_fused_qk_norm_rope_gate + and self.indexer.use_fused_pre_indexer + and ( + self.kv_cache_torch_dtype != torch.uint8 + or current_platform.has_device_capability(89) + ) ) self.fuse_indexer_projection = vllm_config.lora_config is None if self.fuse_indexer_projection: From 2cc2e68d1c5652a5bc7ae10422697308e2193359 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 15:58:59 +0800 Subject: [PATCH 475/952] [Core][Mamba] Remove leftovers of mamba_cache_mode="all" (#60043) Signed-off-by: zjy0516 Signed-off-by: Jiangyun Zhu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/disaggregated.yaml | 2 +- docs/design/hybrid_kv_cache_manager.md | 2 +- docs/features/automatic_prefix_caching.md | 7 +- .../test_kimi_linear_context_parallel.py | 1 - .../configs/Qwen3.8-Flash-Next-FP8-rocm.yaml | 2 +- tests/kernels/mamba/test_mamba_mixer2.py | 10 +- .../test_mamba_eagle_resume_checkpoint.py | 4 +- tests/v1/core/test_contiguous_kv_packing.py | 9 +- tests/v1/core/test_prefix_caching.py | 90 +++------- .../v1/e2e/general/test_mamba_prefix_cache.py | 2 - tests/v1/e2e/test_hybrid_chunked_prefill.py | 2 - tests/v1/e2e/test_replayssm_decode.py | 4 - .../run_mamba_prefix_cache_test.sh | 2 - .../unit/test_offloading_connector.py | 1 - .../unit/test_remote_prefill_lifecycle.py | 3 + tests/v1/simple_kv_offload/test_scheduler.py | 74 ++++++-- vllm/config/cache.py | 28 ++- vllm/config/model.py | 4 - vllm/config/vllm.py | 10 +- vllm/engine/arg_utils.py | 18 +- .../layers/mamba/mamba_mixer2.py | 6 +- .../layers/mamba/ops/causal_conv1d.py | 164 +----------------- vllm/model_executor/models/config.py | 2 +- vllm/model_executor/models/falcon_h1.py | 2 - .../model_executor/models/granitemoehybrid.py | 2 - vllm/model_executor/models/interfaces.py | 25 +-- vllm/model_executor/models/jamba.py | 2 - vllm/model_executor/models/nemotron_h.py | 2 - vllm/model_executor/models/registry.py | 3 - vllm/model_executor/models/zamba2.py | 8 +- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/core/single_type_kv_cache_manager.py | 3 + 33 files changed, 161 insertions(+), 344 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d815ee9268b2..7f0a351524e1 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3523,7 +3523,7 @@ steps: - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' - export PREFILL_BLOCK_SIZE='2048' - export DECODE_BLOCK_SIZE='2048' - - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b9bd59583b82..4821fac0a20b 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -426,7 +426,7 @@ steps: # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. PREFILL_BLOCK_SIZE: "2048" DECODE_BLOCK_SIZE: "2048" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 9891c82f97c9..3984000c7d8f 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -204,7 +204,7 @@ The second question is the cache eviction policy. For now, we use one LRU queue ### Case 3: mamba models -The prefix caching support of the mamba model is work in progress. Once implemented, models with mamba layer + full attention layer can be supported via the full attention + X algorithm in case 2. +Models with mamba layers + full attention layers support prefix caching via the full attention + X algorithm in case 2, with the Mamba state cached at block boundaries. See [Automatic Prefix Caching](../features/automatic_prefix_caching.md). ## Implementation diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md index cfc8c19acb2b..bfd87e7a2c57 100644 --- a/docs/features/automatic_prefix_caching.md +++ b/docs/features/automatic_prefix_caching.md @@ -15,23 +15,22 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl ## Hybrid Mamba models -Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. +With prefix caching enabled, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold: -- `--mamba-cache-mode align` +- prefix caching enabled - EAGLE/MTP speculative decoding on the Mamba group - `--prefix-match-unit` smaller than the Mamba block size - the model does not use multi-module MTP ```bash vllm serve \ - --mamba-cache-mode align \ --prefix-match-unit 64 \ --enable-mamba-shared-prefix-checkpoint ``` -`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect. +`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. For hybrid Mamba models that is the block size itself, so no sub-block boundary exists and the flag has no effect. Choose a value that divides the block size of every prefix-cacheable KV cache group, and that is a multiple of the per-state compression ratio for models that use one, such as sparse MLA. vLLM validates both at startup and names the offending sizes in the error. Read the served block size from the startup log. 64 is a reasonable starting point. diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index 06348074bb14..8bddf6e609c0 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -180,7 +180,6 @@ def _run_k3_partial_prefix_reuse( block_size=block_size, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", ) as runner: assert ( runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml index db43869eb4ae..c35614cd5f8c 100644 --- a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml +++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml @@ -23,4 +23,4 @@ generation_config: extra_body: chat_template_kwargs: enable_thinking: true -server_args: "--tensor-parallel-size 4 --mamba-cache-mode align" +server_args: "--tensor-parallel-size 4" diff --git a/tests/kernels/mamba/test_mamba_mixer2.py b/tests/kernels/mamba/test_mamba_mixer2.py index 973e7885c680..db2274c3c933 100644 --- a/tests/kernels/mamba/test_mamba_mixer2.py +++ b/tests/kernels/mamba/test_mamba_mixer2.py @@ -96,15 +96,17 @@ def mixer2_gated_norm_tensor_parallel( gate_states = torch.randn(batch_size, seq_len, hidden_size) # create gated-norm with TP - mixer = Mixer2RMSNormGated( - full_hidden_size=hidden_size, - full_n_groups=n_groups, - ) + with ensure_current_vllm_config(): + mixer = Mixer2RMSNormGated( + full_hidden_size=hidden_size, + full_n_groups=n_groups, + ) mixer.weight.weight_loader(mixer.weight, weight) # load # create gated-norm without TP to compute reference # - utilize mock patching to disable TP when with ( + ensure_current_vllm_config(), unittest.mock.patch( "vllm.model_executor.layers.mamba.mamba_mixer2." "get_tensor_model_parallel_world_size", diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 324aef425b89..54ae43e8e232 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -38,9 +38,7 @@ def _manager( num_prefill_lookahead=0, ): init_none_hash(sha256) - config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] - ) + config = _make_hybrid_kv_cache_config(block_size, num_blocks, ["full", "mamba"]) if eagle_group is not None: groups = list(config.kv_cache_groups) groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 64cd55054015..f71fe160d83d 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -278,7 +278,14 @@ def test_scratch_group_survives_computed_block_truncation(self): neither divides the hit length nor bounds the (empty) block list.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True - groups = get_kv_cache_groups(config, _make_csa_linear_specs()) + config.cache_config.mamba_cache_mode = "align" + specs = { + name: replace(spec, mamba_cache_mode="align") + if isinstance(spec, MambaSpec) + else spec + for name, spec in _make_csa_linear_specs().items() + } + groups = get_kv_cache_groups(config, specs) kv_cache_config = get_kv_cache_config_from_groups( config, groups, available_memory=BYTES_PER_BLOCK * 64 ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 836286d127fa..52455c608305 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1617,6 +1617,7 @@ def make_kv_cache_config_hybrid_model( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) return KVCacheConfig( @@ -1652,6 +1653,7 @@ def make_kv_cache_config_three_types( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) elif third_spec_type == "sliding_window": third_spec = SlidingWindowSpec( @@ -1732,6 +1734,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): block_size=block_size, shapes=(1, 1), dtypes=(torch.float32,), + mamba_cache_mode="align", ), ), ], @@ -1757,10 +1760,13 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, len(req0.prompt_token_ids), 0, computed_blocks - ) - assert blocks is not None + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the shared prefix. + for num_new_tokens in (len(common_token_ids), 7): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + req0.num_computed_tokens += num_new_tokens + computed_blocks = None req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) @@ -2285,7 +2291,7 @@ def _make_hybrid_kv_cache_config( - "full": FullAttentionSpec - "sliding_window": SlidingWindowSpec with window=2*block_size - "sliding_window_large": SlidingWindowSpec with window=4*block_size - - "mamba": MambaSpec + - "mamba": MambaSpec in prefix-caching ("align") mode """ spec_map = { @@ -2313,11 +2319,6 @@ def _make_hybrid_kv_cache_config( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), - ), - "mamba_align": lambda: MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), mamba_cache_mode="align", ), } @@ -2427,14 +2428,16 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): assert not computed_blocks.blocks[0] # No cache hit initially assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, 55, len(computed_blocks.blocks[0]) * block_size, computed_blocks - ) - assert blocks is not None - # Should have blocks for all groups - assert len(blocks.get_block_ids()) == num_groups - - manager.new_step_starts() + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the common prefix. + for num_new_tokens in (len(common_token_ids), len(unique_token_ids)): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + # Should have blocks for all groups + assert len(blocks.get_block_ids()) == num_groups + req0.num_computed_tokens += num_new_tokens + computed_blocks = None + manager.new_step_starts() # Second request: should hit cached blocks for common prefix req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn) @@ -2549,7 +2552,7 @@ def test_prefill_hybrid_model_mamba_align(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = make_kv_cache_manager( kv_cache_config, @@ -2585,7 +2588,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): """ block_size = 16 manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, @@ -2660,7 +2663,7 @@ def test_hybrid_model_mamba_align_with_dynamic_draft_tokens(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = KVCacheManager( kv_cache_config, @@ -5896,49 +5899,6 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(): - """Manager-level check of the full wiring: a pinned shared-prefix boundary - (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style - detection) keeps its Mamba state block cached under - ``prefix_cache_retention_interval=0``, which otherwise retains only the - end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request.""" - block_size = 16 - - # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. - token_ids = [i for i in range(16) for _ in range(block_size)] - - def cached_mamba_blocks(shared_prefix_boundary): - # Fresh manager per scenario so cached blocks don't leak between runs. - manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"]), - max_model_len=8192, - enable_caching=True, - hash_block_size=block_size, - retention_interval=0, - ) - req = make_request("r", token_ids, block_size, sha256) - req.shared_prefix_boundary = shared_prefix_boundary - computed_blocks, num_computed, _ = manager.get_computed_blocks(req) - blocks = manager.allocate_slots( - req, len(token_ids), num_computed, computed_blocks - ) - assert blocks is not None - pool = manager.block_pool - return { - i - for i in range(16) - if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1]) - is not None - } - - # Without a pinned boundary, retention=0 keeps only the replay boundary (14). - assert cached_mamba_blocks(0) == {14} - # Pinning the shared prefix at token 96 (state block 5) retains it too, so a - # later request sharing that prefix can hit the Mamba state. - assert cached_mamba_blocks(96) == {5, 14} - - def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under @@ -5949,7 +5909,7 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): def last_req_hit(retention, pin): manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 50f34737085d..6c208e19fb5d 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -898,7 +898,6 @@ def _run_mamba_prefix_cache_mrv1( load_format="dummy", enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, @@ -1158,7 +1157,6 @@ def fake_sample( skip_tokenizer_init=True, enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index 66c8d5b9b76c..c63a4ebc7893 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -88,7 +88,6 @@ def test_mtp_speculative_mixed_batch_short_prefill( trust_remote_code=True, enable_chunked_prefill=True, enable_prefix_caching=enable_prefix_caching, - mamba_cache_mode="align" if enable_prefix_caching else "none", ) as llm: sampling_params = SamplingParams( temperature=0.0, @@ -161,7 +160,6 @@ def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch): enforce_eager=True, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", speculative_config={"method": "mtp", "num_speculative_tokens": 2}, ) as runner: llm = runner.get_llm() diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 89797db4cdd7..f75151414afd 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -60,7 +60,6 @@ def _check_replayssm_parity( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", tensor_parallel_size=tensor_parallel_size, mamba_backend=mamba_backend, ) @@ -130,7 +129,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", speculative_config={ "method": "ngram", "num_speculative_tokens": 3, @@ -182,7 +180,6 @@ def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner) max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", mamba_backend="flashinfer", tensor_parallel_size=2, disable_log_stats=False, @@ -250,7 +247,6 @@ def _check_replayssm_prefix_caching_parity( trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, - mamba_cache_mode="align", disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c6f79f37211c..6615a5c103b1 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -60,7 +60,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & @@ -78,7 +77,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 76d219746913..53e7378a7b5c 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -583,7 +583,6 @@ def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int): kv_transfer_config=kv_transfer_config, language_model_only=True, enable_prefix_caching=True, - mamba_cache_mode="align", # Use lossless state storage so this exact-equality regression isolates # offload boundary selection from low-precision Mamba checkpointing. mamba_cache_dtype="float32", diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index b72de811c0f9..74de3c3a8cf1 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -632,6 +632,8 @@ def test_p_side_chunked_prefill_mamba(mock_platform): block_size=BLOCK_SIZE, ) vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + # mamba_cache_mode "none" only runs without prefix caching. + vllm_config.cache_config.enable_prefix_caching = False kv_cache_config = make_kv_cache_config( block_size=BLOCK_SIZE, @@ -855,6 +857,7 @@ def test_async_load_does_not_double_reserve_mamba_spec_blocks( disable_hybrid_kv_cache_manager=False, ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.enable_prefix_caching = mamba_cache_mode == "align" kv_cache_config = make_kv_cache_config( block_size, mamba_enabled=True, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 9f818c866fff..527f7bee05c2 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -49,6 +49,7 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, KVCacheTensor, MambaSpec, SlidingWindowSpec, @@ -2355,8 +2356,6 @@ def _make_hybrid_attention_mamba_scheduler( enable_kv_cache_events: bool = False, ) -> SchedulerFixture: """Build a scheduler for one attention group plus one Mamba group.""" - scheduler_block_size = scheduler_block_size or block_size - hash_block_size = hash_block_size or block_size attention_spec = FullAttentionSpec( block_size=attention_block_size, num_kv_heads=NUM_KV_HEADS, @@ -2369,9 +2368,38 @@ def _make_hybrid_attention_mamba_scheduler( dtypes=(torch.float32,), mamba_cache_mode=mamba_cache_mode, ) + return _make_two_group_scheduler( + attention_spec, + ("mamba", mamba_spec), + num_cpu_blocks=num_cpu_blocks, + num_gpu_blocks=num_gpu_blocks, + scheduler_block_size=scheduler_block_size or block_size, + hash_block_size=hash_block_size or block_size, + dcp_world_size=dcp_world_size, + lazy=lazy, + mamba_cache_mode=mamba_cache_mode, + enable_kv_cache_events=enable_kv_cache_events, + ) + + +def _make_two_group_scheduler( + attention_spec: FullAttentionSpec, + second_group: tuple[str, KVCacheSpec], + *, + num_cpu_blocks: int, + num_gpu_blocks: int, + scheduler_block_size: int, + hash_block_size: int, + dcp_world_size: int, + lazy: bool = False, + mamba_cache_mode: MambaCacheMode = "align", + enable_kv_cache_events: bool = False, +) -> SchedulerFixture: + """Build a scheduler for one attention group plus one named second group.""" + second_name, second_spec = second_group groups = [ KVCacheGroupSpec(["attention"], attention_spec), - KVCacheGroupSpec(["mamba"], mamba_spec), + KVCacheGroupSpec([second_name], second_spec), ] tensors = [ KVCacheTensor( @@ -2380,7 +2408,7 @@ def _make_hybrid_attention_mamba_scheduler( layer_stride=spec.page_size_bytes * num_gpu_blocks, block_stride=spec.page_size_bytes, ) - for group, spec in zip(groups, (attention_spec, mamba_spec)) + for group, spec in zip(groups, (attention_spec, second_spec)) ] kv_cache_config = KVCacheConfig( num_blocks=num_gpu_blocks, @@ -2428,34 +2456,48 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: silently offloads only a fraction of them. """ attention_block_size = BLOCK_SIZE - mamba_block_size = 4 * BLOCK_SIZE + replicated_block_size = 4 * BLOCK_SIZE dcp_world_size = 2 - fix = _make_hybrid_attention_mamba_scheduler( + fix = _make_two_group_scheduler( + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ), + ( + "replicated_attention", + FullAttentionSpec( + block_size=replicated_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + dcp_sharded=False, + ), + ), num_cpu_blocks=32, num_gpu_blocks=32, - attention_block_size=attention_block_size, - block_size=mamba_block_size, + scheduler_block_size=replicated_block_size, # Every prefix-cacheable group's resolved block size must be a multiple # of the hash block, so it cannot exceed the smallest of them. hash_block_size=BLOCK_SIZE, dcp_world_size=dcp_world_size, - mamba_cache_mode="none", ) sched = fix.scheduler gpu_pool = fix.gpu_block_pool - attention_size, mamba_size = sched.group_block_sizes + attention_size, replicated_size = sched.group_block_sizes assert attention_size == attention_block_size * dcp_world_size - assert mamba_size == mamba_block_size + assert replicated_size == replicated_block_size confirmed = 2 * sched.block_size req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) attention_blocks = _allocate_cp_gpu_blocks( gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 ) - mamba_blocks = _allocate_cp_gpu_blocks( - gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 + replicated_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // replicated_size, replicated_size, group_id=1 ) - kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, replicated_blocks)) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) sched.build_connector_meta( make_scheduler_output( @@ -2476,11 +2518,11 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: # count. assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ confirmed // attention_size, - confirmed // mamba_size, + confirmed // replicated_size, ] assert set(meta.store_gpu_blocks) == { *(b.block_id for b in attention_blocks), - *(b.block_id for b in mamba_blocks), + *(b.block_id for b in replicated_blocks), } diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 5eb26726c718..9bb716621594 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -176,8 +176,7 @@ def device_memory_utilization(self, value: float) -> None: page.""" mamba_block_size: int | None = Field(default=None, gt=0) """Size of a contiguous cache block in number of tokens for mamba cache. - Can be set only when prefix caching is enabled. - Value must be a multiple of 8 to align with causal_conv1d kernel.""" + Can be set only when prefix caching is enabled.""" mamba_cache_dtype: MambaDType = "auto" """The data type to use for the Mamba cache (both the conv as well as the ssm state). If set to 'auto', the data type will be inferred from the model @@ -187,19 +186,19 @@ def device_memory_utilization(self, value: float) -> None: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers: + """The cache strategy for Mamba layers, derived from `enable_prefix_caching`: - - "none": set when prefix caching is disabled. - - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. This is the default when prefix - caching is enabled. + - "none": prefix caching is disabled. + - "align": prefix caching is enabled. Only cache the mamba state of the last + token of each scheduler step and when the token is at position + i * block_size. """ enable_mamba_shared_prefix_checkpoint: bool = False """Also register a Mamba "align" checkpoint at the shared-prefix junction -- where an EAGLE/MTP sibling was observed to resume -- instead of only at the - prompt tail. Off by default; only takes effect with `mamba_cache_mode` - "align", EAGLE on the Mamba group, and a prefix match unit smaller than the - Mamba block size.""" + prompt tail. Off by default; only takes effect with prefix caching enabled, + EAGLE on the Mamba group, and a prefix match unit smaller than the Mamba + block size.""" replayssm_buffer_len: int = Field(default=16, gt=0) """ReplaySSM logical history length B for Mamba2. Triton uses B physical rows and FlashInfer uses B+T, where T is the target verification length. @@ -207,11 +206,10 @@ def device_memory_utilization(self, value: float) -> None: use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. - Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton - or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer - and mamba_cache_mode 'none'. In align mode flushes are most efficient when - mamba_block_size is a multiple of replayssm_buffer_len, but this is not - required.""" + Requires the Triton or FlashInfer mamba backend; FlashInfer does not support + prefix caching yet. Mamba2 speculative decode requires FlashInfer. With + prefix caching, flushes are most efficient when mamba_block_size is a + multiple of replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 87759f46061b..c4016e82bc3d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2038,10 +2038,6 @@ def has_noops(self) -> bool: def has_inner_state(self): return self._model_info.has_inner_state - @property - def supports_mamba_prefix_caching(self) -> bool: - return self._model_info.supports_mamba_prefix_caching - @property def supports_replayssm(self) -> bool: return self._model_info.supports_replayssm diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f27c8418ab4d..d6e834861a19 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3360,7 +3360,7 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": and not self.use_v2_model_runner ): raise ValueError( - "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + "RecoverSSM with prefix caching requires VLLM_USE_V2_MODEL_RUNNER=1" ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError( @@ -3371,8 +3371,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif use_mamba_replayssm_spec: if self.cache_config.mamba_cache_mode != "none": raise ValueError( - "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none" + "FlashInfer ReplaySSM speculative decoding does not support " + "prefix caching; use --no-enable-prefix-caching" ) query_len = 1 + self.num_speculative_tokens if self.cache_config.replayssm_buffer_len < query_len: @@ -3390,8 +3390,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: if self.cache_config.mamba_cache_mode == "align": raise ValueError( - "FlashInfer ReplaySSM does not support " - "--mamba-cache-mode align yet; use none" + "FlashInfer ReplaySSM does not support prefix caching yet; " + "use --no-enable-prefix-caching" ) elif self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 86199e5f007b..c9aa433162d2 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE = ( + "--mamba-cache-mode is deprecated and will be removed in v0.33.0. " + "It has no effect: Mamba prefix caching follows --enable-prefix-caching." +) + + _DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( "--disable-cascade-attn is deprecated and has no effect since cascade " "attention has been removed. It will be removed in v0.33.0." @@ -798,7 +804,7 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") - mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + mamba_cache_mode: MambaCacheMode | None = None enable_mamba_shared_prefix_checkpoint: bool = ( CacheConfig.enable_mamba_shared_prefix_checkpoint ) @@ -1403,7 +1409,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--prefix-match-unit", **cache_kwargs["prefix_match_unit"] ) cache_group.add_argument( - "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] + "--mamba-cache-mode", + **{ + **cache_kwargs["mamba_cache_mode"], + "default": None, + "help": _MAMBA_CACHE_MODE_DEPRECATION_MESSAGE, + }, ) cache_group.add_argument( "--enable-mamba-shared-prefix-checkpoint", @@ -2263,6 +2274,8 @@ def create_engine_config( assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + if self.mamba_cache_mode is not None: + logger.warning_once(_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] @@ -2282,7 +2295,6 @@ def create_engine_config( mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, - mamba_cache_mode=self.mamba_cache_mode, enable_mamba_shared_prefix_checkpoint=( self.enable_mamba_shared_prefix_checkpoint ), diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index bfa1103bb368..14c483122ab1 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -812,10 +812,8 @@ def conv_ssm_forward( # - It will read the initial states for every sequence, # that has "has_initial_states_p" == True, # from "cache_indices", using "state_indices_tensor_p". - # - It updates the "conv_state" cache in positions pointed - # to by "state_indices_tensor_p". - # In particular, it will always write the state at the - # sequence end. + # - It writes the state at the sequence end to the "conv_state" + # cache in positions pointed to by "state_indices_tensor_p". x = hidden_states_B_C_p.transpose( 0, 1 ) # this is the form that causal-conv see diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index e8254e4f34eb..2094df902a27 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -20,17 +20,11 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptr, # (dim, width) bias_ptr, initial_states_ptr, # conv_states_ptr - cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains - # the block indices relevant for each sequence - # plus potential 0-padding at the beginning and at the end + cache_indices_ptr, # (batch,) has_initial_states_ptr, query_start_loc_ptr, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, # (batch,) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) - num_computed_tokens, # (batch,) o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, @@ -46,7 +40,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching stride_cache_indices: tl.constexpr, stride_o_dim: tl.constexpr, stride_o_token: tl.int64, - stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M # others pad_slot_id: tl.constexpr, null_block_id: tl.constexpr, @@ -54,7 +47,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching HAS_BIAS: tl.constexpr, KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, @@ -93,40 +85,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching # find the actual sequence length seqlen = sequence_end_index - sequence_start_index - B_size: tl.constexpr = stride_block_m * BLOCK_M - - if IS_APC_ENABLED: - # Handle the case if prefix caching is enabled. - # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" - - # Get the length of the completed sequence so far and compute the offset. - current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - sequence_completed_index = tl.load(num_computed_tokens + idx_seq) - - # Compute the offset where the first stride_block_m-aligned first full block is - # Value in "token-space" - sequence_completed_offset_token = sequence_completed_index % B_size - seq_completed_offset = B_size - sequence_completed_offset_token - seq_end_offset = (seqlen - seq_completed_offset) % B_size - last_full_block_token_index = sequence_end_index - seq_end_offset - # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one - if seq_end_offset == 0: - last_full_block_token_index = last_full_block_token_index - B_size - - # Get the number of blocks to be filled for the current sequence - # If n_block_to_fill = 0, then only the state at the sequence end is stored - n_block_to_fill = current_last_index - current_first_index - - # Get the index of the init block - conv_state_init_index = tl.load(initial_state_idx + idx_seq) - else: - n_block_to_fill = 0 - current_last_index = 0 - conv_state_init_index = 0 - current_first_index = 0 - last_full_block_token_index = 0 - token_offset = BLOCK_M * chunk_offset segment_len = min(BLOCK_M, seqlen - token_offset) @@ -137,7 +95,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + conv_state_indices_ptr + idx_seq * stride_cache_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -224,20 +182,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching loaded_x = tl.load(x_ptrs, mask_x, 0.0) idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - # Compute the offset where the last block should be written in the conv_states - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_last_index - ).to(tl.int64) - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] tl.debug_barrier() # NOTE: use this due to bug in Triton compiler @@ -346,52 +294,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") - # Store intermediate states aligned with stride_block_m - # The additional states are cached starting from the last stride_block_m. - # For example: - # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. - # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last - # stride_block_m are cached. - # For example chunk_offset = n_block_to_fill stores the state at last_full_block - if (chunk_offset - 1) < n_block_to_fill: - # Store the states at the chunk boundaries from the start of the sequence - idx_tokens_last = ( - last_full_block_token_index - - (n_block_to_fill - chunk_offset) * B_size - - state_len - ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] - x_ptrs = ( - x_ptr - + (idx_tokens_last * stride_x_token)[:, None] - + (idx_feats * stride_x_dim)[None, :] - ) # [BLOCK_M,BLOCK_N,] - - mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ - None, : - ] # token-index # token-index # feature-index - loaded_x = tl.load(x_ptrs, mask_x, 0.0) - idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - - # cache_idx - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_first_index - + (chunk_offset - 1) - ).to(tl.int64) - - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] - - mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] - tl.debug_barrier() # NOTE: use this due to bug in Triton compiler - tl.store(conv_states_ptrs_target, loaded_x, mask) - if HAS_BIAS: bias = bias_ptr + idx_feats mask_bias = idx_feats < dim @@ -489,11 +391,6 @@ def causal_conv1d_fn( activation: str | None = "silu", pad_slot_id: int = PAD_SLOT_ID, null_block_id: int = NULL_BLOCK_ID, - block_idx_first_scheduled_token: torch.Tensor | None = None, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, - num_computed_tokens: torch.Tensor | None = None, - block_size_to_align=0, metadata=None, validate_data=False, ): @@ -536,16 +433,6 @@ def causal_conv1d_fn( for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] in this case, the kernel will not process entries at indices 0 and 3 - block_idx_first_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the first cache block to be filled is located. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into cache_indices, where the cache block containing the initial state is located. - num_computed_tokens: (batch,), dtype int32 - The number of tokens already completed for each sequence - block_size_to_align: int - The block size to align the cached states to out: same shape as `x` """ if isinstance(activation, bool) and activation: @@ -631,12 +518,6 @@ def causal_conv1d_fn( assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" - if block_size_to_align is not None and block_size_to_align > 0: - assert (block_size_to_align % BLOCK_M) == 0, ( - "The mamba block size needs to be divisible by the BLOCK_M" - ) - else: - block_size_to_align = BLOCK_M if metadata is None: @@ -720,10 +601,6 @@ def grid(META): query_start_loc, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, - block_idx_last_scheduled_token, - initial_state_idx, - num_computed_tokens, out, # Matrix dimensions dim, @@ -739,7 +616,6 @@ def grid(META): stride_cache_indices, stride_o_dim, stride_o_token, - block_size_to_align // BLOCK_M, # others pad_slot_id, null_block_id, @@ -747,7 +623,6 @@ def grid(META): HAS_BIAS=bias is not None, KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, HAS_NULL_BLOCK=null_block_id is not None, NP2_STATELEN=np2_statelen, # launch_cooperative_grid=True @@ -769,8 +644,6 @@ def _causal_conv1d_update_kernel( conv_state_indices_ptr, num_accepted_tokens_ptr, query_start_loc_ptr, # (batch + 1) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) o_ptr, # (batch, dim, seqlen) # Matrix dimensions batch: int, @@ -798,7 +671,6 @@ def _causal_conv1d_update_kernel( KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, IS_VARLEN: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, @@ -818,17 +690,9 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - if IS_APC_ENABLED: - # Get the state from the initial_state_idx - conv_state_init = tl.load(initial_state_idx + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - else: - conv_state_init = 0 - current_last_index = 0 - # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + conv_state_indices_ptr + idx_seq * stride_state_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -941,14 +805,9 @@ def _causal_conv1d_update_kernel( new_conv_state = tl.where(mask, conv_state, loaded_x) - # Get the state from the initial_state_idx - # cache_idx - conv_states_offset = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index - ).to(tl.int64) conv_state_ptrs_target = ( conv_state_ptr - + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (conv_states_input_coord * stride_conv_state_seq) # Offset from seq + (idx_feats * stride_conv_state_dim) )[None, :] + ( # [BLOCK_N,] idx_tokens * stride_conv_state_tok @@ -1104,8 +963,6 @@ def causal_conv1d_update( query_start_loc: torch.Tensor | None = None, max_query_len: int = -1, null_block_id: int = NULL_BLOCK_ID, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, validate_data=False, out: torch.Tensor | None = None, ): @@ -1123,10 +980,6 @@ def causal_conv1d_update( If not None, the conv_state is a larger tensor along the batch dim, and we are selecting the batch coords specified by conv_state_indices. Useful for a continuous batching scenario. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into conv_state_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into conv_state_indices, where the cache block containing the initial state is located. num_accepted_tokens: (batch,), dtype int32 If not None, it indicates the number of accepted tokens for each sequence in the batch. @@ -1237,8 +1090,6 @@ def grid(META): conv_state_indices, num_accepted_tokens, query_start_loc, - block_idx_last_scheduled_token, - initial_state_idx, out, # Matrix dimensions batch, @@ -1266,7 +1117,6 @@ def grid(META): KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], IS_VARLEN=query_start_loc is not None, - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, IS_SPEC_DECODING=num_accepted_tokens is not None, NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d64e51bfa8bf..69ca643a3a97 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -692,7 +692,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: model_config.architecture, ) assert vllm_config.scheduler_config.enable_chunked_prefill, ( - "Chunked prefill is required for mamba cache mode 'align'." + "Chunked prefill is required for Mamba prefix caching." ) # By default, mamba block size will be set to max_model_len (see # below). When enabling prefix caching, we align mamba block size diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 203a3f3ec393..7aff819dfcd7 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -499,7 +498,6 @@ class FalconH1ForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 979d272c13d1..040f7453fb5b 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, ) @@ -587,7 +586,6 @@ class GraniteMoeHybridForCausalLM( SupportsPP, IsHybrid, SupportsQuant, - SupportsMambaPrefixCaching, ): packed_modules_mapping: dict[str, list[str]] = { "qkv_proj": [ diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 4dc67ff34d27..a780cbe23cf6 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1134,13 +1134,10 @@ def has_noops( @runtime_checkable class SupportsMambaPrefixCaching(Protocol): - """The interface for models whose mamba layers support prefix caching. - - This is currently experimental. + """The interface for attention-free Mamba models that support prefix + caching. Hybrid models get the same methods from `IsHybrid`. """ - supports_mamba_prefix_caching: ClassVar[Literal[True]] = True - @classmethod def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: """Return copy functions for the model's Mamba states.""" @@ -1156,24 +1153,6 @@ def get_mamba_state_copy_funcs( return {mamba_type: copy_funcs for mamba_type in mamba_types} -@overload -def supports_mamba_prefix_caching( - model: type[object], -) -> TypeIs[type[SupportsMambaPrefixCaching]]: ... - - -@overload -def supports_mamba_prefix_caching( - model: object, -) -> TypeIs[SupportsMambaPrefixCaching]: ... - - -def supports_mamba_prefix_caching( - model: type[object] | object, -) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]: - return getattr(model, "supports_mamba_prefix_caching", False) - - @runtime_checkable class SupportsReplaySSM(Protocol): """The interface for models whose recurrent layers support ReplaySSM diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index fe9349c30b10..315c821d6a0e 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -44,7 +44,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -380,7 +379,6 @@ class JambaForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"}, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ae653950ac48..050a13431ddc 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -74,7 +74,6 @@ SupportsEagle, SupportsEagle3, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, SupportsReplaySSM, @@ -725,7 +724,6 @@ class NemotronHForCausalLM( IsHybrid, SupportsQuant, MixtureOfExperts, - SupportsMambaPrefixCaching, SupportsReplaySSM, ): # Relevant only if self.has_moe is True diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cd0557ac4585..09e5025c4e0e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -50,7 +50,6 @@ is_hybrid, requires_raw_input_tokens, supports_hc_sp, - supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -881,7 +880,6 @@ class _ModelInfo: is_attention_free: bool is_hybrid: bool has_noops: bool - supports_mamba_prefix_caching: bool supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -911,7 +909,6 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), - supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index 2e90da86ea01..35e972818726 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -45,7 +45,7 @@ ) from vllm.sequence import IntermediateTensors -from .interfaces import HasInnerState, IsHybrid, SupportsMambaPrefixCaching +from .interfaces import HasInnerState, IsHybrid from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -833,7 +833,7 @@ def forward( return hidden_states -class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching): +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid): """Zamba2 model with causal language modeling head. This class wraps the core Zamba2 model and adds: @@ -911,10 +911,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: LoRA and scheduler settings prefix: Optional prefix for parameter names - Raises: - AssertionError: If prefix caching is enabled - (not supported by Mamba) - """ config = vllm_config.model_config.hf_config diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 3e7ae905a30d..9820e8660e55 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -245,13 +245,10 @@ def _compute_chunk_metadata( The code below carefully constructs the chunks such that: 1. Chunks contain tokens from a *single* sequence only. - 2. For every sequence, we are guaranteed that we can - retrieve the mamba state *every* chunk_size tokens. + 2. Chunk boundaries fall on multiples of chunk_size in each + sequence's absolute token positions, so the chunking does not + depend on how chunked prefill split the sequence. Constraint (1) dramatically simplifies the mamba kernels. - Constraint (2) dramatically simplifies the implementation - of prefix caching for mamba (wip). We need to take care - of the interaction with chunked prefill in order to - satisfy constraint (2). `checkpoint_offsets_p[i]` forces an extra chunk end that many tokens into row i's query (0 = none), because the SSD scan only materializes diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2e4fe55b96c2..28a7a1106f77 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1426,6 +1426,9 @@ def __init__( ) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode + assert self.mamba_cache_mode == "align" or not self.enable_caching, ( + "Mamba prefix caching requires mamba_cache_mode='align'" + ) self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks self.has_prefill_checkpoint_blocks = ( self.mamba_cache_mode == "align" From 3ad096faf81b4734ae995aa66f4ca0c904cea978 Mon Sep 17 00:00:00 2001 From: zhouyou Date: Tue, 13 Oct 2026 15:59:48 +0800 Subject: [PATCH 476/952] [Bugfix][EC Connector] Don't create hash directories on cache-miss lookups in ECExampleConnector (#60826) Signed-off-by: zhouyou9505 --- .../ec_connector/unit/test_ec_example_connector.py | 12 ++++++++++++ .../ec_transfer/ec_connector/example_connector.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index 725bb99eb055..a940fd4fcf5e 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -246,6 +246,18 @@ def test_has_cache_item_partial_exist( assert result[1] # Second exists assert not result[2] # Third doesn't exist + def test_has_cache_item_miss_creates_no_folder( + self, mock_vllm_config_consumer, temp_storage + ): + """Test a cache miss leaves no empty folder in the shared storage.""" + connector = ECExampleConnector( + vllm_config=mock_vllm_config_consumer, + role=ECConnectorRole.SCHEDULER, + ) + + assert not connector.has_cache_item("missing_hash") + assert os.listdir(temp_storage) == [] + class TestStateManagement: """Test connector state management.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py index e080f09b6bef..97075cdf38c9 100644 --- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/example_connector.py @@ -202,7 +202,7 @@ def request_finished( def _found_match_for_mm_data(self, mm_hash) -> bool: """Check if the cache is hit for the request.""" - filename = self._generate_filename_debug(mm_hash) + filename = self._generate_filename_debug(mm_hash, create_folder=False) return os.path.exists(filename) def _generate_foldername_debug( @@ -219,11 +219,13 @@ def _generate_foldername_debug( os.makedirs(foldername, exist_ok=True) return foldername - def _generate_filename_debug(self, mm_hash: str) -> str: + def _generate_filename_debug( + self, + mm_hash: str, + create_folder: bool = True, + ) -> str: """Return the full path of the safetensors file for this mm_hash. - Ensures the parent directory exists because - `_generate_foldername_debug` is called with its default - (`create_folder=True`). + If `create_folder` is True (default) the parent directory is created. """ - foldername = self._generate_foldername_debug(mm_hash) # <- folder auto-created + foldername = self._generate_foldername_debug(mm_hash, create_folder) return os.path.join(foldername, "encoder_cache.safetensors") From ff305934f1c6f98edf7c499464a3c2be5f6e3f04 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 16:15:47 +0800 Subject: [PATCH 477/952] [Tests][Mamba] Wait for accepted counts in fake MTP proposer (#61342) Signed-off-by: zjy0516 --- tests/v1/e2e/general/test_mamba_prefix_cache.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 6c208e19fb5d..cf8aee141bc5 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -129,6 +129,10 @@ def fake_propose_draft_token_ids_fn( first_token_id_index = ( num_computed_tokens + 1 ) # bonus token isn't considered as computed + # num_accepted_tokens_cpu is filled by a non-blocking D2H copy in the + # same step; wait for it or a lagging GPU yields stale counts. + if self.num_accepted_tokens_event is not None: + self.num_accepted_tokens_event.synchronize() first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item() proposed_draft_token_ids = [ prompt_token_ids[ From 24aa1191f433638722178fed646f0fd52e96a587 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 13 Oct 2026 10:35:57 +0200 Subject: [PATCH 478/952] =?UTF-8?q?[ROCm][Bugfix]=20Run=20the=20AITER=20AR?= =?UTF-8?q?+RMSNorm=20fallback=20test=20on=20the=20platform=E2=80=A6=20(#6?= =?UTF-8?q?1350)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Stefan Koncarevic Signed-off-by: Stefan Koncarevic --- tests/distributed/test_aiter_rms_fallback.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py index 4a0fc67873a1..211ff15e6e0d 100644 --- a/tests/distributed/test_aiter_rms_fallback.py +++ b/tests/distributed/test_aiter_rms_fallback.py @@ -11,6 +11,7 @@ import vllm._aiter_ops as aiter_ops import vllm.distributed as vllm_distributed import vllm.kernels # noqa: F401 +from vllm.platforms import current_platform def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce( @@ -72,9 +73,14 @@ def fake_rmsnorm2d_fwd_with_add( fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add monkeypatch.setitem(sys.modules, "aiter", fake_aiter) - inp = torch.randn(8192, 16, dtype=torch.bfloat16) + # On the platform's device, not the CPU. The fallback leaves Python through + # `torch.ops.vllm_aiter.fused_add_rms_norm`, which is registered for the + # platform's dispatch key alone, so CPU tensors have nothing to dispatch to + # on a GPU build. The kernel behind it is still the fake above. + device = current_platform.device_type + inp = torch.randn(8192, 16, dtype=torch.bfloat16, device=device) residual = torch.randn_like(inp) - weight = torch.randn(16, dtype=torch.bfloat16) + weight = torch.randn(16, dtype=torch.bfloat16, device=device) out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( inp, residual, weight, 1e-6, gemma_norm=True From 6f11c220b5b0f3b45b2a14c49ebbca2df12d1c80 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:41:23 +0100 Subject: [PATCH 479/952] [Transformers backend] Fuse MoE routers that score after the top-k (#57410) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/models/transformers/fusers/test_moe.py | 70 ++++++++++++++++++ .../models/transformers/fusers/moe.py | 73 +++++++++++++++---- .../models/transformers/fx_utils.py | 30 +++++++- .../model_executor/models/transformers/moe.py | 31 ++++---- vllm/transformers_utils/config.py | 36 +++++++++ .../model_arch_config_convertor.py | 33 ++------- 6 files changed, 213 insertions(+), 60 deletions(-) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 7de4589eef91..f0e70c896baa 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -29,6 +29,42 @@ def forward(self, hidden_states): return logits, value, index +class SelectionScoredRouter(TopKRouter): + """Scores the selected logits only (GraniteMoE/Aria style), in fp32.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight).float() + value, index = torch.topk(logits, self.top_k, dim=-1) + value = ( + torch.sigmoid(value) if self.sigmoid else F.softmax(value, dim=-1) + ).type_as(hidden_states) + return logits, value, index + + +class EpsilonRenormRouter(TopKRouter): + """Renormalizes via an epsilon-offset denominator (DeepSeek-V3 style).""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) + value, index = torch.topk(scores, self.top_k, dim=-1) + denominator = value.sum(dim=-1, keepdim=True) + 1e-20 + value /= denominator + return logits, value, index + + +class UntraceableRouter(TopKRouter): + """Scores and selects, then loses the trace before the renormalization.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + if value.sum() > 0: # data-dependent branch: tracing stops here + value = value / value.sum(dim=-1, keepdim=True) + return logits, value, index + + class ScaledRouter(TopKRouter): """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`).""" @@ -268,6 +304,19 @@ def test_moe_fuser_detects_router(sigmoid): assert fuser.shared_name is None and fuser.shared_gate_name is None +@pytest.mark.parametrize("sigmoid", [False, True]) +def test_moe_fuser_matches_router_scoring_the_selection(sigmoid): + """Scoring the selection rather than all logits still allows the fuser to match.""" + with torch.device("meta"): + block = MoEBlock(lambda: SelectionScoredRouter(sigmoid=sigmoid)) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.gate_name == "gate" + assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax") + assert fuser.renormalize is (not sigmoid) + assert fuser.router_dtype == torch.float32 + + def test_moe_fuser_matches_scaled_router(): """Weight scaling after the top-k (DeepSeek style) does not break matching.""" with torch.device("meta"): @@ -275,6 +324,27 @@ def test_moe_fuser_matches_scaled_router(): assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) +@pytest.mark.parametrize( + "router_cls,expected", + [ + (TopKRouter, True), # divides by the sum of the selection + (EpsilonRenormRouter, True), # ... with an epsilon in the denominator + (ScaledRouter, False), # scales the weights instead of renormalizing + (UntraceableRouter, None), # trace stops before it could say + ], +) +def test_moe_fuser_reads_renormalize_from_the_gate(router_cls, expected): + """The router's graph, not the config, decides whether vLLM renormalizes: it + is the code being replaced, and it already bakes in `norm_topk_prob` because + the flag is a Python bool that fx resolves at trace time. Absence of the + division only means "no" on a complete trace.""" + with torch.device("meta"): + block = MoEBlock(router_cls) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.renormalize is expected + + def test_moe_fuser_matches_grouped_router(): """Group masking between the score and the routing top-k still matches: the scorer is anchored on the routing (last) top-k, not the group one.""" diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index af3a101a0ce1..06fdac109ca4 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, + is_complete, is_op, peel, trace, @@ -107,6 +108,24 @@ def _reaches(node: fx.Node, key: str) -> set[fx.Node]: return seen +def _renormalizes(topk: fx.Node) -> bool: + """Are the top-k weights divided by their own sum? + + Every HF router spells renormalization as a division by a sum over the + selection, optionally offset by an epsilon.""" + downstream = _reaches(topk, "users") + for node in downstream: + if not (is_op(node, "div") or is_op(node, "truediv")) or len(node.args) < 2: + continue + denominator = node.args[1] + if not isinstance(denominator, fx.Node): + continue + cone = _reaches(denominator, "all_input_nodes") + if any(is_op(n, "sum") and n in downstream for n in cone): + return True + return False + + class SharedExpertMLP(nn.Module): """Wraps an HF shared expert, applying the output gating it is paired with.""" @@ -147,12 +166,18 @@ class MoEBlockFuser: shared_name: str | None shared_gate_name: str | None router_dtype: torch.dtype | None = None + renormalize: bool | None = None @staticmethod - def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: - """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. - - Returns the scoring function and the dtype the router computes in.""" + def _match_router( + gate: nn.Module, + ) -> tuple[str, torch.dtype | None, bool | None] | None: + """Matches `topk(score(linear(x)))` or `score(topk(linear(x)))`, `score` being + `softmax`/`sigmoid`. + + Returns the scoring function, the dtype the router computes in, and + whether the top-k weights are renormalized (`None` if the trace was cut + short before it could say).""" state = {name for name, _ in named_state(gate)} if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None @@ -164,20 +189,31 @@ def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: if not topks: return None topk = topks[-1] - # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. - scorers = [ - n - for n in _reaches(topk, "all_input_nodes") - if is_op(n, "softmax") or is_op(n, "sigmoid") - ] + is_scorer = lambda n: is_op(n, "softmax") or is_op(n, "sigmoid") + post_topk = False + # Find all scoring functions that reach the top-k node. + scorers = [n for n in _reaches(topk, "all_input_nodes") if is_scorer(n)] + if not scorers: + # Find all scoring functions that are users of the top-k node. + scorers = [n for n in _reaches(topk, "users") if is_scorer(n)] + post_topk = True if len(scorers) != 1: return None scorer = scorers[0] logits_cone = _reaches(scorer, "all_input_nodes") if not any(is_op(n, "linear") for n in logits_cone): return None - scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" - return scoring_func, _forced_dtype(logits_cone) + is_softmax = is_op(scorer, "softmax") + # vLLM always scores every expert, so matching a router that softmaxes + # the selection alone takes renormalizing. Otherwise look for the + # division, reading its absence as "no" only from a complete trace. + renormalize: bool | None + if (post_topk and is_softmax) or _renormalizes(topk): + renormalize = True + else: + renormalize = False if is_complete(graph) else None + scoring_func = "softmax" if is_softmax else "sigmoid" + return scoring_func, _forced_dtype(logits_cone), renormalize @staticmethod def _match_shared_experts( @@ -221,14 +257,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None if _returns_tuple(type(moe_block)): return None # Router: the child that scores + top-k selects. - gate_name = scoring_func = router_dtype = None + gate_name = scoring_func = router_dtype = renormalize = None for name, child in moe_block.named_children(): if ( name != experts_name and (router := cls._match_router(child)) is not None ): gate_name = name - scoring_func, router_dtype = router + scoring_func, router_dtype, renormalize = router break if gate_name is None or scoring_func is None: return None @@ -256,7 +292,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None for name, child in moe_block.named_children(): if name not in accounted and next(named_state(child), None) is not None: return None - return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) + return cls( + gate_name, + scoring_func, + shared_name, + shared_gate_name, + router_dtype, + renormalize, + ) def gate( self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index cb630866a390..818e1ca7bb4f 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -35,6 +35,21 @@ """The unpatched `nn.Module.__call__`. During tracing fx patches it to record `call_module` nodes; meta execution must call modules for real.""" +_MODULE_GETATTR = nn.Module.__getattr__ +"""The unpatched `nn.Module.__getattr__`. During tracing fx patches it to record +`get_attr` nodes; reading one while inferring a `get_attr`'s own value would +recurse until the stack runs out, truncating the graph.""" + + +@contextlib.contextmanager +def _unpatched_modules(): + """Restores the module hooks fx patches, so meta execution records nothing.""" + with ( + mock.patch.object(nn.Module, "__call__", _MODULE_CALL), + mock.patch.object(nn.Module, "__getattr__", _MODULE_GETATTR), + ): + yield + def is_leaf_call(node: object) -> TypeIs[fx.Node]: """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" @@ -132,7 +147,8 @@ def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): ) return _UNKNOWN if kind == "get_attr": - value = operator.attrgetter(str(target))(self.root) + with _unpatched_modules(): + value = operator.attrgetter(str(target))(self.root) if isinstance(value, torch.Tensor): value = torch.empty_like(value, device="meta") return value @@ -159,8 +175,6 @@ def meta_of(arg: object) -> object: if kind == "call_module": # Run the child's forward with all its state on "meta", without # mutating it (at match time params may be meta but buffers real). - # fx patches `nn.Module.__call__` while tracing; restore the real - # one so this execution is not itself recorded. child = self.root.get_submodule(str(target)) state = { name: torch.empty_like(tensor, device="meta") @@ -168,7 +182,7 @@ def meta_of(arg: object) -> object: child.named_parameters(), child.named_buffers() ) } - with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): + with _unpatched_modules(): return torch.func.functional_call(child, state, meta_args, meta_kwargs) return _UNKNOWN @@ -686,6 +700,14 @@ def visit(self, node: ast.AST) -> ast.AST: _Replacer().visit(module) +def is_complete(graph: fx.Graph) -> bool: + """Did the trace run to completion? A partial graph has no `output` node. + + Only a complete graph can be read for the *absence* of an op; in a partial + one the tail is exactly what is missing.""" + return any(node.op == "output" for node in graph.nodes) + + def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: """The first node in `graph` matching `predicate`, or `None`.""" return next((n for n in graph.nodes if predicate(n)), None) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 71f5bd066c8d..953459172d36 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -39,6 +39,7 @@ from vllm.model_executor.models.transformers.fusers.glu import GLUFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix +from vllm.transformers_utils.config import CONFIG_FIELD_ALIASES from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op from .base import Base @@ -180,23 +181,18 @@ def recursive_replace(self): text_config = self.text_config # Positional arguments - num_experts = self.model_config.get_num_experts() - top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None) + num_experts_names = CONFIG_FIELD_ALIASES["num_experts"] + num_experts = getattr_iter(text_config, num_experts_names, None) + top_k_names = CONFIG_FIELD_ALIASES["num_experts_per_token"] + top_k = getattr_iter(text_config, top_k_names, None) assert top_k is not None hidden_size = text_config.hidden_size - intermediate_size = getattr_iter( - text_config, ["moe_intermediate_size", "intermediate_size"], None - ) + intermediate_size_names = CONFIG_FIELD_ALIASES["moe_intermediate_size"] + intermediate_size = getattr_iter(text_config, intermediate_size_names, None) assert intermediate_size is not None - num_shared_experts = getattr_iter( - text_config, - [ - "n_shared_experts", # DeepSeek, Docs, GLM - "moe_num_shared_experts", # Aria, Ernie - ], - 0, - ) + num_shared_experts_names = CONFIG_FIELD_ALIASES["num_shared_experts"] + num_shared_experts = getattr_iter(text_config, num_shared_experts_names, 0) # Common kwargs norm_topk_prob = getattr(text_config, "norm_topk_prob", None) @@ -275,6 +271,9 @@ def _recursive_replace(module: nn.Module, prefix: str): if "bias" in experts_param_name: has_bias = True break + # Whether the expert weights are stored in transposed format. + # Comes from Transformers `use_experts_implementation` decorator. + is_transposed = getattr(experts, "is_transposed", False) # If the config does not specify num_shared_experts, but # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: @@ -308,6 +307,7 @@ def _recursive_replace(module: nn.Module, prefix: str): enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, has_bias=has_bias, + is_fused_checkpoint_transposed=is_transposed, routed_experts_cls=TransformersRoutedExperts, ) fuser = MoEBlockFuser.match(moe_block, experts_name) @@ -346,6 +346,11 @@ def _recursive_replace(module: nn.Module, prefix: str): gate = fuser.gate(moe_block, prefix, router_dtype) kwargs |= dict( scoring_func=fuser.scoring_func, + renormalize=( + layer_renormalize + if fuser.renormalize is None + else fuser.renormalize + ), is_sequence_parallel=( self.parallel_config.use_sequence_parallel_moe ), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index e046c301d23a..1d93a04e662a 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -52,6 +52,42 @@ logger = init_logger(__name__) +CONFIG_FIELD_ALIASES: dict[str, tuple[str, ...]] = { + "total_num_kv_heads": ( + "n_head_kv", # Falcon + "num_kv_heads", + "num_key_value_heads", # LLaMA-2 + "multi_query_group_num", # ChatGLM + "num_attention_groups", # Step3p5 + ), + "num_experts": ( + "num_experts", # Jamba + "moe_num_experts", # Dbrx + "n_routed_experts", # DeepSeek + "num_local_experts", # Mixtral + ), + "num_experts_per_token": ( + "num_experts_per_tok", + "num_experts_per_token", + "top_k_experts", + "moe_topk", # Hunyuan + "moe_top_k", + "top_k", + ), + "moe_intermediate_size": ( + "moe_intermediate_size", + "intermediate_size", + ), + "num_shared_experts": ( + "n_shared_experts", # DeepSeek, Docs, GLM + "moe_num_shared_experts", # Aria, Ernie + ), +} +"""HF config fields that models spell differently, keyed by the name vLLM uses. + +`ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the +raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" + _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 06a937697106..6290f6c6645b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -14,6 +14,7 @@ from vllm.config.utils import getattr_iter from vllm.logger import init_logger from vllm.transformers_utils.config import ( + CONFIG_FIELD_ALIASES, ConfigFormat, get_safetensors_params_metadata, ) @@ -148,23 +149,11 @@ def _get_qk_rope_head_dim(self) -> int: return qk_rope_head_dim def get_total_num_kv_heads(self) -> int: - attributes = [ - # For Falcon: - "n_head_kv", - "num_kv_heads", - # For LLaMA-2: - "num_key_value_heads", - # For ChatGLM: - "multi_query_group_num", - # For Step3p5: - "num_attention_groups", - ] + names = CONFIG_FIELD_ALIASES["total_num_kv_heads"] # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. default_factory = self.get_total_num_attention_heads - return getattr_iter( - self.hf_text_config, attributes, default_factory=default_factory - ) + return getattr_iter(self.hf_text_config, names, default_factory=default_factory) def get_num_experts_from_block_configs(self) -> int: """Check block_configs for heterogeneous models (e.g., NemotronH). @@ -188,13 +177,7 @@ def get_num_experts_from_block_configs(self) -> int: def get_num_experts(self) -> int: """Returns the number of experts in the model.""" - num_expert_names = [ - "num_experts", # Jamba - "moe_num_experts", # Dbrx - "n_routed_experts", # DeepSeek - "num_local_experts", # Mixtral - ] - + num_expert_names = CONFIG_FIELD_ALIASES["num_experts"] num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0) if isinstance(num_experts, list): # Ernie VL's remote code uses list[int]... @@ -206,13 +189,7 @@ def get_num_experts(self) -> int: return num_experts def get_num_experts_per_token(self) -> int: - names = [ - "num_experts_per_tok", - "num_experts_per_token", - "top_k_experts", - "moe_topk", - "moe_top_k", - ] + names = CONFIG_FIELD_ALIASES["num_experts_per_token"] num_experts_per_token = getattr_iter(self.hf_text_config, names, 0) if isinstance(num_experts_per_token, list): return max(num_experts_per_token, default=0) From 183625d4f9f49968218ffbe43242877ead72d569 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:48:13 +0400 Subject: [PATCH 480/952] [Bugfix] Fall back to auto for unquantized MoE layers when --moe-backend is marlin (#60219) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- .../moe/test_unquantized_backend_selection.py | 14 ++++++++++++++ .../layers/fused_moe/oracle/unquantized.py | 7 ++++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index b64e909fbd45..e3a83ec3c5c7 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -626,3 +626,17 @@ def test_select_explicit_triton_backend(is_lora_enabled): assert selected_backend == UnquantizedMoeBackend.TRITON assert experts_cls is not None + + +@skipif_not_cuda_rocm +@pytest.mark.parametrize("moe_backend", ["humming", "marlin"]) +def test_quantization_only_backend_falls_back_to_auto(moe_backend): + """A quantization-only --moe-backend also reaches the unquantized MoE + layers of a quantized model (e.g. a BF16 MTP draft), which must use the + auto-selected backend instead of failing.""" + moe_config = make_dummy_moe_config() + moe_config.moe_backend = moe_backend + + assert select_unquantized_moe_backend( + moe_config=moe_config + ) == select_unquantized_moe_backend(moe_config=make_dummy_moe_config()) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 1af0450dfd13..2f05074c643c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -303,9 +303,10 @@ def _return_or_raise( ) runner_backend = moe_config.moe_backend - # 'humming' is quantization-only; an unquantized layer (e.g. excluded via - # modules_to_not_convert) falls through to auto instead of erroring. - if runner_backend not in ["auto", "humming"]: + # 'humming' and 'marlin' are quantization-only; an unquantized layer (e.g. + # excluded via modules_to_not_convert, or an unquantized MTP draft of a + # quantized model) falls through to auto instead of erroring. + if runner_backend not in ["auto", "humming", "marlin"]: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts From 656d0d4506052a23d100c56e148c7b984828d637 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 13 Oct 2026 17:01:12 +0800 Subject: [PATCH 481/952] [Perf] Tune Triton GDN decode tiles on SM100-family GPUs (#61355) Signed-off-by: Thien Tran Co-authored-by: OpenAI Codex --- .../fla/test_fused_recurrent_packed_decode.py | 56 ++++++++++++++----- .../ops/fused_recurrent.py | 12 ++++ 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/tests/kernels/fla/test_fused_recurrent_packed_decode.py b/tests/kernels/fla/test_fused_recurrent_packed_decode.py index be6597cc865f..5ea66ced2390 100644 --- a/tests/kernels/fla/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/fla/test_fused_recurrent_packed_decode.py @@ -18,17 +18,32 @@ ) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("strided_mixed_qkv", [False, True]) +@pytest.mark.parametrize( + "batch,num_key_heads,num_value_heads,dtype,state_dtype", + [ + (32, 4, 8, torch.float16, torch.float16), + (32, 4, 8, torch.bfloat16, torch.bfloat16), + (32, 4, 8, torch.float32, torch.float32), + (1, 16, 64, torch.bfloat16, torch.float32), + (9, 4, 12, torch.bfloat16, torch.float32), + (1, 4, 12, torch.bfloat16, torch.bfloat16), + ], +) def test_fused_recurrent_packed_decode_matches_reference( - dtype: torch.dtype, strided_mixed_qkv: bool + dtype: torch.dtype, + strided_mixed_qkv: bool, + batch: int, + num_key_heads: int, + num_value_heads: int, + state_dtype: torch.dtype, ): - torch.manual_seed(0) + generator = torch.Generator(device=DEVICE).manual_seed(0) # Small but representative GDN config (Qwen3Next defaults are K=128, V=128). - B = 32 - H = 4 - HV = 8 # grouped value attention: HV must be divisible by H + B = batch + H = num_key_heads + HV = num_value_heads K = 128 V = 128 qkv_dim = 2 * (H * K) + (HV * V) @@ -38,22 +53,33 @@ def test_fused_recurrent_packed_decode_matches_reference( if strided_mixed_qkv: # Simulate a packed view into a larger projection buffer: # mixed_qkv.stride(0) > mixed_qkv.shape[1] - proj = torch.randn((B, qkv_dim + 64), device=device, dtype=dtype) + proj = torch.randn( + (B, qkv_dim + 64), device=device, dtype=dtype, generator=generator + ) mixed_qkv = proj[:, :qkv_dim] else: - mixed_qkv = torch.randn((B, qkv_dim), device=device, dtype=dtype) + mixed_qkv = torch.randn( + (B, qkv_dim), device=device, dtype=dtype, generator=generator + ) - a = torch.randn((B, HV), device=device, dtype=dtype) - b = torch.randn((B, HV), device=device, dtype=dtype) - A_log = torch.randn((HV,), device=device, dtype=dtype) - dt_bias = torch.randn((HV,), device=device, dtype=dtype) + a = torch.randn((B, HV), device=device, dtype=dtype, generator=generator) + b = torch.randn((B, HV), device=device, dtype=dtype, generator=generator) + A_log = torch.randn((HV,), device=device, dtype=dtype, generator=generator) + dt_bias = torch.randn((HV,), device=device, dtype=dtype, generator=generator) # Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is # reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1. ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32) - ssm_state_indices[-3:] = -1 - - state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype) + if B > 3: + ssm_state_indices[-3:] = -1 + ssm_state_indices[-1] = 0 + + state0 = torch.randn( + (B + 1, HV, V, K), + device=device, + dtype=state_dtype, + generator=generator, + ) state_ref = state0.clone() state_packed = state0.clone() diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index eb08b938c2fb..2f3997e995fe 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -10,6 +10,7 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .op import exp, log @@ -439,6 +440,17 @@ def fused_recurrent_gated_delta_rule_packed_decode( f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})." ) BV = min(triton.next_power_of_2(V), 32) + if ( + K == V == 128 + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100, device_id=dev.index) + ): + # Tuned on GB300; validated on GB200 and with 148 SMs. + # MRv2 decode warmup/graph capture compiles the reachable BV variants. + if B * HV <= 64: + BV = 4 + elif initial_state.dtype == torch.float32: + BV = 8 num_stages = 3 num_warps = 1 From 206e45d47009447001e9080d430e64c90fba21a6 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 13 Oct 2026 02:24:06 -0700 Subject: [PATCH 482/952] [CI] Shard V1 Attention automatically instead of by fixed pytest shards (#61200) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/attention.yaml | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index ec47e45fe7a0..5952aa40564c 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -2,8 +2,9 @@ group: Attention depends_on: - image-build steps: -- label: ":nvidia: (H100) V1 Attention Shard %N" +- label: ":nvidia: (H100) V1 Attention" key: v1-attention-h100-mi300 + automatic_shard: true timeout_in_minutes: 85 device: h100 source_file_dependencies: @@ -12,14 +13,18 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + - pytest -v -s v1/attention mirror: amd: label: ":amd: (MI300) V1 Attention Shard %N" dind: false device: mi300_1 timeout_in_minutes: 125 + # automatic_shard covers NVIDIA steps only; the mirror keeps its + # fixed pytest shards. + parallelism: 2 + commands: + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT depends_on: - image-build-amd source_file_dependencies: @@ -35,8 +40,9 @@ steps: - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":nvidia: (B200) V1 Attention Shard %N" +- label: ":nvidia: (B200) V1 Attention" key: v1-attention-b200 + automatic_shard: true timeout_in_minutes: 80 device: b200-k8s source_file_dependencies: @@ -45,14 +51,16 @@ steps: - vllm/v1/attention - tests/v1/attention commands: - - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 2 + - pytest -v -s v1/attention mirror: amd: label: ":amd: (MI355 DPX) V1 Attention Shard %N" dind: false device: mi355_dpx timeout_in_minutes: 125 + # automatic_shard covers NVIDIA steps only; the mirror keeps its + # fixed pytest shards. + parallelism: 2 depends_on: - image-build-amd source_file_dependencies: From c7a305ec10f3ec06c2784f1258162483c5fc8d73 Mon Sep 17 00:00:00 2001 From: sungsoo ha Date: Tue, 13 Oct 2026 03:09:17 -0700 Subject: [PATCH 483/952] [Bugfix][Mamba] Store conv_kernel - 1 inputs in Mamba2 prefill checkpoints (#60797) Signed-off-by: Sungsoo Ha Co-authored-by: Claude Opus 5.5 Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> --- .../mamba/test_mamba2_checkpoint_state.py | 16 +++++++++++----- vllm/model_executor/layers/mamba/mamba_mixer2.py | 1 + .../layers/mamba/ops/ssd_checkpoint.py | 7 +++++-- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/kernels/mamba/test_mamba2_checkpoint_state.py b/tests/kernels/mamba/test_mamba2_checkpoint_state.py index 54fcf80205ce..a06d894792f8 100644 --- a/tests/kernels/mamba/test_mamba2_checkpoint_state.py +++ b/tests/kernels/mamba/test_mamba2_checkpoint_state.py @@ -114,23 +114,27 @@ def test_checkpoint_does_not_perturb_the_final_state(): ) +@pytest.mark.parametrize("spec_slots", [0, 3]) @pytest.mark.parametrize("dim_first", [True, False]) -def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first): +def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first, spec_slots): """The store reads the checkpoint row out of every chunk state. Row 0 checkpoints 32 tokens in, its state being chunk 3 rather than the program's own row, so the store only passes if ``chunk_idx`` is used. Row 1 declines and must write nothing. The conv input is a column slice of a wider projection, the SD conv layout is a transposed view, and SSM rows - are padded pages, matching the strides the mixer passes. + are padded pages, matching the strides the mixer passes. Spec decoding + widens the conv state, and the store must still fill only the first + ``state_len`` slots with the newest inputs. """ dim, state_len, heads, head_dim, d_state = 16, 4, 2, 4, 8 projection = torch.randn(64, dim + 8, device=DEVICE) conv_input = projection[:, :dim] + num_slots = state_len + spec_slots if dim_first: - conv_state = torch.zeros(3, dim, state_len, device=DEVICE) + conv_state = torch.zeros(3, dim, num_slots, device=DEVICE) else: - conv_state = torch.zeros(3, state_len, dim, device=DEVICE).transpose(-1, -2) + conv_state = torch.zeros(3, num_slots, dim, device=DEVICE).transpose(-1, -2) varlen_states = torch.randn(5, heads, head_dim, d_state, device=DEVICE) # Blocks are pages holding more than the SSM state, so rows are padded. row = heads * head_dim * d_state @@ -152,12 +156,14 @@ def test_store_selects_the_checkpoint_from_all_chunk_states(dim_first): varlen_states, ssm_state, torch.tensor([0, 40, 64], dtype=torch.int32, device=DEVICE), + state_len=state_len, ) exact = dict(rtol=0, atol=0) torch.testing.assert_close( - conv_state[2], conv_input[32 - state_len : 32].T, **exact + conv_state[2, :, :state_len], conv_input[32 - state_len : 32].T, **exact ) + assert not conv_state[2, :, state_len:].any(), "wrote into the spec slots" torch.testing.assert_close(ssm_state[2], varlen_states[3], **exact) assert not conv_state[:2].any() and not ssm_storage[:2].any() assert not ssm_storage[:, row:].any(), "wrote into the page padding" diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 14c483122ab1..e310aaf40f12 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -890,6 +890,7 @@ def conv_ssm_forward( varlen_states, ssm_state, query_start_loc_p, + state_len=self.conv_kernel_size - 1, ) ssm_state[state_indices_tensor_p] = final_states if ring_start is not None and self._updates_replayssm_trackers: diff --git a/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py index 25bfcedbdc88..fdd8ac8f0ed5 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_checkpoint.py @@ -69,6 +69,8 @@ def store_prefill_checkpoint( chunk_states: torch.Tensor, ssm_state: torch.Tensor, query_start_loc: torch.Tensor, + *, + state_len: int, ) -> None: """Write each prefill row's mid-prefill checkpoint into its paged block. @@ -78,14 +80,15 @@ def store_prefill_checkpoint( checkpoint: Per-row offsets and destination blocks. chunk_idx: Per-row index into ``chunk_states`` of the checkpoint. x: Pre-conv activations, ``(dim, num_prefill_tokens)``. - conv_state: Paged conv state, ``(num_blocks, dim, state_len)``. + conv_state: Paged conv state, ``(num_blocks, dim, state_len + num_spec)``. chunk_states: SSD states at every chunk end. ssm_state: Paged SSM state. query_start_loc: Prefill query start locations. + state_len: ``conv_kernel_size - 1``, not the spec-widened conv width. + Readers take the first ``state_len`` slots as the newest inputs. """ dim = x.size(0) - state_len = conv_state.size(-1) ssm_row_size = chunk_states[0].numel() block_size = 1024 grid = ( From 877f6df889176f056713ac30919153b9dacc0cf1 Mon Sep 17 00:00:00 2001 From: JackRao123 <112158010+JackRao123@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:15:11 -0700 Subject: [PATCH 484/952] [Bugfix][DSv4] Don't bake the C128 compressor skip into CUDA graphs (#59709) Signed-off-by: Jack Rao Co-authored-by: Claude Opus 5.5 --- tests/kernels/test_compressor_kv_cache.py | 78 +++++++++++++++++++++++ vllm/models/deepseek_v4/compressor.py | 6 +- 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index e9924229055a..237c552fc124 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1205,6 +1205,84 @@ def test_get_c128_boundary(starts, query_start_loc, expected): assert _get_c128_boundary(metadata) is expected +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@torch.inference_mode() +def test_c128_skip_is_not_baked_into_cuda_graph(monkeypatch): + """A graph captured on a batch with no C128 boundary must still launch the + compressor on replay; only eager steps may skip it.""" + from vllm.config import CUDAGraphMode + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v4 import compressor as compressor_mod + from vllm.models.deepseek_v4.compressor import DeepseekCompressor + from vllm.models.deepseek_v4.nvidia.ops import sparse_attn_compress_cutedsl + + launches = torch.zeros(1, dtype=torch.int32, device="cuda") + monkeypatch.setattr( + compressor_mod, "_SAVE_PARTIAL_STATES_KERNEL", lambda **kw: None + ) + monkeypatch.setattr( + sparse_attn_compress_cutedsl, + "_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL", + lambda **kw: launches.add_(1), + ) + + compressor = DeepseekCompressor.__new__(DeepseekCompressor) + torch.nn.Module.__init__(compressor) + compressor.head_dim, compressor.rope_head_dim = 512, 64 + compressor.compress_ratio, compressor.coff, compressor.overlap = 128, 1, False + compressor.use_fp4_cache = False + compressor._use_two_stage_fused_compressor = False + compressor._quant_block = 64 + compressor._token_stride = 576 + compressor._scale_dim = 8 + compressor.rms_norm_eps = 1e-6 + compressor.ape = torch.zeros(128, 512, device="cuda") + compressor.norm = SimpleNamespace(weight=torch.ones(512, device="cuda")) + compressor.state_cache = SimpleNamespace( + prefix="state", kv_cache=torch.zeros(2, 8, 1024, device="cuda") + ) + compressor.k_cache_prefix = "main" + kv_cache = torch.zeros(1, 64, 584, dtype=torch.uint8, device="cuda") + compressor._static_forward_context = {"main": SimpleNamespace(kv_cache=kv_cache)} + + num_tokens = 4 + slots = torch.arange(num_tokens, device="cuda") + metadata = { + "state": SimpleNamespace( + slot_mapping=slots, + block_table=torch.zeros(num_tokens, 1, dtype=torch.int32, device="cuda"), + block_size=8, + token_to_req_indices=torch.arange( + num_tokens, dtype=torch.int32, device="cuda" + ), + c128_boundary=False, + ), + "main": SimpleNamespace(slot_mapping=slots), + } + kv_score = torch.zeros(num_tokens, 1024, device="cuda") + positions = torch.arange(num_tokens, device="cuda") + rotary = SimpleNamespace(cos_sin_cache=torch.zeros(1, 64, device="cuda")) + context = ForwardContext( + {}, metadata, {}, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE + ) + + def run(): + with override_forward_context(context): + compressor(kv_score, positions, rotary) + + run() + torch.accelerator.synchronize() + assert launches.item() == 0, "eager step with no boundary should skip" + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + launches.zero_() + graph.replay() + torch.accelerator.synchronize() + assert launches.item() == 1, "captured graph must not bake in the skip" + + # ── Test A: DeepseekV4 Attention path ────────────────────────────────────────────── diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index ccb864faa274..b1e8b5d5bc8b 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -7,7 +7,7 @@ import torch from torch import nn -from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config +from vllm.config import VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -409,12 +409,12 @@ def forward( pdl_kwargs=pdl_kwargs, ) - # full graph cannot branch on per-step CPU metadata after capture + # A captured graph would replay the capture batch's skip decision. if ( current_platform.is_cuda() and self.head_dim == 512 and self.compress_ratio == 128 - and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL + and not torch.cuda.is_current_stream_capturing() and state_metadata.c128_boundary is False ): return From 7173c64e2ec8763c6d8ddd5c3ff9cb6fc92e2da8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:17:14 +0200 Subject: [PATCH 485/952] [Security] Bound decoded and downloaded media per request (#61226) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: juanpas97 Signed-off-by: Juan Pérez de Algaba Co-authored-by: OpenAI Codex --- docs/usage/security.md | 2 + tests/multimodal/media/test_audio.py | 100 ++++++++++++++++++ tests/multimodal/media/test_image.py | 66 +++++++++++- .../media/test_unprocessable_entity_error.py | 22 +++- tests/multimodal/media/test_video.py | 15 +++ tests/test_audio_media_size_precheck.py | 54 +++++++++- tests/test_connections.py | 51 ++++++++- vllm/connections.py | 65 ++++++++++-- vllm/envs.py | 13 +++ vllm/multimodal/media/__init__.py | 3 +- vllm/multimodal/media/audio.py | 60 ++++++++--- vllm/multimodal/media/base.py | 71 +++++++++++++ vllm/multimodal/media/connector.py | 54 +++++++--- vllm/multimodal/media/image.py | 13 ++- vllm/multimodal/media/video.py | 10 ++ vllm/renderers/chat_utils.py | 16 +++ 16 files changed, 565 insertions(+), 50 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 0d7894a467a2..428e64258e0c 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -94,7 +94,9 @@ denial of service: | Environment Variable | Default | Description | | --- | --- | --- | | `VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB` | `256` | Maximum size in MB for a single remote media response. Oversized responses are rejected while streaming before the full body is materialized in memory. | +| `VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB` | `256` | Maximum total HTTP media bytes downloaded for one request, including concurrent downloads. Enforced while streaming. | | `VLLM_MAX_IMAGE_PIXELS` | `178956970` (~179M pixels) | Maximum decoded image size in pixels. Images exceeding this are rejected before raster memory is allocated. Default matches PIL's built-in 2x decompression-bomb threshold (~680 MB for RGB). | +| `VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES` | `715827880` (~683 MiB) | Maximum aggregate decoded image, video-frame, and audio bytes in one request. Images are charged at four bytes per pixel before decoding; audio is limited before decoding; video files are charged after frame extraction. | | `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum compressed filesize in MB for a single audio file. Enforced on all audio inputs (multimodal chat URLs, speech-to-text uploads, data: URLs, and local file paths) before decoding begins. | | `VLLM_MAX_AUDIO_DECODE_DURATION_S` | `600` | Maximum decoded audio duration in seconds. Prevents compressed audio from expanding into gigabytes of float32 PCM. | | `VLLM_MAX_AUDIO_DECODE_BYTES` | `268435456` (256 MiB) | Maximum float32 PCM bytes that audio decoding may allocate. Guards against sample-rate forgery where an inflated header sample rate bypasses the duration guard while the actual frame count causes a multi-GiB allocation. | diff --git a/tests/multimodal/media/test_audio.py b/tests/multimodal/media/test_audio.py index a75c20979e0c..5728fa9a4660 100644 --- a/tests/multimodal/media/test_audio.py +++ b/tests/multimodal/media/test_audio.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +from concurrent.futures import ThreadPoolExecutor from io import BytesIO from pathlib import Path from unittest.mock import patch @@ -9,6 +11,7 @@ import pytest import soundfile as sf +import vllm.envs as envs from vllm.multimodal.media import AudioMediaIO from vllm.multimodal.media import audio as audio_module from vllm.multimodal.media.audio import ( @@ -16,6 +19,7 @@ load_audio_soundfile, load_audio_torchcodec, ) +from vllm.multimodal.media.base import MediaDecodedBudget from ...conftest import AudioTestAssets @@ -118,6 +122,102 @@ def _make_flac_bytes(frames: int, channels: int, samplerate: int) -> bytes: return buf.getvalue() +@pytest.mark.skip_global_cleanup +def test_request_budget_rejects_second_decoded_audio_clip(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 384) + clip = _make_flac_bytes(frames=64, channels=1, samplerate=16000) + audio_io = AudioMediaIO(decoded_budget=MediaDecodedBudget()) + + audio_io.load_bytes(clip) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(clip) + + +@pytest.mark.skip_global_cleanup +def test_request_budget_allows_parallel_audio_decodes(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 512) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Barrier(2) + + def decode(*args, **kwargs): + assert kwargs["max_decode_bytes"] == 256 + decoding.wait(timeout=5) + return np.zeros(32, dtype=np.float32), 16000 + + with ( + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + results = list(pool.map(audio_io.load_bytes, [b"a", b"b"])) + + assert len(results) == 2 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_waits_for_inflight_audio_decode(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + decoding = threading.Event() + finish_decoding = threading.Event() + second_reservation = threading.Event() + reserve = budget.reserve + decode_calls = 0 + + def track_reservation(max_bytes): + if decoding.is_set(): + second_reservation.set() + return reserve(max_bytes) + + def decode(*args, **kwargs): + nonlocal decode_calls + decode_calls += 1 + decoding.set() + assert finish_decoding.wait(timeout=5) + return np.zeros(64, dtype=np.float32), 16000 + + with ( + patch.object(budget, "reserve", side_effect=track_reservation), + patch.object(audio_module, "load_audio", side_effect=decode), + ThreadPoolExecutor(max_workers=2) as pool, + ): + first = pool.submit(audio_io.load_bytes, b"a") + assert decoding.wait(timeout=5) + second = pool.submit(audio_io.load_bytes, b"b") + try: + assert second_reservation.wait(timeout=5) + finally: + finish_decoding.set() + first.result(timeout=5) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + second.result(timeout=5) + + assert decode_calls == 1 + assert budget.used == 256 + + +@pytest.mark.skip_global_cleanup +def test_request_budget_releases_failed_audio_reservation(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 256) + monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_DECODE_BYTES", 256) + budget = MediaDecodedBudget() + audio_io = AudioMediaIO(decoded_budget=budget) + audio = (np.zeros(64, dtype=np.float32), 16000) + + with patch.object( + audio_module, "load_audio", side_effect=[ValueError("invalid audio"), audio] + ): + with pytest.raises(ValueError, match="invalid audio"): + audio_io.load_bytes(b"a") + assert audio_io.load_bytes(b"b")[0].nbytes == 256 + + assert budget.used == 256 + + def test_small_file_passes_memory_guard(): """A small valid file should pass both duration and memory guards.""" payload = _make_flac_bytes(frames=16000, channels=1, samplerate=16000) diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index 94efb312f272..60cfb86cc0d1 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -11,9 +11,9 @@ ALLOWED_IMAGE_FORMATS, open_image, ) -from vllm.multimodal.media import ImageMediaIO +from vllm.multimodal.media import ImageMediaIO, MediaConnector, MediaDecodedBudget -pytestmark = pytest.mark.cpu_test +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] ASSETS_DIR = Path(__file__).parent.parent / "assets" assert ASSETS_DIR.exists() @@ -288,6 +288,68 @@ def test_image_pixel_limit_disabled(monkeypatch): assert result.media.size == (1000, 1000) +def _png_bytes(width: int, height: int) -> bytes: + from io import BytesIO + + buf = BytesIO() + Image.new("RGB", (width, height), (1, 2, 3)).save(buf, format="PNG") + return buf.getvalue() + + +def test_request_decoded_budget_rejects_aggregate(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + image_io.load_bytes(data) + + +def test_request_decoded_budget_accepts_images_within_limit(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 800) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(10, 10) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 800 + + +def test_request_decoded_budget_disabled(monkeypatch): + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 0) + budget = MediaDecodedBudget() + image_io = ImageMediaIO(decoded_budget=budget) + data = _png_bytes(100, 100) + + image_io.load_bytes(data) + image_io.load_bytes(data) + assert budget.used == 0 + + +def test_media_connector_enforces_request_decoded_budget(monkeypatch): + import pybase64 + + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 600) + b64 = pybase64.b64encode(_png_bytes(10, 10)).decode() + url = f"data:image/png;base64,{b64}" + connector = MediaConnector(decoded_budget=MediaDecodedBudget()) + + connector.fetch_image(url) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + connector.fetch_image(url) + + def test_image_media_io_rejects_postscript_payload(): """EPS/PostScript bytes must not reach an external renderer.""" import pybase64 as base64 diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 6d3c2a0b0aa6..8d7142d019a5 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -13,11 +13,16 @@ import aiohttp import pytest -from vllm.connections import MediaDownloadSizeExceededError +from vllm.connections import ( + MediaDownloadSizeExceededError, + MediaRequestDownloadSizeExceededError, +) from vllm.entrypoints.serve import create_error_response from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + class TestVLLMUnprocessableEntityError: """Tests for VLLMUnprocessableEntityError exception.""" @@ -148,6 +153,21 @@ def test_fetch_image_download_size_limit_error(self): assert "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + def test_fetch_image_request_download_size_limit_error(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = MediaRequestDownloadSizeExceededError( + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + connector.fetch_image("https://example.com/image.jpg") + + assert "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB=256" in str(exc_info.value) + class TestErrorResponse: """Tests for error response creation.""" diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index a2fedb1747f5..e3e2ae66894a 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -9,8 +9,10 @@ import numpy.typing as npt import pybase64 import pytest +import torch from PIL import Image +import vllm.envs as envs from vllm.assets.base import get_vllm_public_assets from vllm.assets.video import ( video_get_metadata, @@ -18,6 +20,7 @@ video_to_pil_images_list, ) from vllm.multimodal.media import ImageMediaIO, VideoMediaIO +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, @@ -47,6 +50,18 @@ def load_bytes( return FAKE_OUTPUT_2 +@pytest.mark.skip_global_cleanup +def test_video_request_budget_counts_torch_frames(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 16) + budget = MediaDecodedBudget() + video_io = VideoMediaIO(ImageMediaIO(decoded_budget=budget)) + frames = torch.zeros((2, 2, 2, 3), dtype=torch.uint8) + video_io.video_loader = SimpleNamespace(load_bytes=lambda *_args, **_kwargs: frames) + + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + video_io.load_bytes(b"video") + + def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch): with monkeypatch.context() as m: m.setenv("VLLM_VIDEO_LOADER_BACKEND", "assert_10_frames_1_fps") diff --git a/tests/test_audio_media_size_precheck.py b/tests/test_audio_media_size_precheck.py index 7635022369d1..e2f3c57a0f75 100644 --- a/tests/test_audio_media_size_precheck.py +++ b/tests/test_audio_media_size_precheck.py @@ -2,15 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import AsyncIterator, Iterator +from io import BytesIO from pathlib import Path from unittest.mock import patch +import numpy as np import pytest +from PIL import Image import vllm.envs as envs -from vllm.connections import HTTPConnection +from vllm.connections import HTTPConnection, MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError -from vllm.multimodal.media import AudioMediaIO, MediaConnector +from vllm.multimodal.media import AudioMediaIO, ImageMediaIO, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.utils.mem_constants import KiB_bytes, MiB_bytes pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] @@ -87,6 +91,46 @@ def test_audio_load_bytes_rejects_oversized(monkeypatch: pytest.MonkeyPatch): AudioMediaIO().load_bytes(b"\x00" * (MiB_bytes + 1)) +def test_request_audio_budget_limits_multiple_clips_before_decode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + clip = np.zeros(4, dtype=np.float32) + + def decode(_source, **kwargs): + if kwargs["max_decode_bytes"] < clip.nbytes: + raise ValueError("VLLM_MAX_AUDIO_DECODE_BYTES") + return clip, 16000 + + with patch("vllm.multimodal.media.audio.load_audio", side_effect=decode) as load: + audio_io = AudioMediaIO(decoded_budget=budget) + audio_io.load_bytes(b"first") + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"): + audio_io.load_bytes(b"second") + + assert budget.used == clip.nbytes + assert load.call_count == 2 + + +def test_request_decoded_budget_is_shared_by_image_and_audio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", 24) + budget = MediaDecodedBudget() + image_bytes = BytesIO() + Image.new("RGB", (2, 3)).save(image_bytes, format="PNG") + ImageMediaIO(decoded_budget=budget).load_bytes(image_bytes.getvalue()) + + with ( + patch("vllm.multimodal.media.audio.load_audio") as decode, + pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES"), + ): + AudioMediaIO(decoded_budget=budget).load_bytes(b"audio") + + decode.assert_not_called() + + def test_audio_load_file_rejects_oversized( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): @@ -138,12 +182,16 @@ async def test_audio_connector_threads_byte_limit_to_http_reader( monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1) + budget = MediaRequestDownloadBudget() class _Connection: async def async_get_bytes(self, _url: str, **kwargs): assert kwargs["max_bytes"] == MiB_bytes + assert kwargs["request_budget"] is budget raise VLLMValidationError("Maximum file size exceeded") - connector = MediaConnector(connection=_Connection()) + connector = MediaConnector( # type: ignore[arg-type] + connection=_Connection(), download_budget=budget + ) with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"): await connector.fetch_audio_async("https://example.test/audio") diff --git a/tests/test_connections.py b/tests/test_connections.py index 983e20f6c12a..6385e372e45e 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -9,7 +9,13 @@ import pytest import vllm.envs as envs -from vllm.connections import HTTPConnection, HTTPResponseSizeExceededError +from vllm.connections import ( + HTTPConnection, + HTTPResponseSizeExceededError, + MediaRequestDownloadBudget, +) + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] _ONE_MIB = 1024 * 1024 _SMALL_BODY = b"a" * (_ONE_MIB // 2) @@ -150,6 +156,49 @@ async def test_async_get_bytes_with_limit_streams_instead_of_materializing() -> assert response.content.iterated == 2 +def test_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _SyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _SyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object(connection, "get_response", side_effect=[first, second]): + connection.get_bytes("http://example.com/one", request_budget=budget) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + connection.get_bytes("http://example.com/two", request_budget=budget) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.iterated == 1 + + +@pytest.mark.asyncio +async def test_async_request_download_budget_counts_across_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + first = _AsyncResponse([b"a" * (3 * _ONE_MIB // 4)]) + second = _AsyncResponse([b"b" * (_ONE_MIB // 2)]) + + with patch.object( + connection, "get_async_response", new=AsyncMock(side_effect=[first, second]) + ): + await connection.async_get_bytes( + "http://example.com/one", request_budget=budget + ) + with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB"): + await connection.async_get_bytes( + "http://example.com/two", request_budget=budget + ) + + assert budget.used == 3 * _ONE_MIB // 4 + assert second.content.iterated == 1 + + def test_get_bytes_rejects_content_length_over_limit( monkeypatch: pytest.MonkeyPatch, local_media_server: str ) -> None: diff --git a/vllm/connections.py b/vllm/connections.py index b06d2d563fa3..ded8da430f64 100644 --- a/vllm/connections.py +++ b/vllm/connections.py @@ -3,6 +3,7 @@ import asyncio import functools +import threading import time from collections.abc import Callable, Coroutine, Mapping, MutableMapping from pathlib import Path @@ -60,6 +61,33 @@ def __init__(self, max_bytes: int, received_bytes: int | None = None) -> None: super().__init__(message) +class MediaRequestDownloadSizeExceededError(ValueError): + """Raised when all HTTP media downloads in one request exceed the limit.""" + + +class MediaRequestDownloadBudget: + """Thread-safe total of HTTP media bytes received for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._lock = threading.Lock() + + def charge(self, size: int) -> None: + limit_mb = envs.VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB + if limit_mb <= 0: + return + limit = limit_mb * MiB_bytes + with self._lock: + received = self.used + size + if received > limit: + raise MediaRequestDownloadSizeExceededError( + "Request media downloads exceed " + f"VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB={limit_mb:g} " + f"({received} bytes received)." + ) + self.used = received + + def _get_media_download_limit_bytes() -> int | None: max_size_mb = envs.VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB if max_size_mb <= 0: @@ -105,19 +133,26 @@ def _append_chunk( received_bytes: int, max_bytes: int | None, media_download_limit: bool, + request_budget: MediaRequestDownloadBudget | None = None, ) -> int: if not chunk: return received_bytes received_bytes += len(chunk) if max_bytes is not None and received_bytes > max_bytes: _raise_size_exceeded(max_bytes, received_bytes, media_download_limit) + if request_budget is not None: + request_budget.charge(len(chunk)) chunks.append(chunk) return received_bytes -def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> bytes: +def _read_response_bytes( + response: requests.Response, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, +) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return response.content _check_content_length( @@ -129,16 +164,23 @@ def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> received_bytes = 0 for chunk in response.iter_content(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) async def _async_read_response_bytes( - response: aiohttp.ClientResponse, max_bytes: int | None + response: aiohttp.ClientResponse, + max_bytes: int | None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: max_bytes, media_download_limit = _get_response_size_limit(max_bytes) - if max_bytes is None: + if max_bytes is None and request_budget is None: return await response.read() _check_content_length(response.content_length, max_bytes, media_download_limit) @@ -146,7 +188,12 @@ async def _async_read_response_bytes( received_bytes = 0 async for chunk in response.content.iter_chunked(_RESPONSE_READ_CHUNK_SIZE): received_bytes = _append_chunk( - chunks, chunk, received_bytes, max_bytes, media_download_limit + chunks, + chunk, + received_bytes, + max_bytes, + media_download_limit, + request_budget, ) return b"".join(chunks) @@ -407,6 +454,7 @@ def get_bytes( timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: with self.get_response( url, @@ -416,7 +464,7 @@ def get_bytes( ) as r: r.raise_for_status() - return _read_response_bytes(r, max_bytes) + return _read_response_bytes(r, max_bytes, request_budget) @_async_retry async def async_get_bytes( @@ -426,6 +474,7 @@ async def async_get_bytes( timeout: float | None = None, allow_redirects: bool = True, max_bytes: int | None = None, + request_budget: MediaRequestDownloadBudget | None = None, ) -> bytes: async with await self.get_async_response( url, @@ -434,7 +483,7 @@ async def async_get_bytes( ) as r: r.raise_for_status() - return await _async_read_response_bytes(r, max_bytes) + return await _async_read_response_bytes(r, max_bytes, request_budget) def get_text(self, url: str, *, timeout: float | None = None) -> str: with self.get_response(url, timeout=timeout) as r: diff --git a/vllm/envs.py b/vllm/envs.py index bb1c03a9dabd..dbbc49ff59a8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -78,6 +78,7 @@ VLLM_MEDIA_CACHE_TTL_HOURS: float = 24 VLLM_MEDIA_FETCH_MAX_RETRIES: int = 3 VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB: int = 256 + VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB: int = 256 VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 @@ -86,6 +87,7 @@ VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) VLLM_MAX_EMBED_DECODE_BYTES: int = 2_147_483_648 VLLM_MAX_IMAGE_PIXELS: int = 178_956_970 + VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES: int = 715_827_880 VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" VLLM_TARGET_DEVICE: str = "cuda" @@ -1063,6 +1065,15 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_MAX_IMAGE_PIXELS": lambda: int( os.getenv("VLLM_MAX_IMAGE_PIXELS", "178956970") ), + # Maximum aggregate decoded bytes across media in one request. + # The default allows one image at the per-image pixel limit at 4 bytes + # per pixel. Set to 0 to disable. + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", "715827880") + ), + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB": lambda: int( + os.getenv("VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "256") + ), # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. @@ -2412,6 +2423,8 @@ def compile_factors() -> dict[str, object]: "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", "VLLM_MAX_EMBED_DECODE_BYTES", "VLLM_MAX_IMAGE_PIXELS", + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES", + "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", "VLLM_VIDEO_LOADER_BACKEND", "VLLM_MEDIA_CONNECTOR", "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", diff --git a/vllm/multimodal/media/__init__.py b/vllm/multimodal/media/__init__.py index 94b08c484882..84fc7d6ed88c 100644 --- a/vllm/multimodal/media/__init__.py +++ b/vllm/multimodal/media/__init__.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .connector import MEDIA_CONNECTOR_REGISTRY, MediaConnector from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VIDEO_LOADER_REGISTRY, VideoMediaIO __all__ = [ "MediaIO", + "MediaDecodedBudget", "MediaWithBytes", "AudioEmbeddingMediaIO", "AudioMediaIO", diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index c05f4e4dcc8f..d36b4a984e66 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -23,7 +23,7 @@ safe_to_dense, ) -from .base import MediaIO +from .base import MediaDecodedBudget, MediaIO logger = init_logger(__name__) @@ -527,8 +527,11 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]): error handling. """ - def __init__(self, **kwargs) -> None: + def __init__( + self, *, decoded_budget: MediaDecodedBudget | None = None, **kwargs + ) -> None: super().__init__() + self.decoded_budget = decoded_budget # Allow per-server override of the audio decoding backend, e.g.: # --media-io-kwargs '{"audio": {"audio_backend": "torchcodec"}}' @@ -564,13 +567,44 @@ def _validate_encoded_size(self, size: int) -> None: def load_bytes(self, data: bytes) -> tuple[npt.NDArray, float]: self._validate_encoded_size(len(data)) - return load_audio( - BytesIO(data), - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(BytesIO(data)) + + def _load_with_budget(self, source: BytesIO | Path) -> tuple[npt.NDArray, float]: + budget = self.decoded_budget + if budget is None: + return load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, + backend=self.audio_backend, + ) + + max_decode_bytes = envs.VLLM_MAX_AUDIO_DECODE_BYTES + reserved = budget.reserve(max_decode_bytes) + decoded_bytes = 0 + try: + try: + audio = load_audio( + source, + sr=None, + max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, + max_decode_bytes=reserved, + backend=self.audio_backend, + ) + except ValueError as exc: + if reserved < max_decode_bytes and "VLLM_MAX_AUDIO_DECODE_BYTES" in str( + exc + ): + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) from exc + raise + decoded_bytes = audio[0].nbytes + return audio + finally: + budget.settle(reserved, decoded_bytes) def load_base64( self, @@ -588,13 +622,7 @@ def load_base64( def load_file(self, filepath: Path) -> tuple[npt.NDArray, float]: self._validate_encoded_size(filepath.stat().st_size) - return load_audio( - filepath, - sr=None, - max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, - max_decode_bytes=envs.VLLM_MAX_AUDIO_DECODE_BYTES, - backend=self.audio_backend, - ) + return self._load_with_budget(filepath) def encode_base64( self, diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index d3e622156fc8..952b92e96939 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading from abc import ABC, abstractmethod from collections.abc import Iterable, Iterator from dataclasses import dataclass, field @@ -9,9 +10,79 @@ import numpy as np +import vllm.envs as envs + _T = TypeVar("_T") +class MediaDecodedBudget: + """Thread-safe decoded-media byte accumulator for one request.""" + + def __init__(self) -> None: + self.used = 0 + self._reserved = 0 + self._condition = threading.Condition() + + def reserve(self, max_bytes: int) -> int: + """Reserve decoder capacity, waiting for in-flight decodes if needed.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return max_bytes + max_bytes = max(0, max_bytes) + with self._condition: + while True: + available = limit - self.used - self._reserved + if available >= max_bytes or self._reserved == 0: + if available <= 0: + raise ValueError( + "Request decoded media exceeds " + "VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES." + ) + reserved = min(max_bytes, available) + self._reserved += reserved + return reserved + self._condition.wait() + + def settle(self, reserved: int, decoded_bytes: int) -> None: + """Replace a decoder reservation with its actual decoded size.""" + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0: + return + with self._condition: + self._reserved -= reserved + try: + used = self.used + decoded_bytes + if decoded_bytes > reserved or used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + finally: + self._condition.notify_all() + + def charge(self, decoded_bytes: int) -> None: + limit = envs.VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES + if limit <= 0 or decoded_bytes <= 0: + return + with self._condition: + while ( + self.used + self._reserved + decoded_bytes > limit + and self._reserved > 0 + and self.used + decoded_bytes <= limit + ): + self._condition.wait() + used = self.used + decoded_bytes + if used > limit: + raise ValueError( + "Request decoded media exceeds " + f"VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES={limit} " + f"({used} bytes)." + ) + self.used = used + + @dataclass class MediaWithBytes(Generic[_T]): """Wrapper that couples a media object with its original encoded bytes. diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 297dc91af232..e5a52417fee3 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -25,6 +25,8 @@ from vllm.connections import ( HTTPConnection, MediaDownloadSizeExceededError, + MediaRequestDownloadBudget, + MediaRequestDownloadSizeExceededError, global_http_connection, ) from vllm.exceptions import VLLMUnprocessableEntityError, VLLMValidationError @@ -33,7 +35,7 @@ from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VideoEmbeddingMediaIO, VideoMediaIO @@ -109,7 +111,9 @@ def _wrap_media_fetch_error( value=url, ) - if isinstance(exc, MediaDownloadSizeExceededError): + if isinstance( + exc, (MediaDownloadSizeExceededError, MediaRequestDownloadSizeExceededError) + ): return VLLMUnprocessableEntityError( f"Failed to fetch media from URL: {exc}", parameter="image_url", @@ -162,6 +166,8 @@ def __init__( *, allowed_local_media_path: str = "", allowed_media_domains: list[str] | None = None, + decoded_budget: MediaDecodedBudget | None = None, + download_budget: MediaRequestDownloadBudget | None = None, ) -> None: """Args: media_io_kwargs: Additional args passed to process media @@ -201,6 +207,8 @@ def __init__( if allowed_media_domains is None: allowed_media_domains = [] self.allowed_media_domains = allowed_media_domains + self.decoded_budget = decoded_budget + self.download_budget = download_budget # Media download cache (opt-in via VLLM_MEDIA_CACHE) self._media_cache_dir: str | None = None @@ -230,6 +238,12 @@ def __init__( media_cache, ) + def _image_io(self, image_mode: str | None = "RGB") -> ImageMediaIO: + return ImageMediaIO( + **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})), + decoded_budget=self.decoded_budget, + ) + def _get_cached_bytes(self, url: str) -> bytes | None: """Return cached bytes for a URL, or None if not cached/expired.""" if not self._media_cache_dir: @@ -385,12 +399,18 @@ def load_from_url( return media_io.load_bytes(cached) connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = connection.get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -438,12 +458,18 @@ async def load_from_url_async( return await future connection = self.connection + budget_kwargs = ( + {"request_budget": self.download_budget} + if self.download_budget is not None + else {} + ) try: data = await connection.async_get_bytes( url_spec.url, timeout=fetch_timeout, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, max_bytes=max_bytes, + **budget_kwargs, ) except Exception as e: wrapped = _wrap_media_fetch_error(url, e) @@ -470,7 +496,9 @@ def fetch_audio( audio_url: str, ) -> tuple[np.ndarray, int | float]: """Load audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return self.load_from_url( audio_url, @@ -483,7 +511,9 @@ async def fetch_audio_async( audio_url: str, ) -> tuple[np.ndarray, int | float]: """Asynchronously fetch audio from a URL.""" - audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {})) + audio_io = AudioMediaIO( + decoded_budget=self.decoded_budget, **self.media_io_kwargs.get("audio", {}) + ) return await self.load_from_url_async( audio_url, @@ -503,9 +533,7 @@ def fetch_image( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return self.load_from_url( @@ -529,9 +557,7 @@ async def fetch_image_async( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) try: return await self.load_from_url_async( @@ -551,9 +577,7 @@ def fetch_video( video_processor: str | None = None, ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """Load video from an HTTP or base64 data URL.""" - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) @@ -580,9 +604,7 @@ async def fetch_video_async( `media_io_kwargs={"image": {"image_mode": None}}` to keep the original image mode (e.g. preserving the alpha channel). """ - image_io = ImageMediaIO( - **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {})) - ) + image_io = self._image_io(image_mode) video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) if "video_backend" not in video_io_kwargs and ( video_backend := get_video_loader_backend_for_processor(video_processor) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index fb05f4ed741d..1e2e106189b4 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -17,7 +17,7 @@ ) from ..image import convert_image_mode, normalize_image, open_image, rgba_to_rgb -from .base import MediaIO, MediaWithBytes +from .base import MediaDecodedBudget, MediaIO, MediaWithBytes MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0 @@ -28,12 +28,19 @@ class ImageMediaIO(MediaIO[Image.Image]): error handling. """ - def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None: + def __init__( + self, + image_mode: str | None = "RGB", + *, + decoded_budget: MediaDecodedBudget | None = None, + **kwargs, + ) -> None: super().__init__() # Target mode for loaded images; `None` keeps the original mode # (i.e. no conversion, alpha channel is preserved as-is). self.image_mode = image_mode + self.decoded_budget = decoded_budget # `kwargs` contains custom arguments from # --media-io-kwargs for this modality, merged with # per-request runtime media_io_kwargs via merge_kwargs(). @@ -87,6 +94,8 @@ def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]: f"the maximum of {max_pixels} pixels. Set " f"VLLM_MAX_IMAGE_PIXELS to increase this limit." ) + if self.decoded_budget is not None: + self.decoded_budget.charge(w * h * 4) image = normalize_image(image) image.load() converted = self._convert_image_mode(image) diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 7528331cdfd0..569ef3dc93bf 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -122,6 +122,16 @@ def load_bytes( video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) + frames = video[0] if isinstance(video, tuple) else video + budget = getattr(self.image_io, "decoded_budget", None) + shape = getattr(frames, "shape", None) + if budget is not None and shape is not None and len(shape) >= 3: + decoded_bytes = ( + frames.numel() * frames.element_size() + if isinstance(frames, torch.Tensor) + else frames.nbytes + ) + budget.charge(decoded_bytes) return MediaWithBytes(video, data) def load_base64( diff --git a/vllm/renderers/chat_utils.py b/vllm/renderers/chat_utils.py index f18407e535e9..540564baab8d 100644 --- a/vllm/renderers/chat_utils.py +++ b/vllm/renderers/chat_utils.py @@ -53,6 +53,7 @@ from vllm import envs from vllm.config import ModelConfig +from vllm.connections import MediaRequestDownloadBudget from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from vllm.logger import init_logger @@ -67,6 +68,7 @@ VisionChunkVideo, ) from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector +from vllm.multimodal.media.base import MediaDecodedBudget from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.embed_utils import ( safe_load_prompt_embeds, @@ -609,6 +611,8 @@ def __init__( self._items_by_modality = defaultdict[str, list[_T]](list) # Track original modality for each vision_chunk item (image or video) self._modality_order = defaultdict[str, list[str]](list) + self.decoded_budget = MediaDecodedBudget() + self.download_budget = MediaRequestDownloadBudget() @cached_property def use_unified_vision_chunk_modality(self) -> bool: @@ -1052,6 +1056,8 @@ def _connector(self) -> MediaConnector: media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1138,6 +1144,10 @@ def parse_audio_embeds( def parse_image_pil( self, image_pil: Image.Image | None, uuid: str | None = None ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add("image", (image_pil, uuid)) self._add_placeholder("image", placeholder) @@ -1233,6 +1243,8 @@ def _connector(self) -> MediaConnector: media_io_kwargs=self._tracker.media_io_kwargs, allowed_local_media_path=self._tracker.allowed_local_media_path, allowed_media_domains=self._tracker.allowed_media_domains, + decoded_budget=self._tracker.decoded_budget, + download_budget=self._tracker.download_budget, ) @property @@ -1351,6 +1363,10 @@ def parse_image_pil( image_pil: Image.Image | None, uuid: str | None = None, ) -> None: + if image_pil is not None: + self._tracker.decoded_budget.charge( + image_pil.size[0] * image_pil.size[1] * 4 + ) placeholder = self._tracker.add( "image", partial(self._item_with_uuid_async, image_pil, uuid) ) From 54db759eea66729faa6d41e8cdaceb02833cef7c Mon Sep 17 00:00:00 2001 From: samayp01 <67844360+samayp01@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:27:34 -0400 Subject: [PATCH 486/952] [Bugfix][Structured Output] Reject structural_tag and unhandled specs in backend validation (#58862) Signed-off-by: Samay Co-authored-by: Artem Perevedentsev --- tests/v1/structured_output/test_validation.py | 57 +++++++++++++++++++ .../backend_lm_format_enforcer.py | 16 +++++- vllm/v1/structured_output/backend_outlines.py | 14 ++++- vllm/v1/structured_output/backend_xgrammar.py | 14 ++++- 4 files changed, 94 insertions(+), 7 deletions(-) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 276056b501cd..41166c260280 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -17,6 +17,13 @@ ) from vllm.tokenizers import mistral as mistral_tokenizers from vllm.v1.structured_output import backend_guidance +from vllm.v1.structured_output.backend_lm_format_enforcer import ( + validate_structured_output_request_lm_format_enforcer, +) +from vllm.v1.structured_output.backend_outlines import ( + validate_structured_output_request_outlines, +) +from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar pytestmark = pytest.mark.cpu_test @@ -213,6 +220,9 @@ def test_schema_at_nesting_limit_accepted(field): INVALID_JSON_SCHEMA = {"type": "object", "properties": {"name": {"type": "str"}}} +STRUCTURAL_TAG = ( + '{"type": "structural_tag", "format": {"type": "const_string", "value": "hi"}}' +) @pytest.mark.parametrize( @@ -227,6 +237,9 @@ def test_schema_at_nesting_limit_accepted(field): ("lm-format-enforcer", StructuredOutputsParams(grammar="not a grammar")), ("outlines", StructuredOutputsParams(regex="(")), ("guidance", StructuredOutputsParams(structural_tag='{"nope": 1}')), + ("outlines", StructuredOutputsParams(structural_tag=STRUCTURAL_TAG)), + ("lm-format-enforcer", StructuredOutputsParams(structural_tag=STRUCTURAL_TAG)), + ("outlines", StructuredOutputsParams(json_object=True)), ], ) def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): @@ -241,6 +254,50 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): ) +@pytest.mark.parametrize( + "backend, structured_outputs", + [ + ("xgrammar", StructuredOutputsParams(json_object=True)), + ("xgrammar", StructuredOutputsParams(regex="[0-9]+")), + ("xgrammar", StructuredOutputsParams(regex="")), + ("xgrammar", StructuredOutputsParams(json={})), + ("lm-format-enforcer", StructuredOutputsParams(json_object=True)), + ("lm-format-enforcer", StructuredOutputsParams(regex="")), + ("lm-format-enforcer", StructuredOutputsParams(json={})), + ("outlines", StructuredOutputsParams(regex="")), + ("outlines", StructuredOutputsParams(json={})), + ], +) +def test_supported_spec_is_not_rejected(backend, structured_outputs): + """Supported specs, including the falsy but valid regex="" and json={}, must + not reach the validators' final rejection.""" + params = SamplingParams(structured_outputs=structured_outputs) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend=backend), + tokenizer=object(), + ) + + +@pytest.mark.parametrize( + "validate", + [ + validate_structured_output_request_outlines, + validate_structured_output_request_lm_format_enforcer, + validate_xgrammar_grammar, + ], + ids=["outlines", "lm-format-enforcer", "xgrammar"], +) +def test_unhandled_spec_is_a_validation_error(validate): + """json_object=False is rejected before backend selection, so only a direct + validator call reaches the final branch; it must still be a 400.""" + params = SamplingParams( + structured_outputs=StructuredOutputsParams(json_object=False) + ) + with pytest.raises(VLLMValidationError, match="does not support this"): + validate(params) + + @pytest.mark.parametrize( "schema, expected_backend", [ diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 9a16ddc1fe38..b48cdc791673 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -159,7 +159,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams so_params = params.structured_outputs - if so_params.regex: + if so_params.regex is not None: try: compile_regex_with_timeout( lmformatenforcer.RegexParser, @@ -170,7 +170,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams f"Failed to compile regex for lm-format-enforcer: {err}" ) from err return - elif so_params.json: + elif so_params.json is not None: if isinstance(so_params.json, str): try: # make sure schema is valid json @@ -185,10 +185,20 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams f"Error serializing structured outputs jsonschema: {e}" ) from e return - elif so_params.choice: + elif so_params.choice or so_params.json_object: return elif so_params.grammar: raise VLLMValidationError( "LM Format Enforcer structured outputs backend " "does not support grammar specifications" ) + elif so_params.structural_tag: + raise VLLMValidationError( + "LM Format Enforcer structured outputs backend " + "does not support structural_tag specifications" + ) + else: + raise VLLMValidationError( + "LM Format Enforcer structured outputs backend " + "does not support this structured outputs specification" + ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 176d1a46e684..09588a680ab9 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -198,9 +198,9 @@ def validate_structured_output_request_outlines(params: SamplingParams): so_params = params.structured_outputs - if so_params.regex: + if so_params.regex is not None: validate_regex_is_buildable(so_params.regex) - elif so_params.json: + elif so_params.json is not None: if isinstance(so_params.json, str): try: # make sure schema is valid json @@ -236,6 +236,16 @@ def validate_structured_output_request_outlines(params: SamplingParams): "Outlines structured outputs backend " "does not support grammar specifications" ) + elif so_params.structural_tag: + raise VLLMValidationError( + "Outlines structured outputs backend " + "does not support structural_tag specifications" + ) + else: + raise VLLMValidationError( + "Outlines structured outputs backend " + "does not support this structured outputs specification" + ) def _prefix_needs_context(parsed) -> bool: diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 0199c5f75d6f..9631766d9267 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -414,7 +414,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params = sampling_params.structured_outputs - if so_params.regex: + if so_params.regex is not None: # A NUL byte is never meaningful in a regex pattern and is not handled # by xgrammar's native regex converter. Reject it here, before the # pattern reaches that native code; the try/except below does not cover @@ -432,6 +432,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: raise VLLMValidationError( f"Failed to transform regex into a grammar: {err}" ) from err + return if so_params.choice: choice_grammar = choice_as_grammar(so_params.choice) @@ -445,7 +446,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params.grammar = choice_grammar return - if so_params.json: + if so_params.json is not None: check_json_nesting(so_params.json) if isinstance(so_params.json, str): try: @@ -470,6 +471,9 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: ) from err return + if so_params.json_object: + return + if so_params.grammar: # Parse the grammar with the same syntax `compile_grammar` will use, # but don't compile it. The grammar is passed on unchanged. @@ -502,3 +506,9 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: xgr.Grammar.from_structural_tag(so_params.structural_tag) except Exception as e: raise VLLMValidationError("Invalid structural tag specification.") from e + return + + raise VLLMValidationError( + "XGrammar structured outputs backend " + "does not support this structured outputs specification" + ) From 426f3bd36b941c34ad183e10cf2c30ebd62a4755 Mon Sep 17 00:00:00 2001 From: noone-zhhh Date: Tue, 13 Oct 2026 18:49:39 +0800 Subject: [PATCH 487/952] [Bugfix][Model][Rust Frontend] Qwen3_ASR: Handle ragged Qwen audio batches from Rust frontend (#60979) Signed-off-by: zh98 Co-authored-by: Isotr0py --- vllm/model_executor/models/qwen2_5_omni_thinker.py | 12 +++++++----- vllm/model_executor/models/qwen3_asr.py | 6 ++++-- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 986de3e09636..c09ae1a87e79 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -132,11 +132,11 @@ def _get_feature_tensor(feature: MultiModalFeatureSpec, key: str) -> torch.Tenso def unpad_and_flat_audio_features( - input_audio_features: torch.Tensor, + input_audio_features: torch.Tensor | Sequence[torch.Tensor], audio_feature_lengths: torch.Tensor, ) -> torch.Tensor: - """Unpad and flatten batched audio features.""" - if len(audio_feature_lengths) != input_audio_features.size(0): + """Unpad and flatten stacked or ragged per-audio features.""" + if len(audio_feature_lengths) != len(input_audio_features): raise ValueError( "Length of audio_feature_lengths must match " "the batch size of input_audio_features." @@ -920,8 +920,10 @@ def _parse_and_validate_audio_input( # inputs features from rust frontend is batched and padded # with shape [batch_size, n_mels, padded_seq_len], different - # from python's shape [n_mels, batch_size * seq_len] - if ( + # from python's shape [n_mels, batch_size * seq_len]. + # Requests with different padded lengths are batched as a + # list of [n_mels, padded_seq_len] tensors. + if isinstance(input_audio_features, (list, tuple)) or ( isinstance(input_audio_features, torch.Tensor) and input_audio_features.dim() == 3 ): diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index 161d1efacd06..0283fd814d8d 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -478,8 +478,10 @@ def _parse_and_validate_audio_input( # inputs features from rust frontend is batched and padded # with shape [batch_size, n_mels, padded_seq_len], different - # from python's shape [n_mels, batch_size * seq_len] - if ( + # from python's shape [n_mels, batch_size * seq_len]. + # Requests with different padded lengths are batched as a + # list of [n_mels, padded_seq_len] tensors. + if isinstance(input_audio_features, (list, tuple)) or ( isinstance(input_audio_features, torch.Tensor) and input_audio_features.dim() == 3 ): From e6e84f350c647c6e76dad2e39b08cb85e17f1272 Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Tue, 13 Oct 2026 04:16:21 -0700 Subject: [PATCH 488/952] [Feature][Pooling] Add offline fixed-length late chunking (#60164) Signed-off-by: Linze-Shi --- docs/models/pooling_models/token_embed.md | 34 ++ .../token_embed/late_chunking_offline.py | 72 ++++ .../entrypoints/pooling/test_late_chunking.py | 314 ++++++++++++++++++ .../pooling/token_embed/test_offline.py | 221 +++++++++++- .../layers/test_pooler_heads.py | 1 + .../layers/test_pooler_methods.py | 143 +++++++- tests/test_pooling_params.py | 92 +++++ .../entrypoints/pooling/embed/io_processor.py | 74 +++++ vllm/entrypoints/pooling/late_chunking.py | 92 +++++ vllm/entrypoints/pooling/offline.py | 2 +- vllm/entrypoints/pooling/typing.py | 15 +- vllm/model_executor/layers/pooler/special.py | 5 + .../layers/pooler/tokwise/heads.py | 23 ++ vllm/outputs.py | 17 + vllm/pooling_params.py | 49 ++- 15 files changed, 1147 insertions(+), 7 deletions(-) create mode 100644 examples/pooling/token_embed/late_chunking_offline.py create mode 100644 tests/entrypoints/pooling/test_late_chunking.py create mode 100644 vllm/entrypoints/pooling/late_chunking.py diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 0c2a322e80fc..2bc51f4bec91 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -90,6 +90,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported. ```python --8<-- "vllm/pooling_params.py:common-pooling-params" --8<-- "vllm/pooling_params.py:embed-pooling-params" +--8<-- "vllm/pooling_params.py:token-embed-pooling-params" ``` ### `LLM.encode` @@ -108,6 +109,39 @@ data = output.outputs.data print(f"Data: {data!r}") ``` +### Fixed-length late chunking + +Late chunking encodes the full input before averaging contextual token representations within each chunk. The chunk means pass through the model's embedding head, including its projection and optional normalization. + +This mode is currently available through the offline `LLM.encode` API for plain-text input. With an instance configured for `token_embed`, enable it per request: + +```python +from vllm import PoolingParams +from vllm.pooling_params import LateChunkingParams + +outputs = llm.encode( + [document], + pooling_task="token_embed", + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=256) + ), +) +vectors = outputs[0].outputs.data # [num_chunks, embedding_dimension] +chunks = outputs[0].late_chunking.chunks +``` + +For `N` model-input tokens and chunk size `C`, the output contains `ceil(N / C)` vectors. The last chunk uses its actual length. Special tokens and instruction prefixes count toward both chunk boundaries and means. Omitting `late_chunking_params` preserves ordinary token embeddings. + +Each vector row corresponds to one entry in `chunks`. Its `token_range` and `char_range` are zero-based, half-open ranges. Character ranges index the submitted Python string; use `document[start:end]` to retrieve the source text. A chunk containing only special tokens has `char_range=None`. + +Use nonempty text within the model's context window and a fast tokenizer that provides aligned source offsets. Input truncation, explicit padding and renderer text normalization are not supported. The custom document/query outputs of `JinaForRanking` are also unsupported. Late chunking does not extend the model's context window or reduce its attention computation. + +The [offline example](../../../examples/pooling/token_embed/late_chunking_offline.py) loads an E5 model and prints the vector shape, token ranges and source excerpts for two documents: + +```bash +python examples/pooling/token_embed/late_chunking_offline.py --chunk-size 16 +``` + ### `LLM.score` The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs. diff --git a/examples/pooling/token_embed/late_chunking_offline.py b/examples/pooling/token_embed/late_chunking_offline.py new file mode 100644 index 000000000000..4f82b17cceff --- /dev/null +++ b/examples/pooling/token_embed/late_chunking_offline.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from argparse import Namespace + +from vllm import LLM, EngineArgs, PoolingParams +from vllm.config import PoolerConfig +from vllm.pooling_params import LateChunkingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def parse_args(): + parser = FlexibleArgumentParser() + parser = EngineArgs.add_cli_args(parser) + parser.add_argument( + "--chunk-size", + type=int, + default=16, + help="Number of model-input tokens per chunk, including special tokens.", + ) + parser.set_defaults( + model="intfloat/multilingual-e5-small", + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + max_model_len=512, + enforce_eager=True, + ) + return parser.parse_args() + + +def main(args: Namespace): + pooling_params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=args.chunk_size) + ) + engine_args = vars(args).copy() + del engine_args["chunk_size"] + llm = LLM(**engine_args) + + # E5 document inputs use the "passage: " prefix; it also counts toward chunks. + documents = [ + "passage: Berlin is the capital of Germany. The city has many museums. " + "They display art and historical collections from around the world.", + "passage: Kyoto is a city in Japan. Its historic temples attract visitors " + "throughout the year.", + ] + outputs = llm.encode( + documents, + pooling_task="token_embed", + pooling_params=pooling_params, + ) + + for doc_index, (document, output) in enumerate(zip(documents, outputs)): + vectors = output.outputs.data + metadata = output.late_chunking + assert metadata is not None + print(f"\nDocument {doc_index}: {metadata.input_tokens} input tokens") + print(f"Chunk vectors: {tuple(vectors.shape)}") + for index, chunk in enumerate(metadata.chunks): + # vectors[index] corresponds to this chunk and its source-text range. + if chunk.char_range is None: + excerpt = "" + else: + start, end = chunk.char_range + excerpt = document[start:end] + print( + f" Chunk {index}: tokens={chunk.token_range}, chars={chunk.char_range}" + ) + print(f" {excerpt!r}") + + +if __name__ == "__main__": + main(parse_args()) diff --git a/tests/entrypoints/pooling/test_late_chunking.py b/tests/entrypoints/pooling/test_late_chunking.py new file mode 100644 index 000000000000..eaec62ac196c --- /dev/null +++ b/tests/entrypoints/pooling/test_late_chunking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest import mock +from unittest.mock import Mock + +import pytest +import torch + +from vllm import PoolingParams +from vllm.config import PoolerConfig +from vllm.entrypoints.pooling.embed.io_processor import ( + JinaRankingTokenEmbedIOProcessor, + TokenEmbedIOProcessor, +) +from vllm.entrypoints.pooling.late_chunking import build_late_chunking_metadata +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.pooling.typing import ( + OfflineEncodeInputsContext, + OfflineOutputsContext, +) +from vllm.exceptions import VLLMValidationError +from vllm.outputs import ( + LateChunk, + PoolingOutput, + PoolingRequestOutput, + RequestError, +) +from vllm.pooling_params import LateChunkingParams +from vllm.renderers import TokenizeParams + + +@pytest.fixture +def late_chunk_processor(): + processor = object.__new__(TokenEmbedIOProcessor) + processor.model_config = SimpleNamespace( + is_encoder_decoder=False, + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + processor.vllm_config = SimpleNamespace( + model_config=processor.model_config, + cache_config=SimpleNamespace(enable_prefix_caching=False), + scheduler_config=SimpleNamespace(enable_chunked_prefill=False), + lora_config=None, + ) + processor.renderer = Mock() + processor.renderer.default_cmpl_tok_params = TokenizeParams(max_total_tokens=32) + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 1), (2, 3), (0, 0)], + } + ] + return processor + + +def _late_chunk_context(prompts="a b", **kwargs): + return OfflineEncodeInputsContext( + pooling_task="token_embed", + tokenization_kwargs=kwargs or None, + lora_request=None, + priorities=None, + prompts=prompts, + pooling_params=PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size=2) + ), + ) + + +def test_late_chunk_render_requests_offsets_once(late_chunk_processor): + processor = late_chunk_processor + ctx = _late_chunk_context() + factory, count = processor.get_request_factory_offline(ctx) + render_params = next(factory()) + result = processor.render(render_params) + assert count == 1 + processor.renderer.render_cmpl.assert_called_once() + assert processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets + assert not render_params["tok_params"].return_token_offsets + assert "prompt_token_offsets" not in result["prompts"] + assert "late_chunking" not in result + assert [c.char_range for c in ctx.late_chunking[0].metadata.chunks] == [ + (0, 1), + (2, 3), + ] + assert ctx.late_chunking[0].metadata.input_tokens == 4 + + +@pytest.mark.parametrize("as_sequence", [False, True]) +@pytest.mark.parametrize("chunk_size", [None, 2]) +def test_jina_ranking_rejects_chunking_before_tokenization( + late_chunk_processor, as_sequence, chunk_size +): + processor = object.__new__(JinaRankingTokenEmbedIOProcessor) + processor.__dict__.update(late_chunk_processor.__dict__) + ctx = _late_chunk_context(prompts=["document", "query"]) + params = PoolingParams( + late_chunking_params=LateChunkingParams(chunk_size) if chunk_size else None + ) + ctx.pooling_params = [params] if as_sequence else params + factory, count = processor.get_request_factory_offline(ctx) + assert count == 1 + request = next(factory()) + if chunk_size is not None: + with pytest.raises(VLLMValidationError, match="JinaForRanking.*late chunking"): + processor.render(request) + late_chunk_processor.renderer.render_cmpl.assert_not_called() + else: + result = processor.render(request) + assert result["prompts"]["type"] == "token" + assert result["prompts"]["prompt_token_ids"] == [101, 1, 2, 102] + late_chunk_processor.renderer.render_cmpl.assert_called_once() + assert not request["tok_params"].return_token_offsets + + +@pytest.mark.parametrize( + "prompts, kwargs, error", + [ + pytest.param("", {}, "plain-text", id="empty"), + pytest.param([1, 2], {}, "plain-text", id="token-ids"), + pytest.param( + {"prompt": "abc", "multi_modal_data": {}}, + {}, + "plain-text", + id="multimodal", + ), + ] + + [ + pytest.param("abc", {key: value}, "does not support", id=key) + for key, value in ( + ("truncate_prompt_tokens", 2), + ("truncation", True), + ("pad_prompt_tokens", 8), + ("padding", True), + ("do_lower_case", True), + ) + ], +) +def test_late_chunking_rejects_unsupported_input_before_render( + late_chunk_processor, prompts, kwargs, error +): + processor = late_chunk_processor + with pytest.raises(VLLMValidationError, match=error): + factory, _ = processor.get_request_factory_offline( + _late_chunk_context(prompts, **kwargs) + ) + processor.render(next(factory())) + processor.renderer.render_cmpl.assert_not_called() + + +def test_late_chunk_ranges_keep_unicode_overlaps_and_special_only_chunks(): + text = "中 😀 e\u0301" + offsets = [(0, 0), (0, 1), (2, 3), (2, 3), (4, 6), (0, 0)] + metadata = build_late_chunking_metadata(text, len(offsets), offsets, 1) + assert [c.char_range for c in metadata.chunks] == [ + None, + (0, 1), + (2, 3), + (2, 3), + (4, 6), + None, + ] + assert [ + text[slice(*c.char_range)] if c.char_range else None for c in metadata.chunks + ] == [None, "中", "😀", "😀", "e\u0301", None] + chunks = build_late_chunking_metadata(text, len(offsets), offsets, 4).chunks + assert [(c.token_range, c.char_range) for c in chunks] == [ + ((0, 4), (0, 3)), + ((4, 6), (4, 6)), + ] + + +@pytest.mark.parametrize( + "offsets", [None, [(0, 1)], [(-1, 1), (0, 0)], [(0, 4), (0, 0)], [(2, 3), (0, 1)]] +) +def test_late_chunk_ranges_reject_missing_or_invalid_offsets(offsets): + with pytest.raises(VLLMValidationError, match="offsets"): + build_late_chunking_metadata("abc", 2, offsets, 2) + + +def _mock_chunk_llm(processor, outputs): + llm = mock.Mock(spec=PoolingOfflineMixin) + llm.encode = PoolingOfflineMixin.encode.__get__(llm) + llm._run_tiling_engine = PoolingOfflineMixin._run_tiling_engine.__get__(llm) + llm.pooling_io_processors = {"token_embed": processor} + llm._executor = SimpleNamespace(map=map) + llm.llm_engine = mock.Mock() + llm.llm_engine.vllm_config.scheduler_config.max_num_seqs = 2 + llm.llm_engine.has_unfinished_requests.return_value = False + llm.llm_engine.step.side_effect = outputs + llm._render_and_add_requests = mock.Mock( + side_effect=lambda **kw: [ + str(i) + "-internal" for i in range(len(kw["params"])) + ] + ) + return llm + + +def _chunk_output(request_id, rows=2, **kwargs): + return PoolingRequestOutput( + str(request_id), + PoolingOutput(torch.ones(rows, 4)), + [101, 1, 2, 102], + 0, + True, + **kwargs, + ) + + +def test_late_chunk_mapping_preserves_order_mixed_requests_and_errors( + late_chunk_processor, +): + error = RequestError("test_error", "original error") + failed = _chunk_output(3, rows=0, error=error) + llm = _mock_chunk_llm(late_chunk_processor, RuntimeError("step failed")) + params = [ + PoolingParams(late_chunking_params=LateChunkingParams(2)), + PoolingParams(), + PoolingParams(late_chunking_params=LateChunkingParams(1)), + PoolingParams(late_chunking_params=LateChunkingParams(2)), + ] + # A failed encode call must not leave metadata for the following batch. + with pytest.raises(RuntimeError, match="step failed"): + llm.encode( + ["a b"] * 2, + pooling_task="token_embed", + pooling_params=params[:2], + use_tqdm=False, + ) + llm.llm_engine.step.side_effect = [ + [ + _chunk_output(2, rows=4), + failed, + _chunk_output(1, rows=4), + _chunk_output(0), + ] + ] + outputs = llm.encode( + ["a b"] * 4, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + assert outputs[0].late_chunking.chunks == [ + LateChunk((0, 2), (0, 1)), + LateChunk((2, 4), (2, 3)), + ] + assert outputs[1].late_chunking is None + assert len(outputs[2].late_chunking.chunks) == 4 + assert outputs[3].late_chunking is None + assert outputs[3].error is error + + +def test_late_chunk_mapping_rejects_successful_output_with_wrong_row_count( + late_chunk_processor, +): + llm = _mock_chunk_llm(late_chunk_processor, [[_chunk_output(0, rows=0)]]) + with pytest.raises(ValueError, match="does not match"): + llm.encode( + "a b", + pooling_task="token_embed", + pooling_params=PoolingParams(late_chunking_params=LateChunkingParams(2)), + use_tqdm=False, + ) + # All engine requests have finished before post-processing validates the output. + llm.llm_engine.abort_request.assert_not_called() + + +def test_late_chunk_contexts_remain_isolated_when_rendering_interleaves( + late_chunk_processor, +): + processor = late_chunk_processor + first, second = _late_chunk_context("xy z"), _late_chunk_context() + shared = PoolingParams(late_chunking_params=LateChunkingParams(2)) + first.pooling_params = second.pooling_params = shared + first_factory, _ = processor.get_request_factory_offline(first) + second_factory, _ = processor.get_request_factory_offline(second) + first_request, second_request = next(first_factory()), next(second_factory()) + processor.render(second_request) + assert first.late_chunking[0].metadata is None + processor.renderer.render_cmpl.side_effect = lambda **_: [ + { + "type": "token", + "prompt_token_ids": [101, 1, 2, 102], + "prompt_token_offsets": [(0, 0), (0, 2), (3, 4), (0, 0)], + } + ] + processor.render(first_request) + first_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=first.late_chunking) + )[0] + second_output = processor.post_process_offline( + OfflineOutputsContext([_chunk_output(0)], late_chunking=second.late_chunking) + )[0] + assert first_output.late_chunking.chunks[0].char_range == (0, 2) + assert second_output.late_chunking.chunks[0].char_range == (0, 1) + + +def test_plain_token_embed_skips_chunk_metadata(late_chunk_processor): + ctx = _late_chunk_context() + ctx.pooling_params = PoolingParams() + factory, _ = late_chunk_processor.get_request_factory_offline(ctx) + request = next(factory()) + result = late_chunk_processor.render(request) + assert not ctx.late_chunking + assert result["params"].late_chunking_params is None + assert not late_chunk_processor.renderer.render_cmpl.call_args.kwargs[ + "tok_params" + ].return_token_offsets diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index e0bda6542284..b65ef6ff8bed 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -3,9 +3,13 @@ import weakref import pytest +import torch -from vllm import LLM, PoolingRequestOutput +from tests.models.utils import check_embeddings_close +from vllm import LLM, PoolingParams, PoolingRequestOutput from vllm.config import PoolerConfig +from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -61,3 +65,218 @@ def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) + + +def test_bge_m3_late_chunking_keeps_special_tokens(vllm_runner): + with vllm_runner( + "BAAI/bge-m3", + hf_overrides={"architectures": ["BgeM3EmbeddingModel"]}, + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + # Keep encoder batch size fixed for the numerical reference. + max_num_seqs=1, + gpu_memory_utilization=0.5, + enforce_eager=True, + ) as runner: + llm = runner.llm + text = "A document about Berlin and its museums." + sizes = (1, 3, 128, None) + tokens, *outputs = llm.encode( + [text] * len(sizes), + pooling_task="token_embed", + pooling_params=[ + PoolingParams( + use_activation=size not in (1, None), + late_chunking_params=LateChunkingParams(size) if size else None, + ) + for size in sizes + ], + use_tqdm=False, + ) + assert tokens.outputs.data.shape[0] == len(tokens.prompt_token_ids) + assert tokens.late_chunking.chunks[0].char_range is None + for size, output in zip(sizes[1:], outputs): + if size is None: + # Ordinary BGE-M3 token embeddings still drop BOS only. + expected = tokens.outputs.data[1:] + assert output.late_chunking is None + else: + # BGE-M3's projector is linear, so its unnormalized token + # outputs also provide a reference for projected chunk means. + expected = torch.nn.functional.normalize( + torch.stack([p.mean(0) for p in tokens.outputs.data.split(size)]), + dim=-1, + ) + assert len(output.late_chunking.chunks) == len(expected) + torch.testing.assert_close( + output.outputs.data, expected, atol=2e-5, rtol=2e-4 + ) + + +def test_nomic_late_chunking_without_special_tokens(vllm_runner): + """Keep source ranges without special tokens and reject empty tokenization.""" + with vllm_runner( + "nomic-ai/nomic-embed-text-v1", + revision="720244025c1a7e15661a174c63cce63c8218e52b", + trust_remote_code=True, + runner="pooling", + pooler_config=PoolerConfig(task="token_embed"), + dtype="float32", + max_model_len=128, + max_num_seqs=2, + gpu_memory_utilization=0.3, + enforce_eager=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + ) as runner: + llm = runner.llm + text = "search_document: 中文 😀 café e\u0301" + params = PoolingParams( + late_chunking_params=LateChunkingParams(3), use_activation=False + ) + kwargs = {"add_special_tokens": False} + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + tokenization_kwargs=kwargs, + )[0] + actual = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + )[0] + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(3)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + encoded = llm.get_tokenizer()(text, return_offsets_mapping=True, **kwargs) + assert actual.prompt_token_ids == encoded["input_ids"] + for chunk in actual.late_chunking.chunks: + start, end = chunk.token_range + offsets = [(a, b) for a, b in encoded["offset_mapping"][start:end] if a < b] + assert chunk.char_range == (offsets[0][0], offsets[-1][1]) + for rejected in ("long " * 256, " "): + with pytest.raises(VLLMValidationError): + llm.encode( + rejected, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=params, + tokenization_kwargs=kwargs, + ) + assert ( + llm.encode(text, pooling_task="token_embed", use_tqdm=False)[ + 0 + ].late_chunking + is None + ) + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("normalize", [False, True]) +def test_e5_late_chunking_mixed_requests(llm: LLM, normalize: bool): + texts = [prompt, "Short text.", "中文 😀 café"] + raw = llm.encode( + texts, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + ) + sizes = [3, None, 100] + params = [ + PoolingParams( + late_chunking_params=LateChunkingParams(size) if size is not None else None, + use_activation=normalize, + ) + for size in sizes + ] + actual = llm.encode( + texts, pooling_task="token_embed", pooling_params=params, use_tqdm=False + ) + for original, result, size in zip(raw, actual, sizes): + expected = original.outputs.data + if size is not None: + expected = torch.stack( + [part.float().mean(0) for part in expected.split(size)] + ) + assert result.late_chunking is not None + assert len(result.late_chunking.chunks) == len(expected) + assert result.late_chunking.input_tokens == len(result.prompt_token_ids) + else: + assert result.late_chunking is None + if normalize: + expected = torch.nn.functional.normalize(expected, dim=-1) + check_embeddings_close( + embeddings_0_lst=result.outputs.data.tolist(), + embeddings_1_lst=expected.tolist(), + name_0="late_chunking", + name_1="token_reference", + ) + torch.testing.assert_close( + result.outputs.data.norm(dim=-1), + expected.norm(dim=-1), + atol=1e-3, + rtol=1e-3, + ) + + +@pytest.mark.parametrize("use_v2_runner", [False, True]) +def test_late_chunking_recomputes_cached_prefix_with_chunked_prefill( + vllm_runner, monkeypatch, use_v2_runner +): + # Dummy weights suffice for testing scheduling and reduction of the same states. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_runner else "0") + with vllm_runner( + "openai-community/gpt2", + runner="pooling", + convert="embed", + load_format="dummy", + pooler_config=PoolerConfig( + task="token_embed", seq_pooling_type="LAST", tok_pooling_type="ALL" + ), + dtype="float32", + max_model_len=128, + max_num_batched_tokens=16, + max_num_seqs=2, + enable_prefix_caching=True, + enable_chunked_prefill=True, + enforce_eager=True, + gpu_memory_utilization=0.3, + seed=0, + ) as runner: + text = "The document has several important words. " * 6 + llm = runner.llm + raw = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams(use_activation=False), + )[0] + assert len(raw.prompt_token_ids) > 16 + cached = llm.encode( + text, + pooling_task="token_embed", + use_tqdm=False, + pooling_params=PoolingParams( + use_activation=False, skip_reading_prefix_cache=False + ), + )[0] + assert cached.num_cached_tokens > 0 + params = PoolingParams( + late_chunking_params=LateChunkingParams(7), + use_activation=False, + skip_reading_prefix_cache=False, + ) + actual = llm.encode( + text, pooling_task="token_embed", pooling_params=params, use_tqdm=False + )[0] + assert actual.num_cached_tokens == 0 + assert actual.prompt_token_ids == raw.prompt_token_ids + expected = torch.stack([part.mean(0) for part in raw.outputs.data.split(7)]) + torch.testing.assert_close(actual.outputs.data, expected, atol=1e-5, rtol=1e-4) + assert actual.late_chunking.input_tokens == len(raw.prompt_token_ids) + assert params.skip_reading_prefix_cache is False + assert params.late_chunking_params is not None diff --git a/tests/model_executor/layers/test_pooler_heads.py b/tests/model_executor/layers/test_pooler_heads.py index a80c5b8344f5..00091ec0c292 100644 --- a/tests/model_executor/layers/test_pooler_heads.py +++ b/tests/model_executor/layers/test_pooler_heads.py @@ -308,6 +308,7 @@ def test_passthrough(self): x = torch.randn(5, _HIDDEN) param = PoolingParams(task="token_embed") out = head.forward_chunk(x, param) + assert out is x assert torch.equal(out, x) def test_none_chunked_prefill(self): diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py index 75bbed4f37fe..d6b03dee1954 100644 --- a/tests/model_executor/layers/test_pooler_methods.py +++ b/tests/model_executor/layers/test_pooler_methods.py @@ -13,19 +13,22 @@ from vllm.config import PoolerConfig, set_current_vllm_config from vllm.model_executor.layers.pooler import PoolingParamsUpdate +from vllm.model_executor.layers.pooler.activations import PoolerNormalize from vllm.model_executor.layers.pooler.seqwise.methods import ( CLSPool, LastPool, MeanPool, get_seq_pooling_method, ) -from vllm.model_executor.layers.pooler.special import DispatchPooler +from vllm.model_executor.layers.pooler.special import BOSEOSFilter, DispatchPooler +from vllm.model_executor.layers.pooler.tokwise.heads import TokenEmbeddingPoolerHead from vllm.model_executor.layers.pooler.tokwise.methods import ( AllPool, StepPool, get_tok_pooling_method, ) -from vllm.pooling_params import PoolingParams +from vllm.model_executor.layers.pooler.tokwise.poolers import TokenPooler +from vllm.pooling_params import LateChunkingParams, PoolingParams from vllm.tasks import PoolingTask from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates @@ -606,3 +609,139 @@ def test_step(self): def test_unknown_raises(self): with pytest.raises(NotImplementedError, match="UNKNOWN"): get_tok_pooling_method("UNKNOWN") + + +@pytest.mark.parametrize( + ("chunk_size", "dtype", "head_dtype"), + [ + (1, torch.float32, None), + (2, torch.float16, None), + (2, torch.bfloat16, torch.float32), + (7, torch.float32, None), + (10, torch.float32, None), + (2**63 - 1, torch.float32, None), + ], +) +def test_late_chunk_pool_matches_fp32_means_before_projection_and_normalization( + chunk_size, dtype, head_dtype +): + # Unequal norms expose the incorrect normalize-then-average implementation. + hidden = (torch.arange(28).reshape(7, 4) - 10).to(dtype) + projector = torch.nn.Linear(4, 3, bias=True).to(head_dtype or dtype) + head = TokenEmbeddingPoolerHead(head_dtype, projector, PoolerNormalize()) + pooler = TokenPooler(TestAllPool._make_all_pool(), head) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=chunk_size), + use_activation=True, + ) + metadata = _make_metadata([7], pooling_params=[params]) + actual = pooler(hidden, metadata)[0] + means = torch.stack( + [hidden[a : a + chunk_size].float().mean(0) for a in range(0, 7, chunk_size)] + ).to(head_dtype or dtype) + expected = torch.nn.functional.normalize(projector(means), p=2, dim=-1) + torch.testing.assert_close(actual, expected) + assert actual.dtype == (head_dtype or dtype) + + +def test_late_chunk_pool_reduces_before_nonlinear_projector(): + hidden = torch.tensor([[2.0, 0.0], [0.0, 1.0]]) + projector = torch.nn.Tanh() + head = TokenEmbeddingPoolerHead(projector=projector, activation=PoolerNormalize()) + params = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=2), + use_activation=True, + ) + actual = head.forward_chunk(hidden, params) + expected = torch.nn.functional.normalize(projector(hidden.mean(0, keepdim=True))) + wrong_order = torch.nn.functional.normalize(projector(hidden).mean(0, keepdim=True)) + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, wrong_order) + + +def test_late_chunk_pool_isolates_mixed_requests_and_owns_finished_storage(): + hidden = torch.arange(36, dtype=torch.float32).reshape(9, 4) + params = [ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=2) + ), + PoolingParams(task="token_embed"), + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=1) + ), + ] + pooler = TokenPooler( + TestAllPool._make_all_pool(async_scheduling=True), TokenEmbeddingPoolerHead() + ) + outputs = pooler(hidden, _make_metadata([3, 4, 2], pooling_params=params)) + expected = [ + torch.stack([hidden[:2].mean(0), hidden[2]]), + hidden[3:7].clone(), + hidden[7:].clone(), + ] + hidden.fill_(-999) + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual, reference) + + +@pytest.mark.parametrize("chunk_size", [1, 2, 8]) +@pytest.mark.parametrize("embedding_size", [1, 4]) +def test_bos_eos_filter_preserves_chunk_vectors_in_mixed_batch( + chunk_size, embedding_size +): + hidden = torch.arange(10 * embedding_size, dtype=torch.float32).reshape( + 10, embedding_size + ) + pooler = BOSEOSFilter( + TokenPooler(TestAllPool._make_all_pool(), TokenEmbeddingPoolerHead()), + bos_token_id=101, + eos_token_id=102, + ) + metadata = _make_metadata( + [5, 5], + token_ids=[[101, 1, 2, 3, 102]] * 2, + pooling_params=[ + PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size) + ), + PoolingParams(task="token_embed"), + ], + ) + chunked, ordinary = pooler(hidden, metadata) + expected = torch.stack([part.mean(0) for part in hidden[:5].split(chunk_size)]) + torch.testing.assert_close(chunked, expected) + torch.testing.assert_close(ordinary, hidden[6:9].squeeze(-1)) + + +def test_bos_eos_filter_waits_for_complete_chunk_vectors(): + pooler = BOSEOSFilter( + TokenPooler( + TestAllPool._make_all_pool(chunked=True), TokenEmbeddingPoolerHead() + ), + bos_token_id=101, + ) + params = [ + PoolingParams(task="token_embed", late_chunking_params=LateChunkingParams(3)) + ] + first = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[2], + ) + hidden = torch.arange(8, dtype=torch.float32).reshape(4, 2) + assert pooler(hidden[:2], first) == [None] + last = _make_metadata( + [4], + token_ids=[[101, 1, 2, 102]], + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[4], + ) + last.pooling_states = first.pooling_states + actual = pooler(hidden[2:], last)[0] + torch.testing.assert_close(actual, torch.stack([hidden[:3].mean(0), hidden[3]])) + assert not last.pooling_states[0].hidden_states_cache diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index f34270d0da52..81af02b502ab 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from types import SimpleNamespace from typing import Any +import msgspec import pytest from pydantic import TypeAdapter, ValidationError @@ -13,6 +15,7 @@ from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest from vllm.exceptions import VLLMValidationError +from vllm.pooling_params import LateChunkingParams EMBEDDING_MODELS = [ EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False), @@ -211,3 +214,92 @@ def test_token_classify(pooling_type: str): with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) + + +@pytest.mark.parametrize("value", [0, -1, True, False, 1.5, "2"]) +def test_late_chunk_size_rejects_non_positive_integers(value): + with pytest.raises(VLLMValidationError, match="positive integer"): + PoolingParams(late_chunking_params=LateChunkingParams(chunk_size=value)) + + +def _late_chunking_model_config(): + return SimpleNamespace( + architecture="NomicBertModel", + model_impl="auto", + is_matryoshka=False, + hf_config=SimpleNamespace(), + pooler_config=PoolerConfig(seq_pooling_type="MEAN", tok_pooling_type="ALL"), + ) + + +def test_late_chunking_params_preserve_defaults_clone_and_wire_format(): + model_config = _late_chunking_model_config() + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(chunk_size=3) + ) + params.verify(model_config) + assert params.skip_reading_prefix_cache is True + assert params.use_activation is True + # Chunk means must include cached prefixes even when the user opts into reads. + explicit_cache_flag = PoolingParams( + task="token_embed", + late_chunking_params=LateChunkingParams(chunk_size=3), + skip_reading_prefix_cache=False, + ) + explicit_cache_flag.verify(model_config) + assert explicit_cache_flag.skip_reading_prefix_cache is True + clone = params.clone() + assert clone.late_chunking_params is not None + clone.late_chunking_params.chunk_size = 7 + assert params.late_chunking_params is not None + assert params.late_chunking_params.chunk_size == 3 + # Workers receive only the fixed chunk size, never source ranges. + assert msgspec.msgpack.decode( + msgspec.msgpack.encode(params.late_chunking_params) + ) == [3] + assert ( + msgspec.msgpack.decode(msgspec.msgpack.encode(params), type=PoolingParams) + == params + ) + # A message written before the appended field still uses the old default. + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(PoolingParams())) + assert msgspec.convert(wire[:-1], type=PoolingParams).late_chunking_params is None + + +@pytest.mark.parametrize( + "task", ["embed", "classify", "token_classify", "plugin", None] +) +def test_late_chunking_rejects_other_tasks(task): + with pytest.raises(VLLMValidationError, match="requires token_embed"): + PoolingParams( + task=task, late_chunking_params=LateChunkingParams(chunk_size=2) + ).verify(_late_chunking_model_config()) + + +def test_late_chunking_uses_existing_dimension_validation(): + config = MockMatryoshkaModelConfig( + pooler_config=PoolerConfig(tok_pooling_type="ALL", use_activation=False), + ) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2), dimensions=4 + ) + params.verify(config) + assert params.dimensions == 4 + assert params.use_activation is False + params.dimensions = config.embedding_size + 1 + with pytest.raises(VLLMValidationError, match="dimensions in range"): + params.verify(config) + + +@pytest.mark.parametrize("from_config", [False, True]) +def test_late_chunking_rejects_token_filtering_after_resolving_defaults(from_config): + config = MockModelConfig(pooler_config=PoolerConfig(tok_pooling_type="STEP")) + params = PoolingParams( + task="token_embed", late_chunking_params=LateChunkingParams(2) + ) + if from_config: + config.pooler_config.step_tag_id = 42 + else: + params.step_tag_id = 42 + with pytest.raises(VLLMValidationError, match="all token rows"): + params.verify(config) diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 94f84f09b678..b37a009c925a 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -12,6 +12,7 @@ from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams +from vllm.exceptions import VLLMValidationError from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -26,18 +27,26 @@ from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor +from ..late_chunking import ( + attach_late_chunking_metadata, + build_late_chunking_metadata, + prepare_late_chunking_input, +) from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( AnyOfflineInputsContext, AnyRenderParam, ChunkedEmbeddingMetadata, EncodeChatRenderParams, + LateChunkingSlot, OfflineEncodeInputsContext, + OfflineOutputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, PoolingEngineInput, PoolingServeContext, RequestFactory, + RequestGenerator, ) from .protocol import ( CohereEmbedContent, @@ -657,10 +666,75 @@ def _enforce_cohere_max_tokens(self, ctx: PoolingServeContext) -> None: class TokenEmbedIOProcessor(PoolingIOProcessor): name = "token_embed" + def get_request_factory_offline( + self, ctx: AnyOfflineInputsContext + ) -> tuple[RequestFactory, int]: + factory, num_requests = super().get_request_factory_offline(ctx) + assert isinstance(ctx, OfflineEncodeInputsContext) + if ctx.pooling_params is None or not any( + p.late_chunking_params is not None + for p in self._params_to_seq(ctx.pooling_params, num_requests) + ): + return factory, num_requests + + if (ctx.tokenization_kwargs or {}).get("padding") not in ( + None, + False, + "do_not_pad", + ): + raise VLLMValidationError("Late chunking does not support input padding") + + def request_factory() -> RequestGenerator: + for request in factory(): + slot = None + if request["params"].late_chunking_params is not None: + slot = LateChunkingSlot() + request["late_chunking"] = slot + ctx.late_chunking.append(slot) + yield request + + return request_factory, num_requests + + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + late_chunking = render_params["params"].late_chunking_params + if late_chunking is None: + return super().render(render_params) + + text, tok_params = prepare_late_chunking_input(self.vllm_config, render_params) + render_params = render_params.copy() + render_params["tok_params"] = tok_params + result = super().render(render_params) + prompt = result["prompts"] + if prompt["type"] != "token": + raise VLLMValidationError("Late chunking requires tokenized text") + render_params["late_chunking"].metadata = build_late_chunking_metadata( + text, + len(prompt["prompt_token_ids"]), + prompt.pop("prompt_token_offsets", None), + late_chunking.chunk_size, + ) + return result + + def post_process_offline(self, ctx: OfflineOutputsContext): + if ctx.late_chunking: + for output, slot in zip(ctx.outputs, ctx.late_chunking, strict=True): + if slot is not None: + assert slot.metadata is not None + attach_late_chunking_metadata(output, slot.metadata) + return ctx.outputs + class JinaRankingTokenEmbedIOProcessor( TokenEmbedIOProcessor, JinaRankingIOProcessorMixin ): + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + if render_params["params"].late_chunking_params is not None: + raise VLLMValidationError( + "JinaForRanking returns document/query vectors and does not support " + "late chunking" + ) + return super().render(render_params) + def get_request_factory_online( self, ctx: PoolingServeContext ) -> Sequence[AnyRenderParam]: diff --git a/vllm/entrypoints/pooling/late_chunking.py b/vllm/entrypoints/pooling/late_chunking.py new file mode 100644 index 000000000000..d848c3257f20 --- /dev/null +++ b/vllm/entrypoints/pooling/late_chunking.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from dataclasses import replace +from typing import cast + +from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError +from vllm.outputs import LateChunk, LateChunkingMetadata, PoolingRequestOutput +from vllm.renderers import TokenizeParams + +from .typing import AnyRenderParam, EncodeCMPLRenderParams + + +def prepare_late_chunking_input( + config: VllmConfig, render_params: AnyRenderParam +) -> tuple[str, TokenizeParams]: + """Validate full-document text input before requesting tokenizer offsets.""" + params = render_params["params"] + params.verify(config.model_config) + if "prompts" not in render_params: + raise VLLMValidationError("Late chunking requires plain-text input") + prompt = cast(EncodeCMPLRenderParams, render_params)["prompts"] + if ( + not isinstance(prompt, dict) + or set(prompt) - {"prompt", "cache_salt"} + or not isinstance(text := prompt.get("prompt"), str) + or not text + ): + raise VLLMValidationError("Late chunking requires nonempty plain-text input") + tok_params = render_params["tok_params"] + if ( + tok_params.pad_prompt_tokens is not None + or tok_params.truncate_prompt_tokens is not None + or tok_params.do_lower_case + ): + raise VLLMValidationError( + "Late chunking does not support input padding, truncation or " + "renderer text normalization" + ) + return text, replace(tok_params, return_token_offsets=True) + + +def build_late_chunking_metadata( + text: str, + num_tokens: int, + offsets: Sequence[tuple[int, int]] | None, + chunk_size: int, +) -> LateChunkingMetadata: + """Map chunks to original text, keeping special-token-only chunks.""" + if num_tokens <= 0 or offsets is None or len(offsets) != num_tokens: + raise VLLMValidationError( + "Late chunking requires a fast tokenizer with offsets aligned to " + "all input tokens" + ) + previous_start = previous_end = 0 + for start, end in offsets: + if not 0 <= start <= end <= len(text): + raise VLLMValidationError("Token offsets are not aligned to the input text") + if start < end: + if start < previous_start or end < previous_end: + raise VLLMValidationError("Token offsets are not in source-text order") + previous_start, previous_end = start, end + + chunks = [] + for start in range(0, num_tokens, chunk_size): + end = min(start + chunk_size, num_tokens) + source = [(a, b) for a, b in offsets[start:end] if a < b] + char_range = ( + (min(a for a, _ in source), max(b for _, b in source)) if source else None + ) + chunks.append(LateChunk(token_range=(start, end), char_range=char_range)) + return LateChunkingMetadata( + chunk_size=chunk_size, input_tokens=num_tokens, chunks=chunks + ) + + +def attach_late_chunking_metadata( + output: PoolingRequestOutput, metadata: LateChunkingMetadata +) -> None: + """Attach ranges only to successful, complete chunk-vector results.""" + if not output.finished or output.error is not None: + return + data = output.outputs.data + if ( + data.ndim != 2 + or data.shape[0] != len(metadata.chunks) + or len(output.prompt_token_ids) != metadata.input_tokens + ): + raise ValueError("Late-chunking output does not match its token/source ranges") + output.late_chunking = metadata diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index 19be1eddde91..edf92791a86a 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -129,7 +129,7 @@ def encode( io_processor, request_factory, num_requests, use_tqdm=use_tqdm ) outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs) + ctx=OfflineOutputsContext(outputs=outputs, late_chunking=ctx.late_chunking) ) return outputs diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index a299fc657d54..33aa65998978 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -3,7 +3,7 @@ import time from collections.abc import AsyncGenerator, Callable, Generator, Sequence from dataclasses import dataclass, field -from typing import Any, Generic, TypeAlias, TypedDict, TypeVar +from typing import Any, Generic, NotRequired, TypeAlias, TypedDict, TypeVar from fastapi import Request from pydantic import ConfigDict @@ -11,6 +11,7 @@ from vllm import PoolingParams, PoolingRequestOutput, PromptType from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest +from vllm.outputs import LateChunkingMetadata from vllm.renderers import ChatParams, TokenizeParams from vllm.renderers.chat_utils import ChatCompletionMessageParam from vllm.renderers.inputs import DictPrompt @@ -80,6 +81,13 @@ class ChunkedEmbeddingMetadata: chunk_index: int +@dataclass +class LateChunkingSlot: + """Per-document source ranges shared by frontend rendering and output contexts.""" + + metadata: LateChunkingMetadata | None = None + + @dataclass(kw_only=True) class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) @@ -124,6 +132,9 @@ class OfflineInputsContext: tokenization_kwargs: dict[str, Any] | None lora_request: Sequence[LoRARequest | None] | None priorities: int | Sequence[int] | None + late_chunking: list[LateChunkingSlot | None] = field( + default_factory=list, kw_only=True + ) @dataclass @@ -158,6 +169,7 @@ class OfflineOutputsContext: ## for bi-encoder & late-interaction n_queries: int | None = None + late_chunking: list[LateChunkingSlot | None] = field(default_factory=list) class RenderParams(TypedDict): @@ -168,6 +180,7 @@ class RenderParams(TypedDict): params: PoolingParams lora_requests: LoRARequest | None priorities: int + late_chunking: NotRequired[LateChunkingSlot] class EncodeCMPLRenderParams(RenderParams): diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 23354d66f3c4..f464f28300a7 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -198,6 +198,11 @@ def forward( for i, (prompt_len, token_ids) in enumerate( zip(pooling_metadata.prompt_lens, prompt_token_ids) ): + # Late chunking returns one row per chunk, with special tokens + # already included in the means. Skip BOS/EOS row removal for these + # requests; ordinary token outputs still use the filtering below. + if pooling_metadata.pooling_params[i].late_chunking_params is not None: + continue pooled_data = pooled_outputs[i] assert ( isinstance(pooled_data, torch.Tensor) diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index 78ec82a5384e..95242ee470a7 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -17,6 +17,23 @@ TokenPoolerHeadOutputItem: TypeAlias = torch.Tensor | None +def _mean_pool_chunks(data: torch.Tensor, chunk_size: int) -> torch.Tensor: + """Mean-pool complete contextual states, accumulating in FP32.""" + chunk_size = min(chunk_size, len(data)) + num_full, remainder = divmod(len(data), chunk_size) + full_end = num_full * chunk_size + full_chunks = data[:full_end].reshape(num_full, chunk_size, data.shape[-1]) + if not remainder: + return full_chunks.mean(dim=1, dtype=torch.float32) + + return torch.cat( + [ + full_chunks.mean(dim=1, dtype=torch.float32), + data[full_end:].mean(dim=0, keepdim=True, dtype=torch.float32), + ] + ) + + class TokenPoolerHead(nn.Module, ABC): @abstractmethod def get_supported_tasks(self) -> Set[PoolingTask]: @@ -80,6 +97,12 @@ def forward_chunk( if pooled_data is None: return None + if (late_chunking := pooling_param.late_chunking_params) is not None: + means = _mean_pool_chunks(pooled_data, late_chunking.chunk_size) + pooled_data = ( + means.to(pooled_data.dtype) if self.head_dtype is None else means + ) + if self.head_dtype is not None: pooled_data = pooled_data.to(self.head_dtype) # pooled_data shape: [n_tokens, hidden_size] diff --git a/vllm/outputs.py b/vllm/outputs.py index 145b7d510f71..887ebc5933de 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -277,6 +277,20 @@ def __repr__(self) -> str: _O = TypeVar("_O", default=PoolingOutput) +@dataclass(frozen=True) +class LateChunk: + token_range: tuple[int, int] + char_range: tuple[int, int] | None + + +@dataclass(frozen=True) +class LateChunkingMetadata: + chunk_size: int + input_tokens: int + chunks: list[LateChunk] + strategy: str = "fixed_tokens" + + class PoolingRequestOutput(Generic[_O]): """The output data of a pooling request to the LLM. @@ -287,6 +301,7 @@ class PoolingRequestOutput(Generic[_O]): num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. error: Structured request-level error information, if the request failed. + late_chunking: Token and source ranges for fixed-length late chunking. """ @@ -299,8 +314,10 @@ def __init__( finished: bool, *, error: RequestError | None = None, + late_chunking: LateChunkingMetadata | None = None, ): self.request_id = request_id + self.late_chunking = late_chunking self.prompt_token_ids = prompt_token_ids self.num_cached_tokens = num_cached_tokens self.finished = finished diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 319b77d5986a..45c34afbe3fd 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -36,6 +36,28 @@ class LateInteractionParams( query_uses: int | None = None +class LateChunkingParams( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + array_like=True, +): # type: ignore[call-arg] + """Fixed-length mean pooling before the token embedding head. + + Requests recompute all prompt tokens, bypassing prefix cache reads. + Special tokens count toward chunk boundaries and are included in chunk means, + including the BOS token normally filtered from BGE-M3 token embeddings. + """ + + chunk_size: int + + def __post_init__(self) -> None: + self.verify() + + def verify(self) -> None: + if type(self.chunk_size) is not int or self.chunk_size <= 0: + raise VLLMValidationError("chunk_size must be a positive integer") + + class PoolingParams( msgspec.Struct, omit_defaults=True, # type: ignore[call-arg] @@ -48,6 +70,8 @@ class PoolingParams( `None` uses the pooler's default, which is `True` in most cases. dimensions: Reduce the dimensions of embeddings if model support matryoshka representation. + late_chunking_params: Mean-pool contextual token states into fixed-size + chunks for the token embedding task. """ @@ -72,16 +96,22 @@ class PoolingParams( extra_kwargs: dict[str, Any] | None = None output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY + ## for token embedding models + # --8<-- [start:token-embed-pooling-params] + late_chunking_params: LateChunkingParams | None = None + """Mean-pool contextual tokens into fixed-size chunks for `token_embed`.""" + # --8<-- [end:token-embed-pooling-params] + @property def all_parameters(self) -> list[str]: - return ["dimensions", "use_activation"] + return ["dimensions", "use_activation", "late_chunking_params"] @property def valid_parameters(self): return { "embed": ["dimensions", "use_activation"], "classify": ["use_activation"], - "token_embed": ["dimensions", "use_activation"], + "token_embed": ["dimensions", "use_activation", "late_chunking_params"], "token_classify": ["use_activation"], } @@ -90,6 +120,15 @@ def clone(self) -> "PoolingParams": return deepcopy(self) def verify(self, model_config: ModelConfig) -> None: + if self.late_chunking_params is not None: + self.late_chunking_params.verify() + if self.task != "token_embed": + raise VLLMValidationError("Late chunking requires token_embed") + if self.late_interaction_params is not None: + raise VLLMValidationError( + "Late chunking cannot be combined with late-interaction scoring" + ) + # plugin task uses io_processor.parse_data to verify inputs, # skipping PoolingParams verify if self.task == "plugin": @@ -107,6 +146,11 @@ def verify(self, model_config: ModelConfig) -> None: self._merge_default_parameters(model_config) self._set_default_parameters(model_config) self._verify_valid_parameters() + if self.late_chunking_params is not None: + if self.step_tag_id is not None: + raise VLLMValidationError("Late chunking requires all token rows") + # KV cache hits omit hidden states needed for complete chunk means. + self.skip_reading_prefix_cache = True def _merge_default_parameters(self, model_config: ModelConfig) -> None: pooler_config = model_config.pooler_config @@ -228,6 +272,7 @@ def __repr__(self) -> str: f"requires_token_ids={self.requires_token_ids}, " f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, " f"late_interaction_params={self.late_interaction_params}, " + f"late_chunking_params={self.late_chunking_params}, " f"extra_kwargs={self.extra_kwargs})" ) From cd7517fb17bc26f5e63b348a73cc0dab70be8d48 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:31:53 +0400 Subject: [PATCH 489/952] [Bugfix] Fix ReplaySSM ring shape assertion when TP does not divide Mamba2 n_groups (#60704) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/kernels/mamba/test_ssu_dispatch.py | 24 +++++++++++++++++++ .../layers/mamba/mamba_utils.py | 3 ++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 3b3dc2dc6bec..e12369ef8a04 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -637,3 +637,27 @@ def test_replayssm_physical_ring_shape( (8, expected_ring_len), (2, expected_ring_len, 16), ) + + +def test_replayssm_ring_shape_extends_groups_for_head_shards(): + """A single group under TP is replicated to every head shard, as in + ``mamba2_state_shape``, so each rank's B_cache holds one group.""" + base_shapes = MambaStateShapeCalculator.mamba2_state_shape( + tp_world_size=2, + intermediate_size=256, + n_groups=1, + num_heads=4, + head_dim=64, + state_size=128, + conv_kernel=4, + ) + + shapes = MambaStateShapeCalculator.append_replayssm_ring( + base_shapes, + n_groups=1, + tp_world_size=2, + logical_window=16, + backend=MambaBackendEnum.TRITON, + ) + + assert shapes[4] == (1, 16, 128) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 31e89ffe55b5..5c423739b7a0 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -238,13 +238,14 @@ def append_replayssm_ring( """Append the physical ReplaySSM ring shapes. ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; - B_cache uses the un-extended ``n_groups``. + B_cache extends ``n_groups`` for head shards as in ``mamba2_state_shape``. """ ring_buffer_len = logical_window if backend == MambaBackendEnum.FLASHINFER: # FlashInfer keeps the live window and current verify window together. ring_buffer_len += 1 + num_speculative_tokens local_nheads, head_dim, state_size = base_shapes[1] + n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size) local_ngroups = divide(n_groups, tp_world_size) return ( *base_shapes, From 818f70c295137ddf9d2ba63262adb6f12e954b58 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:38:23 +0400 Subject: [PATCH 490/952] [Bugfix] Mask special added tokens in xgrammar JSON, regex and EBNF grammars (#60749) Signed-off-by: Vadim Gimpelson Signed-off-by: Artem Perevedentsev Co-authored-by: Claude Co-authored-by: Artem Perevedentsev --- .../test_backend_xgrammar_stop_tokens.py | 93 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 86 +++++++++++++++-- 2 files changed, 170 insertions(+), 9 deletions(-) diff --git a/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py index e910d29aa38c..5b0baf44fd3c 100644 --- a/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py +++ b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py @@ -8,8 +8,15 @@ bitmask while the FSM is still mid-object and truncate structured output. ``compile_grammar`` now forwards ``all_stop_token_ids`` to the matcher as ``override_stop_tokens`` so xgrammar masks them until the grammar completes. + +Issue #60580: xgrammar only treats tokens that decode to an empty string as +special, so a tokenizer's control tokens (e.g. ``<|im_start|>``) matched +grammar text and could be sampled inside a JSON string. Structural tags and +Lark grammars, which can name such tokens, must still be able to emit them. """ +import json + import pytest from transformers import AutoTokenizer @@ -25,6 +32,14 @@ QUOTE = 1 # standalone `"`; opens then closes the JSON string LETTER = 55 # `X`: valid string content, not a special/stop token by default +QWEN_TOKENIZER = "Qwen/Qwen3-0.6B" +QWEN_VOCAB_SIZE = 151936 + +# Qwen3 special added tokens. +ENDOFTEXT = 151643 # <|endoftext|>: a generation_config stop token +IM_START = 151644 # <|im_start|>: never a stop token +IM_END = 151645 # <|im_end|>: the tokenizer's eos + def _token_allowed(row, token_id: int) -> bool: word = int(row[token_id // 32].item()) & 0xFFFFFFFF @@ -80,3 +95,81 @@ def test_request_stop_tokens_gated_to_grammar_terminal(backend: XgrammarBackend) assert _token_allowed(bm_override[0], LETTER) assert _token_allowed(bm_default[0], EOS) assert _token_allowed(bm_override[0], EOS) + + +@pytest.fixture(scope="module") +def qwen_tokenizer(): + return AutoTokenizer.from_pretrained(QWEN_TOKENIZER) + + +@pytest.fixture(scope="module") +def qwen_backend(qwen_tokenizer) -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") + ) + return XgrammarBackend( + vllm_config, tokenizer=qwen_tokenizer, vocab_size=QWEN_VOCAB_SIZE + ) + + +@pytest.mark.parametrize( + "request_type, grammar_spec", + [ + (StructuredOutputOptions.JSON, '{"type": "string"}'), + (StructuredOutputOptions.REGEX, '"[^"]*"'), + (StructuredOutputOptions.GRAMMAR, 'root ::= "\\"" [^"]* "\\""'), + ], +) +def test_special_tokens_masked_until_grammar_ends( + qwen_backend: XgrammarBackend, qwen_tokenizer, request_type, grammar_spec +): + grammar = qwen_backend.compile_grammar( + request_type, grammar_spec, stop_token_ids={IM_END, ENDOFTEXT} + ) + bitmask = qwen_backend.allocate_token_bitmask(1) + + text = qwen_tokenizer.encode('"The company is', add_special_tokens=False) + assert grammar.accept_tokens("req", text) + grammar.fill_bitmask(bitmask, 0) + for token_id in (IM_START, IM_END, ENDOFTEXT): + assert not _token_allowed(bitmask[0], token_id) + + quote = qwen_tokenizer.encode('"', add_special_tokens=False) + assert grammar.accept_tokens("req", quote) + grammar.fill_bitmask(bitmask, 0) + assert not _token_allowed(bitmask[0], IM_START) + assert _token_allowed(bitmask[0], IM_END) + assert _token_allowed(bitmask[0], ENDOFTEXT) + + +@pytest.mark.parametrize( + "request_type, grammar_spec", + [ + ( + StructuredOutputOptions.STRUCTURAL_TAG, + json.dumps( + { + "type": "structural_tag", + "format": { + "type": "tag", + "begin": "<|im_start|>", + "content": { + "type": "json_schema", + "json_schema": {"type": "string"}, + }, + "end": "\n", + }, + } + ), + ), + (StructuredOutputOptions.GRAMMAR, f'start: <[{IM_START}]> "assistant"'), + ], +) +def test_token_naming_grammars_can_emit_special_tokens( + qwen_backend: XgrammarBackend, request_type, grammar_spec +): + grammar = qwen_backend.compile_grammar(request_type, grammar_spec) + bitmask = qwen_backend.allocate_token_bitmask(1) + grammar.fill_bitmask(bitmask, 0) + assert _token_allowed(bitmask[0], IM_START) + assert grammar.accept_tokens("req", [IM_START]) diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 9631766d9267..94ccc55b8b98 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -26,6 +26,8 @@ if TYPE_CHECKING: import xgrammar as xgr + + from vllm.tokenizers import TokenizerLike else: xgr = LazyLoader("xgr", globals(), "xgrammar") @@ -43,6 +45,7 @@ def __post_init__(self): model_config is not None and model_config.hf_config.model_type == "plamo3" ) + full_vocab_tokenizer_info = None if is_mistral_tokenizer(self.tokenizer): # NOTE: ideally, xgrammar should handle this accordingly. # refer to https://github.com/mlc-ai/xgrammar/blob/d77c0a0173ef14779c918e3be7966ba852f7910f/python/xgrammar/tokenizer_info.py#L98 @@ -66,15 +69,46 @@ def __post_init__(self): ): tokenizer_info, _ = init_xgrammar() else: - tokenizer_info = xgr.TokenizerInfo.from_huggingface( + # tokenizer for structural tags + # accepts some of the special tokens + full_vocab_tokenizer_info = xgr.TokenizerInfo.from_huggingface( self.tokenizer, vocab_size=self.vocab_size, ) - self.compiler = xgr.GrammarCompiler( - tokenizer_info, - max_threads=8, - cache_enabled=True, - cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024, + # tokenizer for json + # does NOT accept any special tokens + tokenizer_info = self._filter_special_tokens( + self.tokenizer, full_vocab_tokenizer_info + ) + + # FIXME(arpera): + # Here we create not one but TWO compilers + # because xgrammar currently does not properly handle + # special tokens. + # Proper special tokens handling depends on structured output + # request type: json, structural_tag, etc. + # Structural tag may accept some of the special tokens + # At the same time json must NOT accept any special tokens + # Currently xgrammar does not properly handle all the cases + # So, we have to implement additional logic on our side + # We maintain TWO compilers for different structural output types + # If in future xgrammar will fix this problem, + # then we can consider to remove this + # Also, note that other backends like guidance + # do NOT have this problem (as far as I see, I may be wrong). + # So, if someone wants to go and fix the problem on xgrammar side + # you should first definitely check how other backends deal with it + # Some work on xgrammar side has been started: + # https://github.com/mlc-ai/xgrammar/issues/832 + + # Compiler for json. Must NOT accept any special tokens. + self.compiler = self._create_compiler(tokenizer_info) + + # Compiler for structural tag. May accept some of the special tokens. + self.full_vocab_compiler = ( + self.compiler + if full_vocab_tokenizer_info is None + else self._create_compiler(full_vocab_tokenizer_info) ) self.num_speculative_tokens = 0 @@ -122,7 +156,7 @@ def compile_grammar( ) elif request_type == StructuredOutputOptions.GRAMMAR: if grammar_is_likely_lark(grammar_spec): - ctx = self.compiler.compile_lark(grammar_spec) + ctx = self.full_vocab_compiler.compile_lark(grammar_spec) else: ctx = self.compiler.compile_grammar(grammar_spec) elif request_type == StructuredOutputOptions.REGEX: @@ -142,9 +176,11 @@ def compile_grammar( ) for s in s_tag["structures"] ] - ctx = self.compiler.compile_structural_tag(tags, s_tag["triggers"]) + ctx = self.full_vocab_compiler.compile_structural_tag( + tags, s_tag["triggers"] + ) else: - ctx = self.compiler.compile_structural_tag(grammar_spec) + ctx = self.full_vocab_compiler.compile_structural_tag(grammar_spec) else: logger.error( "Validation should have already occurred. Please file an issue." @@ -168,6 +204,38 @@ def allocate_token_bitmask(self, max_num_seqs: int): def destroy(self): del self.compiler + del self.full_vocab_compiler + + @staticmethod + def _create_compiler( + tokenizer_info: "xgr.TokenizerInfo", + ) -> "xgr.GrammarCompiler": + return xgr.GrammarCompiler( + tokenizer_info, + max_threads=8, + cache_enabled=True, + cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024, + ) + + @staticmethod + def _filter_special_tokens( + tokenizer: "TokenizerLike", tokenizer_info: "xgr.TokenizerInfo" + ) -> "xgr.TokenizerInfo": + """Filter ALL special tokens from tokenizer_info""" + special_token_ids = { + token_id + for token_id, token in getattr( + tokenizer, "added_tokens_decoder", {} + ).items() + if token.special + } + encoded_vocab = [""] * tokenizer_info.vocab_size + for token, token_id in tokenizer.get_vocab().items(): + if token_id < len(encoded_vocab) and token_id not in special_token_ids: + encoded_vocab[token_id] = token + return xgr.TokenizerInfo.from_vocab_and_metadata( + encoded_vocab, tokenizer_info.dump_metadata() + ) @dataclass From d9a543c500e9014045d7ae769df0ca7782e31870 Mon Sep 17 00:00:00 2001 From: Aditya Ganesh Kumar Date: Tue, 13 Oct 2026 17:31:59 +0530 Subject: [PATCH 491/952] [Bugfix][Frontend] Forward --enable-log-outputs to the Cohere Chat v2 (#61089) Signed-off-by: Aditya Ganesh Kumar --- vllm/entrypoints/cohere/serving.py | 4 ++++ vllm/entrypoints/generate/api_router.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index ec6dd89c3ef8..93f6d14a3b9c 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -254,6 +254,8 @@ def __init__( tool_strict_level: str = "auto", enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_log_outputs: bool = False, + enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, is_reasoning_model: bool = True, ) -> None: @@ -272,6 +274,8 @@ def __init__( tool_parser=tool_parser, enable_prompt_tokens_details=enable_prompt_tokens_details, enable_force_include_usage=enable_force_include_usage, + enable_log_outputs=enable_log_outputs, + enable_log_deltas=enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, ) # Controls how the assistant's chain-of-thought is surfaced on diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 9bbe4e879b34..942a1c3e78c2 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -208,6 +208,8 @@ async def init_generate_state( reasoning_parser=args.structured_outputs_config.reasoning_parser, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, default_chat_template_kwargs=default_chat_template_kwargs, is_reasoning_model=args.cohere_is_reasoning_model, ) From 253e02b95e9c9deaebdc507f2eb26b29daf3fdf3 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Tue, 13 Oct 2026 14:17:48 +0200 Subject: [PATCH 492/952] [Perf] Warm up DeepGEMM MegaMoE kernels before serving (#60468) Signed-off-by: Simon Veitner Co-authored-by: Claude Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> --- tests/model_executor/test_deep_gemm_warmup.py | 51 ++++++ .../model_executor/warmup/deep_gemm_warmup.py | 152 +++++++++++++++++- vllm/models/kimi_k3/nvidia/model.py | 1 + 3 files changed, 203 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_deep_gemm_warmup.py b/tests/model_executor/test_deep_gemm_warmup.py index e71061e73317..c2dc12a27d97 100644 --- a/tests/model_executor/test_deep_gemm_warmup.py +++ b/tests/model_executor/test_deep_gemm_warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import sys from types import SimpleNamespace import pytest @@ -88,3 +89,53 @@ def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: provider = getattr(layer, "deep_gemm_warmup_provider", None) assert provider is (None if is_bmm else kernel) + + +def test_mega_moe_predicate_matches_only_native_mega_moe_layers(monkeypatch) -> None: + """Without the MegaMoE model modules loaded, nothing matches and nothing is + imported; with them loaded, only DeepSeek-V4 and Kimi K3 MegaMoE layers match.""" + name = "vllm.models.deepseek_v4.nvidia.model" + kimi_name = "vllm.models.kimi_k3.nvidia.model" + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.delitem(sys.modules, kimi_name, raising=False) + other = SimpleNamespace(use_native_mega_moe=True, use_mega_moe=True) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other) + assert name not in sys.modules and kimi_name not in sys.modules + + moe_cls = type("DeepseekV4MoE", (), {}) + monkeypatch.setitem(sys.modules, name, SimpleNamespace(DeepseekV4MoE=moe_cls)) + native, fused = moe_cls(), moe_cls() + native.use_native_mega_moe, fused.use_native_mega_moe = True, False + assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(native) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(fused) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other) + + kimi_cls = type("KimiMoE", (), {}) + monkeypatch.setitem(sys.modules, kimi_name, SimpleNamespace(KimiMoE=kimi_cls)) + kimi_mega, kimi_fused = kimi_cls(), kimi_cls() + kimi_mega.use_mega_moe, kimi_fused.use_mega_moe = True, False + assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_mega) + assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_fused) + + +def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None: + """Each bf16_mega_gate config is warmed with an M above the GateLinear + threshold (1 for 128 experts), so the forward actually reaches the kernel.""" + dg = SimpleNamespace( + get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)}, + get_block_m_for_mega_moe=lambda *_: 16, + ) + monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg) + monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True) + monkeypatch.setattr( + deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8) + ) + buffer = SimpleNamespace(num_max_tokens_per_rank=64) + module = SimpleNamespace( + gate=SimpleNamespace(input_size=4096, output_size=128), + experts=SimpleNamespace( + top_k=6, num_experts=128, get_symm_buffer=lambda: buffer + ), + use_sequence_parallel=False, + ) + assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40] diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index a4775ce22e45..309c66f07e25 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -5,11 +5,19 @@ be used during model execution beforehand. """ +import sys +from typing import TYPE_CHECKING, TypeGuard + import torch from tqdm import tqdm import vllm.envs as envs -from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank +from vllm.distributed.parallel_state import ( + get_dp_group, + get_ep_group, + is_global_first_rank, +) +from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( DeepGemmFp8BlockScaledMMKernel, ) @@ -23,6 +31,7 @@ ) from vllm.tracing import instrument from vllm.utils.deep_gemm import ( + _import_deep_gemm, fp8_gemm_nt, get_mk_alignment_for_contiguous_layout, m_grouped_fp8_gemm_nt_contiguous, @@ -32,6 +41,10 @@ from vllm.utils.platform_utils import num_compute_units from vllm.v1.worker.workspace import current_workspace_manager +if TYPE_CHECKING: + from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MoE + from vllm.models.kimi_k3.nvidia.model import KimiMoE + def _generate_optimal_warmup_m_values( max_tokens: int, n: int, device: torch.device @@ -144,6 +157,22 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool: return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts)) +# Model construction already imports the model modules. Avoid importing them here. +def _mega_moe_uses_mega_gate(module: torch.nn.Module) -> "TypeGuard[DeepseekV4MoE]": + # DeepSeek-V4 routes with bf16_mega_gate, Kimi K3 with its GateLinear. + dsv4 = sys.modules.get("vllm.models.deepseek_v4.nvidia.model") + return dsv4 is not None and isinstance(module, dsv4.DeepseekV4MoE) + + +def _mega_moe_may_use_deep_gemm( + module: torch.nn.Module, +) -> "TypeGuard[DeepseekV4MoE | KimiMoE]": + if _mega_moe_uses_mega_gate(module): + return module.use_native_mega_moe + kimi = sys.modules.get("vllm.models.kimi_k3.nvidia.model") + return kimi is not None and isinstance(module, kimi.KimiMoE) and module.use_mega_moe + + FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set() @@ -305,6 +334,110 @@ def _warmup(w: torch.Tensor, w_scale: torch.Tensor): GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size()) +FP8_FP4_MEGA_MOE_WARMUP_CACHE: set[tuple] = set() + + +def _get_fp8_fp4_mega_moe_warmup_key(module: "DeepseekV4MoE | KimiMoE") -> tuple: + gate = module.gate + return ( + gate.output_size, + getattr(gate, "tid2eid", None) is not None, + gate.e_score_correction_bias is not None, + getattr(module, "image_sentinel_lo", 0), + module.experts.has_fused_shared_experts, + ) + + +def _get_fp8_fp4_mega_moe_m_values( + module: "DeepseekV4MoE | KimiMoE", max_tokens: int +) -> list[int]: + """Get the largest M per distinct fp8_fp4_mega_moe config and, where used, + bf16_mega_gate config.""" + dg = _import_deep_gemm() + gate, experts = module.gate, module.experts + uses_mega_gate = _mega_moe_uses_mega_gate(module) + num_max_tokens_per_rank = experts.get_symm_buffer().num_max_tokens_per_rank + if module.use_sequence_parallel: + max_tokens = cdiv(max_tokens, module.tp_size) + + m_values: dict[tuple, int] = {} + for num_tokens in range(1, max_tokens + 1): + gate_config = ( + dg.get_bf16_mega_gate_config( + num_tokens, gate.input_size, gate.output_size, experts.top_k + ) + if uses_mega_gate + else {} + ) + block_m = dg.get_block_m_for_mega_moe( + get_ep_group().world_size, + experts.num_experts, + num_max_tokens_per_rank, + num_tokens, + experts.top_k, + "fp8xfp4", + ) + # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a + # small token threshold, and a config's largest M is always above it. + m_values[(tuple(gate_config.items()), block_m)] = num_tokens + return sorted(m_values.values()) + + +def _deepgemm_fp8_fp4_mega_moe_warmup( + module: "DeepseekV4MoE | KimiMoE", + max_tokens: int, + pbar: tqdm | None = None, +): + from vllm.models.deepseek_v4.nvidia.model import ( + prepare_mega_gate_routing_metadata, + ) + + key = _get_fp8_fp4_mega_moe_warmup_key(module) + if key in FP8_FP4_MEGA_MOE_WARMUP_CACHE: + return + + m_values = _get_fp8_fp4_mega_moe_m_values(module, max_tokens) + max_m = m_values[-1] # sorted + + device = module.experts.get_symm_buffer().x.device + input_ids = torch.zeros(max_m, device=device, dtype=torch.int64) + hidden_states = torch.zeros( + (max_m, module.gate.input_size), device=device, dtype=torch.bfloat16 + ) + is_padding = torch.zeros(max_m + 1, device=device, dtype=torch.bool) + # With padding skip on, sequence parallel gives ranks > 0 a padding mask + # slice at any alignment. + sp_masked = module.use_sequence_parallel and envs.VLLM_MOE_SKIP_PADDING + offsets = (0, 1) if sp_masked else (0,) + + # Run the whole layer so the gate sees the same optional inputs as serving. + for num_tokens in m_values: + inputs: tuple = (hidden_states[:num_tokens],) + if _mega_moe_uses_mega_gate(module): + ids = input_ids[:num_tokens] + inputs += ( + ids, + prepare_mega_gate_routing_metadata( + ids, + has_hash_routing=module.gate.tid2eid is not None, + image_sentinel_base_id=module.image_sentinel_lo or None, + ), + ) + for offset in offsets: + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + is_padding=is_padding[offset : offset + num_tokens], + ) + with override_forward_context(forward_context): + module(*inputs) + if pbar is not None: + pbar.update(1) + + FP8_FP4_MEGA_MOE_WARMUP_CACHE.add(key) + + def deepgemm_fp8_gemm_nt_warmup( model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None ): @@ -330,11 +463,21 @@ def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ) +@torch.inference_mode() +def deepgemm_fp8_fp4_mega_moe_warmup( + model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None +): + for m in model.modules(): + if _mega_moe_may_use_deep_gemm(m): + _deepgemm_fp8_fp4_mega_moe_warmup(m, max_tokens, pbar=pbar) + + def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE) seen_grouped_sizes: set[torch.Size] = set( GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE ) + seen_mega_moe_keys: set[tuple] = set(FP8_FP4_MEGA_MOE_WARMUP_CACHE) total = 0 for m in model.modules(): @@ -355,6 +498,11 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: if w2.size() not in seen_grouped_sizes: total += n_values seen_grouped_sizes.add(w2.size()) + elif _mega_moe_may_use_deep_gemm(m): + key = _get_fp8_fp4_mega_moe_warmup_key(m) + if key not in seen_mega_moe_keys: + total += len(_get_fp8_fp4_mega_moe_m_values(m, max_tokens)) + seen_mega_moe_keys.add(key) return total @@ -369,6 +517,8 @@ def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int): with tqdm(total=total, desc="DeepGEMM warmup") as pbar: deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar) deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar) + deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, pbar) else: deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None) deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None) + deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, None) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index eb2ac88aa436..ae9b7da3cd11 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -566,6 +566,7 @@ def __init__( ) self.latent_moe_use_norm = config.latent_moe_use_norm self.tp_size = get_tensor_model_parallel_world_size() + self.use_sequence_parallel = use_sequence_parallel self.routed_scaling_factor = config.routed_scaling_factor self.moe_renormalize = moe_renormalize self.use_grouped_topk = config.use_grouped_topk From 7342dc0317c02af21aa15109f147f98982cb8c4f Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:19:53 +0800 Subject: [PATCH 493/952] [MoRIIO] Require pending remote prefill for READ cache hits (#60638) Signed-off-by: whx-sjtu --- .../unit/test_moriio_hma_scheduler.py | 40 ++++++++++++++++++- .../v1/moriio/moriio_connector.py | 6 +-- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 4bd32f786a28..c6df14d7a0fe 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -806,12 +806,50 @@ def test_truncate_mamba_request_noop_for_single_token_prompt(): # -------------------------------------------------------------------------- # get_num_new_matched_tokens (hybrid N-1 accounting) # -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "params,expected", + [ + (None, (0, False)), + ({}, (0, False)), + ({"do_remote_prefill": False}, (0, False)), + ({"do_remote_prefill": True}, (31, False)), + ], +) +def test_read_matches_only_requests_with_pending_remote_prefill(params, expected): + scheduler = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ) + request = SimpleNamespace(num_prompt_tokens=32, kv_transfer_params=params) + + assert scheduler.get_num_new_matched_tokens(request, 0) == expected + + +def test_read_rematch_after_alloc_keeps_queued_transfer(): + """Clearing the flag blocks another external hit, not the queued READ.""" + sched = _make_read_scheduler() + request = _make_read_request([[10, 11], [91]]) + request.num_prompt_tokens = 32 + assert sched.get_num_new_matched_tokens(request, 0) == (31, False) + + sched.update_state_after_alloc( + request, _FakeBlocks(all_groups=([100, 101], [200])), num_external_tokens=31 + ) + + assert sched._reqs_need_recv["req"][1] == [[100, 101], [200]] + assert request.kv_transfer_params["do_remote_prefill"] is False + + meta = sched.build_connector_meta(SchedulerOutput.make_empty()) + queued = meta.reqs_to_recv["req"] + assert queued.transfer_id == "tx" + assert queued.local_block_ids == [[100, 101], [200]] + assert queued.remote_block_ids == [[10, 11], [91]] + assert sched.get_num_new_matched_tokens(request, 0) == (0, False) + + def test_get_num_new_matched_tokens_read_recomputes_last_token(): sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ, _has_mamba=True) req = SimpleNamespace( num_prompt_tokens=10, prompt_token_ids=list(range(10)), - kv_transfer_params=None, + kv_transfer_params={"do_remote_prefill": True}, ) n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0) # READ always recomputes the final token locally: N-1 - computed. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index c398cb20a8b0..0e6eccd8ea91 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -776,10 +776,10 @@ def get_num_new_matched_tokens( num_prompt_tokens = request.num_prompt_tokens num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0) + params = request.kv_transfer_params or {} + if not params.get("do_remote_prefill"): + return 0, False if self.mode == MoRIIOMode.WRITE: - params = request.kv_transfer_params or {} - if not params.get("do_remote_prefill"): - return 0, False # MoriiO in write mode, no remote prefill. Hybrid models never get # here: register_kv_caches refuses WRITE mode for them, so there is # no recurrent-state accounting to do. From 0838cd6593ba751aa215c9ecd44048bb29dcff1b Mon Sep 17 00:00:00 2001 From: Nathan Sala <33957625+natsala13@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:42:29 +0300 Subject: [PATCH 494/952] [Bugfix] Load and apply Nemotron Super final vision LayerNorm (#61401) Signed-off-by: Shaun Kotek Co-authored-by: Shaun Kotek - Nvidia --- .../multimodal/test_nano_nemotron_vl.py | 57 +++++++++++++++++++ .../model_executor/models/nano_nemotron_vl.py | 35 +++++++++++- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 75900cde53a3..41f9428e6d73 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -97,6 +97,7 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", _MissingMultiModalModule()) + object.__setattr__(model, "vision_final_layernorm", _MissingMultiModalModule()) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -106,6 +107,8 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): ("mlp1.0.weight", object()), ("vision_model.radio_model.encoder.weight", object()), ("sound_encoder.encoder.weight", object()), + ("vision_projector.vision_final_layernorm.weight", object()), + ("vision_projector.vision_final_layernorm.bias", object()), ] ) @@ -120,6 +123,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -145,12 +149,65 @@ def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) with pytest.raises(AssertionError): model.load_weights([("sound_encoder.encoder.weight", object())]) +def _vision_norm_model(use_layernorm=True): + model = object.__new__(NemotronH_Nano_VL_V2) + torch.nn.Module.__init__(model) + model.model_config = _ImageOnlyModelConfig() + model.language_model = _LanguageModel() + model.vision_model = _VisionModel() + model.sound_encoder = None + model.mlp1 = torch.nn.Identity() + model.vision_final_layernorm = ( + torch.nn.LayerNorm(4, eps=1e-6, dtype=torch.bfloat16) if use_layernorm else None + ) + model.patch_size = 2 + model.downsample_ratio = 0.5 + model.ps_version = "v2" + model.video_temporal_patch_size = 2 + return model + + +@pytest.mark.parametrize("mode", ["image", "dynamic", "video"]) +@pytest.mark.parametrize("use_layernorm", [True, False]) +def test_vision_layernorm_is_loaded_and_applied_before_pixel_shuffle( + mode, use_layernorm +): + model = _vision_norm_model(use_layernorm) + weight = torch.tensor([2.0, 3.0, 4.0, 5.0], dtype=torch.bfloat16) + bias = torch.tensor([0.5, -0.5, 1.0, -1.0], dtype=torch.bfloat16) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + if use_layernorm + else [] + ) + features = torch.arange(16, dtype=torch.bfloat16).reshape(1, 4, 4) + expected = ( + torch.nn.functional.layer_norm(features, (4,), weight, bias, eps=1e-6) + if use_layernorm + else features + ) + expected = model.pixel_shuffle(expected.reshape(1, 2, 2, 4)).reshape(1, 1, 16) + model.vision_model = lambda *args, **kwargs: (None, features) + pixels = torch.zeros(2 if mode == "video" else 1, 3, 4, 4) + if mode == "dynamic": + output = model.extract_feature_dynamic(pixels, imgs_sizes=[(4, 4)]) + else: + output = model.extract_feature( + pixels, num_frames=2 if mode == "video" else None + ) + torch.testing.assert_close(output, expected) + + def _make_mm_items_with_video_bytes( video_bytes: bytes, ) -> MultiModalDataItems: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 074a89abb93f..b77206295d5c 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1001,6 +1001,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vision_projection_hidden_size = config.projector_hidden_size llm_hidden_size = config.text_config.hidden_size + # Match the HF vision projector for checkpoints trained with MTP. + self.vision_final_layernorm = ( + nn.LayerNorm( + vit_hidden_size, + eps=getattr(vision_config, "layer_norm_eps", 1e-6), + ).to(llm_dtype) + if (getattr(config.text_config, "num_nextn_predict_layers", 0) or 0) > 0 + else None + ) + mlp1 = nn.Sequential( RMSNorm( hidden_size=vit_hidden_size @@ -1095,6 +1105,8 @@ def extract_feature_dynamic( """Dynamic resolution extract_feature for images.""" _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = self.pixel_shuffle_dynamic_res(vit_embeds, imgs_sizes=imgs_sizes) vit_embeds = self.mlp1(vit_embeds) return vit_embeds @@ -1127,6 +1139,8 @@ def extract_feature( else: _, vit_embeds = self.vision_model(chunk) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = vit_embeds.reshape( vit_embeds.shape[0], H_patches, W_patches, -1 ) @@ -1575,7 +1589,7 @@ def get_mm_mapping(self) -> MultiModelKeys: """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", - connector=["mlp1", "sound_encoder.projection"], + connector=["mlp1", "sound_encoder.projection", "vision_final_layernorm"], tower_model=["vision_model", "sound_encoder.encoder"], ) @@ -1593,6 +1607,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): for modality in ("image", "video", "audio") ) adapter_dict = dict(self.mlp1.named_parameters()) + vision_norm_params = ( + dict(self.vision_final_layernorm.named_parameters()) + if load_multimodal_weights and self.vision_final_layernorm is not None + else {} + ) def is_llm(name: str) -> bool: return name.startswith("language_model") @@ -1606,6 +1625,11 @@ def is_vision_weights(name: str) -> bool: def is_sound_weights(name: str) -> bool: return name.startswith("sound") + vision_norm_prefix = "vision_projector.vision_final_layernorm." + + def is_vision_norm_weights(name: str) -> bool: + return name.startswith(vision_norm_prefix) + # LLM weights (the bulk of the model) are streamed lazily through a # generator so each tensor is copied into its parameter before the # iterator advances, avoiding stale-reference corruption with @@ -1637,6 +1661,15 @@ def llm_weights_gen(): continue assert self.sound_encoder is not None sound_weights.append((name, w.detach().clone())) + elif is_vision_norm_weights(name): + if not load_multimodal_weights: + continue + param_name = name.removeprefix(vision_norm_prefix) + assert param_name in vision_norm_params, ( + f"Unexpected vision LayerNorm weight: {name}" + ) + with torch.no_grad(): + default_weight_loader(vision_norm_params[param_name], w) # Fully drain the generator so every mm tensor is buffered, even if # the LLM loader stops iterating early. From 31e90ac3705e7415e4061d72c640c6c01737af15 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Tue, 13 Oct 2026 15:31:37 +0200 Subject: [PATCH 495/952] [Attention] Enable FA4 sparse MQA decode on Blackwell (#55866) Signed-off-by: Simon Veitner Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Co-authored-by: Lucas Wilkinson Co-authored-by: Matthew Bonanni --- cmake/external_projects/vllm_flash_attn.cmake | 2 +- docs/design/attention_backends.md | 2 +- .../attention/test_flashinfer_mla_decode.py | 10 +- ...st_flashmla_nope_sm90_backend_selection.py | 34 ++ .../v1/attention/test_sparse_mla_backends.py | 475 ++++++++++++++---- .../test_sparse_mla_kv_cache_layout.py | 27 +- tests/v1/e2e/general/test_hisparse.py | 12 +- .../layers/attention/mla_attention.py | 6 +- .../layers/attention/sparse_mla_attention.py | 159 +++++- vllm/platforms/cuda.py | 3 +- .../backends/mla/flashattn_mla_sparse.py | 313 +++++++----- .../backends/mla/flashinfer_mla_sparse.py | 172 ++----- vllm/v1/attention/backends/registry.py | 4 + vllm/vllm_flash_attn/flash_attn_interface.py | 24 +- 14 files changed, 860 insertions(+), 383 deletions(-) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a26023f740b..bac2afc828b3 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f + GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d175add40b84..21cde2fbed16 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -147,7 +147,7 @@ Priority is **1 = highest** (tried first). --8<-- "gen:priority-mla" -> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. +> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASH_ATTN_MLA_SPARSE_FA4` (then `FLASHINFER_MLA_SPARSE`) is preferred for low query-head counts (<= 16) on SM 10.x, while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. diff --git a/tests/kernels/attention/test_flashinfer_mla_decode.py b/tests/kernels/attention/test_flashinfer_mla_decode.py index c3f7b4c96f78..fc34b8725746 100644 --- a/tests/kernels/attention/test_flashinfer_mla_decode.py +++ b/tests/kernels/attention/test_flashinfer_mla_decode.py @@ -175,7 +175,15 @@ def unexpected_helper(*args, **kwargs): backend, "prepare_sparse_mla_safe_lengths", unexpected_helper ) - out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts) + out, lse = impl._forward_mqa_kernel( + query.squeeze(1), + kv_cache, + indices, + counts, + layer=None, + block_size=block_size, + is_decode=True, + ) reference = trtllm_batch_decode_with_kv_cache_mla( query=query, kv_cache=kv_cache.unsqueeze(1), diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 8532d2f43aee..8f31bd488aa4 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py @@ -238,6 +238,7 @@ def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100): ) sparse = { "FLASH_ATTN_MLA_SPARSE", + "FLASH_ATTN_MLA_SPARSE_FA4", "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM90", @@ -276,3 +277,36 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla(): for name in order[: order.index("FLASHMLA_SPARSE")]: backend_cls = AttentionBackendEnum[name].get_class() assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends") +def test_sm100_576_bf16_low_heads_falls_back_from_fa4(monkeypatch): + """With FA4 unusable, low-head bf16 selection lands on FlashInfer.""" + import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: "no FA4") + + def reasons(name): + backend_cls = AttentionBackendEnum[name].get_class() + return backend_cls.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=SM100, + attn_type="decoder", + ) + + order = _sparse_order("auto", 576, num_heads=16) + with set_current_vllm_config(VllmConfig()): + # The selector takes the first candidate with no rejection reasons. + selected = next(name for name in order if not reasons(name)) + assert order[0] == "FLASH_ATTN_MLA_SPARSE_FA4", order + assert selected == "FLASHINFER_MLA_SPARSE" diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index f7c47fe6a3ec..a2e216994328 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -34,6 +34,9 @@ ) from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention +from vllm.model_executor.layers.attention import ( + sparse_mla_attention as sparse_mla_module, +) from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, @@ -46,6 +49,7 @@ ) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability # This module's end-to-end sparse-backend correctness test supports CUDA and # ROCm. The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs @@ -71,8 +75,11 @@ flashattn_mla_sparse as flashattn_sparse_module, ) from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( + FlashAttnMLASparseFA4Backend, + FlashAttnMLASparseFA4Impl, FlashAttnMLASparseImpl, FlashAttnMLASparseMetadataBuilder, + _fa4_cute_mla_available, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseImpl, @@ -227,6 +234,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch): req_id_per_token=torch.zeros(1, dtype=torch.int32), block_table=torch.zeros((1, 1), dtype=torch.int32), block_size=32, + num_decode_tokens=1, ) recorded_scale = None @@ -1401,6 +1409,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect ("FLASHMLA_SPARSE", 8, 112 * 1024), ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024), ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024), + ("FLASH_ATTN_MLA_SPARSE_FA4", 4, 36 * 1024), ], ) def test_masked_mha_workspace_guards_long_routing_policy( @@ -3569,56 +3578,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes(): assert builder.reorder_batch_threshold == 4 -def test_flashinfer_hisparse_decode_runs_batched_attention(): - device = torch.device("cpu") - num_tokens = 6 - q = torch.randn(num_tokens, 2, 4, device=device) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device) - physical_topk = topk + 10 - valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device) - kernel_shapes: list[torch.Size] = [] - - def convert_decode(self, *args, **kwargs): # noqa: ARG001 - return physical_topk, valid_counts - - def prepare_kernel(self, *args, **kwargs): # noqa: ARG001 - pass - - def run_kernel(self, q, cache, indices, counts): # noqa: ARG001 - kernel_shapes.append(q.shape) - return q[..., :1], None - - cache_handle = SimpleNamespace( - runtime=SimpleNamespace( - hot=SimpleNamespace(attention_cache=torch.empty(1, device=device)) - ), - ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [cache_handle] - index_group.convert_logical_to_physical_topk = MethodType( - convert_decode, index_group - ) - impl = object.__new__(FlashInferMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl) - impl._run_mqa_kernel = MethodType(run_kernel, impl) - metadata = SimpleNamespace(num_decode_tokens=num_tokens) - - output, lse = FlashInferMLASparseImpl.forward_mqa( - impl, - q, - torch.empty(1, device=device), - metadata, - SimpleNamespace(), - ) - - assert kernel_shapes == [q.shape] - assert output.shape == (num_tokens, 2, 1) - assert lse is None - - @pytest.mark.skipif( not current_platform.is_device_capability_family(90), reason="FlashAttention MLA requires Hopper", @@ -3690,7 +3649,7 @@ def test_flashattn_shared_indices_match_independent_conversion( convert = triton_convert_req_index_to_global_index convert_spy = MagicMock(wraps=convert) - for module in (index_group_module, flashattn_sparse_module): + for module in (index_group_module, sparse_mla_module): monkeypatch.setattr( module, "triton_convert_req_index_to_global_index", convert_spy ) @@ -3735,48 +3694,19 @@ def run(): return_valid_counts=True, ) for impl, cache, actual in zip(impls, caches, outputs): - expected = impl._run_mqa_kernel( - q_nope, q_rope, cache, physical, counts, block_size + expected, _ = impl._forward_mqa_kernel( + (q_nope, q_rope), + cache, + physical, + counts, + layer=None, + block_size=block_size, + is_decode=num_decode_tokens == num_tokens, ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) stream.synchronize() -def test_flashattn_hisparse_decode_uses_index_group(): - num_tokens = 4 - q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) - q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE) - topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE) - physical = topk.clone() - counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.convert_logical_to_physical_topk = MagicMock( - return_value=(physical, counts) - ) - index_group.physical_kv_cache = MagicMock( - return_value=torch.empty(1, device=DEVICE_TYPE) - ) - impl = object.__new__(FlashAttnMLASparseImpl) - impl.topk_indices_buffer = topk - impl.index_group = index_group - impl.index_group_index = 0 - impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1]) - metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64) - - output, lse = FlashAttnMLASparseImpl.forward_mqa( - impl, - (q_nope, q_rope), - torch.empty(1, device=DEVICE_TYPE), - metadata, - SimpleNamespace(), - ) - - assert output.shape == (num_tokens, 2, 1) - assert lse is None - index_group.convert_logical_to_physical_topk.assert_called_once() - impl._run_mqa_kernel.assert_called_once() - - def test_flashinfer_sm120_hisparse_decode_uses_index_group(): num_tokens = 4 q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE) @@ -3811,6 +3741,131 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group(): impl._run_mqa_kernel.assert_called_once() +def test_sparse_mqa_plain_lane_wiring(monkeypatch): + """Under DCP the plain lane converts with the DCP filter at the flat stride.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6 + # Flat row stride (8) differs from block size (4): the wrong one misaddresses. + kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size] + q = torch.zeros(num_tokens, 2, head_dim) + metadata = SimpleNamespace( + num_decode_tokens=num_tokens, + block_size=block_size, + block_table=torch.zeros((num_tokens, 1), dtype=torch.int32), + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32), + cp_kv_cache_interleave_size=2, + ) + index_filter = MagicMock(return_value=(None, None)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", index_filter) + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = MagicMock() + impl.index_group_index = 0 + impl.dcp_world_size = 2 + impl.dcp_rank = 1 + impl._convert_logical_to_physical_topk = MagicMock() + impl._forward_mqa_kernel = MagicMock(return_value=(q, None)) + + impl.forward_mqa(q, kv_cache, metadata, None) + + impl._convert_logical_to_physical_topk.assert_not_called() + assert index_filter.call_args.kwargs["BLOCK_STRIDE_ROWS"] == 8 + impl._forward_mqa_kernel.assert_called_once() + + +@pytest.mark.parametrize( + "num_decode_tokens,all_resident,split_q", + [ + pytest.param(0, True, False, id="resident"), + pytest.param(0, False, False, id="staged"), + pytest.param(2, True, True, id="mixed"), + pytest.param(5, True, False, id="decode"), + ], +) +def test_sparse_mqa_hisparse_lanes( + monkeypatch, num_decode_tokens, all_resident, split_q +): + """HiSparse splits a batch into at most one decode and one prefill call.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + num_tokens, topk_tokens, block_size, head_dim = 5, 8, 4, 6 + num_prefill = num_tokens - num_decode_tokens + # Distinct block counts, so the KV tensor a call is handed names its lane. + host = torch.zeros(2, block_size, head_dim, dtype=torch.bfloat16) + hot = torch.zeros(6, block_size, head_dim, dtype=torch.bfloat16) + staged = torch.zeros(3, block_size, head_dim, dtype=torch.bfloat16) + q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1) + metadata = SimpleNamespace( + num_decode_tokens=num_decode_tokens, block_size=block_size + ) + + def converted(rows, fill): + return ( + torch.full((rows, topk_tokens), fill, dtype=torch.int32), + torch.full((rows,), rows, dtype=torch.int32), + ) + + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.physical_kv_cache = MagicMock(return_value=hot) + index_group.cache = MagicMock( + return_value=SimpleNamespace(all_context_pages_resident=all_resident) + ) + # The fill value names the converter that produced a call's indices. + index_group.convert_logical_to_physical_topk = MagicMock( + return_value=converted(num_decode_tokens, 1) + if num_decode_tokens + else converted(num_tokens, 2) + ) + index_group.stage_prefill_rows = MagicMock(return_value=(staged, None, None)) + calls: list[tuple] = [] + + def convert_staged(*args, **kwargs): + # Decode rows are consumed before a later conversion can reuse their buffer. + assert len(calls) == min(num_decode_tokens, 1) + return converted(num_prefill, 3) + + monkeypatch.setattr( + sparse_mla, + "triton_convert_req_index_to_global_index", + MagicMock(side_effect=convert_staged), + ) + + def hook(rows, cache, indices, counts, *, is_decode, **kwargs): + if split_q: + assert rows[0].tolist() == rows[1].tolist() + rows = rows[0] + calls.append((rows.flatten().tolist(), cache.shape[0], indices[0, 0].item())) + return torch.full((rows.shape[0], 1, 1), 1.0 if is_decode else 2.0), None + + impl = object.__new__(FlashInferMLASparseImpl) + impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32) + impl.index_group = index_group + impl.index_group_index = 0 + impl.dcp_world_size = 1 + impl._forward_mqa_kernel = hook + + out, lse = impl.forward_mqa((q, q) if split_q else q, host, metadata, None) + + q_rows = q.flatten().tolist() + expected = [] + if num_decode_tokens: + # Decode reads the hot cache through the group's conversion. + expected.append((q_rows[:num_decode_tokens], hot.shape[0], 1)) + if num_decode_tokens == 0 and all_resident: + expected.append((q_rows, hot.shape[0], 2)) + elif num_prefill: + expected.append((q_rows[num_decode_tokens:], staged.shape[0], 3)) + assert calls == expected + if not num_prefill: + # A decode-only batch returns before it looks up the prefill cache. + index_group.cache.assert_not_called() + assert lse is None + # A lane handed the whole batch would concatenate to more rows than this. + assert out.flatten().tolist() == [1.0] * num_decode_tokens + [2.0] * num_prefill + + def test_hisparse_resident_prefill_uses_attention_block_stride(): """A batch beyond the swap-row capacity must bypass swap_in. @@ -3887,6 +3942,250 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only(): assert not cache.decode_batch +def _fa4_impl(topk_indices_buffer, **fields): + impl = object.__new__(FlashAttnMLASparseFA4Impl) + impl.__dict__.update( + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + scale=1.0, + kv_cache_dtype="auto", + dcp_world_size=1, + dcp_rank=0, + need_to_return_lse_for_decode=False, + is_nope_mla=False, + _workspace_buffer=None, + bmm1_scale=None, + bmm2_scale=None, + # A non-index-producing layer does not wait on the indexer's ready event. + index_group=SparseMLAIndexGroupBuilder(topk_indices_buffer).register_layer( + False + )[0], + index_group_index=0, + ) + impl.__dict__.update(fields) + # topk_indices_buffer is a data descriptor, so it needs its own setter. + impl.topk_indices_buffer = topk_indices_buffer + return impl + + +def _fa4_inputs(counts, *, num_heads=16, topk=128, num_blocks=4, device=DEVICE_TYPE): + """Decode inputs whose token ``i`` has ``counts[i]`` valid rows, -1 padded.""" + device = torch.device(device) + block_size, kv_lora_rank, rope_dim = 64, 512, 64 + num_tokens = len(counts) + torch.manual_seed(0) + + def values(*shape): + return torch.rand(shape, dtype=torch.bfloat16, device=device) - 0.5 + + topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device) + for tok, count in enumerate(counts): + topk_indices[tok, :count] = torch.randperm( + num_blocks * block_size, device=device + )[:count].to(torch.int32) + + return SimpleNamespace( + kv_lora_rank=kv_lora_rank, + rope_dim=rope_dim, + block_size=block_size, + topk=topk, + kv_cache=values(num_blocks, block_size, kv_lora_rank + rope_dim), + topk_indices=topk_indices, + valid_counts=torch.tensor(counts, dtype=torch.int32, device=device), + ql_nope=values(num_tokens, num_heads, kv_lora_rank), + q_pe=values(num_tokens, num_heads, rope_dim), + scale=(kv_lora_rank + rope_dim) ** -0.5, + metadata=SimpleNamespace( + req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device), + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, num_blocks + ), + block_size=block_size, + num_decode_tokens=num_tokens, + cp_kv_cache_interleave_size=1, + ), + ) + + +def _record_fa4_kernels(monkeypatch, num_heads, *, device): + """Replace both lanes with recorders whose outputs identify the lane.""" + import flashinfer.decode + + fa4_calls: list[dict] = [] + trtllm_calls: list[dict] = [] + + def lse(rows): + return torch.arange(rows * num_heads, dtype=torch.float32, device=device).view( + rows, num_heads + ) + + def fake_fa4(**kwargs): + fa4_calls.append(kwargs) + rows = kwargs["q"].shape[0] + out = torch.ones(rows, num_heads, 512, dtype=torch.bfloat16, device=device) + return (out, lse(rows)) if kwargs.get("return_softmax_lse") else out + + def fake_trtllm(**kwargs): + trtllm_calls.append(kwargs) + rows = kwargs["query"].shape[0] + out = torch.full( + (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device + ) + return (out, lse(rows)) if kwargs.get("return_lse") else out + + monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4) + monkeypatch.setattr( + flashinfer.decode, "trtllm_batch_decode_with_kv_cache_mla", fake_trtllm + ) + return SimpleNamespace(fa4=fa4_calls, trtllm=trtllm_calls, lse=lse) + + +@pytest.mark.parametrize( + "num_heads,return_lse", [(128, False), (16, True)], ids=["heads128", "heads16_lse"] +) +def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse): + if not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") + if (reason := _fa4_cute_mla_available()) is not None: + pytest.skip(reason) + counts = [0, 1, 127, 128, 129, 2047, 2048] + inputs = _fa4_inputs(counts, num_heads=num_heads, topk=2048, num_blocks=64) + impl = _fa4_impl( + topk_indices_buffer=inputs.topk_indices, + scale=inputs.scale, + need_to_return_lse_for_decode=return_lse, + ) + q = (inputs.ql_nope, inputs.q_pe) + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + if not return_lse: + assert lse is None + else: + out_fused, lse_fused = impl.forward_mqa( + torch.cat(q, dim=-1), inputs.kv_cache, inputs.metadata, None + ) + assert lse.shape == (len(counts), num_heads) + assert torch.equal(out, out_fused) and torch.equal(lse, lse_fused) + + kv_flat = inputs.kv_cache.view(-1, inputs.kv_lora_rank + inputs.rope_dim) + keys = kv_flat[inputs.topk_indices.clamp(min=0).long()].float() + scores = torch.einsum("thd,tkd->thk", torch.cat(q, dim=-1).float(), keys) + scores = (scores * inputs.scale).masked_fill( + (inputs.topk_indices < 0)[:, None, :], float("-inf") + ) + # The all-sentinel row softmaxes to NaN here; the kernel writes (0, -inf). + probs = torch.softmax(scores, dim=-1).nan_to_num() + torch.testing.assert_close( + out.float(), + torch.einsum("thk,tkd->thd", probs, keys[..., : inputs.kv_lora_rank]), + rtol=0.01, + atol=0.01, + ) + if return_lse: + torch.testing.assert_close( + lse, torch.logsumexp(scores, dim=-1), rtol=0, atol=5e-3 + ) + + +@pytest.mark.parametrize("is_decode", [True, False], ids=["decode", "prefill"]) +def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode): + """The hook runs a decode batch on FA4 and every other batch on trtllm-gen.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + counts = [1, 2, 3, 4, 5, 6, 7] + num_tokens, dcp = len(counts), 2 + # Under DCP the hook gets the heads of every DCP rank, not just its own. + kernel_heads = 16 * dcp + inputs = _fa4_inputs(counts, num_heads=kernel_heads, device="cpu") + calls = _record_fa4_kernels(monkeypatch, kernel_heads, device="cpu") + # DCP without PCP (the gate rejects the pair) always fuses the query. + q = torch.cat((inputs.ql_nope, inputs.q_pe), dim=-1) + impl = _fa4_impl( + # One row longer than the batch: the varlen-scalar slices become observable. + topk_indices_buffer=inputs.topk_indices.new_zeros(num_tokens + 1, inputs.topk), + dcp_world_size=dcp, + need_to_return_lse_for_decode=True, + ) + impl._convert_logical_to_physical_topk = MagicMock() + dcp_filter = MagicMock(return_value=(inputs.topk_indices, inputs.valid_counts)) + monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", dcp_filter) + inputs.metadata.num_decode_tokens = num_tokens if is_decode else 1 + + with torch.inference_mode(): + out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None) + + dcp_filter.assert_called_once() + impl._convert_logical_to_physical_topk.assert_not_called() + assert out.shape == (num_tokens, kernel_heads, 512) + if is_decode: + (fa4,) = calls.fa4 + assert calls.trtllm == [] + assert fa4["return_softmax_lse"] is True + assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1)) + assert (out == 1.0).all() + # FA4's LSE is already natural-log. + torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0) + else: + (trtllm,) = calls.trtllm + assert calls.fa4 == [] + assert trtllm["return_lse"] is True + assert (out == 2.0).all() + # trtllm-gen's LSE is base 2; the DCP merge expects natural log. + expected = calls.lse(num_tokens) * math.log(2.0) + torch.testing.assert_close(lse, expected, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize( + "config,expected", + [ + ({}, None), + ({"heads": 128}, None), + ({"heads": 12}, "heads"), + ({"dcp": 2}, None), + ({"dcp": 8}, "heads"), + ({"index_topk": 192}, "index_topk"), + ({"rope": 32}, "(kv_lora_rank, qk_rope_head_dim)"), + ({"flashinfer": False}, "FlashInfer"), + ], +) +def test_fa4_sparse_gate(monkeypatch, config, expected): + """``validate_configuration`` on a fixed SM100: one row per gate reason.""" + has_flashinfer = config.get("flashinfer", True) + monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer) + monkeypatch.setattr( + flashattn_sparse_module, "_fa4_cute_mla_available", lambda: None + ) + dcp = config.get("dcp", 1) + vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp) + hf_text_config = vllm_config.model_config.hf_text_config + hf_text_config.index_topk = config.get("index_topk", 128) + hf_text_config.qk_rope_head_dim = config.get("rope", 64) + + with set_current_vllm_config(vllm_config): + reasons = FlashAttnMLASparseFA4Backend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="auto", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(10, 0), + attn_type="decoder", + use_dcp=dcp > 1, + ) + + if expected is None: + assert reasons == [] + else: + (reason,) = reasons + assert expected in reason + + def test_flashmla_cache_dtype_aliases_use_ds_layout(): # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer # stores kv_cache_dtype, so they cannot bypass the gate. diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index 45ba1ac0dc55..b322acf0d5ee 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -13,15 +13,12 @@ from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.platforms import current_platform from vllm.v1.attention.backend import MultipleOf -from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( - FlashInferMLASparseTRTLLMBackend, -) -from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend from vllm.v1.attention.backends.mla.sparse_utils import ( _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL, flat_kv_row_view, ) +from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs @@ -49,17 +46,27 @@ @pytest.mark.parametrize( "backend,cache_dtype,sm100,stride_alignment", [ - (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152), - (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576), - (FlashMLASparseBackend, "auto", True, 1152), - (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)), - (FlashMLASparseBackend, "fp8_ds_mla", True, 656), - (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "auto", True, 32 * 1152), + (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "fp8", True, 32 * 576), + pytest.param( + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, + "auto", + True, + 64 * 1152, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA only" + ), + ), + (AttentionBackendEnum.FLASHMLA_SPARSE, "auto", True, 1152), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", False, MultipleOf(64)), + (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", True, 656), + (AttentionBackendEnum.FLASHMLA_SPARSE, "nvfp4_ds_mla", True, 352), ], ) def test_allocation_and_warmup_follow_addressing_mode( monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device ): + backend = backend.get_class() monkeypatch.setattr( current_platform, "is_device_capability_family", lambda family: sm100 ) diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 2409b5cb690b..d9e1fabf03f3 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -83,13 +83,20 @@ def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" ) @pytest.mark.parametrize( - "with_offloading", [False, True], ids=["standalone", "offload"] + "with_offloading,attention_backend", + [ + pytest.param(False, None, id="standalone"), + pytest.param(True, None, id="offload"), + # Offloading does not touch the attention kernels: FA4 runs standalone. + pytest.param(False, "FLASH_ATTN_MLA_SPARSE_FA4", id="standalone-fa4"), + ], ) @fork_new_process_for_each_test def test_hisparse_spill_and_prefix_restore( monkeypatch: pytest.MonkeyPatch, vllm_runner: type[VllmRunner], with_offloading: bool, + attention_backend: str | None, ): """Spilled prefixes restore and FULL-graph decode writes reach host KV. @@ -101,6 +108,8 @@ def test_hisparse_spill_and_prefix_restore( capability = current_platform.get_device_capability() if capability is None or capability.major < 9: pytest.skip("Sparse MLA requires Hopper or newer") + if attention_backend and not current_platform.is_device_capability_family(100): + pytest.skip("FA4 sparse MLA requires SM 10.x") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") @@ -142,6 +151,7 @@ def test_hisparse_spill_and_prefix_restore( ) ), kv_transfer_config=kv_transfer_config, + attention_backend=attention_backend, block_size=64, max_model_len=320, max_num_batched_tokens=1024, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index db9e8dbc88ca..806faba50774 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -661,6 +661,7 @@ def __init__( "FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM120", + "FLASH_ATTN_MLA_SPARSE_FA4", "DEEPSEEK_V32_INDEXER", ): from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -751,7 +752,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: def _uses_flat_kv_cache(self) -> bool: backend = self.attn_backend.get_name() - return backend == "FLASHINFER_MLA_SPARSE" or ( + return backend in ("FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4") or ( backend == "FLASHMLA_SPARSE" and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla") ) @@ -1857,6 +1858,9 @@ def _use_masked_mha( seq_len: int, has_context: bool, ) -> bool: + if backend_name == "FLASH_ATTN_MLA_SPARSE_FA4": + # Its prefill batches run on FlashInfer's kernel, so it shares those rows. + backend_name = "FLASHINFER_MLA_SPARSE" thresholds = ( _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {}) .get(backend_name, {}) diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 37883593421a..a8fa74f91f62 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -36,12 +36,18 @@ is_quantized_kv_cache, np_to_pinned_tensor, ) -from vllm.v1.attention.backend import AttentionMetadataBuilder +from vllm.v1.attention.backend import AttentionLayer, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, SparseMLAIndexGroup, SparseMLAIndexGroupBuilder, ) +from vllm.v1.attention.backends.mla.sparse_utils import ( + flat_kv_row_view, + triton_convert_req_index_to_global_index, + triton_filter_and_convert_dcp_index, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.dcp import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states @@ -630,6 +636,9 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic """Sparse MLA base with dense and masked-MHA prefill paths.""" is_sparse = True + # Reuse the index group's top-k conversion only for uniform decode batches + # of multi-layer groups; other batches convert per layer. + group_conversion_decode_only: ClassVar[bool] = False def __init__( self, @@ -728,6 +737,154 @@ def _convert_logical_to_physical_topk( return_valid_counts=return_valid_counts, ) + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: T, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self.topk_indices_buffer is not None + num_tokens = q[0].shape[0] if isinstance(q, tuple) else q.shape[0] + topk = self.topk_indices_buffer[:num_tokens] + num_decode_tokens = attn_metadata.num_decode_tokens + block_size = attn_metadata.block_size + index_group = self.index_group + if not isinstance(index_group, HiSparseMLAIndexGroup): + _, block_stride_rows = flat_kv_row_view(kv_c_and_k_pe_cache, block_size) + if self.dcp_world_size > 1: + physical_topk, valid_counts = triton_filter_and_convert_dcp_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + elif self.group_conversion_decode_only and not ( + index_group is not None + and index_group.num_layers > 1 + and num_decode_tokens >= num_tokens + ): + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_tokens], + attn_metadata.block_table, + topk, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + else: + physical_topk, valid_counts = self._convert_logical_to_physical_topk( + topk, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + return self._forward_mqa_kernel( + q, + kv_c_and_k_pe_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=num_decode_tokens >= num_tokens, + ) + + assert self.dcp_world_size == 1 + layer_index = self.index_group_index + hot_cache = index_group.physical_kv_cache(layer_index).view( + kv_c_and_k_pe_cache.dtype + ) + assert hot_cache.is_contiguous(), "HiSparse hot buffer must be contiguous" + decode_out = decode_lse = None + if num_decode_tokens: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk[:num_decode_tokens], + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + q_decode = ( + (q[0][:num_decode_tokens], q[1][:num_decode_tokens]) + if isinstance(q, tuple) + else q[:num_decode_tokens] + ) + decode_out, decode_lse = self._forward_mqa_kernel( + q_decode, + hot_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=True, + ) + if num_decode_tokens >= num_tokens: + return decode_out, decode_lse + + cache = index_group.cache(layer_index) + if num_decode_tokens == 0 and cache.all_context_pages_resident: + physical_topk, valid_counts = index_group.convert_logical_to_physical_topk( + layer_index, + topk, + attn_metadata, + block_stride_rows=None, + return_valid_counts=True, + ) + prefill_cache = hot_cache + else: + prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( + layer_index, kv_c_and_k_pe_cache, attn_metadata + ) + physical_topk, valid_counts = triton_convert_req_index_to_global_index( + req_ids, + block_table, + topk[num_decode_tokens:], + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=topk.shape[1], + return_valid_counts=True, + ) + q_prefill = ( + (q[0][num_decode_tokens:], q[1][num_decode_tokens:]) + if isinstance(q, tuple) + else q[num_decode_tokens:] + ) + prefill_out, prefill_lse = self._forward_mqa_kernel( + q_prefill, + prefill_cache, + physical_topk, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=False, + ) + if decode_out is None: + return prefill_out, prefill_lse + out = torch.cat((decode_out, prefill_out)) + if decode_lse is None: + return out, None + assert prefill_lse is not None + return out, torch.cat((decode_lse, prefill_lse)) + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + raise NotImplementedError + @staticmethod def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool: """Whether this prefill batch's top-k masks fit the workspace.""" diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 401a98f66feb..7a28dcbc96d7 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -105,9 +105,10 @@ def _get_backend_priorities( ] else: # BF16 KV Cache - # Prefer FlashInfer at low head counts (FlashMLA uses padding) + # Prefer FA4, then FlashInfer at low head counts (FlashMLA uses padding) if num_heads is not None and num_heads <= 16: sparse_backends = [ + AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4, AttentionBackendEnum.FLASHINFER_MLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE, ] diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 0016eafedc8a..47347c8fe84e 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from dataclasses import dataclass +from functools import cache from typing import Any, ClassVar import torch @@ -22,11 +24,14 @@ MultipleOf, ) from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla -from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseImpl, + FlashInferMLASparseTRTLLMBackend, + FlashInferMLASparseTRTLLMMetadataBuilder, +) from vllm.v1.attention.backends.mla.sparse_utils import ( align_blocks_to_rows, flat_kv_row_view, - triton_convert_req_index_to_global_index, ) from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec @@ -153,6 +158,8 @@ def update_draft_decode_metadata(self, metadata: SparseMLACommonMetadata) -> Non class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): + group_conversion_decode_only: ClassVar[bool] = True + def __init__( self, num_heads: int, @@ -205,144 +212,23 @@ def __init__( ) self.supports_quant_query_input = False - def forward_mqa( + def _forward_mqa_kernel( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashAttnMLASparseMetadata, + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: if not isinstance(q, tuple): raise NotImplementedError( "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input." ) q_nope, q_rope = q - num_actual_toks = q_rope.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - outputs = [] - if num_decode_tokens: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[:num_decode_tokens], - q_rope[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - if num_decode_tokens < num_actual_toks: - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - prefill_cache = index_group.physical_kv_cache( - self.index_group_index - ).view(kv_c_and_k_pe_cache.dtype) - else: - prefill_cache, block_table, req_ids = ( - index_group.stage_prefill_rows( - self.index_group_index, - kv_c_and_k_pe_cache, - attn_metadata, - ) - ) - physical_topk, valid_counts = ( - triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - ) - outputs.append( - self._run_mqa_kernel( - q_nope[num_decode_tokens:], - q_rope[num_decode_tokens:], - prefill_cache, - physical_topk, - valid_counts, - attn_metadata.block_size, - ) - ) - return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - if ( - index_group is not None - and index_group.num_layers > 1 - and attn_metadata.num_decode_tokens == num_actual_toks - ): - topk_indices, valid_counts = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - else: - topk_indices, valid_counts = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - return ( - self._run_mqa_kernel( - q_nope, - q_rope, - kv_rows, - topk_indices, - valid_counts, - attn_metadata.block_size, - cache_is_flat=True, - ), - None, - ) - - def _run_mqa_kernel( - self, - q_nope: torch.Tensor, - q_rope: torch.Tensor, - kv_cache: torch.Tensor, - topk_indices: torch.Tensor, - valid_counts: torch.Tensor, - block_size: int, - *, - cache_is_flat: bool = False, - ) -> torch.Tensor: - kv_rows = ( - kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0] - ) + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1] v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1) @@ -370,4 +256,167 @@ def _run_mqa_kernel( causal=True, fa_version=3, ) - return out + return out, None + + +def _fa4_cute_mla_available() -> str | None: + """Reason the FA4 cute-DSL MLA entry point is unusable, or None.""" + from vllm.vllm_flash_attn import flash_attn_interface as fa + + if not fa.is_fa_version_supported(4): + return f"FA4 unavailable: {fa.fa_version_unsupported_reason(4)}" + try: + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd # noqa: F401 + except Exception as e: # cute-DSL / cutlass import chain + return f"FA4 cute-DSL entry point failed to import: {e!r}" + return None + + +class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend): + """FA4 for is_decode kernel calls; trtllm-gen for the rest.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + # The qv kernel asserts every descale is None: BF16 KV cache only. + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"] + + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + return [MultipleOf(64)] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN_MLA_SPARSE_FA4" + + @staticmethod + def get_builder_cls() -> type["FlashAttnMLASparseFA4MetadataBuilder"]: + return FlashAttnMLASparseFA4MetadataBuilder + + @staticmethod + def get_impl_cls() -> type[MLAAttentionImpl[Any]]: + return FlashAttnMLASparseFA4Impl + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # 512 latent + 64 RoPE; split back apart for the qv kernel. + return [576] + + @classmethod + def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None: + from vllm.config import get_current_vllm_config_or_none + from vllm.utils.flashinfer import has_flashinfer + + if not has_flashinfer(): + return "FA4 sparse MLA runs its prefill batches on FlashInfer's kernel" + if (reason := _fa4_cute_mla_available()) is not None: + return reason + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + if vllm_config.model_config is not None: + hf_config = vllm_config.model_config.hf_text_config + topk = getattr(hf_config, "index_topk", None) + # gather_kv_indices' last dim must be a whole number of n-tiles. + if topk is None or topk % 128 != 0: + return f"FA4 sparse MLA requires index_topk % 128 == 0, got {topk}" + dims = ( + getattr(hf_config, "kv_lora_rank", None), + getattr(hf_config, "qk_rope_head_dim", None), + ) + if dims != (512, 64): + return ( + "FA4 sparse MLA requires (kv_lora_rank, qk_rope_head_dim) " + f"== (512, 64), got {dims}" + ) + # Under DCP the query is all-gathered: the kernel sees heads * dcp_size. + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + # 128 gathered heads take FA's prefill kernel, unvalidated under DCP. + head_counts = (8, 16, 32, 64) if dcp_size > 1 else (8, 16, 32, 64, 128) + if num_heads * dcp_size not in head_counts: + return ( + f"FA4 sparse MLA requires {head_counts} gathered query heads, " + f"got num_heads={num_heads} * dcp_size={dcp_size}" + ) + return super().supports_combination(*args, **kwargs) + + +class FlashAttnMLASparseFA4MetadataBuilder(FlashInferMLASparseTRTLLMMetadataBuilder): + """The trtllm-gen builder, for the prefill lane's DCP workspace pre-size. + + Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs + one query row per token and bakes the row count into its varlen scalars. + """ + + +@cache +def _fa4_varlen_scalars( + num_tokens: int, num_kv_rows: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Per-token varlen scalars; never evict, CUDA graphs bake in the pointers.""" + cu_seqlens_q = torch.arange(num_tokens + 1, dtype=torch.int32, device=device) + cu_seqlens_k = torch.zeros(num_tokens + 1, dtype=torch.int32, device=device) + seqused_k = torch.full((num_tokens,), num_kv_rows, dtype=torch.int32, device=device) + return cu_seqlens_q, cu_seqlens_k, seqused_k + + +class FlashAttnMLASparseFA4Impl(FlashInferMLASparseImpl): + # FA4 emits a natural-log LSE; prefill converts trtllm-gen's log2 LSE. + lse_base_on_e: bool = True + + def _forward_mqa_kernel( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + topk_indices: torch.Tensor, + valid_counts: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not is_decode: + out, lse = super()._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + valid_counts, + layer=layer, + block_size=block_size, + is_decode=is_decode, + ) + return out, None if lse is None else lse * math.log(2.0) + if isinstance(q, tuple): + ql_nope, q_pe = q + else: + # 16B-aligned halves of the fused query; FA4 reads them in place. + ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :] + kv_rows, _ = flat_kv_row_view(kv_cache, block_size) + num_tokens = q_pe.shape[0] + num_kv_rows = kv_rows.shape[0] + assert self.topk_indices_buffer is not None + cu_seqlens_q, cu_seqlens_k, seqused_k = _fa4_varlen_scalars( + self.topk_indices_buffer.shape[0], num_kv_rows, kv_rows.device + ) + # An all-sentinel row yields (0, -inf), the DCP merge identity; no masking. + kernel_out = flash_attn_varlen_func( + q=q_pe, + k=kv_rows[:, self.kv_lora_rank :].unsqueeze(1), + v=kv_rows[:, : self.kv_lora_rank].unsqueeze(1), + q_v=ql_nope, + max_seqlen_q=1, + cu_seqlens_q=cu_seqlens_q[: num_tokens + 1], + max_seqlen_k=num_kv_rows, + cu_seqlens_k=cu_seqlens_k[: num_tokens + 1], + seqused_k=seqused_k[:num_tokens], + gather_kv_indices=topk_indices, + # Part of FA4's compile key: pass it unconditionally. + gather_kv_valid_length=valid_counts, + softmax_scale=self.scale, + # Causality is in the top-k list; causal=True would clamp flat rows. + causal=False, + fa_version=4, + return_softmax_lse=self.need_to_return_lse_for_decode, + ) + return kernel_out if self.need_to_return_lse_for_decode else (kernel_out, None) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 21eb66aae490..d673868975ab 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -34,8 +34,6 @@ align_blocks_to_rows, flat_kv_row_view, prepare_sparse_mla_safe_lengths, - triton_convert_req_index_to_global_index, - triton_filter_and_convert_dcp_index, ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec @@ -129,14 +127,12 @@ def supports_combination( and parallel_config.decode_context_parallel_size > 1 ): return ( - "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + f"{cls.get_name()} does not support combined PCP+DCP; " "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " "running the rank-local PCP prefill queries" ) if kv_cache_dtype == "fp8_ds_mla": - return ( - "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" - ) + return f"{cls.get_name()} SM10 does not support fp8_ds_mla kv-cache dtype" # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192]. if vllm_config.model_config is not None: @@ -492,135 +488,6 @@ def __init__( # for query and kv_cache (mixed bf16+fp8 is not supported). self.supports_quant_query_input = True - def forward_mqa( - self, - q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: FlashInferMLASparseMetadata, - layer: AttentionLayer, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - if isinstance(q, tuple): - ql_nope, q_pe = q - if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): - q = ql_nope - else: - q = torch.cat(q, dim=-1) - - num_actual_toks = q.shape[0] - - assert self.topk_indices_buffer is not None - topk_indices = self.topk_indices_buffer[:num_actual_toks] - - self._prepare_mqa_kernel(layer, q.device) - - index_group = self.index_group - if isinstance(index_group, HiSparseMLAIndexGroup): - num_decode_tokens = attn_metadata.num_decode_tokens - decode_out: torch.Tensor | None = None - decode_lse: torch.Tensor | None = None - if num_decode_tokens > 0: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices[:num_decode_tokens], - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - decode_out, decode_lse = self._run_mqa_kernel( - q[:num_decode_tokens], - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - if num_decode_tokens == num_actual_toks: - return decode_out, decode_lse - - cache = index_group.cache(self.index_group_index) - if num_decode_tokens == 0 and cache.all_context_pages_resident: - physical_topk, valid_counts = ( - index_group.convert_logical_to_physical_topk( - self.index_group_index, - topk_indices, - attn_metadata, - block_stride_rows=None, - return_valid_counts=True, - ) - ) - return self._run_mqa_kernel( - q, - index_group.physical_kv_cache(self.index_group_index).view( - kv_c_and_k_pe_cache.dtype - ), - physical_topk, - valid_counts, - ) - - prefill_cache, block_table, req_ids = index_group.stage_prefill_rows( - self.index_group_index, kv_c_and_k_pe_cache, attn_metadata - ) - prefill_indices, prefill_lens = triton_convert_req_index_to_global_index( - req_ids, - block_table, - topk_indices[num_decode_tokens:], - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - prefill_out, prefill_lse = self._run_mqa_kernel( - q[num_decode_tokens:], - prefill_cache, - prefill_indices, - prefill_lens, - ) - if decode_out is None: - return prefill_out, prefill_lse - output = torch.cat((decode_out, prefill_out)) - if decode_lse is None: - return output, None - assert prefill_lse is not None - return output, torch.cat((decode_lse, prefill_lse)) - - kv_rows, block_stride_rows = flat_kv_row_view( - kv_c_and_k_pe_cache, attn_metadata.block_size - ) - - if self.dcp_world_size > 1: - topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - dcp_size=self.dcp_world_size, - dcp_rank=self.dcp_rank, - cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size), - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) - else: - topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk( - topk_indices, - attn_metadata, - block_stride_rows=block_stride_rows, - return_valid_counts=True, - ) - - return self._run_mqa_kernel( - q, - # Block sizes and strides have been aligned to 32 rows. - kv_rows.view( - -1, - math.gcd(block_stride_rows, attn_metadata.block_size, 64), - kv_rows.shape[-1], - ), - topk_indices_physical, - seq_lens, - ) - def _prepare_mqa_kernel( self, layer: AttentionLayer, @@ -648,7 +515,6 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: kv_cache = runtime.hot.attention_cache num_tokens = runtime.max_num_reqs topk_tokens = self.topk_indices_buffer.shape[1] - self._prepare_mqa_kernel(layer, kv_cache.device) q_dtype = ( current_platform.fp8_dtype() @@ -679,22 +545,48 @@ def autotune_hisparse_decode(self, layer: AttentionLayer) -> None: dtype=torch.int32, device=kv_cache.device, ) - self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens) + self._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + seq_lens, + layer=layer, + block_size=runtime.hot.block_size, + is_decode=True, + ) - def _run_mqa_kernel( + def _forward_mqa_kernel( self, - q: torch.Tensor, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_cache: torch.Tensor, topk_indices: torch.Tensor, seq_lens: torch.Tensor, + *, + layer: AttentionLayer, + block_size: int, + is_decode: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(q, tuple): + ql_nope, q_pe = q + if q_pe.shape[-1] == 0 and ql_nope.is_contiguous(): + q = ql_nope + else: + q = torch.cat(q, dim=-1) + + self._prepare_mqa_kernel(layer, q.device) assert self._workspace_buffer is not None assert self.bmm1_scale is not None assert self.bmm2_scale is not None from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla - kv_cache = kv_cache.view(q.dtype) + kv_rows, block_stride_rows = flat_kv_row_view( + kv_cache.view(q.dtype), block_size + ) + # Block sizes and strides have been aligned to 32 rows. + kv_cache = kv_rows.view( + -1, math.gcd(block_stride_rows, block_size, 64), kv_rows.shape[-1] + ) # Single-token sparse decode. trtllm-gen requires the q_len_per_request # dim, but the sparse attention mask is fully per-token (each query token diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 207c112a384c..e1c4c8d06741 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -125,6 +125,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" ) + FLASH_ATTN_MLA_SPARSE_FA4 = ( + "vllm.v1.attention.backends.mla.flashattn_mla_sparse." + "FlashAttnMLASparseFA4Backend" + ) MINIMAX_M3_SPARSE = ( "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend" ) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 3c4c0d8948be..c746dd0a54a8 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -214,6 +214,8 @@ def flash_attn_varlen_func( aux_tensors=None, aux_tensor_leading_dims=None, dynamic_causal: "torch.Tensor | None" = None, + gather_kv_indices=None, + gather_kv_valid_length=None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -262,9 +264,10 @@ def flash_attn_varlen_func( testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per - sequence. Mutually exclusive with cu_seqlens_k, and required when + sequence. Mutually exclusive with cu_seqlens_k except on FA4, which reads + cu_seqlens_k as the KV offset and seqused_k as the length. Required when block_table is given. - q_v: optional value tensor fused with q. FA3 only. + q_v: optional value tensor fused with q. FA3 and FA4 only. block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV cache block table. Requires seqused_k. return_softmax_lse: bool. Whether to also return softmax_lse. @@ -289,18 +292,21 @@ def flash_attn_varlen_func( aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only. aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors. dynamic_causal: optional per-sequence causal offsets. FA4 only. + gather_kv_indices: (total_q, k) int32 KV rows per query token, `-1` padded, + k a multiple of 128. Requires q_v, excludes block_table. FA4 only. + gather_kv_valid_length: (total_q,) int32, real entries per row. FA4 only. Return: out: (total, nheads, headdim). - softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The - logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax - normalization factor). + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen), or + (total_q, nheads) on FA4's q_v path. The logsumexp of each row of the matrix + QK^T * scaling (e.g., log of the softmax normalization factor). """ assert cu_seqlens_k is not None or seqused_k is not None, ( "cu_seqlens_k or seqused_k must be provided" ) - assert cu_seqlens_k is None or seqused_k is None, ( + assert fa_version == 4 or cu_seqlens_k is None or seqused_k is None, ( "cu_seqlens_k and seqused_k cannot be provided at the same time" ) assert block_table is None or seqused_k is not None, ( @@ -311,6 +317,9 @@ def flash_attn_varlen_func( f"Fused FP8 output (output_scale) is only supported by FA4, " f"got fa_version={fa_version}" ) + assert fa_version == 4 or ( + gather_kv_indices is None and gather_kv_valid_length is None + ), f"gather_kv_* is only supported by FA4, got fa_version={fa_version}" if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -446,12 +455,15 @@ def flash_attn_varlen_func( q, k, v, + qv=q_v, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, seqused_k=seqused_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, page_table=block_table, + gather_kv_indices=gather_kv_indices, + gather_kv_valid_length=gather_kv_valid_length, softmax_scale=softmax_scale, causal=causal, dynamic_causal=dynamic_causal, From 0fae39a0142c5d296f7938fa753e8f00b5d80171 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 13 Oct 2026 21:37:10 +0800 Subject: [PATCH 496/952] [Perf] per_token_group_quant_8bit: multiply by reciprocal instead of dividing per element (#61343) Signed-off-by: Thien Tran --- .../quantization/fused_kernels/layernorm_utils.cuh | 10 +++++++--- .../quantization/w8a8/fp8/per_token_group_quant.cu | 14 ++++++++------ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/csrc/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh index d633cb9a6e57..8cab1824f278 100644 --- a/csrc/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -516,6 +516,11 @@ __device__ void norm_and_quant( q8x4_t out; + // Group quant multiplies by the reciprocal of the scale (one div per + // group), matching per_token_group_quant_8bit. Per-token fp8 must keep + // exact division for FBGemm parity, so it is excluded here. + constexpr bool mul_by_recip = (group_size > 0) || is_scale_inverted; + float scale_val; if constexpr (group_size > 0) { @@ -528,14 +533,13 @@ __device__ void norm_and_quant( } else { scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size; } - scale_val = - is_scale_inverted ? 1.0f / scale[scale_idx] : scale[scale_idx]; + scale_val = 1.0f / scale[scale_idx]; } else { scale_val = *scale; } #pragma unroll for (int j = 0; j < VEC_SIZE; ++j) { - out.val[j] = ScaledQuant::quant_fn( + out.val[j] = ScaledQuant::quant_fn( static_cast(x.val[j] * rms) * w.val[j], scale_val); } vec_output[i] = out; diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index d2a888643899..0571cc90778d 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -78,9 +78,13 @@ __device__ __forceinline__ void QuantizeGroup( const float y_s, const float min_8bit, const float max_8bit) { constexpr int vec_size = 16 / sizeof(T); + // One div.rn per group instead of per element; the register-resident + // kernel does the same, keeping both schedules bit-identical. + const float inv_y = 1.0f / y_s; + // quantize shared -> global 8-bit auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { - float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); + float q = fminf(fmaxf(static_cast(src) * inv_y, min_8bit), max_8bit); dst = DST_DTYPE(q); }; @@ -272,16 +276,14 @@ __global__ void per_token_group_quant_8bit_register_kernel( } } - // A UE8M0 scale is an exact power of two, so multiplying by the reciprocal - // matches dividing. Otherwise the reciprocal costs an ulp and can move - // values that land exactly on an 8-bit boundary, so divide instead to stay - // bit-exact with the smem kernel. + // One div.rn per group; the smem path (QuantizeGroup) does the same, + // keeping both schedules bit-identical. const float inv_y = 1.0f / y_s; uint32_t packed[VEC_SIZE / 4] = {}; #pragma unroll for (int i = 0; i < VEC_SIZE; ++i) { const float v = static_cast(regs[i]); - const float scaled = SCALE_UE8M0 ? (v * inv_y) : (v / y_s); + const float scaled = v * inv_y; const float q = fminf(fmaxf(scaled, min_8bit), max_8bit); DST_DTYPE qb = DST_DTYPE(q); const uint8_t byte = *reinterpret_cast(&qb); From de8291e24f1312da0239605e33a89fe9cc901e47 Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:51:05 +0800 Subject: [PATCH 497/952] [Core] Add lifecycle hooks for worker extensions (#59463) Signed-off-by: Dakai An Co-authored-by: Claude Opus 5.5 Co-authored-by: aoshen02 --- tests/v1/worker/test_worker_base.py | 85 +++++++++++++++++++++++++++++ vllm/config/parallel.py | 5 +- vllm/v1/worker/worker_base.py | 13 +++++ 3 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 tests/v1/worker/test_worker_base.py diff --git a/tests/v1/worker/test_worker_base.py b/tests/v1/worker/test_worker_base.py new file mode 100644 index 000000000000..079b643dc84f --- /dev/null +++ b/tests/v1/worker/test_worker_base.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm.v1.worker import worker_base +from vllm.v1.worker.worker_base import WorkerBase, WorkerWrapperBase + +pytestmark = pytest.mark.skip_global_cleanup + + +class ExtensionA: + events: list[str] + + def init_worker_extension(self): + self.events.append("init-a") + + def destroy_worker_extension(self): + self.events.append("destroy-a") + + +class ExtensionB: ... + + +class FailingExtension: + def destroy_worker_extension(self): + raise RuntimeError("destroy failed") + + +@pytest.fixture +def make_wrapper(monkeypatch): + # Fresh worker class per test, since extension injection mutates its bases. + class Worker(WorkerBase): + def __init__(self, **_kwargs): + self.events: list[str] = [] + + def init_device(self): + self.events.append("device") + + def shutdown(self): + self.events.append("shutdown") + + classes = { + cls.__name__: cls for cls in (Worker, ExtensionA, ExtensionB, FailingExtension) + } + monkeypatch.setattr(worker_base, "resolve_obj_by_qualname", classes.__getitem__) + monkeypatch.setattr( + worker_base, + "worker_receiver_cache_from_config", + lambda *_args, **_kwargs: None, + ) + + def make(extension: str) -> WorkerWrapperBase: + config = SimpleNamespace( + parallel_config=SimpleNamespace( + worker_cls="Worker", worker_extension_cls=extension + ), + enable_trace_function_call_for_thread=lambda: None, + ) + wrapper = WorkerWrapperBase() + wrapper.init_worker([{"vllm_config": config}]) + wrapper.init_device() + return wrapper + + return make + + +def test_worker_extension_hooks_follow_configured_class(make_wrapper): + for extension, expected in ( + ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]), + ("ExtensionB", ["device", "shutdown"]), + ("ExtensionA", ["device", "init-a", "destroy-a", "shutdown"]), + ): + wrapper = make_wrapper(extension) + wrapper.shutdown() + assert wrapper.worker.events == expected + + +def test_failing_destroy_hook_is_logged_not_raised(make_wrapper, caplog_vllm): + wrapper = make_wrapper("FailingExtension") + wrapper.shutdown() + assert wrapper.worker.events == ["device", "shutdown"] + assert "Worker extension destroy hook failed" in caplog_vllm.text diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index e103c9ffd73d..73d83442baa7 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -309,7 +309,10 @@ class ParallelConfig: """The full name of the worker extension class to use. The worker extension class is dynamically inherited by the worker class. This is used to inject new attributes and methods to the worker class for use in collective_rpc - calls.""" + calls. It may also define `init_worker_extension(self)`, called after + device and distributed initialization but before model loading, and + `destroy_worker_extension(self)`, called (best-effort) before the worker + shuts down.""" master_addr: str = "127.0.0.1" """distributed master address for multi-node distributed inference when distributed_executor_backend is mp.""" diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 95a0c3475d4d..e1a4eaf29102 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -244,9 +244,17 @@ def __init__( # Initialized after init_worker is called self.worker: WorkerBase self.vllm_config: VllmConfig + self._worker_extension_cls: type | None = None def shutdown(self) -> None: if self.worker is not None: + hook = getattr(self._worker_extension_cls, "destroy_worker_extension", None) + if callable(hook): + try: + hook(self.worker) + except Exception: + # Don't raise: callers still need to run their own teardown. + logger.exception("Worker extension destroy hook failed.") self.worker.shutdown() def update_environment_variables( @@ -293,6 +301,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: "and pass the qualified name of the class as a string." ) + worker_extension_cls: type | None = None if parallel_config.worker_extension_cls: worker_extension_cls = resolve_obj_by_qualname( parallel_config.worker_extension_cls @@ -336,6 +345,7 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: with set_current_vllm_config(self.vllm_config): # To make vLLM config available during worker initialization self.worker = worker_class(**kwargs) + self._worker_extension_cls = worker_extension_cls def initialize_from_config(self, kv_cache_configs: list[Any]) -> None: kv_cache_config = kv_cache_configs[self.global_rank] @@ -348,6 +358,9 @@ def init_device(self): with set_current_vllm_config(self.vllm_config): # To make vLLM config available during device initialization self.worker.init_device() # type: ignore + hook = getattr(self._worker_extension_cls, "init_worker_extension", None) + if callable(hook): + hook(self.worker) def __getattr__(self, attr: str): return getattr(self.worker, attr) From 61d9c777bfdef66efbf017c7c363d2bdca5db8ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 22:42:53 +0800 Subject: [PATCH 498/952] [CI] Run the Clef pooler test on GPU (#61409) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 --- tests/models/multimodal/pooling/test_clef_omni.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index 381ac23d53cc..c17b02222efa 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -15,9 +15,11 @@ JointSchemaHead, JointSchemaPooler, ) +from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates +DEVICE_TYPE = current_platform.device_type HIDDEN, VOCAB = 64, 500 TYPES = [0, 1, 1, 2] COUNTS = [2, 3, 1, 5] @@ -55,7 +57,7 @@ def reference_logits(head, hidden, token_ids, embedding, question_spans, spans): base + head.option_summary_norm(torch.stack(summaries)) + head.global_projection(global_vector) - + head.type_embedding(torch.tensor(TYPES)) + + head.type_embedding(torch.tensor(TYPES, device=hidden.device)) ).unsqueeze(0) for layer in head.layers: fields = layer(fields, memory) @@ -90,7 +92,8 @@ def pooler(): with torch.no_grad(): for p in head.parameters(): p.add_(torch.randn_like(p) * 0.2) - embedding = torch.randn(VOCAB, HIDDEN) + head.to(DEVICE_TYPE) + embedding = torch.randn(VOCAB, HIDDEN, device=DEVICE_TYPE) config = SimpleNamespace( scheduler_config=SimpleNamespace( enable_chunked_prefill=False, async_scheduling=False @@ -132,7 +135,7 @@ def test_pooler_matches_reference(pooler): pooler, embedding = pooler lengths = [120, 90, 100] token_ids = [torch.randint(0, VOCAB, (n,)) for n in lengths] - hidden = torch.randn(sum(lengths), HIDDEN) + hidden = torch.randn(sum(lengths), HIDDEN, device=DEVICE_TYPE) first, first_questions, first_spans = schema(lengths[0]) last, last_questions, last_spans = schema(lengths[2]) params = [ @@ -154,7 +157,7 @@ def test_pooler_matches_reference(pooler): metadata.build_pooling_cursor( torch.tensor(lengths).numpy(), seq_lens_cpu=torch.tensor(lengths), - device=torch.device("cpu"), + device=hidden.device, ) probs, empty, logits = pooler(hidden, metadata) From f5fefbffe51e1cc91e0d6507d272f337473989b3 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 13 Oct 2026 11:01:44 -0400 Subject: [PATCH 499/952] [Build] DeepGEMM migration to torch Stable ABI (#59482) Signed-off-by: Chris Leonard Signed-off-by: Shengqi Chen Co-authored-by: Shengqi Chen --- .buildkite/check-torch-abi.py | 1 - cmake/external_projects/deepgemm.cmake | 4 ++-- tools/install_deepgemm.sh | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py index 37e2b8dadbf8..5bfec8abd077 100644 --- a/.buildkite/check-torch-abi.py +++ b/.buildkite/check-torch-abi.py @@ -17,7 +17,6 @@ # _C_stable_libtorch, but these two modules still bind libtorch directly. "_C.abi3.so", "_rocm_C.abi3.so", - "third_party/deep_gemm/_C*.so", ) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index b796623b7ac1..7f1c0ba5157e 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,8 +32,8 @@ else() # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 - # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") + # Mega MoE. TORCH_STABLE_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "a981d8e3a58e877b7f53e63707de5834312fb710") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index b617e6d5679f..c293bb3962ba 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -10,7 +10,7 @@ DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release # plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU # Mega MoE activations, and NVFP4 Mega MoE. -DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" +DEEPGEMM_GIT_REF="a981d8e3a58e877b7f53e63707de5834312fb710" WHEEL_DIR="" # Parse command line arguments From f395a6b1ad49458ea5b0366805c3814c51d6b02b Mon Sep 17 00:00:00 2001 From: Rishabh Agarwal Date: Tue, 13 Oct 2026 11:03:48 -0400 Subject: [PATCH 500/952] [ROCM][Bugfix] Deferring to TRITON_ATTN backend for Prefix lm models instead of ROCM_ATTN and fixed minor bug in Gemma3 Multimodal code (#46754) Signed-off-by: GodRishUniverse Signed-off-by: Rishabh Agarwal Signed-off-by: GodRishUniverse --- .../test_rocm_attention_backends_selection.py | 25 +++++++++++++++++++ vllm/model_executor/models/gemma3_mm.py | 4 +++ 2 files changed, 29 insertions(+) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index e0247a94a1a3..45580b395f72 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -61,6 +61,31 @@ def test_aiter_unified_attention_uses_dedicated_metadata_builder(): ) +def test_rocm_backends_do_not_support_mm_prefix(): + """Regression test: ROCm backends must not claim mm_prefix support. + + If they do, prefix-LM multimodal models that attend vision tokens + bidirectionally (e.g. Gemma3 4B) will use bidirectional attention + via the ROCm backend instead of falling back to Triton, producing + incorrect results. Remove this test only if support is properly added. + """ + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend + + assert not RocmAttentionBackend.supports_mm_prefix(), ( + "RocmAttentionBackend does not support mm_prefix — " + "defer to Triton unified attention instead for now" + "See https://github.com/vllm-project/vllm/issues/48172" + ) + assert not RocmAiterUnifiedAttentionBackend.supports_mm_prefix(), ( + "RocmAiterUnifiedAttentionBackend does not support mm_prefix — " + "defer to Triton unified attention instead for now" + "See https://github.com/ROCm/aiter/issues/4168" + ) + + def test_segmented_attention_uses_dedicated_backend_components(): from vllm.v1.attention.backends.rocm_segmented_attn import ( RocmSegmentedAttentionBackend, diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index c6166bc2f423..b5b0c0d48bc4 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -557,6 +557,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.multi_modal_projector = Gemma3MultiModalProjector(config) + # https://huggingface.co/blog/gemma3, (image tokens attend to each other) + # explicit mention here as text side in gemma3 standalone does not need + # prefix_lm but multimodal does + config.text_config.is_mm_prefix_lm = True with self._mark_language_model(vllm_config): self.language_model = init_vllm_registered_model( vllm_config=vllm_config, From 669f85e686ae533a761362c919d213a87e0e807c Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 13 Oct 2026 23:08:07 +0800 Subject: [PATCH 501/952] [Misc] Break chat_utils into more maintainable modules. (#60972) Signed-off-by: wang.yuqi --- .../cohere/test_serving_conversion.py | 6 +- .../unit_tests/test_chat_billing_header.py | 4 +- .../multimodal/processing/test_common.py | 2 +- tests/renderers/test_chat_utils.py | 12 +- .../test_chat_utils_prompt_embeds.py | 9 +- tests/renderers/test_warmup.py | 2 +- vllm/entrypoints/cohere/serving.py | 4 +- .../entrypoints/pooling/embed/io_processor.py | 2 +- .../pooling/scoring/io_processor.py | 2 +- vllm/entrypoints/pooling/scoring/utils.py | 4 +- vllm/exceptions.py | 6 + vllm/renderers/base.py | 2 +- vllm/renderers/chat_utils.py | 2349 ----------------- vllm/renderers/chat_utils/__init__.py | 42 + vllm/renderers/chat_utils/chat_message.py | 826 ++++++ vllm/renderers/chat_utils/chat_template.py | 112 + vllm/renderers/chat_utils/content_parser.py | 560 ++++ vllm/renderers/chat_utils/item_tracker.py | 562 ++++ vllm/renderers/chat_utils/tool_call.py | 32 + vllm/renderers/chat_utils/types.py | 367 +++ vllm/renderers/hf.py | 3 +- 21 files changed, 2533 insertions(+), 2375 deletions(-) delete mode 100644 vllm/renderers/chat_utils.py create mode 100644 vllm/renderers/chat_utils/__init__.py create mode 100644 vllm/renderers/chat_utils/chat_message.py create mode 100644 vllm/renderers/chat_utils/chat_template.py create mode 100644 vllm/renderers/chat_utils/content_parser.py create mode 100644 vllm/renderers/chat_utils/item_tracker.py create mode 100644 vllm/renderers/chat_utils/tool_call.py create mode 100644 vllm/renderers/chat_utils/types.py diff --git a/tests/entrypoints/cohere/test_serving_conversion.py b/tests/entrypoints/cohere/test_serving_conversion.py index e6696c26b42c..7136b4acdd95 100644 --- a/tests/entrypoints/cohere/test_serving_conversion.py +++ b/tests/entrypoints/cohere/test_serving_conversion.py @@ -511,7 +511,7 @@ def test_text_only_list_preserves_per_block_text_parts(self): # A list of text blocks maps 1-to-1 onto ``text`` content # parts, mirroring the shape ``_convert_user_message`` emits # for multi-part user content. vLLM's - # ``_parse_chat_message_content`` will collapse this back to + # ``parse_chat_message_content`` will collapse this back to # ``"\n".join(...)`` for the ``ConversationMessage`` handed to # renderers (see the tool-role branch in # ``vllm/entrypoints/chat_utils.py``), so downstream chat @@ -533,7 +533,7 @@ def test_document_blocks_serialized_as_text_parts(self): """Document blocks must not survive on the OpenAI-shaped tool message. - vLLM's shared ``_parse_chat_message_content_parts`` runs a + vLLM's shared ``parse_chat_message_content_parts`` runs a Pydantic ``ValidatorIterator`` over the OpenAI content-part union, which only knows ``text`` / ``image_url`` / etc. and rejects ``{"type": "document", ...}`` blocks outright. We @@ -588,7 +588,7 @@ def test_document_only_content_serialized_as_text_part(self): def test_converted_tool_message_content_survives_openai_validation(self): """Regression for the reported crash. - ``vllm/entrypoints/chat_utils.py:_parse_chat_message_content_parts`` + ``vllm/renders/chat_message_content.py:parse_chat_message_content_parts`` runs a Pydantic ``ValidatorIterator`` over ``list[ChatCompletionContentPartParam]``, which is a discriminated union that only accepts ``text`` / ``image_url`` diff --git a/tests/entrypoints/unit_tests/test_chat_billing_header.py b/tests/entrypoints/unit_tests/test_chat_billing_header.py index 3ad694759760..1848aa568b91 100644 --- a/tests/entrypoints/unit_tests/test_chat_billing_header.py +++ b/tests/entrypoints/unit_tests/test_chat_billing_header.py @@ -6,7 +6,7 @@ import pytest -from vllm.renderers.chat_utils import _parse_chat_message_content +from vllm.renderers.chat_utils import parse_chat_message_content HEADER = ( "x-anthropic-billing-header: cc_version=2.1.220.300; " @@ -17,7 +17,7 @@ def _parse(message, content_format="openai"): tracker = MagicMock() tracker.create_parser.return_value.mm_placeholder_storage.return_value = {} - return _parse_chat_message_content( + return parse_chat_message_content( message, tracker, content_format, interleave_strings=False ) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 820d655e905e..667907c3e55c 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -29,7 +29,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.platforms import current_platform -from vllm.renderers.chat_utils import _get_embeds_data +from vllm.renderers.chat_utils.item_tracker import _get_embeds_data from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config from vllm.utils.mistral import is_mistral_tokenizer diff --git a/tests/renderers/test_chat_utils.py b/tests/renderers/test_chat_utils.py index 47fe49a63e6e..877c4de74b5b 100644 --- a/tests/renderers/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -16,24 +16,24 @@ from vllm.config import ModelConfig from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict +from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY from vllm.multimodal.utils import ( encode_audio_url, encode_image_url, encode_video_url, ) from vllm.renderers.chat_utils import ( - MEDIA_CONNECTOR_REGISTRY, AsyncMultiModalItemTracker, ChatCompletionMessageParam, ConversationMessage, - _load_embeds_dict, - _parse_chat_message_content, - _parse_metadata_array, - _postprocess_messages, + parse_chat_message_content, parse_chat_messages, parse_chat_messages_async, validate_chat_template, ) +from vllm.renderers.chat_utils.chat_message import _postprocess_messages +from vllm.renderers.chat_utils.content_parser import _load_embeds_dict +from vllm.renderers.chat_utils.item_tracker import _parse_metadata_array from vllm.utils.serial_utils import tensor2base64 KIMI_K2_5_MODEL_ID = "moonshotai/Kimi-K2.5" @@ -3205,7 +3205,7 @@ def _parse_one(message) -> ConversationMessage: parser = tracker.create_parser.return_value parser.model_config.enable_prompt_embeds = False parser.mm_placeholder_storage.return_value = {} - result = _parse_chat_message_content( + result = parse_chat_message_content( message, tracker, content_format="string", interleave_strings=False ) assert len(result) == 1 diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index f26680793ee2..d6f1673ca79a 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -18,15 +18,16 @@ from vllm.exceptions import VLLMValidationError from vllm.renderers import TokenizeParams -from vllm.renderers.chat_utils import ( - _ENABLE_PROMPT_EMBEDS_ERROR, +from vllm.renderers.chat_utils import parse_chat_messages, parse_chat_messages_async +from vllm.renderers.chat_utils.chat_message import ( _PROMPT_EMBEDS_MISSING_DATA_ERROR, _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, MM_PARSER_MAP, +) +from vllm.renderers.chat_utils.content_parser import ( + _ENABLE_PROMPT_EMBEDS_ERROR, MODALITY_PLACEHOLDERS_MAP, PROMPT_EMBEDS_PLACEHOLDER_TOKEN, - parse_chat_messages, - parse_chat_messages_async, ) from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, diff --git a/tests/renderers/test_warmup.py b/tests/renderers/test_warmup.py index 09fe047d56ef..0c847e5600f2 100644 --- a/tests/renderers/test_warmup.py +++ b/tests/renderers/test_warmup.py @@ -28,9 +28,9 @@ import pytest +from vllm.exceptions import ChatTemplateResolutionError from vllm.multimodal.processing import BaseMultiModalProcessor from vllm.renderers.base import BaseRenderer -from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.params import ChatParams diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 93f6d14a3b9c..4ce734a8121c 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -502,7 +502,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # Map every content block to a plain ``{"type": "text", "text": # ...}`` part, mirroring the shape :meth:`_convert_user_message` # emits for multi-part user content. vLLM's - # ``_parse_chat_message_content`` collapses a text-only tool + # ``parse_chat_message_content`` collapses a text-only tool # part list back into ``"\n".join(...)`` on the way to the # renderer, so simple chat templates that expect string tool # content still work; keeping the list here just preserves the @@ -512,7 +512,7 @@ def _convert_tool_message(cls, msg: ToolChatMessageV2) -> dict[str, Any]: # Tool message content uses ``ToolMessageV2Content``, a union of # ``TextToolContent`` and ``DocumentToolContent``. The latter # produces ``{"type": "document", "document": {...}}`` blocks, - # which vLLM's shared ``_parse_chat_message_content_parts`` in + # which vLLM's shared ``parse_chat_message_content_parts`` in # ``vllm/renderers/chat_utils.py`` rejects because the OpenAI # ``ChatCompletionContentPartParam`` union only knows text / # image / audio / etc. Emitting them anyway triggers a Pydantic diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index b37a009c925a..4b2ad7a3b8ee 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -9,7 +9,6 @@ ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam, ) -from openai.types.chat.chat_completion_content_part_image_param import ImageURL from vllm import PoolingParams from vllm.exceptions import VLLMValidationError @@ -21,6 +20,7 @@ ChatCompletionContentPartParam, ChatCompletionMessageParam, CustomChatCompletionMessageParam, + ImageURL, ) from vllm.renderers.hf import resolve_chat_template from vllm.utils.collection_utils import chunk_list diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 1e4c638968f8..9ab36a13958c 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -7,9 +7,9 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt +from vllm.exceptions import ChatTemplateResolutionError from vllm.logger import init_logger from vllm.renderers import TokenizeParams -from vllm.renderers.chat_utils import ChatTemplateResolutionError from vllm.renderers.hf import safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 8d953cf6ddd4..b3617ae9c601 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -15,7 +15,7 @@ ChatCompletionContentPartTextParam, ConversationMessage, MultiModalItemTracker, - _parse_chat_message_content_parts, + parse_chat_message_content_parts, ) from .typing import ( @@ -193,7 +193,7 @@ def _parse_score_content( mm_parser = mm_tracker.create_parser() - parse_res = _parse_chat_message_content_parts( + parse_res = parse_chat_message_content_parts( role=role, parts=parts, mm_tracker=mm_tracker, diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 11ca13fe0bc3..85863c2fbf1d 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -62,6 +62,12 @@ class VLLMNotFoundError(VLLMClientError): pass +class ChatTemplateResolutionError(VLLMClientError): + """Raised when chat template resolution fails.""" + + pass + + class LoRAAdapterNotFoundError(VLLMNotFoundError): """Exception raised when a LoRA adapter is not found. diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index e04bfa64eb76..7b0f4debaad8 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -362,7 +362,7 @@ def warmup(self, chat_params: ChatParams) -> None: no-op if the background warmup — or a prior reset_mm_cache that joined it — already ran it. """ - from vllm.renderers.chat_utils import ChatTemplateResolutionError + from vllm.exceptions import ChatTemplateResolutionError # prevent MM processor hangs with set_default_torch_num_threads(1): diff --git a/vllm/renderers/chat_utils.py b/vllm/renderers/chat_utils.py deleted file mode 100644 index 540564baab8d..000000000000 --- a/vllm/renderers/chat_utils.py +++ /dev/null @@ -1,2349 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import asyncio -import json -import types -from abc import ABC, abstractmethod -from collections import Counter, defaultdict -from collections.abc import Awaitable, Callable, Iterable -from dataclasses import dataclass -from functools import cached_property, lru_cache, partial -from itertools import accumulate -from pathlib import Path -from typing import ( - TYPE_CHECKING, - Any, - Final, - Generic, - Literal, - Required, - TypeAlias, - TypeVar, - Union, - cast, - get_args, - get_origin, -) - -from openai.types.chat import ( - ChatCompletionAssistantMessageParam, - ChatCompletionContentPartImageParam, - ChatCompletionContentPartInputAudioParam, - ChatCompletionContentPartRefusalParam, - ChatCompletionContentPartTextParam, - ChatCompletionFunctionToolParam, - ChatCompletionMessageToolCallParam, - ChatCompletionToolMessageParam, -) -from openai.types.chat import ( - ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam, -) -from openai.types.chat import ( - ChatCompletionMessageParam as OpenAIChatCompletionMessageParam, -) -from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio -from openai.types.responses import ResponseInputImageParam -from openai_harmony import Message as OpenAIHarmonyMessage -from PIL import Image -from pydantic import BaseModel, ConfigDict, TypeAdapter - -# pydantic needs the TypedDict from typing_extensions -from typing_extensions import TypedDict, override - -from vllm import envs -from vllm.config import ModelConfig -from vllm.connections import MediaRequestDownloadBudget -from vllm.exceptions import VLLMValidationError -from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict -from vllm.logger import init_logger -from vllm.model_executor.models import SupportsMultiModal -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - MultiModalBatchedField, - MultiModalFlatField, - MultiModalSharedField, - VisionChunk, - VisionChunkImage, - VisionChunkVideo, -) -from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector -from vllm.multimodal.media.base import MediaDecodedBudget -from vllm.multimodal.processing import BaseMultiModalProcessor -from vllm.renderers.embed_utils import ( - safe_load_prompt_embeds, - safe_load_prompt_embeds_async, -) -from vllm.transformers_utils.processor import get_video_processor_cls_name -from vllm.utils import random_uuid -from vllm.utils.collection_utils import is_list_of, is_list_of_numbers -from vllm.utils.import_utils import LazyLoader - -if TYPE_CHECKING: - import torch - import transformers -else: - transformers = LazyLoader("transformers", globals(), "transformers") - torch = LazyLoader("torch", globals(), "torch") - -logger = init_logger(__name__) - - -class ChatTemplateResolutionError(ValueError): - """Raised when chat template resolution fails. - - This is a subclass of ValueError for backward compatibility with - existing exception handlers. - """ - - -MODALITY_PLACEHOLDERS_MAP = { - "image": "<##IMAGE##>", - "audio": "<##AUDIO##>", - "video": "<##VIDEO##>", - "prompt_embeds": "<##PROMPT_EMBEDS##>", -} - - -PROMPT_EMBEDS_PLACEHOLDER_TOKEN: Final[str] = "" -"""The special token used as a placeholder for each embedding -position during chat template rendering. - -Registered as an additional special token when `--enable-prompt-embeds` is set. -See `_ensure_prompt_embeds_placeholder_token` in `vllm/renderers/hf.py`. -""" - - -_REQUIRE_MM_PROCESSOR_ERROR: Final[str] = ( - "Resolving modality {modality!r} requires a multimodal processor " - "but none is available." -) - -_ENABLE_PROMPT_EMBEDS_ERROR: Final[str] = ( - "You must set `--enable-prompt-embeds` to input `prompt_embeds`" -) - -_PROMPT_EMBEDS_MISSING_DATA_ERROR: Final[str] = ( - "prompt_embeds content part requires a non-empty `data` field " - "with base64-encoded tensor bytes." -) - -_RESERVED_PLACEHOLDER_IN_TEXT_ERROR: Final[str] = ( - "Text content may not contain the reserved placeholder {token!r}. " - "This placeholder is used internally to mark `prompt_embeds` splice " - "positions in the tokenized prompt." -) - - -class AudioURL(TypedDict, total=False): - url: Required[str] - """ - Either a URL of the audio or a data URL with base64 encoded audio data. - """ - - -class ChatCompletionContentPartAudioParam(TypedDict, total=False): - audio_url: Required[AudioURL] - - type: Required[Literal["audio_url"]] - """The type of the content part.""" - - -MultiModalEmbedsPayload: TypeAlias = str | dict[str, str | list[int | float]] - - -class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False): - image_embeds: MultiModalEmbedsPayload | None - """ - The image embeddings. It can be either: - - A single base64 string. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["image_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False): - audio_embeds: MultiModalEmbedsPayload | None - """ - The audio embeddings. It can be either: - - A single base64 string representing a serialized torch tensor. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["audio_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartVideoEmbedsParam(TypedDict, total=False): - video_embeds: MultiModalEmbedsPayload | None - """ - The video embeddings. It can be either: - - A single base64 string representing a serialized torch tensor. - - A dictionary of base64 tensors or numeric JSON metadata arrays. - """ - type: Required[Literal["video_embeds"]] - """The type of the content part.""" - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class ChatCompletionContentPartPromptEmbedsParam(TypedDict, total=False): - data: Required[str] - """ - Base64-encoded bytes of a serialized `torch.Tensor` of shape - `(num_tokens, hidden_size)`. The tensor's `dtype` and `hidden_size` must - match the model's input embedding layer. - """ - type: Required[Literal["prompt_embeds"]] - """The type of the content part.""" - - -class VideoURL(TypedDict, total=False): - url: Required[str] - """ - Either a URL of the video or a data URL with base64 encoded video data. - """ - - -class ChatCompletionContentPartVideoParam(TypedDict, total=False): - video_url: Required[VideoURL] - - type: Required[Literal["video_url"]] - """The type of the content part.""" - - -class PILImage(BaseModel): - """A PIL.Image.Image object.""" - - image_pil: Image.Image - model_config = ConfigDict(arbitrary_types_allowed=True) - - -class CustomChatCompletionContentPILImageParam(TypedDict, total=False): - """A simpler version of the param that only accepts a PIL image. - - Example: - { - "image_pil": ImageAsset('cherry_blossom').pil_image - } - - """ - - image_pil: PILImage | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain image_url. - This is supported by OpenAI API, although it is not documented. - - Example: - { - "image_url": "https://example.com/image.jpg" - } - - """ - - image_url: str | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain audio_url. - - Example: - { - "audio_url": "https://example.com/audio.mp3" - } - - """ - - audio_url: str | None - - -class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False): - """A simpler version of the param that only accepts a plain audio_url. - - Example: - { - "video_url": "https://example.com/video.mp4" - } - - """ - - video_url: str | None - uuid: str | None - """ - User-provided UUID of a media. User must guarantee that it is properly - generated and unique for different medias. - """ - - -class CustomThinkCompletionContentParam(TypedDict, total=False): - """A Think Completion Content Param that accepts a plain text and a boolean. - - Example: - { - "thinking": "I am thinking about the answer", - "closed": True, - "type": "thinking" - } - - """ - - thinking: Required[str] - """The thinking content.""" - - closed: bool - """Whether the thinking is closed.""" - - type: Required[Literal["thinking"]] - """The thinking type.""" - - -class CustomChatCompletionContentToolReferenceParam(TypedDict, total=False): - """A tool reference content param that only accepts a plain tool name. - - Example: - { - "name": "get_weather", - "type": "tool_reference" - } - - """ - - name: str - """The name of the tool being referenced.""" - - type: Literal["tool_reference"] - """The content type.""" - - -ChatCompletionContentPartParam: TypeAlias = ( - OpenAIChatCompletionContentPartParam - | ChatCompletionContentPartAudioParam - | ChatCompletionContentPartInputAudioParam - | ChatCompletionContentPartVideoParam - | ChatCompletionContentPartRefusalParam - | CustomChatCompletionContentPILImageParam - | CustomChatCompletionContentSimpleImageParam - | ChatCompletionContentPartImageEmbedsParam - | ChatCompletionContentPartAudioEmbedsParam - | ChatCompletionContentPartVideoEmbedsParam - | ChatCompletionContentPartPromptEmbedsParam - | CustomChatCompletionContentSimpleAudioParam - | CustomChatCompletionContentSimpleVideoParam - | CustomChatCompletionContentToolReferenceParam - | str - | CustomThinkCompletionContentParam -) - - -class CustomChatCompletionMessageParam(TypedDict, total=False): - """Enables custom roles in the Chat Completion API.""" - - role: Required[str] - """The role of the message's author.""" - - content: str | list[ChatCompletionContentPartParam] - """The contents of the message.""" - - name: str - """An optional name for the participant. - - Provides the model information to differentiate between participants of the - same role. - """ - - tool_call_id: str | None - """Tool call that this message is responding to.""" - - tool_calls: list[ChatCompletionMessageToolCallParam] | None - """The tool calls generated by the model, such as function calls.""" - - reasoning: str | None - """The reasoning content for interleaved thinking.""" - - tools: list[ChatCompletionFunctionToolParam] | None - """The tools for developer role.""" - - task: str | None - """Model-specific task marker. Currently passed through for DeepSeek V4.""" - - -ChatCompletionMessageParam: TypeAlias = ( - OpenAIChatCompletionMessageParam - | CustomChatCompletionMessageParam - | OpenAIHarmonyMessage -) - - -# TODO: Make fields ReadOnly once mypy supports it -class ConversationMessage(TypedDict, total=False): - role: Required[str] - """The role of the message's author.""" - - content: str | None | list[dict[str, str]] - """The contents of the message""" - - tool_call_id: str | None - """Tool call that this message is responding to.""" - - name: str | None - """The name of the function to call""" - - tool_calls: list[ChatCompletionMessageToolCallParam] | None - """The tool calls generated by the model, such as function calls.""" - - reasoning: str | None - """The reasoning content for interleaved thinking.""" - - reasoning_content: str | None - """Deprecated: The reasoning content for interleaved thinking.""" - - tools: list[ChatCompletionFunctionToolParam] | None - """The tools for developer role.""" - - task: str | None - """Model-specific task marker. Currently passed through for DeepSeek V4.""" - - -# Passed in by user -ChatTemplateContentFormatOption = Literal["auto", "string", "openai"] - -# After resolving "auto" -ChatTemplateContentFormat = Literal["string", "openai"] - - -ModalityStr = Literal[ - "image", - "audio", - "video", - "image_embeds", - "audio_embeds", - "video_embeds", - "vision_chunk", - "prompt_embeds", -] -_T = TypeVar("_T") -_AsyncMultiModalItem: TypeAlias = Callable[[], Awaitable[tuple[object, str | None]]] - - -# Backward compatibility for single item input -class _BatchedSingleItemField(MultiModalSharedField): - pass - - -def _detect_field( - tensors: list[torch.Tensor], - mm_processor: BaseMultiModalProcessor, -): - first_item = tensors[0] - hidden_size = mm_processor.info.ctx.model_config.get_inputs_embeds_size() - - if ( - len(tensors) == 1 - and first_item.ndim == 3 - and first_item.shape[0] == 1 - and first_item.shape[-1] == hidden_size - ): - logger.warning( - "Batched multi-modal embedding inputs are deprecated for Chat API. " - "Please pass a separate content part for each multi-modal item." - ) - return _BatchedSingleItemField(batch_size=1) - - first_shape = first_item.shape - if all(t.shape == first_shape for t in tensors): - return MultiModalBatchedField() - - size_per_item = [len(tensor) for tensor in tensors] - slice_idxs = [0, *accumulate(size_per_item)] - slices = [ - (slice(slice_idxs[i], slice_idxs[i + 1]),) for i in range(len(size_per_item)) - ] - return MultiModalFlatField(slices=slices) - - -def _merge_embeds( - data_items: list[dict[str, "torch.Tensor"]], - mm_processor: BaseMultiModalProcessor, -): - if not data_items: - return {} - - first_keys = set(data_items[0].keys()) - if any(set(item.keys()) != first_keys for item in data_items[1:]): - raise VLLMValidationError( - "All dictionaries in the list of embeddings must have the same keys." - ) - - fields = { - key: _detect_field([item[key] for item in data_items], mm_processor) - for key in first_keys - } - data_merged = { - key: field._reduce_data([item[key] for item in data_items], pin_memory=False) - for key, field in fields.items() - } - - try: - # TODO: Support per-request mm_processor_kwargs - parsed_configs = mm_processor._get_mm_fields_config( - transformers.BatchFeature(data_merged), - {}, - ) - parsed_fields = {key: parsed_configs[key].field for key in first_keys} - keys_to_update = [ - key - for key in first_keys - if ( - fields[key] != parsed_fields[key] - and not isinstance(fields[key], _BatchedSingleItemField) - ) - ] - - for key in keys_to_update: - data_merged[key] = parsed_fields[key]._reduce_data( - [item[key] for item in data_items], pin_memory=False - ) - except Exception: - logger.exception( - "Error when parsing merged embeddings. " - "Falling back to auto-detected fields." - ) - - return data_merged - - -async def _load_embeds_dict( - data: dict[str, str | list[int | float]], - fetch: Callable[[str], Awaitable["torch.Tensor"]], -) -> dict[str, Any]: - encoded = {key: value for key, value in data.items() if isinstance(value, str)} - tensors = await asyncio.gather(*(fetch(value) for value in encoded.values())) - return {**data, **dict(zip(encoded, tensors))} - - -def _parse_metadata_array(key: str, value: list, metadata_fields: set[str]): - if key not in metadata_fields: - raise VLLMValidationError(f"JSON arrays are only supported for metadata: {key}") - if not is_list_of_numbers(value): - raise VLLMValidationError(f"Metadata {key} must be a finite numeric array.") - dtype = torch.float64 if any(isinstance(v, float) for v in value) else torch.int64 - try: - return torch.tensor(value, dtype=dtype) - except (ValueError, TypeError, OverflowError, RuntimeError) as error: - raise VLLMValidationError(f"Invalid metadata array: {key}") from error - - -def _get_embeds_data( - modality: str, - data_items: list[Any], - mm_processor: BaseMultiModalProcessor, -): - if len(data_items) == 0: - return data_items - - if all(item is None for item in data_items): - return data_items - - if is_list_of(data_items, torch.Tensor): - embeds_key = f"{modality}_embeds" - dict_items = [{embeds_key: item} for item in data_items] - return _merge_embeds(dict_items, mm_processor)[embeds_key] - - if is_list_of(data_items, dict): - metadata_fields = mm_processor.info.data_parser.placeholder_metadata_fields( - modality - ) - data_items = [ - { - key: _parse_metadata_array(key, value, metadata_fields) - if isinstance(value, list) - else value - for key, value in item.items() - } - for item in data_items - ] - return _merge_embeds(data_items, mm_processor) - - raise NotImplementedError(type(data_items)) - - -class BaseMultiModalItemTracker(ABC, Generic[_T]): - """Tracks multi-modal items in a given request and ensures that the number - of multi-modal items in a given request does not exceed the configured - maximum per prompt. - """ - - def __init__( - self, - model_config: ModelConfig, - media_io_kwargs: dict[str, dict[str, Any]] | None = None, - ): - super().__init__() - - self._model_config = model_config - self._media_io_kwargs = media_io_kwargs - - self._items_by_modality = defaultdict[str, list[_T]](list) - # Track original modality for each vision_chunk item (image or video) - self._modality_order = defaultdict[str, list[str]](list) - self.decoded_budget = MediaDecodedBudget() - self.download_budget = MediaRequestDownloadBudget() - - @cached_property - def use_unified_vision_chunk_modality(self) -> bool: - """Check if model uses unified vision_chunk modality for images/videos.""" - return getattr(self._model_config.hf_config, "use_unified_vision_chunk", False) - - @property - def model_config(self) -> ModelConfig: - return self._model_config - - @cached_property - def model_cls(self) -> type[SupportsMultiModal]: - from vllm.model_executor.model_loader import get_model_cls - - model_cls = get_model_cls(self.model_config) - return cast(type[SupportsMultiModal], model_cls) - - @property - def media_io_kwargs(self) -> dict[str, dict[str, Any]] | None: - return self._media_io_kwargs or ( - self._model_config.multimodal_config.media_io_kwargs - if self._model_config.multimodal_config - else None - ) - - @property - def allowed_local_media_path(self): - return self._model_config.allowed_local_media_path - - @property - def allowed_media_domains(self): - return self._model_config.allowed_media_domains - - @property - def mm_registry(self): - return MULTIMODAL_REGISTRY - - @cached_property - def mm_processor(self): - return self.mm_registry.create_processor(self.model_config) - - @property - def video_processor_name(self) -> str | None: - return get_video_processor_cls_name(self.model_config) - - def add(self, modality: ModalityStr, item: _T) -> str | None: - """Add a multi-modal item to the current prompt and returns the - placeholder string to use, if any. - - An optional uuid can be added which serves as a unique identifier of the - media. - - Note: - `prompt_embeds` bypass MM-processor validation because they are - pre-computed embeddings that do not go through any HF processor, encoder, - or model-specific placeholder logic. The corresponding placeholder string is - managed by the parser via `_add_placeholder`, so we return None here. - - """ - add_info = self._validate_add(modality) - if add_info is None: - self._items_by_modality["prompt_embeds"].append(item) - return None - - input_modality, original_modality, use_vision_chunk, num_items = add_info - - # Track original modality for vision_chunk items - if use_vision_chunk: - self._items_by_modality[input_modality].append(item) # type: ignore - self._modality_order["vision_chunk"].append(original_modality) - else: - self._items_by_modality[original_modality].append(item) - - return self.model_cls.get_placeholder_str(modality, num_items) - - def _validate_add(self, modality: ModalityStr) -> tuple[str, str, bool, int] | None: - """Validate that one more item of the modality can be tracked.""" - if modality == "prompt_embeds": - return None - - input_modality = modality.replace("_embeds", "") - original_modality = modality - use_vision_chunk = ( - self.use_unified_vision_chunk_modality - and original_modality in ["video", "image"] - ) - - # If use_unified_vision_chunk_modality is enabled, - # map image/video to vision_chunk - if use_vision_chunk: - # To avoid validation fail - # because models with use_unified_vision_chunk_modality=True - # will only accept vision_chunk modality. - input_modality = "vision_chunk" - num_items = len(self._items_by_modality[input_modality]) + 1 - else: - num_items = len(self._items_by_modality[original_modality]) + 1 - - mm_config = self.model_config.multimodal_config - if ( - mm_config is not None - and mm_config.enable_mm_embeds - and mm_config.get_limit_per_prompt(input_modality) == 0 - and original_modality.endswith("_embeds") - ): - # Skip validation: embeddings bypass limit when enable_mm_embeds=True - pass - else: - self.mm_processor.info.validate_num_items(input_modality, num_items) - - return input_modality, original_modality, use_vision_chunk, num_items - - @abstractmethod - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - raise NotImplementedError - - -def _resolve_vision_chunk_items( - vision_chunk_items: list[tuple[object, str | None]], - mm_processor: BaseMultiModalProcessor, - vision_chunks_modality_order: list[str], -): - # Process vision_chunk items - extract from (data, modality) tuples - # and convert to VisionChunk types with proper UUID handling - vision_chunks_uuids = [uuid for data, uuid in vision_chunk_items] - - assert len(vision_chunk_items) == len(vision_chunks_modality_order), ( - f"vision_chunk items ({len(vision_chunk_items)}) and " - f"modality_order ({len(vision_chunks_modality_order)}) must have same length" - ) - - processed_chunks: list[VisionChunk] = [] - video_idx = 0 - for inner_modality, (data, uuid) in zip( - vision_chunks_modality_order, vision_chunk_items - ): - if inner_modality == "image": - # Cast data to proper type for image - # Use .media (PIL.Image) directly to avoid redundant - # bytes→PIL conversion in media_processor - if hasattr(data, "media"): - image_data = data.media # type: ignore[union-attr] - processed_chunks.append( - VisionChunkImage(type="image", image=image_data, uuid=uuid) - ) - else: - processed_chunks.append(data) # type: ignore[arg-type] - elif inner_modality == "video": - # For video, we may need to split into chunks - # if processor supports it - # For now, just wrap as a video chunk placeholder - if hasattr(mm_processor, "split_video_chunks") and data is not None: - try: - video_uuid = uuid or random_uuid() - # video await result is (video_data, video_meta) tuple - if isinstance(data, tuple) and len(data) >= 1: - video_data = data[0] - else: - video_data = data - video_chunks = mm_processor.split_video_chunks(video_data) - for i, vc in enumerate(video_chunks): - processed_chunks.append( - VisionChunkVideo( - type="video_chunk", - video_chunk=vc["video_chunk"], - uuid=f"{video_uuid}-{i}", - video_idx=video_idx, - prompt=vc["prompt"], - ) - ) - video_idx += 1 - except Exception as e: - logger.warning("Failed to split video chunks: %s", e) - processed_chunks.append(data) # type: ignore[arg-type] - else: - processed_chunks.append(data) # type: ignore[arg-type] - return processed_chunks, vision_chunks_uuids - - -def _resolve_items( - items_by_modality: dict[str, list[tuple[object, str | None]]], - mm_processor: BaseMultiModalProcessor | None, - modality_order: dict[str, list[str]], -) -> tuple[MultiModalDataDict, MultiModalUUIDDict]: - """Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`. - - Note: - `mm_processor` is `None` for text-only models (no registered HF - processor) whose only modality is `prompt_embeds`. Every other - modality requires a processor, enforced by the guard below. - - """ - if "image" in items_by_modality and "image_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw image and embedding inputs is not allowed", - parameter="image_embeds", - ) - if "audio" in items_by_modality and "audio_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw audio and embedding inputs is not allowed", - parameter="audio_embeds", - ) - if "video" in items_by_modality and "video_embeds" in items_by_modality: - raise VLLMValidationError( - "Mixing raw video and embedding inputs is not allowed", - parameter="video_embeds", - ) - # `prompt_embeds` bypasses HF MM processors. Every other modality requires one. - processor_modalities = items_by_modality.keys() - {"prompt_embeds"} - if processor_modalities and mm_processor is None: - raise RuntimeError( - _REQUIRE_MM_PROCESSOR_ERROR.format(modality=processor_modalities) - ) - - mm_data = {} - mm_uuids = {} - if "image_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["image"] = _get_embeds_data( - "image", - [data for data, uuid in items_by_modality["image_embeds"]], - mm_processor, - ) - mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image_embeds"]] - if "image" in items_by_modality: - mm_data["image"] = [data for data, uuid in items_by_modality["image"]] - mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image"]] - if "audio_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["audio"] = _get_embeds_data( - "audio", - [data for data, uuid in items_by_modality["audio_embeds"]], - mm_processor, - ) - mm_uuids["audio"] = [uuid for data, uuid in items_by_modality["audio_embeds"]] - if "audio" in items_by_modality: - mm_data["audio"] = [data for data, uuid in items_by_modality["audio"]] - mm_uuids["audio"] = [uuid for data, uuid in items_by_modality["audio"]] - if "video" in items_by_modality: - mm_data["video"] = [data for data, uuid in items_by_modality["video"]] - mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video"]] - if "video_embeds" in items_by_modality: - assert mm_processor is not None - mm_data["video"] = _get_embeds_data( - "video", - [data for data, uuid in items_by_modality["video_embeds"]], - mm_processor, - ) - mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video_embeds"]] - if "vision_chunk" in items_by_modality: - assert mm_processor is not None - # Process vision_chunk items - extract from (data, modality) tuples - # and convert to VisionChunk types with proper UUID handling - processed_chunks, vision_chunk_uuids = _resolve_vision_chunk_items( - items_by_modality["vision_chunk"], - mm_processor, - modality_order.get("vision_chunk", []), - ) - mm_data["vision_chunk"] = processed_chunks - mm_uuids["vision_chunk"] = vision_chunk_uuids - if "prompt_embeds" in items_by_modality: - mm_data["prompt_embeds"] = [ - data for data, _uuid in items_by_modality["prompt_embeds"] - ] - - return mm_data, mm_uuids - - -class MultiModalItemTracker(BaseMultiModalItemTracker[tuple[object, str | None]]): - def resolve_items( - self, - ) -> tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]: - if not self._items_by_modality: - return None, None - - # Text-only models (`is_multimodal_model=False`) with inputs of - # modality `prompt_embeds` have no MM processor since `prompt_embeds` are - # pre-computed and require no processing, so we pass `None`. - mm_processor = ( - self.mm_processor if self._model_config.is_multimodal_model else None - ) - return _resolve_items( - dict(self._items_by_modality), - mm_processor, - self._modality_order, - ) - - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - return MultiModalContentParser(self, mm_processor_kwargs=mm_processor_kwargs) - - -class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[_AsyncMultiModalItem]): - async def resolve_items( - self, - ) -> tuple[MultiModalDataDict | None, MultiModalUUIDDict | None]: - if not self._items_by_modality: - return None, None - - # Fetch all modalities together. Each tracked item is already an - # independent awaitable, and the async connector offloads blocking - # decode work, so waiting for one modality before starting the next - # needlessly adds their latency. - # Keep the original group and item order when rebuilding the result. - item_groups = list(self._items_by_modality.items()) - items = [item for _, group in item_groups for item in group] - results = await asyncio.gather( - *(item() for item in items), return_exceptions=True - ) - for result in results: - if isinstance(result, BaseException): - # Gathering with return_exceptions=True lets every task finish - # (or itself fail) before we raise, instead of abandoning - # still-in-flight fetches (real network/thread-pool work) the - # moment the first one fails. - raise result - - resolved_items_by_modality: dict[str, list[Any]] = {} - result_idx = 0 - for modality, group in item_groups: - next_result_idx = result_idx + len(group) - resolved_items_by_modality[modality] = results[result_idx:next_result_idx] - result_idx = next_result_idx - - mm_processor = ( - self.mm_processor if self._model_config.is_multimodal_model else None - ) - return _resolve_items( - resolved_items_by_modality, - mm_processor, - self._modality_order, - ) - - def create_parser( - self, mm_processor_kwargs: dict[str, Any] | None = None - ) -> "BaseMultiModalContentParser": - return AsyncMultiModalContentParser( - self, mm_processor_kwargs=mm_processor_kwargs - ) - - -class BaseMultiModalContentParser(ABC): - def __init__(self) -> None: - super().__init__() - - # stores model placeholders list with corresponding - # general MM placeholder: - # { - # "<##IMAGE##>": ["", "", ""], - # "<##AUDIO##>": [" + +## determinism takes effect only on the main thread + +In PyTorch 2.13 (nv26.07) Inductor config overrides are thread-local (ContextVar). `torch.use_deterministic_algorithms(True)` sets `torch._inductor.config.deterministic` on the main thread; on autograd's CUDA device thread it stays False. A `torch.compile`d norm VJP first compiles on the autograd thread during backward, sees `deterministic=False`, and the generated looped-reduction kernel enables `dynamic_scale_rblock`: besides `R0_BLOCK=1024` it adds candidate 512, and the first launch picks by timing. The two reduce in different orders and differ at bit level. Only at the end of the dynamo compile is that thread set to deterministic; later compiles have one fixed config. Result: only backward step 0 is not bitwise reproducible across runs; forward logprob and the RL closed loop are unaffected. + +Apply both fixes: `torch.compile(..., options={"deterministic": True})` at the compile site (environment-independent), and `TORCHINDUCTOR_DETERMINISTIC=1` in the recipe environment (covers any first compile on the autograd thread, including the vLLM forward function compiled inside recompute). Also: verl's full_determinism actually sets `warn_only=True`, and the mamba_ssm SSD backward must enable deterministic mode explicitly. + +## How to break down the cost from BI0 to BI1 + +Draw the actual call chain first: what the default fast kernel, the BI branch and the extra preprocessing/sorting/postprocessing each are. Do not infer "split-K was changed" or "a temporary Triton was used" from function names; compare source, binding, compiled artifacts and real launches. + +You may copy the BI0 implementation into a standalone candidate and add the deterministic/BI conditions one at a time for a bisectable A/B. It need not live in the same file, but the default path must stay unchanged. If the BI0 algorithm cannot meet the required tie/reduction contract, switching algorithms is allowed; state the cost. + +Fairness between the arms comes before any kernel conclusion: same fabric (on DS4 two environment variables closed most of the remaining gap), same node, same image, BI=0 not running the BI kernel, no JIT stall, and the KV zeroing fix in both arms. Nemotron's BI1 vs BI0 E2E throughput loss is the median over 5 pairs, with single pairs reported separately. + +## Minimal record for each candidate + +```text +baseline / candidate source+binary hashes +one intended change; BI-specific or common +same-input bitwise + BI composition + prefill/decode gates +BI0 unchanged + real shape/tie/mask/tail + Graph/workspace reuse +paired latency, launch/allocation counts, repetitions +pass / fail / inconclusive; retained or rejected; remaining scope +``` + +On failure, save the minimal counterexample, then withdraw from the default path. When the user requires all micro gates to pass before E2E, do not skip them; performance numbers do not cover a correctness failure. Do not quietly remove an existing stable-sort patch to reproduce a long-length mismatch; that tests a regression, not a defect of the current best code. diff --git a/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md b/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md new file mode 100644 index 000000000000..76b14e4e9302 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md @@ -0,0 +1,304 @@ +# Checkpoint, resume and memory budget + +Use this guide to get checkpoint save, load and resume right in long RL runs, and to budget host (CPU) memory peaks and growth. For per-phase GPU memory peaks see [Locating GPU OOM and staging peaks by lifecycle](#locating-gpu-oom-and-staging-peaks-by-lifecycle). + +## Find by symptom + +| Symptom | Where to go | +| --- | --- | +| OOM while loading a checkpoint | [Load](#4-load-enough-capacity-does-not-mean-enough-for-the-transient-peak) | +| Train-inference mismatch on the first step after restore | [13](#13-stale-quantization-scale-on-the-first-step-after-restore) | +| Resume is set but training starts from step 0 | [14](#14-resume_modeauto-starts-from-step-0) | +| Host RSS grows every step | [3.3](#33-locating-slowly-grows-every-step) | +| Planning host memory for a checkpoint in the hundreds of GB | [4](#4-capacity-planning) | +| vLLM fails to start when colocated with training (KV budget too small) | [5](#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated) | +| GPU OOM during init / export / wake | [Locating GPU OOM](#locating-gpu-oom-and-staging-peaks-by-lifecycle) | + +## Terms + +- **pinned memory:** page-locked host memory for direct GPU DMA (`cudaHostAlloc` / `cudaHostRegister`), cached by PyTorch's caching host allocator. +- **cgroup:** limit `memory.max`, usage `memory.current`, split in `memory.stat` into anon, file (page cache) and shmem. +- **RSS and PSS:** RSS counts shared pages fully in every process; PSS splits them, so only summed PSS approximates the physical total. +- **DCP / Megatron dist_ckpt:** PyTorch and Megatron distributed checkpoint formats; each rank writes its shard, `.metadata` last. +- **staging / D2H:** the host buffer GPU tensors are copied into before writing to disk. +- **Ray object store (plasma):** Ray's shared memory for passing objects, mapped into each worker and charged via PSS. +- **old-logprob:** the logprob training recomputes for rollout tokens with the pre-update policy. +- **Ray OOM threshold:** `RAY_memory_usage_threshold`, the usage fraction at which Ray kills workers. Default 0.95. + +## Key rules + +- Check file, memory and derived state separately; complete files can still mismatch on the first step after restore. +- Capture the first exception (errno, rank, file, offset), not the last message. +- Completion marker last; clear the old directory before overwriting; all ranks check and fail together; a missing file raises instead of falling back to a fresh optimizer. +- Under EP4 + FSDP2, ranks' expert shards in DCP share key and offset, so restore reads another rank's experts. Put the EP rank in the key (port the reference fix), add a 4-GPU bitwise read-back unit test, and raise on EP size mismatch. +- Rotation deletes only tracked, complete, idle directories; in long runs keep one checkpoint every 20 steps outside rotation. +- Set rollout `gpu_memory_utilization` per training mode ([5](#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated)). +- Save fixes pass [the acceptance in 3.2](#32-reproduce-and-discriminate). Before a long run on a new save path or layout, run to the first save and do one kill→resume. +- Resume is proven by matching per-rank digests in a new process, a byte-identical first export, an aligned first old-logprob and one non-zero update; `resume_mode=auto` alone proves nothing. +- Read host memory as anon/file/shmem; `memory.current` includes page cache, and Ray kills on it. +- Pinned buffers round up to a power of 2 (Nemotron: 146.9 GiB pinned for 88.2 GiB of optimizer state, +67%); allocate large long-lived pinned blocks at exact size. +- The save step is the peak; a higher Ray threshold or smaller microbatch does not fix headroom. +- For per-step growth, use the per-process PSS floor and a small CPU-only repro before changing code. +- Verify per save path that the host cache is released after save. + +## Checkpoint and resume + +### 1. Three independent contracts + +Classify the problem before calling it "the checkpoint is broken": + +| Contract | Typical failure | +| --- | --- | +| File: serialization and publish | Disk write error; an incomplete directory treated as complete | +| Memory: transient peaks of save and load | Killed by Ray after the write finishes; cgroup OOM during load | +| Derived state: the first export after restore | Files complete, but the first step after restore mismatches ([stale quantization scale](#13-stale-quantization-scale-on-the-first-step-after-restore)) | + +The third class is a derived-state lifecycle bug, not file corruption. + +### 2. Which state exact resume needs + +| State | Requirement | +| --- | --- | +| Model parameters or LoRA adapter shards | Saved per rank; strict load that rejects missing/unexpected keys ([training forward](training-forward.md)) | +| Optimizer state, including the CPU-offloaded part | Same step as the model. Nemotron chunked optimizer offload: 88.2 GiB per rank, 94.7 GB per rank on disk; DS4.1 saves LoRA parameters and Adam `exp_avg` / `exp_avg_sq` in FP32 | +| LR scheduler | Check the saved value, not the source default (DS4.1 checks its constant, no-warmup schedule per rank) | +| RNG | Four streams: CPU, CUDA, Python, NumPy | +| Counters and data position | global step, optimizer step, data cursor; DS4.1 also saves the cumulative count of effective-signal steps | +| Run contract | Hashes of code, data and recipe, compared item by item on restore; reject a change in total steps or data by default. To extend training, declare a new phase explicitly; do not edit the old contract file or add a generic ignore switch | +| Completion marker | `.metadata` / tracker (e.g. `latest_checkpointed_iteration.txt`), written last, see 3.4 | + +Not in the checkpoint, rebuilt after restore: + +- **Parameter-derived caches:** quantization scale, packed/transposed/fused weight, version markers (see 5). +- **Records outside training:** the metrics jsonl is overwritten on resume. + +### 3. Save + +#### 3.1 Capture the first error, do not fix the last error message + +Keep the writer callback's original exception / errno, rank, path, offset, requested / returned bytes and exception chain, which DCP and Ray wrappers may drop; also filesystem/mount, free space, quota and cgroup OOM events, not just the final `ActorDiedError`. `unexpected pos` can come from disk space, quota, permissions or I/O failure; the first exception tells them apart. + +#### 3.2 Reproduce and discriminate + +A fast repro keeps the real tensors, serializer and concurrency, skips rollout/backward and repeats only save, for a bounded number of rounds sized to the failure probability; stop at the first failure and keep the scene. DS4 model-only save failed on the 3rd attempt on one path and the 6th on another; Nemotron with host cache release failed on the 2nd-4th save. + +**Acceptance for a save fix: N consecutive saves + kill→resume→save again.** Also run it before the first long run of a new save path, parallel layout or node count; running to the first save point is enough. + +Order of discrimination: + +1. Same buffer/serializer, target filesystem versus node-local disk as the control; read the errno of the first failure. +2. Control concurrency, pinned/pageable and memory peak; if the serializer changed along with the filesystem, the result cannot be attributed to one factor. +3. Check growth of threads, processes and file handles, separately from the write error. +4. Only then read kernel / filesystem client code; a matching historical issue is supporting evidence, not proof of a fix. + +#### 3.3 Reliable publish + +Serialize on node-local disk and close the file; copy to a **temporary name on the target filesystem**, handle short copies, flush/fsync as durability requires, verify integrity, then rename within the same filesystem. Publish metadata/tracker/completion marker only after all ranks finish. A move across filesystems is not an atomic rename. + +#### 3.4 Commit protocol: write the completion marker last + +| Rule | If you do not | +| --- | --- | +| Write `.metadata` / tracker last, after every rank's shards and auxiliary files (optimizer, RNG, scheduler) succeed and are confirmed collectively | A rank fails, the marker already exists, and restore reads incomplete state | +| Before overwriting a step directory, delete all old files first (uncommit) | Old optimizer / RNG silently pair with the new model | +| Before restore, all ranks check file completeness together and fail together | The rank missing a file exits; the others hang until the NCCL timeout | +| Ranks with no shard to write still join the same collective | Collective order diverges across ranks and hangs | +| A requested load with a missing file raises an error, not just a log line | Training silently continues with a fresh optimizer | + +Acceptance: a save failing midway (Nemotron hit a real ENOSPC) makes all ranks exit together without advancing the tracker, and training resumes directly; new code loads an old-code checkpoint. Never kill a job mid-write; verify the checkpoint is complete first. + +#### 3.5 Verify a written checkpoint + +Check the expected shard set, non-empty sizes, leftover temporary files, checksums where needed, and do one real full restore. A correct save digest is not a restore test (DS4.1 checkpoints 20/24 had correct save digests but were never restored in a new process, and the report says so). + +### 4. Load: enough capacity does not mean enough for the transient peak + +Load may copy the full optimizer into anonymous memory, and FSDP CPU offload may create a second set of pinned storage. Accept by loading model + optimizer + RNG/extra at real scale, not one small shard. + +Read a checkpoint with the image that wrote it: DS4.1 host PyTorch 2.11 cannot deserialize the DTensor `_FlatLayout` written by the image's 2.13; it reads only inside the image with `torch.load(weights_only=True, mmap=True, map_location="cpu")`. Do not modify the checkpoint or turn off `weights_only` to make it load. + +### 5. Derived state after restore + +DS4 chain: initial HF load saves FP8 source scales → DCP restores the trained BF16 parameters → the first export still uses the old registry → the actor re-quantizes by parameter version, but rollout uses the old scales. The first resumed step mismatches; once the next optimizer hook clears the cache it matches. The prediction "it recovers on the next step" plus a before/after-fix comparison on the same checkpoint locates the cause. + +Fix it in the lifecycle that owns the state: after load/update, invalidate or rebuild parameter-derived caches (scale, packed/transposed/fused weight, version markers), reusing an existing hook without side effects such as a real optimizer update. Do not hide it by waiting a step, disabling resume or clearing all caches by hand. DS4's minimal patch calls the existing post-update hook ([Megatron-LM #197](https://github.com/ISEEKYAN/Megatron-LM/pull/197)). List every parameter-derived cache and its invalidation code, add the missing calls and one regression test. + +### 6. How to prove resume is correct + +1. **Before launch.** For node-local checkpoints, every rank's node in the new allocation has its shard; the tracker / latest pointer points to a complete checkpoint; the launcher's resolved `resume_from_path` and global step are as expected. Confirm the restored step before the first training log line; `resume_mode=auto` is not evidence. +2. **New-process restore.** All worker PIDs changed; per-rank adapter / optimizer / scheduler / RNG digests match the save-time digests item by item. +3. **First export and old-logprob.** Exported files are byte-identical to those of the saving step; also check the exporter scale / policy fingerprint and the full valid logprob. Checking only that the second step returns to 0 misses the bug in 5. +4. **One non-zero update.** A zero-advantage step only exercises the flow (two DS4.1 resume segments had zero gradient on all ranks at step 2, proving nothing about restored adapter values). +5. **Gate coverage.** step0 cold load, first export/old-logprob after restoring a non-zero step, one later update, address / state validity under Graph and sleep/wake, new code reading an old checkpoint. +6. **What acceptance compares.** Restored per-rank digests equal save-time digests; first exported files are byte-identical to the saving step's; per-step `probs_diff` after restore is 0 under BI=1 (Nemotron: 0 on the first step after kill→`resume_mode=auto`). The resumed trajectory need not match an uninterrupted run bitwise (rollout sampling is not bitwise across processes). There is no "metrics are close" threshold; reward curves are reference only, on the same step range and basis. + +### 7. Rotation, disk and time cost + +| Item | Record | +| --- | --- | +| Size per checkpoint and save time | DS4 full32: 99 files, about 3.2 TiB, save about 269 s. Nemotron-H 30B-A3B, 4×GB200: 353 GB (EP4 + FSDP2) or 429 GB (PP4 dist_opt), save 72-152 s, about 2-3% of step time when saving every 5 steps. Layout-specific; state these for your model in the plan | +| Rotation | The mLite engine ignores `max_actor_ckpt_to_keep` | +| Checkpoint used for eval | Rotation may have deleted the step to evaluate; a checkpoint from another chain cannot substitute | +| Resume on a different node after node failure | Copy checkpoint, dequantization cache, pinned image and dependency cache to the new node before launch; disable automatic requeue; change only the node selection; write a transfer record | +| Paths for two-arm experiments | BI=0 / BI=1 output and checkpoint paths are separate; renaming the run breaks a resume chain in progress ([recipe and acceptance](recipe-and-acceptance.md#rules)) | + +## Memory budget: host and GPU + +Host memory fails in two ways in long RL runs: phase **peaks** (save / load / sleep) hitting the limit, and slow per-step **growth**; they are measured differently (1-4). GPU memory is a third class: with training and rollout colocated, rollout `gpu_memory_utilization` is set per training mode (5). + +### 1. Who holds host memory + +| Source | Mechanism | +| --- | --- | +| Pinned buffer of optimizer offload | The caching host allocator rounds up to a power of 2, so the footprint is tens of percent above the nominal state; summed across ranks it pushes residency before save close to the limit | +| Parameter shards of FSDP CPU offload | Each rank's shard stays resident on CPU; doubles if load creates a second pinned copy | +| D2H staging during save | Stays in the pinned cache pool; whether `torch.accelerator.empty_host_cache()` releases it must be verified per save path | +| CPU backup of serving sleep | L1 sleep keeps one full backup per rank; the pinned cache is not returned after wake | +| Ray object store | plasma mappings count into each process's PSS as shared pages; node shmem does not grow with it (3) | + +`param_offload` trades CPU pressure for GPU memory; optimizer offload is separate. Change one offload switch at a time and measure the phase peak before and after. + +### 2. Limits and who does the killing + +- **cgroup.** The job limit `memory.max` is usually below physical memory; budget against it. A cgroup OOM shows an `oom_kill` event and process SIGKILL, with no Python stack. +- **Ray memory monitor.** Kills workers at `RAY_memory_usage_threshold`=0.95 by default; log line `exceeded threshold of B (95.0% ...)`. Its usage figure includes page cache. Raising the threshold does not fix headroom: Nemotron at 0.985 was still killed during save. +- Distinguish GPU allocation exception, host cgroup OOM and Ray kill. A CUDA allocator allocation-failure / retry warning is not a terminal OOM and the step may complete, but it does show memory pressure. + +### 3. How to measure + +#### 3.1 One record per phase + +Log one line at load, save, sleep, wakeup, reset KV cache and generation start / end, with monotonic time, node / rank / PID and policy version. Per node, record cgroup `memory.current` / peak / events and `memory.stat` anon/file/shmem; sample process RSS/PSS, threads and FDs as needed, not with high-frequency scans of all processes. + +`generate_sequences` includes wake/init. An error inside it does not prove token generation started, let alone a backward/optimizer OOM. + +#### 3.2 Pitfalls in the readings + +| Reading | Problem | +| --- | --- | +| `memory.current` | Includes page cache, which masks per-step growth; but it is what Ray and the save peak see | +| `free`, `psutil.virtual_memory().used` | System-visible memory, not process RSS or job cgroup usage | +| `scontrol` FreeMem | Excludes reclaimable cache | +| `ru_maxrss`, Slurm MaxRSS | Cumulative peak, not current RSS | +| VmPin/VmLck | Cannot alone infer CUDA pinned | +| Sum of per-process RSS | Counts shared mappings repeatedly; use PSS | + +Pinned memory from cudaHostAlloc shows up as shmem in cgroup `memory.stat` (measured on Nemotron). + +#### 3.3 Locating "slowly grows every step" + +1. Use node `MemTotal−MemAvailable` and cgroup anon/file/shmem to find which class grows. +2. Sample PSS from `/proc//smaps_rollup` per process once a minute, recording only comm / argv0; take each step's minimum as the floor and find processes that rise monotonically. Within-phase peaks swing with offload and are useless here. +3. For a rising process read `smaps`. On someone else's job or a running job, only read `/proc`; do not attach or send signals. +4. Reproduce the candidate mechanism in a small CPU-only process, then decide where to change code. + +A floor that drops after resume (Nemotron run 1: about 455 back to about 433 GiB) shows the growth accumulates over the process lifetime. + +### 4. Capacity planning + +Budget the save-step peak as the sum of what is alive at the same time, against cgroup limit × Ray threshold: + +```text +peak ≈ resident pinned optimizer offload (actual footprint after allocator rounding) + + save D2H staging (and the part of it left in the cache pool) + + page cache of the newly written checkpoint files + + CPU backup of serving sleep and unreleased pinned cache + + resident amount of the other processes (rollout, driver, Ray workers, etc.; take the floor per 3.3) +limit = cgroup memory.max × RAY_memory_usage_threshold +``` + +Use measured values wherever possible. + +- **The save-step peak is closest to the limit.** The long-run plan states, per training mode, the measured cgroup peak at the save step and in the kill→resume segment, and which load is still unmeasured. +- **Estimate from headers before running.** Safetensors headers alone give the expanded size per dtype (full DS4.1 checkpoint: about 1421.66 GiB as BF16, including auxiliary content). It is not peak RSS; add serializer copies and mapped source pages. +- **A slope only hints at headroom risk.** Continue from a verified checkpoint and observe long enough; do not extrapolate an OOM step linearly from a short timeline. +- **On OOM, do not shrink the microbatch or raise the guard first.** Find the phase and which objects are alive together. + +### 5. GPU memory: set gpu_memory_utilization per mode when training and rollout are colocated + +The training process keeps part of its GPU memory during rollout. Nemotron single node 4×GB200: + +| Training mode | Held per GPU by training during rollout | rollout `gpu_memory_utilization` | +| --- | --- | --- | +| EP4 + FSDP2 (mLite; BI=1 arm default `ACTOR_OPTIMIZER=fsdp2`) | Not measured separately | 0.65 (verified) | +| dist_opt (BI=1 arm optional `ACTOR_OPTIMIZER=dist_opt`; numbers from PP4 dist_opt measurements) | 54-64 GiB (FP32 quantization master resident on GPU) | 0.7, 0.6 and 0.65 fail to start; measured 0.55 | +| EP4 + Megatron (mbridge) dist_opt + chunked optimizer offload (BI=0 arm, QAT + R3) | Not measured separately | 0.45 (current recipe value) | + +FlashInfer autotune warmup needs about 8 GiB more, another direct cause of failing to start. In a two-arm comparison, differences such as gpu_mem 0.45/0.65 are confounders and are listed first ([recipe and acceptance](recipe-and-acceptance.md)). + +Rules: + +- Write the value into the recipe per training mode, not as one global default; the recipe's hardware comment lists only verified topologies. +- Check the resolved value inside the container and the actual KV budget in the vLLM log, not the outer shell variable. +- `param_offload=True` saves GPU memory but moves parameters, changing the `data_ptr` of the quantization master so the pre-export check errors; as with DS4, enable only optimizer offload. +- What sleep level 1/2 release, and why wake may still find no GPU memory: [rollout engine §5](rollout-engine.md#5-lifecycle-and-gpu-memory-ledger) and [Locating GPU OOM](#locating-gpu-oom-and-staging-peaks-by-lifecycle). + +## Locating GPU OOM and staging peaks by lifecycle + +Use this when the shard estimate fits but init, export or wake still OOMs. Find which step allocates and what is still alive, and change that boundary; do not lower the batch or change sharding first, and never blame update for a failure before update. + +### 1. The minimum to record + +```text +phase / process / rank / topology +requested bytes / allocated / reserved / device free +host RSS / container or Ray limit +last completed phase / first failing call +``` + +Mark missing fields as not obtained; do not guess. If output just stops with no error, check CPU/GPU/I/O, log timestamps and phase (reading weights, compiling, collective, generating); slow is not a hang. + +### 2. The memory peak is set by "the objects alive at the same time" + +Parameters divided by GPU count estimates only part of the steady state and misses these peaks: + +| Phase | Objects and boundaries to look for | Minimal verification | +| --- | --- | --- | +| Init / FSDP wrap | Full materialization before sharding; largest single parameter; CPU-resident/tied parameters | Whether meta → shard → load takes effect; per-rank peak | +| forward / backward | Full module from all-gather, activation, workspace | Measure logprob and actor update separately; frozen parameters may be gathered too | +| Export / merge / quantize | A single full tensor, fused groups, lazy generator references | The actual export set; release timing after yield/exception | +| compile / graph capture | autotune workspace and graph pool | Find the allocation in the log; eager only isolates, recheck the production graph separately | +| KV planning / wake | KV budget, trainer residency, allocator free cache | Resolved GMU/TP and allocator metrics, not the configured value | +| checkpoint / host | Full CPU copy, optimizer gather, shared mappings, Ray guard | Per-process and node usage; RSS does not sum to the physical total | + +A small model with the same special layers, attention period and scale layout verifies the mechanism; only the full model verifies capacity, such as the all-gather peak of a huge single parameter. + +### 3. Reserved is high, so why can wake still have no GPU memory + +After update's temporaries are freed, PyTorch keeps the blocks as reserved-but-unallocated. Another allocator such as CuMem cannot use them. + +`empty_cache()` releases free cache only, not live tensors, and not reserved memory held by active blocks or a graph's private pool. `gc.collect()` helps only when collecting a reference cycle drops a tensor's last owner; high reserved alone does not prove GC is needed. + +Compare allocated/reserved/device-free and the same allocation before and after cleanup. If allocated stays high, find the live owner; if reserved-unallocated is small, empty-cache cannot explain a workspace OOM of tens of GiB. Clean up where the resource is handed back (end of update or before wake), chosen by the real call path and a retest. + +### 4. How to tell a fix from a config downgrade + +- Lifecycle error: fix materialization/release timing, the completion contract, or the wrong gather set. +- A legitimate peak exceeds capacity: pick a topology that carries it or shrink the declared configuration, and record the support boundary. +- Configuration not propagated: confirm the resolved value inside the container; editing the shell variable again is not a new comparison. +- Killed by Ray/container policy: account for actual headroom before judging the limit; do not raise the guard by default. +- Long duration: time by phase and optimize only large shares; GPU utilization is not evidence of completion. + +Do not add a resident full-tensor dump to chase an OOM; use the phase record above to decide the next step. For the full-tensor boundary of export quantization see [precision and quantization](precision-quant.md#41-fsdp-path-verl). + +### Evidence boundary of historical experiments + +- A configured path is not an executed path: record a marker/count and treat zero hits as an invalid verification. +- Before comparing tensors, align request, token position, rank/owner and valid region; equal byte hashes alone do not prove the same objects. +- Verify a root cause with a single-variable intervention and a negative control that actually triggers the failure. +- Keep failed, invalid and superseded records with reasons. A small model locates the mechanism but does not extrapolate to full-model capacity or performance; the profiler explains structure, latency comes from independent timing. +- Bind every result to the actual patch stack and imported-file hashes; overlaying only new files keeps the old overlay, so restore and apply by manifest ([code and delivery](../practices/code-and-delivery.md)). + +## Pitfalls + +### 13. Stale quantization scale on the first step after restore + +- **Symptom:** After DS4 restores from a non-zero step, the first rollout mismatches and the next step matches again by itself. +- **Cause and Fix:** See [5](#5-derived-state-after-restore). +- **How to verify:** On the same checkpoint, compare the first export's scale fingerprint and the old-logprob before and after the fix. + +### 14. resume_mode=auto starts from step 0 + +- **Symptom:** A job with `resume_mode=auto` retrained from step 0; the log said `latest_checkpointed_iteration.txt` was not found. +- **Cause:** The checkpoint is on node-local NVMe and the newly allocated nodes have no tracker / shards. Same class: shards on the NVMe of the original eight nodes with the global `data.pt` on only one, so the launcher must check the exact node set. +- **Fix and verification:** Check before launch per item 1 of [6](#6-how-to-prove-resume-is-correct). diff --git a/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md b/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md new file mode 100644 index 000000000000..45b1395a6e79 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md @@ -0,0 +1,321 @@ +# Fault localization and mismatch bisect + +Use this when training and rollout do not match, or a stage of a multi-node run died. It finds the first divergence with the fewest observations and produces a reviewable root cause, without dumping the whole model and without changing precision first. Performance attribution is in [Performance](performance.md); OOM and staging peaks are in [Checkpoint, resume and memory budget](checkpoint-resume.md). + +## Symptom → first step + +| Symptom | Do first | +| --- | --- | +| Not bitwise identical already at θ0 | [01 Isolate the side first](#isolate-the-side-first): four control groups; whichever two are equal rule out that side | +| Only some sequences are wrong | 01 same-prefix sibling: two samples of the same prompt must be bitwise identical over the shared prefix, otherwise it is the engine side | +| Wrong starting from some token position | [02 Bisect](#bisect-and-first-divergence): single-request prefill, bisect along one axis of layer/module/position | +| The difference grows slowly with training (θ0 is basically right) | [01 Isolate the side first](#isolate-the-side-first): save a small batch of prompt/response, score repeatedly at the same version to get each side's noise floor, then compare early and late adapters | +| Same operator input, still different | [Numerical contracts](../workflow/README.md), [BI kernel appendix table](bi-kernels.md) | +| Multi-node will not start / NCCL errors | [05 Communication fault ladder](#gb200-first-tell-apart-topology-implementation-mode-and-build) | +| Not sure what to save to be able to reproduce | [06 Save and compare protocol](#save-and-compare-protocol) | + +## Key rules + +1. Find the earliest divergence, not the largest layer; all later differences propagate from earlier ones. +2. Isolate the side before bisecting; classify before changing precision: input mismatch → weight mismatch → discrete path mismatch → kernel numerics → precision contract. Precision comes last. +3. Align probes by semantics (LSE in ln or log2, before or after norm, before or after the residual add). Misaligned probes produce false divergences. +4. Stop bisecting at a discrete decision (top-k, index); fix the routing, then continue. +5. A single round does not count; hidden coupling (one switch also changed scheduling) cannot be attributed to the switch; the target path must actually be exercised, otherwise the result is VACUOUS. +6. A two-node test must prove cross-node transfer happened; successful initialization does not mean bandwidth is reasonable; one `max_abs=0` does not prove consistency for arbitrary batches. +7. Looks like precision but is something else: determinism only takes effect on the main thread (switch scope). + +## Isolate the side first + +Applies when the logprob difference between FSDP/LoRA training and vLLM rollout is abnormal or grows with training. + +Let μ be the selected-token logprob saved by rollout and π the real trainer old-logprob. `μ != π` does not by itself show vLLM is wrong. The goal is not unconditional bit-exactness: rule out correctness errors, measure the noise floor, and explain growth beyond it. For quantization mechanisms see the [precision and quantization guide](precision-quant.md). + +### Per-sequence observation patterns + +Record valid/differing token counts, first divergence position, longest run of differences, prompt/response length, UID/group, and packing/rank mapping. Distinguish: a few sequences differ from the first response token; all drift slightly; differences only past a chunk/page; differences only after an update. The pattern chooses experiments; it does not prove root cause. + +Confirm token shift, mask and EOS. ±1/±2 offset comparisons are fine for diagnosis, but a fix cannot be declared through a misaligned comparison. Check that no score from another sample or version was referenced. + +### Existing rescoring tool + +A verified prefill rescoring tool takes the full sample tokens as input with zero new tokens and the current adapter, returns scores from `prompt_len-1`, checks the response token IDs, rejects missing scores, and clears the prefix/KV cache before scoring. + +### The cheapest control: do not dump all layers yet + +Save a small batch of real `prompt_ids + response_ids`, the prompt length and the valid response mask. From then on only score: do not regenerate, and do not rewrite the sequence through decode/re-tokenize. Record the base revision, adapter digest/version, config, code SHA, hardware and kernel backend. Compare training and serving under the same weights, with parameters unchanged during scoring. + +Score the same version repeatedly on each side for its noise floor, then compare early and late adapters. Rising ordinary mismatch (samples differ every round) is a signal, not proof that same-input error grows with training. + +| Path | Method | What it can answer | +| --- | --- | --- | +| Prefill | vLLM full input token IDs + `prompt_logprobs`, slice out the response | Difference between training teacher-forcing and the serving full-sequence forward | +| Decode trace | Normal prompt, force fixed response IDs, keep the real logits/logprob | Additional difference from KV/recurrent state, step-by-step generation, and decode kernels | + +`prompt_logprobs=0` means no extra top-k; target-token scores are still returned. If zero generation is not allowed, generate 1 token, discard it, and read only input scores. The first input token has no conditional score; the first response token is predicted by the logits at the last prompt position, so do not shift by one. Check trace switch, runner support and MTP limits against the actual vLLM version. Do not quietly turn off production MTP for the probe; an explicit no-MTP control is allowed but does not cover the MTP path. Prefill consistency does not imply decode or MTP consistency. + +### Metric measurement basis + +- Prefer raw logprob. When comparing the sampling distribution, unify temperature, penalty, top-k/top-p and similar processing. temperature=0 in the request does not prove which kind of logprob is returned; check the backend semantics. +- `logprob_mae` (token MAE): `sum(mask * abs(lp_train-lp_serve)) / sum(mask)` over valid tokens; also report p95/max and position. +- `probability_mae`: `mean(abs(exp(lp_train)-exp(lp_serve)))` over valid tokens. verl's `training/rollout_probs_diff_mean` is this, not logprob MAE. +- Per-token logprob difference in VERL: set `actor_rollout_ref.rollout.calculate_log_probs=true`; after the trainer computes `old_log_probs`, call `verl.utils.debug.metrics.calculate_log_prob_diff(old_log_probs, rollout_log_probs, response_mask.bool())`. It returns the 1-D absolute difference over valid tokens, not a mean; compute mean and quantiles yourself and keep mask indices for localization. The default debug metric applies `exp` first, so it yields a probability difference. Checked against VERL main at commit 8e4a572 (`verl/utils/debug/metrics.py`); on another version or fork, check return shape, mask and `exp` first, and record the VERL revision. +- Relative L2 and cosine are for tensors; do not mix them with logprob MAE/KL. Flag zero vectors separately. Report full-vocab log-softmax, probability total variation and argmax separately from the token metrics. An overall shift of raw logits is not a change in the probability distribution. There is no cross-model universal threshold. +- Schulman k3: with `d=lp_train-lp_rollout`, take statistics of `exp(d)-1-d`. It is a KL estimate only when samples come from the current rollout distribution; on a fixed old trajectory it is a diagnostic statistic, not the current policy KL. Check for overflow; do not silently clip and then claim the difference is small. + +### Build four control groups + +1. Inference alone: frozen tokens, repeated scoring, batch1/mixed batch, decode/full-prefill. +2. Training alone: A/A and A/C packing replay under the same weights. +3. Cross-stack: μ/π for tokens with the same semantics, plus full logits where needed. +4. Score layer: run the scorers separately on exactly the same logits; check dtype, temperature, logsumexp and the reduction over vocabulary shards. + +If training A/C differ and inference alone is stable, look at the training layout first, but keep the inference verification scope. If both sides are stable alone but differ across stacks, check weights, arithmetic semantics and the scoring path first. + +### Same-prefix sibling + +Same policy, same prompt, same response prefix, same score processing rules: conditional probabilities must be equal. Equal sibling μ with different π is evidence that the training side is sensitive to layout/state. For processed logprobs the penalty history and processing config must also match. + +Compare only the shared prefix, never past the generation fork. Use a clean sibling and repeated execution as negative controls. "One side looks more stable" is not proof of correctness over the full scope. + +### Localizing CUDA async errors + +An async illegal access may be reported after the faulty kernel. Save the first fault, every rank's logs, input/version and execution stage; localize with scoped synchronization, CUDA core dump/cuda-gdb, or a sanitizer. Once CUDA state is corrupted, end that experiment and collect no more numerical evidence from it. Disabling async globally or switching attention backend does not prove the stability problem solved. + +### Validity of controls + +Disabling the cache may execute a prefill that never ran before. Disabling Graph also changes speed and scheduling; to blame Graph itself, fix each step's composition offline and compare. A negative control must verify the target branch's execution count: zero difference on an unexercised path counts as not covered. For experiment design see [04 Ablation method](performance.md#ablation-method-make-ab-really-differ-in-only-one-variable). + +After isolating the side, bisect along one axis ([02](#bisect-and-first-divergence)); for classification and handling order see [03](#classification-and-handling-order). + +## Bisect and first divergence + +Separate input/weight errors, operator implementation differences and quantization error with the smallest experiment, without a model-specific debug framework. Metrics: [Metric measurement basis](#metric-measurement-basis); classification: [03](#classification-and-handling-order). + +### Choose the search axis first + +- Version regression: good/bad revision on the same reproducer, compatible environment, fixed input. +- Config regression: paired single-factor intervention; a 2×2 matrix when factors couple. +- Data trigger: request subset, length, neighbors, boundary remainder. +- Compute path: layer → attention/MoE/norm → internal operator → metadata/launch/instruction order. + +Do not change image, weights, kernel and parallel layout at once. For a dirty good/bad, save the full patch and run files; the same git SHA does not mean the same program. Build failure, dependency incompatibility and failure to reproduce are not a bad numerical revision; mark them skip/inconclusive. + +### Single-request prefill: narrow step by step, do not dump the whole model directly + +1. **Read existing training logs first.** Check the metric definition, policy version and sample source. Samples from different rounds only hint at an anomaly; they cannot prove same-input mismatch grows with training. +2. **Fixed-input scoring.** Save token IDs and position/mask from a single render. Both ends use same-source weights and the same-version adapter, and each scores repeatedly to measure its own noise. Rule out weight, input and score-semantics mismatches first ([01](#isolate-the-side-first)). +3. **Minimal forward.** Single request, one prefill, one output token. Keep the minimal prefix or layer set covering the divergent path and its upstream dependencies, with identical config on both ends and no training update. Confirm no cache hit skipped the prefill, and separate extra calls from warmup/draft/decode. An eager probe does not stand for graph/MTP acceptance. +4. **Find the earliest divergence.** Look first at the embedding, block outputs, and logits at the last valid prompt position. Align layout, valid tokens, shards, padding and deferred residuals by semantics, not by module names or tuple indices. Expand only the earliest block that exceeds its own noise. +5. **Same-input operator replay.** Give both implementations the same logical input and weights to separate this operator's difference from upstream propagation. If needed, replace only one operator's output and compare the final distribution; recurrent state must also satisfy the contract. +6. **Return to the original failing path.** After the fix, re-verify with the real adapter, input and run config. Passing with reduced layers/eager does not replace verifying the full model, decode, or CUDA Graph/MTP. + +### Per-module dump: find the earliest divergence, not the largest layer + +1. Build a semantic mapping table: layer, function, input/output role, module path on both sides, logical shape. FSDP wrapper names and fused vs separate Q/K/V need not match; combine before comparing. +2. Capture block boundaries and final logits first, on a few tokens/samples; refine the earliest significantly divergent block into norm, projection, router, expert, residual etc. +3. Align valid tokens, position, mask, TP/FSDP shards, expert ID, dimension order and padding. Never compare a physical shard with the full tensor. +4. A hook takes a snapshot, not an overwritable reference, and never changes the computation. For FSDP, sample at boundaries already materialized in this forward; do not all-gather the whole model or materialize a huge frozen PLE on GPU to observe. Export bounded data per layer/needed shard and release GPU staging promptly. +5. Record forward/recompute counts so backward recomputation does not overwrite the normal forward. Graph replay does not re-execute Python hooks, so an eager dump speaks only for eager. For graph-only problems use in-graph observation or snapshots of existing buffers; do not clear the graph or recapture to cover the problem. +6. If the operator input already differs, the operator is not yet proven buggy; give both ends the same logical input and check whether it amplifies the difference. When a fused kernel's inside is not visible, mark the observation boundary instead of guessing. + +Minimal artifacts: manifest (run conditions) + bounded tensor dump + comparison table with module/role, shape/dtype, fraction of finite values, max abs, relative L2, cosine, first anomalous position. Save a block's detailed intermediate tensors only after the summary has localized it; do not write all tokens of the whole model by default. Artifact format: [06](#save-and-compare-protocol). + +### Align probes by semantics + +Build a cross-stack boundary table: embedding, residual stream, norm input/output, Q/K/V, compressed rows/index, attention output, router logits/indices, expert input/output, combine, lm_head and scorer. Record tensor meaning, dtype, global token/shard coordinates and valid mask, not just module-name pairs. EP dispatch may reorder, duplicate or pad tokens; align expert tensors by semantic token/expert/shard key, not local offset. After expert GEMMs match, still compare separately the communication reduction, the shared + routed expert combination, and its rounding after compilation. + +If the saved rollout's first response token already differs, recompute it once from the original prompt alone. This rules out later-trajectory divergence as the only cause, but attributing it to prefill, attention or a kernel still needs controls on the same policy, position/mask/scorer and per-layer inputs. + +Compare original-dtype digests per layer or at a few checkpoints first; then dump only the neighboring boundaries and target samples. Keep one clean sibling. Snapshot async buffers before reuse without changing the production stream order in a way that hides the problem. + +If the difference does not move under any knob, suspect a too-coarse boundary first. In an external TPU MiMo case, attention core and o_proj were captured together and the difference stayed at `7.8125e-03` across CP 8/2, segment alignment 128/2048 and FP32 accumulation on/off; split apart, the core was 0 and the whole difference was in o_proj. Give adjacent or fused GEMMs separate boundaries by default. + +### Minimal but credible observations + +- Copy tensors before buffer reuse; save only necessary boundaries and tokens, with a byte cap. Model-specific mappings and artifacts stay in the model working directory. +- The manifest records code/image version, checkpoint and adapter identity, input digest, position/mask, dtype/quantization, score semantics, graph/MTP, rank/layout and intervention options. Same config does not mean same weights. +- One directory per call; write the manifest after completion. Refuse comparison on missing boundaries, repeated calls, or incompatible version/shape; do not silently take the intersection or truncate. Non-finite values fail explicitly; metrics on zero norm are marked undefined. +- Compare max abs, MAE, relative L2, cosine; for final full-vocab logits also log-softmax, probability total variation and argmax. + +### Recommended logs: identity + metric measurement basis + evidence boundary + +Reuse existing logs and JSON artifacts; no extra logging framework. One manifest per run, one detailed report per comparison. The terminal prints summaries, not full tensors. Mark unknown fields unknown; do not backfill them as verified. + +```text +[mismatch] run= step= policy= input= phase= rank= repeat= + scope= intervention= + metric= value= reference= + self_repeat= report= classification= +``` + +`metric` uses the names in [Metric measurement basis](#metric-measurement-basis) (`logprob_mae`, `probability_mae`, KL diagnostic statistic, tensor relative L2, vocab TV), never interchangeably. + +Time load, rollout, score and update separately so cold start is not counted as training. A distributed job's outer console log is only a forwarded copy: when updates stop, check TaskRunner/worker raw logs, checkpoint times and completed steps first. One forward/backward stack does not prove a hang; do not restart a healthy job because log forwarding stopped. + +### Limits of bisect + +A floating-point difference may vanish in later rounding, masking or routing, so "this layer differs" is not monotonic in depth and O(log L) does not always find the first difference. Coarse-scan for a suspect interval, then bisect/fine-scan, verify the boundaries before it, and confirm responsibility with the same operator input. + +A cut point must carry the complete state: residual, MHC mix, SSM/conv state, KV, positions, routing and necessary metadata. A fused upper-layer post and this layer's pre have no natural independent cut point. Prefix replacement that drops hidden state creates new differences. + +Only a same-input operator A/B separates upstream propagation from this operator's error. Injecting reference activations, fixing the routing or replacing the kernel can confirm causality, but they are diagnostic interventions; accept the production fix only after restoring the original path. + +### Form a reviewable root cause + +The result must contain: the minimal input, the last identical/earliest different boundary, the same-input control, the source mechanism, the trigger condition, and the minimal change that makes the failure disappear. A config difference is only a lead: in Nemotron, running both reduction block sizes with the same kernel gave consistent results for both, which overturned the guess that the block size difference was the cause. + +If after three rounds only the final score is known to differ, add observation points or re-check the sample mapping instead of rerunning in full. Keep the failing samples and the ruled-out hypotheses so the next round does not repeat them. + +### Conclusion and stop conditions + +The report answers only: what was observed, what the controls proved, what is not yet proven, and one minimal next step. Correctness errors must be fixed. Normal numerical differences are judged by noise and task impact. Quantization error needs same-input evidence; low-precision inference alone is not grounds to blame it. When improvement is needed, evaluate in order: supported high-precision rollout, QAT matching serving, real low-precision training. Do not cover wrong weights or adapters by changing precision. + +**Replacing with the reference output is diagnosis, not a fix. Lower local error does not mean lower training mismatch; per-operator ablation gains do not add up.** Stop once this round's evidence goal is reached; do not change the training config or start new experiments along the way. + +## Classification and handling order + +After isolating the side ([01](#isolate-the-side-first)) and bisecting ([02](#bisect-and-first-divergence)), classify the difference, then decide what to touch. Precision comes last. Model-specific probes, runtime patches and measurements stay in the model working directory. + +### Classification table + +| Class | Evidence | Handling | +| --- | --- | --- | +| Correctness error | Missing weights, wrong mask/position/layout, wrong scale/block, wrong quantization range, operator semantics violate the contract | Must be fixed and the original failing path re-verified; for a diagnostic task, deliver the reproducer and evidence first, without widening modify permission | +| Normal floating-point difference | Explained by same-version repeatability/reduction/kernel noise, no contract error found | Report the noise floor and task impact; bit-exact is not the only pass condition | +| Error from correct quantization | Mapping and quantization rules are correct, and a same-input control localizes the divergence to the quantization path | Precision ablation in the order below | +| Unlocalized | Insufficient evidence or confounded variables | Keep unknown, continue with minimal probes, do not declare "it is just quantization" | + +Precision ablation order: supported and affordable high-precision rollout → QAT matching the serving rules → real low-precision training (most backend requirements, last). A natively quantized model does not support BF16 just by changing the dtype. Fake quant must match the serving weight/activation rules but does not reproduce its GEMM. QAT is not a workaround for a wrong scale. QAT contract and route selection: [precision guide](precision-quant.md#1-choose-the-precision-route-first). + +Keep a single checkpoint source; open a separate high-precision adapter branch when needed. + +### Handling order + +Input mismatch → weight mismatch → discrete path mismatch (top-k, index) → kernel numerics → precision contract. Rule out each item before moving to the next. + +Acceptance: after the fix, the same failing probe passes, the original path is not broken, and error growth on fixed samples is explained/controlled; then look at training process metrics and the fixed task evaluation. Set tolerances from the model and the baseline; do not copy another model's 0.002. Statistics on a fixed old trajectory are diagnostics, not a KL estimate of the current policy. + +This procedure does not authorize stopping jobs, switching dtype, adding GPUs or changing the production recipe. When the next diagnostic gate needs new permissions/resources, report the smallest next probe and its cost. + +### Delivery + +Follow [06 Save and compare protocol](#save-and-compare-protocol): manifest, bounded dump, comparison table, semantic boundary of the earliest divergence, classification and confidence, one next step; state which paths were actually tested. No separate framework, no cross-model universal threshold. + +A case that looks like precision but is a switch-scope correctness error: [determinism takes effect only on the main thread](bi-kernels.md#determinism-takes-effect-only-on-the-main-thread). + +## GB200: first tell apart topology, implementation, mode and build + +Based on experiments on GB200 four-GPU nodes; future DeepEP/NVSHMEM releases may change the support matrix. Check the source of the target fork/commit and the installed artifacts before relying on it. + +### Do not use one name for four things + +| What to ask | Example record | +| --- | --- | +| Which implementation | AllToAll, some version of DeepEP, HybridEP; note the fork/revision | +| Which path of that implementation | The actual class, API and kernel behind high-throughput / low-latency / v2 | +| How it communicates | NVLink within a node; RDMA across nodes; or the MNNVL path actually enabled | +| What it was compiled into | 4-peer/8-peer macro, target GPU, NVSHMEM dynamic dependency | + +`NVL4 + RDMA fallback` describes the communication topology. **It does not state the DeepEP mode**: do not answer low-latency or high-throughput unless the call site proves it. How `DeepEP v2`, `HybridEP` and HT/LL relate depends on the version; do not infer containment from the names. + +MNNVL being off does not rule out cross-node EP: the verified path may be NVLink within a node and RDMA NICs across nodes. Nor does that path work on every GB200; devices, network, peer build, underlying libraries and the called API all need verification. + +### Get the world size right first + +Write the trainer and rollout process groups separately. Formulas for dense world and expert world, and grouping pitfalls such as "two DP4/EP4 are not one EP8 group", are in [training-side layout and boundaries](training-forward.md). The runtime decides the actual grouping; do not change the GPU count to fit a config string. + +### Same-fabric fairness + +Before any communication-bound A/B, prove both arms use the same transport. In one GB200 NVL case the BI=0 arm's NCCL AllGather/ReduceScatter went `via P2P/MNNVL` (same clique, cliqueSize 8) while the BI=1 arm's DeepEP low-latency a2a was pinned to RoCE (`VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=0`, `NVSHMEM_HCA_PREFIX=rocep`); a gap blamed on non-overlappable communication latency was mostly the wrong NIC, fixed by two environment variables. + +Check steps: + +1. `nvidia-smi -q | grep -A4 Fabric`: the ClusterUUID/CliqueId of the two nodes must match. +2. Run `NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET` once per arm and check whether the transport is P2P/MNNVL or NET/IB. +3. grep the arm environment files for DeepEP/NVSHMEM pins. A hard `export` overrides the value passed by the container `-e`; write it as `${VAR:-default}`. + +Mechanism: `VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=1` → `deep_ep.Buffer(allow_mnnvl=True)`, DeepEP no longer exports `NVSHMEM_DISABLE_MNNVL=1`, and NVSHMEM 3.4.5 uses MNNVL for peers in the same clique automatically. `DEEPEP_MAX_NVL_PEERS` is not a runtime switch (it is DeepEP's compile-time `NUM_MAX_NVL_PEERS`); the peer-4/peer-8 build is selected by switching the overlay. MNNVL is not hybrid-only: the DeepEP v1 LL kernel branches per destination on `nvshmemi_get_p2p_ptr` (direct NVLink write vs IBGDA), and MNNVL only makes cross-tray peers P2P-mappable; "hybrid mode" is a v2/elastic concept. MNNVL was on by default before DeepEP #190 (2025-06-06) and turned off afterwards to save NVSHMEM memory. + +## Save and compare protocol + +Large tensors and checkpoints stay in `results//`; source and small summaries go into version control. Each framework dumps differently, so write these per project following the key points below: logprob comparison, per-layer trace, replay controls, a fixed-token oracle (logprob/argmax of a fixed token set before and after a weight update), GSM8K record comparison before and after an update, a `.meta.json` provenance sidecar for the results directory, and a pre-experiment snapshot commit + clean worktree of uncommitted changes with a check of the container's actual import path and file hashes. + +The manifest records artifact path, size, hash, generating command, code/image/policy version, number of valid samples and hook execution count. Confirm the write succeeded before cleaning up compute processes. Do not auto-delete existing artifacts. + +### Archive experiment material by step + +Build the step index before the experiment and fill in status and evidence after each run, not at the end. The layout below is logical; link existing material by explicit path instead of moving or copying it: + +```text +/ + steps/ + INDEX.md # current step, evidence entry of each step, gates not passed + 01-baseline/ + INDEX.md # tuning comparison table, candidates and review records + / + config/ # resolved config, environment variables, version list + scripts/ # actual server/bench commands or pinned-version entry points + results/ # raw results/logs, or the list of persisted artifacts + summary.md # conclusion, limits, failure cause and next experiment + 05-trainer-forward/ # classified by C0-C10 gate number, created as needed + 10-final-acceptance/ # references the baseline ID and this candidate version + results// # large tensors/checkpoints, not in Git +``` + +- **Every run is traceable.** Save the actual launch command, full effective config, code/image/checkpoint identity, nodes and run time, evidence of the backend that actually ran (a CLI flag is not evidence), exit status and raw results path. A branch name is not provenance: record branch, snapshot commit, base SHA, patch, tree hash and image digest in a `.meta.json` sidecar. Replace secrets with references. Large files go to persistent storage with size and hash, never only a node-local path. +- **Baseline tuning: one row per run**: run ID, candidate config (MoE/communication), actual backend, load and concurrency, repeat rounds, valid input/output tokens, cache hits, throughput/latency, exit status, evidence link; N/A where inapplicable. Keep failed, unfinished, diagnostic and excluded runs with the reason. The summary lists the included run IDs; do not pick the best round or keep only the final mean. +- **A candidate is not a reviewed baseline.** The index marks the recommended candidate and the human-confirmed baseline ID separately; the baseline links to exact config/image/code, result set and review record, or reads "pending review". Later comparisons reference that ID; do not swap the denominator for a same-named script's new content or another backend's old result. +- **Reruns keep history.** A new run gets a new ID and overwrites nothing. Mark an invalidated conclusion with the reason and the replacing run. Continuing an experiment keeps separate attempt/segment and exit records. Reuse evidence across steps by reference, not by copy. +- **On handoff, read the index first.** The step owner updates its record. Reporting, recovery and acceptance check raw evidence through the index; a missing file, mismatched identity or chat-only conclusion is an evidence gap, not a pass. Index historical material; do not delete or bulk-move it. + +### Replayable raw data + +Keep the real framework batch: prompt/response IDs, attention/response masks, positions, lengths, UID/group/sibling, rollout logprobs, necessary reward and routing records, plus the packing mapping, parallel config and how to restore that step's rollout weights. The raw batch is the replay input; exported per-token logprobs are for comparison only. + +Step N's rollout usually comes from the policy before update N, so the post-update checkpoint may not reproduce it. Under partial rollout or async staleness, segments of one sequence may come from different policies (re-prefill after interruption continues under the new policy); record the policy version per segment, or a single-checkpoint replay turns expected policy difference into false inconsistency. Freeze weights, mutable buffers and RNG/dropout state, and state whether the scorer uses eval or training mode. + +### Compare per-token logprobs + +Each side (rollout and training old-logprob, or two replays) exports `{request_id: (logprobs, response_mask)}` as original-dtype tensors with `torch.save`. Do not convert to JSON or rebuild from printed logs. Key points: + +- **Align by logical request, not physical row.** Remove padding and restore the pre-packing UID order. Both request sets must match; a missing request raises an error rather than comparing the intersection. +- **Compare only valid response tokens.** Positions outside the mask are excluded. Every sequence needs at least one valid token; an empty batch is an invalid experiment. +- **Bitwise comparison looks at bytes, not values.** Value comparison treats `+0.0` and `-0.0` as equal and two NaNs as unequal. Compare integer views: `a.view(torch.int32)` (float32) or `a.view(torch.int16)` (bf16). Do not cast first and then claim bitwise identity in the original dtype. +- **Same contract on both sides.** Same policy version, same raw/processed basis (for processed, record temperature, truncation, penalty), same input. Compare input token hashes before scores; differences from different inputs are not numerical differences. +- **Report** valid token count, byte mismatch count, first mismatching (request, position), and max/mean absolute error. Matching sampled-token logprobs do not mean matching full logits. + +```python +def compare_logprobs(a, b): + assert a.keys() == b.keys(), "request sets differ" + n = bad = 0; first = None; max_err = 0.0 + for uid in sorted(a): + (x, mx), (y, my) = a[uid], b[uid] + assert torch.equal(mx, my), f"{uid}: masks differ" + x, y = x[mx], y[my] + ix = x.view(torch.int16 if x.dtype == torch.bfloat16 else torch.int32) + iy = y.view(torch.int16 if y.dtype == torch.bfloat16 else torch.int32) + diff = (ix != iy).nonzero() + if len(diff) and first is None: + first = (uid, int(diff[0])) + n += len(x); bad += len(diff) + max_err = max(max_err, (x.double() - y.double()).abs().max().item()) + return {"tokens": n, "mismatched": bad, "first": first, "max_abs": max_err} +``` + +### Per-layer trace: find the boundary of the first divergence + +When logprobs mismatch, record tensors with forward hooks at the same semantic boundaries on both sides (for example each layer's attention output, MoE routing, residual), find the first differing boundary in execution order, then do a bounded dump there. Key points: + +- **Key by semantics, not rank.** Each record's key is `(request_id, boundary name, logical slice)`, for example `("step7-p2-s3", "layer2.attn.out", "tokens=0:128")`. Stitch TP/CP/EP shards back by global coordinates and remove padding before recording. Naming and order must match on both sides. A missing probe invalidates the whole experiment; it cannot be skipped. +- **Hashes first, tensors on demand.** Take SHA256 of the contiguous tensor's original-dtype bytes (`hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy()).hexdigest()`), with shape, dtype and semantic execution sequence number. A hash shows where it differs, not by how much; save raw tensors only at the first differing boundary. +- **Clone before the buffer is reused.** Inference-side workspace and CUDA Graph output buffers are overwritten by the next call; `clone()` in the hook before hashing or saving. +- **The first difference is not the root cause.** To blame an operator, also prove its input is identical on both sides (record its input hash too). + +### Replay controls + +With fixed rollout and fixed weights, change one variable (packing order, batch composition, parallel layout, kernel switch), replay old-logprob, and use the comparison above to judge whether that variable affects the result. Key points: + +- **One process per variant**, with a synchronous entry point that blocks until the replay finishes. Do not use background submission such as `sbatch`: on timeout or failure the job cannot be reclaimed. Variants that change compilation, libraries or initialization environment need a new process. +- **Each variant writes to a new directory** without overwriting old results; record the actual command, environment, exit code and elapsed time. Preview the command and budget before executing; on failure, stop at that variant. +- **Prove the replay is frozen**: record the weight fingerprint, confirm no generation and no optimizer update, and confirm the packing actually in effect, not just the command-line arguments. +- For the experiment record use [the report template](../workflow/report-template.md). diff --git a/.agents/skills/train-inference-consistency/guides/environment-provenance.md b/.agents/skills/train-inference-consistency/guides/environment-provenance.md new file mode 100644 index 000000000000..c418f16d21a7 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/environment-provenance.md @@ -0,0 +1,106 @@ +# Environment, image and runtime version provenance + +Once the environment is set up (the image pins the native runtime, uv manages training dependencies, source arrives as a traceable snapshot or a minimal overlay), prove that it is what is actually running. Code existing does not mean it is executing. + +## Find by symptom + +| Symptom / question | Where to go | +| --- | --- | +| Not sure which version of source / wheel / .so / image is actually running | [Run manifest](#use-one-run-manifest-to-answer-what-exactly-is-it-right-now) | +| How many commits "latest" actually differs from the target branch | ["Latest" must come with a target branch](#latest-must-come-with-a-target-branch) | +| Overlay, or change the source and rebuild | [Boundary between overlay and native wheel](#boundary-between-overlay-and-native-wheel) | +| Recompiles on every start / not sure where the cache should go | [Compile cache](#compile-cache-where-it-lives-decides-whether-it-can-load) | +| Not sure which path DeepEP actually runs / how it was built | [Four things about the communication implementation](#communication-implementation-do-not-use-one-name-for-four-things) | + +## Key rules + +1. Pin the image digest, CPU architecture and Python version first, then resolve training dependencies. Image tag, PyPI version and upstream PR merge status do not replace one minimal inference run in the image. +2. `--system-site-packages` does not stop uv from installing a same-named package that then imports first; check the install plan, the actual import path and the GPU operators together. +3. Each of the six manifest layers (source, installed packages, native libraries, JIT kernels, container, run) has evidence nothing else replaces: `latest` is not the digest, `pip show` in another venv is not the actual import path, the Python source HEAD is not the loaded .so. +4. Audit inside the worker, interpreter and mount namespace that actually run; the login container shell and the Ray worker may differ. +5. One image tag may point to different digests on different dates/nodes; ARM64 and AMD64 tags may be independent builds. +6. Modify in the host worktree, execute in the container: same-commit worktree plus per-file hash check, never an overlay stack detached from the formal source. +7. Write version pins such as "the Megatron fork is only compatible with Megatron-Bridge v0.6.2" into the run manifest. +8. How the image was produced (container commit exported to sqsh, or Dockerfile rebuild), where it was pushed, and the three-repo SHAs in the tag are part of the run identity. + +## Environment identity: code existing does not mean it is executing + +### Use one run manifest to answer "what exactly is it right now" + +Record recipe, trainer, Megatron/MCore, vLLM, DeepEP, DeepGEMM, FlashAttention/FlashInfer and NVSHMEM separately, each with at least: + +| Layer | Evidence to save | What cannot substitute for it | +| --- | --- | --- | +| Source | remote, target branch, HEAD, dirty patch hash, submodule gitlink/actual HEAD | `latest`, the repo directory name | +| Installed package | Actual interpreter, import path, package version, wheel/build revision | `pip show` of another venv | +| Native library | Path/hash of the `.so` actually loaded, build inputs, target architecture, dynamic dependencies | Python source HEAD, extension file name | +| JIT kernel | DSL/source hash, compiler version, specialization/config, cache origin | `.so` unchanged | +| Container | OCI manifest/platform digest or `.sqsh` SHA256, parent image, build record | A mutable Docker tag, image creation time | +| Run | Resolved config, worker environment, model/weight digest, the real kernel | An export in the head node shell | + +Write `unknown` for unknown items; do not infer the whole source tree from a wheel version string. Audit inside **the worker, interpreter and mount namespace that actually run**. If dynamic linking is suspect, check the process's actual mappings/loader result, not the static prediction of `ldd`. + +ARM64 and AMD64 tags may be independent builds or members of one multi-arch manifest. Compare the target platform's manifest/config/layers; a similar tag or update time does not show it is the same image. Record generic scheduling switches such as `CUDA_DEVICE_MAX_CONNECTIONS=1` per the task convention; do not add them quietly and then credit BI. + +A cross-repo compatibility pin is a manifest row too: state the pin (the Megatron fork works only with Megatron-Bridge v0.6.2) and its basis, and re-verify when either side is upgraded. Do not assemble incompatible releases just so every item reads "latest". + +#### Communication implementation: do not use one name for four things + +See [debugging mismatch](debugging-mismatch.md#do-not-use-one-name-for-four-things). + +### "Latest" must come with a target branch + +Pin each repo's target branch per the [delivery target](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). The reference implementation's PR is a design reference, not automatically the PR base. If the user names another fork/branch, record the override. fetch is not checkout; submodule update follows the parent repo's gitlink, not each submodule's own remote. + +```bash +git -C /explicit/repo status --short --branch +git -C /explicit/repo remote -v +git -C /explicit/repo fetch origin +git -C /explicit/repo rev-list --left-right --count A...B +git -C /explicit/repo log --left-right --oneline A...B +git -C /explicit/repo merge-base A B +git -C /explicit/repo diff --stat A B +git -C /explicit/repo submodule status --recursive +``` + +The left/right counts of `A...B` say how many commits each side has; the merge base and diff say what differs. Diverged branches have no single before/after relation. Commit time and short-hash magnitude do not decide ancestry. If a shallow clone lacks the ancestor, fetch history before claiming the lines are unrelated. + +Keep the recipe's declared compatible combination by default. If the user explicitly asks to update a development branch, record the deviation from the pin and check the actual changes and the ABI. + +### Boundary between overlay and native wheel + +During development, fix the container code-sync and build/test method first and record the tool versions in the run manifest. Modify in the host worktree, execute in the container; do not maintain an overlay stack detached from the formal source: + +1. Confirm the config, model registration and import identity of the native image/checkpoint; pin the image digest, the full source commit and the real install directory. Do not guess a missing implementation from the model name or tag. +2. The sync worktree's `HEAD` equals the image's full commit, and the changes under test are a diff on that baseline. They may be ported from the development branch, but a different HEAD cannot bypass the check. +3. In a container not yet running the workload, sync each changed in-package Python file individually: before copying, check the original file's hash against `HEAD`; after copying, check the hash of the file under test; save the sync manifest. For new files, check no original exists on either side. Never overwrite the whole package or modify a running container. +4. Keep the image's version, generated files and unchanged extensions. Recompile native/binding changes and verify the artifacts; for JIT changes check compile inputs and cache; if the ABI/toolchain changes, rebuild the image. +5. After startup, re-check import/library paths in the actual worker and test. Finally build the unified image from the formal PR's clean commit and rerun acceptance; the user does not redo the development sync. Details: [Implementation and delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). + +Overlays are for development/diagnosis only, not the delivery format. + +- An overlay covers specific files in the container with external copies; a wheel is an installed package that may carry compiled machine code. +- Cover only the necessary files at their exact target paths; do not mount the whole repo over the installed package and shadow its extensions. Verify the target's original hash first, apply once at startup, stop on failure, and save the overlay manifest, final hashes and import paths. +- A Python-only change can still change the binding/JIT/shape contract; it is not automatically exempt from recompiling or testing. +- Changes to CUDA/C++, bindings, compile macros, submodule kernels or the CUDA/Torch ABI need checking and recompiling. Before acting on a "missing extension", rule out the wrong venv, shadowing by a whole-repo mount, search paths and dynamic dependencies. +- If the image already has the correct artifact, verify it and use it. Do not recompile a DeepEP already proven to work, and do not change only the source pin in the recipe expecting the old `.so` to change. +- Do not fix a FlashAttention extension load failure by `unset`ting a variable. Accept a config fix only once it is shown the variable selected an unneeded/incompatible path and the run passes with the user's target path restored. +- Keep observability overlays separate from numerical/performance changes and off by default; when cleaning up the main PR, back up still-useful diagnostic code before deleting it. + +### Delivery gate for immutable images + +1. Commit production changes to their PRs; freeze the clean source commit and all native/JIT build inputs and dependencies; record the base image digest, not just the tag. A dirty patch may identify a diagnostic candidate but cannot be an uncommitted dependency of the final run. +2. Build the final image and check import/dependencies/hashes in a fresh container **with no development source mounts and no inherited JIT cache**, keeping the device, model and data mounts production needs. +3. Run the agreed correctness, Graph/cache and performance gates in that image; if observability was on, accept again with the formal config and no probes. +4. Publish exactly the accepted artifact and verify the registry digest/architecture; if conversion or repackaging changes the content, verify the equivalent closure and run a final smoke. +5. Sync recipe, PR description and image receipt; "image push succeeded" does not cover items that failed. Follow this order when the user asks to push after testing. + +The receipt states how the image was produced: Dockerfile rebuild from a clean commit, or changes installed in a container then committed and exported as `.sqsh` (the only option when the coordinator has no docker permission; fine for development, but formal delivery returns to the Dockerfile); which registry it went to (Nemotron: public Docker Hub); a tag carrying the vLLM/Megatron/verl SHAs, not just the date. Record both the `.sqsh` SHA256 and the registry digest under the same tag. + +Adding a new project source layer to an old OCI base does not bring in another experiment image's vLLM/FA4 changes. A deleted file missing from the top-level tar may still exist in lower layers; deleting a base file needs a correct whiteout or a check in the final rootfs. + +Upgrading FlashInfer/NVSHMEM must not silently upgrade Torch/CUDA. Pin the target interpreter and constraints, compare packages/dynamic dependencies before and after, and confirm the APIs used and build compatibility. A version number is not evidence of an HCA fix or of model correctness. + +## Compile cache: where it lives decides whether it can load + +Put the Triton/Inductor/vLLM compile caches on node-local disk mounted explicitly into the container; a `noexec` mount (such as `/dev/shm` in some containers) cannot load the compiled `.so`, and a per-container tmpfs recompiles on every start. Isolate caches by image, source commit and GPU architecture, and give each rank on a node its own copy, since an existing cache directory does not mean a cache hit. diff --git a/.agents/skills/train-inference-consistency/guides/performance.md b/.agents/skills/train-inference-consistency/guides/performance.md new file mode 100644 index 000000000000..97419e524a8a --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/performance.md @@ -0,0 +1,160 @@ +# Performance and profiling + +Record performance gates and correctness gates independently: low throughput does not mean the numerics are wrong, and correct numerics do not mean performance meets the target. BI kernel choices and costs are in [BI kernels](bi-kernels.md); the statistical protocol for paired acceptance is in [Pairing and computation](recipe-and-acceptance.md#pairing-and-computation). + +## 1. Workload and measurement basis + +### Three throughput measurement bases + +- **Per-request tok/s**: 1/TPOT; decode latency. +- **Engine total tok/s**: completed tokens of all concurrent requests / common timing window. With multiple engines use the full common window and total completed tokens, not a sum of per-engine throughput; with unsynchronized node clocks, do not compute the window from timestamps. +- **RL one-step gen time (makespan)**: one batch of rollouts from the first request in to the last request out. This is what RL cares about; long tails and scheduling can move it opposite to the other two. + +State which basis you report; compare ratios only on the same basis. + +### Translate the rollout target into actual kernel workload + +"batch=128, n=8, input 2K, output 8K" must at least specify: + +- Is 128 prompts, expanded requests, or per-engine concurrency? 128 prompts with 8 siblings each is 1024 requests. +- Is 2K/8K the maximum or a fixed length? Record input token mean/max/percentiles, total output, and completed requests. +- `max_num_seqs` limits simultaneously active sequences; `max_num_batched_tokens` is the token budget per scheduling pass; neither equals the total request count. +- Requests per DP engine, prefill chunk size, and the active decode batch over time; a global 128 is not each rank's kernel M. +- Prefix cache cold, warm, or naturally hit; the n siblings of a prompt affect hits and prefill. +- Natural EOS vs ignore EOS / fixed length, sampling parameters, seed, request order, maximum model length, scheduling policy. + +### replay bench: equal-workload replay of one RL step + +Use the same request count and prompt×n as one RL step, take each request's output length from one real rollout, replay with ignore EOS at fixed lengths, and measure makespan. Measure the repeat noise of one config first, then compare arms; draw no conclusion from a difference smaller than the noise. With the same length table, token count, request count and long tail are identical, so the difference belongs to the execution path. Fixed output suits equal-workload comparison, not quality. + +In natural generation, different numerical paths or scheduling can change token counts even with the same weights and seed; record output token receipts and add a frozen-token replay if needed. A short 2K+64 profile can locate startup/local hotspots but **cannot** represent the long-term decode, cache/page pressure, scheduling tail or full rollout speed of 2K+8K. Use different experiment IDs for the short profile, the full stress test and the training E2E. + +### Metrics are cumulative values + +When one engine runs multiple stages (concurrency or parameter sweep), `vllm:iteration_tokens_total_*`, `vllm:spec_decode_num_*` and `vllm:num_preemptions_total` in `/metrics` are cumulative since process start. Always compute per-stage cost as `after - before`, and note in the script that the first stage is right only because before is all zero, which makes the method look correct. Cross-check the prefill amount with an independent client-side basis (prompt_tokens/cached_tokens in the request records); if the two disagree, the counter basis is wrong. + +When the verification method itself fails (a regex that never matches, a hardcoded version number, an unsubtracted cumulative value), the numbers look plausible but measure something else. Reconcile every new script's first numbers against an independent basis. + +## 2. Fair A/B + +### Do not mix the two kinds of A/B + +- **Causal comparison of the BI switch**: same software/model/scheduling/cache, only BI changes. If the switch selects a different kernel automatically, record the actual branch and add no other optimizations. +- **Each side's best available config**: BI=0 uses the strong default / verified best path for that platform and model; BI=1 uses the fastest path that meets the agreed consistency. Backends may differ; disclose every difference. + +Do not force BI=0 onto FA2, a specific MoE, non-default attention, or Graph off to shrink the gap. Do not call an experiment that locally disables switches inside BI=1 "native BI=1". A replaced path is a separately named candidate that passes the full contract gate. Give common optimizations (cutting shared launches, enabling a supported network path) to both sides and report the common gain and the reduction in BI overhead separately. BI=1 faster than BI=0 is not automatically fake, but it triggers an audit of config, workload, cache and timing. + +### What to fix + +Weights, input, generated tokens/lengths, output cap, concurrency, nodes, software, warmup and timing boundaries. If output receipts and workload differ, report only each side's own throughput. Separate the fresh-process first rollout from warm steady state. + +**Same fabric**: before comparing communication bottlenecks, confirm both arms use the same transport; see [Same-fabric fairness](debugging-mismatch.md#same-fabric-fairness). **Same node**: small gains need same-node pairing, crossed order and multiple rounds; counterbalance across node pools and do not pick the fastest cross-pool combination. **JIT pollution**: a JIT kernel specialized by token count recompiles for every new prefill size under BI=0; verify zero compile lines during the bench. **BI=0 silently running the BI kernel**: a fork whose gate only checks whether the `.so` exists makes BI=0 run the BI kernel; grep `vllm::batch_invariant::` in nsys. + +### Report speed switches and numerical switches separately + +One switch (PC, cudagraph, BI) changes both the numerical path and throughput. Before reporting "probs_diff with BI on/off" or "throughput with BI on/off", list the confounders that change with it, rule out or fix each, then compare; see [Paired acceptance](recipe-and-acceptance.md#paired-acceptance). + +### Computation and reporting + +```text +output tokens/ms = completed output tokens / (wall_seconds * 1000) +BI throughput deficit = 1 - throughput_BI1 / throughput_BI0 +BI time overhead = time_BI1 / time_BI0 - 1 # matched work only +``` + +Latency reduction `1-new/old` and throughput gain `old/new-1` differ; 16.35% lower throughput is not 16.35% more time. Report numerator, denominator, request / completed count, raw receipts, config differences, sample count and variation; take the median over pairs and report single pairs separately. Gains of A and B do not add; measure the combination, and do not claim it beats the best single item from a few observations. List skipping the old-logprob forward separately, only when the RL algorithm allows it; it is not a gain of the BI kernel. Every ratio must be recomputable from explicit artifacts. + +## 3. Breakdown and ablation + +### Two stages: remove structural defects first, then fine-tune + +Once proxy inference consistency passes, check feasibility: sweep token row counts on real operator shapes, count launches and host gaps, look for per-token launches, synchronous readbacks, repeated copies and abnormal workspace, and confirm the target Graph captures/replays. After backward is wired in, measure training forward and forward+backward (by token count) on the main topologies fixed at C1 against the native training path on the same build, and check them against the C1 budget, so a scheme fast at decode but slow in training/prefill is not chosen. After full consistency, fine-tune by E2E profile -> micro benchmark (screen high-share shapes, then sweep all target shapes and boundaries; all must be correct; weight performance by real call frequency) -> multi-round paired E2E on the same workload. New or replacement kernels do not use Triton; use CuTe DSL, CUDA/CUTLASS or an existing implementation; historical Triton data is only a comparison. Final delivery still needs full acceptance, see [validation](../workflow/validation.md). Fix a structural problem first and re-check the numerics; do not turn this stage into endless tuning. + +### Break down the overhead first + +Time model load/JIT, rollout prefill/decode, old-logprob, backward/update and offload/wake separately. Kernel times may overlap and do not add up to wall time. Keep the nsys window, rank/GPU normalization, Graph node tracing and sampling interval consistent. + +Profile the effective config and actual kernels, not just the BI switch. In DS4 the same quantization `.so` loads under both BI0 and BI1, but different MoE layouts select different schedulers. For an NVFP4 MoE model, account for dense linear, expert GEMM, expert routing/indexing, all-to-all and attention separately; one vendor name does not describe the path (one Nemotron run mixed FlashInfer FP8 linear, Humming NVFP4 dense GEMM, Humming indexed experts and FlashInfer NVLink one-sided all-to-all). A run that fails on a backend guard (a BI=0 trial expected Humming but got FlashInfer CuTe DSL) calls for re-auditing the backend and build path and does not enter the performance denominator. + +### Collect enough attribution information in one profile + +1. Verify the actual serving config and input/output statistics. For inference-only use vLLM serve or an existing client, not training; use the mp/Ray executor actually needed. +2. Add low-overhead, off-by-default NVTX at semantic module boundaries: attention, indexer/Top-K, index preparation, MLA, MHC, shared/routed MoE, norm, linear, sampling, communication. Do not apply the DS4 module list to dense models. +3. Use nsys first for the timeline and hotspots, with CUDA Graph node tracing on; graph launches without node details do not support per-kernel comparison. +4. Count calls, GPU time, CPU launch/sync, allocation and shape per path; aggregate by rank, phase and module. +5. Pick representative prefill, steady decode, mixed and tail windows, then ncu a few kernels/shapes for bandwidth, occupancy, memory access and compute; ncu-perturbed wall time is not throughput. +6. Measure throughput on the same full workload without profiler or heavy dumps, and record whether profile and timing used the same code and config. + +A kernel name is not a module identity: copy/fill/cast, `index_select`, generic elementwise, GEMM, and different shapes of one template appear in many layers. NVTX/launch correlation and shape are needed to tell "the indexer's gather" from other gathers. Parent/child NVTX inclusive times do not add; the sum of concurrent kernel GPU durations is not wall time; label multi-GPU sum, per-GPU mean and critical-path time separately. + +When a wrapper/overlay also rewrites the kernel selector, confirm who launched a kernel with **three layers of evidence**: per-worker class, object count and import source at initialization; input shape, module name and dispatch at the real forward callsite; launch correlation and GPU kernel in the final profiler. Overridden linear objects or a Humming shape do not prove what launched at replay, and overlay env vars set where the worker does not inherit the entry point falsely suggest the candidate is enabled. Graph may hide ordinary Python CPU op events, so a `cpu_op` count of 0 alone does not prove the callsite did not execute. When evidence conflicts, run a minimal probe that keeps the numerical path; do not fall back to an unverified default operator for throughput. + +For a Roofline bound, use the peak compute of the actual dtype and the bandwidth of the memory level studied; L2 / shared-memory / non-Tensor-Core kernels cannot borrow the HBM / Tensor Core bound, and do not infer the metric from chart coordinates. + +### Pin "where is it slow" to a process and a call stack + +The frontend and control plane are in scope. Split the makespan difference by process: py-spy the engine, worker and frontend once each and see which function holds the main-loop time; use asyncio loop-lag to see whether the frontend is saturated. With R3, most of the frontend main loop is in `output_handler`; without R3 the frontend is almost idle, which moved suspicion from the GPU to the frontend object count. Optimize only parts with a large share; GPU utilization is not evidence of completion. + +### Single-operator replay and synthetic micro benchmark + +A synthetic shape matrix proves algorithm boundaries; real rollout replay supplies score distribution, ties, masks, non-contiguous layouts, cold/hot buffers and call frequency. Weight by sampled shapes/frequencies and keep high percentiles and rare boundaries; report decode small M and prefill large M separately. Time with and without workspace allocation; if the original allocates per call and the candidate preallocates, do not compare only the kernel body. A micro A/B fixes input and resources and checks the candidate bitwise against the baseline; change one thing at a time, covering at least the batch sizes actually used (1/4/32/128), different neighbor compositions, prefill/decode, Graph, mask/tie/boundaries; mark shapes not run as not covered. When Triton is forbidden, a "temporary benchmark" does not bypass that. + +### Ablation method: make A/B really differ in only one variable + +These rules apply to intermittent numerical leaks (the batch-invariance kind) and to performance attribution; the examples come from the DS4 batch-invariance investigation. Process rules are in [Experiment discipline](../practices/evidence-and-experiments.md#experiment-discipline). + +1. **A single round of 0-diff does not count.** For a leak with a hit rate below one half, a single all-green round has probability above one half. Re-check a statistical conclusion with a deterministic reproducer, especially "turned clean"; 0-diff is necessary but not sufficient. +2. **For "turning off X fixes it", first check whether X's switch also changes another branch.** Example: with cudagraph mode NONE, the budget split function also took a different branch, so the eager experiment moved two variables. Hidden coupling cannot be attributed to the switch. +3. **A switch that changes speed also changes the scheduling composition.** Eager is several times slower than graphs, so batch composition changes. Before a statistical comparison, pin each step's batch composition with offline manual steps; otherwise "clean when off" may just mean it was not hit. Turning off the cache is the same: a prefill that never ran before now runs. +4. **The target path must actually be exercised.** A long-prompt filler looks "heavier", but the step exceeds `max_cudagraph_capture_size` and the graph does not run, so identical results across settings are an illusion. A negative control records the target branch's execution count; zero hits is marked VACUOUS and does not count as verification. + +Before localizing, draw a 2×2 (or higher-dimensional) matrix to separate the variables, then choose which cells to run. + +### Ablations must be able to identify interactions + +With optimizations A and B, measure `base / A / B / A+B`; with network/Graph/cache switches, add the corresponding controls instead of running the full combination mechanically. Reuse A+B data only under identical weights, input, software and timing. The first round after restoring checkpoint N is usually step N+1, and adjacent names do not prove the same policy: fix the checkpoint content digest and export fingerprint, and do not mix BI0/BI1 or R3/no-R3 branches. For rollout-only timing, exit cleanly after completion without optimizer/save, and record that all requests completed. Top-K and index post-processing can usually be ablated orthogonally, but shared buffers, layout or Graph create interactions, so re-accept the combination. + +### Segment-by-segment isolation: add the total difference back up + +Cut a long path into arms at boundaries, each enabled up to one segment, and measure what each segment adds: + +| Arm | Enabled up to | What it answers | +| --- | --- | --- | +| All off | No capture | Baseline | +| A0 | Capture only (hook on, nothing lands on the data plane) | Cost of the capture kernel/sync point | +| A | Capture + worker data plane + packing, not sent to the scheduler | Worker-side processing cost | +| E | worker→engine IPC on, scheduler discards | Inter-process transfer cost | +| B | No capture, return path only (fake data) | Return path cost, independent of capture | +| All on | Current path | Total difference | + +Arm differences should roughly sum to all-on minus all-off; otherwise there is an interaction or a dirty cut. This localized the Nemotron-H R3 return path: A0 and A are almost level with baseline, E adds a little, B is close to all on, so about 80% of the overhead is on the per-step engine→frontend return path, not the capture hook. + +Measure same-config repeat noise first and ignore arm differences below it. + +### Optimization priorities and correctness gates + +1. Remove idle/repeated work: launch grids mismatched with the real workload, repeated moves, invalid tokens. +2. Reduce host scheduling/sync and kernel launches; restore a usable Graph with the output contract unchanged. +3. Fuse boundaries such as gather/dequant/index post-processing, keeping the original rounding, traversal order and padding/sentinel. +4. Only then change the kernel algorithm, chunking or collectives; these may change the reduction order and need new equivalence evidence. + +Example: a quantization kernel launches 1024 blocks per expert while decode has few valid rows. Shrink the grid and the stride together; under Graph capture use a host-known capacity bound, not the device token count. If dequant/gather are element-wise independent, re-chunking may keep bit-level equivalence; fusing into sparse attention must keep the scale rounding and KV traversal/reduction order. + +Each candidate first runs a same-input kernel comparison (real dtype/shape, non-contiguous or masked layout, empty experts, tail, extreme values, ties, reused dirty buffers, Graph replay), then the fixed-token engine gates, then actual rollout/training. Formatting checks do not replace compiling and executing; rerun measurements on the final code, since a JIT refactor that looks behavior-preserving may fail to compile. A candidate with no gain or a regression keeps its conclusion and is withdrawn from the default path. + +### Delivery + +Each item lists the actual kernel/path, correctness before and after, micro latency, rollout and full-step results, independent sample count, Graph/topology scope, raw receipts and unsupported configs. State whether quality evaluation and long-term stability are done; kernel speedup does not replace an end-to-end conclusion. + +## 4. Lever results table + +Measured switches: directions and dead ends not to retry. + +| Model / arm | Winners | Dead ends (measured, do not retry) | Structural floor | +| --- | --- | --- | --- | +| DS4-Flash GB200 DP8/EP8 BI=0 | FlashInfer trtllm MoE (`--moe-backend auto` + new FlashInfer); `cudagraph_mode=FULL_AND_PIECEWISE` | deepep_low_latency, `--api-server-count>1`, NCCL_ALGO/PROTO, max-num-batched-tokens 4k/8k, linear backend fitrt/cudnn (cutlass crashes), DBO (forces the V1 runner, which does not accept PIECEWISE) | MoE decode is weight-bound (gemm1 near HBM roofline); DBO rereads weights and loses; overlap capped at the shared experts' few µs | +| DS4-Flash BI=1 | Deterministic top-k + fusion, two PRs together (each alone has almost no effect); DeepEP LL over MNNVL (`VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=1` + nvl8 overlay) is the biggest lever; after that the shared experts aux stream helps only under MNNVL | DBO, TP2/DP4, TP4/DP2, the skinny bi_gemm, async-scheduling (already default), splitting the shared MLP to fill the combine window | MoE backend locked to DEEPGEMM (the only BI fp8 MoE) -> a2a locked to deepep LL; per-layer dispatch+combine is serial; cross-layer overlap breaks BI; finally the TTFT floor | +| Nemotron-H BI=0 | PC off (when R3 forces PC on, upstream #58851 is best); R3 finish-only + per-step packed IPC (overlay), makespan returns close to the R3-off floor; then packing per-step logprobs into a list for the frontend (overlay) unsaturates the frontend and BI=0 is 16% faster than BI=1 | Raising the zero-copy threshold (no effect); `max_num_seqs` 256 (slower); PR 57968 keys-only as it stands | PC on (Mamba align mode) has a fixed cost: extra eager launches per decode step, 8K prefill split into two segments | +| Nemotron-H 256K training update (both arms) | BI=1 per-GPU per-micro-batch token cap 64K -> 128K: update only about 5% faster per token (13.1 -> 12.4 s / million tokens), actor reserved 99 -> 168 GiB (184 GiB per GPU, cap 170); small gain, thin margin | Less recompute (some layers, only Mamba layers, none): both arms OOM at 32K-64K per GPU. Per-GPU token cap for BI=0 (no difference across 32K-192K): with CP>1 the recipe disables dynamic batching, each micro-batch is one sequence, and the parameter has no effect | Update is 14-18% of the step, rollout about 80% (sandbox execution, slowest trajectories). Compute the share before tuning; confirm in the resolved config that a parameter takes effect | + +BI kernel choice: [BI kernels](bi-kernels.md). diff --git a/.agents/skills/train-inference-consistency/guides/precision-quant.md b/.agents/skills/train-inference-consistency/guides/precision-quant.md new file mode 100644 index 000000000000..ddb57a61d7a1 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/precision-quant.md @@ -0,0 +1,490 @@ +# Precision, quantization and QAT + +The weight conversion chain from optimizer master to vLLM bytes, QAT and low-precision training, and known pitfalls. Numerical contracts without quantization: [numerical-contracts](../workflow/README.md). + +## Find by symptom + +| Symptom / question | Where to go | +| --- | --- | +| Weights unchanged, probs_diff grows after export | [P1](#p1-scale-rule-differs-from-the-checkpoint-requantization-flips-codes), [P10](#p10-both-sender-and-receiver-quantize), [P11](#p11-receiver-picks-quantization-targets-by-name-suffix-and-misses-fused-moe) | +| NVFP4/FP8 scale or packed bytes differ from the checkpoint | [Conversion chain §2](#2-inventory-the-checkpoints-full-quantization-rules-first), [P1-P5](#quantization-rules-and-scale) | +| NVFP4 rollout: QAT or BF16 actor? | [QAT §1](#1-choose-the-precision-route-first), [§4](#4-how-much-mismatch-qat-can-remove) | +| QAT is on but the mismatch did not change | [P7](#p7-qat-switch-is-on-target-modules-were-not-replaced), [P8](#p8-training-grid-differs-from-the-deployment-grid) | +| QAT loss does not drop, deployed weights byte-for-byte unchanged | [QAT §6](#6-why-qat-does-not-train), [P24](#p24-master-starts-on-the-quantization-grid) | +| Whether to add activation fake quantization | [QAT §4.1](#41-activation-fake-quant-aligned-per-layer-worse-end-to-end) | +| Same operator input, inference module output differs | [P13-P15](#training-side-numerics) | +| Fixed logprob offset under FSDP | [P16](#p16-fsdp-bf16-policy-overrides-the-fp32-parameter-contract), [P20](#p20-dividing-by-temperature-in-bf16) | +| Some LoRA/adapter gradients are always zero | [P17](#p17-fake-quant-cuts-off-the-gradient) | +| Adding a new quantization format | [Conversion chain §3](#3-what-each-format-outputs), [§5](#5-minimal-adaptation-for-a-new-quantization-format) | +| Implementing or changing backward/STE | [Backward validation](backward-validation.md) | + +## Key rules + +1. "NVFP4/FP8" is only a format name; a different block scale selection, global scale convention, rounding or static/dynamic activation gives different bytes. Inventory the checkpoint's full rules first. +2. Use the rules of a mainstream implementation (TE, ModelOpt, NeMo-RL) that matches the checkpoint; do not invent your own. Verify from the real BF16 release, not from dequantized weights. +3. Training QAT and the exporter use the same quantization function on the same grid (range, axis, scale, rounding). +4. Compute scale on the full logical tensor. A fusion group shares scale (VERL NVFP4: QKV, gate/up fused global scale). Quantizing each shard separately and concatenating is forbidden. +5. Each tensor family has exactly one conversion owner and is quantized exactly once; post-load reordering does not count as quantization. +6. Activation scales (`input_scale`, `input_amax`, `k_scale`, `v_scale`) are calibration values. Pass them through unchanged; they cannot be derived from weights. Read runtime scales from the actual runner. +7. High-precision training + low-precision rollout is not QAT. Neither fake-QAT nor real low-precision GEMM implies a low-precision optimizer or low-precision gradients. +8. Weight QAT lowers Nemotron-H probs_diff mean from about 0.025 to about 0.010; it reaches 0 only when the training forward uses the same kernels as inference. Whether activation fake-quant helps depends on upstream mismatch; measure first. +9. STE backward has its own acceptance: an independent reference, the same dY, a tolerance agreed in advance. A forward 0 diff does not prove the gradient is right. +10. Check that deployed bytes = `requantize(master)`. +11. Start the master from the real BF16 release, not a dequantized checkpoint, whose values sit on grid points so NVFP4 deployed values only change after thousands of steps at small lr. θ0 bytes then differ from the official checkpoint (Nemotron-H: NVFP4 99.776% identical, FP8 24/46 tensors differ); quantify the cost with proxy SFT and have a human confirm it ([QAT §6.2](#62-fix)). + +## Terms + +- **NVFP4**: FP4 E2M1 values, one FP8 E4M3 block scale per 16 elements, plus one FP32 global scale. +- **block scale / global scale**: per-block scale (`weight_scale`, `scale`) / per-tensor or per-fusion-group scale (`weight_scale_2`, `weight_global_scale`, `scale2`). +- **amax**: maximum absolute value of a tensor (or block), used to compute scale. +- **4over6 / MSE**: TE NVFP4 block scale selection settings; the Nemotron-H checkpoint uses them. +- **E4M3 / E2M1 / UE8M0**: FP8 (4 exponent, 3 mantissa bits) / FP4 (2 exponent, 1 mantissa bit) / unsigned 8-bit exponent (power-of-2 scale). +- **W4A16 / W4A4 / W8A8**: weight bit width / activation bit width. +- **QDQ**: quantize then immediately dequantize, giving a high-precision value on the grid. +- **fake-quant / fake-QAT**: the training forward uses the QDQ value while GEMM stays high precision; not low-precision training. +- **STE**: straight-through estimator; treats the derivative of rounding as 1 (or with a clipping mask). +- **master**: high-precision parameters managed by the optimizer. +- **conversion owner**: the single party that converts a tensor family to the inference format. +- **fusion group**: tensors fused in serving (QKV, gate/up) that must share scale. +- **R3**: MoE route replay; training uses the expert selection recorded by rollout. +- **BI=1 / BI=0**: the `VLLM_BATCH_INVARIANT` switch; 1 is the true on-policy arm where training calls the inference kernels, 0 is the regular baseline. +- **probs_diff**: verl's `training/rollout_probs_diff_*`, the token probability difference between rollout and the training forward. +- **TIS / `rollout_is`**: rollout importance sampling correction. +- **θ0**: policy weights before the first update. +- **oracle**: the reference implementation or value that decides pass or fail. +- **VJP**: the backward function that computes input and parameter gradients from the upstream gradient. + +## Conversion chain: from master weights to the bytes the vLLM kernel uses + +### 1. Four stages + +```text +A. master weights FP32/BF16 optimizer master (plus the FP32 exceptions the model declares) + ↓ forward materialization +B. weights the training forward actually sees + BF16; under QAT it is DQ(Q(W)), the dequantized value on the deployment grid + ↓ export (gather → HF mapping → quantize/pack) +C. inference-format bytes + packed NVFP4 + block scale + global scale / FP8 + scale / + INT4 packed + scale / layers still BF16 in mixed precision + ↓ loader → post-load +D. tensors loaded by vLLM and used by the kernel (including post-load reordering and derived scales) +``` + +| Stage | Invariant | How to verify | +| --- | --- | --- | +| A | Record dtype separately for parameters and buffers; an FP32 master does not mean forward uses FP32 | Read optimizer parameter identity and dtype; see [§6](#6-dtype-contract-six-boundaries) | +| B | Under QAT, `Q` is the same function on the same grid as the export in C | Unit test: weights of B = dequantize(bytes of C), bitwise | +| C | Scale computed on the full logical tensor; fusion group shares the global scale; excluded layers keep their dtype; recomputing on the same master is idempotent | Compare packed bytes, block scale, global scale per tensor; see [§2](#2-inventory-the-checkpoints-full-quantization-rules-first) | +| D | Each family quantized once; post-load only reorders, does not recompute scale; calibration values arrive unchanged | Per tensor, compare exported bytes with the bytes vLLM uses after loading | + +Record the real format and scale of weights, activations and KV (packed UE8M0, FP32 scale, MXFP4/FP8 scope); activation and packing rounding may differ between SM90 and SM100. Before measuring quantization error, confirm the scale comes from the current weights, not an old tensor. + +### 2. Inventory the checkpoint's full quantization rules first + +1. From `hf_quant_config.json` and the `quantization_config` in `config.json`, list each layer's format, which layers carry `input_scale`, and whether the KV cache carries k/v scale. Read layer by layer; do not infer from the model name ([P12](#p12-inferring-a-layers-precision-from-the-models-overall-name)). Nemotron-H Lightning 30B-A3B: `MIXED_PRECISION`; 46 FP8 layers (`mixer.in_proj`/`out_proj` of 23 Mamba layers); 5935 `W4A16_NVFP4` layers (routed/shared experts, lm_head); FP8 KV cache; 46 fp32 scalar `input_scale`, FP8 layers only. +2. Encode the BF16 release with the checkpoint's own scales; 100% identical packed bytes confirm the release matches the checkpoint. +3. From the **real BF16 release**, requantize with the candidate rule and report per tensor the match ratio of packed bytes, block scale and global scale. Dequantized weights already sit on the grid and pass under any rule ([P2](#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). Measured match rates: [P1](#p1-scale-rule-differs-from-the-checkpoint-requantization-flips-codes), [P3](#p3-fp8-per-tensor-scale-convention-is-not-amax448). + +Decisions: + +- Call the rules of a mainstream implementation (TE, ModelOpt, NeMo-RL) directly; do not invent a scale policy. Deviate, and report it, only when the mainstream implementation cannot do it. Nemotron's self-invented "enlarge scale only on overflow" was replaced by an existing TE rule. +- θ0 deployed bytes = the real BF16 release requantized by the rule, which differs from the official checkpoint: on Nemotron-H, NVFP4 bytes are 99.776% identical and FP8 24/46 tensors differ; the proxy SFT cost is about 0.03 nat ([QAT §6.2](#62-fix)). This policy needs human confirmation. +- The BI=1 training forward, BI=0 QAT and the exporter call the same `requantize`; θ0 bytes of the two arms are identical. `input_scale`, `k_scale`, `v_scale` are calibration values passed through from the checkpoint. +- When a single FP8 checkpoint is the source of truth, the training base is a deterministic dequantized derivative, and requantizing need not restore the checkpoint byte for byte. Measure separately: (1) with zero updates, re-exported bytes equal θ0; (2) the precision difference from the original checkpoint. + +### 3. What each format outputs + +| Format | Artifacts | Must check | +| --- | --- | --- | +| block FP8 | FP8 weight + `weight_scale_inv`, scale per block | Commonly 128×128; FP32/UE8M0 branch, backend, exclusion scope | +| per-tensor FP8 | FP8 weight + `weight_scale` | Scale from the amax of the full weight, not of a shard; convention may not be amax/448 | +| MXFP8 | E4M3 weight + UE8M0 `weight_scale_inv` | 1×32 block direction, scale encoding, pad and high-precision exceptions | +| NVFP4 W4A16 (compressed-tensors, existing `fp4` sender) | `.weight_packed` + `.weight_scale` + `.weight_global_scale` | 1×16, two-level scale, shared outer scale for gate/up, activation contract | +| NVFP4 W4A4 (ModelOpt, Inkling) | packed `uint8` under the original parameter name + FP8 `scale` + FP32 `scale2` + BF16 `input_amax` + `original_shape` | Naming not interchangeable with compressed-tensors; needs a separate adapter | +| INT4 compressed-tensors | `weight_packed` + `weight_scale` + `weight_shape` | group_size, symmetric quantization, pack layout, list of tensors actually quantized | +| DeepSeek V4.1 attention FP8 | F8_E4M3 weight + F8_E8M0 `.scale`, 32×32 block | Not rowwise 32-element MXFP8 | + +`weight_scale_inv` is only a loader protocol name; interpret the value per implementation. + +### 4. Who converts: conversion owner + +A tensor family has one conversion owner. Once the sender has packed, the receiver must not quantize again. dtype reinterpret, layout staging and post-load repack are not quantization. + +| Path | Training | Export | Rollout load | owner | +| --- | --- | --- | --- | --- | +| FP8 rollout-only | BF16 | Full BF16 | `quant_weights()` converts to block-FP8 | rollout receiver | +| VERL FSDP QAT NVFP4 | BF16 + fake quant | After gathering the full tensor, `QATQuantizer` outputs packed NVFP4/scale | compressed-tensors | FSDP sender | +| VERL Megatron QAT NVFP4 | BF16 + ModelOpt fake quant | `QATWeightExporter` | ModelOpt | Megatron sender | +| Megatron FP8 E2E | TE FP8 | Bridge-specific FP8 export planning | vLLM FP8 | Bridge / rollout, per the pinned recipe | +| Nemotron-H BI=0 (verl Megatron) | BF16 + ModelOpt weight STE | Requantized by the checkpoint's rules into ModelOpt `MIXED_PRECISION` | ModelOpt | sender (`mixed_precision_export=True`) | +| Pre-quantized checkpoint passthrough | No requantization | Original tensor/scale | Per the checkpoint layout | checkpoint producer | + +- Full merge: convert logically complete HF tensors one at a time by their owner; no whole-model all-gather or per-step disk write. +- QAT: the training forward simulates quantization; export generates the actual packed weights and scales. + +#### 4.1 FSDP path (VERL) + +- With `actor_rollout_ref.actor.fsdp_config.qat.enable=true`, quantizable `nn.Linear` modules become `QATLinear` after the PEFT wrapper and before the FSDP wrap. Forward does NVFP4 W4A16/W4A4 fake quant and ends in `F.linear`; after the optimizer step, `invalidate_all_scales` clears the weight scale cache. [Implementation][qat-linear] +- Export: `FSDPEngine.get_per_tensor_param()` calls `full_tensor()` on each `DTensor`, then `QATQuantizer.quantize_with_fusion()`: + +```text +FSDP local shards → full_tensor() all-gather → full HF weight +→ full-weight amax / global_scale → QKV or gate/up fused_global_scale +→ per-row, per-group block scale → packed NVFP4 weight + scales +``` + + The cost is the all-gather and GPU staging. The quantizer caches to CPU per decoder layer and moves back to GPU, so peak GPU memory is not necessarily one tensor. [Export][qat-export] + +- `get_per_tensor_param_shard()` outputs only the local BF16 shard, with no quantizer and no scale. Shard-local block quantization is an unsupported new protocol and must at least solve: + 1. aligning shard and block boundaries, or handling cross-shard blocks; + 2. a collective for the full-weight amax/global scale; + 3. a consistent fused global scale for QKV and gate/up; + 4. W4A4 explicitly carrying the activation-derived `input_global_scale`, which cannot be rebuilt from weight shards. + + Until these contracts and the oracle exist, quantizing each shard and concatenating is forbidden. +- Non-QAT: export BF16 in HF coordinates; BF16 rollout does not quantize, and `rollout.quantization=fp8` is quantized by the receiver. LoRA merge exception: [P6](#p6-lora-mergetrue-bypasses-the-fsdp-qat-export). + +#### 4.2 Megatron path + +- VERL: ModelOpt installs fake-quant modules; `MegatronEngine.get_per_tensor_param()` exports HF weights through Megatron-Bridge, and `QATWeightExporter` reads the module quantizer's amax/block metadata to produce NVFP4 and scales. [ModelOpt exporter][modelopt-source] +- VERL FP8 E2E: TE handles FP8 forward/backward in training; when the Bridge declares `export_weight_dtype=fp8`, VERL does not cache generic conversion tasks. The rollout format is decided by `rollout.quantization` and the serving checkpoint contract; FP8 training does not mean exported tensors are in vLLM layout. +- Megatron capability cannot be inferred as FSDP capability. +- Distinct artifacts: the offline converter prepares the initial checkpoint (may write to disk); the online exporter streams in memory; post-load reorders without recomputing scale; a training checkpoint is not a servable quantized HF checkpoint. + +#### 4.3 Receiver quantization (rollout-only) + +`rollout.quantization=fp8`: the trainer exports BF16; `load_quanted_weights()` in the vLLM worker calls `quant_weights()`, which uses `scaled_fp8_blockwise()` to produce the FP8 weight and `weight_scale_inv`, then replays the quant method's post-load. There are also a DeepSeek V4-specific entry and an Ascend MXFP8 branch with different scale naming. [receiver][receiver-source] When the checkpoint is already FP8/MXFP4, confirm whether it is requantization or passthrough. vLLM cannot quantize BF16 online into ModelOpt `MIXED_PRECISION`. + +### 5. Minimal adaptation for a new quantization format + +1. From the checkpoint config and the actual loader, list key, shape, dtype, block size, packing and fusion group of weight/scale, and the unquantized tensor families. +2. Trace export -> loader -> post-load and mark the single owner of each family. +3. If a converter is missing, add the minimal conversion at the existing owner's entry; verify on one tensor / fusion group, then cold-load a small model. Changing the dtype is not a conversion. +4. Test zero merge and a non-zero update; get quantize-after-full-gather right before shard optimizations. + +One evidence table is enough: tensor family -> source -> export format -> serving format -> owner -> unit/small-model verification. + +#### Minimal acceptance + +- [ ] Record training dtype, export dtype/layout, rollout dtype/layout and owner. +- [ ] Fixed-token logprob of the zero-adapter path and the base are within the noise baseline of the same backend. +- [ ] After a non-zero update, tokenwise logprob/KL between trainer and rollout stays within bounds. +- [ ] name/shape/dtype/count/bytes of packed weight, scale and ignored families all close. +- [ ] Deployed bytes = `requantize(master)`. +- [ ] FSDP QAT proves scale is computed after `full_tensor()`; with shard export, prove it is not quantized, or provide the four shard-local contracts of §4.1 and a cross-rank oracle. +- [ ] Each family is quantized exactly once. + +### 6. dtype contract: six boundaries + +For each model, record `checkpoint → load/bridge → wrapper/sharding → forward → backward/optimizer → export/serving` separately. + +| Boundary | Megatron | VERL/FSDP | Must check | +| --- | --- | --- | --- | +| Source weights and conversion | Bridge may reorder, split, cast; model may override default conversion | HF loader, dequantized derived checkpoint, init dtype | Same revision; logical value/scale/layout match, not only target dtype | +| Parameter storage | `params_dtype` plus model exceptions; FP32 master/optimizer | HF keep-precision list, PEFT master, FSDP shard | An FP32 master does not mean forward uses FP32 | +| forward parameters | Float16Module converts selectively | FSDP2 `MixedPrecisionPolicy.param_dtype` | Measure inside forward, not only before the wrap | +| Operator computation and output | TE/FLA/TileLang/fused each define product, accumulation and output precision | HF/PEFT kernel and autocast | FP32 output does not prove FP32 product; BF16 output does not prove BF16 accumulation | +| Gradient and optimizer | TP/EP/SP/CP collective, DDP | FSDP reduce dtype, gradient accumulation | FP32 reduction does not restore values rounded in forward | +| Export and serving | Undo the parallel layout, convert adapter/quantization format | gather/export HF adapter | Compare values on the same logical layout | + +FP32 protection takes several forms: mark-to-skip before conversion, reset after declaration (which does not restore values already lost), bypassing the default BF16 downcast, and operator-level selection. It is not "all special parameters are raised in precision". + +FSDP: HF `_keep_in_fp32_modules_strict` is a load contract and does not become a per-parameter cast exception in FSDP ([P16](#p16-fsdp-bf16-policy-overrides-the-fp32-parameter-contract)). Respect the model's existing contract first, and verify any change of computation precision against a same-model reference and actual serving. `param_dtype=None`, `param_dtype=fp32` and FP32 autocast are different things. + +### 7. Code entry points + +FSDP [engine][fsdp-source], [module selection][qat-core], [fake quant][qat-linear], [NVFP4 export][qat-export]; Megatron [ModelOpt exporter][modelopt-source]; rollout [receiver][receiver-source]. Historical PR pointers (not the current state): early return on LoRA merge [#5115](https://github.com/verl-project/verl/pull/5115)/[#5411](https://github.com/verl-project/verl/pull/5411)/[#7014](https://github.com/verl-project/verl/pull/7014); fake quant and export quantization [#7149](https://github.com/verl-project/verl/pull/7149)/[#7224](https://github.com/verl-project/verl/pull/7224)/[#7577](https://github.com/verl-project/verl/pull/7577). + +[fsdp-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/workers/engine/fsdp/transformer_impl.py +[qat-core]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/core.py +[qat-linear]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/linear.py +[qat-export]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/quantizer.py +[modelopt-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/modelopt/qat_weight_exporter.py +[receiver-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/vllm/vllm_quant_utils.py + +## QAT, low-precision training and STE + +### 1. Choose the precision route first + +Rule out input, weight, adapter, layout and operator correctness errors first, then choose precision: + +1. When GPU memory allows and rollout can use BF16, use BF16 for both training and inference; check FP32 exception parameters one by one. +2. When rollout must be low precision (such as NVFP4), use weight QAT or a BF16 actor for training. On Nemotron-H the weight QAT arm is stable. Whether `rollout_is=token` stabilizes the BF16 actor is **not tested**. +3. Consider real low-precision training last: backward/kernel support and maintenance cost more. + +Same dtype is not same computation, and QAT is not bit-exact. QAT and low-precision training shrink the **weight value** gap; the kernel and reduction-order gap goes only when the training forward calls the inference kernels (the BI=1 arm of the [BI switch recipe](recipe-and-acceptance.md)). A natively low-precision model does not necessarily run with `dtype=bf16`; do a capacity pre-check first. + +### 2. What the four kinds of path are + +| Path | Training forward | backward / optimizer | Is it QAT | +| --- | --- | --- | --- | +| High-precision training + low-precision rollout | BF16 weights | High precision | No; quantized only at export | +| fake-QAT | `DQ(Q(W))`, GEMM is still BF16 | STE sends the gradient back to the high-precision master | Yes | +| Real low-precision GEMM (FP8/MXFP8/NVFP4) | TE changes the compute kernel | High-precision master/exception tensors remain | Different from fake-QAT | +| Frozen quantized base + high-precision LoRA | Dequantized base + adapter | Only the adapter is trained | No | + +- Neither fake-QAT nor real low-precision training makes the optimizer/gradients low precision. W4A16 quantizes only weights; W4A4 also quantizes inputs. +- Real low precision: NVFP4 forward is not FP4 backward; configurations include high-precision backward and backward on the dequantized operand. +- Only weight-only formats can compute scale from the full weight; formats with activation quantization also need activation scale/calibration state, which weights cannot provide. QAT simulation and export follow the same format rules, but scales are not necessarily constant during training. +- When LoRA is applied separately, do not fake-quant the merged base+adapter yourself. +- Open-source frameworks (verl [#6992](https://github.com/verl-project/verl/pull/6992), [#7971](https://github.com/verl-project/verl/pull/7971), NeMo-RL) essentially only do W4A16; activation inconsistency only appears at W4A8. + +### 3. Implementing fake-QAT + +Forward form (Nemotron-H): + +```text +w_visible = w + (dequant(requantize(w)) - w).detach() +``` + +- `requantize` is the same function as the exporter; replace the library's built-in quantizer when its rules differ. +- Requantization runs on every forward, and again under full recompute; Nemotron-H step time is about +60%. +- A temporary `QuantModuleRegistry` replacement and the old class-attribute patch give the same probs_diff mean on the same quick workload (0.0095/0.0091/0.0100 vs 0.0096/0.0092/0.0101); module hit check: [P7](#p7-qat-switch-is-on-target-modules-were-not-replaced). + +### 4. How much mismatch QAT can remove + +Nemotron-H (`rollout_probs_diff` mean, NVFP4 mixed-precision deployment, single node 4×GB200): + +| Approach | Training forward | probs_diff mean | +| --- | --- | --- | +| BF16 training + quantized rollout | BF16 weights | About 0.025; no rollout correction | +| Weight QAT (STE) | Dequantized deployed weights, BF16 GEMM | About 0.010 (pearson 0.986→0.997); three 3-step smoke runs of one recipe, 0.0091-0.0102 | +| Weight QAT + activation fake-quant | Same, plus static FP8 activations | About 0.0107, about 10% worse than weights only | +| true on-policy (BI=1) | Deployed bytes + the same kernels as vLLM | max = mean = 0 at every step for 100 steps | + +The QAT arm's probs_diff max is still 0.73-1.0; a small mean does not rule out large per-token deviations. Other models (different metrics, not thresholds): DeepSeek-V4 BF16 train + FP8 rollout has a first-step logprob drift of about 0.023 and still reports rising reward/eval ([DeepSeek][dsblog]); Kimi K3 frozen quantized base + LoRA has a train/rollout KL noise floor of about 0.002, a Schulman KL estimate, not a mean absolute logprob diff ([K3][k3blog]). + +#### 4.1 Activation fake-quant: aligned per layer, worse end to end + +Adding vLLM's static FP8 activation fake-quant to training aligns single-layer output to the BF16 noise floor, but end-to-end probs_diff gets about 10% worse: FP8 rounding is discontinuous, so any element-wise upstream mismatch flips codes and amplifies the difference. To judge, add relative noise δ only to the training-side input and sweep the output error; the crossover was about δ≈1%, and activation quantization hurts when upstream mismatch exceeds it. This is one measurement on one model, not a threshold; for a new model, measure upstream mismatch and repeat the sweep. Nemotron-H did not merge it. + +Rule: the weight-side range/axis/scale/rounding must match serving ([rule 3](#key-rules)); matching the activation side depends on upstream mismatch and is not unconditional. Mainstream frameworks (verl, NeMo-RL, slime) use the same **dynamic** activation rule on both sides and do not reuse an external checkpoint's static activation scales; the NeMo-RL Nemotron recipe quantizes only routed experts and leaves Mamba projections, attention and lm_head in BF16. On DeepSeek V4.1, KV quantization simulation adds only about 1.93% improvement on top of R3. Training-side simulation of the FP8 KV cache is not covered. + +#### 4.3 When QAT is worth it + +Evaluate unified FP8/QAT only when evidence shows the train-inference quantization difference is the main limit; "no score gain" is not a reason to switch precision first. For DSv4.1 rank2 LoRA, the post-merge quantized delta is only 5.1-7.7% of the pre-quantization delta (a norm ratio, not information retention), which only supports treating step size/quantization as controlled variables. The long-running Qwen Flash-Next experiment is same-source BF16/BF16 and does not validate the FP8 rollout route. + +### 5. STE contract + +The STE surrogate (identity or clipped), whether dX uses the master or dequantized value, scale handling and verification are in [Backward validation](backward-validation.md#write-down-the-gradient-contract-first). By default dX uses the dequantized value; when using the master value, quantify the bias. + +### 6. Why QAT does not train + +Train-inference consistency is the strictest QAT: forward uses quantized weights, the high-precision master accumulates updates, STE passes the gradient back. Whether it trains depends only on the master's distance to the nearest decision boundary. In W4A16 (the FP4 code is decoded to BF16 in a register and goes through the BF16 tensor core; the gain is 1/4 the weight bandwidth), W4A4 and W8A8 alike, the computed values are "code × scale". + +#### 6.1 Dead zone + +- Each NVFP4 element has only 8 values, 0/0.5/1/1.5/2/3/4/6 × scale, with RNE rounding. From a dequantized checkpoint the master sits on grid points, half a spacing from every decision boundary, and must travel that half spacing before the code changes. +- At lr 1e-6 one step moves the master about 0.01%; a few hundred steps move it far less than half a spacing, so deployed values stay byte-for-byte unchanged. RL gradients are more random than SFT, closer to a random walk, and need more steps. +- Result: NVFP4-covered parameters (on Nemotron-H, routed/shared experts and lm_head, the bulk) have zero deployed change, and validation NLL is clearly worse than the native BF16 baseline; freezing these weights explains about half the gap. +- Recomputing block scale does not flip codes earlier: from grid points requant is idempotent, the E4M3 scale only changes code after block amax moves 3-6%, and a tiny global scale change only rescales the whole tensor. +- Mainstream QAT (NeMo-RL QARL, slime/Kimi INT4, Jet-RL) starts from the original BF16 weights, which lie between grid points, so any small drift flips a proportional share of codes and there is no dead zone. INT4 with 1×32 groups from a dequantized checkpoint stalls the same way. +- FP8 is barely affected: floating-point relative spacing is 6-12%, small elements have small absolute spacing, and some elements are always near a boundary. DS4 also starts from dequantized FP8 and is affected only mildly. + +#### 6.2 Fix + +Preferred: use the original BF16 weights (the real BF16 release) as the master and requantize every step by the same rule (Nemotron-H: TE 4over6, E4M3 upper bound 256). State two points: + +- **θ0 is not equal to the official checkpoint bytes.** Requantizing from the real BF16 release, 99.776% of Nemotron-H's NVFP4 bytes are identical and 24 of 46 FP8 tensors (ModelOpt calibrated scale) differ. "The official checkpoint was generated by this rule" holds only on dequantized weights, where MSE and MAE cannot tell rules apart on grid points ([P2](#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). Proxy SFT gives the cost: 300-step val NLL A-bf16 3.4685 vs C-dequant 3.4641 / C-bf16 3.4373, about 0.03 nat, present at θ0 and not growing during training. This policy needs human confirmation. +- **Three acceptance items** (all pass on Nemotron-H): (1) with zero updates, re-exported bytes equal θ0 (not the official checkpoint); (2) fast-tier RL θ0/θ1/θ2 is 0 diff and deployed bytes start changing; (3) proxy SFT gap to the native BF16 baseline is within budget. + +Alternative, keeping θ0 bytes: add a fixed small jitter to move the master off grid points (toward the neighboring cell, at most half a step; after requant check per block that bytes are unchanged, and roll back changed blocks). Not recommended: stochastic-rounding requantization (about 1.4× deployment error, needs a self-written quantizer, no mainstream precedent), raising lr (under Adam equivalent to QuRL UAQ, only shortens the dead zone, 10× lr is outside the usual RL range), and frozen base + LoRA (QeRL, gives up full-parameter updates). Whichever you choose, state that from a dequantized checkpoint at lr 1e-6, NVFP4 deployed values do not change for roughly the first 10³-10⁴ steps. See [P24](#p24-master-starts-on-the-quantization-grid). + +[dsblog]: https://www.lmsys.org/blog/2026-04-25-deepseek-v4/ +[k3blog]: https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support/ + +## Known pitfalls + +Ordered along the conversion chain: quantization rules -> export and receive -> training-side numerics -> experiment measurement basis. P24 sits under training-side numerics by content. + +### Quantization rules and scale + +#### P1 scale rule differs from the checkpoint, requantization flips codes + +- Symptom: weights unchanged, rollout shifts after a requantized export. +- Cause: recomputing scale on the dequantized master with the ModelOpt default global = amax/(6·448) differs from the checkpoint's TE 4over6 rule and flips about 40% of Nemotron-H's FP4 codes; recomputing with TE 4over6 is idempotent and byte-for-byte unchanged. +- Fix: use a mainstream rule that matches the checkpoint; training QAT and the exporter call the same `requantize`. +- How to verify: requantizing the same master twice is byte-for-byte identical; requantize from the real BF16 release and compare per tensor with the checkpoint. Nemotron-H: + + | Rule | Match with the checkpoint | + | --- | --- | + | TE NVFP4 4over6 + MSE, E4M3 upper bound 256, global = amax/1536 | packed 99.776%, global 100%; no setting closes the remaining gap | + | ModelOpt default global = amax/(6·448) | packed about 50% | + | Standard FP8 per-tensor amax/448 | See P3 | + +- Source: procedure in [Conversion chain §2](#2-inventory-the-checkpoints-full-quantization-rules-first). + +#### P2 Verifying a rule with dequantized weights gives a false positive + +- Symptom: several candidate rules all report a 100% match. +- Cause: dequantized weights already sit on the grid and do not change under any rule. +- Fix: verify from the real BF16 release. +- How to verify: at least one known-wrong rule fails the test. + +#### P3 FP8 per-tensor scale convention is not amax/448 + +- Symptom: FP8 layers differ from the checkpoint after requantization. +- Cause: Nemotron-H's official Mamba projection scale is about 1.5×amax/448; 24 of 46 FP8 tensors differ from amax/448. +- Fix: derive the convention from the checkpoint, or pass the checkpoint's scale through. +- How to verify: report the scale match rate per tensor. + +#### P4 Computing scale separately on each shard + +- Symptom: packed weights concatenated after an FSDP/TP sharded export differ from quantizing the full tensor. +- Cause: blocks cross shards; global scale and fused global scale need the full tensor. +- Fix: `full_tensor()`/gather, then quantize; shard-local quantization must first meet the four contracts of [Conversion chain §4.1](#41-fsdp-path-verl). +- How to verify: evidence shows scale is computed after the gather. + +#### P5 gate/up not quantized as a pair + +- Symptom: wrong values after serving fuses gate/up, or the exporter reports a missing pair. +- Cause: the two HF tensors are fused in serving and need a shared outer scale. +- Fix: quantize as a pair; raise an error when a complete pair is missing, and test pairing across safetensors shards. +- How to verify: test pair completeness, cross-shard gate/up and a TE bitwise reference on matching hardware; a test existing does not mean it ran and passed. + +#### P6 LoRA merge=true bypasses the FSDP QAT export + +- Symptom: QAT and LoRA merge are on, and the export is not the QAT quantizer's packed weights. +- Cause: VERL FSDP returns `_merged_lora_per_tensor_param()` directly for `merge=true`, skipping `QATQuantizer`; with `rollout.quantization=fp8` the receiver may still requantize. The `_qat_enabled` check still follows the `merge=false` branch. +- Fix: the first merge fallback does gather -> merge -> quantize -> yield -> release one full tensor at a time. +- How to verify: check actual module replacement, naming and export scope; both switches being on does not make the combination usable. +- Source: [VERL engine](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/workers/engine/fsdp/transformer_impl.py). + +#### P7 QAT switch is on, target modules were not replaced + +- Symptom: the QAT arm's probs_diff equals the BF16 arm's. +- Cause: VERL `QATLinear` enters through `nn.Linear` and does not cover three-dimensional expert Parameters or arbitrary PEFT wrappers; the framework's built-in QAT changes only the GPT layer spec, while Mamba hybrid models read the hybrid stack spec. +- Fix: attach the quantizer after building the model by the actual spec; use the exporter's `requantize` as the quantizer. +- How to verify: every rank prints the count and list of replaced modules and fails closed (139 for Nemotron-H, the same on all four ranks). +- Source: [Conversion chain §4.1](#41-fsdp-path-verl), [BI=0 integration pitfalls](recipe-and-acceptance.md#bi0-arm-integration-pitfalls). + +#### P8 Training grid differs from the deployment grid + +- Symptom: QAT is declared, the mismatch remains. +- Cause: the converter and the trainer use different group sizes (for example 32 vs 128). An external TPU case is bitwise 0 on the FP8 per-channel grid, but the native weights are 128×128 block, and the author notes the result may not hold on a different grid. +- Fix: align quantized module scope, block/group size, axis, scale dtype/rounding, global scale scope, gate/up fusion, activation scale, TP/ETP partition and high-precision exceptions. FP4 has 1 mantissa bit, so error comes mainly from block partitioning; block size and direction must match the native weights. MiMo-V2.6 does QDQ on the experts after every update under the rollout MXFP4 GEMM's constraints. +- How to verify: training-side QDQ output = dequantize(deployed bytes), bitwise. +- Source: unified NVFP4 constraints in the [technical blog](https://www.lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl/) (its bit-exact quantizer test verifies the format, not whole-model bit-exactness). + +#### P9 Activation scale derived from weights + +- Symptom: W4A4 export raises, or the exported activation scale differs from serving. +- Cause: activation scale depends on actual input. VERL W4A4 export raises when an initialized `input_global_scale` is missing; Inkling's `input_amax` cannot be derived from BF16 weights. +- Fix: reuse unchanged from the QAT calibration state or rollout checkpoint; pass through `input_scale`, `k_scale`, `v_scale`. +- How to verify: manifest activation scales are byte-for-byte identical to the checkpoint. +- Source: [VERL quantizer](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/quantizer.py). + +### Export and receive + +#### P10 Both sender and receiver quantize + +- Symptom: deviation larger than one quantization; or a packed tensor treated as BF16 and quantized again. +- Cause: a tensor family has two conversion owners; the Bridge already outputs packed tensors but the generic packer does not recognize them. +- Fix: one owner per family; pass packed tensors through. +- How to verify: the evidence table lists the owner per family; count/bytes close. +- Source: [Conversion chain §4](#4-who-converts-conversion-owner). + +#### P11 Receiver picks quantization targets by name suffix and misses fused MoE + +- Symptom: under `rollout.quantization=fp8` some experts are still BF16. +- Cause: `gate_up_proj`/`down_proj` do not end in `.weight` and are skipped; fixed by VERL [#7443](https://github.com/verl-project/verl/pull/7443). +- Fix: check coverage using checkpoint names together with runtime modules/dtype; review the list for each new model. +- How to verify: list the serving dtype of every family. +- Source: [receiver](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/vllm/vllm_quant_utils.py). + +#### P12 Inferring a layer's precision from the model's overall name + +- Symptom: diagnosis goes in the wrong direction. +- Cause: in the Nemotron-H "NVFP4" checkpoint, `backbone.layers.16.mixer.out_proj` has `quant_algo` FP8. +- Fix: read `hf_quant_config.json` layer by layer. +- How to verify: the diagnosis states the actual `quant_algo` of the first differing layer. + +### Training-side numerics + +#### P13 scale taken from the wrong object + +- Symptom: the unit-scale unit test passes; real attention output and old-logprob fail. +- Cause: the training adapter quantizes Q with `k_scale`. On the real runner `q_scale=1.0`, `k_scale=0.03173828125`, `v_scale=0.003326416015625`. +- Fix: read scale from the runner/module actually executed; do not infer a default from a missing checkpoint field or another tensor. +- How to verify: replay with real non-unit scales and frozen Q/K/V; after the Nemotron-H fix, 5,832,704/5,832,704 BF16 outputs are bitwise identical. + +#### P14 Quantization modules with the same name take different dispatch + +- Symptom: calling `ops.scaled_fp8_quant` + `flashinfer_scaled_fp8_mm` directly matches training bitwise, yet the inference linear module differs in 1,361 BF16 elements with the same scale. +- Cause: the module selects the **compiled** `forward_native`, which has 44 more differing FP8 elements than uncompiled native and the CUDA custom op; those two agree. Which compile transformation is responsible is not confirmed. +- Fix: under Nemotron-H BI=1, static per-tensor `QuantFP8` switches to the device custom op (GEMM and format unchanged). +- How to verify: with the same operator input, weights and scale, compare FP8 bytes, GEMM output and module output in turn, calling the native/CUDA/actual module paths separately. When the direct GEMM matches training and the module does not, check quantization dispatch inside the module first; do not change GEMM. + +#### P15 Quantizer is not in the `nn.Module` tree, the patch misses layers + +- Symptom: the patch "ran", the target layer still has 1,361 differing elements. +- Cause: traversing `model.modules()` finds only 6 registered modules and misses quantizers on plain objects such as `linear.quant_method.kernel.quant_fp8`. +- Fix: traverse the registered tree, the quant-method kernel objects and that version's legacy fields together. +- How to verify: record per worker the hit count (52 after the fix), target layer names, actual `_forward_method` and quantized output bytes; fail if expected targets are zero or a key layer is missing. Private field paths are only a versioned case. + +#### P16 FSDP BF16 policy overrides the FP32 parameter contract + +- Symptom: Hy4 four-layer fixed-token logprob MAE 0.022982/0.023052. +- Cause: HC and attention sink, FP32 after HF load, are downcast by the uniform BF16 FSDP policy; a later `.float()` in the operator cannot restore them. The router's FP32 conversion is also overridden by the outer BF16 autocast, which changes the expert set. +- Fix: opt in with `wrap_policy.preserve_original_param_dtype=True` (`param_dtype=None`, BF16 autocast kept). This is a whole-model policy, not a per-parameter whitelist; measure communication/memory impact; it does not raise a BF16 base to FP32. Disable autocast inside the router forward. +- How to verify: with original dtype kept, singleton, two ranks and unwrapped are exactly identical; the router has a regression test. + +#### P17 fake-quant cuts off the gradient + +- Symptom: in DSv4.1 LoRA training, 31 of 160 B matrices are zero, all in `self_attn.kv_proj`. +- Cause: `_fake_quant_fp8_block` goes through an FP8 cast on the autograd path, so small gradients underflow to 0; `_fake_quant_fp4_block` has no STE (integer comparison and LUT gather, so gradient flows only through scale, and the ue8m0 branch cuts even that). +- Fix: keep the quantized forward and use identity for the input gradient: `output.detach() + (x - x.detach())`. +- How to verify: a 1e-6 upstream gradient is preserved exactly; after the fix 160/160 B matrices are non-zero. The loss still rises, so this does not prove training improved. + +#### P18 Training forward and export do not quantize at the same place + +- Symptom: the model already has FP4/FP8 STE, the mismatch remains. +- Cause: DSv4.1's existing STE acts on indexer/compressed/window KV activations, not as weight fake-quant on `W+ΔW`; the training forward is unmerged base + LoRA, while export is an FP32 merge followed by block32 FP8/E8M0 conversion. +- Fix: confirm whether training sees the post-export rounding before deciding on weight QAT. +- How to verify: sample the actual dtype inside forward, not only the code path. + +#### P19 KV cache aligned only on the dtype label + +- Symptom: attention results differ with the same weights. +- Cause: in the old vLLM image, KV is FP8/block64 for the non-RoPE part and BF16 for the RoPE part, and the training side handles values differently. +- Fix: set `kv_cache_quantization="fp8_ds_mla"` explicitly in training, use the same rounding for forward with and without gradients, and keep the gradient with STE. Re-verify after the upstream KV layout changes. +- How to verify: compare attention output with the same weights. + +#### P20 Dividing by temperature in BF16 + +- Symptom: logprob has a fixed offset. +- Cause: the old VERL FSDP path divides by temperature in BF16/FP16, then raises precision. +- Fix: `logits.float()` first, then divide by temperature, even when temperature is 1. +- How to verify: both scorers match on the same logits. + +#### P21 Derived scale after resume still belongs to the old weights + +- Symptom: the first rollout after resume is wrong. +- Cause: BF16 parameters were updated, but the FP8 source-scale registry still trusts the initial weights. +- Fix: invalidate derived caches in the weight restore lifecycle. +- How to verify: first-step probs_diff after resume is on the same measurement basis as before (Nemotron-H kill->resume acceptance: still 0 after resume). +- Source: [checkpoint guide §13](checkpoint-resume.md#13-stale-quantization-scale-on-the-first-step-after-restore). + +#### P24 Master starts on the quantization grid + +- Symptom: QAT loss barely drops and NVFP4 deployed weights are byte-for-byte unchanged while the master moves. +- Cause: a master initialized from a dequantized checkpoint sits on grid points ([QAT §6.1](#61-dead-zone)). +- Fix: master from the real BF16 release, requantized every step by the same rule, cost quantified with proxy SFT; without the original weights, a fixed within-cell jitter. Not stochastic rounding, larger lr, or LoRA. +- How to verify: the three acceptance items of [QAT §6.2](#62-fix). + +### Experiment measurement basis + +#### P22 The quantization oracle itself is wrong + +- Symptom: the GPU block32 quantization gate fails (zero block expected 2.27373675e-13, actual 1.16415322e-10). +- Cause: the oracle assumes Triton's epsilon position; the CUDA UE8M0 path also clamps scale to 1e-10 before taking the power of 2. +- Fix: fix the oracle, not the kernel. CUDA and Triton are not thereby bitwise identical on tiny/zero blocks. +- How to verify: check the oracle's epsilon/floor position against the kernel source before judging a failure. + +#### P23 Attributing every logprob difference to quantization + +- Symptom: the difference persists after changing precision. +- Cause: routing selection, attention/kernel and tokenization/mask all produce differences. R3 constrains routing and does not make FP8 GEMM equal BF16; TIS is a statistical correction and does not fix wrong weights. +- Fix: rule out protocol errors first, then record the quantization/numerical noise floor. +- How to verify: after ruling out protocol items (routing, mask) one by one, measure the quantization noise floor on the same input. diff --git a/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md b/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md new file mode 100644 index 000000000000..00762fe8c4e9 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md @@ -0,0 +1,113 @@ +# Recipe, evaluation and paired acceptance + +Which recipe items must match the reference implementation mathematically (not by name), what puts two evaluations on the same measurement basis, and how to pair two arms so a conclusion can be drawn. The reference recipe for train-inference consistency is [verl#7688](https://github.com/verl-project/verl/pull/7688) (DS4 exact 0-diff true on-policy). Which reference a model aligned with, and the values it ended up using, are per-model facts and are not recorded here. + +## Find by question + +| Question | Where to go | +| --- | --- | +| Does a 5% gap between two arms count as a gap | [Paired acceptance](#paired-acceptance): noise floor, median over multiple pairs, confounders | +| The validation score only takes a few discrete values | [Evaluation contract](#evaluation-contract): under greedy the score is a multiple of 1/number of problems | +| How to design the BI=0 baseline arm | [How the two arms are defined](#how-the-two-arms-are-defined) | + +## Key rules + +1. Align the math, not the names; read the final CLI/resolved config, not the README or the model name. +2. Tabulate every deviation from the reference: item, reference value, our value, why, which conclusion it affects. +3. One recipe, one image, a switch selects the arm (following DS4 `EXACT_ALIGNMENT`, default BI=1; BI=0 goes through Bridge/mcore + R3 + QAT); deviations live only in the switch. +4. If any of the six evaluation items differs, it is not the same score: dataset version, prompt template, sampling parameters, number of repeats, scorer, length cap. +5. Rising training reward, probs_diff of 0, and a held-out score gain are three separate conclusions; none substitutes for another. +6. Before running, fix what is compared, how it is computed, the pass line and the noise floor; plant a known error to confirm the criterion can return FAIL. +7. One artifact answers one question: run and report performance, quality and consistency separately. +8. List confounders first, then compare; each is either aligned or quantified separately ([paired acceptance](#paired-acceptance)). +9. Keep failure records as they are and append new results; keep reruns and changed questions apart. + +## One recipe, one image: deliver true on-policy and the regular baseline with the BI switch + +Following DS4's EXACT_ALIGNMENT: one delivery image, one recipe, and `VLLM_BATCH_INVARIANT`, defined by the recipe, selects the whole arm. The arms may differ in training engine, precision and route replay; the switch is only the single entry point. The pitfalls below come from Nemotron-H Lightning NVFP4 (Mamba hybrid MoE); a new model must re-check them item by item, not copy the config. + +### How the two arms are defined + +BI=1 is true on-policy: the training forward calls the vLLM kernels and reads the deployed bytes, rollout takes the vLLM BI path, and the deployed bytes are sent down unchanged. BI=0 is the regular baseline: training takes the default MCore path (verl Megatron engine/Megatron-Bridge, without the alignment implementation of [C6](../workflow/README.md#gate-table)), and rollout uses vLLM's non-BI kernels with the configuration audited per [C1](../workflow/README.md#gate-table). For BI=0, decide the items below per model and record who approved each: + +| Decision item | When needed | Nemotron's choice (example) | +| --- | --- | --- | +| Training precision: bf16 or weight QAT (STE) | QAT exists only when deployment uses quantized weights | QAT by default, bf16 optional | +| MoE route replay (R3) | MoE only | Always on | +| Requantize from master at export | Quantized deployment only | Requantize by the checkpoint's rules at every export | +| Rollout importance sampling correction (TIS etc.) | Arms whose train-inference mismatch is not 0 | Not enabled | +| rollout MoE backend | When the backend picked by `auto` does not run | Explicit `flashinfer_cutedsl` (`auto` picks Marlin, which reports a PTX version error), matching the BI0 arm of the BI=1 bench | + +Do not build BI=0 as "the BI=1 actor with non-BI kernels swapped in". That requires changing the alignment implementation, and the result is neither the regular baseline nor true on-policy; a human rejected this design when it was tried. + +### Rules + +- **Same image, same recipe.** If the recipe lacks a switch, add it on the delivery branch first, then build one image shared by both arms; no separate image, recipe or source mount for a single experiment. Delete the old standalone baseline recipe. +- **Keep output and checkpoint paths separate.** Suffix only BI=0; BI=1 paths stay as before. With shared paths, `resume_mode=auto` resumes from the other arm's checkpoint. Renaming the run breaks an in-progress resume chain (it cannot find the first-segment checkpoint), so a running chain does not change images. +- **The BI=1 command is byte-for-byte unchanged.** The recipe must offer a mode that only prints the fully resolved training command without starting a job (Nemotron uses `DRY_RUN=1`) and a mode that only runs hydra compose. For each mode × optimizer, print the command from the recipe before and after the change and `diff` (or compare sha256): BI=1 must be identical; BI=0, against the original baseline recipe, may differ only in names, paths and newly added explicit switches. Every compose combination must return rc 0. +- **The header comment states only what was measured.** Reviews have caught untrue claims such as "the BI=0 actor uses vLLM's default kernels". Known rollout config differences between arms (such as R3 forcing prefix caching on) go into the comparison report. +- Delete untested switches (such as unverified mxfp8 and route replay variants) instead of leaving them for users to trip over. + +### BI=0 arm integration pitfalls + +| Problem | Consequence | Handling | +| --- | --- | --- | +| THD packing does not pass `total_tokens` into `PackedSeqParams` | Mamba gets no `seq_idx` and scans the whole pack as one sequence; logprob is silently corrupted | Pass it when CP=1; raise an error for unverified CP>1 | +| Actor needs the BF16 release, rollout needs the quantized checkpoint, both share `model.path` | Actor is built from quantized weights | Separate `hf_weights_path` | +| MTP does not turn off cleanly (the provider appends the pattern back) | Extra untrained layers | Clear both `mtp_num_layers` and the MTP pattern | +| Expert bias is updated during training by default | Routing drifts from inference | `moe_router_bias_update_rate=0` | +| Framework's built-in QAT changes only the GPT layer spec | Hybrid models read the hybrid stack spec, so QAT has no effect | Attach the quantizer after building the model from the TE spec, using the same `requantize` as the exporter | +| QAT library incompatible with the TE grouped linear signature | Startup error or wrong quantization | Quantize weights only; unit test that forward = deployed bytes and the gradient reaches master | +| param + optimizer offload, HF safetensors export | Host OOM on save | Chunked optimizer offload, save only dist ckpt, see [checkpoint and host memory](checkpoint-resume.md#memory-budget-host-and-gpu) | + +R3 checks: R3 in vLLM requires prefix caching and Model Runner V2; the capture is indexed by **absolute layer number** (non-MoE layers are zero columns), and the training side selects the column by global `layer_number`, which must also line up under PP; exclude the MTP router. Known pitfalls: omitting `routed_scaling_factor` causes a drift of about 0.28; a monolithic backend that bypasses top-k capture replays all zeros. + +### Smoke criteria + +- BI=1 quick: `rollout_probs_diff` max = mean = 0 at every step. +- BI=0 quick: mean has the same magnitude as the independent reference (Nemotron QAT about 0.010, BF16 about 0.025), max may approach 1, grad norm is finite. +- When a step's grad norm is exactly 0, check within-group rewards first: if the whole group has the same reward, advantage is 0 and it is not a numerical problem. +- After changing the BI=0 QAT implementation, compare probs_diff against the old implementation on the same quick workload, and print the number of replaced modules (identical on every rank). + +## Evaluation contract and paired acceptance + +One artifact answers one question. Bitwise logprob on fixed tokens proves only the scoring of the same token history; natural generation tasks (GSM8K, AIME) answer quality; fixed-length workloads (Rust Bench) answer throughput. Token sets and timing differ, so none substitutes for another. After the numerical path, deployed weights or backend change, every in-scope quality and speed check needs fresh evidence from the **current candidate**; out-of-scope items are marked not verified, and no acceptance tasks are added on our own. + +### Evaluation contract + +If any of the six items differs, it is not the same score: dataset version (AIME24 ≠ AIME25); prompt/chat template and thinking mode; sampling parameters (temperature/top-p/top-k, stop tokens; greedy n=1, sampled average and pass@8 differ; pass@1 at n=8 is not greedy accuracy); number of repeats; scorer (answer region detection, boxed extraction, tail truncation length, reward extractor); length cap (training and eval caps may differ; 4K vs 8K results are not a pure capability difference). + +- Under greedy, n repeats give the same result and the score is a multiple of 1/number of problems: with 30 problems one problem is 3.33 percentage points, and smaller differences are invisible. +- Rising training reward, probs_diff of 0, and a held-out score gain are three separate conclusions; dynamic sampling changes the sample composition of the training reward. Read the truncation ratio together with accuracy. +- A gain from relaxing the token budget is not a training gain. Charts come from raw metrics; do not fabricate a step 0. +- W&B: the run config behind a curve, the step basis (rollout or optimizer step) and the checkpoint version must be traceable; for cumulative metrics, take per-phase differences. +- Changing the scorer or migrating a reward handler: run both scorers offline on the same responses and inspect only the samples they judge differently; do not switch scorers for the sake of "alignment" and mark another model's valid output wrong over a format difference. + +### Paired acceptance + +**Freeze the criteria first, then run.** Fix the model/checkpoint content, image digest, source/overlay, hardware/nodes, runner/parallelism, weight and KV precision, the actual attention/dense GEMM/MoE/all-to-all backends, and the resolved Graph and cache config; fix what is compared, how it is computed, the pass line and the noise floor; plant a known error to confirm the criterion can return FAIL. Tune BI=0 into a reasonably strong baseline at the agreed precision and have a human audit it; do not weaken it to narrow the gap. When the candidate changes quantization dispatch, attention, MoE, exported weights or image, re-evaluate which old evidence is invalidated ([evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger)). + +**List confounders first, then compare.** Tabulate every difference between the arms; each is either aligned or quantified separately; a delta from several simultaneous changes is not attributed to one switch. Example: the two Nemotron-H arms differed at once in PC on/off, `gpu_memory_utilization`, training backend (Bridge vs mLite) and the R3 return path. + +**Quality pairing**: same problem IDs, template, seed, sampling parameters, max length and scorer, no retries; record per problem whether each side is correct; `both_correct`/`only_bi0`/`only_bi1`/`both_wrong` must sum to the problem count; report each side's score, the absolute percentage-point difference and the agreed threshold. Close scores do not mean identical per-problem results, and passing quality does not mean bitwise-equal logprob. + +**Fixed-length throughput**: the agreed workload is vLLM Rust Bench 8192 input/1024 output, ignore EOS, c=1/8/32/64; record actual input/output tokens, completed requests, error count, warmup, and per concurrency point the wall time, throughput and TTFT/TPOT. List any real rollout workload separately; do not mix fixed-length and natural-EOS results in one ratio. + +### Pairing and computation + +Before testing, a human freezes: number of repeats, pairing order, aggregation formula, the budget/pass rule for each target c, and how an invalid leg is identified and re-tested. The suggested statistic is the median per-round paired loss `median_r(1 - T_BI1[c,r] / T_BI0[c,r])`; if `1 - median_r(T_BI1[c,r]) / median_r(T_BI0[c,r])` is chosen instead, label it "ratio of median throughputs". The two are not equivalent and cannot be swapped after seeing results; without a reviewed statistical basis, report only measured values and do not declare the budget passed. Historical results keep their original formula and are not silently re-judged. + +1. For each leg, record on every worker the image/package hash, actual kernel/backend, resolved config, Graph state, cache hits, and request and token statistics; verify client and server exit normally, separately. +2. Same node, same resources, no interfering load, crossed order (AB/BA/AB), each side warmed up separately every round; do not pair each side's fastest run picked across node pools. +3. Check workload and exit receipts item by item first. A mismatched leg is `inconclusive`; do not compute BI loss from throughputs with different request volumes. +4. Report raw throughput and loss `1 - T_BI1/T_BI0` for every round, then compute the formal result per c with the reviewed formula; keep failed/invalid rounds and the reasons for re-testing. Averaging across c must not hide a required c over the limit; keep single-round and aggregate overruns apart. + +Optimization screening may test only the previously failing concurrency points, but it only decides whether to continue. Take the median over multiple pairs and report single pairs separately: if a candidate round sits right at the budget line, it cannot pass on the old baseline's historical median; re-pair both sides in crossed order on the same node, same image, same time window. When the candidate changes the model output path, the old candidate's quality and post-update alignment evidence do not carry over. Crossed pairing does not guarantee a strong denominator: if the same config and workload show clear fast/slow tiers, keep per-round raw numbers and check tuning receipts, CPU/GPU state, Graph, cache and actual operator selection; an optimization is not effective just because BI=0 landed in the slow tier. Mark unexplained fluctuation's causal attribution `inconclusive`; any formal target point over budget by the predefined rule is still `fail`, with no cherry-picking rounds or switching the aggregation formula afterwards. + +Profile only outside the formal profiler-free E2E measurement; GPU kernel durations may overlap and cannot be summed to stand in for E2E; a kernel name is not a module identity, see [Performance and profiling](performance.md). + +### Failure and recovery + +When the acceptance checker errors, first tell apart producer, numerics, scorer and parser. If the original worker/server receipts and score are sealed and the wrapper parsed them incorrectly after a schema change, keep the original artifacts and read them with a recovery checker that is **separately versioned with its hash recorded**; it must verify all workers, EngineCore, API server, package identity and workload, and keep the original failure status. It can correct only the control-plane conclusion; it cannot fill in missing values, unfinished requests or a failed service exit. + +Keep failure records as they are and append new results; keep reruns and changed questions apart. The final report lists at least, separately: train-inference forward, module backward, real update, quality, performance, clean exit, final immutable artifacts; each cell is `pass`/`fail`/`inconclusive`/`not run`. Example scope: "five steps pass bitwise, GSM8K pairing passes, first-round c=32/64 performance over budget". diff --git a/.agents/skills/train-inference-consistency/guides/rollout-engine.md b/.agents/skills/train-inference-consistency/guides/rollout-engine.md new file mode 100644 index 000000000000..e3974e189ae2 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/rollout-engine.md @@ -0,0 +1,108 @@ +# Rollout engine and sampling + +In RL the rollout engine (vLLM) pauses, sleeps and wakes next to the training process, its sampling randomness must be auditable, prefix cache and speculative decoding affect whether the same token sequence is bitwise identical, and RL-only features such as R3 constrain the engine config. Section 6 is a debugging procedure: save one rollout and replay the old-logprob forward. Measurement bases and A/B methods are in [Performance and profiling](performance.md); BI kernels and cudagraph are in [BI kernels](bi-kernels.md). + +## 1. Sampling seed collision + +**Mechanism.** Each vLLM DP engine resets the global RNG with the same `seed` at the end of warmup. Requests without a seed sample from the global stream, so the n-th request samples the same tokens on every engine. When GRPO spreads the siblings of one prompt across DP engines, many samples in a group repeat: before the fix, Nemotron had only 64-70 distinct responses out of 256 and up to 43% of sibling pairs identical; after the fix, 256/256. The fraction of zero-advantage groups barely changes with the fix; zero advantage mostly comes from task difficulty, not seed collision. + +**Fix.** Reset the global RNG by `(seed, data_parallel_index)`; index 0 is bitwise identical to the old behavior (vllm-project/vllm#59788, merged). Use `data_parallel_index`, not `data_parallel_rank`: dense (non-MoE) models run independent engines under DP, and `reconfigure_for_independent_dp_rank()` resets `data_parallel_rank` to 0 (`vllm/v1/engine/core.py`), so 4 engines get the same seed stream and n=4 gets worse. Nemotron is MoE and goes through `DPEngineCoreProc` with a correct rank, so its 1024/1024 verification on verl holds, but only the index version fixes dense models. + +**How to accept.** In the rollout dump, count distinct samples per group and the fraction of fully identical sibling pairs. With different seeds the baseline is about 0.1% identical sibling pairs, all short answers of <=8 tokens (the model outputs `\boxed{}` directly), and about 1-3% share a prefix of >=16 tokens. Anything above this is RNG correlation. + +**Do not do this.** verl `rollout.full_determinism=True` gives the n responses in a group the same seed. Do not enable it. + +**Greedy cannot verify the fix.** Under greedy, n repeats are identical (Nemotron: 32 of 32), and the score is a multiple of 1/number of problems (AIME with 30 problems: multiples of 1/30; 13/30 -> 19-22/30). A greedy A/B before and after the seed fix shows no difference because greedy does not sample. Rising entropy and the three config curves overlapping are visible only on sampling runs. + +## 2. Prefix cache semantics + +**Premise.** "Enabling PC does not change the numerics" requires the same request to be bitwise identical on hit, miss, and PC off. For pure attention models KV block reuse guarantees this. For stateful layers (Mamba2 conv/ssm state) it depends on which states the engine copies. + +**The Mamba2 pitfall.** vLLM's Mamba align mode copies only the conv and ssm states. Under BI=1, Mamba2 uses exact replay (ported from vllm#55905, which stores three extra chunk buffers x/dt/B) and needs 5 states, so BI=1 Mamba2 originally rejected PC. The fix copies every state in each layer's declared state list; after that, hit/miss/PC off are bitwise identical. Early versions had a convolution window alignment bug in state export (written aligned to the tail, read from the head). + +**Cost.** PC under BI=0 has a measurable cost: in Mamba align mode each decode step adds a batch of eager launches, and an 8K prefill is split into two segments. The approach in upstream #58851 is the best. + +**Confounders between the two arms.** BI=1 Mamba2 rejects PC while the BI=0 R3 AuxOutput forces PC on, so PC on/off is a confounder in the BI=0 vs BI=1 throughput comparison. List confounders first, then compare; see [paired acceptance](recipe-and-acceptance.md#paired-acceptance). + +**Relation to cudagraph.** A PC hit changes the token count of a step. Piecewise cudagraph chooses between the graph and eager fallback by total token count vs `max_cudagraph_capture_size`. Each path is batch-invariant on its own, but the two are not bitwise equal. Under BI, pin `FULL_AND_PIECEWISE`; see [BI kernels](bi-kernels.md). + +**Checks for stateful layers and caches.** Mamba/SSM must also save chunk boundary state and partial input buffers; check aligned/unaligned split, cross-chunk and restore separately. Cover indexing and capacity of long-context KV/page/workspace near the boundaries. + +When testing prefill, confirm whether the prefix cache was hit; during diagnosis, disable or isolate it if needed so the target path executes. Verify the production cache's hit, reuse, and preemption/recompute paths separately. Do not make a comparison pass by clearing the cache and then claim the original lifecycle is fixed. + +For performance fixes on the cache path, use bitwise equality with "cache off" as the gold standard: hit positions outside chunk boundaries, every DP rank, recomputed hits after forced preemption, and chunks truncated by the scheduling budget; confirm in the trace that the target branch executed. Then add a hook that checks the device-side state index against the host prediction at every step, plus a per-problem quality evaluation. Two fixes that each pass alone may silently break when stacked; the combination must pass the gate on its own. + +## 3. MTP / speculative decoding + +- The Nemotron checkpoint has `num_nextn_predict_layers: 1`, and the MTP block is not quantized (bf16). It does not take part in training; the training side must explicitly skip its mapping ([training forward](training-forward.md)). +- Usable with BI=0. BI=1 + MTP requires the logprob from speculative verification to be bitwise identical to the full forward; without verifying this you cannot claim true on-policy. +- If used, use it in both arms; otherwise neither the throughput nor the logprob comparison holds. + +## 4. R3 (routed experts replay) feature constraints + +- **What is returned.** One `(num_moe_layers, top_k)` row of uint8 routed expert ids per token (Nemotron: 52 layers x 6 = 312 B), returned to the trainer with the request's CompletionOutput for route replay. 0 is a valid expert id; check distinctness, not non-zero. +- **PC must be on.** The R3 AuxOutput is keyed by KV block hash; with PC off there is no key. This is why the BI=0 arm in §2 is forced to enable PC; the keys-only form (PR 57968) does not lift it. +- **Why per-step return is expensive.** 312 B exceeds the 256 B msgpack zero-copy threshold, so each request's row per step becomes a separate frame, and the frontend accumulates millions of small objects per request before concatenating them at the end. The overhead is on the engine->frontend return path, not in the capture hook. The fix returns once at request finish (finish-only) and packs each step into one ndarray on worker->engine. `output_processor` already exposed `routed_experts` only at finish, so semantics are unchanged. +- **There is another floor below R3.** msgpack unpacking of three small ndarrays per token for sampling logprobs is a fixed frontend cost (about 25-30 µs/token), larger than everything else left after R3 is fixed. BI=1 does not see it when the engine is the bottleneck; BI=0 has a faster engine and hits it. Two arms tying on e2e is often due to this, not the kernel. Check the frontend asyncio busy fraction before comparing the arms ([Performance and profiling](performance.md)). +- **What does not change.** R3 adds no GPU->CPU sync, async scheduling is kept, capture writes into a statically preallocated buffer, and cudagraph is unaffected. flashinfer_cutedsl and Humming are both modular MoE (they receive the router's top-k), so moving R3 into the kernel saves nothing. +- **Limits of the equivalence check.** BI=1 cannot run the equivalence check (Mamba2 BI rejects PC while AuxOutput forces PC, until the §2 fix lands). BI=0 compares sha with strictly serial single-request greedy, but BI=0 routing is itself not batch-invariant, and an unpatched baseline also repeatedly shows partial inequality against the same reference. So first measure the drift with several baseline-vs-baseline runs; the patch counts as equivalent only if it fully equals one clean reference. Bitwise comparison of concurrent greedy dumps is not evidence. + +## 5. Lifecycle and GPU memory ledger + +**Three pause modes** (`EngineCore.pause_scheduler`): `abort` marks all running requests `FINISHED_ABORTED`, queues new requests, and optionally clears the cache; `wait` queues new requests but keeps stepping until running requests finish (not supported by the inproc engine); `keep` pauses everything, and the Future completes when the output queue is empty. In all three modes the Future returned by pause is the barrier; until it completes, workers still step. In `abort` mode the client receives aborted outputs and must resend them. Requests that arrive during pause are queued and released at resume, using the **new** weights. These are "orphan requests" (pause issued before the rollout tail finishes, or the agent loop sending during pause), and their tokens belong to no single policy version. Before each step's pause, confirm in-flight is 0; after resume, confirm the queue is empty. The HTTP endpoints `/sleep`, `/wake_up`, `/is_sleeping` and pause/resume are dev endpoints and need `VLLM_SERVER_DEV_MODE=1`, otherwise 404. + +**Three sleep levels** (`EngineCore.sleep`, `Worker.sleep`): + +| level | What it does | GPU memory | Host memory | +| --- | --- | --- | --- | +| 0 | Only pauses scheduling | Unchanged | Unchanged | +| 1 | Offloads weights to CPU, drops the KV cache, clears the prefix cache | Frees KV and weights | One weight backup per rank (DS4.1 78.61 GiB); the pinned cache is not returned after wake, see [Checkpoint and resume](checkpoint-resume.md) | +| 2 | Discards all GPU memory, saves only `named_buffers` to CPU | All freed | Small | + +After L2, `wake(weights)` only reallocates GPU memory; the contents are uninitialized. "Reserved is very high but still no GPU memory" after wake is an allocator boundary problem: PyTorch's reserved-but-unallocated blocks cannot be used by another allocation path such as CuMem; see [Checkpoint and resume](checkpoint-resume.md#3-reserved-is-high-so-why-can-wake-still-have-no-gpu-memory). + +**Set gpu_memory_utilization by training mode.** The training process does not fully release GPU memory during rollout. Write the value into the recipe per training mode; check the resolved value inside the container and the KV budget in the vLLM log, not the outer setting. Measured values and rules per mode are in [Checkpoint and resume](checkpoint-resume.md#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated). + +**Metrics are cumulative.** `vllm:*_total` counts since process start; when one engine runs several phases, always use `after − before`. Why the first phase happens to be right and how to cross-check are in [Performance and profiling](performance.md#metrics-are-cumulative-values). + +**DP traffic lands on rank 0.** When the Rust frontend sends one request at a time, all of them land on rank 0 (engine selection is `min_by_key(routing_score)`, which takes the first on equal load). "No events on rank 1" then means no traffic, not broken routing. Use `X-data-parallel-rank` to pin requests to each rank in turn. The two frontends interpret this header differently: Python in hybrid mode validates it as a local index in `[0, dp_local_size)`, Rust uses the global rank; do not send it to a single-rank service. Python hybrid/external LB with a dense model does not start (`DPAsyncMPClient` asserts `stats_update_address`, upstream #32252); test Python hybrid with an MoE model. + +--- + +## 6. Save one rollout and replay the old-logprob forward + +### When to use + +For a difference that shows up once across many RL steps, save the failing input batch and matching policy early. Rerunning the rollout changes tokens, scheduling, lengths, cache, packing and later training data at once, hiding whether an intervention worked. + +1. Save the framework batch and the policy identity after the real rollout ends and before the old-logprob forward; append the trainer scores after the comparison. +2. Verify that the save hook executed, that sequence / valid token counts match the log, and that the data reads back. +3. Use the framework's existing training-only/repeat entry point to read the same batch; restore its weights. +4. Return after old-logprob; forbid the optimizer step and regeneration; record the skipped stages and the weight fingerprints before and after. +5. Run A/A first, then change one layout factor. If A/A itself is unstable, first check randomness, mutable buffers, async streams and runtime state. + +### Minimal matrix + +| Variant | Change | Answers | +| --- | --- | --- | +| A | Original packing | Can the frozen input reproduce it | +| B | Exact repeat of A | Is it run-to-run nondeterministic | +| C | Sequence permutation/packing | Does it depend on layout | +| D | micro-batch size or target position | Does it depend on call shape/neighbors | +| E | Separate process or equivalent node | Does it depend on initialization/device state | + +Run first the variants that distinguish the current hypotheses; the full matrix is not needed every time. After reordering, restore and compare by UID and response token position; the physical row no longer represents the same request. C failing does not prove the packing algorithm wrong; it may trigger a CP, page, tile or expert boundary error. + +### Keep the trigger condition when shrinking the input + +Keep the target sample and its necessary neighbors first, then reduce length and request count. After every reduction, rerun the criterion to confirm the failure remains. A CP problem may depend on the remainder from the end of the previous request to the start of the rank; a Graph problem may depend on short fillers and the capture bucket. Keeping only the failing request often makes the problem disappear. + +Freezing the rollout and the weights does not freeze all runtime state. Training-mode dropout, FP8 amax history, RNG, route replay, caches and workspace must each be kept or reset according to the failure condition. A fresh eval forward is not a substitute for the original old-logprob path. + +A resident model saves checkpoint loading, but state must not leak across variants. First verify A/A between a new process and a reused process; when native libraries, autotune or compile config are involved, restart the corresponding workers. + +### Leaving the replay stage + +Once the original reproducer is closed, return to normal generation, real backward/update and the next rollout step. Training-only success supports only the frozen-policy forward conclusion. If the error appears in backward or in long-running state, save that state and build a minimal reproducer for that stage. + +This rollout replay is a debugging procedure. Nemotron's SSD exact replay is an algorithm in the model implementation that recomputes the partial chunk to match prefill. Do not call them by the same name. diff --git a/.agents/skills/train-inference-consistency/guides/training-forward.md b/.agents/skills/train-inference-consistency/guides/training-forward.md new file mode 100644 index 000000000000..049b474530c6 --- /dev/null +++ b/.agents/skills/train-inference-consistency/guides/training-forward.md @@ -0,0 +1,412 @@ +# Training-side model support + +How to build a training implementation cheaply and prove it computes the same math as inference. Where training weights come from (master -> training forward weights) is in [Precision, quantization and QAT](precision-quant.md). + +## Three routes + +The routes differ in where parallelism and kernels come from. In all three we write the model. + +| | MCore (Megatron-Bridge) | FSDP2 (VERL's HF interface path) | mLite | +| --- | --- | --- | --- | +| Positioning | Mature TP/PP/EP/CP, main production RL path | Cheapest: write only the model, no parallelism | Built for consistency: forward calls vLLM kernels, backward reuses the training implementation | +| Forward source | MCore kernels (TE, grouped GEMM) | Our PyTorch model to vLLM semantics, in the Transformers `PreTrainedModel` interface for the VERL loader and PEFT | vLLM kernels (`visible_forward(visible, native, ...)`) | +| Achievable consistency | BI=0, via R3 + importance sampling | BI=0 | BI=1, logprob raw bytes identical | +| What to write for a new model | Bridge weight mapping + modules MCore lacks (Mamba/SSM, special norms) | Whole model: config class, checkpoint loading (with dequantization), per-layer forward, per-operator precision contract; autograd backward, filled in where missing | One `model//`: config, checkpoint mapping, assembly, protocol; modules from `primitive/modules` | +| Typical pitfalls | Mapping checks values, not semantics; MCore and vLLM math differ (MoE activation clamp) | bf16 autocast around the VERL forward rewrites module precision; `--system-site-packages` ABI; PyTorch reference (Mamba scan) too slow; loading a third-party implementation from the image by mistake | Private modules inside the model directory, bypassing the shared skeleton | + +Choose after the [module inventory](#module-inventory-and-placement-rules), by: is BI=1 needed, how many backwards are missing, what parallelism is needed. Nemotron looked cheap by name; FA4, Mamba scan and W4A16 MoE all lacked a backward. + +## Find by symptom + +| Symptom / question | Where to go | +| --- | --- | +| New model loads and config parses, but unclear whether it can train | [Module inventory](#module-inventory-and-placement-rules) | +| Want to reuse vLLM forward code, but tangled in paged KV / cudagraph / scheduling | [Call kernels only](#reproducing-the-forward-call-kernels-only-do-not-port-the-scheduler) | +| Weight mapping shapes are right, output silently degrades | [Mapping checks values and semantics](#43-mapping-checks-cover-values-and-semantics-not-only-shape) | +| Inference kernel has no backward; should we write one | [Three-case decision](backward-validation.md#three-cases) | +| Wrote a VJP; how to prove it is right | [Module-level comparison steps](backward-validation.md#comparison-steps) | +| Unit tests pass but real training fails / step 0 occasionally not bitwise | [Two test-infrastructure false negatives](backward-validation.md#two-test-infrastructure-false-negatives); [BI kernels](bi-kernels.md) | +| Changes after modifying packing or CP/TP/EP; wrong only on long sequences | [Layout and length](#distributed-layout-length-boundaries-and-training-side-acceptance) | +| `ppo_kl` non-zero after rollout switches parallel layout (e.g. DP+EP to TP1 per GPU) | [Bitwise identity after changing the parallel layout](#bitwise-identity-after-changing-the-parallel-layout) | +| `modified by an inplace operation` after no longer recomputing old-logprob | [Entropy when old-logprob is not recomputed](#entropy-when-old-logprob-is-not-recomputed-bypass) | +| How to hook into VERL / Megatron | [Framework integration](#verl--megatron-integration) | +| Which kernel an operator uses in BI mode and its cost | [BI kernels](bi-kernels.md) | + +## Key rules + +1. Semantics follow the vLLM implementation and the checkpoint. +2. Before starting, list for every module whether its inference kernel has a backward. +3. Call kernels only; do not port the scheduler. mLite (BI=1): FP32 logprob raw bytes vs vLLM-captured trajectories, target 0 diff. FSDP2 / MCore (BI=0): no 0-diff target; judge each side against an independent reference, then look at precision. +4. Mapping checks cover values and semantics, not only shape; fused parameters are overwritten whole. +5. The precision contract lives in the model. Where vLLM computes in FP32 (position-encoding angles, router logits and bias, norm statistics, etc.), the training module sets the dtype and disables the outer autocast. Under VERL's bf16 autocast, casting only the inputs of matmul-type operators (including `F.linear` and `einsum` decomposed into bmm) does not stop the downgrade. Measure output dtype and values under the real autocast. +6. Backward follows the formula, not the forward kernel. Only rounding differs -> reuse. Quantization rounding -> fix the STE contract, reuse BF16 GEMM. No whole-segment backward -> write glue, GEMM still TE. "No backward kernel at the same precision" is not a reason to rewrite. +7. Reuse model-agnostic operators: mLite `primitive/modules` and `ops` first, then MCore/TE/mamba_ssm; only then is it new code. +8. Backward need not be bitwise; report bitwise and tolerance separately; the reference builds its graph independently with ATen autograd; tolerances are fixed before seeing results; synthetic inputs are labeled as not real activations. +9. Do not recompute activations; use the forward's saved intermediates (DS4: gate_up / LSE). +10. The model directory holds only config / checkpoint / model / protocol; reusable parts go into `primitive/`; no same-named `primitive/` under the model directory; impl name `lite`. +11. Long-sequence boundary points (KV page, packing prepend, the 7000/9216 tiers) need dedicated tests. +12. MTP block, EPLB redundant experts and draft are explicitly skipped in mapping and sync. + +## Module inventory and placement rules + +List per module: inference kernel, same-math training implementation, quantization rounding, discrete selection, backward. Then pick the route and each module's place in the code tree. This catches missing backwards before they cost weeks. + +Training support means an implementation we maintain that backpropagates correctly. Config registration, same-family fallback, the vLLM inference implementation and third-party implementations in the image (newer Transformers modeling, NeMo AutoModel) do not count. Write it per [Reproducing the forward](#reproducing-the-forward-call-kernels-only-do-not-port-the-scheduler) into real source files and verify by loading through the actual loader in a new process. A REPL-only implementation or runtime monkeypatch is not evidence. + +### 1. How to build the inventory + +One row per module, every column filled; write unknown, not "should exist". + +| Column | Question to answer | +| --- | --- | +| Inference kernel | The function vLLM actually calls (not the model-file class name), with quantization variants, attention backend, cudagraph branch | +| Same-math implementation in the training framework | mLite `primitive/modules`, `ops`; MCore / TE; mamba_ssm, causal_conv1d. Rounding-only differences count as same math; a different activation clamp does not. Third-party same-model implementations are not entered | +| Quantization rounding | FP8/NVFP4 rounding in the forward; weights and activations separately | +| Discrete selection | top-k expert IDs, sparse indices; IDs fixed and not differentiated, the continuous score needs a derivative | +| backward | Does the library have one; do its intermediates (LSE, rstd, token permutation) match what our forward saves | +| Conclusion | Reuse / glue / new code ([decision rules](backward-validation.md#three-cases)) | + +Nemotron-H: + +| Module | Inference kernel | Same-math training implementation | Rounding | Discrete | backward | Conclusion | +| --- | --- | --- | --- | --- | --- | --- | +| Mamba SSD scan | `mamba_chunk_scan_combined_varlen` | mamba_ssm | None | None | mamba_ssm `_mamba_chunk_scan_combined_bwd` | Reuse | +| conv1d | `causal_conv1d_fn` | causal_conv1d | None | None | Yes; dW needs fixed-order FP32 | Reuse + small glue | +| Attention (FP8 KV) | FA4 fixed split + FP8 Q/K/V | FA4 BF16 | Yes | None | FA4 BF16 backward only | Dequantize, recompute BF16 O/LSE, reuse | +| Mamba in/out_proj (FP8), shared expert (NVFP4) | FlashInfer / NVFP4 GEMM | TE `general_gemm` | Yes | None | TE | Fix STE, reuse | +| Routed experts (NVFP4 W4A16) | CuTe-DSL / Humming MoE | TE grouped GEMM | Yes | Yes | No whole-segment backward | Glue `routed_vjp` | +| RMSNorm / gated RMSNorm, routing score | BI kernel / grouped top-k | PyTorch formula | None | None for the score | autograd | Reuse | +| MTP block, draft | Not part of training | - | - | - | - | See §4 | + +### 3. Placement rules + +#### 3.1 mLite (Megatron `experimental/lite/megatron/lite`) + +Read `experimental/lite/docs/porting.md` and `architecture.md` first. Three layers: `runtime` (lifecycle, training-step orchestration), `model` (registration + model-specific build/load/export protocol), `primitive` (reusable: parallel state, TP layers, checkpoint conversion, MoE utilities, optimizer wrappers). The runtime knows no model details: it gets the protocol from the registry, builds a typed config, and lets the protocol build the model. + +| Rule | Content | +| --- | --- | +| The model directory holds only four things | `model//`: config, checkpoint mapping (WeightSpec), assembly, protocol. Impl name `lite` (`/lite/`) | +| Modules go into primitive | norm / linear / router / experts / dispatcher / attention -> `primitive/modules`; generic VJP bridge -> `primitive/ops`; quantization encoding -> `primitive/quantization`; PP splitting, thd, CP -> `primitive/parallel`; HF read/write via `primitive/ckpt/hf_weights` (`SafeTensorReader`, `load_hf_weights`, `export_hf_weights`) | +| Model-specific math still goes into primitive and is registered | e.g. Mamba2 bridge, FP8-KV attention -> `primitive/modules/.py`; no same-named `primitive/` under the model directory | +| No file for reused backward | Reusing TE / mamba_ssm backward means no `backward.py`; glue goes in the module file | +| Dependencies are one-directional | `model/` imports `primitive/`, never the reverse; top-level checkpoint/quantization does not import the impl subtree | +| Layering guards must be registered | Add the model to `tests/unit/runtime/test_layering_contracts.py` or the guard does not apply | +| No hybrid architecture packages | porting.md: "Keep these out of the first PR: Hybrid model implementation packages"; hybrid (Mamba + attention + MoE) modules are split into primitive | +| State third-party kernel dependencies | Document hard dependencies on vllm, flashinfer private APIs, mamba_ssm, causal_conv1d, and the `init_batch_invariance()` call in build_model | + +DS4 is the positive example: `model.py` (650 lines) assembles te.RMSNorm, CSA, HCA/mHC, thd and the fp8 recipe with three model-specific deviations; `moe.py` (102 lines) = SigmoidTopKRouter + Experts + TokenDispatcher + SwiGLUMLP, hash-route the only model-specific code; `checkpoint.py` is one WeightSpec using shared load/export. + +#### 3.2 MCore / Megatron-Bridge + +- Bridge `models//`: provider + HF<->MCore mapping; modules MCore lacks (Mamba/SSM variants, special norms, HC) go into `megatron/core`, not Bridge. +- Mapping checks values, not only shape ([§4.3](#43-mapping-checks-cover-values-and-semantics-not-only-shape)). +- MCore and vLLM may use different formulas (MoE activation clamp); BI=0 cannot align those either, so list them first. +- The Megatron fork works only with Megatron-Bridge v0.6.2, pinned in the image. + +#### 3.3 FSDP2 (VERL's HF interface path) + +- Write a PyTorch model to vLLM semantics in the Transformers `PreTrainedModel` / config interface so the VERL loader, PEFT and HF save/load use it directly. "HF" is only the interface and weight format. +- Delivery: the model goes into the verl trunk, registered to the Auto classes (`AutoConfig` / `AutoModelForCausalLM`, etc.), and ships with the recipe as one verl PR ([Code and delivery §8](../practices/code-and-delivery.md#verl-a-thin-recipe-not-a-second-training-framework)). The verl FSDP extension guide's advice against copying modeling into `verl/models` targets models already in transformers. Register at the current checkout's FSDP load chain (`get_hf_auto_model_class` -> `from_pretrained`). +- Develop and deliver from the same verl source and registration entry. Do not modify the checkpoint's `config.json` / `auto_map` (the low-precision checkpoint is the single source of truth) or enable `trust_remote_code`. The load log must show our registered class was used. +- VERL's automodel engine (`automodel_actor`) uses NeMo AutoModel and is not this route; routes are distinguished by forward source, not by FSDP2 sharding. +- Under bf16 autocast, FP32 computations (position-encoding angles, router logits and bias, norm statistics, following vLLM) disable autocast inside the module and set the dtype. +- Backward prefers autograd; add operators next to the model file only when a backward is missing or too slow (the 160-line PyTorch Mamba scan reference). +- Common pitfalls: `--system-site-packages` ABI, the actual `use_remove_padding`, PEFT targets not going through nn.Linear. +- Under EP4 + FSDP2, all ranks' DCP expert shards share key and offset, so restore reads another rank's experts; the key must carry the EP rank ([checkpoint guide](checkpoint-resume.md)). + +### 4. Parts that do not take part in training + +Each gets its own inventory row: + +| Part | Mapping | Facts | +| --- | --- | --- | +| MTP block | Not loaded for training; WeightSpec lists it as explicit skip, not missing-key tolerance; `_ignore_unexpected_suffixes` etc. only for listed suffixes | Nemotron `num_nextn_predict_layers: 1`, MTP unquantized bf16; usable with BI=0, then in both arms; BI=1 + MTP requires speculative-verification logprob bitwise equal to the full-segment forward | +| EPLB redundant experts | Training has logical experts only; physical-to-logical mapping is inference-side | Completeness counts logical experts, not physical tensors | +| draft / speculative model | Independent checkpoint and mapping | Shared weights (embedding, lm_head) state their owner | + +Skipped parts appear in the load log and manifest as "explicitly skipped", not "no error reported". + +### 5. Example: reuse audit of Nemotron-H against mLite + +Nemotron-H is 3824 lines. From `primitive/` it reuses only protocol-layer orchestration utilities, `load_hf_config_dict` and `vocab_parallel_entropy`; every modeling module is rewritten under `model/nemotron_h/vllm/primitive/`. + +| File (`N = model/nemotron_h`) | Lines | What it does | Verdict | +| --- | --- | --- | --- | +| `N/config.py` | 155 | HF config dataclass, validation | Keep (model dir) | +| `N/quantization.py` | 207 | NVFP4/FP8 requantize, E2M1 grid, FP8 encoding, hand-read index.json | Numerics -> `primitive/quantization/nvfp4.py` (parameterize E2M1 RNE vs mxfp4 ties-down); IO via `SafeTensorReader` | +| `N/checkpoint.py` | 486 | HF name mapping, Export spec, BF16 master invertibility check, θ0 threshold, own `load_hf_weights` | Ordinary tensors via `load_hf_weights` + WeightSpec; keep master check and θ0 threshold; importing the impl subtree is a layering inversion | +| `N/vllm/model.py` | 188 | Layer/Model assembly, PP stage | Keep; PP split via `build_pipeline_chunk_layout` | +| `N/vllm/protocol.py` | 409 | protocol, contract validation, BI check, FSDP2/dist_opt, post-step requant | Keep; hard-coded 52-layer pattern and (2688,128,1856,6) make it single-checkpoint | +| `N/vllm/primitive/dense.py` | 647 | visible forward + BF16-master VJP bridge; RMSNorm VJP; Humming/CuTe NVFP4 linear; projection factory; FP8/NVFP4 TrainingLinear | Bridge and RMSNorm VJP -> `primitive/ops/aligned_vjp.py`; `native_linear_vjp` duplicates `_te_general_gemm`; serving kernel adaptation stays | +| `N/vllm/primitive/logprob.py` | 36 | Chunked lm_head + vLLM `compute_token_logprobs` + FP32 log_softmax STE | Keep; the docstring's DS4 function does not exist | +| `N/vllm/primitive/attention/*` | 231 | `Fa4Fp8KVAttention` (FlashAttnFixedSplit 6768-token page, FP8 Q/KV), `_Fp8AttentionVJP` | Model-specific; reusable once GQAttention supports `rotary=None, qk_norm=False` | +| `N/vllm/primitive/mamba/*` | 316 | `_PackedConvVJP`, `_PackedScanVJP`, `SSMMeta`, `MambaMixer` | Keep both VJPs; generic Mamba2 -> `primitive/modules/mamba2.py`; derive `SSMMeta` from `PackedSeqParams`/`ThdPackMeta` | +| `N/vllm/primitive/moe/module.py` | 185 | `_FixedRouteVJP` + Router, `_CombineVJP`, ReLU² SharedExperts | Router VJP equals `SigmoidTopKRouter`, reusable, not hooked to router_replay; add an activation argument to `mlp.py` for SharedExperts | +| `N/vllm/primitive/moe/grouped.py` | 705 | FlashInfer W4A16 signature/sha pinning, `CuteDslRoutedExperts`, `Nvfp4ExpertWeights`, `_te_grouped_gemm` rewrite, ReLU² `routed_vjp` | Keep W4A16 forward + grouped VJP; TE grouped GEMM wrapper duplicated; `ep4_routed_experts` hard-codes 4 ranks | +| `N/vllm/primitive/moe/communication.py` | 251 | DeepEP normal-mode EP, `reduce_ep4_parts` (FlashInfer one-sided combine order), `EPRoutedExpertsVJP` | Use `_build_deepep_buffer`; keep serving-order reduction, parameterize EP size/topk | + +Six placement violations: impl name `vllm` not `lite`; hybrid architecture package; same-named `primitive/` in the model layer; circular imports between `checkpoint.py`/`quantization.py` and the impl subtree; `nemotron_h` missing from the layering guard; third-party kernel dependencies and `init_batch_invariance()` undocumented. + +Should have been reused: safetensors index reading (5 copies); HF load of ordinary tensors; DeepEP buffer; TE GEMM wrapper (2 copies); visible forward + reference VJP bridge; Router (`SigmoidTopKRouter` + `group_limited_topk`, forward switched to vLLM grouped_topk, hooked to router_replay); shared expert; E2M1 grid and FP8 encoding; PP splitting; `SSMMeta`; two docstrings citing nonexistent DS4 functions. + +Model-specific, keep: the two Mamba2 VJPs + packed conv/scan; `Fa4Fp8KVAttention` + FP8 attention VJP; W4A16 MoE forward + ReLU² grouped VJP; serving-order reduction (parameterized); Humming/CuTe NVFP4 and FP8 serving linear; vLLM forward-value logprob; BF16-master invertibility check and θ0 threshold; ModelOpt 4over6 requantize rule. + +## Reproducing the forward: call kernels only, do not port the scheduler + +The first deliverable after the inventory is a training forward with vLLM's math. Logprob measurement basis and R3 route replay: [workflow](../workflow/README.md), [rollout engine](rollout-engine.md). + +### 1. First separate the responsibilities of the model implementation and the training framework + +A training framework (FSDP2, MCore, mLite runtime) shards a trainable model; it does not make an inference-only kernel differentiable or write the model. Write a minimal PyTorch model for the current loader interface covering the config, checkpoint loading, forward, gradients and save interface actually used; full Transformers integration, generation and all parallelism are not needed first. The semantic basis is vLLM and the checkpoint, not a third-party implementation in the image. + +| Capability | What the training route still has to do | What in inference / Megatron is not copied | +| --- | --- | --- | +| Model math | Real semantics of attention/SSM, norm, position encoding, routing, residuals, special embeddings | vLLM paged KV, decode scheduling, cudagraph management; MCore provider/ModuleSpec wiring | +| Backpropagation | Input gradients of every training-path operator; trained-parameter gradients; recompute correctness ([Backward validation](backward-validation.md)) | TP/SP/CP-specific collectives and gradient reduction | +| Weights | Names, axes, fused order, expert IDs, sharing, derived state | HF<->Megatron shard layout, PP export protocol | +| LoRA | Target coverage, A/B direction and scaling, non-zero adapter taking effect | Parallel linear wrapper, TP adapter aggregation | +| Memory | Initialization peak, large frozen table residency, save/export peak | Megatron memory saver, PP scheduling | + +dtype, dequantization, QAT and update format: [precision guide](precision-quant.md). Here each module states its input/output precision and differentiable boundary. + +### 2. How to wire kernel-only calls + +Each layer's forward runs once: call the vLLM kernel for the visible output, bitwise identical to inference, and save what the backward needs; see [backward wiring](backward-validation.md#how-to-wire-it-forward-once-backward-uses-what-the-forward-saved). Borrow only the kernel, no vLLM runtime object: + +- No KV paging, decode scheduling or CUDA Graph management. Training runs a full-sequence teacher-forced forward; packing passes boundaries via thd/cu_seqlens. +- Reuse order: model-agnostic differentiable operators -> concise PyTorch reference to vLLM semantics -> ready-made operators with backward -> custom autograd/kernel only if necessary. This order only says where to look first; the training-step cost of each choice is checked at C6 against the C1 budget, not deferred to C10. +- Handle discrete top-k / index selection apart from the selected differentiable tensors. What is trained, frozen or given an auxiliary loss follows the model definition; do not detach to get it running or add router/indexer objectives. +- vLLM layer modules used directly (FusedMoE, ModelOpt linear) need `get_tp/ep_group`, so the training process creates vLLM process groups (Nemotron does; DS4's own modules do not). +- Take the same branch as inference: reproduce the quantization config, attention backend and BI-switch kernel selection, or the comparison is against a different path. Tests use the production `CompilationConfig(custom_ops=...)` ([the two false negatives](backward-validation.md#two-test-infrastructure-false-negatives)). + +"The same kernel" still has pitfalls (Nemotron 7-layer localization): + +1. Direct and compiled paths of the same gated RMS kernel differ; calling it through an opaque custom-op boundary removed the difference. +2. A differing reduction block was wrongly blamed alone; the same captured kernel showed both blocks match. +3. The last 4.768e-7 logprob difference came from the HF forward calling `.float()` after lm_head; keeping the real BF16 head output and sharing the scoring fixed it. This is not truncating an FP32 projection to BF16 to manufacture equality. + +### 3. Read the inference implementation along the full support chain + +Order: model file (layer types and order, residual and norm positions, position encoding, attention visibility, router/indexer, special embeddings) -> each layer's actual kernel and quantization variant -> attention backend and KV dtype -> weight loader (name mapping, fused order, derived scale) -> BI-switch branches. Confirm the parameter/buffer list from the same checkpoint/config. + +Sizes sharing one implementation need no separate port. For ordinary dense models, look for a reusable generic skeleton first; we still write the model. Multimodal is validated per task; diffusion is out of scope. + +### 4. The minimal implementation that really has to be filled in + +#### 4.1 Reproduce the math + +Confirm from the checkpoint/config and inference code: layer types and order, parameters/buffers, residual and norm positions, position encoding, attention visibility, router/indexer, special embeddings. Model-specific math (`sqrt(softplus)` scoring, SwiGLU with clamp, mHC Sinkhorn, ReLU² experts) is in no library; write it from the inference source, not the paper. + +#### 4.2 A frozen base still has to propagate input gradients + +```python +W.requires_grad_(False) +y = linear(x, W) + scale * B(A(x)) # no gradient for W, but still propagate to x +loss(y).backward() +``` + +Wrapping a frozen block in `no_grad()` cuts LoRA gradients of earlier layers; only a frozen table lookup needing no input/parameter derivative is exempt. With initial B=0 a zero gradient for A can be normal; test with a non-zero adapter or a real step. A finite loss does not prove backward. + + + +#### 4.3 Mapping checks cover values and semantics, not only shape + +Per weight class: source name -> training name -> serving name, axis transforms, sharing. Focus on fused QKV, gate/up order, expert numbering, extra norm/bias, tie weights, non-persistent buffers, derived state (FP8 KV scale, NVFP4 alpha). Every missing/extra item needs a reason; quantization-derived items follow the precision guide and need not be byte-identical to the low-precision source. FSDP needs no HF<->Megatron bridge but may need checkpoint->training and export mappings. + +#### 4.4 Packing boundaries for Mamba / stateful modules + +Every stateful operator (conv1d, SSM scan, recurrence, PLE shift) must consume sequence boundaries; a correct attention mask proves nothing about them. Incidents: + +- verl via Megatron-Bridge did not pass boundaries to Mamba; sequence states mixed and logprob was silently wrong, visible only as a larger mismatch. +- Qwen3.8 Flash-Next's PLE shift made the second sample reference the first sample's last token. + +Check: unequal-length A/B run separately and merged; sentinel-initialized outputs to catch missed writes. In mLite, derive boundary metadata from `PackedSeqParams`/`ThdPackMeta`, not a separate `SSMMeta`. See also [Long-sequence boundary points](#long-sequence-boundary-points). + +#### 4.5 Recompute and small-model trimming + +With activation checkpointing on, verify recompute reads the same state and gives reasonable gradients. Trim to the smallest layer set covering every non-equivalent path and its dependencies, usually 4-6 layers; keep mixed attention, cross-layer sharing and layer-specific table lookups. Record kept layer indices and the config transform; a slice proves only covered paths, not full-model mismatch or score gain. + +### 5. How to prove it computes the same math + +Feed the training forward vLLM-captured trajectories (fixed tokens, weight version, position, mask, routing); never compare probability summaries or argmax. mLite (BI=1): FP32 logprob raw bytes, target 0 diff. FSDP2 / MCore (BI=0): judge each side against an independent reference, then check the remaining difference comes only from precision rules. + +- Same policy weights. Differing dtypes are a contract mismatch; do not convert to fake bitwise. +- Reference artifact (Nemotron): frozen trainer image, full actor, 2048 tokens, FP32 0 diff; five real updates, 883,998 tokens of identical FP32 logprob. +- Old-logprob recomputation and the grad-enabled forward are two paths; evidence for each ([dual-path contract](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract)). +- Find the first divergence per [debugging](debugging-mismatch.md): per-module dump, single-request prefill, then the operator. + +### 6. Small tasks handed to agents and stop conditions + +Each outputs a short conclusion, source version, reproducible command and artifact. + +| Task | Minimal acceptance | If it fails | +| --- | --- | --- | +| Implementation inventory | Training model class, operator backward, checkpoint loader; gaps listed | Write the minimal model; stop cold-starting the large model | +| Module comparison | Same weights/inputs; gradients vs a differentiable reference | Locate the earliest difference | +| Few-layer training | Non-zero loss, finite gradients, intended adapter changes | Separate broken path, operator bug, gradient-less target | +| Save and restore | Fresh-process restore: consistent forward and next step | Check state, adapter, optimizer, residency | + +Distribute only after the first two pass; scale up after the few-layer two-step loop passes. + +### 7. How to judge development cost + +Training implementation exists: wire loading, LoRA, sharding. Inference only with differentiable references: fill in assembly, mapping, interfaces. Key operator lacks backward or clear semantics: separate blocker; reference and gradient validation first; no uniform person-day figure. + +## Distributed layout, length boundaries and training-side acceptance + +Record trainer and rollout topologies separately. TP/CP/EP, padding, packing and collectives change data mapping and reduction paths; long sequences hit untested boundaries; forward 0 diff does not vouch for gradients or the optimizer. + +### Compute resources first, then request them + +Write trainer and rollout process groups separately. At four GPUs per node, an inference DP8/EP8 group spans two nodes; two DP4/EP4 replicas are not one EP8 group and skip cross-node EP communication. The runtime decides grouping. + +If dense DP excludes CP: `W = TP × PP × CP × DP`. Expert side: verify `W = ETP × EP × PP × EDP` and the CP mapping separately. **Do not multiply EP into the dense world again.** W=8, TP1/PP2/CP2 gives DP=2; a log calling DP×CP "DP" shows 4. Clarify naming and members; do not change GPU count to match a config string. + +GPU memory follows the layout. With PP1 each rank holds all layers' activations, linear in sequence length. Per-route EP backward state stored for other ranks' slots grows by the average ranks per token; store only own routes. Rollout `gpu_memory_utilization` is set from what training still holds during rollout, per training mode, in the recipe ([Memory budget 02 §5](checkpoint-resume.md#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated)). + +### Packing and CP + +For the target UID reconstruct global request `[seq_start, seq_end)`, rank range `[rank_start, rank_end)`, local rows, cu_seqlens, positions, all compression/reordering mappings. + +The intersection `[max(seq_start, rank_start), min(seq_end, rank_end))` contributes rows only when start < end. The overlap window must not give a disjoint request this rank's output. Host cumulative boundaries equal the rows the device compactor produces. + +Cover: before/after/at the boundary, empty intersection, cross-rank, compression-ratio and tile remainders. Window constant 8: check 1…7, not only multiples. Start from a hand-computable counterexample, then enumerate nearby legal configs. + +Passing `cu_seqlens` does not mean a fallback consumes it ([§4.4](#44-packing-boundaries-for-mamba--stateful-modules)). + +### EP and collectives + +Record expert mapping, tokens per rank, padding capacity, route order, shared/routed combine, quantization scales. Empty ranks may still join collectives; partially written reused buffers leave stale data. Use different sentinels per side; zero-filling both gives a false green. + +Pinning NCCL algorithm/channels may remove reduction differences on one stack and cut throughput significantly; verify it per topology, not as a universal env list, before relaxing config or building a deterministic collective. + +MCore dev's grad-ready assertion takes first-batch hook counts as golden and errors on a different count. Short-sequence micro-batches can make counts data-dependent (Nemotron r10 at step 2; r12 did not reproduce). With dp×cp=1 and expert_dp=1, `overlap_grad_reduce` off avoids it without changing values; root cause unknown. With EP4 + dist_opt and dp>1, multi-micro-batch reduction fires early (reported). + +### Bitwise identity after changing the parallel layout + +**Principle.** Floating-point addition is not associative. Layout sets reduction grouping, order and rounding: TP splits a matmul into per-rank partial sums; EP computes a token's top-k experts on different ranks and merges. Last-bit differences become different tokens after dozens of layers and sampling, so bitwise identity holds only for the verified layout; any TP/EP/DP change must be proven again. + +**Rules.** + +1. The reduction tree is part of the numerical contract: per level, accumulation dtype, order and rounding points. All legal rollout layouts and training walk the same tree. +2. A degenerate layout (TP=1, EP=1) must simulate the tree: per-rank partial sums, same rounding, same merge order. +3. After a layout change, compare per-token logprob raw bytes on fixed inputs before training. Non-zero `ppo_kl`/probs_diff is fixed, not avoided by reverting the layout. +4. Locate: dump per layer, find the first differing layer, confirm its inputs (MoE: router logits too) are bitwise equal; the bug is that layer's reduction ([Fault localization](debugging-mismatch.md#bisect-and-first-divergence)). +5. Measure the cost of simulating the order and report it apart from the layout gain. + +**Examples.** + +| Model | Layout change | Where the difference is | Fix and result | +| --- | --- | --- | --- | +| Nemotron-H (W4A16 MoE, BI=1) | rollout from DP4EP4 per node to TP1 per GPU | EP4: per-rank FP32 sums rounded to BF16, merged in FP32 as `((s0+s1)+(s2+s3))+(s4+s5)`. EP1: six outputs summed sequentially in FP32, no intermediate rounding. Max logprob diff 0.4; greedy tokens diverge at positions 10–129 | Without EP, BI mode computes all experts in one launch and reduces in EP4 order (vLLM `661cde6833`, not in the image). Both layouts match training bitwise; `ppo_kl` 0 over 58 steps; multi-turn replay 404 s vs 485 s, reduction cost about 2.5% | +| DeepSeek V4.1 | training TP=1, inference TP>1 | order of per-TP-rank partial sums | both sides use one fixed reduction tree, TP=1 included (`VLLM_TP_INVARIANT_LEAVES`) | + +### Asynchronous communication and lifetime + +Check streams, events, buffer owners and reuse timing. Passing with a global synchronize only proves a timing issue and is not the fix; find the missing dependency and accept on the same path. + +### Long-sequence boundary points + +Test at real lengths: + +| Boundary | Fact | How to test | +| --- | --- | --- | +| KV page | training attention replicates vLLM's KV page size; page boundaries are numerical boundaries | sparse dY across the page boundary | +| Fixed split attention | fixed split count and keys per split; the resulting `MAX_SEQ_LEN` is an invariance bound, not a performance knob | above: error, not silent truncation; below: two real lengths × random/zero/sparse dY | +| Packing prepend | one prepended short sequence shifts every target row offset | two prepend controls vs none, bitwise | +| Micro-batch fits the longest sequence | verl needs each micro-batch to fit the longest sequence (`ppo_max_token_len_per_gpu`, Nemotron 16384 vs DS4 4096; Nemotron training has no CP) | OOM probe at formal length; do not extrapolate from 8K | +| Embedding backward | CUDA `embedding_dense_backward` accumulates at reduced precision with ≤3072 indices, e.g. a short sequence alone in a micro-batch | length-tiered tolerance | +| CP compressed boundary | see [Packing and CP](#packing-and-cp) | before/after/at the boundary | + + + +### The D term in attention backward + +With P rounded to FP8, visible O = fp8(P)·V ≠ softmax(QKᵀ)·V. If `D = rowsum(dO⊙O)` uses visible O while dS uses recomputed P, `rowsum(dS_i) ≠ 0` and dQ gains a same-sign per-row bias (Nemotron: about 3% of dQ/dK, v2 FAIL at steps 150/250; fixed 0.006–0.009, PASS). Mainstream FP8 attention (TE fp8_dpa, SageBwd) uses visible O, so a human decides after being informed. + +Four checks: + +1. **FA4 interface signature**: `_flash_attn_bwd` derives D in preprocess from the `out` passed; D cannot be passed. The public interface errors under `requires_grad`; the actor calls `_flash_attn_fwd/_bwd`. Fix: recompute BF16 O and LSE from dequantized Q/K/V, pass to bwd (layer backward +4–15%, end to end <0.05%). +2. **varlen LSE layout**: ln vs log2, per head vs per token, softmax scale folded or not. On mismatch, recompute with the training kernel. +3. **GQA 32/2**: dK/dV summed within each Q head group; A/A needs fixed summation order. +4. **When the saved Q/K/V are released under recompute=full**: if only dequantized Q/K/V are saved, recompute must read the same state; early release changes the values. + +### Microbatch split and gradient difference + +Different microbatch splits give about 2e-3 relative gradient difference (loss normalization, accumulation order, reduced-precision embedding backward at ≤3072 indices). Not a bug, but: + +- Criterion denominator = `max(B16 reference floor, noise floor of the implementation under test)`, never the FP32 ideal. +- Both arms use the same split; if it changes, rerun A/A. +- grad_norm is not comparable across different microbatch counts. + +### Control matrix + +With policy/tokens fixed, test original layout, reversed order, target first/last, changed micro-batch, CP-boundary sample, each reconstructed by UID/global token coordinates. Real TP/EP changes may legally change the reduction; the experiment contract decides whether cross-topology bitwise equality is required. + +Safe packing is not a fix for a metadata bug; it may only lower trigger probability. After the fix, the original layout and nearby boundaries must pass. + + + +### Training-side correctness acceptance: gradient comparison criterion v2 and optauth + +Two independent checks beyond forward 0 diff. + +#### Gradient comparison criterion v2 + +- Freeze the criterion (tolerance, denominator, thresholds) before the run, record its sha, never change it after. +- Planted-error check first: gradients ×1.25 and ×1.05 must both FAIL. +- Accept on several snapshots (steps 0/50/150/250); the D-term problem appeared only at step 150. +- On failure, inject reference gradients block by block to find the first module over the limit. +- Extended to FSDP2: EP4 + FSDP2 all four snapshots PASS, worst-class median 1.08–1.39, all floors from B16. +- Tool bug: with `overlap_grad_reduce=False` gradients were reduced again before the EP4 dump (fixed). The export path needs A/A too: ‖g‖ bitwise equal, local shards bitwise equal to the dump, replica diff 0/0. +- Cross-process bitwise backward is not required. An occasional non-bitwise first step comes from the thread-local Inductor deterministic config (the autograd thread's first compile picks 512/1024 blocking by timing); fix in [BI kernels: determinism](bi-kernels.md#determinism-takes-effect-only-on-the-main-thread). + +#### Optimizer acceptance (optauth) + +Verify the optimizer consumes clipped `main_grad`, the update lands on master, and deployed bytes = requantize(master). Freeze the criterion first. Split failures into criterion modeling gaps (e.g. a check inapplicable without replicated parameters; the equal-magnitude first Adam step invalidating γ-type criteria while the per-coordinate bound holds), documented and accepted, and implementation errors, fixed. Cover all parameter group types: a first-step probe without an FP32 wd_mult=0 parameter misses weight decay applied to that group. + +#### Facts outside the two criteria + +- Frozen base with LoRA: base modules still propagate input gradients ([§4.2](#42-a-frozen-base-still-has-to-propagate-input-gradients)). +- The two-step RL loop (θ0→θ1→θ2) accepts the forward; three-arm SFT accepts training effect; report all three separately ([validation](../workflow/validation.md)). + +## VERL / Megatron integration + +Upstream config keys may differ from what was used here. Find the actual call paths of rollout dump, old-logprob and skip/repeat in the current checkout and record the commit; local-patch options are not upstream API. Process content (recipe and PR contract, review material, two-step acceptance, E2E targets) is in [validation.md](../workflow/validation.md#recipe-review-materials-and-two-step-acceptance-agreement); PR ownership in [Code and delivery §8](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). + +### Where the integration layer is + +Megatron's `experimental/lite/examples/verl/verl_mlite/` provides the verl engine adapter (`model_engine=mlite`, `pkg://verl_mlite.config`), installable into the unified image without reimplementing the engine in verl. Check engine registration, model factory, batch conversion, loss normalization and optimizer path; list reusable interfaces and gaps. Consistency uses mLite; the training baseline uses the existing recipe's default MCore path; do not mix them. + +### Capture location + +Save the DataProto (or equivalent) after rollout returns and before the first old-logprob; append trainer scores after. Record policy version, UID, input/response IDs, mask, position IDs, reward, route replay data if needed, and the UID -> packed-row mapping if reordered. Keep a pre-update checkpoint bound to the batch by content fingerprint; after strict load compare structure and keys and reject missing/unexpected keys. + +### Historical DS4 entry point + +DS4 used patched VERL options (`skip.rollout.*`, `DS4_SKIP_UPDATE=1`); confirm in code and logs that any such switch takes effect. + +### Minimal adapter contract + +The synchronous entry point takes batch, checkpoint, variant and output directory, then: + +1. Records actual policy, effective config, libraries and parallel layout. +2. Reads the same rollout, repacks per variant, keeps logical UIDs. +3. Runs the real old-logprob forward and exports per-token logprobs in UID order ([Saving and comparing](debugging-mismatch.md#compare-per-token-logprobs)). +4. Records no generation and no update, and mutable-buffer initialization. +5. Waits for workers/collectives, closes its resources, returns an exit code. + +It keeps native `no_grad`/`forward_only` semantics and replaces neither the [grad-enabled forward comparison](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract) nor update acceptance. A/B/C may share one long-lived worker once weight and state isolation is verified. The runner does not take over Ray/Slurm; the launcher blocks on remote jobs and reclaims on timeout. For injected or overlaid adapters, record the actual class and source file and assert the callsite was hit; `PYTHONPATH` can shadow the overlay. + +### Entropy when old-logprob is not recomputed (bypass) + +Bypass uses rollout logprob as old-logprob: `actor/ppo_kl` then measures rollout vs training forward exactly (0 under BI=1), and entropy comes from the grad-enabled update forward. + +- **Put metric-only quantities under `no_grad`.** With `entropy_coeff=0` the loss still holds `0 * entropy_loss`, so entropy's backward runs and saves a logits copy. +- **Chunked entropy in the verl Megatron engine cannot backpropagate.** With `entropy_from_logits_with_chunking=True`, each chunk view's backward overwrites the shared saved logits in place; the shared version counter then raises `one of the variables needed for gradient computation has been modified by an inplace operation`. Hit whenever chunking meets entropy. Fix: [verl-project/verl#8195](https://github.com/verl-project/verl/pull/8195); until merged, with `entropy_coeff=0` compute entropy under `no_grad` with the same chunking. +- **Re-verify gradients after changing forward computation.** Nemotron chunked logprob: BF16 bitwise equal to unchunked; FP32 max diff about 1e-8 at max gradient about 4. + +### Training extension + +Pairing a shared inference forward with a native VJP: check recompute, CP metadata lifetime, PP payload, parameter registration, through the standard training entry point. After a real update, export a new rollout and trainer scores; old bundles replay only with the old policy, and new-vs-old output differences are expected, not a consistency failure. Modules using vLLM layers need vLLM process groups. `param_offload=True` changes the quantized master's `data_ptr` and fails the pre-export check; use optimizer offload only. diff --git a/.agents/skills/train-inference-consistency/practices/code-and-delivery.md b/.agents/skills/train-inference-consistency/practices/code-and-delivery.md new file mode 100644 index 000000000000..aacb47c55bb4 --- /dev/null +++ b/.agents/skills/train-inference-consistency/practices/code-and-delivery.md @@ -0,0 +1,90 @@ +# Code, experiment snapshots and delivery + +Code ownership across three repos for cross-stack train-inference consistency delivery and "three PRs, one image" (§8), and PRs and cleanup (§9). Lifecycle, cold-start and model acceptance details are not repeated. + +## 8. Three-repo delivery: from experimental implementation to three PRs and one image + +Cross-stack RL delivery for a model is one vLLM PR, one Megatron PR, one verl PR, and one unified image built from the versions the three pin. Default targets are `vllm-project/vllm:main`, `NVIDIA/Megatron-LM:dev` and `verl-project/verl:main`; a fork may host the head but does not replace the main repo's target branch. This section covers change ownership and delivery checks; it does not replace the [ten-step flow](../workflow/README.md), the [numerical gates](../workflow/validation.md) or human recipe review. + +### Before changing code: give a change map first + +Fix the inference-side change locations in step 3 and the training side and integration interfaces in step 5; step 1 need not design every interface. Submit a short table of "feature → existing module/interface → proposed change location → reason for reuse or addition", covering at least model construction, forward/backward, quantization, parameter registration, checkpoint and the build entry. + +Read the implementation and callers in the selected checkout before choosing file names; the paths below are design references, not a requirement to copy DeepSeek's directories. Do not write a standalone engine first and discover at the end that the existing integration layer or the real trainer cannot use it. + +### vLLM: implement in the formal execution path, do not deliver a patch loader + +- Development environment and code sync follow [same-commit worktree, per-file sync and build verification](../guides/environment-provenance.md#boundary-between-overlay-and-native-wheel); development sync is not how the final image is installed and does not change the PR's target branch. +- Model-specific selection goes in the model construction/execution entry; changes to shared kernels, backends and metadata go in the existing owning module, stating which models and callers are affected. +- BI=0 keeps the reviewed native path; BI=1 explicitly selects the aligned implementation. Check the actual backend and fail early when a required BI capability is missing; never silently fall back while still claiming consistency. +- Kernel numerics, binding and build registration ship together; CUDA/C++ changes enter the normal build, and JIT changes such as CuTe DSL record the compiler and specialization. Do not submit only the Python call while depending on an uncommitted `.so`, source file or build script inside the experiment container. +- Verify the default path is unchanged and cover the modified shared interfaces; when only a certain model/runner/topology is supported, restrict it explicitly rather than implying other models support BI. + +Reference points in [vLLM #54955](https://github.com/vllm-project/vllm/pull/54955): the formal kernel integration in `csrc/batch_invariant/`, the DeepGEMM BI capability check, and the stable tie-break/prefill-decode semantics of sparse selection. These address operator numerics, backend authenticity and execution shape change respectively; they are more than a BI switch. + +### Megatron mLite: separate the model, the aligned operators and the runtime + +The train-inference consistency implementation uses mLite; the training baseline uses the default MCore path of the existing verl recipe, without the aligned implementation. Record the training baseline and the vLLM BI=0 inference performance baseline separately; they are not the same comparison. When target-model support has gaps, state them and get them reviewed; do not change the training scheme on your own. + +| Change | Preferred owner | Should not go in | +| --- | --- | --- | +| Model layer structure, parameters and checkpoint mapping | model/protocol/checkpoint modules of `experimental/lite/model//` | a temporary model implementation inside a verl recipe | +| forward aligned with vLLM, autograd/STE, model-specific quantization | that model's aligned implementation and primitive modules | miscellaneous top-level mLite scripts, global monkeypatch | +| Reusable training/communication/optimizer capabilities | existing shared primitive/runtime, changed only for a real interface gap | one model's recipe constants and diagnostic branches | +| verl engine adaptation, batch/loss bridging | existing `experimental/lite/examples/verl/verl_mlite/` interface | another equivalent engine inside verl | + +- Select native vs aligned implementations through the existing protocol/factory/explicit config; do not scatter environment checks across low-level forwards or replace instance methods/global functions at runtime. Small-scope branches can stay; do not copy the whole model for a tidy directory; reuse shared math and backward, and implement independently only primitives whose numerical semantics really differ. +- Do not assemble production operators with `inspect.getsource`, AST rewriting, or string replacement plus `exec`, and do not register models via `sitecustomize` or temporary `PYTHONPATH` hijacking. Import stable interfaces from the real dependency, and commit necessary changes to the owning module. +- Keep the parameter registration, state_dict/checkpoint, distributed reduction and recomputation contracts; do not replace the training master parameters with an inference cache the optimizer does not manage just to align forward. Backward is accepted separately per the [gradient contract](../guides/backward-validation.md#write-down-the-gradient-contract-first); it need not be bitwise identical with inference, and a kernel without autograd does not automatically need STE. + +[Megatron #7050](https://github.com/NVIDIA/Megatron-LM/pull/7050) separates DeepSeek-V4's `lite/` from `vllm/primitive/` and reuses the existing `verl_mlite`, illustrating these boundaries. `verl_mlite` exists in official Megatron dev, not only in a fork; still check the version this task pins and its actual capabilities. + +### verl: a thin recipe, not a second training framework + +What is delivered depends on the route: + +- **mLite route**: for non-observability changes, deliver ideally only one recipe script and one image build script, reusing the real trainer of the current version and the installed `verl_mlite`. +- **FSDP2 Day-0 route**: verl lacks the model, so the deliverable is the self-written model implementation in verl mainline, its Auto class registration and the recipe, in one verl PR; the model implementation itself is the deliverable and needs no separate "generic interface gap" pre-review. See [Training-side model support §3.3](../guides/training-forward.md#33-fsdp2-verls-hf-interface-path). + +On both routes, when there is a real generic interface gap, first give call-chain evidence, the interface to change and the scope of impact, **get human review before implementing**, and make the minimal fix in the owning repo; do not hide code in scripts or move it to the wrong repo to make up the file count. + +The two-file form of [verl #7688](https://github.com/verl-project/verl/pull/7688) and the `verl.trainer.main_ppo`/`model_engine=mlite` chain are the reference; its development mounts and overlay are not how the final image is delivered. For data, config differences and review rules, use the [DAPO reference and review materials](../workflow/validation.md#dapo-reference-and-review-materials); do not invent an "approximate DAPO" script that bypasses the trainer. + +### Step 10: from clean PRs to verifiable delivery + +1. **Three complete diffs.** For each repo, record the target owner/branch/base SHA agreed with the user and the PR head; check the merge-base, commit history and actual diff. Required model support must be in the PR or already merged into the target branch; another unmerged feature branch cannot be the base. Do not carry over a preview PR's history drift or unrelated deletions; pin cross-repo dependencies by commit. After changing the base or rebasing, re-review the diff; evidence from the old head does not transfer. +2. **One rebuildable image.** Build from three clean checkouts, pinning the base image digest, dependencies and necessary native/JIT build inputs, and install all runtime packages normally, including external adapter packages. Training and inference workers use the same immutable platform digest and the same Python/Torch and runtime dependencies; different nodes/processes are fine, but not two independent environments inside the image to dodge compatibility. It must run without host source mounts, startup-time patches or uncommitted overlays. +3. **Re-accept from a new container.** Check the actual worker's interpreter, imports, libraries and versions; verify the proxy, full-model zero-update and the real RL loop of the reviewed recipe per the existing gates; then record performance/quality. A successful build, offline disk replay and historical 0 diff do not replace verifying the final artifact; on a difference, shrink back to a fixed failing replay first instead of relaunching full training to locate it. +4. **One mapping.** Use the [report template](../workflow/report-template.md) to list the three PRs, image digest, recipe review version and acceptance receipts, marking each item pass/fail/inconclusive/not run. Update it after code cleanup, build or dependency changes; an unverified clean implementation does not inherit the overlay prototype's pass. + +When slimming the implementation PR, non-essential diagnostic tools, docs and experiment artifacts stay in the private repo/model workspace; necessary regression tests can go with the PR. Do not delete tests that exist on the target branch, and a slim PR still runs tests. Follow each target repo's contribution requirements, record the regressions and model evaluations actually run, and submit for human review; if the gates do not pass, keep the not-accepted status. + +### Reference versions, not proof of passing + +All three PRs were open when reviewed, and the first two were marked DO NOT MERGE: + +| Reference | Reviewed head SHA | +| --- | --- | +| vLLM #54955 | `4b5c627c23e6690b8987100d1cbf78e34ac14bf9` | +| Megatron #7050 | `327a238239448de6771beea185dc6902a8de3eaa` | +| verl #7688 | `640c5f0a317307e99636924804c40624242fa19b` | + +A PR body may describe a different historical pin; when citing the design, distinguish the reviewed head, the experiment version the body claims and the version measured now, and do not treat the reference as fully correct or merged. Its Triton implementation also does not override this workflow's constraints on added/replaced kernels. + +## 9. PRs and cleanup + +- Three-repo delivery follows the target branches in §8. +- The PR description records: minimal purpose, base and patch stack, related existing PRs/why it is not a duplicate, upstream/private PRs, RL-only content, test commands and results, model impact, rollback method, AI assistance and human review requirements. +- Split a large PR into small independently reviewable PRs and prove equivalence before and after the split with the tree hash; open PRs only from your own fork. +- Resolve conflicts on the head branch of the PR the user specified; do not open a similarly named new PR to pass the original off as fixed. +- Cross-repo order is usually kernel/runtime → framework integration → recipe/pins → final image acceptance, with each later stage locking the actual revisions of the earlier ones. + +After a patch is merged and a new image is built, archive the worktree, delete the temporary overlay, and keep the PR, commit, manifest and verification report. + +## Minimal changes, contribution and delivery + +Change the config along the entry script the user specified, so multiple temporary launchers do not drift apart. End with one launch entry that actually runs, with defaults, necessary overrides and image/code pins traceable; exports that exist only in chat are not delivery. + +A fix lives with the state/feature owner: fix vLLM backend/selector problems in vLLM, not with a global monkeypatch in the training framework; a numerically matched forward must keep the real backward. Call existing reusable functions instead of building synonymous logprob metrics and compatibility switches. + +An authorized independent audit does not transfer responsibility to the auditor. Before declaring completion, check: are all of the user's goals covered? Do the results belong to the final code/image? Does the wording hide any untested item? If the evidence is insufficient, write down the gap; do not wrap it in a percentage or a success marker. diff --git a/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md b/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md new file mode 100644 index 000000000000..224e652d9056 --- /dev/null +++ b/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md @@ -0,0 +1,179 @@ +# Evidence and experiment rules + +Rules for what counts as evidence, how experiments are designed and recorded, and when old evidence stops holding. Mechanisms live in the guides. + +## 1. Evidence labels + +PASS / FAIL / INVALID (did not enter the target path) / VACUOUS (condition had zero hits) / SUPERSEDED / HISTORICAL. These label evidence, not lifecycle states. The tables in [validation.md](../workflow/validation.md) and the [report template](../workflow/report-template.md) have only `pass / fail / inconclusive / not run`. PASS→`pass`, FAIL→`fail`; INVALID and VACUOUS become `inconclusive` with the reason in the notes; SUPERSEDED and HISTORICAL stay out of the current table and live only in the [candidate evidence ledger](#6-candidate-evidence-ledger) and historical reports. One SUPERSEDED piece of evidence makes the gate it covered stale. A successful run, train-inference consistency and a held-out score gain are different conclusions and cannot substitute for each other. + +## 2. Evidence that must not be confused + +- Enabling `VLLM_BATCH_INVARIANT` is not cross-framework bitwise alignment; the same text is not the same logprob. +- Zero difference in probability space, zero KL or identical printed numbers do not replace comparing valid logprobs in the original dtype. +- Replaying fixed inputs is a localization tool; statistics from old trajectories are not an unbiased KL under the current policy. +- The same function name, generated code or environment variables do not prove the same actual kernel, launch config and shared libraries. +- Output replacement, forced routing, disabling Graph and bypassing bad packing are valid causal experiments if you state what they change. +- Similar quality scores are not identical per-problem results; same-problem quality, bitwise logprob and 8K/1K performance are accounted separately. When the acceptance script fails, check the producer receipt first; do not resample or declare a pass automatically. +- Verify quantization rules on the deployed bytes, not dequantized weights ([precision pitfall P2](../guides/precision-quant.md#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). QAT/low-precision training only narrows the weight gap; 0 diff comes only from training calling the same kernel as inference; see [How much QAT can shrink the mismatch](../guides/precision-quant.md#4-how-much-mismatch-qat-can-remove). +- Do not manufacture a pass by changing tokens, weights, model geometry, precision or tolerance. If the comparison baseline really must change, start a new experiment and re-verify quality. +- Do not carry DS4's EP/MoE config over to dense Qwen; do not manufacture a performance gain by weakening the BI=0 config, locally disguising BI as off, or removing existing correctness patches. +- Source tests passing is not the final image passing. Development overlay results belong only to their experiment; the delivered image digest, installed source and build artifacts must match the acceptance object, and the evidence chain is rebuilt after repackaging. + +## 3. Inferences no longer carried forward + +Historical numbers and criteria stay with the run that produced them; a new task does not inherit a past run's model, dataset, topology, performance statistics method or pass status. Lessons from retracted or narrowed claims: + +- A number in a summary is not evidence: find the raw numerator and denominator for the same policy/workload/config, or retract the assertion. +- Execution completion is not numerical acceptance: finished steps, a started job, an existing adapter file or finite nonzero gradients do not show the target class ran or the comparison passed. Record the instantiated class and its import source, and accept non-empty original-dtype token comparison, module VJP against an independent reference, actual update, shutdown and long-term stability separately. +- Evidence does not transfer across candidates: a source pin, the same Docker tag, an earlier proxy, an experiment stack or a direct operator call does not prove the current candidate, image or compiled module dispatch. Verify branch, dirty patch, wheel, `.so`, dependencies and digest separately, and re-establish entry-gate evidence for the current candidate. +- A fixed split count or a pass on one path does not cover other batches, prefill/decode, the first round after resume or cross-node communication; each needs its own evidence. +- When an acceptance wrapper fails, find the failing layer before resampling or rerunning; a control-plane parsing error does not discard sealed numerical evidence. + +Each new case records only "contract → minimal counterexample → discriminating experiment → root cause/minimal fix → regression on the original reproducer → scope/source", with the small raw receipt for disputed numbers. Merge status, recommended dependency versions and best-performing config are volatile; look them up in the target repo. + +## 4. Experiment contract + +### Define the objects of equality first + +State the checkpoint/adapter/policy version, tokenizer and chat template, raw token IDs, position IDs, attention/response mask, padding/packing, and the logits processing on both training and inference sides. + +For response token `y[t]`, compare `log p(y[t] | prompt, y[:t])`. Make explicit the prediction position of the first response token, whether EOS counts, and the exclusion rules for truncation and padding. Do not align logits rows of different positions. + +Distinguish raw logits, raw logprobs, and logprobs after temperature/repetition penalty/top-k/top-p, and record which one the interface returns. The same seed does not guarantee the same scheduling or RNG consumption: freeze the tokens instead of resampling. + +### Record these goals separately + +| Goal | Valid evidence | Cannot infer | +| --- | --- | --- | +| Single-side repeatability | A/A with the same input and state | Consistency across batches or processes | +| batch/layout invariance | Same semantic sample identical under different positions, neighbors and packing | Training and inference are the same | +| prefill/decode parity | Two execution paths with the same forced tokens | backward or the training update is correct | +| Cross-stack forward parity | Real training forward against a frozen rollout | Full-vocabulary logits match, if only sampled tokens are compared | +| Effective training loop | Non-degenerate gradients, parameter change, post-update parity | Convergence or long-term stability | + +When the user requires bitwise, compare original-dtype bytes and record mismatching/valid token counts; do not reach a pass through a float cast, probability-space rounding or a tolerance loosened afterwards. If the goal is an error budget, fix the metric and threshold in advance and measure each side's own noise; generic numerical alignment need not be forced to bitwise. + +A finite logprob is necessary; an empty mask, all-zero placeholders, missed capture, NaN or the wrong weight version cannot pass. A zero may be legitimate; tell it from a placeholder by hook count, score source and a negative control. + +### Recipe and coverage record + +When building a train-inference recipe, also fix these inputs, recovering missing items from the existing config first: + +| Item | What to record | +| --- | --- | +| Workload | Training context length, rollout input/output length distribution, concurrency, packing, training token batch | +| Hardware | GPU architecture/count, GPU memory, interconnect and target TP/CP/EP/PP/DP, trainer and rollout separately | +| Numerics | master/deployed weight, activation, KV, accumulation and gradient dtype; quantization granularity, scale rule and update method | +| Scope | Raw sampled logprob or full-vocabulary logits, applicable hardware/software, Graph/cache path | +| Budget | Throughput, peak GPU memory, quality and gradient criteria, experiment count and stop condition; with an unknown threshold, report the measurement and do not declare a pass | + +Record the measured baseline of the current config, then build a strong performance baseline per the main flow, without exhaustive tuning first. Rust Bench 8192/1024, ignore EOS, c=1–64 is a fixed comparison and does not replace the target context and natural-length workload; a user-specified workload governs, with the difference reported. + +A proxy records "full-model function → proxy position → test case → uncovered items" and keeps the hidden/head/expert/quantization geometry and layout that affect kernel selection. Reuse existing modules and tests, but after changing model, dtype or parallelism re-check the contract; changing the config does not carry over the old acceptance. + +### Environment fingerprint + +Keep the code commit, dirty patch hash, untracked runtime files, image digest, GPU architecture, Torch/Triton/CUDA, actual import path, and native `.so` path and hash. Record TP/PP/CP/EP/DP for trainer and rollout, dtype/quantization/scale, kernel backend, and Graph/compile/cache and autotune config. + +Read the resolved config and actual modules/libraries inside the worker, not launcher settings. For weights, record content version and update status of derived weights/scales, not just the path. Keep immutable known-good and known-bad artifacts before diagnosing. + +### Scope of conclusions + +For each result, record model scale, valid sample count, independent token count, lengths and boundaries, process count, GPU/parallel config and number of updates. Four samples repeated into batch32 are still four independent samples; mixed scheduling a small model never triggers cannot cover the full model. + +Record machine budget, maximum experiment count and stop condition. GPU or system faults give `inconclusive`; a failed numerical comparison gives `fail`; `pass` for a scope requires every declared acceptance item. + +### Experiment discipline + +Write these nine rules into the record before comparing two arms; if any is unmet, the conclusion is "pending re-check". + +| Rule | How to do it | Mechanism and example | +| --- | --- | --- | +| One image per recipe | Both arms use the same image and recipe, differing only in the recipe switch; inject changes with an overlay and record the hash | [One recipe, one image](../guides/recipe-and-acceptance.md#one-recipe-one-image-deliver-true-on-policy-and-the-regular-baseline-with-the-bi-switch) | +| One seed per config | One seed by default; add more only when the user asks for variance, reported separately | Extra seeds spend budget on variance; most problems are systematic | +| Small-workload iteration | Iterate with small batch, short outputs and few steps; run the recipe default workload only for final acceptance | The default workload allows only one or two iterations a day | +| Freeze the criterion first and plant an error | Fix comparison window and tolerance in writing before the run; confirm with a known error that it can judge FAIL; never adjust the threshold after a failure | A threshold chosen afterwards can always pass; [paired acceptance](../guides/recipe-and-acceptance.md#paired-acceptance) | +| No conclusion from a single e2e run | Measure noise first; on the same hardware and image, report the median of multiple pairs and single pairs separately | [Performance and profiling](../guides/performance.md) | +| Same fabric | Before comparing communication bottlenecks, confirm both arms use the same transport and record the relevant environment variables | [Same-fabric fairness](../guides/debugging-mismatch.md#same-fabric-fairness) | +| Same-node comparison | Run paired arms sequentially on the same node or node group, with node names recorded | Memory fragmentation, resident processes and interconnect topology differ between nodes | +| JIT pollution: check `tl_after_bench=0` | After the run, confirm no new compilation after the bench; results with recompilation are void | [Performance and profiling](../guides/performance.md) | +| Take metrics by key name | Read numbers from raw metrics (jsonl/W&B) by full key name and cite the key in the report; never by column position in log lines | An `actor/grad_norm` column once read as entropy reversed a conclusion | + +List confounders before comparing; see [paired acceptance](../guides/recipe-and-acceptance.md#paired-acceptance). + +## 5. Execution state + +For every comparison, record the actual execution state, and keep results from different execution states as separate evidence. + +### Fields that must be recorded + +- Resolved Graph mode, capture sizes, whether the step is prefill/decode/mixed, and kernel call evidence; a successful capture does not mean every step uses the same graph, nor numerics identical to eager. +- Whether the prefix cache hit, and whether it was disabled/isolated during diagnosis; passing with the cache off does not mean the native lifecycle is fixed. +- Weight lifecycle phase: cold load, sleep/wake, Graph replay, reload in a new process, one row each. +- Whether chunk boundary state and partial input buffers of stateful layers such as Mamba/SSM are saved. + +### Keep separate evidence + +Eager, compiled with Graph off, and production Graph each keep their own artifacts; hooks run at capture, not replay, so a capture dump does not represent replay. List separately: diagnostic runtime override vs formal source candidate, Graph vs no-Graph, forward-only vs real `train()`+autograd. A new process passing does not prove the old process's cache/Graph is fixed; two fixes that each pass go through the gate again when stacked. Once CUDA state is corrupted, end that experiment and collect no more numerical evidence from it. + +### Where the mechanisms and methods are + +| Question | Where to go | +| --- | --- | +| Graph capture vs replay, the Graph/eager 2×2 comparison with fixed composition, which bucket the filler falls into | [BI kernel appendix](../guides/bi-kernels.md) (cudagraph row) | +| prefix cache semantics, Mamba chunk boundary state, bitwise comparison with "cache off as the gold standard" | [rollout-engine](../guides/rollout-engine.md) | +| Localizing CUDA asynchronous errors | [debugging-mismatch](../guides/debugging-mismatch.md#isolate-the-side-first) | + +## 6. Candidate evidence ledger + +Goal: keep an old pass from becoming a new pass. + +When an experiment changes model scale, quantization dispatch, attention, MoE, Graph or policy weights, create one ledger row per **candidate execution closure**, then decide what to rerun. + +### What one row of evidence is bound to + +| Field | What to record and check | +| --- | --- | +| Object | Original layer numbers and real adjacency, parameter/quantization geometry, tokenizer, checkpoint content, policy version; a proxy's manually spliced boundaries listed separately | +| Execution on both sides | Trainer/server image digest, source and overlay hashes, actually imported classes and files, resolved config, native/JIT artifacts | +| Path | `train()`/autograd/whether backward runs, actual dispatch of attention/norm/linear/MoE, compile/Graph/cache, TP/EP/PP layout | +| Workload | Content digest of prompt/response tokens, valid mask, length distribution, packing, concurrency, reward/advantage and actual accepted sample count | +| Criterion | Original dtype, finite valid token count, bitwise mismatch, sample count, first difference, and the independent states of gradient/update/normal exit | +| Artifacts | Inputs, score, worker/rank receipt, command, time, parent version, hash and `pass`/`fail`/`inconclusive`/`not run` | + +`policy_manifest_sha256`, the exported `policy_content_id` and the rollout receipt hash answer three separate questions, **manifest version, weight content, run artifact**, and are not comparable with each other. If the exported content ID is the hash of the export receipt, follow the receipt to check shard hashes against the files the server actually loaded; one equal version string is not enough. + +### Entry gates from cheap to expensive + +1. **Module.** Use a real triggering input and a non-unit scale/real layout; A/A first, then the candidate's forward against an independent reference; for gradients, check dX/dW/other differentiable quantities separately with the same dY. A pass covers only the shapes, arithmetic and VJP tested. +2. **Minimal real-adjacency proxy.** Cover the actual order of layer types, not just the operator set. With a fixed 8192+1024 tokens, check inference BI/PD first, then per the [dual-path contract](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract) verify native old-logprob and the grad-enabled actor forward separately; for a fixed-token comparison only, do not add GSM8K/RL to the proxy. Record source closure, actual class, Graph, non-empty valid tokens and original-dtype mismatch. When any item changes, re-examine whether this step still transfers. +3. **Full-model zero-update.** Generate a real-workload rollout with the full model on the **current** serving candidate, saving tokens and raw scores; the actor replays the same policy and history before the optimizer update. Check the earliest failing sample, then expand to the agreed sample size; if unequal, stop here and narrow the boundary instead of working around it with five-step RL or a quality evaluation. +4. **Per-version update.** For each θk, compare on-policy first, then check reward/advantage, backward, master parameters and the optimizer step, then sample θk+1; "the last training step exited" is not "the last version's comparison passed". +5. **Independent acceptance.** Module gradients, whole-model gradient coverage, real RL, quality, performance, clean exit and the final image each get their own result. Numbers from a previous candidate or another execution mode stay as historical comparison and never fill the current candidate's pass cell. + +A fixed-weight serving overlay without RL lifecycle operations such as sleep/wake is only a serving screening candidate, even if BI/PD and 8K/1K paired performance pass; once the lifecycle is restored, rerun the affected proxy, zero-update and post-update comparison. Do not splice historical throughput and another candidate's RL pass into one "completed" version. + +If a phase fails midway, classify the failure as numerical, control, environment or resource plane; keep the sealed rollout and producer hash. A control-plane acceptance bug is handled by a separately versioned recovery acceptor reading the original artifacts; never edit the old receipt, or resample and call it the same trajectory. + +### When evidence must be invalidated and rerun + +| Change | Minimum regression; old conclusions that cannot be carried over | +| --- | --- | +| Proxy's original layer numbers/adjacency or hidden, head, expert geometry change | Coverage table, BI/PD, training forward; the old proxy's 0 diff no longer counts for promotion | +| Quantization algorithm, scale, actual dispatch or JIT/compiler change | Real module input/quantized bytes, proxy, full-model zero-update; a direct-call unit test does not replace the compiled module | +| Attention/MoE operator or its training adapter changes | Real class and source, same-input operator, relevant VJP, proxy and full-model zero-update; a pass on a substitute backend is diagnostic only | +| Graph/cache/packing/parallel layout changes | The corresponding BI/PD, replay and state lifecycle; an eager pass does not cover FULL Graph | +| Optimizer or master→deployment export changes | Gradient and one-step update, export byte hash, post-update comparison; a θ0 forward pass is not enough | +| Final image/overlay/native build changes | Verify the closure in the worker and rerun affected numerics and performance; development-mount evidence does not transfer | + +If you can prove a change does not affect a gate, write the dependency closure and reason in the ledger to skip a mechanical rerun; equivalence cannot be decided from a file name, PR description, same tag or "looks like a cleanup". + +### Separate cells for performance and quality + +Tune BI=0 well under the agreed weight/KV precision and quality requirements, record the real dense GEMM, MoE expert, all-to-all and attention backends, and freeze the denominator after human audit. Pair BI=1 and BI=0 on the same full workload, nodes and candidate closure; Rust Bench 8192/1024, ignore EOS, c=1/8/32/64 is one fixed throughput workload, not natural DAPO lengths or GSM8K quality. Check effective tokens/error count first, then compute the loss per the [pre-reviewed statistical protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); do not reach the budget by slowing BI=0, changing weight precision or changing prefix cache/Graph parameters. + +Quality needs the same problems, prompt, chat template, seed, generation limit, scorer and no-retry policy; report per-problem pairing and the absolute score difference. A quality pass does not transfer to a candidate whose numerical path changed. If the user changes the threshold, the current one governs; keep the old threshold and change date for historical reports. + +## 7. Failure records + +Keep failure records as they are; append new results, with no overwriting, deletion or editing. For W&B, use a separate project per comparison and read back through the API to check no step is missing. For freezing the criterion, planting a known error and "no conclusion from a single e2e run", see [§4 Experiment discipline](#experiment-discipline). diff --git a/.agents/skills/train-inference-consistency/workflow/README.md b/.agents/skills/train-inference-consistency/workflow/README.md new file mode 100644 index 000000000000..44e7b8419456 --- /dev/null +++ b/.agents/skills/train-inference-consistency/workflow/README.md @@ -0,0 +1,122 @@ +# Train-inference consistency workflow: C0–C10 + +Turn a train-inference difference into a repeatable experiment with fixed input and fixed weights, locate the first divergence, fix or optimize, then return to the real training loop. Keep the user's model, precision, parallelism and consistency target. Switches, thresholds, passing scope and implementation paths from historical cases, reference PRs and implementations already in the training framework (such as an existing true-on-policy path) are experience, not defaults for the current model, dataset or hardware; measure their numerics and training-step time before reusing them. + +Forward consistency, backward conformance to the gradient contract, training quality and performance are accepted separately; none replaces another. User config comes first. Recover missing information from the source and resolved config, and ask only for what cannot be recovered and affects the approach. Choose the low-precision recipe by hardware and existing paths, not by a default precision, model or framework. Changing the model or precision requires rechecking coverage and acceptance evidence. + +Work in one directory per task, with state changes in an event log and artifacts carrying provenance sidecars. Details are in [validation.md](validation.md) (entry gate table, dual-path contract, tiered criteria, SFT three-way control, two-step RL sequence) and [numerical contracts](#numerical-contracts-and-shared-forward). + +Archive materials as you go: each run ID links config/scripts, execution identity, raw results and conclusion; the step index states the current candidate, failed/excluded items, unverified items and the human-confirmed version. C1 keeps the full tuning comparison table; C10 references the reviewed baseline ID. Layout: [archive experiment materials by step](../guides/debugging-mismatch.md#archive-experiment-material-by-step). + +## Gate table + +Iterate as "single operator / fixed replay -> proxy -> staged full model"; do not rerun every gate each time, and do not skip the human stop points at C1 and C9. Inference-only tasks apply only the relevant gates and state that the training loop is not accepted. + +| Gate | Input | Action | Pass condition | Artifacts | See | +| --- | --- | --- | --- | --- | --- | +| **C0 Module inventory and route** | checkpoint, vLLM implementation, target (BI=1 / BI=0), hardware | Per module: inference kernel, whether training implements the same math, quantization rounding, discrete choices; choose FSDP2 / Megatron-Bridge / mLite | Inventory covers all modules; route fits "BI=1 needed? how many backwards missing? what parallelism?" | Module inventory, route and rationale | [Training-side model support](../guides/training-forward.md) | +| **C1 Recipe and strong baseline** | Model, target context length, hardware | Decide weight/activation/KV precision, quantization rules, runner, parallelism, execution config; decide the training master start (BF16 original or dequantized weights; do θ0 bytes equal the official checkpoint); agree scoring scope, consistency standard, performance budget (inference throughput, and limits for training forward and forward+backward against the native training path on the main topologies; a human decides which parallelism to enable and which topologies are main); tune BI=0 with **vLLM Rust Bench** at 8192 in / 1024 out / ignore EOS / c=1–64; for MoE, measure MoE / all-to-all backend combinations | Baseline = the formal recipe's BI=0 config (including R3, PC, gpu_mem), not the serving-only optimum; multi-round measurements and backends recorded; **freeze the candidate and submit for human audit** | Tuning table, resolved config, image digest, audit materials | [Precision, quantization and QAT](../guides/precision-quant.md), [Recipe and evaluation](../guides/recipe-and-acceptance.md), [Performance](../guides/performance.md) | +| **C2 Proxy model** | Full model structure | About 4 layers covering all operator types, needed adjacent-layer interactions and state paths; real hidden/head/expert/quantization geometry and needed parallelism constraints; coverage table "full-model feature -> proxy layer/module -> test case" | Uncovered items listed; the proxy is the daily development environment | Coverage table, proxy config | [Training-side model support](../guides/training-forward.md) | +| **C3 Inference BI/PD** | Proxy, fixed weights and tokens | Verify batch and prefill/decode invariance separately, covering chunks, mixed scheduling, cache state; consider kernel suitability for training / large token counts. On failure, coarse-scan layer boundaries, bisect, capture identical operator inputs, fix kernel/metadata/layout/state | Bitwise identical in all three cases under fixed input; **proves only consistency inside inference** | Comparison matrix, minimal failing replay | [BI kernel appendix table](../guides/bi-kernels.md), [Fault localization](../guides/debugging-mismatch.md); before changing inference read [implementation ownership](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image); sync source per [runtime provenance](../guides/environment-provenance.md) (same-commit worktree + per-file hash) | +| **C4 Structural performance** | Proxy, real operator shapes | Profile prefill/decode (and training forward and forward+backward once wired): launches vs tokens, CPU gaps, sync, repeated quantization/copies/allocation, abnormal workspace; target CUDA Graph captures/replays correctly | No per-token Python loop or many-small-kernels defects (autoregressive decode steps do not count); consistency retest passes after the fix | Profile, defect list | [Performance](../guides/performance.md) | +| **C5 Training forward↔inference** | Proxy, fixed weights, input/output tokens, state | Reuse inference primitives and quantization, layout, mask, position, routing, reduction, logprob semantics; keep the training parameter container and native entry point; same-packing A/A and different-packing A/C first; then compare native old-logprob recomputation and the grad-enabled actor forward separately per the [dual-path contract](validation.md#old-logprob-and-training-forward-dual-path-contract) | Both paths: valid logprob mismatch=0 in original dtype; record non-empty valid token count and execution identity. Inference BI/PD, single operator or diagnostic substitute paths do not replace this gate | Dual-path receipt | [Training-side model support](../guides/training-forward.md) | +| **C6 backward / STE** | Forward implementation from C5 | Reuse existing backward, autograd, communication, recompute; where the gradient passes quantization rounding, decide the STE surrogate (identity/clipping, master or dequantized value, scale rule). A kernel without autograd does not force an approximate gradient; no AST rewriting or monkeypatching in the production path | Every module has a "reuse / glue / newly written" verdict and rationale, and its backward uses only what the forward saved; the training step is within budget on the main topologies fixed at C1 | Module-level backward inventory, training-step time table for the main topologies | [Backward validation](../guides/backward-validation.md), [STE contract](../guides/precision-quant.md) | +| **C7 Independent gradient validation / SFT control** | C6 implementation, same input, weights, saved state, same-dtype upstream gradient | Compare against a same-contract reference that does not call the backward under test: dX, dW, other trainable gradients, one optimizer step; fix tolerances first; report bitwise and tolerance separately. When required, compare production, an independent surrogate with the same visible forward, and native training per the [three-way SFT control](../guides/backward-validation.md#layer-3-proxy-sft-three-arm-comparison) | Independent comparison passes; without SFT, mark "quality not covered", a loss decrease is not quality. Forward 0 diff or finite gradients do not replace this gate | Gradient audit, three-way SFT results | Same as above | +| **C8 Proxy RL loop** | Proxy, real framework entry point | Follow the [exact sequence](validation.md#two-step-rl-sequence): θ0 compare -> update to θ1 -> θ1 compare -> update to θ2 -> θ2 compare; check finite, meaningful gradients, actual master change, normal exit | No cross-version comparison; a forward skipped by zero advantage does not count; proxy limited to logprob alignment: mark "loop not verified"; proxy asked to verify forward only: full-model zero-update comparison first, then the training gates; never start multi-step RL directly | Two-step RL receipt | [Two-step RL sequence](validation.md#two-step-rl-sequence) | +| **C9 Recipe review and full-model entry gate** | Launch script, image digest, code version, full resolved config, differences from the reference recipe (verl #7688 etc.), proxy smoke | **Human review; start the full model only after the recipe version is confirmed** (data, length, batch, parallelism, precision/quantization, optimizer); fill in the [entry gate table](validation.md#full-model-entry-gate-pass-proxy-train-inference-forward-first); then add layers, length, concurrency, parallelism one at a time and verify mixed scheduling, cache/Graph, multi-step loop; with checkpoint restore in scope, check alignment after the first restore | No escalation while C5 is fail / inconclusive / without artifact; two-step acceptance is the formal recipe with only fewer steps, same sequence/batch/data/algorithm; on full-model failure, shrink to a minimal replay first | Review materials, entry gate record, two-step results | [Checkpoint, resume and memory](../guides/checkpoint-resume.md), [recipe review materials and two-step agreement](validation.md#recipe-review-materials-and-two-step-acceptance-agreement) | +| **C10 Optimization and final E2E** | Full-model consistency passed, delivery PRs, reviewed recipe | E2E profile -> BI=1 hotspots (shape, dtype, layout, frequency, time share) -> representative micro bench -> correctness and performance on all target shapes -> E2E. **Triton is forbidden**; use CuTe DSL or better, CUDA/CUTLASS, or an existing implementation. Each round regresses inference, training forward and gradients by impact | On the delivery image through the real trainer entry: θ0 compare, update, θ1 compare, update, θ2 compare; 0 diff = equal original-dtype logprob for the same weight version, token history and valid mask; 8K/1K Rust Bench BI=0/BI=1 per the [pairing protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); no added long runs or gates | Delivery manifest, pairing results, loop receipt | [Performance](../guides/performance.md), [Paired acceptance](../guides/recipe-and-acceptance.md), [Delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image) | + +The consistency implementation uses mLite by default: it separates model, operators and runtime, so the inference forward can be reused, backward/STE wired in independently, and the existing `verl_mlite` interface and MCore distributed optimizer wrapper kept. The training baseline follows the default MCore path of the existing verl recipe without the alignment implementation, and is accounted for separately from C1's vLLM BI=0 inference baseline. **mLite does not inherently guarantee 0 diff or better performance**; verify the target model's parallelism, recompute and checkpoint support with the actual recipe. Both arms share one recipe and one image, and `VLLM_BATCH_INVARIANT` selects the arm (see [bi-switch-recipe](../guides/recipe-and-acceptance.md)); the delivery contract is in [validation.md](validation.md#recipe-and-pr-delivery-contract). + +## Run the cheapest failure criterion first + +1. host metadata / shape arithmetic: minimal counterexample of a pure function and boundary parameterization. +2. Single operator: same input, real dtype/layout, non-finite detection, repeat / cross-process / Graph. +3. frozen batch replay: the original failing A/C and A/A, full valid-token comparison. +4. engine: batch composition, prefill/decode, chunk boundaries, cache/Graph, actual parallelism scope. +5. Real training: non-degenerate reward/advantage, finite non-zero gradients, optimizer step, parameter change, frozen-token alignment after the update. +6. Production acceptance: full model, agreed multiple steps, restore lifecycle, quality evaluation and paired performance. + +If a second run adds no localization information, change the probe. Not reproduced, empty artifacts and an unexecuted path are not passes. When the failure is in backward, replaying the forward is not enough: save checkpoint, batch, RNG and state and move to a minimal backward experiment. Details are in [validation.md](validation.md). + +## RL delivery gates + +Training and inference use **the same immutable image digest and the same Python/Torch and runtime dependencies**. They may run on separate nodes and processes, but two environments inside one image cannot be used to dodge compatibility. Complete at least two real updates through the target RL framework's actual entry point and verify each round: original-dtype logprob 0 diff before the update, an actual parameter update, and 0 diff for the new policy after it. Saved rollouts, offline replay, separate images and cold start are diagnostics only. This is a delivery agreement, not a technical claim about all RL systems. + +Deliver one vLLM PR, one Megatron PR, one verl PR and one unified image, targeting `vllm-project/vllm:main`, `NVIDIA/Megatron-LM:dev` and `verl-project/verl:main`. Non-observability changes in verl should be limited to the recipe and the image script where possible. Follow [from clean PR to delivery](../practices/code-and-delivery.md#step-10-from-clean-prs-to-verifiable-delivery): check the target branch and full diff, pin cross-repository commits, build the image from clean source, and fill in the [delivery manifest](report-template.md#delivery-manifest). + +## Baseline and consistency run as two parallel tracks + +- **Baseline track**: C1's BI=0 tuning, backend combinations and multi-round performance evidence. Do not sacrifice default performance for consistency in advance; after reasonable tuning, freeze the candidate and wait for human audit. +- **Consistency track**: do not wait for the baseline; advance through C2–C8 on the proxy. C9 still requires the entry gate; parallel progress is not a reason to go to the full model early. +- Share the model/precision/hardware/image contract; isolate code, environment and artifacts; keep consistency patches out of BI=0. While the baseline is not final, mark the scope consistency covers and re-verify after config changes. +- Do not add nodes beyond the budget. On a single node, GPU experiments time-share; during a formal benchmark, run no GPU, heavy CPU or I/O work from the other track. The final performance gap is accepted against the human-confirmed strong baseline with full consistency evidence. + +## Two human stop points + +**After C1: baseline audit.** Once measurements support "reasonably tuned", freeze the candidate, report and wait; do not keep chasing the performance limit. Materials: model/weight version, hardware and node count, image digest / build time, code and dependency versions, reproducible launch and Rust Bench commands; weight/activation/KV precision, runner, parallelism and scheduling, Graph/cache state, actual operator backends, and a statement that alignment patches do not change BI=0; configs tried and why, valid tokens per concurrency, warmup, repeat count, throughput/latency and variation, failed items and uncovered scope; the config proposed for locking, the [statistical measurement basis](../guides/recipe-and-acceptance.md#pairing-and-computation), repeat count, per-concurrency criteria and budget, and open questions for humans. An unconfirmed candidate baseline must not be the final denominator. Re-audit after changing model/precision/hardware/image; while waiting, hand over or release nodes as needed. + +**Before C9: recipe review.** See the gate table. Changes to workload or numerical/update semantics are reviewed again; pure path or log fixes are not. + +## Find by problem + +| Current problem | Read first | +| --- | --- | +| Setting the consistency target or evidence basis | [Experiment contract](../practices/evidence-and-experiments.md#4-experiment-contract) | +| Reproducible input not yet saved | [artifacts](../guides/debugging-mismatch.md), [rollout-engine](../guides/rollout-engine.md) | +| rollout and old logprob differ: side unknown / stable failing sample / identical operator inputs still differ | [Fault localization](../guides/debugging-mismatch.md), [numerical contracts](#numerical-contracts-and-shared-forward) | +| Results change with packing, CP/TP/EP | [Layout and length](../guides/training-forward.md) | +| Fails only under Graph, compile or cache | [Execution state](../practices/evidence-and-experiments.md#5-execution-state) | +| Wiring in backward/STE, comparing gradients, SFT | [Reuse backward and backward validation (three layers)](../guides/backward-validation.md) | +| Going to the full model, shrinking a full-model failure, declaring completion | [validation.md](validation.md) | +| Candidate/backend/image changed: is old evidence still valid | [Candidate evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger), [inferences no longer carried over](../practices/evidence-and-experiments.md#3-inferences-no-longer-carried-forward) | +| Aligned; recovering performance, profiling, ablation, BI kernels | [Performance](../guides/performance.md), [BI kernel appendix table](../guides/bi-kernels.md) | +| Separating quality, fixed-length Rust Bench and normal exit | [Evaluation and paired acceptance](../guides/recipe-and-acceptance.md) | +| Image, source, wheel, submodule, overlay versions confused | [Runtime provenance](../guides/environment-provenance.md) | +| Three-repo ownership, PRs, unified image | [Three-repository delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image) | +| GB200 devices, NVSHMEM, DeepEP, cross-node initialization failure | [Communication fault ladder](../guides/debugging-mismatch.md) | +| checkpoint save/load OOM, `unexpected pos`, host memory rising, first resume misaligned | [Checkpoint, resume and memory](../guides/checkpoint-resume.md) | +| BI=0 baseline arm, two modes in the same recipe | [bi-switch-recipe](../guides/recipe-and-acceptance.md) | +| Requantization rules, QAT selection, activation fake quantization | [Precision, quantization and QAT](../guides/precision-quant.md) | +| Many repeated samples within a rollout group | [rollout-engine](../guides/rollout-engine.md#1-sampling-seed-collision) | +| Different seeds give the same samples (greedy n repeats are identical; SGLang deterministic mode without a seed always uses `sampling_seed=42`) | [rollout-engine §1](../guides/rollout-engine.md#1-sampling-seed-collision) | + +The ten kinds of evidence that must not be confused are in [evidence and experiment rules §2](../practices/evidence-and-experiments.md#2-evidence-that-must-not-be-confused). Writing comparisons and replays: [save and compare](../guides/debugging-mismatch.md#compare-per-token-logprobs). Record experiments and report with the [report template](report-template.md). This workflow grants no extra GPU, job stop, image publish or cross-repository write permission. + +## Numerical contracts and shared forward + +### Check the arithmetic that actually executes + +The same formula does not guarantee the same floating-point result. Check accumulation dtype, cast position, FMA, reduction tree, chunking/split-K, approximate activations, quantization scale/rounding, epilogue and the input dtype of the final logprob. The same source can also produce different execution configs through autotune, compile specialization or a different GPU. + +A shared kernel still needs its caller, arguments, layout and compile boundary checked; a direct call and a torch.compile fused path may differ. If needed, keep the same forward behind an existing opaque custom-op boundary; do not introduce unproven arithmetic rewrites to imitate the output. + +Quantization execution checks (scale taken from the actual runner, quantization dispatch, quantizer outside the `nn.Module` tree, training vs deployment quantization grid) are in the [precision, quantization and QAT guide](../guides/precision-quant.md#training-side-numerics), the grid in [P8](../guides/precision-quant.md#p8-training-grid-differs-from-the-deployment-grid); format and scale records are in [conversion chain §1](../guides/precision-quant.md#1-four-stages). + +First check that the semantics exist on both sides. When a feature is inside the kernel on one side and external or missing on the other, the difference does not go away with precision settings. Run a two-layer minimal comparison per model feature (attention sink, RoPE variant, final norm, etc.), enabling one at a time. In the external MiMo case SWA alone gave 0 and SWA+sink gave `3.906e-03`, because only one side's kernel included the sink. + +### Discrete decisions come before last-bit error + +After router/indexer logits, also compare the selected set, the order and the tie-break. Top-K with the same set in a different order still changes downstream gather/reduction. A deterministic tie-break defines the index order for equal scores and keeps the invalid/padding index contract. + +For MoE, also fix or verify dispatch, expert capacity, grouped GEMM geometry, quantization and combine. Do not hide a router difference by forcing expert indices; forced routing only isolates the downstream. + +The top-k/top-p truncation set is also a discrete decision: a last-bit logit difference moves tokens near the cutoff in or out of the set, and the normalization denominator of the processed logprob jumps with it. Compare the candidate set before processed logprob. MiMo-V2.6 replays the candidate set together with R3 in production: rollout records a fixed-shape, full-vocabulary bitmap (avoiding host sync and truncation), and training renormalizes within the set. If the contract adopts this replay, state in the report that it masks the set difference, and count the set flip rate separately. + +### Test shape and process boundaries + +For the same victim, construct batch1 / multiple neighbors, different physical positions, lengths on both sides of a boundary, cross-process autotune, prefill/decode and Graph/eager. Different launches when only M changes are not a numerical difference; measure the output. Finding no difference covers only the shapes tested. + +SSM chunk recurrence and per-token recurrence can be two legal but different rounding paths. Unifying them requires matching the chunk grid, boundary state precision and gate order, not just a deterministic switch. + +For history-dependent operators, a single token's scan/conv input and output can make an upstream history difference look like an error in that operator. First find the first differing token at each layer over the **full prompt batch**, then capture norm, projection, conv, scan, gate and output projection in order at that layer. A single-position probe only narrows candidates; it does not prove the first cause. A plain Python hook in a compiled graph may hold stale warmup data; verify capture freshness with a predictable sentinel such as embedding/residual, and if needed use the custom-op boundary that actually executes and a separate buffer. + +### Boundary between shared forward and backward + +With a shared inference forward + native differentiable VJP, verify separately: forward parity, a finite gradient for every trainable parameter, and gradient comparison under a cotangent representable in the same dtype. Operator-level VJP verification does not prove the whole-model gradient is bitwise equal to the original training implementation. + +Do not detach the inference output and use it as a trainable substitute. A finite gradient does not mean effective learning; also check that advantage/reward is not degenerate, the gradient is non-zero, and parameters and optimizer state are updated. + +### Changing the arithmetic baseline + +If you deliberately change both sides' arithmetic baseline to align them, keep the old baseline, build a new frozen rollout, state that "old outputs are not guaranteed to stay bitwise unchanged", and run model quality and performance regressions. Success numbers from the old baseline do not carry over to the new implementation. diff --git a/.agents/skills/train-inference-consistency/workflow/report-template.md b/.agents/skills/train-inference-consistency/workflow/report-template.md new file mode 100644 index 000000000000..e8fffb0e383d --- /dev/null +++ b/.agents/skills/train-inference-consistency/workflow/report-template.md @@ -0,0 +1,119 @@ +# Train–inference consistency report + +## Outcome and scope + +State pass/fail/inconclusive for the requested scope. Record model/layers, independent samples, valid tokens, dtype, topology, Graph/cache paths, update count and untested conditions. + +Link the human-approved acceptance checklist: required vs out-of-scope checks, recipe revision, two-update stop condition, quality/resume coverage and target-repository submission requirements. Out-of-scope checks are unverified, not passed; finishing the short run does not complete missing required checks. + +## Immutable experiment + +Record source commits and dirty patch hashes, loaded libraries, image digest, GPU, checkpoint/policy identity, input bundle digest, score semantics, resolved config, commands, resource budget and artifacts. + +Separate source, installed wheel, loaded native libraries and JIT specializations. State whether the measured runtime and the final delivered image are identical; if not, name the missing verification instead of inheriting the result. + +For recipe work, record training context, rollout length distribution, concurrency, packing/token batch, hardware/interconnect and master/deployment/activation/KV/accumulation/gradient dtypes, plus quantization granularity, scale/update rules and the quality/throughput/memory budgets defined before testing. Unknown budgets stay unaccepted even when measurements exist. + +## Delivery manifest + +For the three-PR / one-image delivery, fill this from the final clean build, not development overlays. Use full SHAs and immutable image identity; mark missing evidence as not run. + +| Deliverable | Review / source identity | Build or acceptance evidence | +| --- | --- | --- | +| vLLM PR | URL; target owner/branch/base SHA; head SHA | Build command and wheel/native/JIT inputs | +| Megatron PR | URL; target owner/branch/base SHA; head SHA | Installed model / primitive / adapter identity | +| verl PR | URL; target owner/branch/base SHA; head SHA | Recipe and image build script paths | +| Unified image | Base and final platform digests; dependency pins | Build log; actual trainer and rollout worker imports/libraries; no development source mounts | +| Approved recipe | Script SHA, resolved config and human approval | Changes from reference recipe; dataset and topology | +| Final validation | Commands and artifact IDs tied to the above image | Proxy / full zero-update / θ0–θ2 / quality / performance status, each reported separately | + +Each PR is complete against the agreed base, not stacked on an unmerged feature branch; cross-repository dependency pins are expected. Every required runtime file outside the PRs must be a pinned existing dependency; an uncommitted production patch is a delivery gap. A changed base, build input or recipe requires reassessing affected evidence. + +## Proxy coverage and gradient contract (when training is in scope) + +Map each full-model feature to its proxy layer/module and test artifact; name omitted depth, geometry, topology and execution paths. Per module, record visible forward semantics, trainable variables, master vs dequantized surrogate values, STE/clipping/scale rules, fixed discrete choices and differentiable scores, VJP accumulation/collectives, and saved/recomputed state with its memory cost. + +Link the independent reference and explain why it does not call the tested backward. Record shape/dtype/layout and random, zero, sparse and captured upstream-gradient coverage; per input/parameter: missing/nonfinite gradients, max absolute error, relative L2, norms and cosine where defined. Include predefined thresholds with rationale, original forward-byte comparison and one-step master-parameter deltas. + +## Symptom and localization + +Record per-request/token differences, A/A and A/C results, side-isolation evidence, last equal / first unequal semantic boundary, identical-input operator comparison, competing explanations and counterexamples. + +## Cause and change + +Explain the violated arithmetic/metadata/state contract and the smallest triggering input. Separate the actual repair from diagnostic interventions. Link code changes and the preserved regression. Write unknown if evidence does not establish a cause. + +## Verification + +Record the proxy upgrade decision before escalating to full-model BI consistency or training. Do not merge inference-only and cross-stack parity into one result; a failed, missing, empty or stale forward comparison blocks escalation. If the user scoped the proxy to fixed-token forward parity, mark SFT/RL not run and list the later training gates. + +| Proxy upgrade gate | Status (pass/fail/inconclusive/not run) | Bound command / artifact and valid-token count | Decision | +| --- | --- | --- | --- | +| Proxy feature/operator coverage | | | | +| Proxy inference BI and prefill/decode parity | | | | +| Proxy inference ↔ native old-logprob recomputation | | | | +| Proxy inference ↔ actual grad-enabled actor forward | | | | +| Backward / SFT / RL, if in proxy scope | | | | +| Upgrade to full-model BI work | | | | + +For each cross-stack row, follow the [dual-path contract](validation.md#old-logprob-and-training-forward-dual-path-contract): model geometry, frozen policy/checkpoint identity, source/image digest, dtype, topology, model/grad modes, output requires_grad, actual backends, raw logprob mismatch count and first mismatch. Keep native no-grad scoring semantics; its pass does not prove grad-enabled forward parity. Capture scores before updates and loss transforms. + +| Gate | Command / artifact | Result and count | Scope limitation | +| --- | --- | --- | --- | +| Original failing replay | | | | +| Operator / metadata regression | | | | +| Engine and production layouts | | | | +| Independent module VJP and one optimizer update | | | | +| Controlled proxy SFT (A/B/C paths) | | | | +| Real backward/update | | | | +| θ0 / θ1 / θ2 parity | | | | +| First post-checkpoint-resume export / parity | | | | +| Model quality | | | | +| Clean shutdown / sustained run | | | | +| Final immutable image, without development overlays | | | | + +For two-update RL, keep a receipt per policy θ0, θ1, θ2: input/mask identity, valid token count, original-byte mismatches, actual trainer-forward execution and update norm/count. Record both forward-path results at matching frozen policy versions. Worker exit is checked at final shutdown, not at each policy transition. + +## SFT controls and quality (when training is in scope) + +Separate SFT-A (production), SFT-B (same visible forward at the same snapshot plus an independent surrogate) and SFT-C (verified native training); these labels are independent of replay packing variants. Link same-checkpoint/same-batch A/B gradients and optimizer deltas; once trajectories diverge, later forward values are not assumed identical. A/C differences include forward/quantization changes. + +Record initial master and optimizer state, data/order, labels/masks, packing, effective-token normalization, batch/accumulation, learning-rate schedule, clipping, precision/loss scaling, RNG/dropout/recompute and intentional differences. Keep raw train, held-out and fixed-probe loss, gradient norm, update RMS, clipping/nonfinite/skipped-step counts, throughput and baseline/seed variation when available. Use one optimizer-step or effective-token axis; report predefined windows/thresholds, missing controls and whether evidence is only a short smoke run. + +## Performance + +Record matched work, warm/cold timing, node pool and paired order, repetitions, latency/throughput, kernel/rollout/full-step results and numerical receipts. Separate measured improvements, rejected candidates and estimates. + +For 8K/1K acceptance, link the [approved aggregation protocol](../guides/recipe-and-acceptance.md#pairing-and-computation): formula, repetitions, per-concurrency thresholds and invalid-leg replacement rules. Include all raw pairs, exclusions with reasons and the approved aggregate per c; median paired loss and loss from median throughputs are different statistics. Other phase timings are diagnostics, not gates. + +Where relevant, include training time, peak memory, launch count vs token batch, CPU gaps and repeated quantization/allocation. Separate early feasibility profiles from final full-model optimization; normal autoregressive decoding is not evidence of an avoidable per-token prefill/training loop. + +State whether this is a BI-only causal comparison or each side's best valid stack, and list all differing flags/backends. Include prompt count vs expanded requests, actual input/output lengths, completed output tokens, wall-time denominator and whether the run is fixed-length synthetic or natural-EOS. Do not sum overlapping GPU ranges or percentages from separate ablations. + +## Evidence corrections and retention + +List superseded claims with the correcting artifact and reason. Record which original failure bundles, checkpoints and small receipts are retained, their storage budget/location, and any missing eval steps. Never fill a curve with a different policy lineage. + +## Remaining work + +List concrete gaps and the smallest next experiment. State whether production hooks/config were restored. Keep raw evidence; do not upgrade historical or partial results into current completion. + +## Experiment record (YAML) + +```yaml +# Train-inference consistency replay experiment record: each variant changes only one variable, everything else is fixed. +purpose: "frozen-policy old-logprob replay" +target: "bitwise selected-token raw logprob" +rollout: "/absolute/path/to/frozen-rollout" # original framework batch, not exported scores +checkpoint: "content id + pre-update recovery path" +environment: "image digest, source+patch hashes, libraries, GPU" +topology: + trainer: "TP/PP/CP/EP/DP" + rollout: "TP/PP/CP/EP/DP" +state: "RNG, mutable buffers, Graph/cache and reset policy" +budget: {max_runs: 3, timeout_seconds: 600} +variants: + - {name: A, change: "none (A/A baseline)", cmd: "", output: "", exit: null, result: ""} + - {name: B, change: "none (A/A repeat)", cmd: "", output: "", exit: null, result: ""} + - {name: C, change: "packing reversed", cmd: "", output: "", exit: null, result: ""} +``` diff --git a/.agents/skills/train-inference-consistency/workflow/validation.md b/.agents/skills/train-inference-consistency/workflow/validation.md new file mode 100644 index 000000000000..9bac973376cb --- /dev/null +++ b/.agents/skills/train-inference-consistency/workflow/validation.md @@ -0,0 +1,169 @@ +# Tiered regression and completion criteria + +## Full-model entry gate: pass proxy train-inference forward first + +When backends, compile paths or delivery images have changed, check against the [candidate evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger) whether old results still apply. A passing old candidate is history only and does not fill the current candidate's entry gate cells. + +Before full-model correctness or training experiments, provide the evidence below. Missing or not executed counts as not passed; repeated full-model launches are no substitute. Proxy inference-internal BI/PD and proxy training forward↔inference logprob comparison are two independent results: the former passing does not mean the latter passes. + +| Gate | Minimal evidence | +| --- | --- | +| Proxy coverage | Inventory of operator types, necessary adjacent-layer interactions, quantization geometry, state and parallel paths; uncovered items stated | +| Inference BI/PD | Batch, prefill/decode and mixed-scheduling comparison on the proxy with fixed weights and tokens; proves only consistency inside inference | +| Training forward↔inference | Same proxy, same policy weights, fixed tokens: verify native old-logprob recomputation and the grad-enabled actor forward used in training separately, each against the target production inference path; keep evidence per the dual-path contract below | +| Backward and quality | [Independent module gradient comparison](../guides/backward-validation.md#layer-1-module-level-vjp-comparison) with reference implementation, tolerances and raw results; controlled SFT per task scope, or an explicit statement that quality is not covered | +| Real RL update | At least two effective updates, θ0/θ1/θ2 same-policy comparison, actual parameter change, normal exit; sequence below | +| Performance feasibility | launch/sync/memory overhead on real shapes and the proxy profile, ruling out bottlenecks such as per-token launches in batchable prefill/training; the target Graph path works | + +This is an escalation gate, not the completion criterion: two-step RL does not prove convergence, and the proxy does not prove full depth, capacity, long context or cross-node correctness. On the full model, extend coverage one item at a time. Inference-only tasks apply only the relevant gates, state that training is not accepted, and add no training work. + +Before escalating, record per gate in the [report template](report-template.md#verification): `pass`/`fail`/`inconclusive`/`not run`, command, artifact path and the bound model geometry, weights, source/image, precision and parallelism. Accept only current-candidate evidence. For training forward↔inference, the bitwise target requires original-dtype mismatch of valid tokens to be 0; tasks with an agreed error budget use the predefined standard and cannot use it to report a bitwise failure as a pass. An unexecuted training forward, an empty mask, or passing only single-operator or inference BI/PD tests must not release to the full model. Keep the first failing sample on the proxy for layer-by-layer localization, and after the fix rerun the original failing case and the affected matrix. + +### old-logprob and training forward dual-path contract + +- **Native old-logprob recomputation:** follow the framework's actual path and its `no_grad`/`forward_only` semantics; do not force autograd on for acceptance. Its result proves only that this scoring path matches rollout. +- **Actual training forward:** with the same frozen θ, token history and valid mask, run the grad-enabled actor branch training actually uses, capture logprob before the loss transform, and compare with rollout independently. Keep the target packing, recompute, operators and parallelism; do not substitute a diagnostic forward. Passing does not replace the later independent backward validation. +- Record per path: model mode, grad mode, output `requires_grad`, operator class / backend / source hash, original dtype, non-empty valid token count and mismatch. Diagnostics (substitute attention/MoE, Graph off) prove only that path; `forward_only` alone is not grounds for rejection, and native old-logprob passing does not prove the grad-enabled branch. +- The comparison happens before the corresponding policy update, for both proxy and full-model θ0 entry. If the formal recipe has several minibatch updates per round, do not require 0 diff between the updated actor and the old-θ rollout, and do not modify the algorithm to manufacture equality; record the real parameter version and compare only within one version. + +When an adapter / overlay is injected dynamically, the receipt must also record the actual object's class, source file and key config and assert that the target callsite was hit. A script directory placed before the overlay in `PYTHONPATH` can silently load an old adapter; a 0 diff from such a run cannot vouch for the new candidate. + +If the user limits the minimal proxy to fixed-token logprob comparison (or forward only), do not require GSM8K, SFT or two-step RL; mark those gates `not run`, not passed. Scale up as agreed only after the proxy train-inference forward passes. Before any multi-step full-model RL, first complete a zero-update same-policy logprob comparison with **full-model θ0 real workload, real training mode, target production inference mode**, then the necessary gradient validation and effective-update checks. A fixed 8K/1K proxy pass does not replace the full-model variable-length workload. When full-model θ0 fails, stop optimizer updates, keep the failing trajectory and shrink to the first divergence. This scope adjustment does not waive the proxy forward gate. + +### Limited exceptions that allow going to the full model early + +- Establishing the BI=0 performance/quality baseline of the full target model is not approval for full-model BI debugging or training. +- When a fault depends on real depth, capacity, cross-layer state or cross-node topology and cannot be reproduced on a semantics-preserving proxy, targeted full-model diagnosis is allowed. Before the experiment, write down why the proxy is insufficient, the hypotheses to distinguish, the minimal scale and the stop condition. +- Scaling up does not scale up resource permissions: no extra GPUs, no stopping other jobs, no bypassing the model/runner scope the user set. + +### Faster iteration and fallback on failure + +Judge each change's impact first: minimal counterexample, single operator or fixed replay, then the relevant proxy gates, then the staged full model; do not repeat all ten steps per commit. Changes to shared arithmetic or gradients regress the affected forward/backward; final delivery still runs the full agreed acceptance. + +On a full-model failure, first save the input before the first divergence, the weight version, metadata, cache/routing and necessary RNG state, and extract a module or few-layer reproducer. Iterate on it only after confirming the fault survives shrinking. If the fault disappears, a passing shrunk test is not a fix, and the full-model reproduction is kept. + +Fixed rollout / module inputs may be cached and replayed in an isolated resident process, bound to fingerprints of weights, input, config and necessary state; a new implementation still reruns the comparison. Each scale-up states the existing evidence and the open question the new scale answers. + +If a candidate fix changes the sampling trajectory or EOS stop point, save the new token history and replay training old logprob on it with the same weight version. The old trajectory remains a fixed-input diagnostic; comparing new inference against the old response length is not an end-to-end pass. List evidence separately for a diagnostic runtime override, a formal source candidate, Graph/no-Graph, forward-only and real `train()`+autograd. An isolation package may be used for localization; final acceptance returns to the delivered code and the default target execution mode. + +## Start from the cheapest failure criterion + +1. host metadata/shape arithmetic: minimal counterexample of a pure function and boundary parameterization; if needed, exhaustively check equivalence before and after the fix. +2. Single operator: same input, real dtype/layout, non-finite detection, repeat/cross-process and Graph. +3. frozen batch replay: the original failing A/C and A/A, full valid-token comparison; UID mapping passes. +4. engine: batch composition, prefill/decode, chunk boundaries, cache/Graph, actual parallelism scope. +5. Real training (when in scope): non-degenerate reward/advantage, finite and meaningful non-zero gradients, optimizer step, actual parameter change, frozen-token alignment after the update. +6. Production acceptance: full model, agreed multiple steps, update/restore lifecycle, quality evaluation and paired performance. + +Not every task runs every tier. Report the tiers completed; when the user asks only for diagnosis, do not start full training. Extend existing test suites first; do not replace behavior tests with assertions on source strings. + +## Preventing false greens + +- Zero valid tokens, comparison not executed, missing rank/sample: invalid experiment. +- All-zero score, mask all false, the same sample repeated: check the producer and coverage; a passing comparison is not enough. +- probability/KL/logPPL is zero: still check the original logprob bytes. +- Zero gradient or identical rewards: maybe no effective update; check advantage, update count and parameter difference. +- Small model / short output passes: confirm the full model's mixed-step and long-context trigger conditions. +- A shrunk proxy's error magnitude is not the full-model baseline: in the external MiMo case the 12-layer proxy was `1.866e-12` and full-depth real weights were `1.143e-02`. The proxy only verifies the execution path. +- In-training mismatch metric is 0: confirm the micro-batch really ran the trainer forward. The framework may skip the forward for a zero-advantage batch and return 0, a trivial match; count only reports with non-zero advantage. +- A success marker written before the process was killed: separate a completed numerical experiment from normal exit / long-term stability. +- probs_diff is 0 but many sibling samples within a group are identical: a sampling diversity defect, independent of the consistency metric. Count distinct rollout samples per group; the mechanism and fix are in [rollout-engine §1](../guides/rollout-engine.md#1-sampling-seed-collision). +- The launcher shows `COMPLETED 0:0`: read the recipe's own exit code. When tests report passed, check the skipped count; GPU cases may be skipped wholesale by a marker or environment gate. + +## Unified measurement basis and matrix for forward consistency + +Fix policy, token history, position, mask and tensor semantics, and record per item the original dtype/shape, valid token count (greater than 0), non-finite count, bitwise mismatch count, maximum absolute difference and first failing position. A dtype difference is a contract inconsistency first; do not convert and then claim a bitwise pass. After aligning semantics, compare contiguous tensor bytes that preserve the bit pattern or an equal-width integer view; floating-point `==` does not distinguish positive and negative zero. NaN/Inf in a valid logprob is a failure even with identical bytes. Sampled-token comparison does not extrapolate to the full vocabulary. + +Keep cases and artifacts per item for the target execution scope, and list unsupported / not run items: + +| Dimension | Comparison under fixed semantics | +| --- | --- | +| A/A | Repeated execution with the same weights, input and state; in a separate process if needed | +| BI | Single request, different batch sizes, companion requests, permutation/position and packing | +| PD | Full prefill, chunked prefill and step-by-step decode on the same history, covering mixed scheduling | +| Length and branches | Real lengths, chunk/causal boundaries, and cases that trigger sparse selection, MoE and special modules; not only random tokens | +| Execution state | Target Graph capture/replay, cache reuse/reset and the actual parallel layout | + +The bitwise target requires mismatch=0; error-budget tasks report against the predefined standard and are never recorded as a bitwise pass. Gradient tolerances are defined independently, see [module VJP](../guides/backward-validation.md#layer-1-module-level-vjp-comparison) and [SFT control](#proxy-sft-three-way-control). + +## Two-step RL sequence + +This sequence is the lower bound for entry; "container started / training script exited 0" is not a pass. + +When the task includes the full train-inference loop, complete at least this sequence; every actor recomputation uses the same policy and token history as its rollout: + +```text +θ0 → rollout0 → actor(θ0) recompute and compare + → backward0 / optimizer update → θ1 +θ1 → rollout1 → actor(θ1) recompute and compare + → backward1 / optimizer update → θ2 +θ2 → rollout2 → actor(θ2) recompute and compare +``` + +The θ2 comparison needs no third optimizer update. To compare outputs across policies, use the same frozen token history; do not rely on sampling happening to match. + +For each stage save the policy content identifier, token/mask, non-empty valid comparison count and raw byte comparison result; for each update save finite loss/gradient, coverage of parameters that should have gradients, master update norm and optimizer step. + +For a bitwise task the θ0, θ1 and θ2 comparisons must all have mismatch=0, with two effective updates and a normal exit actually completed. Skipped updates, both sides reading an old cache, an empty mask or comparing only probability summaries are not allowed. + +A small master update may not cross a quantization bucket, so deployment bytes need not change every step. Judge an update by master update norm and optimizer step, not by step count or byte change. + +## Proxy SFT three-way control + +The three paths, controlled variables, recorded items and localization order are in [Backward validation layer 3](../guides/backward-validation.md#layer-3-proxy-sft-three-arm-comparison). If the module gradient comparison fails, fix the implementation before starting SFT. + +## Performance and model quality + +When forward arithmetic, precision, scorer or inference path changes, rerun the same quality evaluation within the human-reviewed scope and the target repository's requirements. Record model/data, problem count, prompt format, generation settings, statistical variation and unfinished parts; mark out-of-round items unverified, and add no SFT or long runs. A loss decrease or two-step reward change does not prove no quality regression. + +Micro benchmarks, fixed-token parity, generation quality, normal exit and multi-step stability are separate. A run with four non-zero updates and 139,260 bitwise-equal tokens that ended in a DataLoader error supports four-step alignment, not clean shutdown or 50-step stability. + +The full model covers the reviewed scope (real depth / weights, target context, concurrency, packing, parallelism/hardware, cache/Graph) and reruns affected correctness gates; proxy curves do not extrapolate to full-model quality. Final performance is accepted on 8K/1K per the [paired statistics protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); training and full-step times may be listed as diagnostics and add no performance gate. Scope and stop condition of the two-step short run are in the [general two-step acceptance agreement](#general-two-step-acceptance-agreement). + +## Failure classification and delivery + +Use `pass`, `fail`, `inconclusive` with a reason. Keep source/patch/image/bundle hashes, commands, per-item counts, earliest difference, excluded hypotheses, and before/after results on the original reproducer. Cleaned-up code needs re-verification or explicit behavior-equivalence evidence; old tests do not vouch for new code. + +Passed items stand despite unfinished ones, but the final conclusion is limited in scope. Use the [report template](report-template.md). Formal upstream contributions also follow the target repository's duplicate check, human line-by-line review, tests and AI assistance disclosure; historical cases do not mean the review was done on the author's behalf. + +If the delivery includes an image, the gate object is the final immutable artifact, not a development image plus a source mount. Verify the closure per [runtime provenance](../guides/environment-provenance.md); accept save per [01 §3.2](../guides/checkpoint-resume.md#32-reproduce-and-discriminate), load per [§4](../guides/checkpoint-resume.md#4-load-enough-capacity-does-not-mean-enough-for-the-transient-peak), and the first resume per [§6](../guides/checkpoint-resume.md#6-how-to-prove-resume-is-correct). Passing environment checks proves only the preconditions for running, not GPU numerics or performance. + +## Recipe review materials and two-step acceptance agreement + +The integration point is in [training-side model support](../guides/training-forward.md); only process agreements are here. + +### Recipe and PR delivery contract + +Locate the existing integration layer in the selected version first. Megatron's `experimental/lite/examples/verl/verl_mlite/` provides the verl engine adapter; an external integration package installs normally into the unified image, and removing the runtime overlay does not require reimplementing the engine in verl. Check the actual engine registration, model factory, batch conversion, loss normalization and optimizer path, and list reusable interfaces and specific gaps. + +- **verl PR** contains by default the model's DAPO recipe and the unified image Dockerfile. Public interface changes beyond these two: submit the gap and plan first, implement after human review; necessary regression tests may stay. +- **Megatron PR** contains the model, operators and necessary support in the existing `verl_mlite` integration layer; do not move generic code that does not belong in Megatron there just to keep the verl PR at two files. +- **vLLM PR** contains the necessary inference-side model, kernel, state and weight loading changes. +- Per the [general delivery target](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image), submit complete PRs to the three main repositories, including the required model adaptation and not stacked on another unmerged feature PR; pin cross-repository dependency versions and state the build order. Consistency uses mLite and the training baseline uses the existing verl recipe's default MCore path (see [the workflow](README.md#gate-table)); do not merge them into one training scheme. + +### DAPO reference and review materials + +[verl #7688](https://github.com/verl-project/verl/pull/7688), at reviewed version `640c5f0a317307e99636924804c40624242fa19b`, uses Megatron's integration package through `model_engine=mlite` and `pkg://verl_mlite.config`, so the verl diff needs only the recipe and the Dockerfile. Follow its real `verl.trainer.main_ppo` execution chain; do not assume its config keys, DeepSeek precision/topology, mounts or overlay fit the current model. + +The reference case's model, data and evaluation set are not defaults. Choose training data / reward by target model and human agreement, and keep the training set distinct from the quality evaluation set. Finish config resolution and the necessary proxy smoke, then submit for human review: + +- Executable launch script, unified image digest, commit of each repository and full resolved config. +- Parameter differences from the reference recipe with reasons: data/reward, prompt/response length, rollout count, batch/microbatch, parallelism, quantization / KV precision, optimizer. +- Proxy check results, known uncovered items, and the acceptance scope of the full-model zero update and the two real updates. + +Formal full-model E2E starts after a human confirms a specific recipe version and the proxy entry passes. Workload or numerical/update changes are reviewed again; path and log changes that leave input and execution semantics unchanged are only recorded. This gate does not block already authorized builds or proxy development. + +### General two-step acceptance agreement + +- **Freeze versions first.** Record the heads of the three delivery PRs, the unified image digest and the full recipe config. Changes that affect results require a version update and the related verification; old artifacts cannot be reused for acceptance. +- **Review the recipe first** (materials above). Confirming process or scope does not approve an unsubmitted recipe. The review checklist also marks which quality, restore and long-run checks are required or out of scope for this round, and the evidence the target repository requires; out of scope does not mean passed, and the short run ending does not hide missing required items. +- **Run two effective updates by default.** Use the formal recipe and shorten only the execution length to two steps; do not downgrade data, sequence length, rollout count, batch, optimizer, algorithm, parallelism or backend. If the stop step count drives the learning-rate schedule, keep the formal recipe's learning-rate/warmup semantics for the first two steps instead of compressing the full schedule into two; the implementation is reviewed with the recipe. +- **Keep the existing correctness gates.** After the proxy real training forward passes, run the full-model [two-step RL sequence](#two-step-rl-sequence) with actual parameter updates; a fixed-weight replay does not count instead. +- **Performance is accepted only on 8K/1K.** On the frozen version, use Rust Bench at 8192 input / 1024 output, ignore EOS, c=1/8/32/64, against the human-confirmed BI=0 baseline and performance budget. Training forward, backward and full RL step times may be recorded but add no performance gate. +- After the two updates, the θ2 comparison and the agreed 8K/1K performance pass, stop the short run and report. Overall acceptance is still judged item by item against the review checklist; if a required item is missing, state that it is not complete. Do not extend to 100 steps, add an MCore long-run comparison or new quality gates, and do not extrapolate the two-step result to long-run training quality. + +### E2E comparison target + +Run the [two-step RL sequence](#two-step-rl-sequence) through the real trainer entry point. 0 diff means equal original-dtype logprob for the same weight version, token history and valid mask; θ0 and θ1 outputs may differ. Verify the actual master update each round (not step count or byte change). Disk replay is a localization tool and does not replace acceptance through the real trainer entry point. + +Finally re-verify with the image built from the delivery PRs and the reviewed recipe. Record standalone engine smoke, old image numerical results and historical performance separately; they do not replace full evidence from this execution entry point. diff --git a/.claude/skills/train-inference-consistency b/.claude/skills/train-inference-consistency new file mode 120000 index 000000000000..906cd763172c --- /dev/null +++ b/.claude/skills/train-inference-consistency @@ -0,0 +1 @@ +../../.agents/skills/train-inference-consistency \ No newline at end of file From 237d99b4aff67ecfedacca7acb0ee93f299fa54e Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 17 Oct 2026 21:54:30 +0800 Subject: [PATCH 797/952] Revert "[Docs] Add train-inference consistency agent skill" (#61551) (#62186) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../train-inference-consistency/SKILL.md | 31 -- .../agents/openai.yaml | 4 - .../guides/backward-validation.md | 221 -------- .../guides/bi-kernels.md | 143 ----- .../guides/checkpoint-resume.md | 304 ----------- .../guides/debugging-mismatch.md | 321 ------------ .../guides/environment-provenance.md | 106 ---- .../guides/performance.md | 160 ------ .../guides/precision-quant.md | 490 ------------------ .../guides/recipe-and-acceptance.md | 113 ---- .../guides/rollout-engine.md | 108 ---- .../guides/training-forward.md | 412 --------------- .../practices/code-and-delivery.md | 90 ---- .../practices/evidence-and-experiments.md | 179 ------- .../workflow/README.md | 122 ----- .../workflow/report-template.md | 119 ----- .../workflow/validation.md | 169 ------ .claude/skills/train-inference-consistency | 1 - 18 files changed, 3093 deletions(-) delete mode 100644 .agents/skills/train-inference-consistency/SKILL.md delete mode 100644 .agents/skills/train-inference-consistency/agents/openai.yaml delete mode 100644 .agents/skills/train-inference-consistency/guides/backward-validation.md delete mode 100644 .agents/skills/train-inference-consistency/guides/bi-kernels.md delete mode 100644 .agents/skills/train-inference-consistency/guides/checkpoint-resume.md delete mode 100644 .agents/skills/train-inference-consistency/guides/debugging-mismatch.md delete mode 100644 .agents/skills/train-inference-consistency/guides/environment-provenance.md delete mode 100644 .agents/skills/train-inference-consistency/guides/performance.md delete mode 100644 .agents/skills/train-inference-consistency/guides/precision-quant.md delete mode 100644 .agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md delete mode 100644 .agents/skills/train-inference-consistency/guides/rollout-engine.md delete mode 100644 .agents/skills/train-inference-consistency/guides/training-forward.md delete mode 100644 .agents/skills/train-inference-consistency/practices/code-and-delivery.md delete mode 100644 .agents/skills/train-inference-consistency/practices/evidence-and-experiments.md delete mode 100644 .agents/skills/train-inference-consistency/workflow/README.md delete mode 100644 .agents/skills/train-inference-consistency/workflow/report-template.md delete mode 100644 .agents/skills/train-inference-consistency/workflow/validation.md delete mode 120000 .claude/skills/train-inference-consistency diff --git a/.agents/skills/train-inference-consistency/SKILL.md b/.agents/skills/train-inference-consistency/SKILL.md deleted file mode 100644 index 6f47ad8041b6..000000000000 --- a/.agents/skills/train-inference-consistency/SKILL.md +++ /dev/null @@ -1,31 +0,0 @@ ---- -name: train-inference-consistency -description: Build and validate train-inference consistency (true on-policy RL), where vLLM rollout and the training forward give bitwise identical logprobs at high throughput. Use when aligning a training stack with vLLM batch-invariant inference, validating forward, backward and quantization contracts, or diagnosing logprob mismatch, layout and performance regressions between training and rollout. ---- - -# Train-inference consistency skill - -Make the training side and the inference side in RL produce bitwise identical logprobs for the same token (diff = 0), and, while keeping 0 diff, bring performance back close to the high-performance baseline. - -How to enable batch-invariant inference in vLLM is in [batch invariance](../../../docs/features/batch_invariance.md); this skill covers making the training side match it. - -Start here: **[train-inference consistency workflow (C0–C10)](workflow/README.md)**. It answers "what to do next, and what counts as passing"; acceptance details are in [validation.md](workflow/validation.md), and the report template is [report-template.md](workflow/report-template.md). - -Find a guide by problem: - -| Guide | Questions it answers | -| --- | --- | -| [Fault localization and mismatch bisect](guides/debugging-mismatch.md) | Find the side first, bisect, classify before changing precision, save and compare protocol | -| [Training forward](guides/training-forward.md) | Module inventory and route (mLite / MCore / FSDP2), reproducing the inference forward by calling its kernels, layout and length, framework integration | -| [Backward validation](guides/backward-validation.md) | When to reuse backward, gradient contract, module-level VJP comparison, proxy SFT three-way control | -| [BI kernels](guides/bi-kernels.md) | Which kernel each operator uses in BI mode and what it costs | -| [Precision and quantization](guides/precision-quant.md) | The conversion chain master -> training forward weights -> exported bytes -> vLLM bytes; QAT/STE; known pitfalls | -| [Rollout engine and sampling](guides/rollout-engine.md) | Sampling RNG, prefix cache semantics, MTP constraints, R3 constraints, rollout replay | -| [Performance and profiling](guides/performance.md) | Workload and measurement basis, fair A/B, breakdown and ablation | -| [Checkpoint and resume](guides/checkpoint-resume.md) | Exact resume, derived state after restore, memory budget for colocated training and rollout | -| [Environment provenance](guides/environment-provenance.md) | Which version is actually running now; immutable delivery image | -| [Recipe and acceptance](guides/recipe-and-acceptance.md) | BI=1 / BI=0 arms in one recipe and one image; evaluation contract; paired acceptance | - -Process rules: [evidence and experiments](practices/evidence-and-experiments.md), [code and delivery](practices/code-and-delivery.md). - -Note: this skill comes from the DeepSeek-V4 and Nemotron-H train-inference consistency work and covers train-inference consistency only. Day-0 model support, per-model case records and cluster-specific notes are not included. diff --git a/.agents/skills/train-inference-consistency/agents/openai.yaml b/.agents/skills/train-inference-consistency/agents/openai.yaml deleted file mode 100644 index 012b1477efd8..000000000000 --- a/.agents/skills/train-inference-consistency/agents/openai.yaml +++ /dev/null @@ -1,4 +0,0 @@ -interface: - display_name: "Train-Inference Consistency" - short_description: "Make RL training and vLLM rollout logprobs bitwise identical." - default_prompt: "Use the train-inference consistency workflow to align training with vLLM inference for this model." diff --git a/.agents/skills/train-inference-consistency/guides/backward-validation.md b/.agents/skills/train-inference-consistency/guides/backward-validation.md deleted file mode 100644 index 1a1c6248ea9b..000000000000 --- a/.agents/skills/train-inference-consistency/guides/backward-validation.md +++ /dev/null @@ -1,221 +0,0 @@ -# Backward: reuse and validation - -## Reusing backward: when you actually need to write your own - -When to reuse backward and when to write glue. - -### Conclusion - -Backward is determined by the **math formula**, not by the forward kernel. The backward of `y = f(x, W)` needs only the formula and the saved inputs: `dX = dY·W`, `dW = dYᵀ·x`. A vLLM forward kernel does not by itself require a new backward. "No backward operator at the same precision" is not a reason to rewrite: under the high_precision contract the backward runs in BF16/FP32 anyway. - -### How to wire it: forward once, backward uses what the forward saved - -The training forward is the inference forward. Each module is an autograd Function: its forward calls the inference kernel to produce the visible output and saves what its own backward needs (inputs, visible output, LSE, fixed routing IDs, quantized values and scales); its backward uses only those, through a reused library backward or a closed-form VJP. If the inference kernel does not output something the backward needs, change the kernel to output it. - -Delivery requires 0 diff and a training step within budget on the main topologies fixed at C1; C6 checks both together. Nemotron audited this against DS4: attention saves q/k/v, out and LSE; routed MoE saves x, the visible FC1 output and the visible output; linear and norms save their inputs; the router saves logits and fixed IDs. Its open gap is Mamba SSD, because vLLM's visible SSD does not output the chunk state. - -### Three cases - -| Case | Characteristics | Examples | What to do | -| --- | --- | --- | --- | -| A only rounding differs | Same math; summation order/chunking differs; difference in the last BF16 bit | RMSNorm, Mamba SSD scan, attention, conv1d, unquantized linear | Reuse the training framework's backward (MCore/TE/mamba_ssm); write no code | -| B forward has rounding | FP8/NVFP4 quantization rounds; the true derivative is 0 almost everywhere | Quantized linear, FP8 Q/K/V | First fix the STE proxy (rounding passes through as 1; dX uses dequantized values by default; DS4 and Nemotron currently approximate with TE high_precision master values, and the deviation must be quantified per [layer 2](#layer-2-gradient-contract-comparison)), then reuse the ordinary BF16 GEMM backward | -| C no ready-made backward for the whole segment | Composite of several operators with no library backward for the whole segment | Routed experts: sort -> grouped GEMM -> relu² -> restore -> top-k weighted sum | Write glue that chains ready-made operators; GEMM still uses TE grouped GEMM; compare gradients against an independent reference | - -Decision flow: (1) does the forward have rounding/quantization (top-k ID selection does not count; IDs are fixed and not differentiated)? No -> A. Yes -> (2) fix the STE contract -> (3) is there a packaged training backward for the whole segment (TE Linear/GroupedLinear, MCore MoE)? Yes -> B, reuse it and feed master values; no -> C, glue. If A has no same-math training implementation, also write only glue. Look first in mLite `primitive/modules` and `ops`, then in MCore/TE/mamba_ssm. - -### How DS4 did it - -DS4's only hand-written kernel is the small FP8 quantization Triton kernel in the forward. The backward has no hand-written kernel, only PyTorch closed-form formulas and glue, all in FP32. - -| Operator | Case | Backward source | What we wrote | -| --- | --- | --- | --- | -| FP8 dense linear (qkv, shared expert, o-proj) | B | TE `general_gemm`, BF16 master | Nothing | -| Routed experts (grouped) | B+C | TE `general_grouped_gemm`; reuses gate_up, the output of the forward vLLM FC1, without recomputing | Glue: sort, activation derivative, summation | -| Sparse attention | A | cuDNN sparse attention backward, consuming the forward's saved out and LSE | Nothing | -| SwiGLU (with clamp), RMSNorm, q/kv norm, mHC | A | Hand-written closed-form derivative, FP32, `torch.compile` | PyTorch formula | -| Routing score | A (IDs fixed) | FP32 replay of `sqrt(softplus)` + normalization, handed to autograd | PyTorch formula | -| Others | A | Generic bridge: forward via vLLM, backward recomputes the PyTorch graph then `autograd.grad` | One utility function | - -DS4 wrote its own for four reasons, none of them a missing same-precision backward kernel: (1) library backwards need intermediates from the library's own forward (TE RMSNorm needs rstd, MCore MoE the token permutation; the vLLM forward has neither); (2) model-specific math is not in the library (`sqrt(softplus)` scoring, SwiGLU with clamp, mHC Sinkhorn); (3) discrete selection must copy the forward (recomputed top-k may pick a different expert on ties, so store the IDs and differentiate only the continuous score); (4) the generic bridge costs too much memory and time (autograd stores all 20 mHC iterations; use a closed-form derivative fused with compile). A custom backward function has no autograd graph inside, so its activation derivative must be closed-form (`dynamic_clamped_swiglu_vjp`; likewise relu² in Nemotron `routed_vjp`). - -### Three rules when reusing - -1. Intermediate definitions must match: attention backward uses the forward's out and LSE, so LSE base (ln or log2), softmax scale and memory layout must line up. If they do not, recompute the forward with the training kernel inside the backward. -2. Backward need not be bitwise identical: recomputed intermediates differ from the vLLM forward in the last BF16 bit and gradients differ by the same amount, which is allowed. The requirement is that the gradient contract holds and the error is near the noise floor. -3. Reused code must be verified too, per [layer 1](#layer-1-module-level-vjp-comparison): compare dX and dW against an FP32 reference that does not call the backward under test, with the same dY, in multiples of the noise floor. - -### Two test-infrastructure false negatives - -- A unit test that instantiated QuantFP8 directly lacked production training's vLLM config context (`CompilationConfig(custom_ops=['none','+quant_fp8'])`, Humming, PP4), and the first long-sequence test failed before computing anything. Fix: the test wrapper uses the production config; production operators and tolerances are unchanged. -- An import check without GPU visibility makes vLLM select a non-CUDA platform, which cannot import `reshape_and_cache_flash` (exported only on CUDA). Dependency checks need real device visibility. -- Both are test errors, not train-inference differences. - -## Backward validation: is the implementation right, is the approximation reasonable, is training biased - -Forward 0 diff shows the logprobs match, not that gradients are right. Validate the backward in three layers: - -| Layer | Question answered | Method | Cannot replace | -| --- | --- | --- | --- | -| 1. Module-level VJP comparison | Does the implementation match the agreed gradient contract (missing scale, wrong transpose, broken gradient chain) | compare dX, dW and a one-step update of every module against an independent reference that does not call the backward under test | layers 2 and 3 | -| 2. Gradient contract comparison | Is the contract reasonable: does an approximation (e.g. dX from the master value) bias training | same batch, compare gradients under the two conventions per class and per layer | layer 1 | -| 3. Proxy SFT three-arm comparison | Does the whole training chain have a systematic bias | three SFT paths (under test, gradient reference, native baseline) with the same data and hyperparameters | layers 1 and 2 | - -Run 1 → 2 → 3. If layer 1 fails, fix the implementation before any SFT. If layer 3 is abnormal, go back to layers 1 and 2. Passing all three clears the backward; none proves long-term RL stability. - -### Write down the gradient contract first - -#### The three tiers of a weight - -In quantized training one parameter has three values, and the gradient must end up at the one the optimizer manages: - -| Tier | Precision | What it is | -| --- | --- | --- | -| main | FP32 | optimizer state; updates accumulate here | -| master | BF16 | the model parameter, converted from main every step; non-quantized layers compute with it directly | -| deployed value | NVFP4 / FP8 etc. | master quantized per the deployment rule; both the training forward and inference use these bytes | - -Keep the optimizer's parameter identity: when replacing the forward, do not create parameters the optimizer does not know about. The deployed value is not a trainable parameter. - -#### Contract table - -A kernel without autograd does not need a surrogate gradient: non-quantized GEMM, norm and attention reuse the existing backward when the math matches; STE is needed only through discrete operations such as quantization rounding. Before implementing, decide per module class: - -| Item | What must be decided | -| --- | --- | -| Visible forward | the inference implementation; input/output dtype, layout; quantize/dequantize positions, scale granularity, rounding, clamp, accumulation precision | -| Trainable variables | input, master weight, bias, norm, router; whether scale is fixed, stop-gradient or trainable | -| Surrogate function | whether dX uses the master or dequantized value (see below); identity or clipped STE, and the value at the threshold | -| Discrete state | fixed expert IDs, sparse indices, ordering; continuous score/combine weights that still need gradients | -| backward | reused implementation or custom VJP; cast, accumulation/reduction precision, communication and normalization | -| Save/recompute | quantized values, scale, norm/attention intermediates, routing, RNG; memory cost and recompute rules | -| Reference and acceptance | independent reference implementation, target shape/execution path, metrics and tolerances agreed in advance | - -Identity STE sets the surrogate derivative of Q(x) to 1; clipped STE sets 1 inside the range and 0 outside, with the threshold value stated and never switched silently. Declare weight and activation quantization separately; frozen weights do not mean nothing upstream of the activations needs gradients. - -#### Linear: which weight dX uses - -For `Y = X · Wᵀ + b` with upstream gradient `G`: - -```text -dW = Gᵀ · X_ref # returns to master via STE, then the optimizer updates main -db = sum_over_tokens(G) -dX = G · W_ref -``` - -Two choices for `W_ref`: - -- **Dequantized value**: `W_ref = dequant(quant(W_master))`, the value the forward actually used. This is the true derivative of `Y = X · Wqᵀ` with respect to X, and what fake-quant QAT gets from autograd. **Use this by default.** -- **Master value**: `W_ref = W_master` (BF16). Reuses the plain BF16 GEMM backward without dequantizing; the cheapest hookup when the inference kernel takes only packed bytes and has no backward. Cost: dX is no longer the exact forward derivative and ignores the quantization error's effect on the input gradient. - -The two give different gradients and must not be swapped silently. With the master value, the layer-2 comparison is mandatory: quantify and record the bias. Likewise for activation quantization, state whether `X_ref` is the quantized or original activation. - -Before reusing a linear or grouped GEMM backward, check flatten/transpose, padding and valid tokens, bias reduction, saved inputs, scaling and the accumulation/backprop dtype. Matching output shape does not prove reusability. - -### Layer 1: module-level VJP comparison - -#### How to build the reference implementation - -The reference may use PyTorch ops and autograd but **must not call the backward under test**, or both paths share the error. For discrete operations implement the agreed surrogate backward explicitly; the zero gradient of bare rounding is not a reference. - -This bridge takes the output value from the inference implementation and the gradient from an independent surrogate. Use it to build the reference path: - -```python -import torch - - -class VisibleWithSurrogate(torch.autograd.Function): - @staticmethod - def forward(ctx, visible, surrogate): - assert visible.shape == surrogate.shape - assert visible.device == surrogate.device - assert visible.dtype == surrogate.dtype - return visible - - @staticmethod - def backward(ctx, grad_output): - return None, grad_output - - -with torch.no_grad(): - visible = inference_linear(x, deployment_weight, scales) -w_ref = dequantize(deployment_weight, scales) # default contract; use master_weight for the master-value contract -surrogate = torch.nn.functional.linear(x, w_ref, bias) -y = VisibleWithSurrogate.apply(visible.detach(), surrogate) -``` - -- Build the surrogate graph in normal grad mode, outside the custom Function's forward. -- With the dequantized surrogate, `w_ref`'s gradient must reach master via STE: bridge a tensor valued as the dequantized weight whose gradient passes straight to master, then feed it to a plain linear. -- Do not use `visible + surrogate - surrogate.detach()`: floating-point add/subtract and compiler transforms change the result. Do not modify the bridge output in place downstream. -- This bridge is the independent validation reference, not the production wiring; production wiring is in [How to wire it](#how-to-wire-it-forward-once-backward-uses-what-the-forward-saved). -- `visible_forward(visible, native, *inputs)` is another way to write the same bridge: `visible` saves inputs without building a graph, `native` is the training framework's implementation of the same math and reruns the forward on those inputs in backward; precision and speed depend on which `native` is chosen. - -#### Comparison steps - -1. **Fix the state.** Fix inputs, weights, quantization config, mask/routing. Reference and candidate get independent copies; clear accumulated gradients. Check that the visible forward matches inference, then the saved state the backward needs; at integration, confirm the real forward produces that state. -2. **Same upstream gradient.** Both sides use the same external dY, dtype and values. Cover fixed random, zero, and sparse (non-zero at selected positions) gradients, then samples from real training. `y.sum().backward()` alone misses transpose, permutation and broadcast errors. -3. **Compare item by item.** Compare dX, dW and every bias/norm/router/scale gradient separately: missing, shape, finiteness, non-zero where expected. Do not ignore unused differentiable inputs by default; list parameters the branch genuinely does not use. Report max absolute error, relative L2 `‖g − g_ref‖₂ / max(‖g_ref‖₂, eps)`, norm ratio and cosine; direction alone is not enough. With a zero norm the cosine is undefined; use absolute error. -4. **Set tolerances in advance.** Derive atol/rtol and the relative L2 threshold from dtype, reduction length and reference precision; record the reasoning. No single threshold for all modules, no relaxing after a failure, and no relaxing a same-semantics kernel optimization because "STE is an approximation anyway". Cross-framework bitwise identity is not required. -5. **Coverage.** Small M (decode), large M (training/prefill), different context lengths, unaligned tails, padding/packing/mask, zeros and extreme values; both sides of the quantization threshold, the saturation region, different scales. For MoE cover empty experts, repeated routing, shared experts, top-k ties; for shared parameters check gradient accumulation across use sites. Test under the target dtype/layout, recompute, autocast, CUDA Graph, gradient accumulation and the actual parallel reduction; local first, then integrated. -6. **One-step update.** From the same parameters and optimizer state, run one update each and compare the change in master. If gradients are close but updates differ, check scaling, clipping, optimizer state and parameter binding. - -Cross-check smooth non-quantized modules with small, high-precision [gradcheck](https://docs.pytorch.org/docs/stable/generated/torch.autograd.gradcheck.html) or finite differences: fix the discrete branch, avoid non-smooth boundaries, and perturb the weights the forward actually sees. **Do not validate STE with finite differences of the real quantized forward**: that measures the discrete forward, which the STE surrogate derivative deliberately differs from. - -#### What to check per module - -| Module | Reuse first | Semantics to check | -| --- | --- | --- | -| Dense/QKV/LM head | training linear/grouped VJP, gradient returns to master | dX value convention, cast, bias, gradient accumulation of shared parameters | -| Norm/RoPE | existing training VJP, FP32 statistics if needed | epsilon, RMS or LayerNorm, position offset, rotation layout | -| Activation + quantization | the actual composition of activation derivative and STE | saturation region, whether scale is detached, fusion order | -| Attention/sparse attention | training backward with the same mask/index semantics | causal boundary, softmax scale, LSE definition, GQA, order of selected KV; do not mix auxiliary quantities from different forwards | -| MoE | expert linear/grouped VJP and a differentiable combine | inverse dispatch permutation, gradient accumulation of repeated tokens, empty experts, shared experts, top-k weight normalization | -| Router/sparse selection | fix the discrete selection, keep the continuous score's derivative | no continuous gradient for IDs/indices; do not detach the whole router output | -| Linear attention/SSM | training backward with the same chunked arithmetic | chunk boundary, initial state, state precision, conv state | -| Selected logprob | log-softmax/gather VJP matching the definition | vocab-shard reduction, temperature, token mask, loss normalization | - -Recompute must save or deterministically rebuild the scale, routing/indices, RNG and necessary intermediates, and must not switch the surrogate function. - -### Layer 2: gradient contract comparison - -Layer 1 answers "does the implementation equal the contract", not "is the contract reasonable". When the contract contains an approximation (most often dX from the master instead of the dequantized value), quantify its distance from the exact derivative: - -1. On a fixed checkpoint and batch, compute dX (and gradients related to activation quantization) for every quantized layer under both conventions. -2. Report cosine and norm ratio per module class and per layer; inspect outlier layers. -3. Measure each target precision combination separately: changing weight or activation precision changes the bias. - -On Nemotron-H (NVFP4 routed experts, FP8 Mamba projections), master-value and dequantized-value dX have cosine ≈0.99 and norm ratio ≈1.00 per class and layer, with both BF16 and FP8 activations, so this approximation does not cause the training score gap. Remeasure for other precisions, a larger learning rate or longer training. - -### Layer 3: Proxy SFT three-arm comparison - -After the first two layers pass, run a controlled SFT over the whole training chain: - -| Path | Visible forward | backward | Purpose | -| --- | --- | --- | --- | -| Under test | aligned inference numerics | production backward | the scheme to deliver | -| Gradient reference | same as under test at the same parameter snapshot | independent surrogate reference (the bridge above works) | isolate backward implementation errors | -| Native baseline | original training numerical path | verified original training backward | measure the recipe's overall quality change | - -- First compare gradients and one-step update of under-test vs gradient-reference on a fixed checkpoint and batch, then run curves; once trajectories diverge, visible values differ too. The under-test vs native-baseline gap includes quantization and forward-path differences, not only backward. If the gradient reference's full curve is unaffordable, compare on representative batches and state the missing range. -- Controlled variables: initial weights and optimizer state, data and order, labels/mask, packing, loss normalization, global batch and gradient accumulation, learning rate and schedule, clipping, precision and loss scaling, seed, recompute policy. -- Record train loss, fixed-validation-set NLL, gradient norm, update RMS, clipping ratio, non-finite values, throughput, against optimizer step or cumulative effective tokens. Fix the comparison window and acceptable deviation before the run; a short run is only a smoke test. -- Localization: loss differs at the initial weights -> check forward, data, mask; initial loss matches but one-step update differs -> backward, reduction, scale, clipping, optimizer; diverges only after several steps -> save the first abnormal checkpoint, check recompute, quantization refresh, learning rate, RNG, then replay at the same weights. -- Mark unexplained or unbudgeted differences inconclusive; do not cover them with forward 0 diff. - -On Nemotron-H the three paths (aligned trainer with BF16 activations, Megatron QAT, Megatron BF16) ran 40 steps on gsm8k with the same data in the same order: loss 0.3388 / 0.3383 / 0.3374, gradient norm 0.685 / 0.674 / 0.728, curves overlapping. The training loop itself is fine; the RL score gap comes from RL-specific factors. - -### Common false differences - -Each of these made gradients look mismatched when the comparison method was the problem: - -- **Upstream gradient precision differs.** FP32 dY on one side and BF16-rounded dY on the other creates a difference unrelated to the backward. Use the same external dY. -- **Dump timing.** Expert gradients are multiplied by the expert-DP to DP scaling ratio only inside the optimizer step; dumping earlier shows a false multiplicative difference. -- **Unsharded copy under FSDP2.** After forward-only the root module is not resharded, so the parameter fetched by name is the unsharded compute copy, not the real master. Perturbing it or reading its gradient both mislead. -- **"Frozen" is not frozen.** Zeroing the gradient still lets AdamW weight decay change the parameter; a real freeze needs grad None or removal from the optimizer. -- **Test environment differs from production.** A unit test missing production config context, or an import without GPU visibility, selects the wrong implementation or fails to import (see [Two test-infrastructure false negatives](#two-test-infrastructure-false-negatives)). - -### What to write in the report - -Per module class: gradient contract (dX value choice, STE type, scale handling), reference implementation, upstream gradient construction, shapes/paths covered, metrics and tolerances with reasoning. Layer 2: per-class per-layer cosine/norm ratio and outlier layers. Layer 3: config diff table of the three paths, curves, conclusion within the window. List missing coverage explicitly, not as "the rest is analogous". - -For the basic definition of STE see [Bengio et al.](https://arxiv.org/abs/1308.3432); it does not decide the surrogate gradient rule for a specific model. diff --git a/.agents/skills/train-inference-consistency/guides/bi-kernels.md b/.agents/skills/train-inference-consistency/guides/bi-kernels.md deleted file mode 100644 index 4b188b4b2e4c..000000000000 --- a/.agents/skills/train-inference-consistency/guides/bi-kernels.md +++ /dev/null @@ -1,143 +0,0 @@ -# BI kernels: which kernel each operator uses and what it costs - -Read [the numerical contracts](../workflow/README.md) first, then build candidates from this page. Each row gives the kernel used under BI=1, the cost of invariance, known pitfalls and directions already tested dead. Numbers come from Nemotron-H (GB200 single node, 4 GPUs) and DS4-Flash (GB200 DP8/EP8); only conclusions are given. When the user asks to optimize only the explicit Top-K/index overhead, do not expand into whole-model tuning; switching DSL does not relax the contract. - -## By operator - -| Operator | BI=1 kernel / approach | Cost and upper bound | Pitfalls | -| --- | --- | --- | --- | -| Attention (FA4) | `FlashAttnFixedSplit`: 32 splits × 640 keys, `MAX_SEQ_LEN` 16384; with FP8 Q/KV the KV dequantization path is fixed | The 16384 limit exists for invariance; exceeding it must raise, not silently change the split. MLA matrix absorption changes rounding; mathematical equivalence does not give train-inference bitwise equivalence | On SM100, do not declare support by bypassing the selector's "backend does not support BI" check. Implement and verify the numerical contract in the backend first, then let the selector declare it; BI0 keeps the default fast path | -| Dense GEMM (FP8 / NVFP4 linear) | fixed split-K; under Nemotron BI small-M does not use CuTe GEMM, FP8 bmm pinned to cutlass, W4A16 uses Humming | tcgen05 hand-written skinny GEMM: a 64-row UMMA tile under "TileM must not depend on M" caps tensor core utilization at 50%; large M with narrow N is about 2× cuBLAS; loading the skinny bi_gemm `.so` on DS4 BI=1 does not change throughput | The small-M gain comes from cluster split-K over DSMEM (S=8 at K=4096), removing the producer/MMA warp `try_wait` peek (~0.5 µs), and `st.async` sent directly to the owner slot (~0.5 µs at M=64); stage count, spin `test_wait` and S=16 gave nothing, do not retry | -| mHC (DS4 TileLang) | BI=1 pins `split_tokens=1` | no extra cost | Pitfall only under BI=0: `n_splits` varies with `num_tokens` and is a `@tilelang.jit` specialization parameter, so every unseen prefill size recompiles for 10–14 s, all DP ranks stall, and every BI=0 comparison is polluted. BI=0 must load the `fixed_splitk` patch and check `tl_after_bench=0`. `mhc_post` spreading h over 128 blocks is numerically right but 40% slower; gains can only come from fusion | -| MoE (NVFP4 W4A16, Nemotron) | CuTe-DSL W4A16 modular MoE (takes the router's top-k); one CTA runs the whole K (`_W4A16_K_TILE=256` fixed), no split-K, no atomic | Per-row K accumulation order is the same under both tile sets; whether a single FP32 element is bitwise identical across tcgen05 instruction shapes (M128/M256, N8/N128) is undocumented and must be measured | flashinfer_cutedsl and Humming are both modular; moving R3 routing capture into the kernel (monolithic, #44214 style) saves nothing | -| MoE (FP8 block, DS4) | backend locked to DEEPGEMM (open source, built-in BI mode, the deployment kernel; FlashInfer TRT-LLM MoE ships as precompiled cubins that cannot be modified or verified as BI); `flashinfer_cutlass` rejects the block-fp8 QuantKey, `triton` resolves to BATCHED_TRITON which does not declare BI → a2a locked to deepep low-latency | BI=1 vs BI=0 gap floor is 33%, 83% of it non-overlappable deepep LL communication (86 serial a2a round trips per decode step); BI=1 does 14% less GPU work than BI=0 but takes 28% more wall clock (overlap ratio 1.33 vs 1.97) | DBO measured −39% (split ubatches read expert weights twice per layer, and weight reads are already at 84% of HBM roofline); TP2/DP4 −17%, TP4/DP2 −28%, DP8 is optimal for BI=1 too; deterministic top-k + histogram radix-select together +6.7% (alone +2.3% / +0.1%, super-additive, measure together); the biggest lever is DeepEP LL over MNNVL instead of RoCE (`VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=1`, `DEEPEP_MAX_NVL_PEERS=8`): gap 33% (start) → 17% (after kernel PRs) → 4.5% (paired, same fabric), finally BI=1 ahead by 3.2%; transport does not touch BI (output bitwise identical across transports) | -| Stable top-k / routing scores | BI grouped top-k, tie-break by request-local index | Deterministic top-k (radix above 4K) 5.9 µs vs a histogram radix-select 56.7 µs | Full output contract below; "select candidates unordered, then sort" is correct only if the candidate set itself meets the contract for cutoff ties; a final stable sort cannot repair an earlier step that randomly dropped equal-score candidates | -| Mamba2 (SSD scan + conv1d) | exact replay ported from vllm#55905 (three extra chunk buffers x/dt/B); a partial chunk is replayed from the FP32 boundary state through the same SSD path | row-gated emit (emit only outputs needed now, update state at chunk boundary) restores Graph, but measured 0.1874 ms/token vs BI0 0.1726, still 8.57% slower; not default | Under BI Mamba2 declares 5 states (conv + temporal×4) but align-mode PC copied only 2 → after a hit, every prompt longer than one block is wrong; now copies the full per-layer declared state list, with a unit test asserting equal counts. The conv window is written tail-aligned by `conv_state.size(-1)` while the kernel reads `[0, width-1)`; when spec decode widens the cache they misalign; write only `[0, conv_kernel_size-1)`. `test_mamba2_ssd_emit` is occasionally flaky because the reference's autotune picks different blocking; pin the reference to the BI config. Only chunk64/8 heads exposes that the cumsum tile changes FP32 prefix sums (#55895). Gating inside the FP32 scan vs rounding first then gated RMSNorm round differently (#56244). Merging #58851 × #57329 directly: 28-shot GSM8K 432 → 4/500, because later groups reused the first group's checkpoint target block; 438 after adding the per-group gather. Under BI Mamba2 rejects the KV connector; real PD disaggregation is not covered | -| RMSNorm / gated RMSNorm | `cuda_rms_norm` (`CudaRMSNorm`), `grouped_gated_rms_norm` custom op, norm/act+quant fusion off, standalone `+quant_fp8` | Training-side TE fused RMSNorm differs from vLLM native RMSNorm by 1 ulp → only 74.4% of in_proj GEMM inputs bitwise identical and 1.7% of FP8 codes flip; so the training side aligns with the vLLM kernel instead of te.RMSNorm | Direct and compiled paths of the same gated RMS kernel differ; the difference disappears behind an opaque custom-op boundary. Review comment: this is a second BI RMSNorm implementation; DS4 reuses `rms_norm_batch_invariant` | -| FP8 activation quantization | standalone `+quant_fp8` op, `use_cuda_fp8_quant` under BI | Audit scale grouping, amax, rounding, clamp and aggregation order separately; the communication mode does not guarantee quant BI | The DS4 vLLM fork's `is_batch_invariant_quant_kernel_enabled()` only checks that the `.so` exists, not the BI switch, so BI=0 runs the unoptimized BI kernel for nothing (nsys 9.9%, 35.7 µs); with the fixed `.so` + `VLLM_BI_QUANT_TOKEN_BLOCKS=128` it is 5.3 µs, DS4 16703 → 18313 tok/s. Before a BI=0 comparison, grep nsys for `vllm::batch_invariant::`, or set `VLLM_BATCH_INVARIANT_KERNEL_LIB=/nonexistent` to fall back to upstream | -| cudagraph | pinned to FULL under BI (on DS4 FULL_AND_PIECEWISE is fastest: 16409 vs FULL 16068 vs PIECEWISE 10423 tok/s, BI=1 bitwise PASS) | FULL has worse TTFT (1792 vs 764 ms): on DS4 FULL lands as FULL_DECODE_ONLY and mixed prefill is always eager | Piecewise breaks BI through the switch "does this step run the graph or fall back to eager"; the criterion is total token count ≤ `max_cudagraph_capture_size`. Graph and eager paths are each BI but not bitwise identical to each other. Unrelated to padding. It reappears under FULL when the capture limit is set below `max_num_seqs × decode_query_len`. A short-prompt probe passing in all three modes does not prove PIECEWISE safe | -| determinism | `torch.use_deterministic_algorithms(True)` + `CUBLAS_WORKSPACE_CONFIG` (rewritten to `:16:8` by vLLM BI init); the mamba_ssm SSD backward uses atomics for dA/dD/ddt_bias, call `set_deterministic_mode(True)` explicitly | affects only backward step 0 | see the section below | -| KV block zeroing (DS4 multi-segment KV) | unrelated to BI, but both arms need it | The DS4 vLLM fork's `KVBlockZeroer` picks a 16 B program block by the gcd of per-segment page sizes, 1–12 ms per step (2.27 ms/step = 7.6% of wall clock); upstream #52058 uses a 3-D grid + masked tail; chunk mode 60 µs | Visible only as `_zero_kv_blocks*` in nsys; for a memset-like kernel taking ms, first check `blk_size` ≤16 | -| Prefix caching (Mamba align mode, BI=0) | precompute multi-group alignment indices once (#58851 fits best), skip align postprocess without spec decode, skip preprocess/precopy when the host determines nothing crosses a block; on prefill export block-boundary checkpoints inside the scan (#57329) | PC on as is: BI0 throughput −4 to −20%, TPOT +13 to 21%, c=1 TTFT roughly doubles (eager launches per step 9→37, 8K prefill split in two); after the fix BI1 is about on par with PC off | The BI path must export checkpoints separately or BI1 is silently wrong; a chunk truncated by the token budget must also export its last checkpoint (under DP+EP one extra prefill step on any rank stalls global decode) | - -## Reference: DS4-Flash per-operator choices - -The full per-operator table behind the blog's summary table (consistent arm = BI=1). "Same as left" means the training forward calls the same vLLM kernel as rollout. - -### Dense, norm and mHC - -| Operator | rollout (consistent) | Training | What is fixed / given up; cost | -| --- | --- | --- | --- | -| Embedding | Table lookup | Forward: Same as left
Backward: PyTorch | Naturally consistent | -| mHC residual mixing (pre / post / Sinkhorn) | Split TileLang mHC kernels, `split_tokens=1` | Forward: Same as left
Backward: Hand-written VJP, Sinkhorn iterations unrolled, FP32 | Gives up the baseline's fused post+pre kernel (split count varies with token count)
Cost: Both sides lose one fusion | -| RMSNorm | Batch-invariant RMSNorm (one program per row) | Forward: Same as left
Backward: FP32 VJP | Per-row reduction order fixed
Cost: Small | -| FP8 linear layers (q_a/kv, q_b, wo_b, shared expert) | DeepGEMM FP8 GEMM, batch-invariant mode; weights in 128×128 blocks, activations quantized dynamically per token in groups of 128 | Forward: Quantizes BF16 master weights on the fly to the same block FP8, same DeepGEMM
Backward: Transformer Engine (TE) BF16 GEMM w.r.t. BF16 master weights (STE) | Gives up the baseline's FlashInfer switch at M<32; scale is UE8M0 on SM100 (8-bit, exponent only, i.e. a power of 2), FP32 on SM90
Cost: rollout: no faster small-batch kernel; training: FP8 weights repacked every forward (including recomputation) | -| Small linear layers (compressor, indexer weight projection) | FP32-output cuBLAS GEMM / batch-invariant persistent Triton matmul | Forward: Same as left
Backward: BF16 | Summation order along K fixed
Cost: Small | - -### Attention (DSA sparse attention + FP8 KV cache) - -| Operator | rollout (consistent) | Training | What is fixed / given up; cost | -| --- | --- | --- | --- | -| q-norm + RoPE + KV write | Fused kernel quantizes K and writes the FP8 cache | Forward: Same kernel writes a private training-side FP8 cache, then reads it back
Backward: FP32 RoPE / norm in PyTorch (straight-through) | K takes the same FP8 round trip
Cost: Training: one extra cache and dequantization | -| FP8 KV cache | Stored as FP8; prefill and decode dequantize to BF16 before attention | Forward: Same as left
Backward: — | Gives up the default FP8-reading decode kernel (blog 4.2)
Cost: rollout decode | -| CSA compressor | vLLM compression kernel, state reset per request | Forward: Same kernel's values, functional graph's gradients
Backward: autograd | In-request state starts as in prefill
Cost: Training: loop over requests | -| DSA indexer scoring | DeepGEMM FP8 MQA logits; prefill reads contiguous KV, decode reads paged KV | Forward: Same as left, per request
Backward: None (indexer not trained) | Paged and contiguous reads bitwise identical | -| Top-K selection + index arrangement | Batch-invariant Top-K (64-bit key + radix select), sorted descending, invalid entries last | Forward: Same as left
Backward: None | Ties go to the larger in-request index (blog 4.3)
Cost: One extra sort per side, recovered by fusion | -| Sparse MLA | `flash_mla_sparse_fwd`; prefill one segment per request, decode same sparse kernel | Forward: Same as left, once per request
Backward: cuDNN frontend DSA backward | In-request geometry ignores neighbors; decode emits 1 token per request per call
Cost: rollout: multiple prefill launches; training: serial over requests | -| o_proj | SM100: DeepGEMM FP8 einsum; SM90: grouped Triton W8A8 | Forward: Same as left
Backward: TE BF16 grouped GEMM + inverse-RoPE VJP | Aligned per architecture; no bitwise identity across architectures
Cost: Loop over groups on SM90 | - -### MoE - -| Operator | rollout (consistent) | Training | What is fixed / given up; cost | -| --- | --- | --- | --- | -| Router scoring + expert selection | Same FP32-output gate GEMM for all batch sizes; sqrt-softplus + bias selects top-k | Forward: Same as left, routing recomputed bitwise
Backward: Expert IDs fixed, FP32 gradient w.r.t. routing weights | Gives up the small-batch gate kernel; **R3 is not needed** (blog 4.4)
Cost: Small | -| Dispatch | DeepEP low-latency (LL below), fixed per-rank capacity, all ranks segment in sync; FP8 quantization uses elementwise round-to-nearest | Forward: DeepEP normal, plus one dispatch of routing metadata (fingerprint, expert, weight)
Backward: DeepEP backward | Experts get bitwise-identical FP8 values and scales either way
Cost: rollout: LL communication cannot overlap compute; training: metadata dispatch and hash matching | -| SiLU · mul · quantization | Batch-invariant fused kernel: product rounded once to BF16, then quantized in groups of 128 | Forward: Same as left (mandatory)
Backward: FP32 VJP | Rounding position fixed (blog 4.4)
Cost: Small after optimization | -| Routed expert GEMM | DeepGEMM masked grouped GEMM, static capacity as scheduling hint | Forward: DeepGEMM contiguous grouped GEMM, per-expert rows aligned to 128
Backward: TE BF16 grouped GEMM | Gives up the baseline's FlashInfer TRT-LLM MoE (ships as precompiled cubins; cannot be modified or verified as batch-invariant)
Cost: rollout: a main throughput gap; training: alignment padding | -| Combine | DeepEP LL combine, fixed merge order | Forward: FP32 accumulation in top-k slot order on the token's owning rank, then BF16
Backward: Gradient scatter | Merge order by slot, not arrival order or rank
Cost: 43 layers × 2 serial all-to-all per step | -| Shared expert | FP8 linear layer; auxiliary stream under LL + MNNVL (multi-node NVLink) | Forward: Same as left
Backward: BF16 (STE) | Stream changes overlap, not values
Cost: A source of recovered speed | - -### Output layer and infrastructure - -| Item | rollout (consistent) | Training | What is fixed / given up; cost | -| --- | --- | --- | --- | -| lm_head + logprob | Batch-invariant matmul; FP32 batch-invariant log-softmax | Logits per 8192 tokens, vLLM logprob value; backward via PyTorch log-softmax on the same logits (straight-through) | temperature must be a scalar
Cost: Training: each chunk's FP32 log-softmax kept for backward | -| Cross-GPU reduction | TP=1, no all-reduce; NCCL channel count and algorithm pinned as a guard | TP=1; CP exchanges boundaries only, no reduction | TP>1 not supported (blog 4.5)
Cost: Cannot split with TP | -| CUDA Graph | Kept, with the right mode (blog 4.6) | — | Graph and eager paths bitwise identical | -| Weight sync | Reload copies new values to original addresses, refreshes derived weights | Exports block-FP8 weights | Graph-referenced addresses unchanged
Cost: One export and transfer per step | - -Details that back the table: - -- **Scales and precision.** DeepGEMM scales are powers of 2: UE8M0-packed on SM100 (a Triton kernel emits packed weights), FP32 on SM90 (vLLM weight post-processing). Router gate and compressor projection output FP32; indexer weight projection and lm_head are BF16. Before the first update the checkpoint scales are reused; afterwards they are recomputed from amax. FP8 weights are not cached by default, so every forward including recompute repacks them. -- **FP8 KV layout.** Per token per layer: 448 FP8 NoPE dims, 64 BF16 RoPE dims, one UE8M0 scale per 64 elements, 1 alignment byte = 584 bytes (BF16: 1024). The sparse decode kernel must be selected explicitly; the BI switch alone does not change the default FP8-reading FlashMLA decode. The CSA compressor zeroes state at each request start. -- **Top-K key.** 64-bit key: high 32 bits = order-preserving transform of the score, low 32 bits = request-local index; radix select at 8 bits per round. Deterministic Top-K and the Top-K + index-sort fusion gave +2.3% and +0.1% alone but +6.7% together; the interaction is unexplained, so measure them together. Under DCP>1 the cross-GPU candidate merge uses the opposite tie rule; unsupported. The indexer auxiliary loss does not support CP>1. -- **Expert GEMM.** Contiguous layout aligns each expert's rows to 128. LL quantization switched to elementwise round-to-nearest; training quantizes the contiguous layout under the LL dispatch convention. The SiLU·mul·quant kernel has fixed launch geometry, is mandatory on both sides (training refuses to start without it), and uses different SiLU formulations on SM100 and SM90 (no cross-architecture bitwise promise). Backward: TE BF16 grouped GEMM, second GEMM dgrad/wgrad first, SwiGLU recomputed from saved gate_up in FP32. -- **Layouts tried for performance only.** TP2/DP4 17% slower, TP4/DP2 28% slower; DP8 is best for BI=1. -- **Cross-GPU.** Training dispatches routing metadata once more (per route: 16-element token fingerprint, expert index, FP32 routing weight), matches outputs by bitwise hash, and scatters with a deterministic Triton kernel. NCCL channel count, thread count, protocol and algorithm are pinned and symmetric-memory / custom all-reduce are off; under TP=1 this is a guard, not the mechanism. CP pads each sequence, splits contiguously, exchanges boundary hidden states and all-gathers compressed KV rows. -- **CUDA Graph.** Workspaces are allocated in warmup before capture and locked afterwards; the Top-K row-start buffer is a resident zero buffer. Weight-derived mHC parameters are recomputed in place on reload. - -## The full output contract of stable Top-K - -What orders equal scores (e.g. the request-local index)? Take it from the training side or current baseline; do not invent a tie-break. Preserve: - -- Comparison key for equal scores, positive and negative zero, NaN/Inf (if allowed); a global batch offset cannot replace the request-local index. -- Output indices and their order, not just the same set; swapping equal-score indices changes later floating-point reduction order. -- valid length, mask, invalid sentinel, K larger than the valid count, empty rows, padding, tail blocks. -- prefill/decode index coordinates, cache page mapping, the position of the same semantic token. - -Check the candidate for repeated allocation, sort input range, radix/merge passes, workspace, launch/grid/block; prove order preservation before deleting a pass. For any CUB/custom sort or selection, verify the binding and the real lifetime of its temporary storage; a library's stability claim does not replace integration tests. A pre-generated index/workspace saves overhead but must not overwrite in-use content under multiple streams, concurrent requests or Graph replay. - -## Index / elementwise fusion - -Limit scope to confirmed hotspots such as gather lengths, row start/end, sparse index concatenation and invalid-row zeroing. For every row write down input coordinates, whether the interval is half-open (closed left, open right), output capacity, valid range, and who writes the padding. - -Fusion can save a launch or temporary tensor, but can also raise register pressure or hurt memory access. Measure the gain on all real shapes; the number of changed lines says nothing. Shrink the grid and loop stride together so the full output is covered; fill a reused buffer with poison values to catch unwritten tail/empty rows. - -Integer indices are not automatically correct: off-by-one, request offset, CP boundary, overflow and uninitialized rows can bias a whole sequence. If fusion reaches into dequant/norm, keep the original dtype conversions and rounding positions. Unit-test Top-K and index post-processing separately, then combined, then module forward, then engine gates. Record whether the optimization also speeds up BI0: dynamic block/grid selection on a shared path is a general optimization and cannot all be counted as BI-specific overhead reduction. - -## Split-K and the M of grouped GEMM - -A GEMM's M is usually the number of rows processed; masked grouped GEMM computes several experts in groups and uses only each group's valid token rows. `expert_num_tokens.max().item()` pulls the device max token count to the host, commonly as a capacity/scheduling hint; it does not "fix the shape" or "guarantee BI". The readback synchronizes and usually cannot be used in CUDA Graph capture. Before removing it, read the API: is the parameter a correctness bound, a capacity, or a tuning hint? Restore an estimate/safe upper bound the library supports, and cover heavy load and empty experts. Do not claim without source evidence that underestimating M makes the default kernel access memory illegally, and do not set an arbitrarily small value for speed. - -One launch per token avoids some batch-dependent reduction/scheduling but means many launches and low utilization. Batched launch with fixed `n_split` may keep the per-row numerical order, but the K range of each split, the accumulation order within it, the accumulator dtype, partial reduce order, tail/padding and actual kernel specialization must still be fixed or verified. "Fixed K/n_split" is a candidate mechanism, not a proof of prefill=decode: different tiles, epilogues, precision or preprocessing can still differ at the same split count. Test identical operator inputs and the full semantic path. - -## BI coverage of attention - -Attention BI covers QK, online softmax, PV and the merge order of blocks/splits. Beyond the split count, the tile scheduler, Q-stage, cluster/scheduling choice, page layout, mask/sink, score scale and accumulator/rounding can depend on batch/length. Passing small batches does not cover long context, prefill chunk and decode, Graph buckets, mixed batches, page boundaries, prefix cache hit/miss. Graph+prefix enabled does not mean one test executed all these branches; record the coverage. - -## Observation in Graph mode - -Save the resolved Graph mode, capture sizes, prefill/decode/mixed step types and evidence of kernel calls. A successful capture guarantees neither eager-equal numerics nor that every step uses the same kind of graph. A Python hook usually runs at capture, not replay; a capture dump cannot stand for all replays. - -To observe replay, use controlled device-side outputs or static diagnostic buffers, and state whether instrumentation changes the capture. Eager hooks can localize, but the final recheck must return to the original Graph mode. Even if the compiled path with Graph off matches production Graph on fixed samples, eager may still differ; keep artifacts for all three. A hook not firing on the compiled path does not mean the layer did not execute. - -Turning the graph off changes speed and therefore online batching. First fix each step's tokens, requests/neighbors, padding and cache state, then compare the supported modes (eager/FULL/PIECEWISE). Use a 2×2 experiment to rule out the dispatch/split plan the switch changes at the same time. A short filler may put a mixed step into a capture bucket while a long filler goes eager. Prefill/decode overlap differs between the full and a small model; a small model passing does not cover untriggered production mixed steps. For intermittent failures keep multi-round samples and the trigger rate; do not pick only successful rounds. - - - -## determinism takes effect only on the main thread - -In PyTorch 2.13 (nv26.07) Inductor config overrides are thread-local (ContextVar). `torch.use_deterministic_algorithms(True)` sets `torch._inductor.config.deterministic` on the main thread; on autograd's CUDA device thread it stays False. A `torch.compile`d norm VJP first compiles on the autograd thread during backward, sees `deterministic=False`, and the generated looped-reduction kernel enables `dynamic_scale_rblock`: besides `R0_BLOCK=1024` it adds candidate 512, and the first launch picks by timing. The two reduce in different orders and differ at bit level. Only at the end of the dynamo compile is that thread set to deterministic; later compiles have one fixed config. Result: only backward step 0 is not bitwise reproducible across runs; forward logprob and the RL closed loop are unaffected. - -Apply both fixes: `torch.compile(..., options={"deterministic": True})` at the compile site (environment-independent), and `TORCHINDUCTOR_DETERMINISTIC=1` in the recipe environment (covers any first compile on the autograd thread, including the vLLM forward function compiled inside recompute). Also: verl's full_determinism actually sets `warn_only=True`, and the mamba_ssm SSD backward must enable deterministic mode explicitly. - -## How to break down the cost from BI0 to BI1 - -Draw the actual call chain first: what the default fast kernel, the BI branch and the extra preprocessing/sorting/postprocessing each are. Do not infer "split-K was changed" or "a temporary Triton was used" from function names; compare source, binding, compiled artifacts and real launches. - -You may copy the BI0 implementation into a standalone candidate and add the deterministic/BI conditions one at a time for a bisectable A/B. It need not live in the same file, but the default path must stay unchanged. If the BI0 algorithm cannot meet the required tie/reduction contract, switching algorithms is allowed; state the cost. - -Fairness between the arms comes before any kernel conclusion: same fabric (on DS4 two environment variables closed most of the remaining gap), same node, same image, BI=0 not running the BI kernel, no JIT stall, and the KV zeroing fix in both arms. Nemotron's BI1 vs BI0 E2E throughput loss is the median over 5 pairs, with single pairs reported separately. - -## Minimal record for each candidate - -```text -baseline / candidate source+binary hashes -one intended change; BI-specific or common -same-input bitwise + BI composition + prefill/decode gates -BI0 unchanged + real shape/tie/mask/tail + Graph/workspace reuse -paired latency, launch/allocation counts, repetitions -pass / fail / inconclusive; retained or rejected; remaining scope -``` - -On failure, save the minimal counterexample, then withdraw from the default path. When the user requires all micro gates to pass before E2E, do not skip them; performance numbers do not cover a correctness failure. Do not quietly remove an existing stable-sort patch to reproduce a long-length mismatch; that tests a regression, not a defect of the current best code. diff --git a/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md b/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md deleted file mode 100644 index 76b14e4e9302..000000000000 --- a/.agents/skills/train-inference-consistency/guides/checkpoint-resume.md +++ /dev/null @@ -1,304 +0,0 @@ -# Checkpoint, resume and memory budget - -Use this guide to get checkpoint save, load and resume right in long RL runs, and to budget host (CPU) memory peaks and growth. For per-phase GPU memory peaks see [Locating GPU OOM and staging peaks by lifecycle](#locating-gpu-oom-and-staging-peaks-by-lifecycle). - -## Find by symptom - -| Symptom | Where to go | -| --- | --- | -| OOM while loading a checkpoint | [Load](#4-load-enough-capacity-does-not-mean-enough-for-the-transient-peak) | -| Train-inference mismatch on the first step after restore | [13](#13-stale-quantization-scale-on-the-first-step-after-restore) | -| Resume is set but training starts from step 0 | [14](#14-resume_modeauto-starts-from-step-0) | -| Host RSS grows every step | [3.3](#33-locating-slowly-grows-every-step) | -| Planning host memory for a checkpoint in the hundreds of GB | [4](#4-capacity-planning) | -| vLLM fails to start when colocated with training (KV budget too small) | [5](#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated) | -| GPU OOM during init / export / wake | [Locating GPU OOM](#locating-gpu-oom-and-staging-peaks-by-lifecycle) | - -## Terms - -- **pinned memory:** page-locked host memory for direct GPU DMA (`cudaHostAlloc` / `cudaHostRegister`), cached by PyTorch's caching host allocator. -- **cgroup:** limit `memory.max`, usage `memory.current`, split in `memory.stat` into anon, file (page cache) and shmem. -- **RSS and PSS:** RSS counts shared pages fully in every process; PSS splits them, so only summed PSS approximates the physical total. -- **DCP / Megatron dist_ckpt:** PyTorch and Megatron distributed checkpoint formats; each rank writes its shard, `.metadata` last. -- **staging / D2H:** the host buffer GPU tensors are copied into before writing to disk. -- **Ray object store (plasma):** Ray's shared memory for passing objects, mapped into each worker and charged via PSS. -- **old-logprob:** the logprob training recomputes for rollout tokens with the pre-update policy. -- **Ray OOM threshold:** `RAY_memory_usage_threshold`, the usage fraction at which Ray kills workers. Default 0.95. - -## Key rules - -- Check file, memory and derived state separately; complete files can still mismatch on the first step after restore. -- Capture the first exception (errno, rank, file, offset), not the last message. -- Completion marker last; clear the old directory before overwriting; all ranks check and fail together; a missing file raises instead of falling back to a fresh optimizer. -- Under EP4 + FSDP2, ranks' expert shards in DCP share key and offset, so restore reads another rank's experts. Put the EP rank in the key (port the reference fix), add a 4-GPU bitwise read-back unit test, and raise on EP size mismatch. -- Rotation deletes only tracked, complete, idle directories; in long runs keep one checkpoint every 20 steps outside rotation. -- Set rollout `gpu_memory_utilization` per training mode ([5](#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated)). -- Save fixes pass [the acceptance in 3.2](#32-reproduce-and-discriminate). Before a long run on a new save path or layout, run to the first save and do one kill→resume. -- Resume is proven by matching per-rank digests in a new process, a byte-identical first export, an aligned first old-logprob and one non-zero update; `resume_mode=auto` alone proves nothing. -- Read host memory as anon/file/shmem; `memory.current` includes page cache, and Ray kills on it. -- Pinned buffers round up to a power of 2 (Nemotron: 146.9 GiB pinned for 88.2 GiB of optimizer state, +67%); allocate large long-lived pinned blocks at exact size. -- The save step is the peak; a higher Ray threshold or smaller microbatch does not fix headroom. -- For per-step growth, use the per-process PSS floor and a small CPU-only repro before changing code. -- Verify per save path that the host cache is released after save. - -## Checkpoint and resume - -### 1. Three independent contracts - -Classify the problem before calling it "the checkpoint is broken": - -| Contract | Typical failure | -| --- | --- | -| File: serialization and publish | Disk write error; an incomplete directory treated as complete | -| Memory: transient peaks of save and load | Killed by Ray after the write finishes; cgroup OOM during load | -| Derived state: the first export after restore | Files complete, but the first step after restore mismatches ([stale quantization scale](#13-stale-quantization-scale-on-the-first-step-after-restore)) | - -The third class is a derived-state lifecycle bug, not file corruption. - -### 2. Which state exact resume needs - -| State | Requirement | -| --- | --- | -| Model parameters or LoRA adapter shards | Saved per rank; strict load that rejects missing/unexpected keys ([training forward](training-forward.md)) | -| Optimizer state, including the CPU-offloaded part | Same step as the model. Nemotron chunked optimizer offload: 88.2 GiB per rank, 94.7 GB per rank on disk; DS4.1 saves LoRA parameters and Adam `exp_avg` / `exp_avg_sq` in FP32 | -| LR scheduler | Check the saved value, not the source default (DS4.1 checks its constant, no-warmup schedule per rank) | -| RNG | Four streams: CPU, CUDA, Python, NumPy | -| Counters and data position | global step, optimizer step, data cursor; DS4.1 also saves the cumulative count of effective-signal steps | -| Run contract | Hashes of code, data and recipe, compared item by item on restore; reject a change in total steps or data by default. To extend training, declare a new phase explicitly; do not edit the old contract file or add a generic ignore switch | -| Completion marker | `.metadata` / tracker (e.g. `latest_checkpointed_iteration.txt`), written last, see 3.4 | - -Not in the checkpoint, rebuilt after restore: - -- **Parameter-derived caches:** quantization scale, packed/transposed/fused weight, version markers (see 5). -- **Records outside training:** the metrics jsonl is overwritten on resume. - -### 3. Save - -#### 3.1 Capture the first error, do not fix the last error message - -Keep the writer callback's original exception / errno, rank, path, offset, requested / returned bytes and exception chain, which DCP and Ray wrappers may drop; also filesystem/mount, free space, quota and cgroup OOM events, not just the final `ActorDiedError`. `unexpected pos` can come from disk space, quota, permissions or I/O failure; the first exception tells them apart. - -#### 3.2 Reproduce and discriminate - -A fast repro keeps the real tensors, serializer and concurrency, skips rollout/backward and repeats only save, for a bounded number of rounds sized to the failure probability; stop at the first failure and keep the scene. DS4 model-only save failed on the 3rd attempt on one path and the 6th on another; Nemotron with host cache release failed on the 2nd-4th save. - -**Acceptance for a save fix: N consecutive saves + kill→resume→save again.** Also run it before the first long run of a new save path, parallel layout or node count; running to the first save point is enough. - -Order of discrimination: - -1. Same buffer/serializer, target filesystem versus node-local disk as the control; read the errno of the first failure. -2. Control concurrency, pinned/pageable and memory peak; if the serializer changed along with the filesystem, the result cannot be attributed to one factor. -3. Check growth of threads, processes and file handles, separately from the write error. -4. Only then read kernel / filesystem client code; a matching historical issue is supporting evidence, not proof of a fix. - -#### 3.3 Reliable publish - -Serialize on node-local disk and close the file; copy to a **temporary name on the target filesystem**, handle short copies, flush/fsync as durability requires, verify integrity, then rename within the same filesystem. Publish metadata/tracker/completion marker only after all ranks finish. A move across filesystems is not an atomic rename. - -#### 3.4 Commit protocol: write the completion marker last - -| Rule | If you do not | -| --- | --- | -| Write `.metadata` / tracker last, after every rank's shards and auxiliary files (optimizer, RNG, scheduler) succeed and are confirmed collectively | A rank fails, the marker already exists, and restore reads incomplete state | -| Before overwriting a step directory, delete all old files first (uncommit) | Old optimizer / RNG silently pair with the new model | -| Before restore, all ranks check file completeness together and fail together | The rank missing a file exits; the others hang until the NCCL timeout | -| Ranks with no shard to write still join the same collective | Collective order diverges across ranks and hangs | -| A requested load with a missing file raises an error, not just a log line | Training silently continues with a fresh optimizer | - -Acceptance: a save failing midway (Nemotron hit a real ENOSPC) makes all ranks exit together without advancing the tracker, and training resumes directly; new code loads an old-code checkpoint. Never kill a job mid-write; verify the checkpoint is complete first. - -#### 3.5 Verify a written checkpoint - -Check the expected shard set, non-empty sizes, leftover temporary files, checksums where needed, and do one real full restore. A correct save digest is not a restore test (DS4.1 checkpoints 20/24 had correct save digests but were never restored in a new process, and the report says so). - -### 4. Load: enough capacity does not mean enough for the transient peak - -Load may copy the full optimizer into anonymous memory, and FSDP CPU offload may create a second set of pinned storage. Accept by loading model + optimizer + RNG/extra at real scale, not one small shard. - -Read a checkpoint with the image that wrote it: DS4.1 host PyTorch 2.11 cannot deserialize the DTensor `_FlatLayout` written by the image's 2.13; it reads only inside the image with `torch.load(weights_only=True, mmap=True, map_location="cpu")`. Do not modify the checkpoint or turn off `weights_only` to make it load. - -### 5. Derived state after restore - -DS4 chain: initial HF load saves FP8 source scales → DCP restores the trained BF16 parameters → the first export still uses the old registry → the actor re-quantizes by parameter version, but rollout uses the old scales. The first resumed step mismatches; once the next optimizer hook clears the cache it matches. The prediction "it recovers on the next step" plus a before/after-fix comparison on the same checkpoint locates the cause. - -Fix it in the lifecycle that owns the state: after load/update, invalidate or rebuild parameter-derived caches (scale, packed/transposed/fused weight, version markers), reusing an existing hook without side effects such as a real optimizer update. Do not hide it by waiting a step, disabling resume or clearing all caches by hand. DS4's minimal patch calls the existing post-update hook ([Megatron-LM #197](https://github.com/ISEEKYAN/Megatron-LM/pull/197)). List every parameter-derived cache and its invalidation code, add the missing calls and one regression test. - -### 6. How to prove resume is correct - -1. **Before launch.** For node-local checkpoints, every rank's node in the new allocation has its shard; the tracker / latest pointer points to a complete checkpoint; the launcher's resolved `resume_from_path` and global step are as expected. Confirm the restored step before the first training log line; `resume_mode=auto` is not evidence. -2. **New-process restore.** All worker PIDs changed; per-rank adapter / optimizer / scheduler / RNG digests match the save-time digests item by item. -3. **First export and old-logprob.** Exported files are byte-identical to those of the saving step; also check the exporter scale / policy fingerprint and the full valid logprob. Checking only that the second step returns to 0 misses the bug in 5. -4. **One non-zero update.** A zero-advantage step only exercises the flow (two DS4.1 resume segments had zero gradient on all ranks at step 2, proving nothing about restored adapter values). -5. **Gate coverage.** step0 cold load, first export/old-logprob after restoring a non-zero step, one later update, address / state validity under Graph and sleep/wake, new code reading an old checkpoint. -6. **What acceptance compares.** Restored per-rank digests equal save-time digests; first exported files are byte-identical to the saving step's; per-step `probs_diff` after restore is 0 under BI=1 (Nemotron: 0 on the first step after kill→`resume_mode=auto`). The resumed trajectory need not match an uninterrupted run bitwise (rollout sampling is not bitwise across processes). There is no "metrics are close" threshold; reward curves are reference only, on the same step range and basis. - -### 7. Rotation, disk and time cost - -| Item | Record | -| --- | --- | -| Size per checkpoint and save time | DS4 full32: 99 files, about 3.2 TiB, save about 269 s. Nemotron-H 30B-A3B, 4×GB200: 353 GB (EP4 + FSDP2) or 429 GB (PP4 dist_opt), save 72-152 s, about 2-3% of step time when saving every 5 steps. Layout-specific; state these for your model in the plan | -| Rotation | The mLite engine ignores `max_actor_ckpt_to_keep` | -| Checkpoint used for eval | Rotation may have deleted the step to evaluate; a checkpoint from another chain cannot substitute | -| Resume on a different node after node failure | Copy checkpoint, dequantization cache, pinned image and dependency cache to the new node before launch; disable automatic requeue; change only the node selection; write a transfer record | -| Paths for two-arm experiments | BI=0 / BI=1 output and checkpoint paths are separate; renaming the run breaks a resume chain in progress ([recipe and acceptance](recipe-and-acceptance.md#rules)) | - -## Memory budget: host and GPU - -Host memory fails in two ways in long RL runs: phase **peaks** (save / load / sleep) hitting the limit, and slow per-step **growth**; they are measured differently (1-4). GPU memory is a third class: with training and rollout colocated, rollout `gpu_memory_utilization` is set per training mode (5). - -### 1. Who holds host memory - -| Source | Mechanism | -| --- | --- | -| Pinned buffer of optimizer offload | The caching host allocator rounds up to a power of 2, so the footprint is tens of percent above the nominal state; summed across ranks it pushes residency before save close to the limit | -| Parameter shards of FSDP CPU offload | Each rank's shard stays resident on CPU; doubles if load creates a second pinned copy | -| D2H staging during save | Stays in the pinned cache pool; whether `torch.accelerator.empty_host_cache()` releases it must be verified per save path | -| CPU backup of serving sleep | L1 sleep keeps one full backup per rank; the pinned cache is not returned after wake | -| Ray object store | plasma mappings count into each process's PSS as shared pages; node shmem does not grow with it (3) | - -`param_offload` trades CPU pressure for GPU memory; optimizer offload is separate. Change one offload switch at a time and measure the phase peak before and after. - -### 2. Limits and who does the killing - -- **cgroup.** The job limit `memory.max` is usually below physical memory; budget against it. A cgroup OOM shows an `oom_kill` event and process SIGKILL, with no Python stack. -- **Ray memory monitor.** Kills workers at `RAY_memory_usage_threshold`=0.95 by default; log line `exceeded threshold of B (95.0% ...)`. Its usage figure includes page cache. Raising the threshold does not fix headroom: Nemotron at 0.985 was still killed during save. -- Distinguish GPU allocation exception, host cgroup OOM and Ray kill. A CUDA allocator allocation-failure / retry warning is not a terminal OOM and the step may complete, but it does show memory pressure. - -### 3. How to measure - -#### 3.1 One record per phase - -Log one line at load, save, sleep, wakeup, reset KV cache and generation start / end, with monotonic time, node / rank / PID and policy version. Per node, record cgroup `memory.current` / peak / events and `memory.stat` anon/file/shmem; sample process RSS/PSS, threads and FDs as needed, not with high-frequency scans of all processes. - -`generate_sequences` includes wake/init. An error inside it does not prove token generation started, let alone a backward/optimizer OOM. - -#### 3.2 Pitfalls in the readings - -| Reading | Problem | -| --- | --- | -| `memory.current` | Includes page cache, which masks per-step growth; but it is what Ray and the save peak see | -| `free`, `psutil.virtual_memory().used` | System-visible memory, not process RSS or job cgroup usage | -| `scontrol` FreeMem | Excludes reclaimable cache | -| `ru_maxrss`, Slurm MaxRSS | Cumulative peak, not current RSS | -| VmPin/VmLck | Cannot alone infer CUDA pinned | -| Sum of per-process RSS | Counts shared mappings repeatedly; use PSS | - -Pinned memory from cudaHostAlloc shows up as shmem in cgroup `memory.stat` (measured on Nemotron). - -#### 3.3 Locating "slowly grows every step" - -1. Use node `MemTotal−MemAvailable` and cgroup anon/file/shmem to find which class grows. -2. Sample PSS from `/proc//smaps_rollup` per process once a minute, recording only comm / argv0; take each step's minimum as the floor and find processes that rise monotonically. Within-phase peaks swing with offload and are useless here. -3. For a rising process read `smaps`. On someone else's job or a running job, only read `/proc`; do not attach or send signals. -4. Reproduce the candidate mechanism in a small CPU-only process, then decide where to change code. - -A floor that drops after resume (Nemotron run 1: about 455 back to about 433 GiB) shows the growth accumulates over the process lifetime. - -### 4. Capacity planning - -Budget the save-step peak as the sum of what is alive at the same time, against cgroup limit × Ray threshold: - -```text -peak ≈ resident pinned optimizer offload (actual footprint after allocator rounding) - + save D2H staging (and the part of it left in the cache pool) - + page cache of the newly written checkpoint files - + CPU backup of serving sleep and unreleased pinned cache - + resident amount of the other processes (rollout, driver, Ray workers, etc.; take the floor per 3.3) -limit = cgroup memory.max × RAY_memory_usage_threshold -``` - -Use measured values wherever possible. - -- **The save-step peak is closest to the limit.** The long-run plan states, per training mode, the measured cgroup peak at the save step and in the kill→resume segment, and which load is still unmeasured. -- **Estimate from headers before running.** Safetensors headers alone give the expanded size per dtype (full DS4.1 checkpoint: about 1421.66 GiB as BF16, including auxiliary content). It is not peak RSS; add serializer copies and mapped source pages. -- **A slope only hints at headroom risk.** Continue from a verified checkpoint and observe long enough; do not extrapolate an OOM step linearly from a short timeline. -- **On OOM, do not shrink the microbatch or raise the guard first.** Find the phase and which objects are alive together. - -### 5. GPU memory: set gpu_memory_utilization per mode when training and rollout are colocated - -The training process keeps part of its GPU memory during rollout. Nemotron single node 4×GB200: - -| Training mode | Held per GPU by training during rollout | rollout `gpu_memory_utilization` | -| --- | --- | --- | -| EP4 + FSDP2 (mLite; BI=1 arm default `ACTOR_OPTIMIZER=fsdp2`) | Not measured separately | 0.65 (verified) | -| dist_opt (BI=1 arm optional `ACTOR_OPTIMIZER=dist_opt`; numbers from PP4 dist_opt measurements) | 54-64 GiB (FP32 quantization master resident on GPU) | 0.7, 0.6 and 0.65 fail to start; measured 0.55 | -| EP4 + Megatron (mbridge) dist_opt + chunked optimizer offload (BI=0 arm, QAT + R3) | Not measured separately | 0.45 (current recipe value) | - -FlashInfer autotune warmup needs about 8 GiB more, another direct cause of failing to start. In a two-arm comparison, differences such as gpu_mem 0.45/0.65 are confounders and are listed first ([recipe and acceptance](recipe-and-acceptance.md)). - -Rules: - -- Write the value into the recipe per training mode, not as one global default; the recipe's hardware comment lists only verified topologies. -- Check the resolved value inside the container and the actual KV budget in the vLLM log, not the outer shell variable. -- `param_offload=True` saves GPU memory but moves parameters, changing the `data_ptr` of the quantization master so the pre-export check errors; as with DS4, enable only optimizer offload. -- What sleep level 1/2 release, and why wake may still find no GPU memory: [rollout engine §5](rollout-engine.md#5-lifecycle-and-gpu-memory-ledger) and [Locating GPU OOM](#locating-gpu-oom-and-staging-peaks-by-lifecycle). - -## Locating GPU OOM and staging peaks by lifecycle - -Use this when the shard estimate fits but init, export or wake still OOMs. Find which step allocates and what is still alive, and change that boundary; do not lower the batch or change sharding first, and never blame update for a failure before update. - -### 1. The minimum to record - -```text -phase / process / rank / topology -requested bytes / allocated / reserved / device free -host RSS / container or Ray limit -last completed phase / first failing call -``` - -Mark missing fields as not obtained; do not guess. If output just stops with no error, check CPU/GPU/I/O, log timestamps and phase (reading weights, compiling, collective, generating); slow is not a hang. - -### 2. The memory peak is set by "the objects alive at the same time" - -Parameters divided by GPU count estimates only part of the steady state and misses these peaks: - -| Phase | Objects and boundaries to look for | Minimal verification | -| --- | --- | --- | -| Init / FSDP wrap | Full materialization before sharding; largest single parameter; CPU-resident/tied parameters | Whether meta → shard → load takes effect; per-rank peak | -| forward / backward | Full module from all-gather, activation, workspace | Measure logprob and actor update separately; frozen parameters may be gathered too | -| Export / merge / quantize | A single full tensor, fused groups, lazy generator references | The actual export set; release timing after yield/exception | -| compile / graph capture | autotune workspace and graph pool | Find the allocation in the log; eager only isolates, recheck the production graph separately | -| KV planning / wake | KV budget, trainer residency, allocator free cache | Resolved GMU/TP and allocator metrics, not the configured value | -| checkpoint / host | Full CPU copy, optimizer gather, shared mappings, Ray guard | Per-process and node usage; RSS does not sum to the physical total | - -A small model with the same special layers, attention period and scale layout verifies the mechanism; only the full model verifies capacity, such as the all-gather peak of a huge single parameter. - -### 3. Reserved is high, so why can wake still have no GPU memory - -After update's temporaries are freed, PyTorch keeps the blocks as reserved-but-unallocated. Another allocator such as CuMem cannot use them. - -`empty_cache()` releases free cache only, not live tensors, and not reserved memory held by active blocks or a graph's private pool. `gc.collect()` helps only when collecting a reference cycle drops a tensor's last owner; high reserved alone does not prove GC is needed. - -Compare allocated/reserved/device-free and the same allocation before and after cleanup. If allocated stays high, find the live owner; if reserved-unallocated is small, empty-cache cannot explain a workspace OOM of tens of GiB. Clean up where the resource is handed back (end of update or before wake), chosen by the real call path and a retest. - -### 4. How to tell a fix from a config downgrade - -- Lifecycle error: fix materialization/release timing, the completion contract, or the wrong gather set. -- A legitimate peak exceeds capacity: pick a topology that carries it or shrink the declared configuration, and record the support boundary. -- Configuration not propagated: confirm the resolved value inside the container; editing the shell variable again is not a new comparison. -- Killed by Ray/container policy: account for actual headroom before judging the limit; do not raise the guard by default. -- Long duration: time by phase and optimize only large shares; GPU utilization is not evidence of completion. - -Do not add a resident full-tensor dump to chase an OOM; use the phase record above to decide the next step. For the full-tensor boundary of export quantization see [precision and quantization](precision-quant.md#41-fsdp-path-verl). - -### Evidence boundary of historical experiments - -- A configured path is not an executed path: record a marker/count and treat zero hits as an invalid verification. -- Before comparing tensors, align request, token position, rank/owner and valid region; equal byte hashes alone do not prove the same objects. -- Verify a root cause with a single-variable intervention and a negative control that actually triggers the failure. -- Keep failed, invalid and superseded records with reasons. A small model locates the mechanism but does not extrapolate to full-model capacity or performance; the profiler explains structure, latency comes from independent timing. -- Bind every result to the actual patch stack and imported-file hashes; overlaying only new files keeps the old overlay, so restore and apply by manifest ([code and delivery](../practices/code-and-delivery.md)). - -## Pitfalls - -### 13. Stale quantization scale on the first step after restore - -- **Symptom:** After DS4 restores from a non-zero step, the first rollout mismatches and the next step matches again by itself. -- **Cause and Fix:** See [5](#5-derived-state-after-restore). -- **How to verify:** On the same checkpoint, compare the first export's scale fingerprint and the old-logprob before and after the fix. - -### 14. resume_mode=auto starts from step 0 - -- **Symptom:** A job with `resume_mode=auto` retrained from step 0; the log said `latest_checkpointed_iteration.txt` was not found. -- **Cause:** The checkpoint is on node-local NVMe and the newly allocated nodes have no tracker / shards. Same class: shards on the NVMe of the original eight nodes with the global `data.pt` on only one, so the launcher must check the exact node set. -- **Fix and verification:** Check before launch per item 1 of [6](#6-how-to-prove-resume-is-correct). diff --git a/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md b/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md deleted file mode 100644 index 45b1395a6e79..000000000000 --- a/.agents/skills/train-inference-consistency/guides/debugging-mismatch.md +++ /dev/null @@ -1,321 +0,0 @@ -# Fault localization and mismatch bisect - -Use this when training and rollout do not match, or a stage of a multi-node run died. It finds the first divergence with the fewest observations and produces a reviewable root cause, without dumping the whole model and without changing precision first. Performance attribution is in [Performance](performance.md); OOM and staging peaks are in [Checkpoint, resume and memory budget](checkpoint-resume.md). - -## Symptom → first step - -| Symptom | Do first | -| --- | --- | -| Not bitwise identical already at θ0 | [01 Isolate the side first](#isolate-the-side-first): four control groups; whichever two are equal rule out that side | -| Only some sequences are wrong | 01 same-prefix sibling: two samples of the same prompt must be bitwise identical over the shared prefix, otherwise it is the engine side | -| Wrong starting from some token position | [02 Bisect](#bisect-and-first-divergence): single-request prefill, bisect along one axis of layer/module/position | -| The difference grows slowly with training (θ0 is basically right) | [01 Isolate the side first](#isolate-the-side-first): save a small batch of prompt/response, score repeatedly at the same version to get each side's noise floor, then compare early and late adapters | -| Same operator input, still different | [Numerical contracts](../workflow/README.md), [BI kernel appendix table](bi-kernels.md) | -| Multi-node will not start / NCCL errors | [05 Communication fault ladder](#gb200-first-tell-apart-topology-implementation-mode-and-build) | -| Not sure what to save to be able to reproduce | [06 Save and compare protocol](#save-and-compare-protocol) | - -## Key rules - -1. Find the earliest divergence, not the largest layer; all later differences propagate from earlier ones. -2. Isolate the side before bisecting; classify before changing precision: input mismatch → weight mismatch → discrete path mismatch → kernel numerics → precision contract. Precision comes last. -3. Align probes by semantics (LSE in ln or log2, before or after norm, before or after the residual add). Misaligned probes produce false divergences. -4. Stop bisecting at a discrete decision (top-k, index); fix the routing, then continue. -5. A single round does not count; hidden coupling (one switch also changed scheduling) cannot be attributed to the switch; the target path must actually be exercised, otherwise the result is VACUOUS. -6. A two-node test must prove cross-node transfer happened; successful initialization does not mean bandwidth is reasonable; one `max_abs=0` does not prove consistency for arbitrary batches. -7. Looks like precision but is something else: determinism only takes effect on the main thread (switch scope). - -## Isolate the side first - -Applies when the logprob difference between FSDP/LoRA training and vLLM rollout is abnormal or grows with training. - -Let μ be the selected-token logprob saved by rollout and π the real trainer old-logprob. `μ != π` does not by itself show vLLM is wrong. The goal is not unconditional bit-exactness: rule out correctness errors, measure the noise floor, and explain growth beyond it. For quantization mechanisms see the [precision and quantization guide](precision-quant.md). - -### Per-sequence observation patterns - -Record valid/differing token counts, first divergence position, longest run of differences, prompt/response length, UID/group, and packing/rank mapping. Distinguish: a few sequences differ from the first response token; all drift slightly; differences only past a chunk/page; differences only after an update. The pattern chooses experiments; it does not prove root cause. - -Confirm token shift, mask and EOS. ±1/±2 offset comparisons are fine for diagnosis, but a fix cannot be declared through a misaligned comparison. Check that no score from another sample or version was referenced. - -### Existing rescoring tool - -A verified prefill rescoring tool takes the full sample tokens as input with zero new tokens and the current adapter, returns scores from `prompt_len-1`, checks the response token IDs, rejects missing scores, and clears the prefix/KV cache before scoring. - -### The cheapest control: do not dump all layers yet - -Save a small batch of real `prompt_ids + response_ids`, the prompt length and the valid response mask. From then on only score: do not regenerate, and do not rewrite the sequence through decode/re-tokenize. Record the base revision, adapter digest/version, config, code SHA, hardware and kernel backend. Compare training and serving under the same weights, with parameters unchanged during scoring. - -Score the same version repeatedly on each side for its noise floor, then compare early and late adapters. Rising ordinary mismatch (samples differ every round) is a signal, not proof that same-input error grows with training. - -| Path | Method | What it can answer | -| --- | --- | --- | -| Prefill | vLLM full input token IDs + `prompt_logprobs`, slice out the response | Difference between training teacher-forcing and the serving full-sequence forward | -| Decode trace | Normal prompt, force fixed response IDs, keep the real logits/logprob | Additional difference from KV/recurrent state, step-by-step generation, and decode kernels | - -`prompt_logprobs=0` means no extra top-k; target-token scores are still returned. If zero generation is not allowed, generate 1 token, discard it, and read only input scores. The first input token has no conditional score; the first response token is predicted by the logits at the last prompt position, so do not shift by one. Check trace switch, runner support and MTP limits against the actual vLLM version. Do not quietly turn off production MTP for the probe; an explicit no-MTP control is allowed but does not cover the MTP path. Prefill consistency does not imply decode or MTP consistency. - -### Metric measurement basis - -- Prefer raw logprob. When comparing the sampling distribution, unify temperature, penalty, top-k/top-p and similar processing. temperature=0 in the request does not prove which kind of logprob is returned; check the backend semantics. -- `logprob_mae` (token MAE): `sum(mask * abs(lp_train-lp_serve)) / sum(mask)` over valid tokens; also report p95/max and position. -- `probability_mae`: `mean(abs(exp(lp_train)-exp(lp_serve)))` over valid tokens. verl's `training/rollout_probs_diff_mean` is this, not logprob MAE. -- Per-token logprob difference in VERL: set `actor_rollout_ref.rollout.calculate_log_probs=true`; after the trainer computes `old_log_probs`, call `verl.utils.debug.metrics.calculate_log_prob_diff(old_log_probs, rollout_log_probs, response_mask.bool())`. It returns the 1-D absolute difference over valid tokens, not a mean; compute mean and quantiles yourself and keep mask indices for localization. The default debug metric applies `exp` first, so it yields a probability difference. Checked against VERL main at commit 8e4a572 (`verl/utils/debug/metrics.py`); on another version or fork, check return shape, mask and `exp` first, and record the VERL revision. -- Relative L2 and cosine are for tensors; do not mix them with logprob MAE/KL. Flag zero vectors separately. Report full-vocab log-softmax, probability total variation and argmax separately from the token metrics. An overall shift of raw logits is not a change in the probability distribution. There is no cross-model universal threshold. -- Schulman k3: with `d=lp_train-lp_rollout`, take statistics of `exp(d)-1-d`. It is a KL estimate only when samples come from the current rollout distribution; on a fixed old trajectory it is a diagnostic statistic, not the current policy KL. Check for overflow; do not silently clip and then claim the difference is small. - -### Build four control groups - -1. Inference alone: frozen tokens, repeated scoring, batch1/mixed batch, decode/full-prefill. -2. Training alone: A/A and A/C packing replay under the same weights. -3. Cross-stack: μ/π for tokens with the same semantics, plus full logits where needed. -4. Score layer: run the scorers separately on exactly the same logits; check dtype, temperature, logsumexp and the reduction over vocabulary shards. - -If training A/C differ and inference alone is stable, look at the training layout first, but keep the inference verification scope. If both sides are stable alone but differ across stacks, check weights, arithmetic semantics and the scoring path first. - -### Same-prefix sibling - -Same policy, same prompt, same response prefix, same score processing rules: conditional probabilities must be equal. Equal sibling μ with different π is evidence that the training side is sensitive to layout/state. For processed logprobs the penalty history and processing config must also match. - -Compare only the shared prefix, never past the generation fork. Use a clean sibling and repeated execution as negative controls. "One side looks more stable" is not proof of correctness over the full scope. - -### Localizing CUDA async errors - -An async illegal access may be reported after the faulty kernel. Save the first fault, every rank's logs, input/version and execution stage; localize with scoped synchronization, CUDA core dump/cuda-gdb, or a sanitizer. Once CUDA state is corrupted, end that experiment and collect no more numerical evidence from it. Disabling async globally or switching attention backend does not prove the stability problem solved. - -### Validity of controls - -Disabling the cache may execute a prefill that never ran before. Disabling Graph also changes speed and scheduling; to blame Graph itself, fix each step's composition offline and compare. A negative control must verify the target branch's execution count: zero difference on an unexercised path counts as not covered. For experiment design see [04 Ablation method](performance.md#ablation-method-make-ab-really-differ-in-only-one-variable). - -After isolating the side, bisect along one axis ([02](#bisect-and-first-divergence)); for classification and handling order see [03](#classification-and-handling-order). - -## Bisect and first divergence - -Separate input/weight errors, operator implementation differences and quantization error with the smallest experiment, without a model-specific debug framework. Metrics: [Metric measurement basis](#metric-measurement-basis); classification: [03](#classification-and-handling-order). - -### Choose the search axis first - -- Version regression: good/bad revision on the same reproducer, compatible environment, fixed input. -- Config regression: paired single-factor intervention; a 2×2 matrix when factors couple. -- Data trigger: request subset, length, neighbors, boundary remainder. -- Compute path: layer → attention/MoE/norm → internal operator → metadata/launch/instruction order. - -Do not change image, weights, kernel and parallel layout at once. For a dirty good/bad, save the full patch and run files; the same git SHA does not mean the same program. Build failure, dependency incompatibility and failure to reproduce are not a bad numerical revision; mark them skip/inconclusive. - -### Single-request prefill: narrow step by step, do not dump the whole model directly - -1. **Read existing training logs first.** Check the metric definition, policy version and sample source. Samples from different rounds only hint at an anomaly; they cannot prove same-input mismatch grows with training. -2. **Fixed-input scoring.** Save token IDs and position/mask from a single render. Both ends use same-source weights and the same-version adapter, and each scores repeatedly to measure its own noise. Rule out weight, input and score-semantics mismatches first ([01](#isolate-the-side-first)). -3. **Minimal forward.** Single request, one prefill, one output token. Keep the minimal prefix or layer set covering the divergent path and its upstream dependencies, with identical config on both ends and no training update. Confirm no cache hit skipped the prefill, and separate extra calls from warmup/draft/decode. An eager probe does not stand for graph/MTP acceptance. -4. **Find the earliest divergence.** Look first at the embedding, block outputs, and logits at the last valid prompt position. Align layout, valid tokens, shards, padding and deferred residuals by semantics, not by module names or tuple indices. Expand only the earliest block that exceeds its own noise. -5. **Same-input operator replay.** Give both implementations the same logical input and weights to separate this operator's difference from upstream propagation. If needed, replace only one operator's output and compare the final distribution; recurrent state must also satisfy the contract. -6. **Return to the original failing path.** After the fix, re-verify with the real adapter, input and run config. Passing with reduced layers/eager does not replace verifying the full model, decode, or CUDA Graph/MTP. - -### Per-module dump: find the earliest divergence, not the largest layer - -1. Build a semantic mapping table: layer, function, input/output role, module path on both sides, logical shape. FSDP wrapper names and fused vs separate Q/K/V need not match; combine before comparing. -2. Capture block boundaries and final logits first, on a few tokens/samples; refine the earliest significantly divergent block into norm, projection, router, expert, residual etc. -3. Align valid tokens, position, mask, TP/FSDP shards, expert ID, dimension order and padding. Never compare a physical shard with the full tensor. -4. A hook takes a snapshot, not an overwritable reference, and never changes the computation. For FSDP, sample at boundaries already materialized in this forward; do not all-gather the whole model or materialize a huge frozen PLE on GPU to observe. Export bounded data per layer/needed shard and release GPU staging promptly. -5. Record forward/recompute counts so backward recomputation does not overwrite the normal forward. Graph replay does not re-execute Python hooks, so an eager dump speaks only for eager. For graph-only problems use in-graph observation or snapshots of existing buffers; do not clear the graph or recapture to cover the problem. -6. If the operator input already differs, the operator is not yet proven buggy; give both ends the same logical input and check whether it amplifies the difference. When a fused kernel's inside is not visible, mark the observation boundary instead of guessing. - -Minimal artifacts: manifest (run conditions) + bounded tensor dump + comparison table with module/role, shape/dtype, fraction of finite values, max abs, relative L2, cosine, first anomalous position. Save a block's detailed intermediate tensors only after the summary has localized it; do not write all tokens of the whole model by default. Artifact format: [06](#save-and-compare-protocol). - -### Align probes by semantics - -Build a cross-stack boundary table: embedding, residual stream, norm input/output, Q/K/V, compressed rows/index, attention output, router logits/indices, expert input/output, combine, lm_head and scorer. Record tensor meaning, dtype, global token/shard coordinates and valid mask, not just module-name pairs. EP dispatch may reorder, duplicate or pad tokens; align expert tensors by semantic token/expert/shard key, not local offset. After expert GEMMs match, still compare separately the communication reduction, the shared + routed expert combination, and its rounding after compilation. - -If the saved rollout's first response token already differs, recompute it once from the original prompt alone. This rules out later-trajectory divergence as the only cause, but attributing it to prefill, attention or a kernel still needs controls on the same policy, position/mask/scorer and per-layer inputs. - -Compare original-dtype digests per layer or at a few checkpoints first; then dump only the neighboring boundaries and target samples. Keep one clean sibling. Snapshot async buffers before reuse without changing the production stream order in a way that hides the problem. - -If the difference does not move under any knob, suspect a too-coarse boundary first. In an external TPU MiMo case, attention core and o_proj were captured together and the difference stayed at `7.8125e-03` across CP 8/2, segment alignment 128/2048 and FP32 accumulation on/off; split apart, the core was 0 and the whole difference was in o_proj. Give adjacent or fused GEMMs separate boundaries by default. - -### Minimal but credible observations - -- Copy tensors before buffer reuse; save only necessary boundaries and tokens, with a byte cap. Model-specific mappings and artifacts stay in the model working directory. -- The manifest records code/image version, checkpoint and adapter identity, input digest, position/mask, dtype/quantization, score semantics, graph/MTP, rank/layout and intervention options. Same config does not mean same weights. -- One directory per call; write the manifest after completion. Refuse comparison on missing boundaries, repeated calls, or incompatible version/shape; do not silently take the intersection or truncate. Non-finite values fail explicitly; metrics on zero norm are marked undefined. -- Compare max abs, MAE, relative L2, cosine; for final full-vocab logits also log-softmax, probability total variation and argmax. - -### Recommended logs: identity + metric measurement basis + evidence boundary - -Reuse existing logs and JSON artifacts; no extra logging framework. One manifest per run, one detailed report per comparison. The terminal prints summaries, not full tensors. Mark unknown fields unknown; do not backfill them as verified. - -```text -[mismatch] run= step= policy= input= phase= rank= repeat= - scope= intervention= - metric= value= reference= - self_repeat= report= classification= -``` - -`metric` uses the names in [Metric measurement basis](#metric-measurement-basis) (`logprob_mae`, `probability_mae`, KL diagnostic statistic, tensor relative L2, vocab TV), never interchangeably. - -Time load, rollout, score and update separately so cold start is not counted as training. A distributed job's outer console log is only a forwarded copy: when updates stop, check TaskRunner/worker raw logs, checkpoint times and completed steps first. One forward/backward stack does not prove a hang; do not restart a healthy job because log forwarding stopped. - -### Limits of bisect - -A floating-point difference may vanish in later rounding, masking or routing, so "this layer differs" is not monotonic in depth and O(log L) does not always find the first difference. Coarse-scan for a suspect interval, then bisect/fine-scan, verify the boundaries before it, and confirm responsibility with the same operator input. - -A cut point must carry the complete state: residual, MHC mix, SSM/conv state, KV, positions, routing and necessary metadata. A fused upper-layer post and this layer's pre have no natural independent cut point. Prefix replacement that drops hidden state creates new differences. - -Only a same-input operator A/B separates upstream propagation from this operator's error. Injecting reference activations, fixing the routing or replacing the kernel can confirm causality, but they are diagnostic interventions; accept the production fix only after restoring the original path. - -### Form a reviewable root cause - -The result must contain: the minimal input, the last identical/earliest different boundary, the same-input control, the source mechanism, the trigger condition, and the minimal change that makes the failure disappear. A config difference is only a lead: in Nemotron, running both reduction block sizes with the same kernel gave consistent results for both, which overturned the guess that the block size difference was the cause. - -If after three rounds only the final score is known to differ, add observation points or re-check the sample mapping instead of rerunning in full. Keep the failing samples and the ruled-out hypotheses so the next round does not repeat them. - -### Conclusion and stop conditions - -The report answers only: what was observed, what the controls proved, what is not yet proven, and one minimal next step. Correctness errors must be fixed. Normal numerical differences are judged by noise and task impact. Quantization error needs same-input evidence; low-precision inference alone is not grounds to blame it. When improvement is needed, evaluate in order: supported high-precision rollout, QAT matching serving, real low-precision training. Do not cover wrong weights or adapters by changing precision. - -**Replacing with the reference output is diagnosis, not a fix. Lower local error does not mean lower training mismatch; per-operator ablation gains do not add up.** Stop once this round's evidence goal is reached; do not change the training config or start new experiments along the way. - -## Classification and handling order - -After isolating the side ([01](#isolate-the-side-first)) and bisecting ([02](#bisect-and-first-divergence)), classify the difference, then decide what to touch. Precision comes last. Model-specific probes, runtime patches and measurements stay in the model working directory. - -### Classification table - -| Class | Evidence | Handling | -| --- | --- | --- | -| Correctness error | Missing weights, wrong mask/position/layout, wrong scale/block, wrong quantization range, operator semantics violate the contract | Must be fixed and the original failing path re-verified; for a diagnostic task, deliver the reproducer and evidence first, without widening modify permission | -| Normal floating-point difference | Explained by same-version repeatability/reduction/kernel noise, no contract error found | Report the noise floor and task impact; bit-exact is not the only pass condition | -| Error from correct quantization | Mapping and quantization rules are correct, and a same-input control localizes the divergence to the quantization path | Precision ablation in the order below | -| Unlocalized | Insufficient evidence or confounded variables | Keep unknown, continue with minimal probes, do not declare "it is just quantization" | - -Precision ablation order: supported and affordable high-precision rollout → QAT matching the serving rules → real low-precision training (most backend requirements, last). A natively quantized model does not support BF16 just by changing the dtype. Fake quant must match the serving weight/activation rules but does not reproduce its GEMM. QAT is not a workaround for a wrong scale. QAT contract and route selection: [precision guide](precision-quant.md#1-choose-the-precision-route-first). - -Keep a single checkpoint source; open a separate high-precision adapter branch when needed. - -### Handling order - -Input mismatch → weight mismatch → discrete path mismatch (top-k, index) → kernel numerics → precision contract. Rule out each item before moving to the next. - -Acceptance: after the fix, the same failing probe passes, the original path is not broken, and error growth on fixed samples is explained/controlled; then look at training process metrics and the fixed task evaluation. Set tolerances from the model and the baseline; do not copy another model's 0.002. Statistics on a fixed old trajectory are diagnostics, not a KL estimate of the current policy. - -This procedure does not authorize stopping jobs, switching dtype, adding GPUs or changing the production recipe. When the next diagnostic gate needs new permissions/resources, report the smallest next probe and its cost. - -### Delivery - -Follow [06 Save and compare protocol](#save-and-compare-protocol): manifest, bounded dump, comparison table, semantic boundary of the earliest divergence, classification and confidence, one next step; state which paths were actually tested. No separate framework, no cross-model universal threshold. - -A case that looks like precision but is a switch-scope correctness error: [determinism takes effect only on the main thread](bi-kernels.md#determinism-takes-effect-only-on-the-main-thread). - -## GB200: first tell apart topology, implementation, mode and build - -Based on experiments on GB200 four-GPU nodes; future DeepEP/NVSHMEM releases may change the support matrix. Check the source of the target fork/commit and the installed artifacts before relying on it. - -### Do not use one name for four things - -| What to ask | Example record | -| --- | --- | -| Which implementation | AllToAll, some version of DeepEP, HybridEP; note the fork/revision | -| Which path of that implementation | The actual class, API and kernel behind high-throughput / low-latency / v2 | -| How it communicates | NVLink within a node; RDMA across nodes; or the MNNVL path actually enabled | -| What it was compiled into | 4-peer/8-peer macro, target GPU, NVSHMEM dynamic dependency | - -`NVL4 + RDMA fallback` describes the communication topology. **It does not state the DeepEP mode**: do not answer low-latency or high-throughput unless the call site proves it. How `DeepEP v2`, `HybridEP` and HT/LL relate depends on the version; do not infer containment from the names. - -MNNVL being off does not rule out cross-node EP: the verified path may be NVLink within a node and RDMA NICs across nodes. Nor does that path work on every GB200; devices, network, peer build, underlying libraries and the called API all need verification. - -### Get the world size right first - -Write the trainer and rollout process groups separately. Formulas for dense world and expert world, and grouping pitfalls such as "two DP4/EP4 are not one EP8 group", are in [training-side layout and boundaries](training-forward.md). The runtime decides the actual grouping; do not change the GPU count to fit a config string. - -### Same-fabric fairness - -Before any communication-bound A/B, prove both arms use the same transport. In one GB200 NVL case the BI=0 arm's NCCL AllGather/ReduceScatter went `via P2P/MNNVL` (same clique, cliqueSize 8) while the BI=1 arm's DeepEP low-latency a2a was pinned to RoCE (`VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=0`, `NVSHMEM_HCA_PREFIX=rocep`); a gap blamed on non-overlappable communication latency was mostly the wrong NIC, fixed by two environment variables. - -Check steps: - -1. `nvidia-smi -q | grep -A4 Fabric`: the ClusterUUID/CliqueId of the two nodes must match. -2. Run `NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,NET` once per arm and check whether the transport is P2P/MNNVL or NET/IB. -3. grep the arm environment files for DeepEP/NVSHMEM pins. A hard `export` overrides the value passed by the container `-e`; write it as `${VAR:-default}`. - -Mechanism: `VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=1` → `deep_ep.Buffer(allow_mnnvl=True)`, DeepEP no longer exports `NVSHMEM_DISABLE_MNNVL=1`, and NVSHMEM 3.4.5 uses MNNVL for peers in the same clique automatically. `DEEPEP_MAX_NVL_PEERS` is not a runtime switch (it is DeepEP's compile-time `NUM_MAX_NVL_PEERS`); the peer-4/peer-8 build is selected by switching the overlay. MNNVL is not hybrid-only: the DeepEP v1 LL kernel branches per destination on `nvshmemi_get_p2p_ptr` (direct NVLink write vs IBGDA), and MNNVL only makes cross-tray peers P2P-mappable; "hybrid mode" is a v2/elastic concept. MNNVL was on by default before DeepEP #190 (2025-06-06) and turned off afterwards to save NVSHMEM memory. - -## Save and compare protocol - -Large tensors and checkpoints stay in `results//`; source and small summaries go into version control. Each framework dumps differently, so write these per project following the key points below: logprob comparison, per-layer trace, replay controls, a fixed-token oracle (logprob/argmax of a fixed token set before and after a weight update), GSM8K record comparison before and after an update, a `.meta.json` provenance sidecar for the results directory, and a pre-experiment snapshot commit + clean worktree of uncommitted changes with a check of the container's actual import path and file hashes. - -The manifest records artifact path, size, hash, generating command, code/image/policy version, number of valid samples and hook execution count. Confirm the write succeeded before cleaning up compute processes. Do not auto-delete existing artifacts. - -### Archive experiment material by step - -Build the step index before the experiment and fill in status and evidence after each run, not at the end. The layout below is logical; link existing material by explicit path instead of moving or copying it: - -```text -/ - steps/ - INDEX.md # current step, evidence entry of each step, gates not passed - 01-baseline/ - INDEX.md # tuning comparison table, candidates and review records - / - config/ # resolved config, environment variables, version list - scripts/ # actual server/bench commands or pinned-version entry points - results/ # raw results/logs, or the list of persisted artifacts - summary.md # conclusion, limits, failure cause and next experiment - 05-trainer-forward/ # classified by C0-C10 gate number, created as needed - 10-final-acceptance/ # references the baseline ID and this candidate version - results// # large tensors/checkpoints, not in Git -``` - -- **Every run is traceable.** Save the actual launch command, full effective config, code/image/checkpoint identity, nodes and run time, evidence of the backend that actually ran (a CLI flag is not evidence), exit status and raw results path. A branch name is not provenance: record branch, snapshot commit, base SHA, patch, tree hash and image digest in a `.meta.json` sidecar. Replace secrets with references. Large files go to persistent storage with size and hash, never only a node-local path. -- **Baseline tuning: one row per run**: run ID, candidate config (MoE/communication), actual backend, load and concurrency, repeat rounds, valid input/output tokens, cache hits, throughput/latency, exit status, evidence link; N/A where inapplicable. Keep failed, unfinished, diagnostic and excluded runs with the reason. The summary lists the included run IDs; do not pick the best round or keep only the final mean. -- **A candidate is not a reviewed baseline.** The index marks the recommended candidate and the human-confirmed baseline ID separately; the baseline links to exact config/image/code, result set and review record, or reads "pending review". Later comparisons reference that ID; do not swap the denominator for a same-named script's new content or another backend's old result. -- **Reruns keep history.** A new run gets a new ID and overwrites nothing. Mark an invalidated conclusion with the reason and the replacing run. Continuing an experiment keeps separate attempt/segment and exit records. Reuse evidence across steps by reference, not by copy. -- **On handoff, read the index first.** The step owner updates its record. Reporting, recovery and acceptance check raw evidence through the index; a missing file, mismatched identity or chat-only conclusion is an evidence gap, not a pass. Index historical material; do not delete or bulk-move it. - -### Replayable raw data - -Keep the real framework batch: prompt/response IDs, attention/response masks, positions, lengths, UID/group/sibling, rollout logprobs, necessary reward and routing records, plus the packing mapping, parallel config and how to restore that step's rollout weights. The raw batch is the replay input; exported per-token logprobs are for comparison only. - -Step N's rollout usually comes from the policy before update N, so the post-update checkpoint may not reproduce it. Under partial rollout or async staleness, segments of one sequence may come from different policies (re-prefill after interruption continues under the new policy); record the policy version per segment, or a single-checkpoint replay turns expected policy difference into false inconsistency. Freeze weights, mutable buffers and RNG/dropout state, and state whether the scorer uses eval or training mode. - -### Compare per-token logprobs - -Each side (rollout and training old-logprob, or two replays) exports `{request_id: (logprobs, response_mask)}` as original-dtype tensors with `torch.save`. Do not convert to JSON or rebuild from printed logs. Key points: - -- **Align by logical request, not physical row.** Remove padding and restore the pre-packing UID order. Both request sets must match; a missing request raises an error rather than comparing the intersection. -- **Compare only valid response tokens.** Positions outside the mask are excluded. Every sequence needs at least one valid token; an empty batch is an invalid experiment. -- **Bitwise comparison looks at bytes, not values.** Value comparison treats `+0.0` and `-0.0` as equal and two NaNs as unequal. Compare integer views: `a.view(torch.int32)` (float32) or `a.view(torch.int16)` (bf16). Do not cast first and then claim bitwise identity in the original dtype. -- **Same contract on both sides.** Same policy version, same raw/processed basis (for processed, record temperature, truncation, penalty), same input. Compare input token hashes before scores; differences from different inputs are not numerical differences. -- **Report** valid token count, byte mismatch count, first mismatching (request, position), and max/mean absolute error. Matching sampled-token logprobs do not mean matching full logits. - -```python -def compare_logprobs(a, b): - assert a.keys() == b.keys(), "request sets differ" - n = bad = 0; first = None; max_err = 0.0 - for uid in sorted(a): - (x, mx), (y, my) = a[uid], b[uid] - assert torch.equal(mx, my), f"{uid}: masks differ" - x, y = x[mx], y[my] - ix = x.view(torch.int16 if x.dtype == torch.bfloat16 else torch.int32) - iy = y.view(torch.int16 if y.dtype == torch.bfloat16 else torch.int32) - diff = (ix != iy).nonzero() - if len(diff) and first is None: - first = (uid, int(diff[0])) - n += len(x); bad += len(diff) - max_err = max(max_err, (x.double() - y.double()).abs().max().item()) - return {"tokens": n, "mismatched": bad, "first": first, "max_abs": max_err} -``` - -### Per-layer trace: find the boundary of the first divergence - -When logprobs mismatch, record tensors with forward hooks at the same semantic boundaries on both sides (for example each layer's attention output, MoE routing, residual), find the first differing boundary in execution order, then do a bounded dump there. Key points: - -- **Key by semantics, not rank.** Each record's key is `(request_id, boundary name, logical slice)`, for example `("step7-p2-s3", "layer2.attn.out", "tokens=0:128")`. Stitch TP/CP/EP shards back by global coordinates and remove padding before recording. Naming and order must match on both sides. A missing probe invalidates the whole experiment; it cannot be skipped. -- **Hashes first, tensors on demand.** Take SHA256 of the contiguous tensor's original-dtype bytes (`hashlib.sha256(t.contiguous().view(torch.uint8).cpu().numpy()).hexdigest()`), with shape, dtype and semantic execution sequence number. A hash shows where it differs, not by how much; save raw tensors only at the first differing boundary. -- **Clone before the buffer is reused.** Inference-side workspace and CUDA Graph output buffers are overwritten by the next call; `clone()` in the hook before hashing or saving. -- **The first difference is not the root cause.** To blame an operator, also prove its input is identical on both sides (record its input hash too). - -### Replay controls - -With fixed rollout and fixed weights, change one variable (packing order, batch composition, parallel layout, kernel switch), replay old-logprob, and use the comparison above to judge whether that variable affects the result. Key points: - -- **One process per variant**, with a synchronous entry point that blocks until the replay finishes. Do not use background submission such as `sbatch`: on timeout or failure the job cannot be reclaimed. Variants that change compilation, libraries or initialization environment need a new process. -- **Each variant writes to a new directory** without overwriting old results; record the actual command, environment, exit code and elapsed time. Preview the command and budget before executing; on failure, stop at that variant. -- **Prove the replay is frozen**: record the weight fingerprint, confirm no generation and no optimizer update, and confirm the packing actually in effect, not just the command-line arguments. -- For the experiment record use [the report template](../workflow/report-template.md). diff --git a/.agents/skills/train-inference-consistency/guides/environment-provenance.md b/.agents/skills/train-inference-consistency/guides/environment-provenance.md deleted file mode 100644 index c418f16d21a7..000000000000 --- a/.agents/skills/train-inference-consistency/guides/environment-provenance.md +++ /dev/null @@ -1,106 +0,0 @@ -# Environment, image and runtime version provenance - -Once the environment is set up (the image pins the native runtime, uv manages training dependencies, source arrives as a traceable snapshot or a minimal overlay), prove that it is what is actually running. Code existing does not mean it is executing. - -## Find by symptom - -| Symptom / question | Where to go | -| --- | --- | -| Not sure which version of source / wheel / .so / image is actually running | [Run manifest](#use-one-run-manifest-to-answer-what-exactly-is-it-right-now) | -| How many commits "latest" actually differs from the target branch | ["Latest" must come with a target branch](#latest-must-come-with-a-target-branch) | -| Overlay, or change the source and rebuild | [Boundary between overlay and native wheel](#boundary-between-overlay-and-native-wheel) | -| Recompiles on every start / not sure where the cache should go | [Compile cache](#compile-cache-where-it-lives-decides-whether-it-can-load) | -| Not sure which path DeepEP actually runs / how it was built | [Four things about the communication implementation](#communication-implementation-do-not-use-one-name-for-four-things) | - -## Key rules - -1. Pin the image digest, CPU architecture and Python version first, then resolve training dependencies. Image tag, PyPI version and upstream PR merge status do not replace one minimal inference run in the image. -2. `--system-site-packages` does not stop uv from installing a same-named package that then imports first; check the install plan, the actual import path and the GPU operators together. -3. Each of the six manifest layers (source, installed packages, native libraries, JIT kernels, container, run) has evidence nothing else replaces: `latest` is not the digest, `pip show` in another venv is not the actual import path, the Python source HEAD is not the loaded .so. -4. Audit inside the worker, interpreter and mount namespace that actually run; the login container shell and the Ray worker may differ. -5. One image tag may point to different digests on different dates/nodes; ARM64 and AMD64 tags may be independent builds. -6. Modify in the host worktree, execute in the container: same-commit worktree plus per-file hash check, never an overlay stack detached from the formal source. -7. Write version pins such as "the Megatron fork is only compatible with Megatron-Bridge v0.6.2" into the run manifest. -8. How the image was produced (container commit exported to sqsh, or Dockerfile rebuild), where it was pushed, and the three-repo SHAs in the tag are part of the run identity. - -## Environment identity: code existing does not mean it is executing - -### Use one run manifest to answer "what exactly is it right now" - -Record recipe, trainer, Megatron/MCore, vLLM, DeepEP, DeepGEMM, FlashAttention/FlashInfer and NVSHMEM separately, each with at least: - -| Layer | Evidence to save | What cannot substitute for it | -| --- | --- | --- | -| Source | remote, target branch, HEAD, dirty patch hash, submodule gitlink/actual HEAD | `latest`, the repo directory name | -| Installed package | Actual interpreter, import path, package version, wheel/build revision | `pip show` of another venv | -| Native library | Path/hash of the `.so` actually loaded, build inputs, target architecture, dynamic dependencies | Python source HEAD, extension file name | -| JIT kernel | DSL/source hash, compiler version, specialization/config, cache origin | `.so` unchanged | -| Container | OCI manifest/platform digest or `.sqsh` SHA256, parent image, build record | A mutable Docker tag, image creation time | -| Run | Resolved config, worker environment, model/weight digest, the real kernel | An export in the head node shell | - -Write `unknown` for unknown items; do not infer the whole source tree from a wheel version string. Audit inside **the worker, interpreter and mount namespace that actually run**. If dynamic linking is suspect, check the process's actual mappings/loader result, not the static prediction of `ldd`. - -ARM64 and AMD64 tags may be independent builds or members of one multi-arch manifest. Compare the target platform's manifest/config/layers; a similar tag or update time does not show it is the same image. Record generic scheduling switches such as `CUDA_DEVICE_MAX_CONNECTIONS=1` per the task convention; do not add them quietly and then credit BI. - -A cross-repo compatibility pin is a manifest row too: state the pin (the Megatron fork works only with Megatron-Bridge v0.6.2) and its basis, and re-verify when either side is upgraded. Do not assemble incompatible releases just so every item reads "latest". - -#### Communication implementation: do not use one name for four things - -See [debugging mismatch](debugging-mismatch.md#do-not-use-one-name-for-four-things). - -### "Latest" must come with a target branch - -Pin each repo's target branch per the [delivery target](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). The reference implementation's PR is a design reference, not automatically the PR base. If the user names another fork/branch, record the override. fetch is not checkout; submodule update follows the parent repo's gitlink, not each submodule's own remote. - -```bash -git -C /explicit/repo status --short --branch -git -C /explicit/repo remote -v -git -C /explicit/repo fetch origin -git -C /explicit/repo rev-list --left-right --count A...B -git -C /explicit/repo log --left-right --oneline A...B -git -C /explicit/repo merge-base A B -git -C /explicit/repo diff --stat A B -git -C /explicit/repo submodule status --recursive -``` - -The left/right counts of `A...B` say how many commits each side has; the merge base and diff say what differs. Diverged branches have no single before/after relation. Commit time and short-hash magnitude do not decide ancestry. If a shallow clone lacks the ancestor, fetch history before claiming the lines are unrelated. - -Keep the recipe's declared compatible combination by default. If the user explicitly asks to update a development branch, record the deviation from the pin and check the actual changes and the ABI. - -### Boundary between overlay and native wheel - -During development, fix the container code-sync and build/test method first and record the tool versions in the run manifest. Modify in the host worktree, execute in the container; do not maintain an overlay stack detached from the formal source: - -1. Confirm the config, model registration and import identity of the native image/checkpoint; pin the image digest, the full source commit and the real install directory. Do not guess a missing implementation from the model name or tag. -2. The sync worktree's `HEAD` equals the image's full commit, and the changes under test are a diff on that baseline. They may be ported from the development branch, but a different HEAD cannot bypass the check. -3. In a container not yet running the workload, sync each changed in-package Python file individually: before copying, check the original file's hash against `HEAD`; after copying, check the hash of the file under test; save the sync manifest. For new files, check no original exists on either side. Never overwrite the whole package or modify a running container. -4. Keep the image's version, generated files and unchanged extensions. Recompile native/binding changes and verify the artifacts; for JIT changes check compile inputs and cache; if the ABI/toolchain changes, rebuild the image. -5. After startup, re-check import/library paths in the actual worker and test. Finally build the unified image from the formal PR's clean commit and rerun acceptance; the user does not redo the development sync. Details: [Implementation and delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). - -Overlays are for development/diagnosis only, not the delivery format. - -- An overlay covers specific files in the container with external copies; a wheel is an installed package that may carry compiled machine code. -- Cover only the necessary files at their exact target paths; do not mount the whole repo over the installed package and shadow its extensions. Verify the target's original hash first, apply once at startup, stop on failure, and save the overlay manifest, final hashes and import paths. -- A Python-only change can still change the binding/JIT/shape contract; it is not automatically exempt from recompiling or testing. -- Changes to CUDA/C++, bindings, compile macros, submodule kernels or the CUDA/Torch ABI need checking and recompiling. Before acting on a "missing extension", rule out the wrong venv, shadowing by a whole-repo mount, search paths and dynamic dependencies. -- If the image already has the correct artifact, verify it and use it. Do not recompile a DeepEP already proven to work, and do not change only the source pin in the recipe expecting the old `.so` to change. -- Do not fix a FlashAttention extension load failure by `unset`ting a variable. Accept a config fix only once it is shown the variable selected an unneeded/incompatible path and the run passes with the user's target path restored. -- Keep observability overlays separate from numerical/performance changes and off by default; when cleaning up the main PR, back up still-useful diagnostic code before deleting it. - -### Delivery gate for immutable images - -1. Commit production changes to their PRs; freeze the clean source commit and all native/JIT build inputs and dependencies; record the base image digest, not just the tag. A dirty patch may identify a diagnostic candidate but cannot be an uncommitted dependency of the final run. -2. Build the final image and check import/dependencies/hashes in a fresh container **with no development source mounts and no inherited JIT cache**, keeping the device, model and data mounts production needs. -3. Run the agreed correctness, Graph/cache and performance gates in that image; if observability was on, accept again with the formal config and no probes. -4. Publish exactly the accepted artifact and verify the registry digest/architecture; if conversion or repackaging changes the content, verify the equivalent closure and run a final smoke. -5. Sync recipe, PR description and image receipt; "image push succeeded" does not cover items that failed. Follow this order when the user asks to push after testing. - -The receipt states how the image was produced: Dockerfile rebuild from a clean commit, or changes installed in a container then committed and exported as `.sqsh` (the only option when the coordinator has no docker permission; fine for development, but formal delivery returns to the Dockerfile); which registry it went to (Nemotron: public Docker Hub); a tag carrying the vLLM/Megatron/verl SHAs, not just the date. Record both the `.sqsh` SHA256 and the registry digest under the same tag. - -Adding a new project source layer to an old OCI base does not bring in another experiment image's vLLM/FA4 changes. A deleted file missing from the top-level tar may still exist in lower layers; deleting a base file needs a correct whiteout or a check in the final rootfs. - -Upgrading FlashInfer/NVSHMEM must not silently upgrade Torch/CUDA. Pin the target interpreter and constraints, compare packages/dynamic dependencies before and after, and confirm the APIs used and build compatibility. A version number is not evidence of an HCA fix or of model correctness. - -## Compile cache: where it lives decides whether it can load - -Put the Triton/Inductor/vLLM compile caches on node-local disk mounted explicitly into the container; a `noexec` mount (such as `/dev/shm` in some containers) cannot load the compiled `.so`, and a per-container tmpfs recompiles on every start. Isolate caches by image, source commit and GPU architecture, and give each rank on a node its own copy, since an existing cache directory does not mean a cache hit. diff --git a/.agents/skills/train-inference-consistency/guides/performance.md b/.agents/skills/train-inference-consistency/guides/performance.md deleted file mode 100644 index 97419e524a8a..000000000000 --- a/.agents/skills/train-inference-consistency/guides/performance.md +++ /dev/null @@ -1,160 +0,0 @@ -# Performance and profiling - -Record performance gates and correctness gates independently: low throughput does not mean the numerics are wrong, and correct numerics do not mean performance meets the target. BI kernel choices and costs are in [BI kernels](bi-kernels.md); the statistical protocol for paired acceptance is in [Pairing and computation](recipe-and-acceptance.md#pairing-and-computation). - -## 1. Workload and measurement basis - -### Three throughput measurement bases - -- **Per-request tok/s**: 1/TPOT; decode latency. -- **Engine total tok/s**: completed tokens of all concurrent requests / common timing window. With multiple engines use the full common window and total completed tokens, not a sum of per-engine throughput; with unsynchronized node clocks, do not compute the window from timestamps. -- **RL one-step gen time (makespan)**: one batch of rollouts from the first request in to the last request out. This is what RL cares about; long tails and scheduling can move it opposite to the other two. - -State which basis you report; compare ratios only on the same basis. - -### Translate the rollout target into actual kernel workload - -"batch=128, n=8, input 2K, output 8K" must at least specify: - -- Is 128 prompts, expanded requests, or per-engine concurrency? 128 prompts with 8 siblings each is 1024 requests. -- Is 2K/8K the maximum or a fixed length? Record input token mean/max/percentiles, total output, and completed requests. -- `max_num_seqs` limits simultaneously active sequences; `max_num_batched_tokens` is the token budget per scheduling pass; neither equals the total request count. -- Requests per DP engine, prefill chunk size, and the active decode batch over time; a global 128 is not each rank's kernel M. -- Prefix cache cold, warm, or naturally hit; the n siblings of a prompt affect hits and prefill. -- Natural EOS vs ignore EOS / fixed length, sampling parameters, seed, request order, maximum model length, scheduling policy. - -### replay bench: equal-workload replay of one RL step - -Use the same request count and prompt×n as one RL step, take each request's output length from one real rollout, replay with ignore EOS at fixed lengths, and measure makespan. Measure the repeat noise of one config first, then compare arms; draw no conclusion from a difference smaller than the noise. With the same length table, token count, request count and long tail are identical, so the difference belongs to the execution path. Fixed output suits equal-workload comparison, not quality. - -In natural generation, different numerical paths or scheduling can change token counts even with the same weights and seed; record output token receipts and add a frozen-token replay if needed. A short 2K+64 profile can locate startup/local hotspots but **cannot** represent the long-term decode, cache/page pressure, scheduling tail or full rollout speed of 2K+8K. Use different experiment IDs for the short profile, the full stress test and the training E2E. - -### Metrics are cumulative values - -When one engine runs multiple stages (concurrency or parameter sweep), `vllm:iteration_tokens_total_*`, `vllm:spec_decode_num_*` and `vllm:num_preemptions_total` in `/metrics` are cumulative since process start. Always compute per-stage cost as `after - before`, and note in the script that the first stage is right only because before is all zero, which makes the method look correct. Cross-check the prefill amount with an independent client-side basis (prompt_tokens/cached_tokens in the request records); if the two disagree, the counter basis is wrong. - -When the verification method itself fails (a regex that never matches, a hardcoded version number, an unsubtracted cumulative value), the numbers look plausible but measure something else. Reconcile every new script's first numbers against an independent basis. - -## 2. Fair A/B - -### Do not mix the two kinds of A/B - -- **Causal comparison of the BI switch**: same software/model/scheduling/cache, only BI changes. If the switch selects a different kernel automatically, record the actual branch and add no other optimizations. -- **Each side's best available config**: BI=0 uses the strong default / verified best path for that platform and model; BI=1 uses the fastest path that meets the agreed consistency. Backends may differ; disclose every difference. - -Do not force BI=0 onto FA2, a specific MoE, non-default attention, or Graph off to shrink the gap. Do not call an experiment that locally disables switches inside BI=1 "native BI=1". A replaced path is a separately named candidate that passes the full contract gate. Give common optimizations (cutting shared launches, enabling a supported network path) to both sides and report the common gain and the reduction in BI overhead separately. BI=1 faster than BI=0 is not automatically fake, but it triggers an audit of config, workload, cache and timing. - -### What to fix - -Weights, input, generated tokens/lengths, output cap, concurrency, nodes, software, warmup and timing boundaries. If output receipts and workload differ, report only each side's own throughput. Separate the fresh-process first rollout from warm steady state. - -**Same fabric**: before comparing communication bottlenecks, confirm both arms use the same transport; see [Same-fabric fairness](debugging-mismatch.md#same-fabric-fairness). **Same node**: small gains need same-node pairing, crossed order and multiple rounds; counterbalance across node pools and do not pick the fastest cross-pool combination. **JIT pollution**: a JIT kernel specialized by token count recompiles for every new prefill size under BI=0; verify zero compile lines during the bench. **BI=0 silently running the BI kernel**: a fork whose gate only checks whether the `.so` exists makes BI=0 run the BI kernel; grep `vllm::batch_invariant::` in nsys. - -### Report speed switches and numerical switches separately - -One switch (PC, cudagraph, BI) changes both the numerical path and throughput. Before reporting "probs_diff with BI on/off" or "throughput with BI on/off", list the confounders that change with it, rule out or fix each, then compare; see [Paired acceptance](recipe-and-acceptance.md#paired-acceptance). - -### Computation and reporting - -```text -output tokens/ms = completed output tokens / (wall_seconds * 1000) -BI throughput deficit = 1 - throughput_BI1 / throughput_BI0 -BI time overhead = time_BI1 / time_BI0 - 1 # matched work only -``` - -Latency reduction `1-new/old` and throughput gain `old/new-1` differ; 16.35% lower throughput is not 16.35% more time. Report numerator, denominator, request / completed count, raw receipts, config differences, sample count and variation; take the median over pairs and report single pairs separately. Gains of A and B do not add; measure the combination, and do not claim it beats the best single item from a few observations. List skipping the old-logprob forward separately, only when the RL algorithm allows it; it is not a gain of the BI kernel. Every ratio must be recomputable from explicit artifacts. - -## 3. Breakdown and ablation - -### Two stages: remove structural defects first, then fine-tune - -Once proxy inference consistency passes, check feasibility: sweep token row counts on real operator shapes, count launches and host gaps, look for per-token launches, synchronous readbacks, repeated copies and abnormal workspace, and confirm the target Graph captures/replays. After backward is wired in, measure training forward and forward+backward (by token count) on the main topologies fixed at C1 against the native training path on the same build, and check them against the C1 budget, so a scheme fast at decode but slow in training/prefill is not chosen. After full consistency, fine-tune by E2E profile -> micro benchmark (screen high-share shapes, then sweep all target shapes and boundaries; all must be correct; weight performance by real call frequency) -> multi-round paired E2E on the same workload. New or replacement kernels do not use Triton; use CuTe DSL, CUDA/CUTLASS or an existing implementation; historical Triton data is only a comparison. Final delivery still needs full acceptance, see [validation](../workflow/validation.md). Fix a structural problem first and re-check the numerics; do not turn this stage into endless tuning. - -### Break down the overhead first - -Time model load/JIT, rollout prefill/decode, old-logprob, backward/update and offload/wake separately. Kernel times may overlap and do not add up to wall time. Keep the nsys window, rank/GPU normalization, Graph node tracing and sampling interval consistent. - -Profile the effective config and actual kernels, not just the BI switch. In DS4 the same quantization `.so` loads under both BI0 and BI1, but different MoE layouts select different schedulers. For an NVFP4 MoE model, account for dense linear, expert GEMM, expert routing/indexing, all-to-all and attention separately; one vendor name does not describe the path (one Nemotron run mixed FlashInfer FP8 linear, Humming NVFP4 dense GEMM, Humming indexed experts and FlashInfer NVLink one-sided all-to-all). A run that fails on a backend guard (a BI=0 trial expected Humming but got FlashInfer CuTe DSL) calls for re-auditing the backend and build path and does not enter the performance denominator. - -### Collect enough attribution information in one profile - -1. Verify the actual serving config and input/output statistics. For inference-only use vLLM serve or an existing client, not training; use the mp/Ray executor actually needed. -2. Add low-overhead, off-by-default NVTX at semantic module boundaries: attention, indexer/Top-K, index preparation, MLA, MHC, shared/routed MoE, norm, linear, sampling, communication. Do not apply the DS4 module list to dense models. -3. Use nsys first for the timeline and hotspots, with CUDA Graph node tracing on; graph launches without node details do not support per-kernel comparison. -4. Count calls, GPU time, CPU launch/sync, allocation and shape per path; aggregate by rank, phase and module. -5. Pick representative prefill, steady decode, mixed and tail windows, then ncu a few kernels/shapes for bandwidth, occupancy, memory access and compute; ncu-perturbed wall time is not throughput. -6. Measure throughput on the same full workload without profiler or heavy dumps, and record whether profile and timing used the same code and config. - -A kernel name is not a module identity: copy/fill/cast, `index_select`, generic elementwise, GEMM, and different shapes of one template appear in many layers. NVTX/launch correlation and shape are needed to tell "the indexer's gather" from other gathers. Parent/child NVTX inclusive times do not add; the sum of concurrent kernel GPU durations is not wall time; label multi-GPU sum, per-GPU mean and critical-path time separately. - -When a wrapper/overlay also rewrites the kernel selector, confirm who launched a kernel with **three layers of evidence**: per-worker class, object count and import source at initialization; input shape, module name and dispatch at the real forward callsite; launch correlation and GPU kernel in the final profiler. Overridden linear objects or a Humming shape do not prove what launched at replay, and overlay env vars set where the worker does not inherit the entry point falsely suggest the candidate is enabled. Graph may hide ordinary Python CPU op events, so a `cpu_op` count of 0 alone does not prove the callsite did not execute. When evidence conflicts, run a minimal probe that keeps the numerical path; do not fall back to an unverified default operator for throughput. - -For a Roofline bound, use the peak compute of the actual dtype and the bandwidth of the memory level studied; L2 / shared-memory / non-Tensor-Core kernels cannot borrow the HBM / Tensor Core bound, and do not infer the metric from chart coordinates. - -### Pin "where is it slow" to a process and a call stack - -The frontend and control plane are in scope. Split the makespan difference by process: py-spy the engine, worker and frontend once each and see which function holds the main-loop time; use asyncio loop-lag to see whether the frontend is saturated. With R3, most of the frontend main loop is in `output_handler`; without R3 the frontend is almost idle, which moved suspicion from the GPU to the frontend object count. Optimize only parts with a large share; GPU utilization is not evidence of completion. - -### Single-operator replay and synthetic micro benchmark - -A synthetic shape matrix proves algorithm boundaries; real rollout replay supplies score distribution, ties, masks, non-contiguous layouts, cold/hot buffers and call frequency. Weight by sampled shapes/frequencies and keep high percentiles and rare boundaries; report decode small M and prefill large M separately. Time with and without workspace allocation; if the original allocates per call and the candidate preallocates, do not compare only the kernel body. A micro A/B fixes input and resources and checks the candidate bitwise against the baseline; change one thing at a time, covering at least the batch sizes actually used (1/4/32/128), different neighbor compositions, prefill/decode, Graph, mask/tie/boundaries; mark shapes not run as not covered. When Triton is forbidden, a "temporary benchmark" does not bypass that. - -### Ablation method: make A/B really differ in only one variable - -These rules apply to intermittent numerical leaks (the batch-invariance kind) and to performance attribution; the examples come from the DS4 batch-invariance investigation. Process rules are in [Experiment discipline](../practices/evidence-and-experiments.md#experiment-discipline). - -1. **A single round of 0-diff does not count.** For a leak with a hit rate below one half, a single all-green round has probability above one half. Re-check a statistical conclusion with a deterministic reproducer, especially "turned clean"; 0-diff is necessary but not sufficient. -2. **For "turning off X fixes it", first check whether X's switch also changes another branch.** Example: with cudagraph mode NONE, the budget split function also took a different branch, so the eager experiment moved two variables. Hidden coupling cannot be attributed to the switch. -3. **A switch that changes speed also changes the scheduling composition.** Eager is several times slower than graphs, so batch composition changes. Before a statistical comparison, pin each step's batch composition with offline manual steps; otherwise "clean when off" may just mean it was not hit. Turning off the cache is the same: a prefill that never ran before now runs. -4. **The target path must actually be exercised.** A long-prompt filler looks "heavier", but the step exceeds `max_cudagraph_capture_size` and the graph does not run, so identical results across settings are an illusion. A negative control records the target branch's execution count; zero hits is marked VACUOUS and does not count as verification. - -Before localizing, draw a 2×2 (or higher-dimensional) matrix to separate the variables, then choose which cells to run. - -### Ablations must be able to identify interactions - -With optimizations A and B, measure `base / A / B / A+B`; with network/Graph/cache switches, add the corresponding controls instead of running the full combination mechanically. Reuse A+B data only under identical weights, input, software and timing. The first round after restoring checkpoint N is usually step N+1, and adjacent names do not prove the same policy: fix the checkpoint content digest and export fingerprint, and do not mix BI0/BI1 or R3/no-R3 branches. For rollout-only timing, exit cleanly after completion without optimizer/save, and record that all requests completed. Top-K and index post-processing can usually be ablated orthogonally, but shared buffers, layout or Graph create interactions, so re-accept the combination. - -### Segment-by-segment isolation: add the total difference back up - -Cut a long path into arms at boundaries, each enabled up to one segment, and measure what each segment adds: - -| Arm | Enabled up to | What it answers | -| --- | --- | --- | -| All off | No capture | Baseline | -| A0 | Capture only (hook on, nothing lands on the data plane) | Cost of the capture kernel/sync point | -| A | Capture + worker data plane + packing, not sent to the scheduler | Worker-side processing cost | -| E | worker→engine IPC on, scheduler discards | Inter-process transfer cost | -| B | No capture, return path only (fake data) | Return path cost, independent of capture | -| All on | Current path | Total difference | - -Arm differences should roughly sum to all-on minus all-off; otherwise there is an interaction or a dirty cut. This localized the Nemotron-H R3 return path: A0 and A are almost level with baseline, E adds a little, B is close to all on, so about 80% of the overhead is on the per-step engine→frontend return path, not the capture hook. - -Measure same-config repeat noise first and ignore arm differences below it. - -### Optimization priorities and correctness gates - -1. Remove idle/repeated work: launch grids mismatched with the real workload, repeated moves, invalid tokens. -2. Reduce host scheduling/sync and kernel launches; restore a usable Graph with the output contract unchanged. -3. Fuse boundaries such as gather/dequant/index post-processing, keeping the original rounding, traversal order and padding/sentinel. -4. Only then change the kernel algorithm, chunking or collectives; these may change the reduction order and need new equivalence evidence. - -Example: a quantization kernel launches 1024 blocks per expert while decode has few valid rows. Shrink the grid and the stride together; under Graph capture use a host-known capacity bound, not the device token count. If dequant/gather are element-wise independent, re-chunking may keep bit-level equivalence; fusing into sparse attention must keep the scale rounding and KV traversal/reduction order. - -Each candidate first runs a same-input kernel comparison (real dtype/shape, non-contiguous or masked layout, empty experts, tail, extreme values, ties, reused dirty buffers, Graph replay), then the fixed-token engine gates, then actual rollout/training. Formatting checks do not replace compiling and executing; rerun measurements on the final code, since a JIT refactor that looks behavior-preserving may fail to compile. A candidate with no gain or a regression keeps its conclusion and is withdrawn from the default path. - -### Delivery - -Each item lists the actual kernel/path, correctness before and after, micro latency, rollout and full-step results, independent sample count, Graph/topology scope, raw receipts and unsupported configs. State whether quality evaluation and long-term stability are done; kernel speedup does not replace an end-to-end conclusion. - -## 4. Lever results table - -Measured switches: directions and dead ends not to retry. - -| Model / arm | Winners | Dead ends (measured, do not retry) | Structural floor | -| --- | --- | --- | --- | -| DS4-Flash GB200 DP8/EP8 BI=0 | FlashInfer trtllm MoE (`--moe-backend auto` + new FlashInfer); `cudagraph_mode=FULL_AND_PIECEWISE` | deepep_low_latency, `--api-server-count>1`, NCCL_ALGO/PROTO, max-num-batched-tokens 4k/8k, linear backend fitrt/cudnn (cutlass crashes), DBO (forces the V1 runner, which does not accept PIECEWISE) | MoE decode is weight-bound (gemm1 near HBM roofline); DBO rereads weights and loses; overlap capped at the shared experts' few µs | -| DS4-Flash BI=1 | Deterministic top-k + fusion, two PRs together (each alone has almost no effect); DeepEP LL over MNNVL (`VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL=1` + nvl8 overlay) is the biggest lever; after that the shared experts aux stream helps only under MNNVL | DBO, TP2/DP4, TP4/DP2, the skinny bi_gemm, async-scheduling (already default), splitting the shared MLP to fill the combine window | MoE backend locked to DEEPGEMM (the only BI fp8 MoE) -> a2a locked to deepep LL; per-layer dispatch+combine is serial; cross-layer overlap breaks BI; finally the TTFT floor | -| Nemotron-H BI=0 | PC off (when R3 forces PC on, upstream #58851 is best); R3 finish-only + per-step packed IPC (overlay), makespan returns close to the R3-off floor; then packing per-step logprobs into a list for the frontend (overlay) unsaturates the frontend and BI=0 is 16% faster than BI=1 | Raising the zero-copy threshold (no effect); `max_num_seqs` 256 (slower); PR 57968 keys-only as it stands | PC on (Mamba align mode) has a fixed cost: extra eager launches per decode step, 8K prefill split into two segments | -| Nemotron-H 256K training update (both arms) | BI=1 per-GPU per-micro-batch token cap 64K -> 128K: update only about 5% faster per token (13.1 -> 12.4 s / million tokens), actor reserved 99 -> 168 GiB (184 GiB per GPU, cap 170); small gain, thin margin | Less recompute (some layers, only Mamba layers, none): both arms OOM at 32K-64K per GPU. Per-GPU token cap for BI=0 (no difference across 32K-192K): with CP>1 the recipe disables dynamic batching, each micro-batch is one sequence, and the parameter has no effect | Update is 14-18% of the step, rollout about 80% (sandbox execution, slowest trajectories). Compute the share before tuning; confirm in the resolved config that a parameter takes effect | - -BI kernel choice: [BI kernels](bi-kernels.md). diff --git a/.agents/skills/train-inference-consistency/guides/precision-quant.md b/.agents/skills/train-inference-consistency/guides/precision-quant.md deleted file mode 100644 index ddb57a61d7a1..000000000000 --- a/.agents/skills/train-inference-consistency/guides/precision-quant.md +++ /dev/null @@ -1,490 +0,0 @@ -# Precision, quantization and QAT - -The weight conversion chain from optimizer master to vLLM bytes, QAT and low-precision training, and known pitfalls. Numerical contracts without quantization: [numerical-contracts](../workflow/README.md). - -## Find by symptom - -| Symptom / question | Where to go | -| --- | --- | -| Weights unchanged, probs_diff grows after export | [P1](#p1-scale-rule-differs-from-the-checkpoint-requantization-flips-codes), [P10](#p10-both-sender-and-receiver-quantize), [P11](#p11-receiver-picks-quantization-targets-by-name-suffix-and-misses-fused-moe) | -| NVFP4/FP8 scale or packed bytes differ from the checkpoint | [Conversion chain §2](#2-inventory-the-checkpoints-full-quantization-rules-first), [P1-P5](#quantization-rules-and-scale) | -| NVFP4 rollout: QAT or BF16 actor? | [QAT §1](#1-choose-the-precision-route-first), [§4](#4-how-much-mismatch-qat-can-remove) | -| QAT is on but the mismatch did not change | [P7](#p7-qat-switch-is-on-target-modules-were-not-replaced), [P8](#p8-training-grid-differs-from-the-deployment-grid) | -| QAT loss does not drop, deployed weights byte-for-byte unchanged | [QAT §6](#6-why-qat-does-not-train), [P24](#p24-master-starts-on-the-quantization-grid) | -| Whether to add activation fake quantization | [QAT §4.1](#41-activation-fake-quant-aligned-per-layer-worse-end-to-end) | -| Same operator input, inference module output differs | [P13-P15](#training-side-numerics) | -| Fixed logprob offset under FSDP | [P16](#p16-fsdp-bf16-policy-overrides-the-fp32-parameter-contract), [P20](#p20-dividing-by-temperature-in-bf16) | -| Some LoRA/adapter gradients are always zero | [P17](#p17-fake-quant-cuts-off-the-gradient) | -| Adding a new quantization format | [Conversion chain §3](#3-what-each-format-outputs), [§5](#5-minimal-adaptation-for-a-new-quantization-format) | -| Implementing or changing backward/STE | [Backward validation](backward-validation.md) | - -## Key rules - -1. "NVFP4/FP8" is only a format name; a different block scale selection, global scale convention, rounding or static/dynamic activation gives different bytes. Inventory the checkpoint's full rules first. -2. Use the rules of a mainstream implementation (TE, ModelOpt, NeMo-RL) that matches the checkpoint; do not invent your own. Verify from the real BF16 release, not from dequantized weights. -3. Training QAT and the exporter use the same quantization function on the same grid (range, axis, scale, rounding). -4. Compute scale on the full logical tensor. A fusion group shares scale (VERL NVFP4: QKV, gate/up fused global scale). Quantizing each shard separately and concatenating is forbidden. -5. Each tensor family has exactly one conversion owner and is quantized exactly once; post-load reordering does not count as quantization. -6. Activation scales (`input_scale`, `input_amax`, `k_scale`, `v_scale`) are calibration values. Pass them through unchanged; they cannot be derived from weights. Read runtime scales from the actual runner. -7. High-precision training + low-precision rollout is not QAT. Neither fake-QAT nor real low-precision GEMM implies a low-precision optimizer or low-precision gradients. -8. Weight QAT lowers Nemotron-H probs_diff mean from about 0.025 to about 0.010; it reaches 0 only when the training forward uses the same kernels as inference. Whether activation fake-quant helps depends on upstream mismatch; measure first. -9. STE backward has its own acceptance: an independent reference, the same dY, a tolerance agreed in advance. A forward 0 diff does not prove the gradient is right. -10. Check that deployed bytes = `requantize(master)`. -11. Start the master from the real BF16 release, not a dequantized checkpoint, whose values sit on grid points so NVFP4 deployed values only change after thousands of steps at small lr. θ0 bytes then differ from the official checkpoint (Nemotron-H: NVFP4 99.776% identical, FP8 24/46 tensors differ); quantify the cost with proxy SFT and have a human confirm it ([QAT §6.2](#62-fix)). - -## Terms - -- **NVFP4**: FP4 E2M1 values, one FP8 E4M3 block scale per 16 elements, plus one FP32 global scale. -- **block scale / global scale**: per-block scale (`weight_scale`, `scale`) / per-tensor or per-fusion-group scale (`weight_scale_2`, `weight_global_scale`, `scale2`). -- **amax**: maximum absolute value of a tensor (or block), used to compute scale. -- **4over6 / MSE**: TE NVFP4 block scale selection settings; the Nemotron-H checkpoint uses them. -- **E4M3 / E2M1 / UE8M0**: FP8 (4 exponent, 3 mantissa bits) / FP4 (2 exponent, 1 mantissa bit) / unsigned 8-bit exponent (power-of-2 scale). -- **W4A16 / W4A4 / W8A8**: weight bit width / activation bit width. -- **QDQ**: quantize then immediately dequantize, giving a high-precision value on the grid. -- **fake-quant / fake-QAT**: the training forward uses the QDQ value while GEMM stays high precision; not low-precision training. -- **STE**: straight-through estimator; treats the derivative of rounding as 1 (or with a clipping mask). -- **master**: high-precision parameters managed by the optimizer. -- **conversion owner**: the single party that converts a tensor family to the inference format. -- **fusion group**: tensors fused in serving (QKV, gate/up) that must share scale. -- **R3**: MoE route replay; training uses the expert selection recorded by rollout. -- **BI=1 / BI=0**: the `VLLM_BATCH_INVARIANT` switch; 1 is the true on-policy arm where training calls the inference kernels, 0 is the regular baseline. -- **probs_diff**: verl's `training/rollout_probs_diff_*`, the token probability difference between rollout and the training forward. -- **TIS / `rollout_is`**: rollout importance sampling correction. -- **θ0**: policy weights before the first update. -- **oracle**: the reference implementation or value that decides pass or fail. -- **VJP**: the backward function that computes input and parameter gradients from the upstream gradient. - -## Conversion chain: from master weights to the bytes the vLLM kernel uses - -### 1. Four stages - -```text -A. master weights FP32/BF16 optimizer master (plus the FP32 exceptions the model declares) - ↓ forward materialization -B. weights the training forward actually sees - BF16; under QAT it is DQ(Q(W)), the dequantized value on the deployment grid - ↓ export (gather → HF mapping → quantize/pack) -C. inference-format bytes - packed NVFP4 + block scale + global scale / FP8 + scale / - INT4 packed + scale / layers still BF16 in mixed precision - ↓ loader → post-load -D. tensors loaded by vLLM and used by the kernel (including post-load reordering and derived scales) -``` - -| Stage | Invariant | How to verify | -| --- | --- | --- | -| A | Record dtype separately for parameters and buffers; an FP32 master does not mean forward uses FP32 | Read optimizer parameter identity and dtype; see [§6](#6-dtype-contract-six-boundaries) | -| B | Under QAT, `Q` is the same function on the same grid as the export in C | Unit test: weights of B = dequantize(bytes of C), bitwise | -| C | Scale computed on the full logical tensor; fusion group shares the global scale; excluded layers keep their dtype; recomputing on the same master is idempotent | Compare packed bytes, block scale, global scale per tensor; see [§2](#2-inventory-the-checkpoints-full-quantization-rules-first) | -| D | Each family quantized once; post-load only reorders, does not recompute scale; calibration values arrive unchanged | Per tensor, compare exported bytes with the bytes vLLM uses after loading | - -Record the real format and scale of weights, activations and KV (packed UE8M0, FP32 scale, MXFP4/FP8 scope); activation and packing rounding may differ between SM90 and SM100. Before measuring quantization error, confirm the scale comes from the current weights, not an old tensor. - -### 2. Inventory the checkpoint's full quantization rules first - -1. From `hf_quant_config.json` and the `quantization_config` in `config.json`, list each layer's format, which layers carry `input_scale`, and whether the KV cache carries k/v scale. Read layer by layer; do not infer from the model name ([P12](#p12-inferring-a-layers-precision-from-the-models-overall-name)). Nemotron-H Lightning 30B-A3B: `MIXED_PRECISION`; 46 FP8 layers (`mixer.in_proj`/`out_proj` of 23 Mamba layers); 5935 `W4A16_NVFP4` layers (routed/shared experts, lm_head); FP8 KV cache; 46 fp32 scalar `input_scale`, FP8 layers only. -2. Encode the BF16 release with the checkpoint's own scales; 100% identical packed bytes confirm the release matches the checkpoint. -3. From the **real BF16 release**, requantize with the candidate rule and report per tensor the match ratio of packed bytes, block scale and global scale. Dequantized weights already sit on the grid and pass under any rule ([P2](#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). Measured match rates: [P1](#p1-scale-rule-differs-from-the-checkpoint-requantization-flips-codes), [P3](#p3-fp8-per-tensor-scale-convention-is-not-amax448). - -Decisions: - -- Call the rules of a mainstream implementation (TE, ModelOpt, NeMo-RL) directly; do not invent a scale policy. Deviate, and report it, only when the mainstream implementation cannot do it. Nemotron's self-invented "enlarge scale only on overflow" was replaced by an existing TE rule. -- θ0 deployed bytes = the real BF16 release requantized by the rule, which differs from the official checkpoint: on Nemotron-H, NVFP4 bytes are 99.776% identical and FP8 24/46 tensors differ; the proxy SFT cost is about 0.03 nat ([QAT §6.2](#62-fix)). This policy needs human confirmation. -- The BI=1 training forward, BI=0 QAT and the exporter call the same `requantize`; θ0 bytes of the two arms are identical. `input_scale`, `k_scale`, `v_scale` are calibration values passed through from the checkpoint. -- When a single FP8 checkpoint is the source of truth, the training base is a deterministic dequantized derivative, and requantizing need not restore the checkpoint byte for byte. Measure separately: (1) with zero updates, re-exported bytes equal θ0; (2) the precision difference from the original checkpoint. - -### 3. What each format outputs - -| Format | Artifacts | Must check | -| --- | --- | --- | -| block FP8 | FP8 weight + `weight_scale_inv`, scale per block | Commonly 128×128; FP32/UE8M0 branch, backend, exclusion scope | -| per-tensor FP8 | FP8 weight + `weight_scale` | Scale from the amax of the full weight, not of a shard; convention may not be amax/448 | -| MXFP8 | E4M3 weight + UE8M0 `weight_scale_inv` | 1×32 block direction, scale encoding, pad and high-precision exceptions | -| NVFP4 W4A16 (compressed-tensors, existing `fp4` sender) | `.weight_packed` + `.weight_scale` + `.weight_global_scale` | 1×16, two-level scale, shared outer scale for gate/up, activation contract | -| NVFP4 W4A4 (ModelOpt, Inkling) | packed `uint8` under the original parameter name + FP8 `scale` + FP32 `scale2` + BF16 `input_amax` + `original_shape` | Naming not interchangeable with compressed-tensors; needs a separate adapter | -| INT4 compressed-tensors | `weight_packed` + `weight_scale` + `weight_shape` | group_size, symmetric quantization, pack layout, list of tensors actually quantized | -| DeepSeek V4.1 attention FP8 | F8_E4M3 weight + F8_E8M0 `.scale`, 32×32 block | Not rowwise 32-element MXFP8 | - -`weight_scale_inv` is only a loader protocol name; interpret the value per implementation. - -### 4. Who converts: conversion owner - -A tensor family has one conversion owner. Once the sender has packed, the receiver must not quantize again. dtype reinterpret, layout staging and post-load repack are not quantization. - -| Path | Training | Export | Rollout load | owner | -| --- | --- | --- | --- | --- | -| FP8 rollout-only | BF16 | Full BF16 | `quant_weights()` converts to block-FP8 | rollout receiver | -| VERL FSDP QAT NVFP4 | BF16 + fake quant | After gathering the full tensor, `QATQuantizer` outputs packed NVFP4/scale | compressed-tensors | FSDP sender | -| VERL Megatron QAT NVFP4 | BF16 + ModelOpt fake quant | `QATWeightExporter` | ModelOpt | Megatron sender | -| Megatron FP8 E2E | TE FP8 | Bridge-specific FP8 export planning | vLLM FP8 | Bridge / rollout, per the pinned recipe | -| Nemotron-H BI=0 (verl Megatron) | BF16 + ModelOpt weight STE | Requantized by the checkpoint's rules into ModelOpt `MIXED_PRECISION` | ModelOpt | sender (`mixed_precision_export=True`) | -| Pre-quantized checkpoint passthrough | No requantization | Original tensor/scale | Per the checkpoint layout | checkpoint producer | - -- Full merge: convert logically complete HF tensors one at a time by their owner; no whole-model all-gather or per-step disk write. -- QAT: the training forward simulates quantization; export generates the actual packed weights and scales. - -#### 4.1 FSDP path (VERL) - -- With `actor_rollout_ref.actor.fsdp_config.qat.enable=true`, quantizable `nn.Linear` modules become `QATLinear` after the PEFT wrapper and before the FSDP wrap. Forward does NVFP4 W4A16/W4A4 fake quant and ends in `F.linear`; after the optimizer step, `invalidate_all_scales` clears the weight scale cache. [Implementation][qat-linear] -- Export: `FSDPEngine.get_per_tensor_param()` calls `full_tensor()` on each `DTensor`, then `QATQuantizer.quantize_with_fusion()`: - -```text -FSDP local shards → full_tensor() all-gather → full HF weight -→ full-weight amax / global_scale → QKV or gate/up fused_global_scale -→ per-row, per-group block scale → packed NVFP4 weight + scales -``` - - The cost is the all-gather and GPU staging. The quantizer caches to CPU per decoder layer and moves back to GPU, so peak GPU memory is not necessarily one tensor. [Export][qat-export] - -- `get_per_tensor_param_shard()` outputs only the local BF16 shard, with no quantizer and no scale. Shard-local block quantization is an unsupported new protocol and must at least solve: - 1. aligning shard and block boundaries, or handling cross-shard blocks; - 2. a collective for the full-weight amax/global scale; - 3. a consistent fused global scale for QKV and gate/up; - 4. W4A4 explicitly carrying the activation-derived `input_global_scale`, which cannot be rebuilt from weight shards. - - Until these contracts and the oracle exist, quantizing each shard and concatenating is forbidden. -- Non-QAT: export BF16 in HF coordinates; BF16 rollout does not quantize, and `rollout.quantization=fp8` is quantized by the receiver. LoRA merge exception: [P6](#p6-lora-mergetrue-bypasses-the-fsdp-qat-export). - -#### 4.2 Megatron path - -- VERL: ModelOpt installs fake-quant modules; `MegatronEngine.get_per_tensor_param()` exports HF weights through Megatron-Bridge, and `QATWeightExporter` reads the module quantizer's amax/block metadata to produce NVFP4 and scales. [ModelOpt exporter][modelopt-source] -- VERL FP8 E2E: TE handles FP8 forward/backward in training; when the Bridge declares `export_weight_dtype=fp8`, VERL does not cache generic conversion tasks. The rollout format is decided by `rollout.quantization` and the serving checkpoint contract; FP8 training does not mean exported tensors are in vLLM layout. -- Megatron capability cannot be inferred as FSDP capability. -- Distinct artifacts: the offline converter prepares the initial checkpoint (may write to disk); the online exporter streams in memory; post-load reorders without recomputing scale; a training checkpoint is not a servable quantized HF checkpoint. - -#### 4.3 Receiver quantization (rollout-only) - -`rollout.quantization=fp8`: the trainer exports BF16; `load_quanted_weights()` in the vLLM worker calls `quant_weights()`, which uses `scaled_fp8_blockwise()` to produce the FP8 weight and `weight_scale_inv`, then replays the quant method's post-load. There are also a DeepSeek V4-specific entry and an Ascend MXFP8 branch with different scale naming. [receiver][receiver-source] When the checkpoint is already FP8/MXFP4, confirm whether it is requantization or passthrough. vLLM cannot quantize BF16 online into ModelOpt `MIXED_PRECISION`. - -### 5. Minimal adaptation for a new quantization format - -1. From the checkpoint config and the actual loader, list key, shape, dtype, block size, packing and fusion group of weight/scale, and the unquantized tensor families. -2. Trace export -> loader -> post-load and mark the single owner of each family. -3. If a converter is missing, add the minimal conversion at the existing owner's entry; verify on one tensor / fusion group, then cold-load a small model. Changing the dtype is not a conversion. -4. Test zero merge and a non-zero update; get quantize-after-full-gather right before shard optimizations. - -One evidence table is enough: tensor family -> source -> export format -> serving format -> owner -> unit/small-model verification. - -#### Minimal acceptance - -- [ ] Record training dtype, export dtype/layout, rollout dtype/layout and owner. -- [ ] Fixed-token logprob of the zero-adapter path and the base are within the noise baseline of the same backend. -- [ ] After a non-zero update, tokenwise logprob/KL between trainer and rollout stays within bounds. -- [ ] name/shape/dtype/count/bytes of packed weight, scale and ignored families all close. -- [ ] Deployed bytes = `requantize(master)`. -- [ ] FSDP QAT proves scale is computed after `full_tensor()`; with shard export, prove it is not quantized, or provide the four shard-local contracts of §4.1 and a cross-rank oracle. -- [ ] Each family is quantized exactly once. - -### 6. dtype contract: six boundaries - -For each model, record `checkpoint → load/bridge → wrapper/sharding → forward → backward/optimizer → export/serving` separately. - -| Boundary | Megatron | VERL/FSDP | Must check | -| --- | --- | --- | --- | -| Source weights and conversion | Bridge may reorder, split, cast; model may override default conversion | HF loader, dequantized derived checkpoint, init dtype | Same revision; logical value/scale/layout match, not only target dtype | -| Parameter storage | `params_dtype` plus model exceptions; FP32 master/optimizer | HF keep-precision list, PEFT master, FSDP shard | An FP32 master does not mean forward uses FP32 | -| forward parameters | Float16Module converts selectively | FSDP2 `MixedPrecisionPolicy.param_dtype` | Measure inside forward, not only before the wrap | -| Operator computation and output | TE/FLA/TileLang/fused each define product, accumulation and output precision | HF/PEFT kernel and autocast | FP32 output does not prove FP32 product; BF16 output does not prove BF16 accumulation | -| Gradient and optimizer | TP/EP/SP/CP collective, DDP | FSDP reduce dtype, gradient accumulation | FP32 reduction does not restore values rounded in forward | -| Export and serving | Undo the parallel layout, convert adapter/quantization format | gather/export HF adapter | Compare values on the same logical layout | - -FP32 protection takes several forms: mark-to-skip before conversion, reset after declaration (which does not restore values already lost), bypassing the default BF16 downcast, and operator-level selection. It is not "all special parameters are raised in precision". - -FSDP: HF `_keep_in_fp32_modules_strict` is a load contract and does not become a per-parameter cast exception in FSDP ([P16](#p16-fsdp-bf16-policy-overrides-the-fp32-parameter-contract)). Respect the model's existing contract first, and verify any change of computation precision against a same-model reference and actual serving. `param_dtype=None`, `param_dtype=fp32` and FP32 autocast are different things. - -### 7. Code entry points - -FSDP [engine][fsdp-source], [module selection][qat-core], [fake quant][qat-linear], [NVFP4 export][qat-export]; Megatron [ModelOpt exporter][modelopt-source]; rollout [receiver][receiver-source]. Historical PR pointers (not the current state): early return on LoRA merge [#5115](https://github.com/verl-project/verl/pull/5115)/[#5411](https://github.com/verl-project/verl/pull/5411)/[#7014](https://github.com/verl-project/verl/pull/7014); fake quant and export quantization [#7149](https://github.com/verl-project/verl/pull/7149)/[#7224](https://github.com/verl-project/verl/pull/7224)/[#7577](https://github.com/verl-project/verl/pull/7577). - -[fsdp-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/workers/engine/fsdp/transformer_impl.py -[qat-core]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/core.py -[qat-linear]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/linear.py -[qat-export]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/quantizer.py -[modelopt-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/modelopt/qat_weight_exporter.py -[receiver-source]: https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/vllm/vllm_quant_utils.py - -## QAT, low-precision training and STE - -### 1. Choose the precision route first - -Rule out input, weight, adapter, layout and operator correctness errors first, then choose precision: - -1. When GPU memory allows and rollout can use BF16, use BF16 for both training and inference; check FP32 exception parameters one by one. -2. When rollout must be low precision (such as NVFP4), use weight QAT or a BF16 actor for training. On Nemotron-H the weight QAT arm is stable. Whether `rollout_is=token` stabilizes the BF16 actor is **not tested**. -3. Consider real low-precision training last: backward/kernel support and maintenance cost more. - -Same dtype is not same computation, and QAT is not bit-exact. QAT and low-precision training shrink the **weight value** gap; the kernel and reduction-order gap goes only when the training forward calls the inference kernels (the BI=1 arm of the [BI switch recipe](recipe-and-acceptance.md)). A natively low-precision model does not necessarily run with `dtype=bf16`; do a capacity pre-check first. - -### 2. What the four kinds of path are - -| Path | Training forward | backward / optimizer | Is it QAT | -| --- | --- | --- | --- | -| High-precision training + low-precision rollout | BF16 weights | High precision | No; quantized only at export | -| fake-QAT | `DQ(Q(W))`, GEMM is still BF16 | STE sends the gradient back to the high-precision master | Yes | -| Real low-precision GEMM (FP8/MXFP8/NVFP4) | TE changes the compute kernel | High-precision master/exception tensors remain | Different from fake-QAT | -| Frozen quantized base + high-precision LoRA | Dequantized base + adapter | Only the adapter is trained | No | - -- Neither fake-QAT nor real low-precision training makes the optimizer/gradients low precision. W4A16 quantizes only weights; W4A4 also quantizes inputs. -- Real low precision: NVFP4 forward is not FP4 backward; configurations include high-precision backward and backward on the dequantized operand. -- Only weight-only formats can compute scale from the full weight; formats with activation quantization also need activation scale/calibration state, which weights cannot provide. QAT simulation and export follow the same format rules, but scales are not necessarily constant during training. -- When LoRA is applied separately, do not fake-quant the merged base+adapter yourself. -- Open-source frameworks (verl [#6992](https://github.com/verl-project/verl/pull/6992), [#7971](https://github.com/verl-project/verl/pull/7971), NeMo-RL) essentially only do W4A16; activation inconsistency only appears at W4A8. - -### 3. Implementing fake-QAT - -Forward form (Nemotron-H): - -```text -w_visible = w + (dequant(requantize(w)) - w).detach() -``` - -- `requantize` is the same function as the exporter; replace the library's built-in quantizer when its rules differ. -- Requantization runs on every forward, and again under full recompute; Nemotron-H step time is about +60%. -- A temporary `QuantModuleRegistry` replacement and the old class-attribute patch give the same probs_diff mean on the same quick workload (0.0095/0.0091/0.0100 vs 0.0096/0.0092/0.0101); module hit check: [P7](#p7-qat-switch-is-on-target-modules-were-not-replaced). - -### 4. How much mismatch QAT can remove - -Nemotron-H (`rollout_probs_diff` mean, NVFP4 mixed-precision deployment, single node 4×GB200): - -| Approach | Training forward | probs_diff mean | -| --- | --- | --- | -| BF16 training + quantized rollout | BF16 weights | About 0.025; no rollout correction | -| Weight QAT (STE) | Dequantized deployed weights, BF16 GEMM | About 0.010 (pearson 0.986→0.997); three 3-step smoke runs of one recipe, 0.0091-0.0102 | -| Weight QAT + activation fake-quant | Same, plus static FP8 activations | About 0.0107, about 10% worse than weights only | -| true on-policy (BI=1) | Deployed bytes + the same kernels as vLLM | max = mean = 0 at every step for 100 steps | - -The QAT arm's probs_diff max is still 0.73-1.0; a small mean does not rule out large per-token deviations. Other models (different metrics, not thresholds): DeepSeek-V4 BF16 train + FP8 rollout has a first-step logprob drift of about 0.023 and still reports rising reward/eval ([DeepSeek][dsblog]); Kimi K3 frozen quantized base + LoRA has a train/rollout KL noise floor of about 0.002, a Schulman KL estimate, not a mean absolute logprob diff ([K3][k3blog]). - -#### 4.1 Activation fake-quant: aligned per layer, worse end to end - -Adding vLLM's static FP8 activation fake-quant to training aligns single-layer output to the BF16 noise floor, but end-to-end probs_diff gets about 10% worse: FP8 rounding is discontinuous, so any element-wise upstream mismatch flips codes and amplifies the difference. To judge, add relative noise δ only to the training-side input and sweep the output error; the crossover was about δ≈1%, and activation quantization hurts when upstream mismatch exceeds it. This is one measurement on one model, not a threshold; for a new model, measure upstream mismatch and repeat the sweep. Nemotron-H did not merge it. - -Rule: the weight-side range/axis/scale/rounding must match serving ([rule 3](#key-rules)); matching the activation side depends on upstream mismatch and is not unconditional. Mainstream frameworks (verl, NeMo-RL, slime) use the same **dynamic** activation rule on both sides and do not reuse an external checkpoint's static activation scales; the NeMo-RL Nemotron recipe quantizes only routed experts and leaves Mamba projections, attention and lm_head in BF16. On DeepSeek V4.1, KV quantization simulation adds only about 1.93% improvement on top of R3. Training-side simulation of the FP8 KV cache is not covered. - -#### 4.3 When QAT is worth it - -Evaluate unified FP8/QAT only when evidence shows the train-inference quantization difference is the main limit; "no score gain" is not a reason to switch precision first. For DSv4.1 rank2 LoRA, the post-merge quantized delta is only 5.1-7.7% of the pre-quantization delta (a norm ratio, not information retention), which only supports treating step size/quantization as controlled variables. The long-running Qwen Flash-Next experiment is same-source BF16/BF16 and does not validate the FP8 rollout route. - -### 5. STE contract - -The STE surrogate (identity or clipped), whether dX uses the master or dequantized value, scale handling and verification are in [Backward validation](backward-validation.md#write-down-the-gradient-contract-first). By default dX uses the dequantized value; when using the master value, quantify the bias. - -### 6. Why QAT does not train - -Train-inference consistency is the strictest QAT: forward uses quantized weights, the high-precision master accumulates updates, STE passes the gradient back. Whether it trains depends only on the master's distance to the nearest decision boundary. In W4A16 (the FP4 code is decoded to BF16 in a register and goes through the BF16 tensor core; the gain is 1/4 the weight bandwidth), W4A4 and W8A8 alike, the computed values are "code × scale". - -#### 6.1 Dead zone - -- Each NVFP4 element has only 8 values, 0/0.5/1/1.5/2/3/4/6 × scale, with RNE rounding. From a dequantized checkpoint the master sits on grid points, half a spacing from every decision boundary, and must travel that half spacing before the code changes. -- At lr 1e-6 one step moves the master about 0.01%; a few hundred steps move it far less than half a spacing, so deployed values stay byte-for-byte unchanged. RL gradients are more random than SFT, closer to a random walk, and need more steps. -- Result: NVFP4-covered parameters (on Nemotron-H, routed/shared experts and lm_head, the bulk) have zero deployed change, and validation NLL is clearly worse than the native BF16 baseline; freezing these weights explains about half the gap. -- Recomputing block scale does not flip codes earlier: from grid points requant is idempotent, the E4M3 scale only changes code after block amax moves 3-6%, and a tiny global scale change only rescales the whole tensor. -- Mainstream QAT (NeMo-RL QARL, slime/Kimi INT4, Jet-RL) starts from the original BF16 weights, which lie between grid points, so any small drift flips a proportional share of codes and there is no dead zone. INT4 with 1×32 groups from a dequantized checkpoint stalls the same way. -- FP8 is barely affected: floating-point relative spacing is 6-12%, small elements have small absolute spacing, and some elements are always near a boundary. DS4 also starts from dequantized FP8 and is affected only mildly. - -#### 6.2 Fix - -Preferred: use the original BF16 weights (the real BF16 release) as the master and requantize every step by the same rule (Nemotron-H: TE 4over6, E4M3 upper bound 256). State two points: - -- **θ0 is not equal to the official checkpoint bytes.** Requantizing from the real BF16 release, 99.776% of Nemotron-H's NVFP4 bytes are identical and 24 of 46 FP8 tensors (ModelOpt calibrated scale) differ. "The official checkpoint was generated by this rule" holds only on dequantized weights, where MSE and MAE cannot tell rules apart on grid points ([P2](#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). Proxy SFT gives the cost: 300-step val NLL A-bf16 3.4685 vs C-dequant 3.4641 / C-bf16 3.4373, about 0.03 nat, present at θ0 and not growing during training. This policy needs human confirmation. -- **Three acceptance items** (all pass on Nemotron-H): (1) with zero updates, re-exported bytes equal θ0 (not the official checkpoint); (2) fast-tier RL θ0/θ1/θ2 is 0 diff and deployed bytes start changing; (3) proxy SFT gap to the native BF16 baseline is within budget. - -Alternative, keeping θ0 bytes: add a fixed small jitter to move the master off grid points (toward the neighboring cell, at most half a step; after requant check per block that bytes are unchanged, and roll back changed blocks). Not recommended: stochastic-rounding requantization (about 1.4× deployment error, needs a self-written quantizer, no mainstream precedent), raising lr (under Adam equivalent to QuRL UAQ, only shortens the dead zone, 10× lr is outside the usual RL range), and frozen base + LoRA (QeRL, gives up full-parameter updates). Whichever you choose, state that from a dequantized checkpoint at lr 1e-6, NVFP4 deployed values do not change for roughly the first 10³-10⁴ steps. See [P24](#p24-master-starts-on-the-quantization-grid). - -[dsblog]: https://www.lmsys.org/blog/2026-04-25-deepseek-v4/ -[k3blog]: https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support/ - -## Known pitfalls - -Ordered along the conversion chain: quantization rules -> export and receive -> training-side numerics -> experiment measurement basis. P24 sits under training-side numerics by content. - -### Quantization rules and scale - -#### P1 scale rule differs from the checkpoint, requantization flips codes - -- Symptom: weights unchanged, rollout shifts after a requantized export. -- Cause: recomputing scale on the dequantized master with the ModelOpt default global = amax/(6·448) differs from the checkpoint's TE 4over6 rule and flips about 40% of Nemotron-H's FP4 codes; recomputing with TE 4over6 is idempotent and byte-for-byte unchanged. -- Fix: use a mainstream rule that matches the checkpoint; training QAT and the exporter call the same `requantize`. -- How to verify: requantizing the same master twice is byte-for-byte identical; requantize from the real BF16 release and compare per tensor with the checkpoint. Nemotron-H: - - | Rule | Match with the checkpoint | - | --- | --- | - | TE NVFP4 4over6 + MSE, E4M3 upper bound 256, global = amax/1536 | packed 99.776%, global 100%; no setting closes the remaining gap | - | ModelOpt default global = amax/(6·448) | packed about 50% | - | Standard FP8 per-tensor amax/448 | See P3 | - -- Source: procedure in [Conversion chain §2](#2-inventory-the-checkpoints-full-quantization-rules-first). - -#### P2 Verifying a rule with dequantized weights gives a false positive - -- Symptom: several candidate rules all report a 100% match. -- Cause: dequantized weights already sit on the grid and do not change under any rule. -- Fix: verify from the real BF16 release. -- How to verify: at least one known-wrong rule fails the test. - -#### P3 FP8 per-tensor scale convention is not amax/448 - -- Symptom: FP8 layers differ from the checkpoint after requantization. -- Cause: Nemotron-H's official Mamba projection scale is about 1.5×amax/448; 24 of 46 FP8 tensors differ from amax/448. -- Fix: derive the convention from the checkpoint, or pass the checkpoint's scale through. -- How to verify: report the scale match rate per tensor. - -#### P4 Computing scale separately on each shard - -- Symptom: packed weights concatenated after an FSDP/TP sharded export differ from quantizing the full tensor. -- Cause: blocks cross shards; global scale and fused global scale need the full tensor. -- Fix: `full_tensor()`/gather, then quantize; shard-local quantization must first meet the four contracts of [Conversion chain §4.1](#41-fsdp-path-verl). -- How to verify: evidence shows scale is computed after the gather. - -#### P5 gate/up not quantized as a pair - -- Symptom: wrong values after serving fuses gate/up, or the exporter reports a missing pair. -- Cause: the two HF tensors are fused in serving and need a shared outer scale. -- Fix: quantize as a pair; raise an error when a complete pair is missing, and test pairing across safetensors shards. -- How to verify: test pair completeness, cross-shard gate/up and a TE bitwise reference on matching hardware; a test existing does not mean it ran and passed. - -#### P6 LoRA merge=true bypasses the FSDP QAT export - -- Symptom: QAT and LoRA merge are on, and the export is not the QAT quantizer's packed weights. -- Cause: VERL FSDP returns `_merged_lora_per_tensor_param()` directly for `merge=true`, skipping `QATQuantizer`; with `rollout.quantization=fp8` the receiver may still requantize. The `_qat_enabled` check still follows the `merge=false` branch. -- Fix: the first merge fallback does gather -> merge -> quantize -> yield -> release one full tensor at a time. -- How to verify: check actual module replacement, naming and export scope; both switches being on does not make the combination usable. -- Source: [VERL engine](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/workers/engine/fsdp/transformer_impl.py). - -#### P7 QAT switch is on, target modules were not replaced - -- Symptom: the QAT arm's probs_diff equals the BF16 arm's. -- Cause: VERL `QATLinear` enters through `nn.Linear` and does not cover three-dimensional expert Parameters or arbitrary PEFT wrappers; the framework's built-in QAT changes only the GPT layer spec, while Mamba hybrid models read the hybrid stack spec. -- Fix: attach the quantizer after building the model by the actual spec; use the exporter's `requantize` as the quantizer. -- How to verify: every rank prints the count and list of replaced modules and fails closed (139 for Nemotron-H, the same on all four ranks). -- Source: [Conversion chain §4.1](#41-fsdp-path-verl), [BI=0 integration pitfalls](recipe-and-acceptance.md#bi0-arm-integration-pitfalls). - -#### P8 Training grid differs from the deployment grid - -- Symptom: QAT is declared, the mismatch remains. -- Cause: the converter and the trainer use different group sizes (for example 32 vs 128). An external TPU case is bitwise 0 on the FP8 per-channel grid, but the native weights are 128×128 block, and the author notes the result may not hold on a different grid. -- Fix: align quantized module scope, block/group size, axis, scale dtype/rounding, global scale scope, gate/up fusion, activation scale, TP/ETP partition and high-precision exceptions. FP4 has 1 mantissa bit, so error comes mainly from block partitioning; block size and direction must match the native weights. MiMo-V2.6 does QDQ on the experts after every update under the rollout MXFP4 GEMM's constraints. -- How to verify: training-side QDQ output = dequantize(deployed bytes), bitwise. -- Source: unified NVFP4 constraints in the [technical blog](https://www.lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl/) (its bit-exact quantizer test verifies the format, not whole-model bit-exactness). - -#### P9 Activation scale derived from weights - -- Symptom: W4A4 export raises, or the exported activation scale differs from serving. -- Cause: activation scale depends on actual input. VERL W4A4 export raises when an initialized `input_global_scale` is missing; Inkling's `input_amax` cannot be derived from BF16 weights. -- Fix: reuse unchanged from the QAT calibration state or rollout checkpoint; pass through `input_scale`, `k_scale`, `v_scale`. -- How to verify: manifest activation scales are byte-for-byte identical to the checkpoint. -- Source: [VERL quantizer](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/qat/quantizer.py). - -### Export and receive - -#### P10 Both sender and receiver quantize - -- Symptom: deviation larger than one quantization; or a packed tensor treated as BF16 and quantized again. -- Cause: a tensor family has two conversion owners; the Bridge already outputs packed tensors but the generic packer does not recognize them. -- Fix: one owner per family; pass packed tensors through. -- How to verify: the evidence table lists the owner per family; count/bytes close. -- Source: [Conversion chain §4](#4-who-converts-conversion-owner). - -#### P11 Receiver picks quantization targets by name suffix and misses fused MoE - -- Symptom: under `rollout.quantization=fp8` some experts are still BF16. -- Cause: `gate_up_proj`/`down_proj` do not end in `.weight` and are skipped; fixed by VERL [#7443](https://github.com/verl-project/verl/pull/7443). -- Fix: check coverage using checkpoint names together with runtime modules/dtype; review the list for each new model. -- How to verify: list the serving dtype of every family. -- Source: [receiver](https://github.com/verl-project/verl/blob/753aed3e1c286ba6825a74342b28669e72c083ea/verl/utils/vllm/vllm_quant_utils.py). - -#### P12 Inferring a layer's precision from the model's overall name - -- Symptom: diagnosis goes in the wrong direction. -- Cause: in the Nemotron-H "NVFP4" checkpoint, `backbone.layers.16.mixer.out_proj` has `quant_algo` FP8. -- Fix: read `hf_quant_config.json` layer by layer. -- How to verify: the diagnosis states the actual `quant_algo` of the first differing layer. - -### Training-side numerics - -#### P13 scale taken from the wrong object - -- Symptom: the unit-scale unit test passes; real attention output and old-logprob fail. -- Cause: the training adapter quantizes Q with `k_scale`. On the real runner `q_scale=1.0`, `k_scale=0.03173828125`, `v_scale=0.003326416015625`. -- Fix: read scale from the runner/module actually executed; do not infer a default from a missing checkpoint field or another tensor. -- How to verify: replay with real non-unit scales and frozen Q/K/V; after the Nemotron-H fix, 5,832,704/5,832,704 BF16 outputs are bitwise identical. - -#### P14 Quantization modules with the same name take different dispatch - -- Symptom: calling `ops.scaled_fp8_quant` + `flashinfer_scaled_fp8_mm` directly matches training bitwise, yet the inference linear module differs in 1,361 BF16 elements with the same scale. -- Cause: the module selects the **compiled** `forward_native`, which has 44 more differing FP8 elements than uncompiled native and the CUDA custom op; those two agree. Which compile transformation is responsible is not confirmed. -- Fix: under Nemotron-H BI=1, static per-tensor `QuantFP8` switches to the device custom op (GEMM and format unchanged). -- How to verify: with the same operator input, weights and scale, compare FP8 bytes, GEMM output and module output in turn, calling the native/CUDA/actual module paths separately. When the direct GEMM matches training and the module does not, check quantization dispatch inside the module first; do not change GEMM. - -#### P15 Quantizer is not in the `nn.Module` tree, the patch misses layers - -- Symptom: the patch "ran", the target layer still has 1,361 differing elements. -- Cause: traversing `model.modules()` finds only 6 registered modules and misses quantizers on plain objects such as `linear.quant_method.kernel.quant_fp8`. -- Fix: traverse the registered tree, the quant-method kernel objects and that version's legacy fields together. -- How to verify: record per worker the hit count (52 after the fix), target layer names, actual `_forward_method` and quantized output bytes; fail if expected targets are zero or a key layer is missing. Private field paths are only a versioned case. - -#### P16 FSDP BF16 policy overrides the FP32 parameter contract - -- Symptom: Hy4 four-layer fixed-token logprob MAE 0.022982/0.023052. -- Cause: HC and attention sink, FP32 after HF load, are downcast by the uniform BF16 FSDP policy; a later `.float()` in the operator cannot restore them. The router's FP32 conversion is also overridden by the outer BF16 autocast, which changes the expert set. -- Fix: opt in with `wrap_policy.preserve_original_param_dtype=True` (`param_dtype=None`, BF16 autocast kept). This is a whole-model policy, not a per-parameter whitelist; measure communication/memory impact; it does not raise a BF16 base to FP32. Disable autocast inside the router forward. -- How to verify: with original dtype kept, singleton, two ranks and unwrapped are exactly identical; the router has a regression test. - -#### P17 fake-quant cuts off the gradient - -- Symptom: in DSv4.1 LoRA training, 31 of 160 B matrices are zero, all in `self_attn.kv_proj`. -- Cause: `_fake_quant_fp8_block` goes through an FP8 cast on the autograd path, so small gradients underflow to 0; `_fake_quant_fp4_block` has no STE (integer comparison and LUT gather, so gradient flows only through scale, and the ue8m0 branch cuts even that). -- Fix: keep the quantized forward and use identity for the input gradient: `output.detach() + (x - x.detach())`. -- How to verify: a 1e-6 upstream gradient is preserved exactly; after the fix 160/160 B matrices are non-zero. The loss still rises, so this does not prove training improved. - -#### P18 Training forward and export do not quantize at the same place - -- Symptom: the model already has FP4/FP8 STE, the mismatch remains. -- Cause: DSv4.1's existing STE acts on indexer/compressed/window KV activations, not as weight fake-quant on `W+ΔW`; the training forward is unmerged base + LoRA, while export is an FP32 merge followed by block32 FP8/E8M0 conversion. -- Fix: confirm whether training sees the post-export rounding before deciding on weight QAT. -- How to verify: sample the actual dtype inside forward, not only the code path. - -#### P19 KV cache aligned only on the dtype label - -- Symptom: attention results differ with the same weights. -- Cause: in the old vLLM image, KV is FP8/block64 for the non-RoPE part and BF16 for the RoPE part, and the training side handles values differently. -- Fix: set `kv_cache_quantization="fp8_ds_mla"` explicitly in training, use the same rounding for forward with and without gradients, and keep the gradient with STE. Re-verify after the upstream KV layout changes. -- How to verify: compare attention output with the same weights. - -#### P20 Dividing by temperature in BF16 - -- Symptom: logprob has a fixed offset. -- Cause: the old VERL FSDP path divides by temperature in BF16/FP16, then raises precision. -- Fix: `logits.float()` first, then divide by temperature, even when temperature is 1. -- How to verify: both scorers match on the same logits. - -#### P21 Derived scale after resume still belongs to the old weights - -- Symptom: the first rollout after resume is wrong. -- Cause: BF16 parameters were updated, but the FP8 source-scale registry still trusts the initial weights. -- Fix: invalidate derived caches in the weight restore lifecycle. -- How to verify: first-step probs_diff after resume is on the same measurement basis as before (Nemotron-H kill->resume acceptance: still 0 after resume). -- Source: [checkpoint guide §13](checkpoint-resume.md#13-stale-quantization-scale-on-the-first-step-after-restore). - -#### P24 Master starts on the quantization grid - -- Symptom: QAT loss barely drops and NVFP4 deployed weights are byte-for-byte unchanged while the master moves. -- Cause: a master initialized from a dequantized checkpoint sits on grid points ([QAT §6.1](#61-dead-zone)). -- Fix: master from the real BF16 release, requantized every step by the same rule, cost quantified with proxy SFT; without the original weights, a fixed within-cell jitter. Not stochastic rounding, larger lr, or LoRA. -- How to verify: the three acceptance items of [QAT §6.2](#62-fix). - -### Experiment measurement basis - -#### P22 The quantization oracle itself is wrong - -- Symptom: the GPU block32 quantization gate fails (zero block expected 2.27373675e-13, actual 1.16415322e-10). -- Cause: the oracle assumes Triton's epsilon position; the CUDA UE8M0 path also clamps scale to 1e-10 before taking the power of 2. -- Fix: fix the oracle, not the kernel. CUDA and Triton are not thereby bitwise identical on tiny/zero blocks. -- How to verify: check the oracle's epsilon/floor position against the kernel source before judging a failure. - -#### P23 Attributing every logprob difference to quantization - -- Symptom: the difference persists after changing precision. -- Cause: routing selection, attention/kernel and tokenization/mask all produce differences. R3 constrains routing and does not make FP8 GEMM equal BF16; TIS is a statistical correction and does not fix wrong weights. -- Fix: rule out protocol errors first, then record the quantization/numerical noise floor. -- How to verify: after ruling out protocol items (routing, mask) one by one, measure the quantization noise floor on the same input. diff --git a/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md b/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md deleted file mode 100644 index 00762fe8c4e9..000000000000 --- a/.agents/skills/train-inference-consistency/guides/recipe-and-acceptance.md +++ /dev/null @@ -1,113 +0,0 @@ -# Recipe, evaluation and paired acceptance - -Which recipe items must match the reference implementation mathematically (not by name), what puts two evaluations on the same measurement basis, and how to pair two arms so a conclusion can be drawn. The reference recipe for train-inference consistency is [verl#7688](https://github.com/verl-project/verl/pull/7688) (DS4 exact 0-diff true on-policy). Which reference a model aligned with, and the values it ended up using, are per-model facts and are not recorded here. - -## Find by question - -| Question | Where to go | -| --- | --- | -| Does a 5% gap between two arms count as a gap | [Paired acceptance](#paired-acceptance): noise floor, median over multiple pairs, confounders | -| The validation score only takes a few discrete values | [Evaluation contract](#evaluation-contract): under greedy the score is a multiple of 1/number of problems | -| How to design the BI=0 baseline arm | [How the two arms are defined](#how-the-two-arms-are-defined) | - -## Key rules - -1. Align the math, not the names; read the final CLI/resolved config, not the README or the model name. -2. Tabulate every deviation from the reference: item, reference value, our value, why, which conclusion it affects. -3. One recipe, one image, a switch selects the arm (following DS4 `EXACT_ALIGNMENT`, default BI=1; BI=0 goes through Bridge/mcore + R3 + QAT); deviations live only in the switch. -4. If any of the six evaluation items differs, it is not the same score: dataset version, prompt template, sampling parameters, number of repeats, scorer, length cap. -5. Rising training reward, probs_diff of 0, and a held-out score gain are three separate conclusions; none substitutes for another. -6. Before running, fix what is compared, how it is computed, the pass line and the noise floor; plant a known error to confirm the criterion can return FAIL. -7. One artifact answers one question: run and report performance, quality and consistency separately. -8. List confounders first, then compare; each is either aligned or quantified separately ([paired acceptance](#paired-acceptance)). -9. Keep failure records as they are and append new results; keep reruns and changed questions apart. - -## One recipe, one image: deliver true on-policy and the regular baseline with the BI switch - -Following DS4's EXACT_ALIGNMENT: one delivery image, one recipe, and `VLLM_BATCH_INVARIANT`, defined by the recipe, selects the whole arm. The arms may differ in training engine, precision and route replay; the switch is only the single entry point. The pitfalls below come from Nemotron-H Lightning NVFP4 (Mamba hybrid MoE); a new model must re-check them item by item, not copy the config. - -### How the two arms are defined - -BI=1 is true on-policy: the training forward calls the vLLM kernels and reads the deployed bytes, rollout takes the vLLM BI path, and the deployed bytes are sent down unchanged. BI=0 is the regular baseline: training takes the default MCore path (verl Megatron engine/Megatron-Bridge, without the alignment implementation of [C6](../workflow/README.md#gate-table)), and rollout uses vLLM's non-BI kernels with the configuration audited per [C1](../workflow/README.md#gate-table). For BI=0, decide the items below per model and record who approved each: - -| Decision item | When needed | Nemotron's choice (example) | -| --- | --- | --- | -| Training precision: bf16 or weight QAT (STE) | QAT exists only when deployment uses quantized weights | QAT by default, bf16 optional | -| MoE route replay (R3) | MoE only | Always on | -| Requantize from master at export | Quantized deployment only | Requantize by the checkpoint's rules at every export | -| Rollout importance sampling correction (TIS etc.) | Arms whose train-inference mismatch is not 0 | Not enabled | -| rollout MoE backend | When the backend picked by `auto` does not run | Explicit `flashinfer_cutedsl` (`auto` picks Marlin, which reports a PTX version error), matching the BI0 arm of the BI=1 bench | - -Do not build BI=0 as "the BI=1 actor with non-BI kernels swapped in". That requires changing the alignment implementation, and the result is neither the regular baseline nor true on-policy; a human rejected this design when it was tried. - -### Rules - -- **Same image, same recipe.** If the recipe lacks a switch, add it on the delivery branch first, then build one image shared by both arms; no separate image, recipe or source mount for a single experiment. Delete the old standalone baseline recipe. -- **Keep output and checkpoint paths separate.** Suffix only BI=0; BI=1 paths stay as before. With shared paths, `resume_mode=auto` resumes from the other arm's checkpoint. Renaming the run breaks an in-progress resume chain (it cannot find the first-segment checkpoint), so a running chain does not change images. -- **The BI=1 command is byte-for-byte unchanged.** The recipe must offer a mode that only prints the fully resolved training command without starting a job (Nemotron uses `DRY_RUN=1`) and a mode that only runs hydra compose. For each mode × optimizer, print the command from the recipe before and after the change and `diff` (or compare sha256): BI=1 must be identical; BI=0, against the original baseline recipe, may differ only in names, paths and newly added explicit switches. Every compose combination must return rc 0. -- **The header comment states only what was measured.** Reviews have caught untrue claims such as "the BI=0 actor uses vLLM's default kernels". Known rollout config differences between arms (such as R3 forcing prefix caching on) go into the comparison report. -- Delete untested switches (such as unverified mxfp8 and route replay variants) instead of leaving them for users to trip over. - -### BI=0 arm integration pitfalls - -| Problem | Consequence | Handling | -| --- | --- | --- | -| THD packing does not pass `total_tokens` into `PackedSeqParams` | Mamba gets no `seq_idx` and scans the whole pack as one sequence; logprob is silently corrupted | Pass it when CP=1; raise an error for unverified CP>1 | -| Actor needs the BF16 release, rollout needs the quantized checkpoint, both share `model.path` | Actor is built from quantized weights | Separate `hf_weights_path` | -| MTP does not turn off cleanly (the provider appends the pattern back) | Extra untrained layers | Clear both `mtp_num_layers` and the MTP pattern | -| Expert bias is updated during training by default | Routing drifts from inference | `moe_router_bias_update_rate=0` | -| Framework's built-in QAT changes only the GPT layer spec | Hybrid models read the hybrid stack spec, so QAT has no effect | Attach the quantizer after building the model from the TE spec, using the same `requantize` as the exporter | -| QAT library incompatible with the TE grouped linear signature | Startup error or wrong quantization | Quantize weights only; unit test that forward = deployed bytes and the gradient reaches master | -| param + optimizer offload, HF safetensors export | Host OOM on save | Chunked optimizer offload, save only dist ckpt, see [checkpoint and host memory](checkpoint-resume.md#memory-budget-host-and-gpu) | - -R3 checks: R3 in vLLM requires prefix caching and Model Runner V2; the capture is indexed by **absolute layer number** (non-MoE layers are zero columns), and the training side selects the column by global `layer_number`, which must also line up under PP; exclude the MTP router. Known pitfalls: omitting `routed_scaling_factor` causes a drift of about 0.28; a monolithic backend that bypasses top-k capture replays all zeros. - -### Smoke criteria - -- BI=1 quick: `rollout_probs_diff` max = mean = 0 at every step. -- BI=0 quick: mean has the same magnitude as the independent reference (Nemotron QAT about 0.010, BF16 about 0.025), max may approach 1, grad norm is finite. -- When a step's grad norm is exactly 0, check within-group rewards first: if the whole group has the same reward, advantage is 0 and it is not a numerical problem. -- After changing the BI=0 QAT implementation, compare probs_diff against the old implementation on the same quick workload, and print the number of replaced modules (identical on every rank). - -## Evaluation contract and paired acceptance - -One artifact answers one question. Bitwise logprob on fixed tokens proves only the scoring of the same token history; natural generation tasks (GSM8K, AIME) answer quality; fixed-length workloads (Rust Bench) answer throughput. Token sets and timing differ, so none substitutes for another. After the numerical path, deployed weights or backend change, every in-scope quality and speed check needs fresh evidence from the **current candidate**; out-of-scope items are marked not verified, and no acceptance tasks are added on our own. - -### Evaluation contract - -If any of the six items differs, it is not the same score: dataset version (AIME24 ≠ AIME25); prompt/chat template and thinking mode; sampling parameters (temperature/top-p/top-k, stop tokens; greedy n=1, sampled average and pass@8 differ; pass@1 at n=8 is not greedy accuracy); number of repeats; scorer (answer region detection, boxed extraction, tail truncation length, reward extractor); length cap (training and eval caps may differ; 4K vs 8K results are not a pure capability difference). - -- Under greedy, n repeats give the same result and the score is a multiple of 1/number of problems: with 30 problems one problem is 3.33 percentage points, and smaller differences are invisible. -- Rising training reward, probs_diff of 0, and a held-out score gain are three separate conclusions; dynamic sampling changes the sample composition of the training reward. Read the truncation ratio together with accuracy. -- A gain from relaxing the token budget is not a training gain. Charts come from raw metrics; do not fabricate a step 0. -- W&B: the run config behind a curve, the step basis (rollout or optimizer step) and the checkpoint version must be traceable; for cumulative metrics, take per-phase differences. -- Changing the scorer or migrating a reward handler: run both scorers offline on the same responses and inspect only the samples they judge differently; do not switch scorers for the sake of "alignment" and mark another model's valid output wrong over a format difference. - -### Paired acceptance - -**Freeze the criteria first, then run.** Fix the model/checkpoint content, image digest, source/overlay, hardware/nodes, runner/parallelism, weight and KV precision, the actual attention/dense GEMM/MoE/all-to-all backends, and the resolved Graph and cache config; fix what is compared, how it is computed, the pass line and the noise floor; plant a known error to confirm the criterion can return FAIL. Tune BI=0 into a reasonably strong baseline at the agreed precision and have a human audit it; do not weaken it to narrow the gap. When the candidate changes quantization dispatch, attention, MoE, exported weights or image, re-evaluate which old evidence is invalidated ([evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger)). - -**List confounders first, then compare.** Tabulate every difference between the arms; each is either aligned or quantified separately; a delta from several simultaneous changes is not attributed to one switch. Example: the two Nemotron-H arms differed at once in PC on/off, `gpu_memory_utilization`, training backend (Bridge vs mLite) and the R3 return path. - -**Quality pairing**: same problem IDs, template, seed, sampling parameters, max length and scorer, no retries; record per problem whether each side is correct; `both_correct`/`only_bi0`/`only_bi1`/`both_wrong` must sum to the problem count; report each side's score, the absolute percentage-point difference and the agreed threshold. Close scores do not mean identical per-problem results, and passing quality does not mean bitwise-equal logprob. - -**Fixed-length throughput**: the agreed workload is vLLM Rust Bench 8192 input/1024 output, ignore EOS, c=1/8/32/64; record actual input/output tokens, completed requests, error count, warmup, and per concurrency point the wall time, throughput and TTFT/TPOT. List any real rollout workload separately; do not mix fixed-length and natural-EOS results in one ratio. - -### Pairing and computation - -Before testing, a human freezes: number of repeats, pairing order, aggregation formula, the budget/pass rule for each target c, and how an invalid leg is identified and re-tested. The suggested statistic is the median per-round paired loss `median_r(1 - T_BI1[c,r] / T_BI0[c,r])`; if `1 - median_r(T_BI1[c,r]) / median_r(T_BI0[c,r])` is chosen instead, label it "ratio of median throughputs". The two are not equivalent and cannot be swapped after seeing results; without a reviewed statistical basis, report only measured values and do not declare the budget passed. Historical results keep their original formula and are not silently re-judged. - -1. For each leg, record on every worker the image/package hash, actual kernel/backend, resolved config, Graph state, cache hits, and request and token statistics; verify client and server exit normally, separately. -2. Same node, same resources, no interfering load, crossed order (AB/BA/AB), each side warmed up separately every round; do not pair each side's fastest run picked across node pools. -3. Check workload and exit receipts item by item first. A mismatched leg is `inconclusive`; do not compute BI loss from throughputs with different request volumes. -4. Report raw throughput and loss `1 - T_BI1/T_BI0` for every round, then compute the formal result per c with the reviewed formula; keep failed/invalid rounds and the reasons for re-testing. Averaging across c must not hide a required c over the limit; keep single-round and aggregate overruns apart. - -Optimization screening may test only the previously failing concurrency points, but it only decides whether to continue. Take the median over multiple pairs and report single pairs separately: if a candidate round sits right at the budget line, it cannot pass on the old baseline's historical median; re-pair both sides in crossed order on the same node, same image, same time window. When the candidate changes the model output path, the old candidate's quality and post-update alignment evidence do not carry over. Crossed pairing does not guarantee a strong denominator: if the same config and workload show clear fast/slow tiers, keep per-round raw numbers and check tuning receipts, CPU/GPU state, Graph, cache and actual operator selection; an optimization is not effective just because BI=0 landed in the slow tier. Mark unexplained fluctuation's causal attribution `inconclusive`; any formal target point over budget by the predefined rule is still `fail`, with no cherry-picking rounds or switching the aggregation formula afterwards. - -Profile only outside the formal profiler-free E2E measurement; GPU kernel durations may overlap and cannot be summed to stand in for E2E; a kernel name is not a module identity, see [Performance and profiling](performance.md). - -### Failure and recovery - -When the acceptance checker errors, first tell apart producer, numerics, scorer and parser. If the original worker/server receipts and score are sealed and the wrapper parsed them incorrectly after a schema change, keep the original artifacts and read them with a recovery checker that is **separately versioned with its hash recorded**; it must verify all workers, EngineCore, API server, package identity and workload, and keep the original failure status. It can correct only the control-plane conclusion; it cannot fill in missing values, unfinished requests or a failed service exit. - -Keep failure records as they are and append new results; keep reruns and changed questions apart. The final report lists at least, separately: train-inference forward, module backward, real update, quality, performance, clean exit, final immutable artifacts; each cell is `pass`/`fail`/`inconclusive`/`not run`. Example scope: "five steps pass bitwise, GSM8K pairing passes, first-round c=32/64 performance over budget". diff --git a/.agents/skills/train-inference-consistency/guides/rollout-engine.md b/.agents/skills/train-inference-consistency/guides/rollout-engine.md deleted file mode 100644 index e3974e189ae2..000000000000 --- a/.agents/skills/train-inference-consistency/guides/rollout-engine.md +++ /dev/null @@ -1,108 +0,0 @@ -# Rollout engine and sampling - -In RL the rollout engine (vLLM) pauses, sleeps and wakes next to the training process, its sampling randomness must be auditable, prefix cache and speculative decoding affect whether the same token sequence is bitwise identical, and RL-only features such as R3 constrain the engine config. Section 6 is a debugging procedure: save one rollout and replay the old-logprob forward. Measurement bases and A/B methods are in [Performance and profiling](performance.md); BI kernels and cudagraph are in [BI kernels](bi-kernels.md). - -## 1. Sampling seed collision - -**Mechanism.** Each vLLM DP engine resets the global RNG with the same `seed` at the end of warmup. Requests without a seed sample from the global stream, so the n-th request samples the same tokens on every engine. When GRPO spreads the siblings of one prompt across DP engines, many samples in a group repeat: before the fix, Nemotron had only 64-70 distinct responses out of 256 and up to 43% of sibling pairs identical; after the fix, 256/256. The fraction of zero-advantage groups barely changes with the fix; zero advantage mostly comes from task difficulty, not seed collision. - -**Fix.** Reset the global RNG by `(seed, data_parallel_index)`; index 0 is bitwise identical to the old behavior (vllm-project/vllm#59788, merged). Use `data_parallel_index`, not `data_parallel_rank`: dense (non-MoE) models run independent engines under DP, and `reconfigure_for_independent_dp_rank()` resets `data_parallel_rank` to 0 (`vllm/v1/engine/core.py`), so 4 engines get the same seed stream and n=4 gets worse. Nemotron is MoE and goes through `DPEngineCoreProc` with a correct rank, so its 1024/1024 verification on verl holds, but only the index version fixes dense models. - -**How to accept.** In the rollout dump, count distinct samples per group and the fraction of fully identical sibling pairs. With different seeds the baseline is about 0.1% identical sibling pairs, all short answers of <=8 tokens (the model outputs `\boxed{}` directly), and about 1-3% share a prefix of >=16 tokens. Anything above this is RNG correlation. - -**Do not do this.** verl `rollout.full_determinism=True` gives the n responses in a group the same seed. Do not enable it. - -**Greedy cannot verify the fix.** Under greedy, n repeats are identical (Nemotron: 32 of 32), and the score is a multiple of 1/number of problems (AIME with 30 problems: multiples of 1/30; 13/30 -> 19-22/30). A greedy A/B before and after the seed fix shows no difference because greedy does not sample. Rising entropy and the three config curves overlapping are visible only on sampling runs. - -## 2. Prefix cache semantics - -**Premise.** "Enabling PC does not change the numerics" requires the same request to be bitwise identical on hit, miss, and PC off. For pure attention models KV block reuse guarantees this. For stateful layers (Mamba2 conv/ssm state) it depends on which states the engine copies. - -**The Mamba2 pitfall.** vLLM's Mamba align mode copies only the conv and ssm states. Under BI=1, Mamba2 uses exact replay (ported from vllm#55905, which stores three extra chunk buffers x/dt/B) and needs 5 states, so BI=1 Mamba2 originally rejected PC. The fix copies every state in each layer's declared state list; after that, hit/miss/PC off are bitwise identical. Early versions had a convolution window alignment bug in state export (written aligned to the tail, read from the head). - -**Cost.** PC under BI=0 has a measurable cost: in Mamba align mode each decode step adds a batch of eager launches, and an 8K prefill is split into two segments. The approach in upstream #58851 is the best. - -**Confounders between the two arms.** BI=1 Mamba2 rejects PC while the BI=0 R3 AuxOutput forces PC on, so PC on/off is a confounder in the BI=0 vs BI=1 throughput comparison. List confounders first, then compare; see [paired acceptance](recipe-and-acceptance.md#paired-acceptance). - -**Relation to cudagraph.** A PC hit changes the token count of a step. Piecewise cudagraph chooses between the graph and eager fallback by total token count vs `max_cudagraph_capture_size`. Each path is batch-invariant on its own, but the two are not bitwise equal. Under BI, pin `FULL_AND_PIECEWISE`; see [BI kernels](bi-kernels.md). - -**Checks for stateful layers and caches.** Mamba/SSM must also save chunk boundary state and partial input buffers; check aligned/unaligned split, cross-chunk and restore separately. Cover indexing and capacity of long-context KV/page/workspace near the boundaries. - -When testing prefill, confirm whether the prefix cache was hit; during diagnosis, disable or isolate it if needed so the target path executes. Verify the production cache's hit, reuse, and preemption/recompute paths separately. Do not make a comparison pass by clearing the cache and then claim the original lifecycle is fixed. - -For performance fixes on the cache path, use bitwise equality with "cache off" as the gold standard: hit positions outside chunk boundaries, every DP rank, recomputed hits after forced preemption, and chunks truncated by the scheduling budget; confirm in the trace that the target branch executed. Then add a hook that checks the device-side state index against the host prediction at every step, plus a per-problem quality evaluation. Two fixes that each pass alone may silently break when stacked; the combination must pass the gate on its own. - -## 3. MTP / speculative decoding - -- The Nemotron checkpoint has `num_nextn_predict_layers: 1`, and the MTP block is not quantized (bf16). It does not take part in training; the training side must explicitly skip its mapping ([training forward](training-forward.md)). -- Usable with BI=0. BI=1 + MTP requires the logprob from speculative verification to be bitwise identical to the full forward; without verifying this you cannot claim true on-policy. -- If used, use it in both arms; otherwise neither the throughput nor the logprob comparison holds. - -## 4. R3 (routed experts replay) feature constraints - -- **What is returned.** One `(num_moe_layers, top_k)` row of uint8 routed expert ids per token (Nemotron: 52 layers x 6 = 312 B), returned to the trainer with the request's CompletionOutput for route replay. 0 is a valid expert id; check distinctness, not non-zero. -- **PC must be on.** The R3 AuxOutput is keyed by KV block hash; with PC off there is no key. This is why the BI=0 arm in §2 is forced to enable PC; the keys-only form (PR 57968) does not lift it. -- **Why per-step return is expensive.** 312 B exceeds the 256 B msgpack zero-copy threshold, so each request's row per step becomes a separate frame, and the frontend accumulates millions of small objects per request before concatenating them at the end. The overhead is on the engine->frontend return path, not in the capture hook. The fix returns once at request finish (finish-only) and packs each step into one ndarray on worker->engine. `output_processor` already exposed `routed_experts` only at finish, so semantics are unchanged. -- **There is another floor below R3.** msgpack unpacking of three small ndarrays per token for sampling logprobs is a fixed frontend cost (about 25-30 µs/token), larger than everything else left after R3 is fixed. BI=1 does not see it when the engine is the bottleneck; BI=0 has a faster engine and hits it. Two arms tying on e2e is often due to this, not the kernel. Check the frontend asyncio busy fraction before comparing the arms ([Performance and profiling](performance.md)). -- **What does not change.** R3 adds no GPU->CPU sync, async scheduling is kept, capture writes into a statically preallocated buffer, and cudagraph is unaffected. flashinfer_cutedsl and Humming are both modular MoE (they receive the router's top-k), so moving R3 into the kernel saves nothing. -- **Limits of the equivalence check.** BI=1 cannot run the equivalence check (Mamba2 BI rejects PC while AuxOutput forces PC, until the §2 fix lands). BI=0 compares sha with strictly serial single-request greedy, but BI=0 routing is itself not batch-invariant, and an unpatched baseline also repeatedly shows partial inequality against the same reference. So first measure the drift with several baseline-vs-baseline runs; the patch counts as equivalent only if it fully equals one clean reference. Bitwise comparison of concurrent greedy dumps is not evidence. - -## 5. Lifecycle and GPU memory ledger - -**Three pause modes** (`EngineCore.pause_scheduler`): `abort` marks all running requests `FINISHED_ABORTED`, queues new requests, and optionally clears the cache; `wait` queues new requests but keeps stepping until running requests finish (not supported by the inproc engine); `keep` pauses everything, and the Future completes when the output queue is empty. In all three modes the Future returned by pause is the barrier; until it completes, workers still step. In `abort` mode the client receives aborted outputs and must resend them. Requests that arrive during pause are queued and released at resume, using the **new** weights. These are "orphan requests" (pause issued before the rollout tail finishes, or the agent loop sending during pause), and their tokens belong to no single policy version. Before each step's pause, confirm in-flight is 0; after resume, confirm the queue is empty. The HTTP endpoints `/sleep`, `/wake_up`, `/is_sleeping` and pause/resume are dev endpoints and need `VLLM_SERVER_DEV_MODE=1`, otherwise 404. - -**Three sleep levels** (`EngineCore.sleep`, `Worker.sleep`): - -| level | What it does | GPU memory | Host memory | -| --- | --- | --- | --- | -| 0 | Only pauses scheduling | Unchanged | Unchanged | -| 1 | Offloads weights to CPU, drops the KV cache, clears the prefix cache | Frees KV and weights | One weight backup per rank (DS4.1 78.61 GiB); the pinned cache is not returned after wake, see [Checkpoint and resume](checkpoint-resume.md) | -| 2 | Discards all GPU memory, saves only `named_buffers` to CPU | All freed | Small | - -After L2, `wake(weights)` only reallocates GPU memory; the contents are uninitialized. "Reserved is very high but still no GPU memory" after wake is an allocator boundary problem: PyTorch's reserved-but-unallocated blocks cannot be used by another allocation path such as CuMem; see [Checkpoint and resume](checkpoint-resume.md#3-reserved-is-high-so-why-can-wake-still-have-no-gpu-memory). - -**Set gpu_memory_utilization by training mode.** The training process does not fully release GPU memory during rollout. Write the value into the recipe per training mode; check the resolved value inside the container and the KV budget in the vLLM log, not the outer setting. Measured values and rules per mode are in [Checkpoint and resume](checkpoint-resume.md#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated). - -**Metrics are cumulative.** `vllm:*_total` counts since process start; when one engine runs several phases, always use `after − before`. Why the first phase happens to be right and how to cross-check are in [Performance and profiling](performance.md#metrics-are-cumulative-values). - -**DP traffic lands on rank 0.** When the Rust frontend sends one request at a time, all of them land on rank 0 (engine selection is `min_by_key(routing_score)`, which takes the first on equal load). "No events on rank 1" then means no traffic, not broken routing. Use `X-data-parallel-rank` to pin requests to each rank in turn. The two frontends interpret this header differently: Python in hybrid mode validates it as a local index in `[0, dp_local_size)`, Rust uses the global rank; do not send it to a single-rank service. Python hybrid/external LB with a dense model does not start (`DPAsyncMPClient` asserts `stats_update_address`, upstream #32252); test Python hybrid with an MoE model. - ---- - -## 6. Save one rollout and replay the old-logprob forward - -### When to use - -For a difference that shows up once across many RL steps, save the failing input batch and matching policy early. Rerunning the rollout changes tokens, scheduling, lengths, cache, packing and later training data at once, hiding whether an intervention worked. - -1. Save the framework batch and the policy identity after the real rollout ends and before the old-logprob forward; append the trainer scores after the comparison. -2. Verify that the save hook executed, that sequence / valid token counts match the log, and that the data reads back. -3. Use the framework's existing training-only/repeat entry point to read the same batch; restore its weights. -4. Return after old-logprob; forbid the optimizer step and regeneration; record the skipped stages and the weight fingerprints before and after. -5. Run A/A first, then change one layout factor. If A/A itself is unstable, first check randomness, mutable buffers, async streams and runtime state. - -### Minimal matrix - -| Variant | Change | Answers | -| --- | --- | --- | -| A | Original packing | Can the frozen input reproduce it | -| B | Exact repeat of A | Is it run-to-run nondeterministic | -| C | Sequence permutation/packing | Does it depend on layout | -| D | micro-batch size or target position | Does it depend on call shape/neighbors | -| E | Separate process or equivalent node | Does it depend on initialization/device state | - -Run first the variants that distinguish the current hypotheses; the full matrix is not needed every time. After reordering, restore and compare by UID and response token position; the physical row no longer represents the same request. C failing does not prove the packing algorithm wrong; it may trigger a CP, page, tile or expert boundary error. - -### Keep the trigger condition when shrinking the input - -Keep the target sample and its necessary neighbors first, then reduce length and request count. After every reduction, rerun the criterion to confirm the failure remains. A CP problem may depend on the remainder from the end of the previous request to the start of the rank; a Graph problem may depend on short fillers and the capture bucket. Keeping only the failing request often makes the problem disappear. - -Freezing the rollout and the weights does not freeze all runtime state. Training-mode dropout, FP8 amax history, RNG, route replay, caches and workspace must each be kept or reset according to the failure condition. A fresh eval forward is not a substitute for the original old-logprob path. - -A resident model saves checkpoint loading, but state must not leak across variants. First verify A/A between a new process and a reused process; when native libraries, autotune or compile config are involved, restart the corresponding workers. - -### Leaving the replay stage - -Once the original reproducer is closed, return to normal generation, real backward/update and the next rollout step. Training-only success supports only the frozen-policy forward conclusion. If the error appears in backward or in long-running state, save that state and build a minimal reproducer for that stage. - -This rollout replay is a debugging procedure. Nemotron's SSD exact replay is an algorithm in the model implementation that recomputes the partial chunk to match prefill. Do not call them by the same name. diff --git a/.agents/skills/train-inference-consistency/guides/training-forward.md b/.agents/skills/train-inference-consistency/guides/training-forward.md deleted file mode 100644 index 049b474530c6..000000000000 --- a/.agents/skills/train-inference-consistency/guides/training-forward.md +++ /dev/null @@ -1,412 +0,0 @@ -# Training-side model support - -How to build a training implementation cheaply and prove it computes the same math as inference. Where training weights come from (master -> training forward weights) is in [Precision, quantization and QAT](precision-quant.md). - -## Three routes - -The routes differ in where parallelism and kernels come from. In all three we write the model. - -| | MCore (Megatron-Bridge) | FSDP2 (VERL's HF interface path) | mLite | -| --- | --- | --- | --- | -| Positioning | Mature TP/PP/EP/CP, main production RL path | Cheapest: write only the model, no parallelism | Built for consistency: forward calls vLLM kernels, backward reuses the training implementation | -| Forward source | MCore kernels (TE, grouped GEMM) | Our PyTorch model to vLLM semantics, in the Transformers `PreTrainedModel` interface for the VERL loader and PEFT | vLLM kernels (`visible_forward(visible, native, ...)`) | -| Achievable consistency | BI=0, via R3 + importance sampling | BI=0 | BI=1, logprob raw bytes identical | -| What to write for a new model | Bridge weight mapping + modules MCore lacks (Mamba/SSM, special norms) | Whole model: config class, checkpoint loading (with dequantization), per-layer forward, per-operator precision contract; autograd backward, filled in where missing | One `model//`: config, checkpoint mapping, assembly, protocol; modules from `primitive/modules` | -| Typical pitfalls | Mapping checks values, not semantics; MCore and vLLM math differ (MoE activation clamp) | bf16 autocast around the VERL forward rewrites module precision; `--system-site-packages` ABI; PyTorch reference (Mamba scan) too slow; loading a third-party implementation from the image by mistake | Private modules inside the model directory, bypassing the shared skeleton | - -Choose after the [module inventory](#module-inventory-and-placement-rules), by: is BI=1 needed, how many backwards are missing, what parallelism is needed. Nemotron looked cheap by name; FA4, Mamba scan and W4A16 MoE all lacked a backward. - -## Find by symptom - -| Symptom / question | Where to go | -| --- | --- | -| New model loads and config parses, but unclear whether it can train | [Module inventory](#module-inventory-and-placement-rules) | -| Want to reuse vLLM forward code, but tangled in paged KV / cudagraph / scheduling | [Call kernels only](#reproducing-the-forward-call-kernels-only-do-not-port-the-scheduler) | -| Weight mapping shapes are right, output silently degrades | [Mapping checks values and semantics](#43-mapping-checks-cover-values-and-semantics-not-only-shape) | -| Inference kernel has no backward; should we write one | [Three-case decision](backward-validation.md#three-cases) | -| Wrote a VJP; how to prove it is right | [Module-level comparison steps](backward-validation.md#comparison-steps) | -| Unit tests pass but real training fails / step 0 occasionally not bitwise | [Two test-infrastructure false negatives](backward-validation.md#two-test-infrastructure-false-negatives); [BI kernels](bi-kernels.md) | -| Changes after modifying packing or CP/TP/EP; wrong only on long sequences | [Layout and length](#distributed-layout-length-boundaries-and-training-side-acceptance) | -| `ppo_kl` non-zero after rollout switches parallel layout (e.g. DP+EP to TP1 per GPU) | [Bitwise identity after changing the parallel layout](#bitwise-identity-after-changing-the-parallel-layout) | -| `modified by an inplace operation` after no longer recomputing old-logprob | [Entropy when old-logprob is not recomputed](#entropy-when-old-logprob-is-not-recomputed-bypass) | -| How to hook into VERL / Megatron | [Framework integration](#verl--megatron-integration) | -| Which kernel an operator uses in BI mode and its cost | [BI kernels](bi-kernels.md) | - -## Key rules - -1. Semantics follow the vLLM implementation and the checkpoint. -2. Before starting, list for every module whether its inference kernel has a backward. -3. Call kernels only; do not port the scheduler. mLite (BI=1): FP32 logprob raw bytes vs vLLM-captured trajectories, target 0 diff. FSDP2 / MCore (BI=0): no 0-diff target; judge each side against an independent reference, then look at precision. -4. Mapping checks cover values and semantics, not only shape; fused parameters are overwritten whole. -5. The precision contract lives in the model. Where vLLM computes in FP32 (position-encoding angles, router logits and bias, norm statistics, etc.), the training module sets the dtype and disables the outer autocast. Under VERL's bf16 autocast, casting only the inputs of matmul-type operators (including `F.linear` and `einsum` decomposed into bmm) does not stop the downgrade. Measure output dtype and values under the real autocast. -6. Backward follows the formula, not the forward kernel. Only rounding differs -> reuse. Quantization rounding -> fix the STE contract, reuse BF16 GEMM. No whole-segment backward -> write glue, GEMM still TE. "No backward kernel at the same precision" is not a reason to rewrite. -7. Reuse model-agnostic operators: mLite `primitive/modules` and `ops` first, then MCore/TE/mamba_ssm; only then is it new code. -8. Backward need not be bitwise; report bitwise and tolerance separately; the reference builds its graph independently with ATen autograd; tolerances are fixed before seeing results; synthetic inputs are labeled as not real activations. -9. Do not recompute activations; use the forward's saved intermediates (DS4: gate_up / LSE). -10. The model directory holds only config / checkpoint / model / protocol; reusable parts go into `primitive/`; no same-named `primitive/` under the model directory; impl name `lite`. -11. Long-sequence boundary points (KV page, packing prepend, the 7000/9216 tiers) need dedicated tests. -12. MTP block, EPLB redundant experts and draft are explicitly skipped in mapping and sync. - -## Module inventory and placement rules - -List per module: inference kernel, same-math training implementation, quantization rounding, discrete selection, backward. Then pick the route and each module's place in the code tree. This catches missing backwards before they cost weeks. - -Training support means an implementation we maintain that backpropagates correctly. Config registration, same-family fallback, the vLLM inference implementation and third-party implementations in the image (newer Transformers modeling, NeMo AutoModel) do not count. Write it per [Reproducing the forward](#reproducing-the-forward-call-kernels-only-do-not-port-the-scheduler) into real source files and verify by loading through the actual loader in a new process. A REPL-only implementation or runtime monkeypatch is not evidence. - -### 1. How to build the inventory - -One row per module, every column filled; write unknown, not "should exist". - -| Column | Question to answer | -| --- | --- | -| Inference kernel | The function vLLM actually calls (not the model-file class name), with quantization variants, attention backend, cudagraph branch | -| Same-math implementation in the training framework | mLite `primitive/modules`, `ops`; MCore / TE; mamba_ssm, causal_conv1d. Rounding-only differences count as same math; a different activation clamp does not. Third-party same-model implementations are not entered | -| Quantization rounding | FP8/NVFP4 rounding in the forward; weights and activations separately | -| Discrete selection | top-k expert IDs, sparse indices; IDs fixed and not differentiated, the continuous score needs a derivative | -| backward | Does the library have one; do its intermediates (LSE, rstd, token permutation) match what our forward saves | -| Conclusion | Reuse / glue / new code ([decision rules](backward-validation.md#three-cases)) | - -Nemotron-H: - -| Module | Inference kernel | Same-math training implementation | Rounding | Discrete | backward | Conclusion | -| --- | --- | --- | --- | --- | --- | --- | -| Mamba SSD scan | `mamba_chunk_scan_combined_varlen` | mamba_ssm | None | None | mamba_ssm `_mamba_chunk_scan_combined_bwd` | Reuse | -| conv1d | `causal_conv1d_fn` | causal_conv1d | None | None | Yes; dW needs fixed-order FP32 | Reuse + small glue | -| Attention (FP8 KV) | FA4 fixed split + FP8 Q/K/V | FA4 BF16 | Yes | None | FA4 BF16 backward only | Dequantize, recompute BF16 O/LSE, reuse | -| Mamba in/out_proj (FP8), shared expert (NVFP4) | FlashInfer / NVFP4 GEMM | TE `general_gemm` | Yes | None | TE | Fix STE, reuse | -| Routed experts (NVFP4 W4A16) | CuTe-DSL / Humming MoE | TE grouped GEMM | Yes | Yes | No whole-segment backward | Glue `routed_vjp` | -| RMSNorm / gated RMSNorm, routing score | BI kernel / grouped top-k | PyTorch formula | None | None for the score | autograd | Reuse | -| MTP block, draft | Not part of training | - | - | - | - | See §4 | - -### 3. Placement rules - -#### 3.1 mLite (Megatron `experimental/lite/megatron/lite`) - -Read `experimental/lite/docs/porting.md` and `architecture.md` first. Three layers: `runtime` (lifecycle, training-step orchestration), `model` (registration + model-specific build/load/export protocol), `primitive` (reusable: parallel state, TP layers, checkpoint conversion, MoE utilities, optimizer wrappers). The runtime knows no model details: it gets the protocol from the registry, builds a typed config, and lets the protocol build the model. - -| Rule | Content | -| --- | --- | -| The model directory holds only four things | `model//`: config, checkpoint mapping (WeightSpec), assembly, protocol. Impl name `lite` (`/lite/`) | -| Modules go into primitive | norm / linear / router / experts / dispatcher / attention -> `primitive/modules`; generic VJP bridge -> `primitive/ops`; quantization encoding -> `primitive/quantization`; PP splitting, thd, CP -> `primitive/parallel`; HF read/write via `primitive/ckpt/hf_weights` (`SafeTensorReader`, `load_hf_weights`, `export_hf_weights`) | -| Model-specific math still goes into primitive and is registered | e.g. Mamba2 bridge, FP8-KV attention -> `primitive/modules/.py`; no same-named `primitive/` under the model directory | -| No file for reused backward | Reusing TE / mamba_ssm backward means no `backward.py`; glue goes in the module file | -| Dependencies are one-directional | `model/` imports `primitive/`, never the reverse; top-level checkpoint/quantization does not import the impl subtree | -| Layering guards must be registered | Add the model to `tests/unit/runtime/test_layering_contracts.py` or the guard does not apply | -| No hybrid architecture packages | porting.md: "Keep these out of the first PR: Hybrid model implementation packages"; hybrid (Mamba + attention + MoE) modules are split into primitive | -| State third-party kernel dependencies | Document hard dependencies on vllm, flashinfer private APIs, mamba_ssm, causal_conv1d, and the `init_batch_invariance()` call in build_model | - -DS4 is the positive example: `model.py` (650 lines) assembles te.RMSNorm, CSA, HCA/mHC, thd and the fp8 recipe with three model-specific deviations; `moe.py` (102 lines) = SigmoidTopKRouter + Experts + TokenDispatcher + SwiGLUMLP, hash-route the only model-specific code; `checkpoint.py` is one WeightSpec using shared load/export. - -#### 3.2 MCore / Megatron-Bridge - -- Bridge `models//`: provider + HF<->MCore mapping; modules MCore lacks (Mamba/SSM variants, special norms, HC) go into `megatron/core`, not Bridge. -- Mapping checks values, not only shape ([§4.3](#43-mapping-checks-cover-values-and-semantics-not-only-shape)). -- MCore and vLLM may use different formulas (MoE activation clamp); BI=0 cannot align those either, so list them first. -- The Megatron fork works only with Megatron-Bridge v0.6.2, pinned in the image. - -#### 3.3 FSDP2 (VERL's HF interface path) - -- Write a PyTorch model to vLLM semantics in the Transformers `PreTrainedModel` / config interface so the VERL loader, PEFT and HF save/load use it directly. "HF" is only the interface and weight format. -- Delivery: the model goes into the verl trunk, registered to the Auto classes (`AutoConfig` / `AutoModelForCausalLM`, etc.), and ships with the recipe as one verl PR ([Code and delivery §8](../practices/code-and-delivery.md#verl-a-thin-recipe-not-a-second-training-framework)). The verl FSDP extension guide's advice against copying modeling into `verl/models` targets models already in transformers. Register at the current checkout's FSDP load chain (`get_hf_auto_model_class` -> `from_pretrained`). -- Develop and deliver from the same verl source and registration entry. Do not modify the checkpoint's `config.json` / `auto_map` (the low-precision checkpoint is the single source of truth) or enable `trust_remote_code`. The load log must show our registered class was used. -- VERL's automodel engine (`automodel_actor`) uses NeMo AutoModel and is not this route; routes are distinguished by forward source, not by FSDP2 sharding. -- Under bf16 autocast, FP32 computations (position-encoding angles, router logits and bias, norm statistics, following vLLM) disable autocast inside the module and set the dtype. -- Backward prefers autograd; add operators next to the model file only when a backward is missing or too slow (the 160-line PyTorch Mamba scan reference). -- Common pitfalls: `--system-site-packages` ABI, the actual `use_remove_padding`, PEFT targets not going through nn.Linear. -- Under EP4 + FSDP2, all ranks' DCP expert shards share key and offset, so restore reads another rank's experts; the key must carry the EP rank ([checkpoint guide](checkpoint-resume.md)). - -### 4. Parts that do not take part in training - -Each gets its own inventory row: - -| Part | Mapping | Facts | -| --- | --- | --- | -| MTP block | Not loaded for training; WeightSpec lists it as explicit skip, not missing-key tolerance; `_ignore_unexpected_suffixes` etc. only for listed suffixes | Nemotron `num_nextn_predict_layers: 1`, MTP unquantized bf16; usable with BI=0, then in both arms; BI=1 + MTP requires speculative-verification logprob bitwise equal to the full-segment forward | -| EPLB redundant experts | Training has logical experts only; physical-to-logical mapping is inference-side | Completeness counts logical experts, not physical tensors | -| draft / speculative model | Independent checkpoint and mapping | Shared weights (embedding, lm_head) state their owner | - -Skipped parts appear in the load log and manifest as "explicitly skipped", not "no error reported". - -### 5. Example: reuse audit of Nemotron-H against mLite - -Nemotron-H is 3824 lines. From `primitive/` it reuses only protocol-layer orchestration utilities, `load_hf_config_dict` and `vocab_parallel_entropy`; every modeling module is rewritten under `model/nemotron_h/vllm/primitive/`. - -| File (`N = model/nemotron_h`) | Lines | What it does | Verdict | -| --- | --- | --- | --- | -| `N/config.py` | 155 | HF config dataclass, validation | Keep (model dir) | -| `N/quantization.py` | 207 | NVFP4/FP8 requantize, E2M1 grid, FP8 encoding, hand-read index.json | Numerics -> `primitive/quantization/nvfp4.py` (parameterize E2M1 RNE vs mxfp4 ties-down); IO via `SafeTensorReader` | -| `N/checkpoint.py` | 486 | HF name mapping, Export spec, BF16 master invertibility check, θ0 threshold, own `load_hf_weights` | Ordinary tensors via `load_hf_weights` + WeightSpec; keep master check and θ0 threshold; importing the impl subtree is a layering inversion | -| `N/vllm/model.py` | 188 | Layer/Model assembly, PP stage | Keep; PP split via `build_pipeline_chunk_layout` | -| `N/vllm/protocol.py` | 409 | protocol, contract validation, BI check, FSDP2/dist_opt, post-step requant | Keep; hard-coded 52-layer pattern and (2688,128,1856,6) make it single-checkpoint | -| `N/vllm/primitive/dense.py` | 647 | visible forward + BF16-master VJP bridge; RMSNorm VJP; Humming/CuTe NVFP4 linear; projection factory; FP8/NVFP4 TrainingLinear | Bridge and RMSNorm VJP -> `primitive/ops/aligned_vjp.py`; `native_linear_vjp` duplicates `_te_general_gemm`; serving kernel adaptation stays | -| `N/vllm/primitive/logprob.py` | 36 | Chunked lm_head + vLLM `compute_token_logprobs` + FP32 log_softmax STE | Keep; the docstring's DS4 function does not exist | -| `N/vllm/primitive/attention/*` | 231 | `Fa4Fp8KVAttention` (FlashAttnFixedSplit 6768-token page, FP8 Q/KV), `_Fp8AttentionVJP` | Model-specific; reusable once GQAttention supports `rotary=None, qk_norm=False` | -| `N/vllm/primitive/mamba/*` | 316 | `_PackedConvVJP`, `_PackedScanVJP`, `SSMMeta`, `MambaMixer` | Keep both VJPs; generic Mamba2 -> `primitive/modules/mamba2.py`; derive `SSMMeta` from `PackedSeqParams`/`ThdPackMeta` | -| `N/vllm/primitive/moe/module.py` | 185 | `_FixedRouteVJP` + Router, `_CombineVJP`, ReLU² SharedExperts | Router VJP equals `SigmoidTopKRouter`, reusable, not hooked to router_replay; add an activation argument to `mlp.py` for SharedExperts | -| `N/vllm/primitive/moe/grouped.py` | 705 | FlashInfer W4A16 signature/sha pinning, `CuteDslRoutedExperts`, `Nvfp4ExpertWeights`, `_te_grouped_gemm` rewrite, ReLU² `routed_vjp` | Keep W4A16 forward + grouped VJP; TE grouped GEMM wrapper duplicated; `ep4_routed_experts` hard-codes 4 ranks | -| `N/vllm/primitive/moe/communication.py` | 251 | DeepEP normal-mode EP, `reduce_ep4_parts` (FlashInfer one-sided combine order), `EPRoutedExpertsVJP` | Use `_build_deepep_buffer`; keep serving-order reduction, parameterize EP size/topk | - -Six placement violations: impl name `vllm` not `lite`; hybrid architecture package; same-named `primitive/` in the model layer; circular imports between `checkpoint.py`/`quantization.py` and the impl subtree; `nemotron_h` missing from the layering guard; third-party kernel dependencies and `init_batch_invariance()` undocumented. - -Should have been reused: safetensors index reading (5 copies); HF load of ordinary tensors; DeepEP buffer; TE GEMM wrapper (2 copies); visible forward + reference VJP bridge; Router (`SigmoidTopKRouter` + `group_limited_topk`, forward switched to vLLM grouped_topk, hooked to router_replay); shared expert; E2M1 grid and FP8 encoding; PP splitting; `SSMMeta`; two docstrings citing nonexistent DS4 functions. - -Model-specific, keep: the two Mamba2 VJPs + packed conv/scan; `Fa4Fp8KVAttention` + FP8 attention VJP; W4A16 MoE forward + ReLU² grouped VJP; serving-order reduction (parameterized); Humming/CuTe NVFP4 and FP8 serving linear; vLLM forward-value logprob; BF16-master invertibility check and θ0 threshold; ModelOpt 4over6 requantize rule. - -## Reproducing the forward: call kernels only, do not port the scheduler - -The first deliverable after the inventory is a training forward with vLLM's math. Logprob measurement basis and R3 route replay: [workflow](../workflow/README.md), [rollout engine](rollout-engine.md). - -### 1. First separate the responsibilities of the model implementation and the training framework - -A training framework (FSDP2, MCore, mLite runtime) shards a trainable model; it does not make an inference-only kernel differentiable or write the model. Write a minimal PyTorch model for the current loader interface covering the config, checkpoint loading, forward, gradients and save interface actually used; full Transformers integration, generation and all parallelism are not needed first. The semantic basis is vLLM and the checkpoint, not a third-party implementation in the image. - -| Capability | What the training route still has to do | What in inference / Megatron is not copied | -| --- | --- | --- | -| Model math | Real semantics of attention/SSM, norm, position encoding, routing, residuals, special embeddings | vLLM paged KV, decode scheduling, cudagraph management; MCore provider/ModuleSpec wiring | -| Backpropagation | Input gradients of every training-path operator; trained-parameter gradients; recompute correctness ([Backward validation](backward-validation.md)) | TP/SP/CP-specific collectives and gradient reduction | -| Weights | Names, axes, fused order, expert IDs, sharing, derived state | HF<->Megatron shard layout, PP export protocol | -| LoRA | Target coverage, A/B direction and scaling, non-zero adapter taking effect | Parallel linear wrapper, TP adapter aggregation | -| Memory | Initialization peak, large frozen table residency, save/export peak | Megatron memory saver, PP scheduling | - -dtype, dequantization, QAT and update format: [precision guide](precision-quant.md). Here each module states its input/output precision and differentiable boundary. - -### 2. How to wire kernel-only calls - -Each layer's forward runs once: call the vLLM kernel for the visible output, bitwise identical to inference, and save what the backward needs; see [backward wiring](backward-validation.md#how-to-wire-it-forward-once-backward-uses-what-the-forward-saved). Borrow only the kernel, no vLLM runtime object: - -- No KV paging, decode scheduling or CUDA Graph management. Training runs a full-sequence teacher-forced forward; packing passes boundaries via thd/cu_seqlens. -- Reuse order: model-agnostic differentiable operators -> concise PyTorch reference to vLLM semantics -> ready-made operators with backward -> custom autograd/kernel only if necessary. This order only says where to look first; the training-step cost of each choice is checked at C6 against the C1 budget, not deferred to C10. -- Handle discrete top-k / index selection apart from the selected differentiable tensors. What is trained, frozen or given an auxiliary loss follows the model definition; do not detach to get it running or add router/indexer objectives. -- vLLM layer modules used directly (FusedMoE, ModelOpt linear) need `get_tp/ep_group`, so the training process creates vLLM process groups (Nemotron does; DS4's own modules do not). -- Take the same branch as inference: reproduce the quantization config, attention backend and BI-switch kernel selection, or the comparison is against a different path. Tests use the production `CompilationConfig(custom_ops=...)` ([the two false negatives](backward-validation.md#two-test-infrastructure-false-negatives)). - -"The same kernel" still has pitfalls (Nemotron 7-layer localization): - -1. Direct and compiled paths of the same gated RMS kernel differ; calling it through an opaque custom-op boundary removed the difference. -2. A differing reduction block was wrongly blamed alone; the same captured kernel showed both blocks match. -3. The last 4.768e-7 logprob difference came from the HF forward calling `.float()` after lm_head; keeping the real BF16 head output and sharing the scoring fixed it. This is not truncating an FP32 projection to BF16 to manufacture equality. - -### 3. Read the inference implementation along the full support chain - -Order: model file (layer types and order, residual and norm positions, position encoding, attention visibility, router/indexer, special embeddings) -> each layer's actual kernel and quantization variant -> attention backend and KV dtype -> weight loader (name mapping, fused order, derived scale) -> BI-switch branches. Confirm the parameter/buffer list from the same checkpoint/config. - -Sizes sharing one implementation need no separate port. For ordinary dense models, look for a reusable generic skeleton first; we still write the model. Multimodal is validated per task; diffusion is out of scope. - -### 4. The minimal implementation that really has to be filled in - -#### 4.1 Reproduce the math - -Confirm from the checkpoint/config and inference code: layer types and order, parameters/buffers, residual and norm positions, position encoding, attention visibility, router/indexer, special embeddings. Model-specific math (`sqrt(softplus)` scoring, SwiGLU with clamp, mHC Sinkhorn, ReLU² experts) is in no library; write it from the inference source, not the paper. - -#### 4.2 A frozen base still has to propagate input gradients - -```python -W.requires_grad_(False) -y = linear(x, W) + scale * B(A(x)) # no gradient for W, but still propagate to x -loss(y).backward() -``` - -Wrapping a frozen block in `no_grad()` cuts LoRA gradients of earlier layers; only a frozen table lookup needing no input/parameter derivative is exempt. With initial B=0 a zero gradient for A can be normal; test with a non-zero adapter or a real step. A finite loss does not prove backward. - - - -#### 4.3 Mapping checks cover values and semantics, not only shape - -Per weight class: source name -> training name -> serving name, axis transforms, sharing. Focus on fused QKV, gate/up order, expert numbering, extra norm/bias, tie weights, non-persistent buffers, derived state (FP8 KV scale, NVFP4 alpha). Every missing/extra item needs a reason; quantization-derived items follow the precision guide and need not be byte-identical to the low-precision source. FSDP needs no HF<->Megatron bridge but may need checkpoint->training and export mappings. - -#### 4.4 Packing boundaries for Mamba / stateful modules - -Every stateful operator (conv1d, SSM scan, recurrence, PLE shift) must consume sequence boundaries; a correct attention mask proves nothing about them. Incidents: - -- verl via Megatron-Bridge did not pass boundaries to Mamba; sequence states mixed and logprob was silently wrong, visible only as a larger mismatch. -- Qwen3.8 Flash-Next's PLE shift made the second sample reference the first sample's last token. - -Check: unequal-length A/B run separately and merged; sentinel-initialized outputs to catch missed writes. In mLite, derive boundary metadata from `PackedSeqParams`/`ThdPackMeta`, not a separate `SSMMeta`. See also [Long-sequence boundary points](#long-sequence-boundary-points). - -#### 4.5 Recompute and small-model trimming - -With activation checkpointing on, verify recompute reads the same state and gives reasonable gradients. Trim to the smallest layer set covering every non-equivalent path and its dependencies, usually 4-6 layers; keep mixed attention, cross-layer sharing and layer-specific table lookups. Record kept layer indices and the config transform; a slice proves only covered paths, not full-model mismatch or score gain. - -### 5. How to prove it computes the same math - -Feed the training forward vLLM-captured trajectories (fixed tokens, weight version, position, mask, routing); never compare probability summaries or argmax. mLite (BI=1): FP32 logprob raw bytes, target 0 diff. FSDP2 / MCore (BI=0): judge each side against an independent reference, then check the remaining difference comes only from precision rules. - -- Same policy weights. Differing dtypes are a contract mismatch; do not convert to fake bitwise. -- Reference artifact (Nemotron): frozen trainer image, full actor, 2048 tokens, FP32 0 diff; five real updates, 883,998 tokens of identical FP32 logprob. -- Old-logprob recomputation and the grad-enabled forward are two paths; evidence for each ([dual-path contract](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract)). -- Find the first divergence per [debugging](debugging-mismatch.md): per-module dump, single-request prefill, then the operator. - -### 6. Small tasks handed to agents and stop conditions - -Each outputs a short conclusion, source version, reproducible command and artifact. - -| Task | Minimal acceptance | If it fails | -| --- | --- | --- | -| Implementation inventory | Training model class, operator backward, checkpoint loader; gaps listed | Write the minimal model; stop cold-starting the large model | -| Module comparison | Same weights/inputs; gradients vs a differentiable reference | Locate the earliest difference | -| Few-layer training | Non-zero loss, finite gradients, intended adapter changes | Separate broken path, operator bug, gradient-less target | -| Save and restore | Fresh-process restore: consistent forward and next step | Check state, adapter, optimizer, residency | - -Distribute only after the first two pass; scale up after the few-layer two-step loop passes. - -### 7. How to judge development cost - -Training implementation exists: wire loading, LoRA, sharding. Inference only with differentiable references: fill in assembly, mapping, interfaces. Key operator lacks backward or clear semantics: separate blocker; reference and gradient validation first; no uniform person-day figure. - -## Distributed layout, length boundaries and training-side acceptance - -Record trainer and rollout topologies separately. TP/CP/EP, padding, packing and collectives change data mapping and reduction paths; long sequences hit untested boundaries; forward 0 diff does not vouch for gradients or the optimizer. - -### Compute resources first, then request them - -Write trainer and rollout process groups separately. At four GPUs per node, an inference DP8/EP8 group spans two nodes; two DP4/EP4 replicas are not one EP8 group and skip cross-node EP communication. The runtime decides grouping. - -If dense DP excludes CP: `W = TP × PP × CP × DP`. Expert side: verify `W = ETP × EP × PP × EDP` and the CP mapping separately. **Do not multiply EP into the dense world again.** W=8, TP1/PP2/CP2 gives DP=2; a log calling DP×CP "DP" shows 4. Clarify naming and members; do not change GPU count to match a config string. - -GPU memory follows the layout. With PP1 each rank holds all layers' activations, linear in sequence length. Per-route EP backward state stored for other ranks' slots grows by the average ranks per token; store only own routes. Rollout `gpu_memory_utilization` is set from what training still holds during rollout, per training mode, in the recipe ([Memory budget 02 §5](checkpoint-resume.md#5-gpu-memory-set-gpu_memory_utilization-per-mode-when-training-and-rollout-are-colocated)). - -### Packing and CP - -For the target UID reconstruct global request `[seq_start, seq_end)`, rank range `[rank_start, rank_end)`, local rows, cu_seqlens, positions, all compression/reordering mappings. - -The intersection `[max(seq_start, rank_start), min(seq_end, rank_end))` contributes rows only when start < end. The overlap window must not give a disjoint request this rank's output. Host cumulative boundaries equal the rows the device compactor produces. - -Cover: before/after/at the boundary, empty intersection, cross-rank, compression-ratio and tile remainders. Window constant 8: check 1…7, not only multiples. Start from a hand-computable counterexample, then enumerate nearby legal configs. - -Passing `cu_seqlens` does not mean a fallback consumes it ([§4.4](#44-packing-boundaries-for-mamba--stateful-modules)). - -### EP and collectives - -Record expert mapping, tokens per rank, padding capacity, route order, shared/routed combine, quantization scales. Empty ranks may still join collectives; partially written reused buffers leave stale data. Use different sentinels per side; zero-filling both gives a false green. - -Pinning NCCL algorithm/channels may remove reduction differences on one stack and cut throughput significantly; verify it per topology, not as a universal env list, before relaxing config or building a deterministic collective. - -MCore dev's grad-ready assertion takes first-batch hook counts as golden and errors on a different count. Short-sequence micro-batches can make counts data-dependent (Nemotron r10 at step 2; r12 did not reproduce). With dp×cp=1 and expert_dp=1, `overlap_grad_reduce` off avoids it without changing values; root cause unknown. With EP4 + dist_opt and dp>1, multi-micro-batch reduction fires early (reported). - -### Bitwise identity after changing the parallel layout - -**Principle.** Floating-point addition is not associative. Layout sets reduction grouping, order and rounding: TP splits a matmul into per-rank partial sums; EP computes a token's top-k experts on different ranks and merges. Last-bit differences become different tokens after dozens of layers and sampling, so bitwise identity holds only for the verified layout; any TP/EP/DP change must be proven again. - -**Rules.** - -1. The reduction tree is part of the numerical contract: per level, accumulation dtype, order and rounding points. All legal rollout layouts and training walk the same tree. -2. A degenerate layout (TP=1, EP=1) must simulate the tree: per-rank partial sums, same rounding, same merge order. -3. After a layout change, compare per-token logprob raw bytes on fixed inputs before training. Non-zero `ppo_kl`/probs_diff is fixed, not avoided by reverting the layout. -4. Locate: dump per layer, find the first differing layer, confirm its inputs (MoE: router logits too) are bitwise equal; the bug is that layer's reduction ([Fault localization](debugging-mismatch.md#bisect-and-first-divergence)). -5. Measure the cost of simulating the order and report it apart from the layout gain. - -**Examples.** - -| Model | Layout change | Where the difference is | Fix and result | -| --- | --- | --- | --- | -| Nemotron-H (W4A16 MoE, BI=1) | rollout from DP4EP4 per node to TP1 per GPU | EP4: per-rank FP32 sums rounded to BF16, merged in FP32 as `((s0+s1)+(s2+s3))+(s4+s5)`. EP1: six outputs summed sequentially in FP32, no intermediate rounding. Max logprob diff 0.4; greedy tokens diverge at positions 10–129 | Without EP, BI mode computes all experts in one launch and reduces in EP4 order (vLLM `661cde6833`, not in the image). Both layouts match training bitwise; `ppo_kl` 0 over 58 steps; multi-turn replay 404 s vs 485 s, reduction cost about 2.5% | -| DeepSeek V4.1 | training TP=1, inference TP>1 | order of per-TP-rank partial sums | both sides use one fixed reduction tree, TP=1 included (`VLLM_TP_INVARIANT_LEAVES`) | - -### Asynchronous communication and lifetime - -Check streams, events, buffer owners and reuse timing. Passing with a global synchronize only proves a timing issue and is not the fix; find the missing dependency and accept on the same path. - -### Long-sequence boundary points - -Test at real lengths: - -| Boundary | Fact | How to test | -| --- | --- | --- | -| KV page | training attention replicates vLLM's KV page size; page boundaries are numerical boundaries | sparse dY across the page boundary | -| Fixed split attention | fixed split count and keys per split; the resulting `MAX_SEQ_LEN` is an invariance bound, not a performance knob | above: error, not silent truncation; below: two real lengths × random/zero/sparse dY | -| Packing prepend | one prepended short sequence shifts every target row offset | two prepend controls vs none, bitwise | -| Micro-batch fits the longest sequence | verl needs each micro-batch to fit the longest sequence (`ppo_max_token_len_per_gpu`, Nemotron 16384 vs DS4 4096; Nemotron training has no CP) | OOM probe at formal length; do not extrapolate from 8K | -| Embedding backward | CUDA `embedding_dense_backward` accumulates at reduced precision with ≤3072 indices, e.g. a short sequence alone in a micro-batch | length-tiered tolerance | -| CP compressed boundary | see [Packing and CP](#packing-and-cp) | before/after/at the boundary | - - - -### The D term in attention backward - -With P rounded to FP8, visible O = fp8(P)·V ≠ softmax(QKᵀ)·V. If `D = rowsum(dO⊙O)` uses visible O while dS uses recomputed P, `rowsum(dS_i) ≠ 0` and dQ gains a same-sign per-row bias (Nemotron: about 3% of dQ/dK, v2 FAIL at steps 150/250; fixed 0.006–0.009, PASS). Mainstream FP8 attention (TE fp8_dpa, SageBwd) uses visible O, so a human decides after being informed. - -Four checks: - -1. **FA4 interface signature**: `_flash_attn_bwd` derives D in preprocess from the `out` passed; D cannot be passed. The public interface errors under `requires_grad`; the actor calls `_flash_attn_fwd/_bwd`. Fix: recompute BF16 O and LSE from dequantized Q/K/V, pass to bwd (layer backward +4–15%, end to end <0.05%). -2. **varlen LSE layout**: ln vs log2, per head vs per token, softmax scale folded or not. On mismatch, recompute with the training kernel. -3. **GQA 32/2**: dK/dV summed within each Q head group; A/A needs fixed summation order. -4. **When the saved Q/K/V are released under recompute=full**: if only dequantized Q/K/V are saved, recompute must read the same state; early release changes the values. - -### Microbatch split and gradient difference - -Different microbatch splits give about 2e-3 relative gradient difference (loss normalization, accumulation order, reduced-precision embedding backward at ≤3072 indices). Not a bug, but: - -- Criterion denominator = `max(B16 reference floor, noise floor of the implementation under test)`, never the FP32 ideal. -- Both arms use the same split; if it changes, rerun A/A. -- grad_norm is not comparable across different microbatch counts. - -### Control matrix - -With policy/tokens fixed, test original layout, reversed order, target first/last, changed micro-batch, CP-boundary sample, each reconstructed by UID/global token coordinates. Real TP/EP changes may legally change the reduction; the experiment contract decides whether cross-topology bitwise equality is required. - -Safe packing is not a fix for a metadata bug; it may only lower trigger probability. After the fix, the original layout and nearby boundaries must pass. - - - -### Training-side correctness acceptance: gradient comparison criterion v2 and optauth - -Two independent checks beyond forward 0 diff. - -#### Gradient comparison criterion v2 - -- Freeze the criterion (tolerance, denominator, thresholds) before the run, record its sha, never change it after. -- Planted-error check first: gradients ×1.25 and ×1.05 must both FAIL. -- Accept on several snapshots (steps 0/50/150/250); the D-term problem appeared only at step 150. -- On failure, inject reference gradients block by block to find the first module over the limit. -- Extended to FSDP2: EP4 + FSDP2 all four snapshots PASS, worst-class median 1.08–1.39, all floors from B16. -- Tool bug: with `overlap_grad_reduce=False` gradients were reduced again before the EP4 dump (fixed). The export path needs A/A too: ‖g‖ bitwise equal, local shards bitwise equal to the dump, replica diff 0/0. -- Cross-process bitwise backward is not required. An occasional non-bitwise first step comes from the thread-local Inductor deterministic config (the autograd thread's first compile picks 512/1024 blocking by timing); fix in [BI kernels: determinism](bi-kernels.md#determinism-takes-effect-only-on-the-main-thread). - -#### Optimizer acceptance (optauth) - -Verify the optimizer consumes clipped `main_grad`, the update lands on master, and deployed bytes = requantize(master). Freeze the criterion first. Split failures into criterion modeling gaps (e.g. a check inapplicable without replicated parameters; the equal-magnitude first Adam step invalidating γ-type criteria while the per-coordinate bound holds), documented and accepted, and implementation errors, fixed. Cover all parameter group types: a first-step probe without an FP32 wd_mult=0 parameter misses weight decay applied to that group. - -#### Facts outside the two criteria - -- Frozen base with LoRA: base modules still propagate input gradients ([§4.2](#42-a-frozen-base-still-has-to-propagate-input-gradients)). -- The two-step RL loop (θ0→θ1→θ2) accepts the forward; three-arm SFT accepts training effect; report all three separately ([validation](../workflow/validation.md)). - -## VERL / Megatron integration - -Upstream config keys may differ from what was used here. Find the actual call paths of rollout dump, old-logprob and skip/repeat in the current checkout and record the commit; local-patch options are not upstream API. Process content (recipe and PR contract, review material, two-step acceptance, E2E targets) is in [validation.md](../workflow/validation.md#recipe-review-materials-and-two-step-acceptance-agreement); PR ownership in [Code and delivery §8](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image). - -### Where the integration layer is - -Megatron's `experimental/lite/examples/verl/verl_mlite/` provides the verl engine adapter (`model_engine=mlite`, `pkg://verl_mlite.config`), installable into the unified image without reimplementing the engine in verl. Check engine registration, model factory, batch conversion, loss normalization and optimizer path; list reusable interfaces and gaps. Consistency uses mLite; the training baseline uses the existing recipe's default MCore path; do not mix them. - -### Capture location - -Save the DataProto (or equivalent) after rollout returns and before the first old-logprob; append trainer scores after. Record policy version, UID, input/response IDs, mask, position IDs, reward, route replay data if needed, and the UID -> packed-row mapping if reordered. Keep a pre-update checkpoint bound to the batch by content fingerprint; after strict load compare structure and keys and reject missing/unexpected keys. - -### Historical DS4 entry point - -DS4 used patched VERL options (`skip.rollout.*`, `DS4_SKIP_UPDATE=1`); confirm in code and logs that any such switch takes effect. - -### Minimal adapter contract - -The synchronous entry point takes batch, checkpoint, variant and output directory, then: - -1. Records actual policy, effective config, libraries and parallel layout. -2. Reads the same rollout, repacks per variant, keeps logical UIDs. -3. Runs the real old-logprob forward and exports per-token logprobs in UID order ([Saving and comparing](debugging-mismatch.md#compare-per-token-logprobs)). -4. Records no generation and no update, and mutable-buffer initialization. -5. Waits for workers/collectives, closes its resources, returns an exit code. - -It keeps native `no_grad`/`forward_only` semantics and replaces neither the [grad-enabled forward comparison](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract) nor update acceptance. A/B/C may share one long-lived worker once weight and state isolation is verified. The runner does not take over Ray/Slurm; the launcher blocks on remote jobs and reclaims on timeout. For injected or overlaid adapters, record the actual class and source file and assert the callsite was hit; `PYTHONPATH` can shadow the overlay. - -### Entropy when old-logprob is not recomputed (bypass) - -Bypass uses rollout logprob as old-logprob: `actor/ppo_kl` then measures rollout vs training forward exactly (0 under BI=1), and entropy comes from the grad-enabled update forward. - -- **Put metric-only quantities under `no_grad`.** With `entropy_coeff=0` the loss still holds `0 * entropy_loss`, so entropy's backward runs and saves a logits copy. -- **Chunked entropy in the verl Megatron engine cannot backpropagate.** With `entropy_from_logits_with_chunking=True`, each chunk view's backward overwrites the shared saved logits in place; the shared version counter then raises `one of the variables needed for gradient computation has been modified by an inplace operation`. Hit whenever chunking meets entropy. Fix: [verl-project/verl#8195](https://github.com/verl-project/verl/pull/8195); until merged, with `entropy_coeff=0` compute entropy under `no_grad` with the same chunking. -- **Re-verify gradients after changing forward computation.** Nemotron chunked logprob: BF16 bitwise equal to unchunked; FP32 max diff about 1e-8 at max gradient about 4. - -### Training extension - -Pairing a shared inference forward with a native VJP: check recompute, CP metadata lifetime, PP payload, parameter registration, through the standard training entry point. After a real update, export a new rollout and trainer scores; old bundles replay only with the old policy, and new-vs-old output differences are expected, not a consistency failure. Modules using vLLM layers need vLLM process groups. `param_offload=True` changes the quantized master's `data_ptr` and fails the pre-export check; use optimizer offload only. diff --git a/.agents/skills/train-inference-consistency/practices/code-and-delivery.md b/.agents/skills/train-inference-consistency/practices/code-and-delivery.md deleted file mode 100644 index aacb47c55bb4..000000000000 --- a/.agents/skills/train-inference-consistency/practices/code-and-delivery.md +++ /dev/null @@ -1,90 +0,0 @@ -# Code, experiment snapshots and delivery - -Code ownership across three repos for cross-stack train-inference consistency delivery and "three PRs, one image" (§8), and PRs and cleanup (§9). Lifecycle, cold-start and model acceptance details are not repeated. - -## 8. Three-repo delivery: from experimental implementation to three PRs and one image - -Cross-stack RL delivery for a model is one vLLM PR, one Megatron PR, one verl PR, and one unified image built from the versions the three pin. Default targets are `vllm-project/vllm:main`, `NVIDIA/Megatron-LM:dev` and `verl-project/verl:main`; a fork may host the head but does not replace the main repo's target branch. This section covers change ownership and delivery checks; it does not replace the [ten-step flow](../workflow/README.md), the [numerical gates](../workflow/validation.md) or human recipe review. - -### Before changing code: give a change map first - -Fix the inference-side change locations in step 3 and the training side and integration interfaces in step 5; step 1 need not design every interface. Submit a short table of "feature → existing module/interface → proposed change location → reason for reuse or addition", covering at least model construction, forward/backward, quantization, parameter registration, checkpoint and the build entry. - -Read the implementation and callers in the selected checkout before choosing file names; the paths below are design references, not a requirement to copy DeepSeek's directories. Do not write a standalone engine first and discover at the end that the existing integration layer or the real trainer cannot use it. - -### vLLM: implement in the formal execution path, do not deliver a patch loader - -- Development environment and code sync follow [same-commit worktree, per-file sync and build verification](../guides/environment-provenance.md#boundary-between-overlay-and-native-wheel); development sync is not how the final image is installed and does not change the PR's target branch. -- Model-specific selection goes in the model construction/execution entry; changes to shared kernels, backends and metadata go in the existing owning module, stating which models and callers are affected. -- BI=0 keeps the reviewed native path; BI=1 explicitly selects the aligned implementation. Check the actual backend and fail early when a required BI capability is missing; never silently fall back while still claiming consistency. -- Kernel numerics, binding and build registration ship together; CUDA/C++ changes enter the normal build, and JIT changes such as CuTe DSL record the compiler and specialization. Do not submit only the Python call while depending on an uncommitted `.so`, source file or build script inside the experiment container. -- Verify the default path is unchanged and cover the modified shared interfaces; when only a certain model/runner/topology is supported, restrict it explicitly rather than implying other models support BI. - -Reference points in [vLLM #54955](https://github.com/vllm-project/vllm/pull/54955): the formal kernel integration in `csrc/batch_invariant/`, the DeepGEMM BI capability check, and the stable tie-break/prefill-decode semantics of sparse selection. These address operator numerics, backend authenticity and execution shape change respectively; they are more than a BI switch. - -### Megatron mLite: separate the model, the aligned operators and the runtime - -The train-inference consistency implementation uses mLite; the training baseline uses the default MCore path of the existing verl recipe, without the aligned implementation. Record the training baseline and the vLLM BI=0 inference performance baseline separately; they are not the same comparison. When target-model support has gaps, state them and get them reviewed; do not change the training scheme on your own. - -| Change | Preferred owner | Should not go in | -| --- | --- | --- | -| Model layer structure, parameters and checkpoint mapping | model/protocol/checkpoint modules of `experimental/lite/model//` | a temporary model implementation inside a verl recipe | -| forward aligned with vLLM, autograd/STE, model-specific quantization | that model's aligned implementation and primitive modules | miscellaneous top-level mLite scripts, global monkeypatch | -| Reusable training/communication/optimizer capabilities | existing shared primitive/runtime, changed only for a real interface gap | one model's recipe constants and diagnostic branches | -| verl engine adaptation, batch/loss bridging | existing `experimental/lite/examples/verl/verl_mlite/` interface | another equivalent engine inside verl | - -- Select native vs aligned implementations through the existing protocol/factory/explicit config; do not scatter environment checks across low-level forwards or replace instance methods/global functions at runtime. Small-scope branches can stay; do not copy the whole model for a tidy directory; reuse shared math and backward, and implement independently only primitives whose numerical semantics really differ. -- Do not assemble production operators with `inspect.getsource`, AST rewriting, or string replacement plus `exec`, and do not register models via `sitecustomize` or temporary `PYTHONPATH` hijacking. Import stable interfaces from the real dependency, and commit necessary changes to the owning module. -- Keep the parameter registration, state_dict/checkpoint, distributed reduction and recomputation contracts; do not replace the training master parameters with an inference cache the optimizer does not manage just to align forward. Backward is accepted separately per the [gradient contract](../guides/backward-validation.md#write-down-the-gradient-contract-first); it need not be bitwise identical with inference, and a kernel without autograd does not automatically need STE. - -[Megatron #7050](https://github.com/NVIDIA/Megatron-LM/pull/7050) separates DeepSeek-V4's `lite/` from `vllm/primitive/` and reuses the existing `verl_mlite`, illustrating these boundaries. `verl_mlite` exists in official Megatron dev, not only in a fork; still check the version this task pins and its actual capabilities. - -### verl: a thin recipe, not a second training framework - -What is delivered depends on the route: - -- **mLite route**: for non-observability changes, deliver ideally only one recipe script and one image build script, reusing the real trainer of the current version and the installed `verl_mlite`. -- **FSDP2 Day-0 route**: verl lacks the model, so the deliverable is the self-written model implementation in verl mainline, its Auto class registration and the recipe, in one verl PR; the model implementation itself is the deliverable and needs no separate "generic interface gap" pre-review. See [Training-side model support §3.3](../guides/training-forward.md#33-fsdp2-verls-hf-interface-path). - -On both routes, when there is a real generic interface gap, first give call-chain evidence, the interface to change and the scope of impact, **get human review before implementing**, and make the minimal fix in the owning repo; do not hide code in scripts or move it to the wrong repo to make up the file count. - -The two-file form of [verl #7688](https://github.com/verl-project/verl/pull/7688) and the `verl.trainer.main_ppo`/`model_engine=mlite` chain are the reference; its development mounts and overlay are not how the final image is delivered. For data, config differences and review rules, use the [DAPO reference and review materials](../workflow/validation.md#dapo-reference-and-review-materials); do not invent an "approximate DAPO" script that bypasses the trainer. - -### Step 10: from clean PRs to verifiable delivery - -1. **Three complete diffs.** For each repo, record the target owner/branch/base SHA agreed with the user and the PR head; check the merge-base, commit history and actual diff. Required model support must be in the PR or already merged into the target branch; another unmerged feature branch cannot be the base. Do not carry over a preview PR's history drift or unrelated deletions; pin cross-repo dependencies by commit. After changing the base or rebasing, re-review the diff; evidence from the old head does not transfer. -2. **One rebuildable image.** Build from three clean checkouts, pinning the base image digest, dependencies and necessary native/JIT build inputs, and install all runtime packages normally, including external adapter packages. Training and inference workers use the same immutable platform digest and the same Python/Torch and runtime dependencies; different nodes/processes are fine, but not two independent environments inside the image to dodge compatibility. It must run without host source mounts, startup-time patches or uncommitted overlays. -3. **Re-accept from a new container.** Check the actual worker's interpreter, imports, libraries and versions; verify the proxy, full-model zero-update and the real RL loop of the reviewed recipe per the existing gates; then record performance/quality. A successful build, offline disk replay and historical 0 diff do not replace verifying the final artifact; on a difference, shrink back to a fixed failing replay first instead of relaunching full training to locate it. -4. **One mapping.** Use the [report template](../workflow/report-template.md) to list the three PRs, image digest, recipe review version and acceptance receipts, marking each item pass/fail/inconclusive/not run. Update it after code cleanup, build or dependency changes; an unverified clean implementation does not inherit the overlay prototype's pass. - -When slimming the implementation PR, non-essential diagnostic tools, docs and experiment artifacts stay in the private repo/model workspace; necessary regression tests can go with the PR. Do not delete tests that exist on the target branch, and a slim PR still runs tests. Follow each target repo's contribution requirements, record the regressions and model evaluations actually run, and submit for human review; if the gates do not pass, keep the not-accepted status. - -### Reference versions, not proof of passing - -All three PRs were open when reviewed, and the first two were marked DO NOT MERGE: - -| Reference | Reviewed head SHA | -| --- | --- | -| vLLM #54955 | `4b5c627c23e6690b8987100d1cbf78e34ac14bf9` | -| Megatron #7050 | `327a238239448de6771beea185dc6902a8de3eaa` | -| verl #7688 | `640c5f0a317307e99636924804c40624242fa19b` | - -A PR body may describe a different historical pin; when citing the design, distinguish the reviewed head, the experiment version the body claims and the version measured now, and do not treat the reference as fully correct or merged. Its Triton implementation also does not override this workflow's constraints on added/replaced kernels. - -## 9. PRs and cleanup - -- Three-repo delivery follows the target branches in §8. -- The PR description records: minimal purpose, base and patch stack, related existing PRs/why it is not a duplicate, upstream/private PRs, RL-only content, test commands and results, model impact, rollback method, AI assistance and human review requirements. -- Split a large PR into small independently reviewable PRs and prove equivalence before and after the split with the tree hash; open PRs only from your own fork. -- Resolve conflicts on the head branch of the PR the user specified; do not open a similarly named new PR to pass the original off as fixed. -- Cross-repo order is usually kernel/runtime → framework integration → recipe/pins → final image acceptance, with each later stage locking the actual revisions of the earlier ones. - -After a patch is merged and a new image is built, archive the worktree, delete the temporary overlay, and keep the PR, commit, manifest and verification report. - -## Minimal changes, contribution and delivery - -Change the config along the entry script the user specified, so multiple temporary launchers do not drift apart. End with one launch entry that actually runs, with defaults, necessary overrides and image/code pins traceable; exports that exist only in chat are not delivery. - -A fix lives with the state/feature owner: fix vLLM backend/selector problems in vLLM, not with a global monkeypatch in the training framework; a numerically matched forward must keep the real backward. Call existing reusable functions instead of building synonymous logprob metrics and compatibility switches. - -An authorized independent audit does not transfer responsibility to the auditor. Before declaring completion, check: are all of the user's goals covered? Do the results belong to the final code/image? Does the wording hide any untested item? If the evidence is insufficient, write down the gap; do not wrap it in a percentage or a success marker. diff --git a/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md b/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md deleted file mode 100644 index 224e652d9056..000000000000 --- a/.agents/skills/train-inference-consistency/practices/evidence-and-experiments.md +++ /dev/null @@ -1,179 +0,0 @@ -# Evidence and experiment rules - -Rules for what counts as evidence, how experiments are designed and recorded, and when old evidence stops holding. Mechanisms live in the guides. - -## 1. Evidence labels - -PASS / FAIL / INVALID (did not enter the target path) / VACUOUS (condition had zero hits) / SUPERSEDED / HISTORICAL. These label evidence, not lifecycle states. The tables in [validation.md](../workflow/validation.md) and the [report template](../workflow/report-template.md) have only `pass / fail / inconclusive / not run`. PASS→`pass`, FAIL→`fail`; INVALID and VACUOUS become `inconclusive` with the reason in the notes; SUPERSEDED and HISTORICAL stay out of the current table and live only in the [candidate evidence ledger](#6-candidate-evidence-ledger) and historical reports. One SUPERSEDED piece of evidence makes the gate it covered stale. A successful run, train-inference consistency and a held-out score gain are different conclusions and cannot substitute for each other. - -## 2. Evidence that must not be confused - -- Enabling `VLLM_BATCH_INVARIANT` is not cross-framework bitwise alignment; the same text is not the same logprob. -- Zero difference in probability space, zero KL or identical printed numbers do not replace comparing valid logprobs in the original dtype. -- Replaying fixed inputs is a localization tool; statistics from old trajectories are not an unbiased KL under the current policy. -- The same function name, generated code or environment variables do not prove the same actual kernel, launch config and shared libraries. -- Output replacement, forced routing, disabling Graph and bypassing bad packing are valid causal experiments if you state what they change. -- Similar quality scores are not identical per-problem results; same-problem quality, bitwise logprob and 8K/1K performance are accounted separately. When the acceptance script fails, check the producer receipt first; do not resample or declare a pass automatically. -- Verify quantization rules on the deployed bytes, not dequantized weights ([precision pitfall P2](../guides/precision-quant.md#p2-verifying-a-rule-with-dequantized-weights-gives-a-false-positive)). QAT/low-precision training only narrows the weight gap; 0 diff comes only from training calling the same kernel as inference; see [How much QAT can shrink the mismatch](../guides/precision-quant.md#4-how-much-mismatch-qat-can-remove). -- Do not manufacture a pass by changing tokens, weights, model geometry, precision or tolerance. If the comparison baseline really must change, start a new experiment and re-verify quality. -- Do not carry DS4's EP/MoE config over to dense Qwen; do not manufacture a performance gain by weakening the BI=0 config, locally disguising BI as off, or removing existing correctness patches. -- Source tests passing is not the final image passing. Development overlay results belong only to their experiment; the delivered image digest, installed source and build artifacts must match the acceptance object, and the evidence chain is rebuilt after repackaging. - -## 3. Inferences no longer carried forward - -Historical numbers and criteria stay with the run that produced them; a new task does not inherit a past run's model, dataset, topology, performance statistics method or pass status. Lessons from retracted or narrowed claims: - -- A number in a summary is not evidence: find the raw numerator and denominator for the same policy/workload/config, or retract the assertion. -- Execution completion is not numerical acceptance: finished steps, a started job, an existing adapter file or finite nonzero gradients do not show the target class ran or the comparison passed. Record the instantiated class and its import source, and accept non-empty original-dtype token comparison, module VJP against an independent reference, actual update, shutdown and long-term stability separately. -- Evidence does not transfer across candidates: a source pin, the same Docker tag, an earlier proxy, an experiment stack or a direct operator call does not prove the current candidate, image or compiled module dispatch. Verify branch, dirty patch, wheel, `.so`, dependencies and digest separately, and re-establish entry-gate evidence for the current candidate. -- A fixed split count or a pass on one path does not cover other batches, prefill/decode, the first round after resume or cross-node communication; each needs its own evidence. -- When an acceptance wrapper fails, find the failing layer before resampling or rerunning; a control-plane parsing error does not discard sealed numerical evidence. - -Each new case records only "contract → minimal counterexample → discriminating experiment → root cause/minimal fix → regression on the original reproducer → scope/source", with the small raw receipt for disputed numbers. Merge status, recommended dependency versions and best-performing config are volatile; look them up in the target repo. - -## 4. Experiment contract - -### Define the objects of equality first - -State the checkpoint/adapter/policy version, tokenizer and chat template, raw token IDs, position IDs, attention/response mask, padding/packing, and the logits processing on both training and inference sides. - -For response token `y[t]`, compare `log p(y[t] | prompt, y[:t])`. Make explicit the prediction position of the first response token, whether EOS counts, and the exclusion rules for truncation and padding. Do not align logits rows of different positions. - -Distinguish raw logits, raw logprobs, and logprobs after temperature/repetition penalty/top-k/top-p, and record which one the interface returns. The same seed does not guarantee the same scheduling or RNG consumption: freeze the tokens instead of resampling. - -### Record these goals separately - -| Goal | Valid evidence | Cannot infer | -| --- | --- | --- | -| Single-side repeatability | A/A with the same input and state | Consistency across batches or processes | -| batch/layout invariance | Same semantic sample identical under different positions, neighbors and packing | Training and inference are the same | -| prefill/decode parity | Two execution paths with the same forced tokens | backward or the training update is correct | -| Cross-stack forward parity | Real training forward against a frozen rollout | Full-vocabulary logits match, if only sampled tokens are compared | -| Effective training loop | Non-degenerate gradients, parameter change, post-update parity | Convergence or long-term stability | - -When the user requires bitwise, compare original-dtype bytes and record mismatching/valid token counts; do not reach a pass through a float cast, probability-space rounding or a tolerance loosened afterwards. If the goal is an error budget, fix the metric and threshold in advance and measure each side's own noise; generic numerical alignment need not be forced to bitwise. - -A finite logprob is necessary; an empty mask, all-zero placeholders, missed capture, NaN or the wrong weight version cannot pass. A zero may be legitimate; tell it from a placeholder by hook count, score source and a negative control. - -### Recipe and coverage record - -When building a train-inference recipe, also fix these inputs, recovering missing items from the existing config first: - -| Item | What to record | -| --- | --- | -| Workload | Training context length, rollout input/output length distribution, concurrency, packing, training token batch | -| Hardware | GPU architecture/count, GPU memory, interconnect and target TP/CP/EP/PP/DP, trainer and rollout separately | -| Numerics | master/deployed weight, activation, KV, accumulation and gradient dtype; quantization granularity, scale rule and update method | -| Scope | Raw sampled logprob or full-vocabulary logits, applicable hardware/software, Graph/cache path | -| Budget | Throughput, peak GPU memory, quality and gradient criteria, experiment count and stop condition; with an unknown threshold, report the measurement and do not declare a pass | - -Record the measured baseline of the current config, then build a strong performance baseline per the main flow, without exhaustive tuning first. Rust Bench 8192/1024, ignore EOS, c=1–64 is a fixed comparison and does not replace the target context and natural-length workload; a user-specified workload governs, with the difference reported. - -A proxy records "full-model function → proxy position → test case → uncovered items" and keeps the hidden/head/expert/quantization geometry and layout that affect kernel selection. Reuse existing modules and tests, but after changing model, dtype or parallelism re-check the contract; changing the config does not carry over the old acceptance. - -### Environment fingerprint - -Keep the code commit, dirty patch hash, untracked runtime files, image digest, GPU architecture, Torch/Triton/CUDA, actual import path, and native `.so` path and hash. Record TP/PP/CP/EP/DP for trainer and rollout, dtype/quantization/scale, kernel backend, and Graph/compile/cache and autotune config. - -Read the resolved config and actual modules/libraries inside the worker, not launcher settings. For weights, record content version and update status of derived weights/scales, not just the path. Keep immutable known-good and known-bad artifacts before diagnosing. - -### Scope of conclusions - -For each result, record model scale, valid sample count, independent token count, lengths and boundaries, process count, GPU/parallel config and number of updates. Four samples repeated into batch32 are still four independent samples; mixed scheduling a small model never triggers cannot cover the full model. - -Record machine budget, maximum experiment count and stop condition. GPU or system faults give `inconclusive`; a failed numerical comparison gives `fail`; `pass` for a scope requires every declared acceptance item. - -### Experiment discipline - -Write these nine rules into the record before comparing two arms; if any is unmet, the conclusion is "pending re-check". - -| Rule | How to do it | Mechanism and example | -| --- | --- | --- | -| One image per recipe | Both arms use the same image and recipe, differing only in the recipe switch; inject changes with an overlay and record the hash | [One recipe, one image](../guides/recipe-and-acceptance.md#one-recipe-one-image-deliver-true-on-policy-and-the-regular-baseline-with-the-bi-switch) | -| One seed per config | One seed by default; add more only when the user asks for variance, reported separately | Extra seeds spend budget on variance; most problems are systematic | -| Small-workload iteration | Iterate with small batch, short outputs and few steps; run the recipe default workload only for final acceptance | The default workload allows only one or two iterations a day | -| Freeze the criterion first and plant an error | Fix comparison window and tolerance in writing before the run; confirm with a known error that it can judge FAIL; never adjust the threshold after a failure | A threshold chosen afterwards can always pass; [paired acceptance](../guides/recipe-and-acceptance.md#paired-acceptance) | -| No conclusion from a single e2e run | Measure noise first; on the same hardware and image, report the median of multiple pairs and single pairs separately | [Performance and profiling](../guides/performance.md) | -| Same fabric | Before comparing communication bottlenecks, confirm both arms use the same transport and record the relevant environment variables | [Same-fabric fairness](../guides/debugging-mismatch.md#same-fabric-fairness) | -| Same-node comparison | Run paired arms sequentially on the same node or node group, with node names recorded | Memory fragmentation, resident processes and interconnect topology differ between nodes | -| JIT pollution: check `tl_after_bench=0` | After the run, confirm no new compilation after the bench; results with recompilation are void | [Performance and profiling](../guides/performance.md) | -| Take metrics by key name | Read numbers from raw metrics (jsonl/W&B) by full key name and cite the key in the report; never by column position in log lines | An `actor/grad_norm` column once read as entropy reversed a conclusion | - -List confounders before comparing; see [paired acceptance](../guides/recipe-and-acceptance.md#paired-acceptance). - -## 5. Execution state - -For every comparison, record the actual execution state, and keep results from different execution states as separate evidence. - -### Fields that must be recorded - -- Resolved Graph mode, capture sizes, whether the step is prefill/decode/mixed, and kernel call evidence; a successful capture does not mean every step uses the same graph, nor numerics identical to eager. -- Whether the prefix cache hit, and whether it was disabled/isolated during diagnosis; passing with the cache off does not mean the native lifecycle is fixed. -- Weight lifecycle phase: cold load, sleep/wake, Graph replay, reload in a new process, one row each. -- Whether chunk boundary state and partial input buffers of stateful layers such as Mamba/SSM are saved. - -### Keep separate evidence - -Eager, compiled with Graph off, and production Graph each keep their own artifacts; hooks run at capture, not replay, so a capture dump does not represent replay. List separately: diagnostic runtime override vs formal source candidate, Graph vs no-Graph, forward-only vs real `train()`+autograd. A new process passing does not prove the old process's cache/Graph is fixed; two fixes that each pass go through the gate again when stacked. Once CUDA state is corrupted, end that experiment and collect no more numerical evidence from it. - -### Where the mechanisms and methods are - -| Question | Where to go | -| --- | --- | -| Graph capture vs replay, the Graph/eager 2×2 comparison with fixed composition, which bucket the filler falls into | [BI kernel appendix](../guides/bi-kernels.md) (cudagraph row) | -| prefix cache semantics, Mamba chunk boundary state, bitwise comparison with "cache off as the gold standard" | [rollout-engine](../guides/rollout-engine.md) | -| Localizing CUDA asynchronous errors | [debugging-mismatch](../guides/debugging-mismatch.md#isolate-the-side-first) | - -## 6. Candidate evidence ledger - -Goal: keep an old pass from becoming a new pass. - -When an experiment changes model scale, quantization dispatch, attention, MoE, Graph or policy weights, create one ledger row per **candidate execution closure**, then decide what to rerun. - -### What one row of evidence is bound to - -| Field | What to record and check | -| --- | --- | -| Object | Original layer numbers and real adjacency, parameter/quantization geometry, tokenizer, checkpoint content, policy version; a proxy's manually spliced boundaries listed separately | -| Execution on both sides | Trainer/server image digest, source and overlay hashes, actually imported classes and files, resolved config, native/JIT artifacts | -| Path | `train()`/autograd/whether backward runs, actual dispatch of attention/norm/linear/MoE, compile/Graph/cache, TP/EP/PP layout | -| Workload | Content digest of prompt/response tokens, valid mask, length distribution, packing, concurrency, reward/advantage and actual accepted sample count | -| Criterion | Original dtype, finite valid token count, bitwise mismatch, sample count, first difference, and the independent states of gradient/update/normal exit | -| Artifacts | Inputs, score, worker/rank receipt, command, time, parent version, hash and `pass`/`fail`/`inconclusive`/`not run` | - -`policy_manifest_sha256`, the exported `policy_content_id` and the rollout receipt hash answer three separate questions, **manifest version, weight content, run artifact**, and are not comparable with each other. If the exported content ID is the hash of the export receipt, follow the receipt to check shard hashes against the files the server actually loaded; one equal version string is not enough. - -### Entry gates from cheap to expensive - -1. **Module.** Use a real triggering input and a non-unit scale/real layout; A/A first, then the candidate's forward against an independent reference; for gradients, check dX/dW/other differentiable quantities separately with the same dY. A pass covers only the shapes, arithmetic and VJP tested. -2. **Minimal real-adjacency proxy.** Cover the actual order of layer types, not just the operator set. With a fixed 8192+1024 tokens, check inference BI/PD first, then per the [dual-path contract](../workflow/validation.md#old-logprob-and-training-forward-dual-path-contract) verify native old-logprob and the grad-enabled actor forward separately; for a fixed-token comparison only, do not add GSM8K/RL to the proxy. Record source closure, actual class, Graph, non-empty valid tokens and original-dtype mismatch. When any item changes, re-examine whether this step still transfers. -3. **Full-model zero-update.** Generate a real-workload rollout with the full model on the **current** serving candidate, saving tokens and raw scores; the actor replays the same policy and history before the optimizer update. Check the earliest failing sample, then expand to the agreed sample size; if unequal, stop here and narrow the boundary instead of working around it with five-step RL or a quality evaluation. -4. **Per-version update.** For each θk, compare on-policy first, then check reward/advantage, backward, master parameters and the optimizer step, then sample θk+1; "the last training step exited" is not "the last version's comparison passed". -5. **Independent acceptance.** Module gradients, whole-model gradient coverage, real RL, quality, performance, clean exit and the final image each get their own result. Numbers from a previous candidate or another execution mode stay as historical comparison and never fill the current candidate's pass cell. - -A fixed-weight serving overlay without RL lifecycle operations such as sleep/wake is only a serving screening candidate, even if BI/PD and 8K/1K paired performance pass; once the lifecycle is restored, rerun the affected proxy, zero-update and post-update comparison. Do not splice historical throughput and another candidate's RL pass into one "completed" version. - -If a phase fails midway, classify the failure as numerical, control, environment or resource plane; keep the sealed rollout and producer hash. A control-plane acceptance bug is handled by a separately versioned recovery acceptor reading the original artifacts; never edit the old receipt, or resample and call it the same trajectory. - -### When evidence must be invalidated and rerun - -| Change | Minimum regression; old conclusions that cannot be carried over | -| --- | --- | -| Proxy's original layer numbers/adjacency or hidden, head, expert geometry change | Coverage table, BI/PD, training forward; the old proxy's 0 diff no longer counts for promotion | -| Quantization algorithm, scale, actual dispatch or JIT/compiler change | Real module input/quantized bytes, proxy, full-model zero-update; a direct-call unit test does not replace the compiled module | -| Attention/MoE operator or its training adapter changes | Real class and source, same-input operator, relevant VJP, proxy and full-model zero-update; a pass on a substitute backend is diagnostic only | -| Graph/cache/packing/parallel layout changes | The corresponding BI/PD, replay and state lifecycle; an eager pass does not cover FULL Graph | -| Optimizer or master→deployment export changes | Gradient and one-step update, export byte hash, post-update comparison; a θ0 forward pass is not enough | -| Final image/overlay/native build changes | Verify the closure in the worker and rerun affected numerics and performance; development-mount evidence does not transfer | - -If you can prove a change does not affect a gate, write the dependency closure and reason in the ledger to skip a mechanical rerun; equivalence cannot be decided from a file name, PR description, same tag or "looks like a cleanup". - -### Separate cells for performance and quality - -Tune BI=0 well under the agreed weight/KV precision and quality requirements, record the real dense GEMM, MoE expert, all-to-all and attention backends, and freeze the denominator after human audit. Pair BI=1 and BI=0 on the same full workload, nodes and candidate closure; Rust Bench 8192/1024, ignore EOS, c=1/8/32/64 is one fixed throughput workload, not natural DAPO lengths or GSM8K quality. Check effective tokens/error count first, then compute the loss per the [pre-reviewed statistical protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); do not reach the budget by slowing BI=0, changing weight precision or changing prefix cache/Graph parameters. - -Quality needs the same problems, prompt, chat template, seed, generation limit, scorer and no-retry policy; report per-problem pairing and the absolute score difference. A quality pass does not transfer to a candidate whose numerical path changed. If the user changes the threshold, the current one governs; keep the old threshold and change date for historical reports. - -## 7. Failure records - -Keep failure records as they are; append new results, with no overwriting, deletion or editing. For W&B, use a separate project per comparison and read back through the API to check no step is missing. For freezing the criterion, planting a known error and "no conclusion from a single e2e run", see [§4 Experiment discipline](#experiment-discipline). diff --git a/.agents/skills/train-inference-consistency/workflow/README.md b/.agents/skills/train-inference-consistency/workflow/README.md deleted file mode 100644 index 44e7b8419456..000000000000 --- a/.agents/skills/train-inference-consistency/workflow/README.md +++ /dev/null @@ -1,122 +0,0 @@ -# Train-inference consistency workflow: C0–C10 - -Turn a train-inference difference into a repeatable experiment with fixed input and fixed weights, locate the first divergence, fix or optimize, then return to the real training loop. Keep the user's model, precision, parallelism and consistency target. Switches, thresholds, passing scope and implementation paths from historical cases, reference PRs and implementations already in the training framework (such as an existing true-on-policy path) are experience, not defaults for the current model, dataset or hardware; measure their numerics and training-step time before reusing them. - -Forward consistency, backward conformance to the gradient contract, training quality and performance are accepted separately; none replaces another. User config comes first. Recover missing information from the source and resolved config, and ask only for what cannot be recovered and affects the approach. Choose the low-precision recipe by hardware and existing paths, not by a default precision, model or framework. Changing the model or precision requires rechecking coverage and acceptance evidence. - -Work in one directory per task, with state changes in an event log and artifacts carrying provenance sidecars. Details are in [validation.md](validation.md) (entry gate table, dual-path contract, tiered criteria, SFT three-way control, two-step RL sequence) and [numerical contracts](#numerical-contracts-and-shared-forward). - -Archive materials as you go: each run ID links config/scripts, execution identity, raw results and conclusion; the step index states the current candidate, failed/excluded items, unverified items and the human-confirmed version. C1 keeps the full tuning comparison table; C10 references the reviewed baseline ID. Layout: [archive experiment materials by step](../guides/debugging-mismatch.md#archive-experiment-material-by-step). - -## Gate table - -Iterate as "single operator / fixed replay -> proxy -> staged full model"; do not rerun every gate each time, and do not skip the human stop points at C1 and C9. Inference-only tasks apply only the relevant gates and state that the training loop is not accepted. - -| Gate | Input | Action | Pass condition | Artifacts | See | -| --- | --- | --- | --- | --- | --- | -| **C0 Module inventory and route** | checkpoint, vLLM implementation, target (BI=1 / BI=0), hardware | Per module: inference kernel, whether training implements the same math, quantization rounding, discrete choices; choose FSDP2 / Megatron-Bridge / mLite | Inventory covers all modules; route fits "BI=1 needed? how many backwards missing? what parallelism?" | Module inventory, route and rationale | [Training-side model support](../guides/training-forward.md) | -| **C1 Recipe and strong baseline** | Model, target context length, hardware | Decide weight/activation/KV precision, quantization rules, runner, parallelism, execution config; decide the training master start (BF16 original or dequantized weights; do θ0 bytes equal the official checkpoint); agree scoring scope, consistency standard, performance budget (inference throughput, and limits for training forward and forward+backward against the native training path on the main topologies; a human decides which parallelism to enable and which topologies are main); tune BI=0 with **vLLM Rust Bench** at 8192 in / 1024 out / ignore EOS / c=1–64; for MoE, measure MoE / all-to-all backend combinations | Baseline = the formal recipe's BI=0 config (including R3, PC, gpu_mem), not the serving-only optimum; multi-round measurements and backends recorded; **freeze the candidate and submit for human audit** | Tuning table, resolved config, image digest, audit materials | [Precision, quantization and QAT](../guides/precision-quant.md), [Recipe and evaluation](../guides/recipe-and-acceptance.md), [Performance](../guides/performance.md) | -| **C2 Proxy model** | Full model structure | About 4 layers covering all operator types, needed adjacent-layer interactions and state paths; real hidden/head/expert/quantization geometry and needed parallelism constraints; coverage table "full-model feature -> proxy layer/module -> test case" | Uncovered items listed; the proxy is the daily development environment | Coverage table, proxy config | [Training-side model support](../guides/training-forward.md) | -| **C3 Inference BI/PD** | Proxy, fixed weights and tokens | Verify batch and prefill/decode invariance separately, covering chunks, mixed scheduling, cache state; consider kernel suitability for training / large token counts. On failure, coarse-scan layer boundaries, bisect, capture identical operator inputs, fix kernel/metadata/layout/state | Bitwise identical in all three cases under fixed input; **proves only consistency inside inference** | Comparison matrix, minimal failing replay | [BI kernel appendix table](../guides/bi-kernels.md), [Fault localization](../guides/debugging-mismatch.md); before changing inference read [implementation ownership](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image); sync source per [runtime provenance](../guides/environment-provenance.md) (same-commit worktree + per-file hash) | -| **C4 Structural performance** | Proxy, real operator shapes | Profile prefill/decode (and training forward and forward+backward once wired): launches vs tokens, CPU gaps, sync, repeated quantization/copies/allocation, abnormal workspace; target CUDA Graph captures/replays correctly | No per-token Python loop or many-small-kernels defects (autoregressive decode steps do not count); consistency retest passes after the fix | Profile, defect list | [Performance](../guides/performance.md) | -| **C5 Training forward↔inference** | Proxy, fixed weights, input/output tokens, state | Reuse inference primitives and quantization, layout, mask, position, routing, reduction, logprob semantics; keep the training parameter container and native entry point; same-packing A/A and different-packing A/C first; then compare native old-logprob recomputation and the grad-enabled actor forward separately per the [dual-path contract](validation.md#old-logprob-and-training-forward-dual-path-contract) | Both paths: valid logprob mismatch=0 in original dtype; record non-empty valid token count and execution identity. Inference BI/PD, single operator or diagnostic substitute paths do not replace this gate | Dual-path receipt | [Training-side model support](../guides/training-forward.md) | -| **C6 backward / STE** | Forward implementation from C5 | Reuse existing backward, autograd, communication, recompute; where the gradient passes quantization rounding, decide the STE surrogate (identity/clipping, master or dequantized value, scale rule). A kernel without autograd does not force an approximate gradient; no AST rewriting or monkeypatching in the production path | Every module has a "reuse / glue / newly written" verdict and rationale, and its backward uses only what the forward saved; the training step is within budget on the main topologies fixed at C1 | Module-level backward inventory, training-step time table for the main topologies | [Backward validation](../guides/backward-validation.md), [STE contract](../guides/precision-quant.md) | -| **C7 Independent gradient validation / SFT control** | C6 implementation, same input, weights, saved state, same-dtype upstream gradient | Compare against a same-contract reference that does not call the backward under test: dX, dW, other trainable gradients, one optimizer step; fix tolerances first; report bitwise and tolerance separately. When required, compare production, an independent surrogate with the same visible forward, and native training per the [three-way SFT control](../guides/backward-validation.md#layer-3-proxy-sft-three-arm-comparison) | Independent comparison passes; without SFT, mark "quality not covered", a loss decrease is not quality. Forward 0 diff or finite gradients do not replace this gate | Gradient audit, three-way SFT results | Same as above | -| **C8 Proxy RL loop** | Proxy, real framework entry point | Follow the [exact sequence](validation.md#two-step-rl-sequence): θ0 compare -> update to θ1 -> θ1 compare -> update to θ2 -> θ2 compare; check finite, meaningful gradients, actual master change, normal exit | No cross-version comparison; a forward skipped by zero advantage does not count; proxy limited to logprob alignment: mark "loop not verified"; proxy asked to verify forward only: full-model zero-update comparison first, then the training gates; never start multi-step RL directly | Two-step RL receipt | [Two-step RL sequence](validation.md#two-step-rl-sequence) | -| **C9 Recipe review and full-model entry gate** | Launch script, image digest, code version, full resolved config, differences from the reference recipe (verl #7688 etc.), proxy smoke | **Human review; start the full model only after the recipe version is confirmed** (data, length, batch, parallelism, precision/quantization, optimizer); fill in the [entry gate table](validation.md#full-model-entry-gate-pass-proxy-train-inference-forward-first); then add layers, length, concurrency, parallelism one at a time and verify mixed scheduling, cache/Graph, multi-step loop; with checkpoint restore in scope, check alignment after the first restore | No escalation while C5 is fail / inconclusive / without artifact; two-step acceptance is the formal recipe with only fewer steps, same sequence/batch/data/algorithm; on full-model failure, shrink to a minimal replay first | Review materials, entry gate record, two-step results | [Checkpoint, resume and memory](../guides/checkpoint-resume.md), [recipe review materials and two-step agreement](validation.md#recipe-review-materials-and-two-step-acceptance-agreement) | -| **C10 Optimization and final E2E** | Full-model consistency passed, delivery PRs, reviewed recipe | E2E profile -> BI=1 hotspots (shape, dtype, layout, frequency, time share) -> representative micro bench -> correctness and performance on all target shapes -> E2E. **Triton is forbidden**; use CuTe DSL or better, CUDA/CUTLASS, or an existing implementation. Each round regresses inference, training forward and gradients by impact | On the delivery image through the real trainer entry: θ0 compare, update, θ1 compare, update, θ2 compare; 0 diff = equal original-dtype logprob for the same weight version, token history and valid mask; 8K/1K Rust Bench BI=0/BI=1 per the [pairing protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); no added long runs or gates | Delivery manifest, pairing results, loop receipt | [Performance](../guides/performance.md), [Paired acceptance](../guides/recipe-and-acceptance.md), [Delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image) | - -The consistency implementation uses mLite by default: it separates model, operators and runtime, so the inference forward can be reused, backward/STE wired in independently, and the existing `verl_mlite` interface and MCore distributed optimizer wrapper kept. The training baseline follows the default MCore path of the existing verl recipe without the alignment implementation, and is accounted for separately from C1's vLLM BI=0 inference baseline. **mLite does not inherently guarantee 0 diff or better performance**; verify the target model's parallelism, recompute and checkpoint support with the actual recipe. Both arms share one recipe and one image, and `VLLM_BATCH_INVARIANT` selects the arm (see [bi-switch-recipe](../guides/recipe-and-acceptance.md)); the delivery contract is in [validation.md](validation.md#recipe-and-pr-delivery-contract). - -## Run the cheapest failure criterion first - -1. host metadata / shape arithmetic: minimal counterexample of a pure function and boundary parameterization. -2. Single operator: same input, real dtype/layout, non-finite detection, repeat / cross-process / Graph. -3. frozen batch replay: the original failing A/C and A/A, full valid-token comparison. -4. engine: batch composition, prefill/decode, chunk boundaries, cache/Graph, actual parallelism scope. -5. Real training: non-degenerate reward/advantage, finite non-zero gradients, optimizer step, parameter change, frozen-token alignment after the update. -6. Production acceptance: full model, agreed multiple steps, restore lifecycle, quality evaluation and paired performance. - -If a second run adds no localization information, change the probe. Not reproduced, empty artifacts and an unexecuted path are not passes. When the failure is in backward, replaying the forward is not enough: save checkpoint, batch, RNG and state and move to a minimal backward experiment. Details are in [validation.md](validation.md). - -## RL delivery gates - -Training and inference use **the same immutable image digest and the same Python/Torch and runtime dependencies**. They may run on separate nodes and processes, but two environments inside one image cannot be used to dodge compatibility. Complete at least two real updates through the target RL framework's actual entry point and verify each round: original-dtype logprob 0 diff before the update, an actual parameter update, and 0 diff for the new policy after it. Saved rollouts, offline replay, separate images and cold start are diagnostics only. This is a delivery agreement, not a technical claim about all RL systems. - -Deliver one vLLM PR, one Megatron PR, one verl PR and one unified image, targeting `vllm-project/vllm:main`, `NVIDIA/Megatron-LM:dev` and `verl-project/verl:main`. Non-observability changes in verl should be limited to the recipe and the image script where possible. Follow [from clean PR to delivery](../practices/code-and-delivery.md#step-10-from-clean-prs-to-verifiable-delivery): check the target branch and full diff, pin cross-repository commits, build the image from clean source, and fill in the [delivery manifest](report-template.md#delivery-manifest). - -## Baseline and consistency run as two parallel tracks - -- **Baseline track**: C1's BI=0 tuning, backend combinations and multi-round performance evidence. Do not sacrifice default performance for consistency in advance; after reasonable tuning, freeze the candidate and wait for human audit. -- **Consistency track**: do not wait for the baseline; advance through C2–C8 on the proxy. C9 still requires the entry gate; parallel progress is not a reason to go to the full model early. -- Share the model/precision/hardware/image contract; isolate code, environment and artifacts; keep consistency patches out of BI=0. While the baseline is not final, mark the scope consistency covers and re-verify after config changes. -- Do not add nodes beyond the budget. On a single node, GPU experiments time-share; during a formal benchmark, run no GPU, heavy CPU or I/O work from the other track. The final performance gap is accepted against the human-confirmed strong baseline with full consistency evidence. - -## Two human stop points - -**After C1: baseline audit.** Once measurements support "reasonably tuned", freeze the candidate, report and wait; do not keep chasing the performance limit. Materials: model/weight version, hardware and node count, image digest / build time, code and dependency versions, reproducible launch and Rust Bench commands; weight/activation/KV precision, runner, parallelism and scheduling, Graph/cache state, actual operator backends, and a statement that alignment patches do not change BI=0; configs tried and why, valid tokens per concurrency, warmup, repeat count, throughput/latency and variation, failed items and uncovered scope; the config proposed for locking, the [statistical measurement basis](../guides/recipe-and-acceptance.md#pairing-and-computation), repeat count, per-concurrency criteria and budget, and open questions for humans. An unconfirmed candidate baseline must not be the final denominator. Re-audit after changing model/precision/hardware/image; while waiting, hand over or release nodes as needed. - -**Before C9: recipe review.** See the gate table. Changes to workload or numerical/update semantics are reviewed again; pure path or log fixes are not. - -## Find by problem - -| Current problem | Read first | -| --- | --- | -| Setting the consistency target or evidence basis | [Experiment contract](../practices/evidence-and-experiments.md#4-experiment-contract) | -| Reproducible input not yet saved | [artifacts](../guides/debugging-mismatch.md), [rollout-engine](../guides/rollout-engine.md) | -| rollout and old logprob differ: side unknown / stable failing sample / identical operator inputs still differ | [Fault localization](../guides/debugging-mismatch.md), [numerical contracts](#numerical-contracts-and-shared-forward) | -| Results change with packing, CP/TP/EP | [Layout and length](../guides/training-forward.md) | -| Fails only under Graph, compile or cache | [Execution state](../practices/evidence-and-experiments.md#5-execution-state) | -| Wiring in backward/STE, comparing gradients, SFT | [Reuse backward and backward validation (three layers)](../guides/backward-validation.md) | -| Going to the full model, shrinking a full-model failure, declaring completion | [validation.md](validation.md) | -| Candidate/backend/image changed: is old evidence still valid | [Candidate evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger), [inferences no longer carried over](../practices/evidence-and-experiments.md#3-inferences-no-longer-carried-forward) | -| Aligned; recovering performance, profiling, ablation, BI kernels | [Performance](../guides/performance.md), [BI kernel appendix table](../guides/bi-kernels.md) | -| Separating quality, fixed-length Rust Bench and normal exit | [Evaluation and paired acceptance](../guides/recipe-and-acceptance.md) | -| Image, source, wheel, submodule, overlay versions confused | [Runtime provenance](../guides/environment-provenance.md) | -| Three-repo ownership, PRs, unified image | [Three-repository delivery](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image) | -| GB200 devices, NVSHMEM, DeepEP, cross-node initialization failure | [Communication fault ladder](../guides/debugging-mismatch.md) | -| checkpoint save/load OOM, `unexpected pos`, host memory rising, first resume misaligned | [Checkpoint, resume and memory](../guides/checkpoint-resume.md) | -| BI=0 baseline arm, two modes in the same recipe | [bi-switch-recipe](../guides/recipe-and-acceptance.md) | -| Requantization rules, QAT selection, activation fake quantization | [Precision, quantization and QAT](../guides/precision-quant.md) | -| Many repeated samples within a rollout group | [rollout-engine](../guides/rollout-engine.md#1-sampling-seed-collision) | -| Different seeds give the same samples (greedy n repeats are identical; SGLang deterministic mode without a seed always uses `sampling_seed=42`) | [rollout-engine §1](../guides/rollout-engine.md#1-sampling-seed-collision) | - -The ten kinds of evidence that must not be confused are in [evidence and experiment rules §2](../practices/evidence-and-experiments.md#2-evidence-that-must-not-be-confused). Writing comparisons and replays: [save and compare](../guides/debugging-mismatch.md#compare-per-token-logprobs). Record experiments and report with the [report template](report-template.md). This workflow grants no extra GPU, job stop, image publish or cross-repository write permission. - -## Numerical contracts and shared forward - -### Check the arithmetic that actually executes - -The same formula does not guarantee the same floating-point result. Check accumulation dtype, cast position, FMA, reduction tree, chunking/split-K, approximate activations, quantization scale/rounding, epilogue and the input dtype of the final logprob. The same source can also produce different execution configs through autotune, compile specialization or a different GPU. - -A shared kernel still needs its caller, arguments, layout and compile boundary checked; a direct call and a torch.compile fused path may differ. If needed, keep the same forward behind an existing opaque custom-op boundary; do not introduce unproven arithmetic rewrites to imitate the output. - -Quantization execution checks (scale taken from the actual runner, quantization dispatch, quantizer outside the `nn.Module` tree, training vs deployment quantization grid) are in the [precision, quantization and QAT guide](../guides/precision-quant.md#training-side-numerics), the grid in [P8](../guides/precision-quant.md#p8-training-grid-differs-from-the-deployment-grid); format and scale records are in [conversion chain §1](../guides/precision-quant.md#1-four-stages). - -First check that the semantics exist on both sides. When a feature is inside the kernel on one side and external or missing on the other, the difference does not go away with precision settings. Run a two-layer minimal comparison per model feature (attention sink, RoPE variant, final norm, etc.), enabling one at a time. In the external MiMo case SWA alone gave 0 and SWA+sink gave `3.906e-03`, because only one side's kernel included the sink. - -### Discrete decisions come before last-bit error - -After router/indexer logits, also compare the selected set, the order and the tie-break. Top-K with the same set in a different order still changes downstream gather/reduction. A deterministic tie-break defines the index order for equal scores and keeps the invalid/padding index contract. - -For MoE, also fix or verify dispatch, expert capacity, grouped GEMM geometry, quantization and combine. Do not hide a router difference by forcing expert indices; forced routing only isolates the downstream. - -The top-k/top-p truncation set is also a discrete decision: a last-bit logit difference moves tokens near the cutoff in or out of the set, and the normalization denominator of the processed logprob jumps with it. Compare the candidate set before processed logprob. MiMo-V2.6 replays the candidate set together with R3 in production: rollout records a fixed-shape, full-vocabulary bitmap (avoiding host sync and truncation), and training renormalizes within the set. If the contract adopts this replay, state in the report that it masks the set difference, and count the set flip rate separately. - -### Test shape and process boundaries - -For the same victim, construct batch1 / multiple neighbors, different physical positions, lengths on both sides of a boundary, cross-process autotune, prefill/decode and Graph/eager. Different launches when only M changes are not a numerical difference; measure the output. Finding no difference covers only the shapes tested. - -SSM chunk recurrence and per-token recurrence can be two legal but different rounding paths. Unifying them requires matching the chunk grid, boundary state precision and gate order, not just a deterministic switch. - -For history-dependent operators, a single token's scan/conv input and output can make an upstream history difference look like an error in that operator. First find the first differing token at each layer over the **full prompt batch**, then capture norm, projection, conv, scan, gate and output projection in order at that layer. A single-position probe only narrows candidates; it does not prove the first cause. A plain Python hook in a compiled graph may hold stale warmup data; verify capture freshness with a predictable sentinel such as embedding/residual, and if needed use the custom-op boundary that actually executes and a separate buffer. - -### Boundary between shared forward and backward - -With a shared inference forward + native differentiable VJP, verify separately: forward parity, a finite gradient for every trainable parameter, and gradient comparison under a cotangent representable in the same dtype. Operator-level VJP verification does not prove the whole-model gradient is bitwise equal to the original training implementation. - -Do not detach the inference output and use it as a trainable substitute. A finite gradient does not mean effective learning; also check that advantage/reward is not degenerate, the gradient is non-zero, and parameters and optimizer state are updated. - -### Changing the arithmetic baseline - -If you deliberately change both sides' arithmetic baseline to align them, keep the old baseline, build a new frozen rollout, state that "old outputs are not guaranteed to stay bitwise unchanged", and run model quality and performance regressions. Success numbers from the old baseline do not carry over to the new implementation. diff --git a/.agents/skills/train-inference-consistency/workflow/report-template.md b/.agents/skills/train-inference-consistency/workflow/report-template.md deleted file mode 100644 index e8fffb0e383d..000000000000 --- a/.agents/skills/train-inference-consistency/workflow/report-template.md +++ /dev/null @@ -1,119 +0,0 @@ -# Train–inference consistency report - -## Outcome and scope - -State pass/fail/inconclusive for the requested scope. Record model/layers, independent samples, valid tokens, dtype, topology, Graph/cache paths, update count and untested conditions. - -Link the human-approved acceptance checklist: required vs out-of-scope checks, recipe revision, two-update stop condition, quality/resume coverage and target-repository submission requirements. Out-of-scope checks are unverified, not passed; finishing the short run does not complete missing required checks. - -## Immutable experiment - -Record source commits and dirty patch hashes, loaded libraries, image digest, GPU, checkpoint/policy identity, input bundle digest, score semantics, resolved config, commands, resource budget and artifacts. - -Separate source, installed wheel, loaded native libraries and JIT specializations. State whether the measured runtime and the final delivered image are identical; if not, name the missing verification instead of inheriting the result. - -For recipe work, record training context, rollout length distribution, concurrency, packing/token batch, hardware/interconnect and master/deployment/activation/KV/accumulation/gradient dtypes, plus quantization granularity, scale/update rules and the quality/throughput/memory budgets defined before testing. Unknown budgets stay unaccepted even when measurements exist. - -## Delivery manifest - -For the three-PR / one-image delivery, fill this from the final clean build, not development overlays. Use full SHAs and immutable image identity; mark missing evidence as not run. - -| Deliverable | Review / source identity | Build or acceptance evidence | -| --- | --- | --- | -| vLLM PR | URL; target owner/branch/base SHA; head SHA | Build command and wheel/native/JIT inputs | -| Megatron PR | URL; target owner/branch/base SHA; head SHA | Installed model / primitive / adapter identity | -| verl PR | URL; target owner/branch/base SHA; head SHA | Recipe and image build script paths | -| Unified image | Base and final platform digests; dependency pins | Build log; actual trainer and rollout worker imports/libraries; no development source mounts | -| Approved recipe | Script SHA, resolved config and human approval | Changes from reference recipe; dataset and topology | -| Final validation | Commands and artifact IDs tied to the above image | Proxy / full zero-update / θ0–θ2 / quality / performance status, each reported separately | - -Each PR is complete against the agreed base, not stacked on an unmerged feature branch; cross-repository dependency pins are expected. Every required runtime file outside the PRs must be a pinned existing dependency; an uncommitted production patch is a delivery gap. A changed base, build input or recipe requires reassessing affected evidence. - -## Proxy coverage and gradient contract (when training is in scope) - -Map each full-model feature to its proxy layer/module and test artifact; name omitted depth, geometry, topology and execution paths. Per module, record visible forward semantics, trainable variables, master vs dequantized surrogate values, STE/clipping/scale rules, fixed discrete choices and differentiable scores, VJP accumulation/collectives, and saved/recomputed state with its memory cost. - -Link the independent reference and explain why it does not call the tested backward. Record shape/dtype/layout and random, zero, sparse and captured upstream-gradient coverage; per input/parameter: missing/nonfinite gradients, max absolute error, relative L2, norms and cosine where defined. Include predefined thresholds with rationale, original forward-byte comparison and one-step master-parameter deltas. - -## Symptom and localization - -Record per-request/token differences, A/A and A/C results, side-isolation evidence, last equal / first unequal semantic boundary, identical-input operator comparison, competing explanations and counterexamples. - -## Cause and change - -Explain the violated arithmetic/metadata/state contract and the smallest triggering input. Separate the actual repair from diagnostic interventions. Link code changes and the preserved regression. Write unknown if evidence does not establish a cause. - -## Verification - -Record the proxy upgrade decision before escalating to full-model BI consistency or training. Do not merge inference-only and cross-stack parity into one result; a failed, missing, empty or stale forward comparison blocks escalation. If the user scoped the proxy to fixed-token forward parity, mark SFT/RL not run and list the later training gates. - -| Proxy upgrade gate | Status (pass/fail/inconclusive/not run) | Bound command / artifact and valid-token count | Decision | -| --- | --- | --- | --- | -| Proxy feature/operator coverage | | | | -| Proxy inference BI and prefill/decode parity | | | | -| Proxy inference ↔ native old-logprob recomputation | | | | -| Proxy inference ↔ actual grad-enabled actor forward | | | | -| Backward / SFT / RL, if in proxy scope | | | | -| Upgrade to full-model BI work | | | | - -For each cross-stack row, follow the [dual-path contract](validation.md#old-logprob-and-training-forward-dual-path-contract): model geometry, frozen policy/checkpoint identity, source/image digest, dtype, topology, model/grad modes, output requires_grad, actual backends, raw logprob mismatch count and first mismatch. Keep native no-grad scoring semantics; its pass does not prove grad-enabled forward parity. Capture scores before updates and loss transforms. - -| Gate | Command / artifact | Result and count | Scope limitation | -| --- | --- | --- | --- | -| Original failing replay | | | | -| Operator / metadata regression | | | | -| Engine and production layouts | | | | -| Independent module VJP and one optimizer update | | | | -| Controlled proxy SFT (A/B/C paths) | | | | -| Real backward/update | | | | -| θ0 / θ1 / θ2 parity | | | | -| First post-checkpoint-resume export / parity | | | | -| Model quality | | | | -| Clean shutdown / sustained run | | | | -| Final immutable image, without development overlays | | | | - -For two-update RL, keep a receipt per policy θ0, θ1, θ2: input/mask identity, valid token count, original-byte mismatches, actual trainer-forward execution and update norm/count. Record both forward-path results at matching frozen policy versions. Worker exit is checked at final shutdown, not at each policy transition. - -## SFT controls and quality (when training is in scope) - -Separate SFT-A (production), SFT-B (same visible forward at the same snapshot plus an independent surrogate) and SFT-C (verified native training); these labels are independent of replay packing variants. Link same-checkpoint/same-batch A/B gradients and optimizer deltas; once trajectories diverge, later forward values are not assumed identical. A/C differences include forward/quantization changes. - -Record initial master and optimizer state, data/order, labels/masks, packing, effective-token normalization, batch/accumulation, learning-rate schedule, clipping, precision/loss scaling, RNG/dropout/recompute and intentional differences. Keep raw train, held-out and fixed-probe loss, gradient norm, update RMS, clipping/nonfinite/skipped-step counts, throughput and baseline/seed variation when available. Use one optimizer-step or effective-token axis; report predefined windows/thresholds, missing controls and whether evidence is only a short smoke run. - -## Performance - -Record matched work, warm/cold timing, node pool and paired order, repetitions, latency/throughput, kernel/rollout/full-step results and numerical receipts. Separate measured improvements, rejected candidates and estimates. - -For 8K/1K acceptance, link the [approved aggregation protocol](../guides/recipe-and-acceptance.md#pairing-and-computation): formula, repetitions, per-concurrency thresholds and invalid-leg replacement rules. Include all raw pairs, exclusions with reasons and the approved aggregate per c; median paired loss and loss from median throughputs are different statistics. Other phase timings are diagnostics, not gates. - -Where relevant, include training time, peak memory, launch count vs token batch, CPU gaps and repeated quantization/allocation. Separate early feasibility profiles from final full-model optimization; normal autoregressive decoding is not evidence of an avoidable per-token prefill/training loop. - -State whether this is a BI-only causal comparison or each side's best valid stack, and list all differing flags/backends. Include prompt count vs expanded requests, actual input/output lengths, completed output tokens, wall-time denominator and whether the run is fixed-length synthetic or natural-EOS. Do not sum overlapping GPU ranges or percentages from separate ablations. - -## Evidence corrections and retention - -List superseded claims with the correcting artifact and reason. Record which original failure bundles, checkpoints and small receipts are retained, their storage budget/location, and any missing eval steps. Never fill a curve with a different policy lineage. - -## Remaining work - -List concrete gaps and the smallest next experiment. State whether production hooks/config were restored. Keep raw evidence; do not upgrade historical or partial results into current completion. - -## Experiment record (YAML) - -```yaml -# Train-inference consistency replay experiment record: each variant changes only one variable, everything else is fixed. -purpose: "frozen-policy old-logprob replay" -target: "bitwise selected-token raw logprob" -rollout: "/absolute/path/to/frozen-rollout" # original framework batch, not exported scores -checkpoint: "content id + pre-update recovery path" -environment: "image digest, source+patch hashes, libraries, GPU" -topology: - trainer: "TP/PP/CP/EP/DP" - rollout: "TP/PP/CP/EP/DP" -state: "RNG, mutable buffers, Graph/cache and reset policy" -budget: {max_runs: 3, timeout_seconds: 600} -variants: - - {name: A, change: "none (A/A baseline)", cmd: "", output: "", exit: null, result: ""} - - {name: B, change: "none (A/A repeat)", cmd: "", output: "", exit: null, result: ""} - - {name: C, change: "packing reversed", cmd: "", output: "", exit: null, result: ""} -``` diff --git a/.agents/skills/train-inference-consistency/workflow/validation.md b/.agents/skills/train-inference-consistency/workflow/validation.md deleted file mode 100644 index 9bac973376cb..000000000000 --- a/.agents/skills/train-inference-consistency/workflow/validation.md +++ /dev/null @@ -1,169 +0,0 @@ -# Tiered regression and completion criteria - -## Full-model entry gate: pass proxy train-inference forward first - -When backends, compile paths or delivery images have changed, check against the [candidate evidence ledger](../practices/evidence-and-experiments.md#6-candidate-evidence-ledger) whether old results still apply. A passing old candidate is history only and does not fill the current candidate's entry gate cells. - -Before full-model correctness or training experiments, provide the evidence below. Missing or not executed counts as not passed; repeated full-model launches are no substitute. Proxy inference-internal BI/PD and proxy training forward↔inference logprob comparison are two independent results: the former passing does not mean the latter passes. - -| Gate | Minimal evidence | -| --- | --- | -| Proxy coverage | Inventory of operator types, necessary adjacent-layer interactions, quantization geometry, state and parallel paths; uncovered items stated | -| Inference BI/PD | Batch, prefill/decode and mixed-scheduling comparison on the proxy with fixed weights and tokens; proves only consistency inside inference | -| Training forward↔inference | Same proxy, same policy weights, fixed tokens: verify native old-logprob recomputation and the grad-enabled actor forward used in training separately, each against the target production inference path; keep evidence per the dual-path contract below | -| Backward and quality | [Independent module gradient comparison](../guides/backward-validation.md#layer-1-module-level-vjp-comparison) with reference implementation, tolerances and raw results; controlled SFT per task scope, or an explicit statement that quality is not covered | -| Real RL update | At least two effective updates, θ0/θ1/θ2 same-policy comparison, actual parameter change, normal exit; sequence below | -| Performance feasibility | launch/sync/memory overhead on real shapes and the proxy profile, ruling out bottlenecks such as per-token launches in batchable prefill/training; the target Graph path works | - -This is an escalation gate, not the completion criterion: two-step RL does not prove convergence, and the proxy does not prove full depth, capacity, long context or cross-node correctness. On the full model, extend coverage one item at a time. Inference-only tasks apply only the relevant gates, state that training is not accepted, and add no training work. - -Before escalating, record per gate in the [report template](report-template.md#verification): `pass`/`fail`/`inconclusive`/`not run`, command, artifact path and the bound model geometry, weights, source/image, precision and parallelism. Accept only current-candidate evidence. For training forward↔inference, the bitwise target requires original-dtype mismatch of valid tokens to be 0; tasks with an agreed error budget use the predefined standard and cannot use it to report a bitwise failure as a pass. An unexecuted training forward, an empty mask, or passing only single-operator or inference BI/PD tests must not release to the full model. Keep the first failing sample on the proxy for layer-by-layer localization, and after the fix rerun the original failing case and the affected matrix. - -### old-logprob and training forward dual-path contract - -- **Native old-logprob recomputation:** follow the framework's actual path and its `no_grad`/`forward_only` semantics; do not force autograd on for acceptance. Its result proves only that this scoring path matches rollout. -- **Actual training forward:** with the same frozen θ, token history and valid mask, run the grad-enabled actor branch training actually uses, capture logprob before the loss transform, and compare with rollout independently. Keep the target packing, recompute, operators and parallelism; do not substitute a diagnostic forward. Passing does not replace the later independent backward validation. -- Record per path: model mode, grad mode, output `requires_grad`, operator class / backend / source hash, original dtype, non-empty valid token count and mismatch. Diagnostics (substitute attention/MoE, Graph off) prove only that path; `forward_only` alone is not grounds for rejection, and native old-logprob passing does not prove the grad-enabled branch. -- The comparison happens before the corresponding policy update, for both proxy and full-model θ0 entry. If the formal recipe has several minibatch updates per round, do not require 0 diff between the updated actor and the old-θ rollout, and do not modify the algorithm to manufacture equality; record the real parameter version and compare only within one version. - -When an adapter / overlay is injected dynamically, the receipt must also record the actual object's class, source file and key config and assert that the target callsite was hit. A script directory placed before the overlay in `PYTHONPATH` can silently load an old adapter; a 0 diff from such a run cannot vouch for the new candidate. - -If the user limits the minimal proxy to fixed-token logprob comparison (or forward only), do not require GSM8K, SFT or two-step RL; mark those gates `not run`, not passed. Scale up as agreed only after the proxy train-inference forward passes. Before any multi-step full-model RL, first complete a zero-update same-policy logprob comparison with **full-model θ0 real workload, real training mode, target production inference mode**, then the necessary gradient validation and effective-update checks. A fixed 8K/1K proxy pass does not replace the full-model variable-length workload. When full-model θ0 fails, stop optimizer updates, keep the failing trajectory and shrink to the first divergence. This scope adjustment does not waive the proxy forward gate. - -### Limited exceptions that allow going to the full model early - -- Establishing the BI=0 performance/quality baseline of the full target model is not approval for full-model BI debugging or training. -- When a fault depends on real depth, capacity, cross-layer state or cross-node topology and cannot be reproduced on a semantics-preserving proxy, targeted full-model diagnosis is allowed. Before the experiment, write down why the proxy is insufficient, the hypotheses to distinguish, the minimal scale and the stop condition. -- Scaling up does not scale up resource permissions: no extra GPUs, no stopping other jobs, no bypassing the model/runner scope the user set. - -### Faster iteration and fallback on failure - -Judge each change's impact first: minimal counterexample, single operator or fixed replay, then the relevant proxy gates, then the staged full model; do not repeat all ten steps per commit. Changes to shared arithmetic or gradients regress the affected forward/backward; final delivery still runs the full agreed acceptance. - -On a full-model failure, first save the input before the first divergence, the weight version, metadata, cache/routing and necessary RNG state, and extract a module or few-layer reproducer. Iterate on it only after confirming the fault survives shrinking. If the fault disappears, a passing shrunk test is not a fix, and the full-model reproduction is kept. - -Fixed rollout / module inputs may be cached and replayed in an isolated resident process, bound to fingerprints of weights, input, config and necessary state; a new implementation still reruns the comparison. Each scale-up states the existing evidence and the open question the new scale answers. - -If a candidate fix changes the sampling trajectory or EOS stop point, save the new token history and replay training old logprob on it with the same weight version. The old trajectory remains a fixed-input diagnostic; comparing new inference against the old response length is not an end-to-end pass. List evidence separately for a diagnostic runtime override, a formal source candidate, Graph/no-Graph, forward-only and real `train()`+autograd. An isolation package may be used for localization; final acceptance returns to the delivered code and the default target execution mode. - -## Start from the cheapest failure criterion - -1. host metadata/shape arithmetic: minimal counterexample of a pure function and boundary parameterization; if needed, exhaustively check equivalence before and after the fix. -2. Single operator: same input, real dtype/layout, non-finite detection, repeat/cross-process and Graph. -3. frozen batch replay: the original failing A/C and A/A, full valid-token comparison; UID mapping passes. -4. engine: batch composition, prefill/decode, chunk boundaries, cache/Graph, actual parallelism scope. -5. Real training (when in scope): non-degenerate reward/advantage, finite and meaningful non-zero gradients, optimizer step, actual parameter change, frozen-token alignment after the update. -6. Production acceptance: full model, agreed multiple steps, update/restore lifecycle, quality evaluation and paired performance. - -Not every task runs every tier. Report the tiers completed; when the user asks only for diagnosis, do not start full training. Extend existing test suites first; do not replace behavior tests with assertions on source strings. - -## Preventing false greens - -- Zero valid tokens, comparison not executed, missing rank/sample: invalid experiment. -- All-zero score, mask all false, the same sample repeated: check the producer and coverage; a passing comparison is not enough. -- probability/KL/logPPL is zero: still check the original logprob bytes. -- Zero gradient or identical rewards: maybe no effective update; check advantage, update count and parameter difference. -- Small model / short output passes: confirm the full model's mixed-step and long-context trigger conditions. -- A shrunk proxy's error magnitude is not the full-model baseline: in the external MiMo case the 12-layer proxy was `1.866e-12` and full-depth real weights were `1.143e-02`. The proxy only verifies the execution path. -- In-training mismatch metric is 0: confirm the micro-batch really ran the trainer forward. The framework may skip the forward for a zero-advantage batch and return 0, a trivial match; count only reports with non-zero advantage. -- A success marker written before the process was killed: separate a completed numerical experiment from normal exit / long-term stability. -- probs_diff is 0 but many sibling samples within a group are identical: a sampling diversity defect, independent of the consistency metric. Count distinct rollout samples per group; the mechanism and fix are in [rollout-engine §1](../guides/rollout-engine.md#1-sampling-seed-collision). -- The launcher shows `COMPLETED 0:0`: read the recipe's own exit code. When tests report passed, check the skipped count; GPU cases may be skipped wholesale by a marker or environment gate. - -## Unified measurement basis and matrix for forward consistency - -Fix policy, token history, position, mask and tensor semantics, and record per item the original dtype/shape, valid token count (greater than 0), non-finite count, bitwise mismatch count, maximum absolute difference and first failing position. A dtype difference is a contract inconsistency first; do not convert and then claim a bitwise pass. After aligning semantics, compare contiguous tensor bytes that preserve the bit pattern or an equal-width integer view; floating-point `==` does not distinguish positive and negative zero. NaN/Inf in a valid logprob is a failure even with identical bytes. Sampled-token comparison does not extrapolate to the full vocabulary. - -Keep cases and artifacts per item for the target execution scope, and list unsupported / not run items: - -| Dimension | Comparison under fixed semantics | -| --- | --- | -| A/A | Repeated execution with the same weights, input and state; in a separate process if needed | -| BI | Single request, different batch sizes, companion requests, permutation/position and packing | -| PD | Full prefill, chunked prefill and step-by-step decode on the same history, covering mixed scheduling | -| Length and branches | Real lengths, chunk/causal boundaries, and cases that trigger sparse selection, MoE and special modules; not only random tokens | -| Execution state | Target Graph capture/replay, cache reuse/reset and the actual parallel layout | - -The bitwise target requires mismatch=0; error-budget tasks report against the predefined standard and are never recorded as a bitwise pass. Gradient tolerances are defined independently, see [module VJP](../guides/backward-validation.md#layer-1-module-level-vjp-comparison) and [SFT control](#proxy-sft-three-way-control). - -## Two-step RL sequence - -This sequence is the lower bound for entry; "container started / training script exited 0" is not a pass. - -When the task includes the full train-inference loop, complete at least this sequence; every actor recomputation uses the same policy and token history as its rollout: - -```text -θ0 → rollout0 → actor(θ0) recompute and compare - → backward0 / optimizer update → θ1 -θ1 → rollout1 → actor(θ1) recompute and compare - → backward1 / optimizer update → θ2 -θ2 → rollout2 → actor(θ2) recompute and compare -``` - -The θ2 comparison needs no third optimizer update. To compare outputs across policies, use the same frozen token history; do not rely on sampling happening to match. - -For each stage save the policy content identifier, token/mask, non-empty valid comparison count and raw byte comparison result; for each update save finite loss/gradient, coverage of parameters that should have gradients, master update norm and optimizer step. - -For a bitwise task the θ0, θ1 and θ2 comparisons must all have mismatch=0, with two effective updates and a normal exit actually completed. Skipped updates, both sides reading an old cache, an empty mask or comparing only probability summaries are not allowed. - -A small master update may not cross a quantization bucket, so deployment bytes need not change every step. Judge an update by master update norm and optimizer step, not by step count or byte change. - -## Proxy SFT three-way control - -The three paths, controlled variables, recorded items and localization order are in [Backward validation layer 3](../guides/backward-validation.md#layer-3-proxy-sft-three-arm-comparison). If the module gradient comparison fails, fix the implementation before starting SFT. - -## Performance and model quality - -When forward arithmetic, precision, scorer or inference path changes, rerun the same quality evaluation within the human-reviewed scope and the target repository's requirements. Record model/data, problem count, prompt format, generation settings, statistical variation and unfinished parts; mark out-of-round items unverified, and add no SFT or long runs. A loss decrease or two-step reward change does not prove no quality regression. - -Micro benchmarks, fixed-token parity, generation quality, normal exit and multi-step stability are separate. A run with four non-zero updates and 139,260 bitwise-equal tokens that ended in a DataLoader error supports four-step alignment, not clean shutdown or 50-step stability. - -The full model covers the reviewed scope (real depth / weights, target context, concurrency, packing, parallelism/hardware, cache/Graph) and reruns affected correctness gates; proxy curves do not extrapolate to full-model quality. Final performance is accepted on 8K/1K per the [paired statistics protocol](../guides/recipe-and-acceptance.md#pairing-and-computation); training and full-step times may be listed as diagnostics and add no performance gate. Scope and stop condition of the two-step short run are in the [general two-step acceptance agreement](#general-two-step-acceptance-agreement). - -## Failure classification and delivery - -Use `pass`, `fail`, `inconclusive` with a reason. Keep source/patch/image/bundle hashes, commands, per-item counts, earliest difference, excluded hypotheses, and before/after results on the original reproducer. Cleaned-up code needs re-verification or explicit behavior-equivalence evidence; old tests do not vouch for new code. - -Passed items stand despite unfinished ones, but the final conclusion is limited in scope. Use the [report template](report-template.md). Formal upstream contributions also follow the target repository's duplicate check, human line-by-line review, tests and AI assistance disclosure; historical cases do not mean the review was done on the author's behalf. - -If the delivery includes an image, the gate object is the final immutable artifact, not a development image plus a source mount. Verify the closure per [runtime provenance](../guides/environment-provenance.md); accept save per [01 §3.2](../guides/checkpoint-resume.md#32-reproduce-and-discriminate), load per [§4](../guides/checkpoint-resume.md#4-load-enough-capacity-does-not-mean-enough-for-the-transient-peak), and the first resume per [§6](../guides/checkpoint-resume.md#6-how-to-prove-resume-is-correct). Passing environment checks proves only the preconditions for running, not GPU numerics or performance. - -## Recipe review materials and two-step acceptance agreement - -The integration point is in [training-side model support](../guides/training-forward.md); only process agreements are here. - -### Recipe and PR delivery contract - -Locate the existing integration layer in the selected version first. Megatron's `experimental/lite/examples/verl/verl_mlite/` provides the verl engine adapter; an external integration package installs normally into the unified image, and removing the runtime overlay does not require reimplementing the engine in verl. Check the actual engine registration, model factory, batch conversion, loss normalization and optimizer path, and list reusable interfaces and specific gaps. - -- **verl PR** contains by default the model's DAPO recipe and the unified image Dockerfile. Public interface changes beyond these two: submit the gap and plan first, implement after human review; necessary regression tests may stay. -- **Megatron PR** contains the model, operators and necessary support in the existing `verl_mlite` integration layer; do not move generic code that does not belong in Megatron there just to keep the verl PR at two files. -- **vLLM PR** contains the necessary inference-side model, kernel, state and weight loading changes. -- Per the [general delivery target](../practices/code-and-delivery.md#8-three-repo-delivery-from-experimental-implementation-to-three-prs-and-one-image), submit complete PRs to the three main repositories, including the required model adaptation and not stacked on another unmerged feature PR; pin cross-repository dependency versions and state the build order. Consistency uses mLite and the training baseline uses the existing verl recipe's default MCore path (see [the workflow](README.md#gate-table)); do not merge them into one training scheme. - -### DAPO reference and review materials - -[verl #7688](https://github.com/verl-project/verl/pull/7688), at reviewed version `640c5f0a317307e99636924804c40624242fa19b`, uses Megatron's integration package through `model_engine=mlite` and `pkg://verl_mlite.config`, so the verl diff needs only the recipe and the Dockerfile. Follow its real `verl.trainer.main_ppo` execution chain; do not assume its config keys, DeepSeek precision/topology, mounts or overlay fit the current model. - -The reference case's model, data and evaluation set are not defaults. Choose training data / reward by target model and human agreement, and keep the training set distinct from the quality evaluation set. Finish config resolution and the necessary proxy smoke, then submit for human review: - -- Executable launch script, unified image digest, commit of each repository and full resolved config. -- Parameter differences from the reference recipe with reasons: data/reward, prompt/response length, rollout count, batch/microbatch, parallelism, quantization / KV precision, optimizer. -- Proxy check results, known uncovered items, and the acceptance scope of the full-model zero update and the two real updates. - -Formal full-model E2E starts after a human confirms a specific recipe version and the proxy entry passes. Workload or numerical/update changes are reviewed again; path and log changes that leave input and execution semantics unchanged are only recorded. This gate does not block already authorized builds or proxy development. - -### General two-step acceptance agreement - -- **Freeze versions first.** Record the heads of the three delivery PRs, the unified image digest and the full recipe config. Changes that affect results require a version update and the related verification; old artifacts cannot be reused for acceptance. -- **Review the recipe first** (materials above). Confirming process or scope does not approve an unsubmitted recipe. The review checklist also marks which quality, restore and long-run checks are required or out of scope for this round, and the evidence the target repository requires; out of scope does not mean passed, and the short run ending does not hide missing required items. -- **Run two effective updates by default.** Use the formal recipe and shorten only the execution length to two steps; do not downgrade data, sequence length, rollout count, batch, optimizer, algorithm, parallelism or backend. If the stop step count drives the learning-rate schedule, keep the formal recipe's learning-rate/warmup semantics for the first two steps instead of compressing the full schedule into two; the implementation is reviewed with the recipe. -- **Keep the existing correctness gates.** After the proxy real training forward passes, run the full-model [two-step RL sequence](#two-step-rl-sequence) with actual parameter updates; a fixed-weight replay does not count instead. -- **Performance is accepted only on 8K/1K.** On the frozen version, use Rust Bench at 8192 input / 1024 output, ignore EOS, c=1/8/32/64, against the human-confirmed BI=0 baseline and performance budget. Training forward, backward and full RL step times may be recorded but add no performance gate. -- After the two updates, the θ2 comparison and the agreed 8K/1K performance pass, stop the short run and report. Overall acceptance is still judged item by item against the review checklist; if a required item is missing, state that it is not complete. Do not extend to 100 steps, add an MCore long-run comparison or new quality gates, and do not extrapolate the two-step result to long-run training quality. - -### E2E comparison target - -Run the [two-step RL sequence](#two-step-rl-sequence) through the real trainer entry point. 0 diff means equal original-dtype logprob for the same weight version, token history and valid mask; θ0 and θ1 outputs may differ. Verify the actual master update each round (not step count or byte change). Disk replay is a localization tool and does not replace acceptance through the real trainer entry point. - -Finally re-verify with the image built from the delivery PRs and the reviewed recipe. Record standalone engine smoke, old image numerical results and historical performance separately; they do not replace full evidence from this execution entry point. diff --git a/.claude/skills/train-inference-consistency b/.claude/skills/train-inference-consistency deleted file mode 120000 index 906cd763172c..000000000000 --- a/.claude/skills/train-inference-consistency +++ /dev/null @@ -1 +0,0 @@ -../../.agents/skills/train-inference-consistency \ No newline at end of file From c61a1e73c664ef59e4dd496e6b32eb4ba0c95b9a Mon Sep 17 00:00:00 2001 From: siyu Date: Sat, 17 Oct 2026 23:10:04 +0800 Subject: [PATCH 798/952] [Feature] Skip MLA weight absorption when weights are pre-processed (#57320) Co-authored-by: Xun Sun --- vllm/models/kimi_k3/nvidia/mla.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 18fd93579be7..83a7a205ba00 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -75,7 +75,7 @@ resolve_quant_method, ) from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import is_weights_pre_processed, replace_parameter from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.kimi_k3.nvidia.low_latency_gemm import try_low_latency_gemm from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( @@ -440,18 +440,21 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: projected into latent space by ``W_UK_T`` and the attention output is projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. """ - W_UK, W_UV = split_kv_b_proj( - self.kv_b_proj, - act_dtype, - self.kv_lora_rank, - self.num_local_heads, - self.qk_nope_head_dim, - self.v_head_dim, - ) - # (L, N, V) -> (N, L, V) - replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) - # (L, N, P) -> (N, P, L) - replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) + # The weight cache daemon exports the absorbed weights. Repeating the + # absorption would mutate its shared CUDA allocation. + if not (is_weights_pre_processed() and getattr(self, "W_UV", None) is not None): + W_UK, W_UV = split_kv_b_proj( + self.kv_b_proj, + act_dtype, + self.kv_lora_rank, + self.num_local_heads, + self.qk_nope_head_dim, + self.v_head_dim, + ) + # (L, N, V) -> (N, L, V) + replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) + # (L, N, P) -> (N, P, L) + replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) quant_method = ( resolve_quant_method(self.quant_config, self, prefix=self.layer_name) From c0b6f74bee518aeaf5d72b200a0216e72fbb5e7c Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 17 Oct 2026 08:27:11 -0700 Subject: [PATCH 799/952] [Benchmark] Add per-case kernel profiles to vllm bench forward (#61692) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .agents/skills/model-microbenchmark/SKILL.md | 40 ++++ docs/cli/README.md | 2 +- tests/benchmarks/test_forward_cli.py | 37 +++ vllm/benchmarks/forward/capture.py | 18 ++ vllm/benchmarks/forward/main.py | 24 ++ vllm/benchmarks/forward/profile.py | 236 +++++++++++++++++++ 6 files changed, 356 insertions(+), 1 deletion(-) create mode 100644 vllm/benchmarks/forward/profile.py diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index 7f18c2bb870c..f630f698a141 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -51,6 +51,46 @@ Key flags (plus all `EngineArgs`): Read `median_us` (each replay is already the max over ranks). +## Kernel breakdown + +Add `--profile-dir DIR` to see where the time goes. After timing, each case is +replayed 3 more times under `torch.profiler`; every rank writes +`DIR/.rank.kernels.json` (GPU time per kernel per forward, sorted) +and `DIR/.rank.trace.json.gz` (open in Perfetto). Rank 0's top +kernels are printed and included in `--output-json`. + +```bash +python -m vllm.benchmarks.forward.profile DIR/decode.rank0.kernels.json +# Before/after a change, same case: per-kernel deltas, largest first. +python -m vllm.benchmarks.forward.profile base/decode.rank0.kernels.json new/decode.rank0.kernels.json +``` + +- Kernels are grouped by name without template or function arguments; the + JSON keeps one full name per group in `example`. +- Only the timed forward is included: the in-graph KV restore and rank + rendezvous are excluded (`pre_forward_us`). `span_us` (first to last + forward kernel) should be close to `median_us`, a bit higher from profiler + overhead; `busy_us` can exceed it when streams overlap. +- Compare the same rank across runs: under EP, ranks differ (e.g. collective + wait time shows up inside communication kernels). + +For Nsight Systems (NVTX, NCCL, CPU side, GPU metrics), skip `--profile-dir`: +the two can't share CUPTI. Each case's 3 profiled replays run between +`cudaProfilerStart`/`Stop` in an NVTX range `forward/`, so a capture range +records only them, one report per case (and rank): + +```bash +nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat \ + --cuda-graph-trace=node -o prof vllm bench forward ... # prof.1, prof.2, ... +torchrun --nproc-per-node=4 --no-python nsys profile ... -o prof_rank%q{RANK} \ + python -m vllm.entrypoints.cli.main bench forward ... +nsys stats --report cuda_gpu_kern_sum --format csv -o kern prof.1.nsys-rep +``` + +Keep `--capture-range-end=repeat`: nsys's default stops after the first case +and leaves the bench hanging. `nsys stats` also counts the KV restore copies +that the JSON summary excludes. + ## Caveats - **Dummy weights lie for data-dependent kernels.** Dummy gates route every diff --git a/docs/cli/README.md b/docs/cli/README.md index c2b27bed26fd..3cb6ae68d80f 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -136,7 +136,7 @@ vllm bench {forward, latency, serve, throughput} ### forward -Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; sampling and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP. CUDA only for now. +Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; sampling and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP. CUDA only for now. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case. ```bash vllm bench forward \ diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py index d44c6b946665..07b216e341cc 100644 --- a/tests/benchmarks/test_forward_cli.py +++ b/tests/benchmarks/test_forward_cli.py @@ -3,6 +3,7 @@ import json import subprocess from dataclasses import replace +from pathlib import Path from types import SimpleNamespace import numpy as np @@ -10,6 +11,7 @@ from vllm.benchmarks.forward.batch import build_batch from vllm.benchmarks.forward.main import run_case +from vllm.benchmarks.forward.profile import summarize from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case from vllm.platforms import current_platform @@ -97,6 +99,35 @@ def test_bench_forward_rejects_oversized_verify(capacity): run_case(runner, case, SimpleNamespace()) +def test_profile_summary(): + """Each replay's leading restore/rendezvous activities are left out, and + kernels are grouped by name without template or function arguments.""" + + def act(name, ts, dur, cat="kernel"): + return {"cat": cat, "name": name, "ts": ts, "dur": dur} + + replay = [ + ("Memcpy DtoD (Device -> Device)", "gpu_memcpy", 5), + ("void ns::gemm<128>(Params)", "kernel", 10), + ("void ns::gemm<64>(Params)", "kernel", 20), + ("std::enable_if<(1>0), void>::type ns::norm(float*)", "kernel", 4), + ] + events = [{"cat": "cpu_op", "name": "x", "ts": 0, "dur": 1}] + for r in range(2): + for i, (name, cat, dur) in enumerate(replay): + events.append(act(name, 1000 * r + 100 * i, dur, cat)) + summary = summarize(events, replays=2, skip=1) + assert summary["pre_forward_us"] == 5 + assert summary["busy_us"] == 34 + assert summary["span_us"] == 204 + assert [(k["name"], k["us"], k["calls"]) for k in summary["kernels"]] == [ + ("ns::gemm", 30, 2), + ("ns::norm", 4, 1), + ] + with pytest.raises(ValueError, match="replays"): + summarize(events[:-1], replays=2, skip=1) + + @pytest.mark.benchmark # TODO(woosuk): Support ROCm and re-enable this test there. @pytest.mark.skipif( @@ -126,6 +157,8 @@ def test_bench_forward(tmp_path): "2", "--output-json", str(output), + "--profile-dir", + str(tmp_path / "profile"), ] result = subprocess.run(command, capture_output=True, text=True) print(result.stdout) @@ -139,3 +172,7 @@ def test_bench_forward(tmp_path): assert len(rank0["batch_order"]) == len(row["requests"]) assert rank0["finite"] and rank0["graph_deterministic"] assert row["median_us"] > 0 + profile = rank0["profile"] + assert profile["top_kernels"] and profile["span_us"] > 0 + summary = json.loads(Path(profile["summary"]).read_text()) + assert Path(summary["trace"]).exists() diff --git a/vllm/benchmarks/forward/capture.py b/vllm/benchmarks/forward/capture.py index 8b378bb18cf5..8f4f3396ae45 100644 --- a/vllm/benchmarks/forward/capture.py +++ b/vllm/benchmarks/forward/capture.py @@ -55,6 +55,8 @@ def __init__( iters: int, samples: int, state_views: list[torch.Tensor], + name: str, + profile_prefix: str | None = None, ): self.runner = runner self.batch = batch @@ -62,6 +64,8 @@ def __init__( self.iters = iters self.samples = samples self.state_views = state_views + self.name = name + self.profile_prefix = profile_prefix self.phase = "warmup" self.result: dict[str, Any] = {} @@ -193,6 +197,20 @@ def restore() -> None: ), ) self.result["latency_us"] = self._time(graph) + from vllm.benchmarks.forward.profile import profile_graph + + # Each replay restores the state views, then (multi-rank) rendezvouses, + # before the timed forward. + skip = len(self.state_views) + (world.world_size > 1) + summary = profile_graph(graph, self.name, self.profile_prefix, skip) + if summary is not None: + self.result["profile"] = { + "summary": f"{self.profile_prefix}.kernels.json", + "span_us": summary["span_us"], + "busy_us": summary["busy_us"], + "activities": summary["activities"], + "top_kernels": summary["kernels"][:10], + } del graph return output diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py index d7841c641e0b..7525c93d9f64 100644 --- a/vllm/benchmarks/forward/main.py +++ b/vllm/benchmarks/forward/main.py @@ -74,6 +74,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None: parser.add_argument("--iters", type=int, default=20, help="Replays per sample.") parser.add_argument("--samples", type=int, default=5) parser.add_argument("--output-json", help="Write results to this file.") + parser.add_argument( + "--profile-dir", + help="After timing, trace a few replays of each case and write " + "`.rank.trace.json.gz` (Perfetto) and `.rank.kernels.json` " + "(GPU time per kernel per forward) here. Summarize or compare with " + "`python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]`.", + ) EngineArgs.add_cli_args(parser) parser.set_defaults(load_format="dummy", enable_prefix_caching=False) @@ -146,6 +153,13 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any dcp_size, ) zero_kv_blocks(context, kv_config.kv_cache_groups, batch, kv_config.num_blocks) + profile_prefix = None + if args.profile_dir: + from vllm.distributed.parallel_state import get_world_group + + os.makedirs(args.profile_dir, exist_ok=True) + rank = get_world_group().rank + profile_prefix = os.path.join(args.profile_dir, f"{case.name}.rank{rank}") bench = ForwardBench( runner, batch, @@ -154,6 +168,8 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any touched_kv_views( context, kv_config.kv_cache_groups, batch, kv_config.num_blocks, dcp_size ), + case.name, + profile_prefix, ) routing, balanced = install_routing( runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed @@ -221,6 +237,14 @@ def main(args: argparse.Namespace) -> None: f"eager rel rms {rank0['graph_vs_eager']['relative_rms']:.1e}", flush=True, ) + if profile := rank0.get("profile"): + from vllm.benchmarks.forward.profile import format_summary + + print( + f"[forward] {case.name} rank 0 kernels ({profile['summary']}):\n" + + format_summary({**profile, "kernels": profile["top_kernels"]}), + flush=True, + ) if latency["sample_spread_pct"] > 0.5: logger.warning( "%s: unstable timing (sample medians spread %.2f%%).", diff --git a/vllm/benchmarks/forward/profile.py b/vllm/benchmarks/forward/profile.py new file mode 100644 index 000000000000..d06b50aba0eb --- /dev/null +++ b/vllm/benchmarks/forward/profile.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU kernel profiles of `vllm bench forward` cases. + +With `--profile-dir`, each case's graph is replayed a few more times under +`torch.profiler` after timing. Each rank writes a Perfetto trace and a kernel +summary (`.rank.kernels.json`): per-kernel time per forward, sorted. + +Summarize one file or compare two (e.g. before/after a change): +`python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]`. + +The same replays run between `cudaProfilerStart`/`Stop` in an NVTX range per +case (with or without `--profile-dir`), so Nsight Systems can record just them: +`nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat +--cuda-graph-trace=node vllm bench forward ...`. +""" + +import argparse +import gzip +import json +import statistics +from collections import defaultdict +from typing import Any + +import regex as re +import torch + +_GPU_CATEGORIES = ("kernel", "gpu_memcpy", "gpu_memset") +# Kernels named after a generic launcher are identified by their first +# template argument instead. +_GENERIC = re.compile(r"(^|::)(device_kernel|Kernel|kernel|KernelWrapper)$") +# CuTe DSL kernel names embed per-shape layouts after these markers. +_CUTE_DSL_SUFFIX = re.compile(r"_object_at_|_tensorptr") + + +def _match_back(text: str, end: int, open_ch: str, close_ch: str) -> int: + """Index of the bracket opening the one closing at `end`, skipping + brackets of the other kind nested inside parentheses.""" + depth, parens = 0, 0 + for i in range(end, -1, -1): + ch = text[i] + if open_ch == "<" and ch in "()": + parens += 1 if ch == ")" else -1 + continue + if parens: + continue + if ch == close_ch: + depth += 1 + elif ch == open_ch: + depth -= 1 + if depth == 0: + return i + return -1 + + +def short_name(name: str) -> str: + """`void ns::foo(args)` -> `ns::foo`: return type, template and + function arguments dropped. Generic launchers (e.g. + `cutlass::device_kernel>`) keep their first template argument.""" + if name.startswith("kernel_cutlass_"): + return _CUTE_DSL_SUFFIX.split(name, maxsplit=1)[0] + head = name.strip() + if head.endswith(")"): + head = head[: max(_match_back(head, len(head) - 1, "(", ")"), 0)].rstrip() + template = "" + if head.endswith(">"): + start = _match_back(head, len(head) - 1, "<", ">") + if start > 0: + head, template = head[:start], head[start + 1 : -1] + base = head.split()[-1] if head.split() else name + if _GENERIC.search(base) and template: + arg = re.split(r"[<>(,]", template, maxsplit=1)[0].strip() + if arg: + return f"{base}<{arg.removeprefix('void ')}>" + return base + + +def summarize(events: list[dict], replays: int, skip: int) -> dict[str, Any]: + """Per-forward kernel times from the GPU activities of `replays` replays. + + Each replay first runs `skip` activities that aren't part of the timed + forward (the KV state restore and the rank rendezvous); they are reported + separately as `pre_forward_us`. + """ + acts = sorted( + (e for e in events if e.get("cat") in _GPU_CATEGORIES and "dur" in e), + key=lambda e: e["ts"], + ) + if not acts: + raise ValueError( + "No GPU activity recorded. Another profiler (e.g. nsys) may hold " + "CUPTI; don't combine --profile-dir with it." + ) + if len(acts) % replays: + raise ValueError( + f"{len(acts)} GPU activities don't split into {replays} replays" + ) + per_replay = len(acts) // replays + if per_replay <= skip: + raise ValueError(f"{per_replay} activities per replay, {skip} to skip") + kernels: dict[str, dict[str, Any]] = defaultdict( + lambda: {"total_us": 0.0, "calls": 0, "example": ""} + ) + spans, busy, pre = [], [], [] + for r in range(replays): + chunk = acts[r * per_replay : (r + 1) * per_replay] + pre.append(sum(e["dur"] for e in chunk[:skip])) + forward = chunk[skip:] + spans.append( + max(e["ts"] + e["dur"] for e in forward) - min(e["ts"] for e in forward) + ) + busy.append(sum(e["dur"] for e in forward)) + for e in forward: + name = short_name(e["name"]) if e["cat"] == "kernel" else e["name"] + k = kernels[name] + k["total_us"] += e["dur"] + k["calls"] += 1 + k["example"] = k["example"] or e["name"][:300] + total = sum(k["total_us"] for k in kernels.values()) + rows = [ + { + "name": name, + "us": round(k["total_us"] / replays, 2), + "pct": round(100 * k["total_us"] / total, 2), + "calls": k["calls"] // replays, + "example": k["example"], + } + for name, k in kernels.items() + ] + rows.sort(key=lambda row: -row["us"]) + return { + "replays": replays, + # First to last forward activity; busy can exceed it when streams overlap. + "span_us": round(statistics.median(spans), 2), + "busy_us": round(statistics.median(busy), 2), + "pre_forward_us": round(statistics.median(pre), 2), + "activities": per_replay - skip, + "kernels": rows, + } + + +def profile_graph( + graph: torch.cuda.CUDAGraph, + name: str, + prefix: str | None, + skip: int, + replays: int = 3, +) -> dict[str, Any] | None: + """Replay `replays` times between `cudaProfilerStart`/`Stop`, in an NVTX + range named after the case, for `nsys --capture-range=cudaProfilerApi`. + + With `prefix`, also trace them with `torch.profiler` and write + `prefix.trace.json.gz` and `prefix.kernels.json`. Every rank must call + this (replays may contain collectives). + """ + + def run() -> None: + torch.cuda.profiler.start() + with torch.cuda.nvtx.range(f"forward/{name}"): + for _ in range(replays): + graph.replay() + torch.accelerator.synchronize() + torch.cuda.profiler.stop() + + if prefix is None: + run() + return None + activities = [torch.profiler.ProfilerActivity.CUDA] + with torch.profiler.profile(activities=activities) as prof: + run() + trace = f"{prefix}.trace.json.gz" + prof.export_chrome_trace(trace) + with gzip.open(trace, "rt") as f: + events = json.load(f)["traceEvents"] + summary = {"trace": trace, **summarize(events, replays, skip)} + with open(f"{prefix}.kernels.json", "w") as f: + json.dump(summary, f, indent=1) + return summary + + +def format_summary(summary: dict[str, Any], top: int = 20) -> str: + lines = [ + f"span {summary['span_us']:.1f} us, busy {summary['busy_us']:.1f} us, " + f"{summary['activities']} activities/forward", + f"{'us':>10} {'%':>6} {'calls':>6} kernel", + ] + for row in summary["kernels"][:top]: + lines.append( + f"{row['us']:10.1f} {row['pct']:6.2f} {row['calls']:6d} {row['name']}" + ) + return "\n".join(lines) + + +def format_diff(a: dict[str, Any], b: dict[str, Any], top: int = 20) -> str: + rows_a = {row["name"]: row for row in a["kernels"]} + rows_b = {row["name"]: row for row in b["kernels"]} + deltas = [] + for name in rows_a.keys() | rows_b.keys(): + ra, rb = rows_a.get(name), rows_b.get(name) + us_a, us_b = ra["us"] if ra else 0.0, rb["us"] if rb else 0.0 + calls_a, calls_b = ra["calls"] if ra else 0, rb["calls"] if rb else 0 + deltas.append((us_b - us_a, us_a, us_b, calls_a, calls_b, name)) + deltas.sort(key=lambda d: -abs(d[0])) + lines = [ + f"span {a['span_us']:.1f} -> {b['span_us']:.1f} us " + f"({b['span_us'] - a['span_us']:+.1f}), busy {a['busy_us']:.1f} -> " + f"{b['busy_us']:.1f} us ({b['busy_us'] - a['busy_us']:+.1f})", + f"{'delta us':>10} {'A us':>10} {'B us':>10} {'A n':>5} {'B n':>5} kernel", + ] + for delta, us_a, us_b, calls_a, calls_b, name in deltas[:top]: + lines.append( + f"{delta:+10.1f} {us_a:10.1f} {us_b:10.1f} {calls_a:5d} {calls_b:5d} " + f"{name}" + ) + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("summaries", nargs="+", help="One or two *.kernels.json.") + parser.add_argument("--top", type=int, default=20) + args = parser.parse_args() + if len(args.summaries) > 2: + parser.error("Pass one summary, or two to compare.") + loaded = [] + for path in args.summaries: + with open(path) as f: + loaded.append(json.load(f)) + if len(loaded) == 1: + print(format_summary(loaded[0], args.top)) + else: + print(format_diff(loaded[0], loaded[1], args.top)) + + +if __name__ == "__main__": + main() From e744eacf682f2b1c7f91d2dcf348a015046e56b9 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:10:40 -0700 Subject: [PATCH 800/952] [Rust Frontend] Validate bad_words limits during request lowering (#61933) Signed-off-by: Alec Flowers Co-authored-by: Claude Opus 5.5 --- rust/src/text/src/lower.rs | 41 ++++++++++++++++++++++++++++ rust/src/text/src/lower/token_ids.rs | 35 ++++++++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 9fa8c78f4f91..319c9953698a 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -1338,6 +1338,47 @@ mod tests { )); } + #[test] + fn lower_sampling_params_enforces_bad_words_limits() { + let tokenizer = TestTokenizer::new() + .with_regular_token("blocked", 300) + .with_regular_token(" blocked", 301); + let lower = |bad_words: Option>, bad_words_token_ids: Vec>| { + lower_sampling_params( + SamplingParams { + bad_words, + bad_words_token_ids: Some(bad_words_token_ids), + ..Default::default() + }, + SamplingHints::default(), + sample_sampling_limits(), + 3, + &tokenizer, + ) + }; + + lower(None, vec![vec![7; 8]; 128]).unwrap(); + for (bad_words, raw, expected) in [ + ( + None, + vec![vec![7]; 129], + "Too many bad words: 129. The max number is 128.", + ), + ( + Some(vec!["blocked".to_string()]), + vec![vec![7]; 127], + "Too many bad words: 129. The max number is 128.", + ), + ( + None, + vec![vec![7; 1025]], + "Too many total bad word tokens: 1025. The max is 1024.", + ), + ] { + assert_eq!(lower(bad_words, raw).unwrap_err().to_string(), expected); + } + } + #[test] fn lower_sampling_params_rejects_out_of_vocab_logit_bias() { let error = lower_sampling_params_with_limits( diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index 51329f9290ed..b7f160b919a4 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -8,12 +8,21 @@ use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::SamplingLimits; +const MAX_NUM_BAD_WORDS_ENV: &str = "VLLM_MAX_NUM_BAD_WORDS"; +const DEFAULT_MAX_NUM_BAD_WORDS: usize = 128; +const MAX_BAD_WORDS_TOTAL_TOKENS_ENV: &str = "VLLM_MAX_BAD_WORDS_TOTAL_TOKENS"; +const DEFAULT_MAX_BAD_WORDS_TOTAL_TOKENS: usize = 1024; + #[derive(Debug, Error)] pub enum TokenIdsError { #[error("allowed_token_ids should not be empty")] EmptyAllowedTokenIds, #[error("bad_words_token_ids must not contain empty sequences")] EmptyBadWordSequence, + #[error("Too many bad words: {count}. The max number is {max}.")] + TooManyBadWords { count: usize, max: usize }, + #[error("Too many total bad word tokens: {count}. The max is {max}.")] + TooManyBadWordTokens { count: usize, max: usize }, #[error( "token_id(s) {token_ids:?} in {parameter} are out of vocabulary. \ Vocabulary size: {vocab_size}" @@ -45,6 +54,31 @@ fn validate_param( }) } +fn env_limit(name: &str, default: usize) -> usize { + std::env::var(name).ok().and_then(|value| value.parse().ok()).unwrap_or(default) +} + +/// Mirrors Model Runner V2's per-request bad-words buffer limits: +/// +fn validate_bad_words_limits(bad_words_token_ids: &[Vec]) -> Result<(), TokenIdsError> { + let count = bad_words_token_ids.len(); + let max = env_limit(MAX_NUM_BAD_WORDS_ENV, DEFAULT_MAX_NUM_BAD_WORDS); + if count > max { + return Err(TokenIdsError::TooManyBadWords { count, max }); + } + + let count = bad_words_token_ids.iter().map(Vec::len).sum(); + let max = env_limit( + MAX_BAD_WORDS_TOTAL_TOKENS_ENV, + DEFAULT_MAX_BAD_WORDS_TOTAL_TOKENS, + ); + if count > max { + return Err(TokenIdsError::TooManyBadWordTokens { count, max }); + } + + Ok(()) +} + /// Validate that pre-tokenized prompt IDs are within the engine-visible prompt /// vocabulary range. pub(crate) fn validate_prompt_token_ids( @@ -101,6 +135,7 @@ pub(crate) fn validate_vocab_range( if bad_words_token_ids.iter().any(Vec::is_empty) { return Err(TokenIdsError::EmptyBadWordSequence); } + validate_bad_words_limits(bad_words_token_ids)?; validate_param( "bad_words", bad_words_token_ids.iter().flatten().copied(), From 763e3f7c911308fa609fb9437642929efe955798 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sun, 18 Oct 2026 00:41:04 +0800 Subject: [PATCH 801/952] [Bugfix][Spec Decode] Disable dynamic speculative decoding for fixed-K proposers instead of crashing (#46399) Signed-off-by: Ting Sun Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/config/vllm.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53b9784dadb1..050cc561194c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1056,6 +1056,28 @@ def _maybe_override_dynamic_sd_cudagraph_mode(self) -> None: ) self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + def _maybe_disable_dynamic_sd_for_v1(self) -> None: + speculative_config = self.speculative_config + # V2 truncates drafts centrally. V1 requires proposer support. + if ( + speculative_config is None + or not speculative_config.uses_dynamic_speculative_decoding() + or self.use_v2_model_runner + or speculative_config.method + in ("ngram", "eagle", "eagle3", "mtp", "draft_model", "dflash") + ): + return + + logger.warning_once( + "Dynamic speculative decoding is not supported with the '%s' " + "speculative method on Model Runner V1. " + "Disabling num_speculative_tokens_per_batch_size " + "and falling back to static num_speculative_tokens=%d.", + speculative_config.method, + speculative_config.num_speculative_tokens, + ) + speculative_config.num_speculative_tokens_per_batch_size = None + def _maybe_disable_dynamic_sd_for_data_parallel(self) -> None: speculative_config = self.speculative_config if ( @@ -1829,6 +1851,7 @@ def has_blocked_weights(): ) self._maybe_disable_dynamic_sd_for_data_parallel() + self._maybe_disable_dynamic_sd_for_v1() self._maybe_override_dynamic_sd_cudagraph_mode() if ( From 0c5dd18e7a0ace18d99438fb16b8375572ea76b5 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sat, 17 Oct 2026 10:27:57 -0700 Subject: [PATCH 802/952] [MRV2] Collect KV transfer results before the drafter runs (#62041) Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 (1M context) --- tests/v1/worker/test_gpu_kv_connector.py | 28 +++++++++++++++++++++--- vllm/v1/worker/gpu/kv_connector.py | 23 +++++++++++++++---- vllm/v1/worker/gpu/model_runner.py | 11 +++++++--- 3 files changed, 52 insertions(+), 10 deletions(-) diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index 78fe3c185bd9..dc46c22c7cd1 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/tests/v1/worker/test_gpu_kv_connector.py @@ -22,7 +22,12 @@ def _make_connector( backend.bind_connector_metadata.side_effect = lambda _: events.append("bind") backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start") backend.wait_for_save.side_effect = lambda: events.append("wait") - backend.get_transfer_results.return_value = KVConnectorTransferResults() + + def get_transfer_results(_finished_req_ids): + events.append("results") + return KVConnectorTransferResults() + + backend.get_transfer_results.side_effect = get_transfer_results backend.get_block_ids_with_load_errors.return_value = set() backend.get_kv_connector_stats.return_value = None backend.get_kv_connector_kv_cache_events.return_value = None @@ -73,7 +78,7 @@ def test_load_start_phase( ) connector.post_forward(set()) - assert events == ["handle", "bind", "start", "wait", "clear"] + assert events == ["handle", "bind", "start", "wait", "results", "clear"] kwargs = connector.kv_connector.start_load_kv.call_args.kwargs assert kwargs["request_state_indices"] is request_indices @@ -94,4 +99,21 @@ def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch): connector.no_forward(_scheduler_output(False)) - assert events == ["handle", "bind", "start", "wait", "clear"] + assert events == ["handle", "bind", "start", "wait", "results", "clear"] + + +def test_deferred_finalize_collects_results_before_saving( + monkeypatch: pytest.MonkeyPatch, +): + """With a drafter, the transfer results are collected right after the + target forward, and the saves and metadata are finalized after the + drafter writes its KV.""" + events: list[str] = [] + connector = _make_connector(monkeypatch, events) + connector.pre_forward(_scheduler_output(False)) + + connector.post_forward(set(), defer_finalize=True) + assert events == ["handle", "bind", "start", "results"] + + connector.finalize_step() + assert events == ["handle", "bind", "start", "results", "wait", "clear"] diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index e9b71dae3e8c..9edcc2831a8a 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -41,9 +41,14 @@ def pre_forward( def finish_forward(self) -> None: pass - def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: + def post_forward( + self, finished_req_ids: set[str], defer_finalize: bool = False + ) -> KVConnectorOutput | None: return None + def finalize_step(self) -> None: + pass + def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: return EMPTY_MODEL_RUNNER_OUTPUT @@ -128,7 +133,9 @@ def finish_forward(self) -> None: def reset_capture_state(self) -> None: self.kv_connector.reset_capture_state() - def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: + def post_forward( + self, finished_req_ids: set[str], defer_finalize: bool = False + ) -> KVConnectorOutput | None: if self._disabled: return None @@ -136,7 +143,8 @@ def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: self._start_load_kv() output = KVConnectorOutput() - self.kv_connector.wait_for_save() + if not defer_finalize: + self.kv_connector.wait_for_save() transfer_results = self.kv_connector.get_transfer_results(finished_req_ids) output.finished_sending = transfer_results.finished_sending or None output.finished_recving = transfer_results.finished_recving or None @@ -147,9 +155,16 @@ def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None: output.kv_connector_worker_meta = ( self.kv_connector.build_connector_worker_meta() ) - self.kv_connector.clear_connector_metadata() + if not defer_finalize: + self.kv_connector.clear_connector_metadata() return output + def finalize_step(self) -> None: + if self._disabled: + return + self.kv_connector.wait_for_save() + self.kv_connector.clear_connector_metadata() + def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput: if self._disabled: return EMPTY_MODEL_RUNNER_OUTPUT diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1b75d7349395..b8ad557b68a8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2207,6 +2207,12 @@ def sample_tokens( pending_aux_output=pending_aux_output, ) + # Collect KV transfer results now; finalize saves after the drafter runs. + defer_kv_finalize = self.speculator is not None + model_runner_output.kv_connector_output = self.kv_connector.post_forward( + finished_req_ids, defer_finalize=defer_kv_finalize + ) + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None # The encoder runner exists only on the first PP rank, so later ranks # have no cached embeddings to gather. @@ -2288,9 +2294,8 @@ def sample_tokens( self.req_states.draft_tokens, input_batch ) - # Post-step KV connector related operations. - kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - model_runner_output.kv_connector_output = kv_connector_output + if defer_kv_finalize: + self.kv_connector.finalize_step() model_runner_output.ec_connector_output = ec_connector_output return async_output From bafb263664e87daea420f47bd07f4831b7b31bcc Mon Sep 17 00:00:00 2001 From: Tyrone <71038642+TyroneNel@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:47:52 +0200 Subject: [PATCH 803/952] [Bugfix][Benchmarks] Name the real cause when the /tokenize alignment probe fails (#59888) --- vllm/benchmarks/serve.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 1a1107d3e406..4c7f4a2facb3 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -122,8 +122,22 @@ async def _detokenize(tokens: list[int]) -> str: try: first_tokens = await _tokenize(input_requests[0].prompt) - except Exception: - print("WARNING: /tokenize unavailable, skipping alignment.") + except aiohttp.ClientConnectionError as e: + print(f"WARNING: {base_url} unreachable ({e!r}), skipping alignment.") + return input_requests + except aiohttp.ClientResponseError as e: + if e.status == 404: + hint = ( + "404 Not Found: either this server has no /tokenize route," + f" or it does not serve a model named `{model_id}`" + " (its served names are listed by GET /v1/models)" + ) + else: + hint = f"HTTP {e.status}" + print(f"WARNING: /tokenize unavailable ({hint}), skipping alignment.") + return input_requests + except Exception as e: + print(f"WARNING: /tokenize probe failed ({e!r}), skipping alignment.") return input_requests expected = input_requests[0].prompt_len From 5c1c903501927ea7c6c699aabd4c271015e66e71 Mon Sep 17 00:00:00 2001 From: Afonso Ribeiro Date: Sat, 17 Oct 2026 19:49:06 +0100 Subject: [PATCH 804/952] [Perf][Parser] Stream Qwen3 tool argument values incrementally (12 KB call: 2.1 s to 11 ms) (#62027) Signed-off-by: Afonso Ribeiro Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/test_qwen3.py | 131 ++++++++++++++++++++++++++++++ vllm/parser/mimo.py | 5 ++ vllm/parser/qwen3.py | 60 ++++++++++++++ 3 files changed, 196 insertions(+) diff --git a/tests/parser/engine/test_qwen3.py b/tests/parser/engine/test_qwen3.py index fda047d42245..be4b33e3861f 100644 --- a/tests/parser/engine/test_qwen3.py +++ b/tests/parser/engine/test_qwen3.py @@ -6,6 +6,7 @@ Qwen3 XML-style tool calls. """ +import dataclasses import json from unittest.mock import MagicMock @@ -808,6 +809,136 @@ def test_partial_value_with_angle_bracket_and_complete_param(self): assert result == {"city": "Tokyo", "expr": "x<5"} +WRITE_CALL_HEAD = ( + "\n\n" + "\n/repo/app.py\n\n" + "\n644\n\n" + "\n" +) +WRITE_CALL_TAIL = "\n\n\n" +# Escapes, non-ASCII, blank lines and tag-like text that must stay inside +# the value, including prefixes of the closing and next-parameter tags. +CONTENT_LINE = 'if a < b > c: print("x\\y") #

int: + calls = 0 + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + calls += partial + return converter(raw_args, partial) + + parser = ParserEngine( + mock_tokenizer, + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + text = WRITE_CALL_HEAD + CONTENT_LINE * lines + WRITE_CALL_TAIL + simulate_tool_streaming( + parser, + mock_request, + [text[i : i + 4] for i in range(0, len(text), 4)], + ) + return calls + + assert streaming_conversions(lines=20) == streaming_conversions(lines=200) + + def test_extender_follows_newline_trimming_and_split_end(self): + from vllm.parser.qwen3 import _qwen3_arg_value_extender + + extend = _qwen3_arg_value_extender("") + assert extend is not None + deltas = ["\n", 'a"\n', "b", "\n", ""] + expected = ["", 'a\\"', "\\nb", "", "\\nx") + assert extend is not None + assert [extend(d) for d in ["< para", "meter ="]] == ["< para", None] + + @pytest.mark.parametrize( + "raw_args", + [ + "x", + "123", + ], + ids=["closed_value", "open_name", "repeated_key_not_last"], + ) + def test_no_extender_unless_open_value_is_last(self, raw_args): + from vllm.parser.qwen3 import _qwen3_arg_value_extender + + assert _qwen3_arg_value_extender(raw_args) is None + + class TestSchemaAwareTypeCoercion: """Verify that _fix_arg_types corrects miscoerced values using the tool schema.""" diff --git a/vllm/parser/mimo.py b/vllm/parser/mimo.py index b1c37dd74644..3704907efe56 100644 --- a/vllm/parser/mimo.py +++ b/vllm/parser/mimo.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MiMo's compact XML parameters contain verbatim string values.""" +import functools import json from dataclasses import replace @@ -10,6 +11,7 @@ _PARTIAL_PARAM_RE, CHATML_TURN_BOUNDARIES, Qwen3Parser, + _param_value_extender, qwen3_config, ) @@ -37,6 +39,9 @@ def __init__(self, tokenizer, tools=None, **kwargs): turn_boundary_tokens=CHATML_TURN_BOUNDARIES, ), arg_converter=_mimo_arg_converter, + arg_value_extender=functools.partial( + _param_value_extender, trim_newlines=False + ), ), ) super().__init__(tokenizer, tools, **kwargs) diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index d9a71f2a46e2..577e01edf6ed 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -25,6 +25,7 @@ from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -86,6 +87,64 @@ def _qwen3_arg_converter(raw_args: str, partial: bool) -> str: return json.dumps(params, ensure_ascii=False) +# What ends an open value in ``_PARAM_RE``: a closing tag or the next opener. +_PARAM_VALUE_END_RE = re.compile(r"<\s*/\s*parameter\s*>|<\s*parameter\s*=") + + +def _param_value_end_window(text: str) -> str: + """Return the suffix of *text* that may still grow into a value end.""" + start = text.rfind("<") + if start == -1: + return "" + tail = text[start:] + return tail if _PARAM_VALUE_END_RE.fullmatch(tail, partial=True) else "" + + +def _param_value_extender( + raw_args: str, trim_newlines: bool +) -> ArgValueExtender | None: + """Extend the open trailing ```` value of *raw_args*. + + The extender streams exactly what the converter would for the same text. + With *trim_newlines* it follows ``_trim_wrapping_newlines``: it drops a + leading newline of the value and holds back a trailing one. + """ + match = _PARTIAL_PARAM_RE.search(_PARAM_RE.sub("", raw_args)) + if match is None: + return None + name, value = match.group(1), match.group(2) + # A repeated key keeps its first position, so its value is not last. + if any(m.group(1) == name for m in _PARAM_RE.finditer(raw_args)): + return None + + window = _param_value_end_window(value) + started = bool(value) + held = "\n" if trim_newlines and value.endswith("\n") and value != "\n" else "" + + def extend(delta: str) -> str | None: + nonlocal window, started, held + text = window + delta + if _PARAM_VALUE_END_RE.search(text): + return None + window = _param_value_end_window(text) + if trim_newlines: + text = held + delta + if not started and text: + started = True + text = text.removeprefix("\n") + held = "\n" if text.endswith("\n") else "" + text = text[: len(text) - len(held)] + else: + text = delta + return json.dumps(text, ensure_ascii=False)[1:-1] + + return extend + + +def _qwen3_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + return _param_value_extender(raw_args, trim_newlines=True) + + @functools.cache def qwen3_config( thinking: bool = True, @@ -223,6 +282,7 @@ def qwen3_config( ), }, arg_converter=_qwen3_arg_converter, + arg_value_extender=_qwen3_arg_value_extender, stream_arg_deltas=True, strip_trailing_reasoning_whitespace=False, tool_args_json=False, From 03e8300d2cac77b24f7ab6013a803369ee64c19e Mon Sep 17 00:00:00 2001 From: yifanFengg Date: Sat, 17 Oct 2026 11:53:06 -0700 Subject: [PATCH 805/952] [Bugfix] Enable inductor batch-invariant mode under VLLM_BATCH_INVARIANT (#61118) Signed-off-by: Yifan Feng Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- vllm/model_executor/determinism/batch_invariant.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index e4b04793feed..933e5ecef6ea 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -1424,6 +1424,9 @@ def override_envs_for_invariance(): # torch.compile settings os.environ["VLLM_USE_AOT_COMPILE"] = "0" + import torch._inductor.config as inductor_config + + inductor_config.batch_invariant = True def init_batch_invariance(): From 255f14599ec0544c35539414810abdfa7e4e1c86 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 17 Oct 2026 11:56:28 -0700 Subject: [PATCH 806/952] [Bugfix][Rust Frontend] Fix macOS build: read inherited listener family from its address (#62177) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- rust/src/engine-core-client/src/transport.rs | 29 +++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index e453cafdc8fb..dc1c743ede6c 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -405,7 +405,9 @@ fn adopt_listener(fd: RawFd) -> Result { // transferring ownership to the listener constructed below. let fd = unsafe { OwnedFd::from_raw_fd(fd) }; let socket: socket2::Socket = fd.into(); - let domain = socket.domain()?; + // `Socket::domain` (SO_DOMAIN) exists only on Linux, Android, FreeBSD and + // Fuchsia; the bound address carries the family on every platform. + let domain = socket.local_addr()?.domain(); if domain == socket2::Domain::IPV4 || domain == socket2::Domain::IPV6 { let listener: std::net::TcpListener = socket.into(); @@ -672,4 +674,29 @@ mod tests { assert_eq!(endpoint.to_string(), format!("tcp://{expected}")); assert!(router.close().await.is_empty()); } + + #[tokio::test] + async fn inherited_unix_listener_is_adopted_without_rebinding() { + let path = std::env::temp_dir().join(format!( + "vllm-inherited-listener-{}-{}.sock", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("clock after epoch") + .as_nanos() + )); + let listener = std::os::unix::net::UnixListener::bind(&path).expect("bind listener"); + let mut router = RouterSocket::new(); + + let endpoint = router + .bind_listener( + adopt_listener(listener.into_raw_fd()).expect("adopt inherited listener fd"), + ) + .await + .expect("bind inherited listener"); + + assert!(endpoint.to_string().starts_with("ipc://"), "{endpoint}"); + assert!(router.close().await.is_empty()); + let _ = std::fs::remove_file(&path); + } } From 7879f3687b5b6a4be4754242e70b9c21cf7b4183 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Sat, 17 Oct 2026 14:24:35 -0700 Subject: [PATCH 807/952] [NIXL] Drop unreachable producer-side notif counting in push mode (#62210) Signed-off-by: David Cheung Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../unit/test_nixl_push_connector.py | 21 ++---- .../kv_connector/v1/nixl/push_worker.py | 64 ++++++------------- 2 files changed, 26 insertions(+), 59 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index d92857c34186..e2f9e99b494f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -929,14 +929,8 @@ def test_get_new_notifs_processes_forwarded_completion_notif(self): notif_msg = f"{request_id}:1".encode() w._pending_completion_notifs.put(notif_msg) - # transfer_topo is consulted only for the producer-side path; we - # make it a MagicMock because the D-side branch returns early. - w.transfer_topo = MagicMock() - - notified = w._get_new_notifs() + w._get_new_notifs() - # Notif consumed; D-side just touches _recving_transfers. - assert notified == set() assert request_id in w._recving_transfers def test_get_transfer_results_evicts_completed_state(self): @@ -1352,12 +1346,10 @@ def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), ) w = _StubWriterWorker.fresh() - w.transfer_topo = MagicMock() # Forward a completion notif for an unknown request_id. w._pending_completion_notifs.put(b"never-heard-of-you:1") - notified = w._get_new_notifs() - assert notified == set() + w._get_new_notifs() # Did not register anywhere. assert "never-heard-of-you" not in w._recving_transfers assert w.xfer_stats.data["num_notifications_after_expiry"] == [1] @@ -1381,7 +1373,6 @@ def test_get_new_notifs_extends_lease_on_heartbeat(self): leases of tracked P-side requests on the engine main thread, and ignore request IDs that aren't being tracked.""" w = _StubWriterWorker.fresh() - w.transfer_topo = MagicMock() # _handle_heartbeat reads ``self._lease_extension`` (set in the # real ``__init__``). w._lease_extension = 10 @@ -1395,8 +1386,7 @@ def test_get_new_notifs_extends_lease_on_heartbeat(self): # and another tracked one. w._pending_completion_notifs.put(b"HB:req-a,req-unknown,req-b") - notified = w._get_new_notifs() - assert notified == set() + w._get_new_notifs() # Tracked leases were renewed strictly forward in time. now = time.perf_counter() @@ -1421,20 +1411,19 @@ def test_completion_waits_for_one_notif_per_pp_stage(self): """Each PP stage WRITEs its own layers and sends one notif; D must collect pp_size notifs before reporting the recv done.""" w = _StubWriterWorker.fresh() - w.transfer_topo = MagicMock() request_id = "req-pp-2" w._recving_metadata[request_id] = MagicMock(pp_size=2) notif = f"{request_id}:1".encode() # First stage: counted, not yet done. w._pending_completion_notifs.put(notif) - assert w._get_new_notifs() == set() + w._get_new_notifs() assert request_id not in w._recving_transfers assert w.consumer_notification_counts_by_req[request_id] == 1 # Second (final) stage: now reported done. w._pending_completion_notifs.put(notif) - assert w._get_new_notifs() == set() + w._get_new_notifs() assert request_id in w._recving_transfers assert request_id not in w.consumer_notification_counts_by_req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 6d82fbcc0ac7..38e7d0fbb78d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -748,11 +748,9 @@ def _get_new_notifs(self) -> set[str]: """Drain HB / completion notifs forwarded by the writer thread. The writer owns ``nixl_wrapper.get_new_notifs`` for push; PUSH_REG - notifs are handled there. Everything else is forwarded here for - existing accounting. + notifs are handled there. P learns of finished pushes from its WRITE + handles, so this never reports finished sends. """ - assert self.transfer_topo is not None - notified_req_ids: set[str] = set() while True: try: notif = self._pending_completion_notifs.get_nowait() @@ -766,46 +764,27 @@ def _get_new_notifs(self) -> set[str]: req_id, tp_size = msg.rsplit(":", 1) - # Not tracked as a P-side send/process for this notif. - if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process: - if (meta := self._recving_metadata.get(req_id)) is not None: - # Consumer waits for one notif per producer rank writing - # here: pp_size stages * producers-per-consumer (>1 when - # producer TP > consumer TP; tp_size is the producer TP). - producers_per_consumer = max(1, int(tp_size) // self.world_size) - expected_notifs = meta.pp_size * producers_per_consumer - self.consumer_notification_counts_by_req[req_id] += 1 - notifs = self.consumer_notification_counts_by_req[req_id] - if notifs < expected_notifs: - continue - del self.consumer_notification_counts_by_req[req_id] - # P drove the transfer (we own no NIXL handle), so - # materialise an empty ``_recving_transfers`` entry for - # ``_pop_done_transfers`` to report done. - self._recving_transfers.setdefault(req_id, []) - else: - # Not tracked on either side (lease may have expired - # before the notif arrived). Log and skip. - self.xfer_stats.record_notification_after_expiry() - logger.error( - "Unrecognized request %s notif (may have expired).", - req_id, - ) + if (meta := self._recving_metadata.get(req_id)) is None: + # The lease may have expired before the notif arrived. + self.xfer_stats.record_notification_after_expiry() + logger.error( + "Unrecognized request %s notif (may have expired).", req_id + ) continue - - n_consumers = int(tp_size) - tp_ratio = self.transfer_topo.tp_ratio(n_consumers) - consumers_per_producer = -tp_ratio if n_consumers > self.world_size else 1 + # Consumer waits for one notif per producer rank writing + # here: pp_size stages * producers-per-consumer (>1 when + # producer TP > consumer TP; tp_size is the producer TP). + producers_per_consumer = max(1, int(tp_size) // self.world_size) + expected_notifs = meta.pp_size * producers_per_consumer self.consumer_notification_counts_by_req[req_id] += 1 - if ( - self.consumer_notification_counts_by_req[req_id] - == consumers_per_producer - ): - notified_req_ids.add(req_id) - del self.consumer_notification_counts_by_req[req_id] - self._reqs_to_process.remove(req_id) - self._reqs_to_send.pop(req_id, None) - return notified_req_ids + if self.consumer_notification_counts_by_req[req_id] < expected_notifs: + continue + del self.consumer_notification_counts_by_req[req_id] + # P drove the transfer (we own no NIXL handle), so + # materialise an empty ``_recving_transfers`` entry for + # ``_pop_done_transfers`` to report done. + self._recving_transfers.setdefault(req_id, []) + return set() def get_transfer_results(self) -> KVConnectorTransferResults: # Engine main thread asking for completions: also wake the writer @@ -835,7 +814,6 @@ def get_transfer_results(self) -> KVConnectorTransferResults: for req_id in done_pushing: self._reqs_to_send.pop(req_id, None) self._reqs_to_process.discard(req_id) - self.consumer_notification_counts_by_req.pop(req_id, None) done_sending.add(req_id) # Tell the writer to drop any state it still holds for any From 9c778e9ff8d0c78a241d7f530c441c2dc1d67533 Mon Sep 17 00:00:00 2001 From: djramic Date: Sun, 18 Oct 2026 01:18:02 +0200 Subject: [PATCH 808/952] [Bugfix][EPLB] Wait for pending expert buffer copies before NIXL READs (#62081) Signed-off-by: Djordje Ramic Co-authored-by: Andreas Karatzas --- vllm/distributed/eplb/eplb_communicator.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 473058d5258e..409a031c03dc 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -354,6 +354,7 @@ def __init__( # Per-rank expert_id -> physical row; set by set_transfer_context. self._expert_to_src_row: list[dict[int, int]] | None = None self._layer_idx: int | None = None + self.set_stream(None) nixl_agent_config = nixl_utils.nixl_agent_config config = ( @@ -403,7 +404,7 @@ def _make_agent_name(self) -> str: return f"eplb-{self._rank}{pp_suffix}-{uid}" def set_stream(self, stream: torch.Stream | None) -> None: - pass + self._stream = stream def add_send( self, @@ -421,6 +422,13 @@ def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: f"pending transfers from layer {self._layer_idx}; " f"execute() was not called after previous add_recv() calls" ) + # NIXL READs bypass device streams; wait for pending copies out of + # the receive buffer before add_recv overwrites it. + if current_platform.is_rocm(): + with gpu_sync_allowed(): + torch.accelerator.current_stream().synchronize() + if self._stream is not None: + self._stream.synchronize() self._layer_idx = layer_idx n = self._num_local_experts rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n) From 03ee2e134d6cc60c4cc84096c8492fcc971f0526 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 17 Oct 2026 18:10:08 -0700 Subject: [PATCH 809/952] [CI][ROCm] Fix order-dependent test_rocm_aiter_sampler_defers_import (CPU tensors hit Triton) (#62213) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- tests/v1/sample/test_topk_topp_sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 0cd44fe213bb..93b71a5d9503 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -110,8 +110,8 @@ def guard_aiter_sampling_import(name, *args, **kwargs): monkeypatch.setattr("builtins.__import__", guard_aiter_sampling_import) sampler = topk_topp_sampler.TopKTopPSampler() - logits = torch.randn(2, 8) - k = torch.full((2,), 2, dtype=torch.int32) + logits = torch.randn(2, 8, device=DEVICE_TYPE) + k = torch.full((2,), 2, dtype=torch.int32, device=DEVICE_TYPE) generators = {0: torch.Generator(device=logits.device).manual_seed(0)} token_ids, logits_to_return = sampler(logits, generators, k, None) From b2456c622b48ae85af5e55406fb138d9a23364be Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Sat, 17 Oct 2026 21:43:59 -0500 Subject: [PATCH 810/952] [ROCm][CI] Make ROCm 10.0 the default image; keep ROCm 7.2 as -rocm72 (#58761) Signed-off-by: Rohan Potdar Signed-off-by: Rohan138 Signed-off-by: Micah Williamson Co-authored-by: Claude Opus 5.5 Co-authored-by: Micah Williamson Co-authored-by: Kevin H. Luu Co-authored-by: Andreas Karatzas --- .buildkite/release-pipeline.yaml | 667 ++++++------------ .buildkite/scripts/annotate-rocm-release.sh | 108 ++- .buildkite/scripts/cache-rocm-base-wheels.sh | 62 +- .buildkite/scripts/ci-bake-rocm.sh | 61 +- .buildkite/scripts/cleanup-nightly-builds.sh | 7 +- .buildkite/scripts/generate-nightly-index.py | 72 +- .buildkite/scripts/publish-release-images.sh | 44 +- .../scripts/push-nightly-builds-rocm.sh | 107 ++- .buildkite/scripts/rocm/build-config.sh | 4 +- .buildkite/scripts/rocm/stack.sh | 78 ++ .buildkite/scripts/upload-rocm-wheels.sh | 88 ++- docker/Dockerfile.rock_base | 453 ------------ docker/Dockerfile.rocm | 82 +-- .../{Dockerfile.rock => Dockerfile.rocm_72} | 98 +-- docker/Dockerfile.rocm_72_base | 479 +++++++++++++ docker/Dockerfile.rocm_base | 539 +++++++------- docker/ci-rocm.hcl | 34 +- docker/docker-bake-rocm.hcl | 4 - .../installation/gpu.rocm.inc.md | 100 +-- setup.py | 29 +- .../pooling/test_token_classification.py | 2 +- .../generate-rocm-wheels-root-index.sh | 10 +- tools/vllm-rocm/therock_wheels.py | 234 ++++++ vllm/platforms/rocm.py | 6 + 24 files changed, 1738 insertions(+), 1630 deletions(-) create mode 100644 .buildkite/scripts/rocm/stack.sh delete mode 100644 docker/Dockerfile.rock_base rename docker/{Dockerfile.rock => Dockerfile.rocm_72} (92%) create mode 100644 docker/Dockerfile.rocm_72_base create mode 100644 tools/vllm-rocm/therock_wheels.py diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 6064f93104f5..84987524b23d 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -703,547 +703,280 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" + - trigger: "perf-eval" + label: "Trigger perf-eval (CUDA)" + key: trigger-perf-eval-cuda + if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" + depends_on: + - build-release-image-x86 + async: true + soft_fail: true + build: + branch: "main" + message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (CUDA)" + env: + NIGHTLY: 1 # run all perf-eval workloads + VLLM_COMMIT: "$BUILDKITE_COMMIT" + VLLM_IMAGE_CUDA: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64" + + - trigger: "perf-eval" + label: "Trigger perf-eval (ROCm)" + key: trigger-perf-eval-rocm + if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" + depends_on: + - build-rocm100-release-image + async: true + soft_fail: true + build: + branch: "main" + message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (ROCm)" + env: + NIGHTLY: 1 # run all perf-eval workloads + VLLM_COMMIT: "$BUILDKITE_COMMIT" + VLLM_IMAGE_ROCM: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm100" + # ============================================================================= - # ROCm Release Pipeline (x86_64 only) + # ROCm Release Pipelines (x86_64 only) # ============================================================================= # - # vLLM version is determined by the Buildkite checkout (like CUDA pipeline). - # To build a specific version, trigger the build from that branch/tag. + # One group per ROCm stack, identified by its base Dockerfile (ROCM_BASE_DOCKERFILE). + # The scripts derive the ROCm version, wheel variant, cache keys and image tags from + # it (.buildkite/scripts/rocm/stack.sh), so both groups share one set of step + # definitions, anchored on the ROCm 10.0 group. # - # Environment variables for ROCm builds (set via Buildkite UI or schedule): - # - # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base + # - ROCm 10.0 (TheRock, docker/Dockerfile.rocm{,_base}) is the default: it owns the + # :nightly / :latest images and the root wheel index. + # - ROCm 7.2 (docker/Dockerfile.rocm_72{,_base}) is kept for a transition period as + # rocm72 images and wheels. # + # vLLM version is determined by the Buildkite checkout (like the CUDA pipeline). # ============================================================================= - - group: "Build ROCm Wheel / Image " - key: "build-rocm-wheel-image" + - group: "Build ROCm 10.0 Wheel / Image" + key: "build-rocm100" depends_on: ~ steps: - # ROCm Job 1: Build ROCm Base Wheels (with S3 caching) - - label: ":rocm: Build ROCm Base Image & Wheels" - id: build-rocm-base-wheels + - label: ":rocm: Build ROCm 10.0 Base Image & Wheels" + id: build-rocm100-base depends_on: ~ - agents: + agents: &rocm_build_agents queue: cpu_queue_release - commands: + timeout_in_minutes: 360 + env: &rocm100_env + ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_base" + ROCM_DEFAULT_ARG: "--default" + DOCKER_BUILDKIT: "1" + S3_BUCKET: "vllm-wheels" + DOCKERHUB_USERNAME: "vllmbot" + commands: &rocm_base_commands - | set -euo pipefail - - # Generate cache key - CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base" - - echo "========================================" - echo "ROCm Base Build Configuration" - echo "========================================" - echo " CACHE_KEY: $${CACHE_KEY}" - echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}" - echo "========================================" - - # Login to ECR + source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" + SCCACHE_ARGS=(--build-arg USE_SCCACHE=1 --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache + --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0) aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - IMAGE_EXISTS=false - WHEELS_EXIST=false - - # Check ECR for Docker image - - if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then - IMAGE_EXISTS=true - echo "ECR image cache HIT" - fi - - # Check S3 for wheels - WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check) - if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then - WHEELS_EXIST=true - echo "S3 wheels cache HIT" - fi - - # Scenario 1: Both cached (best case) - if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then - echo "" - echo "FULL CACHE HIT - Reusing both image and wheels" - echo "" - - # Download wheels - .buildkite/scripts/cache-rocm-base-wheels.sh download - - # Save ECR tag for downstream jobs - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - - # Scenario 2: Full rebuild needed + # Base image, cached in ECR by the base Dockerfile's hash + if docker manifest inspect "$${ROCM_STACK_ECR_BASE}" > /dev/null 2>&1; then + echo "ECR image cache HIT: $${ROCM_STACK_ECR_BASE}" else - echo "" - echo " CACHE MISS - Building from scratch..." - echo "" - - # Build full base image and push to ECR - DOCKER_BUILDKIT=1 docker buildx build \ - --file docker/Dockerfile.rocm_base \ - --tag "$${ECR_CACHE_TAG}" \ - --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --push \ - . - - # Build wheel extraction stage - DOCKER_BUILDKIT=1 docker buildx build \ - --file docker/Dockerfile.rocm_base \ - --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \ - --target debs_wheel_release \ - --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --load \ - . - - # Extract and upload wheels - mkdir -p artifacts/rocm-base-wheels - cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) - docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/ - docker rm $${cid} - - .buildkite/scripts/cache-rocm-base-wheels.sh upload - - # Cache base docker image to ECR - docker push "$${ECR_CACHE_TAG}" - - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - - echo "" - echo " Build complete - Image and wheels cached" + docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${ROCM_STACK_ECR_BASE}" \ + "$${SCCACHE_ARGS[@]}" --push . fi - artifact_paths: - - "artifacts/rocm-base-wheels/*.whl" - env: - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" + # Base wheels, cached in S3 (see cache-rocm-base-wheels.sh) + CACHE=.buildkite/scripts/cache-rocm-base-wheels.sh + if [ "$$($${CACHE} "$${ROCM_BASE_DOCKERFILE}" check)" = "hit" ]; then + $${CACHE} "$${ROCM_BASE_DOCKERFILE}" download + else + DEBS="rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}-$${ROCM_STACK_VARIANT}" + docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${DEBS}" \ + --target debs_wheel_release "$${SCCACHE_ARGS[@]}" --load . + mkdir -p "$${ROCM_STACK_BASE_WHEELS_DIR}" + cid=$$(docker create "$${DEBS}") + docker cp "$${cid}":/app/debs/. "$${ROCM_STACK_BASE_WHEELS_DIR}/" + docker rm "$${cid}" + $${CACHE} "$${ROCM_BASE_DOCKERFILE}" upload + fi + buildkite-agent artifact upload "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" - # ROCm Job 2: Build vLLM ROCm Wheel - - label: ":python: Build vLLM ROCm Wheel - x86_64" - id: build-rocm-vllm-wheel - depends_on: - - step: build-rocm-base-wheels - allow_failure: false - agents: - queue: cpu_queue_release + - label: ":python: Build vLLM ROCm 10.0 Wheel - x86_64" + id: build-rocm100-vllm-wheel + depends_on: build-rocm100-base + agents: *rocm_build_agents timeout_in_minutes: 180 - commands: - # Download artifacts and prepare Docker image + env: *rocm100_env + commands: &rocm_vllm_wheel_commands - | set -euo pipefail - - # Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags) - # This fixes version detection when tags are moved/force-pushed - echo "Fetching latest tags from origin..." + source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" + # Buildkite's default fetch doesn't update moved tags; version detection needs them git fetch --tags --force origin - - # Log tag information for debugging version detection - echo "========================================" - echo "Git Tag Verification" - echo "========================================" - echo "Current HEAD: $(git rev-parse HEAD)" - echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')" - echo "" - echo "Recent tags (pointing to commits near HEAD):" - git tag -l --sort=-creatordate | head -5 - echo "setuptools_scm version detection:" - pip install -q setuptools_scm 2>/dev/null || true - python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)" - echo "========================================" - - # Download wheel artifacts from current build - echo "Downloading wheel artifacts from current build" - buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . - - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" - echo "This should have been set by the build-rocm-base-wheels job" - exit 1 - fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Login to ECR + buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" . aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - - # Prepare base wheels for Docker build context + docker pull "$${ROCM_STACK_ECR_BASE}" mkdir -p docker/context/base-wheels touch docker/context/base-wheels/.keep - cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/ - echo "Base wheels for vLLM build:" - ls -lh docker/context/base-wheels/ - - echo "========================================" - echo "Building vLLM wheel with:" - echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" - echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" - echo "========================================" - - # Build vLLM wheel using local checkout (REMOTE_VLLM=0) - DOCKER_BUILDKIT=1 docker build \ - --file docker/Dockerfile.rocm \ - --target export_vllm_wheel_release \ + cp "$${ROCM_STACK_BASE_WHEELS_DIR}"/*.whl docker/context/base-wheels/ + docker build --file "$${ROCM_STACK_DOCKERFILE}" --target export_vllm_wheel_release \ --output type=local,dest=rocm-dist \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ - --build-arg REMOTE_VLLM=0 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg USE_SCCACHE=1 \ + --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" --build-arg REMOTE_VLLM=0 \ + --build-arg GIT_REPO_CHECK=1 --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - . - echo "Built vLLM wheel:" - ls -lh rocm-dist/*.whl - # Copy wheel to artifacts directory - mkdir -p artifacts/rocm-vllm-wheel - cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/ - echo "Final vLLM wheel:" - ls -lh artifacts/rocm-vllm-wheel/ - artifact_paths: - - "artifacts/rocm-vllm-wheel/*.whl" - env: - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" - - # ROCm Job 3: Upload Wheels to S3 - - label: ":s3: Upload ROCm Wheels to S3" - id: upload-rocm-wheels - depends_on: - - step: build-rocm-vllm-wheel - allow_failure: false - agents: - queue: cpu_queue_release + --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 . + mkdir -p "$${ROCM_STACK_VLLM_WHEEL_DIR}" + cp rocm-dist/*.whl "$${ROCM_STACK_VLLM_WHEEL_DIR}/" + buildkite-agent artifact upload "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl" + + - label: ":s3: Upload ROCm 10.0 Wheels to S3" + id: upload-rocm100-wheels + depends_on: build-rocm100-vllm-wheel + agents: *rocm_build_agents timeout_in_minutes: 60 - commands: - # Download all wheel artifacts and run upload + env: *rocm100_env + commands: &rocm_upload_commands - | set -euo pipefail - - # Download artifacts from current build - echo "Downloading artifacts from current build" - buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . - buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" . - - # # Run upload script - bash .buildkite/scripts/upload-rocm-wheels.sh - env: - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" - - # ROCm Job 4: Annotate ROCm Wheel Release - - label: ":memo: Annotate ROCm wheel release" - id: annotate-rocm-release - depends_on: - - upload-rocm-wheels - agents: - queue: cpu_queue_release - commands: - - "bash .buildkite/scripts/annotate-rocm-release.sh" - env: - S3_BUCKET: "vllm-wheels" - - # ROCm Job 5: Generate Root Index for ROCm Wheels (for release only) - # This is the job to create https://wheels.vllm.ai/rocm/ index allowing - # users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/` + source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" + buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" . + buildkite-agent artifact download "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl" . + bash .buildkite/scripts/upload-rocm-wheels.sh "$${ROCM_BASE_DOCKERFILE}" + + - label: ":memo: Annotate ROCm 10.0 wheel release" + id: annotate-rocm100-release + depends_on: upload-rocm100-wheels + agents: *rocm_build_agents + env: *rocm100_env + commands: &rocm_annotate_commands + - bash .buildkite/scripts/annotate-rocm-release.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG} + + # Generate https://wheels.vllm.ai/rocm/ (the default stack's variant), so users can + # `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/` (release only) - block: "Generate Root Index for ROCm Wheels for Release" key: block-generate-root-index-rocm-wheels - depends_on: upload-rocm-wheels + depends_on: + - upload-rocm100-wheels + - upload-rocm72-wheels - label: ":package: Generate Root Index for ROCm Wheels for Release" depends_on: block-generate-root-index-rocm-wheels id: generate-root-index-rocm-wheels - agents: - queue: cpu_queue_release + agents: *rocm_build_agents + env: *rocm100_env commands: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" - env: - S3_BUCKET: "vllm-wheels" - VARIANT: "rocm723" - # ROCm Job 6: Build ROCm Release Docker Image - - label: ":docker: Build release image - x86_64 - ROCm" - id: build-rocm-release-image + - label: ":docker: Build release image - x86_64 - ROCm 10.0" + id: build-rocm100-release-image depends_on: - step: block-build-release-images allow_failure: true - - step: build-rocm-base-wheels + - step: build-rocm100-base allow_failure: false - agents: - queue: cpu_queue_release - timeout_in_minutes: 60 - commands: + agents: *rocm_build_agents + timeout_in_minutes: 180 + env: *rocm100_env + commands: &rocm_release_image_commands - | set -euo pipefail - - # Login to ECR + source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" - echo "This should have been set by the build-rocm-base-wheels job" - exit 1 - fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - - # Pass the base image ECR tag to downstream steps (nightly publish) - buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}" - - echo "========================================" - echo "Building vLLM ROCm release image with:" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" - echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" - echo "========================================" - - # Build vLLM ROCm release image using cached base - DOCKER_BUILDKIT=1 docker build \ - --build-arg max_jobs=16 \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ - --build-arg USE_SCCACHE=1 \ + docker pull "$${ROCM_STACK_ECR_BASE}" + docker build --file "$${ROCM_STACK_DOCKERFILE}" --target vllm-openai \ + --tag "$${ROCM_STACK_ECR_IMAGE}" --progress plain \ + --build-arg max_jobs=16 --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" \ + --build-arg INSTALL_LMCACHE=true --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --build-arg INSTALL_LMCACHE=true \ - --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile.rocm . - - # Push to ECR - docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm - - echo "" - echo " Successfully built and pushed ROCm release image" - echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" - echo "" - env: - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" - - - label: "Publish nightly XPU image to DockerHub" - depends_on: - - create-manifest-xpu - if: build.env("NIGHTLY") == "1" - agents: - queue: small_cpu_queue_release - commands: - - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" - - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" - plugins: - - docker-login#v3.0.0: - username: vllmbot - password-env: DOCKERHUB_TOKEN - env: - DOCKER_BUILDKIT: "1" - DOCKERHUB_USERNAME: "vllmbot" + --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 . + docker push "$${ROCM_STACK_ECR_IMAGE}" - - label: "Publish nightly ROCm image to DockerHub" - depends_on: - - build-rocm-release-image + - label: "Publish nightly ROCm 10.0 image to DockerHub" + id: publish-rocm100-nightly + depends_on: build-rocm100-release-image if: build.env("NIGHTLY") == "1" - agents: + agents: &rocm_publish_agents queue: small_cpu_queue_release - commands: - - "bash .buildkite/scripts/push-nightly-builds-rocm.sh" - # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm" - - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm" - plugins: + env: *rocm100_env + commands: &rocm_publish_nightly_commands + - bash .buildkite/scripts/push-nightly-builds-rocm.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG} + plugins: &rocm_dockerhub_login - docker-login#v3.0.0: username: vllmbot password-env: DOCKERHUB_TOKEN - env: - DOCKER_BUILDKIT: "1" - DOCKERHUB_USERNAME: "vllmbot" - - - trigger: "perf-eval" - label: "Trigger perf-eval (CUDA)" - key: trigger-perf-eval-cuda - if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" - depends_on: - - build-release-image-x86 - async: true - soft_fail: true - build: - branch: "main" - message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (CUDA)" - env: - NIGHTLY: 1 # run all perf-eval workloads - VLLM_COMMIT: "$BUILDKITE_COMMIT" - VLLM_IMAGE_CUDA: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64" - - - trigger: "perf-eval" - label: "Trigger perf-eval (ROCm)" - key: trigger-perf-eval-rocm - if: build.env("NIGHTLY") == "1" || build.env("RUN_PERF_EVAL") == "1" - depends_on: - - build-rocm-release-image - async: true - soft_fail: true - build: - branch: "main" - message: "(NIGHTLY=${NIGHTLY:-0}) Evaluate vLLM $BUILDKITE_COMMIT (ROCm)" - env: - NIGHTLY: 1 # run all perf-eval workloads - VLLM_COMMIT: "$BUILDKITE_COMMIT" - VLLM_IMAGE_ROCM: "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-rocm" - - # ============================================================================= - # ROCk Release Pipeline (x86_64 only) - # ============================================================================= - # - # Same shape as the ROCm pipeline above (Jobs 1 and 6 only), but built from - # docker/Dockerfile.rock_base and docker/Dockerfile.rock, which install the - # ROCm SDK from TheRock wheels instead of the ROCm apt repos. - # - # Note: ROCm SDK version is determined by ROCM_SDK_VERSION in - # docker/Dockerfile.rock_base - # - # ============================================================================= - - group: "Build ROCk Image" - key: "build-rock-image" + - group: "Build ROCm 7.2 Wheel / Image" + key: "build-rocm72" depends_on: ~ steps: - # ROCk Job 1: Build ROCk Base Image (with ECR caching) - - label: ":rocm: Build ROCk Base Image" - id: build-rock-base-image + - label: ":rocm: Build ROCm 7.2 Base Image & Wheels" + id: build-rocm72-base depends_on: ~ - agents: - queue: cpu_queue_release + agents: *rocm_build_agents timeout_in_minutes: 360 - commands: - - | - set -euo pipefail - - # The SDK version, torch wheels and component branches all live in - # the Dockerfile, so its hash is the cache key. - CACHE_KEY=$$(sha256sum docker/Dockerfile.rock_base | cut -c1-16) - ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rock-base" - - echo "========================================" - echo "ROCk Base Build Configuration" - echo "========================================" - echo " CACHE_KEY: $${CACHE_KEY}" - echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}" - echo "========================================" - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then - echo "ECR image cache HIT" - else - echo "CACHE MISS - Building from scratch..." - - DOCKER_BUILDKIT=1 docker buildx build \ - --file docker/Dockerfile.rock_base \ - --tag "$${ECR_CACHE_TAG}" \ - --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --push \ - . - fi - - # Save ECR tag for downstream jobs - buildkite-agent meta-data set "rock-base-image-tag" "$${ECR_CACHE_TAG}" - env: + env: &rocm72_env + ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_72_base" + ROCM_DEFAULT_ARG: "" DOCKER_BUILDKIT: "1" + S3_BUCKET: "vllm-wheels" + DOCKERHUB_USERNAME: "vllmbot" + commands: *rocm_base_commands + + - label: ":python: Build vLLM ROCm 7.2 Wheel - x86_64" + id: build-rocm72-vllm-wheel + depends_on: build-rocm72-base + agents: *rocm_build_agents + timeout_in_minutes: 180 + env: *rocm72_env + commands: *rocm_vllm_wheel_commands - # ROCk Job 2: Build ROCk Release Docker Image - - label: ":docker: Build release image - x86_64 - ROCk" - id: build-rock-release-image + - label: ":s3: Upload ROCm 7.2 Wheels to S3" + id: upload-rocm72-wheels + depends_on: build-rocm72-vllm-wheel + agents: *rocm_build_agents + timeout_in_minutes: 60 + env: *rocm72_env + commands: *rocm_upload_commands + + - label: ":memo: Annotate ROCm 7.2 wheel release" + id: annotate-rocm72-release + depends_on: upload-rocm72-wheels + agents: *rocm_build_agents + env: *rocm72_env + commands: *rocm_annotate_commands + + - label: ":docker: Build release image - x86_64 - ROCm 7.2" + id: build-rocm72-release-image depends_on: - step: block-build-release-images allow_failure: true - - step: build-rock-base-image + - step: build-rocm72-base allow_failure: false - agents: - queue: cpu_queue_release + agents: *rocm_build_agents timeout_in_minutes: 180 - commands: - - | - set -euo pipefail - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rock-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rock-base-image-tag metadata not found" - echo "This should have been set by the build-rock-base-image job" - exit 1 - fi - - docker pull "$${ECR_IMAGE_TAG}" - - # Pass the base image ECR tag to downstream steps (nightly publish) - buildkite-agent meta-data set "rock-base-ecr-tag" "$${ECR_IMAGE_TAG}" - - echo "========================================" - echo "Building vLLM ROCk release image with:" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" - echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" - echo "========================================" - - DOCKER_BUILDKIT=1 docker build \ - --build-arg max_jobs=16 \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ - --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 \ - --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --build-arg INSTALL_LMCACHE=true \ - --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rock \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile.rock . + env: *rocm72_env + commands: *rocm_release_image_commands - docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rock - env: - DOCKER_BUILDKIT: "1" + - label: "Publish nightly ROCm 7.2 image to DockerHub" + id: publish-rocm72-nightly + depends_on: build-rocm72-release-image + if: build.env("NIGHTLY") == "1" + agents: *rocm_publish_agents + env: *rocm72_env + commands: *rocm_publish_nightly_commands + plugins: *rocm_dockerhub_login - - label: "Publish nightly ROCk image to DockerHub" + - label: "Publish nightly XPU image to DockerHub" depends_on: - - build-rock-release-image + - create-manifest-xpu if: build.env("NIGHTLY") == "1" agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/push-nightly-builds-rocm.sh rock rocm100" - # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly-rocm100- vllm/vllm-openai-rocm" - - "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly-rocm100- vllm/vllm-openai-rocm" + - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" plugins: - docker-login#v3.0.0: username: vllmbot @@ -1264,7 +997,8 @@ steps: - create-multi-arch-manifest-ubuntu2404 - create-multi-arch-manifest-cuda-12-9-ubuntu2404 - create-manifest-xpu - - build-rocm-release-image + - build-rocm100-release-image + - build-rocm72-release-image - input-release-version # Wait for CPU builds if their block steps were unblocked, so publish # doesn't race the in-progress CPU build. allow_failure lets publish @@ -1289,6 +1023,7 @@ steps: - "cuda-13-0-ubuntu-24-04" - "cuda-12-9-ubuntu-24-04" - "rocm" + - "rocm72" - "xpu" - "cpu" plugins: diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh index d66129722748..1889db43b5ad 100755 --- a/.buildkite/scripts/annotate-rocm-release.sh +++ b/.buildkite/scripts/annotate-rocm-release.sh @@ -2,15 +2,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Generate Buildkite annotation for ROCm wheel release +# Generate Buildkite annotation for a ROCm stack's wheel release. +# +# Usage: annotate-rocm-release.sh [--default] +# base-dockerfile e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base +# --default the stack also owns the plain :latest / :v Docker Hub tags set -ex -# Extract build configuration from Dockerfile.rocm_base (single source of truth) -# Extract ROCm version dynamically from Dockerfile.rocm_base -# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0" -ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown") -PYTHON_VERSION=$(grep '^ARG PYTHON_VERSION=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTHON_VERSION=//') -PYTORCH_ROCM_ARCH=$(grep '^ARG PYTORCH_ROCM_ARCH=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTORCH_ROCM_ARCH=//') +# shellcheck source=.buildkite/scripts/rocm/stack.sh +source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 [--default]}" +TAG_VARIANTS=("-${ROCM_STACK_VARIANT}") +if [[ "${2:-}" == "--default" ]]; then + TAG_VARIANTS=("" "-${ROCM_STACK_VARIANT}") +fi +ROCM_VERSION="$ROCM_STACK_VERSION" +VARIANT="$ROCM_STACK_VARIANT" +PYTHON_VERSION=$(sed -nE 's/^ARG PYTHON_VERSION="?([^" ]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE") +PYTORCH_ROCM_ARCH=$(sed -nE 's/^ARG PYTORCH_ROCM_ARCH="?([^"]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE") # Get release version, default to 1.0.0.dev for nightly/per-commit builds RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "") @@ -18,22 +26,32 @@ if [ -z "${RELEASE_VERSION}" ]; then RELEASE_VERSION="1.0.0.dev" fi -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - # S3 URLs S3_BUCKET="${S3_BUCKET:-vllm-wheels}" S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}" -S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com" - -# Format ROCm version for path (e.g., "7.1" -> "rocm710") -ROCM_VERSION_PATH="rocm$(echo "${ROCM_VERSION}" | tr -d '.')" -ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}" -buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF -## ROCm Wheel and Docker Image Releases +S3_HOST="${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com" +S3_URL="http://${S3_HOST}" +WHEEL_PATH="rocm/${BUILDKITE_COMMIT}/${VARIANT}-wheels/" + +DOCKER_CMDS="docker pull ${ROCM_STACK_ECR_BASE} +docker pull ${ROCM_STACK_ECR_IMAGE}" +for suffix in "${TAG_VARIANTS[@]}"; do + for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do + DOCKER_CMDS+=" +docker tag ${ROCM_STACK_ECR_BASE} vllm/vllm-openai-rocm:${tag}-base +docker tag ${ROCM_STACK_ECR_IMAGE} vllm/vllm-openai-rocm:${tag} +docker push vllm/vllm-openai-rocm:${tag}-base +docker push vllm/vllm-openai-rocm:${tag}" + done +done + +buildkite-agent annotate --style 'success' --context "rocm-release-workflow-${VARIANT}" << EOF +## ROCm ${ROCM_VERSION} Wheel and Docker Image Releases ### Build Configuration | Setting | Value | |---------|-------| | **ROCm Version** | ${ROCM_VERSION} | +| **Index Variant** | ${VARIANT} | | **Python Version** | ${PYTHON_VERSION} | | **GPU Architectures** | ${PYTORCH_ROCM_ARCH} | | **Branch** | \`${BUILDKITE_BRANCH}\` | @@ -41,49 +59,20 @@ buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << ### :package: Installation -**Install from this build (by commit):** - -\`\`\`bash -pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com - -# Example for ROCm ${ROCM_VERSION}: -pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com -\`\`\` - -**Install from nightly (if published):** - \`\`\`bash -pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com +# This build (by commit) +pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${VARIANT}/ --trusted-host ${S3_HOST} +# Nightly (if published) +pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/${VARIANT}/ --trusted-host ${S3_HOST} \`\`\` ### :floppy_disk: Download Wheels Directly \`\`\`bash -# List all ROCm wheels -aws s3 ls s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/ -# Download specific wheels -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/vllm-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torch-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-kernels-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl . +aws s3 ls s3://${S3_BUCKET}/${WHEEL_PATH} +aws s3 cp --recursive --exclude "*" --include "*.whl" s3://${S3_BUCKET}/${WHEEL_PATH} . \`\`\` -### :gear: Included Packages -- **vllm**: vLLM with ROCm support -- **torch**: PyTorch built for ROCm ${ROCM_VERSION} -- **triton**: Triton -- **triton-kernels**: Triton kernels -- **torchvision**: TorchVision for ROCm PyTorch -- **torchaudio**: Torchaudio for ROCm PyTorch -- **amdsmi**: AMD SMI Python bindings -- **amd_aiter**: Aiter for ROCm -- **flash-attn**: Flash Attention for ROCm - ### :warning: Notes - These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs - Supported GPU architectures: ${PYTORCH_ROCM_ARCH} @@ -91,23 +80,8 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash- ### :package: Docker Image Release -To download and upload the image: - \`\`\` -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base -docker push vllm/vllm-openai-rocm:latest-base -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest -docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker push vllm/vllm-openai-rocm:latest -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} +${DOCKER_CMDS} \`\`\` EOF diff --git a/.buildkite/scripts/cache-rocm-base-wheels.sh b/.buildkite/scripts/cache-rocm-base-wheels.sh index dc76544d3603..02fbbf52cf54 100755 --- a/.buildkite/scripts/cache-rocm-base-wheels.sh +++ b/.buildkite/scripts/cache-rocm-base-wheels.sh @@ -5,39 +5,33 @@ # Cache helper for ROCm base wheels # # This script manages caching of pre-built ROCm base wheels (torch, triton, etc.) -# to avoid rebuilding them when Dockerfile.rocm_base hasn't changed. +# to avoid rebuilding them when the stack's base Dockerfile hasn't changed. # -# Usage: -# cache-rocm-base-wheels.sh check - Check if cache exists, outputs "hit" or "miss" -# cache-rocm-base-wheels.sh upload - Upload wheels to cache -# cache-rocm-base-wheels.sh download - Download wheels from cache -# cache-rocm-base-wheels.sh key - Output the cache key +# Usage: cache-rocm-base-wheels.sh +# base-dockerfile e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base +# Commands: +# check - Check if cache exists, outputs "hit" or "miss" +# upload - Upload wheels to cache +# download - Download wheels from cache +# key - Output the cache key +# path - Output the full S3 cache path # # Environment variables: # S3_BUCKET - S3 bucket name (default: vllm-wheels) # -# Note: ROCm version is determined by BASE_IMAGE in Dockerfile.rocm_base, -# so changes to ROCm version are captured by the Dockerfile hash. +# The cache key hashes the base Dockerfile and the repo files it COPYs in, which +# capture the ROCm version and every component pin (see rocm/stack.sh). set -euo pipefail BUCKET="${S3_BUCKET:-vllm-wheels}" -DOCKERFILE="docker/Dockerfile.rocm_base" CACHE_PREFIX="rocm/cache" -# Generate hash from Dockerfile content + build args -generate_cache_key() { - # Include Dockerfile content - if [[ ! -f "$DOCKERFILE" ]]; then - echo "ERROR: Dockerfile not found: $DOCKERFILE" >&2 - exit 1 - fi - local dockerfile_hash=$(sha256sum "$DOCKERFILE" | cut -c1-16) - - echo "${dockerfile_hash}" -} - -CACHE_KEY=$(generate_cache_key) +# shellcheck source=.buildkite/scripts/rocm/stack.sh +source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 }" +shift +CACHE_KEY="$ROCM_STACK_WHEEL_KEY" +WHEELS_DIR="$ROCM_STACK_BASE_WHEELS_DIR" CACHE_PATH="s3://${BUCKET}/${CACHE_PREFIX}/${CACHE_KEY}/" case "${1:-}" in @@ -67,19 +61,19 @@ case "${1:-}" in echo "Cache path: ${CACHE_PATH}" echo "" - if [[ ! -d "artifacts/rocm-base-wheels" ]]; then - echo "ERROR: artifacts/rocm-base-wheels directory not found" >&2 + if [[ ! -d "${WHEELS_DIR}" ]]; then + echo "ERROR: ${WHEELS_DIR} directory not found" >&2 exit 1 fi - WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) + WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) if [[ "$WHEEL_COUNT" -eq 0 ]]; then - echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2 + echo "ERROR: No wheels found in ${WHEELS_DIR}/" >&2 exit 1 fi echo "Uploading $WHEEL_COUNT wheels..." - aws s3 cp --recursive artifacts/rocm-base-wheels/ "${CACHE_PATH}" + aws s3 cp --recursive "${WHEELS_DIR}/" "${CACHE_PATH}" echo "" echo "Cache upload complete!" @@ -93,17 +87,17 @@ case "${1:-}" in echo "Cache key: ${CACHE_KEY}" echo "Cache path: ${CACHE_PATH}" echo "" - mkdir -p artifacts/rocm-base-wheels + mkdir -p "${WHEELS_DIR}" # Use sync with include/exclude to only download .whl files - aws s3 sync "${CACHE_PATH}" artifacts/rocm-base-wheels/ \ + aws s3 sync "${CACHE_PATH}" "${WHEELS_DIR}/" \ --exclude "*" \ --include "*.whl" echo "" echo "Downloaded wheels:" - find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' -exec ls -lh {} \; - WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) + find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' -exec ls -lh {} \; + WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) echo "" echo "Total: $WHEEL_COUNT wheels" echo "========================================" @@ -118,12 +112,12 @@ case "${1:-}" in ;; *) - echo "Usage: $0 {check|upload|download|key|path}" >&2 + echo "Usage: $0 {check|upload|download|key|path}" >&2 echo "" >&2 echo "Commands:" >&2 echo " check - Check if cache exists, outputs 'hit' or 'miss'" >&2 - echo " upload - Upload wheels from artifacts/rocm-base-wheels/ to cache" >&2 - echo " download - Download wheels from cache to artifacts/rocm-base-wheels/" >&2 + echo " upload - Upload wheels from ${WHEELS_DIR}/ to cache" >&2 + echo " download - Download wheels from cache to ${WHEELS_DIR}/" >&2 echo " key - Output the cache key" >&2 echo " path - Output the full S3 cache path" >&2 exit 1 diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index ae81f59e843c..5935938616bc 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_rocshmem build_deepep mori_base ci_base" +DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="3" # ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected # local-source stages rather than unreachable remote-fetch alternatives. @@ -555,9 +555,6 @@ configure_custom_rocm_stages() { CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${stages}}" ROCM_CSRC_DOCKERFILE_STAGES="${ROCM_CSRC_DOCKERFILE_STAGES:-${stages}}" ROCM_RUST_DOCKERFILE_STAGES="${ROCM_RUST_DOCKERFILE_STAGES:-${stages}}" - if [[ " ${stages} " != *" build_rocshmem "* ]]; then - unset ROCSHMEM_BRANCH ROCSHMEM_CACHE_KEY DEEPEP_CACHE_KEY - fi } compose_dependency_cache_key() { @@ -1724,14 +1721,11 @@ ci_base_metadata_pairs() { metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}" metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")" metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}" - metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")" - metadata_pair "vllm.rocm.rocshmem_commit" "${ROCSHMEM_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_BRANCH")}" metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")" metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}" - metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}" metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" } @@ -2314,7 +2308,7 @@ extract_dependency_pins() { return 0 fi - for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do + for var in NIXL_BRANCH UCX_BRANCH DEEPEP_BRANCH; do if [[ -n "${!var:-}" ]]; then echo "Using provided ${var}: ${!var}" continue @@ -2336,17 +2330,14 @@ compute_dependency_cache_keys() { local dockerfile_rocm="" local nixl_branch="" local ucx_branch="" - local rocshmem_branch="" local deepep_branch="" local nixl_material="" - local rocshmem_material="" local deepep_material="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH") ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") - rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH") if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then @@ -2360,22 +2351,11 @@ compute_dependency_cache_keys() { echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}" fi - if [[ -n "${rocshmem_branch}" ]]; then - rocshmem_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem") - ROCSHMEM_CACHE_KEY=$( - compose_dependency_cache_key \ - "${rocshmem_branch}" \ - "${rocshmem_material}" - ) - export ROCSHMEM_CACHE_KEY - echo "ROCShmem dependency cache key: ${ROCSHMEM_CACHE_KEY}" - fi - - if [[ -n "${deepep_branch}" && -n "${rocshmem_branch}" ]]; then - deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem build_deepep") + if [[ -n "${deepep_branch}" ]]; then + deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_deepep") DEEPEP_CACHE_KEY=$( compose_dependency_cache_key \ - "${deepep_branch}-rocshmem-${rocshmem_branch}" \ + "${deepep_branch}" \ "${deepep_material}" ) export DEEPEP_CACHE_KEY @@ -2415,18 +2395,11 @@ dependency_cache_ref_for_target() { printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}" fi ;; - rocshmem-rocm-ci) - if [[ -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then - printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" - elif [[ -n "${ROCSHMEM_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_BRANCH}" - fi - ;; deepep-rocm-ci) if [[ -n "${DEEPEP_CACHE_KEY:-}" ]]; then printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_CACHE_KEY}" elif [[ -n "${DEEPEP_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}" + printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}" fi ;; esac @@ -2444,7 +2417,6 @@ add_dependency_cache_target() { resolve_ci_base_dependency_targets() { local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}" local nixl_ref="" - local rocshmem_ref="" local deepep_ref="" [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0 @@ -2452,7 +2424,7 @@ resolve_ci_base_dependency_targets() { case "${mode}" in always) echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially" - for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do + for target in nixl-rocm-ci deepep-rocm-ci; do if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then add_dependency_cache_target "${target}" fi @@ -2482,16 +2454,6 @@ resolve_ci_base_dependency_targets() { fi fi - if [[ "${mode}" != "always" && -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then - rocshmem_ref=$(dependency_cache_ref_for_target "rocshmem-rocm-ci") - if dependency_cache_ref_exists "${rocshmem_ref}"; then - echo "ROCShmem dependency cache exists: ${rocshmem_ref}" - else - echo "ROCShmem dependency cache missing; will seed: ${rocshmem_ref}" - add_dependency_cache_target "rocshmem-rocm-ci" - fi - fi - if [[ "${mode}" != "always" && -n "${DEEPEP_CACHE_KEY:-}" ]]; then deepep_ref=$(dependency_cache_ref_for_target "deepep-rocm-ci") if dependency_cache_ref_exists "${deepep_ref}"; then @@ -2502,15 +2464,6 @@ resolve_ci_base_dependency_targets() { fi fi - # DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too - # so the pair stays consistent for future ci_base rebuilds. - if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \ - && ! printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "deepep-rocm-ci" \ - && [[ -n "${DEEPEP_BRANCH:-}" ]]; then - echo "ROCShmem cache is missing; also seeding DeepEP cache" - add_dependency_cache_target "deepep-rocm-ci" - fi - BAKE_TARGETS=("ci-base-rocm-ci") if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then echo "All dependency caches exist; building ci_base without dependency cache exports" diff --git a/.buildkite/scripts/cleanup-nightly-builds.sh b/.buildkite/scripts/cleanup-nightly-builds.sh index 41b04996e455..925ae66dd6b5 100755 --- a/.buildkite/scripts/cleanup-nightly-builds.sh +++ b/.buildkite/scripts/cleanup-nightly-builds.sh @@ -5,14 +5,17 @@ set -ex # Clean up old nightly builds from DockerHub, keeping only the last 14 builds # This script uses DockerHub API to list and delete old tags with specified prefix # Tags starting with "nightly-dev" are always excluded from deletion -# Usage: cleanup-nightly-builds.sh [TAG_PREFIX] [REPO] +# Usage: cleanup-nightly-builds.sh [TAG_PREFIX] [REPO] [EXCLUDE_PREFIX] # Example: cleanup-nightly-builds.sh "nightly-" # Example: cleanup-nightly-builds.sh "cu130-nightly-" # Example: cleanup-nightly-builds.sh "nightly-" "vllm/vllm-openai-rocm" +# Example: cleanup-nightly-builds.sh "nightly-" "vllm/vllm-openai-rocm" "nightly-rocm" +# (EXCLUDE_PREFIX skips tags owned by another variant sharing TAG_PREFIX) # Get tag prefix and repo from arguments TAG_PREFIX="${1:-nightly-}" REPO="${2:-vllm/vllm-openai}" +EXCLUDE_PREFIX="${3:-}" echo "Cleaning up tags with prefix: $TAG_PREFIX in repository: $REPO" @@ -57,7 +60,7 @@ get_all_tags() { # Get both last_updated timestamp and tag name, separated by | # Exclude mutable architecture tags and nightly-dev tags from cleanup - local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" '.results[] | select(.name | startswith($prefix)) | select(.name != ($prefix + "x86_64") and .name != ($prefix + "aarch64")) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"') + local tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" --arg exclude "$EXCLUDE_PREFIX" '.results[] | select(.name | startswith($prefix)) | select($exclude == "" or (.name | startswith($exclude) | not)) | select(.name != ($prefix + "x86_64") and .name != ($prefix + "aarch64")) | select(.name | startswith("nightly-dev") | not) | "\(.last_updated)|\(.name)"') if [ -z "$tags" ]; then break diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/scripts/generate-nightly-index.py index 8f92996ff9ab..d7bb1c24e8c3 100644 --- a/.buildkite/scripts/generate-nightly-index.py +++ b/.buildkite/scripts/generate-nightly-index.py @@ -12,7 +12,7 @@ from datetime import datetime from pathlib import Path from typing import Any -from urllib.parse import quote +from urllib.parse import quote, unquote, urlparse import regex as re @@ -51,6 +51,8 @@ class WheelFileInfo: platform_tag: str variant: str | None filename: str + # Absolute URL for files hosted elsewhere (see --external-links) + url: str | None = None def parse_from_filename(file: str) -> WheelFileInfo: @@ -110,6 +112,29 @@ def parse_from_filename(file: str) -> WheelFileInfo: ) +def parse_external_url(url: str) -> WheelFileInfo: + """Parse an absolute wheel or sdist URL hosted on another index.""" + filename = unquote(urlparse(url).path.rsplit("/", 1)[-1]) + if filename.endswith(".whl"): + info = parse_from_filename(filename) + else: + match = re.match(r"^(?P.+)-(?P[^-]+)\.(tar\.gz|zip)$", filename) + if not match: + raise ValueError(f"Unsupported external file: {filename}") + info = WheelFileInfo( + package_name=match.group("name"), + version=match.group("version"), + build_tag=None, + python_tag="source", + abi_tag="none", + platform_tag="any", + variant=None, + filename=filename, + ) + info.url = url + return info + + def generate_project_list(subdir_names: list[str], comment: str = "") -> str: """Generate project list HTML content linking to each project & variant subdirectory.""" href_tags = [] @@ -129,14 +154,19 @@ def generate_package_index_and_metadata( href_tags = [] metadata = [] for file in sorted(wheel_files, key=lambda x: x.filename): - relative_path = ( - wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename - ) - # handle with '+' in URL, and avoid double-encoding '/' and already-encoded '%2B' - # NOTE: this is AWS S3 specific behavior! - file_path_quoted = quote(relative_path.as_posix(), safe=":%/") + if file.url: + file_path_quoted = file.url + else: + relative_path = ( + wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename + ) + # handle with '+' in URL, and avoid double-encoding '/' and already-encoded '%2B' + # NOTE: this is AWS S3 specific behavior! + file_path_quoted = quote(relative_path.as_posix(), safe=":%/") href_tags.append(f' {file.filename}
') file_meta = asdict(file) + if file_meta["url"] is None: + del file_meta["url"] file_meta["path"] = file_path_quoted metadata.append(file_meta) index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment) @@ -151,6 +181,7 @@ def generate_index_and_metadata( default_variant: str | None = None, alias_to_default: str | None = None, comment: str = "", + external_urls: list[str] | None = None, ): """Generate index for all wheel files. @@ -160,6 +191,8 @@ def generate_index_and_metadata( index_base_dir (Path): Base directory to store index files. default_variant (str | None): The default variant name, if any. alias_to_default (str | None): Alias variant name for the default variant, if any. + external_urls (list[str] | None): Absolute URLs of files hosted on another index, + listed in the generated indices without being copied. comment (str | None): Optional comment to include in the generated HTML files. First, parse all wheel files to extract metadata. @@ -214,6 +247,7 @@ def generate_index_and_metadata( """ parsed_files = [parse_from_filename(f) for f in whl_files] + parsed_files += [parse_external_url(u) for u in external_urls or []] if not parsed_files: print("No wheel files found, skipping index generation.") @@ -232,10 +266,14 @@ def generate_index_and_metadata( print(f"Detected ROCm variant from vllm: {rocm_variant}") break - # Apply ROCm variant to all wheels without a variant + # Apply ROCm variant to all wheels without a variant. Dependencies with their + # own ROCm local version (e.g. TheRock's torch-2.12.0+rocm10.0.0) belong to + # the same variant rather than one of their own. if rocm_variant: for file in parsed_files: - if file.variant is None: + if file.variant is None or ( + file.variant.startswith("rocm") and file.variant != rocm_variant + ): file.variant = rocm_variant print(f"Inherited variant '{rocm_variant}' for {file.filename}") @@ -366,6 +404,13 @@ def generate_index_and_metadata( default=None, help="Alias variant name for the default variant", ) + parser.add_argument( + "--external-links", + type=str, + default=None, + help="Optional file with one absolute URL per line for packages hosted on " + "another index (listed alongside, but not copied into, this index)", + ) parser.add_argument( "--comment", type=str, @@ -463,5 +508,14 @@ def generate_index_and_metadata( default_variant=None, alias_to_default=args.alias_to_default, comment=args.comment.strip(), + external_urls=( + [ + line.strip() + for line in Path(args.external_links).read_text().splitlines() + if line.strip() and not line.startswith("#") + ] + if args.external_links + else None + ), ) print(f"Successfully generated index and metadata in {output_dir}") diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index dec366390b00..81dee9a4df5f 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -11,9 +11,9 @@ set -euo pipefail TARGET="${1:-all}" case "${TARGET}" in cuda-13-0 | cuda-12-9 | cuda-13-0-ubuntu-24-04 | \ - cuda-12-9-ubuntu-24-04 | rocm | xpu | cpu | all) ;; + cuda-12-9-ubuntu-24-04 | rocm | rocm72 | xpu | cpu | all) ;; *) - echo "Usage: $0 {cuda-13-0|cuda-12-9|cuda-13-0-ubuntu-24-04|cuda-12-9-ubuntu-24-04|rocm|xpu|cpu|all}" + echo "Usage: $0 {cuda-13-0|cuda-12-9|cuda-13-0-ubuntu-24-04|cuda-12-9-ubuntu-24-04|rocm|rocm72|xpu|cpu|all}" exit 2 ;; esac @@ -143,23 +143,35 @@ if target_enabled cuda-12-9-ubuntu-24-04; then fi # ---- ROCm ---- +# Each ROCm stack is identified by its base Dockerfile (see rocm/stack.sh). The +# default stack also owns the plain :latest / :v tags. + +publish_rocm_stack() { + local base_dockerfile="$1" default="$2" suffix + ( + # shellcheck source=.buildkite/scripts/rocm/stack.sh + source .buildkite/scripts/rocm/stack.sh "$base_dockerfile" + local suffixes=("-${ROCM_STACK_VARIANT}") + [[ "$default" == "1" ]] && suffixes=("" "-${ROCM_STACK_VARIANT}") + docker pull "$ROCM_STACK_ECR_IMAGE" + docker pull "$ROCM_STACK_ECR_BASE" + for suffix in "${suffixes[@]}"; do + for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do + docker tag "$ROCM_STACK_ECR_IMAGE" "vllm/vllm-openai-rocm:${tag}" + docker tag "$ROCM_STACK_ECR_BASE" "vllm/vllm-openai-rocm:${tag}-base" + docker push "vllm/vllm-openai-rocm:${tag}" + docker push "vllm/vllm-openai-rocm:${tag}-base" + done + done + ) +} if target_enabled rocm; then - ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - echo "ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" - - docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" - docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" - - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" vllm/vllm-openai-rocm:latest - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm" "vllm/vllm-openai-rocm:v${RELEASE_VERSION}" - docker push vllm/vllm-openai-rocm:latest - docker push "vllm/vllm-openai-rocm:v${RELEASE_VERSION}" + publish_rocm_stack docker/Dockerfile.rocm_base 1 +fi - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" vllm/vllm-openai-rocm:latest-base - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base" "vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base" - docker push vllm/vllm-openai-rocm:latest-base - docker push "vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base" +if target_enabled rocm72; then + publish_rocm_stack docker/Dockerfile.rocm_72_base 0 fi # ---- XPU ---- diff --git a/.buildkite/scripts/push-nightly-builds-rocm.sh b/.buildkite/scripts/push-nightly-builds-rocm.sh index 2a1a0b080491..8b120883eff4 100644 --- a/.buildkite/scripts/push-nightly-builds-rocm.sh +++ b/.buildkite/scripts/push-nightly-builds-rocm.sh @@ -2,42 +2,36 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Push a ROCm-family nightly base image and nightly image from ECR to Docker Hub -# under vllm/vllm-openai-rocm, as base-nightly[-], nightly[-] -# and their - forms. -# Run when NIGHTLY=1 after the matching build-*-release-image step has pushed to ECR. +# Push a ROCm stack's nightly base image and nightly image from ECR to Docker Hub +# under vllm/vllm-openai-rocm, as base-nightly-, nightly- and their +# - forms, then prune old nightly tags. With --default, also push the plain +# base-nightly and nightly tags for the same images. +# Run when NIGHTLY=1 after the stack's release image has been pushed to ECR. # -# Usage: push-nightly-builds-rocm.sh [ECR_TAG_SUFFIX TAG_VARIANT] -# Pass both arguments or neither. With none, the default ROCm apt build is -# published as :nightly. The TheRock build passes "rock" (the source ECR tag -# suffix) and "rocm714" (the Docker Hub tag flavor), giving :nightly-rocm714. +# Usage: push-nightly-builds-rocm.sh [--default] +# e.g. "docker/Dockerfile.rocm_base --default" (ROCm 10.0, also :nightly) +# and "docker/Dockerfile.rocm_72_base" (ROCm 7.2, :nightly-rocm72). # # Local testing (no push to Docker Hub): -# BASE_ECR_IMAGE= \ -# BUILDKITE_COMMIT= DRY_RUN=1 \ -# bash .buildkite/scripts/push-nightly-builds-rocm.sh +# BUILDKITE_COMMIT= DRY_RUN=1 \ +# bash .buildkite/scripts/push-nightly-builds-rocm.sh docker/Dockerfile.rocm_base --default # Requires: AWS CLI configured (for ECR public login), Docker. For full run: Docker Hub login. set -euxo pipefail usage() { - echo "Usage: $0 [ECR_TAG_SUFFIX TAG_VARIANT]" >&2 + echo "Usage: $0 [--default]" >&2 } -DOCKERHUB_REPO="vllm/vllm-openai-rocm" - +PUSH_DEFAULT=0 case "$#" in - 0) - ECR_TAG_SUFFIX="rocm" - TAG_VARIANT="" - ;; + 1) ;; 2) - if [[ -z "$1" || -z "$2" ]]; then + if [[ "$2" != "--default" ]]; then usage exit 2 fi - ECR_TAG_SUFFIX="$1" - TAG_VARIANT="$2" + PUSH_DEFAULT=1 ;; *) usage @@ -49,61 +43,44 @@ esac BUILDKITE_COMMIT="${BUILDKITE_COMMIT:?Set BUILDKITE_COMMIT to the commit SHA that has the image in ECR (e.g. from a previous release pipeline run)}" DRY_RUN="${DRY_RUN:-0}" -ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-release-repo" +# shellcheck source=.buildkite/scripts/rocm/stack.sh +source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "$1" +DOCKERHUB_REPO="vllm/vllm-openai-rocm" +VARIANT="$ROCM_STACK_VARIANT" -# Get the base image ECR tag (set by the build-*-release-image pipeline step) -BASE_METADATA_KEY="${ECR_TAG_SUFFIX}-base-ecr-tag" -if [[ -n "${BASE_ECR_IMAGE:-}" ]]; then - if [[ "$DRY_RUN" != "1" ]]; then - echo "ERROR: BASE_ECR_IMAGE may only be used with DRY_RUN=1" >&2 - exit 1 - fi - BASE_ORIG_TAG="$BASE_ECR_IMAGE" -elif ! BASE_ORIG_TAG="$(buildkite-agent meta-data get "$BASE_METADATA_KEY")"; then - echo "ERROR: Failed to read required Buildkite metadata '$BASE_METADATA_KEY'" >&2 - exit 1 +BASE_TAGS=("base-nightly-${VARIANT}" "base-nightly-${VARIANT}-${BUILDKITE_COMMIT}") +TAGS=("nightly-${VARIANT}" "nightly-${VARIANT}-${BUILDKITE_COMMIT}") +if [[ "$PUSH_DEFAULT" == "1" ]]; then + BASE_TAGS+=("base-nightly" "base-nightly-${BUILDKITE_COMMIT}") + TAGS+=("nightly" "nightly-${BUILDKITE_COMMIT}") fi -if [[ -z "$BASE_ORIG_TAG" ]]; then - echo "ERROR: Required Buildkite metadata '$BASE_METADATA_KEY' is empty" >&2 - echo "Set BASE_ECR_IMAGE to the full ECR base image reference for local testing" >&2 - exit 1 -fi - -ORIG_TAG="${BUILDKITE_COMMIT}-${ECR_TAG_SUFFIX}" -VARIANT_SUFFIX="" -[[ -n "$TAG_VARIANT" ]] && VARIANT_SUFFIX="-$TAG_VARIANT" -BASE_TAG_NAME="base-nightly${VARIANT_SUFFIX}" -TAG_NAME="nightly${VARIANT_SUFFIX}" -BASE_TAG_NAME_COMMIT="base-nightly${VARIANT_SUFFIX}-${BUILDKITE_COMMIT}" -TAG_NAME_COMMIT="nightly${VARIANT_SUFFIX}-${BUILDKITE_COMMIT}" -echo "Pushing base image from ECR: $BASE_ORIG_TAG" -echo "Pushing release image from ECR tag: $ORIG_TAG to Docker Hub as $DOCKERHUB_REPO:$TAG_NAME and $DOCKERHUB_REPO:$TAG_NAME_COMMIT" -[[ "$DRY_RUN" == "1" ]] && echo "[DRY_RUN] Skipping push to Docker Hub" +echo "Pushing base image $ROCM_STACK_ECR_BASE as: ${BASE_TAGS[*]}" +echo "Pushing release image $ROCM_STACK_ECR_IMAGE as: ${TAGS[*]}" -# Login to ECR and pull the images built by build-*-release-image +# Login to ECR and pull the stack's images aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 -docker pull "$BASE_ORIG_TAG" -docker pull "${ECR_REPO}:${ORIG_TAG}" +docker pull "$ROCM_STACK_ECR_BASE" +docker pull "$ROCM_STACK_ECR_IMAGE" -# Tag for Docker Hub (base-nightly and base-nightly-, nightly and nightly-) -docker tag "$BASE_ORIG_TAG" "$DOCKERHUB_REPO":"$BASE_TAG_NAME" -docker tag "$BASE_ORIG_TAG" "$DOCKERHUB_REPO":"$BASE_TAG_NAME_COMMIT" -docker tag "${ECR_REPO}:${ORIG_TAG}" "$DOCKERHUB_REPO":"$TAG_NAME" -docker tag "${ECR_REPO}:${ORIG_TAG}" "$DOCKERHUB_REPO":"$TAG_NAME_COMMIT" +for tag in "${BASE_TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_BASE" "$DOCKERHUB_REPO:$tag"; done +for tag in "${TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_IMAGE" "$DOCKERHUB_REPO:$tag"; done if [[ "$DRY_RUN" == "1" ]]; then - echo "[DRY_RUN] Would push $DOCKERHUB_REPO:$BASE_TAG_NAME and $DOCKERHUB_REPO:$BASE_TAG_NAME_COMMIT" - echo "[DRY_RUN] Would push $DOCKERHUB_REPO:$TAG_NAME and $DOCKERHUB_REPO:$TAG_NAME_COMMIT" echo "[DRY_RUN] Local tags created. Exiting without push." exit 0 fi # Push to Docker Hub (docker-login plugin runs before this step in CI) -docker push "$DOCKERHUB_REPO":"$BASE_TAG_NAME" -docker push "$DOCKERHUB_REPO":"$BASE_TAG_NAME_COMMIT" -docker push "$DOCKERHUB_REPO":"$TAG_NAME" -docker push "$DOCKERHUB_REPO":"$TAG_NAME_COMMIT" +for tag in "${BASE_TAGS[@]}" "${TAGS[@]}"; do docker push "$DOCKERHUB_REPO:$tag"; done +echo "Pushed $DOCKERHUB_REPO: ${BASE_TAGS[*]} ${TAGS[*]}" -echo "Pushed $DOCKERHUB_REPO:$BASE_TAG_NAME and $DOCKERHUB_REPO:$BASE_TAG_NAME_COMMIT" -echo "Pushed $DOCKERHUB_REPO:$TAG_NAME and $DOCKERHUB_REPO:$TAG_NAME_COMMIT" +# Keep only the last 14 builds of each tag family. The plain nightly- families +# exclude the variant ones sharing their prefix; commit SHAs never start with "rocm". +CLEANUP=.buildkite/scripts/cleanup-nightly-builds.sh +bash "$CLEANUP" "nightly-${VARIANT}-" "$DOCKERHUB_REPO" +bash "$CLEANUP" "base-nightly-${VARIANT}-" "$DOCKERHUB_REPO" +if [[ "$PUSH_DEFAULT" == "1" ]]; then + bash "$CLEANUP" "nightly-" "$DOCKERHUB_REPO" "nightly-rocm" + bash "$CLEANUP" "base-nightly-" "$DOCKERHUB_REPO" "base-nightly-rocm" +fi diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh index ad0179b16359..f6c6ff56ae38 100644 --- a/.buildkite/scripts/rocm/build-config.sh +++ b/.buildkite/scripts/rocm/build-config.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash # Shared Dockerfile selection for the normal AMD CI build chain. # CI_ROCM_DOCKERFILE_BASE and CI_ROCM_DOCKERFILE accept repository-relative paths. -# For Rock, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rock_base and -# CI_ROCM_DOCKERFILE=docker/Dockerfile.rock. +# For the legacy ROCm 7.2 stack, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rocm_72_base +# and CI_ROCM_DOCKERFILE=docker/Dockerfile.rocm_72. # Custom recipes must preserve the CI targets and artifact layout. Additional # COPY inputs also need the corresponding *_CONTENT_FILES overrides. diff --git a/.buildkite/scripts/rocm/stack.sh b/.buildkite/scripts/rocm/stack.sh new file mode 100644 index 000000000000..14bfeb48a017 --- /dev/null +++ b/.buildkite/scripts/rocm/stack.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Derive everything the release pipeline needs about a ROCm stack from its base +# Dockerfile, so the stack's image is the single input to every ROCm script. +# +# Usage: source .buildkite/scripts/rocm/stack.sh +# e.g. docker/Dockerfile.rocm_base (ROCm 10.0) or docker/Dockerfile.rocm_72_base +# +# Exports: +# ROCM_STACK_BASE_DOCKERFILE the base Dockerfile +# ROCM_STACK_DOCKERFILE the matching vLLM Dockerfile (base name minus _base) +# ROCM_STACK_THEROCK 1 if the ROCm SDK comes from TheRock wheels +# ROCM_STACK_VERSION ROCm version, e.g. 10.0.0 or 7.2.3 +# ROCM_STACK_VARIANT wheel variant and Docker Hub tag flavor from major.minor, +# same rule as setup.py (rocm100, rocm72) +# ROCM_STACK_IMAGE_KEY hash of the base Dockerfile (base image cache key) +# ROCM_STACK_WHEEL_KEY hash of the base Dockerfile plus files it COPYs from +# the repo (base wheel cache key) +# ROCM_STACK_ECR_BASE ECR tag of the cached base image +# ROCM_STACK_ECR_IMAGE ECR tag of this commit's release image +# ROCM_STACK_BASE_WHEELS_DIR local dir for base wheels (also the artifact path) +# ROCM_STACK_VLLM_WHEEL_DIR local dir for the vLLM wheel (also the artifact path) + +rocm_stack_init() { + local base="$1" + if [[ ! -f "$base" ]]; then + echo "ROCm base Dockerfile not found: $base" >&2 + return 2 + fi + ROCM_STACK_BASE_DOCKERFILE="$base" + ROCM_STACK_DOCKERFILE="${base%_base}" + + local sdk + sdk=$(sed -nE 's/^ARG ROCM_SDK_VERSION=([0-9.]+).*/\1/p' "$base" | head -1) + if [[ -n "$sdk" ]]; then + ROCM_STACK_THEROCK=1 + ROCM_STACK_VERSION="$sdk" + else + ROCM_STACK_THEROCK=0 + # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.2.3-complete -> 7.2.3 + ROCM_STACK_VERSION=$(sed -nE 's/^ARG BASE_IMAGE=.*:([0-9]+(\.[0-9]+)+).*/\1/p' "$base" | head -1) + fi + if [[ -z "$ROCM_STACK_VERSION" ]]; then + echo "Could not determine the ROCm version from $base" >&2 + return 2 + fi + + local major minor + IFS=. read -r major minor _ <<< "$ROCM_STACK_VERSION" + ROCM_STACK_VARIANT="rocm${major}${minor}" + + ROCM_STACK_IMAGE_KEY=$(sha256sum "$base" | cut -c1-16) + local inputs + inputs=$(grep -E '^COPY ' "$base" | grep -v -- '--from' | awk '{print $2}' || true) + if [[ -z "$inputs" ]]; then + ROCM_STACK_WHEEL_KEY="$ROCM_STACK_IMAGE_KEY" + else + # shellcheck disable=SC2086 + ROCM_STACK_WHEEL_KEY=$(sha256sum "$base" $inputs | sha256sum | cut -c1-16) + fi + + local ecr="public.ecr.aws/q9t5s3a7/vllm-release-repo" + ROCM_STACK_ECR_BASE="${ecr}:${ROCM_STACK_IMAGE_KEY}-${ROCM_STACK_VARIANT}-base" + ROCM_STACK_ECR_IMAGE="${ecr}:${BUILDKITE_COMMIT:-unknown}-${ROCM_STACK_VARIANT}" + ROCM_STACK_BASE_WHEELS_DIR="artifacts/${ROCM_STACK_VARIANT}-base-wheels" + ROCM_STACK_VLLM_WHEEL_DIR="artifacts/${ROCM_STACK_VARIANT}-vllm-wheel" + + export ROCM_STACK_BASE_DOCKERFILE ROCM_STACK_DOCKERFILE ROCM_STACK_THEROCK \ + ROCM_STACK_VERSION ROCM_STACK_VARIANT ROCM_STACK_IMAGE_KEY \ + ROCM_STACK_WHEEL_KEY ROCM_STACK_ECR_BASE ROCM_STACK_ECR_IMAGE \ + ROCM_STACK_BASE_WHEELS_DIR ROCM_STACK_VLLM_WHEEL_DIR +} + +if [[ $# -gt 0 ]]; then + rocm_stack_init "$1" +fi diff --git a/.buildkite/scripts/upload-rocm-wheels.sh b/.buildkite/scripts/upload-rocm-wheels.sh index 65fca944c7d5..5cf32bcfb358 100755 --- a/.buildkite/scripts/upload-rocm-wheels.sh +++ b/.buildkite/scripts/upload-rocm-wheels.sh @@ -4,21 +4,31 @@ # # Upload ROCm wheels to S3 with proper index generation # +# Usage: upload-rocm-wheels.sh +# base-dockerfile the ROCm stack, e.g. docker/Dockerfile.rocm_base (see rocm/stack.sh) +# # Required environment variables: # AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (or IAM role) # S3_BUCKET (default: vllm-wheels) # # S3 path structure: -# s3://vllm-wheels/rocm/{commit}/ - All wheels for this commit -# s3://vllm-wheels/rocm/nightly/ - Index pointing to latest nightly -# s3://vllm-wheels/rocm/{version}/ - Index for release versions +# s3://vllm-wheels/rocm/{commit}/-wheels/ - This stack's wheels for this commit +# s3://vllm-wheels/rocm/{commit}// - Index per ROCm variant (e.g. rocm100) +# s3://vllm-wheels/rocm/nightly// - Index pointing to latest nightly +# s3://vllm-wheels/rocm/{version}// - Index for release versions +# The top-level index.html at each prefix lists every variant present, since +# several ROCm stacks publish into the same prefixes. set -ex # ======== Configuration ======== BUCKET="${S3_BUCKET:-vllm-wheels}" +# shellcheck source=.buildkite/scripts/rocm/stack.sh +source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 }" ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}" -S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/" +WHEEL_DIRS="$ROCM_STACK_BASE_WHEELS_DIR $ROCM_STACK_VLLM_WHEEL_DIR" +WHEEL_SUBPATH="$ROCM_SUBPATH/${ROCM_STACK_VARIANT}-wheels" +S3_WHEEL_PREFIX="s3://$BUCKET/$WHEEL_SUBPATH/" INDICES_OUTPUT_DIR="rocm-indices" echo "========================================" @@ -49,9 +59,11 @@ source .buildkite/scripts/lib/manylinux.sh # ======== Part 1: Collect and prepare wheels ======== # Collect all wheels +rm -rf all-rocm-wheels "$INDICES_OUTPUT_DIR" mkdir -p all-rocm-wheels -cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true -cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true +for dir in $WHEEL_DIRS; do + cp "$dir"/*.whl all-rocm-wheels/ 2>/dev/null || true +done WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) echo "Total wheels to upload: $WHEEL_COUNT" @@ -84,9 +96,9 @@ ls -lh all-rocm-wheels/ # ======== Part 2: Upload wheels to S3 ======== echo "" -echo "Uploading wheels to $S3_COMMIT_PREFIX" +echo "Uploading wheels to $S3_WHEEL_PREFIX" for wheel in all-rocm-wheels/*.whl; do - aws s3 cp "$wheel" "$S3_COMMIT_PREFIX" + aws s3 cp "$wheel" "$S3_WHEEL_PREFIX" done # ======== Part 3: Generate and upload indices ======== @@ -95,28 +107,63 @@ done echo "" echo "Generating indices..." obj_json="rocm-objects.json" -aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$ROCM_SUBPATH/" --delimiter / --output json > "$obj_json" +aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$WHEEL_SUBPATH/" --delimiter / --output json > "$obj_json" mkdir -p "$INDICES_OUTPUT_DIR" # Use the existing generate-nightly-index.py # HACK: Replace regex module with stdlib re (same as CUDA script) sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py - +sed -i 's/import regex as re/import re/g' tools/vllm-rocm/therock_wheels.py +sed -i 's/import pybase64 as base64/import base64/g' tools/vllm-rocm/therock_wheels.py + +INDEX_ARGS=(--version "$ROCM_SUBPATH" --wheel-dir "${WHEEL_SUBPATH#rocm/}") +# TheRock's SDK, device kernels and torchvision/torchaudio are linked, not hosted +if [[ "$ROCM_STACK_THEROCK" == "1" ]]; then + $PYTHON tools/vllm-rocm/therock_wheels.py external-links \ + --dockerfile "$ROCM_STACK_BASE_DOCKERFILE" \ + --torch-wheel "$(ls "$ROCM_STACK_BASE_WHEELS_DIR"/torch-*.whl)" > therock-external-links.txt + cat therock-external-links.txt + INDEX_ARGS+=(--external-links therock-external-links.txt) +fi $PYTHON .buildkite/scripts/generate-nightly-index.py \ - --version "$ROCM_SUBPATH" \ + "${INDEX_ARGS[@]}" \ --current-objects "$obj_json" \ --output-dir "$INDICES_OUTPUT_DIR" \ --comment "ROCm commit $BUILDKITE_COMMIT" -# Upload indices to commit directory -echo "Uploading indices to $S3_COMMIT_PREFIX" -aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX" +# Upload this build's variant indices under an S3 prefix (e.g. rocm/nightly), then +# rebuild that prefix's top-level project list from every variant now present. +publish_indices() { + local prefix="$1" + local listing variants merged + echo "Uploading indices to s3://$BUCKET/$prefix/" + aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$prefix/" --exclude "index.html" + listing=$(aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$prefix/" \ + --delimiter / --query 'CommonPrefixes[].Prefix' --output text) + variants=$(tr '\t' '\n' <<< "$listing" | sed -n "s|^$prefix/\(rocm[0-9]*\)/$|\1|p" | sort -u) + merged=$(mktemp) + { + echo '' + echo '' + echo ' ' + echo ' ' + for variant in $variants; do + echo " $variant/
" + done + echo ' ' + echo '' + } > "$merged" + echo "Variants under $prefix/: $(echo "$variants" | tr '\n' ' ')" + aws s3 cp "$merged" "s3://$BUCKET/$prefix/index.html" + rm -f "$merged" +} + +publish_indices "$ROCM_SUBPATH" # Only scheduled nightly builds should update the moving nightly index. if [[ "${NIGHTLY:-0}" == "1" ]]; then - echo "Updating rocm/nightly/ index..." - aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/" + publish_indices "rocm/nightly" fi # Extract version from vLLM wheel and update version-specific index @@ -129,8 +176,7 @@ if [ -n "$VLLM_WHEEL" ]; then echo "Pure version: $PURE_VERSION" if [[ "$VERSION" != *"dev"* ]]; then - echo "Updating rocm/$PURE_VERSION/ index..." - aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/$PURE_VERSION/" + publish_indices "rocm/$PURE_VERSION" fi fi @@ -142,14 +188,14 @@ echo "ROCm Wheel Upload Complete!" echo "========================================" echo "" echo "Wheels available at:" -echo " s3://$BUCKET/$ROCM_SUBPATH/" +echo " s3://$BUCKET/$WHEEL_SUBPATH/" echo "" echo "Install command (by commit):" -echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/" +echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/$ROCM_STACK_VARIANT/" echo "" if [[ "${NIGHTLY:-0}" == "1" ]]; then echo "Install command (nightly):" - echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/" + echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/$ROCM_STACK_VARIANT/" fi echo "" echo "Wheel count: $WHEEL_COUNT" diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base deleted file mode 100644 index 57bf6ae6febe..000000000000 --- a/docker/Dockerfile.rock_base +++ /dev/null @@ -1,453 +0,0 @@ -ARG BASE_IMAGE=ubuntu:22.04 -ARG PYTHON_VERSION=3.12 -ARG TORCH_VERSION=2.12.0+rocm10.0.0 -ARG TORCHVISION_VERSION=0.27.0+rocm10.0.0 -ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0 -ARG TRITON_BRANCH="669b31ac" # release/internal/3.8.x as of 09/29/2026 -ARG TRITON_VERSION=3.8.0+git669b31ac -ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG FA_BRANCH="0e60e394" -ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.25" -ARG AITER_REPO="https://github.com/ROCm/aiter.git" -ARG MORI_BRANCH="v1.2.3.post1" -ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" -ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2" - - -# TheRock configuration -ARG SITE_PACKAGES=/usr/local/lib/python${PYTHON_VERSION}/dist-packages -ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL="https://stable.repo.amd.com/rocm/whl-next/" # For pre-built GPU images, packages, etc. -ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL="https://nightly.repo.amd.com/rocm/whl-next/" # For torch, triton, etc. wheels with latest fixes. -ARG ROCM_SDK_VERSION=10.0.0 - -# Build configuration -ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" -ARG AITER_ROCM_ARCH="gfx942;gfx950" - -# Sccache configuration (only used in release pipeline) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - -FROM ${BASE_IMAGE} AS base - -ARG PYTORCH_ROCM_ARCH -ENV PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" -ARG AITER_ROCM_ARCH -ENV AITER_ROCM_ARCH="${AITER_ROCM_ARCH}" - -ARG SITE_PACKAGES -ENV SITE_PACKAGES=${SITE_PACKAGES} -ENV SDK_DEV="$SITE_PACKAGES/_rocm_sdk_devel" \ - SDK_CORE="$SITE_PACKAGES/_rocm_sdk_core" -ENV ROCM_PATH="$SDK_DEV" \ - ROCM_HOME="$SDK_DEV" \ - HIP_PATH="$SDK_DEV" \ - HIP_CLANG_PATH="$SDK_DEV/lib/llvm/bin" \ - HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode" -ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib" -ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH" - -# Required for RCCL in ROCm7.1 -ENV HSA_NO_SCRATCH_RECLAIM=1 - -ARG PYTHON_VERSION -ENV PYTHON_VERSION=${PYTHON_VERSION} - -RUN mkdir -p /app -WORKDIR /app -ENV DEBIAN_FRONTEND=noninteractive - -# Install Python and other dependencies -RUN apt-get update -y \ - && apt-get install -y software-properties-common build-essential git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config libdrm-dev libnuma-dev kmod \ - && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ - done \ - && apt-get update -y \ - && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ - && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ - && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ - && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version - -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython -RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* - -# Install ROCm SDK via TheRock -ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL -ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL -ARG ROCM_SDK_VERSION -RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ - --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ - "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \ - && rocm-sdk init - -# Install PyTorch wheels (nightly for latest torch, release as fallback) -ARG TORCH_VERSION -ARG TORCHAUDIO_VERSION -ARG TORCHVISION_VERSION -RUN pip install \ - --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ - --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ - "torch==${TORCH_VERSION}" \ - "torchaudio==${TORCHAUDIO_VERSION}" \ - "torchvision==${TORCHVISION_VERSION}" - -# Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH: - -# Possible values that can be in PYTORCH_ROCM_ARCH: -# -# gfx90a (MI210/250), gfx942 (MI300), gfx950 (MI350), gfx1100 (7900XTX), -# gfx1101 (7800XT), gfx1200 (9070XT), gfx1201 (9070), gfx1150/1151 (Strix APU) - -# For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that -# PyTorch needs to run on that specific chip. -SHELL ["/bin/bash", "-c"] -RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ - && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \ - echo "Installing device packages for ${arch}" \ - && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ - --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ - amd-torch-device-${arch}==${TORCH_VER} \ - rocm-sdk-device-${arch}==${ROCM_SDK_VERSION}; \ - done - -# Remove gfx90a xnack kpacks (ASAN artifacts that cause hipErrorInvalidImage on MI250) -RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \ - "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack" - -RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install . -ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi - -# TheRock runtime fixes -RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ - sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ - fi \ - && if [ -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" ] \ - && [ ! -e "${ROCM_PATH}/lib/librocm_smi64.so.7" ]; then \ - ln -sfn "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" "${ROCM_PATH}/lib/librocm_smi64.so.7"; \ - fi \ - && printf '%s\n%s\n' "${ROCM_PATH}/lib" "${SITE_PACKAGES}/_rocm_sdk_core/lib" > /etc/ld.so.conf.d/rocm-pip.conf \ - && for f in ${SITE_PACKAGES}/_rocm_sdk_core/lib/libamdhip64.so*; do \ - bn=$(basename "$f"); \ - if [ -f "${ROCM_PATH}/lib/${bn}" ] && [ -f "$f" ] && [ ! -L "$f" ]; then \ - rm -f "$f"; \ - ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ - fi; \ - done \ - && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ - librocprofiler-sdk-roctx.so.1; do \ - if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ - && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ - rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ - ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ - "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ - fi; \ - done \ - && ldconfig - -# Install sccache if USE_SCCACHE is enabled (for release builds) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME -ARG SCCACHE_REGION_NAME -ARG SCCACHE_S3_NO_CREDENTIALS -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi - -# Setup sccache for HIP compilation via HIP_CLANG_PATH -# This creates wrapper scripts in a separate directory and points HIP to use them -# This avoids modifying the original ROCm binaries which can break detection -# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) -# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ - fi - -ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} -ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} -ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} -ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} - - -### -### MORI Build -### -FROM base AS build_mori -ARG MORI_BRANCH -ARG MORI_REPO -RUN git clone ${MORI_REPO} -RUN cd mori \ - && git checkout ${MORI_BRANCH} \ - && git submodule update --init --recursive \ - && ln -sfn ${ROCM_PATH} /opt/rocm \ - && pip install "setuptools_scm>=6.2" \ - && CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \ - python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl -RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install - -### -### FlashAttention Build -### -FROM base AS build_fa -ARG FA_BRANCH -ARG FA_REPO -ARG USE_SCCACHE -RUN git clone ${FA_REPO} -RUN cd flash-attention \ - && git checkout ${FA_BRANCH} \ - && git submodule update --init \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install - - -### -### AITER Build -### -FROM base AS build_aiter -ARG AITER_BRANCH -ARG AITER_REPO -ARG USE_SCCACHE -RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} -RUN cd aiter \ - && git submodule update --init --recursive \ - && pip install -r requirements.txt -# AITER_USE_SYSTEM_TRITON keeps setup.py from resolving its own triton, which would -# displace the pinned one already installed in base. -RUN pip install pyyaml && cd aiter \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && ls /app/aiter/dist/*.whl -RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install - - -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && if [ ! -f setup.py ]; then cd python; fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install - - -### -### ROCr + CLR build -### -# Brings in multi-stream perf improvements; drop after these arrive in an official release -FROM base AS build_rocm_runtime -ARG ROCM_SYSTEMS_REPO -ARG ROCM_RUNTIME_COMMIT -# libgl-dev: rocclr does find_package(OpenGL REQUIRED) unconditionally and the wheel image has -# no GL headers. The link must not pick up a libGL DT_NEEDED, which is asserted after the build. -RUN apt-get update \ - && apt-get install -y --no-install-recommends g++ libdw-dev xxd libgl-dev libglx-dev \ - && rm -rf /var/lib/apt/lists/* -RUN pip install --no-cache-dir CppHeaderParser -RUN git init -q /src && cd /src \ - && git remote add origin ${ROCM_SYSTEMS_REPO} \ - && git config core.sparseCheckout true \ - && git config remote.origin.promisor true \ - && git config remote.origin.partialclonefilter blob:none \ - && git sparse-checkout init --cone \ - && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \ - && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \ - && git checkout -q FETCH_HEAD -# libelf is taken from the SDK's own vendored sysdeps rather than apt -- the wheel image has no -# /usr/include/libelf.h -- so ROCr links against the copy the wheel itself was built against. -ENV SYSDEPS="$SDK_DEV/lib/rocm_sysdeps" -# ROCr first, and installed into $SDK_DEV, because the clr build below needs its headers. -RUN cd /src/projects/rocr-runtime \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ - -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ - -DLIBELF_INCLUDE_DIRS="$SYSDEPS/include" \ - -DLIBELF_LIBRARIES="$SYSDEPS/lib/libelf.so" \ - && cmake --build build --parallel 16 \ - && cmake --install build \ - && cmake --install build --prefix /rocr-install --strip -RUN cd /src/projects/clr \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCLR_BUILD_HIP=ON \ - -DCLR_BUILD_OCL=OFF \ - -DROCM_KPACK_ENABLED=ON \ - -DHIP_COMMON_DIR=/src/projects/hip \ - -DROCM_PATH="$SDK_DEV" \ - -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ - -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ - -DHIP_LLVM_ROOT="$SDK_DEV/lib/llvm" \ - && cmake --build build --parallel 16 \ - && cmake --install build --prefix /clr-install --strip -RUN mkdir -p /staging/lib \ - && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \ - && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \ - && ls -l /staging/lib - - -### -### Final Build -### - -# Wheel release stage - -# only includes dependencies used by wheel release pipeline -FROM base AS debs_wheel_release -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs - -FROM base AS debs -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs - -FROM base AS final -ARG TRITON_VERSION - -# torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA -# teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix. -ENV ROCPROFILER_QUEUE_INTERPOSITION=0 - -COPY --from=build_rocm_runtime /staging/lib/ ${SDK_DEV}/lib/ -# Hard link the patched library onto every stock file name, in both SDK trees, so that the -# SONAME lookup and the hard coded paths all resolve to it. Hard links rather than symlinks -# because ldconfig rewrites SONAME symlinks but cannot re-point an inode. $SDK_CORE keeps its -# own libhsa-runtime64 copy, which is on LD_LIBRARY_PATH and in ld.so.conf; relinking it is -# belt and braces rather than strictly required, but it leaves no stock ROCr in the image. -RUN set -eux; \ - cd "$SDK_DEV/lib"; \ - ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \ - ours_hip="$(readlink -f libamdhip64.so.7)"; \ - for f in "$SDK_DEV"/lib/libhsa-runtime64.so.1* "$SDK_CORE"/lib/libhsa-runtime64.so.1* \ - "$SDK_DEV"/lib/libamdhip64.so.7* "$SDK_CORE"/lib/libamdhip64.so.7*; do \ - [ -f "$f" ] && [ ! -L "$f" ] || continue; \ - case "${f##*/}" in libhsa*) ours="$ours_hsa";; *) ours="$ours_hip";; esac; \ - [ "$f" -ef "$ours" ] || ln -f "$ours" "$f"; \ - done - -RUN ldconfig - -# This is to prevent resolvers from checking METADATA and RECORD and clobbering -# the Triton source build we installed. Constraints files are unable to help -# with this. However, these entries are optional and can be blanked which -# solves the problem. -RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \ - && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \ - && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \ - && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD" - -RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - pip install /install/*.whl \ - && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \ - && test "$(pip show triton | awk '/^Version:/ {print $2}')" = "${TRITON_VERSION}" \ - # Create a constraints file to protect the versions we need, especially - # triton, requirements will try to reinstall repeatedly. - && pip list --format=freeze \ - | grep -E '^(torch|triton|torchaudio|torchvision|pybind11)==' > /etc/rocm-constraints.txt \ - && echo "ROCm constraints:" && cat /etc/rocm-constraints.txt \ - && test "$(wc -l < /etc/rocm-constraints.txt)" -eq 5 \ - && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \ - && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \ - && python3 -c "import triton; print(f'triton={triton.__version__}')" \ - && pip list 2>/dev/null | grep -iE '^(torch|triton) ' - -ARG BASE_IMAGE -ARG PYTHON_VERSION -ARG ROCM_SDK_VERSION -ARG TORCH_VERSION -ARG TORCHAUDIO_VERSION -ARG TORCHVISION_VERSION -ARG PYTORCH_ROCM_ARCH -ARG AITER_ROCM_ARCH -ARG TRITON_BRANCH -ARG TRITON_REPO -ARG FA_BRANCH -ARG FA_REPO -ARG AITER_BRANCH -ARG AITER_REPO -ARG MORI_BRANCH -ARG MORI_REPO -ARG ROCM_SYSTEMS_REPO -ARG ROCM_RUNTIME_COMMIT -ARG TRITON_VERSION -# Record hashes for Triton and pybind11 so we can verify later. -RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \ - && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ - && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ - && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ - && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ - && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \ - && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \ - && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ - && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ - && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ - && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ - && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ - && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ - && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ - && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ - # Fingerprint our triton so Dockerfile.rocm can prove it survived the build. - && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ - && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ - -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ - >> /app/versions.txt \ - && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ - && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \ - $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \ - >> /app/versions.txt \ - && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \ - -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ - >> /app/versions.txt diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 36aa0aa4df72..029d003b5542 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -280,8 +280,9 @@ COPY --from=fetch_vllm \ ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \ requirements/ +# Use constraints file to protect triton, pybind11, and torch dependencies. RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -r requirements/rocm.txt + uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt # Fetch only the pinned ROCm Triton kernels package. Keeping this before the # vLLM source COPYs makes the download reusable when native sources change. @@ -322,8 +323,9 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rocm ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && export CCACHE_BASEDIR="$PWD" \ && echo "=== ccache stats before ROCm native build ===" \ && (ccache --show-stats || true) \ @@ -366,7 +368,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools /tools +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ @@ -386,7 +389,6 @@ ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96" ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git" ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194" ARG UCX_REPO="https://github.com/openucx/ucx.git" -ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV NIXL_HOME=/usr/local/nixl @@ -541,7 +543,7 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \ test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \ expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \ - actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \ + actual="$(llvm-objdump --offloading "${c_ops}" \ | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \ if [ "${actual}" != "${expected}" ]; then \ printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \ @@ -549,32 +551,10 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ exit 1; \ fi -# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not -# invalidate the slow ROCShmem build. -FROM base AS build_rocshmem -ARG ROCSHMEM_BRANCH="f0acb0c6" -ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so -# it can be linked against DeepEP without arch mismatches. -ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" -ENV ROCM_PATH=/opt/rocm -ENV ROCSHMEM_DIR=/opt/rocshmem - -RUN --mount=type=cache,target=/root/.cache/ccache \ - git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ - && cd rocm-systems \ - && git sparse-checkout set --cone projects/rocshmem \ - && git checkout ${ROCSHMEM_BRANCH} \ - && mkdir -p projects/rocshmem/build \ - && cd projects/rocshmem/build \ - && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ - bash ../scripts/build_configs/all_backends \ - -DROCM_PATH=${ROCM_PATH} \ - -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ - -DUSE_EXTERNAL_MPI=OFF - -# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. -FROM build_rocshmem AS build_deepep +# DeepEP build stage. ROCShmem is provided by TheRock SDK +# ROCSHMEM_DIR set to the SDK install path so DeepEP finds the headers, etc. +FROM base AS build_deepep +ENV ROCSHMEM_DIR=${ROCM_PATH} ARG DEEPEP_BRANCH="0f63d3e9" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" @@ -693,8 +673,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Install the packaging helpers before modifying requirements/rocm.txt. RUN --mount=type=cache,target=/root/.cache/uv \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ vllm/cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rocm ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && cd vllm \ && uv pip install --system setuptools_scm regex @@ -750,7 +731,8 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requir COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools /tools +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ @@ -795,9 +777,6 @@ RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /nixl_install/*.whl /deep_install/*.whl -# Copy ROCShmem runtime libraries. -COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem - # RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \ -o Acquire::Retries=3 \ @@ -830,7 +809,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Pre-install shared ROCm runtime dependencies. COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + uv pip install --system -c /etc/rocm-constraints.txt \ + -r /tmp/ci-base-requirements/rocm.txt \ && rm -rf /tmp/ci-base-requirements # Enable fast and less brittle model downloads in tests. @@ -908,8 +888,6 @@ RUN mkdir src && mv vllm src/vllm # Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. ENV VLLM_GPU_SYNC_CHECK=error -# Run the structural smoke test inside BuildKit, where the test image layers -# are already available. Export only a build-scoped success marker. FROM test AS test_smoke ARG ROCM_SMOKE_ID RUN --network=none \ @@ -956,15 +934,37 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \ - && uv pip install --system -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ + && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ && pip uninstall -y vllm \ - && uv pip install --system -c /tmp/otel-constraints.txt *.whl \ + && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt *.whl \ && rm /tmp/otel-constraints.txt # Install NIXL ROCm wheel RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ uv pip install --system /nixl_install/*.whl +# Recompute the base's triton fingerprint to detect resolver interference. +RUN EXPECTED=$(grep '^TRITON_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ + && ACTUAL=$(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ + -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ + && echo "triton sha256 expected=${EXPECTED} actual=${ACTUAL}" \ + && test -n "${EXPECTED}" \ + && test "${EXPECTED}" = "${ACTUAL}" \ + && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" + +# Recompute the base's pybind11 fingerprint to detect resolver interference. +RUN EXPECTED=$(grep '^PYBIND11_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ + && EXPECTED_ABI=$(grep '^PYBIND11_INTERNALS_VERSION: ' /app/versions.txt | cut -d' ' -f2) \ + && INC=$(python3 -c 'import pybind11; print(pybind11.get_include())') \ + && ACTUAL=$(find "${INC}" -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ + && ACTUAL_ABI=$(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' "${INC}" | awk '{print $NF}') \ + && echo "pybind11 $(pip show pybind11 | awk '/^Version:/ {print $2}')" \ + "sha256 expected=${EXPECTED} actual=${ACTUAL}" \ + "internals expected=${EXPECTED_ABI} actual=${ACTUAL_ABI}" \ + && test -n "${EXPECTED}" \ + && test "${EXPECTED}" = "${ACTUAL}" \ + && test "${EXPECTED_ABI}" = "${ACTUAL_ABI}" + ARG COMMON_WORKDIR ARG BASE_IMAGE ARG NIC_BACKEND diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rocm_72 similarity index 92% rename from docker/Dockerfile.rock rename to docker/Dockerfile.rocm_72 index ed66b8e85e96..a24b65458ab5 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rocm_72 @@ -3,7 +3,7 @@ ARG REMOTE_VLLM="0" ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" ARG VLLM_BRANCH="main" ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=rocm/vllm-dev:base +ARG BASE_IMAGE=vllm/vllm-openai-rocm:base-nightly-rocm72 ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 ARG INSTALL_LMCACHE=false @@ -46,7 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # Note: mold is installed but not set as the system default linker because # some packages use JIT compilation at runtime with flags mold does not support. # Build stages opt in via LDFLAGS="-fuse-ld=mold". -# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) +# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_72_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ apt-get purge -y sccache || true; \ @@ -280,9 +280,8 @@ COPY --from=fetch_vllm \ ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \ requirements/ -# Use constraints file to protect triton, pybind11, and torch dependencies. RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt + uv pip install --system -r requirements/rocm.txt # Fetch only the pinned ROCm Triton kernels package. Keeping this before the # vLLM source COPYs makes the download reusable when native sources change. @@ -323,9 +322,8 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ cmake/external_projects/triton_kernels.cmake \ - || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && export CCACHE_BASEDIR="$PWD" \ && echo "=== ccache stats before ROCm native build ===" \ && (ccache --show-stats || true) \ @@ -368,12 +366,11 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools /tools COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72 /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl @@ -389,6 +386,7 @@ ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96" ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git" ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194" ARG UCX_REPO="https://github.com/openucx/ucx.git" +ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV NIXL_HOME=/usr/local/nixl @@ -543,7 +541,7 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \ test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \ expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \ - actual="$(llvm-objdump --offloading "${c_ops}" \ + actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \ | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \ if [ "${actual}" != "${expected}" ]; then \ printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \ @@ -551,10 +549,32 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ exit 1; \ fi -# DeepEP build stage. ROCShmem is provided by TheRock SDK -# ROCSHMEM_DIR set to the SDK install path so DeepEP finds the headers, etc. -FROM base AS build_deepep -ENV ROCSHMEM_DIR=${ROCM_PATH} +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem +ARG ROCSHMEM_BRANCH="f0acb0c6" +ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +ENV ROCM_PATH=/opt/rocm +ENV ROCSHMEM_DIR=/opt/rocshmem + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF + +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep ARG DEEPEP_BRANCH="0f63d3e9" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" @@ -571,8 +591,8 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ && git checkout ${DEEPEP_BRANCH} \ && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install -# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do -# not force users to rebuild the long-lived Dockerfile.rocm_base image. +# MoRI runtime dependencies live in Dockerfile.rocm_72 so NIC backend changes do +# not force users to rebuild the long-lived Dockerfile.rocm_72_base image. FROM base AS mori_base ARG NIC_BACKEND ARG AINIC_VERSION @@ -673,9 +693,8 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Install the packaging helpers before modifying requirements/rocm.txt. RUN --mount=type=cache,target=/root/.cache/uv \ - { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ vllm/cmake/external_projects/triton_kernels.cmake \ - || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && cd vllm \ && uv pip install --system setuptools_scm regex @@ -731,12 +750,11 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requir COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools /tools COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72 /docker/ +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl @@ -777,6 +795,9 @@ RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /nixl_install/*.whl /deep_install/*.whl +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem + # RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \ -o Acquire::Retries=3 \ @@ -809,8 +830,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Pre-install shared ROCm runtime dependencies. COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -c /etc/rocm-constraints.txt \ - -r /tmp/ci-base-requirements/rocm.txt \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ && rm -rf /tmp/ci-base-requirements # Enable fast and less brittle model downloads in tests. @@ -888,6 +908,8 @@ RUN mkdir src && mv vllm src/vllm # Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. ENV VLLM_GPU_SYNC_CHECK=error +# Run the structural smoke test inside BuildKit, where the test image layers +# are already available. Export only a build-scoped success marker. FROM test AS test_smoke ARG ROCM_SMOKE_ID RUN --network=none \ @@ -934,37 +956,15 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \ - && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ + && uv pip install --system -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ && pip uninstall -y vllm \ - && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt *.whl \ + && uv pip install --system -c /tmp/otel-constraints.txt *.whl \ && rm /tmp/otel-constraints.txt # Install NIXL ROCm wheel RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ uv pip install --system /nixl_install/*.whl -# Recompute the base's triton fingerprint to detect resolver interference. -RUN EXPECTED=$(grep '^TRITON_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ - && ACTUAL=$(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ - -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ - && echo "triton sha256 expected=${EXPECTED} actual=${ACTUAL}" \ - && test -n "${EXPECTED}" \ - && test "${EXPECTED}" = "${ACTUAL}" \ - && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" - -# Recompute the base's pybind11 fingerprint to detect resolver interference. -RUN EXPECTED=$(grep '^PYBIND11_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ - && EXPECTED_ABI=$(grep '^PYBIND11_INTERNALS_VERSION: ' /app/versions.txt | cut -d' ' -f2) \ - && INC=$(python3 -c 'import pybind11; print(pybind11.get_include())') \ - && ACTUAL=$(find "${INC}" -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ - && ACTUAL_ABI=$(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' "${INC}" | awk '{print $NF}') \ - && echo "pybind11 $(pip show pybind11 | awk '/^Version:/ {print $2}')" \ - "sha256 expected=${EXPECTED} actual=${ACTUAL}" \ - "internals expected=${EXPECTED_ABI} actual=${ACTUAL_ABI}" \ - && test -n "${EXPECTED}" \ - && test "${EXPECTED}" = "${ACTUAL}" \ - && test "${EXPECTED_ABI}" = "${ACTUAL_ABI}" - ARG COMMON_WORKDIR ARG BASE_IMAGE ARG NIC_BACKEND diff --git a/docker/Dockerfile.rocm_72_base b/docker/Dockerfile.rocm_72_base new file mode 100644 index 000000000000..b70434f09012 --- /dev/null +++ b/docker/Dockerfile.rocm_72_base @@ -0,0 +1,479 @@ +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete +ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17 +ARG TRITON_REPO="https://github.com/ROCm/triton.git" +ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17 +ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" +ARG PYTORCH_VISION_BRANCH="v0.28.0" +ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" +ARG PYTORCH_AUDIO_BRANCH="v2.11.0" +ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" +ARG FA_BRANCH="0e60e394" +ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" +ARG AITER_BRANCH="v0.1.25" +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG MORI_BRANCH="v1.2.3.post1" +ARG MORI_REPO="https://github.com/ROCm/mori.git" +ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" +ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350" +ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521" + +# Sccache configuration (only used in release pipeline) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + + +FROM ${BASE_IMAGE} AS base + +ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENV ROCM_PATH=/opt/rocm +ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: + +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV AITER_ROCM_ARCH=gfx942;gfx950 +# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime + +# Required for RCCL in ROCm7.1 +ENV HSA_NO_SCRATCH_RECLAIM=1 + +ARG PYTHON_VERSION=3.12 +ENV PYTHON_VERSION=${PYTHON_VERSION} + +RUN mkdir -p /app +WORKDIR /app +ENV DEBIAN_FRONTEND=noninteractive + +# Install Python and other dependencies (plus gcc-13, selected later in build_pytorch). +RUN apt-get update -y \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ + && python3 --version && python3 -m pip --version + +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython +RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +# Setup sccache for HIP compilation via HIP_CLANG_PATH +# This creates wrapper scripts in a separate directory and points HIP to use them +# This avoids modifying the original ROCm binaries which can break detection +# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm_72) +# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + + +### +### Triton Build +### +FROM base AS build_triton +ARG TRITON_BRANCH +ARG TRITON_REPO +RUN git clone ${TRITON_REPO} +RUN cd triton \ + && git checkout ${TRITON_BRANCH} \ + && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ + && if [ ! -f setup.py ]; then cd python; fi \ + && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \ + && mkdir -p /app/install && cp dist/*.whl /app/install + + +### +### AMD SMI Build +### +FROM base AS build_amdsmi +RUN cd /opt/rocm/share/amd_smi \ + && pip wheel . --wheel-dir=dist +RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install + + +### +### Pytorch build +### +FROM base AS build_pytorch +ARG PYTORCH_BRANCH +ARG PYTORCH_REPO +ARG USE_SCCACHE + +RUN apt-get update && apt-get install -y pkg-config liblzma-dev +# torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13. +RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 +RUN git clone ${PYTORCH_REPO} pytorch +RUN cd pytorch && git checkout ${PYTORCH_BRANCH} +RUN cd pytorch \ + && pip install -r requirements.txt && git submodule update --init --recursive +RUN cd pytorch && python3 tools/amd_build/build_amd.py \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ + && sccache --show-stats; \ + fi \ + && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi +RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install + + +### +### Shared PyTorch runtime for dependent builds +### +FROM base AS build_pytorch_runtime +RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install,ro \ + pip install /install/*.whl + + +### +### TorchVision build +### +FROM build_pytorch_runtime AS build_torchvision +ARG PYTORCH_VISION_BRANCH +ARG PYTORCH_VISION_REPO +ARG USE_SCCACHE + +RUN git clone ${PYTORCH_VISION_REPO} vision +RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi +RUN mkdir -p /app/install && cp /app/vision/dist/*.whl /app/install + + +### +### TorchAudio build +### +FROM build_pytorch_runtime AS build_torchaudio +ARG PYTORCH_AUDIO_BRANCH +ARG PYTORCH_AUDIO_REPO +ARG USE_SCCACHE + +RUN git clone ${PYTORCH_AUDIO_REPO} audio +RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ + && git submodule update --init --recursive +RUN cd audio && pip install -r requirements.txt +RUN cd audio \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi +RUN mkdir -p /app/install && cp /app/audio/dist/*.whl /app/install + + +### +### MORI Build +### +FROM base AS build_mori +ARG MORI_BRANCH +ARG MORI_REPO +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl +RUN git clone ${MORI_REPO} +RUN cd mori \ + && git checkout ${MORI_BRANCH} \ + && git submodule update --init --recursive \ + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl +RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install + + +### +### FlashAttention Build +### +FROM base AS build_fa +ARG FA_BRANCH +ARG FA_REPO +ARG USE_SCCACHE +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl +RUN git clone ${FA_REPO} +RUN cd flash-attention \ + && git checkout ${FA_BRANCH} \ + && git submodule update --init \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi +RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install + + +### +### AITER Build +### +FROM base AS build_aiter +ARG AITER_BRANCH +ARG AITER_REPO +ARG USE_SCCACHE +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} +RUN cd aiter \ + && git submodule update --init --recursive \ + && pip install -r requirements.txt +RUN pip install pyyaml && cd aiter \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && ls /app/aiter/dist/*.whl +RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install + + +### +### Final Build +### + +# Wheel release stage - +# only includes dependencies used by wheel release pipeline +FROM base AS debs_wheel_release +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs + +# Full debs stage - includes Mori (used by Docker releases) +FROM base AS debs +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs + +### +### ROCr + CLR build +### +# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond. +FROM base AS build_rocm_runtime +ARG ROCM_SYSTEMS_REPO +ARG ROCM_RUNTIME_COMMIT +RUN apt-get update \ + && apt-get install -y --no-install-recommends g++ \ + libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \ + && rm -rf /var/lib/apt/lists/* +RUN pip install --no-cache-dir CppHeaderParser +RUN git init -q /src && cd /src \ + && git remote add origin ${ROCM_SYSTEMS_REPO} \ + && git config core.sparseCheckout true \ + && git config remote.origin.promisor true \ + && git config remote.origin.partialclonefilter blob:none \ + && git sparse-checkout init --cone \ + && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \ + && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \ + && git checkout -q FETCH_HEAD +# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get +# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so +# supply one pointing at them. +RUN mkdir -p /opt/rocm-cmake-shim \ + && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \ + && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake +# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers +RUN cd /src/projects/rocr-runtime \ + && cmake -B build -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH=/opt/rocm \ + -DCMAKE_INSTALL_PREFIX=/opt/rocm \ + -DClang_DIR=/opt/rocm-cmake-shim \ + -DLLVM_DIR=/opt/rocm-cmake-shim \ + && cmake --build build --parallel 16 \ + && cmake --install build \ + && cmake --install build --prefix /rocr-install --strip +RUN cd /src/projects/clr \ + && cmake -B build -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCLR_BUILD_HIP=ON \ + -DCLR_BUILD_OCL=OFF \ + -DHIP_COMMON_DIR=/src/projects/hip \ + -DROCM_PATH=/opt/rocm \ + -DCMAKE_PREFIX_PATH=/opt/rocm \ + -DCMAKE_INSTALL_PREFIX=/opt/rocm \ + -DLLVM_DIR=/opt/rocm-cmake-shim \ + -DHIP_LLVM_ROOT=/opt/rocm/llvm \ + && cmake --build build --parallel 16 \ + && cmake --install build --prefix /clr-install --strip +RUN mkdir -p /staging/lib \ + && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \ + && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \ + && ls -l /staging/lib + +# rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3 +FROM base AS build_rocprofiler_sdk +ARG ROCM_SYSTEMS_REPO +ARG ROCPROFILER_SDK_COMMIT +ARG USE_SCCACHE +RUN apt-get update \ + && apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \ + && rm -rf /var/lib/apt/lists/* +RUN git init -q /src && cd /src \ + && git remote add origin ${ROCM_SYSTEMS_REPO} \ + && git config core.sparseCheckout true \ + && git config remote.origin.promisor true \ + && git config remote.origin.partialclonefilter blob:none \ + && git sparse-checkout init --cone \ + && git sparse-checkout set projects/rocprofiler-sdk shared cmake \ + && git fetch --filter=blob:none --depth 1 origin ${ROCPROFILER_SDK_COMMIT} \ + && git checkout -q FETCH_HEAD +# Patch the ROCm 7.2.3 build of rocprofiler-sdk to include queue-interposition fixes that are +# absent from the pinned 7.14 build. This resolves hangs on torch.profiler on ROCm. +# TODO: Drop these once ${ROCPROFILER_SDK_COMMIT} is bumped to a build that already +# contains them, such as ROCm 10.0. See https://github.com/vllm-project/vllm/issues/54087. +ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a +ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653 +RUN cd /src \ + && HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \ + && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \ + && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \ + && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \ + && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch +RUN cd /src/projects/rocprofiler-sdk \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ + && cmake -B build -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH=/opt/rocm \ + -DCMAKE_INSTALL_PREFIX=/opt/rocm \ + -DROCPROFILER_BUILD_TESTS=OFF \ + -DROCPROFILER_BUILD_SAMPLES=OFF \ + -DROCPROFILER_BUILD_DOCS=OFF \ + && cmake --build build --parallel "$(nproc)" \ + && cmake --install build --prefix /staging --strip + +FROM base AS final +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ + pip install /install/*.whl + +COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/ +ENV ROCPROFILER_QUEUE_INTERPOSITION=0 + +COPY --from=build_rocm_runtime /staging/ /opt/rocm/ +# Hard link the patched library onto every stock file name so that both the +# SONAME lookup and the hard coded paths resolve to it. +RUN set -eux; \ + cd /opt/rocm/lib; \ + ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \ + ours_hip="$(readlink -f libamdhip64.so.7)"; \ + for f in libhsa-runtime64.so.1.*; do \ + [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \ + done; \ + for f in libamdhip64.so.7.*; do \ + [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \ + done + +RUN ldconfig + +ARG BASE_IMAGE +ARG TRITON_BRANCH +ARG TRITON_REPO +ARG PYTORCH_BRANCH +ARG PYTORCH_VISION_BRANCH +ARG PYTORCH_REPO +ARG PYTORCH_VISION_REPO +ARG PYTORCH_AUDIO_BRANCH +ARG PYTORCH_AUDIO_REPO +ARG FA_BRANCH +ARG FA_REPO +ARG AITER_BRANCH +ARG AITER_REPO +ARG MORI_BRANCH +ARG MORI_REPO +ARG ROCM_SYSTEMS_REPO +ARG ROCPROFILER_SDK_COMMIT +ARG ROCM_RUNTIME_COMMIT +RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ + && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ + && echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \ + && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ + && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ + && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ + && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ + && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ + && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ + && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 33dd0f6b10fb..933f112e8ecb 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,12 +1,11 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17 +ARG BASE_IMAGE=ubuntu:22.04 +ARG PYTHON_VERSION=3.12 +ARG TORCH_VERSION=2.12.0+rocm10.0.0 +ARG TORCHVISION_VERSION=0.27.0+rocm10.0.0 +ARG TORCHAUDIO_VERSION=2.11.0+rocm10.0.0 +ARG TRITON_BRANCH="669b31ac" # release/internal/3.8.x as of 09/29/2026 +ARG TRITON_VERSION=3.8.0+git669b31ac ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17 -ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.28.0" -ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.11.0" -ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" ARG AITER_BRANCH="v0.1.25" @@ -14,8 +13,18 @@ ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.2.3.post1" ARG MORI_REPO="https://github.com/ROCm/mori.git" ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" -ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350" -ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521" +ARG ROCM_RUNTIME_COMMIT="25e14349f2606a6b95c4ccd45eb92a1835646cf2" + + +# TheRock configuration +ARG SITE_PACKAGES=/usr/local/lib/python${PYTHON_VERSION}/dist-packages +ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL="https://stable.repo.amd.com/rocm/whl-next/" # For pre-built GPU images, packages, etc. +ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL="https://nightly.repo.amd.com/rocm/whl-next/" # For torch, triton, etc. wheels with latest fixes. +ARG ROCM_SDK_VERSION=10.0.0 + +# Build configuration +ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" +ARG AITER_ROCM_ARCH="gfx942;gfx950" # Sccache configuration (only used in release pipeline) ARG USE_SCCACHE @@ -25,47 +34,128 @@ ARG SCCACHE_BUCKET_NAME=vllm-build-sccache ARG SCCACHE_REGION_NAME=us-west-2 ARG SCCACHE_S3_NO_CREDENTIALS=0 - FROM ${BASE_IMAGE} AS base -ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -ENV ROCM_PATH=/opt/rocm -ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: - -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 -ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950 -# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime +ARG PYTORCH_ROCM_ARCH +ENV PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH}" +ARG AITER_ROCM_ARCH +ENV AITER_ROCM_ARCH="${AITER_ROCM_ARCH}" + +ARG SITE_PACKAGES +ENV SITE_PACKAGES=${SITE_PACKAGES} +ENV SDK_DEV="$SITE_PACKAGES/_rocm_sdk_devel" \ + SDK_CORE="$SITE_PACKAGES/_rocm_sdk_core" +ENV ROCM_PATH="$SDK_DEV" \ + ROCM_HOME="$SDK_DEV" \ + HIP_PATH="$SDK_DEV" \ + HIP_CLANG_PATH="$SDK_DEV/lib/llvm/bin" \ + HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode" +ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib" +ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH" # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 -ARG PYTHON_VERSION=3.12 +ARG PYTHON_VERSION ENV PYTHON_VERSION=${PYTHON_VERSION} RUN mkdir -p /app WORKDIR /app ENV DEBIAN_FRONTEND=noninteractive -# Install Python and other dependencies (plus gcc-13, selected later in build_pytorch). +# Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ + && apt-get install -y software-properties-common build-essential git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config libdrm-dev libnuma-dev kmod \ && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ done \ && apt-get update -y \ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ && python3 --version && python3 -m pip --version -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* +# Install ROCm SDK via TheRock +ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL +ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL +ARG ROCM_SDK_VERSION +RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ + "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \ + && rocm-sdk init + +# Install PyTorch wheels (nightly for latest torch, release as fallback) +ARG TORCH_VERSION +ARG TORCHAUDIO_VERSION +ARG TORCHVISION_VERSION +RUN pip install \ + --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ + "torch==${TORCH_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" + +# Install the AMD-compiled device kernels corresponding to PYTORCH_ROCM_ARCH: + +# Possible values that can be in PYTORCH_ROCM_ARCH: +# +# gfx90a (MI210/250), gfx942 (MI300), gfx950 (MI350), gfx1100 (7900XTX), +# gfx1101 (7800XT), gfx1200 (9070XT), gfx1201 (9070), gfx1150/1151 (Strix APU) + +# For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that +# PyTorch needs to run on that specific chip. +SHELL ["/bin/bash", "-c"] +RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ + && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \ + echo "Installing device packages for ${arch}" \ + && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ + "torch[device-${arch}]==${TORCH_VER}" \ + "torchvision[device-${arch}]==${TORCHVISION_VERSION}" \ + "rocm[device-${arch}]==${ROCM_SDK_VERSION}"; \ + done + +# Remove gfx90a xnack kpacks (ASAN artifacts that cause hipErrorInvalidImage on MI250) +RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \ + "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack" + +RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install . +ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi + +# TheRock runtime fixes +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi \ + && if [ -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" ] \ + && [ ! -e "${ROCM_PATH}/lib/librocm_smi64.so.7" ]; then \ + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_core/lib/librocm_smi64.so.1" "${ROCM_PATH}/lib/librocm_smi64.so.7"; \ + fi \ + && printf '%s\n%s\n' "${ROCM_PATH}/lib" "${SITE_PACKAGES}/_rocm_sdk_core/lib" > /etc/ld.so.conf.d/rocm-pip.conf \ + && for f in ${SITE_PACKAGES}/_rocm_sdk_core/lib/libamdhip64.so*; do \ + bn=$(basename "$f"); \ + if [ -f "${ROCM_PATH}/lib/${bn}" ] && [ -f "$f" ] && [ ! -L "$f" ]; then \ + rm -f "$f"; \ + ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ + fi; \ + done \ + && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ + librocprofiler-sdk-roctx.so.1; do \ + if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ + && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ + rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ + "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + fi; \ + done \ + && ldconfig + # Install sccache if USE_SCCACHE is enabled (for release builds) ARG USE_SCCACHE ARG SCCACHE_DOWNLOAD_URL @@ -94,141 +184,35 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ RUN if [ "$USE_SCCACHE" = "1" ]; then \ echo "Setting up sccache wrappers for HIP compilation..." \ && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \ && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \ && chmod +x /opt/sccache-wrappers/clang \ && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi -# Set sccache environment variables only when USE_SCCACHE=1 -# This prevents S3 config from leaking into images when sccache is not used -ARG USE_SCCACHE ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && if [ ! -f setup.py ]; then cd python; fi \ - && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install - - -### -### AMD SMI Build -### -FROM base AS build_amdsmi -RUN cd /opt/rocm/share/amd_smi \ - && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install - - -### -### Pytorch build -### -FROM base AS build_pytorch -ARG PYTORCH_BRANCH -ARG PYTORCH_REPO -ARG USE_SCCACHE - -RUN apt-get update && apt-get install -y pkg-config liblzma-dev -# torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13. -RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 -RUN git clone ${PYTORCH_REPO} pytorch -RUN cd pytorch && git checkout ${PYTORCH_BRANCH} -RUN cd pytorch \ - && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch && python3 tools/amd_build/build_amd.py \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ - && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install - - -### -### Shared PyTorch runtime for dependent builds -### -FROM base AS build_pytorch_runtime -RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install,ro \ - pip install /install/*.whl - - -### -### TorchVision build -### -FROM build_pytorch_runtime AS build_torchvision -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_VISION_REPO -ARG USE_SCCACHE - -RUN git clone ${PYTORCH_VISION_REPO} vision -RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/vision/dist/*.whl /app/install - - -### -### TorchAudio build -### -FROM build_pytorch_runtime AS build_torchaudio -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO -ARG USE_SCCACHE - -RUN git clone ${PYTORCH_AUDIO_REPO} audio -RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ - && git submodule update --init --recursive -RUN cd audio && pip install -r requirements.txt -RUN cd audio \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/audio/dist/*.whl /app/install - - ### ### MORI Build ### FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone ${MORI_REPO} RUN cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl + && ln -sfn ${ROCM_PATH} /opt/rocm \ + && pip install "setuptools_scm>=6.2" \ + && CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \ + python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install - ### ### FlashAttention Build ### @@ -236,8 +220,6 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone ${FA_REPO} RUN cd flash-attention \ && git checkout ${FA_BRANCH} \ @@ -258,12 +240,12 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ && git submodule update --init --recursive \ && pip install -r requirements.txt +# AITER_USE_SYSTEM_TRITON keeps setup.py from resolving its own triton, which would +# displace the pinned one already installed in base. RUN pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ @@ -276,58 +258,30 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install ### -### Final Build +### Triton Build ### +FROM base AS build_triton +ARG TRITON_BRANCH +ARG TRITON_REPO +RUN git clone ${TRITON_REPO} +RUN cd triton \ + && git checkout ${TRITON_BRANCH} \ + && if [ ! -f setup.py ]; then cd python; fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && mkdir -p /app/install && cp dist/*.whl /app/install -# Wheel release stage - -# only includes dependencies used by wheel release pipeline -FROM base AS debs_wheel_release -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs - -# Full debs stage - includes Mori (used by Docker releases) -FROM base AS debs -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs ### ### ROCr + CLR build ### -# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond. +# Brings in multi-stream perf improvements; drop after these arrive in an official release FROM base AS build_rocm_runtime ARG ROCM_SYSTEMS_REPO ARG ROCM_RUNTIME_COMMIT +# libgl-dev: rocclr does find_package(OpenGL REQUIRED) unconditionally and the wheel image has +# no GL headers. The link must not pick up a libGL DT_NEEDED, which is asserted after the build. RUN apt-get update \ - && apt-get install -y --no-install-recommends g++ \ - libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \ + && apt-get install -y --no-install-recommends g++ libdw-dev xxd libgl-dev libglx-dev \ && rm -rf /var/lib/apt/lists/* RUN pip install --no-cache-dir CppHeaderParser RUN git init -q /src && cd /src \ @@ -339,20 +293,17 @@ RUN git init -q /src && cd /src \ && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \ && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \ && git checkout -q FETCH_HEAD -# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get -# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so -# supply one pointing at them. -RUN mkdir -p /opt/rocm-cmake-shim \ - && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \ - && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake -# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers +# libelf is taken from the SDK's own vendored sysdeps rather than apt -- the wheel image has no +# /usr/include/libelf.h -- so ROCr links against the copy the wheel itself was built against. +ENV SYSDEPS="$SDK_DEV/lib/rocm_sysdeps" +# ROCr first, and installed into $SDK_DEV, because the clr build below needs its headers. RUN cd /src/projects/rocr-runtime \ && cmake -B build -G Ninja \ -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DClang_DIR=/opt/rocm-cmake-shim \ - -DLLVM_DIR=/opt/rocm-cmake-shim \ + -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ + -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ + -DLIBELF_INCLUDE_DIRS="$SYSDEPS/include" \ + -DLIBELF_LIBRARIES="$SYSDEPS/lib/libelf.so" \ && cmake --build build --parallel 16 \ && cmake --install build \ && cmake --install build --prefix /rocr-install --strip @@ -361,12 +312,12 @@ RUN cd /src/projects/clr \ -DCMAKE_BUILD_TYPE=Release \ -DCLR_BUILD_HIP=ON \ -DCLR_BUILD_OCL=OFF \ + -DROCM_KPACK_ENABLED=ON \ -DHIP_COMMON_DIR=/src/projects/hip \ - -DROCM_PATH=/opt/rocm \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DLLVM_DIR=/opt/rocm-cmake-shim \ - -DHIP_LLVM_ROOT=/opt/rocm/llvm \ + -DROCM_PATH="$SDK_DEV" \ + -DCMAKE_PREFIX_PATH="$SDK_DEV;$SYSDEPS" \ + -DCMAKE_INSTALL_PREFIX="$SDK_DEV" \ + -DHIP_LLVM_ROOT="$SDK_DEV/lib/llvm" \ && cmake --build build --parallel 16 \ && cmake --install build --prefix /clr-install --strip RUN mkdir -p /staging/lib \ @@ -374,82 +325,115 @@ RUN mkdir -p /staging/lib \ && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \ && ls -l /staging/lib -# rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3 -FROM base AS build_rocprofiler_sdk -ARG ROCM_SYSTEMS_REPO -ARG ROCPROFILER_SDK_COMMIT -ARG USE_SCCACHE -RUN apt-get update \ - && apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \ - && rm -rf /var/lib/apt/lists/* -RUN git init -q /src && cd /src \ - && git remote add origin ${ROCM_SYSTEMS_REPO} \ - && git config core.sparseCheckout true \ - && git config remote.origin.promisor true \ - && git config remote.origin.partialclonefilter blob:none \ - && git sparse-checkout init --cone \ - && git sparse-checkout set projects/rocprofiler-sdk shared cmake \ - && git fetch --filter=blob:none --depth 1 origin ${ROCPROFILER_SDK_COMMIT} \ - && git checkout -q FETCH_HEAD -# Patch the ROCm 7.2.3 build of rocprofiler-sdk to include queue-interposition fixes that are -# absent from the pinned 7.14 build. This resolves hangs on torch.profiler on ROCm. -# TODO: Drop these once ${ROCPROFILER_SDK_COMMIT} is bumped to a build that already -# contains them, such as ROCm 10.0. See https://github.com/vllm-project/vllm/issues/54087. -ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a -ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653 -RUN cd /src \ - && HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \ - && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \ - && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \ - && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \ - && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch -RUN cd /src/projects/rocprofiler-sdk \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DROCPROFILER_BUILD_TESTS=OFF \ - -DROCPROFILER_BUILD_SAMPLES=OFF \ - -DROCPROFILER_BUILD_DOCS=OFF \ - && cmake --build build --parallel "$(nproc)" \ - && cmake --install build --prefix /staging --strip + +### +### Final Build +### + +# Wheel release stage - +# only includes dependencies used by wheel release pipeline +FROM base AS debs_wheel_release +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +# The wheel index also serves a torch repacked to require our triton instead of +# the one TheRock built it with, and amdsmi, which TheRock ships only as sources +# inside rocm-sdk-core. +ARG TORCH_VERSION +ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL +ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL +COPY tools/vllm-rocm/therock_wheels.py /tmp/therock_wheels.py +# regex and pybase64 aren't installed here; fall back to the stdlib modules. +RUN sed -i -e 's/import regex as re/import re/' \ + -e 's/import pybase64 as base64/import base64/' /tmp/therock_wheels.py \ + && pip download --no-deps --dest /tmp/torch-dl \ + --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ + --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ + "torch==${TORCH_VERSION}" \ + && TRITON_VER=$(basename /app/debs/triton-*.whl | cut -d- -f2) \ + && python3 /tmp/therock_wheels.py repack-torch /tmp/torch-dl/torch-*.whl \ + --triton-version "${TRITON_VER}" --out-dir /app/debs \ + && python3 /tmp/therock_wheels.py patch-amdsmi "${SDK_CORE}/share/amd_smi" \ + --out-dir /tmp/amd_smi \ + && pip wheel --no-deps --wheel-dir /app/debs /tmp/amd_smi \ + && rm -rf /tmp/torch-dl /tmp/amd_smi /tmp/therock_wheels.py \ + && ls -l /app/debs + +FROM base AS debs +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs FROM base AS final -RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - pip install /install/*.whl +ARG TRITON_VERSION -COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/ +# torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA +# teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix. ENV ROCPROFILER_QUEUE_INTERPOSITION=0 -COPY --from=build_rocm_runtime /staging/ /opt/rocm/ -# Hard link the patched library onto every stock file name so that both the -# SONAME lookup and the hard coded paths resolve to it. +COPY --from=build_rocm_runtime /staging/lib/ ${SDK_DEV}/lib/ +# Hard link the patched library onto every stock file name, in both SDK trees, so that the +# SONAME lookup and the hard coded paths all resolve to it. Hard links rather than symlinks +# because ldconfig rewrites SONAME symlinks but cannot re-point an inode. $SDK_CORE keeps its +# own libhsa-runtime64 copy, which is on LD_LIBRARY_PATH and in ld.so.conf; relinking it is +# belt and braces rather than strictly required, but it leaves no stock ROCr in the image. RUN set -eux; \ - cd /opt/rocm/lib; \ + cd "$SDK_DEV/lib"; \ ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \ ours_hip="$(readlink -f libamdhip64.so.7)"; \ - for f in libhsa-runtime64.so.1.*; do \ - [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \ - done; \ - for f in libamdhip64.so.7.*; do \ - [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \ + for f in "$SDK_DEV"/lib/libhsa-runtime64.so.1* "$SDK_CORE"/lib/libhsa-runtime64.so.1* \ + "$SDK_DEV"/lib/libamdhip64.so.7* "$SDK_CORE"/lib/libamdhip64.so.7*; do \ + [ -f "$f" ] && [ ! -L "$f" ] || continue; \ + case "${f##*/}" in libhsa*) ours="$ours_hsa";; *) ours="$ours_hip";; esac; \ + [ "$f" -ef "$ours" ] || ln -f "$ours" "$f"; \ done RUN ldconfig +# This is to prevent resolvers from checking METADATA and RECORD and clobbering +# the Triton source build we installed. Constraints files are unable to help +# with this. However, these entries are optional and can be blanked which +# solves the problem. +RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \ + && grep -q '^Requires-Dist: triton' "$TORCH_DIST_INFO/METADATA" \ + && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \ + && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD" + +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ + pip install /install/*.whl \ + && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \ + && test "$(pip show triton | awk '/^Version:/ {print $2}')" = "${TRITON_VERSION}" \ + # Create a constraints file to protect the versions we need, especially + # triton, requirements will try to reinstall repeatedly. + && pip list --format=freeze \ + | grep -E '^(torch|triton|torchaudio|torchvision|pybind11)==' > /etc/rocm-constraints.txt \ + && echo "ROCm constraints:" && cat /etc/rocm-constraints.txt \ + && test "$(wc -l < /etc/rocm-constraints.txt)" -eq 5 \ + && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \ + && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \ + && python3 -c "import triton; print(f'triton={triton.__version__}')" \ + && pip list 2>/dev/null | grep -iE '^(torch|triton) ' + ARG BASE_IMAGE +ARG PYTHON_VERSION +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHAUDIO_VERSION +ARG TORCHVISION_VERSION +ARG PYTORCH_ROCM_ARCH +ARG AITER_ROCM_ARCH ARG TRITON_BRANCH ARG TRITON_REPO -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH @@ -457,23 +441,36 @@ ARG AITER_REPO ARG MORI_BRANCH ARG MORI_REPO ARG ROCM_SYSTEMS_REPO -ARG ROCPROFILER_SDK_COMMIT ARG ROCM_RUNTIME_COMMIT +ARG TRITON_VERSION +# Record hashes for Triton and pybind11 so we can verify later. RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ - && echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \ - && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ + && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}" >> /app/versions.txt \ + && echo "AITER_ROCM_ARCH: ${AITER_ROCM_ARCH}" >> /app/versions.txt \ && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ + && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ + && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ + # Fingerprint our triton so Dockerfile.rocm can prove it survived the build. + && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ + && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ + -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ + >> /app/versions.txt \ + && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ + && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \ + $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \ + >> /app/versions.txt \ + && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \ + -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ + >> /app/versions.txt diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl index 1794b8468a16..c251641452f6 100644 --- a/docker/ci-rocm.hcl +++ b/docker/ci-rocm.hcl @@ -83,10 +83,6 @@ variable "UCX_BRANCH" { default = "" } -variable "ROCSHMEM_BRANCH" { - default = "" -} - variable "DEEPEP_BRANCH" { default = "" } @@ -95,10 +91,6 @@ variable "NIXL_CACHE_KEY" { default = "" } -variable "ROCSHMEM_CACHE_KEY" { - default = "" -} - variable "DEEPEP_CACHE_KEY" { default = "" } @@ -236,7 +228,7 @@ function "get_cache_to_rocm_rust" { ]) } -# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP). +# Cache functions for upstream dependency stages (NIXL/UCX, DeepEP). # These stages are pinned to specific upstream commit hashes, so cache keys use # those hashes rather than the Buildkite commit. This means the cache persists # across all vLLM commits as long as the upstream dependency pins don't change. @@ -245,8 +237,7 @@ function "get_cache_from_rocm_deps" { params = [] result = compact([ NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), - ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""), - DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""), + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}" : ""), ]) } @@ -257,17 +248,10 @@ function "get_cache_to_rocm_nixl" { ]) } -function "get_cache_to_rocm_rocshmem" { - params = [] - result = compact([ - ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY},mode=min" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH},mode=min" : ""), - ]) -} - function "get_cache_to_rocm_deepep" { params = [] result = compact([ - DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH},mode=min" : ""), + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH},mode=min" : ""), ]) } @@ -385,14 +369,6 @@ target "nixl-rocm-ci" { output = ["type=cacheonly"] } -target "rocshmem-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "build_rocshmem" - cache-from = get_cache_from_rocm_deps() - cache-to = get_cache_to_rocm_rocshmem() - output = ["type=cacheonly"] -} - target "deepep-rocm-ci" { inherits = ["_common-rocm", "_ci-rocm"] target = "build_deepep" @@ -418,7 +394,7 @@ target "ci-base-rocm-ci" { CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "", CI_BASE_TRUSTED_CONTENT_REF != "" ? "type=registry,ref=${CI_BASE_TRUSTED_CONTENT_REF}" : "", ]), - # Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages + # Import upstream dependency caches so NIXL/DeepEP stages # are cache hits even when ci_base itself needs rebuilding. get_cache_from_rocm_deps(), ) @@ -431,5 +407,5 @@ target "ci-base-rocm-ci" { # Group for ci_base builds -- exports dependency stage caches alongside the # ci_base image so future rebuilds can reuse them independently. group "ci-base-rocm-ci-with-deps" { - targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] + targets = ["nixl-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] } diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl index ab523a296a93..2eb1d5a43076 100644 --- a/docker/docker-bake-rocm.hcl +++ b/docker/docker-bake-rocm.hcl @@ -65,10 +65,6 @@ variable "UCX_BRANCH" { default = "" } -variable "ROCSHMEM_BRANCH" { - default = "" -} - variable "DEEPEP_BRANCH" { default = "" } diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index 28e74b670ca7..d8debbab3ea3 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -1,7 +1,7 @@ --8<-- [start:installation] -vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available for ROCm 7.0 and ROCm 7.2.1. +vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available for ROCm 10.0 (the default) and ROCm 7.2. #### Prebuilt Wheels @@ -9,6 +9,8 @@ vLLM supports AMD GPUs with ROCm 6.3 or above. Pre-built wheels are available fo | ------------ | -------------- | ------------ | ----------------- | ------------------ | | `rocm700` | 3.12 | 7.0 | >= 2.35 | `0.14.0` to `0.18.0` | | `rocm721` | 3.12 | 7.2.1 | >= 2.35 | Nightly releases after commit `171775f306a333a9cf105bfd533bf3e113d401d9` | +| `rocm72` | 3.12 | 7.2 | >= 2.39 | Nightly releases on the legacy ROCm 7.2 stack | +| `rocm100` | 3.12 | 10.0 | >= 2.35 | Nightly releases (default) | --8<-- [end:installation] --8<-- [start:requirements] @@ -40,12 +42,17 @@ If you need a different ROCm version or want to use an existing PyTorch installa source .venv/bin/activate ``` -To install the latest version of vLLM for Python 3.12, ROCm 7.0 and `glibc >= 2.35`. +To install the latest version of vLLM for Python 3.12, ROCm 10.0 and `glibc >= 2.35`: ```bash uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/ --upgrade ``` +!!! tip + The ROCm 10.0 wheels install the ROCm SDK from pip ([TheRock](https://github.com/ROCm/TheRock)), including the GPU kernels for every supported architecture. + + Wheels for the legacy ROCm 7.2 stack, which use a system ROCm installation, remain available at `https://wheels.vllm.ai/rocm/${VLLM_VERSION}/rocm72`. + !!! tip You can find out about which ROCm version the latest vLLM supports by checking the `vllm` package in index in extra-index-url at [https://wheels.vllm.ai/rocm/vllm](https://wheels.vllm.ai/rocm/vllm) . @@ -195,9 +202,9 @@ install it without re-running the full installation: and place it in `vllm/vllm-rs`. If you see an error about a wheel not being found, the wheel for your base -commit and ROCm patch version might not be available. Check the available -variants under `https://wheels.vllm.ai/rocm//`. For example, ROCm 7.2.1 -uses the `rocm721` variant. +commit and ROCm version might not be available. Check the available +variants under `https://wheels.vllm.ai/rocm//`. Variants are named by ROCm +major and minor version, e.g. ROCm 7.2.3 uses the `rocm72` variant. There are more environment variables to control the behavior of Python-only build: @@ -208,7 +215,7 @@ build: - `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the full commit hash used to download the pre-compiled wheel. - `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the ROCm variant subdirectory, e.g., - `rocm700` or `rocm721`. If not specified, the variant is auto-detected based + `rocm72` or `rocm100`. If not specified, the variant is auto-detected based on your system's ROCm version. An explicitly specified variant must match the detected environment. @@ -226,23 +233,34 @@ You can find more information about vLLM's wheels in #### Full build (with compilation) {#full-build} !!! tip - - If you found that the following installation step does not work for you, please refer to [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). Dockerfile is a form of installation steps. - -0. Install prerequisites (skip if you are already in an environment/docker with the following installed): + - The steps below follow [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base), which pins the validated version of every component. If a step does not work for you, refer to it; a Dockerfile is a form of installation steps. + - The legacy ROCm 7.2 stack uses a system ROCm installation instead; its steps are in [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base). - - [ROCm](https://rocm.docs.amd.com/en/latest/deploy/linux/index.html) - - [PyTorch](https://pytorch.org/) +0. Install the ROCm SDK and PyTorch (skip if you are already in an environment or docker image with them installed, e.g. `vllm/vllm-openai-rocm:base-nightly`). - For installing PyTorch, you can start from a fresh docker image, e.g, `rocm/pytorch:rocm7.0_ubuntu22.04_py3.10_pytorch_release_2.8.0`, `rocm/pytorch-nightly`. If you are using docker image, you can skip to Step 3. - - Alternatively, you can install PyTorch using PyTorch wheels. You can check PyTorch installation guide in PyTorch [Getting Started](https://pytorch.org/get-started/locally/). Example: + ROCm 10.0 is installed from pip ([TheRock](https://github.com/ROCm/TheRock)), together with PyTorch and the GPU kernels for your GPU's architecture. In a Python 3.12 environment: ```bash - # Install PyTorch - pip uninstall torch -y - pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm7.0 + # Your GPU's architecture, e.g. gfx942 for MI300 or gfx950 for MI350 + export GPU_ARCH=$(uvx --from rocm-bootstrap rocm-bootstrap-detect --unique) + export ROCM_INDEX=https://stable.repo.amd.com/rocm/whl-next/ + + pip install --index-url ${ROCM_INDEX} \ + "rocm[libraries,devel,device-${GPU_ARCH}]==10.0.0" \ + "torch[device-${GPU_ARCH}]==2.12.0+rocm10.0.0" \ + "torchvision[device-${GPU_ARCH}]==0.27.0+rocm10.0.0" + rocm-sdk init + + # Point builds at the pip-installed ROCm SDK + export ROCM_PATH=$(rocm-sdk path --root) + export ROCM_HOME=${ROCM_PATH} HIP_PATH=${ROCM_PATH} + export PATH=$(rocm-sdk path --bin):${PATH} ``` + !!! note + - The validated `ROCM_SDK_VERSION`, `TORCH_VERSION` and `TORCHVISION_VERSION` can be found in [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). + - To support several GPU architectures, add a `device-` extra for each of them. + 1. Install [Triton for ROCm](https://github.com/ROCm/triton.git) Install ROCm's Triton following the instructions from [ROCm/triton](https://github.com/ROCm/triton.git) @@ -252,8 +270,7 @@ You can find more information about vLLM's wheels in pip uninstall -y triton git clone https://github.com/ROCm/triton.git cd triton - # git checkout $TRITON_BRANCH - git checkout f9e5bf54 + git checkout $TRITON_BRANCH if [ ! -f setup.py ]; then cd python; fi python3 setup.py install cd ../.. @@ -265,17 +282,14 @@ You can find more information about vLLM's wheels in 2. Optionally, if you choose to use CK flash attention, you can install [flash attention for ROCm](https://github.com/Dao-AILab/flash-attention.git) - Install ROCm's flash attention (v2.8.0) following the instructions from [ROCm/flash-attention](https://github.com/Dao-AILab/flash-attention#amd-rocm-support) - - For example, for ROCm 7.0, suppose your gfx arch is `gfx942`. To get your gfx architecture, run `rocminfo |grep gfx`. + Install ROCm's flash attention following the instructions from [ROCm/flash-attention](https://github.com/Dao-AILab/flash-attention#amd-rocm-support) ```bash git clone https://github.com/Dao-AILab/flash-attention.git cd flash-attention - # git checkout $FA_BRANCH - git checkout 0e60e394 + git checkout $FA_BRANCH git submodule update --init - GPU_ARCHS="gfx942" python3 setup.py install + GPU_ARCHS="${GPU_ARCH}" python3 setup.py install cd .. ``` @@ -290,7 +304,8 @@ You can find more information about vLLM's wheels in cd aiter git checkout $AITER_BRANCH_OR_COMMIT git submodule sync; git submodule update --init --recursive - python3 setup.py develop + # Keep the Triton installed in step 1 + AITER_USE_SYSTEM_TRITON=1 python3 setup.py develop ``` !!! note @@ -300,32 +315,34 @@ You can find more information about vLLM's wheels in 4. Optionally, if you want to use MORI for EP or PD disaggregation, you can install [MORI](https://github.com/ROCm/mori) using the following steps: ```bash + pip install "setuptools_scm>=6.2" git clone https://github.com/ROCm/mori.git cd mori git checkout $MORI_BRANCH_OR_COMMIT git submodule sync; git submodule update --init --recursive - MORI_GPU_ARCHS="gfx942;gfx950" python3 setup.py install + # The ROCm SDK keeps its CMake configs for system dependencies (e.g. NUMA) under rocm_sysdeps + CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \ + MORI_GPU_ARCHS="gfx942;gfx950" python3 setup.py install ``` !!! note - You will need to config the `$MORI_BRANCH_OR_COMMIT` for your purpose. - The validated `$MORI_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). -5. Build vLLM. For example, vLLM on ROCM 7.0 can be built with the following steps: +5. Build vLLM. For example, vLLM on ROCm 10.0 can be built with the following steps: ???+ console "Commands" ```bash pip install --upgrade pip - # Build & install AMD SMI - pip install /opt/rocm/share/amd_smi + # Install AMD SMI, which the ROCm SDK ships as sources + SDK_CORE=$(python3 -c "import _rocm_sdk_core, os; print(os.path.dirname(_rocm_sdk_core.__file__))") + pip install ${SDK_CORE}/share/amd_smi + # It loads the SDK's libamd_smi relative to its own directory + export PYTHONPATH=${SDK_CORE}/share/amd_smi # Install dependencies - pip install --upgrade numba \ - scipy \ - huggingface-hub \ - setuptools_scm pip install -r requirements/rocm.txt # To build for a single architecture (e.g., MI300) for faster installation (recommended): @@ -352,8 +369,9 @@ You can find more information about vLLM's wheels in vLLM offers official Docker images for deployment. The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags). -- `vllm/vllm-openai-rocm:latest` — stable release -- `vllm/vllm-openai-rocm:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes +- `vllm/vllm-openai-rocm:latest` — stable release, built on ROCm 10.0 (also tagged `latest-rocm100`) +- `vllm/vllm-openai-rocm:nightly` — preview build from the latest development branch on ROCm 10.0 (also tagged `nightly-rocm100`), use this if you want the latest features and fixes +- `vllm/vllm-openai-rocm:latest-rocm72` / `nightly-rocm72` — the same builds on the legacy ROCm 7.2 stack, kept for a transition period ```bash docker run --rm \ @@ -409,7 +427,7 @@ You can build and run vLLM from source via the provided [docker/Dockerfile.rocm] ??? info "(Optional) Build an image with ROCm software stack" Build a docker image from [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base) which setup ROCm software stack needed by the vLLM. - **This step is optional as this rocm_base image is usually prebuilt and store at [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev) under tag `rocm/vllm-dev:base` to speed up user experience.** + **This step is optional as this rocm_base image is usually prebuilt and store at [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev) under tag `rocm/vllm-dev:base` to speed up user experience.** For the legacy ROCm 7.2 stack, the base image is built from [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base) and published as `vllm/vllm-openai-rocm:base-nightly-rocm72`. If you choose to build this rocm_base image yourself, the steps are as follows. It is important that the user kicks off the docker build using buildkit. Either the user put `DOCKER_BUILDKIT=1` as environment variable when calling docker build command, or the user needs to set up buildkit in the docker daemon configuration `/etc/docker/daemon.json` as follows and restart the daemon: @@ -422,7 +440,7 @@ You can build and run vLLM from source via the provided [docker/Dockerfile.rocm] } ``` - To build vllm on ROCm 7.0 for MI200 and MI300 series, you can use the default: + To build vllm on ROCm 10.0 for all supported GPUs, you can use the default: ```bash DOCKER_BUILDKIT=1 docker build \ @@ -441,15 +459,15 @@ It is important that the user kicks off the docker build using buildkit. Either } ``` -[docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm) uses ROCm 7.0 by default, but also supports ROCm 5.7, 6.0, 6.1, 6.2, 6.3, and 6.4, in older vLLM branches. +[docker/Dockerfile.rocm](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm) uses ROCm 10.0 (installed from TheRock wheels) by default. The legacy ROCm 7.2 stack is available as [docker/Dockerfile.rocm_72](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72) with [docker/Dockerfile.rocm_72_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72_base); older vLLM branches support ROCm 5.7 through 7.0. It provides flexibility to customize the build of docker image using the following arguments: -- `BASE_IMAGE`: specifies the base image used when running `docker build`. The default value `rocm/vllm-dev:base` is an image published and maintained by AMD. It is being built using [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base) +- `BASE_IMAGE`: specifies the base image used when running `docker build`. The default value `rocm/vllm-dev:base` is an image published and maintained by AMD. It is being built using [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). For [docker/Dockerfile.rocm_72](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_72), the default is `vllm/vllm-openai-rocm:base-nightly-rocm72` - `ARG_PYTORCH_ROCM_ARCH`: Allows to override the gfx architecture values from the base docker image Their values can be passed in when running `docker build` with `--build-arg` options. -To build vllm on ROCm 7.0 for MI200 and MI300 series, you can use the default (which build a docker image with `vllm serve` as entrypoint): +To build vllm on ROCm 10.0 for all supported GPUs, you can use the default (which build a docker image with `vllm serve` as entrypoint): ```bash DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm/vllm-openai-rocm . diff --git a/setup.py b/setup.py index 83b59e303a6e..cf99ec3aaaf4 100644 --- a/setup.py +++ b/setup.py @@ -645,8 +645,10 @@ def detect_system_cuda_variant() -> str: @staticmethod def rocm_version_to_variant(rocm_version: str) -> str: - """Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723.""" - return "rocm" + rocm_version.replace(".", "") + """Convert a ROCm version string to a wheel variant from its major and + minor version, e.g. 7.2.3 -> rocm72, 10.0.0 -> rocm100. + """ + return "rocm" + "".join(rocm_version.split(".")[:2]) @staticmethod def detect_system_rocm_variant() -> str | None: @@ -1382,7 +1384,9 @@ def get_vllm_version() -> str: # Get the Rocm Version rocm_version = get_rocm_version() or torch.version.hip if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: - version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" + version += sep + precompiled_wheel_utils.rocm_version_to_variant( + rocm_version + ) elif _is_tpu(): version += f"{sep}tpu" elif _is_cpu(): @@ -1512,6 +1516,7 @@ def _read_requirements(filename: str) -> list[str]: requirements = _read_requirements("rocm.txt") if "develop" in sys.argv[1:] and "--no-deps" not in sys.argv[1:]: _check_requirements_preinstalled(requirements) + requirements += get_rocm_device_requirements() elif _is_tpu(): requirements = _read_requirements("tpu.txt") elif _is_cpu(): @@ -1523,6 +1528,24 @@ def _read_requirements(filename: str) -> list[str]: return requirements +def get_rocm_device_requirements() -> list[str]: + """GPU kernel packages for TheRock builds. + + TheRock ships GPU kernels as separate per-arch wheels selected through + ``[device-]`` extras on torch, torchvision and rocm. Require them for + every arch this wheel was compiled for, so ``pip install vllm`` works on + any supported GPU. + """ + if not _is_hip() or importlib.util.find_spec("rocm_sdk") is None: + return [] + arches = [a for a in os.getenv("PYTORCH_ROCM_ARCH", "").split(";") if a] + return [ + f"{pkg}[device-{arch}]" + for pkg in ("torch", "torchvision", "rocm") + for arch in arches + ] + + ext_modules = [] if _is_cuda() or _is_hip(): diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index c205a3c42a8a..9a8ac646a741 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -110,7 +110,7 @@ def test_modernbert_models( for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): hf_output = hf_output.detach().clone().cpu().float() vllm_output = vllm_output.detach().clone().cpu().float() - torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3) + torch.testing.assert_close(hf_output, vllm_output, atol=6e-2, rtol=1e-3) @pytest.mark.parametrize("model", ["Davlan/xlm-roberta-base-ner-hrl"]) diff --git a/tools/vllm-rocm/generate-rocm-wheels-root-index.sh b/tools/vllm-rocm/generate-rocm-wheels-root-index.sh index 650a71937899..8ccfab17f79f 100755 --- a/tools/vllm-rocm/generate-rocm-wheels-root-index.sh +++ b/tools/vllm-rocm/generate-rocm-wheels-root-index.sh @@ -17,14 +17,20 @@ # # Environment variables: # S3_BUCKET - Bucket name (default: vllm-wheels) -# VARIANT - ROCm variant (default: rocm721) +# ROCM_BASE_DOCKERFILE - default ROCm stack whose variant the root index serves +# (default: docker/Dockerfile.rocm_base) +# VARIANT - override the ROCm variant derived from ROCM_BASE_DOCKERFILE # DRY_RUN - Set to 1 for preview mode (same as --dry-run) set -euo pipefail # ======== Configuration ======== BUCKET="${S3_BUCKET:-vllm-wheels}" -VARIANT="${VARIANT:-rocm721}" +if [[ -z "${VARIANT:-}" ]]; then + # shellcheck source=.buildkite/scripts/rocm/stack.sh + source .buildkite/scripts/rocm/stack.sh "${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}" + VARIANT="$ROCM_STACK_VARIANT" +fi DRY_RUN="${DRY_RUN:-0}" FORCE_VERSION="" diff --git a/tools/vllm-rocm/therock_wheels.py b/tools/vllm-rocm/therock_wheels.py new file mode 100644 index 000000000000..3728ab4d08fa --- /dev/null +++ b/tools/vllm-rocm/therock_wheels.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Adapt TheRock (ROCm 10) wheels for the vLLM ROCm wheel index. + +repack-torch: TheRock's torch wheel hard-pins the triton it was built with. + vLLM ships a newer source-built triton, so rewrite that one Requires-Dist + and the METADATA hash in RECORD. The version and filename stay the same, + because torchvision and the amd-torch-device-* wheels pin torch exactly. + +external-links: list the TheRock wheels the index links to instead of hosting + (ROCm SDK, device kernels, torchvision/torchaudio, rocm-bootstrap), with + versions and GPU arches read from docker/Dockerfile.rocm_base. torch's + device- extras can also pull in a GPU-family wheel (e.g. + amd-torch-device-gfx110x), so those are read from the torch wheel itself. + +patch-amdsmi: rocm-sdk-core ships the amdsmi bindings under + _rocm_sdk_core/share/amd_smi, and the loader finds libamd_smi relative to + that directory. Once installed as its own package it would miss the SDK + library, so also look it up through the installed _rocm_sdk_core package. +""" + +import argparse +import hashlib +import shutil +import sys +import tempfile +import urllib.request +import zipfile +from html.parser import HTMLParser +from pathlib import Path +from urllib.parse import unquote, urljoin + +import pybase64 as base64 +import regex as re + + +def _record_hash(data: bytes) -> str: + digest = base64.urlsafe_b64encode(hashlib.sha256(data).digest()).rstrip(b"=") + return f"sha256={digest.decode()}" + + +def repack_torch(wheel: Path, triton_version: str, out_dir: Path) -> Path: + with zipfile.ZipFile(wheel) as src: + infos = src.infolist() + metadata = next(i for i in infos if i.filename.endswith(".dist-info/METADATA")) + record = next(i for i in infos if i.filename.endswith(".dist-info/RECORD")) + + text = src.read(metadata).decode() + new_text, count = re.subn( + r"^Requires-Dist: triton\s*==\s*\S+\s*$", + f"Requires-Dist: triton=={triton_version}", + text, + flags=re.MULTILINE, + ) + if count != 1: + raise SystemExit( + f"expected exactly one pinned triton requirement, found {count}" + ) + new_metadata = new_text.encode() + + lines = src.read(record).decode().splitlines() + prefix = metadata.filename + "," + replaced = [ + f"{metadata.filename},{_record_hash(new_metadata)},{len(new_metadata)}" + if line.startswith(prefix) + else line + for line in lines + ] + if replaced == lines: + raise SystemExit(f"{metadata.filename} not found in RECORD") + new_record = ("\n".join(replaced) + "\n").encode() + + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / wheel.name + with zipfile.ZipFile(out, "w") as dst: + for info in infos: + if info.filename == metadata.filename: + data = new_metadata + elif info.filename == record.filename: + data = new_record + else: + data = src.read(info) + dst.writestr(info, data, compress_type=info.compress_type) + return out + + +_AMDSMI_ANCHOR = " possible_locations.append(libamd_smi_path)\n" +_AMDSMI_PATCH = """\ + # vLLM: installed as its own package, find the lib in TheRock's rocm-sdk-core. + try: + import importlib.util + _sdk_core = importlib.util.find_spec("_rocm_sdk_core") + if _sdk_core is not None and _sdk_core.origin: + possible_locations.append( + Path(_sdk_core.origin).parent / "lib/libamd_smi.so.27") + except Exception: + pass +""" + + +def patch_amdsmi(src_dir: Path, out_dir: Path) -> Path: + out_dir = out_dir.resolve() + if out_dir.exists(): + shutil.rmtree(out_dir) + shutil.copytree(src_dir, out_dir, ignore=shutil.ignore_patterns("tests")) + wrapper = out_dir / "amdsmi" / "amdsmi_wrapper.py" + text = wrapper.read_text() + if text.count(_AMDSMI_ANCHOR) != 1: + raise SystemExit(f"amdsmi loader anchor not found in {wrapper}") + wrapper.write_text(text.replace(_AMDSMI_ANCHOR, _AMDSMI_ANCHOR + _AMDSMI_PATCH)) + return out_dir + + +def dockerfile_args(dockerfile: Path) -> dict[str, str]: + args = {} + for m in re.finditer(r'^ARG (\w+)="?([^"\s]*)"?', dockerfile.read_text(), re.M): + args.setdefault(m.group(1), m.group(2)) + return args + + +class _Links(HTMLParser): + def __init__(self) -> None: + super().__init__() + self.hrefs: list[str] = [] + + def handle_starttag(self, tag, attrs): + if tag == "a": + self.hrefs += [v for k, v in attrs if k == "href" and v] + + +def _project_files(index_url: str, project: str) -> list[str]: + page = urljoin(index_url, f"{project}/") + with urllib.request.urlopen(page, timeout=60) as resp: + parser = _Links() + parser.feed(resp.read().decode()) + return [urljoin(page, href.split("#")[0]) for href in parser.hrefs] + + +def _usable(url: str, version: str | None) -> bool: + name = unquote(url.rsplit("/", 1)[-1]) + if ( + version is not None + and f"-{version}-" not in name + and f"-{version}." not in name + ): + return False + if name.endswith(".tar.gz"): + return True + tags = name[: -len(".whl")].split("-")[-3:] + return ( + name.endswith(".whl") + and tags[0] in ("cp312", "py3") + and (tags[2] == "any" or ("x86_64" in tags[2] and "win" not in tags[2])) + ) + + +def torch_device_requirements( + torch_wheel: Path, arches: list[str] +) -> list[tuple[str, str]]: + """(project, version) pairs that torch's device- extras require.""" + with zipfile.ZipFile(torch_wheel) as zf: + metadata = next(n for n in zf.namelist() if n.endswith(".dist-info/METADATA")) + text = zf.read(metadata).decode() + wanted = {f"device-{arch}" for arch in arches} + reqs = [] + for m in re.finditer( + r'^Requires-Dist: ([\w.-]+)\s*==\s*(\S+)\s*;.*extra == "([^"]+)"', text, re.M + ): + if m.group(3) in wanted and (m.group(1), m.group(2)) not in reqs: + reqs.append((m.group(1), m.group(2))) + missing = wanted - set(re.findall(r'extra == "(device-[^"]+)"', text)) + if missing: + raise SystemExit(f"torch has no extras for {sorted(missing)}") + return reqs + + +def external_links(dockerfile: Path, torch_wheel: Path) -> list[str]: + args = dockerfile_args(dockerfile) + index = args["ROCM_RELEASE_WHEELS_MULTIARCH_URL"] + sdk = args["ROCM_SDK_VERSION"] + vision_v, audio_v = args["TORCHVISION_VERSION"], args["TORCHAUDIO_VERSION"] + arches = [a for a in args["PYTORCH_ROCM_ARCH"].split(";") if a] + wanted: list[tuple[str, str | None]] = [ + ("rocm", sdk), + ("rocm-sdk-core", sdk), + ("rocm-sdk-libraries", sdk), + ("torchvision", vision_v), + ("torchaudio", audio_v), + ("rocm-bootstrap", None), + ] + for arch in arches: + wanted += [ + (f"rocm-sdk-device-{arch}", sdk), + (f"amd-torchvision-device-{arch}", vision_v), + ] + wanted += torch_device_requirements(torch_wheel, arches) + urls = [] + for project, version in wanted: + found = [u for u in _project_files(index, project) if _usable(u, version)] + if not found: + raise SystemExit(f"no usable {project}=={version} files on {index}") + urls += found + return urls + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("repack-torch") + p.add_argument("wheel", type=Path) + p.add_argument("--triton-version", required=True) + p.add_argument("--out-dir", type=Path, required=True) + p = sub.add_parser("external-links") + p.add_argument( + "--dockerfile", type=Path, default=Path("docker/Dockerfile.rocm_base") + ) + p.add_argument("--torch-wheel", type=Path, required=True) + p = sub.add_parser("patch-amdsmi") + p.add_argument("src_dir", type=Path) + p.add_argument("--out-dir", type=Path, default=None) + args = parser.parse_args() + + if args.cmd == "repack-torch": + print(repack_torch(args.wheel, args.triton_version, args.out_dir)) + elif args.cmd == "external-links": + print("\n".join(external_links(args.dockerfile, args.torch_wheel))) + else: + out = args.out_dir or Path(tempfile.mkdtemp()) / "amd_smi" + print(patch_amdsmi(args.src_dir, out)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index ec3d26c92c6f..4d0103f51d31 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -405,6 +405,12 @@ def get_rocm_version() -> tuple[int, ...] | None: ): os.environ["HIP_ONLINE_TUNING"] = "1" +# RCCL 2.30's RAS thread can abort the process at exit ("terminate called +# without an active exception") when its shutdown notification fails. +# https://github.com/ROCm/rocm-systems/issues/13572 +if torch.version.hip: + os.environ.setdefault("NCCL_RAS_ENABLE", "0") + @cache def use_rocm_custom_paged_attention( From da9134bc89364465f68aa679564435936d157a14 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sat, 17 Oct 2026 21:49:27 -0500 Subject: [PATCH 811/952] [ROCm][CI] Give the MoE parallel-launch watchdog room for AITER JIT builds (#62161) Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Claude --- tests/kernels/moe/modular_kernel_tools/parallel_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index 235d1b81a493..b6cafe4a4152 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -18,13 +18,16 @@ init_distributed_environment, initialize_model_parallel, ) +from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port ## Parallel Processes Utils P = ParamSpec("P") -LAUNCH_TIMEOUT_S = 900 +# On ROCm the first AITER CK MoE launch JIT-compiles two stage modules that the +# AITER wheel does not prebuild, ~10 min each (ROCm/aiter#5766). +LAUNCH_TIMEOUT_S = 1800 if current_platform.is_rocm() else 900 @dataclasses.dataclass From 7c9042c05addf034366b9a1da5b1d3e3687c052b Mon Sep 17 00:00:00 2001 From: Junru Chen <114817987+chenjr97@users.noreply.github.com> Date: Sun, 18 Oct 2026 14:51:32 +0800 Subject: [PATCH 812/952] [Bugfix][MRV2] Fix speculative warmup deadlock with pipeline parallelism (#62129) Signed-off-by: chenjr97 Co-authored-by: Nick Hill Co-authored-by: vllm-agent --- tests/v1/worker/test_gpu_warmup_blocks.py | 1 + tests/v1/worker/test_mixed_warmup_gate.py | 25 +++++++++++++++++++++++ vllm/v1/worker/gpu/warmup.py | 3 ++- 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index b48f8249e66f..6a0c59874b96 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -88,6 +88,7 @@ def _make_runner( """Stub model runner exposing only what the warmup entry points read.""" return SimpleNamespace( num_speculative_steps=num_spec_steps, + speculative_config=SimpleNamespace() if num_spec_steps else None, adaptive_verification=None, rejection_sampler=None, decode_query_len=num_spec_steps + 1, diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 199ce3b54a37..9d3f60d6f494 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -74,6 +74,7 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_ runner = SimpleNamespace( adaptive_verification=manager, rejection_sampler=rejection_sampler, + speculative_config=SimpleNamespace(), ) def run_steps(model_runner, execute, sample, greedy=False): @@ -91,3 +92,27 @@ def run_steps(model_runner, execute, sample, greedy=False): assert runner.adaptive_verification is manager assert manager.cost_tables is None assert rejection_sampler.enable_adaptive_verification + + +@pytest.mark.parametrize("speculative", [False, True]) +def test_kernel_warmup_runs_same_rounds_on_every_pp_stage(monkeypatch, speculative): + """PP receivers need a matching sender for every speculative warmup round.""" + rounds = [] + monkeypatch.setattr( + warmup, + "_warmup_kernels", + lambda runner, execute, sample, greedy=False: rounds.append(greedy), + ) + for is_last_pp_rank in (False, True): + rounds.clear() + runner = SimpleNamespace( + adaptive_verification=None, + speculative_config=SimpleNamespace() if speculative else None, + rejection_sampler=( + SimpleNamespace(enable_adaptive_verification=False) + if is_last_pp_rank and speculative + else None + ), + ) + warmup.warmup_kernels(runner, _fail, _fail) + assert rounds == ([False, True] if speculative else [False]) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 54181475008e..848d984295fd 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -238,7 +238,8 @@ def warmup_kernels( rejection_sampler.enable_adaptive_verification = False try: _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens) - if rejection_sampler is not None: + # All PP ranks must run the same forwards to match sends and receives. + if model_runner.speculative_config is not None: # Greedy batches skip logits processing, so rejection kernels # also see model-dtype logits instead of FP32. _warmup_kernels( From b510857d37640209a4a3e01a13260ac26f2a93f3 Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:14:06 +0800 Subject: [PATCH 813/952] [Bugfix][Build] Build fused GDN MTP decode for SM80-SM90 with CUDA < 13 (#60896) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: Shengqi Chen Co-authored-by: kimseunghyun --- CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index e67221978c7b..c84e2b9eed43 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1119,6 +1119,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "9.0a;10.0f;12.0f" "${CUDA_ARCHS}") cuda_archs_loose_intersection(FUSED_GDN_DECODE_ARCHS "8.0;8.6;8.9;9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + else() + # The fused GDN decode kernel needs only cp.async and BF16 math. Only the + # family-specific targets (10.0f, 11.0f, 12.0f) need CUDA >= 13.0. + cuda_archs_loose_intersection(FUSED_GDN_DECODE_ARCHS + "8.0;8.6;8.9;9.0a" "${CUDA_ARCHS}") endif() if(FUSED_KDA_DECODE_ARCHS) set(FUSED_KDA_DECODE_SRC From 83c9a547bf9a3b4f2653fb5766886bdc5dba35be Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 18 Oct 2026 00:48:00 -0700 Subject: [PATCH 814/952] [CI][Release] Fix release wheel builds: stop excluding tracked recipes/ from Docker context (#62221) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .dockerignore | 2 - .pre-commit-config.yaml | 6 +++ .../pre_commit/check_dockerignore_tracked.py | 48 +++++++++++++++++++ 3 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 tools/pre_commit/check_dockerignore_tracked.py diff --git a/.dockerignore b/.dockerignore index cdf2c0431a59..e0284129eefa 100644 --- a/.dockerignore +++ b/.dockerignore @@ -6,8 +6,6 @@ dist vllm/*.so vllm/vllm-rs -# Model recipes are rendered by vllm-project/recipes, never part of the image. -/recipes # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 78a672c1104b..d1715b15c13c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -277,6 +277,12 @@ repos: files: ^docker/(Dockerfile|versions\.json)$ pass_filenames: false additional_dependencies: [dockerfile-parse] + - id: check-dockerignore-tracked + name: Check .dockerignore does not exclude tracked files + entry: python tools/pre_commit/check_dockerignore_tracked.py + language: python + always_run: true + pass_filenames: false - id: check-boolean-context-manager name: Check for boolean ops in with-statements entry: python tools/pre_commit/check_boolean_context_manager.py diff --git a/tools/pre_commit/check_dockerignore_tracked.py b/tools/pre_commit/check_dockerignore_tracked.py new file mode 100644 index 000000000000..9aeae9c88e61 --- /dev/null +++ b/tools/pre_commit/check_dockerignore_tracked.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check that no root-anchored `.dockerignore` entry excludes tracked files. + +Docker builds bind-mount `.git` but copy the working tree through the build +context, so git sees any tracked file that `.dockerignore` drops as deleted. +That dirty tree fails `tools/check_repo.sh` in release builds +(`GIT_REPO_CHECK=1`) and adds a `.dYYYYMMDD` dirty suffix to the wheel version. + +Usage: + python tools/pre_commit/check_dockerignore_tracked.py +""" + +import subprocess +import sys +from pathlib import Path + +DOCKERIGNORE = Path(".dockerignore") + + +def main() -> int: + rooted = [ + line.strip().strip("/") + for line in DOCKERIGNORE.read_text(encoding="utf-8").splitlines() + if line.startswith("/") + ] + if not rooted: + return 0 + tracked = subprocess.run( + ["git", "ls-files", "--", *rooted], + check=True, + stdout=subprocess.PIPE, + text=True, + ).stdout.split() + if not tracked: + return 0 + print( + f"{DOCKERIGNORE} excludes {len(tracked)} tracked file(s). Docker builds " + "would see them as deleted and the git tree as dirty, failing " + "tools/check_repo.sh in release builds. First few:" + ) + for path in tracked[:10]: + print(f" {path}") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From 1b2a666391f2a24228a1436b64d547c34cc6b6ce Mon Sep 17 00:00:00 2001 From: Louie Tsai Date: Sun, 18 Oct 2026 01:46:56 -0700 Subject: [PATCH 815/952] [Doc][Recipe] Fix for redhat fp8 xeon support issue (#62245) --- recipes/models/Qwen/Qwen2.5-VL-7B-Instruct.yaml | 1 + recipes/models/Qwen/Qwen3-4B.yaml | 1 + recipes/models/meta-llama/Llama-3.1-8B-Instruct.yaml | 3 +++ recipes/models/meta-llama/Llama-3.3-70B-Instruct.yaml | 3 +++ recipes/models/meta-llama/Llama-4-Scout-17B-16E-Instruct.yaml | 3 +++ 5 files changed, 11 insertions(+) diff --git a/recipes/models/Qwen/Qwen2.5-VL-7B-Instruct.yaml b/recipes/models/Qwen/Qwen2.5-VL-7B-Instruct.yaml index 4b21c7d3a63a..5f27444e466f 100644 --- a/recipes/models/Qwen/Qwen2.5-VL-7B-Instruct.yaml +++ b/recipes/models/Qwen/Qwen2.5-VL-7B-Instruct.yaml @@ -60,6 +60,7 @@ variants: precision: fp8 vram_minimum_gb: 9 description: "RedHatAI FP8 (E4M3) with dynamic per-token activation quantization" + supported_hardware: [trillium] extra_args: - "--kv-cache-dtype" - "fp8" diff --git a/recipes/models/Qwen/Qwen3-4B.yaml b/recipes/models/Qwen/Qwen3-4B.yaml index 8ef62812aeb9..1d2445470243 100644 --- a/recipes/models/Qwen/Qwen3-4B.yaml +++ b/recipes/models/Qwen/Qwen3-4B.yaml @@ -66,6 +66,7 @@ variants: label: "FP8 Dynamic (RedHat)" vram_minimum_gb: 5 description: "RedHatAI FP8 (E4M3) with dynamic per-token activation quantization" + supported_hardware: [trillium] extra_args: - "--kv-cache-dtype" - "fp8" diff --git a/recipes/models/meta-llama/Llama-3.1-8B-Instruct.yaml b/recipes/models/meta-llama/Llama-3.1-8B-Instruct.yaml index 04450187b129..b5d0ffaaabcb 100644 --- a/recipes/models/meta-llama/Llama-3.1-8B-Instruct.yaml +++ b/recipes/models/meta-llama/Llama-3.1-8B-Instruct.yaml @@ -88,6 +88,7 @@ variants: label: "FP8 Dynamic (RedHat)" vram_minimum_gb: 10 description: "RedHatAI FP8 (E4M3) with dynamic per-token activation quantization" + supported_hardware: [h100, h200] extra_args: - "--kv-cache-dtype" - "fp8" @@ -98,6 +99,7 @@ variants: label: "FP8 (RedHat)" vram_minimum_gb: 10 description: "RedHatAI FP8 quantization" + supported_hardware: [h100, h200] extra_args: - "--kv-cache-dtype" - "fp8" @@ -108,6 +110,7 @@ variants: label: "FP8 Block (RedHat)" vram_minimum_gb: 10 description: "RedHatAI FP8 block quantization" + supported_hardware: [h100, h200, xeon6] extra_args: - "--kv-cache-dtype" - "fp8" diff --git a/recipes/models/meta-llama/Llama-3.3-70B-Instruct.yaml b/recipes/models/meta-llama/Llama-3.3-70B-Instruct.yaml index ce0c3e1528e0..8390b98fcdbf 100644 --- a/recipes/models/meta-llama/Llama-3.3-70B-Instruct.yaml +++ b/recipes/models/meta-llama/Llama-3.3-70B-Instruct.yaml @@ -75,6 +75,7 @@ variants: label: "FP8 (RedHat)" vram_minimum_gb: 84 description: "RedHatAI FP8 (E4M3) with dynamic per-token activation quantization" + supported_hardware: [h100, h200, b200, gb200] extra_args: - "--kv-cache-dtype" - "fp8" @@ -95,6 +96,7 @@ variants: label: "FP8 Block (RedHat)" vram_minimum_gb: 84 description: "RedHatAI FP8 block quantization" + supported_hardware: [h100, h200, b200, gb200, xeon6] extra_args: - "--kv-cache-dtype" - "fp8" @@ -105,6 +107,7 @@ variants: label: "NVFP4 (RedHat)" vram_minimum_gb: 46 description: "RedHatAI NVFP4 quantization" + supported_hardware: [b200, gb200] compatible_strategies: - single_node_tp diff --git a/recipes/models/meta-llama/Llama-4-Scout-17B-16E-Instruct.yaml b/recipes/models/meta-llama/Llama-4-Scout-17B-16E-Instruct.yaml index de845242cf3c..3826170e8040 100644 --- a/recipes/models/meta-llama/Llama-4-Scout-17B-16E-Instruct.yaml +++ b/recipes/models/meta-llama/Llama-4-Scout-17B-16E-Instruct.yaml @@ -67,6 +67,7 @@ variants: label: "FP8 (RedHat)" vram_minimum_gb: 131 description: "RedHatAI FP8 (E4M3) with dynamic per-token activation quantization" + supported_hardware: [h100, b200, gb200] extra_args: - "--kv-cache-dtype" - "fp8" @@ -77,6 +78,7 @@ variants: label: "NVFP4 (RedHat)" vram_minimum_gb: 65 description: "RedHatAI compressed-tensors NVFP4 quantization for Blackwell GPUs" + supported_hardware: [b200, gb200] extra_args: - "--kv-cache-dtype" - "fp8" @@ -89,6 +91,7 @@ variants: label: "FP8 Block (RedHat)" vram_minimum_gb: 131 description: "RedHatAI FP8 block quantization" + supported_hardware: [h100, b200, gb200] extra_args: - "--kv-cache-dtype" - "fp8" From 2365177ddda7714502e4df7a1669cc79c0a4dfd0 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 18 Oct 2026 17:26:56 +0800 Subject: [PATCH 816/952] [Doc][RL] Document the weight-transfer HTTP protocol for trainers without vLLM (#62267) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- docs/training/weight_transfer/README.md | 9 +- docs/training/weight_transfer/ipc.md | 38 +++++ docs/training/weight_transfer/nccl.md | 46 ++++++ examples/rl/rlhf_http_ipc_no_vllm.py | 168 +++++++++++++++++++ examples/rl/rlhf_http_nccl_no_vllm.py | 174 ++++++++++++++++++++ tools/pre_commit/check_forbidden_imports.py | 1 + 6 files changed, 435 insertions(+), 1 deletion(-) create mode 100644 examples/rl/rlhf_http_ipc_no_vllm.py create mode 100644 examples/rl/rlhf_http_nccl_no_vllm.py diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index c8382e2a062d..98666e7b285a 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -145,7 +145,7 @@ engine.send_weights() # called on every rank ## API Endpoints -When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first four for you. +When running vLLM as an HTTP server, the following endpoints are available for weight transfer. `HTTPVLLMWeightSyncClient` speaks the first five for you. A trainer may also implement the protocol directly without importing vLLM: see the wire contracts in [NCCL](nccl.md#trainer-without-vllm-over-http) and [IPC](ipc.md#trainer-without-vllm-over-http). This keeps the trainer decoupled from vLLM: it needs no vLLM install or version pin, only this documented HTTP and wire contract, so trainer and server can upgrade independently (non-PyTorch trainers such as JAX can use it too). | Endpoint | Method | Description | | -------- | ------ | ----------- | @@ -153,6 +153,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w | `/start_weight_update` | POST | Start a weight update | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | | `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | +| `/shutdown_weight_transfer_engine` | POST | Release the transfer channel without stopping inference | | `/update_weight_version` | POST | Update `weight_version` without changing model weights | | `/weight_info` | GET | Get the latest committed weight version | | `/pause` | POST | Pause generation before weight sync to handle inflight requests | @@ -162,6 +163,12 @@ When running vLLM as an HTTP server, the following endpoints are available for w !!! note The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set. +!!! warning + If `/update_weights` fails, the model is left mid-reload while `/health` + stays 200, and `/finish_weight_update` is rejected. Before generating again, + resend a full `/start_weight_update` → `/update_weights` → + `/finish_weight_update` (see [#57854](https://github.com/vllm-project/vllm/pull/57854)). + The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, sparse-NCCL, or sharded-RDT transport. ## Extending the System diff --git a/docs/training/weight_transfer/ipc.md b/docs/training/weight_transfer/ipc.md index b71568c3df28..7dab95d66dff 100644 --- a/docs/training/weight_transfer/ipc.md +++ b/docs/training/weight_transfer/ipc.md @@ -110,6 +110,43 @@ mode carries a per-chunk barrier across ranks: without it, a rank could overwrit its buffer while its colocated worker is still reading the current chunk. This is handled inside `send_weights()`. +## Trainer Without vLLM over HTTP + +A colocated trainer can speak the [HTTP endpoints](README.md#api-endpoints) +directly, with only PyTorch and an HTTP client. The server needs +`VLLM_ALLOW_INSECURE_SERIALIZATION=1`; without it, init and start return 200 and +the first `/update_weights` returns 500. Reference: +[`rlhf_http_ipc_no_vllm.py`](../../../examples/rl/rlhf_http_ipc_no_vllm.py). + +`POST /init_weight_transfer_engine` body: `{"init_info": {"packed": false}}`. + +`POST /update_weights` body `{"update_info": {...}}`: + +| Field | Meaning | +| ----- | ------- | +| `names`, `dtype_names`, `shapes` | Checkpoint names, `torch` dtype attribute names (`"bfloat16"`), full shapes | +| `ipc_handles_pickled` | `base64(pickle.dumps(handles))`. Unpacked: a list with one dict per name. Packed: one dict for the shared buffer | +| `tensor_sizes` | Packed only: byte size of each tensor in the chunk, in `names` order | + +- A handle dict is `{gpu_uuid: reduce_tensor(t)[1]}` with one entry per GPU the + server uses (`torch.multiprocessing.reductions.reduce_tensor`). + `gpu_uuid = str(torch.cuda.get_device_properties(i).uuid)`, **without** the + `GPU-` prefix that `nvidia-smi` prints. The worker overwrites `args[6]` + (device index) with its own. +- `ipc_handles` (unpickled) is not JSON-serializable; over HTTP use + `ipc_handles_pickled`. +- Packed: export one `uint8` buffer once and reuse it for every chunk. Copy + `t.contiguous().view(-1).view(torch.uint8)` in `names` order and close a chunk + **before** the next tensor would overflow it (unlike NCCL, which closes after). + A tensor larger than the buffer is an error. Synchronize before each + `/update_weights`; do not overwrite the buffer until the call returns. +- Keep every exported tensor alive until `/finish_weight_update` returns, then + drop the references and call `torch.cuda.ipc_collect()`. +- Multi-rank trainers: each rank exports handles for its own GPU, the ranks + gather the `{uuid: args}` dicts to rank 0 over their own `torch.distributed` + group and merge them per index, and rank 0 sends all HTTP. In packed mode, + barrier before and after each `/update_weights`. + ## Rank-Local Updates For different parameter subsets per worker, `update_info` may be a list indexed @@ -119,4 +156,5 @@ worker to participate. ## Examples - [RLHF with IPC weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - **Start here.** Server and training model share a single GPU; HTTP control plane, CUDA IPC data plane. Launches and tears down its own server +- [IPC weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_ipc_no_vllm.py) - `torch` + `requests` only, `ipc_handles_pickled` over HTTP; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http) - [RLHF with IPC + FSDP2 and expert parallelism](../../../examples/rl/rlhf_ipc_fsdp_ep.py) - Multi-rank trainer colocated with a `--data-parallel-size 4` server on the same 4 GPUs: every FSDP rank builds an engine and joins the handle all-gather, with packed chunking and sleep/wake around the transfer diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index 3a3e207fdc0f..f7c73c1ea422 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/docs/training/weight_transfer/nccl.md @@ -130,6 +130,51 @@ pass — this is the invariant to test first. IPC does not read `metadata()` at all: it derives the update info from iteration as it goes, so it cannot observe a divergence. +## Trainer Without vLLM over HTTP + +A trainer can speak the [HTTP endpoints](README.md#api-endpoints) directly, with +only PyTorch and an HTTP client. It mints the `ncclUniqueId` itself and is NCCL +rank 0. Reference: [`rlhf_http_nccl_no_vllm.py`](../../../examples/rl/rlhf_http_nccl_no_vllm.py). + +`POST /init_weight_transfer_engine` body `{"init_info": {...}}`: + +| Field | Default | Meaning | +| ----- | ------- | ------- | +| `nccl_unique_id_b64` | — | Standard (not URL-safe) base64 of the 128-byte `ncclUniqueId`. Replaces `master_address` / `master_port` | +| `rank_offset` | — | First worker rank; `1` for one trainer rank | +| `world_size` | — | Trainer ranks + `GET /get_world_size` | +| `packed` | `false` | Packed broadcasting | +| `packed_buffer_size_bytes` | 1 GiB | Must match the trainer's chunking | +| `packed_num_buffers` | 2 | Worker-side receive buffers | + +- Worker rank = `rank_offset + dp_index * (TP * PP) + rank`. +- `/init_weight_transfer_engine` returns only after all ranks join: send it in + the background while the trainer calls `ncclCommInitRank(world_size, uid, 0)`. + There is no timeout; if the trainer never joins, the workers wait forever. +- Right after init every worker issues a warm-up `all_reduce` (sum) of + `torch.zeros(1)` (one float32 element). The trainer must issue the same one + before any other collective. + +Each sync: `POST /start_weight_update`, then `POST /update_weights` with +`{"update_info": {"names", "dtype_names", "shapes"}}` **in the background** +while rank 0 broadcasts in `names` order, then `POST /finish_weight_update`. +`dtype_names` are `torch` attribute names (`"bfloat16"`); `shapes` are full +checkpoint shapes. + +- Unpacked: one broadcast per tensor. +- Packed: concatenate `t.contiguous().view(-1).view(torch.uint8)` in `names` + order; broadcast a chunk as soon as its size **exceeds** + `packed_buffer_size_bytes` (the crossing tensor stays in the chunk), then the + remainder. + +Trainer-side communicators without vLLM: + +| Communicator | Status | +| ------------ | ------ | +| `torch.distributed.ProcessGroupNCCL` on a private `TCPStore` + `eager_connect_single_device` | Works; the unique id is read from store key `"0"`, a torch implementation detail (used by the example) | +| `ctypes` against torch's bundled `libnccl.so.2` | Works | +| `torch.cuda.nccl.init_rank` | Fails on Python 3.10+ (`PY_SSIZE_T_CLEAN` `SystemError`) | + ## Sparse NCCL Sparse, flat-index weight patches use @@ -196,6 +241,7 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still ## Examples - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server +- [NCCL weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_nccl_no_vllm.py) - `torch` + `requests` only, `nccl_unique_id_b64` rendezvous; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http) - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model; uses an in-process `AsyncLLMEngine` with `RayVLLMWeightSyncClient` diff --git a/examples/rl/rlhf_http_ipc_no_vllm.py b/examples/rl/rlhf_http_ipc_no_vllm.py new file mode 100644 index 000000000000..20a92a149c2a --- /dev/null +++ b/examples/rl/rlhf_http_ipc_no_vllm.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA IPC weight sync over HTTP from a trainer that never imports vLLM. + +The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the +stdlib, and speaks the wire protocol in +`docs/training/weight_transfer/ipc.md` ("Trainer Without vLLM over HTTP"): +it exports its tensors with `torch.multiprocessing.reductions.reduce_tensor` +and sends the handles, pickled and base64-encoded, in `ipc_handles_pickled`. + +IPC requires colocation: this process holds a full copy of the weights on every +GPU the server uses, and keys each handle by that GPU's UUID. A multi-process +trainer would instead gather the per-rank `{uuid: args}` dicts to one rank. + +Start a server; IPC handles are pickled, so it needs +`VLLM_ALLOW_INSECURE_SERIALIZATION=1` (the weights it starts with are discarded): + + VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \ + CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-0.6B \ + --load-format dummy --enforce-eager --gpu-memory-utilization 0.4 \ + --weight-transfer-config '{"backend": "ipc"}' + +Then run the trainer on the same GPU(s) (`--packed` for chunked transfer): + + CUDA_VISIBLE_DEVICES=0 python examples/rl/rlhf_http_ipc_no_vllm.py \ + --checkpoint /path/to/Qwen3-0.6B + +For a TP=2 server on GPUs 0-1, use `CUDA_VISIBLE_DEVICES=0,1` on both sides and +pass `--num-gpus 2` to the trainer. +""" + +import sys + + +class _BlockVLLM: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] == "vllm": + raise ImportError(f"this trainer must not import {name!r}") + return None + + +sys.meta_path.insert(0, _BlockVLLM()) + +import argparse # noqa: E402 +import glob # noqa: E402 +import pickle # noqa: E402 + +import pybase64 as base64 # noqa: E402 +import requests # noqa: E402 +import torch # noqa: E402 +from safetensors.torch import load_file # noqa: E402 +from torch.multiprocessing.reductions import reduce_tensor # noqa: E402 + +parser = argparse.ArgumentParser() +parser.add_argument("--url", default="http://localhost:8000") +parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors") +parser.add_argument("--num-gpus", type=int, default=1) +parser.add_argument("--packed", action="store_true") +parser.add_argument("--packed-buffer-size-bytes", type=int, default=512 << 20) +args = parser.parse_args() + + +def post(path: str, body: dict | None = None) -> dict: + r = requests.post(args.url + path, json=body or {}, timeout=600) + r.raise_for_status() + return r.json() + + +def generate(prompts: list[str]) -> list[str]: + model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"] + body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0} + choices = post("/v1/completions", body)["choices"] + return [c["text"] for c in sorted(choices, key=lambda c: c["index"])] + + +# Weights are {gpu_uuid: {name: tensor}}: one full copy per colocated GPU. The +# UUID is str(torch.cuda.get_device_properties(i).uuid), with no "GPU-" prefix. +Weights = dict[str, dict[str, torch.Tensor]] + + +def update_weights( + weights: Weights, + names: list[str], + ipc_handles: list[dict[str, tuple]] | dict[str, tuple], + tensor_sizes: list[int] | None = None, +) -> None: + w = next(iter(weights.values())) + update_info = { + "names": names, + "dtype_names": [str(w[n].dtype).removeprefix("torch.") for n in names], + "shapes": [list(w[n].shape) for n in names], + "ipc_handles_pickled": base64.b64encode(pickle.dumps(ipc_handles)).decode(), + } + if tensor_sizes is not None: + update_info["tensor_sizes"] = tensor_sizes + post("/update_weights", {"update_info": update_info}) + + +def send_unpacked(weights: Weights, names: list[str]) -> list[torch.Tensor]: + """One /update_weights call; one {uuid: args} dict per parameter.""" + copies = {u: {n: w[n].contiguous() for n in names} for u, w in weights.items()} + handles = [{u: reduce_tensor(copies[u][n])[1] for u in copies} for n in names] + update_weights(weights, names, handles) + return [t for c in copies.values() for t in c.values()] + + +def send_packed(weights: Weights, names: list[str], cap: int) -> list[torch.Tensor]: + """One /update_weights call per chunk; one {uuid: args} dict per chunk.""" + buffers = { + u: torch.empty(cap, dtype=torch.uint8, device=w[names[0]].device) + for u, w in weights.items() + } + handle = {u: reduce_tensor(b)[1] for u, b in buffers.items()} # export once + chunk: list[str] = [] + sizes: list[int] = [] + + def flush() -> None: + for b in buffers.values(): + torch.accelerator.synchronize(b.device) # copies must land before the read + update_weights(weights, chunk, handle, sizes) + chunk.clear() + sizes.clear() + + for n in names: + nbytes = next(iter(weights.values()))[n].nbytes + assert nbytes <= cap, f"{n} ({nbytes} B) exceeds the buffer" + if chunk and sum(sizes) + nbytes > cap: # close *before* overflowing + flush() + offset = sum(sizes) + for u, w in weights.items(): + flat = w[n].contiguous().view(-1).view(torch.uint8) + buffers[u][offset : offset + nbytes].copy_(flat) + chunk.append(n) + sizes.append(nbytes) + if chunk: + flush() + return list(buffers.values()) + + +def main() -> None: + weights: Weights = {} + for i in range(args.num_gpus): + uuid = str(torch.cuda.get_device_properties(i).uuid) + weights[uuid] = {} + for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")): + weights[uuid].update(load_file(path, device=f"cuda:{i}")) + names = sorted(next(iter(weights.values()))) + + prompts = ["The capital of France is", "def fibonacci(n):"] + print("before:", generate(prompts)) + + post("/init_weight_transfer_engine", {"init_info": {"packed": args.packed}}) + post("/start_weight_update") + if args.packed: + exported = send_packed(weights, names, args.packed_buffer_size_bytes) + else: + exported = send_unpacked(weights, names) + post("/finish_weight_update", {"weight_version": None}) + # The worker reads the exported memory until finish; only then release it. + del exported + torch.cuda.ipc_collect() + + print("after: ", generate(prompts)) + assert not [m for m in sys.modules if m.split(".")[0] == "vllm"] + + +if __name__ == "__main__": + main() diff --git a/examples/rl/rlhf_http_nccl_no_vllm.py b/examples/rl/rlhf_http_nccl_no_vllm.py new file mode 100644 index 000000000000..c9d8abad365b --- /dev/null +++ b/examples/rl/rlhf_http_nccl_no_vllm.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NCCL weight sync over HTTP from a trainer that never imports vLLM. + +The trainer uses only `torch`, `requests`, `safetensors`, `pybase64` and the +stdlib, and speaks the wire protocol in +`docs/training/weight_transfer/nccl.md` ("Trainer Without vLLM over HTTP"): +`nccl_unique_id_b64` rendezvous, a one-element warm-up `all_reduce`, then +`/start_weight_update` -> `/update_weights` (concurrent with the rank-0 +broadcast) -> `/finish_weight_update`. + +The NCCL communicator is a plain `torch.distributed.ProcessGroupNCCL` over a +private `TCPStore` that only the trainer joins. Its `ncclUniqueId` is read back +from store key "0", which is a torch implementation detail. + +Start a server on GPUs 0-1 (the weights it starts with are discarded): + + VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen3-0.6B \ + --tensor-parallel-size 2 --load-format dummy --enforce-eager \ + --weight-transfer-config '{"backend": "nccl"}' + +Then run the trainer on GPU 2 (`--packed` for packed broadcasting): + + CUDA_VISIBLE_DEVICES=2 python examples/rl/rlhf_http_nccl_no_vllm.py \ + --checkpoint /path/to/Qwen3-0.6B +""" + +import sys + + +class _BlockVLLM: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] == "vllm": + raise ImportError(f"this trainer must not import {name!r}") + return None + + +sys.meta_path.insert(0, _BlockVLLM()) + +import argparse # noqa: E402 +import datetime # noqa: E402 +import glob # noqa: E402 +import threading # noqa: E402 + +import pybase64 as base64 # noqa: E402 +import requests # noqa: E402 +import torch # noqa: E402 +import torch.distributed as dist # noqa: E402 +from safetensors.torch import load_file # noqa: E402 + +parser = argparse.ArgumentParser() +parser.add_argument("--url", default="http://localhost:8000") +parser.add_argument("--checkpoint", required=True, help="dir with *.safetensors") +parser.add_argument("--packed", action="store_true") +parser.add_argument("--packed-buffer-size-bytes", type=int, default=64 << 20) +parser.add_argument("--store-port", type=int, default=29612) +args = parser.parse_args() + + +def post(path: str, body: dict | None = None) -> dict: + r = requests.post(args.url + path, json=body or {}, timeout=600) + r.raise_for_status() + return r.json() + + +class PostInBackground(threading.Thread): + """An HTTP call that only returns once this process joins the collective.""" + + def __init__(self, path: str, body: dict) -> None: + super().__init__() + self.path, self.body = path, body + self.error: BaseException | None = None + self.start() + + def run(self) -> None: + try: + post(self.path, self.body) + except BaseException as e: + self.error = e + + def result(self) -> None: + self.join() + if self.error is not None: + raise self.error + + +def generate(prompts: list[str]) -> list[str]: + model = requests.get(args.url + "/v1/models", timeout=60).json()["data"][0]["id"] + body = {"model": model, "prompt": prompts, "max_tokens": 16, "temperature": 0} + choices = post("/v1/completions", body)["choices"] + return [c["text"] for c in sorted(choices, key=lambda c: c["index"])] + + +def main() -> None: + device = torch.device("cuda:0") + torch.accelerator.set_device_index(device) + weights: dict[str, torch.Tensor] = {} + for path in sorted(glob.glob(f"{args.checkpoint}/*.safetensors")): + weights.update(load_file(path, device="cuda:0")) + names = sorted(weights) + + prompts = ["The capital of France is", "def fibonacci(n):"] + print("before:", generate(prompts)) + + # Trainer is NCCL rank 0; vLLM worker i is rank 1 + i. + world_size = 1 + requests.get(args.url + "/get_world_size").json()["world_size"] + store = dist.TCPStore( + "127.0.0.1", args.store_port, is_master=True, wait_for_workers=False + ) + pg = dist.ProcessGroupNCCL(store, 0, world_size, dist.ProcessGroupNCCL.Options()) + connect = threading.Thread(target=pg.eager_connect_single_device, args=(device,)) + connect.start() # publishes the unique id, then waits in ncclCommInitRank + reader = dist.TCPStore("127.0.0.1", args.store_port, is_master=False) + reader.wait(["0"], datetime.timedelta(seconds=60)) + unique_id = reader.get("0") + assert len(unique_id) == 128 + + init = PostInBackground( + "/init_weight_transfer_engine", + { + "init_info": { + "nccl_unique_id_b64": base64.b64encode(unique_id).decode(), + "rank_offset": 1, + "world_size": world_size, + "packed": args.packed, + "packed_buffer_size_bytes": args.packed_buffer_size_bytes, + } + }, + ) + connect.join() + pg.allreduce([torch.zeros(1, device=device)]).wait() # vLLM's warm-up + init.result() + + def broadcast(t: torch.Tensor) -> None: + opts = dist.BroadcastOptions() + opts.rootRank = 0 + pg.broadcast([t], opts).wait() + + post("/start_weight_update") + update_info = { + "names": names, + "dtype_names": [str(weights[n].dtype).removeprefix("torch.") for n in names], + "shapes": [list(weights[n].shape) for n in names], + } + update = PostInBackground("/update_weights", {"update_info": update_info}) + if not args.packed: + for n in names: + broadcast(weights[n].contiguous()) + else: + # Same chunking as the worker: a chunk closes right after the tensor + # that pushes it past the buffer size. + chunk: list[torch.Tensor] = [] + size = 0 + for n in names: + chunk.append(weights[n].contiguous().view(-1).view(torch.uint8)) + size += chunk[-1].numel() + if size > args.packed_buffer_size_bytes: + broadcast(torch.cat(chunk)) + chunk, size = [], 0 + if chunk: + broadcast(torch.cat(chunk)) + torch.accelerator.synchronize() + update.result() + post("/finish_weight_update", {"weight_version": None}) + + print("after: ", generate(prompts)) + post("/shutdown_weight_transfer_engine") + pg.shutdown() + + assert not [m for m in sys.modules if m.split(".")[0] == "vllm"] + + +if __name__ == "__main__": + main() diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 740043bfa05c..e5812e667358 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -59,6 +59,7 @@ class ForbiddenImport: "vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", + "examples/rl/rlhf_http_ipc_no_vllm.py", "vllm/model_executor/model_loader/weight_cache/protocol.py", "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", From 3f1b04e045beb841329519bacafa9b2ded86a762 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 18 Oct 2026 18:37:39 +0800 Subject: [PATCH 817/952] [Bugfix] Fix Opencv video backend frames count missing for WebM video inputs (#61759) Signed-off-by: Isotr0py --- tests/multimodal/test_video.py | 36 +++++++++++++- tests/multimodal/utils.py | 63 ++++++++++++++++++++++-- vllm/multimodal/video_decoders/opencv.py | 32 ++++++++++-- 3 files changed, 122 insertions(+), 9 deletions(-) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 1c7730a37ef8..16334fdb0dc7 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -49,7 +49,9 @@ from .utils import ( create_edit_list_trimmed_video, + create_live_webm_video, create_long_gop_video, + create_raw_h264_video, create_video_from_image, ) @@ -1261,8 +1263,8 @@ def test_qwen3_vl_unknown_source_fps_matches_hf( the metadata carries the fps HF assumes, which the frame timestamps use.""" get_video_metadata = OpenCVVideoBackendMixin.get_video_metadata - def unknown_fps(cap) -> VideoSourceMetadata: - return get_video_metadata(cap)._replace(original_fps=0.0, duration=0.0) + def unknown_fps(data) -> VideoSourceMetadata: + return get_video_metadata(data)._replace(original_fps=0.0, duration=0.0) monkeypatch.setattr( OpenCVVideoBackendMixin, "get_video_metadata", staticmethod(unknown_fps) @@ -1476,6 +1478,36 @@ def test_video_backend_handles_edit_list_trimmed_video( assert qwen_frames.shape[0] >= 4 +def test_opencv_demux_counts_frames_when_video_header_lacks_metadata( + monkeypatch: pytest.MonkeyPatch, +): + """Streamed WebM and raw H.264 carry no frame count/duration in their + headers, so OpenCV reports a bogus negative ``CAP_PROP_FRAME_COUNT``. + The loader must count the frames via a demux-only pass, without + decoding. + """ + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") + loader = VIDEO_LOADER_REGISTRY.load("opencv") + + webm_data = create_live_webm_video(num_frames=20, fps=10) + frames, metadata = loader.load_bytes(webm_data, backend="opencv") + assert frames.shape[0] == 20 + assert metadata["total_num_frames"] == 20 + assert metadata["duration"] == pytest.approx(2.0) + + h264_data = create_raw_h264_video(num_frames=20, fps=25) + frames, metadata = loader.load_bytes(h264_data, backend="opencv") + assert frames.shape[0] == 20 + assert metadata["total_num_frames"] == 20 + assert metadata["duration"] == pytest.approx(0.8) + + # The Qwen samplers used to raise on the bogus negative frame count. + qwen_frames, qwen_metadata = Qwen2VLVideoBackend.load_bytes(webm_data) + assert qwen_metadata["total_num_frames"] == 20 + assert qwen_frames.shape[0] >= 4 + + # ============================================================================ # Frame Recovery Tests # ============================================================================ diff --git a/tests/multimodal/utils.py b/tests/multimodal/utils.py index 23a05a094422..d74f68453146 100644 --- a/tests/multimodal/utils.py +++ b/tests/multimodal/utils.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io + import cv2 import numpy as np import numpy.typing as npt @@ -77,8 +79,6 @@ def create_long_gop_video( The marker lets a test recover which frame the decoder actually returned, independent of any metadata label. """ - import io - import av buf = io.BytesIO() @@ -103,6 +103,63 @@ def create_long_gop_video( return buf.getvalue() +def create_live_webm_video( + num_frames: int = 20, + fps: int = 10, + width: int = 64, + height: int = 64, +) -> bytes: + """Encode a WebM stream without container duration or frame-count metadata. + + The ``live`` muxing option mimics writing to a pipe: the Segment size is + unknown and no Duration or Cues elements are written, so decoder backends + cannot learn the frame count from the header. + """ + import av + + buf = io.BytesIO() + with av.open(buf, mode="w", format="webm", options={"live": "1"}) as container: + stream = container.add_stream("libvpx", rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + for i in range(num_frames): + img = np.zeros((height, width, 3), dtype=np.uint8) + img[:, :, 1] = i % 256 + frame = av.VideoFrame.from_ndarray(img, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return buf.getvalue() + + +def create_raw_h264_video( + num_frames: int = 20, + fps: int = 25, + width: int = 64, + height: int = 64, +) -> bytes: + """Encode a raw H.264 Annex B elementary stream (no container metadata).""" + import av + + buf = io.BytesIO() + with av.open(buf, mode="w", format="h264") as container: + stream = container.add_stream("h264", rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + for i in range(num_frames): + img = np.zeros((height, width, 3), dtype=np.uint8) + img[:, :, 1] = i % 256 + frame = av.VideoFrame.from_ndarray(img, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return buf.getvalue() + + def create_edit_list_trimmed_video( num_frames: int = 90, trim_start_frame: int = 60, @@ -117,8 +174,6 @@ def create_edit_list_trimmed_video( trailing frames are presentable — the shape produced by lossless trims (e.g. ``ffmpeg -ss ... -c copy``). Returns (video_bytes, visible_frames). """ - import io - import av src = create_long_gop_video(num_frames=num_frames, fps=fps) diff --git a/vllm/multimodal/video_decoders/opencv.py b/vllm/multimodal/video_decoders/opencv.py index 3ca141c738c4..88d7fb837551 100644 --- a/vllm/multimodal/video_decoders/opencv.py +++ b/vllm/multimodal/video_decoders/opencv.py @@ -38,7 +38,9 @@ def decode_opencv( int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), ) - source = loader_cls._prepare_source(OpenCVVideoBackendMixin.get_video_metadata(cap)) + source = loader_cls._prepare_source( + OpenCVVideoBackendMixin.get_video_metadata(data) + ) frame_idx = loader_cls.compute_frames_index_to_sample( source=source, target=target, **sampling_kwargs ) @@ -74,10 +76,34 @@ def open_video_capture(cls, data: bytes) -> "cv2.VideoCapture": raise ValueError("Could not open video stream") return cap + @classmethod + def demux_count_frames(cls, data: bytes) -> int: + """Count frames with OpenCV's demuxer-only mode.""" + # NOTE: buf must stay referenced until after release(). + # OpenCV's Python stream binding does not safely own the stream object + buf = BytesIO(data) + demux_cap = cv2.VideoCapture( + buf, cls.get_cv2_video_api(), [cv2.CAP_PROP_FORMAT, -1] + ) + count = 0 + while demux_cap.grab(): + count += 1 + demux_cap.release() + if count == 0: + raise ValueError("Could not determine the number of video frames") + return count + @staticmethod - def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: + def get_video_metadata(data: bytes) -> VideoSourceMetadata: + cap = OpenCVVideoBackendMixin.open_video_capture(data) total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) original_fps = cap.get(cv2.CAP_PROP_FPS) + if total_frames_num <= 0: + # Streamed WebM/MKV and raw H.264 miss frame count in headers; + # We need to use demux frames count for retrieval + total_frames_num = OpenCVVideoBackendMixin.demux_count_frames(data) + duration = total_frames_num / original_fps if original_fps > 0 else 0 + return VideoSourceMetadata(total_frames_num, original_fps, duration) # CAP_PROP_FRAME_COUNT counts every physical sample in the container, # overstating the presentable frames when an mp4 edit list hides the # decode lead-in (e.g. a lossless `ffmpeg -ss ... -c copy` trim). @@ -87,7 +113,7 @@ def get_video_metadata(cap: "cv2.VideoCapture") -> VideoSourceMetadata: # may land mid-stream on corrupt files, so only trust it when the # stream really ends there. A one-frame tolerance keeps mere header # miscounts from being "corrected". - if total_frames_num > 0 and cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): + if cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1.0): visible_frames = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) at_stream_end = not cap.grab() cap.set(cv2.CAP_PROP_POS_FRAMES, 0) From 953fa442fafedbf51c3310d577474f80e457964a Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 18 Oct 2026 04:00:16 -0700 Subject: [PATCH 818/952] [CI] Per-test timeout for Language Models (Extended Pooling) (#62279) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_language.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index be9f6927b314..94c16ba118e4 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -210,7 +210,10 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language/pooling commands: - - pytest -v -s models/language/pooling -m 'not core_model' + # Per-test timeout: a hung test fails with a stack dump instead of holding + # the job until Buildkite cancels it, which also blocks the full run's + # coverage collection. A whole shard takes 12-21 minutes. + - pytest -v -s models/language/pooling -m 'not core_model' --timeout=900 --timeout-method=thread mirror: amd: parallelism: 4 From 3c13653738dad0acd877c23373816fefffbdd373 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Delhaye?= Date: Sun, 18 Oct 2026 13:30:12 +0200 Subject: [PATCH 819/952] [Frontend] Score with the chat template declared by Sentence Transformers CrossEncoders (#60877) Signed-off-by: rdlhio --- docs/models/pooling_models/scoring.md | 25 ++++++ docs/models/pooling_models/specific_models.md | 8 +- .../pooling/scoring/test_io_processor_unit.py | 2 + .../pooling_mteb_test/test_nemotron.py | 8 ++ .../pooling/test_llama_nemotron_vl.py | 26 ++---- .../pooling/scoring/io_processor.py | 64 +++++++++++++- vllm/transformers_utils/config.py | 86 +++++++++++++++++++ 7 files changed, 192 insertions(+), 27 deletions(-) diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 7ff186b5cdb2..3bda12a09c89 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -466,6 +466,31 @@ This approach is more robust than index-based access (`messages[0]`, `messages[1 Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) +#### Templates declared by Sentence Transformers checkpoints + +vLLM does not apply a checkpoint's chat template by default, since many rerankers inherit an unrelated template from their base LLM. The exception is a Sentence Transformers (v5.4+) `CrossEncoder` with a single `sequence-classification` Transformer module whose `sentence_bert_config.json` declares the `message` modality. Sentence Transformers formats the pairs of such checkpoints as `query`/`document` messages, so vLLM does the same with the saved template: + +```json +{ + "transformer_task": "sequence-classification", + "modality_config": { + "text": {"method": "forward", "method_output_name": "logits"}, + "message": {"method": "forward", "method_output_name": "logits", "format": "flat"} + }, + "module_output_name": "scores", + "processing_kwargs": {"chat_template": {"add_generation_prompt": true}} +} +``` + +- The default chat template is used, unless `processing_kwargs.chat_template.chat_template` names another template saved with the tokenizer (e.g. `additional_chat_templates/score.jinja` for `"score"`). +- The other `processing_kwargs.chat_template` settings, such as `add_generation_prompt`, are passed to the template; the request's `chat_template_kwargs` override them. +- Message contents are passed as strings, as with `--chat-template`. +- vLLM does not restore the template suffix after truncation (`restore_suffix`). Use `max_tokens_per_doc` rather than `truncate_prompt_tokens` to truncate long documents. +- `--chat-template` (or the `chat_template` argument of `LLM.score`) takes precedence over the saved template and its settings. +- Models with a built-in score template keep using it. + +The template in use is logged at startup. + ### Enable/disable activation You can enable or disable activation via `use_activation` only works for cross-encoder models. diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 0c7b0a1bf3fc..5b2871ca923f 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -313,14 +313,12 @@ Start the server: ```shell vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \ --runner pooling \ - --trust-remote-code \ - --chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja + --trust-remote-code ``` !!! note - The chat template bundled with this checkpoint's tokenizer is not suitable - for the Score/Rerank APIs. Use the provided override template when serving: - `examples/pooling/score/template/nemotron-vl-rerank.jinja`. + The checkpoint declares its score template in its Sentence Transformers config, so vLLM applies it automatically (see [Templates declared by Sentence Transformers checkpoints](scoring.md#templates-declared-by-sentence-transformers-checkpoints)). + Revisions older than October 2026 need `--chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja`. Score a text query against an image document: diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 7f2c157f2516..d5a585f5796a 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -65,6 +65,8 @@ def llm_reranker_processor() -> CrossEncoderIOProcessor: processor.supports_score_template = False processor.use_sep_token = False processor.model = None + processor.saved_chat_template = None + processor.saved_chat_template_kwargs = {} return processor diff --git a/tests/models/language/pooling_mteb_test/test_nemotron.py b/tests/models/language/pooling_mteb_test/test_nemotron.py index 711224a1a29f..3c3618de6510 100644 --- a/tests/models/language/pooling_mteb_test/test_nemotron.py +++ b/tests/models/language/pooling_mteb_test/test_nemotron.py @@ -39,6 +39,14 @@ is_prefix_caching_supported=False, is_chunked_prefill_supported=False, ), + # Declares its score template in its Sentence Transformers config, so no + # chat template is passed. + RerankModelInfo( + "nvidia/llama-nemotron-rerank-vl-1b-v2", + architecture="LlamaNemotronVLForSequenceClassification", + mteb_score=0.33919, + mteb_tol=1e-2, + ), ] diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index 3cbeda2e49e1..227542023a5c 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -10,7 +10,6 @@ from collections.abc import Sequence from io import BytesIO -from pathlib import Path import pybase64 as base64 import pytest @@ -202,13 +201,6 @@ def test_models( RERANKER_MODELS = ["nvidia/llama-nemotron-rerank-vl-1b-v2"] -# The tokenizer's built-in chat template is not suitable for the Score/Rerank -# APIs (it's inherited from the base LLM). We must use the provided override. -_RERANKER_SCORE_TEMPLATE = ( - Path(__file__).parents[4] - / "examples/pooling/score/template/nemotron-vl-rerank.jinja" -).read_text() - RERANKER_TEXT_QUERY = "How is AI improving the intelligence and capabilities of robots?" RERANKER_TEXT_DOCS = [ "AI enables robots to perceive, plan, and act autonomously.", @@ -281,7 +273,11 @@ def _run_vllm_reranker( dtype: str, input_cases: Sequence[RerankerCase], ) -> list[list[float]]: - """Run all vLLM reranker cases in one model lifecycle.""" + """Run all vLLM reranker cases in one model lifecycle. + + The checkpoint declares its score template in its Sentence Transformers + config, so no chat template is passed. + """ with vllm_runner( model, runner="pooling", @@ -298,11 +294,7 @@ def _run_vllm_reranker( if not has_images: queries = [query] * len(docs) doc_texts = [doc_text for doc_text, _ in docs] - outputs = vllm_model.score( - queries, - doc_texts, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + outputs = vllm_model.score(queries, doc_texts) else: query_params = [ ScoreMultiModalParam( @@ -334,11 +326,7 @@ def _run_vllm_reranker( ) doc_params.append(ScoreMultiModalParam(content=content)) - raw_outputs = vllm_model.llm.score( - query_params, - doc_params, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + raw_outputs = vllm_model.llm.score(query_params, doc_params) outputs = [output.outputs.score for output in raw_outputs] scores_per_case.append(outputs) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 7c4a1fe9170f..2613a4073a8e 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -9,14 +9,17 @@ from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.exceptions import ChatTemplateResolutionError from vllm.logger import init_logger -from vllm.renderers import TokenizeParams -from vllm.renderers.hf import safe_apply_chat_template +from vllm.renderers import TokenizeParams, merge_kwargs +from vllm.renderers.hf import resolve_chat_template, safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, parse_model_prompt, prompt_to_seq, ) from vllm.tasks import PoolingTask +from vllm.transformers_utils.config import ( + get_sentence_transformers_chat_template_kwargs, +) from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor @@ -457,9 +460,22 @@ def __init__(self, *args, **kwargs): self.model = model if self.supports_score_template else None self.use_sep_token = self.model_config.use_sep_token + self.saved_chat_template: str | None = None + self.saved_chat_template_kwargs: dict[str, Any] = {} + if not self.supports_score_template: + try: + self.saved_chat_template, self.saved_chat_template_kwargs = ( + self._get_saved_chat_template() + ) + except ValueError: + # An explicit chat template takes precedence anyway. + if self.chat_template is None: + raise + if ( getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False) and self.chat_template is None + and self.saved_chat_template is None ): suggested_template = ( "examples/pooling/score/template/qwen3_vl_reranker.jinja" @@ -475,6 +491,41 @@ def __init__(self, *args, **kwargs): suggested_template, ) + def _get_saved_chat_template(self) -> tuple[str | None, dict[str, Any]]: + """Get the chat template and kwargs declared by the Sentence + Transformers config, if any.""" + saved_kwargs = get_sentence_transformers_chat_template_kwargs( + self.model_config.model, self.model_config.revision + ) + if saved_kwargs is None: + return None, {} + + saved_kwargs = dict(saved_kwargs) + name = saved_kwargs.pop("chat_template", None) + if name == "default": + name = None + template = resolve_chat_template( + self.tokenizer, + chat_template=name, + tools=None, + model_config=self.model_config, + ) + # Hugging Face returns an unknown template name unchanged. + if template is None or template == name: + raise ValueError( + "The Sentence Transformers config declares message inputs, but " + f"the checkpoint has no {name or 'default'!r} chat template." + ) + + if self.chat_template is None: + logger.info_once( + "Scoring with the %r chat template declared by the Sentence " + "Transformers config. Its suffix is not restored after " + "truncation; use max_tokens_per_doc to truncate documents.", + name or "default", + ) + return template, saved_kwargs + ####################################### # online APIs @@ -648,6 +699,12 @@ def get_score_prompt( model_config = self.model_config tokenizer = self.tokenizer + if chat_template is None and self.saved_chat_template is not None: + chat_template = self.saved_chat_template + chat_template_kwargs = merge_kwargs( + self.saved_chat_template_kwargs, chat_template_kwargs + ) + prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data( data_1, data_2, @@ -703,7 +760,8 @@ def default_tokenizer_encode(): prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) return full_prompt, prompt_inputs - # FIXME: For now, we only apply a template when one is explicitly provided. + # FIXME: For now, we only apply a template when one is explicitly provided, + # either directly or by the Sentence Transformers config. # We cannot rely on the tokenizer's chat template because many models # inherit junk templates from their base LLM, which breaks both the models # and the tests that use them. diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c67d468a13b3..8587c02b680c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -89,6 +89,10 @@ `ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" +_ST_TRANSFORMER_MODULE_TYPES = { + "sentence_transformers.models.Transformer", + "sentence_transformers.base.modules.transformer.Transformer", +} _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", @@ -1354,6 +1358,88 @@ def try_get_dense_modules( return None +@cache +def get_sentence_transformers_chat_template_kwargs( + model: str | Path, + revision: str | None = None, +) -> dict[str, Any] | None: + """Get the saved chat template settings of a Sentence Transformers + CrossEncoder whose single Transformer module declares the `message` + modality. + + Sentence Transformers formats the pairs of such checkpoints as + `query`/`document` messages with the saved chat template. + + Args: + model: Local checkpoint path or Hugging Face model ID. + revision: Model revision to inspect. + + Returns: + The saved `processing_kwargs.chat_template` (possibly empty), or + `None` if the checkpoint does not declare the `message` modality. + + Raises: + ValueError: If the saved chat template settings are unsupported. + + """ + try: + st_config = get_hf_file_to_dict( + "config_sentence_transformers.json", model, revision + ) + if not isinstance(st_config, dict) or ( + st_config.get("model_type") != "CrossEncoder" + ): + return None + modules = get_hf_file_to_dict("modules.json", model, revision) + transformer_config = get_hf_file_to_dict( + "sentence_bert_config.json", model, revision + ) + except Exception: + logger.debug("Failed to read the Sentence Transformers config", exc_info=True) + return None + + if isinstance(modules, dict): + modules = modules.get("modules") + if not ( + isinstance(modules, list) + and len(modules) == 1 + and isinstance(modules[0], dict) + and modules[0].get("type") in _ST_TRANSFORMER_MODULE_TYPES + and not modules[0].get("path") + and isinstance(transformer_config, dict) + and transformer_config.get("transformer_task") == "sequence-classification" + and isinstance( + modality_config := transformer_config.get("modality_config"), dict + ) + and "message" in modality_config + ): + return None + + processing_kwargs = transformer_config.get("processing_kwargs") or {} + chat_template_kwargs = ( + processing_kwargs.get("chat_template") or {} + if isinstance(processing_kwargs, dict) + else None + ) + if not isinstance(chat_template_kwargs, dict): + raise ValueError( + "The Sentence Transformers processing_kwargs.chat_template must be " + "a JSON object." + ) + chat_template_kwargs = dict(chat_template_kwargs) + # vLLM does not restore the template suffix after truncation. + chat_template_kwargs.pop("restore_suffix", None) + if not isinstance(chat_template_kwargs.get("chat_template", ""), str): + raise ValueError( + "The Sentence Transformers chat template name must be a string." + ) + if reserved := {"tokenize", "tools"} & chat_template_kwargs.keys(): + raise ValueError( + f"Unsupported Sentence Transformers chat template settings: {reserved}" + ) + return chat_template_kwargs + + def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]: return { param_name: info From 494663940b9b44f667c4f61c8c2407ba8e725ceb Mon Sep 17 00:00:00 2001 From: Harry Huang Date: Sun, 18 Oct 2026 19:38:21 +0800 Subject: [PATCH 820/952] [Qwen4Exp] Clean up HC and PLE implementations for Qwen3.8-Flash-Next (#61454) Signed-off-by: huanghaoyan.hhy Co-authored-by: Thien Tran --- tests/models/qwen4_exp/test_ple.py | 130 +++++++++++++----- vllm/models/qwen4_exp/amd/hyperconnection.py | 30 ++-- vllm/models/qwen4_exp/amd/ple_layer.py | 31 ++++- .../qwen4_exp/common/ngram_embedding.py | 73 +++++----- vllm/models/qwen4_exp/common/ops/ple.py | 37 +---- vllm/models/qwen4_exp/common/ple.py | 114 ++++++++------- vllm/models/qwen4_exp/common/qsa_cache.py | 8 +- .../qwen4_exp/nvidia/hyperconnection.py | 26 ++-- .../qwen4_exp/nvidia/ngram_embedding.py | 7 +- vllm/models/qwen4_exp/nvidia/ops/hc.py | 114 +-------------- .../qwen4_exp/nvidia/ops/qsa_indexer.py | 3 +- 11 files changed, 262 insertions(+), 311 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 119ec3b1ac8f..61a020da2576 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging from dataclasses import dataclass -from functools import partial from itertools import accumulate from types import SimpleNamespace import pytest import torch +from compressed_tensors.quantization import QuantizationArgs from torch import nn from torch.nn import functional as F from transformers import Qwen4ExpTextConfig @@ -19,6 +20,10 @@ from vllm.config import SpeculativeConfig, VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, +) from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.modelopt import ( @@ -33,8 +38,7 @@ ) from vllm.models.qwen4_exp.common.ple import ( PLEShardOverlap, - compute_ple_shard_overlap, - copy_ple_embedding_shard_, + PLEVocabParallelEmbedding, ) from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( Qwen4ExpNGramEmbedding, @@ -210,12 +214,23 @@ def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: return module +def _make_test_ple_embedding(tp_start: int, tp_end: int) -> PLEVocabParallelEmbedding: + """Create an embedding with only the state needed for shard loading.""" + embedding = PLEVocabParallelEmbedding.__new__(PLEVocabParallelEmbedding) + nn.Module.__init__(embedding) + embedding.shard_indices = SimpleNamespace( + org_vocab_start_index=tp_start, + org_vocab_end_index=tp_end, + ) + return embedding + + def _set_test_embedding_weight_loader(embedding) -> None: - embedding.weight.weight_loader = partial( - copy_ple_embedding_shard_, + loader = _make_test_ple_embedding( tp_start=embedding.shard_indices.org_vocab_start_index, tp_end=embedding.shard_indices.org_vocab_end_index, ) + embedding.weight.weight_loader = loader.weight_loader def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: @@ -244,19 +259,16 @@ def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: def test_ple_shard_overlap_and_copy() -> None: - overlap = compute_ple_shard_overlap( - checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8 - ) + embedding = _make_test_ple_embedding(tp_start=4, tp_end=8) + overlap = embedding._compute_shard_overlap(checkpoint_start=2, checkpoint_rows=5) assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3) destination = torch.full((4, 2), -1.0) loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( destination, loaded, checkpoint_start=2, - tp_start=4, - tp_end=8, ) assert copied == 3 @@ -265,13 +277,12 @@ def test_ple_shard_overlap_and_copy() -> None: def test_ple_shard_copy_is_a_noop_without_overlap() -> None: + embedding = _make_test_ple_embedding(tp_start=4, tp_end=8) destination = torch.ones(4, 2) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( destination, torch.zeros(2, 2), checkpoint_start=10, - tp_start=4, - tp_end=8, ) assert copied == 0 @@ -704,17 +715,44 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None: ) -def test_ple_embedding_rejects_unsupported_quantization_configs() -> None: +@pytest.mark.parametrize( + "quant_config", + [ + pytest.param(ModelOptNvFp4Config(exclude_modules=[]), id="modelopt-nvfp4"), + pytest.param( + OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor")), + id="online-fp8", + ), + pytest.param( + CompressedTensorsConfig( + target_scheme_map={ + "re:.*ngram_embedding": { + "weights": QuantizationArgs(num_bits=8, type="int"), + }, + }, + ignore=[], + quant_format="int-quantized", + ), + id="compressed-tensors", + ), + ], +) +def test_ple_embedding_falls_back_for_unsupported_quantization_configs( + quant_config: QuantizationConfig, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Warn when global quantization falls back to an unquantized PLE table.""" prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" - nvfp4_config = ModelOptNvFp4Config(exclude_modules=[]) - with pytest.raises(NotImplementedError, match="ModelOptNvFp4Config"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(nvfp4_config, prefix) - - online_fp8_config = OnlineQuantizationConfig( - QuantizationConfigArgs(linear="fp8_per_tensor") + test_logger = logging.getLogger("test.ple_fallback") + monkeypatch.setattr(ngram_embedding_module, "logger", test_logger) + caplog.set_level(logging.WARNING, logger=test_logger.name) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, ) - with pytest.raises(NotImplementedError, match="OnlineQuantizationConfig"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(online_fp8_config, prefix) + assert caplog.records[0].levelno == logging.WARNING + assert "Falling back to unquantized PLE" in caplog.text def test_ple_embedding_respects_modelopt_exclusion() -> None: @@ -741,8 +779,10 @@ def test_ple_embedding_respects_inc_layer_config() -> None: ) quant_config.extra_config = None - with pytest.raises(NotImplementedError, match="INCConfig"): - Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) def test_ple_embedding_is_unquantized_under_quark() -> None: @@ -926,12 +966,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva( loaded_weight = ( torch.arange(12, dtype=torch.float32).reshape(4, 3).to(storage_dtype) ) - copied = copy_ple_embedding_shard_( + copied = embedding._copy_embedding_shard_( embedding.weight, loaded_weight, checkpoint_start=0, - tp_start=0, - tp_end=4, ) input_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") output = embedding._lookup(input_ids) @@ -988,12 +1026,10 @@ def test_pinned_embedding_start_prefetch_allocates_lazily( loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3) loaded_weight = loaded_weight.to(torch.bfloat16) - copy_ple_embedding_shard_( + embedding._copy_embedding_shard_( embedding.weight, loaded_weight, checkpoint_start=0, - tp_start=0, - tp_end=4, ) ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") @@ -1012,8 +1048,18 @@ def test_pinned_embedding_start_prefetch_allocates_lazily( torch.testing.assert_close(output, expected, rtol=0, atol=0) -def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: +@pytest.mark.parametrize( + "exclude_pattern", + [ + "model.language_model.layers.1.ple.ple_embedding.ngram_embedding", + "*.ngram_embedding", + ], +) +def test_ple_fp8_embedding_supports_mixed_precision_config( + exclude_pattern: str, +) -> None: prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + nvfp4_prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding" quant_config = ModelOptMixedPrecisionConfig.from_config( { "quantization": { @@ -1022,6 +1068,7 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: "group_size": 16, "quantized_layers": { prefix: {"quant_algo": "FP8"}, + nvfp4_prefix: {"quant_algo": "NVFP4"}, "model.language_model.layers.2.moe.gate_proj": { "quant_algo": "MXFP8" }, @@ -1034,13 +1081,30 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), Qwen4ExpPLEFp8EmbeddingMethod, ) - # A quant_algo the PLE table cannot load must not fall back to unquantized. + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config( + quant_config, + "model.language_model.layers.3.ple.ple_embedding.ngram_embedding", + ), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, nvfp4_prefix), + Qwen4ExpPLENvFp4EmbeddingMethod, + ) + # An explicitly unsupported PLE format must not fall back to unquantized. with pytest.raises(ValueError, match="Unsupported quant_algo 'MXFP8'"): Qwen4ExpPLEEmbeddingMethod.from_quant_config( quant_config, "model.language_model.layers.2.moe.gate_proj", ) + quant_config.exclude_modules = [exclude_pattern] + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), + Qwen4ExpPLEUnquantizedEmbeddingMethod, + ) + def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None: conv_state_len = 6 diff --git a/vllm/models/qwen4_exp/amd/hyperconnection.py b/vllm/models/qwen4_exp/amd/hyperconnection.py index 087f03fdcf5f..b6376ac238ff 100644 --- a/vllm/models/qwen4_exp/amd/hyperconnection.py +++ b/vllm/models/qwen4_exp/amd/hyperconnection.py @@ -70,7 +70,7 @@ def __init__( ) -> None: super().__init__() self.config = config - self.lora_rank = config.hc_lowrank + self.hc_lowrank = config.hc_lowrank self.hc_count = config.hc_count self.hidden_size = config.hidden_size self.use_combine = use_combine @@ -91,11 +91,11 @@ def __init__( # -- vLLM Linear weights -------------------------------------------- # The merged skinny-GEMM shape is physically padded to 16 rows for # alignment and efficient backend dispatch. - self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0 if use_combine: self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( self.hyper_hidden_size, - [self.lora_rank, self.hc_count] + [self.hc_lowrank, self.hc_count] + ([self.pad_size] if self.pad_size else []), bias=False, params_dtype=config.params_dtype, @@ -107,7 +107,7 @@ def __init__( else: self.input_mix_weight_down = ReplicatedLinear( self.hyper_hidden_size, - self.lora_rank, + self.hc_lowrank, bias=False, params_dtype=config.params_dtype, quant_config=None, @@ -115,7 +115,7 @@ def __init__( return_bias=False, ) self.input_mix_weight_up = ReplicatedLinear( - self.lora_rank, + self.hc_lowrank, self.hyper_hidden_size, bias=False, params_dtype=config.params_dtype, @@ -136,15 +136,15 @@ def mix( if self.use_combine: # produce injection logits for combine - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] down_and_injection = self.input_mix_weight_down_block_inject(xn) - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) else: - lora = self.input_mix_weight_down(xn) + lowrank = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank = hc_silu(lowrank, self.hc_count) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection @@ -172,15 +172,15 @@ def combine_and_mix( if self.use_combine: # produce injection logits for combine - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] down_and_injection = self.input_mix_weight_down_block_inject(xn) - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) else: - lora = self.input_mix_weight_down(xn) + lowrank = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank = hc_silu(lowrank, self.hc_count) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index c1cfc0f24627..478b0ae97215 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -29,6 +29,7 @@ sp_reduce_scatter, sp_shard, ) +from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( @@ -242,7 +243,7 @@ def __init__( persistent=True, ) divisor = int(config.make_ngram_vocab_size_divisible_by) - padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor + padded_vocab_size = round_up(total_vocab_size, divisor) if params_dtype is None: params_dtype = torch.get_default_dtype() embedding_prefix = f"{prefix}.ngram_embedding" @@ -280,9 +281,11 @@ def __init__( reduce_results=not self.use_reduce_scatter, ) logger.info( - "Initialized AMD PLE embedding %s: quantization_method=%s, " + "Initialized AMD PLE embedding %s: " + "quantization_config=%s, quantization_method=%s, " "weight_dtype=%s, weight_device=%s, pinned=%s", embedding_prefix, + type(quant_config).__name__, type(embedding_quant_method).__name__, self.ngram_embedding.weight.dtype, self.ngram_embedding.weight.device, @@ -742,7 +745,7 @@ def forward( ) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) - gated_output, conv_input = ple_gate( + gated_output, conv_input = torch.ops.vllm.qwen4_exp_ple_gate( key, value, hidden_states, @@ -848,6 +851,28 @@ def qwen4_exp_ple_short_conv( layer._short_conv(inputs, residual, outer_residual) +def _ple_gate_fake( + key: torch.Tensor, + value: torch.Tensor, + hidden: torch.Tensor, + norm_key_w: torch.Tensor, + norm_query_w: torch.Tensor, + norm_conv_w: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return the gate output shapes for AMD graph tracing.""" + return torch.empty_like(hidden), torch.empty_like(hidden) + + +# Keep the AMD gate opaque to Dynamo while reusing the shared Triton kernel. +direct_register_custom_op( + op_name="qwen4_exp_ple_gate", + op_func=ple_gate, + mutates_args=[], + fake_impl=_ple_gate_fake, +) + + direct_register_custom_op( op_name="qwen4_exp_amd_ple_ngram_ids", op_func=qwen4_exp_amd_ple_ngram_ids, diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 0348b242245a..d5bd9f689270 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -194,22 +194,21 @@ def from_quant_config( ) if quant_config is None: return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance(quant_config, ModelOptMixedPrecisionConfig): - quant_algo = quant_config._resolve_quant_algo(prefix) - if quant_algo == "FP8": - return Qwen4ExpPLEFp8EmbeddingMethod() - if quant_algo == "NVFP4": - return Qwen4ExpPLENvFp4EmbeddingMethod() - if quant_algo is not None: - raise ValueError( - f"Unsupported quant_algo {quant_algo!r} for the PLE embedding " - f"{prefix!r}; expected FP8 or NVFP4" - ) - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance( - quant_config, ModelOptQuantConfigBase - ) and quant_config.is_layer_excluded(prefix): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance(quant_config, ModelOptQuantConfigBase): + if quant_config.is_layer_excluded(prefix): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance(quant_config, ModelOptMixedPrecisionConfig): + quant_algo = quant_config._resolve_quant_algo(prefix) + if quant_algo == "FP8": + return Qwen4ExpPLEFp8EmbeddingMethod() + if quant_algo == "NVFP4": + return Qwen4ExpPLENvFp4EmbeddingMethod() + if quant_algo is not None: + raise ValueError( + f"Unsupported quant_algo {quant_algo!r} for the PLE embedding " + f"{prefix!r}; expected FP8 or NVFP4" + ) + return Qwen4ExpPLEUnquantizedEmbeddingMethod() if ( isinstance(quant_config, CompressedTensorsConfig) and quant_config.get_scheme_dict(None, layer_name=prefix) is None @@ -223,29 +222,32 @@ def from_quant_config( # Quark quantizes only Linear and MoE layers; PLE tables stay BF16. if isinstance(quant_config, QuarkConfig): return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not isinstance(quant_config, Fp8Config): + if isinstance(quant_config, Fp8Config): + ignored_layers = quant_config.ignored_layers + if is_layer_skipped( + prefix, + ignored_layers, + quant_config.packed_modules_mapping, + match_mode=quant_config.ignored_layers_match_mode, + ): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + # PLE checkpoint shards form one runtime embedding parameter. + shard_prefix = f"{prefix}.shard_" + if any(name.startswith(shard_prefix) for name in ignored_layers): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if quant_config.is_checkpoint_fp8_serialized: + return Qwen4ExpPLEFp8EmbeddingMethod() raise NotImplementedError( - "Qwen4Exp PLE embedding does not support quantization config " - f"{type(quant_config).__name__}" + "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" ) - ignored_layers = quant_config.ignored_layers - if is_layer_skipped( + logger.warning( + "Falling back to unquantized PLE embedding for %s " + "with quantization config %s.", prefix, - ignored_layers, - quant_config.packed_modules_mapping, - match_mode=quant_config.ignored_layers_match_mode, - ): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - # PLE checkpoint shards form one runtime embedding parameter. - shard_prefix = f"{prefix}.shard_" - if any(name.startswith(shard_prefix) for name in ignored_layers): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not quant_config.is_checkpoint_fp8_serialized: - raise NotImplementedError( - "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" - ) - return Qwen4ExpPLEFp8EmbeddingMethod() + type(quant_config).__name__, + ) + return Qwen4ExpPLEUnquantizedEmbeddingMethod() def apply( self, @@ -773,6 +775,7 @@ def start_prefetch(self, ngram_ids: torch.Tensor) -> None: ) active_output = buffer[: gathered_ids.shape[0]] prefetch_stream.wait_stream(torch.cuda.current_stream()) + # Prevent ID storage reuse until the prefetch stream finishes reading it. gathered_ids.record_stream(prefetch_stream) with torch.cuda.stream(prefetch_stream): self._lookup(gathered_ids, output=active_output) diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py index 41ce082c9e70..994b799ecdcd 100644 --- a/vllm/models/qwen4_exp/common/ops/ple.py +++ b/vllm/models/qwen4_exp/common/ops/ple.py @@ -16,7 +16,6 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -240,7 +239,7 @@ def _ple_gate_kernel( tl.store(normed_ptr + t * HC * H + offs, normed, mask=mask) -def _ple_gate( +def ple_gate( key: torch.Tensor, value: torch.Tensor, hidden: torch.Tensor, @@ -281,40 +280,6 @@ def _ple_gate( return gated, normed -def _ple_gate_fake( - key: torch.Tensor, - value: torch.Tensor, - hidden: torch.Tensor, - norm_key_w: torch.Tensor, - norm_query_w: torch.Tensor, - norm_conv_w: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.empty_like(hidden), torch.empty_like(hidden) - - -direct_register_custom_op( - op_name="qwen4_exp_ple_gate", - op_func=_ple_gate, - mutates_args=[], - fake_impl=_ple_gate_fake, -) - - -def ple_gate( - key: torch.Tensor, - value: torch.Tensor, - hidden: torch.Tensor, - norm_key_w: torch.Tensor, - norm_query_w: torch.Tensor, - norm_conv_w: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.qwen4_exp_ple_gate( - key, value, hidden, norm_key_w, norm_query_w, norm_conv_w, eps - ) - - # Dilated short convolution @triton.jit(do_not_specialize=["num_reqs", "bs_iters", "has_token_map"]) def _ple_conv_kernel( diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py index 72fc494b6037..e15e803b6946 100644 --- a/vllm/models/qwen4_exp/common/ple.py +++ b/vllm/models/qwen4_exp/common/ple.py @@ -20,65 +20,63 @@ class PLEShardOverlap: row_count: int -def compute_ple_shard_overlap( - *, - checkpoint_start: int, - checkpoint_rows: int, - tp_start: int, - tp_end: int, -) -> PLEShardOverlap | None: - """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" - if checkpoint_start < 0 or checkpoint_rows < 0: - raise ValueError("checkpoint shard bounds must be non-negative") - if tp_start < 0 or tp_end < tp_start: - raise ValueError("invalid TP vocabulary range") - checkpoint_end = checkpoint_start + checkpoint_rows - overlap_start = max(checkpoint_start, tp_start) - overlap_end = min(checkpoint_end, tp_end) - if overlap_start >= overlap_end: - return None - return PLEShardOverlap( - source_start=overlap_start - checkpoint_start, - destination_start=overlap_start - tp_start, - row_count=overlap_end - overlap_start, - ) - +class PLEVocabParallelEmbedding(VocabParallelEmbedding): + """Vocab-parallel embedding that accepts checkpoint row shards.""" -def copy_ple_embedding_shard_( - destination: torch.Tensor, - loaded_weight: torch.Tensor, - *, - checkpoint_start: int, - tp_start: int, - tp_end: int, -) -> int: - """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" - if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: - raise ValueError("destination and loaded weight must have matching ranks") - if destination.shape[1:] != loaded_weight.shape[1:]: - raise ValueError( - "embedding shard dimensions do not match: " - f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" + def _compute_shard_overlap( + self, + *, + checkpoint_start: int, + checkpoint_rows: int, + ) -> PLEShardOverlap | None: + """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" + tp_start = self.shard_indices.org_vocab_start_index + tp_end = self.shard_indices.org_vocab_end_index + if checkpoint_start < 0 or checkpoint_rows < 0: + raise ValueError("checkpoint shard bounds must be non-negative") + if tp_start < 0 or tp_end < tp_start: + raise ValueError("invalid TP vocabulary range") + checkpoint_end = checkpoint_start + checkpoint_rows + overlap_start = max(checkpoint_start, tp_start) + overlap_end = min(checkpoint_end, tp_end) + if overlap_start >= overlap_end: + return None + return PLEShardOverlap( + source_start=overlap_start - checkpoint_start, + destination_start=overlap_start - tp_start, + row_count=overlap_end - overlap_start, ) - if destination.shape[0] < tp_end - tp_start: - raise ValueError("destination does not cover the requested TP range") - overlap = compute_ple_shard_overlap( - checkpoint_start=checkpoint_start, - checkpoint_rows=loaded_weight.shape[0], - tp_start=tp_start, - tp_end=tp_end, - ) - if overlap is None: - return 0 - source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) - target = destination.narrow(0, overlap.destination_start, overlap.row_count) - with torch.no_grad(): - target.copy_(source.to(device=target.device, dtype=target.dtype)) - return overlap.row_count - -class PLEVocabParallelEmbedding(VocabParallelEmbedding): - """Vocab-parallel embedding that accepts checkpoint row shards.""" + def _copy_embedding_shard_( + self, + destination: torch.Tensor, + loaded_weight: torch.Tensor, + *, + checkpoint_start: int, + ) -> int: + """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" + tp_start = self.shard_indices.org_vocab_start_index + tp_end = self.shard_indices.org_vocab_end_index + if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: + raise ValueError("destination and loaded weight must have matching ranks") + if destination.shape[1:] != loaded_weight.shape[1:]: + raise ValueError( + "embedding shard dimensions do not match: " + f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" + ) + if destination.shape[0] < tp_end - tp_start: + raise ValueError("destination does not cover the requested TP range") + overlap = self._compute_shard_overlap( + checkpoint_start=checkpoint_start, + checkpoint_rows=loaded_weight.shape[0], + ) + if overlap is None: + return 0 + source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) + target = destination.narrow(0, overlap.destination_start, overlap.row_count) + with torch.no_grad(): + target.copy_(source.to(device=target.device, dtype=target.dtype)) + return overlap.row_count def weight_loader( self, @@ -89,10 +87,8 @@ def weight_loader( if checkpoint_start is None: super().weight_loader(param, loaded_weight) return - copy_ple_embedding_shard_( + self._copy_embedding_shard_( param, loaded_weight, checkpoint_start=checkpoint_start, - tp_start=self.shard_indices.org_vocab_start_index, - tp_end=self.shard_indices.org_vocab_end_index, ) diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index f960dea07194..5eba092eec25 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_up from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -818,9 +818,7 @@ def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: key_head_size = int(kwargs.pop("head_size")) self.key_head_size = key_head_size self.cache_rope_positions = bool(cache_rope_positions) - self.rope_position_offset = ( - (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 - ) * self._BF16_PER_INT64 + self.rope_position_offset = round_up(key_head_size, self._BF16_PER_INT64) storage_head_size = key_head_size if self.cache_rope_positions: storage_head_size = self.rope_position_offset + ( @@ -848,7 +846,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # the scheduler block size). Anything narrower lets a rejected draft row # overwrite a committed key the next step needs to close the group. span = self.compress_ratio + vllm_config.num_speculative_tokens - capacity = self.compress_ratio * cdiv(span, self.compress_ratio) + capacity = round_up(span, self.compress_ratio) assert self.cache_config.block_size % capacity == 0, ( f"QSA ring capacity {capacity} must divide the attention block " f"size {self.cache_config.block_size}" diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index dcf8daa956ae..3844f6de32bc 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -74,7 +74,7 @@ def __init__( ) -> None: super().__init__() self.config = config - self.lora_rank = config.hc_lowrank + self.hc_lowrank = config.hc_lowrank self.hc_count = config.hc_count self.hidden_size = config.hidden_size self.use_combine = use_combine @@ -95,11 +95,11 @@ def __init__( # -- vLLM Linear weights -------------------------------------------- # The merged skinny-GEMM shape is physically padded to 16 rows to ensure # good alignment and performant implementation chosen by CuBLAS heuristics. - self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + self.pad_size = (-(self.hc_lowrank + self.hc_count)) % 16 if use_combine else 0 if use_combine: self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( self.hyper_hidden_size, - [self.lora_rank, self.hc_count] + [self.hc_lowrank, self.hc_count] + ([self.pad_size] if self.pad_size else []), bias=False, params_dtype=config.params_dtype, @@ -117,7 +117,7 @@ def __init__( else: self.input_mix_weight_down = ReplicatedLinear( self.hyper_hidden_size, - self.lora_rank, + self.hc_lowrank, bias=False, params_dtype=config.params_dtype, quant_config=None, @@ -125,7 +125,7 @@ def __init__( return_bias=False, ) self.input_mix_weight_up = ReplicatedLinear( - self.lora_rank, + self.hc_lowrank, self.hyper_hidden_size, bias=False, params_dtype=config.params_dtype, @@ -150,13 +150,13 @@ def _down_and_inject( return hc_down_silu( xn, self.input_mix_weight_down_block_inject.weight, - self.lora_rank, + self.hc_lowrank, self.hc_count, ) down_and_injection = self.input_mix_weight_down_block_inject(xn) - split_sizes = [self.lora_rank, self.hc_count, self.pad_size] - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) - return hc_silu(lora, self.hc_count), injection + split_sizes = [self.hc_lowrank, self.hc_count, self.pad_size] + lowrank, injection, _ = down_and_injection.split(split_sizes, dim=-1) + return hc_silu(lowrank, self.hc_count), injection def mix( self, hidden_states: torch.Tensor @@ -168,8 +168,8 @@ def mix( self.hc_count, ) - lora, injection = self._down_and_inject(xn) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank, injection = self._down_and_inject(xn) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection @@ -196,8 +196,8 @@ def combine_and_mix( self.hc_count, ) - lora, injection = self._down_and_inject(xn) - gate = self.input_mix_weight_up(lora) # [M, D] + lowrank, injection = self._down_and_inject(xn) + gate = self.input_mix_weight_up(lowrank) # [M, D] block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 1a3655104690..7b0d9ed8c2f7 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard +from vllm.utils.math_utils import round_up from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, @@ -203,7 +204,7 @@ def __init__( persistent=True, ) divisor = int(config.make_ngram_vocab_size_divisible_by) - padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor + padded_vocab_size = round_up(total_vocab_size, divisor) embedding_prefix = f"{prefix}.ngram_embedding" embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( quant_config, @@ -243,9 +244,11 @@ def __init__( ) weight = self.ngram_embedding.weight logger.info( - "Initialized PLE embedding %s: quantization_method=%s, " + "Initialized PLE embedding %s: " + "quantization_config=%s, quantization_method=%s, " "weight_dtype=%s, weight_device=%s, pinned=%s", embedding_prefix, + type(quant_config).__name__, type(embedding_quant_method).__name__, weight.dtype, weight.device, diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py index 36bc8cefbc4e..3dee9db436b7 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/hc.py +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -6,7 +6,6 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op @triton.jit @@ -52,7 +51,7 @@ def _grouped_gemma_rmsnorm_kernel( tl.store(y_ptr + row * stride_y + offsets, y, mask) -def _grouped_gemma_rmsnorm( +def grouped_gemma_rmsnorm( x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int ) -> torch.Tensor: N, DIM = x.shape @@ -105,7 +104,7 @@ def _hc_silu_kernel( tl.store(y_ptr + row * stride_y + offs, y, mask) -def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: num_tokens, DIM = x.shape assert x.stride(1) == 1 @@ -160,7 +159,7 @@ def _hc_gate_mix_kernel( tl.store(y_ptr + row * stride_y + offs_inner, acc, mask) -def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: N, DIM = gate.shape assert x.shape == gate.shape assert DIM % hc_count == 0 @@ -232,7 +231,7 @@ def _hc_combine_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask) -def _hc_combine( +def hc_combine( residual: torch.Tensor, block_output: torch.Tensor, injection_logits: torch.Tensor | None, @@ -345,7 +344,7 @@ def _hc_combine_norm_kernel( tl.store(y_ptr + row * stride_y + offs, y, mask_inner) -def _hc_combine_norm( +def hc_combine_norm( residual: torch.Tensor, block_output: torch.Tensor, injection_logits: torch.Tensor | None, @@ -392,109 +391,6 @@ def _hc_combine_norm( return out, y -def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor: - return x.new_empty(x.shape) - - -def _hc_gate_mix_fake( - x: torch.Tensor, gate: torch.Tensor, hc_count: int -) -> torch.Tensor: - del gate - return x.new_empty((x.shape[0], x.shape[1] // hc_count)) - - -def _hc_combine_fake( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - hc_count: int, -) -> torch.Tensor: - del block_output, injection_logits, hc_count - return residual.new_empty(residual.shape) - - -def _hc_combine_norm_fake( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, -) -> tuple[torch.Tensor, torch.Tensor]: - del block_output, injection_logits, norm_weight, eps, hc_count - return residual.new_empty(residual.shape), residual.new_empty(residual.shape) - - -direct_register_custom_op( - op_name="qwen4_exp_grouped_gemma_rmsnorm", - op_func=_grouped_gemma_rmsnorm, - fake_impl=_same_shape_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_silu", - op_func=_hc_silu, - fake_impl=_same_shape_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_gate_mix", - op_func=_hc_gate_mix, - fake_impl=_hc_gate_mix_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_combine", - op_func=_hc_combine, - fake_impl=_hc_combine_fake, -) -direct_register_custom_op( - op_name="qwen4_exp_hc_combine_norm", - op_func=_hc_combine_norm, - fake_impl=_hc_combine_norm_fake, -) - - -def grouped_gemma_rmsnorm( - x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int -) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups) - - -def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count) - - -def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count) - - -def hc_combine( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - hc_count: int, -) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_combine( - residual, block_output, injection_logits, hc_count - ) - - -def hc_combine_norm( - residual: torch.Tensor, - block_output: torch.Tensor, - injection_logits: torch.Tensor | None, - norm_weight: torch.Tensor, - eps: float, - hc_count: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.qwen4_exp_hc_combine_norm( - residual, - block_output, - injection_logits, - norm_weight, - eps, - hc_count, - ) - - __all__ = [ "grouped_gemma_rmsnorm", "hc_combine", diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 9b28ed6dcaa3..4aa85f79f3b0 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -10,6 +10,7 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up _TOPK_WORKSPACE_BYTES = 1024 * 1024 _DECODE_BLOCK_N = 64 @@ -611,7 +612,7 @@ def qsa_select_paged_prefill( # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed # columns. Round up to 64 to keep the logits row stride # cooperative_topk-compatible. - logits_width = triton.cdiv(triton.cdiv(max_seq_len, compress_ratio), 64) * 64 + logits_width = round_up(triton.cdiv(max_seq_len, compress_ratio), 64) logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1]) # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB From 6b24947628b6e0f3363540f3fbdc5c14f7d99b94 Mon Sep 17 00:00:00 2001 From: Muhammad Ibrahim Date: Sun, 18 Oct 2026 14:12:34 +0200 Subject: [PATCH 821/952] [CI/Build] Build DeepEP for SM103 (B300/GB300) in the Docker image (#61829) Signed-off-by: Muhammad Ibrahim Co-authored-by: Muhammad Ibrahim Co-authored-by: Shengqi Chen --- docker/Dockerfile | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index e29142dd963a..cc87f9a21fe7 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -502,7 +502,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \ > /tmp/nccl-override.txt && \ export UV_OVERRIDE=/tmp/nccl-override.txt; \ fi && \ - export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \ + DEEPEP_ARCH_LIST='9.0a 10.0a' && \ + if [ "$(printf '%s\n' 12.9 "$CUDA_VERSION" | sort -V | head -n1)" = 12.9 ]; then \ + DEEPEP_ARCH_LIST="$DEEPEP_ARCH_LIST 10.3a"; \ + fi && \ + export TORCH_CUDA_ARCH_LIST="$DEEPEP_ARCH_LIST" && \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ --mode wheel \ From 072c0886d53d933a7e1d3da7e8a3a91670cf181d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20Lustosa?= <30708955+andre-motta@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:13:17 -0400 Subject: [PATCH 822/952] [Build] Keep plain CUDA arch when its suffixed variant is also requested (#56549) Signed-off-by: Andre Lustosa Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Shengqi Chen --- cmake/utils.cmake | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 6ded936213c6..83eed340eb7d 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -425,6 +425,7 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR if("${_base}" IN_LIST _TGT_CUDA_ARCHS) list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") + list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a" "${_base}f") list(APPEND _CUDA_ARCHS "${_arch}") endif() endif() @@ -464,8 +465,10 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR string(REGEX REPLACE "[af]$" "" _base "${_arch}") if ("${_base}" IN_LIST _SRC_CUDA_ARCHS) list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_arch}") - list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}") list(APPEND _CUDA_ARCHS "${_arch}") + if (NOT "${_base}" IN_LIST _TGT_CUDA_ARCHS) + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}") + endif() endif() endif() endforeach() From 371745ebd6c12f314402c6c319de2e7ef82c2981 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 18 Oct 2026 21:11:45 +0800 Subject: [PATCH 823/952] [Doc][RL] Document routed experts return for MoE routing replay (R3) (#62302) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- docs/training/rlhf.md | 2 ++ docs/training/routed_experts.md | 54 +++++++++++++++++++++++++++++++++ mkdocs.yaml | 1 + 3 files changed, 57 insertions(+) create mode 100644 docs/training/routed_experts.md diff --git a/docs/training/rlhf.md b/docs/training/rlhf.md index 2aae267e1d85..ccce0f37dd2d 100644 --- a/docs/training/rlhf.md +++ b/docs/training/rlhf.md @@ -18,6 +18,8 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines. To verify that an update landed on every rank, see the [Weight Checker](weight_checker.md). +For MoE policies, see [Routed Experts Return](routed_experts.md) (`--enable-return-routed-experts`) to get the per-token expert routing of each rollout for routing replay (R3) in the trainer. + For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight. See the following notebooks showing how to use vLLM for GRPO: diff --git a/docs/training/routed_experts.md b/docs/training/routed_experts.md new file mode 100644 index 000000000000..f531a0ebbd8c --- /dev/null +++ b/docs/training/routed_experts.md @@ -0,0 +1,54 @@ +# Routed Experts Return (MoE Routing Replay / R3) + +`--enable-return-routed-experts` makes vLLM return the routed experts (the expert IDs each MoE layer's router selected) for every token of a request. RL trainers use it for MoE routing replay (R3): the training forward pass reuses the rollout's routing. + +```bash +vllm serve Qwen/Qwen3-30B-A3B --enable-return-routed-experts +# offline: LLM("Qwen/Qwen3-30B-A3B", enable_return_routed_experts=True) +``` + +The flag sets `AuxOutputConfig.enable_return_routed_experts` (full form: `--aux-output-config '{"enable_return_routed_experts": true, "max_bytes": N}'`). There is no per-request on/off switch. + +## Requirements + +Startup fails with a `ValueError` unless: + +- **Model Runner V2** is used (`VLLM_USE_V2_MODEL_RUNNER=1` if vLLM falls back to V1); +- the runner is `generate` and the model is **MoE**; +- **prefix caching is enabled** (the default). `--no-enable-prefix-caching` fails; the old v0.21 advice to disable it is wrong. + +Rejected at startup: adaptive speculative verification, PP > 1, DCP/PCP > 1, and the `NixlConnector`, `NixlPullConnector`, `NixlPushConnector`, `MoRIIOConnector`, `MooncakeConnector` KV connectors. Regular speculative decoding works. + +### How it works: aux output connector + +Routing is stored by the aux output connector (`vllm/distributed/aux_output_connector/`, configured via `--aux-output-config`), keyed by KV block hash and shared with prefix caching, which is why prefix caching is required. `max_bytes` is its host-memory LRU capacity (default: derived from KV cache capacity). If a prefix-cache hit needs evicted rows, the worker raises `BlockObjectStoreError`; increase `max_bytes`. + +## Output format + +| API | Field | +| --- | --- | +| Offline (`LLM`, `AsyncLLM`) | `CompletionOutput.routed_experts`: `np.ndarray` or `None` | +| `/v1/completions`, `/v1/chat/completions` | `choices[i].routed_experts`: base64 `.npy`, **non-streaming only** | +| `/inference/v1/generate` | `routed_experts` per choice (base64 `.npy`); with `stream: true`, on the final chunk | + +- **Shape:** `(num_tokens - 1 - routed_experts_prompt_start, num_moe_layers, num_experts_per_tok)`, `num_tokens` = prompt + output. The last sampled token has no row. +- **Rows:** row `i` is position `routed_experts_prompt_start + i` of `prompt_token_ids + output_token_ids`. +- **dtype:** `uint8` for at most 256 experts, else `uint16`; logical expert IDs. +- Filled only when the request finishes. `None` if the flag is off or the request aborted before any forward pass. + +```python +import base64, io +import numpy as np + +experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"]))) +``` + +## `routed_experts_prompt_start` + +Per-request, default `0`: drops the first N prompt-token rows (e.g. routing already received in earlier turns). Must be `0 <= value <= prompt_len`, else HTTP 400. Pass it top-level on `/v1/*completions`, inside `sampling_params` on `/inference/v1/generate`, or as `SamplingParams(routed_experts_prompt_start=...)`. + +## Limitations + +- The Rust HTTP frontend never returns `routed_experts` and always sends `routed_experts_prompt_start = 0`. Use the Python API server. + +Example: [examples/rl/routed_experts_e2e.py](../../examples/rl/routed_experts_e2e.py). diff --git a/mkdocs.yaml b/mkdocs.yaml index bed8318998c3..d591df7b1435 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -106,6 +106,7 @@ plugins: examples/features/pause_resume.md: examples/rl/pause_resume.md serving/online_serving/decisions.md: models/structured_decisions_models.md serving/online_serving/structured_decisions.md: models/structured_decisions_models.md + training/routed_experts_replay.md: training/routed_experts.md markdown_extensions: - attr_list From e80b6c642f21fd5b24439825592f5b9ced36add8 Mon Sep 17 00:00:00 2001 From: Taoyu Zhu Date: Sun, 18 Oct 2026 22:39:38 +0800 Subject: [PATCH 824/952] [ROCm][DSv4] Overlap the unfused shared expert with AITER MegaMoE (#61819) Signed-off-by: Taoyu Zhu Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../models/test_deepseek_v4_aiter_mega_moe.py | 10 ++++- vllm/models/deepseek_v4/amd/model.py | 45 ++++++++++++++++--- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index 1a5347669401..c3ffded51296 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/tests/models/test_deepseek_v4_aiter_mega_moe.py @@ -15,6 +15,7 @@ ) from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE from vllm.platforms import current_platform +from vllm.utils.torch_utils import aux_stream HIDDEN = 64 INTER = 64 @@ -156,7 +157,10 @@ def test_shared_expert_route_appended(monkeypatch): @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm") @pytest.mark.parametrize("hash_routing", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 37]) -def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_tokens): +@pytest.mark.parametrize("shared_stream", [False, True]) +def test_mega_moe_routing_matches_reference( + monkeypatch, hash_routing, num_tokens, shared_stream +): num_experts, top_k, vocab = 16, 4, 64 scaling = 2.5 device = "cuda" @@ -195,6 +199,10 @@ def experts(x, topk_weights, topk_ids): moe.gate = Gate() moe.experts = experts moe.shared_experts = lambda x: 3 * x + moe.shared_experts_stream = aux_stream() if shared_stream else None + moe.shared_experts_events = ( + (torch.cuda.Event(), torch.cuda.Event()) if shared_stream else None + ) moe.scoring_func = "sqrtsoftplus" moe.n_activated_experts = top_k moe.renormalize = True diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index d1a8ca1d9d02..569cc0c9f5c1 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -82,6 +82,8 @@ from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950, on_gfx1250 from vllm.sequence import IntermediateTensors +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID @@ -713,6 +715,18 @@ def __init__( ), prefix=f"{prefix}.experts", ) + self.shared_experts_stream = ( + aux_stream() + if self.shared_experts is not None + and not self.experts.fuse_shared_expert + and not envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM + else None + ) + self.shared_experts_events = ( + (torch.cuda.Event(), torch.cuda.Event()) + if self.shared_experts_stream is not None + else None + ) return fuse_shared_into_routed = ( @@ -777,6 +791,30 @@ def forward( def _forward_mega_moe( self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None + ) -> torch.Tensor: + shared_experts = self.shared_experts + if shared_experts is None: + return self._forward_mega_moe_routed(hidden_states, input_ids) + events = self.shared_experts_events + if ( + events is None + or hidden_states.shape[0] > envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + ): + return self._forward_mega_moe_routed( + hidden_states, input_ids + ) + shared_experts(hidden_states) + # The unfused shared expert overlaps router + MegaMoE in small steps. + routed, shared = maybe_execute_in_parallel( + fn0=lambda: self._forward_mega_moe_routed(hidden_states, input_ids), + fn1=lambda: shared_experts(hidden_states), + event0=events[0], + event1=events[1], + aux_stream=self.shared_experts_stream, + ) + return routed + shared + + def _forward_mega_moe_routed( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None ) -> torch.Tensor: router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( @@ -792,12 +830,7 @@ def _forward_mega_moe( bias_vl=self.gate.bias_vl, image_sentinel_lo=self.image_sentinel_lo, ) - final_hidden_states = self.experts(hidden_states, topk_weights, topk_ids) - if self.shared_experts is not None: - final_hidden_states = final_hidden_states + self.shared_experts( - hidden_states - ) - return final_hidden_states + return self.experts(hidden_states, topk_weights, topk_ids) # Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm. From fe8f2140eeb6103e2ac01fbe2ce40dcc9dba19de Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 18 Oct 2026 23:26:36 +0800 Subject: [PATCH 825/952] [Doc] Document /inference/v1/generate request/response fields and DP rank header (#62305) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- .../online_serving/token_in_token_out.md | 46 ++++++++++++++++++- docs/training/sampling_mask.md | 5 +- examples/scale_out/token_generation_client.py | 2 + 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index eade608bc451..6279fbd78b18 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/docs/serving/online_serving/token_in_token_out.md @@ -1,8 +1,50 @@ # Tokens In <> Tokens Out API -`/inference/v1/generate` takes prompt token IDs and returns generated token IDs. It is the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)). +`POST /inference/v1/generate` is the token-in/token-out (TITO) HTTP endpoint: it takes prompt token IDs and returns generated token IDs, logprobs keyed by integer token ID, and `finish_reason`/`stop_reason`, with no detokenization. Use it when the caller does its own tokenization and chat templating, such as an RL rollout worker. It is also the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)). -The endpoint is registered on `vllm serve` with `--enable-scale-out`, and always with `--tokens-only`. `--tokens-only` also skips loading the tokenizer, so the server can't turn token IDs back into text. +## Enabling the endpoint (`--enable-scale-out`, `--tokens-only`) + +- `--enable-scale-out` registers it, with `/render`, `/derender` and `/inference/v1/abort_requests`. +- `--tokens-only` always registers it, skips loading the tokenizer and forces `detokenize=False`. Stop strings (`sampling_params.stop`) are then rejected with a 400; use `stop_token_ids`. +- With neither flag, it returns 403 `Scale-out endpoints are disabled. Restart vllm serve with --enable-scale-out to enable them.` + +```bash +vllm serve Qwen/Qwen3-0.6B --enable-scale-out + +curl -X POST http://localhost:8000/inference/v1/generate \ + -H "Content-Type: application/json" \ + -d '{"token_ids": [151644, 872, 198], "sampling_params": {"max_tokens": 32, "logprobs": 1}}' +``` + +A full client is in [examples/scale_out/token_generation_client.py](../../../examples/scale_out/token_generation_client.py). + +## Request fields + +| Field | Notes | +| --- | --- | +| `token_ids` | Required `list[int]`: the already-templated prompt. Non-empty, no negative IDs. | +| `sampling_params` | Required. Any [SamplingParams](../../api/README.md#inference-parameters) field, e.g. `max_tokens`, `temperature`, `logprobs`, `stop_token_ids`. Without `max_tokens`, the server applies its own default (as on the OpenAI endpoints), not 16. | +| `output_mode` | `"tokens"` (default) or `"text"`, see [Output modes](#output-modes). | +| `stream`, `stream_options` | SSE chunks ending with `data: [DONE]`. | +| `return_token_ids` | Echo the final prompt IDs as `prompt_token_ids`. | +| `features` / `content_parts` | Multimodal input, mutually exclusive. | +| `request_id`, `model`, `cache_salt`, `priority`, `kv_transfer_params`, `ec_transfer_params` | Optional. | + +## Response fields + +| `choices[i]` field | Notes | +| --- | --- | +| `token_ids` | Generated token IDs. | +| `finish_reason` | `stop`, `length` or `abort`. | +| `stop_reason` | Stop token ID or string; `null` on EOS, length and abort. | +| `logprobs` | `{"content": [{"token_id", "logprob", "rank", "top_logprobs": [...]}]}`, one entry per token, sampled token first. `token_id` is an integer, not the older `"token_id:"` string. Needs `sampling_params.logprobs`. | +| `weight_versions`, `routed_experts`, `sampling_mask` | Optional, see [Weight versions](#weight-versions); the last two need `--enable-return-routed-experts` / `--return-sampling-mask`. | + +The top level adds `request_id`, `usage`, `prompt_token_ids` and `prompt_logprobs`. Stream chunks carry the same per-choice fields for the new tokens. + +## Targeting a data-parallel rank (`X-data-parallel-rank`) + +Send the header `X-data-parallel-rank: ` to pin a request to one data-parallel (DP) engine rank; there is no body field. Without the header, or with a non-integer value, the request is load balanced. A rank outside `[0, data_parallel_size)` is rejected. ## Output modes diff --git a/docs/training/sampling_mask.md b/docs/training/sampling_mask.md index 7de236680070..25dfb414d025 100644 --- a/docs/training/sampling_mask.md +++ b/docs/training/sampling_mask.md @@ -139,8 +139,9 @@ Each position then carries the sampled token's logprob and the top-128 drop them. The finite entries are the sampling mask when it has at most `k` tokens. - The sampled token is always the first entry and may repeat inside the top-k. -- `/inference/v1/generate` encodes IDs as `token_id:`; the OpenAI - endpoints need `return_tokens_as_token_ids`. +- The `token_id:` strings come from the OpenAI endpoints with + `return_tokens_as_token_ids`; `/inference/v1/generate` returns an integer + `token_id` instead. ## Limitations diff --git a/examples/scale_out/token_generation_client.py b/examples/scale_out/token_generation_client.py index a4792a68f9a3..063e8ae39b59 100644 --- a/examples/scale_out/token_generation_client.py +++ b/examples/scale_out/token_generation_client.py @@ -3,6 +3,8 @@ """Call a token-in/token-out server started with: vllm serve Qwen/Qwen3-0.6B --enable-scale-out + +Fields and DP rank routing: docs/serving/online_serving/token_in_token_out.md """ import httpx From ff01a6b14f329ccd4162fce309913542e2abcd5b Mon Sep 17 00:00:00 2001 From: alp Date: Sun, 18 Oct 2026 09:49:42 -0700 Subject: [PATCH 826/952] [Bugfix][Frontend] Stop beam search on generation config EOS tokens (#62072) Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> Co-authored-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> --- tests/samplers/test_beam_search.py | 1 + tests/samplers/test_beam_search_online.py | 71 ++++++++++++++++++- .../generate/beam_search/offline.py | 14 ++-- .../generate/beam_search/online.py | 16 +++-- .../entrypoints/generate/beam_search/utils.py | 24 +++++-- 5 files changed, 110 insertions(+), 16 deletions(-) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 823202ba32f1..4e00c415593f 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -99,6 +99,7 @@ def run_requests(prompts, **kwargs): llm = LLM.__new__(LLM) llm.llm_engine = Mock() + llm.llm_engine.input_processor.generation_config_fields = {} tokenizer = SimpleNamespace( eos_token_id=0, decode=lambda tokens, skip_special_tokens=False: " ".join( diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index a6e4967620b2..b8134e6c8d7e 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -67,6 +67,7 @@ def _check_supports_watermarking(self, params: BeamSearchParams) -> bool: class _InputProcessor: vllm_config = _VllmConfig() + generation_config_fields: dict[str, Any] = {} def resolve_watermarking(self, params): return self.vllm_config._check_supports_watermarking(params) @@ -128,7 +129,7 @@ def _beam_search_step( self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable[..., Any], @@ -140,6 +141,74 @@ def _beam_search_step( return True +@pytest.mark.asyncio +@pytest.mark.parametrize("offline", [False, True]) +@pytest.mark.parametrize("config_eos", [None, 11, [0, 11]]) +@pytest.mark.parametrize("ignore_eos", [False, True]) +async def test_beam_search_respects_generation_config_eos( + monkeypatch, offline: bool, config_eos: int | list[int] | None, ignore_eos: bool +) -> None: + """Model-specific end-of-turn tokens must finish beams like tokenizer EOS.""" + monkeypatch.setattr( + _InputProcessor, "generation_config_fields", {"eos_token_id": config_eos} + ) + prompt: TokensInput = { + "type": "token", + "prompt_token_ids": [1], + } + params = BeamSearchParams( + beam_width=1, + max_tokens=3, + ignore_eos=ignore_eos, + include_stop_str_in_output=True, + watermarking=False, + ) + result = await anext( + _EngineClient().generate(prompt, SamplingParams(watermarking=False)) + ) + result.outputs[0].logprobs = [{11: Logprob(-0.1)}] + should_stop = config_eos is not None and not ignore_eos + expected_tokens = [11] if should_stop else [11, 11, 11] + engine_params: list[SamplingParams] = [] + + if offline: + offline_serving = _OfflineServing() + monkeypatch.setattr( + _OfflineServing, + "_beam_search_step", + BeamSearchOfflineMixin._beam_search_step, + ) + + def render_and_run_requests(*, params, **kwargs): + engine_params.extend(params) + return [result] + + monkeypatch.setattr( + offline_serving, "_render_and_run_requests", render_and_run_requests + ) + beam = offline_serving.beam_search(cast(list, [prompt]), params)[0].sequences[0] + assert beam.tokens == [1, *expected_tokens] + else: + + async def generate(engine_prompt, sampling_params, *args, **kwargs): + engine_params.append(sampling_params) + yield result + + online_serving = _AsyncServing() + monkeypatch.setattr(online_serving.engine_client, "generate", generate) + output = ( + await anext(online_serving.beam_search(prompt, "request", params)) + ).outputs[0] + assert output.token_ids == expected_tokens + assert output.finish_reason == ("stop" if should_stop else "length") + assert output.stop_reason == (11 if should_stop else None) + + # The engine validates stop_token_ids against the vocab before it applies + # the generation config, so the config EOS ids must not be copied into them. + assert engine_params + assert all(not p.stop_token_ids for p in engine_params) + + @pytest.mark.asyncio async def test_beam_search_handles_extra_logprob_candidates() -> None: prompt: TokensInput = { diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 441475cafd60..53297a4736c4 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -27,6 +27,7 @@ BeamSearchOutput, BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) logger = init_logger(__name__) @@ -86,8 +87,6 @@ def beam_search( self.llm_engine.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) engine_inputs = self._preprocess_cmpl(prompts) lora_requests = self._lora_request_to_seq(lora_request, len(engine_inputs)) @@ -129,6 +128,11 @@ def beam_search( params.structured_outputs, tokenizer ) + eos_token_ids = get_beam_search_eos_token_ids( + self.llm_engine.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) instances: list[BeamSearchInstance] = [] for lora_req, prompt in zip(lora_requests, engine_inputs): @@ -169,7 +173,7 @@ def beam_search( should_stop = self._beam_search_step( instances_batch=instances_batch, base_sampling_params=base_sampling_params, - eos_token_id=eos_token_id, + eos_token_ids=eos_token_ids, ignore_eos=ignore_eos, beam_width=beam_width, sort_beams_key=sort_beams_key, @@ -204,7 +208,7 @@ def _beam_search_step( self, instances_batch: list[BeamSearchInstance], base_sampling_params: SamplingParams, - eos_token_id: int | None, + eos_token_ids: set[int], ignore_eos: bool, beam_width: int, sort_beams_key: Callable, @@ -350,7 +354,7 @@ def _beam_search_step( cum_logprob=current_beam.cum_logprob + logprob_obj.logprob, ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: instance.completed.append(new_beam) else: instance_new_beams.append(new_beam) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 9cbe4ea2588b..1df1be014541 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -20,6 +20,7 @@ from .utils import ( BeamSearchSequence, create_sort_beams_key_function, + get_beam_search_eos_token_ids, ) @@ -47,8 +48,6 @@ async def beam_search( self.engine_client.input_processor.resolve_watermarking(params) tokenizer = self.renderer.get_tokenizer() - eos_token_id = tokenizer.eos_token_id - sort_beams_key = create_sort_beams_key_function(eos_token_id, length_penalty) if prompt["type"] == "embeds": raise NotImplementedError("Embedding prompt not supported for beam search") @@ -74,6 +73,11 @@ async def beam_search( sampling_params, self.engine_client.input_processor.vllm_config ) allowed = set(params.allowed_token_ids) if params.allowed_token_ids else None + eos_token_ids = get_beam_search_eos_token_ids( + self.engine_client.input_processor.generation_config_fields, + tokenizer.eos_token_id, + ) + sort_beams_key = create_sort_beams_key_function(eos_token_ids, length_penalty) all_beams = [ BeamSearchSequence( orig_prompt=prompt, @@ -159,17 +163,17 @@ async def beam_search( candidate_logprob = ( current_beam.cum_logprob + logprob_obj.logprob ) - if token_id == eos_token_id and not ignore_eos: + if token_id in eos_token_ids and not ignore_eos: completed.append( BeamSearchSequence( orig_prompt=prompt, - tokens=current_beam.tokens + [eos_token_id] + tokens=current_beam.tokens + [token_id] if include_stop_str_in_output else current_beam.tokens, logprobs=current_beam.logprobs + [logprobs], cum_logprob=candidate_logprob, finish_reason="stop", - stop_reason=eos_token_id, + stop_reason=token_id, ) ) else: @@ -219,7 +223,7 @@ async def beam_search( best_beams = sorted_completed[:beam_width] for beam in best_beams: - if beam.tokens[-1] == eos_token_id and not ignore_eos: + if beam.tokens[-1] in eos_token_ids and not ignore_eos: # Skip the eos token in the text. tokens = beam.tokens[tokenized_length:-1] else: diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 61f97963b6f3..b0312b75f464 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass +from typing import Any from vllm.inputs import ( DecoderOnlyEngineInput, @@ -13,6 +14,7 @@ ) from vllm.logprobs import Logprob from vllm.lora.request import LoRARequest +from vllm.sampling_params import SamplingParams @dataclass @@ -137,7 +139,7 @@ def __init__( def get_beam_search_score( tokens: list[int], cumulative_logprob: float, - eos_token_id: int, + eos_token_ids: set[int], length_penalty: float = 1.0, ) -> float: """Calculate the beam search score with length penalty. @@ -147,17 +149,31 @@ def get_beam_search_score( https://github.com/huggingface/transformers/blob/ccb92be23def445f2afdea94c31286f84b89eb5b/src/transformers/generation/beam_search.py#L938 """ seq_len = len(tokens) - if tokens[-1] == eos_token_id: + if tokens[-1] in eos_token_ids: seq_len -= 1 # An aborted beam may contain only an EOS prompt token. return cumulative_logprob / (max(seq_len, 1) ** length_penalty) -def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): +def create_sort_beams_key_function(eos_token_ids: set[int], length_penalty: float): def sort_beams_key(x: BeamSearchSequence) -> float: return get_beam_search_score( - x.tokens, x.cum_logprob, eos_token_id, length_penalty + x.tokens, x.cum_logprob, eos_token_ids, length_penalty ) return sort_beams_key + + +def get_beam_search_eos_token_ids( + generation_config: dict[str, Any], eos_token_id: int | None +) -> set[int]: + """Return the token ids that end generation: the tokenizer EOS plus the + generation config's ``eos_token_id``. + + This uses a separate SamplingParams so the params sent to the engine keep + their own ``stop_token_ids`` and are validated as before. + """ + params = SamplingParams() + params.update_from_generation_config(generation_config, eos_token_id) + return params.all_stop_token_ids From 6540d5c8f8c73bb2dd6194e56ed002fc32a31837 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 18 Oct 2026 14:29:21 -0400 Subject: [PATCH 827/952] [Bugfix][Spec Decode][Structured Output] Drive grammar masks from GPU logit counts (#52477) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: vllm-agent --- .../spec_decode/test_adaptive_verification.py | 50 ------------------- .../spec_decode/test_mtp_structured_output.py | 2 +- .../test_structured_output_manager.py | 11 ++-- tests/v1/worker/test_gpu_batch_shard.py | 13 ++--- vllm/v1/structured_output/__init__.py | 14 ++++-- vllm/v1/structured_output/utils.py | 17 ++++--- vllm/v1/worker/gpu/model_runner.py | 1 - vllm/v1/worker/gpu/sample/batch_shard.py | 33 ++++-------- vllm/v1/worker/gpu/structured_outputs.py | 49 ++++-------------- vllm/v1/worker/gpu/warmup.py | 4 +- 10 files changed, 54 insertions(+), 140 deletions(-) diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 73c478a10405..84ba8da55ce9 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -16,7 +16,6 @@ maybe_create_adaptive_verification_manager, resolve_adaptive_cudagraph_mode, ) -from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping def make_manager( @@ -307,52 +306,3 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits(): assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype # The prefill keeps its scheduled tokens; only drafts are dropped. assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) - - -def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): - # The scheduler sizes the grammar bitmask from the *scheduled* drafts - # (len(drafts) + 1 rows per request), but a zero budget rewrites - # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping - # from those rewritten offsets drops rows and trips the - # `num_masks == len(mapping)` assert in apply_grammar_bitmask. - manager = make_manager( - np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), - np.ones(64), - ) - manager.req_states.req_id_to_index["prefill"] = 2 - manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) - manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) - manager._max_total_logits = 2 # < 3 requests * 1 bonus token - - scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]} - manager.get_num_tokens( - {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens - ) - assert manager._batch_budget is not None - assert manager._batch_budget[2] == 0 - - req_ids = ["low", "high", "prefill"] - num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) - _, cu_num_logits_np = manager.compact_batch( - num_draft_tokens_per_req, - np.array([3, 3, 40], dtype=np.int32), - np.array([0, 3, 6, 7], dtype=np.int32), - ) - - mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens - mapping = _build_grammar_mapping( - req_ids, - req_ids, - cu_num_logits_np, - num_draft_tokens_per_req, - manager.num_bonus_tokens, - mask_stride, - ) - - num_bitmask_rows = sum( - len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids - ) - assert len(mapping) == num_bitmask_rows - # (request, position) keys, so the kernel can mask rows the compacted - # device layout no longer has room for. - assert mapping == [0, 1, 2, 3, 4, 5, 6] diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 72fa3e7ad7d7..3cdc39fb9864 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -157,7 +157,7 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend): ) assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 + assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 assert (bitmask[0] == -1).all() assert (bitmask[1] == -1).all() assert not (bitmask[2] == -1).all() diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 0cb051f0fddc..b08c11829c22 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -219,8 +219,9 @@ def _run_real_flow( scheduled_spec_decode_tokens={request.request_id: padded}, ) assert bitmask is not None - assert bitmask.shape[0] == len(raw_drafts) + 1 - assert _row_pattern(bitmask, masked_ids or ()) == expected_row_pattern + assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 + row_pattern = _row_pattern(bitmask[: len(raw_drafts) + 1], masked_ids or ()) + assert row_pattern == expected_row_pattern # grammar_bitmask() must rollback any speculative state it advanced. assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated @@ -247,7 +248,7 @@ def _run_real_flow( scheduled_spec_decode_tokens={}, ) assert post_accept_bitmask is not None - assert post_accept_bitmask.shape[0] == 1 + assert post_accept_bitmask.shape[0] == NUM_SPEC_TOKENS + 1 assert (post_accept_bitmask[0] == expected_post_accept_row).all() assert structured_req.reasoning_ended == expected_reasoning assert grammar.is_terminated() is expect_terminated @@ -496,7 +497,7 @@ def test_reasoning_rows_mask_grammar_invalid_implicit_end( scheduled_spec_decode_tokens={request.request_id: drafts}, ) assert bitmask is not None - assert _row_pattern(bitmask, [marker]) == expected_row_pattern + assert _row_pattern(bitmask[: len(drafts) + 1], [marker]) == expected_row_pattern @pytest.mark.parametrize("backend", BACKENDS) @@ -642,7 +643,7 @@ def test_rejected_draft_keeps_later_rows_constrained( scheduled_spec_decode_tokens={request.request_id: drafts}, ) assert bitmask is not None - assert _row_pattern(bitmask) == "CCCC" + assert _row_pattern(bitmask[: len(drafts) + 1]) == "CCCC" assert (bitmask[2] == bitmask[1]).all() assert (bitmask[3] == bitmask[1]).all() diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 60cf6479a70d..ad1d412a33d3 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -292,25 +292,22 @@ def test_shard_grammar_output(tp_size: int): max_num_reqs = 32 num_reqs = 12 slots = rng.choice(np.arange(max_num_reqs), size=num_reqs, replace=False) - k = rng.integers(0, 4, size=num_reqs) - cu = np.zeros(num_reqs + 1, dtype=np.int32) - np.cumsum(1 + k, out=cu[1:]) + mask_stride = 4 owner = slots % tp_size # Grammar requests: every other request, in a scrambled order. req_ids = [f"req-{i}" for i in range(num_reqs)] grammar_idx = list(rng.permutation(np.arange(0, num_reqs, 2))) grammar_req_ids = [req_ids[i] for i in grammar_idx] - num_grammar_rows = sum(int(cu[i + 1] - cu[i]) for i in grammar_idx) + num_grammar_rows = len(grammar_idx) * mask_stride bitmask = rng.integers(0, 2**31, size=(num_grammar_rows, 4), dtype=np.int32) - input_batch = SimpleNamespace(req_ids=req_ids, cu_num_logits_np=cu) # Global row offset of each grammar request within the bitmask. row_offsets = {} cursor = 0 for i in grammar_idx: row_offsets[i] = cursor - cursor += int(cu[i + 1] - cu[i]) + cursor += mask_stride grammar_output = GrammarOutput( structured_output_request_ids=grammar_req_ids, @@ -320,7 +317,7 @@ def test_shard_grammar_output(tp_size: int): kept_total = 0 for rank in range(tp_size): local_req_ids = [req_ids[i] for i in range(num_reqs) if owner[i] == rank] - local = _shard_grammar_output(grammar_output, input_batch, local_req_ids) + local = _shard_grammar_output(grammar_output, local_req_ids) expected_idx = [i for i in grammar_idx if owner[i] == rank] if not expected_idx: assert local is None @@ -329,7 +326,7 @@ def test_shard_grammar_output(tp_size: int): assert local.structured_output_request_ids == [req_ids[i] for i in expected_idx] expected_rows = np.concatenate( [ - np.arange(row_offsets[i], row_offsets[i] + cu[i + 1] - cu[i]) + np.arange(row_offsets[i], row_offsets[i] + mask_stride) for i in expected_idx ] ) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 7de9dc22165e..dc3a236aa7cc 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -375,6 +375,8 @@ def grammar_bitmask( # Covers both speculative decoding and diffusion LLMs (canvas_length). max_num_spec_tokens = self.vllm_config.num_speculative_tokens + is_diffusion = self.vllm_config.model_config.is_diffusion + mask_stride = max_num_spec_tokens + (0 if is_diffusion else 1) if self._grammar_bitmask is None: assert self.backend is not None @@ -390,7 +392,7 @@ def grammar_bitmask( # Generate a batched bitmask for all structured output requests. # When speculative decoding is enabled, we need to include multiple # masks for each request, one for each possible bonus token position. - # These are stored inline in the tensor and unpacked by the gpu runner. + # Each request owns a fixed-width block of `mask_stride` rows. cumulative_index = 0 # Optimized parallel filling of bitmasks for @@ -430,7 +432,8 @@ def grammar_bitmask( promise.result() else: # Fallback to serial filling of bitmasks for small-batch-size cases - for req_id in structured_output_request_ids: + for grammar_idx, req_id in enumerate(structured_output_request_ids): + cumulative_index = grammar_idx * mask_stride request = requests[req_id] structured_output_request = request.structured_output_request @@ -484,7 +487,7 @@ def grammar_bitmask( # Diffusion LLMs don't sample a bonus token after the # scheduled positions, so skip its bitmask in that case. - if not (self.vllm_config.model_config.is_diffusion and req_tokens): + if not (is_diffusion and req_tokens): if failed_index is not None: bitmask[cumulative_index].copy_(bitmask[failed_index]) else: @@ -502,8 +505,9 @@ def grammar_bitmask( grammar.rollback(state_advancements) bitmask_tensor = self._grammar_bitmask - if cumulative_index < bitmask_tensor.shape[0]: - bitmask_tensor = bitmask_tensor[:cumulative_index] + num_masks = len(structured_output_request_ids) * mask_stride + if num_masks < bitmask_tensor.shape[0]: + bitmask_tensor = bitmask_tensor[:num_masks] # After finishing with the xgrammar operations, we convert to # np.ndarray, because that is much more efficient for serialization diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 5a0b95c17bdb..936ec4f9af2a 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -136,7 +136,7 @@ def apply_grammar_bitmask( if req_id in struct_out_req_ids: struct_out_req_batch_indices[req_id] = logit_index - out_indices = [] + out_indices: list[int] = [] # Reorder the bitmask to match the order of the requests in the batch. sorted_bitmask_tensor = torch.full( @@ -146,15 +146,18 @@ def apply_grammar_bitmask( pin_memory=PIN_MEMORY, ) sorted_bitmask = sorted_bitmask_tensor.numpy() + mask_stride = grammar_bitmask.shape[0] // len( + grammar_output.structured_output_request_ids + ) cumulative_index = 0 for req_id in grammar_output.structured_output_request_ids: - num_spec_tokens = len(spec_tokens.get(req_id, ())) + num_rows = 1 + len(spec_tokens.get(req_id, ())) if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None: - for i in range(1 + num_spec_tokens): - bitmask_index = logit_idx + i - sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i] - out_indices.append(bitmask_index) - cumulative_index += 1 + num_spec_tokens + sorted_bitmask[logit_idx : logit_idx + num_rows] = grammar_bitmask[ + cumulative_index : cumulative_index + num_rows + ] + out_indices.extend(range(logit_idx, logit_idx + num_rows)) + cumulative_index += mask_stride # Copy async to device. grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b8ad557b68a8..bb45af43f48b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -527,7 +527,6 @@ def load_model( vocab_size=self.vocab_size, device=self.device, mask_stride=self.decode_query_len, - num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index 695b1d543563..d8cecb7f87b7 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -275,9 +275,7 @@ def shard_sampler_inputs( ) local_grammar_output = None if grammar_output is not None: - local_grammar_output = _shard_grammar_output( - grammar_output, input_batch, local_req_ids - ) + local_grammar_output = _shard_grammar_output(grammar_output, local_req_ids) metadata = BatchShardMetadata( tp_size=tp_size, num_logits_per_rank=num_logits_per_rank_np.tolist(), @@ -314,34 +312,25 @@ def all_to_all_logits( def _shard_grammar_output( grammar_output: GrammarOutput, - input_batch: InputBatch, local_req_ids: list[str], ) -> GrammarOutput | None: owned = set(local_req_ids) - req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)} - cu_num_logits_np = input_batch.cu_num_logits_np - - local_ids: list[str] = [] - keep_indices: list[int] = [] - cursor = 0 - for req_id in grammar_output.structured_output_request_ids: - req_idx = req_id_to_idx[req_id] - num_req_logits = int(cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]) - if req_id in owned: - local_ids.append(req_id) - keep_indices.extend(range(cursor, cursor + num_req_logits)) - cursor += num_req_logits - if not local_ids: + grammar_req_ids = grammar_output.structured_output_request_ids + keep = [i for i, req_id in enumerate(grammar_req_ids) if req_id in owned] + if not keep: return None + local_ids = [grammar_req_ids[i] for i in keep] + # Each request owns a fixed-width block of rows. + bitmask = grammar_output.grammar_bitmask + blocks = bitmask.reshape(len(grammar_req_ids), -1, bitmask.shape[-1]) + local_bitmask = blocks[keep].reshape(-1, bitmask.shape[-1]) # num_acceptable_drafts is ordered with structured_output_request_ids. num_acceptable = grammar_output.num_acceptable_drafts if num_acceptable is not None: - owned = set(local_ids) - ids = grammar_output.structured_output_request_ids - num_acceptable = [n for i, n in zip(ids, num_acceptable) if i in owned] + num_acceptable = [num_acceptable[i] for i in keep] return GrammarOutput( structured_output_request_ids=local_ids, - grammar_bitmask=grammar_output.grammar_bitmask[keep_indices], + grammar_bitmask=local_bitmask, num_acceptable_drafts=num_acceptable, ) diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 336bc85861d9..017d7b53a709 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -9,32 +9,6 @@ from vllm.v1.worker.gpu.input_batch import InputBatch -def _build_grammar_mapping( - req_ids: list[str], - grammar_req_ids: list[str], - cu_num_logits_np: np.ndarray, - num_draft_tokens_per_req: np.ndarray | None, - num_bonus_tokens: int, - mask_stride: int, -) -> list[int]: - mapping: list[int] = [] - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - if num_draft_tokens_per_req is None: - num_positions = int( - cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] - ) - else: - # Grammar masks follow the scheduled layout even when adaptive - # verification compacts the actual CPU logit offsets to bonus-only. - num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens - mapping.extend( - req_idx * mask_stride + position for position in range(num_positions) - ) - return mapping - - def grammar_invalid_drafts( input_batch: InputBatch, grammar_req_ids: list[str], @@ -79,7 +53,6 @@ def __init__( vocab_size: int, device: torch.device, mask_stride: int, - num_bonus_tokens: int, ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device @@ -90,7 +63,6 @@ def __init__( self.device = device self.copy_stream = torch.cuda.Stream() self.mask_stride = mask_stride - self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -108,18 +80,15 @@ def apply_grammar_bitmask( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) - # Construct bitmask -> logits mapping - # Key by (request, position) rather than absolute logit index: - # adaptive verification finalizes per-request logit offsets on - # device, so the kernel resolves them from the GPU cu_num_logits. - mapping = _build_grammar_mapping( - input_batch.req_ids, - grammar_req_ids, - input_batch.cu_num_logits_np, - input_batch.num_draft_tokens_per_req, - self.num_bonus_tokens, - self.mask_stride, - ) + # Each grammar request owns `mask_stride` bitmask rows. Key them by + # (request, position); the kernel skips positions beyond the device + # cu_num_logits, which adaptive verification may trim. + req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)} + mapping = [ + req_id_to_idx[req_id] * self.mask_stride + position + for req_id in grammar_req_ids + for position in range(self.mask_stride) + ] # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream): diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 848d984295fd..9b54319a879b 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -384,7 +384,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: vocab_size = model_runner.model_config.get_vocab_size() bitmask_width = (vocab_size + 31) // 32 grammar_bitmask = np.full( - (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 + (len(req_ids) * model_runner.decode_query_len, bitmask_width), + fill_value=-1, + dtype=np.int32, ) grammar_output = GrammarOutput( structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask From 82a24bb82dd9344cce7a43195bb51a2390adf00a Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 18 Oct 2026 11:45:08 -0700 Subject: [PATCH 828/952] [Bugfix][ROCm] Only set inductor batch_invariant when torch has it (#62274) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- vllm/model_executor/determinism/batch_invariant.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 933e5ecef6ea..589452c75141 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -1426,7 +1426,8 @@ def override_envs_for_invariance(): os.environ["VLLM_USE_AOT_COMPILE"] = "0" import torch._inductor.config as inductor_config - inductor_config.batch_invariant = True + if hasattr(inductor_config, "batch_invariant"): + inductor_config.batch_invariant = True def init_batch_invariance(): From e0adacc9a67a5f2b2d8ba508f7f522bc1ff83c6d Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:29:43 -0700 Subject: [PATCH 829/952] [Bugfix][Rust Frontend] Use xgrammar for gRPC structured outputs (#61892) Signed-off-by: Alec Flowers Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: Bugen Zhao Co-authored-by: vllm-agent --- .../src/protocol/structured_outputs.rs | 58 ++++++- .../protocol/structured_outputs/xgrammar.rs | 142 ++++++++++++++++++ rust/src/server/src/grpc/convert.rs | 51 +++++-- rust/src/server/src/grpc/tests.rs | 51 +++++++ rust/src/text/src/error.rs | 3 +- rust/src/text/src/lower.rs | 31 ++++ .../test_structured_output_manager.py | 19 +++ vllm/v1/structured_output/backend_xgrammar.py | 6 + 8 files changed, 343 insertions(+), 18 deletions(-) create mode 100644 rust/src/engine-core-client/src/protocol/structured_outputs/xgrammar.rs diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs.rs b/rust/src/engine-core-client/src/protocol/structured_outputs.rs index 99fa66ae2b02..003d150cac52 100644 --- a/rust/src/engine-core-client/src/protocol/structured_outputs.rs +++ b/rust/src/engine-core-client/src/protocol/structured_outputs.rs @@ -7,11 +7,14 @@ use serde_json::Value; use crate::error::{Error, Result}; +mod xgrammar; + /// Structured-output backend selected for EngineCore grammar compilation. /// /// Python vLLM stores this in `StructuredOutputsParams._backend` after request -/// validation. The Rust frontend currently always lowers structured-output -/// requests to guidance, while ignoring any user-supplied `_backend` value. +/// validation. The Rust frontend picks the backend itself and ignores any +/// user-supplied `_backend` value: gRPC constraints and parser-built output +/// grammars use xgrammar, other HTTP constraints use guidance. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub enum StructuredOutputBackend { @@ -118,18 +121,27 @@ impl StructuredOutputsParams { message: "structured_outputs.json cannot be an empty string".to_string(), }); } - serde_json::from_str::(value).map_err(|_| { + let schema = serde_json::from_str::(value).map_err(|_| { Error::InvalidStructuredOutputsParams { message: "Invalid JSON grammar specification.".to_string(), } })?; - Ok(()) + self.validate_json_schema(&schema) } + StructuredOutputConstraint::Json(schema) => self.validate_json_schema(schema), StructuredOutputConstraint::Choice(choice) if choice.is_empty() => { Err(Error::InvalidStructuredOutputsParams { message: "structured_outputs.choice cannot be empty".to_string(), }) } + // NUL is never meaningful in a regex and the regex-to-grammar + // conversion does not handle it; Python rejects it for every backend. + StructuredOutputConstraint::Regex(regex) if regex.contains('\0') => { + Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.regex must not contain a NUL character ('\\x00')" + .to_string(), + }) + } StructuredOutputConstraint::StructuralTag(tag) if tag.trim().is_empty() => { Err(Error::InvalidStructuredOutputsParams { message: "structured_outputs.structural_tag cannot be an empty string" @@ -139,6 +151,20 @@ impl StructuredOutputsParams { _ => Ok(()), } } + + /// Reject a schema the selected backend would accept but enforce only in + /// part, as Python does when validating against xgrammar. + fn validate_json_schema(&self, schema: &Value) -> Result<()> { + if self.backend == StructuredOutputBackend::Xgrammar + && xgrammar::has_unsupported_json_features(schema) + { + return Err(Error::InvalidStructuredOutputsParams { + message: "The provided JSON schema contains features not supported by xgrammar." + .to_string(), + }); + } + Ok(()) + } } /// Wire-compatible structured-output payload used by Python engine-core. @@ -364,6 +390,30 @@ mod tests { assert!(error.to_string().contains("choice cannot be empty")); } + #[test] + fn validate_rejects_nul_in_regex() { + let error = StructuredOutputsParams::regex("a\0b").validate().unwrap_err(); + assert!(error.to_string().contains("must not contain a NUL character")); + } + + #[test] + fn validate_rejects_json_features_xgrammar_ignores() { + let schema = serde_json::json!({ "type": "array", "uniqueItems": true }); + for params in [ + StructuredOutputsParams::json(schema.clone()), + StructuredOutputsParams::json(Value::String(schema.to_string())), + ] { + assert!(params.validate().is_ok(), "guidance enforces {schema}"); + let error = StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..params + } + .validate() + .unwrap_err(); + assert!(error.to_string().contains("not supported by xgrammar")); + } + } + #[test] fn validate_rejects_empty_structural_tag() { let error = StructuredOutputsParams::structural_tag("\n\t").validate().unwrap_err(); diff --git a/rust/src/engine-core-client/src/protocol/structured_outputs/xgrammar.rs b/rust/src/engine-core-client/src/protocol/structured_outputs/xgrammar.rs new file mode 100644 index 000000000000..6dba08720a55 --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/structured_outputs/xgrammar.rs @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Request-side xgrammar handling mirrored from Python's xgrammar backend. + +use serde_json::{Map, Value}; + +const STRING_SUPPORTED_FORMATS: &[&str] = &[ + "email", + "date", + "time", + "date-time", + "duration", + "ipv4", + "ipv6", + "hostname", + "uuid", + "uri", + "uri-reference", + "uri-template", + "json-pointer", + "relative-json-pointer", +]; + +/// Port of Python's `has_xgrammar_unsupported_json_features`. xgrammar +/// ignores these constraints rather than rejecting the schema. +pub(super) fn has_unsupported_json_features(schema: &Value) -> bool { + let Some(obj) = schema.as_object() else { + return false; + }; + let types = schema_types(obj); + let has_type = |name: &str| types.contains(&name); + + if (has_type("integer") || has_type("number")) && obj.contains_key("multipleOf") { + return true; + } + if has_type("array") + && ["uniqueItems", "contains", "minContains", "maxContains"] + .iter() + .any(|key| obj.contains_key(*key)) + { + return true; + } + if has_type("string") { + let unsupported_format = obj.get("format").is_some_and(|format| { + !format.as_str().is_some_and(|f| STRING_SUPPORTED_FORMATS.contains(&f)) + }); + if unsupported_format || has_pattern_and_length_bounds(obj) { + return true; + } + } + if has_type("object") { + if let Some(property_names) = obj.get("propertyNames") + && (property_names.as_object().is_some_and(has_pattern_and_length_bounds) + || obj.contains_key("patternProperties") + || obj.contains_key("properties") + || obj + .get("unevaluatedProperties") + .is_some_and(|value| *value != Value::Bool(true))) + { + return true; + } + if let Some(Value::Object(pattern_properties)) = obj.get("patternProperties") + && (obj.contains_key("properties") || pattern_properties.len() > 1) + { + return true; + } + } + if obj + .get("allOf") + .and_then(Value::as_array) + .is_some_and(|branches| branches.len() >= 2) + { + return true; + } + + // serde_json's parser recursion limit bounds this recursion. + obj.values().any(|value| match value { + Value::Object(_) => has_unsupported_json_features(value), + Value::Array(items) => items.iter().any(has_unsupported_json_features), + _ => false, + }) +} + +fn schema_types(obj: &Map) -> Vec<&str> { + match obj.get("type") { + Some(Value::String(name)) => vec![name.as_str()], + Some(Value::Array(names)) => names.iter().filter_map(Value::as_str).collect(), + _ => Vec::new(), + } +} + +fn has_pattern_and_length_bounds(obj: &Map) -> bool { + (obj.contains_key("pattern") || obj.contains_key("format")) + && (obj.contains_key("minLength") || obj.contains_key("maxLength")) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn unsupported_json_features_match_python() { + let unsupported = [ + json!({"type": "integer", "multipleOf": 2}), + json!({"type": ["number", "null"], "multipleOf": 0.5}), + json!({"type": "array", "uniqueItems": true}), + json!({"type": "array", "contains": {"type": "string"}}), + json!({"type": "string", "format": "binary"}), + json!({"type": "string", "format": 7}), + json!({"type": "string", "pattern": "a+", "maxLength": 3}), + json!({"type": "object", "propertyNames": {"pattern": "a", "minLength": 1}}), + json!({"type": "object", "propertyNames": {}, "properties": {}}), + json!({"type": "object", "propertyNames": {}, "unevaluatedProperties": false}), + json!({"type": "object", "patternProperties": {"a": {}, "b": {}}}), + json!({"type": "object", "patternProperties": {"a": {}}, "properties": {}}), + json!({"allOf": [{"type": "string"}, {"enum": ["yes", "no"]}]}), + json!({"type": "object", "properties": {"x": {"type": "array", "uniqueItems": true}}}), + json!({"anyOf": [{"type": "integer", "multipleOf": 3}]}), + ]; + for schema in unsupported { + assert!(has_unsupported_json_features(&schema), "{schema}"); + } + + let supported = [ + json!({"type": "object", "properties": {"x": {"type": "string"}}}), + json!({"type": "string", "format": "email"}), + json!({"type": "string", "pattern": "a+"}), + json!({"type": "array", "items": {"type": "integer"}, "minItems": 1}), + json!({"multipleOf": 2}), + json!({"type": "object", "propertyNames": {"pattern": "a"}}), + json!({"type": "object", "propertyNames": {}, "unevaluatedProperties": true}), + json!({"allOf": [{"type": "string"}]}), + json!(true), + ]; + for schema in supported { + assert!(!has_unsupported_json_features(&schema), "{schema}"); + } + } +} diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 1eec1d0fe320..a5354da2501b 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -4,13 +4,14 @@ //! Conversion between gRPC protobuf types and internal `vllm-text` //! request/response types. -use thiserror_ext::AsReport as _; use tonic::Status; use uuid::Uuid; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::output::StopReason; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; -use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; +use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputBackend, StructuredOutputsParams, +}; use vllm_text::{ DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, PromptTruncation, PromptTruncationLimit, SamplingParams, TextDecodeOptions, TextRequest, TruncationSide, @@ -275,6 +276,8 @@ fn candidate_logprob_spec( } } +/// EngineCore compiles every request with the backend of the first one it sees, +/// so all gRPC constraints use xgrammar. fn convert_structured_output( d: &pb::DecodingParameters, ) -> Result, Status> { @@ -283,7 +286,7 @@ fn convert_structured_output( Some(so) => so, }; use pb::decoding_parameters::StructuredOutput; - let params = match so { + let mut params = match so { StructuredOutput::Json(schema) => { let json: serde_json::Value = serde_json::from_str(schema) .map_err(|e| Status::invalid_argument(format!("invalid json schema: {e}")))?; @@ -300,9 +303,7 @@ fn convert_structured_output( StructuredOutputsParams::structural_tag(tag.clone()) } }; - params - .validate() - .map_err(|error| Status::invalid_argument(error.to_report_string()))?; + params.backend = StructuredOutputBackend::Xgrammar; Ok(Some(params)) } @@ -574,7 +575,8 @@ mod tests { use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; use super::{ - ResponseOpts, json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request, + ResponseOpts, Status, StructuredOutputBackend, StructuredOutputsParams, + json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request, }; fn base_request() -> pb::GenerateRequest { @@ -759,18 +761,41 @@ mod tests { assert_eq!(text.sampling_params.temperature, Some(0.0)); } - #[test] - fn grpc_rejects_empty_grammar_before_engine() { - use super::pb::decoding_parameters::StructuredOutput; + fn lower_structured_output( + so: super::pb::decoding_parameters::StructuredOutput, + ) -> Result { let req = pb::GenerateRequest { decoding: Some(pb::DecodingParameters { - structured_output: Some(StructuredOutput::Grammar(" ".to_string())), + structured_output: Some(so), ..Default::default() }), ..base_request() }; - let err = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); - assert!(err.message().contains("grammar cannot be an empty string")); + Ok(to_text_request(req, false, &["test-model".to_string()])? + .sampling_params + .structured_outputs + .expect("structured outputs set")) + } + + #[test] + fn grpc_structured_outputs_use_xgrammar() { + use super::pb::decoding_parameters::StructuredOutput; + for so in [ + StructuredOutput::Json(r#"{"type": "object"}"#.to_string()), + StructuredOutput::JsonObject(true), + StructuredOutput::Regex("[a-z]+".to_string()), + StructuredOutput::Grammar(r#"root ::= "a""#.to_string()), + StructuredOutput::Choice(super::pb::decoding_parameters::StringChoices { + choices: vec!["a".to_string()], + }), + StructuredOutput::StructuralTag( + r#"{"type": "structural_tag", "format": {"type": "const_string", "value": "x"}}"# + .to_string(), + ), + ] { + let params = lower_structured_output(so.clone()).expect("convert ok"); + assert_eq!(params.backend, StructuredOutputBackend::Xgrammar, "{so:?}"); + } } #[test] diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index fa041bdf33d2..98d6579bce75 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1388,6 +1388,57 @@ async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_invalid_structured_outputs_returns_invalid_argument() { + use pb::decoding_parameters::{StringChoices, StructuredOutput}; + + let (mut client, server_task, _engine_task) = grpc_test_server( + b"engine-grpc-invalid-structured-outputs", + default_stream_output_specs(), + ) + .await; + + for (structured_output, message) in [ + ( + StructuredOutput::Grammar(" ".to_string()), + "grammar cannot be an empty string", + ), + ( + StructuredOutput::Choice(StringChoices { choices: vec![] }), + "choice cannot be empty", + ), + ( + StructuredOutput::Regex("a\0b".to_string()), + "must not contain a NUL character", + ), + // xgrammar would accept this schema but ignore `uniqueItems`. + ( + StructuredOutput::Json(r#"{"type": "array", "uniqueItems": true}"#.to_string()), + "not supported by xgrammar", + ), + ] { + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-structured-outputs".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + decoding: Some(pb::DecodingParameters { + structured_output: Some(structured_output), + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("invalid structured outputs must be rejected before the engine"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains(message), "{}", status.message()); + } + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index b917fe57cb80..c3e8ccc11d5d 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -77,7 +77,8 @@ impl Error { | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) | Self::Llm(LlmError::EngineCoreClient( EngineCoreError::InvalidDataParallelRank { .. }, - )) => true, + )) + | Self::EngineCore(EngineCoreError::InvalidStructuredOutputsParams { .. }) => true, _ => false, } } diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 319c9953698a..edae40317b62 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -136,6 +136,11 @@ pub fn lower_sampling_params( sampling_limits, )?; validate_repetition_detection(repetition_detection.as_ref())?; + // Validate the final constraint: chat may have replaced the user's + // constraint with a parser-built output grammar. + if let Some(structured_outputs) = &structured_outputs { + structured_outputs.validate()?; + } // Mirrors the model-generation-config inheritance used by vLLM's OpenAI chat // path: https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/openai/chat_completion/protocol.py#L424-L450 @@ -346,10 +351,14 @@ mod tests { use std::num::NonZeroU32; use serial_test::file_serial; + use thiserror_ext::AsReport as _; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use vllm_engine_core_client::protocol::multimodal::{ MmFeatureSpec, MmModality, PlaceholderRange, }; + use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputBackend, StructuredOutputsParams, + }; use vllm_engine_core_client::protocol::tensor::WireNdArray; use vllm_tokenizer::test_utils::TestTokenizer; @@ -480,6 +489,28 @@ mod tests { assert_eq!(params.min_tokens, 0); } + #[test] + fn lower_sampling_params_rejects_schema_xgrammar_cannot_enforce() { + let structured_outputs = StructuredOutputsParams { + backend: StructuredOutputBackend::Xgrammar, + ..StructuredOutputsParams::json(serde_json::json!({ + "type": "integer", + "multipleOf": 2, + })) + }; + let error = lower_sampling_params_with_limits( + SamplingParams { + structured_outputs: Some(structured_outputs), + ..Default::default() + }, + sample_sampling_limits(), + ) + .unwrap_err(); + + assert!(error.is_request_validation_error()); + assert!(error.to_report_string().contains("not supported by xgrammar")); + } + #[test] fn lower_sampling_params_validates_repetition_detection() { let lower = |repetition_detection| { diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index b08c11829c22..8d3b98be7a5b 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -25,6 +25,7 @@ from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.structured_output.backend_outlines import OutlinesBackend, OutlinesGrammar +from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend from vllm.v1.structured_output.request import get_structured_output_key TOKENIZER = "gpt2" @@ -694,6 +695,24 @@ def test_outlines_choice_with_non_bmp_characters(tokenizer): assert grammar.accept_tokens("", tokenizer.encode(choices[0])) +def test_xgrammar_compiles_choice_without_rewrite(tokenizer): + """Frontends that skip Python request validation send `choice` as is.""" + choices = ["yes", "no\nplease"] + backend = XgrammarBackend( + VllmConfig(), tokenizer=tokenizer, vocab_size=len(tokenizer) + ) + request_type, grammar_spec = get_structured_output_key( + StructuredOutputsParams(choice=choices) + ) + + def accepts(text: str) -> bool: + grammar = backend.compile_grammar(request_type, grammar_spec) + return grammar.accept_tokens("", tokenizer.encode(text)) + + assert accepts(choices[1]) + assert not accepts("maybe") + + def test_clear_backend_waits_for_running_grammar_and_cancels_pending( tokenizer, monkeypatch ): diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 94ccc55b8b98..aae72873b2e6 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -164,6 +164,12 @@ def compile_grammar( self.compiler.compile_regex, grammar_spec, ) + elif request_type == StructuredOutputOptions.CHOICE: + # xgrammar has no choice type. Python request validation rewrites + # `choice` into this grammar, but other frontends send it as is. + ctx = self.compiler.compile_grammar( + choice_as_grammar(json.loads(grammar_spec)) + ) elif request_type == StructuredOutputOptions.STRUCTURAL_TAG: s_tag = json.loads(grammar_spec) if "structures" in s_tag: From 36e0ae2df2882de8027bb484da9179c10646ccee Mon Sep 17 00:00:00 2001 From: Anirudh Ravichandran Date: Sun, 18 Oct 2026 15:47:48 -0700 Subject: [PATCH 830/952] [Bugfix][MoE] Keep DeepEP high-throughput buffers alive across DBO streams (#60735) Signed-off-by: Anirudh Ravichandran Signed-off-by: Andreas Karatzas Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- .../fused_moe/prepare_finalize/deepep_ht.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index 035d5552693d..e5bae6247d5a 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input from vllm.platforms import current_platform from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import current_stream from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, dbo_enabled, @@ -26,6 +27,23 @@ ) +def _dbo_record_stream(*tensors: torch.Tensor | None) -> None: + """Mark tensors allocated on the other DBO stream as used by this one. + + Under DBO, DeepEP allocates its outputs on the comm stream while the + experts read them on the compute stream, and the reverse holds for the + dispatch and combine inputs. Without this, the caching allocator can hand + a freed block to the other ubatch while kernels on this stream still read + it. + """ + if not dbo_enabled(): + return + stream = current_stream() + for t in tensors: + if t is not None: + t.record_stream(stream) + + class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): """Prepare/Finalize using DeepEP High-Throughput kernels.""" @@ -124,6 +142,9 @@ def _do_dispatch( # kernel will block the CPU so we want to queue up all the compute # for the other ubatch before the dispatch kernel starts. dbo_yield_and_switch_from_compute_to_comm() + # The ROCm host sync after dispatch already orders these reads. + if not self.sync_dbo_comm: + _dbo_record_stream(tokens, token_scales, rank_topk_ids, rank_topk_weights) ( num_tokens_per_rank, @@ -209,6 +230,9 @@ def _receiver( expert_x, expert_x_scale = token_data else: expert_x, expert_x_scale = token_data, None + _dbo_record_stream( + expert_x, expert_x_scale, expert_topk_ids, expert_topk_weights + ) # The existing MOE kernels assume that all entries of topk_ids are # valid. To that effect, set the -1s in expert_topk_ids to some expert @@ -370,6 +394,8 @@ def _finalize( ) previous_event = dbo_get_previous_event(self.buffer.capture) dbo_yield_and_switch_from_compute_to_comm() + if not self.sync_dbo_comm: + _dbo_record_stream(fused_expert_output) assert fused_expert_output.dtype == torch.bfloat16, ( f"Expected fused_expert_output bfloat16, got {fused_expert_output.dtype}" ) From 8f17dbf707459686eb7e500748552be98612218c Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 19:07:55 -0500 Subject: [PATCH 831/952] [Bugfix][KV Connector] Pair non-divisible hybrid MLA block sizes in NixlConnector (#61002) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- .../unit/test_nixl_desc_geometry.py | 84 +++++++++++++++---- .../kv_connector/v1/nixl/base_worker.py | 55 ++++++++++++ 2 files changed, 123 insertions(+), 16 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index b22d7c9485d0..057209a4f495 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -11,6 +11,7 @@ mid-decode (silent corruption of an unrelated request). """ +import math from collections import defaultdict from threading import Event, Lock from types import SimpleNamespace @@ -573,8 +574,11 @@ def test_registration_ignores_a_sub_block_padding_tail(page_covers_view): assert _descriptor_geometry(padded) == _descriptor_geometry(unpadded) -def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks): - """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout.""" +def _make_mla_hybrid_worker( + local_block_size, kernel_block_size, num_logical_blocks, worker_cls=None +): + """Build a real worker (pull by default) with a hybrid MLA + 2xKDA HMA + layout.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( base_worker as bw, ) @@ -655,7 +659,9 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo ), set_current_vllm_config(vllm_config), ): - worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + worker = (worker_cls or NixlConnectorWorker)( + vllm_config, "local-engine", kv_cache_config + ) worker.use_mla = True # Attention caches are kernel-block granular on dim 0, as the @@ -667,13 +673,15 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo ) for _ in range(2) ] + # MLA caches are [B, 1, N, C] on device, tokens on dim 2. + mla_views = [t.view(t.shape[0], 1, kernel_block_size, -1) for t in tensors] worker.register_kv_caches( { "kda_a.0": tensors[0], - "mla.0": tensors[0], + "mla.0": mla_views[0], "kda_b.0": tensors[0], "kda_a.1": tensors[1], - "mla.1": tensors[1], + "mla.1": mla_views[1], "kda_b.1": tensors[1], } ) @@ -736,11 +744,11 @@ def _make_remote_meta( ) remote_ppl = remote_block_size // remote_kernel_block_size - # Kernel-granularity pages are TP-independent for MLA hybrids and must - # match the local ones for the handshake to pass, scaled down by the - # block-size ratio when the remote's kernel block is smaller. - block_size_ratio = worker.block_size // remote_kernel_block_size - kernel_page = worker.block_len_per_layer[0] // block_size_ratio + # Per-token page bytes are TP-independent for MLA hybrids and must match + # the local ones for the handshake to pass. + kernel_page = ( + worker.block_len_per_layer[0] * remote_kernel_block_size // worker.block_size + ) return NixlAgentMetadata( engine_id="remote-engine", agent_metadata=b"remote-agent-meta", @@ -904,7 +912,8 @@ def _run_hetero_case( ) remote_kernel = remote_kernel or kernel - block_size_ratio = kernel // remote_kernel + # Transfer unit: the remote kernel block when it divides the local one. + unit = math.gcd(kernel, remote_kernel) remote_ppl = remote_block // remote_kernel matched = num_tokens - 1 # mamba N-1 rule n_local = -(-num_tokens // local_block) @@ -967,9 +976,8 @@ def _run_hetero_case( remote_bases = [0x10_000_000, 0x20_000_000] local_unified = worker._test_unified_page remote_unified = (local_unified // local_block) * remote_block - # With block_size_ratio > 1 the local page is split into ratio sub-descs, - # each the size of a whole remote kernel page. - desc_page = worker.block_len_per_layer[0] // block_size_ratio + # The local page is split into sub-descs of one transfer unit each. + desc_page = worker.block_len_per_layer[0] * unit // kernel meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified covered_tokens = set() for op, lh, lids, rh, rids in nixl.xfers: @@ -1013,8 +1021,8 @@ def _run_hetero_case( covered_tokens.add(ltok) # Invariant 3: full coverage of the matched tokens, at the finest - # transfer granularity (the remote kernel block). - needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)} + # transfer granularity. + needed = {t for t in range(0, matched - matched % unit, unit)} missing = needed - covered_tokens assert not missing, ( f"tokens never transferred: {sorted(missing)[:8]} " @@ -1075,6 +1083,50 @@ def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens) ) +@pytest.mark.cpu_test +@pytest.mark.parametrize("local_block,remote_block", [(36, 20), (20, 36)]) +@pytest.mark.parametrize( + "num_tokens", [2, 4, 5, 19, 20, 21, 36, 37, 40, 41, 72, 73, 80, 81] +) +def test_hetero_non_divisible_block_sizes(local_block, remote_block, num_tokens): + """Hybrid MLA blocks follow the TP-sharded SSM page, so on backends that + take any kernel block size P and D may have blocks neither of which + divides the other; they pair in gcd-sized units.""" + _run_hetero_case( + local_block, + kernel=local_block, + remote_block=remote_block, + remote_kernel=remote_block, + num_tokens=num_tokens, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "worker_cls,local_block,remote_block", + [(NixlPullConnectorWorker, 257, 256), (NixlPushConnectorWorker, 36, 20)], +) +def test_hetero_block_split_is_declined(worker_cls, local_block, remote_block): + """Pull pairs whose gcd would split each block into very many transfers, + and every push pair, fail the handshake instead of being split.""" + worker = _make_mla_hybrid_worker( + local_block_size=local_block, + kernel_block_size=local_block, + num_logical_blocks=8, + worker_cls=worker_cls, + ) + meta_r = _make_remote_meta( + worker, + remote_block_size=remote_block, + remote_kernel_block_size=remote_block, + remote_num_logical=8, + remote_ssm_sizes=(24, 32), + ) + with pytest.raises(AssertionError, match="not divisible"): + _register_remote_agents(worker, meta_r, tp_size=2) + assert meta_r.block_size == remote_block + + @pytest.mark.cpu_test @pytest.mark.parametrize( "num_tokens", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 4e86000bb11b..c604b210cfc6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -107,6 +107,7 @@ # invalidated, eg due to a persistent RDMA fault. _INVALID_PEER_MIN_BACKOFF_S = 1.0 _INVALID_PEER_MAX_BACKOFF_S = 60.0 +_MAX_HYBRID_BLOCK_SPLIT = 128 def _region_sort_key(layer_name: str) -> tuple[tuple[int, int | str], ...]: @@ -2246,6 +2247,57 @@ def register_local_xfer_handler( ) return handle, blocks_data + def _split_remote_blocks(self, meta: NixlAgentMetadata) -> None: + """View remote blocks as gcd-sized kernel blocks when the local block + size is not a multiple of the remote one. + + Hybrid blocks are sized to cover the TP-sharded SSM page, so P and D + differ with TP. Backends that accept any kernel block size keep them + whole, and non-multiples cannot otherwise be paired. + """ + if self.block_size % meta.block_size == 0: + return + unit = math.gcd(self.block_size, meta.block_size) + split = meta.block_size // unit + if max(split, self.block_size // unit) > _MAX_HYBRID_BLOCK_SPLIT: + logger.warning( + "Not pairing remote %d-token blocks with local %d-token blocks: " + "their %d-token common unit would split each block into more " + "than %d transfers.", + meta.block_size, + self.block_size, + unit, + _MAX_HYBRID_BLOCK_SPLIT, + ) + return + if ( + meta.block_strides != meta.block_lens + or meta.ple_block_len is not None + or any(block_len % split for block_len in meta.block_lens) + ): + logger.warning( + "Not pairing remote %d-token blocks with local %d-token blocks: " + "the remote KV layout cannot be split into %d-token units.", + meta.block_size, + self.block_size, + unit, + ) + return + logger.info_once( + "Transferring remote %d-token blocks as %d-token kernel blocks to " + "pair them with local %d-token blocks.", + meta.block_size, + unit, + self.block_size, + ) + meta.block_size //= split + meta.physical_blocks_per_logical_kv_block *= split + meta.num_blocks *= split + meta.block_lens = [block_len // split for block_len in meta.block_lens] + meta.block_strides = list(meta.block_lens) + if meta.region_num_blocks is not None: + meta.region_num_blocks = [n * split for n in meta.region_num_blocks] + def add_remote_agent( self, nixl_agent_meta: NixlAgentMetadata, @@ -2355,6 +2407,9 @@ def add_remote_agent( nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] + # The push WRITE path is not validated with split remote blocks. + if self._has_mamba and self.use_mla and self._TRANSFER_MODE == "pull": + self._split_remote_blocks(nixl_agent_meta) if ( self._packed_member_layouts From 52e88e25ec7476a85a964f0e58cf3e8dbeb4977c Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Sun, 18 Oct 2026 17:22:13 -0700 Subject: [PATCH 832/952] [Core] Remove the obsolete VLLM_USE_RAY_V2_EXECUTOR_BACKEND env var (#62228) Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 (1M context) --- vllm/envs.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 0492c5a59d94..f11d4faa8d90 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -169,7 +169,6 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_MLA_DISABLE: bool = False - VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = True VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" VLLM_CUDART_SO_PATH: str | None = None @@ -1480,12 +1479,6 @@ def _resolve_rust_cli_path() -> str | None: # which indices are used for the Ray bundle, for every worker. # Format: comma-separated list of integers, e.g. "0,1,2,3" "VLLM_RAY_BUNDLE_INDICES": lambda: os.getenv("VLLM_RAY_BUNDLE_INDICES", ""), - # Deprecated and ignored: RayExecutorV2 is the only Ray executor. Kept so - # out-of-tree platforms that still read it keep working; will be removed - # in v0.34. - "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool( - int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "1")) - ), # In some system, find_loaded_library() may not work. So we allow users to # specify the path through environment variable VLLM_CUDART_SO_PATH. "VLLM_CUDART_SO_PATH": lambda: os.getenv("VLLM_CUDART_SO_PATH", None), From 8e5cae09f16d779a43524fd2af1142bdce2a592c Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 19:29:16 -0500 Subject: [PATCH 833/952] [ROCm][CI] Pin Inductor pointwise configs in the MoRIIO dummy-weight P/D test (#62280) Signed-off-by: Andreas Karatzas Co-authored-by: Claude Co-authored-by: vllm-agent --- .../test_write_dummy_e2e_rocm.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py index 74c9a956980a..f208f93f363e 100644 --- a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py +++ b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py @@ -26,9 +26,9 @@ a whole number of chunks plus one token. With the split matched, a correct transfer reproduces the direct serve's first-token log-probabilities exactly, so the comparison needs no noise allowance. The -servers compile with Inductor's deterministic mode, since otherwise each one -benchmarks its own reduction launch configs and separately started servers -can drift apart by a few bf16 logit steps. Guards: +servers compile with Inductor's deterministic mode and without pointwise +autotuning, since otherwise each one benchmarks its own launch configs and +separately started servers can drift apart by a few bf16 logit steps. Guards: * decode must report an external prefix-cache hit for every transferred token of every prompt, so the match is not a local recompute; @@ -303,12 +303,17 @@ def id(self) -> str: def _graph_args(mode: str) -> list[str]: if mode == "NONE": return ["--enforce-eager"] - # Inductor otherwise picks reduction launch configs by timing them, so each - # separately compiled server may sum in a different order and drift by a - # few bf16 logit steps; deterministic mode fixes the choice. + # Inductor otherwise picks launch configs by timing them, so separately + # compiled servers may drift by a few bf16 logit steps: reductions sum in a + # different order, and on ROCm the pointwise block size changes which + # multiply-adds are fused into FMAs. Deterministic mode fixes only the + # reduction configs, so pointwise autotuning is off as well. config = { "cudagraph_mode": mode, - "inductor_compile_config": {"deterministic": True}, + "inductor_compile_config": { + "deterministic": True, + "triton.autotune_pointwise": False, + }, } return ["--compilation-config", json.dumps(config)] From 0ee62b99b663b76ebd4fb5a4ef63740932cf8b5e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 19:31:13 -0500 Subject: [PATCH 834/952] [ROCm][CI] Contain the MI355 DP + EP exit-time abort (#62286) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- .buildkite/test_areas/distributed.yaml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 5a982efc3ce5..4f885974b0ef 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -425,7 +425,6 @@ steps: amd: label: ":amd: (MI355) Distributed DP + EP" dind: false - optional: false device: mi355_8 timeout_in_minutes: 50 depends_on: @@ -439,6 +438,14 @@ steps: - vllm/v1/executor/uniproc_executor.py - vllm/v1/worker/gpu_worker.py - vllm/platforms/rocm.py + commands: + - export NCCL_CUMEM_HOST_ENABLE=0 + # TODO(akaratza): drop NCCL_RAS_ENABLE once RCCL's RAS thread no longer + # aborts at exit: https://github.com/ROCm/rocm-systems/issues/13572 + - export NCCL_RAS_ENABLE=0 + # Logs each process group's teardown, to locate the exit-time abort. + - export TORCH_CPP_LOG_LEVEL=INFO + - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep - label: ":nvidia: (A100) Distributed" key: distributed-tests-4xa100 From 689effee6d07bfeb0afe42d44046141657238c23 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 19:39:14 -0500 Subject: [PATCH 835/952] [Bugfix][Attention] Fix Triton MLA decode int32 overflow and logits workspace accounting (#61938) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- .../attention/test_triton_decode_attention.py | 70 ++++++++++++++++++ .../layers/attention/mla_attention.py | 1 + vllm/v1/attention/backend.py | 5 ++ vllm/v1/attention/backends/mla/triton_mla.py | 72 +++++++++++-------- .../attention/ops/triton_decode_attention.py | 12 ++-- 5 files changed, 126 insertions(+), 34 deletions(-) diff --git a/tests/kernels/attention/test_triton_decode_attention.py b/tests/kernels/attention/test_triton_decode_attention.py index 5ac1d6cdcbf8..4de0436e0e16 100644 --- a/tests/kernels/attention/test_triton_decode_attention.py +++ b/tests/kernels/attention/test_triton_decode_attention.py @@ -325,3 +325,73 @@ def run(k_buffer, v_buffer): # Same data and same compute order; only addressing differs. assert torch.equal(o_ref, o_xl) assert torch.equal(lse_ref, lse_xl) + + +@pytest.mark.parametrize( + "B,H_Q,H_KV,D_QK,D_V,is_mla", + [ + (256, 32, 1, 576, 512, True), # grouped kernel + (1024, 32, 32, 128, 128, False), # normal kernel + ], +) +def test_decode_attention_int64_logits_offset(B, H_Q, H_KV, D_QK, D_V, is_mla): + """attn_logits offsets must not overflow int32 when B * H_Q * num_kv_splits + * (D_V + 1) exceeds 2**31 (long max_model_len, large decode batch).""" + PAGE_SIZE = 16 + seq_len = 33 + num_kv_splits = 512 + dtype = torch.bfloat16 + sm_scale = 1.0 / (D_QK**0.5) + assert B * H_Q * num_kv_splits * (D_V + 1) > 2**31 + + num_pages_per_batch = cdiv(seq_len, PAGE_SIZE) + req_to_page = torch.arange( + B * num_pages_per_batch, device=DEVICE_TYPE, dtype=torch.int32 + ).view(B, num_pages_per_batch) + k_buffer = torch.randn( + B * num_pages_per_batch, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE + ) + v_buffer = ( + k_buffer[..., :D_V] + if is_mla + else torch.randn( + B * num_pages_per_batch, + PAGE_SIZE, + H_KV, + D_V, + dtype=dtype, + device=DEVICE_TYPE, + ) + ) + q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE) + b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE) + + def run(rows): + n = rows.stop - rows.start + o = torch.zeros(n, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE) + lse = torch.zeros(n, H_Q, dtype=dtype, device=DEVICE_TYPE) + attn_logits = torch.empty( + (n, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE + ) + decode_attention_fwd( + q[rows], + k_buffer, + v_buffer, + o, + lse, + req_to_page[rows], + b_seq_len[rows], + attn_logits, + num_kv_splits, + sm_scale, + PAGE_SIZE, + is_mla=is_mla, + ) + return o, lse + + o, lse = run(slice(0, B)) + o_last, lse_last = run(slice(B - 1, B)) + + # Same compute order; the last request only differs in its logits offset. + assert torch.equal(o[-1:], o_last) + assert torch.equal(lse[-1:], lse_last) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 865913b78245..07130abd7f75 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -936,6 +936,7 @@ def forward_impl( device=k_c_normed.device, dtype=k_c_normed.dtype, ) + self.impl.reserve_workspace() # The zero fill is required when used with DP + EP # to ensure all ranks within a DP group compute the diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d3c8d7ea2433..7982ae220765 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -1085,6 +1085,11 @@ def forward_mqa( """MQA-style decode forward pass.""" raise NotImplementedError + def reserve_workspace(self) -> None: + """Called in the memory profiling run to claim persistent workspace + that is otherwise reserved after profiling.""" + return + def fused_output_quant_supported(self, quant_key: "QuantKey"): """Does this attention implementation support fused output quantization. Since MLA quantization is done manually in forward_impl (common code), diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 950d798f2210..077fe7b5927e 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -6,6 +6,7 @@ import torch import vllm.envs as envs +from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( @@ -24,6 +25,7 @@ AttentionLayer, AttentionType, MultipleOf, + max_decode_query_len, ) from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd from vllm.v1.worker.workspace import ( @@ -48,6 +50,27 @@ def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int: return min(ideal_splits, max_splits) +def _reserve_attn_logits_workspace( + vllm_config: VllmConfig, max_decode_rows: int, q_num_heads: int, kv_lora_rank: int +) -> None: + """Pre-size the shared workspace for the decode split-KV attn logits. + + Reserving at the worst case (max_model_len -> max num_kv_splits, + max_decode_rows) before warmup/cudagraph capture means the per-call + ``get_simultaneous`` in ``forward_mqa`` never has to grow the buffer at + runtime (which would raise once the workspace is locked). + """ + if not is_workspace_manager_initialized(): + return + max_splits = _compute_num_kv_splits( + vllm_config.model_config.max_model_len, + current_platform.num_compute_units(), + ) + current_workspace_manager().get_simultaneous( + ((max_decode_rows, q_num_heads, max_splits, kv_lora_rank + 1), torch.float32), + ) + + class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): # forward_mqa flattens a uniform multi-token block to one decode row per # query token, so causal and non-causal blocks both take the decode path. @@ -59,38 +82,19 @@ def __init__(self, kv_cache_spec, layer_names, vllm_config, device): super().__init__(kv_cache_spec, layer_names, vllm_config, device) # DCP local sequence lengths are not advanced between draft steps. self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 - self._reserve_attn_logits_workspace() + # forward_mqa flattens each request's block to query_len decode rows, + # and query_len is bounded by the reorder threshold. DCP all-gathers + # the query heads before forward_mqa. + _reserve_attn_logits_workspace( + vllm_config, + vllm_config.scheduler_config.max_num_seqs * self.reorder_batch_threshold, + self.num_heads * self.dcp_world_size, + self.mla_dims.kv_lora_rank, + ) def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None: pass - def _reserve_attn_logits_workspace(self) -> None: - """Pre-size the shared workspace for the decode split-KV attn logits. - - Reserving at the worst case (max_model_len -> max num_kv_splits, - max_num_seqs decode tokens) before warmup/cudagraph capture means the - per-call ``get_simultaneous`` in ``forward_mqa`` never has to grow the - buffer at runtime (which would raise once the workspace is locked). - """ - if not is_workspace_manager_initialized(): - return - # forward_mqa flattens each request's block to query_len decode rows, - # and query_len is bounded by the reorder threshold. - B = ( - self.vllm_config.scheduler_config.max_num_seqs - * self.reorder_batch_threshold - ) - # DCP all-gathers the query heads before forward_mqa. - q_num_heads = self.num_heads * self.dcp_world_size - max_splits = _compute_num_kv_splits( - self.model_config.max_model_len, - current_platform.num_compute_units(), - ) - lse_dim = self.mla_dims.kv_lora_rank + 1 - current_workspace_manager().get_simultaneous( - ((B, q_num_heads, max_splits, lse_dim), torch.float32), - ) - class TritonMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -226,6 +230,18 @@ def __init__( self._sm_count = current_platform.num_compute_units() + def reserve_workspace(self) -> None: + # The builder reserves after memory profiling; reserve its upper bound + # here so the KV cache is sized around it. + vllm_config = get_current_vllm_config() + _reserve_attn_logits_workspace( + vllm_config, + vllm_config.scheduler_config.max_num_seqs + * max_decode_query_len(vllm_config), + self.num_heads * self.dcp_world_size, + self.kv_lora_rank, + ) + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py index 130760d1ded8..f59a874a8104 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -181,7 +181,7 @@ def _fwd_kernel_stage1( e_max = n_e_max offs_mid_o = ( - cur_batch * stride_mid_ob + cur_batch.to(tl.int64) * stride_mid_ob + cur_head * stride_mid_oh + split_kv_id * stride_mid_os + offs_dv @@ -194,7 +194,7 @@ def _fwd_kernel_stage1( ) offs_mid_o_1 = ( - cur_batch * stride_mid_ob + cur_batch.to(tl.int64) * stride_mid_ob + cur_head * stride_mid_oh + split_kv_id * stride_mid_os + Lv @@ -440,7 +440,7 @@ def _fwd_grouped_kernel_stage1( e_max = n_e_max offs_mid_o = ( - cur_batch * stride_mid_ob + cur_batch.to(tl.int64) * stride_mid_ob + cur_head[:, None] * stride_mid_oh + split_kv_id * stride_mid_os + offs_dv[None, :] @@ -453,7 +453,7 @@ def _fwd_grouped_kernel_stage1( ) offs_mid_o_1 = ( - cur_batch * stride_mid_ob + cur_batch.to(tl.int64) * stride_mid_ob + cur_head * stride_mid_oh + split_kv_id * stride_mid_os + Lv @@ -600,8 +600,8 @@ def _fwd_kernel_stage2( e_max = -float("inf") acc = tl.zeros([BLOCK_DV], dtype=tl.float32) - offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d - offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv + offs_v = cur_batch.to(tl.int64) * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = cur_batch.to(tl.int64) * stride_mid_ob + cur_head * stride_mid_oh + Lv for split_kv_id in range(0, NUM_KV_SPLITS): kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) From c07c59f8211acb5283324a3644250d1672ce2e1a Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 20:17:15 -0500 Subject: [PATCH 836/952] [ROCm][Bugfix] Release vLLM's communicators on external-launcher shutdown (#62241) Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Claude --- vllm/v1/executor/uniproc_executor.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py index da1f2c588b6e..229879c0866c 100644 --- a/vllm/v1/executor/uniproc_executor.py +++ b/vllm/v1/executor/uniproc_executor.py @@ -205,3 +205,15 @@ def determine_available_memory(self) -> list[int]: # in bytes memory_tensor = torch.tensor([memory], device="cpu", dtype=torch.int64) dist.all_reduce(memory_tensor, group=cpu_group, op=dist.ReduceOp.MIN) return [memory_tensor.item()] + + def shutdown(self) -> None: + super().shutdown() + # Release vLLM's communicators (TP/PP/DP/EP groups) now, as + # MultiprocExecutor's workers do. Left alive, they are torn down by + # static destructors at interpreter exit, where newer RCCL aborts the + # process ("terminate called without an active exception"). The world + # group and the default process group may belong to the launcher, so + # they stay. + from vllm.distributed.parallel_state import destroy_model_parallel + + destroy_model_parallel() From b0d61e5d26561f865bf89688eb0aea5000133a42 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sun, 18 Oct 2026 18:37:04 -0700 Subject: [PATCH 837/952] [Model] Add Reflection Beam (#62074) Signed-off-by: Woosuk Kwon Co-authored-by: Lain Co-authored-by: kaiwu-reflection Co-authored-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Cursor Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_basic.yaml | 7 +- docs/models/supported_models.md | 1 + tests/models/registry.py | 12 + tests/models/test_beam_construction.py | 212 +++++ tests/models/test_beam_numerics.py | 481 +++++++++++ tests/models/test_beam_weight_mapping.py | 206 +++++ tests/test_config.py | 5 +- tests/transformers_utils/test_beam_config.py | 83 ++ tests/v1/spec_decode/test_dflash2.py | 399 ++++++++- vllm/config/vllm.py | 9 +- vllm/model_executor/models/registry.py | 2 + vllm/models/beam/__init__.py | 8 + vllm/models/beam/nvidia/__init__.py | 3 + vllm/models/beam/nvidia/activation.py | 59 ++ vllm/models/beam/nvidia/attention.py | 219 +++++ vllm/models/beam/nvidia/attention_ops.py | 250 ++++++ .../beam/nvidia/attention_ops_cutedsl.py | 766 ++++++++++++++++++ vllm/models/beam/nvidia/dflash2.py | 623 ++++++++++++++ .../models/beam/nvidia/dflash2_context_ops.py | 127 +++ vllm/models/beam/nvidia/dflash2_conv_ops.py | 94 +++ vllm/models/beam/nvidia/dflash2_ops.py | 85 ++ vllm/models/beam/nvidia/layernorm.py | 184 +++++ vllm/models/beam/nvidia/model.py | 411 ++++++++++ vllm/models/beam/nvidia/moe.py | 103 +++ vllm/models/beam/nvidia/projection.py | 61 ++ vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + vllm/transformers_utils/configs/beam.py | 111 +++ vllm/v1/worker/gpu/spec_decode/__init__.py | 4 +- 29 files changed, 4482 insertions(+), 46 deletions(-) create mode 100644 tests/models/test_beam_construction.py create mode 100644 tests/models/test_beam_numerics.py create mode 100644 tests/models/test_beam_weight_mapping.py create mode 100644 tests/transformers_utils/test_beam_config.py create mode 100644 vllm/models/beam/__init__.py create mode 100644 vllm/models/beam/nvidia/__init__.py create mode 100644 vllm/models/beam/nvidia/activation.py create mode 100644 vllm/models/beam/nvidia/attention.py create mode 100644 vllm/models/beam/nvidia/attention_ops.py create mode 100644 vllm/models/beam/nvidia/attention_ops_cutedsl.py create mode 100644 vllm/models/beam/nvidia/dflash2.py create mode 100644 vllm/models/beam/nvidia/dflash2_context_ops.py create mode 100644 vllm/models/beam/nvidia/dflash2_conv_ops.py create mode 100644 vllm/models/beam/nvidia/dflash2_ops.py create mode 100644 vllm/models/beam/nvidia/layernorm.py create mode 100644 vllm/models/beam/nvidia/model.py create mode 100644 vllm/models/beam/nvidia/moe.py create mode 100644 vllm/models/beam/nvidia/projection.py create mode 100644 vllm/transformers_utils/configs/beam.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index e76a99d059c3..191e71b76067 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -11,9 +11,12 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/test_initialization.py - tests/models/registry.py + - tests/models/test_beam_construction.py commands: # Run a subset of model initialization tests - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset + # Exercise Beam construction and generation without a checkpoint download. + - pytest -v -s models/test_beam_construction.py mirror: amd: label: ":amd: (MI355 DPX) Basic Models (Initialization)" @@ -258,6 +261,8 @@ steps: - tests/models/test_qwen3_5_mtp_config.py - tests/models/test_alice_ai.py - tests/models/test_alice_ai_config.py + - tests/models/test_beam_weight_mapping.py + - tests/models/test_beam_numerics.py - tests/models/transformers/fusers/ - tests/models/transformers/test_dynamic_rope.py - tests/model_executor/layers/test_activation.py @@ -268,7 +273,7 @@ steps: - tests/models/transformers/test_layer_registry.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0b69ac8f4753..2166c3a8eb08 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -379,6 +379,7 @@ th { | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ | | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ | | `BailingMoeV3ForCausalLM` | Ling | `inclusionAI/Ling-3.0-flash` | | ✅︎ | +| `BeamForCausalLM` | Beam | `reflection/Beam-1.0-501B-A23B-NVFP4`, etc. | | ✅︎ | | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ | | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index f9dfafb40479..d27aeba00f70 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -216,6 +216,12 @@ def check_available_online( "BailingMoeV3ForCausalLM": _HfExamplesInfo( "inclusionAI/Ling-3.0-flash", trust_remote_code=True ), + "BeamForCausalLM": _HfExamplesInfo( + "reflection/Beam-1.0-501B-A23B-NVFP4", + is_available_online=False, + # Include MoE in the reduced one-layer initialization test. + hf_overrides={"mlp_layer_types": ["sparse"]}, + ), "BloomForCausalLM": _HfExamplesInfo( "bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"} ), @@ -1490,6 +1496,12 @@ def check_available_online( max_model_len=8192, max_num_seqs=32, ), + "DFlash2BeamDraftModel": _HfExamplesInfo( + "reflection/Beam-1.0-501B-A23B-NVFP4", + speculative_model="reflection/Beam-1.0-501B-A23B-DFlash2", + is_available_online=False, + use_original_num_layers=True, + ), "DFlashLagunaForCausalLM": _HfExamplesInfo( "poolside/Laguna-XS-2.1-NVFP4", speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4", diff --git a/tests/models/test_beam_construction.py b/tests/models/test_beam_construction.py new file mode 100644 index 000000000000..1010e4bc59f6 --- /dev/null +++ b/tests/models/test_beam_construction.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Construct Beam from a synthetic config and generate with dummy weights. + +Exercises NoPE + SSMax, hybrid sliding-window attention, the sigmoid output +gate and its per-head bias, sandwich norm, depth-scaled residuals, grouped +sigmoid routing, and logit softcapping without downloading a checkpoint. +""" + +import json +from operator import attrgetter +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm import LLM, SamplingParams +from vllm.models.beam.nvidia import model as beam +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.beam import BeamConfig +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + set_eagle3_aux_hidden_state_layers, +) + +# `moe_intermediate_size` must be a multiple of 128 for the BF16 MoE kernel. +TINY_CONFIG = { + "architectures": ["BeamForCausalLM"], + "model_type": "beam", + "vocab_size": 1024, + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "head_dim": 16, + "max_position_embeddings": 512, + "rms_norm_eps": 1e-6, + "sliding_window": 128, + "moe_intermediate_size": 256, + "num_experts_per_tok": 2, + "num_experts": 8, + "n_group": 2, + "topk_group": 1, + "routed_scaling_factor": 1.0, + "norm_topk_prob": True, + # Layer 1 is NoPE, so it takes the SSMax path and full attention while + # layer 0 stays RoPE + sliding. Exercises both branches in one model. + "layer_types": ["sliding_attention", "full_attention"], + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + "mlp_layer_types": ["dense", "sparse"], + "final_logit_softcapping": 30.0, + "torch_dtype": "bfloat16", + "bos_token_id": 1, + "eos_token_id": 2, +} + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Beam targets CUDA") +@pytest.mark.usefixtures("enable_pickle") +@pytest.mark.parametrize( + "enforce_eager,tp_size,num_experts,enable_ep", + [ + (True, 1, 8, False), + (False, 1, 8, False), + (True, 4, 2, False), + (True, 4, 8, True), + ], +) +def test_beam_builds_and_generates( + tmp_path, enforce_eager, tp_size, num_experts, enable_ep +): + """Release-style metadata must load and generate in both execution modes.""" + if torch.accelerator.device_count() < tp_size: + pytest.skip(f"Requires {tp_size} GPUs") + config = { + **TINY_CONFIG, + "num_experts": num_experts, + "moe_intermediate_size": max(256, 128 * tp_size), + "n_group": 1 if num_experts == 2 else 2, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + + llm = LLM( + model=str(tmp_path), + load_format="dummy", + skip_tokenizer_init=True, + max_model_len=256, + max_num_seqs=1, + max_num_batched_tokens=256, + kv_cache_memory_bytes=16 * 1024**2, + gpu_memory_utilization=0.25, + enforce_eager=enforce_eager, + tensor_parallel_size=tp_size, + enable_expert_parallel=enable_ep, + ) + outputs = llm.generate( + {"prompt_token_ids": [5, 6, 7, 8]}, + SamplingParams(max_tokens=8, ignore_eos=True), + ) + + assert len(outputs[0].outputs[0].token_ids) == 8 + # TP shards every expert; only EP partitions the expert set across ranks. + local_experts = num_experts // tp_size if enable_ep else num_experts + assert llm.apply_model(attrgetter("num_local_physical_experts")) == ( + [local_experts] * tp_size + ) + + +def test_eplb_is_rejected_until_validated(monkeypatch): + """Expert parallelism is validated; expert rebalancing is not.""" + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=True, use_sequence_parallel_moe=False + ), + ) + with pytest.raises(NotImplementedError, match="EPLB"): + beam.BeamForCausalLM(vllm_config=config) + + +def test_sequence_parallel_moe_is_rejected_until_validated(): + """DP attention with EP MoE takes the sequence-parallel path; unvalidated.""" + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=False, use_sequence_parallel_moe=True + ), + ) + with pytest.raises(NotImplementedError, match="sequence-parallel"): + beam.BeamForCausalLM(vllm_config=config) + + +def test_dense_only_pipeline_stage_initializes(monkeypatch): + """A stage with dense layers and remote MoE layers needs no local experts.""" + layer = beam.BeamDecoderLayer.__new__(beam.BeamDecoderLayer) + nn.Module.__init__(layer) + layer.mlp = nn.Identity() + stage = nn.Module() + stage.layers = nn.ModuleList([layer, beam.PPMissingLayer()]) + stage.make_empty_intermediate_tensors = lambda *args, **kwargs: None + monkeypatch.setattr(beam, "BeamModel", lambda **kwargs: stage) + monkeypatch.setattr( + beam, "get_pp_group", lambda: SimpleNamespace(is_last_rank=False) + ) + monkeypatch.setattr(beam, "LogitsProcessor", lambda *args: nn.Identity()) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=BeamConfig()), + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=False, use_sequence_parallel_moe=False + ), + ) + + model = beam.BeamForCausalLM(vllm_config=config) + + assert model.moe_layers == [] + assert model.num_moe_layers == 0 + assert model.num_logical_experts == 0 + assert model.num_local_physical_experts == 0 + + +class _AddOneLayer(nn.Module): + """Adds one to the residual stream and hands the next layer double that.""" + + def __init__(self): + super().__init__() + self.input_layernorm = nn.Identity() + + def forward(self, positions, hidden_states, normalized_states, next_norm_weight): + hidden_states = hidden_states + 1 + return hidden_states, hidden_states * 2 + + +@pytest.mark.cpu_test +def test_dflash_target_layers_capture_layer_outputs(monkeypatch): + """DFlash `target_layer_ids` name the layer whose output is captured.""" + model = beam.BeamModel.__new__(beam.BeamModel) + nn.Module.__init__(model) + model.start_layer, model.end_layer = 0, 3 + model.layers = nn.ModuleList([_AddOneLayer() for _ in range(3)]) + model.norm = nn.Identity() + target = beam.BeamForCausalLM.__new__(beam.BeamForCausalLM) + nn.Module.__init__(target) + target.model = model + monkeypatch.setattr( + beam, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + spec_config = SimpleNamespace( + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(dflash_config={"target_layer_ids": [0, 1]}) + ) + ) + + set_eagle3_aux_hidden_state_layers(target, spec_config) + output, aux_hidden_states = target( + input_ids=None, + positions=torch.tensor([0]), + inputs_embeds=torch.tensor([[1.0, 2.0]]), + ) + + torch.testing.assert_close(output, torch.tensor([[8.0, 10.0]])) + assert len(aux_hidden_states) == 2 + torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[2.0, 3.0]])) + torch.testing.assert_close(aux_hidden_states[1], torch.tensor([[3.0, 4.0]])) diff --git a/tests/models/test_beam_numerics.py b/tests/models/test_beam_numerics.py new file mode 100644 index 000000000000..22fa15ef6d92 --- /dev/null +++ b/tests/models/test_beam_numerics.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regressions for Beam normalization and activation rounding.""" + +import pytest +import torch + + +def _torch_rms_norm(x, weight, eps): + y = x.float() + y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + eps) * weight.float() + return y.to(x.dtype) + + +def _torch_boundary(x, residual, post_weight, post_eps, next_weight, next_eps, scales): + x = _torch_rms_norm(x, post_weight, post_eps) * scales[0] + hidden = x + residual.to(x.dtype) * scales[1] + normalized = ( + _torch_rms_norm(hidden, next_weight, next_eps) if next_eps is not None else None + ) + return hidden, normalized + + +def _assert_boundary_close(actual, expected, *, relative_tolerance=1e-4): + for value, reference in zip(actual, expected): + if reference is None: + assert value is None + continue + bf16 = reference.dtype == torch.bfloat16 + torch.testing.assert_close( + value, + reference, + rtol=0.016 if bf16 else 2e-5, + atol=0.016 if bf16 else 2e-6, + ) + # Allow isolated BF16 rounding differences, but catch omitted casts. + if reference.numel(): + error = (value.float() - reference.float()).norm() + assert error <= reference.float().norm() * ( + relative_tolerance if bf16 else 2e-6 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 8, 128]) +@pytest.mark.parametrize("cap", [7.3, 30.0]) +@torch.inference_mode() +def test_softcap_rounds_into_original_logits_without_touching_padding(dtype, rows, cap): + from vllm.models.beam.nvidia.activation import beam_softcap_logits + + torch.manual_seed(42) + storage = torch.randn(rows, 264, device="cuda", dtype=dtype) * (cap * 3) + logits = storage[:, :257] + if rows: + logits[0, :6] = torch.tensor( + [0.0, -0.0, float("inf"), -float("inf"), float("nan"), 1e-38], + device=logits.device, + dtype=dtype, + ) + original = storage.clone() + expected = (torch.tanh(logits.float() / cap) * cap).to(dtype) + actual = beam_softcap_logits(logits, cap) + + assert actual is logits + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + torch.testing.assert_close( + storage[:, 257:], original[:, 257:], rtol=0, atol=0, equal_nan=True + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "rows,width", [(0, 2560), (1, 5120), (8, 2560), (33, 513), (2048, 2560)] +) +@torch.inference_mode() +def test_dense_activation_rounds_after_product(default_vllm_config, dtype, rows, width): + """BF16 I/O must not round the SiLU result before multiplication.""" + import vllm._custom_ops # noqa: F401 + from vllm.models.beam.nvidia.activation import beam_silu_mul + + torch.manual_seed(42) + x = torch.randn(rows, 2 * width, device="cuda", dtype=dtype) * 4 + expected = torch.empty(rows, width, device="cuda", dtype=torch.float32) + if rows: + torch.ops._C.silu_and_mul(expected, x.float()) + actual = beam_silu_mul(x) + torch.testing.assert_close( + actual, + expected.to(dtype), + rtol=0.008 if dtype == torch.bfloat16 else 1e-6, + atol=1e-7, + ) + if rows: + error = (actual.float() - expected.to(dtype).float()).norm() + assert error <= expected.norm() * 1e-5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(33, 513), (17, 6, 128)]) +@torch.inference_mode() +def test_fused_norm_preserves_fp32_affine_for_strided_rows(dtype, shape): + """QKV slices and near-epsilon values retain the reference's FP32 affine.""" + from vllm.models.beam.nvidia.layernorm import beam_rms_norm + + generator = torch.Generator(device="cuda").manual_seed(42) + storage = torch.randn( + (*shape[:-1], shape[-1] + 64), + device="cuda", + dtype=dtype, + generator=generator, + ) + x = storage[..., 1 : shape[-1] + 1] + x[::2].mul_(1e-3) + weight = torch.randn(shape[-1], device="cuda", dtype=dtype, generator=generator) + expected = _torch_rms_norm(x, weight, 1e-6) + actual = beam_rms_norm(x, weight, 1e-6) + # FP32 reduction trees may differ; allow one BF16 ULP after final rounding. + torch.testing.assert_close( + actual, + expected, + rtol=0.008 if dtype == torch.bfloat16 else 1e-6, + atol=1e-7, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("shape", [(0, 5120), (33, 513)]) +@pytest.mark.parametrize("next_", [False, True]) +@torch.inference_mode() +def test_norm_residual_norm_preserves_intermediate_rounding(dtype, shape, next_): + """Both residual and next input must retain BF16 casts between fused stages.""" + from vllm.models.beam.nvidia.layernorm import beam_norm_residual_norm + + torch.manual_seed(42) + x = torch.randn(shape[0], shape[1] + 32, device="cuda", dtype=dtype)[ + :, 1 : shape[1] + 1 + ] + residual = torch.randn_like(x) + x[::2].mul_(1e-3) + post_weight = torch.randn(shape[1], device="cuda", dtype=dtype) + next_weight = torch.randn(shape[1], device="cuda", dtype=dtype) if next_ else None + scales = (0.125, 0.9682458365518543) + original_x, original_residual = x.clone(), residual.clone() + args = ( + x, + residual, + post_weight, + 1e-6, + next_weight, + 1e-6 if next_ else None, + scales, + ) + expected = _torch_boundary(*args) + actual = beam_norm_residual_norm(*args) + _assert_boundary_close(actual, expected) + torch.testing.assert_close(x, original_x, rtol=0, atol=0) + torch.testing.assert_close(residual, original_residual, rtol=0, atol=0) + + +def _separate_qk_norm_rope( + q, + k, + qw, + kw, + eps, + positions, + cache, + head_dim, + rotary_dim=0, + neox=True, + ssmax=False, +): + from vllm import _custom_ops as ops + + q = _torch_rms_norm( + q.reshape(-1, q.shape[-1] // head_dim, head_dim), qw, eps + ).reshape_as(q) + k = _torch_rms_norm( + k.reshape(-1, k.shape[-1] // head_dim, head_dim), kw, eps + ).reshape_as(k) + if rotary_dim and q.shape[0]: + qf, kf = q.float(), k.float() + ops.rotary_embedding(positions, qf, kf, head_dim, cache, neox) + q, k = qf.to(q.dtype), kf.to(k.dtype) + elif ssmax: + scale = (0.1 * torch.log1p((positions.float() + 1) / 2048)).exp() + q = (q * scale[:, None]).to(q.dtype) + return q, k + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 17, 512, 2049]) +@pytest.mark.parametrize("packed", [False, True]) +@torch.inference_mode() +def test_output_gate_preserves_activation_dtype_rounding(dtype, rows, packed): + from vllm.models.beam.nvidia.attention_ops import beam_output_gate + + torch.manual_seed(42) + x = torch.randn(rows, 768, device="cuda", dtype=dtype) + gate = torch.randn(rows, 1792 if packed else 768, device="cuda", dtype=dtype) * 8 + gate = gate[:, -768:] + weight = torch.randn(6, device="cuda", dtype=dtype) + # Preserve each activation-dtype operation in the parent implementation. + g = gate.view(rows, 6, 128) + weight[None, :, None] + expected = x * torch.sigmoid(g.reshape_as(gate)) + actual = beam_output_gate(x, gate, weight, num_heads=6, head_dim=128) + torch.testing.assert_close(actual, expected, rtol=2e-6, atol=1e-7) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rotary_dim,neox", [(0, True), (128, True), (64, False)]) +@pytest.mark.parametrize("head_dim,q_heads,k_heads", [(128, 6, 1), (80, 4, 2)]) +@torch.inference_mode() +def test_qk_norm_rope_keeps_norm_rounding_and_fp32_rotation( + dtype, rotary_dim, neox, head_dim, q_heads, k_heads +): + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + rows = 17 + rotary_dim = min(rotary_dim, head_dim) + sizes = (q_heads * head_dim, k_heads * head_dim, k_heads * head_dim) + qkv = torch.randn(rows, sum(sizes) + 1, device="cuda", dtype=dtype)[:, 1:] + q, k, _ = qkv.split(sizes, -1) + q[::2].mul_(1e-3) + qw = torch.randn(head_dim, device="cuda", dtype=dtype) + kw = torch.randn(head_dim, device="cuda", dtype=dtype) + positions = torch.arange(rows, device="cuda") * 3 + angle = torch.randn(max(1, rows * 3), rotary_dim // 2, device="cuda") + cache = torch.cat((angle.cos(), angle.sin()), -1) + args = (q, k, qw, kw, 1e-6, positions, cache, head_dim, rotary_dim, neox) + expected = _separate_qk_norm_rope(*args) + actual = beam_qk_norm_rope(*args) + for value, reference in zip(actual, expected): + torch.testing.assert_close( + value, reference, rtol=0.016 if dtype == torch.bfloat16 else 2e-5, atol=2e-6 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 8, 513, 2049]) +@pytest.mark.parametrize("kv_heads", [1, 2]) +@pytest.mark.parametrize("nhd", [False, True]) +@pytest.mark.parametrize("rotary_dim,neox", [(128, True), (64, False), (0, True)]) +@torch.inference_mode() +def test_qk_norm_rope_inserts_only_mapped_kv(rows, kv_heads, nhd, rotary_dim, neox): + """Cache insertion preserves Q/K and follows live slot mappings on replay.""" + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + sizes = (6 * kv_heads * 128, kv_heads * 128, kv_heads * 128) + qkv = torch.randn(rows, sum(sizes) + 1, device="cuda", dtype=torch.bfloat16)[:, 1:] + q, k, v = qkv.split(sizes, -1) + weight = torch.randn(128, device="cuda") + positions = torch.arange(rows, device="cuda") + rope = torch.randn(rows, rotary_dim, device="cuda") if rotary_dim else None + args = (q, k, weight, weight, 1e-6, positions, rope, 128, rotary_dim, neox) + pages = (rows + 127) // 128 + 1 + shape = (pages, 128, kv_heads, 256) if nhd else (pages, kv_heads, 128, 256) + cache = torch.zeros(shape, device="cuda", dtype=q.dtype) + if nhd: + cache = cache.transpose(1, 2) + live = max(1, rows - 1) + slots = torch.randperm(pages * 128, device="cuda")[:live].clone() + if live > 1: + slots[-1] = -1 + + def fused(): + return beam_qk_norm_rope( + *args, ssmax=True, value=v, kv_cache=cache, slot_mapping=slots + ) + + fused() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fused() + # Mutate both data and metadata after capture, including a previously skipped slot. + qkv.add_(0.125) + slots.copy_(torch.randperm(pages * 128, device="cuda")[:live]) + if live > 1: + slots[0] = -1 + cache.zero_() + graph.replay() + expected = beam_qk_norm_rope(*args, ssmax=True) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, rtol=0, atol=0) + valid = slots >= 0 + mapped = slots[valid] + reference_cache = torch.zeros_like(cache) + reference_cache[mapped // 128, :, mapped % 128, :128] = expected[1][:live][ + valid + ].reshape(-1, kv_heads, 128) + reference_cache[mapped // 128, :, mapped % 128, 128:] = v[:live][valid].reshape( + -1, kv_heads, 128 + ) + torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@torch.inference_mode() +def test_ssmax_fuses_after_norm_rounding_and_scales_only_queries(dtype): + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + positions = torch.tensor( + [0, 1, 15, 2047, 2048, 16384, 100000, 1048575], device="cuda" + ) + qkv = torch.randn(len(positions), 1025, device="cuda", dtype=dtype)[:, 1:] + q, k, _ = qkv.split((768, 128, 128), -1) + qw, kw = [torch.randn(128, device="cuda", dtype=dtype) for _ in range(2)] + args = (q, k, qw, kw, 1e-6, positions, None, 128, 0, True) + actual = beam_qk_norm_rope(*args, ssmax=True) + expected = _separate_qk_norm_rope(*args, ssmax=True) + for value, reference in zip(actual, expected): + # Norm and position scaling each round to BF16; allow two rounding steps. + torch.testing.assert_close( + value, + reference, + rtol=0.016 if dtype == torch.bfloat16 else 2e-6, + atol=2e-6, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rotary_dim", [0, 128]) +@torch.inference_mode() +def test_attention_fusions_replay_with_changed_positions_and_gates(rotary_dim): + from vllm.models.beam.nvidia.attention_ops import ( + beam_output_gate, + beam_qk_norm_rope, + ) + + torch.manual_seed(42) + qkv = torch.randn(17, 1025, device="cuda", dtype=torch.bfloat16)[:, 1:] + q, k, _ = qkv.split((768, 128, 128), -1) + positions = torch.arange(17, device="cuda") + angle = torch.randn(128, 64, device="cuda") + cache = torch.cat((angle.cos(), angle.sin()), -1) + gate = torch.randn_like(q) + weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(6, device="cuda", dtype=torch.bfloat16) + + def run(): + nq, nk = beam_qk_norm_rope( + q, + k, + weight, + weight, + 1e-6, + positions, + cache, + 128, + rotary_dim, + ssmax=not rotary_dim, + ) + return beam_output_gate(nq, gate, bias, num_heads=6, head_dim=128), nk + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = run() + graph.replay() + torch.accelerator.synchronize() + positions.add_(31) + gate.neg_() + qkv.mul_(0.5) + expected = run() + graph.replay() + torch.accelerator.synchronize() + for actual, reference in zip(captured, expected): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.accelerator.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rotary_dim", [0, 128]) +@pytest.mark.parametrize("insert_kv", [False, True]) +@torch.inference_mode() +def test_fused_kernels_reuse_specializations_across_token_counts( + monkeypatch, rotary_dim, insert_kv +): + """Warm tile buckets once, then vary batch sizes, positions offsets and scales.""" + from triton import knobs + + from vllm.models.beam.nvidia.activation import beam_softcap_logits + from vllm.models.beam.nvidia.attention_ops import ( + beam_output_gate, + beam_qk_norm_rope, + ) + from vllm.models.beam.nvidia.layernorm import beam_norm_residual_norm, beam_rms_norm + + device = "cuda" + dtype = torch.bfloat16 + rows = 8193 + qkv = torch.randn(rows, 1024, device=device, dtype=dtype) + kv_cache = torch.zeros((65, 1, 128, 256), device=device, dtype=dtype) + slots = torch.arange(rows, device=device) + positions = torch.arange(rows + 16, device=device) + cache = torch.randn(rows + 16, 128, device=device) + weight = torch.ones(128, device=device, dtype=dtype) + bias = torch.zeros(6, device=device, dtype=dtype) + x = torch.randn(rows, 5120, device=device, dtype=dtype) + residual = torch.randn_like(x) + norm_weight = torch.ones(5120, device=device, dtype=dtype) + + def run(tokens, offset, scales): + q, k = beam_qk_norm_rope( + qkv[:tokens, :768], + qkv[:tokens, 768:896], + weight, + weight, + 1e-6, + positions[offset : offset + tokens], + cache if rotary_dim else None, + 128, + rotary_dim, + ssmax=not rotary_dim, + value=qkv[:tokens, 896:] if insert_kv else None, + kv_cache=kv_cache if insert_kv else None, + slot_mapping=slots[: tokens - 1] if insert_kv else None, + ) + beam_output_gate(q, q, bias, num_heads=6, head_dim=128) + beam_rms_norm(x[:tokens], norm_weight, 1e-6) + beam_softcap_logits(x[:tokens], 30.0) + beam_norm_residual_norm( + x[:tokens], + residual[:tokens], + norm_weight, + 1e-6, + norm_weight, + 1e-6, + scales, + ) + + for tokens in (1, 512, 2049): + run(tokens, 0, (0.125, 0.9682458365518543)) + + def unexpected_compile(**kwargs): + if kwargs["fn"].module.startswith("vllm.models.beam.nvidia"): + pytest.fail(f"Unexpected shape specialization: {kwargs['fn'].name}") + + with monkeypatch.context() as patch: + patch.setattr(knobs.runtime, "jit_cache_hook", unexpected_compile) + token_counts = (2, 7, 17, 511, 513, 2048, 2050, 8193) + for i, tokens in enumerate(token_counts): + run(tokens, i % 16, (0.125 + i / 1000, 0.968 - i / 1000)) + torch.accelerator.synchronize() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 2, 4, 8, 12, 16, 24, 64]) +@pytest.mark.parametrize("qkv_rows,gate_rows", [(2048, 1536), (1024, 768)]) +@torch.inference_mode() +def test_packed_attention_projection(rows, qkv_rows, gate_rows): + """The shared GEMMs preserve QKV/gate ordering and BF16 output accuracy.""" + from vllm.models.beam.nvidia.projection import ( + packed_projection, + use_low_latency_gemm, + ) + + torch.manual_seed(42) + x = torch.randn(rows, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(qkv_rows + gate_rows, 5120, device="cuda", dtype=x.dtype) * 0.01 + packed = packed_projection(x, w, use_low_latency_gemm(w)) + for actual, weight in zip( + packed.split([qkv_rows, gate_rows], -1), + w.split([qkv_rows, gate_rows]), + ): + expected = (x.double() @ weight.double().T).to(x.dtype) + torch.testing.assert_close(actual, expected, rtol=0.008, atol=1e-5) diff --git a/tests/models/test_beam_weight_mapping.py b/tests/models/test_beam_weight_mapping.py new file mode 100644 index 000000000000..2d6a60a1d6b5 --- /dev/null +++ b/tests/models/test_beam_weight_mapping.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""`BeamModel.hf_to_vllm_mapper` must fuse only the dense projections. + +Beam's checkpoints carry two `gate_proj` names that are *not* the SwiGLU +gate -- the per-expert `mlp.experts.N.gate_proj` (owned by +`MoERunner.load_weights`) and the attention output gate `self_attn.gate_proj`. +Folding either into `mlp.gate_up_proj` silently corrupts the weights, so pin the +mapping. The attention gate is additionally renamed to `output_gate_proj` so +that `packed_modules_mapping` consumers (LoRA) cannot match it as a +`gate_up_proj` shard either. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.beam.nvidia import model as beam +from vllm.models.beam.nvidia.model import BeamModel +from vllm.transformers_utils.configs.beam import BeamConfig + +MAPPER = BeamModel.hf_to_vllm_mapper + + +def _map(name: str) -> tuple[str, int | str | None]: + weight = torch.empty(0) + ((mapped, out),) = MAPPER.apply([(name, weight)]) + return mapped, getattr(out, "shard_id", None) + + +@pytest.mark.parametrize( + ("checkpoint_name", "expected_name", "expected_shard"), + [ + ( + "model.layers.3.self_attn.q_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "q", + ), + ( + "model.layers.3.self_attn.k_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "k", + ), + ( + "model.layers.3.self_attn.v_proj.weight", + "model.layers.3.self_attn.qkv_proj.weight", + "v", + ), + ( + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 0, + ), + ( + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 1, + ), + ], +) +def test_dense_projections_are_fused(checkpoint_name, expected_name, expected_shard): + assert _map(checkpoint_name) == (expected_name, expected_shard) + + +@pytest.mark.parametrize( + "checkpoint_name", + [ + # Owned by MoERunner.load_weights, which is also what keeps expert + # weights correct under EP and redundant experts. + "model.layers.1.mlp.experts.7.gate_proj.weight", + "model.layers.1.mlp.experts.7.up_proj.weight", + "model.layers.1.mlp.experts.7.down_proj.weight", + "model.layers.1.self_attn.gate_bias", + # The MoE router. + "model.layers.1.mlp.gate.weight", + "model.layers.1.mlp.gate.e_score_correction_bias", + ], +) +def test_gates_and_experts_pass_through(checkpoint_name): + assert _map(checkpoint_name) == (checkpoint_name, None) + + +def test_attention_output_gate_is_renamed_not_fused(): + """The attention gate must be renamed, and must not pick up a shard id.""" + assert _map("model.layers.1.self_attn.gate_proj.weight") == ( + "model.layers.1.self_attn.output_gate_proj.weight", + None, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@torch.inference_mode() +def test_checkpoint_norm_names_match_sublayer_order(default_vllm_config, monkeypatch): + """Load reference names directly and use each weight on its intended branch.""" + + class Attention(nn.Module): + def forward(self, positions, hidden_states): + return hidden_states + + monkeypatch.setattr(beam, "BeamAttention", lambda *args, **kwargs: Attention()) + monkeypatch.setattr(beam, "BeamMLP", lambda **kwargs: nn.Identity()) + config = BeamConfig(hidden_size=8, num_hidden_layers=1, mlp_layer_types=["dense"]) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + cache_config=None, + quant_config=None, + ) + model = BeamModel.__new__(BeamModel) + nn.Module.__init__(model) + model.config = config + layer = beam.BeamDecoderLayer(vllm_config, prefix="layers.0") + model.layers = nn.ModuleList([layer]) + model.cuda() + names = [ + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", + ] + weights = { + f"layers.0.{name}.weight": torch.arange(1, 9, device="cuda").float().roll(i) / 8 + for i, name in enumerate(names) + } + assert model.load_weights(weights.items()) == set(weights) + x = torch.arange(1, 17, device="cuda").float().view(2, 8) / 16 + + def norm(value, name): + return ( + value + * torch.rsqrt(value.square().mean(-1, keepdim=True) + 1e-6) + * weights[f"layers.0.{name}.weight"] + ) + + x_scale, r_scale = beam._depth_scaled_coeffs(0, is_attention=True) + expected = x * r_scale + norm(norm(x, names[0]), names[1]) * x_scale + x_scale, r_scale = beam._depth_scaled_coeffs(0, is_attention=False) + expected = expected * r_scale + norm(norm(expected, names[2]), names[3]) * x_scale + actual, _ = layer(torch.arange(2, device="cuda"), x, None) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tp_size", [4, 8]) +@torch.inference_mode() +def test_checkpoint_loads_into_merged_qkv_gate( + default_vllm_config, monkeypatch, tp_size +): + """All four checkpoint projections load directly into the allocated TP shard.""" + from vllm.models.beam.nvidia import attention + from vllm.models.beam.nvidia.projection import packed_projection + + for module in ( + "vllm.model_executor.layers.linear", + "vllm.model_executor.parameter", + ): + monkeypatch.setattr( + f"{module}.get_tensor_model_parallel_rank", lambda: tp_size - 1 + ) + monkeypatch.setattr( + f"{module}.get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr( + attention, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + attention_stub = nn.Identity() + attention_stub.attn_backend = SimpleNamespace(get_name=lambda: "TEST") + monkeypatch.setattr(attention, "Attention", lambda *args, **kwargs: attention_stub) + config = BeamConfig( + hidden_size=64, num_attention_heads=48, num_key_value_heads=8, head_dim=8 + ) + attn = ( + attention.BeamAttention(config, None, None, prefix="layers.0.self_attn") + .bfloat16() + .cuda() + ) + model = BeamModel.__new__(BeamModel) + nn.Module.__init__(model) + model.config = config + model.hf_to_vllm_mapper = model.get_weights_mapper(config, tp_size) + layer = nn.Module() + layer.self_attn = attn + model.layers = nn.ModuleList([layer]) + param = attn.qkv_gate_proj.weight + initial_pointer = param.data_ptr() + for _ in range(2): + weights = [ + torch.randn(n, 64, device="cuda", dtype=torch.bfloat16) + for n in (384, 64, 64, 384) + ] + loaded = model.load_weights( + (f"layers.0.self_attn.{name}.weight", w) + for name, w in zip(("q_proj", "k_proj", "v_proj", "gate_proj"), weights) + ) + assert loaded == {"layers.0.self_attn.qkv_gate_proj.weight"} + assert param.data_ptr() == initial_pointer + expected_weight = torch.cat([w.chunk(tp_size)[-1] for w in weights]) + torch.testing.assert_close(param, expected_weight, rtol=0, atol=0) + x = torch.randn(3, 64, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close( + packed_projection(x, param, False), + torch.nn.functional.linear(x, expected_weight), + rtol=0, + atol=0, + ) diff --git a/tests/test_config.py b/tests/test_config.py index c68ecca605ac..b1f977bd6e0b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -924,7 +924,10 @@ def test_v2_model_runner_supports_custom_logits_processors(): assert config._get_v2_model_runner_unsupported_features() == [] -@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"]) +@pytest.mark.parametrize( + "architecture", + ["DFlash2DraftModel", "DFlash2BeamDraftModel", "LiLiCorrDraftModel"], +) def test_dflash_candidate_draft_forces_v2_model_runner(architecture): """A DFlash2 draft must reach the V2 speculator, the only one that runs its candidate selector; on V1 it would draft as DFlash1 without raising.""" diff --git a/tests/transformers_utils/test_beam_config.py b/tests/transformers_utils/test_beam_config.py new file mode 100644 index 000000000000..d9477c9481c9 --- /dev/null +++ b/tests/transformers_utils/test_beam_config.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test that Beam's released config schema loads and round-trips.""" + +import json + +import pytest + +from vllm.transformers_utils.config import get_config +from vllm.transformers_utils.configs.beam import BeamConfig + +_NUM_LAYERS = 6 +# Shaped like reflection/Beam-1.0-501B-A23B's config.json, scaled down. +_RELEASE_CONFIG = { + "architectures": ["BeamForCausalLM"], + "model_type": "beam", + "attention_bias": False, + "final_logit_softcapping": 30.0, + "head_dim": 16, + "hidden_size": 64, + "intermediate_size": 128, + "layer_types": ["sliding_attention"] * 5 + ["full_attention"], + "max_position_embeddings": 1048576, + "mlp_layer_types": ["dense"] * 3 + ["sparse"] * 3, + "moe_intermediate_size": 32, + "n_group": 8, + "norm_topk_prob": True, + "num_attention_heads": 4, + "num_experts": 16, + "num_experts_per_tok": 4, + "num_hidden_layers": _NUM_LAYERS, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-6, + "rope_parameters": { + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + "rope_theta": None, + "routed_scaling_factor": 2.5, + "sliding_window": 512, + "tie_word_embeddings": False, + "topk_group": 4, + "vocab_size": 256, +} + + +def test_release_config_loads_and_round_trips(tmp_path): + (tmp_path / "config.json").write_text(json.dumps(_RELEASE_CONFIG)) + config = get_config(str(tmp_path), trust_remote_code=False) + assert isinstance(config, BeamConfig) + for resolved in (config, BeamConfig(**config.to_dict())): + # The NoPE full-attention entry carries no RoPE parameters. + assert resolved.rope_parameters == { + "rope_type": "default", + "rope_theta": 10000.0, + } + assert resolved.layer_types == _RELEASE_CONFIG["layer_types"] + assert resolved.mlp_layer_types == _RELEASE_CONFIG["mlp_layer_types"] + assert resolved.routed_scaling_factor == 2.5 + assert resolved.final_logit_softcapping == 30.0 + + +def test_mismatched_rope_types_are_rejected(): + rope = {"rope_type": "default", "rope_theta": 10000.0} + with pytest.raises(ValueError, match="exactly one type"): + BeamConfig( + num_hidden_layers=2, + rope_parameters={"sliding_attention": rope, "full_attention": rope}, + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {"layer_types": ["sliding_attention"]}, + {"layer_types": ["sliding_attention", "chunked_attention"]}, + {"mlp_layer_types": ["dense"]}, + {"mlp_layer_types": ["dense", "shared"]}, + ], +) +def test_invalid_layer_schedules_are_rejected(overrides): + with pytest.raises(ValueError, match="one"): + BeamConfig(num_hidden_layers=2, **overrides) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 558bb38e6bc9..e94ae3591fc3 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -39,6 +39,50 @@ def test_grouped_conv_matches_reference(block_size: int): torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize( + "rows,width,query,group,taps", + [ + (8, 5120, 8, 16, 2), + (64, 5120, 8, 16, 2), + (256, 5120, 8, 16, 2), + (15, 143, 5, 11, 3), + ], +) +@pytest.mark.parametrize("warps", [None, 4, 8, 16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_beam_conv_residual_norm(rows, width, query, group, taps, warps, dtype): + """Dtype rounding and request boundaries survive residual/norm fusion.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_conv_ops import conv_residual_norm + + torch.manual_seed(42) + hidden = torch.randn(rows, width, device="cuda", dtype=dtype) + residual = torch.randn_like(hidden) + weight = torch.randn(width, device="cuda", dtype=hidden.dtype) + base = torch.randn(taps, width, device="cuda", dtype=hidden.dtype) + delta = torch.randn( + rows, 2, taps, width // group, device="cuda", dtype=hidden.dtype + )[:, 1] + expected = _grouped_conv(hidden, delta, base, query, width // group, group, taps) + expected_residual = residual.clone() + ops.fused_add_rms_norm(expected, expected_residual, weight, 1e-6) + actual, actual_residual = conv_residual_norm( + hidden, delta, base, residual, weight, 1e-6, query, group, num_warps=warps + ) + tolerance = {torch.bfloat16: 0.015625, torch.float16: 0.002, torch.float32: 3e-6}[ + dtype + ] + residual_tolerance = 0 if dtype == torch.bfloat16 else tolerance + torch.testing.assert_close( + actual_residual, + expected_residual, + rtol=residual_tolerance, + atol=residual_tolerance, + ) + torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize( @@ -106,6 +150,121 @@ def test_draft_quant_exclusions_include_global_layer_indices(): assert "lilicorr.layers.88.mlp.0" not in quant_config.exclude_modules +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("rows,heads", [(8, 1), (64, 2), (256, 1)]) +def test_beam_fused_context_kv(rows, heads): + """All layer pointers and strided cache pages match the staged context path.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_context_ops import store_context_kv + + torch.manual_seed(42) + layers, dim, page = 5, 128, 16 + kv = torch.randn(rows, layers, 2, heads, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(layers, dim, device="cuda", dtype=kv.dtype) + positions = torch.arange(rows, device="cuda", dtype=torch.int64) + angles = torch.randn(rows, dim // 2, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(kv.dtype) + staged = kv.permute(2, 1, 0, 3, 4).contiguous() + keys = torch.empty_like(staged[0]) + ops.rms_norm(keys, staged[0], weight, 1e-6) + ops.rotary_embedding( + positions.repeat(layers), keys.view(layers * rows, -1), None, dim, rope, True + ) + values = staged[1] + storage = torch.randn( + (rows + page - 1) // page, + layers, + heads, + page, + 2 * dim, + device="cuda", + dtype=kv.dtype, + ) + expected = storage.clone() + slots = positions.flip(0)[:-1].contiguous() + slots[::3] = -1 + valid = slots >= 0 + ids = slots[valid] + for layer in range(layers): + expected[ids // page, layer, :, ids % page, :dim] = keys[layer, :-1][valid] + expected[ids // page, layer, :, ids % page, dim:] = values[layer, :-1][valid] + caches = storage.unbind(1) + pointers = torch.tensor( + [c.data_ptr() for c in caches], device="cuda", dtype=torch.int64 + ) + store_context_kv(kv, weight, positions, rope, pointers, slots, caches[0], 1e-6) + torch.testing.assert_close(storage, expected, rtol=0.01, atol=0.015625) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("rows", [1, 8, 65, 2048]) +@pytest.mark.parametrize("cache_dtype", ["fp8_e4m3", "fp8_e5m2"]) +def test_beam_fp8_context_cache_matches_backend_and_replays(rows, cache_dtype): + """Quantize rounded context K/V with each layer's live scales and slot mask.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.dflash2_context_ops import store_context_kv + + torch.manual_seed(42) + layers, heads, dim, page = 3, 2, 128, 16 + kv = torch.randn(rows, layers, 2, heads, dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(layers, dim, device="cuda", dtype=kv.dtype) + positions = torch.arange(rows, device="cuda") + angles = torch.randn(rows, dim // 2, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(kv.dtype) + shape = ((rows + page - 1) // page, layers, heads, page, 2 * dim) + bf16 = torch.zeros(shape, device="cuda", dtype=kv.dtype) + actual = torch.zeros(shape, device="cuda", dtype=torch.uint8) + expected = torch.zeros_like(actual) + dtype = torch.float8_e4m3fn if cache_dtype == "fp8_e4m3" else torch.float8_e5m2 + scales = torch.tensor([[0.01, 0.02], [0.5, 1.5], [2.0, 3.0]], device="cuda") + scale_ptrs = torch.tensor( + [s.data_ptr() for pair in scales.unbind() for s in pair.unbind()], + device="cuda", + ) + + def pointers(storage): + return torch.tensor([c.data_ptr() for c in storage.unbind(1)], device="cuda") + + slots = positions.flip(0).contiguous() + slots[1::3] = -1 + args = (kv, weight, positions, rope) + bf16_ptrs, fp8_ptrs = pointers(bf16), pointers(actual) + + def run(): + store_context_kv( + *args, + fp8_ptrs, + slots, + actual[:, 0].view(dtype), + 1e-6, + scale_ptrs, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for factor in (1.0, 0.5): + scales.mul_(factor) + kv.mul_(factor) + store_context_kv(*args, bf16_ptrs, positions, bf16[:, 0], 1e-6) + for layer in range(layers): + source = bf16[:, layer].transpose(1, 2).reshape(-1, heads, 2 * dim) + cache = expected[:, layer].transpose(1, 2) + ops.reshape_and_cache_flash( + source[:rows, :, :dim].contiguous(), + source[:rows, :, dim:].contiguous(), + cache[..., :dim], + cache[..., dim:], + slots, + cache_dtype, + scales[layer, 0], + scales[layer, 1], + ) + graph.replay() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_selector_edges_match_sequential_reference(): torch.manual_seed(1) batch, steps, top_k, rank = 2, 4, 3, 5 @@ -142,6 +301,78 @@ def test_selector_edges_match_sequential_reference(): torch.testing.assert_close(actual, expected) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize( + "batch,steps,top_k,rank", + [(1, 1, 3, 5), (1, 7, 16, 256), (8, 7, 16, 256), (32, 7, 16, 256), (3, 5, 17, 65)], +) +@pytest.mark.parametrize("unary_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("anchor_dtype", [torch.int32, torch.int64]) +def test_beam_fused_candidate_edges( + batch, steps, top_k, rank, unary_dtype, anchor_dtype +): + """Keep the anchor, predecessor mapping, and BF16 intermediate rounding.""" + from vllm.models.beam.nvidia.dflash2_ops import candidate_edges + + torch.manual_seed(1) + vocab = 257 + predecessors = torch.randn(vocab, rank, device="cuda", dtype=torch.bfloat16) + successors = torch.randn_like(predecessors) + candidates = torch.randint(vocab, (batch, steps, top_k), device="cuda") + unary = torch.randn(batch, steps, top_k, device="cuda", dtype=unary_dtype) + hidden = torch.randn(batch, steps, rank, device="cuda", dtype=torch.bfloat16) + # Production anchors are a strided view of the query token buffer. + anchors = torch.randint(vocab, (batch * 8,), device="cuda", dtype=anchor_dtype)[::8] + args = (predecessors, successors, candidates, unary, hidden, anchors, top_k) + expected = _score_edges(*args) + actual = candidate_edges(*args) + torch.testing.assert_close(actual, expected, rtol=0.01, atol=0.015625) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("rows", [8, 64, 256]) +def test_beam_draft_qk_norm_rope_cache(rows): + """Fused query preparation preserves BF16 rounding and ignores padded slots.""" + from vllm import _custom_ops as ops + from vllm.models.beam.nvidia.attention_ops import beam_qk_norm_rope + + torch.manual_seed(42) + qkv = torch.randn(rows, 896, device="cuda", dtype=torch.bfloat16) + q, k, v = qkv.split((640, 128, 128), -1) + qw, kw = torch.randn(2, 128, device="cuda", dtype=q.dtype) + positions = torch.arange(rows, device="cuda", dtype=torch.int64) + angles = torch.randn(rows, 64, device="cuda") + rope = torch.cat((angles.cos(), angles.sin()), -1).to(q.dtype) + oq, ok = torch.empty_like(q), torch.empty_like(k) + ops.rms_norm(oq.view(-1, 5, 128), q.view(-1, 5, 128), qw, 1e-6) + ops.rms_norm(ok.view(-1, 1, 128), k.view(-1, 1, 128), kw, 1e-6) + ops.rotary_embedding(positions, oq, ok, 128, rope, True) + cache = torch.randn((rows + 15) // 16, 1, 16, 256, device="cuda", dtype=q.dtype) + expected_cache = cache.clone() + slots = positions.flip(0).contiguous() + slots[::3] = -1 + valid = slots >= 0 + slot_ids = slots[valid] + expected_cache[slot_ids // 16, 0, slot_ids % 16, :128] = ok[valid] + expected_cache[slot_ids // 16, 0, slot_ids % 16, 128:] = v[valid] + actual = beam_qk_norm_rope( + q, + k, + qw, + kw, + 1e-6, + positions, + rope, + 128, + 128, + value=v, + kv_cache=cache, + slot_mapping=slots, + ) + torch.testing.assert_close(actual, (oq, ok), rtol=0.01, atol=0.015625) + torch.testing.assert_close(cache, expected_cache, rtol=0.01, atol=0.015625) + + def _stub_base(monkeypatch, draft_logits): """A DFlashSpeculator.__init__ that allocates only what the base class would. @@ -189,20 +420,8 @@ def test_selector_asks_for_fp32_proposal_logits(): assert fill == float("-inf") -@pytest.mark.skip_global_cleanup -@pytest.mark.parametrize("variant", ["dflash2", "lilicorr", "lilicorr_plain"]) -def test_candidate_model_decoder_layer_cls(monkeypatch, variant): - from types import SimpleNamespace - - from vllm.config import set_current_vllm_config - from vllm.model_executor.models.lilicorr import LiLiCorr - from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer - from vllm.model_executor.models.qwen3_dflash2 import ( - DFlash2Qwen3DecoderLayer, - DFlash2Qwen3Model, - ) - - # 1. Mock get_current_vllm_config and TP groups +def _mock_draft_configs(monkeypatch, hf_config): + """The current and draft vLLM configs for building a draft model on CPU.""" mock_current_vllm_config = SimpleNamespace( cache_config=SimpleNamespace( block_size=16, @@ -234,13 +453,12 @@ def test_candidate_model_decoder_layer_cls(monkeypatch, variant): dtype=torch.float32, is_mm_prefix_lm=False, rswa_window=None, + head_dtype=None, ), kernel_config=SimpleNamespace( linear_backend="auto", ), ) - from vllm.platforms import current_platform - monkeypatch.setattr( current_platform, "get_attn_backend_cls", @@ -258,7 +476,41 @@ class MockGroup: MockGroup(), ) - # 2. Mock vllm_config + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace( + hf_config=hf_config, + quantization=None, + ), + num_speculative_tokens=4, + enable_adaptive_verification=False, + ), + model_config=SimpleNamespace( + dtype=torch.float32, + is_mm_prefix_lm=False, + get_total_num_hidden_layers=lambda: 2, + ), + load_config=SimpleNamespace( + quantization=None, + quantization_param_path=None, + ), + ) + mock_current_vllm_config.speculative_config = vllm_config.speculative_config + vllm_config.compilation_config = mock_current_vllm_config.compilation_config + return mock_current_vllm_config, vllm_config + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("variant", ["dflash2", "lilicorr", "lilicorr_plain"]) +def test_candidate_model_decoder_layer_cls(monkeypatch, variant): + from vllm.config import set_current_vllm_config + from vllm.model_executor.models.lilicorr import LiLiCorr + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer + from vllm.model_executor.models.qwen3_dflash2 import ( + DFlash2Qwen3DecoderLayer, + DFlash2Qwen3Model, + ) + hf_config = SimpleNamespace( vocab_size=1000, hidden_size=256, @@ -287,33 +539,13 @@ class MockGroup: "use_aux_hidden_state": False, }, ) - vllm_config = SimpleNamespace( - speculative_config=SimpleNamespace( - draft_model_config=SimpleNamespace( - hf_config=hf_config, - quantization=None, - ), - num_speculative_tokens=4, - enable_adaptive_verification=False, - ), - model_config=SimpleNamespace( - dtype=torch.float32, - is_mm_prefix_lm=False, - ), - load_config=SimpleNamespace( - quantization=None, - quantization_param_path=None, - ), - ) - mock_current_vllm_config.speculative_config = vllm_config.speculative_config - vllm_config.compilation_config = mock_current_vllm_config.compilation_config + current_config, vllm_config = _mock_draft_configs(monkeypatch, hf_config) - # 3. Instantiate the model under meta device to avoid parameter allocation issues - with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"): + # Instantiate the model under meta device to avoid parameter allocation issues + with set_current_vllm_config(current_config), torch.device("meta"): model_cls = DFlash2Qwen3Model if variant == "dflash2" else LiLiCorr model = model_cls(vllm_config=vllm_config) - # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass) assert len(model.layers) == 2 expected = ( DFlashQwen3DecoderLayer @@ -323,6 +555,93 @@ class MockGroup: assert type(model.layers[0]) is expected +@pytest.mark.skip_global_cleanup +def test_beam_dflash2_loads_every_checkpoint_tensor(monkeypatch): + """The exported checkpoint's names load into every parameter but the + embedding and LM head, which alias the target's.""" + from vllm.config import set_current_vllm_config + from vllm.models.beam import DFlash2BeamForCausalLM + + hidden, heads, kv_heads, head_dim, ffn, vocab, rank = 64, 4, 2, 16, 128, 100, 8 + taps, groups = 2, 4 + hf_config = SimpleNamespace( + vocab_size=vocab, + hidden_size=hidden, + num_hidden_layers=2, + num_attention_heads=heads, + num_key_value_heads=kv_heads, + head_dim=head_dim, + max_position_embeddings=2048, + rms_norm_eps=1e-6, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + intermediate_size=ffn, + hidden_act="silu", + is_causal=False, + layer_types=["sliding_attention"] * 2, + sliding_window=64, + dflash_config={ + "block_size": 16, + "conv_group_size": hidden // groups, + "conv_kernel_size": taps, + "final_logit_softcapping": 30, + "mask_token_id": 99, + "selector_rank": rank, + "selector_top_k": 3, + "target_layer_ids": [0, 1], + }, + ) + current_config, vllm_config = _mock_draft_configs(monkeypatch, hf_config) + with set_current_vllm_config(current_config): + model = DFlash2BeamForCausalLM(vllm_config=vllm_config) + + shapes = { + "fc.weight": (hidden, 2 * hidden), + "hidden_norm.weight": (hidden,), + "norm.weight": (hidden,), + "candidate_selector.hidden_projection.weight": (rank, hidden), + "candidate_selector.predecessor_codebook": (vocab, rank), + "candidate_selector.successor_codebook": (vocab, rank), + } + for layer in range(2): + for conv in ("attention_conv", "mlp_conv"): + shapes[f"layers.{layer}.{conv}.base_kernel"] = (2, taps, hidden) + shapes[f"layers.{layer}.{conv}.kernel_projection.weight"] = ( + 2 * taps * groups, + hidden, + ) + shapes |= { + f"layers.{layer}.input_layernorm.weight": (hidden,), + f"layers.{layer}.post_attention_layernorm.weight": (hidden,), + f"layers.{layer}.self_attn.q_proj.weight": (heads * head_dim, hidden), + f"layers.{layer}.self_attn.k_proj.weight": (kv_heads * head_dim, hidden), + f"layers.{layer}.self_attn.v_proj.weight": (kv_heads * head_dim, hidden), + f"layers.{layer}.self_attn.o_proj.weight": (hidden, heads * head_dim), + f"layers.{layer}.self_attn.q_norm.weight": (head_dim,), + f"layers.{layer}.self_attn.k_norm.weight": (head_dim,), + f"layers.{layer}.mlp.gate_proj.weight": (ffn, hidden), + f"layers.{layer}.mlp.up_proj.weight": (ffn, hidden), + f"layers.{layer}.mlp.down_proj.weight": (hidden, ffn), + } + torch.manual_seed(0) + checkpoint = {name: torch.randn(shape) for name, shape in shapes.items()} + + loaded = model.load_weights(iter(checkpoint.items())) + + assert loaded == {name for name, _ in model.named_parameters()} + context_kv = [ + checkpoint[f"layers.{layer}.self_attn.{proj}_proj.weight"] + for layer in range(2) + for proj in ("k", "v") + ] + torch.testing.assert_close(model.model._context_kv_weight, torch.cat(context_kv)) + # Numbered after the target's two layers, so the KV cache keys stay distinct. + assert model.get_draft_kv_cache_layer_names() == [ + "model.layers.2.self_attn.attn", + "model.layers.3.self_attn.attn", + ] + assert model.get_draft_attn_causal() == [False, False] + + def test_conv_projections_use_draft_quant_config(monkeypatch): from torch import nn diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 050cc561194c..0a68f585a4a1 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -74,6 +74,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { "AliceAIForCausalLM", + "BeamForCausalLM", "DeepseekV32MTPModel", "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", @@ -749,9 +750,11 @@ def _is_dflash_candidate_draft(self) -> bool: if draft_config is None: return False return bool( - {"DFlash2DraftModel", "LiLiCorrDraftModel"}.intersection( - draft_config.architectures or [] - ) + { + "DFlash2DraftModel", + "DFlash2BeamDraftModel", + "LiLiCorrDraftModel", + }.intersection(draft_config.architectures or []) ) def _dflash_needs_multi_kv_group(self) -> bool: diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 50280061e257..0e4235ae6685 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -79,6 +79,7 @@ "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"), "BailingMoeV3ForCausalLM": ("bailing_moe_v3", "BailingMoeV3ForCausalLM"), + "BeamForCausalLM": ("vllm.models.beam", "BeamForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), @@ -637,6 +638,7 @@ "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "LiLiCorrDraftModel": ("lilicorr", "LiLiCorrForCausalLM"), "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"), + "DFlash2BeamDraftModel": ("vllm.models.beam", "DFlash2BeamForCausalLM"), # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash # implementation. EAGLEConfig rewrites a dflash draft's architecture to # DFlash{arch} unless it already starts or ends with "DFlash" (see diff --git a/vllm/models/beam/__init__.py b/vllm/models/beam/__init__.py new file mode 100644 index 000000000000..faed0717d1a7 --- /dev/null +++ b/vllm/models/beam/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam model entry point.""" + +from .nvidia.dflash2 import DFlash2BeamForCausalLM +from .nvidia.model import BeamForCausalLM + +__all__ = ["BeamForCausalLM", "DFlash2BeamForCausalLM"] diff --git a/vllm/models/beam/nvidia/__init__.py b/vllm/models/beam/nvidia/__init__.py new file mode 100644 index 000000000000..d2d40aaa7d8d --- /dev/null +++ b/vllm/models/beam/nvidia/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA-specific Beam model implementation.""" diff --git a/vllm/models/beam/nvidia/activation.py b/vllm/models/beam/nvidia/activation.py new file mode 100644 index 000000000000..9afc18a0b464 --- /dev/null +++ b/vllm/models/beam/nvidia/activation.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + + +@triton.jit(do_not_specialize=["size"]) +def _silu_mul_kernel(X, Y, size, width: tl.constexpr, BLOCK: tl.constexpr): + i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + src = i // width * (2 * width) + i % width + gate = tl.load(X + src, i < size, other=0).to(tl.float32) + up = tl.load(X + src + width, i < size, other=0).to(tl.float32) + act = tl.div_rn(gate, 1.0 + libdevice.exp(-gate)) + tl.store(Y + i, act * up, i < size) + + +def beam_silu_mul(x: torch.Tensor) -> torch.Tensor: + """Compute SiLU and multiplication in FP32, rounding only the product.""" + width = x.shape[-1] // 2 + output = torch.empty((*x.shape[:-1], width), dtype=x.dtype, device=x.device) + if output.numel(): + _silu_mul_kernel[(triton.cdiv(output.numel(), 1024),)]( + x, output, output.numel(), width, 1024, enable_fp_fusion=False + ) + return output + + +@triton.jit +def _softcap_logits_kernel( + X, + WIDTH: tl.constexpr, + STRIDE: tl.constexpr, + CAP: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(X + row * STRIDE + col, col < WIDTH, 0).to(tl.float32) + y = libdevice.tanh(x * (1.0 / CAP)) * CAP + tl.store(X + row * STRIDE + col, y, col < WIDTH) + + +def beam_softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: + """Softcap in FP32 and round back into the original logits buffer.""" + assert logits.ndim == 2 and logits.stride(1) == 1 + if logits.numel(): + _softcap_logits_kernel[(logits.shape[0], triton.cdiv(logits.shape[1], 1024))]( + logits, + logits.shape[1], + logits.stride(0), + cap, + 1024, + enable_fp_fusion=False, + enable_reflect_ftz=False, + ) + return logits diff --git a/vllm/models/beam/nvidia/attention.py b/vllm/models/beam/nvidia/attention.py new file mode 100644 index 000000000000..f685129ffce7 --- /dev/null +++ b/vllm/models/beam/nvidia/attention.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.config import CacheConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import set_weight_attrs +from vllm.transformers_utils.configs.beam import BeamConfig +from vllm.v1.attention.backend import AttentionType + +from .attention_ops import beam_output_gate, beam_qk_norm_rope +from .layernorm import RMSNorm +from .projection import ( + packed_projection, + use_low_latency_gemm, + warmup_packed_projection, +) + + +def use_merged_qkv_gate(config: BeamConfig, tp_size: int) -> bool: + return config.num_key_value_heads >= tp_size + + +class BeamAttention(nn.Module): + """GQA with per-head Q/K RMSNorm, then RoPE (sliding-window layers) or + SSMax (NoPE full-attention layers). The attention output is scaled by + ``sigmoid(gate + bias)`` before ``o_proj``. + """ + + def __init__( + self, + config: BeamConfig, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = extract_layer_index(prefix) + hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.merged_qkv_gate = use_merged_qkv_gate(config, tp_size) + if self.merged_qkv_gate: + self.qkv_gate_proj = MergedColumnParallelLinear( + hidden_size, + [ + self.total_num_heads * self.head_dim, + self.total_num_kv_heads * self.head_dim, + self.total_num_kv_heads * self.head_dim, + self.total_num_heads * self.head_dim, + ], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_gate_proj", + ) + else: + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.output_gate_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rms_norm_eps = config.rms_norm_eps + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + # The output gate has a separate bias per head. + self.gate_bias = nn.Parameter(torch.zeros(self.num_heads)) + set_weight_attrs(self.gate_bias, {"weight_loader": sharded_weight_loader(0)}) + + is_sliding = config.layer_types[self.layer_id] == "sliding_attention" + self.disable_rope = not is_sliding + if self.disable_rope: + self.rotary_emb = None + else: + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + is_neox_style=True, + rope_parameters=config.rope_parameters, + dtype=torch.float32, + ) + + # Backends interpret `sliding_window` as `(q - k) <= sliding_window - 1`. + per_layer_sw = config.sliding_window if is_sliding else None + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=per_layer_sw, + attn_type=AttentionType.DECODER, + prefix=f"{prefix}.attn", + ) + + self.fuse_kv_insert = ( + self.attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + and self.attn.impl.dcp_world_size == 1 + and self.attn.kv_sharing_target_layer_name is None + ) + self.low_latency_projection = False + + def warmup_projection(self) -> None: + if not self.merged_qkv_gate: + return + projection = self.qkv_gate_proj + self.low_latency_projection = isinstance( + projection.quant_method, UnquantizedLinearMethod + ) and use_low_latency_gemm(projection.weight) + if self.low_latency_projection: + warmup_packed_projection(projection.weight) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + if self.merged_qkv_gate: + if self.low_latency_projection: + packed = packed_projection( + hidden_states, self.qkv_gate_proj.weight, True + ) + else: + packed, _ = self.qkv_gate_proj(hidden_states) + qkv, gate = packed.split( + [self.q_size + 2 * self.kv_size, self.q_size], dim=-1 + ) + else: + qkv, _ = self.qkv_proj(hidden_states) + gate, _ = self.output_gate_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + kv_cache, slot_mapping = None, None + if ( + self.fuse_kv_insert + and self.attn.kv_cache.dtype == q.dtype == torch.bfloat16 + ): + slots = get_forward_context().slot_mapping + if isinstance(slots, dict): + slot_mapping = slots.get(self.attn.layer_name) + if slot_mapping is not None: + kv_cache = self.attn.kv_cache + + rope = self.rotary_emb + q, k = beam_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + self.rms_norm_eps, + positions, + rope.cos_sin_cache if rope is not None else None, + self.head_dim, + rope.rotary_dim if rope is not None else 0, + rope.is_neox_style if rope is not None else True, + ssmax=self.disable_rope, + value=v, + kv_cache=kv_cache, + slot_mapping=slot_mapping, + ) + + # Both backends read the populated cache when K/V are omitted. + attn_output = ( + self.attn(q, k, v) if kv_cache is None else self.attn(q, None, None) + ) + attn_output = beam_output_gate( + attn_output, + gate, + self.gate_bias, + num_heads=self.num_heads, + head_dim=self.head_dim, + ) + output, _ = self.o_proj(attn_output) + return output diff --git a/vllm/models/beam/nvidia/attention_ops.py b/vllm/models/beam/nvidia/attention_ops.py new file mode 100644 index 000000000000..1e30f11a4aae --- /dev/null +++ b/vllm/models/beam/nvidia/attention_ops.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Q/K normalization, position transforms, and attention output gating.""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + + +@triton.jit(do_not_specialize=["SIZE"]) +def _output_gate_kernel( + X, + G, + B, + Y, + SIZE, + G_STRIDE: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + dtype = G.dtype.element_ty + g_idx = i // (HEADS * DIM) * G_STRIDE + i % (HEADS * DIM) + g = tl.load(G + g_idx, i < SIZE, other=0).to(tl.float32) + b = tl.load(B + i // DIM % HEADS).to(tl.float32) + g = (g + b).to(dtype).to(tl.float32) + sigmoid = (1.0 / (1.0 + tl.exp(-g))).to(dtype).to(tl.float32) + x = tl.load(X + i, i < SIZE, other=0).to(tl.float32) + tl.store(Y + i, x * sigmoid, i < SIZE) + + +def beam_output_gate( + attn_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + *, + num_heads: int, + head_dim: int, +) -> torch.Tensor: + """``attn_output * sigmoid(gate + gate_bias)`` with activation-dtype rounding + after the bias add and the sigmoid.""" + output = torch.empty_like(attn_output) + if output.numel(): + if ( + attn_output.dtype == gate.dtype == gate_bias.dtype == torch.bfloat16 + and current_platform.is_device_capability_family(100) + ): + from .attention_ops_cutedsl import run + + run(attn_output, gate, gate_bias, num_heads, head_dim, output) + return output + _output_gate_kernel[(triton.cdiv(output.numel(), 256),)]( + attn_output, + gate, + gate_bias, + output, + output.numel(), + gate.stride(0), + num_heads, + head_dim, + 256, + enable_fp_fusion=False, + ) + return output + + +@triton.jit( + do_not_specialize=["M", "NUM_SLOTS"], do_not_specialize_on_alignment=["POS"] +) +def _qk_norm_rope_kernel( + Q, + K, + QW, + KW, + POS, + CACHE, + OQ, + OK, + M, + V, + KV_CACHE, + SLOTS, + NUM_SLOTS, + WRITE_CACHE: tl.constexpr, + V_STRIDE: tl.constexpr, + CACHE_BLOCK_STRIDE: tl.constexpr, + CACHE_HEAD_STRIDE: tl.constexpr, + CACHE_TOKEN_STRIDE: tl.constexpr, + PAGE_SIZE: tl.constexpr, + TOKENS: tl.constexpr, + Q_HEADS: tl.constexpr, + K_HEADS: tl.constexpr, + DIM: tl.constexpr, + Q_STRIDE: tl.constexpr, + K_STRIDE: tl.constexpr, + EPS: tl.constexpr, + ROTARY: tl.constexpr, + NEOX: tl.constexpr, + SSMAX: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) * TOKENS + tl.arange(0, TOKENS) + head = tl.program_id(1) + c = tl.arange(0, BLOCK) + mask = (token[:, None] < M) & (c[None, :] < DIM) + if head < Q_HEADS: + x = tl.load( + Q + token[:, None] * Q_STRIDE + head * DIM + c[None, :], mask, 0 + ).to(tl.float32) + if BLOCK == 128: + squares = (x * x).reshape((TOKENS, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(x * x, 1) + y = x * tl.rsqrt(total / DIM + EPS)[:, None] + y = y * tl.load(QW + c, c < DIM, 0).to(tl.float32)[None, :] + out = OQ + (token[:, None] * Q_HEADS + head) * DIM + else: + h = head - Q_HEADS + x = tl.load(K + token[:, None] * K_STRIDE + h * DIM + c[None, :], mask, 0).to( + tl.float32 + ) + if BLOCK == 128: + squares = (x * x).reshape((TOKENS, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(x * x, 1) + y = x * tl.rsqrt(total / DIM + EPS)[:, None] + y = y * tl.load(KW + c, c < DIM, 0).to(tl.float32)[None, :] + out = OK + (token[:, None] * K_HEADS + h) * DIM + y = y.to(Q.dtype.element_ty).to(tl.float32) + if ROTARY: + half: tl.constexpr = ROTARY // 2 + if NEOX: + partner = tl.where(c < half, c + half, c - half) + second = c >= half + offset = c % half + else: + partner = c ^ 1 + second = c % 2 != 0 + offset = c // 2 + partner = tl.minimum(tl.maximum(partner, 0), BLOCK - 1) + other = tl.gather(y, tl.broadcast_to(partner[None, :], (TOKENS, BLOCK)), axis=1) + pos = tl.load(POS + token, token < M, 0).to(tl.int64) + rope_mask = (token[:, None] < M) & (c[None, :] < ROTARY) + cos = tl.load(CACHE + pos[:, None] * ROTARY + offset[None, :], rope_mask, 0) + sin = tl.load( + CACHE + pos[:, None] * ROTARY + half + offset[None, :], rope_mask, 0 + ) + # Match the FP32 multiply/FMA ordering of the CUDA RoPE kernel. + rotated = tl.where( + second[None, :], + tl.fma(other, sin, y * cos), + tl.fma(y, cos, -other * sin), + ) + y = tl.where(c[None, :] < ROTARY, rotated, y) + elif SSMAX: + if head < Q_HEADS: + # (1 + (position + 1) / 2048) ** 0.1 + p = tl.load(POS + token, token < M, 0).to(tl.float32) + 1 + scale = libdevice.exp(0.1 * libdevice.log1p(p / 2048.0)) + y = y * scale[:, None] + tl.store(out + c[None, :], y, mask) + if WRITE_CACHE and head >= Q_HEADS: + slot = tl.load(SLOTS + token, token < NUM_SLOTS, -1).to(tl.int64) + cache_mask = mask & (token[:, None] < NUM_SLOTS) & (slot[:, None] >= 0) + address = ( + slot[:, None] // PAGE_SIZE * CACHE_BLOCK_STRIDE + + slot[:, None] % PAGE_SIZE * CACHE_TOKEN_STRIDE + + (head - Q_HEADS) * CACHE_HEAD_STRIDE + + c[None, :] + ) + tl.store(KV_CACHE + address, y, cache_mask) + value = tl.load( + V + token[:, None] * V_STRIDE + (head - Q_HEADS) * DIM + c[None, :], + cache_mask, + 0, + ) + tl.store(KV_CACHE + address + DIM, value, cache_mask) + + +def beam_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor | None, + head_dim: int, + rotary_dim: int = 0, + is_neox_style: bool = True, + *, + ssmax: bool = False, + value: torch.Tensor | None = None, + kv_cache: torch.Tensor | None = None, + slot_mapping: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Round norm before position transforms and optionally insert BF16 K/V.""" + if kv_cache is not None: + assert value is not None and slot_mapping is not None + assert kv_cache.dtype == q.dtype == torch.bfloat16 + assert kv_cache.ndim == 4 and kv_cache.stride(-1) == 1 + q_out, k_out = torch.empty_like(q), torch.empty_like(k) + q_heads, k_heads = q.shape[-1] // head_dim, k.shape[-1] // head_dim + if q.shape[0]: + if rotary_dim or q.shape[0] < 512: + tile, warps = 4, 4 + elif q.shape[0] <= 2048: + tile, warps = 16, 8 + else: + tile, warps = 32, 4 + _qk_norm_rope_kernel[(triton.cdiv(q.shape[0], tile), q_heads + k_heads)]( + q, + k, + q_weight, + k_weight, + positions, + cos_sin_cache, + q_out, + k_out, + q.shape[0], + value, + kv_cache, + slot_mapping, + slot_mapping.numel() if slot_mapping is not None else 0, + kv_cache is not None, + value.stride(0) if value is not None else 0, + kv_cache.stride(0) if kv_cache is not None else 0, + kv_cache.stride(1) if kv_cache is not None else 0, + kv_cache.stride(2) if kv_cache is not None else 0, + kv_cache.shape[2] if kv_cache is not None else 1, + tile, + q_heads, + k_heads, + head_dim, + q.stride(0), + k.stride(0), + eps, + rotary_dim, + is_neox_style, + ssmax, + triton.next_power_of_2(head_dim), + num_warps=warps, + enable_fp_fusion=False, + ) + return q_out, k_out diff --git a/vllm/models/beam/nvidia/attention_ops_cutedsl.py b/vllm/models/beam/nvidia/attention_ops_cutedsl.py new file mode 100644 index 000000000000..fd051569e9ca --- /dev/null +++ b/vllm/models/beam/nvidia/attention_ops_cutedsl.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bit-exact BF16 output gate in CuTe DSL.""" + +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Int32, Int64, Uint32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.utils.torch_utils import current_stream + +_LDQ = "ld.global.nc.L1::no_allocate" + + +def _asm(res_ty, operands, text, constraints, loc=None, ip=None, side_effects=False): + return llvm.inline_asm( + res_ty, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_v4(addr: Int64, pf="", *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + pf + ".v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def make_policy(kind: str, *, loc=None, ip=None) -> Int64: + return Int64( + _asm( + T.i64(), + [], + f"createpolicy.fractional.L2::{kind}.b64 $0, 1.0;", + "=l", + loc, + ip, + ) + ) + + +@dsl_user_op +def st_v4_hint(addr: Int64, words, pol: Int64, *, loc=None, ip=None): + ops = [Int64(addr).ir_value(loc=loc, ip=ip)] + [ + Uint32(w).ir_value(loc=loc, ip=ip) for w in words + ] + ops.append(Int64(pol).ir_value(loc=loc, ip=ip)) + _asm( + None, + ops, + "st.global.L2::cache_hint.v4.b32 [$0], {$1, $2, $3, $4}, $5;", + "l,r,r,r,r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def ld_v2(addr: Int64, pf=".L2::128B", *, loc=None, ip=None): + res = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + pf + ".v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(2) + ] + + +@dsl_user_op +def ld_b32(addr: Int64, *, loc=None, ip=None): + return [ + Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + _LDQ + ".L2::128B.b32 $0, [$1];", + "=r,l", + loc, + ip, + ) + ) + ] + + +@dsl_user_op +def ld_b16(addr: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\nld.global.nc.b16 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def ld_b16_l2_128(addr: Int64, pf=".L2::128B", *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\n" + _LDQ + pf + ".b16 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def ld_bias_b16x2(addr: Int64, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 h;\nld.global.nc.b16 h, [$1];\nmov.b32 $0, {h, h};\n}", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def st_vec(addr: Int64, words, *, loc=None, ip=None): + n = len(words) + ops = [Int64(addr).ir_value(loc=loc, ip=ip)] + [ + Uint32(w).ir_value(loc=loc, ip=ip) for w in words + ] + if n == 4: + text = "st.global.v4.b32 [$0], {$1, $2, $3, $4};" + elif n == 2: + text = "st.global.v2.b32 [$0], {$1, $2};" + else: + text = "st.global.b32 [$0], $1;" + _asm(None, ops, text, "l" + ",r" * n, loc, ip, side_effects=True) + + +@dsl_user_op +def st_b16(addr: Int64, word, *, loc=None, ip=None): + _asm( + None, + [Int64(addr).ir_value(loc=loc, ip=ip), Uint32(word).ir_value(loc=loc, ip=ip)], + "{\n.reg .b16 l, h;\nmov.b32 {l, h}, $1;\nst.global.b16 [$0], l;\n}", + "l,r", + loc, + ip, + side_effects=True, + ) + + +def _gate_pair_ptx(): + return "\n".join( + [ + "{", + ".reg .b32 w, lo, hi;", + ".reg .b64 p, kl, k1;", + ".reg .f32 a0, a1;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + "add.rn.bf16x2 w, $1, $3;", + "shl.b32 lo, w, 16;", + "and.b32 hi, w, 0xffff0000;", + "mov.b64 p, {lo, hi};", + "mul.rn.f32x2 p, p, kl;", + "mov.b64 {a0, a1}, p;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "mov.b64 p, {a0, a1};", + "add.rn.f32x2 p, p, k1;", + "mov.b64 {a0, a1}, p;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + "mul.rn.bf16x2 $0, w, $2;", + "}", + ] + ) + + +@dsl_user_op +def gate_pair(g, x, b, *, loc=None, ip=None) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + ] + return Uint32(_asm(T.i32(), ops, _gate_pair_ptx(), "=r,r,r,r", loc, ip)) + + +def _gate_vec_exp_interleaved_ptx(npairs): + """Two BF16x2 pairs with independent MUFU operations issued together.""" + assert npairs == 2 + s = [ + "{", + ".reg .b32 w0, w1, lo, hi;", + ".reg .b64 p0, p1, kl, k1;", + ".reg .f32 a0, a1, a2, a3;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + ] + s += [ + "add.rn.bf16x2 w0, $2, $6;", + "add.rn.bf16x2 w1, $3, $6;", + "shl.b32 lo, w0, 16;", + "and.b32 hi, w0, 0xffff0000;", + "mov.b64 p0, {lo, hi};", + "mul.rn.f32x2 p0, p0, kl;", + "mov.b64 {a0, a1}, p0;", + "shl.b32 lo, w1, 16;", + "and.b32 hi, w1, 0xffff0000;", + "mov.b64 p1, {lo, hi};", + "mul.rn.f32x2 p1, p1, kl;", + "mov.b64 {a2, a3}, p1;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "ex2.approx.ftz.f32 a2, a2;", + "ex2.approx.ftz.f32 a3, a3;", + "mov.b64 p0, {a0, a1};", + "mov.b64 p1, {a2, a3};", + "add.rn.f32x2 p0, p0, k1;", + "add.rn.f32x2 p1, p1, k1;", + "mov.b64 {a0, a1}, p0;", + "mov.b64 {a2, a3}, p1;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "rcp.approx.ftz.f32 a2, a2;", + "rcp.approx.ftz.f32 a3, a3;", + "cvt.rn.bf16x2.f32 w0, a1, a0;", + "cvt.rn.bf16x2.f32 w1, a3, a2;", + "mul.rn.bf16x2 $0, w0, $4;", + "mul.rn.bf16x2 $1, w1, $5;", + "}", + ] + return "\n".join(s) + + +@dsl_user_op +def gate_vec_exp_interleaved(gs, xs, b, *, loc=None, ip=None): + n = len(gs) + ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in gs + xs + [b]] + res = _asm( + llvm.StructType.get_literal([T.i32()] * n), + ops, + _gate_vec_exp_interleaved_ptx(n), + ",".join(["=r"] * n + ["r"] * len(ops)), + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n) + ] + + +def _gate_vec_ptx(npairs): + s = [ + "{", + ".reg .b32 w, lo, hi;", + ".reg .b64 p, kl, k1;", + ".reg .f32 a0, a1;", + "mov.b64 kl, 0xBFB8AA3BBFB8AA3B;", + "mov.b64 k1, 0x3F8000003F800000;", + ] + for j in range(npairs): + g = f"${npairs + j}" + x = f"${2 * npairs + j}" + s += [ + f"add.rn.bf16x2 w, {g}, ${3 * npairs};", + "shl.b32 lo, w, 16;", + "and.b32 hi, w, 0xffff0000;", + "mov.b64 p, {lo, hi};", + "mul.rn.f32x2 p, p, kl;", + "mov.b64 {a0, a1}, p;", + "ex2.approx.ftz.f32 a0, a0;", + "ex2.approx.ftz.f32 a1, a1;", + "mov.b64 p, {a0, a1};", + "add.rn.f32x2 p, p, k1;", + "mov.b64 {a0, a1}, p;", + "rcp.approx.ftz.f32 a0, a0;", + "rcp.approx.ftz.f32 a1, a1;", + "cvt.rn.bf16x2.f32 w, a1, a0;", + f"mul.rn.bf16x2 ${j}, w, {x};", + ] + s.append("}") + return "\n".join(s) + + +@dsl_user_op +def gate_vec(gs, xs, b, *, loc=None, ip=None): + """Gate 2 or 4 packed BF16x2 pairs sharing one bias word.""" + n = len(gs) + ops = [Uint32(v).ir_value(loc=loc, ip=ip) for v in gs + xs + [b]] + res = _asm( + llvm.StructType.get_literal([T.i32()] * n), + ops, + _gate_vec_ptx(n), + ",".join(["=r"] * n + ["r"] * len(ops)), + loc, + ip, + ) + return [ + Uint32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n) + ] + + +def _gate_one_ptx(): + return "\n".join( + [ + "{", + ".reg .b16 g, b, x, w;", + ".reg .f32 a, r;", + "cvt.u16.u32 g, $1;", + "cvt.u16.u32 x, $2;", + "cvt.u16.u32 b, $3;", + "add.rn.bf16 w, g, b;", + "cvt.f32.bf16 a, w;", + "mul.rn.f32 a, a, 0fBFB8AA3B;", + "ex2.approx.ftz.f32 a, a;", + "add.rn.f32 a, a, 0f3F800000;", + "rcp.approx.ftz.f32 r, a;", + "cvt.rn.bf16.f32 w, r;", + "mul.rn.bf16 w, w, x;", + "mov.b32 $0, {w, w};", + "}", + ] + ) + + +@dsl_user_op +def gate_one(g, x, b, *, loc=None, ip=None) -> Uint32: + ops = [ + Uint32(g).ir_value(loc=loc, ip=ip), + Uint32(x).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + ] + return Uint32(_asm(T.i32(), ops, _gate_one_ptx(), "=r,r,r,r", loc, ip)) + + +def _ld_words(addr, VEC, PF256): + if const_expr(VEC == 8): + return ld_v4(addr, ".L2::256B" if PF256 else "") + elif const_expr(VEC == 4): + return ld_v2(addr, ".L2::256B" if PF256 else ".L2::128B") + elif const_expr(VEC == 2): + return ld_b32(addr) + return [ld_b16_l2_128(addr, ".L2::256B" if PF256 else ".L2::128B")] + + +@cute.kernel +def _gate2d_kernel( + x_addr: Int64, + g_addr: Int64, + b_addr: Int64, + o_addr: Int64, + n_rows: Int32, + W: cutlass.Constexpr, + D: cutlass.Constexpr, + G_STRIDE: cutlass.Constexpr, + VEC: cutlass.Constexpr, + E: cutlass.Constexpr, + R: cutlass.Constexpr, + BLOCK: cutlass.Constexpr, + BIAS_MODE: cutlass.Constexpr, + ROW_EXACT: cutlass.Constexpr, + ONE_ROW: cutlass.Constexpr, + ROWS_ON_X: cutlass.Constexpr, + PF256: cutlass.Constexpr, + INTERLEAVED8: cutlass.Constexpr, +): + tidx, _, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() + if const_expr(not ROWS_ON_X): + bx, by = by, bx + cta_cols = BLOCK * VEC * E + col_exact = W % cta_cols == 0 + nwords = (VEC + 1) // 2 + row0 = Uint32(0) if const_expr(ONE_ROW) else Uint32(bx) * Uint32(R) + colbase = Uint32(by) * Uint32(cta_cols) + Uint32(tidx) * Uint32(VEC) + + cols = [] + colcs = [] + bws = [] + for e in cutlass.range_constexpr(E): + col = colbase + Uint32(e * BLOCK * VEC) + colc = col if const_expr(col_exact) else cutlass.min(col, Uint32(W - VEC)) + cols.append(col) + colcs.append(colc) + if const_expr(BIAS_MODE == 1): + bw = [ld_bias_b16x2(b_addr + Int64(colc // Uint32(D)) * 2)] * nwords + else: + bw = [] + for j in cutlass.range_constexpr(nwords): + lo = ld_b16(b_addr + Int64((colc + Uint32(2 * j)) // Uint32(D)) * 2) + if const_expr(VEC > 1): + hi = ld_b16( + b_addr + Int64((colc + Uint32(2 * j + 1)) // Uint32(D)) * 2 + ) + bw.append(lo | (hi << 16)) + else: + bw.append(lo) + bws.append(bw) + + rows = [] + xs = [] + gs = [] + for r in cutlass.range_constexpr(R): + row = Uint32(0) if const_expr(ONE_ROW) else row0 + Uint32(r) + rowc = ( + row + if const_expr(ONE_ROW or ROW_EXACT) + else cutlass.min(row, Uint32(n_rows) - Uint32(1)) + ) + rows.append(row) + if const_expr(ONE_ROW): + x_row = x_addr + g_row = g_addr + else: + x_row = x_addr + Int64(rowc) * (W * 2) + g_row = g_addr + Int64(rowc) * (G_STRIDE * 2) + xr = [] + gr = [] + for e in cutlass.range_constexpr(E): + xr.append(_ld_words(x_row + Int64(colcs[e]) * 2, VEC, PF256)) + gr.append(_ld_words(g_row + Int64(colcs[e]) * 2, VEC, PF256)) + xs.append(xr) + gs.append(gr) + + outs = [] + if const_expr(VEC == 1): + flat_g = [] + flat_x = [] + flat_b = [] + for r in cutlass.range_constexpr(R): + for e in cutlass.range_constexpr(E): + flat_g.append(gs[r][e][0]) + flat_x.append(xs[r][e][0]) + flat_b.append(bws[e][0]) + n = R * E + flat_o = [] + for i in cutlass.range_constexpr(0, n, 2): + if const_expr(i + 1 < n): + rr = gate_pair( + flat_g[i] | (flat_g[i + 1] << 16), + flat_x[i] | (flat_x[i + 1] << 16), + (flat_b[i] & Uint32(0xFFFF)) | (flat_b[i + 1] << 16), + ) + flat_o.append(rr) + flat_o.append(rr >> 16) + else: + flat_o.append(gate_one(flat_g[i], flat_x[i], flat_b[i])) + for r in cutlass.range_constexpr(R): + outs.append([[flat_o[r * E + e]] for e in range(E)]) + else: + for r in cutlass.range_constexpr(R): + orow = [] + for e in cutlass.range_constexpr(E): + if const_expr(VEC >= 4 and BIAS_MODE == 1): + if const_expr(VEC == 4): + ow = gate_vec_exp_interleaved(gs[r][e], xs[r][e], bws[e][0]) + elif const_expr(INTERLEAVED8): + ow = gate_vec_exp_interleaved( + gs[r][e][0:2], xs[r][e][0:2], bws[e][0] + ) + gate_vec_exp_interleaved( + gs[r][e][2:4], xs[r][e][2:4], bws[e][0] + ) + else: + ow = gate_vec(gs[r][e], xs[r][e], bws[e][0]) + else: + ow = [ + gate_pair(gs[r][e][j], xs[r][e][j], bws[e][j]) + for j in cutlass.range_constexpr(nwords) + ] + orow.append(ow) + outs.append(orow) + + st_pol = make_policy("evict_last") if const_expr(VEC == 8) else None + for r in cutlass.range_constexpr(R): + o_row = o_addr if const_expr(ONE_ROW) else o_addr + Int64(rows[r]) * (W * 2) + for e in cutlass.range_constexpr(E): + addr = o_row + Int64(cols[e]) * 2 + if const_expr(ROW_EXACT and col_exact): + if const_expr(VEC == 1): + st_b16(addr, outs[r][e][0]) + elif const_expr(st_pol is not None): + st_v4_hint(addr, outs[r][e], st_pol) + else: + st_vec(addr, outs[r][e]) + else: + if const_expr(ROW_EXACT): + ok = cols[e] < Uint32(W) + elif const_expr(col_exact): + ok = rows[r] < Uint32(n_rows) + else: + ok = (rows[r] < Uint32(n_rows)) & (cols[e] < Uint32(W)) + if ok: + if const_expr(VEC == 1): + st_b16(addr, outs[r][e][0]) + elif const_expr(st_pol is not None): + st_v4_hint(addr, outs[r][e], st_pol) + else: + st_vec(addr, outs[r][e]) + + +@cute.jit +def _gate2d_launch( + x_addr: Int64, + g_addr: Int64, + b_addr: Int64, + o_addr: Int64, + n_rows: Int32, + stream: cuda.CUstream, + W: cutlass.Constexpr, + D: cutlass.Constexpr, + G_STRIDE: cutlass.Constexpr, + VEC: cutlass.Constexpr, + E: cutlass.Constexpr, + R: cutlass.Constexpr, + BLOCK: cutlass.Constexpr, + BIAS_MODE: cutlass.Constexpr, + ROW_EXACT: cutlass.Constexpr, + ONE_ROW: cutlass.Constexpr, + ROWS_ON_X: cutlass.Constexpr, + CLUSTER_LAUNCH: cutlass.Constexpr, + PF256: cutlass.Constexpr, + INTERLEAVED8: cutlass.Constexpr, +): + ncb = (W + BLOCK * VEC * E - 1) // (BLOCK * VEC * E) + nrb = (n_rows + (R - 1)) // R + grid = [nrb, ncb, 1] if const_expr(ROWS_ON_X) else [ncb, nrb, 1] + kernel = _gate2d_kernel( + x_addr, + g_addr, + b_addr, + o_addr, + n_rows, + W, + D, + G_STRIDE, + VEC, + E, + R, + BLOCK, + BIAS_MODE, + ROW_EXACT, + ONE_ROW, + ROWS_ON_X, + PF256, + INTERLEAVED8, + ) + # Small grids: launching as (trivial 1x1x1) clusters measured ~0.1 us lower end-to- + # end + # latency on B300, and adding the cooperative (co-resident) attribute on top of that + # a + # further 1-3 timer ticks; large grids lose CTA launch throughput this way. + if const_expr(CLUSTER_LAUNCH == 2): + kernel.launch( + grid=grid, + block=[BLOCK, 1, 1], + cluster=[1, 1, 1], + cooperative=True, + stream=stream, + preferred_smem_carveout=0, + min_blocks_per_mp=1, + ) + elif const_expr(CLUSTER_LAUNCH == 1): + kernel.launch( + grid=grid, + block=[BLOCK, 1, 1], + cluster=[1, 1, 1], + stream=stream, + preferred_smem_carveout=0, + ) + else: + kernel.launch( + grid=grid, block=[BLOCK, 1, 1], stream=stream, preferred_smem_carveout=0 + ) + + +_cache: dict[tuple, Callable] = {} +_no_coop = set() +_CLUSTER_MAX_CTAS = 768 +# Cooperative 1x1x1-cluster launches were measured to be refused above ~1-1.5K CTAs on +# B300 +# even when plain occupancy would allow more, so cap them here (and fall back if +# refused). +_COOP_MAX_CTAS = 1024 +_sm_count: list[int] = [] + + +def _num_sms(): + if not _sm_count: + _sm_count.append( + torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + ) + return _sm_count[0] + + +def _max_vec(W, gs, ptrs): + for vec in (8, 4, 2): + if W % vec == 0 and gs % vec == 0 and all(p % (2 * vec) == 0 for p in ptrs): + return vec + return 1 + + +def _coop_fits(ncta, block, regs): + """Conservative co-residency bound (no shared memory; regs rounded up per warp).""" + per_warp = -(-regs * 32 // 256) * 256 + by_regs = 65536 // (per_warp * -(-block // 32)) + per_sm = min(32, 2048 // block, by_regs) + return ncta <= min(_COOP_MAX_CTAS, _num_sms() * per_sm) + + +def _choose(T_, W, maxvec, contiguous_gate): + """Return launch parameters and an optional cooperative register bound.""" + n = T_ * W + if n <= 512 * 1024: + # Decode/small-prefill sizes. Scalar lanes in 96-ish + # CTAs + # up to 64K elements, then 8-byte packets with ~1.5-3K warps in flight. + pf256 = False + if n <= 16 * 1024: + vec, e, block_cap = 1, 1, 128 + elif n <= 32 * 1024: + # A 256-thread scalar CTA spans 512 B: request 256 B per L2 miss. + vec, e, block_cap, pf256 = 1, 1, 256, True + elif n <= 64 * 1024: + # 512-thread scalar CTAs each cover 1 KB of a row: fetch 256 B per L2 miss. + vec, e, block_cap, pf256 = 1, 1, 512, True + elif n <= 128 * 1024: + vec, e, block_cap = 4, 1, 256 + elif n <= 256 * 1024: + vec, e, block_cap = 4, 1, 512 + else: + vec, e, block_cap = 4, 1, 384 + vec = min(vec, maxvec) + if vec != 1: + pf256 = False + need = -(-W // (vec * e)) + return vec, e, 1, min(block_cap, max(32, -(-need // 32) * 32)), pf256, None + if n <= 32 * 1024 * 1024: + vec, e, r = 8, 1, 1 + else: + vec, e, r = 8, 1, 1 if contiguous_gate else 2 + vec = min(vec, maxvec) + if vec == 8 and n > 1024 * 1024: + # Two rows per 192-thread CTA keep the T=512 primary grid co-resident. + ncb = -(-W // (192 * 8)) + if ncb * 192 * 8 - W < 192 * 8 // 2 and _coop_fits(ncb * -(-T_ // 2), 192, 32): + return 8, 1, 2, 192, False, 32 + need = -(-W // (vec * e)) + block_cap = 256 if vec == 8 and 512 * 1024 <= n <= 1024 * 1024 else 128 + if ( + vec == 8 + and n > 1024 * 1024 + and -(-T_ // r) * -(-W // (256 * vec * e)) <= _CLUSTER_MAX_CTAS + ): + block_cap = 256 # keeps ~1-1.5M-element grids within the cluster-launch tier + block = min(block_cap, max(32, -(-need // 32) * 32)) + return vec, e, r, block, False, None + + +def run(attn_output, gate, gate_bias, num_heads, head_dim, output): + T_ = attn_output.shape[0] + H = int(num_heads) + D = int(head_dim) + W = H * D + if T_ == 0 or W == 0: + return + gs = gate.stride(0) if T_ > 1 else max(gate.stride(0), W) + ptrs = [attn_output.data_ptr(), gate.data_ptr(), output.data_ptr()] + + vec, e, r, block, pf256, coop_regs = _choose(T_, W, _max_vec(W, gs, ptrs), gs == W) + bias_mode = 1 if D % vec == 0 else 2 + nrb = -(-T_ // r) + compile_stride = W if T_ == 1 else gs + ncta = nrb * -(-W // (block * vec * e)) + launch_mode = 0 + if coop_regs is not None: + launch_mode = 2 + elif ncta <= _CLUSTER_MAX_CTAS: + # Cooperative launch needs every CTA resident at once; only claim it at <= half + # the + # SM's thread capacity so it cannot fail for any register footprint used here. + launch_mode = 2 if ncta <= _num_sms() * max(1, 1024 // block) else 1 + # T=512's two-row cooperative tile exposes the final packets' math at the one-wave + # tail. + # Interleave independent exponential sigmoid operations for these packets. + interleaved8 = vec == 8 and r == 2 and launch_mode == 2 and bias_mode == 1 + key = ( + W, + D, + compile_stride, + vec, + e, + r, + block, + bias_mode, + T_ % r == 0, + T_ == 1, + nrb > 65535, + launch_mode, + pf256, + interleaved8, + ) + if key in _no_coop: + key = key[:11] + (1, pf256, interleaved8) + args = ( + Int64(ptrs[0]), + Int64(ptrs[1]), + Int64(gate_bias.data_ptr()), + Int64(ptrs[2]), + Int32(T_), + cuda.CUstream(current_stream().cuda_stream), + ) + fn = _cache.get(key) + if fn is None: + fn = cute.compile(_gate2d_launch, *args, *key) + _cache[key] = fn + if launch_mode == 2 and coop_regs is not None: + try: + fn(*args) + except ( + Exception + ) as exc: # cooperative grid refused by the driver: launch clustered + if "COOPERATIVE_LAUNCH_TOO_LARGE" not in str(exc): + raise + _no_coop.add(key) + key = key[:11] + (1, pf256, interleaved8) + fn = _cache.get(key) + if fn is None: + fn = cute.compile(_gate2d_launch, *args, *key) + _cache[key] = fn + fn(*args) + return + fn(*args) diff --git a/vllm/models/beam/nvidia/dflash2.py b/vllm/models/beam/nvidia/dflash2.py new file mode 100644 index 000000000000..cf9874ab72be --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2.py @@ -0,0 +1,623 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DFlash2 draft model for Beam targets. + +A DFlash2 drafter predicts a block of tokens at once from hidden states the +Beam target captures at `dflash_config.target_layer_ids`, then a low-rank +candidate selector rescores the top-k tokens at each block position against +the token before it. The drafter shares the target's embeddings and LM head. + +Nothing here is torch-compiled. The DFlash speculator captures the whole draft +step (decoder layers, candidate head and selector) in its own CUDA graphs, and +the context K/V precompute runs eagerly as a few fused kernels. +""" + +from collections.abc import Iterable + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm import _custom_ops as ops +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.qwen3_dflash2 import DFlashGroupedConv +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) + +from .dflash2_ops import candidate_edges + +logger = init_logger(__name__) + + +class DFlash2BeamAttention(nn.Module): + """Query-token attention; the context K/V is already in the KV cache.""" + + def __init__( + self, + config, + layer_idx: int, + cache_config: CacheConfig | None, + prefix: str, + ) -> None: + super().__init__() + hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + total_num_heads = config.num_attention_heads + total_num_kv_heads = config.num_key_value_heads + assert total_num_heads % tp_size == 0 + if total_num_kv_heads >= tp_size: + assert total_num_kv_heads % tp_size == 0 + else: + assert tp_size % total_num_kv_heads == 0 + self.num_heads = total_num_heads // tp_size + self.num_kv_heads = max(1, total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + # Every layer is non-causal sliding-window attention over the block. + assert config.layer_types[layer_idx] == "sliding_attention" + assert not config.is_causal + self.causal = False + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + total_num_heads, + total_num_kv_heads, + bias=False, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + total_num_heads * self.head_dim, + hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + is_neox_style=True, + rope_parameters=config.rope_parameters, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.head_dim**-0.5, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + per_layer_sliding_window=config.sliding_window, + prefix=f"{prefix}.attn", + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + from .attention_ops import beam_qk_norm_rope + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + cache, slots = None, None + mapping = get_forward_context().slot_mapping + assert isinstance(mapping, dict) + if mapping and self.attn.kv_cache.dtype == q.dtype == torch.bfloat16: + slots = mapping[self.attn.layer_name] + cache = self.attn.kv_cache + assert self.attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + assert self.attn.impl.dcp_world_size == 1 + assert self.attn.kv_sharing_target_layer_name is None + assert cache.dtype == q.dtype == torch.bfloat16 + assert cache.ndim == 4 + rope = self.rotary_emb + q, k = beam_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + positions, + rope.cos_sin_cache, + self.head_dim, + rope.rotary_dim, + rope.is_neox_style, + value=v, + kv_cache=cache, + slot_mapping=slots, + ) + output = self.attn(q, k, v) if cache is None else self.attn(q, None, None) + output, _ = self.o_proj(output) + return output + + +class DFlash2BeamMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, prefix: str) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.gate_up_proj(x) + x, _ = self.down_proj(self.act_fn(x)) + return x + + +class DFlash2BeamDecoderLayer(nn.Module): + """Takes the normalized input and returns the next layer's normalized input, + fusing each grouped convolution's second half into the residual norm.""" + + def __init__( + self, + vllm_config: VllmConfig, + config, + layer_idx: int, + prefix: str, + ) -> None: + super().__init__() + self.self_attn = DFlash2BeamAttention( + config, layer_idx, vllm_config.cache_config, f"{prefix}.self_attn" + ) + self.mlp = DFlash2BeamMLP( + config.hidden_size, config.intermediate_size, f"{prefix}.mlp" + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + draft_config = config.dflash_config + assert vllm_config.speculative_config is not None + conv_args = dict( + hidden_size=config.hidden_size, + taps=int(draft_config["conv_kernel_size"]), + group_size=int(draft_config["conv_group_size"]), + # Query tokens per request: the bonus token plus the mask tokens. + block_size=1 + vllm_config.speculative_config.num_speculative_tokens, + params_dtype=vllm_config.model_config.dtype, + ) + self.attention_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "attention_conv") + ) + self.mlp_conv = DFlashGroupedConv( + **conv_args, prefix=maybe_prefix(prefix, "mlp_conv") + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor, + next_norm: RMSNorm, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .dflash2_conv_ops import conv_residual_norm + + conv = self.attention_conv + hidden_states, coefficients = conv.prepare(hidden_states) + hidden_states = self.self_attn(positions, hidden_states) + hidden_states, residual = conv_residual_norm( + hidden_states, + coefficients, + conv.base_kernel[1], + residual, + self.post_attention_layernorm.weight, + self.post_attention_layernorm.variance_epsilon, + conv.block_size, + conv.group_size, + ) + conv = self.mlp_conv + hidden_states, coefficients = conv.prepare(hidden_states) + hidden_states = self.mlp(hidden_states) + return conv_residual_norm( + hidden_states, + coefficients, + conv.base_kernel[1], + residual, + next_norm.weight, + next_norm.variance_epsilon, + conv.block_size, + conv.group_size, + ) + + +class DFlash2BeamCandidateSelector(nn.Module): + """Scores each candidate token against every candidate one position back. + + For block position `p`, candidate `c` and predecessor `q`, the score is the + candidate's own logit plus ``, where + `h_p` is the drafter's hidden state projected to the codebook rank. The + predecessor of the first position is the anchor token. + """ + + def __init__( + self, + hidden_size: int, + vocab_size: int, + rank: int, + top_k: int, + params_dtype: torch.dtype, + prefix: str, + ) -> None: + super().__init__() + self.top_k = top_k + self.predecessor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.successor_codebook = nn.Parameter( + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False + ) + self.hidden_projection = ReplicatedLinear( + hidden_size, + rank, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "hidden_projection"), + return_bias=False, + ) + + def forward( + self, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden_states: torch.Tensor, + anchor_token_ids: torch.Tensor, + ) -> torch.Tensor: + """Edge scores `[requests, steps, top_k (predecessor), top_k]`.""" + assert hidden_states.dtype == torch.bfloat16 + return candidate_edges( + self.predecessor_codebook, + self.successor_codebook, + candidate_ids, + unary_logits, + self.hidden_projection(hidden_states), + anchor_token_ids, + self.top_k, + ) + + +class DFlash2BeamModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + }, + ) + + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + draft_config = config.dflash_config + params_dtype = vllm_config.model_config.dtype + + # Replaced by the target's embedding when the speculator loads the draft. + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + # Numbered after the target's layers so the attention layer names, which + # key the KV cache, do not collide with the target's. + current_vllm_config = get_current_vllm_config() + self.layers = nn.ModuleList( + [ + DFlash2BeamDecoderLayer( + current_vllm_config, + config, + layer_idx, + prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + + self.fc = ReplicatedLinear( + config.hidden_size * len(draft_config["target_layer_ids"]), + config.hidden_size, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "fc"), + return_bias=False, + ) + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self._context_cache_ptrs: torch.Tensor | None = None + self._context_scale_ptrs: torch.Tensor | None = None + self._context_cache_addresses: tuple[int, ...] | None = None + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.candidate_selector = DFlash2BeamCandidateSelector( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + rank=int(draft_config["selector_rank"]), + top_k=int(draft_config["selector_top_k"]), + params_dtype=params_dtype, + prefix=maybe_prefix(prefix, "candidate_selector"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def build_context_kv_weights(self) -> None: + """Stack every layer's K/V projection and K-norm weight. + + The context K/V for all layers then comes from one GEMM and one grouped + norm. Must run after the weights are loaded, since it copies them. + """ + attns = [layer.self_attn for layer in self.layers] + self._context_kv_weight = torch.cat( + [attn.qkv_proj.weight[attn.q_size :] for attn in attns] + ) + self._context_k_norm_weight = torch.stack( + [attn.k_norm.weight for attn in attns] + ).contiguous() + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project the combined target states to each layer's K/V and cache them. + + The drafter's queries attend to these context keys, which carry the + layer's K-norm and RoPE like the queries' own keys. Profiling uses a + `None` slot mapping to run the projection without writing the cache. + """ + if not hasattr(self, "_context_kv_weight"): + # Dummy weights skip `load_weights`. + logger.warning_once( + "The Beam DFlash2 context K/V weights were built lazily; this is " + "expected only with dummy weights." + ) + self.build_context_kv_weights() + + attn0 = self.layers[0].self_attn + num_layers = len(self.layers) + num_ctx = context_states.shape[0] + num_kv_heads, head_dim = attn0.num_kv_heads, attn0.head_dim + + normed = torch.empty_like(context_states) + ops.rms_norm( + normed, + context_states, + self.hidden_norm.weight, + self.hidden_norm.variance_epsilon, + ) + kv = F.linear(normed, self._context_kv_weight).view( + num_ctx, num_layers, 2, num_kv_heads, head_dim + ) + self._store_fused_context_kv(kv, context_positions, context_slot_mapping) + + def _store_fused_context_kv( + self, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + from .dflash2_context_ops import store_context_kv + + if slot_mapping is None: + return + if isinstance(slot_mapping, (list, tuple)): + slots = slot_mapping[0] + if slots is None: + assert all(s is None for s in slot_mapping) + return + assert all( + s is not None + and s.data_ptr() == slots.data_ptr() + and s.shape == slots.shape + and s.stride() == slots.stride() + for s in slot_mapping + ) + else: + slots = slot_mapping + first = self.layers[0].self_attn + rotary = first.rotary_emb + assert kv.dtype == torch.bfloat16 + assert rotary.rotary_dim == first.head_dim + assert rotary.cos_sin_cache.dtype == kv.dtype + assert slots.stride(0) == 1 + attns = [layer.self_attn.attn for layer in self.layers] + caches = [attn.kv_cache for attn in attns] + cache = caches[0] + cache_dtype = attns[0].kv_cache_dtype + fp8_dtype = { + "fp8": torch.float8_e4m3fn, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + }.get(cache_dtype) + for attn, c in zip(attns, caches): + assert attn.attn_backend.get_name() in ("FLASH_ATTN", "FLASHINFER") + assert attn.impl.dcp_world_size == 1 + assert attn.kv_sharing_target_layer_name is None + assert attn.kv_cache_dtype == cache_dtype + assert c.dtype == cache.dtype and c.device == kv.device + assert c.dtype == kv.dtype or fp8_dtype is not None + assert c.ndim == 4 and c.stride(-1) == 1 + assert c.shape[1] == kv.shape[3] + assert c.shape[-1] == 2 * kv.shape[-1] + assert c.shape == cache.shape and c.stride() == cache.stride() + cache_pointers = tuple(c.data_ptr() for c in caches) + scale_pointers = ( + tuple( + s.data_ptr() for attn in attns for s in (attn._k_scale, attn._v_scale) + ) + if fp8_dtype is not None + else () + ) + pointers = cache_pointers + scale_pointers + if self._context_cache_addresses != pointers: + self._context_cache_ptrs = torch.tensor( + cache_pointers, dtype=torch.int64, device=kv.device + ) + self._context_scale_ptrs = ( + torch.tensor(scale_pointers, dtype=torch.int64, device=kv.device) + if scale_pointers + else None + ) + self._context_cache_addresses = pointers + assert self._context_cache_ptrs is not None + assert self._context_cache_ptrs.device == kv.device + store_context_kv( + kv, + self._context_k_norm_weight, + positions, + rotary.cos_sin_cache, + self._context_cache_ptrs, + slots, + cache.view(fp8_dtype) if fp8_dtype is not None else cache, + first.k_norm.variance_epsilon, + self._context_scale_ptrs, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + residual = inputs_embeds + hidden_states = self.layers[0].input_layernorm(inputs_embeds) + for i, layer in enumerate(self.layers): + next_norm = ( + self.layers[i + 1].input_layernorm + if i + 1 < len(self.layers) + else self.norm + ) + hidden_states, residual = layer( + positions, hidden_states, residual, next_norm + ) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class DFlash2BeamForCausalLM(nn.Module): + # The checkpoint ships neither, so the speculator always aliases the target's. + has_own_embed_tokens = False + has_own_lm_head = False + # Full vocabulary: draft ids are target ids. + draft_id_to_target_id = None + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.model = DFlash2BeamModel( + vllm_config=vllm_config, + start_layer_id=vllm_config.model_config.get_total_num_hidden_layers(), + prefix=maybe_prefix(prefix, "model"), + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + draft_config = self.config.dflash_config + self.logits_processor = LogitsProcessor( + self.config.vocab_size, + soft_cap=float(draft_config["final_logit_softcapping"]), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project the concatenated target-layer states to the drafter's width.""" + expected = self.model.fc.input_size + if hidden_states.shape[-1] != expected: + raise ValueError( + f"The DFlash2 drafter expects {expected} concatenated target " + f"features but received {hidden_states.shape[-1]}; its " + "target_layer_ids do not match what the target captures." + ) + return self.model.fc(hidden_states) + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.attn.layer_name for layer in self.model.layers] + + def get_draft_attn_causal(self) -> list[bool]: + """Per-layer causality, aligned with `get_draft_kv_cache_layer_names`.""" + return [layer.self_attn.causal for layer in self.model.layers] + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def compute_candidates( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """The selector's top-k candidate ids and their logits at each position.""" + return self.logits_processor.get_top_k_tokens( + self.lm_head, hidden_states, self.model.candidate_selector.top_k + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded = {f"model.{name}" for name in self.model.load_weights(weights)} + self.model.build_context_kv_weights() + # Aliased to the target's after loading. + return loaded | {"model.embed_tokens.weight", "lm_head.weight"} diff --git a/vllm/models/beam/nvidia/dflash2_context_ops.py b/vllm/models/beam/nvidia/dflash2_context_ops.py new file mode 100644 index 000000000000..ea7ce0c016e5 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_context_ops.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["ROWS", "NUM_SLOTS"]) +def _context_kv_kernel( + KV, + WEIGHT, + POSITIONS, + ROPE, + CACHE_PTRS, + CACHE, + SCALE_PTRS, + SLOTS, + ROWS, + NUM_SLOTS, + LAYERS: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + BLOCK_STRIDE: tl.constexpr, + HEAD_STRIDE: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + PAGE: tl.constexpr, + EPS: tl.constexpr, + FP8: tl.constexpr, + FP8_MAX: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) * 4 + tl.arange(0, 4) + layer_head = tl.program_id(1) + layer = layer_head // HEADS + head = layer_head % HEADS + col = tl.arange(0, BLOCK) + slot = tl.load(SLOTS + row, (row < ROWS) & (row < NUM_SLOTS), -1).to(tl.int64) + mask = (row[:, None] < ROWS) & (slot[:, None] >= 0) & (col[None, :] < DIM) + offsets = ( + row[:, None] * (LAYERS * 2 * HEADS * DIM) + + (layer * 2 * HEADS + head) * DIM + + col[None, :] + ) + key = tl.load(KV + offsets, mask, 0).to(tl.float32) + if BLOCK == 128: + squares = (key * key).reshape((4, 4, 32)) + total = tl.sum(tl.sum(squares, 2), 1) + else: + total = tl.sum(key * key, 1) + weight = tl.load(WEIGHT + layer * DIM + col, col < DIM, 0).to(tl.float32) + key = key * tl.rsqrt(total / DIM + EPS)[:, None] * weight[None, :] + key = key.to(KV.dtype.element_ty).to(tl.float32) + half: tl.constexpr = DIM // 2 + partner = (col + half) % DIM + second = col >= half + rope_col = col % half + other = tl.gather(key, tl.broadcast_to(partner[None, :], (4, BLOCK)), 1) + position = tl.load(POSITIONS + row, (row < ROWS) & (slot >= 0), 0).to(tl.int64) + cosine = tl.load(ROPE + position[:, None] * DIM + rope_col[None, :], mask, 0) + sine = tl.load(ROPE + position[:, None] * DIM + half + rope_col[None, :], mask, 0) + key = tl.where( + second[None, :], + tl.fma(other, sine, key * cosine), + tl.fma(key, cosine, -other * sine), + ) + cache = tl.load(CACHE_PTRS + layer).to(tl.pointer_type(CACHE.dtype.element_ty)) + address = ( + slot[:, None] // PAGE * BLOCK_STRIDE + + slot[:, None] % PAGE * TOKEN_STRIDE + + head * HEAD_STRIDE + + col[None, :] + ) + value = tl.load(KV + offsets + HEADS * DIM, mask, 0).to(tl.float32) + if FP8: + # Preserve BF16 rounding after RoPE. + key = key.to(KV.dtype.element_ty).to(tl.float32) + k_scale = tl.load(SCALE_PTRS + layer * 2).to(tl.pointer_type(tl.float32)) + v_scale = tl.load(SCALE_PTRS + layer * 2 + 1).to(tl.pointer_type(tl.float32)) + key = tl.minimum(tl.maximum(key / tl.load(k_scale), -FP8_MAX), FP8_MAX) + value = tl.minimum(tl.maximum(value / tl.load(v_scale), -FP8_MAX), FP8_MAX) + tl.store(cache + address, key, mask) + tl.store(cache + address + DIM, value, mask) + + +def store_context_kv( + kv: torch.Tensor, + weight: torch.Tensor, + positions: torch.Tensor, + rope: torch.Tensor, + cache_ptrs: torch.Tensor, + slots: torch.Tensor, + cache: torch.Tensor, + eps: float, + scale_ptrs: torch.Tensor | None = None, +) -> None: + """Normalize, apply NeoX RoPE, and scatter all draft layers' context K/V.""" + rows, layers, _, heads, dim = kv.shape + fp8 = cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + assert cache.dtype == kv.dtype or fp8 + assert not fp8 or scale_ptrs is not None + if rows: + _context_kv_kernel[(triton.cdiv(rows, 4), layers * heads)]( + kv, + weight, + positions, + rope, + cache_ptrs, + cache, + scale_ptrs, + slots, + rows, + slots.numel(), + layers, + heads, + dim, + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.shape[2], + eps, + fp8, + torch.finfo(cache.dtype).max if fp8 else 0.0, + triton.next_power_of_2(dim), + num_warps=4, + enable_fp_fusion=False, + ) diff --git a/vllm/models/beam/nvidia/dflash2_conv_ops.py b/vllm/models/beam/nvidia/dflash2_conv_ops.py new file mode 100644 index 000000000000..4ce6508b7040 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_conv_ops.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _conv_residual_norm_kernel( + X, + DELTA, + BASE, + RESIDUAL, + WEIGHT, + OUTPUT, + X_STRIDE: tl.constexpr, + DELTA_ROW_STRIDE: tl.constexpr, + DELTA_TAP_STRIDE: tl.constexpr, + RESIDUAL_STRIDE: tl.constexpr, + WIDTH: tl.constexpr, + GROUP: tl.constexpr, + TAPS: tl.constexpr, + QUERY: tl.constexpr, + EPS: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + mask = col < WIDTH + group = col // GROUP + coefficient = tl.load(BASE + col, mask, 0).to(tl.float32) + tl.load( + DELTA + row * DELTA_ROW_STRIDE + group, mask, 0 + ).to(tl.float32) + value = tl.load(X + row * X_STRIDE + col, mask, 0).to(tl.float32) + result = coefficient * value + for tap in tl.static_range(1, TAPS): + valid = mask & (row % QUERY >= tap) + coefficient = tl.load(BASE + tap * WIDTH + col, valid, 0).to( + tl.float32 + ) + tl.load( + DELTA + row * DELTA_ROW_STRIDE + tap * DELTA_TAP_STRIDE + group, valid, 0 + ).to(tl.float32) + value = tl.load(X + (row - tap) * X_STRIDE + col, valid, 0).to(tl.float32) + result = tl.fma(coefficient, value, result) + result = result.to(X.dtype.element_ty).to(tl.float32) + residual = tl.load(RESIDUAL + row * RESIDUAL_STRIDE + col, mask, 0).to(tl.float32) + result = (result + residual).to(RESIDUAL.dtype.element_ty).to(tl.float32) + tl.store(RESIDUAL + row * RESIDUAL_STRIDE + col, result, mask) + variance = tl.sum(result * result, 0) / WIDTH + weight = tl.load(WEIGHT + col, mask, 0).to(tl.float32) + normalized = result * tl.rsqrt(variance + EPS) * weight + tl.store(OUTPUT + row * WIDTH + col, normalized, mask) + + +def conv_residual_norm( + hidden: torch.Tensor, + coefficients: torch.Tensor, + base: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + query_len: int, + group_size: int, + *, + num_warps: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Finish grouped convolution and update the residual before RMSNorm.""" + rows, width = hidden.shape + if num_warps is None: + num_warps = 16 if rows <= 256 else (8 if rows < 4096 else 4) + output = torch.empty_like(hidden) + if rows: + _conv_residual_norm_kernel[(rows,)]( + hidden, + coefficients, + base, + residual, + weight, + output, + hidden.stride(0), + coefficients.stride(0), + coefficients.stride(1), + residual.stride(0), + width, + group_size, + base.shape[0], + query_len, + eps, + triton.next_power_of_2(width), + num_warps=num_warps, + enable_fp_fusion=False, + ) + return output, residual diff --git a/vllm/models/beam/nvidia/dflash2_ops.py b/vllm/models/beam/nvidia/dflash2_ops.py new file mode 100644 index 000000000000..ed1dbfce5c01 --- /dev/null +++ b/vllm/models/beam/nvidia/dflash2_ops.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _candidate_edges_kernel( + PREDECESSORS, + SUCCESSORS, + IDS, + UNARY, + HIDDEN, + ANCHORS, + OUTPUT, + ANCHOR_STRIDE: tl.constexpr, + STEPS: tl.constexpr, + TOP_K: tl.constexpr, + RANK: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_R: tl.constexpr, +): + position = tl.program_id(0) + predecessor = tl.program_id(1) + step = position % STEPS + if step == 0: + pred_id = tl.load(ANCHORS + position // STEPS * ANCHOR_STRIDE).to(tl.int64) + else: + pred_id = tl.load(IDS + (position - 1) * TOP_K + predecessor).to(tl.int64) + r = tl.arange(0, BLOCK_R) + c = tl.arange(0, BLOCK_K) + pred = tl.load(PREDECESSORS + pred_id.to(tl.int64) * RANK + r, r < RANK, 0) + hidden = tl.load(HIDDEN + position * RANK + r, r < RANK, 0) + # Preserve the materialized product and GEMM's output rounding. + weighted = (pred.to(tl.float32) * hidden.to(tl.float32)).to(hidden.dtype) + ids = tl.load(IDS + position * TOP_K + c, c < TOP_K, 0).to(tl.int64) + successors = tl.load( + SUCCESSORS + ids[:, None] * RANK + r[None, :], + (c[:, None] < TOP_K) & (r[None, :] < RANK), + 0, + ) + scores = tl.sum(successors.to(tl.float32) * weighted[None, :].to(tl.float32), 1) + scores = scores.to(hidden.dtype).to(tl.float32) + unary = tl.load(UNARY + position * TOP_K + c, c < TOP_K, 0).to(tl.float32) + tl.store( + OUTPUT + (position * TOP_K + predecessor) * TOP_K + c, scores + unary, c < TOP_K + ) + + +def candidate_edges( + predecessor_table: torch.Tensor, + successor_table: torch.Tensor, + candidate_ids: torch.Tensor, + unary_logits: torch.Tensor, + hidden: torch.Tensor, + anchor_token_ids: torch.Tensor, + top_k: int, +) -> torch.Tensor: + """Gather codebooks and score candidate edges without intermediate tensors.""" + batch, steps, _ = candidate_ids.shape + rank = hidden.shape[-1] + output = torch.empty( + (batch, steps, top_k, top_k), dtype=unary_logits.dtype, device=hidden.device + ) + if batch * steps: + _candidate_edges_kernel[(batch * steps, top_k)]( + predecessor_table, + successor_table, + candidate_ids, + unary_logits, + hidden, + anchor_token_ids, + output, + anchor_token_ids.stride(0), + steps, + top_k, + rank, + triton.next_power_of_2(top_k), + triton.next_power_of_2(rank), + num_warps=4, + enable_fp_fusion=False, + ) + return output diff --git a/vllm/models/beam/nvidia/layernorm.py b/vllm/models/beam/nvidia/layernorm.py new file mode 100644 index 000000000000..4abab83c62f5 --- /dev/null +++ b/vllm/models/beam/nvidia/layernorm.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam kernels preserving FP32 normalization and affine arithmetic.""" + +import torch +from torch import nn + +from vllm.triton_utils import tl, triton + + +class RMSNorm(nn.Module): + """Parameters for normalization performed directly by the model kernels.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, *args, **kwargs): + raise RuntimeError("Call the model's fused normalization kernels directly") + + +@triton.jit +def _rms_norm_kernel( + X, + W, + Y, + N: tl.constexpr, + HEADS: tl.constexpr, + STRIDE_TOKEN: tl.constexpr, + STRIDE_HEAD: tl.constexpr, + EPS: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + idx = tl.arange(0, BLOCK) + token = row // HEADS + head = row % HEADS + x = tl.load( + X + token * STRIDE_TOKEN + head * STRIDE_HEAD + idx, idx < N, other=0 + ).to(tl.float32) + mean = tl.sum(x * x, 0) / N + y = x * tl.rsqrt(mean + EPS) + w = tl.load(W + idx, idx < N, other=0).to(tl.float32) + y = y * w + tl.store(Y + row * N + idx, y, idx < N) + + +def beam_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Normalize strided token/head rows, rounding only after the affine multiply.""" + assert x.ndim in (2, 3) and x.stride(-1) == 1 + n = x.shape[-1] + heads = x.shape[1] if x.ndim == 3 else 1 + out = torch.empty(x.shape, device=x.device, dtype=x.dtype) + if out.numel() == 0: + return out + _rms_norm_kernel[(x.numel() // n,)]( + x, + weight, + out, + n, + heads, + x.stride(0), + x.stride(-2) if x.ndim == 3 else n, + eps, + triton.next_power_of_2(n), + num_warps=4 if n <= 1024 else 8, + enable_fp_fusion=False, + ) + return out + + +@triton.jit +def _boundary_math( + x, + r, + POST_W, + NEXT_W, + col, + dtype: tl.constexpr, + N: tl.constexpr, + POST_EPS: tl.constexpr, + NEXT_EPS: tl.constexpr, + X_SCALE, + R_SCALE, + NEXT_NORM: tl.constexpr, +): + x = x * tl.rsqrt(tl.sum(x * x, 0) / N + POST_EPS) + w = tl.load(POST_W + col, col < N, other=0).to(tl.float32) + x = (x * w).to(dtype).to(tl.float32) + x = (x * X_SCALE).to(dtype).to(tl.float32) + r = (r * R_SCALE).to(dtype).to(tl.float32) + h = (x + r).to(dtype) + if NEXT_NORM: + h_fp32 = tl.where(col < N, h.to(tl.float32), 0) + y = h_fp32 * tl.rsqrt(tl.sum(h_fp32 * h_fp32, 0) / N + NEXT_EPS) + w = tl.load(NEXT_W + col, col < N, other=0).to(tl.float32) + y = y * w + else: + y = h + return h, y + + +@triton.jit +def _norm_residual_norm_kernel( + X, + R, + POST_W, + NEXT_W, + H, + Y, + N: tl.constexpr, + X_STRIDE: tl.constexpr, + R_STRIDE: tl.constexpr, + POST_EPS: tl.constexpr, + NEXT_EPS: tl.constexpr, + X_SCALE, + R_SCALE, + NEXT_NORM: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + dtype = X.dtype.element_ty + x = tl.load(X + row * X_STRIDE + col, col < N, other=0).to(tl.float32) + r = tl.load(R + row * R_STRIDE + col, col < N, other=0).to(dtype).to(tl.float32) + h, y = _boundary_math( + x, + r, + POST_W, + NEXT_W, + col, + dtype, + N, + POST_EPS, + NEXT_EPS, + X_SCALE, + R_SCALE, + NEXT_NORM, + ) + tl.store(H + row * N + col, h, col < N) + if NEXT_NORM: + tl.store(Y + row * N + col, y, col < N) + + +def beam_norm_residual_norm( + x: torch.Tensor, + residual: torch.Tensor, + post_weight: torch.Tensor, + post_eps: float, + next_weight: torch.Tensor | None, + next_eps: float | None, + scales: tuple[float, float], +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Fuse post-norm, scaled residual add and next pre-norm, keeping each cast. + + Returns the new residual and, when ``next_weight`` is given, its + normalization for the next sublayer. + """ + assert x.ndim == residual.ndim == 2 and x.shape == residual.shape + assert x.stride(-1) == residual.stride(-1) == 1 + assert (next_weight is None) == (next_eps is None) + hidden = torch.empty(x.shape, device=x.device, dtype=x.dtype) + normalized = torch.empty_like(hidden) if next_weight is not None else None + n = x.shape[-1] + if hidden.numel(): + _norm_residual_norm_kernel[(x.shape[0],)]( + x, + residual, + post_weight, + next_weight, + hidden, + normalized, + n, + x.stride(0), + residual.stride(0), + post_eps, + next_eps or 0.0, + *scales, + next_weight is not None, + triton.next_power_of_2(n), + num_warps=4 if n <= 1024 else 8, + enable_fp_fusion=False, + ) + return hidden, normalized diff --git a/vllm/models/beam/nvidia/model.py b/vllm/models/beam/nvidia/model.py new file mode 100644 index 000000000000..b19887f2145e --- /dev/null +++ b/vllm/models/beam/nvidia/model.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Beam model compatible with HuggingFace weights. + +Beam is a decoder-only sparse MoE causal LM. Its config-gated features +are NoPE layers with SSMax query scaling, hybrid sliding-window attention, +per-head Q/K RMSNorm, sandwich norm, depth-scaled residuals, and sigmoid +attention output gating. The MoE router uses grouped sigmoid ``noaux_tc`` +selection. +""" + +import math +from collections.abc import Iterable +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.model_executor.layers.fused_moe import MoERunner +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + MixtureOfExperts, + SupportsEagle3, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.beam import BeamConfig + +from .activation import beam_softcap_logits +from .attention import BeamAttention, use_merged_qkv_gate +from .layernorm import RMSNorm, beam_norm_residual_norm, beam_rms_norm +from .moe import BeamMLP, BeamSparseMoeBlock + + +def _depth_scaled_coeffs(layer_id: int, is_attention: bool) -> tuple[float, float]: + """Depth-scaled residual coefficients ``(x_scale, r_scale)`` for one sublayer. + + ``num_adds`` counts residual adds up to and including this one, with a + 1-based layer index ``L = layer_id + 1``: ``2L`` for the attention add, + ``2L + 1`` for the MLP add. + """ + one_based = layer_id + 1 + num_adds = 2 * one_based if is_attention else 2 * one_based + 1 + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +class BeamDecoderLayer(nn.Module): + """Sandwich-norm block carrying residual and next normalized input separately.""" + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + config = cast(BeamConfig, vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.layer_id = extract_layer_index(prefix) + self.rms_norm_eps = config.rms_norm_eps + + self.self_attn = BeamAttention( + config, cache_config, quant_config, prefix=f"{prefix}.self_attn" + ) + + if config.mlp_layer_types[self.layer_id] == "dense": + self.mlp: nn.Module = BeamMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = BeamSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + + self.input_layernorm = RMSNorm(config.hidden_size) + self.pre_feedforward_layernorm = RMSNorm(config.hidden_size) + self.post_attention_layernorm = RMSNorm(config.hidden_size) + self.post_feedforward_layernorm = RMSNorm(config.hidden_size) + self.attn_scales = _depth_scaled_coeffs(self.layer_id, is_attention=True) + self.mlp_scales = _depth_scaled_coeffs(self.layer_id, is_attention=False) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + normalized_states: torch.Tensor | None = None, + next_norm_weight: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + x = ( + beam_rms_norm(hidden_states, self.input_layernorm.weight, self.rms_norm_eps) + if normalized_states is None + else normalized_states + ) + x = self.self_attn(positions, x) + hidden_states, x = beam_norm_residual_norm( + x, + hidden_states, + self.post_attention_layernorm.weight, + self.rms_norm_eps, + self.pre_feedforward_layernorm.weight, + self.rms_norm_eps, + self.attn_scales, + ) + assert x is not None + x = self.mlp(x) + return beam_norm_residual_norm( + x, + hidden_states, + self.post_feedforward_layernorm.weight, + self.rms_norm_eps, + next_norm_weight, + self.rms_norm_eps if next_norm_weight is not None else None, + self.mlp_scales, + ) + + +class BeamModel(nn.Module, EagleModelMixin): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + # Renamed so it cannot be matched as a `gate_up_proj` shard by + # `packed_modules_mapping` consumers such as LoRA. + ".self_attn.gate_proj": ".self_attn.output_gate_proj", + }, + orig_to_new_stacked={ + # weight_name: (param_name, shard_id) + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + # `.mlp.experts.N.gate_proj` does not contain `.mlp.gate_proj`, so + # expert weights fall through to `MoERunner.load_weights` (which is + # also what keeps them correct under EP and redundant experts). + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + }, + ) + + @classmethod + def get_weights_mapper(cls, config: BeamConfig, tp_size: int) -> WeightsMapper: + mapper = cls.hf_to_vllm_mapper + if not use_merged_qkv_gate(config, tp_size): + return mapper + return WeightsMapper( + orig_to_new_substr=mapper.orig_to_new_substr, + orig_to_new_stacked={ + **mapper.orig_to_new_stacked, + **{ + f".self_attn.{name}": (".self_attn.qkv_gate_proj", shard) + for shard, name in enumerate( + ("q_proj", "k_proj", "v_proj", "output_gate_proj") + ) + }, + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.hf_to_vllm_mapper = self.get_weights_mapper( + config, get_tensor_model_parallel_world_size() + ) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + # Beam drives SWA per layer off `layer_types`, passing an explicit + # `per_layer_sliding_window` (None on the full-attention layers). + # Without this, those layers would pick up `cache_config.sliding_window` + # as a fallback in `Attention.__init__` and silently get a window. + if cache_config is not None: + cache_config.sliding_window = None + + self.config = config + self.quant_config = quant_config + self.vocab_size = config.vocab_size + self.embed_tokens: VocabParallelEmbedding | PPMissingLayer + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: BeamDecoderLayer(vllm_config=vllm_config, prefix=prefix), + prefix=f"{prefix}.layers", + ) + self.norm = ( + RMSNorm(config.hidden_size) + if get_pp_group().is_last_rank + else PPMissingLayer() + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + assert input_ids is not None + hidden_states = self.embed_input_ids(input_ids) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, None + ) + normalized_states = None + for index in range(self.start_layer, self.end_layer): + next_norm = ( + self.layers[index + 1].input_layernorm + if index + 1 < self.end_layer + else self.norm + ) + hidden_states, normalized_states = self.layers[index]( + positions, + hidden_states, + normalized_states, + next_norm.weight if isinstance(next_norm, RMSNorm) else None, + ) + self._maybe_add_hidden_state( + aux_hidden_states, index + 1, hidden_states, None + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + hidden_states = ( + beam_rms_norm(hidden_states, self.norm.weight, self.config.rms_norm_eps) + if normalized_states is None + else normalized_states + ) + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class BeamForCausalLM( + nn.Module, + SupportsPP, + MixtureOfExperts, + SupportsQuant, + SupportsEagle3, +): + """Beam causal LM. LoRA support is deferred to a follow-up change.""" + + hf_to_vllm_mapper = BeamModel.hf_to_vllm_mapper + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "qkv_gate_proj": ["q_proj", "k_proj", "v_proj", "output_gate_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + fall_back_to_pt_during_load = False + + def process_weights_after_loading(self) -> None: + for module in self.model.modules(): + if isinstance(module, BeamAttention): + module.warmup_projection() + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + if parallel_config.enable_eplb: + # Expert parallelism is validated; expert rebalancing is not. + raise NotImplementedError( + "Beam has not been validated with EPLB; run without --enable-eplb." + ) + if parallel_config.use_sequence_parallel_moe: + # Data-parallel attention with expert-parallel MoE routes tokens + # through the sequence-parallel path, which is not yet validated. + raise NotImplementedError( + "Beam has not been validated with sequence-parallel MoE " + "(data parallelism with expert parallelism)." + ) + self.config = config + self.quant_config = quant_config + + self.model = BeamModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if config.tie_word_embeddings: + raise NotImplementedError("Beam requires tie_word_embeddings=False.") + # Untied LM head, in the model dtype. An fp32 head is available the + # standard way, via `head_dtype` on the config (see + # `LogitsProcessor._apply_head`), rather than by pinning the parameter + # dtype here. + self.lm_head: ParallelLMHead | PPMissingLayer + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + # Apply the softcap in fp32 while keeping the head in the model dtype. + self.final_logit_softcapping = config.final_logit_softcapping + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + # MoE metadata for the MixtureOfExperts interface. + self.moe_layers: list[MoERunner] = [] + example_layer: BeamSparseMoeBlock | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + assert isinstance(layer, BeamDecoderLayer) + if isinstance(layer.mlp, BeamSparseMoeBlock): + example_layer = layer.mlp + self.moe_layers.append(layer.mlp.experts) + if example_layer is None: + # A pipeline stage may contain only the initial dense layers. + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_shared_experts = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_redundant_experts = 0 + return + + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = config.n_group or 1 + self.num_shared_experts = 0 + moe_config = example_layer.experts.moe_config + self.num_logical_experts = moe_config.num_logical_experts + self.num_physical_experts = moe_config.num_experts + self.num_local_physical_experts = moe_config.num_local_experts + self.num_routed_experts = config.num_experts + self.num_redundant_experts = ( + self.num_physical_experts - self.num_logical_experts + ) + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.moe_layers: + layer.update_expert_map() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is not None and self.final_logit_softcapping is not None: + logits = beam_softcap_logits(logits, self.final_logit_softcapping) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm/models/beam/nvidia/moe.py b/vllm/models/beam/nvidia/moe.py new file mode 100644 index 000000000000..b079b7b9ccf9 --- /dev/null +++ b/vllm/models/beam/nvidia/moe.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig + +from .activation import beam_silu_mul + + +class BeamMLP(nn.Module): + """SwiGLU MLP for the ``dense`` entries of ``mlp_layer_types``.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported." + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + # Beam rounds once after the fp32 SiLU and gate/up multiplication. + out = beam_silu_mul(gate_up) + out, _ = self.down_proj(out) + return out + + +class BeamSparseMoeBlock(nn.Module): + """Grouped sigmoid ``noaux_tc`` MoE block.""" + + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.gate = GateLinear( + config.hidden_size, + config.num_experts, + bias=False, + params_dtype=torch.float32, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + e_score_correction_bias = nn.Parameter( + torch.zeros(config.num_experts, dtype=torch.float32) + ) + self.gate.register_parameter("e_score_correction_bias", e_score_correction_bias) + + self.experts = FusedMoEFactory( + gate=self.gate, + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + use_grouped_topk=( + config.n_group is not None and config.topk_group is not None + ), + num_expert_group=config.n_group, + topk_group=config.topk_group, + scoring_func="sigmoid", + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=config.routed_scaling_factor, + apply_routed_scale_to_output=True, + router_logits_dtype=self.gate.out_dtype, + enable_eplb=parallel_config.enable_eplb, + num_redundant_experts=parallel_config.eplb_config.num_redundant_experts, + prefix=f"{prefix}.experts", + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.experts(hidden_states=hidden_states, router_logits=hidden_states) diff --git a/vllm/models/beam/nvidia/projection.py b/vllm/models/beam/nvidia/projection.py new file mode 100644 index 000000000000..ce28ccda67dd --- /dev/null +++ b/vllm/models/beam/nvidia/projection.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed attention projections using shared low-latency GEMMs.""" + +import functools + +import torch +import torch.nn.functional as F + +from vllm.platforms import current_platform + + +@functools.cache +def _tinygemm(): + from flashinfer.jit import gen_tinygemm2_module + + return gen_tinygemm2_module().build_and_load() + + +def use_low_latency_gemm(weight: torch.Tensor) -> bool: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import is_available + + return ( + weight.dtype == torch.bfloat16 + and weight.device.type == "cuda" + and current_platform.is_device_capability_family(100) + and weight.shape in ((1792, 5120), (3584, 5120)) + and is_available() + ) + + +def warmup_packed_projection(weight: torch.Tensor) -> None: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + ll_bf16_gemm_c1_pdl_kernel, + ll_bf16_gemm_kernel, + ) + + shapes = [(weight.shape[1], weight.shape[0])] + ll_bf16_gemm_c1_pdl_kernel.warmup( + shapes=shapes, m_values=[1], output_dtype=torch.bfloat16 + ) + ll_bf16_gemm_kernel.warmup( + shapes=shapes, m_values=[2, 16], output_dtype=torch.bfloat16 + ) + _tinygemm() + + +def packed_projection( + x: torch.Tensor, weight: torch.Tensor, low_latency: bool +) -> torch.Tensor: + rows = x.shape[0] + if not low_latency or not 1 <= rows <= 16 or not x.is_contiguous(): + return F.linear(x, weight) + tiny_max = 8 if weight.shape[0] == 1792 else 4 + if 3 <= rows <= tiny_max: + output = x.new_empty((rows, weight.shape[0])) + _tinygemm().tinygemm2_nobias_op(x, weight, output, False) + return output + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm + + return ll_bf16_gemm(x, weight, output_dtype=torch.bfloat16) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 8587c02b680c..69350a53e9ca 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -122,6 +122,7 @@ def __getitem__(self, key): _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( bagel="BagelConfig", bailing_moe_v3_vl="BailingMoeV3VLConfig", + beam="BeamConfig", chatglm="ChatGLMConfig", modernvbert="ColModernVBertConfig", colpali="ColPaliConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 68f82ead6b39..e674dd5a2e85 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -18,6 +18,7 @@ "BailingMoeV3TextConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VisionConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", "BailingMoeV3VLConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl", + "BeamConfig": "vllm.transformers_utils.configs.beam", "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm", "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert", "ColPaliConfig": "vllm.transformers_utils.configs.colpali", @@ -90,6 +91,7 @@ "BailingMoeV3TextConfig", "BailingMoeV3VisionConfig", "BailingMoeV3VLConfig", + "BeamConfig", "ChatGLMConfig", "ColModernVBertConfig", "ColPaliConfig", diff --git a/vllm/transformers_utils/configs/beam.py b/vllm/transformers_utils/configs/beam.py new file mode 100644 index 000000000000..2e1241ae590d --- /dev/null +++ b/vllm/transformers_utils/configs/beam.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Beam model configuration.""" + +from transformers.configuration_utils import PreTrainedConfig + + +class BeamConfig(PreTrainedConfig): + """Configuration for the Beam decoder-only MoE causal LM. + + Full-attention layers are NoPE with SSMax query scaling; sliding-window + layers use RoPE. ``mlp_layer_types`` marks each layer's feed-forward as + ``dense`` or ``sparse`` (grouped sigmoid ``noaux_tc`` MoE). + """ + + model_type = "beam" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 2048, + intermediate_size: int = 6144, + num_hidden_layers: int = 24, + num_attention_heads: int = 32, + num_key_value_heads: int = 4, + head_dim: int | None = None, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + rms_norm_eps: float = 1e-6, + tie_word_embeddings: bool = False, + sliding_window: int = 4096, + layer_types: list[str] | None = None, + mlp_layer_types: list[str] | None = None, + moe_intermediate_size: int = 768, + num_experts_per_tok: int = 8, + num_experts: int = 128, + norm_topk_prob: bool = True, + n_group: int | None = None, + topk_group: int | None = None, + routed_scaling_factor: float = 1.0, + final_logit_softcapping: float | None = None, + pad_token_id: int | None = None, + bos_token_id: int | None = None, + eos_token_id: int | None = None, + **kwargs, + ): + if "architectures" not in kwargs: + kwargs["architectures"] = ["BeamForCausalLM"] + if layer_types is None: + layer_types = ["sliding_attention"] * num_hidden_layers + if len(layer_types) != num_hidden_layers or any( + t not in ("full_attention", "sliding_attention") for t in layer_types + ): + raise ValueError( + "layer_types must contain one valid attention type per layer." + ) + if mlp_layer_types is None: + mlp_layer_types = ["sparse"] * num_hidden_layers + if len(mlp_layer_types) != num_hidden_layers or any( + t not in ("dense", "sparse") for t in mlp_layer_types + ): + raise ValueError( + "mlp_layer_types must contain one 'dense' or 'sparse' entry per layer." + ) + rope_parameters = kwargs.setdefault( + "rope_parameters", {"rope_type": "default", "rope_theta": 10000.0} + ) + if set(rope_parameters) <= { + "sliding_attention", + "full_attention", + }: + # NoPE layers have no rotary parameters. Canonicalize before HF's + # generic RoPE validation, which cannot process null entries. + rotary_parameters = [p for p in rope_parameters.values() if p is not None] + if len(rotary_parameters) != 1: + raise ValueError("Beam expects RoPE parameters for exactly one type.") + kwargs["rope_parameters"] = rotary_parameters[0] + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.sliding_window = sliding_window + self.layer_types = layer_types + self.mlp_layer_types = mlp_layer_types + self.moe_intermediate_size = moe_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.final_logit_softcapping = final_logit_softcapping + + kwargs.update( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + ) + super().__init__(**kwargs) + + +__all__ = ["BeamConfig"] diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index fc028941cf0d..e1f97d2affcf 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -31,7 +31,9 @@ def init_speculator( ) return LiLiCorrSpeculator(vllm_config, device) - if "DFlash2DraftModel" in speculative_config.draft_model_config.architectures: + if {"DFlash2DraftModel", "DFlash2BeamDraftModel"}.intersection( + speculative_config.draft_model_config.architectures + ): from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import ( DFlash2Speculator, ) From 3497d72a72cc4145b5e5600c7569b71b9578584d Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Mon, 19 Oct 2026 10:13:24 +0800 Subject: [PATCH 838/952] [Frontend] Add Muse Spark Open reasoning and tool parsers (#62207) Signed-off-by: Jee Jee Li Signed-off-by: khluu Co-authored-by: Bugen Zhao Co-authored-by: Jiahan Chang (Cyrus) <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Isotr0py Co-authored-by: Woosuk Kwon Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com> Co-authored-by: qizixi <22851944+zixi-qi@users.noreply.github.com> Co-authored-by: haic0 <149741444+haic0@users.noreply.github.com> Co-authored-by: Tianyu Guo Co-authored-by: Shangdi Yu Co-authored-by: JohnQinAMD Co-authored-by: Duncan Moss <6038874+djmmoss@users.noreply.github.com> Co-authored-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- docs/features/tool_calling.md | 19 + rust/src/chat/src/lib.rs | 4 +- rust/src/chat/src/parser/unified.rs | 42 +- rust/src/parser/src/unified/mod.rs | 2 + rust/src/parser/src/unified/muse.rs | 1166 +++++++++++++++++ tests/tool_use/test_muse.py | 615 +++++++++ tests/tool_use/test_muse_glimmer.py | 19 + tests/tool_use/test_muse_parse_delta.py | 228 ++++ vllm/reasoning/__init__.py | 9 + .../muse_glimmer_reasoning_parser.py | 325 +---- vllm/reasoning/muse_reasoning_parser.py | 428 ++++++ vllm/tool_parsers/__init__.py | 9 + vllm/tool_parsers/muse_glimmer_tool_parser.py | 507 +------ vllm/tool_parsers/muse_tool_parser.py | 539 ++++++++ 14 files changed, 3087 insertions(+), 825 deletions(-) create mode 100644 rust/src/parser/src/unified/muse.rs create mode 100644 tests/tool_use/test_muse.py create mode 100644 tests/tool_use/test_muse_parse_delta.py create mode 100644 vllm/reasoning/muse_reasoning_parser.py create mode 100644 vllm/tool_parsers/muse_tool_parser.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 9980429324ce..00115e5b76a6 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -567,6 +567,25 @@ Supported models: Flags: `--tool-call-parser apertus` +### Muse Models (`muse`) + +The Muse parser supports the ATEM tool-call format used by Muse Spark Open. +Enable tool calling with `--enable-auto-tool-choice`, +`--tool-call-parser muse`, and `--reasoning-parser muse`. `musespark1x` is an +alias for both parsers. The `muse_glimmer` entry points share the Muse implementation while retaining +their original channel semantics, including recipient-less tool-parser content. + +The parsers route `to=self` and recipient-less messages to reasoning, +`to=user` and `to=commentary` messages to content, and tool recipients to +ATEM tool-call extraction. `to=developer` messages are omitted. Framing +special tokens are preserved for parsing and removed from surfaced content. + +The Rust frontend also provides a Muse unified parser under `muse` and +`musespark1x`, and selects it automatically for Muse Spark model names. +Existing automatic parser selection for Muse Glimmer is unchanged. + +Parser availability alone does not enable Muse Spark model execution. + ### Models with Pythonic Tool Calls (`pythonic`) A growing number of models output a python list to represent tool calls instead of using JSON. This has the advantage of inherently supporting parallel tool calls and removing ambiguity around the JSON schema required for tool calls. The `pythonic` tool parser can support such models. diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 5abdd36cc9fd..4c31d8c8e6f5 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -434,7 +434,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: beam, deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: beam, deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, muse, musespark1x, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -445,6 +445,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: beam, cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: beam, cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, muse, musespark1x, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 705133bb0a71..32cdd109f8d4 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -7,7 +7,8 @@ use std::sync::LazyLock; pub use vllm_parser::unified::{ BeamUnifiedParser, Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, - HyV4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser, + HyV4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, MuseUnifiedParser, + ResponseTemplate, UnifiedParser, }; use vllm_tokenizer::DynTokenizer; @@ -21,6 +22,9 @@ pub mod names { pub const HY_V4: &str = "hy_v4"; pub const INKLING: &str = "inkling"; pub const KIMI_K3: &str = "kimi_k3"; + pub const MUSE: &str = "muse"; + /// Alias of [`MUSE`]: the parser name of the original Muse Spark plugin. + pub const MUSE_SPARK_1X: &str = "musespark1x"; /// Executes the checkpoint's Hugging Face `response_template`. Registered for /// name resolution only: it is built from the model's template, never by its @@ -56,6 +60,8 @@ impl UnifiedParserFactory { factory.register_parser::(names::HY_V4); factory.register_parser::(names::INKLING); factory.register_parser::(names::KIMI_K3); + factory.register_parser::(names::MUSE); + factory.register_parser::(names::MUSE_SPARK_1X); // Opt-in only: no model patterns map to it. factory.register_parser::(names::HF); factory.register_parser::(names::BEAM); @@ -67,7 +73,8 @@ impl UnifiedParserFactory { .register_pattern("hy4", names::HY_V4) .register_pattern("inkling", names::INKLING) .register_pattern("kimi-k3", names::KIMI_K3) - .register_pattern("kimi_k3", names::KIMI_K3); + .register_pattern("kimi_k3", names::KIMI_K3) + .register_pattern("muse-spark", names::MUSE); factory } @@ -163,6 +170,37 @@ mod tests { factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap(); } + #[test] + fn factory_registers_muse_spark() { + let factory = UnifiedParserFactory::new(); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|start|>", 200022) + .with_regular_token("<|message|>", 200023) + .with_regular_token("<|eom|>", 200007) + .with_regular_token("<|eot|>", 200008); + + let model_id = "meta-models/Muse-Spark-open-816B-A42B-NVFP4"; + assert_eq!(factory.resolve_name_for_model(model_id), Some(names::MUSE)); + assert_eq!( + crate::ParserSelection::Auto.resolve_reasoning_name(model_id), + Some(names::MUSE) + ); + assert_eq!( + crate::ParserSelection::Auto.resolve_tool_name(model_id), + Some(names::MUSE) + ); + // Do not change automatic parser selection for existing Glimmer deployments. + assert_eq!( + factory.resolve_name_for_model("meta-models/Muse-Glimmer-30B"), + None + ); + for name in [names::MUSE, names::MUSE_SPARK_1X] { + factory.create(name, &[], Arc::new(tokenizer.clone())).unwrap(); + let explicit = crate::ParserSelection::Explicit(name.to_owned()); + crate::parser::validate_parser_overrides(&explicit, &explicit).unwrap(); + } + } + #[test] fn factory_registers_hy_v4() { let factory = UnifiedParserFactory::new(); diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 585f28cfdebd..8219f6fcb7c8 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -10,6 +10,7 @@ mod hf; mod hy; mod inkling; mod kimi_k3; +mod muse; pub use beam::BeamUnifiedParser; pub use combined::CombinedParser; @@ -18,6 +19,7 @@ pub use hf::{HfPosition, HfTemplateError, HfUnifiedParser, ResponseTemplate}; pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser}; pub use inkling::InklingUnifiedParser; pub use kimi_k3::KimiK3UnifiedParser; +pub use muse::MuseUnifiedParser; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::{DecodedText, DynTokenizer}; diff --git a/rust/src/parser/src/unified/muse.rs b/rust/src/parser/src/unified/muse.rs new file mode 100644 index 000000000000..c04d519b2552 --- /dev/null +++ b/rust/src/parser/src/unified/muse.rs @@ -0,0 +1,1166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for the Muse chat format (Muse Spark and Muse Glimmer). +//! +//! The generation prompt ends with `<|start|>assistant`; the completion body is +//! a sequence of recipient-routed segments forming one assistant turn: +//! +//! ```text +//! to=self<|message|>REASONING<|eom|><|start|>assistant +//! to=commentary<|message|>VISIBLE<|eom|><|start|>assistant +//! to=FUNCTION<|message|> +//! +//! +//! VALUE +//! +//! <|eom|><|start|>assistant +//! to=user<|message|>FINAL<|eot|> +//! ``` +//! +//! `self` bodies are reasoning, `commentary`/`user` bodies are visible text, +//! and any other recipient names a tool whose body is one ATEM block with one +//! or more `` calls. Parameters declared as strings retain their +//! raw text; other values are JSON-decoded when possible and kept as raw +//! strings otherwise. Each invoke's arguments are reassembled as a compact +//! JSON object preserving parameter order. +//! +//! The model sometimes opens the next segment without closing the current one +//! with `<|eom|>`, e.g. writing `to=FUNCTION<|message|>` straight after its +//! reasoning when calling a tool with empty arguments. Since `<|start|>` and +//! `<|message|>` never occur inside a body, either one also ends the current +//! segment and starts the next. + +use std::collections::{HashMap, HashSet}; + +use serde_json::{Map, Value}; +use vllm_tokenizer::{DecodedText, DynTokenizer}; +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; +use winnow::combinator::{alt, delimited, eof, peek, preceded, repeat, seq}; +use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::{literal, rest, take_till, take_until, take_while}; + +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::{Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +const START: &str = "<|start|>"; +const MESSAGE: &str = "<|message|>"; +const EOM: &str = "<|eom|>"; +const EOT: &str = "<|eot|>"; + +const ASSISTANT_START: &str = "<|start|>assistant"; +const RECIPIENT_PREFIX: &str = "to="; +const ATEM_OPEN: &str = ""; +const ATEM_CLOSE: &str = ""; +const INVOKE_OPEN: &str = "` is missing. +const BODY_MARKERS: &[&str] = &[EOM, EOT, START, MESSAGE, RECIPIENT_PREFIX]; + +/// Recipient whose segments are dropped: `developer` is a prompt role, not a +/// valid assistant recipient. +const DEVELOPER_RECIPIENT: &str = "developer"; + +type MuseInput<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct MuseToolCall { + name: String, + parameters: Vec<(String, String)>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum MuseEvent { + Text, + Reasoning, + /// Structural noise consumed without emitting anything. + Skip, + /// A complete segment routing header; `None` for a bare `<|message|>` + /// without recipient. + SegmentHeader { + recipient: Option, + }, + /// `<|eom|>`: the current segment ends and another segment may follow. + SegmentEnd, + /// `<|eot|>`: the assistant turn ends. + TurnEnd, + /// One complete ATEM block parsed into raw calls. + ToolCalls { + calls: Vec, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum MuseMode { + /// Awaiting the ` to=RECIPIENT<|message|>` routing header of the next + /// segment; the full `<|start|>assistant to=…` form is also accepted. + Header { scan: MarkerScanState }, + /// Inside a `to=self` segment (private reasoning), or a bare `<|message|>` + /// segment without recipient. + Reasoning, + /// Inside a `to=commentary`/`to=user` segment (visible text). + Text, + /// Inside a tool segment, buffering one ATEM block through its close tag. + Tool { scan: MarkerScanState }, + /// Dropping text until the next segment boundary: noise after a complete + /// ATEM block, or the body of a `to=developer` segment. + Skip, + /// After `<|eot|>`: ignore the rest (EOS leakage guard). + Done, +} + +impl Default for MuseMode { + fn default() -> Self { + Self::Header { + scan: MarkerScanState::default(), + } + } +} + +/// Unified parser for Muse recipient routing and ATEM tool calls. +pub struct MuseUnifiedParser { + buffer: DecodedText, + mode: MuseMode, + /// Names of the tools registered on the request. + tool_names: HashSet, + string_parameters: HashMap>, + /// Number of calls emitted in the current response. + emitted_call_count: usize, +} + +impl MuseUnifiedParser { + /// Create a Muse parser. + pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result { + token_id(tokenizer.as_ref(), START)?; + token_id(tokenizer.as_ref(), MESSAGE)?; + token_id(tokenizer.as_ref(), EOM)?; + token_id(tokenizer.as_ref(), EOT)?; + + Ok(Self { + buffer: DecodedText::default(), + mode: MuseMode::default(), + tool_names: tools.iter().map(|tool| tool.name.clone()).collect(), + string_parameters: tools + .iter() + .map(|tool| { + let strings = tool + .parameters + .get("properties") + .and_then(Value::as_object) + .into_iter() + .flatten() + .filter(|(_, schema)| { + schema.get("type").and_then(Value::as_str) == Some("string") + }) + .map(|(name, _)| name.clone()) + .collect(); + (tool.name.clone(), strings) + }) + .collect(), + emitted_call_count: 0, + }) + } + + fn apply_event( + &mut self, + event: MuseEvent, + piece: DecodedText, + output: &mut UnifiedParserOutput, + ) { + match event { + MuseEvent::Text => output.push_text(piece.text), + MuseEvent::Reasoning => output.push_reasoning(piece), + MuseEvent::Skip => {} + MuseEvent::SegmentHeader { recipient } => { + self.mode = match recipient.as_deref() { + None | Some("self") => MuseMode::Reasoning, + Some("commentary" | "user") => MuseMode::Text, + Some(DEVELOPER_RECIPIENT) => MuseMode::Skip, + Some(_) => MuseMode::Tool { + scan: MarkerScanState::default(), + }, + }; + } + MuseEvent::SegmentEnd => { + self.mode = MuseMode::Header { + scan: MarkerScanState::default(), + }; + } + MuseEvent::TurnEnd => self.mode = MuseMode::Done, + MuseEvent::ToolCalls { calls } => { + for call in calls { + let name = normalize_tool_name(call.name, &self.tool_names); + let strings = self.string_parameters.get(&name); + let arguments: Map = call + .parameters + .into_iter() + .map(|(key, raw)| { + let value = if strings.is_some_and(|strings| strings.contains(&key)) { + Value::String(raw) + } else { + decode_parameter_value(&raw) + }; + (key, value) + }) + .collect(); + let arguments = Value::Object(arguments).to_string(); + let tool_index = self.emitted_call_count; + self.emitted_call_count += 1; + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments, + }); + } + self.mode = MuseMode::Skip; + } + } + } +} + +impl UnifiedParser for MuseUnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, _prompt_token_ids: &[u32]) -> Result<()> { + // The generation prompt always ends with `<|start|>assistant`, so the + // completion starts with the bare ` to=RECIPIENT<|message|>` header + // continuation; no prefill channel detection is needed. + self.reset(); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + // TODO: add a `structural_tag_builder` covering the full output envelope + // (self/commentary/user recipients and ATEM tool calls) for strict tool + // calling and JSON response formats. + + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.append(delta); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer.text, |input| { + parse_next_muse_event(input, &mut self.mode) + })? { + let piece = self.buffer.drain_prefix(consumed_len); + self.apply_event(event, piece, output); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + // A truncated header (or marker-free plain text) falls back to + // visible text, matching how other unified parsers treat + // incomplete output. + MuseMode::Header { .. } | MuseMode::Text => { + output.push_text(self.buffer.take().text); + } + MuseMode::Reasoning => output.push_reasoning(self.buffer.take()), + MuseMode::Skip | MuseMode::Done => self.buffer.clear(), + MuseMode::Tool { .. } => { + return Err(parsing_failed!("incomplete Muse tool call")); + } + } + + self.mode = MuseMode::default(); + Ok(output) + } + + fn reset(&mut self) -> String { + self.mode = MuseMode::default(); + self.emitted_call_count = 0; + self.buffer.take().text + } +} + +/// Map an emitted ATEM invoke name back to a registered tool name. +/// +/// The Muse Glimmer template lists a registered bare name `NAME` as the +/// namespace recipient `"NAME.*"`, so the model may call it as `NAME.NAME`. +/// Collapse only that doubled form, and only when `NAME` is registered and the +/// emitted name is not: matching on the trailing segment alone could dispatch +/// `weather.get` to a registered `calendar.get`. +fn normalize_tool_name(name: String, tool_names: &HashSet) -> String { + if tool_names.contains(&name) { + return name; + } + match name.split_once('.') { + Some((head, tail)) if head == tail && tool_names.contains(head) => head.to_string(), + _ => name, + } +} + +/// Parse one Muse event from buffered streaming input. +fn parse_next_muse_event(input: &mut MuseInput<'_>, mode: &mut MuseMode) -> ModalResult { + match mode { + MuseMode::Header { scan } => alt(( + literal(EOT).value(MuseEvent::TurnEnd), + |input: &mut MuseInput<'_>| header_event(input, scan), + )) + .parse_next(input), + MuseMode::Reasoning => body_event(input, MuseEvent::Reasoning), + MuseMode::Text => body_event(input, MuseEvent::Text), + MuseMode::Tool { scan } => tool_block_event(input, scan), + MuseMode::Skip => body_event(input, MuseEvent::Skip), + MuseMode::Done => parse_done_event(input), + } +} + +/// Parse one segment routing header up to and including `<|message|>`. +fn header_event<'i>( + input: &mut MuseInput<'i>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (header,) = seq!( + take_until_marker(MESSAGE, scan), + _: literal(MESSAGE), + ) + .parse_next(input)?; + let recipient = parse_segment_recipient(header)?; + + Ok(MuseEvent::SegmentHeader { recipient }) +} + +/// Parse an event inside a segment body, emitting body text as `body`. +fn body_event(input: &mut MuseInput<'_>, body: MuseEvent) -> ModalResult { + alt(( + literal(EOM).value(MuseEvent::SegmentEnd), + literal(EOT).value(MuseEvent::TurnEnd), + inline_header_event, + // A `to=` not followed by `NAME<|message|>` is ordinary body text. + literal(RECIPIENT_PREFIX).value(body.clone()), + safe_body_event.value(body), + )) + .parse_next(input) +} + +/// Parse safe body text before the next body marker. +fn safe_body_event(input: &mut MuseInput<'_>) -> ModalResult { + safe_text_len_mul(input, BODY_MARKERS) +} + +/// Parse the header of a segment opened inside a body whose `<|eom|>` is +/// missing: `<|start|>assistant to=NAME<|message|>`, a bare +/// `to=NAME<|message|>`, or a recipient-less `<|message|>`. +fn inline_header_event(input: &mut MuseInput<'_>) -> ModalResult { + alt(( + preceded(peek(literal(START)), |input: &mut MuseInput<'_>| { + header_event(input, &mut MarkerScanState::default()) + }), + delimited( + literal(RECIPIENT_PREFIX), + take_while(1.., is_recipient_char), + literal(MESSAGE), + ) + .map(|recipient: &str| MuseEvent::SegmentHeader { + recipient: Some(recipient.to_string()), + }), + literal(MESSAGE).value(MuseEvent::SegmentHeader { recipient: None }), + )) + .parse_next(input) +} + +/// Parse a buffered ATEM block through `` into completed calls. +// TODO: calls commit at block granularity; the canonical wire shape is one +// invoke per block (the template renders each tool call as its own segment), +// so this matches call-at-a-time emission. Revisit per-invoke emission only if +// multi-invoke blocks ever matter. +fn tool_block_event( + input: &mut MuseInput<'_>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (body,) = seq!( + take_until_marker(ATEM_CLOSE, scan), + _: literal(ATEM_CLOSE), + ) + .parse_next(input)?; + let calls = parse_atem_block(body)?; + + Ok(MuseEvent::ToolCalls { calls }) +} + +/// Ignore everything after the assistant turn ended. +fn parse_done_event(input: &mut MuseInput<'_>) -> ModalResult { + rest.value(MuseEvent::Skip).parse_next(input) +} + +/// Parse a complete segment routing header into its recipient name, or `None` +/// for a header without recipient. +/// +/// Accepted shapes: ` to=NAME` (generation-prompt continuation) and +/// `<|start|>assistant to=NAME` (full segment start), each optionally without +/// the recipient. +fn parse_segment_recipient(header: &str) -> ModalResult> { + let header = header.strip_prefix(ASSISTANT_START).unwrap_or(header).trim(); + if header.is_empty() { + return Ok(None); + } + + let mut input = header; + let parsed: ModalResult<(&str,)> = seq!( + _: literal(RECIPIENT_PREFIX), + take_while(1.., is_recipient_char), + _: eof, + ) + .parse_next(&mut input); + let (recipient,) = parsed.map_err(|_| muse_error("Muse segment header"))?; + + Ok(Some(recipient.to_string())) +} + +/// One recipient name character, per the Python parser's `[A-Za-z0-9_.-]+`. +fn is_recipient_char(char: char) -> bool { + char.is_ascii_alphanumeric() || matches!(char, '_' | '.' | '-') +} + +/// Parse one complete ATEM block body (up to ``) into +/// raw calls, one per ``. +fn parse_atem_block(body: &str) -> ModalResult> { + let mut input = body; + delimited( + (ws0, literal(ATEM_OPEN)), + repeat(1.., atem_invoke), + (ws0, eof), + ) + .parse_next(&mut input) + .map_err(|_| muse_error("Muse ATEM block")) +} + +/// Parse one `` element into its name and raw parameters. +fn atem_invoke(input: &mut &str) -> ModalResult { + let (name, parameters): (String, Vec<(String, String)>) = seq!( + _: (ws0, literal(INVOKE_OPEN), ws1, literal("name"), literal("=")), + atem_attr_value, + _: (take_until(0.., ">"), literal(">")), + repeat(0.., atem_parameter), + _: (ws0, literal(INVOKE_CLOSE)), + ) + .parse_next(input)?; + + Ok(MuseToolCall { name, parameters }) +} + +/// Parse one `` element into its key/value pair. +fn atem_parameter(input: &mut &str) -> ModalResult<(String, String)> { + let (name, raw) = seq!( + _: (ws0, literal(PARAMETER_OPEN), ws1, literal("name"), literal("=")), + atem_attr_value, + _: (take_until(0.., ">"), literal(">")), + take_until(0.., PARAMETER_CLOSE), + _: literal(PARAMETER_CLOSE), + ) + .parse_next(input)?; + + Ok((name, raw.to_string())) +} + +/// Parse one ATEM attribute value, double/single-quoted or bare. +fn atem_attr_value(input: &mut &str) -> ModalResult { + alt(( + delimited(literal("\""), take_till(0.., '"'), literal("\"")), + delimited(literal("'"), take_till(0.., '\''), literal("'")), + take_while(1.., is_recipient_char), + )) + .map(str::to_string) + .parse_next(input) +} + +/// Decode one ATEM parameter value: JSON when it parses, raw string otherwise. +/// +/// The raw value is used verbatim — whitespace inside string values is +/// significant ("spaces for string values are not stripped"). +fn decode_parameter_value(raw: &str) -> Value { + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// Build a cut error for determinably malformed Muse structure. +fn muse_error(label: &'static str) -> ErrMode { + let mut error = ContextError::new(); + error.push(StrContext::Label(label)); + ErrMode::Cut(error) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use vllm_tokenizer::{DecodedText, test_utils::TestTokenizer}; + + use super::{EOM, EOT, MESSAGE, MuseUnifiedParser, START}; + use crate::tool::{Tool, ToolCallDelta}; + use crate::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + }; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(START, 200022) + .with_special_token(MESSAGE, 200023) + .with_special_token(EOM, 200007) + .with_special_token(EOT, 200008) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for MuseUnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(DecodedText::unattributed(chunk), &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_parser() -> MuseUnifiedParser { + MuseUnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() + } + + fn collect_stream(parser: &mut MuseUnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + /// Split `text` into small chunks to stress marker-split handling. + fn char_chunks(text: &str, size: usize) -> Vec { + let chars: Vec = text.chars().collect(); + chars.chunks(size).map(|chunk| chunk.iter().collect()).collect() + } + + /// Every two-way char-boundary split of `text`. + fn two_way_splits(text: &str) -> Vec<(String, String)> { + text.char_indices() + .map(|(index, _)| index) + .chain([text.len()]) + .map(|index| (text[..index].to_string(), text[index..].to_string())) + .collect() + } + + /// A `<|start|>assistant`-prefixed segment header followed by its body and + /// the `<|eom|>` continuation boundary. + fn segment(recipient: &str, body: &str) -> String { + format!("{START}assistant to={recipient}{MESSAGE}{body}{EOM}") + } + + /// One ATEM block with one invoke per (name, params) entry. + fn atem(invokes: &[(&str, &[(&str, &str)])]) -> String { + let mut block = String::from("\n"); + for (name, params) in invokes { + block.push_str(&format!("\n")); + for (key, value) in *params { + block.push_str(&format!( + "{value}\n" + )); + } + block.push_str("\n"); + } + block.push_str(""); + block + } + + /// A completion exercising every recipient kind and both boundary markers. + fn full_completion() -> String { + let commentary = segment("commentary", "working"); + let calls = atem(&[ + ( + "get_weather", + &[("location", "Shanghai"), ("days", "1.5")][..], + ), + ("get_time", &[][..]), + ]); + format!( + " to=self{MESSAGE}plan it{EOM}{commentary}\ + {START}assistant to=get_weather{MESSAGE}{calls}{EOM}\ + {START}assistant to=user{MESSAGE}done{EOT}" + ) + } + + #[test] + fn muse_create_requires_structural_tokens() { + let error = match MuseUnifiedParser::new(&[], Arc::new(TestTokenizer::new())) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == START + )); + } + + #[test] + fn muse_preserves_special_tokens() { + assert!(test_parser().preserve_special_tokens()); + } + + #[test] + fn muse_parses_self_commentary_user_segments() { + let text = format!( + " to=self{MESSAGE}think step{EOM}{}\ + {START}assistant to=user{MESSAGE}final answer{EOT}", + segment("commentary", "intermediate"), + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "think step"); + // Visible segments concatenate without any implicit separator. + assert_eq!(output.normal_text(), "intermediatefinal answer"); + assert!(output.calls().is_empty()); + } + + #[test] + fn muse_tolerates_full_header_on_first_segment() { + let text = format!("{START}assistant to=user{MESSAGE}answer{EOT}"); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn muse_eot_ends_reasoning_segment() { + let text = format!(" to=self{MESSAGE}only reasoning{EOT}"); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "only reasoning"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn muse_parses_tool_call_with_typed_and_string_parameters() { + let text = format!( + " to=get_weather{MESSAGE}{}{EOT}", + atem( + &[( + "get_weather", + &[ + ("location", "Shanghai"), + ("days", "1.5"), + ("detailed", "true"), + ("filters", r#"{"kind":"rain"}"#), + ("hours", "[8,20]"), + ("note", "not json {"), + ][..], + )][..] + ) + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&calls[0].arguments).unwrap(), + json!({ + "location": "Shanghai", + "days": 1.5, + "detailed": true, + "filters": { "kind": "rain" }, + "hours": [8, 20], + "note": "not json {", + }) + ); + } + + #[test] + fn muse_preserves_schema_string_parameters_in_batch_and_streaming() { + let tools = [Tool { + name: "echo".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": { + "text": {"type": "string"}, + "count": {"type": "integer"}, + "items": {"type": "array", "items": {"type": "integer"}} + } + }), + strict: None, + defer_loading: None, + }]; + for value in [ + "25", + "true", + "null", + "[1,2]", + r#"{"x":1}"#, + "3%", + "1,500", + " 25 ", + ] { + let text = format!( + " to=echo.echo{MESSAGE}{}{EOT}", + atem(&[( + "echo.echo", + &[("text", value), ("count", "25"), ("items", "[1,2]")] + )]) + ); + let mut parser = MuseUnifiedParser::new(&tools, Arc::new(tokenizer())).unwrap(); + let batch = parser.parse_complete(&text).unwrap(); + let calls = batch.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name.as_deref(), Some("echo")); + assert_eq!( + serde_json::from_str::(&calls[0].arguments).unwrap(), + json!({"text": value, "count": 25, "items": [1, 2]}) + ); + for size in [1, 17, text.len()] { + parser.reset(); + let chunks = char_chunks(&text, size); + let chunks: Vec<_> = chunks.iter().map(String::as_str).collect(); + assert_eq!(collect_stream(&mut parser, &chunks), batch); + } + } + } + + #[test] + fn muse_arguments_preserve_order_and_compact_number_formatting() { + let text = format!( + " to=add{MESSAGE}{}{EOT}", + atem( + &[( + "add", + &[("y", "1.0"), ("x", "2"), ("items", r#"["left","right"]"#)][..], + )][..] + ) + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!( + output.calls()[0].arguments, + r#"{"y":1.0,"x":2,"items":["left","right"]}"# + ); + } + + #[test] + fn muse_parses_multiple_invokes_in_one_block() { + let text = format!( + " to=multi{MESSAGE}{}{EOT}", + atem(&[ + ("first", &[("a", "1")][..]), + ("second", &[][..]), + ("third", &[("b", "x")][..]), + ]) + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 3); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("first")); + assert_eq!(calls[0].arguments, r#"{"a":1}"#); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("second")); + assert_eq!(calls[1].arguments, "{}"); + assert_eq!(calls[2].tool_index, 2); + assert_eq!(calls[2].name.as_deref(), Some("third")); + } + + #[test] + fn muse_parses_consecutive_tool_segments_across_eom() { + let text = format!( + "{}{START}assistant to=add{MESSAGE}{}{EOT}", + segment( + "get_weather", + &atem(&[("get_weather", &[("location", "SF")][..])]) + ), + atem(&[("add", &[("x", "2")][..])]), + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("add")); + } + + #[test] + fn muse_streaming_matches_batch_for_every_marker_split() { + assert_streaming_matches_batch(&full_completion()); + } + + /// Stream `text` under every two-way split and several fixed chunk sizes, + /// asserting each run emits the same events as one batch parse. + fn assert_streaming_matches_batch(text: &str) { + let batch = test_parser().parse_complete(text).unwrap(); + + // Exhaustive two-way splits cover every marker cut at every boundary. + for (first, second) in two_way_splits(text) { + let output = collect_stream(&mut test_parser(), &[&first, &second]); + assert_eq!(output.events, batch.events, "split: {first:?} + {second:?}"); + } + + // Small fixed-size chunks split longer markers across 3+ chunks. + for size in 1..=9 { + let chunks = char_chunks(text, size); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let output = collect_stream(&mut test_parser(), &chunk_refs); + assert_eq!(output.events, batch.events, "chunk size {size}"); + } + } + + #[test] + fn muse_streaming_emits_text_incrementally() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(" to=user").unwrap(); + assert!(output.events.is_empty()); + + output.append(parser.parse_chunk(&format!("{MESSAGE}Hel")).unwrap()); + assert_eq!(output.normal_text(), "Hel"); + + output.append(parser.parse_chunk(&format!("lo{EOT}")).unwrap()); + assert_eq!(output.normal_text(), "Hello"); + } + + #[test] + fn muse_ignores_output_after_turn_end() { + let text = format!(" to=user{MESSAGE}answer{EOT}junk{EOM}more"); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn muse_bare_eot_between_segments_ends_turn() { + let text = format!("{}{EOT}", segment("commentary", "note")); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "note"); + } + + #[test] + fn muse_finish_flushes_unclosed_reasoning() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!(" to=self{MESSAGE}still thinking")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "still thinking"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn muse_finish_flushes_unclosed_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!(" to=user{MESSAGE}partial answer")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "partial answer"); + } + + #[test] + fn muse_finish_flushes_marker_free_text_as_text() { + let output = test_parser().parse_complete("plain answer").unwrap(); + + assert_eq!(output.normal_text(), "plain answer"); + } + + #[test] + fn muse_finish_after_closed_tool_block_keeps_calls() { + // The terminal `<|eot|>` may already be stripped by the serving layer + // (e.g. stop-token handling); a complete ATEM block still commits. + let mut parser = test_parser(); + let mut output = parser + .parse_chunk(&format!( + " to=calc{MESSAGE}{}", + atem(&[("calc", &[("x", "42")][..])]) + )) + .unwrap(); + output.append(parser.finish().unwrap()); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name.as_deref(), Some("calc")); + assert_eq!(calls[0].arguments, r#"{"x":42}"#); + } + + #[test] + fn muse_finish_fails_mid_tool_block() { + let mut parser = test_parser(); + parser + .parse_chunk(&format!( + " to=calc{MESSAGE}\ngarbage" + )) + .unwrap_err(); + + assert!(error.to_report_string().contains("Muse ATEM block")); + } + + #[test] + fn muse_reset_returns_buffered_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!(" to=user{MESSAGE}answer<|eo")).unwrap(); + assert_eq!(output.normal_text(), "answer"); + + let raw = parser.reset(); + + assert_eq!(raw, "<|eo"); + + output = parser.parse_chunk(&format!(" to=user{MESSAGE}again")).unwrap(); + output.append(parser.finish().unwrap()); + assert_eq!(output.normal_text(), "again"); + } + + #[test] + fn muse_bare_tool_header_closes_reasoning_without_eom() { + // The model may skip `<|eom|>` and write the tool header right after + // its reasoning, e.g. for a tool call with empty arguments. + let text = format!( + " to=self{MESSAGE}I should check the time. to=get_time{MESSAGE}{}{EOM}\ + {START}assistant to=user{MESSAGE}It is noon.{EOT}", + atem(&[("get_time", &[][..])]), + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "I should check the time. "); + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name.as_deref(), Some("get_time")); + assert_eq!(calls[0].arguments, "{}"); + assert_eq!(output.normal_text(), "It is noon."); + assert_streaming_matches_batch(&text); + } + + #[test] + fn muse_full_header_closes_segments_without_eom() { + let text = format!( + " to=self{MESSAGE}plan{START}assistant to=calc{MESSAGE}{}\ + {START}assistant to=user{MESSAGE}done{EOT}", + atem(&[("calc", &[("x", "1")][..])]), + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "plan"); + assert_eq!(output.calls()[0].name.as_deref(), Some("calc")); + assert_eq!(output.normal_text(), "done"); + assert_streaming_matches_batch(&text); + } + + #[test] + fn muse_recipient_mentions_in_body_stay_text() { + let text = format!( + " to=self{MESSAGE}send it to=user later, too{EOM}\ + {START}assistant to=user{MESSAGE}auto to=x{EOT}" + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "send it to=user later, too"); + assert_eq!(output.normal_text(), "auto to=x"); + assert_streaming_matches_batch(&text); + } + + #[test] + fn muse_recipient_less_segment_is_reasoning() { + let text = format!( + "{MESSAGE}thinking{EOM}{START}assistant to=user{MESSAGE}answer{MESSAGE}more{EOT}" + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "thinkingmore"); + assert_eq!(output.normal_text(), "answer"); + assert_streaming_matches_batch(&text); + } + + #[test] + fn muse_mixed_channels_match_python_routing() { + let text = concat!( + "<|message|>Think.<|eom|>", + "<|start|>assistant to=commentary<|message|>Checking. ", + "to=developer<|message|>internal<|eom|>", + "<|start|>assistant to=weather.get<|message|>", + "", + "Paris", + "<|eom|>", + "<|start|>assistant to=self<|message|>Recheck.<|eom|>", + "<|start|>assistant to=user<|message|>Done.<|eot|>", + ); + let output = test_parser().parse_complete(text).unwrap(); + expect_test::expect![[r#" + UnifiedParserOutput { + events: [ + Reasoning( + DecodedText { + text: "Think.", + attributions: [], + }, + ), + Text( + "Checking. ", + ), + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "weather.get", + ), + arguments: "{\"city\":\"Paris\"}", + }, + ), + Reasoning( + DecodedText { + text: "Recheck.", + attributions: [], + }, + ), + Text( + "Done.", + ), + ], + } + "#]] + .assert_debug_eq(&output); + assert_streaming_matches_batch(text); + } + + #[test] + fn muse_drops_developer_segment() { + let text = format!( + " to=developer{MESSAGE}not for the user{EOM}\ + {START}assistant to=user{MESSAGE}answer{EOT}" + ); + + let output = test_parser().parse_complete(&text).unwrap(); + + assert!(output.reasoning_text().is_empty()); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn muse_collapses_doubled_namespace_tool_name() { + let tools = ["get_weather", "calendar.get"].map(|name| Tool { + name: name.to_string(), + description: None, + parameters: json!({"type": "object"}), + strict: None, + defer_loading: None, + }); + let mut parser = MuseUnifiedParser::new(&tools, Arc::new(tokenizer())).unwrap(); + let text = format!( + " to=get_weather.get_weather{MESSAGE}{}", + atem(&[ + ("get_weather.get_weather", &[][..]), + ("calendar.get", &[][..]), + ("weather.get", &[][..]), + ("calendar.calendar", &[][..]), + ]) + ); + + let output = parser.parse_complete(&text).unwrap(); + + let names: Vec<_> = output.calls().into_iter().filter_map(|call| call.name).collect(); + assert_eq!( + names, + [ + "get_weather", + "calendar.get", + "weather.get", + "calendar.calendar" + ] + ); + } +} diff --git a/tests/tool_use/test_muse.py b/tests/tool_use/test_muse.py new file mode 100644 index 000000000000..91f9481636d8 --- /dev/null +++ b/tests/tool_use/test_muse.py @@ -0,0 +1,615 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the Muse ATEM tool parser and reasoning parser. + +Muse writes every turn as a sequence of channel-scoped messages rather +than JSON, so the two parsers are tested together: the reasoning parser strips +the reasoning span and forwards the remaining channels as content, and the tool +parser reads ATEM markup out of those channels. + +Five areas, in order: + + 1. non-streaming tool-call extraction, including channel scoping (an + ```` echoed inside reasoning must never become a call); + 2. the reasoning -> tool-parser handoff, which regressed once by returning + ``content=None`` and starving the tool parser; + 3. channel routing: ``to=commentary`` is content, a recipient-less + ``<|message|>`` is reasoning, ``to=developer`` is dropped, and a tool + header written without the closing ``<|eom|>`` still ends the reasoning; + 4. streaming, where markers routinely straddle chunk boundaries, plus + truncation isolation for an unterminated ``to=self`` block; + 5. tool-name normalization against the tools registered on the request. + +These drive the parsers directly and need no checkpoint. The tests that require +a real tokenizer live in ``test_muse_parse_delta.py``. +""" + +import json +from contextlib import suppress +from types import SimpleNamespace +from typing import cast + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.reasoning.muse_reasoning_parser import MuseReasoningParser +from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser +from vllm.tool_parsers.muse_tool_parser import MuseToolParser + +R: MuseReasoningParser +T: MuseToolParser + + +@pytest.fixture(autouse=True) +def _fresh_parsers(): + """Give each test request-scoped parser state through the real constructors.""" + global R, T + R = MuseReasoningParser(object()) + T = MuseToolParser(object()) + + +# Any framing token that must NEVER appear in surfaced reasoning/content. +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "to=read.read", + " ChatCompletionRequest: + return cast(ChatCompletionRequest, _FakeReq()) + + +# The parsers accept a missing request at runtime. +_NO_REQ = cast(ChatCompletionRequest, None) + + +def _req(*names: str) -> ChatCompletionRequest: + """A request with ``names`` registered as tools.""" + return cast( + ChatCompletionRequest, + SimpleNamespace( + tools=[SimpleNamespace(function=SimpleNamespace(name=n)) for n in names] + ), + ) + + +def _call(name): + """One tool-call turn invoking ``name``.""" + return ( + f"<|start|>assistant to={name}<|message|>" + f'\n\n' + f'Paris\n' + f"\n" + ) + + +# ---------------------------------------------------------------- tool calls + + +def test_single_tool_call_after_reasoning(): + raw = ( + "to=self<|message|>Let me check the weather.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + 'celsius\n' + "\n<|eot|>" + ) + out = MuseToolParser.extract_tool_calls(T, raw, _NO_REQ) + assert out.tools_called and len(out.tool_calls) == 1 + assert out.tool_calls[0].function.name == "weather.get" + assert json.loads(out.tool_calls[0].function.arguments) == { + "city": "Paris", + "units": "celsius", + } + + +def test_parallel_calls_across_eom_boundaries(): + raw = ( + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n' + '2\n' + "\n<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n' + '4\n' + "\n<|eot|>" + ) + out = MuseToolParser.extract_tool_calls(T, raw, _NO_REQ) + assert out.tools_called and len(out.tool_calls) == 2, len(out.tool_calls) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"] + # JSON-typed values decode to ints + assert json.loads(out.tool_calls[0].function.arguments) == {"a": 1, "b": 2} + + +def test_echoed_invoke_in_reasoning_is_not_parsed(): + """Channel scoping: an invoke quoted inside reasoning is not a call.""" + raw = ( + 'to=self<|message|>I could call ' + '1 ' + "but I will not.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + out = MuseToolParser.extract_tool_calls(T, raw, _NO_REQ) + assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!" + assert out.content == "The answer is 42.", repr(out.content) + + +def test_plain_answer_yields_no_tool_calls(): + out = MuseToolParser.extract_tool_calls( + T, "to=user<|message|>Just a plain answer.<|eot|>", _NO_REQ + ) + assert not out.tools_called + + +def test_json_object_array_and_bool_params_decode(): + raw = ( + "<|start|>assistant to=api.call<|message|>" + '\n\n' + '{"nested": [1, 2, 3]}\n' + 'true\n' + "\n<|eot|>" + ) + out = MuseToolParser.extract_tool_calls(T, raw, _NO_REQ) + assert json.loads(out.tool_calls[0].function.arguments) == { + "payload": {"nested": [1, 2, 3]}, + "flag": True, + } + + +# ------------------------------------------------- reasoning -> tool handoff + + +def test_reasoning_to_toolcall_handoff(): + """The regression: content=None here starved the tool parser.""" + raw = ( + " to=self<|message|>Let me call the tool.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + "\n" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning == "Let me call the tool.", repr(reasoning) + assert content is not None and "thinking<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning == "thinking", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert not MuseToolParser.extract_tool_calls(T, content, _NO_REQ).tools_called + + +def test_plain_content_without_framing_passes_through(): + reasoning, content = MuseReasoningParser.extract_reasoning( + R, "Just a direct answer.", _NO_REQ + ) + assert reasoning is None and content == "Just a direct answer.", ( + reasoning, + content, + ) + + +def test_reasoning_then_parallel_calls(): + raw = ( + " to=self<|message|>need two calls<|eom|>" + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n\n' + "<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n\n' + "<|eot|>" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning == "need two calls", repr(reasoning) + assert content is not None + out = MuseToolParser.extract_tool_calls(T, content, _NO_REQ) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"], ( + out.tool_calls + ) + + +def test_full_block_inside_reasoning_is_not_parsed_end_to_end(): + """Serving runs the reasoning parser first; a full block in to=self stays inert.""" + raw = ( + 'to=self<|message|>' + '1' + "<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning is not None and 'name="evil.fn"' in reasoning + assert content == "The answer is 42.", repr(content) + out = MuseToolParser.extract_tool_calls(T, content, _NO_REQ) + assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!" + + +# ---------------------------------------------------------- channel routing + +RAW_COMMENTARY_TOOLCALL = ( + " to=self<|message|>Plan the lookup.<|eom|>" + "<|start|>assistant to=commentary<|message|>Checking the weather.<|eom|>" + + _call("weather.get") +) + +RAW_COMMENTARY_ANSWER = ( + " to=self<|message|>Plan.<|eom|>" + "<|start|>assistant to=commentary<|message|>Working on it. <|eom|>" + "<|start|>assistant to=user<|message|>Done.<|eot|>" +) + +# The model skips <|eom|> and writes the tool header right after its reasoning. +RAW_MISSING_EOM_TOOLCALL = ( + " to=self<|message|>Need the time. to=time.now<|message|>" + '\n\n\n' + "<|eom|>" + "<|start|>assistant to=user<|message|>It is noon.<|eot|>" +) + + +def test_commentary_before_tool_call_is_content(): + reasoning, content = MuseReasoningParser.extract_reasoning( + R, RAW_COMMENTARY_TOOLCALL, _NO_REQ + ) + assert reasoning == "Plan the lookup.", repr(reasoning) + assert content is not None + out = MuseToolParser.extract_tool_calls(T, content, _NO_REQ) + assert [t.function.name for t in out.tool_calls] == ["weather.get"] + assert out.content == "Checking the weather.", repr(out.content) + + +def test_commentary_and_answer_join_as_content(): + reasoning, content = MuseReasoningParser.extract_reasoning( + R, RAW_COMMENTARY_ANSWER, _NO_REQ + ) + assert reasoning == "Plan.", repr(reasoning) + assert content == "Working on it. Done.", repr(content) + + +def test_recipient_less_message_is_reasoning(): + raw = ( + "<|message|>Thinking.<|eom|><|start|>assistant to=user<|message|>Answer.<|eot|>" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning == "Thinking.", repr(reasoning) + assert content == "Answer.", repr(content) + + +def test_developer_channel_is_dropped(): + raw = ( + " to=developer<|message|>internal note<|eom|>" + "<|start|>assistant to=user<|message|>Answer.<|eot|>" + ) + reasoning, content = MuseReasoningParser.extract_reasoning(R, raw, _NO_REQ) + assert reasoning is None and content == "Answer.", (reasoning, content) + out = MuseToolParser.extract_tool_calls(T, raw, _NO_REQ) + assert not out.tools_called and out.content == "Answer.", out + + +def test_missing_eom_before_tool_header_keeps_tool_call(): + reasoning, content = MuseReasoningParser.extract_reasoning( + R, RAW_MISSING_EOM_TOOLCALL, _NO_REQ + ) + assert reasoning == "Need the time. ", repr(reasoning) + assert content is not None + out = MuseToolParser.extract_tool_calls(T, content, _NO_REQ) + assert [t.function.name for t in out.tool_calls] == ["time.now"] + assert out.content == "It is noon.", repr(out.content) + + +# ----------------------------------------------------------------- streaming + + +def _stream(raw: str, chunk: int): + """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming + parsers; return (reasoning, content, tool_calls).""" + reasoning, content, toolcalls = [], [], [] + prev = "" + i = 0 + while i < len(raw): + cur = raw[: i + chunk] + delta = cur[len(prev) :] + dm = MuseReasoningParser.extract_reasoning_streaming( + R, prev, cur, delta, [], [], [] + ) + if dm is not None: + if dm.reasoning: + reasoning.append(dm.reasoning) + content_delta = getattr(dm, "content", None) + # Tool-channel content is an internal handoff to the tool parser, + # not client-visible content from the unified parser. + if content_delta and "" not in content_delta: + content.append(content_delta) + dt = MuseToolParser.extract_tool_calls_streaming( + T, prev, cur, delta, [], [], [], _fake_req() + ) + if dt is not None and dt.tool_calls: + toolcalls.extend(dt.tool_calls) + prev = cur + i += chunk + return "".join(reasoning), "".join(content), toolcalls + + +def _fn_of(tc): + fn = tc.function + if isinstance(fn, dict): + return fn.get("name"), fn.get("arguments") + return fn.name, fn.arguments + + +RAW_TOOLCALL = ( + " to=self<|message|>I should read the hostname file to answer.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n" +) + +RAW_ANSWER = ( + " to=self<|message|>Think about it.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" +) + +# NO closing <|eom|> -> truncated CoT +RAW_TRUNCATED = ( + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait" +) + + +# 1 is the worst case: markers arrive one char at a time (mid-marker deltas). +@pytest.mark.parametrize("chunk", [1, 3, 17]) +def test_streaming_toolcall(chunk): + reasoning, content, tcs = _stream(RAW_TOOLCALL, chunk) + # (a) no framing token leaks into reasoning or content + for f in _FRAMING: + assert f not in reasoning, ( + f"framing {f!r} leaked into reasoning (chunk={chunk})" + ) + assert f not in content, f"framing {f!r} leaked into content (chunk={chunk})" + # (b) exactly one tool_call with correct name + args + assert len(tcs) == 1, f"expected 1 tool_call, got {len(tcs)} (chunk={chunk})" + name, args = _fn_of(tcs[0]) + assert name == "read.read", name + assert json.loads(args) == {"path": "/etc/hostname"}, args + assert tcs[0].index == 0 and tcs[0].type == "function" and tcs[0].id + # (c) reasoning captured separately and clean + assert reasoning == "I should read the hostname file to answer.", repr(reasoning) + assert content == "", repr(content) + + +def test_streaming_reasoning_then_content(): + reasoning, content, tcs = _stream(RAW_ANSWER, 3) + for f in _FRAMING: + assert f not in reasoning and f not in content, f + assert reasoning == "Think about it.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tcs == [] + + +def test_streaming_commentary_and_answer(): + reasoning, content, tcs = _stream(RAW_COMMENTARY_ANSWER, 3) + for f in _FRAMING: + assert f not in reasoning and f not in content, f + assert reasoning == "Plan.", repr(reasoning) + assert content == "Working on it. Done.", repr(content) + assert tcs == [] + + +def test_streaming_missing_eom_before_tool_header(): + reasoning, _, tcs = _stream(RAW_MISSING_EOM_TOOLCALL, 1) + assert reasoning == "Need the time. ", repr(reasoning) + assert [_fn_of(tc)[0] for tc in tcs] == ["time.now"], tcs + + +def test_truncated_cot_no_toolcall_nonstreaming(): + out = MuseToolParser.extract_tool_calls(T, RAW_TRUNCATED, _fake_req()) + assert not out.tools_called and out.tool_calls == [] + # partial reasoning must still be recovered by the reasoning parser + reasoning, _ = MuseReasoningParser.extract_reasoning(R, RAW_TRUNCATED, _fake_req()) + assert reasoning and "Maybe I should call" in reasoning, repr(reasoning) + + +def test_truncated_cot_no_toolcall_streaming(): + _, _, tcs = _stream(RAW_TRUNCATED, 3) + assert tcs == [], f"truncated CoT invoke leaked as streaming tool call: {tcs}" + + +# ------------------------------------------------------- name normalization +# +# Muse emits `get_weather.get_weather` for a bare-registered +# `get_weather`, and `weather.get` verbatim for a namespaced one. The parser +# normalizes against the tools actually registered on the request. + + +def test_doubled_bare_name_collapses(): + out = MuseToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), _req("get_weather") + ) + assert out.tools_called and out.tool_calls[0].function.name == "get_weather", ( + out.tool_calls[0].function.name + ) + + +def test_namespaced_name_preserved(): + out = MuseToolParser.extract_tool_calls( + T, _call("weather.get"), _req("weather.get") + ) + assert out.tool_calls[0].function.name == "weather.get" + + +def test_unregistered_namespace_is_preserved(): + # Suffix-only matching can silently dispatch a tool from the wrong namespace. + out = MuseToolParser.extract_tool_calls( + T, _call("foo.get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "foo.get_weather" + + +def test_trailing_segment_ambiguous_left_alone(): + # two registered tools share leaf 'get' -> ambiguous -> do NOT rewrite + out = MuseToolParser.extract_tool_calls( + T, _call("x.get"), _req("weather.get", "time.get") + ) + assert out.tool_calls[0].function.name == "x.get" + + +def test_no_registered_tools_passthrough(): + out = MuseToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), _NO_REQ + ) + assert out.tool_calls[0].function.name == "get_weather.get_weather" + + +def test_exact_match_kept(): + out = MuseToolParser.extract_tool_calls( + T, _call("get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "get_weather" + + +def test_adjust_request_forces_skip_special_tokens_off(): + req = SimpleNamespace(tools=None, tool_choice=None, skip_special_tokens=True) + MuseToolParser.adjust_request(T, cast(ChatCompletionRequest, req)) + assert req.skip_special_tokens is False + + +def test_named_custom_tool_choice_skips_structured_outputs(): + # Custom tools ride the function shim; a named custom choice must skip + # the base adjust_request like a named function choice does. + request = ResponsesRequest( + input="hi", + tools=[{"type": "custom", "name": "emit_command"}], + tool_choice={"type": "custom", "name": "emit_command"}, + ) + T.adjust_request(request) + assert request.structured_outputs is None + + +@pytest.mark.parametrize( + ("raw", "reasoning", "content"), + [ + ("<|message|>Thinking.<|eot|>", "Thinking.", None), + ("to=developer<|message|>internal note<|eom|>", None, None), + ( + "to=commentary<|message|>Checking. " + "to=developer<|message|>internal note<|eom|>", + None, + "Checking. ", + ), + ( + "<|message|>Think.<|start|>assistant to=user<|message|>Done.<|eot|>", + "Think.", + "Done.", + ), + ( + "to=self<|message|>First.<|eom|>" + "<|start|>assistant<|message|>Second.<|eot|>", + "First.Second.", + None, + ), + ], +) +def test_channel_boundaries_match_streaming(raw, reasoning, content): + """Hidden channels and message boundaries must not leak in either path.""" + assert R.extract_reasoning(raw, _NO_REQ) == (reasoning, content) + for chunk in (1, 3, len(raw)): + parser = MuseReasoningParser(object()) + reasons, contents = [], [] + for end in range(chunk, len(raw) + chunk, chunk): + previous = raw[: max(0, end - chunk)] + current = raw[:end] + delta = parser.extract_reasoning_streaming( + previous, current, current[len(previous) :], [], [], [] + ) + if delta: + reasons.append(delta.reasoning or "") + contents.append(delta.content or "") + assert "".join(reasons) == (reasoning or "") + assert "".join(contents) == (content or "") + + +@pytest.mark.parametrize( + "value", ["25", "true", "null", "[1,2]", '{"x":1}', "3%", "1,500", " 25 "] +) +@pytest.mark.parametrize("request_kind", ["chat", "responses"]) +@pytest.mark.parametrize("parser_class", [MuseToolParser, MuseGlimmerToolParser]) +def test_string_parameter_schema_preserves_raw_text(value, request_kind, parser_class): + """String schemas disambiguate ATEM values; Glimmer keeps legacy decoding.""" + parameters = { + "type": "object", + "properties": { + "text": {"type": "string"}, + "count": {"type": "integer"}, + "items": {"type": "array", "items": {"type": "integer"}}, + }, + } + function = {"name": "echo", "parameters": parameters} + if request_kind == "chat": + request = ChatCompletionRequest( + model="test", + messages=[], + tools=[{"type": "function", "function": function}], + ) + else: + request = ResponsesRequest(input="hi", tools=[{"type": "function", **function}]) + raw = ( + 'to=echo.echo<|message|>' + f'{value}' + '25' + '[1,2]' + "<|eot|>" + ) + expected_text = value + if parser_class is MuseGlimmerToolParser: + with suppress(ValueError): + expected_text = json.loads(value) + expected = {"text": expected_text, "count": 25, "items": [1, 2]} + parser = parser_class(object()) + calls = parser.extract_tool_calls(raw, request).tool_calls + assert len(calls) == 1 + assert calls[0].function.name == "echo" + assert json.loads(calls[0].function.arguments) == expected + + for chunk in (1, 17, len(raw)): + parser = parser_class(object()) + calls = [] + for end in range(chunk, len(raw) + chunk, chunk): + previous, current = raw[: end - chunk], raw[:end] + delta = parser.extract_tool_calls_streaming( + previous, current, current[len(previous) :], [], [], [], request + ) + if delta and delta.tool_calls: + calls.extend(delta.tool_calls) + assert len(calls) == 1 + name, arguments = _fn_of(calls[0]) + assert name == "echo" + assert json.loads(arguments) == expected diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py index 2fbf3cff49c0..c986b3e2fc8d 100644 --- a/tests/tool_use/test_muse_glimmer.py +++ b/tests/tool_use/test_muse_glimmer.py @@ -413,3 +413,22 @@ def test_named_custom_tool_choice_skips_structured_outputs(): ) T.adjust_request(request) assert request.structured_outputs is None + + +@pytest.mark.parametrize("chunk", [1, 3, 100]) +def test_recipient_less_content_keeps_glimmer_behavior(chunk): + """The shared Muse implementation must not reclassify Glimmer content.""" + raw = "<|message|>Hello.<|eot|>" + assert T.extract_tool_calls(raw, _NO_REQ).content == "Hello." + assert R.extract_reasoning(raw, _NO_REQ) == (None, raw) + content, reasoning = [], [] + for end in range(chunk, len(raw) + chunk, chunk): + previous, current = raw[: max(0, end - chunk)], raw[:end] + delta = T.extract_tool_calls_streaming( + previous, current, current[len(previous) :], [], [], [], _fake_req() + ) + if delta: + content.append(delta.content or "") + reasoning.append(delta.reasoning or "") + assert "".join(content) == "Hello." + assert "".join(reasoning) == "" diff --git a/tests/tool_use/test_muse_parse_delta.py b/tests/tool_use/test_muse_parse_delta.py new file mode 100644 index 000000000000..0b6875a5cd94 --- /dev/null +++ b/tests/tool_use/test_muse_parse_delta.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""parse_delta-level regression tests for the Muse parsers. + +These drive the UNIFIED parser-engine streaming API the serving layer actually +uses (``DelegatingParser.parse_delta``), NOT the ``extract_*_streaming`` methods +directly. This is the coverage that was missing when a live no-tools streaming +request leaked raw harmony framing into ``content``: the phase machine marked +``reasoning_ended=True`` at the prompt boundary (because ``is_reasoning_end`` on +a prompt with no ``to=self`` wrongly returned True), skipping the reasoning +phase for the whole generation. + +Requires a real Muse tokenizer; skipped if the checkpoint is unavailable. +""" + +import json +import os + +import pytest + +CKPT = os.environ.get("MUSE_CKPT", "") + +pytestmark = pytest.mark.skipif( + not os.path.isdir(CKPT), reason=f"Muse checkpoint not found at {CKPT}" +) + +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "= len(gen_ids)), + ) + if dm is None: + continue + if getattr(dm, "reasoning", None): + reasoning.append(dm.reasoning) + if getattr(dm, "content", None): + content.append(dm.content) + for tc in getattr(dm, "tool_calls", None) or []: + fn = tc.function + name = fn.get("name") if isinstance(fn, dict) else getattr(fn, "name", None) + args = ( + fn.get("arguments") + if isinstance(fn, dict) + else getattr(fn, "arguments", None) + ) + tools.append((tc.index, name, args)) + return "".join(reasoning), "".join(content), tools + + +def _assert_no_framing(text): + for f in _FRAMING: + assert f not in text, f"framing {f!r} leaked: {text!r}" + + +def test_parse_delta_no_tools_reasoning_then_answer(parser_cls, tok): + # The core regression: reasoning phase must stay active for a no-tools turn. + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>Let me think step by step about the sum.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>", + ) + _assert_no_framing(content) + assert reasoning == "Let me think step by step about the sum.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tools == [] + + +def test_parse_delta_content_only(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=user<|message|>Just a direct answer.<|eot|>", + ) + _assert_no_framing(content) + assert content == "Just a direct answer.", repr(content) + assert tools == [] + + +def test_parse_delta_tool_call(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>I should read the hostname.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n", + ) + _assert_no_framing(content) + assert reasoning == "I should read the hostname.", repr(reasoning) + assert len(tools) == 1, tools + idx, name, args = tools[0] + assert idx == 0 and name == "read.read" + assert json.loads(args) == {"path": "/etc/hostname"} + + +def test_parse_delta_truncated_cot_no_toolcall(parser_cls, tok): + # Contemplated invoke inside an unterminated to=self block: no tool call, + # no framing leak into content, partial reasoning recovered. + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait", + ) + _assert_no_framing(content) + assert tools == [], f"contemplated invoke leaked as tool call: {tools}" + assert "Maybe I should call" in reasoning, repr(reasoning) + + +def test_parse_delta_reasoning_suppressed_when_not_requested(parser_cls, tok): + class _NoReasonReq: + tools = None + tool_choice = None + include_reasoning = False + + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>secret thoughts<|eom|>" + "<|start|>assistant to=user<|message|>Public answer.<|eot|>", + req=_NoReasonReq(), + ) + assert reasoning == "", repr(reasoning) # suppressed + assert content == "Public answer.", repr(content) + _assert_no_framing(content) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) + + +@pytest.mark.parametrize("chunk", [1, 3, 10000]) +@pytest.mark.parametrize( + ("raw", "reasoning", "content"), + [ + ("<|message|>Think.<|eot|>", "Think.", ""), + ("to=developer<|message|>internal<|eot|>", "", ""), + ( + "to=commentary<|message|>Working. to=developer<|message|>internal" + "<|eom|><|start|>assistant to=user<|message|>Done.<|eot|>", + "", + "Working. Done.", + ), + ], +) +def test_parse_delta_channel_boundaries( + parser_cls, tok, raw, reasoning, content, chunk +): + assert _drive(parser_cls, tok, raw, chunk=chunk) == (reasoning, content, []) + + +@pytest.mark.parametrize("chunk", [1, 3, 10000]) +def test_parse_delta_commentary_tool_and_answer(parser_cls, tok, chunk): + raw = ( + "<|message|>Think.<|eom|>" + "<|start|>assistant to=commentary<|message|>Checking. <|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '' + 'Paris' + "<|eom|>" + "<|start|>assistant to=self<|message|>Recheck.<|eom|>" + "<|start|>assistant to=user<|message|>Done.<|eot|>" + ) + reasoning, content, tools = _drive(parser_cls, tok, raw, chunk=chunk) + assert reasoning == "Think.Recheck." + assert content == "Checking. Done." + assert len(tools) == 1 + assert tools[0][1] == "weather.get" + assert json.loads(tools[0][2]) == {"city": "Paris"} diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 796fc3ba3f37..5eac114850b2 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -152,10 +152,19 @@ "plamo3_engine_reasoning_parser", "Plamo3ParserReasoningAdapter", ), + "muse": ( + "muse_reasoning_parser", + "MuseReasoningParser", + ), "muse_glimmer": ( "muse_glimmer_reasoning_parser", "MuseGlimmerReasoningParser", ), + # Alias of "muse". + "musespark1x": ( + "muse_reasoning_parser", + "MuseReasoningParser", + ), "qwen3": ( "qwen3_engine_reasoning_parser", "Qwen3ParserReasoningAdapter", diff --git a/vllm/reasoning/muse_glimmer_reasoning_parser.py b/vllm/reasoning/muse_glimmer_reasoning_parser.py index 5991eb4d2d4c..a0ee8940cd1b 100644 --- a/vllm/reasoning/muse_glimmer_reasoning_parser.py +++ b/vllm/reasoning/muse_glimmer_reasoning_parser.py @@ -1,327 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reasoning-content parser for MuseGlimmer. -Port of the ``reasoning_content`` rule from the HuggingFace MuseGlimmer -``MUSE_GLIMMER_RESPONSE_SCHEMA`` (synced with internal master). MuseGlimmer emits -chain-of-thought in ``to=self`` channels delimited by ``<|message|>`` ... ``<|eom|>``: - to=self<|message|>...reasoning...<|eom|> -A turn may contain several ``to=self`` blocks interleaved with tool calls, and a -tool call or final answer follows in its own channel. -Because MuseGlimmer's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed -to be single vocab tokens across every checkpoint's tokenizer, this parser works -on the decoded text with regexes rather than the single start/end-token base class. -Usage: ``--reasoning-parser muse_glimmer`` -""" -from __future__ import annotations +"""Compatibility entry point for the Muse Glimmer reasoning parser.""" -from collections.abc import Iterable, Sequence +from vllm.reasoning.muse_reasoning_parser import MuseReasoningParser -import regex as re -from vllm.entrypoints.generate.base.protocol import DeltaMessage -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser - -_EOM = "<|eom|>" -_EOT = "<|eot|>" -_FUNCTION_CALLS_OPEN = "" -_REASONING_OPEN = "to=self<|message|>" -_ASSISTANT_TURN_OPEN = "<|start|>assistant" -# A channel header: ``to=<|message|>`` where recipient is ``self`` -# (reasoning), ``user`` (final answer) or ``[.]`` (tool call). -_CHANNEL_HEADER_RE = re.compile(r"to=(?P[^\s<]+)<\|message\|>") -_HEADER_PAT = r"to=[^\s<]+<\|message\|>" -# Collapse the gap between reasoning blocks so multiple to=self spans join. -_COLLAPSE_RE = re.compile( - r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL -) -_REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL) -_CONTENT_RE = re.compile( - r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL -) -# Strip a CLOSED reasoning span (header .. <|eom|>). -_STRIP_REASONING_RE = re.compile( - r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL -) -# An UNTERMINATED trailing reasoning span. The model sometimes leaves the -# analysis channel WITHOUT emitting <|eom|>, writing a bare -# ``to=<|message|>`` header instead (observed deterministically for a call -# with EMPTY arguments on a tool that has optional parameters; reproduced on -# other engines too, so it is a model-side defect, not engine-specific). -# -# These two patterns MUST therefore stop at the next channel header rather than -# running to end-of-text. An unbounded ``...$`` version consumes the real tool -# call along with the reasoning: `is_reasoning_end` then never fires, the parser -# never leaves the reasoning phase, the tool parser is never invoked, and the -# entire generation is dropped (empty reasoning, empty content, no tool call). -_STRIP_OPEN_REASONING_RE = re.compile( - r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>" - r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*" - r"(?=" + _HEADER_PAT + r"|$)", - re.DOTALL, -) -_OPEN_REASONING_RE = re.compile( - r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)" - r"(?=" + _HEADER_PAT + r"|$)", - re.DOTALL, -) -# Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body. -_HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>") -# A trailing fragment that could still grow into a channel header (" t", " to", -# " to=", " to=skill"). Without this the recipient name leaks into reasoning and -# then has to be un-emitted once ``<|message|>`` arrives. -_OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$") - - -def _current_assistant_turn(text: str) -> str: - """Return only the text generated in the current assistant turn. - ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start, - and an MuseGlimmer prompt legitimately contains ATEM markers (``render_tool_defs`` - writes a literal ```` example into the system message, - and prior assistant turns may carry real tool calls). Anchoring on the last - channel-open keeps prompt text from deciding the phase. - """ - idx = text.rfind(_ASSISTANT_TURN_OPEN) - return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text - - -def _trim_open_body(body: str) -> str: - """Hold back any tail of a still-growing body that could still be framing. - Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs - the partial-marker trim (``<``) before the partial-header trim can see - `" to=skill"`. Trimming only once leaks the recipient name as reasoning. - """ - while True: - trimmed = body - for marker in _HOLDBACK_MARKERS: - for k in range(min(len(marker) - 1, len(trimmed)), 0, -1): - if trimmed.endswith(marker[:k]): - trimmed = trimmed[:-k] - break - else: - continue - break - header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed) - if header_tail is not None: - trimmed = trimmed[: header_tail.start()] - if trimmed == body: - return body - body = trimmed - - -class MuseGlimmerReasoningParser(ReasoningParser): - def __init__(self, tokenizer, *args, **kwargs) -> None: - super().__init__(tokenizer, *args, **kwargs) - # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified - # `previous_text` is unsafe: a classified body legitimately shrinks when - # a partial header becomes recognisable, and diffing against the shrunken - # value re-emits text that already went out. - self._emitted_reasoning: str = "" - self._emitted_content: str = "" - self._tool_handoff_done: bool = False - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - """Preserve MuseGlimmer's ATEM framing tokens in the decoded output. - vLLM's serving default is ``skip_special_tokens=True``, which strips - ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the - parsers run, collapsing reasoning into content and breaking channel - scoping. Unlike the base tool-parser hook we do NOT touch - ``structured_outputs`` -- MuseGlimmer emits native ATEM, not JSON. - """ - request.skip_special_tokens = False - return request - - def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: - """Whether the model has left reasoning and opened a TOOL channel. - A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this - parser surfaces that content itself. Only a real tool channel switches - the ``DelegatingParser`` phase machine over to the tool parser. - Both closed and unterminated reasoning spans are stripped before the - check, so an ```` the model merely echoes inside its CoT - never flips the phase. - """ - try: - text = self.model_tokenizer.decode(input_ids) - except Exception: - return False - remainder = self._tool_channel_remainder(text) - return _FUNCTION_CALLS_OPEN in remainder or " bool: - return self.is_reasoning_end(list(input_ids)) - - def extract_content_ids(self, input_ids: list[int]) -> list[int]: - # Content-id slicing is unreliable for multi-token markers; the serving - # path uses extract_reasoning() for the final split. - return [] - - @classmethod - def _scoped_turn(cls, text: str) -> str: - """Current assistant turn with reasoning spans removed.""" - scoped = _current_assistant_turn(text) - scoped = _STRIP_REASONING_RE.sub("", scoped) - return _STRIP_OPEN_REASONING_RE.sub("", scoped) - - @classmethod - def _tool_channel_remainder(cls, text: str) -> str: - """Text from the first tool-channel header onward, framing INCLUDED. - ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever - this parser returns as ``.content`` on the transition delta and commits - it; anything not returned is destroyed. It must start AT the - ``to=<|message|>`` header -- handing over the text after the header - loses the recipient, and the tool parser then sees a bare ``<|message|>``, - classifies it as the content channel, and leaks the ATEM markup. - """ - scoped = cls._scoped_turn(text) - for match in _CHANNEL_HEADER_RE.finditer(scoped): - if match.group("recipient") not in ("self", "user"): - return scoped[match.start() :] - return "" - - @staticmethod - def _classify_bodies(text: str) -> tuple[str, str]: - """Split ``text`` into (reasoning_body, content_body), channel-aware. - Framing markers and tool channels contribute nothing -- the tool parser - owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel - header, or at end-of-text (an OPEN body, which is held back). - """ - reasoning_parts: list[str] = [] - content_parts: list[str] = [] - pos = 0 - n = len(text) - while pos < n: - match = _CHANNEL_HEADER_RE.search(text, pos) - if not match: - break - recipient = match.group("recipient") - body_start = match.end() - eom = text.find(_EOM, body_start) - eot = text.find(_EOT, body_start) - terminators = [p for p in (eom, eot) if p != -1] - next_header = _CHANNEL_HEADER_RE.search(text, body_start) - if next_header is not None: - terminators.append(next_header.start()) - body_end = min(terminators) if terminators else n - body = text[body_start:body_end] - if not terminators: - body = _trim_open_body(body) - if recipient == "self": - reasoning_parts.append(body) - elif ( - # Never surface tool XML echoed into a user channel. - recipient == "user" - and _FUNCTION_CALLS_OPEN not in body - and " str | None: - """Promote un-surfaced content when the stream ends mid-reasoning. - ``DelegatingParser.finalize_generation`` calls this when - ``reasoning_ended`` is still False. Returns only the channel-classified - ``to=user`` body, and only the portion not already streamed. - """ - _, content_body = self._classify_bodies(previous_text) - remainder = content_body[len(self._emitted_content) :] - return remainder or None - - def extract_reasoning( - self, - model_output: str, - request: ChatCompletionRequest | ResponsesRequest, - ) -> tuple[str | None, str | None]: - collapsed = _COLLAPSE_RE.sub("\n", model_output) - matches = _REASONING_RE.findall(collapsed) - reasoning = "\n".join(matches) if matches else None - # Truncation fallback: generation stopped inside a to=self block, so - # there is no closing <|eom|>. Bounded at the next channel header so a - # real tool call that follows a header-less channel switch is not - # absorbed into the reasoning field. - open_match = _OPEN_REASONING_RE.search(model_output) - if open_match and open_match.group(1): - partial = open_match.group(1) - reasoning = f"{reasoning}\n{partial}" if reasoning else partial - # Content is everything that is not a reasoning block. In a - # reasoning+tool-call turn there is no to=user answer, but the tool - # channels MUST be forwarded -- the unified parser runs the tool parser - # on this returned `content`, not on the original model_output. - remainder = _STRIP_REASONING_RE.sub("", model_output) - remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder) - if " DeltaMessage | None: - """Channel-aware streaming split of reasoning vs content. - Classifies the full ``current_text`` and emits only what has not been - emitted yet, so no framing token is ever surfaced and a delta straddling - a channel boundary only contributes the portion inside a real body. - """ - curr_reason, curr_content = self._classify_bodies(current_text) - reasoning_delta = "" - if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len( - self._emitted_reasoning - ): - reasoning_delta = curr_reason[len(self._emitted_reasoning) :] - self._emitted_reasoning = curr_reason - content_delta = "" - if curr_content.startswith(self._emitted_content) and len(curr_content) > len( - self._emitted_content - ): - content_delta = curr_content[len(self._emitted_content) :] - self._emitted_content = curr_content - # Hand the tool channel to the tool parser exactly once, starting at its - # header. parse_delta discards anything not returned here. - # - # This MUST fire on the same delta where is_reasoning_end() flips, i.e. - # only once the tool channel actually contains ATEM. Emitting it earlier - # -- when only the bare `to=<|message|>` header has arrived -- keeps - # the parser in the reasoning phase, so parse_delta never replaces this - # DeltaMessage with the tool parser's and the header is delivered to the - # client as visible content. - handoff = "" - if not self._tool_handoff_done: - remainder = self._tool_channel_remainder(current_text) - if _FUNCTION_CALLS_OPEN in remainder or "`` ... ``<|eom|>``: + to=self<|message|>...reasoning...<|eom|> +Spark also routes recipient-less messages to reasoning, commentary to content, +and drops developer messages. The Glimmer compatibility profile retains its +original channel semantics. +Because Muse's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed +to be single vocab tokens across every checkpoint's tokenizer, this parser works +on the decoded text with regexes rather than the single start/end-token base class. +Usage: ``--reasoning-parser muse`` +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator, Sequence + +import regex as re + +from vllm.entrypoints.generate.base.protocol import DeltaMessage +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser + +_EOM = "<|eom|>" +_EOT = "<|eot|>" +_FUNCTION_CALLS_OPEN = "" +_REASONING_OPEN = "to=self<|message|>" +_ASSISTANT_TURN_OPEN = "<|start|>assistant" +# A channel header: ``to=<|message|>`` where recipient is ``self`` +# (reasoning), ``user`` (final answer) or ``[.]`` (tool call). +_CHANNEL_HEADER_RE = re.compile(r"to=(?P[^\s<]+)<\|message\|>") +_HEADER_PAT = r"to=[^\s<]+<\|message\|>" +# Collapse the gap between reasoning blocks so multiple to=self spans join. +_COLLAPSE_RE = re.compile( + r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL +) +_REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL) +_CONTENT_RE = re.compile( + r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL +) +# Strip a CLOSED reasoning span (header .. <|eom|>). +_STRIP_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL +) +# An UNTERMINATED trailing reasoning span. The model sometimes leaves the +# analysis channel WITHOUT emitting <|eom|>, writing a bare +# ``to=<|message|>`` header instead (observed deterministically for a call +# with EMPTY arguments on a tool that has optional parameters; reproduced on +# other engines too, so it is a model-side defect, not engine-specific). +# +# These two patterns MUST therefore stop at the next channel header rather than +# running to end-of-text. An unbounded ``...$`` version consumes the real tool +# call along with the reasoning: `is_reasoning_end` then never fires, the parser +# never leaves the reasoning phase, the tool parser is never invoked, and the +# entire generation is dropped (empty reasoning, empty content, no tool call). +_STRIP_OPEN_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>" + r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +_OPEN_REASONING_RE = re.compile( + r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +# Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body. +_HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>") +# A trailing fragment that could still grow into a channel header (" t", " to", +# " to=", " to=skill"). Without this the recipient name leaks into reasoning and +# then has to be un-emitted once ``<|message|>`` arrives. +_OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$") + + +_MUSE_HEADER_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?(?:to=(?P[^\s<]+))?<\|message\|>" +) +_MUSE_BOUNDARY_RE = re.compile(r"<\|eom\|>|<\|eot\|>|<\|start\|>") +_MUSE_OTHER_ROLE_RE = re.compile(r"<\|start\|>(?!assistant(?:\s|<|$))[^\s<]+") +_MUSE_NON_TOOL_RECIPIENTS = (None, "self", "user", "commentary", "developer") + + +def _muse_messages(text: str) -> Iterator[tuple[str | None, str, bool, str]]: + """Yield recipient, body, completion status, and framed text per message.""" + headers = list(_MUSE_HEADER_RE.finditer(text)) + for i, header in enumerate(headers): + end = headers[i + 1].start() if i + 1 < len(headers) else len(text) + boundary = _MUSE_BOUNDARY_RE.search(text, header.end(), end) + body_end = boundary.start() if boundary else end + closed = boundary is not None or i + 1 < len(headers) + yield ( + header.group("recipient"), + text[header.end() : body_end], + closed, + text[header.start() : end], + ) + + +def _muse_current_turn(text: str) -> str: + # Prompt examples and earlier turns must not trigger a tool handoff. + roles = list(_MUSE_OTHER_ROLE_RE.finditer(text)) + if roles: + start = text.find(_ASSISTANT_TURN_OPEN, roles[-1].end()) + return text[start:] if start != -1 else "" + return text + + +def _current_assistant_turn(text: str) -> str: + """Return only the text generated in the current assistant turn. + ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start, + and an Muse prompt legitimately contains ATEM markers (``render_tool_defs`` + writes a literal ```` example into the system message, + and prior assistant turns may carry real tool calls). Anchoring on the last + channel-open keeps prompt text from deciding the phase. + """ + idx = text.rfind(_ASSISTANT_TURN_OPEN) + return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text + + +def _trim_open_body(body: str) -> str: + """Hold back any tail of a still-growing body that could still be framing. + Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs + the partial-marker trim (``<``) before the partial-header trim can see + `" to=skill"`. Trimming only once leaks the recipient name as reasoning. + """ + while True: + trimmed = body + for marker in _HOLDBACK_MARKERS: + for k in range(min(len(marker) - 1, len(trimmed)), 0, -1): + if trimmed.endswith(marker[:k]): + trimmed = trimmed[:-k] + break + else: + continue + break + header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed) + if header_tail is not None: + trimmed = trimmed[: header_tail.start()] + if trimmed == body: + return body + body = trimmed + + +class MuseReasoningParser(ReasoningParser): + _glimmer_compat = False + + def __init__(self, tokenizer, *args, **kwargs) -> None: + super().__init__(tokenizer, *args, **kwargs) + # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified + # `previous_text` is unsafe: a classified body legitimately shrinks when + # a partial header becomes recognisable, and diffing against the shrunken + # value re-emits text that already went out. + self._emitted_reasoning: str = "" + self._emitted_content: str = "" + self._tool_handoff_done: bool = False + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Preserve Muse's ATEM framing tokens in the decoded output. + vLLM's serving default is ``skip_special_tokens=True``, which strips + ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the + parsers run, collapsing reasoning into content and breaking channel + scoping. Unlike the base tool-parser hook we do NOT touch + ``structured_outputs`` -- Muse emits native ATEM, not JSON. + """ + request.skip_special_tokens = False + return request + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + """Whether the model has left reasoning and opened a TOOL channel. + A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this + parser surfaces that content itself. Only a real tool channel switches + the ``DelegatingParser`` phase machine over to the tool parser. + Both closed and unterminated reasoning spans are stripped before the + check, so an ```` the model merely echoes inside its CoT + never flips the phase. + """ + try: + text = self.model_tokenizer.decode(input_ids) + except Exception: + return False + remainder = self._tool_channel_remainder(text) + return _FUNCTION_CALLS_OPEN in remainder or " bool: + return self.is_reasoning_end(list(input_ids)) + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + # Content-id slicing is unreliable for multi-token markers; the serving + # path uses extract_reasoning() for the final split. + return [] + + @classmethod + def _scoped_turn(cls, text: str) -> str: + """Current assistant turn with reasoning spans removed.""" + scoped = _current_assistant_turn(text) + scoped = _STRIP_REASONING_RE.sub("", scoped) + return _STRIP_OPEN_REASONING_RE.sub("", scoped) + + @classmethod + def _tool_channel_remainder(cls, text: str) -> str: + """Text from the first tool-channel header onward, framing INCLUDED. + ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever + this parser returns as ``.content`` on the transition delta and commits + it; anything not returned is destroyed. It must start AT the + ``to=<|message|>`` header -- handing over the text after the header + loses the recipient, and the tool parser then sees a bare ``<|message|>``, + classifies it as the content channel, and leaks the ATEM markup. + """ + if not cls._glimmer_compat: + scoped = _muse_current_turn(text) + headers = list(_MUSE_HEADER_RE.finditer(scoped)) + for header, (recipient, body, _, _) in zip(headers, _muse_messages(scoped)): + if recipient not in _MUSE_NON_TOOL_RECIPIENTS and ( + _FUNCTION_CALLS_OPEN in body or " tuple[str, str]: + """Split ``text`` into (reasoning_body, content_body), channel-aware. + Framing markers and tool channels contribute nothing -- the tool parser + owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel + header, or at end-of-text (an OPEN body, which is held back). + """ + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + if not cls._glimmer_compat: + for recipient, body, closed, _ in _muse_messages(text): + if not closed: + body = _trim_open_body(body) + if recipient in (None, "self"): + reasoning_parts.append(body) + elif recipient in ("user", "commentary"): + content_parts.append(body) + return "".join(reasoning_parts), "".join(content_parts) + pos = 0 + n = len(text) + while pos < n: + match = _CHANNEL_HEADER_RE.search(text, pos) + if not match: + break + recipient = match.group("recipient") + body_start = match.end() + eom = text.find(_EOM, body_start) + eot = text.find(_EOT, body_start) + terminators = [p for p in (eom, eot) if p != -1] + next_header = _CHANNEL_HEADER_RE.search(text, body_start) + if next_header is not None: + terminators.append(next_header.start()) + body_end = min(terminators) if terminators else n + body = text[body_start:body_end] + if not terminators: + body = _trim_open_body(body) + if recipient == "self": + reasoning_parts.append(body) + elif ( + # Never surface tool XML echoed into a user channel. + recipient == "user" + and _FUNCTION_CALLS_OPEN not in body + and " str | None: + """Promote un-surfaced content when the stream ends mid-reasoning. + ``DelegatingParser.finalize_generation`` calls this when + ``reasoning_ended`` is still False. Returns only the channel-classified + ``to=user`` body, and only the portion not already streamed. + """ + _, content_body = self._classify_bodies(previous_text) + remainder = content_body[len(self._emitted_content) :] + return remainder or None + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + content: str | None + if not self._glimmer_compat: + messages = list(_muse_messages(model_output)) + if not messages: + return None, model_output or None + reasoning = ( + "".join( + body + for recipient, body, _, _ in messages + if recipient in (None, "self") + ) + or None + ) + has_tool_call = any( + recipient not in _MUSE_NON_TOOL_RECIPIENTS + and (_FUNCTION_CALLS_OPEN in body or ". Bounded at the next channel header so a + # real tool call that follows a header-less channel switch is not + # absorbed into the reasoning field. + open_match = _OPEN_REASONING_RE.search(model_output) + if open_match and open_match.group(1): + partial = open_match.group(1) + reasoning = f"{reasoning}\n{partial}" if reasoning else partial + # Content is everything that is not a reasoning block. In a + # reasoning+tool-call turn there is no to=user answer, but the tool + # channels MUST be forwarded -- the unified parser runs the tool parser + # on this returned `content`, not on the original model_output. + remainder = _STRIP_REASONING_RE.sub("", model_output) + remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder) + if " DeltaMessage | None: + """Channel-aware streaming split of reasoning vs content. + Classifies the full ``current_text`` and emits only what has not been + emitted yet, so no framing token is ever surfaced and a delta straddling + a channel boundary only contributes the portion inside a real body. + """ + if not self._glimmer_compat and not self._tool_handoff_done: + remainder = self._tool_channel_remainder(current_text) + if remainder: + prefix = current_text[: -len(remainder)] + reasoning, content = self._classify_bodies(prefix) + reasoning = reasoning[len(self._emitted_reasoning) :] + content = content[len(self._emitted_content) :] + # A transition batch may contain commentary and several messages. + # Let the tool parser emit all pending channels in that batch. + pending = "" + if reasoning: + pending += f"to=self<|message|>{reasoning}<|eom|>" + if content: + pending += f"to=commentary<|message|>{content}<|eom|>" + self._tool_handoff_done = True + return DeltaMessage(content=pending + remainder) + curr_reason, curr_content = self._classify_bodies(current_text) + reasoning_delta = "" + if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len( + self._emitted_reasoning + ): + reasoning_delta = curr_reason[len(self._emitted_reasoning) :] + self._emitted_reasoning = curr_reason + content_delta = "" + if curr_content.startswith(self._emitted_content) and len(curr_content) > len( + self._emitted_content + ): + content_delta = curr_content[len(self._emitted_content) :] + self._emitted_content = curr_content + # Hand the tool channel to the tool parser exactly once, starting at its + # header. parse_delta discards anything not returned here. + # + # This MUST fire on the same delta where is_reasoning_end() flips, i.e. + # only once the tool channel actually contains ATEM. Emitting it earlier + # -- when only the bare `to=<|message|>` header has arrived -- keeps + # the parser in the reasoning phase, so parse_delta never replaces this + # DeltaMessage with the tool parser's and the header is delivered to the + # client as visible content. + handoff = "" + if not self._tool_handoff_done: + remainder = self._tool_channel_remainder(current_text) + if _FUNCTION_CALLS_OPEN in remainder or "assistant to=self<|message|>...reasoning...<|eom|> - <|start|>assistant to=.<|message|> - - - value - - <|eom|> # non-final call - <|start|>assistant to=user<|message|>...final answer...<|eot|> -Channel scoping is essential: an ```` echoed inside a ``to=self`` -reasoning block or a ``to=user`` final answer must NOT be parsed as a real tool -call. - -Rather than *subtracting* reasoning/answer spans with regex substitutions (the -approach the HF ``response_schema`` uses, which is safe only on a complete, -well-formed turn), this parser *segments* the output into messages and then -*selects* the tool-channel bodies. On a complete turn the two are equivalent; -on a truncated or damaged turn, subtraction can delete a valid tool call -(an unterminated ``to=self`` block makes the non-greedy strip run to the next -``<|eom|>``, which belongs to the tool-call message) whereas selection cannot. - -Usage: ``--enable-auto-tool-choice --tool-call-parser muse_glimmer`` -""" - -from __future__ import annotations - -import json -from collections.abc import Iterator, Sequence - -import regex as re -from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction -from transformers import PreTrainedTokenizerBase - -from vllm.entrypoints.generate.base.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionNamedToolChoiceParam, - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.logger import init_logger -from vllm.renderers.chat_utils import make_tool_call_id -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) - -logger = init_logger(__name__) - -# --- Message framing ------------------------------------------------------- -# An assistant message header. All three parts are optional except the -# <|message|> terminator: -# "<|start|>assistant to=get_weather<|message|>" -- after an <|eom|> boundary -# " to=self<|message|>" -- first message of a turn -# (the prompt already ended -# with "<|start|>assistant") -# "<|message|>" -- bare recipient (public CoT -# / untagged content) -_MSG_HEADER_RE = re.compile( - r"(?:<\|start\|>\s*assistant)?[^\S\n]*(?:to=(?P[A-Za-z0-9_.\-]+))?<\|message\|>" -) -_MSG_END_RE = re.compile(r"<\|eom\|>|<\|eot\|>") - -# Recipients whose bodies are NOT tool calls. -_REASONING_RECIPIENT = "self" -_USER_RECIPIENT = "user" - -# Structural markers that must never reach the client. A streamed body is held -# back by up to len(marker)-1 characters so a marker split across two chunks is -# not emitted as content. -_STRUCTURAL_MARKERS = ("<|eom|>", "<|eot|>", "<|start|>", "<|message|>") -_MAX_MARKER_LEN = max(len(m) for m in _STRUCTURAL_MARKERS) -# A trailing " to=NAME" that could still grow into a bare message header (the -# first message of a turn has no <|start|> prefix -- the prompt ends with -# "<|start|>assistant", so the model's first emitted text is " to=self<|message|>"). -_OPEN_TAIL_TO_RE = re.compile(r"[^\S\n]+to=[A-Za-z0-9_.\-]*$") - -# --- Tool-call extraction (unchanged from MUSE_GLIMMER_RESPONSE_SCHEMA) ------------- -_INVOKE_RE = re.compile(r"()", re.DOTALL) -_NAME_RE = re.compile(r']*?\bname="([^"]+)"') -_PARAM_RE = re.compile( - r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)', - re.DOTALL, -) -_FUNCTION_CALLS_OPEN = "" - - -def _decode_value(raw: str): - """JSON-decode a parameter value when possible, else keep the raw string. - - Mirrors the schema's ``x-parser: json`` with ``allow_non_json: True``. - """ - try: - return json.loads(raw) - except (json.JSONDecodeError, ValueError): - return raw - - -def _iter_messages(text: str) -> Iterator[tuple[str | None, str, bool]]: - """Segment *text* into assistant messages. - - Yields ``(recipient, body, closed)`` per message, where ``recipient`` is - ``None`` for a bare ``<|message|>`` header and ``closed`` is False for a - message that has not (yet) seen ``<|eom|>`` / ``<|eot|>``. - - A message is also terminated by the start of the NEXT header. Without that, - a reasoning block whose ``<|eom|>`` is missing (truncation, or a chunk - dropped at the reasoning -> tool transition) would absorb the tool-call - message that follows it and the call would be lost -- the same defect the - subtractive regexes have. - """ - pos = 0 - while pos < len(text): - header = _MSG_HEADER_RE.search(text, pos) - if header is None: - return - body_start = header.end() - end = _MSG_END_RE.search(text, body_start) - nxt = _MSG_HEADER_RE.search(text, body_start) - body_end = end.start() if end is not None else len(text) - closed = end is not None - if nxt is not None and nxt.start() < body_end: - body_end = nxt.start() - closed = False - next_pos = nxt.start() - else: - next_pos = end.end() if end is not None else len(text) - body = text[body_start:body_end] - # A body can never legitimately contain <|start|>. Seeing one means the - # next header is only partially generated (its <|message|> has not - # arrived), so the regex above could not recognise it yet. Cut there, - # otherwise the streamed body would grow to include the next header and - # then shrink back once it completes. - start_tok = body.find("<|start|>") - if start_tok != -1: - body = body[:start_tok] - closed = False - yield header.group("rcpt"), body, closed - pos = next_pos - - -def _trailing_partial_marker_len(text: str) -> int: - """Length of the longest suffix of *text* that prefixes a structural marker.""" - max_overlap = min(len(text), _MAX_MARKER_LEN - 1) - for overlap in range(max_overlap, 0, -1): - suffix = text[-overlap:] - if any(marker.startswith(suffix) for marker in _STRUCTURAL_MARKERS): - return overlap - return 0 - - -def _safe_open_body(body: str) -> str: - """Trim the tail of a still-growing body to what is safe to emit now. - - Holds back anything that could still turn out to be structural, so the - emitted prefix only ever grows. Chunks under speculative decoding are large - enough that markers routinely straddle them. - """ - tail_to = _OPEN_TAIL_TO_RE.search(body) - if tail_to is not None: - return body[: tail_to.start()] - partial = _trailing_partial_marker_len(body) - return body[: len(body) - partial] if partial else body - - -class MuseGlimmerToolParser(ToolParser): - # MuseGlimmer emits ATEM markup around tool-call arguments. The generic - # named/required tool_choice path in vllm/parser/abstract_parser.py assigns - # the raw model text straight to FunctionCall.arguments, which leaks that - # framing and yields invalid JSON; the "required" branch then silently - # swallows the ValidationError and returns tool_calls=null. Opting out - # routes named/required through extract_tool_calls / - # extract_tool_calls_streaming -- the same path "auto" already uses. - supports_required_and_named = False - - def __init__( - self, - tokenizer: PreTrainedTokenizerBase, - tools: list[Tool] | None = None, - ) -> None: - super().__init__(tokenizer, tools) - # Streaming cursors. vLLM constructs one ToolParser per request, so - # instance state is per-stream. - self._streamed_content_len: int = 0 - self._streamed_reasoning_len: int = 0 - self._emitted_tool_calls: int = 0 - - def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest - ) -> ChatCompletionRequest | ResponsesRequest: - """Force special tokens through, and keep JSON guided decoding off. - - ``skip_special_tokens`` defaults to True on both ChatCompletionRequest - and ResponsesRequest. Every rule in this parser keys off ``<|message|>`` - / ``<|eom|>`` / ``<|eot|>`` / ``<|start|>``, so with the default the - channel framing is stripped before we see it: message segmentation - finds nothing, reasoning-channel invokes are indistinguishable from real - ones, and the raw ATEM markup falls through to the client as content. - Set it unconditionally and FIRST -- not only for tools requests, and not - relying on the reasoning parser to have set it. - - For required/named ``tool_choice`` the base hook installs a JSON schema - constraint (ToolParser.adjust_request -> get_json_schema_from_tools). - MuseGlimmer emits ATEM XML, so under that constraint it writes JSON *inside* - the tool channel -- ``[{"name": ...`` -- with no - ```` for this parser to find. Skip the base hook so those - choices decode natively, the same as "auto". - """ - request.skip_special_tokens = False - - tool_choice = getattr(request, "tool_choice", None) - if request.tools and ( - tool_choice == "required" - or isinstance( - tool_choice, - ( - ChatCompletionNamedToolChoiceParam, - ToolChoiceFunction, - ToolChoiceCustom, - ), - ) - ): - return request - return super().adjust_request(request) - - # ---------------- channel selection ---------------- - - @classmethod - def _tool_channel_text(cls, text: str) -> str: - """Concatenate the bodies of messages addressed to a tool. - - Falls back to the whole text when no message header is present at all -- - that means the framing never reached us (``skip_special_tokens`` was on, - or the chunk carrying the header was dropped upstream), and scanning - everything is strictly better than returning nothing. - """ - bodies = [ - body - for rcpt, body, _closed in _iter_messages(text) - if rcpt is not None - and rcpt != _REASONING_RECIPIENT - and rcpt != _USER_RECIPIENT - ] - if bodies: - return "\n".join(bodies) - if _MSG_HEADER_RE.search(text) is None and ( - _FUNCTION_CALLS_OPEN in text or " tuple[str, str, bool, bool]: - """Return ``(content, reasoning, content_open, reasoning_open)``. - - The ``*_open`` flags say whether that channel's LAST message is still - being generated; only then must the caller hold back a partial - structural marker. Tracking them per channel matters: a closed - reasoning block whose text happens to end in ``<`` would otherwise stay - permanently truncated while a later content message is open. - """ - content_parts: list[str] = [] - reasoning_parts: list[str] = [] - content_open = False - reasoning_open = False - for rcpt, body, closed in _iter_messages(text): - if rcpt == _REASONING_RECIPIENT: - reasoning_parts.append(body) - reasoning_open = not closed - elif rcpt is None or rcpt == _USER_RECIPIENT: - content_parts.append(body) - content_open = not closed - return ( - "".join(content_parts), - "".join(reasoning_parts), - content_open, - reasoning_open, - ) - - # ---------------- tool name binding ---------------- - - @staticmethod - def _registered_names(request: ChatCompletionRequest | None) -> set[str]: - """Names of the tools the client registered on this request.""" - names: set[str] = set() - tools = getattr(request, "tools", None) if request is not None else None - for t in tools or []: - fn = getattr(t, "function", None) or t - name = getattr(fn, "name", None) - if name is None and isinstance(fn, dict): - name = fn.get("name") - if name: - names.add(name) - return names - - @staticmethod - def _normalize_name(emitted: str, registered: set[str]) -> str: - """Map an emitted ATEM invoke name back to a registered tool name. - - When a client registers a BARE name (e.g. ``get_weather``) the shipped - chat template renders the valid recipient as ``"get_weather.*"``, and - the model duly emits ``get_weather.get_weather``. Collapsing that - doubled form is safe: head and tail are identical and the collapsed - name is registered. - - Anything else is passed through unchanged. Matching on the trailing - segment alone is NOT safe -- an emitted ``weather.get`` against a - registered ``{calendar.get}`` has a unique leaf match and would silently - dispatch the wrong tool. - """ - if not registered or emitted in registered: - return emitted - head, sep, tail = emitted.partition(".") - if sep and head == tail and head in registered: - return head - logger.warning( - "MuseGlimmer: emitted tool name %r does not match any registered tool; " - "passing through unchanged.", - emitted, - ) - return emitted - - @classmethod - def _parse_tool_calls( - cls, text: str, registered: set[str] | None = None - ) -> list[ToolCall]: - registered = registered or set() - scoped = cls._tool_channel_text(text) - tool_calls: list[ToolCall] = [] - for invoke in _INVOKE_RE.findall(scoped): - name_m = _NAME_RE.search(invoke) - if not name_m: - continue - name = cls._normalize_name(name_m.group(1), registered) - args: dict = {} - for pm in _PARAM_RE.finditer(invoke): - args[pm.group("key")] = _decode_value(pm.group("value")) - tool_calls.append( - ToolCall( - function=FunctionCall( - name=name, - arguments=json.dumps(args, ensure_ascii=False), - ) - ) - ) - return tool_calls - - @classmethod - def _extract_content(cls, text: str) -> str | None: - """Return the user-facing body, or the raw text when unframed.""" - content, reasoning, _c_open, _r_open = cls._visible_channels(text) - if content: - return content - # No framing at all -> the whole thing is plain content. - if not reasoning and _MSG_HEADER_RE.search(text) is None: - return text or None - return None - - # ---------------- non-streaming ---------------- - - def extract_tool_calls( - self, model_output: str, request: ChatCompletionRequest - ) -> ExtractedToolCallInformation: - if ( - _FUNCTION_CALLS_OPEN not in model_output - and "... parsed -- typically a truncated - # call (finish_reason='length'/abort). Log it: silently - # returning "no tool call" here is indistinguishable from the - # model choosing not to call one, which makes this failure mode - # invisible in production. - logger.warning( - "MuseGlimmer: tool channel opened but no complete " - "parsed (truncated tool call?); returning content only." - ) - return ExtractedToolCallInformation( - tools_called=False, - tool_calls=[], - content=self._extract_content(model_output), - ) - return ExtractedToolCallInformation( - tools_called=True, - tool_calls=tool_calls, - content=self._extract_content(model_output), - ) - except Exception: - logger.exception("Error extracting MuseGlimmer ATEM tool calls.") - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - # ---------------- streaming ---------------- - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - """Incremental ATEM streaming for tool calls AND content. - - This parser owns every delta once reasoning has ended: in - vllm/parser/abstract_parser.py::parse_delta the "pass through as - content" fallback is guarded by ``not self._in_tool_call_phase(state)``, - and ``_in_tool_call_phase`` is simply ``tool_parser is not None and - state.reasoning_ended``. So with a tool parser loaded that fallback is - dead code, and returning None here DISCARDS the delta. Anything we do - not emit -- including the ``to=user`` final answer -- never reaches the - client. Hence content is emitted here, not left to the reasoning parser. - - Tool calls are surfaced only when an ```` block becomes - complete: the XML is opaque until closed and MuseGlimmer parameters are not - incremental JSON, so there is nothing meaningful to stream before then. - """ - if not previous_text: - # First delta of the tool phase (parse_delta resets previous_text - # to "" when it hands the stream over). Reset the cursors. - self._streamed_content_len = 0 - self._streamed_reasoning_len = 0 - self._emitted_tool_calls = 0 - - try: - registered = self._registered_names(request) - calls = self._parse_tool_calls(current_text, registered) - content, reasoning, content_open, reasoning_open = self._visible_channels( - current_text - ) - - # Trim the tail of a channel that is still growing, so the emitted - # prefix never shrinks between deltas. - if content_open: - content = _safe_open_body(content) - if reasoning_open: - reasoning = _safe_open_body(reasoning) - - content_delta = content[self._streamed_content_len :] - reasoning_delta = reasoning[self._streamed_reasoning_len :] - - tool_deltas: list[DeltaToolCall] = [] - for i in range(self._emitted_tool_calls, len(calls)): - fn = calls[i].function - tool_deltas.append( - DeltaToolCall( - index=i, - type="function", - id=make_tool_call_id(), - function=DeltaFunctionCall( - name=fn.name, - arguments=fn.arguments, - ).model_dump(exclude_none=True), - ) - ) - - if not content_delta and not reasoning_delta and not tool_deltas: - return None - - self._streamed_content_len = len(content) - self._streamed_reasoning_len = len(reasoning) - self._emitted_tool_calls = len(calls) - - message = DeltaMessage() - if content_delta: - message.content = content_delta - if reasoning_delta: - message.reasoning = reasoning_delta - if tool_deltas: - message.tool_calls = tool_deltas - return message - except Exception: - logger.exception("Error extracting MuseGlimmer ATEM streaming tool calls.") - return None +class MuseGlimmerToolParser(MuseToolParser): + _preserve_string_parameters = False + _reasoning_recipients = ("self",) + _content_recipients = (None, "user") + _non_tool_recipients = ("self", "user") diff --git a/vllm/tool_parsers/muse_tool_parser.py b/vllm/tool_parsers/muse_tool_parser.py new file mode 100644 index 000000000000..8b684889f94b --- /dev/null +++ b/vllm/tool_parsers/muse_tool_parser.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ATEM tool-call parser for Muse (Muse Spark and Muse Glimmer). + +Port of the ``response_schema`` tool-call contract from the HuggingFace Muse +Glimmer export (``convert_muse_glimmer_weights_to_hf.py``: +``MUSE_GLIMMER_RESPONSE_SCHEMA``), which Muse Spark shares. + +Muse emits tool calls in an XML-ish ATEM format inside channel-scoped messages: + + <|start|>assistant to=self<|message|>...reasoning...<|eom|> + <|start|>assistant to=.<|message|> + + + value + + <|eom|> # non-final call + <|start|>assistant to=user<|message|>...final answer...<|eot|> + +``to=commentary`` messages are user-visible updates between tool calls, and +like ``to=user`` messages they are content; ``to=developer`` messages are +dropped. Channel scoping is essential: an ```` echoed inside a +``to=self`` reasoning block or a ``to=user`` final answer must NOT be parsed as +a real tool call. + +Rather than *subtracting* reasoning/answer spans with regex substitutions (the +approach the HF ``response_schema`` uses, which is safe only on a complete, +well-formed turn), this parser *segments* the output into messages and then +*selects* the tool-channel bodies. On a complete turn the two are equivalent; +on a truncated or damaged turn, subtraction can delete a valid tool call +(an unterminated ``to=self`` block makes the non-greedy strip run to the next +``<|eom|>``, which belongs to the tool-call message) whereas selection cannot. + +Usage: ``--enable-auto-tool-choice --tool-call-parser muse`` +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Sequence + +import regex as re +from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction +from transformers import PreTrainedTokenizerBase + +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.logger import init_logger +from vllm.renderers.chat_utils import make_tool_call_id +from vllm.tool_parsers.abstract_tool_parser import ( + Tool, + ToolParser, +) + +logger = init_logger(__name__) + +# --- Message framing ------------------------------------------------------- +# An assistant message header. All three parts are optional except the +# <|message|> terminator: +# "<|start|>assistant to=get_weather<|message|>" -- after an <|eom|> boundary +# " to=self<|message|>" -- first message of a turn +# (the prompt already ended +# with "<|start|>assistant") +# "<|message|>" -- bare recipient (CoT, +# treated as reasoning) +_MSG_HEADER_RE = re.compile( + r"(?:<\|start\|>\s*assistant)?[^\S\n]*(?:to=(?P[A-Za-z0-9_.\-]+))?<\|message\|>" +) +_MSG_END_RE = re.compile(r"<\|eom\|>|<\|eot\|>") + +# Structural markers that must never reach the client. A streamed body is held +# back by up to len(marker)-1 characters so a marker split across two chunks is +# not emitted as content. +_STRUCTURAL_MARKERS = ("<|eom|>", "<|eot|>", "<|start|>", "<|message|>") +_MAX_MARKER_LEN = max(len(m) for m in _STRUCTURAL_MARKERS) +# A trailing " to=NAME" that could still grow into a bare message header (the +# first message of a turn has no <|start|> prefix -- the prompt ends with +# "<|start|>assistant", so the model's first emitted text is " to=self<|message|>"). +_OPEN_TAIL_TO_RE = re.compile(r"[^\S\n]+to=[A-Za-z0-9_.\-]*$") + +# --- Tool-call extraction (unchanged from MUSE_GLIMMER_RESPONSE_SCHEMA) ------------- +_INVOKE_RE = re.compile(r"()", re.DOTALL) +_NAME_RE = re.compile(r']*?\bname="([^"]+)"') +_PARAM_RE = re.compile( + r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)', + re.DOTALL, +) +_FUNCTION_CALLS_OPEN = "" + + +def _decode_value(raw: str): + """JSON-decode a parameter value when possible, else keep the raw string. + + Mirrors the schema's ``x-parser: json`` with ``allow_non_json: True``. + """ + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw + + +def _iter_messages(text: str) -> Iterator[tuple[str | None, str, bool]]: + """Segment *text* into assistant messages. + + Yields ``(recipient, body, closed)`` per message, where ``recipient`` is + ``None`` for a bare ``<|message|>`` header and ``closed`` is False for a + message that has not (yet) seen ``<|eom|>`` / ``<|eot|>``. + + A message is also terminated by the start of the NEXT header. Without that, + a reasoning block whose ``<|eom|>`` is missing (truncation, or a chunk + dropped at the reasoning -> tool transition) would absorb the tool-call + message that follows it and the call would be lost -- the same defect the + subtractive regexes have. + """ + pos = 0 + while pos < len(text): + header = _MSG_HEADER_RE.search(text, pos) + if header is None: + return + body_start = header.end() + end = _MSG_END_RE.search(text, body_start) + nxt = _MSG_HEADER_RE.search(text, body_start) + body_end = end.start() if end is not None else len(text) + closed = end is not None + if nxt is not None and nxt.start() < body_end: + body_end = nxt.start() + closed = False + next_pos = nxt.start() + else: + next_pos = end.end() if end is not None else len(text) + body = text[body_start:body_end] + # A body can never legitimately contain <|start|>. Seeing one means the + # next header is only partially generated (its <|message|> has not + # arrived), so the regex above could not recognise it yet. Cut there, + # otherwise the streamed body would grow to include the next header and + # then shrink back once it completes. + start_tok = body.find("<|start|>") + if start_tok != -1: + body = body[:start_tok] + closed = False + yield header.group("rcpt"), body, closed + pos = next_pos + + +def _trailing_partial_marker_len(text: str) -> int: + """Length of the longest suffix of *text* that prefixes a structural marker.""" + max_overlap = min(len(text), _MAX_MARKER_LEN - 1) + for overlap in range(max_overlap, 0, -1): + suffix = text[-overlap:] + if any(marker.startswith(suffix) for marker in _STRUCTURAL_MARKERS): + return overlap + return 0 + + +def _safe_open_body(body: str) -> str: + """Trim the tail of a still-growing body to what is safe to emit now. + + Holds back anything that could still turn out to be structural, so the + emitted prefix only ever grows. Chunks under speculative decoding are large + enough that markers routinely straddle them. + """ + tail_to = _OPEN_TAIL_TO_RE.search(body) + if tail_to is not None: + return body[: tail_to.start()] + partial = _trailing_partial_marker_len(body) + return body[: len(body) - partial] if partial else body + + +class MuseToolParser(ToolParser): + _preserve_string_parameters = True + _reasoning_recipients: tuple[str | None, ...] = (None, "self") + _content_recipients: tuple[str | None, ...] = ("user", "commentary") + _non_tool_recipients: tuple[str, ...] = ("self", "user", "commentary", "developer") + + # Muse emits ATEM markup around tool-call arguments. The generic + # named/required tool_choice path in vllm/parser/abstract_parser.py assigns + # the raw model text straight to FunctionCall.arguments, which leaks that + # framing and yields invalid JSON; the "required" branch then silently + # swallows the ValidationError and returns tool_calls=null. Opting out + # routes named/required through extract_tool_calls / + # extract_tool_calls_streaming -- the same path "auto" already uses. + supports_required_and_named = False + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + tools: list[Tool] | None = None, + ) -> None: + super().__init__(tokenizer, tools) + # Streaming cursors. vLLM constructs one ToolParser per request, so + # instance state is per-stream. + self._streamed_content_len: int = 0 + self._streamed_reasoning_len: int = 0 + self._emitted_tool_calls: int = 0 + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Force special tokens through, and keep JSON guided decoding off. + + ``skip_special_tokens`` defaults to True on both ChatCompletionRequest + and ResponsesRequest. Every rule in this parser keys off ``<|message|>`` + / ``<|eom|>`` / ``<|eot|>`` / ``<|start|>``, so with the default the + channel framing is stripped before we see it: message segmentation + finds nothing, reasoning-channel invokes are indistinguishable from real + ones, and the raw ATEM markup falls through to the client as content. + Set it unconditionally and FIRST -- not only for tools requests, and not + relying on the reasoning parser to have set it. + + For required/named ``tool_choice`` the base hook installs a JSON schema + constraint (ToolParser.adjust_request -> get_json_schema_from_tools). + Muse emits ATEM XML, so under that constraint it writes JSON *inside* + the tool channel -- ``[{"name": ...`` -- with no + ```` for this parser to find. Skip the base hook so those + choices decode natively, the same as "auto". + """ + request.skip_special_tokens = False + + tool_choice = getattr(request, "tool_choice", None) + if request.tools and ( + tool_choice == "required" + or isinstance( + tool_choice, + ( + ChatCompletionNamedToolChoiceParam, + ToolChoiceFunction, + ToolChoiceCustom, + ), + ) + ): + return request + return super().adjust_request(request) + + # ---------------- channel selection ---------------- + + @classmethod + def _tool_channel_text(cls, text: str) -> str: + """Concatenate the bodies of messages addressed to a tool. + + Falls back to the whole text when no message header is present at all -- + that means the framing never reached us (``skip_special_tokens`` was on, + or the chunk carrying the header was dropped upstream), and scanning + everything is strictly better than returning nothing. + """ + bodies = [ + body + for rcpt, body, _closed in _iter_messages(text) + if rcpt is not None and rcpt not in cls._non_tool_recipients + ] + if bodies: + return "\n".join(bodies) + if _MSG_HEADER_RE.search(text) is None and ( + _FUNCTION_CALLS_OPEN in text or " tuple[str, str, bool, bool]: + """Return ``(content, reasoning, content_open, reasoning_open)``. + + The ``*_open`` flags say whether that channel's LAST message is still + being generated; only then must the caller hold back a partial + structural marker. Tracking them per channel matters: a closed + reasoning block whose text happens to end in ``<`` would otherwise stay + permanently truncated while a later content message is open. + """ + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + content_open = False + reasoning_open = False + for rcpt, body, closed in _iter_messages(text): + if rcpt in cls._reasoning_recipients: + reasoning_parts.append(body) + reasoning_open = not closed + elif rcpt in cls._content_recipients: + content_parts.append(body) + content_open = not closed + return ( + "".join(content_parts), + "".join(reasoning_parts), + content_open, + reasoning_open, + ) + + # ---------------- tool name binding ---------------- + + @staticmethod + def _registered_names(request: ChatCompletionRequest | None) -> set[str]: + """Names of the tools the client registered on this request.""" + names: set[str] = set() + tools = getattr(request, "tools", None) if request is not None else None + for t in tools or []: + fn = getattr(t, "function", None) or t + name = getattr(fn, "name", None) + if name is None and isinstance(fn, dict): + name = fn.get("name") + if name: + names.add(name) + return names + + @classmethod + def _string_parameters( + cls, request: ChatCompletionRequest | ResponsesRequest | None + ) -> dict[str, set[str]]: + if not cls._preserve_string_parameters: + return {} + strings: dict[str, set[str]] = {} + for tool in getattr(request, "tools", None) or []: + function = getattr(tool, "function", None) or tool + if isinstance(function, dict): + name = function.get("name") + parameters = function.get("parameters") + else: + name = getattr(function, "name", None) + parameters = getattr(function, "parameters", None) + if name and isinstance(parameters, dict): + strings[name] = { + key + for key, schema in parameters.get("properties", {}).items() + if isinstance(schema, dict) and schema.get("type") == "string" + } + return strings + + @staticmethod + def _normalize_name(emitted: str, registered: set[str]) -> str: + """Map an emitted ATEM invoke name back to a registered tool name. + + When a client registers a BARE name (e.g. ``get_weather``) the shipped + chat template renders the valid recipient as ``"get_weather.*"``, and + the model duly emits ``get_weather.get_weather``. Collapsing that + doubled form is safe: head and tail are identical and the collapsed + name is registered. + + Anything else is passed through unchanged. Matching on the trailing + segment alone is NOT safe -- an emitted ``weather.get`` against a + registered ``{calendar.get}`` has a unique leaf match and would silently + dispatch the wrong tool. + """ + if not registered or emitted in registered: + return emitted + head, sep, tail = emitted.partition(".") + if sep and head == tail and head in registered: + return head + logger.warning( + "Muse: emitted tool name %r does not match any registered tool; " + "passing through unchanged.", + emitted, + ) + return emitted + + @classmethod + def _parse_tool_calls( + cls, + text: str, + registered: set[str] | None = None, + strings: dict[str, set[str]] | None = None, + ) -> list[ToolCall]: + registered = registered or set() + scoped = cls._tool_channel_text(text) + tool_calls: list[ToolCall] = [] + for invoke in _INVOKE_RE.findall(scoped): + name_m = _NAME_RE.search(invoke) + if not name_m: + continue + name = cls._normalize_name(name_m.group(1), registered) + string_parameters = (strings or {}).get(name, set()) + args: dict = {} + for pm in _PARAM_RE.finditer(invoke): + key, raw = pm.group("key"), pm.group("value") + args[key] = raw if key in string_parameters else _decode_value(raw) + tool_calls.append( + ToolCall( + function=FunctionCall( + name=name, + arguments=json.dumps(args, ensure_ascii=False), + ) + ) + ) + return tool_calls + + @classmethod + def _extract_content(cls, text: str) -> str | None: + """Return the user-facing body, or the raw text when unframed.""" + content, reasoning, _c_open, _r_open = cls._visible_channels(text) + if content: + return content + # No framing at all -> the whole thing is plain content. + if not reasoning and _MSG_HEADER_RE.search(text) is None: + return text or None + return None + + # ---------------- non-streaming ---------------- + + def extract_tool_calls( + self, model_output: str, request: ChatCompletionRequest + ) -> ExtractedToolCallInformation: + if ( + _FUNCTION_CALLS_OPEN not in model_output + and "... parsed -- typically a truncated + # call (finish_reason='length'/abort). Log it: silently + # returning "no tool call" here is indistinguishable from the + # model choosing not to call one, which makes this failure mode + # invisible in production. + logger.warning( + "Muse: tool channel opened but no complete " + "parsed (truncated tool call?); returning content only." + ) + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._extract_content(model_output), + ) + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=self._extract_content(model_output), + ) + except Exception: + logger.exception("Error extracting Muse ATEM tool calls.") + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + # ---------------- streaming ---------------- + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + """Incremental ATEM streaming for tool calls AND content. + + This parser owns every delta once reasoning has ended: in + vllm/parser/abstract_parser.py::parse_delta the "pass through as + content" fallback is guarded by ``not self._in_tool_call_phase(state)``, + and ``_in_tool_call_phase`` is simply ``tool_parser is not None and + state.reasoning_ended``. So with a tool parser loaded that fallback is + dead code, and returning None here DISCARDS the delta. Anything we do + not emit -- including the ``to=user`` final answer -- never reaches the + client. Hence content is emitted here, not left to the reasoning parser. + + Tool calls are surfaced only when an ```` block becomes + complete: the XML is opaque until closed and Muse parameters are not + incremental JSON, so there is nothing meaningful to stream before then. + """ + if not previous_text: + # First delta of the tool phase (parse_delta resets previous_text + # to "" when it hands the stream over). Reset the cursors. + self._streamed_content_len = 0 + self._streamed_reasoning_len = 0 + self._emitted_tool_calls = 0 + + try: + registered = self._registered_names(request) + calls = self._parse_tool_calls( + current_text, registered, self._string_parameters(request) + ) + content, reasoning, content_open, reasoning_open = self._visible_channels( + current_text + ) + + # Trim the tail of a channel that is still growing, so the emitted + # prefix never shrinks between deltas. + if content_open: + content = _safe_open_body(content) + if reasoning_open: + reasoning = _safe_open_body(reasoning) + + content_delta = content[self._streamed_content_len :] + reasoning_delta = reasoning[self._streamed_reasoning_len :] + + tool_deltas: list[DeltaToolCall] = [] + for i in range(self._emitted_tool_calls, len(calls)): + fn = calls[i].function + tool_deltas.append( + DeltaToolCall( + index=i, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=fn.name, + arguments=fn.arguments, + ).model_dump(exclude_none=True), + ) + ) + + if not content_delta and not reasoning_delta and not tool_deltas: + return None + + self._streamed_content_len = len(content) + self._streamed_reasoning_len = len(reasoning) + self._emitted_tool_calls = len(calls) + + message = DeltaMessage() + if content_delta: + message.content = content_delta + if reasoning_delta: + message.reasoning = reasoning_delta + if tool_deltas: + message.tool_calls = tool_deltas + return message + except Exception: + logger.exception("Error extracting Muse ATEM streaming tool calls.") + return None From 83db1ab4967759df85d1915e27d1974764f869e0 Mon Sep 17 00:00:00 2001 From: Lele Date: Mon, 19 Oct 2026 10:18:33 +0800 Subject: [PATCH 839/952] [Bugfix][LoRA] Keep failed activations retryable (#61618) Signed-off-by: LeslieWylie <112309453+LeslieWylie@users.noreply.github.com> Co-authored-by: linitra24 --- vllm/lora/model_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index a7f9f02e77af..7b8b27c47302 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -361,12 +361,10 @@ def activate_adapter( if first_free_slot is None: raise ValueError("No free lora slots") index, _ = first_free_slot - self._active_adapters[lora_id] = None lora_model = self._registered_adapters[lora_id] logger.debug( "Activating LoRA. int id: %d, slot index: %d", lora_model.id, index ) - self.lora_index_to_id[index] = lora_model.id num_lora_weights_applied = 0 for module_name, module in self.modules.items(): if isinstance(module, ClassificationHeadWithLoRA): @@ -400,6 +398,8 @@ def activate_adapter( "This may be expected with pipeline or expert parallelism.", lora_model.id, ) + self._active_adapters[lora_id] = None + self.lora_index_to_id[index] = lora_model.id return True def _deactivate_adapter(self, lora_id: int): From e719b41e2233db39041ba78a70ac8f425a7f03e1 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:20:56 -0400 Subject: [PATCH 840/952] [ROCm][Perf][DSv4.1] Mono decode for the DSpark draft layers (#62131) Signed-off-by: Fangzhou Ai Co-authored-by: Claude Sonnet 5.5 --- tests/models/test_dsv41_mono_decode.py | 56 ++++++++- vllm/models/deepseek_v41/amd/mono/README.md | 4 +- .../deepseek_v41/amd/mono/attention/front.py | 6 +- vllm/models/deepseek_v41/amd/mono/layer.py | 64 ++++++++-- vllm/models/deepseek_v41/amd/mono/runner.py | 60 ++++++--- vllm/models/deepseek_v41/amd/mono_decode.py | 117 ++++++++++++++---- 6 files changed, 251 insertions(+), 56 deletions(-) diff --git a/tests/models/test_dsv41_mono_decode.py b/tests/models/test_dsv41_mono_decode.py index 309d278bcb42..45e59984dc55 100644 --- a/tests/models/test_dsv41_mono_decode.py +++ b/tests/models/test_dsv41_mono_decode.py @@ -22,9 +22,10 @@ class _Runner: def __init__(self): self.calls = [] + self.draft = None def forward(self, *args, **kwargs): - self.calls.append((args, kwargs)) + self.calls.append((args, {**kwargs, "draft_runner": self.draft})) return ("mono outputs",) def first_pre_mix(self, rows): @@ -86,6 +87,7 @@ def _run( ffn_only=False, entry=None, layer=None, + mono=None, **inputs, ): if metadata is None: @@ -97,10 +99,15 @@ def _run( runner = _Runner() monkeypatch.setattr(md, "is_forward_context_available", lambda: True) monkeypatch.setattr(md, "get_forward_context", lambda: fc) - monkeypatch.setattr(md, "_mono_runner", lambda device: runner) + + def pick(device, draft=None): + runner.draft = draft + return runner + + monkeypatch.setattr(md, "_mono_runner", pick) monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None) args = {**_inputs(), **inputs} - mono = md.MonoDecodeLayer(ffn_only) + mono = mono or md.MonoDecodeLayer(ffn_only) if (entry or ("ffn" if ffn_only else "layer")) == "ffn": # x stands for wo_b's unreduced output; the first layer's attention # has no compressed cache @@ -135,6 +142,49 @@ def test_first_layer_runs_mono_from_the_embedding(run): assert args[2].shape == (M, 4, 5120) and not kwargs["fold"] +DRAFT_WIDTH = 192 # the DSpark draft's non-causal SWA rows: window 128 + 5 drafts + + +def _draft(**overrides): + """A DSpark draft layer's mono path, and the metadata its step carries.""" + layer = _layer(comp=None) + layer.attn.compress_ratio = 0 + swa = _swa( + decode_swa_width=DRAFT_WIDTH, + decode_swa_indices=torch.zeros(64, 1, DRAFT_WIDTH, dtype=torch.int32), + ) + mono = md.MonoDecodeLayer(draft=True, experts=128, topk=3, swa_width=DRAFT_WIDTH) + return dict(layer=layer, mono=mono, metadata={"swa": swa}, **overrides) + + +def test_draft_layer_runs_mono_on_its_own_runner(run): + out, calls = run(**_draft()) + assert out == ("mono outputs",) and len(calls) == 1 + assert calls[0][0][9].stride(0) == DRAFT_WIDTH # the non-causal SWA rows + assert calls[0][1]["draft_runner"] is not None # not the target's runner + + +def test_draft_first_layer_takes_the_broadcast_embedding(run): + """The draft hands its first layer the embedding already repeated to the + hc copies: K1 takes it as the settled residual.""" + x = torch.zeros(M, 4, 5120, dtype=torch.bfloat16) + out, calls = run( + **_draft(residual=None, post_mix=None, res_mix=None, pre_mix=None, x=x) + ) + assert out == ("mono outputs",) + args, kwargs = calls[0] + assert args[1].shape == (M, 5120) and args[2] is x and not kwargs["fold"] + + +def test_draft_and_target_layers_take_only_their_own_swa_rows(run): + causal = {"swa": _swa(), "comp": SimpleNamespace(block_size=128)} + out, calls = run(**{**_draft(), "metadata": causal}) + assert out is None and not calls + target = {"swa": _draft()["metadata"]["swa"], "comp": causal["comp"]} + out, calls = run(metadata=target) + assert out is None and not calls + + def test_index_source_selects_between_the_kernels(run): """An index source hands the runner its compressor and indexer to run between K1 and K2.""" diff --git a/vllm/models/deepseek_v41/amd/mono/README.md b/vllm/models/deepseek_v41/amd/mono/README.md index 462f7bdecd28..2f9be2e15ef4 100644 --- a/vllm/models/deepseek_v41/amd/mono/README.md +++ b/vllm/models/deepseek_v41/amd/mono/README.md @@ -20,7 +20,9 @@ new, reading vLLM's KV records and weights as loaded. their compressor; 24, 28, 32, 36) K1 also hands out its normed input and bf16 qr, vLLM's compressor and indexer run between K1 and K2, and K2 builds the key table from that top-k. The FFN launch serves a layer outside these (an - MXFP8 KV cache, another indexer); the DSpark draft layers keep vLLM's path. + MXFP8 KV cache, another indexer). The DSpark draft's three layers run the same + K1 / K2 on a runner of their own (128 experts / top-3, window-only attention over + the block's non-causal SWA rows). The paths interleave freely (same tensors at the layer boundary). - Steps: decode only (FULL graph or eager), M <= 48 rows (8 requests x 6 DSpark tokens), TP2 / TP4, `fp8_ds_mla` KV records (ROCm), text tokens (the diff --git a/vllm/models/deepseek_v41/amd/mono/attention/front.py b/vllm/models/deepseek_v41/amd/mono/attention/front.py index 23c009739dc2..311a5de79305 100644 --- a/vllm/models/deepseek_v41/amd/mono/attention/front.py +++ b/vllm/models/deepseek_v41/amd/mono/attention/front.py @@ -106,6 +106,7 @@ class FrontBuild: tp: int ratio: int # the layer's compress ratio (0: no top-k keys) timeline: bool = False + swa_width: int = WINDOW # the SWA index rows' width (the DSpark draft's: wider) def plus(t0, v): @@ -202,8 +203,8 @@ def stage_kt(c, t): kv_len = ntopk + nswa for i in range_constexpr((KEYS + THREADS - 1) // THREADS): k = fresh(tid) + THREADS * i - sidx = fx.min(fx.max(k - ntopk, 0), WINDOW - 1) - s_slot = ld_i32(a["swa_idx"], t * WINDOW + sidx) + sidx = fx.min(fx.max(k - ntopk, 0), c["swa_width"] - 1) + s_slot = ld_i32(a["swa_idx"], t * c["swa_width"] + sidx) val = (k < kv_len).select(s_slot, fx.Int32(-1)) if const_expr(ratio > 0): local = ld_i32(a["topk"], t * TOPK + fx.min(k, TOPK - 1)) @@ -531,6 +532,7 @@ def front_context(key, lds, args: dict, scratch, tag, layout=None) -> dict: "mb": Mailbox(tag), "d": d, "ratio": key.ratio, + "swa_width": key.swa_width, "ldbf": lambda ptr, i: fx.Float32( fx.BFloat16(bo.buffer_load(rsrc(ptr), i, vec_width=1, dtype=T.bf16)) ), diff --git a/vllm/models/deepseek_v41/amd/mono/layer.py b/vllm/models/deepseek_v41/amd/mono/layer.py index dcd39ae2b754..a334bedf90e1 100644 --- a/vllm/models/deepseek_v41/amd/mono/layer.py +++ b/vllm/models/deepseek_v41/amd/mono/layer.py @@ -25,7 +25,7 @@ from .attention import back as mb_back from .attention import front as mb_front from .attention.device import BLOCKS, CM_DEV, THREADS, gstore, kernel_symbol, rsrc -from .attention.plan import HIDDEN, Dims, back_scratch, front_scratch +from .attention.plan import HIDDEN, WINDOW, Dims, back_scratch, front_scratch from .common.mx import clamp_fp8 from .common.ops import butterfly, fp8_pack4, peer_bases, traced from .common.plan import WAVES, first_task, key_tuple @@ -33,7 +33,7 @@ from .sources import SOURCES from .stages import moe, seam from .stages.dims import Dims as StageDims -from .stages.moe_shape import SORT_NETS, MoeBuild, route_shape +from .stages.moe_shape import EXPERTS, SORT_NETS, TOPK, MoeBuild, route_shape _STREAM = fx.Stream(None) MAX_TOKENS = 48 @@ -64,21 +64,49 @@ class MonoBuild: # the layer selects its own top-k (an index source): K1 hands its normed # input and qr out for vLLM's indexer, K2 builds the key table index: bool = False + # a layer of the DSpark draft: its routed experts and top-k, the FP4 groups + # of an ug task, and the width of its (non-causal) SWA index rows + experts: int = EXPERTS + topk: int = TOPK + ug_groups: int = 2 + swa_width: int = WINDOW def _moe_key(key: MonoBuild) -> MoeBuild: - return MoeBuild(tokens=key.tokens, tp=key.tp) + return MoeBuild( + tokens=key.tokens, + experts=key.experts, + topk=key.topk, + tp=key.tp, + ug_groups=key.ug_groups, + ) def _front_key(key: MonoBuild): - return mb_front.FrontBuild(key.tokens, key.tp, key.ratio, key.timeline) + return mb_front.FrontBuild( + key.tokens, key.tp, key.ratio, key.timeline, key.swa_width + ) + + +def _symbol_params(key: MonoBuild) -> dict: + """The kernel symbol's draft-only parameters (none for a target layer).""" + out = {} + if key.experts != EXPERTS: + out["e"] = key.experts + if key.topk != TOPK: + out["k"] = key.topk + if key.ug_groups != 2: + out["g"] = key.ug_groups + if key.swa_width != WINDOW: + out["w"] = key.swa_width + return out def _back_key(key: MonoBuild): return mb_back.BackBuild(key.tokens, key.tp, key.ratio, key.timeline) -def scratch_layout(s: int, tp: int) -> dict: +def scratch_layout(s: int, tp: int, experts: int = EXPERTS, topk: int = TOPK) -> dict: """Every region of step width s's launches -> (byte offset, bytes), disjoint: K1's front and seam, K2's back, the MoE's regions, the normed rows and their flags. Each width has a scratch of its own (``DSV41MonoLayer.scratch``): @@ -101,7 +129,9 @@ def place(regions): place(front_scratch(s, d)) place(seam.scratch_layout(s)) place(back_scratch(s, d, start=0)) - regions = moe.scratch_layout(_moe_key(MonoBuild(s, tp, 1))).items() + regions = moe.scratch_layout( + _moe_key(MonoBuild(s, tp, 1, experts=experts, topk=topk)) + ).items() mr = {n: v for n, v in regions if n not in COUNTERS} lo = min(o for o, _ in mr.values()) place({n: (o - lo, n_) for n, (o, n_) in mr.items()}) @@ -110,8 +140,10 @@ def place(regions): return out -def scratch_bytes(s: int = MAX_TOKENS, tp: int = 2) -> int: - return max(o + n for o, n in scratch_layout(s, tp).values()) +def scratch_bytes( + s: int = MAX_TOKENS, tp: int = 2, experts: int = EXPERTS, topk: int = TOPK +) -> int: + return max(o + n for o, n in scratch_layout(s, tp, experts, topk).values()) def peer_half_bytes(tp: int) -> int: @@ -164,7 +196,7 @@ def publish_x8(c, t): def build_mono_k1(key: MonoBuild): s, tp = key.tokens, key.tp assert 1 <= s <= MAX_TOKENS - layout = scratch_layout(s, tp) + layout = scratch_layout(s, tp, key.experts, key.topk) fkey = _front_key(key) FrontLds = mb_front.front_lds(s, Dims(tp)) NORM0 = SLICES # the norms on the CTAs past the slices: the wqkv waits on them @@ -190,6 +222,7 @@ class K1Lds: tl=key.timeline, f=key.fold, i=key.index, + **_symbol_params(key), ) keyed = key_tuple(key, SOURCES) @@ -458,7 +491,7 @@ def run_ffn(key: MonoBuild, lds, a: dict, amb, bases, own, rank, ep, part=None): the residual), the MoE and its all-reduce into ``a["out"]``, then the launch pair's epoch moved on.""" s, tp = key.tokens, key.tp - layout = scratch_layout(s, tp) + layout = scratch_layout(s, tp, key.experts, key.topk) mkey = _moe_key(key) rs = route_shape(mkey) GATE0 = SLICES @@ -554,7 +587,7 @@ def reset(): def build_mono_k2(key: MonoBuild): s, tp = key.tokens, key.tp assert 1 <= s <= MAX_TOKENS - layout = scratch_layout(s, tp) + layout = scratch_layout(s, tp, key.experts, key.topk) bkey = _back_key(key) mkey = _moe_key(key) rs = route_shape(mkey) @@ -574,7 +607,13 @@ class K2Lds: moe: MoeLds # type: ignore[valid-type] name = kernel_symbol( - "dsv41_mono_k2", s=s, tp=tp, r=key.ratio, tl=key.timeline, i=key.index + "dsv41_mono_k2", + s=s, + tp=tp, + r=key.ratio, + tl=key.timeline, + i=key.index, + **_symbol_params(key), ) keyed = key_tuple(key, SOURCES) @@ -688,6 +727,7 @@ def push(t, col, v0, v1): "tid": c["tid"], "dev": True, "ratio": key.ratio, + "swa_width": key.swa_width, "args": { "slot": slot, "swa_idx": swa_idx, diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py index ef9a1978e422..4d64b15633ff 100644 --- a/vllm/models/deepseek_v41/amd/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -22,7 +22,7 @@ import torch -from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, Dims +from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, WINDOW, Dims from .common.plan import BLOCKS from .layer import ( EPOCH_WORDS, @@ -35,7 +35,7 @@ scratch_bytes, ) from .stages.dims import Dims as MoeDims -from .stages.moe_shape import EXPERTS +from .stages.moe_shape import EXPERTS, TOPK __all__ = ["BLOCKS", "MAX_TOKENS", "AttnWeights", "DSV41MonoLayer", "MonoLayerWeights"] @@ -125,10 +125,10 @@ class MonoLayerWeights: sw2: torch.Tensor # shared down [5120, inter] e4m3 sw2_s: torch.Tensor - def check(self, tp: int) -> None: - """The MoE's tensors: 384 experts, TP-sharded intermediates (the routed - one padded as the loader pads it), AITER's A8W4 layout, 32 x 32 E8M0 - blocks for the shared expert.""" + def check(self, tp: int, experts: int = EXPERTS) -> None: + """The MoE's tensors: ``experts`` routed experts (384; the DSpark draft's + 128), TP-sharded intermediates (the routed one padded as the loader pads + it), AITER's A8W4 layout, 32 x 32 E8M0 blocks for the shared expert.""" d, fp4, e4m3, u8 = ( MoeDims(tp), torch.float4_e2m1fn_x2, @@ -136,12 +136,12 @@ def check(self, tp: int) -> None: torch.uint8, ) want = { - "gate_w": ((EXPERTS, HIDDEN), torch.bfloat16), - "bias": ((EXPERTS,), torch.float32), - "w13": ((EXPERTS, 2 * d.inter, HIDDEN // 2), fp4), - "w13_s": ((EXPERTS * 2 * d.inter, HIDDEN // 32), u8), - "w2": ((EXPERTS, HIDDEN, d.inter // 2), fp4), - "w2_s": ((EXPERTS * HIDDEN, d.down_scale_cols), u8), + "gate_w": ((experts, HIDDEN), torch.bfloat16), + "bias": ((experts,), torch.float32), + "w13": ((experts, 2 * d.inter, HIDDEN // 2), fp4), + "w13_s": ((experts * 2 * d.inter, HIDDEN // 32), u8), + "w2": ((experts, HIDDEN, d.inter // 2), fp4), + "w2_s": ((experts * HIDDEN, d.down_scale_cols), u8), "sgu": ((2 * d.sh_inter, HIDDEN), e4m3), "sgu_s": ((2 * d.sh_inter // 32, HIDDEN // 32), u8), "sw2": ((HIDDEN, d.sh_inter), e4m3), @@ -154,11 +154,25 @@ class DSV41MonoLayer: """The mono decode runner of one TP rank (``group``: its TP group, for the peer-memory handle exchange; a gloo / CPU group).""" - def __init__(self, tp: int, rank: int, group, device: torch.device | str = "cuda"): + def __init__( + self, + tp: int, + rank: int, + group, + device: torch.device | str = "cuda", + experts: int = EXPERTS, + topk: int = TOPK, + ug_groups: int = 2, + swa_width: int = WINDOW, + ): _ensure_writable_flydsl_cache() from .common.peer_memory import PeerBuffer self.tp, self.rank = tp, rank + # the layers it runs: the target's (384 experts, top-6, 128-wide causal SWA + # rows) or the DSpark draft's, which has a runner of its own + self.experts, self.topk, self.ug_groups = experts, topk, ug_groups + self.swa_width = swa_width self.d = Dims(tp) dev = self.device = torch.device(device) self._scratch: dict[int, torch.Tensor] = {} @@ -205,7 +219,9 @@ def scratch(self, tokens: int) -> torch.Tensor: "a CUDA graph capture" ) buf = torch.zeros( - scratch_bytes(tokens, self.tp), dtype=torch.uint8, device=self.device + scratch_bytes(tokens, self.tp, self.experts, self.topk), + dtype=torch.uint8, + device=self.device, ) self._scratch[tokens] = buf return buf @@ -220,7 +236,17 @@ def supports(tokens: int) -> bool: def kernels(self, tokens: int, ratio: int, fold: bool = True, index: bool = False): key = (tokens, ratio, fold, index) if key not in self._kernels: - b = MonoBuild(tokens, self.tp, ratio, fold=fold, index=index) + b = MonoBuild( + tokens, + self.tp, + ratio, + fold=fold, + index=index, + experts=self.experts, + topk=self.topk, + ug_groups=self.ug_groups, + swa_width=self.swa_width, + ) self._kernels[key] = (build_mono_k1(b), build_mono_k2(b)) return self._kernels[key] @@ -282,7 +308,9 @@ def forward( assert t_.is_contiguous() assert positions.dtype == slot_mapping.dtype == torch.int64 assert swa_indices.dtype == swa_lens.dtype == token_to_req.dtype == torch.int32 - assert swa_indices.stride(0) == 128 and swa_cache.dtype == torch.uint8 + assert ( + swa_indices.stride(0) == self.swa_width and swa_cache.dtype == torch.uint8 + ) assert swa_cache.shape[-1] == 584 and swa_cache.stride(1) == 584 swa_block = swa_cache.shape[1] if ratio: diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py index 2dd17cc6d21c..6eb1615a7fd6 100644 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -13,9 +13,10 @@ and attention and runs the rest as one launch, the FFN launch, from wo_b's unreduced output. A step takes the kernels when it is decode only, at most ``MAX_ROWS`` rows (eager or a FULL CUDA graph: the model is not torch.compiled); -a whole layer also needs causal SWA windows. Anything else runs the layer as -before; the paths share every tensor at the layer boundary, so they interleave -freely. +a whole layer also needs its SWA rows as the kernels read them (causal windows; +for the DSpark draft's layers, the block's non-causal rows). Anything else runs +the layer as before; the paths share every tensor at the layer boundary, so they +interleave freely. Every TP rank must take the same path at every layer (the kernels wait on each other's pushes): the conditions read only metadata every rank holds alike. @@ -47,17 +48,41 @@ RECORD = 584 # an fp8_ds_mla KV record: 576 data bytes, 8 scale bytes SWA_WIDTH = 128 # a causal decode token's window slots MAX_ROWS = 48 # the kernels' step rows at most: 8 requests x (1 + 5 DSpark drafts) +# FP4 groups of an ug task: the draft's 128-expert unions are small +DRAFT_UG_GROUPS = 1 # the kernels' ISA: CDNA4 (gfx950) scaled MFMA and MX formats CDNA_VERSIONS = (4,) _runner = None - - -def _mono_runner(device: torch.device): - """The process's mono runner (one a TP rank, shared by every layer). Its - peer-memory handshake is collective over the TP group: every rank builds it - at the same layer of the same step.""" - global _runner +_draft_runner = None + + +def _mono_runner(device: torch.device, draft: "MonoDecodeLayer | None" = None): + """The process's mono runner (one a TP rank, shared by every layer of the + target; the DSpark draft's layers share one of their own, as its routing and + SWA rows differ). Its peer-memory handshake is collective over the TP group: + every rank builds it at the same layer of the same step.""" + global _runner, _draft_runner + if draft is not None: + if _draft_runner is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "--monokernel's draft runner was first reached inside a CUDA " + "graph capture, before an eager draft step." + ) + from vllm.models.deepseek_v41.amd.mono.runner import DSV41MonoLayer + + _draft_runner = DSV41MonoLayer( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + get_tp_group().cpu_group, + device, + experts=draft.experts, + topk=draft.topk, + ug_groups=DRAFT_UG_GROUPS, + swa_width=draft.swa_width, + ) + return _draft_runner if _runner is None: if torch.cuda.is_current_stream_capturing(): # its peer memory is allocated and exchanged eagerly; vLLM warms @@ -149,8 +174,22 @@ class MonoDecodeLayer: launch after vLLM's attention (``ffn``). Its weights are taken from the layer at the first call, after loading.""" - def __init__(self, ffn_only: bool = False) -> None: + def __init__( + self, + ffn_only: bool = False, + draft: bool = False, + experts: int = 384, + topk: int = 6, + swa_width: int = SWA_WIDTH, + ) -> None: self.ffn_only = ffn_only + # a layer of the DSpark draft: its own routing and (non-causal) SWA rows + self.draft, self.experts, self.topk, self.swa_width = ( + draft, + experts, + topk, + swa_width, + ) self._weights: MonoLayerWeights | None = None @staticmethod @@ -184,23 +223,30 @@ def create( # the layer: the kernels serve the backbone's seams and MoE config = vllm_config.model_config.hf_config attn, ffn = layer.attn, layer.ffn + draft = attn.layer_id >= config.num_hidden_layers + # a draft layer's routing is the DSpark expert set (128 / 3 on V4.1-Flash) + experts, topk = 384, 6 + if draft: + experts = getattr(config, "dspark_n_routed_experts", 0) or experts + topk = getattr(config, "dspark_num_experts_per_tok", 0) or topk if layer.use_sequence_parallel: why = "sequence-parallel layer" elif not layer.fuse_seam_norm: why = "needs aiter's fused mHC seam" - elif attn.layer_id >= config.num_hidden_layers: - why = "draft layer" elif ( ffn.shared_experts is None or ffn.gate.tid2eid is not None - or ffn.n_routed_experts != 384 - or ffn.n_activated_experts != 6 + or ffn.n_routed_experts != experts + or ffn.n_activated_experts != topk or ffn.routed_scaling_factor != 1.5 or ffn.swiglu_limit != 10.0 or ffn.scoring_func != "sqrtsoftplus" or not ffn.renormalize ): - why = "MoE routing outside the kernels' (384 / 6, sqrtsoftplus, x1.5)" + why = ( + "MoE routing outside the kernels' " + f"({experts} / {topk}, sqrtsoftplus, x1.5)" + ) if why is not None: logger.info_once("DSv4.1 mono decode off for some layers: %s", why) return None @@ -218,11 +264,31 @@ def create( why = f"compress ratio {attn.compress_ratio}" elif attn.kv_mxfp8 or attn.kv_cache_dtype != "fp8_ds_mla": why = f"a {attn.kv_cache_dtype} KV cache" + elif draft and (attn.compress_ratio or attn.indexer is not None): + why = "a draft layer with compressed keys" + if why is not None and draft: + logger.info_once("DSv4.1 mono decode off for the draft layers: %s", why) + return None if why is not None: logger.info_once( "DSv4.1 mono decode: the FFN launch for layers with %s", why ) return MonoDecodeLayer(ffn_only=True) + if draft: + from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, + ) + + spec = vllm_config.speculative_config + logger.info_once("DSv4.1 mono decode: the DSpark draft layers") + return MonoDecodeLayer( + draft=True, + experts=experts, + topk=topk, + swa_width=get_dspark_swa_index_width( + config.sliding_window, spec.num_speculative_tokens + ), + ) logger.info_once("DSv4.1 mono decode: the whole standard backbone layers") return MonoDecodeLayer() @@ -255,7 +321,7 @@ def weights(self, layer: "DeepseekV4DecoderLayer"): sw2=sh.down_proj.weight, sw2_s=_scale_bytes(sh.down_proj.weight_scale), ) - self._weights.check(get_tensor_model_parallel_world_size()) + self._weights.check(get_tensor_model_parallel_world_size(), self.experts) return self._weights @staticmethod @@ -341,7 +407,9 @@ def __call__( if self.ffn_only: return None first = residual is None - if first and x.dim() != 2: + # the draft's first layer takes the embedding already broadcast to the + # hc copies ([rows, hc, dim]), the target's the 2-D embedding + if first and x.dim() != (3 if self.draft else 2): return self._skip("first layer of a pipeline stage") why, swa, comp = self._decode_step(layer, x.shape[0]) attn = layer.attn @@ -350,7 +418,7 @@ def __call__( # the kernels' attention: causal SWA windows, the compressed cache if comp is None and attn.compress_ratio: why = "layer metadata missing" - elif swa.decode_swa_width != SWA_WIDTH: + elif swa.decode_swa_width != self.swa_width: why = f"SWA width {swa.decode_swa_width}" elif ( swa.decode_swa_indices is None @@ -363,15 +431,20 @@ def __call__( if why is not None: return self._skip(why) assert swa is not None - runner = _mono_runner(x.device) + runner = _mono_runner(x.device, self if self.draft else None) logger.info_once( - "DSv4.1 mono decode layer active (decode steps of <= %d rows)", MAX_ROWS + "DSv4.1 mono decode %slayer active (decode steps of <= %d rows)", + "draft " if self.draft else "", + MAX_ROWS, ) # K1's seam folds the owed post into the residual, except where the # residual arrives settled: the embedding, or after the Engram write fold = True if first: - residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() + if self.draft: + residual, x = x, x[:, 0].contiguous() + else: + residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous() pre_mix = runner.first_pre_mix(x.shape[0]) post_mix = res_mix = pre_mix # unread without the fold fold = False From 38127f629025ff4df7c5844f7c0cf94571c8613e Mon Sep 17 00:00:00 2001 From: Suraj Date: Sun, 18 Oct 2026 22:36:29 -0400 Subject: [PATCH 841/952] [Multimodal][Docs] State that retried media downloads count toward the request budget (#62204) Signed-off-by: surajm20061998 Co-authored-by: Claude Opus 5.5 --- docs/usage/security.md | 2 +- tests/test_connections.py | 59 +++++++++++++++++++++++++++++++++++++-- vllm/connections.py | 2 +- 3 files changed, 58 insertions(+), 5 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 4a704578fa93..589e562c7097 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -94,7 +94,7 @@ denial of service: | Environment Variable | Default | Description | | --- | --- | --- | | `VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB` | `256` | Maximum size in MB for a single remote media response. Oversized responses are rejected while streaming before the full body is materialized in memory. | -| `VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB` | `256` | Maximum total HTTP media bytes downloaded for one request, including concurrent downloads. Enforced while streaming. | +| `VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB` | `256` | Maximum total HTTP media bytes downloaded for one request, including concurrent downloads and retried attempts. Enforced while streaming. | | `VLLM_MAX_IMAGE_PIXELS` | `178956970` (~179M pixels) | Maximum decoded image size in pixels. Images exceeding this are rejected before raster memory is allocated. Default matches PIL's built-in 2x decompression-bomb threshold (~680 MB for RGB). | | `VLLM_MAX_MEDIA_REQUEST_DECODE_BYTES` | `715827880` (~683 MiB) | Maximum aggregate decoded image, video-frame, and audio bytes in one request. Images are charged at four bytes per pixel before decoding; audio is limited before decoding; video files are charged after frame extraction. | | `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` | `25` | Maximum compressed filesize in MB for a single audio file. Enforced on all audio inputs (multimodal chat URLs, speech-to-text uploads, data: URLs, and local file paths) before decoding begins. | diff --git a/tests/test_connections.py b/tests/test_connections.py index 6385e372e45e..69df261c86a2 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -6,7 +6,9 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import AsyncMock, patch +import aiohttp import pytest +import requests import vllm.envs as envs from vllm.connections import ( @@ -56,7 +58,7 @@ def local_media_server() -> Generator[str, None, None]: class _SyncResponse: - def __init__(self, chunks: list[bytes]) -> None: + def __init__(self, chunks: list[bytes | Exception]) -> None: self.headers: dict[str, str] = {} self._chunks = chunks self.content_accessed = False @@ -71,6 +73,8 @@ def iter_content(self, chunk_size: int): del chunk_size for chunk in self._chunks: self.iterated += 1 + if isinstance(chunk, Exception): + raise chunk yield chunk def raise_for_status(self) -> None: @@ -84,7 +88,7 @@ def __exit__(self, exc_type, exc_value, traceback): class _AsyncContent: - def __init__(self, chunks: list[bytes]) -> None: + def __init__(self, chunks: list[bytes | Exception]) -> None: self._chunks = chunks self.iterated = 0 @@ -92,11 +96,13 @@ async def iter_chunked(self, chunk_size: int): del chunk_size for chunk in self._chunks: self.iterated += 1 + if isinstance(chunk, Exception): + raise chunk yield chunk class _AsyncResponse: - def __init__(self, chunks: list[bytes]) -> None: + def __init__(self, chunks: list[bytes | Exception]) -> None: self.content = _AsyncContent(chunks) self.content_length = None self.read = AsyncMock( @@ -199,6 +205,53 @@ async def test_async_request_download_budget_counts_across_responses( assert second.content.iterated == 1 +def test_request_download_budget_counts_retried_attempts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + monkeypatch.setattr("vllm.connections.time.sleep", lambda _: None) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + body = b"a" * (3 * _ONE_MIB // 4) + interrupted = _SyncResponse([body[: _ONE_MIB // 2], requests.ConnectionError()]) + + with ( + patch.object( + connection, + "get_response", + side_effect=[interrupted, _SyncResponse([body])], + ), + pytest.raises(ValueError, match="including retried attempts"), + ): + connection.get_bytes("http://example.com/media", request_budget=budget) + + +@pytest.mark.asyncio +async def test_async_request_download_budget_counts_retried_attempts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB", 1) + monkeypatch.setattr("vllm.connections.asyncio.sleep", AsyncMock()) + budget = MediaRequestDownloadBudget() + connection = HTTPConnection() + body = b"a" * (3 * _ONE_MIB // 4) + interrupted = _AsyncResponse( + [body[: _ONE_MIB // 2], aiohttp.ClientConnectionError()] + ) + + with ( + patch.object( + connection, + "get_async_response", + new=AsyncMock(side_effect=[interrupted, _AsyncResponse([body])]), + ), + pytest.raises(ValueError, match="including retried attempts"), + ): + await connection.async_get_bytes( + "http://example.com/media", request_budget=budget + ) + + def test_get_bytes_rejects_content_length_over_limit( monkeypatch: pytest.MonkeyPatch, local_media_server: str ) -> None: diff --git a/vllm/connections.py b/vllm/connections.py index ded8da430f64..2913a640e0cb 100644 --- a/vllm/connections.py +++ b/vllm/connections.py @@ -83,7 +83,7 @@ def charge(self, size: int) -> None: raise MediaRequestDownloadSizeExceededError( "Request media downloads exceed " f"VLLM_MAX_MEDIA_REQUEST_DOWNLOAD_SIZE_MB={limit_mb:g} " - f"({received} bytes received)." + f"({received} bytes received, including retried attempts)." ) self.used = received From 37891a98ea0789607b9ec1c5e896f837b6072ca2 Mon Sep 17 00:00:00 2001 From: YAASHU DAVE <107920191+YaashuDave@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:52:41 -0400 Subject: [PATCH 842/952] [Bugfix] Reject LoRA at config time for models without SupportsLoRA (#61884) Signed-off-by: Yaashu Dave Co-authored-by: vllm-agent --- tests/test_config.py | 22 ++++++++++++++++++++++ vllm/config/lora.py | 8 ++++++++ vllm/config/model.py | 4 ++++ vllm/model_executor/models/mamba2.py | 3 --- vllm/model_executor/models/registry.py | 3 +++ vllm/model_executor/models/zamba2.py | 3 --- 6 files changed, 37 insertions(+), 6 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index b1f977bd6e0b..985af6284561 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -29,6 +29,7 @@ HiSparseConfig, KernelConfig, KVTransferConfig, + LoRAConfig, ModelConfig, ObservabilityConfig, ParallelConfig, @@ -284,6 +285,27 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) +@pytest.mark.parametrize("supported", [True, False]) +def test_lora_rejected_at_config_time_for_unsupported_model(supported): + """Models without SupportsLoRA must fail during config validation with a + clear message, not with a bare assert inside the model constructor.""" + model_config = SimpleNamespace( + is_lora_supported=supported, + architecture="Zamba2ForCausalLM", + architectures=["Zamba2ForCausalLM"], + model="Zyphra/Zamba2-1.2B-instruct", + dtype=torch.bfloat16, + ) + lora_config = LoRAConfig() + if supported: + lora_config.verify_with_model_config(model_config) + else: + with pytest.raises( + ValueError, match="LoRA is not supported for Zamba2ForCausalLM" + ): + lora_config.verify_with_model_config(model_config) + + def test_mamba_cache_mode_all_is_rejected(): """The removed 'all' mode must fail validation instead of being ignored.""" with pytest.raises(ValidationError, match="mamba_cache_mode"): diff --git a/vllm/config/lora.py b/vllm/config/lora.py index a77660f3b130..5eaca2b4930b 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -135,6 +135,14 @@ def _validate_lora_config(self) -> Self: return self def verify_with_model_config(self, model_config: ModelConfig): + if not model_config.is_lora_supported: + raise ValueError( + f"LoRA is not supported for {model_config.architecture} " + f"({model_config.model}). Supported models implement the " + "`SupportsLoRA` interface. Run this model without " + "`enable_lora`." + ) + if self.lora_dtype in (None, "auto"): self.lora_dtype = model_config.dtype elif isinstance(self.lora_dtype, str): diff --git a/vllm/config/model.py b/vllm/config/model.py index 6dc2d16d529e..4620213e529c 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2015,6 +2015,10 @@ def score_type(self) -> ScoreType: def is_pp_supported(self) -> bool: return self._model_info.supports_pp + @property + def is_lora_supported(self) -> bool: + return self._model_info.supports_lora + @property def is_attention_free(self) -> bool: return self._model_info.is_attention_free diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py index cb8b2b481631..7788fe22c3b2 100644 --- a/vllm/model_executor/models/mamba2.py +++ b/vllm/model_executor/models/mamba2.py @@ -101,9 +101,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config - lora_config = vllm_config.lora_config - is_lora_enabled = bool(lora_config) - assert not is_lora_enabled self.config = config diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 0e4235ae6685..31ab5f350b04 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -51,6 +51,7 @@ requires_raw_input_tokens, supports_hc_sp, supports_kv_cache_grouping, + supports_lora, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -877,6 +878,7 @@ class _ModelInfo: requires_raw_input_tokens: bool supports_multimodal_encoder_tp_data: bool supports_pp: bool + supports_lora: bool supports_hc_sp: bool has_inner_state: bool is_attention_free: bool @@ -908,6 +910,7 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": model ), supports_pp=supports_pp(model), + supports_lora=supports_lora(model), supports_hc_sp=supports_hc_sp(model), has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index 35e972818726..1a173a326334 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -714,9 +714,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config - lora_config = vllm_config.lora_config - is_lora_enabled = bool(lora_config) - assert not is_lora_enabled self.config = config From dcf8158cdc3843cda9830936ca9cadf477f1c3dc Mon Sep 17 00:00:00 2001 From: Dilber <44312867+dilberx@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:24:00 +0530 Subject: [PATCH 843/952] [Bugfix][Benchmark] Don't count content-free chunks as streamed outputs in chat (#62341) Signed-off-by: Dilber P Shakir --- benchmarks/backend_request_func.py | 22 +- .../test_endpoint_request_func_timing.py | 190 +++++++++++++++++- vllm/benchmarks/lib/endpoint_request_func.py | 21 +- 3 files changed, 229 insertions(+), 4 deletions(-) diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py index b8e5f5be4c13..0f9977f77fc0 100644 --- a/benchmarks/backend_request_func.py +++ b/benchmarks/backend_request_func.py @@ -8,6 +8,7 @@ import time import traceback from dataclasses import dataclass, field +from typing import Any import aiohttp import huggingface_hub.constants @@ -363,6 +364,23 @@ async def async_request_openai_completions( return output +def _choice_has_streamed_output(choice: dict[str, Any]) -> bool: + delta = choice.get("delta") + if not isinstance(delta, dict): + return False + if ( + delta.get("content") + or delta.get("reasoning") + or delta.get("tool_calls") + or delta.get("function_call") + ): + return True + # A generated token can decode to an empty string, e.g. a special token + # under skip_special_tokens. Those chunks carry `content` and no `role`. + # The opening metadata chunk carries `role`. + return delta.get("content") == "" and "role" not in delta + + async def async_request_openai_chat_completions( request_func_input: RequestFuncInput, pbar: tqdm | None = None, @@ -439,7 +457,9 @@ async def async_request_openai_chat_completions( timestamp = time.perf_counter() data = json.loads(chunk) - if choices := data.get("choices"): + if ( + choices := data.get("choices") + ) and _choice_has_streamed_output(choices[0]): content = choices[0]["delta"].get("content") # First token if not first_chunk_received: diff --git a/tests/benchmarks/test_endpoint_request_func_timing.py b/tests/benchmarks/test_endpoint_request_func_timing.py index d400e6743a3b..efea304f75a3 100644 --- a/tests/benchmarks/test_endpoint_request_func_timing.py +++ b/tests/benchmarks/test_endpoint_request_func_timing.py @@ -4,8 +4,10 @@ serve.py derives TPOT as (latency - ttft) / (output_len - 1), which only means "average inter-token latency" if latency - ttft == sum(itl). These tests pin -that identity, and pin that the trailing choice-less usage chunk does not -extend the measured end of the request. +that identity, pin that the trailing choice-less usage chunk does not extend +the measured end of the request, and pin that chat chunks carrying no token +(the opening role chunk, an empty delta) are skipped for TTFT and ITL while a +token that renders as an empty string is still counted. Driven against a fake session and a scripted clock: no GPU, model, or socket. """ @@ -86,6 +88,20 @@ def _usage_chunk(n_tokens: int) -> bytes: ) +def _role_chunk(role: str = "assistant") -> bytes: + choice = {"index": 0, "delta": {"role": role, "content": ""}} + return _sse( + {"id": "cmpl-test", "object": "chat.completion.chunk", "choices": [choice]} + ) + + +def _empty_chunk() -> bytes: + choice = {"index": 0, "delta": {}} + return _sse( + {"id": "cmpl-test", "object": "chat.completion.chunk", "choices": [choice]} + ) + + def _build_stream( n_tokens: int, *, @@ -508,3 +524,173 @@ def test_legacy_usage_only_stream_is_not_reported_as_success( assert not output.success assert output.itl == [] + + +def _run_chat( + monkeypatch: pytest.MonkeyPatch, + client: str, + chunks: list[tuple[float, bytes]], + n_tokens: int, + *, + clock: _ScriptedClock | None = None, + expect_success: bool = True, +): + if client == "packaged": + if clock is not None: + monkeypatch.setattr(request_func_module, "time", _FakeTime(clock)) + return _run( + async_request_openai_chat_completions, + "http://test/v1/chat/completions", + "chat", + chunks, + n_tokens, + expect_success=expect_success, + ) + elif client == "legacy": + output = _run_legacy( + monkeypatch, + "async_request_openai_chat_completions", + "http://test/v1/chat/completions", + "chat", + chunks, + n_tokens, + clock=clock, + ) + if expect_success: + assert output.success, output.error + return output + raise ValueError(f"Unknown client: {client}") + + +@pytest.mark.parametrize("client", ["packaged", "legacy"]) +def test_chat_role_chunk_is_not_streamed_output( + monkeypatch: pytest.MonkeyPatch, client: str +) -> None: + """Opening role chunk is not a streamed output: TTFT is measured to tok0.""" + tick = 1.0 + clock = _ScriptedClock(tick=tick) + n_tokens = 5 + chunks = [ + (0.0, _role_chunk()), + *[(0.0, _token_chunk(i, chat=True)) for i in range(n_tokens)], + (0.0, _usage_chunk(n_tokens)), + (0.0, b"data: [DONE]\n\n"), + ] + output = _run_chat(monkeypatch, client, chunks, n_tokens, clock=clock) + + assert output.success, output.error + assert len(output.itl) == n_tokens - 1 + assert output.output_tokens == n_tokens + + # Clock read 0 is start_time (1000.0). + # Clock read 1 is role chunk (1001.0); skipped because delta has no content. + # Clock read 2 is tok0 chunk (1002.0); sets TTFT = 1002.0 - 1000.0 = 2 * tick. + expected_ttft = 2.0 * tick + assert output.ttft == pytest.approx(expected_ttft, abs=tick / 1000), ( + f"{client}: TTFT {output.ttft} should be measured to the first content chunk " + f"({expected_ttft}), not the opening role chunk ({1.0 * tick})" + ) + + residual = (output.latency - output.ttft) - sum(output.itl) + assert residual == pytest.approx(0.0, abs=tick / 1000), ( + f"{client}: (latency - ttft) - sum(itl) = {residual!r}" + ) + + +@pytest.mark.parametrize("client", ["packaged", "legacy"]) +def test_chat_midstream_content_free_chunk_adds_no_interval( + monkeypatch: pytest.MonkeyPatch, client: str +) -> None: + """Mid-stream content-free chunk adds no ITL and does not advance timestamp.""" + tick = 1.0 + clock = _ScriptedClock(tick=tick) + n_tokens = 3 + chunks = [ + (0.0, _role_chunk()), + (0.0, _token_chunk(0, chat=True)), + (0.0, _token_chunk(1, chat=True)), + (0.0, _empty_chunk()), + (0.0, _token_chunk(2, chat=True)), + (0.0, _usage_chunk(n_tokens)), + (0.0, b"data: [DONE]\n\n"), + ] + output = _run_chat(monkeypatch, client, chunks, n_tokens, clock=clock) + + assert output.success, output.error + assert len(output.itl) == n_tokens - 1 + # tok0 to tok1 is 1 tick. empty_chunk is at tick 4 but does not advance + # most_recent_timestamp (remains tick 3 from tok1). tok2 is at tick 5, so + # interval is 5 - 3 = 2 ticks. + assert output.itl == pytest.approx([1.0 * tick, 2.0 * tick], abs=tick / 1000), ( + f"{client}: empty chunk should not add interval or advance " + f"most_recent_timestamp: {output.itl}" + ) + + residual = (output.latency - output.ttft) - sum(output.itl) + assert residual == pytest.approx(0.0, abs=tick / 1000), ( + f"{client}: (latency - ttft) - sum(itl) = {residual!r}" + ) + + +@pytest.mark.parametrize("client", ["packaged", "legacy"]) +def test_chat_tokenless_stream_is_rejected( + monkeypatch: pytest.MonkeyPatch, client: str +) -> None: + """A stream with no streamed output is rejected (#54708).""" + chunks = [ + (0.0, _role_chunk()), + (0.0, _empty_chunk()), + (0.0, _usage_chunk(0)), + (0.0, b"data: [DONE]\n\n"), + ] + output = _run_chat(monkeypatch, client, chunks, 0, expect_success=False) + + assert not output.success, f"{client}: tokenless stream should have failed" + assert output.itl == [] + assert "TTFT" in output.error + + +@pytest.mark.parametrize("client", ["packaged", "legacy"]) +def test_chat_empty_string_token_is_counted( + monkeypatch: pytest.MonkeyPatch, client: str +) -> None: + """A token that renders as an empty string is counted as streamed output.""" + tick = 1.0 + clock = _ScriptedClock(tick=tick) + n_tokens = 3 + + def _chunk(content: str) -> bytes: + choice = {"index": 0, "delta": {"content": content}} + return _sse( + {"id": "cmpl-test", "object": "chat.completion.chunk", "choices": [choice]} + ) + + chunks = [ + (0.0, _role_chunk()), + (0.0, _chunk("")), + (0.0, _chunk("a")), + (0.0, _chunk("b")), + (0.0, _usage_chunk(n_tokens)), + (0.0, b"data: [DONE]\n\n"), + ] + output = _run_chat(monkeypatch, client, chunks, n_tokens, clock=clock) + + assert output.success, output.error + assert len(output.itl) == 2 + assert output.output_tokens == n_tokens + + # Clock read 0 is start_time (1000.0). + # Clock read 1 is role chunk (1001.0); skipped because delta carries role. + # Clock read 2 is empty token chunk (1002.0); sets TTFT = 2 * tick. + expected_ttft = 2.0 * tick + assert output.ttft == pytest.approx(expected_ttft, abs=tick / 1000), ( + f"{client}: TTFT {output.ttft} should be measured to the empty token chunk " + f"({expected_ttft}), not the next chunk ({3.0 * tick})" + ) + + assert output.itl == pytest.approx([1.0 * tick, 1.0 * tick], abs=tick / 1000) + + residual = (output.latency - output.ttft) - sum(output.itl) + assert residual == pytest.approx(0.0, abs=tick / 1000), ( + f"{client}: (latency - ttft) - sum(itl) = {residual!r}" + ) diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index a4b55736a951..d6f6c2f5930a 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -368,6 +368,23 @@ def _get_chat_messages( ] +def _choice_has_streamed_output(choice: dict[str, Any]) -> bool: + delta = choice.get("delta") + if not isinstance(delta, dict): + return False + if ( + delta.get("content") + or delta.get("reasoning") + or delta.get("tool_calls") + or delta.get("function_call") + ): + return True + # A generated token can decode to an empty string, e.g. a special token + # under skip_special_tokens. Those chunks carry `content` and no `role`. + # The opening metadata chunk carries `role`. + return delta.get("content") == "" and "role" not in delta + + async def async_request_openai_chat_completions( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, @@ -428,7 +445,9 @@ async def async_request_openai_chat_completions( timestamp = time.perf_counter() data = json.loads(chunk) - if choices := data.get("choices"): + if ( + choices := data.get("choices") + ) and _choice_has_streamed_output(choices[0]): content = choices[0]["delta"].get("content") # First token if not first_chunk_received: From ebf5d31e15b0f02eaa09c9fe85c51f0afab48b1f Mon Sep 17 00:00:00 2001 From: Athrael Soju Date: Mon, 19 Oct 2026 04:09:26 +0100 Subject: [PATCH 844/952] [Model] Add ColUltraVBERT late-interaction retrieval model (#62237) Signed-off-by: Athrael Soju --- docs/models/pooling_models/specific_models.md | 25 + docs/models/pooling_models/token_embed.md | 1 + .../multimodal/pooling/test_colultravbert.py | 126 +++++ tests/models/registry.py | 1 + vllm/model_executor/models/colultravbert.py | 533 ++++++++++++++++++ vllm/model_executor/models/registry.py | 1 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + .../configs/colultravbert.py | 104 ++++ 9 files changed, 794 insertions(+) create mode 100644 tests/models/multimodal/pooling/test_colultravbert.py create mode 100644 vllm/model_executor/models/colultravbert.py create mode 100644 vllm/transformers_utils/configs/colultravbert.py diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 5b2871ca923f..87fed54139aa 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -236,6 +236,31 @@ curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{ An example can be found here: [examples/pooling/score/colqwen3_5_rerank_online.py](../../../examples/pooling/score/colqwen3_5_rerank_online.py) +## ColUltraVBERT Multi-Modal Late Interaction Models + +ColUltraVBERT pairs the UltraBERT text encoder with a SigLIP vision encoder and returns per-token vectors for MaxSim scoring. + +| Architecture | Backbone | Example HF Models | +| - | - | - | +| `ColUltraVBertForRetrieval` | UltraVBERT | `athrael-soju/ColUltraVBERT` | + +Clients send the prompts the model was trained with: a page goes in as `User:Describe the image.\nAssistant:` with its image, and a query as its text followed by ten `` tokens. + +```python +from PIL import Image + +from vllm import LLM + +llm = LLM(model="athrael-soju/ColUltraVBERT", runner="pooling") +page = { + "prompt": "User:Describe the image.\nAssistant:", + "multi_modal_data": {"image": Image.open("page.png").convert("RGB")}, +} +query = "What was the total revenue in 2023?" + "" * 10 +page_output, query_output = llm.encode([page, query], pooling_task="token_embed") +page_vectors, query_vectors = page_output.outputs.data, query_output.outputs.data +``` + ## Llama Nemotron Multimodal ### Embedding Model diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 2bc51f4bec91..a8bbd4de0032 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -59,6 +59,7 @@ Models of any architecture can be converted into embedding models using `--conve | Architecture | Models | Inputs | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | | ------------ | ------ | ----- | ----------------- | ------------------------------ | ------------------------------------------ | | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | | +| `ColUltraVBertForRetrieval` | ColUltraVBERT | T / I | `athrael-soju/ColUltraVBERT` | | | | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | | | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | | | `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | | diff --git a/tests/models/multimodal/pooling/test_colultravbert.py b/tests/models/multimodal/pooling/test_colultravbert.py new file mode 100644 index 000000000000..81deedcb0334 --- /dev/null +++ b/tests/models/multimodal/pooling/test_colultravbert.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +from transformers import Idefics3ImageProcessor + +from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score +from vllm.model_executor.models.colultravbert import get_capped_longest_edge + +MODEL_NAME = "athrael-soju/ColUltraVBERT" +DTYPE = "bfloat16" +PAGE_PROMPT = "User:Describe the image.\nAssistant:" +IMAGE_SEQ_LEN = 64 +EMBEDDING_DIM = 768 +LONGEST_EDGE = 3072 +TILE_SIZE = 512 +MAX_IMAGE_TILES = 27 + + +def page_grid(width: int, height: int) -> tuple[int, int]: + edge = get_capped_longest_edge( + height, width, LONGEST_EDGE, TILE_SIZE, MAX_IMAGE_TILES + ) + processor = Idefics3ImageProcessor( + size={"longest_edge": LONGEST_EDGE}, + max_image_size={"longest_edge": TILE_SIZE}, + ) + _, rows, cols = processor.get_number_of_image_patches( + height, width, {"size": {"longest_edge": edge}} + ) + return rows, cols + + +# Grids from the checkpoint's own UltraVBertImageProcessor. +@pytest.mark.parametrize( + ("width", "height", "grid"), + [ + (1700, 2200, (5, 4)), + (2200, 1700, (4, 5)), + (2000, 2000, (5, 5)), + (4000, 1000, (2, 6)), + (1000, 4000, (6, 2)), + (300, 200, (4, 6)), + (3000, 2400, (4, 5)), + ], +) +def test_colultravbert_tile_cap(width, height, grid): + assert page_grid(width, height) == grid + + +@pytest.fixture(scope="module") +def colultravbert_model(vllm_runner): + with vllm_runner( + MODEL_NAME, + runner="pooling", + dtype=DTYPE, + max_model_len=2048, + enforce_eager=True, + ) as vllm_model: + yield vllm_model + + +def test_colultravbert_text_token_embed(colultravbert_model): + outputs = colultravbert_model.token_embed(["What is machine learning?"]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == EMBEDDING_DIM + assert emb.shape[0] > 1 + + +def test_colultravbert_text_relevance_ordering(colultravbert_model): + query = "What is machine learning?" + documents = [ + "Machine learning is a subset of artificial intelligence.", + "The weather in Paris is mild in spring.", + ] + + scores = colultravbert_model.score(query, documents) + + assert len(scores) == 2 + assert scores[0] > scores[1] + + +def test_colultravbert_text_late_interaction(colultravbert_model): + query = "What is the capital of France?" + doc = "The capital of France is Paris." + + q_out = colultravbert_model.token_embed([query]) + d_out = colultravbert_model.token_embed([doc]) + + q_emb = torch.tensor(q_out[0]) + d_emb = torch.tensor(d_out[0]) + manual_score = compute_maxsim_score(q_emb, d_emb).item() + + vllm_scores = colultravbert_model.score(query, doc) + + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + + +def test_colultravbert_image_token_embed(colultravbert_model, image_assets): + image = image_assets[0].pil_image + inputs = colultravbert_model.get_inputs( + [PAGE_PROMPT], + images=[image], + ) + req_outputs = colultravbert_model.llm.encode( + inputs, + pooling_task="token_embed", + ) + outputs = [req_output.outputs.data for req_output in req_outputs] + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == EMBEDDING_DIM + + rows, cols = page_grid(*image.size) + tokenizer = colultravbert_model.llm.get_tokenizer() + prompt_ids = req_outputs[0].prompt_token_ids + image_tokens = prompt_ids.count(tokenizer.convert_tokens_to_ids("")) + assert image_tokens == (rows * cols + 1) * IMAGE_SEQ_LEN + assert emb.shape[0] == len(prompt_ids) diff --git a/tests/models/registry.py b/tests/models/registry.py index d27aeba00f70..658754188af1 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -678,6 +678,7 @@ def check_available_online( "ColModernVBertForRetrieval": _HfExamplesInfo( "ModernVBERT/colmodernvbert-merged", ), + "ColUltraVBertForRetrieval": _HfExamplesInfo("athrael-soju/ColUltraVBERT"), "ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"), "ColQwen3": _HfExamplesInfo( "TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True diff --git a/vllm/model_executor/models/colultravbert.py b/vllm/model_executor/models/colultravbert.py new file mode 100644 index 000000000000..ba7091dcae8a --- /dev/null +++ b/vllm/model_executor/models/colultravbert.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Iterable, Mapping + +import torch +from torch import nn +from transformers import BatchFeature, Idefics3ImageProcessor, Idefics3Processor +from transformers.models.idefics3.image_processing_idefics3 import ( + MAX_IMAGE_SIZE, + _resize_output_size_rescale_to_max_len, + _resize_output_size_scale_below_upper_bound, +) + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.config.multimodal import MultiModalDummyOptions +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.inputs import MultiModalDataDict +from vllm.model_executor.layers.activation import get_act_and_mul_fn +from vllm.model_executor.layers.attention import EncoderOnlyAttention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.parse import ( + ImageProcessorItems, + ImageSize, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.colultravbert import ( + ColUltraVBertConfig, + UltraBertConfig, +) +from vllm.transformers_utils.processor import cached_get_image_processor + +from .idefics3 import ( + Idefics3DummyInputsBuilder, + Idefics3ImagePixelInputs, + Idefics3MultiModalProcessor, + Idefics3ProcessingInfo, +) +from .interfaces import ( + MultiModalEmbeddings, + SupportsLateInteraction, + SupportsMultiModal, +) +from .interfaces_base import default_pooling_type +from .siglip import SiglipVisionModel +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix + + +class UltraBertAttention(nn.Module): + def __init__( + self, + config: UltraBertConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + layer_type = config.layer_types[layer_id] + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters[layer_type], + ) + # `sliding_window` is an inclusive radius; vLLM's window counts the + # query token, so radius r becomes r + 1. + sliding_window = ( + config.sliding_window + 1 if layer_type == "sliding_attention" else None + ) + self.attn = EncoderOnlyAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + scale=self.head_dim**-0.5, + num_kv_heads=self.num_kv_heads, + quant_config=quant_config, + per_layer_sliding_window=sliding_window, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + output, _ = self.o_proj(self.attn(q, k, v)) + return output + + +class UltraBertMLP(nn.Module): + def __init__( + self, + config: UltraBertConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [config.intermediate_size] * 2, + bias=config.mlp_bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=config.mlp_bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = get_act_and_mul_fn(config.hidden_act) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(hidden_states) + output, _ = self.down_proj(self.act_fn(gate_up)) + return output + + +class UltraBertLayer(nn.Module): + def __init__( + self, + config: UltraBertConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = UltraBertAttention( + config, layer_id, quant_config=quant_config, prefix=f"{prefix}.self_attn" + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = UltraBertMLP( + config, quant_config=quant_config, prefix=f"{prefix}.mlp" + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states + self.self_attn( + positions, self.input_layernorm(hidden_states) + ) + return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) + + +@support_torch_compile +class UltraBertModel(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: UltraBertConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.layers = nn.ModuleList( + [ + UltraBertLayer( + config, + layer_id, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_id}", + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # UltraBERT feeds token embeddings to its first layer unnormalized. + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer(positions, hidden_states) + return self.norm(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("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_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class ColUltraVBertConnector(nn.Module): + def __init__( + self, + config: ColUltraVBertConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.factor = config.pixel_shuffle_factor + hidden_size = config.text_config.hidden_size + self.modality_projection = ReplicatedLinear( + config.vision_config.hidden_size * self.factor**2, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "modality_projection"), + ) + self.norm = ( + RMSNorm(hidden_size, eps=config.text_config.rms_norm_eps) + if config.connector_norm + else nn.Identity() + ) + + def pixel_shuffle(self, x: torch.Tensor) -> torch.Tensor: + batch, length, dim = x.shape + side, f = math.isqrt(length), self.factor + x = x.reshape(batch, side // f, f, side // f, f, dim) + x = x.permute(0, 1, 3, 2, 4, 5) + return x.reshape(batch, length // (f * f), dim * f * f) + + def forward(self, image_hidden_states: torch.Tensor) -> torch.Tensor: + projected, _ = self.modality_projection(self.pixel_shuffle(image_hidden_states)) + return self.norm(projected) + + +def get_capped_longest_edge( + height: int, width: int, longest_edge: int, tile_size: int, max_image_tiles: int +) -> int: + # Counts tiles with the tiling's own resize arithmetic, as transformers' + # UltraVBertImageProcessor does; a closed-form count disagrees at tile + # boundaries. + for edge in range(longest_edge, tile_size - 1, -tile_size): + h, w = _resize_output_size_rescale_to_max_len(height, width, max_len=edge) + h, w = _resize_output_size_scale_below_upper_bound(h, w, max_len=MAX_IMAGE_SIZE) + aspect_ratio = w / h + if w >= h: + cols = math.ceil(w / tile_size) + rows = math.ceil(int(cols * tile_size / aspect_ratio) / tile_size) + else: + rows = math.ceil(h / tile_size) + cols = math.ceil(int(rows * tile_size * aspect_ratio) / tile_size) + if rows * cols <= max_image_tiles: + return edge + return tile_size + + +class ColUltraVBertProcessingInfo(Idefics3ProcessingInfo): + def get_hf_config(self) -> ColUltraVBertConfig: + return self.ctx.get_hf_config(ColUltraVBertConfig) + + def get_image_processor(self) -> Idefics3ImageProcessor: + model_config = self.ctx.model_config + return cached_get_image_processor( + model_config.model, + revision=model_config.revision, + trust_remote_code=model_config.trust_remote_code, + processor_cls_overrides=Idefics3ImageProcessor, + ) + + def get_hf_processor(self, **kwargs: object) -> Idefics3Processor: + return self.ctx.init_processor( + Idefics3Processor, + image_processor=self.get_image_processor(), + tokenizer=self.get_tokenizer(), + image_seq_len=self.get_hf_config().image_seq_len, + ) + + def get_longest_edge(self, *, image_width: int, image_height: int) -> int: + image_processor = self.get_image_processor() + edge = image_processor.size["longest_edge"] + max_image_tiles = getattr(image_processor, "max_image_tiles", None) + if not max_image_tiles: + return edge + return get_capped_longest_edge( + image_height, + image_width, + edge, + image_processor.max_image_size["longest_edge"], + max_image_tiles, + ) + + def _get_image_feature_grid_size( + self, + *, + image_width: int, + image_height: int, + processor: Idefics3Processor, + mm_kwargs: Mapping[str, object], + ) -> tuple[int, int, int]: + edge = self.get_longest_edge(image_width=image_width, image_height=image_height) + return processor.image_processor.get_number_of_image_patches( + image_height, image_width, {"size": {"longest_edge": edge}} + ) + + def get_image_size_with_most_features(self) -> ImageSize: + processor = self.get_hf_processor() + edge = processor.image_processor.size["longest_edge"] + tile = processor.image_processor.max_image_size["longest_edge"] + sizes = [ImageSize(width=edge, height=h) for h in range(tile, edge + 1, tile)] + return max( + sizes, + key=lambda s: self.get_num_image_tokens( + image_width=s.width, + image_height=s.height, + processor=processor, + mm_kwargs={}, + ), + ) + + +class ColUltraVBertDummyInputsBuilder(Idefics3DummyInputsBuilder): + info: ColUltraVBertProcessingInfo + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> MultiModalDataDict: + size = self.info.get_image_size_with_most_features() + return { + "image": self._get_dummy_images( + width=size.width, + height=size.height, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), + ) + } + + +class ColUltraVBertMultiModalProcessor(Idefics3MultiModalProcessor): + info: ColUltraVBertProcessingInfo + + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Each image tiles at its own capped longest edge, which the HF + # processor only takes per call. + images = self.info.parse_mm_data( + {"image": hf_data["images"]}, validate=False + ).get_items("image", ImageProcessorItems) + outputs = [] + for i in range(len(images)): + size = images.get_image_size(i) + edge = self.info.get_longest_edge( + image_width=size.width, image_height=size.height + ) + outputs.append( + super()._call_hf_processor( + {**hf_data, "images": [images.get(i)]}, + {**hf_kwargs, "size": {"longest_edge": edge}}, + ) + ) + return BatchFeature( + { + key: torch.cat([output[key] for output in outputs], dim=1) + for key in ("pixel_values", "pixel_attention_mask") + } + ) + + +@MULTIMODAL_REGISTRY.register_processor( + ColUltraVBertMultiModalProcessor, + info=ColUltraVBertProcessingInfo, + dummy_inputs=ColUltraVBertDummyInputsBuilder, +) +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") +class ColUltraVBertForRetrieval(nn.Module, SupportsMultiModal, SupportsLateInteraction): + is_pooling_model = True + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + # transformers >= 5.10 saves SigLIP without its inner `vision_model.` + # prefix; the longer prefix must stay first. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "vision_model.vision_model.": "vision_model.vision_model.", + "vision_model.": "vision_model.vision_model.", + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: ColUltraVBertConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + + with self._mark_tower_model(vllm_config, "image"): + self.vision_model = SiglipVisionModel( + config.vision_config, + quant_config, + prefix=maybe_prefix(prefix, "vision_model"), + ) + self.connector = ColUltraVBertConnector( + config, quant_config, prefix=maybe_prefix(prefix, "connector") + ) + + with self._mark_language_model(vllm_config): + self.text_model = UltraBertModel( + vllm_config=vllm_config.with_hf_config(config.text_config), + prefix=maybe_prefix(prefix, "text_model"), + ) + + pooler_config = vllm_config.model_config.pooler_config + assert pooler_config is not None + self.pooler = pooler_for_token_embed(pooler_config) + + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> Idefics3ImagePixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + if pixel_values is None: + return None + + size = self.config.vision_config.image_size + return Idefics3ImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + pixel_attention_mask=kwargs.pop("pixel_attention_mask"), + num_patches=kwargs.pop("num_patches"), + resolve_bindings={"h": size, "w": size}, + ) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None: + return [] + + pixel_values = image_input["pixel_values"] + features = self.connector( + self.vision_model(pixel_values.to(dtype=self.vision_model.dtype)) + ) + num_patches = image_input["num_patches"].tolist() + return [f.flatten(0, 1) for f in features.split(num_patches)] + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.text_model( + input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds + ) + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + projector = getattr(self.pooler.head, "projector", None) + if isinstance(projector, nn.Module): + loaded |= { + f"pooler.head.projector.{name}" + for name, _ in projector.named_parameters() + } + return loaded diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 31ab5f350b04..fc8f062c3a2e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -275,6 +275,7 @@ "JinaForRanking": ("jina", "JinaForRanking"), # [Multimodal] "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"), + "ColUltraVBertForRetrieval": ("colultravbert", "ColUltraVBertForRetrieval"), "ColPaliForRetrieval": ("colpali", "ColPaliModel"), "ColQwen3": ("colqwen3", "ColQwen3Model"), "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 69350a53e9ca..f0ec18bde155 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -125,6 +125,7 @@ def __getitem__(self, key): beam="BeamConfig", chatglm="ChatGLMConfig", modernvbert="ColModernVBertConfig", + ultravbert="ColUltraVBertConfig", colpali="ColPaliConfig", colqwen3="ColQwen3Config", ops_colqwen3="OpsColQwen3Config", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index e674dd5a2e85..cb6867faab5d 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -21,6 +21,7 @@ "BeamConfig": "vllm.transformers_utils.configs.beam", "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm", "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert", + "ColUltraVBertConfig": "vllm.transformers_utils.configs.colultravbert", "ColPaliConfig": "vllm.transformers_utils.configs.colpali", "ColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3", @@ -94,6 +95,7 @@ "BeamConfig", "ChatGLMConfig", "ColModernVBertConfig", + "ColUltraVBertConfig", "ColPaliConfig", "ColQwen3Config", "OpsColQwen3Config", diff --git a/vllm/transformers_utils/configs/colultravbert.py b/vllm/transformers_utils/configs/colultravbert.py new file mode 100644 index 000000000000..5688b7131e29 --- /dev/null +++ b/vllm/transformers_utils/configs/colultravbert.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from transformers import PreTrainedConfig, SiglipVisionConfig + + +class UltraBertConfig(PreTrainedConfig): + model_type = "ultrabert" + + def __init__( + self, + vocab_size: int = 49280, + hidden_size: int = 768, + intermediate_size: int = 1024, + num_hidden_layers: int = 24, + num_attention_heads: int = 6, + num_key_value_heads: int | None = None, + head_dim: int | None = None, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + rms_norm_eps: float = 1e-6, + attention_bias: bool = False, + mlp_bias: bool = False, + sliding_window: int = 128, + layer_types: list[str] | None = None, + rope_parameters: dict[str, Any] | None = None, + **kwargs, + ): + if layer_types is None: + layer_types = [ + "sliding_attention" if (i + 1) % 3 else "full_attention" + for i in range(num_hidden_layers) + ] + if rope_parameters is None: + rope_parameters = { + "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0}, + "full_attention": {"rope_type": "default", "rope_theta": 2_560_000.0}, + } + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads or num_attention_heads + self.head_dim = head_dim or hidden_size // num_attention_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.attention_bias = attention_bias + self.mlp_bias = mlp_bias + self.sliding_window = sliding_window + self.layer_types = layer_types + super().__init__(**kwargs) + # Set after the base init: transformers < 5.19 adds a shared + # `rope_theta` key next to the per-layer-type dicts, which 5.17 and + # 5.18 then reject. + self.rope_parameters = rope_parameters + + +class ColUltraVBertConfig(PreTrainedConfig): + model_type = "ultravbert" + + def __init__( + self, + text_config: dict | UltraBertConfig | None = None, + vision_config: dict | SiglipVisionConfig | None = None, + image_token_id: int = 49191, + pixel_shuffle_factor: int = 4, + connector_norm: bool = True, + **kwargs, + ): + if not isinstance(text_config, UltraBertConfig): + text_config = { + k: v + for k, v in (text_config or {}).items() + if k not in ("model_type", "auto_map") + } + text_config = UltraBertConfig(**text_config) + if not isinstance(vision_config, SiglipVisionConfig): + vision_config = { + k: v for k, v in (vision_config or {}).items() if k != "model_type" + } + vision_config = SiglipVisionConfig(**vision_config) + + self.text_config = text_config + self.vision_config: SiglipVisionConfig = vision_config + self.image_token_id = image_token_id + self.pixel_shuffle_factor = pixel_shuffle_factor + self.connector_norm = connector_norm + self.hidden_size = text_config.hidden_size + + kwargs.setdefault("architectures", ["ColUltraVBertForRetrieval"]) + super().__init__(**kwargs) + + @property + def image_seq_len(self) -> int: + grid = self.vision_config.image_size // self.vision_config.patch_size + return grid * grid // self.pixel_shuffle_factor**2 + + def get_text_config(self, **kwargs): + return self.text_config From 7dc1705433e01e5c465547234a0d908d74055756 Mon Sep 17 00:00:00 2001 From: Beichen Huang <71347503+BeichenHuang@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:10:35 -0500 Subject: [PATCH 845/952] [Bugfix][LoRA] Apply text-only LoRAs to the LM of multimodal wrappers (#62224) Signed-off-by: BeichenHuang Co-authored-by: fooSynaptic <19420328+fooSynaptic@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: linitra24 --- tests/lora/test_lora_manager.py | 41 ++++++++++++++++++++++++ vllm/lora/model_manager.py | 56 ++++++++++++++++++++++++++++----- 2 files changed, 89 insertions(+), 8 deletions(-) diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index ffb064ebfbdc..0d1ffdd61f51 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -30,6 +30,7 @@ from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager from vllm.model_executor.layers.fused_moe import GateLinear +from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.platforms import current_platform from .utils import create_peft_lora @@ -138,6 +139,46 @@ def test_from_lora_tensors(qwen3_lora_files, device): assert lora.lora_a.shape[0] == 8 +@pytest.mark.parametrize("is_pooling_model", [False, True]) +@pytest.mark.parametrize( + "lm_prefix,module_name,text_name", + [ + ( + "language_model", + "language_model.model.layers.0.self_attn.o_proj", + "model.layers.0.self_attn.o_proj", + ), + ("language_model", "language_model.lm_head", "lm_head"), + ( + "model.text_model", + "model.text_model.layers.0.self_attn.o_proj", + "model.layers.0.self_attn.o_proj", + ), + ], +) +def test_text_only_lora_binds_to_multimodal_language_model( + lm_prefix, module_name, text_name, is_pooling_model +): + """An adapter trained on the text-only *ForCausalLM must apply to the + language model of a multimodal wrapper instead of silently doing nothing.""" + manager = LoRAModelManager.__new__(LoRAModelManager) + manager.is_pooling_model = is_pooling_model + manager.supports_mm = True + manager.mm_mapping = MultiModelKeys.from_string_field(language_model=lm_prefix) + + def weights(name): + return LoRALayerWeights(name, 8, 16, torch.zeros(8, 4), torch.zeros(4, 8)) + + text_lora = weights(text_name) + adapter = LoRAModel(1, 8, {text_name: text_lora}) + assert manager._get_lora_layer_weights(adapter, module_name) is text_lora + assert manager._get_lora_layer_weights(adapter, "visual.blocks.0.proj") is None + + exact_lora = weights(module_name) + adapter = LoRAModel(1, 8, {text_name: text_lora, module_name: exact_lora}) + assert manager._get_lora_layer_weights(adapter, module_name) is exact_lora + + def create_lora( lora_id: int, model: nn.Module, sub_modules: list[str], device: torch.device ) -> LoRAModel: diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 7b8b27c47302..bd6e48926de3 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -871,8 +871,7 @@ def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: module_name, ) ) - for module in packed_module_names: - lora_model.loras.pop(module, None) + self._pop_lora_weights(lora_model, packed_module_names) continue if self._is_non_gated_moe and len(replacement_loras) > 0: replacement_loras = self._pad_lora_pairs_to_triplets( @@ -891,8 +890,7 @@ def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: # entries that were loaded but did not contribute to the # packed result. Without this they would keep extra CPU # memory alive after pack_moe. - for module in packed_module_names: - lora_model.loras.pop(module, None) + self._pop_lora_weights(lora_model, packed_module_names) for lora in lora_model.loras.values(): lora.optimize() @@ -1128,7 +1126,7 @@ def _convert_3d_to_2d_moe_lora( down_proj_lora.lora_b = [w1_b, down_b, w3_b] # Drop the redundant base_layer entry to avoid double pin_memory # and to keep the activation path looking up only the wrapper key. - lora_model.loras.pop(module_name + ".base_layer", None) + self._pop_lora_weights(lora_model, [module_name + ".base_layer"]) def _slice_moe_lora_ep( self, @@ -1224,15 +1222,57 @@ def _get_lora_layer_weights( if self.is_pooling_model and not lora_model.check_lora_name(module_name): # If it's a pool model, and the layer name is not found, # remove the prefix 'model.' and search again. - module_name = module_name.removeprefix("model.") - if lora_model.check_lora_name(module_name): - org_module_name = module_name + stripped_name = module_name.removeprefix("model.") + if lora_model.check_lora_name(stripped_name): + org_module_name = stripped_name logger.info_once( "For the pool model, successfully loaded the LoRA weights " "after removing the prefix 'model.'." ) + if not lora_model.check_lora_name(org_module_name): + text_name = self._get_text_lora_name(lora_model, module_name) + if text_name is not None: + org_module_name = text_name return lora_model.get_lora(org_module_name) + def _pop_lora_weights(self, lora_model: LoRAModel, module_names: list[str]) -> None: + for module in module_names: + lora_model.loras.pop(module, None) + if text_name := self._get_text_lora_name(lora_model, module): + lora_model.loras.pop(text_name) + + def _get_text_lora_name( + self, lora_model: LoRAModel, module_name: str + ) -> str | None: + for name in self._get_text_lora_names(module_name): + if lora_model.check_lora_name(name): + logger.info_once( + "Loaded text-only LoRA weights for the language model of a " + "multimodal model." + ) + return name + return None + + def _get_text_lora_names(self, module_name: str) -> list[str]: + """Names a text-only adapter uses for a language-model module. + + Adapters trained on the text-only `*ForCausalLM` (e.g. PEFT on + `Qwen3_5ForCausalLM`) name modules `model.layers.*` / `lm_head`, while + multimodal wrappers mount the language model under a prefix such as + `language_model.` or `llm.`. If the prefix holds a `*ForCausalLM`, the + text name is the remainder; if it holds a bare decoder, it is the + remainder under `model.`. + """ + if not getattr(self, "supports_mm", False): + return [] + lm_prefix = self.mm_mapping.language_model[0].rstrip(".") + "." + if not module_name.startswith(lm_prefix): + return [] + name = module_name.removeprefix(lm_prefix) + if name.startswith(("model.", "lm_head")): + return [name] + return [f"model.{name}", name] + def _get_module_to_save_weights( self, lora_model: LoRAModel, module_name: str ) -> LoRAFullModuleWeights | None: From 4e35590dff6d6adfd318486c83efc4a84e749e40 Mon Sep 17 00:00:00 2001 From: Akshat Anand Date: Mon, 19 Oct 2026 08:44:30 +0530 Subject: [PATCH 846/952] [Frontend] Serve Gemma 4 on /v1/systemone and /v1/decisions (#62303) Signed-off-by: cipheraxat Co-authored-by: cipheraxat Co-authored-by: Chauncey --- docs/models/structured_decisions_models.md | 19 ++++++++------ .../structured_decisions/test_strategies.py | 26 +++++++++++++++++++ .../openai/decisions/strategies.py | 3 +++ 3 files changed, 40 insertions(+), 8 deletions(-) diff --git a/docs/models/structured_decisions_models.md b/docs/models/structured_decisions_models.md index 3b53150a5fc4..10e6ca498497 100644 --- a/docs/models/structured_decisions_models.md +++ b/docs/models/structured_decisions_models.md @@ -11,6 +11,7 @@ used for classification and routing. | Architecture | ReadStrategy | Task | Head | Example HF Models | |---------------------------|---------------------|------------------|-------------------|-------------------------| | `ClefFlashForJointSchema` | `ClefFlashStrategy` | `token_classify` | `JointSchemaHead` | `Cloudflare/clef-flash` | +| `Gemma4ForCausalLM` | `NextTokenStrategy` | `generate` | N/A | `google/gemma-4-E2B-it` | | `Qwen3ForCausalLM` | `NextTokenStrategy` | `generate` | N/A | `Qwen/Qwen3-0.6B`, etc. | ### Multimodal Models @@ -18,14 +19,16 @@ used for classification and routing. !!! note For more information about multimodal models inputs, see [this page](./supported_models.md#list-of-multimodal-language-models). -| Architecture | Inputs | ReadStrategy | Task | Head | Example HF Models | -|--------------------------------------|-------------------------------------|--------------------------------|------------------|-------------------|------------------------------------------| -| `ClefOmniForJointSchema` | T + I + V | `JointSchemaStrategy` | `token_classify` | `JointSchemaHead` | `Cloudflare/clef-omni` | -| `D3ForSequenceClassification` | T + I + V | `D3Strategy` | `classify` | Linear (`score`) | `vllm-sr/d3`, `vllm-sr/d3-mini`, etc. | -| `Lfm2VlForConditionalGeneration` | T + I+ | `LiquidStrategy` | `generate` | N/A | `LiquidAI/d1-3B` | -| `Qwen3_5ForConditionalGeneration` | T + IE+ + VE+ | `NextTokenStrategy` | `generate` | N/A | `Qwen/Qwen3.5-0.8B`, etc. | -| `Qwen3_5MoeForConditionalGeneration` | T + IE+ + VE+ | `NextTokenStrategy` | `generate` | N/A | `Qwen/Qwen3.5-35B-A3B`, etc. | -| `DiffusionGemmaForBlockDiffusion` | T + I+ + V+ | `DiffusionGemmaCanvasStrategy` | `generate` | N/A | `google/diffusiongemma-26B-A4B-it`, etc. | +| Architecture | Inputs | ReadStrategy | Task | Head | Example HF Models | +|-----------------------------------------|---------------------------------------|--------------------------------|------------------|-------------------|------------------------------------------| +| `ClefOmniForJointSchema` | T + I + V | `JointSchemaStrategy` | `token_classify` | `JointSchemaHead` | `Cloudflare/clef-omni` | +| `D3ForSequenceClassification` | T + I + V | `D3Strategy` | `classify` | Linear (`score`) | `vllm-sr/d3`, `vllm-sr/d3-mini`, etc. | +| `Gemma4ForConditionalGeneration` | T + I+ + V + A* | `NextTokenStrategy` | `generate` | N/A | `google/gemma-4-E2B-it`, etc. | +| `Gemma4UnifiedForConditionalGeneration` | T + I+ + V + A | `NextTokenStrategy` | `generate` | N/A | `google/gemma-4-12B-it`, etc. | +| `Lfm2VlForConditionalGeneration` | T + I+ | `LiquidStrategy` | `generate` | N/A | `LiquidAI/d1-3B` | +| `Qwen3_5ForConditionalGeneration` | T + IE+ + VE+ | `NextTokenStrategy` | `generate` | N/A | `Qwen/Qwen3.5-0.8B`, etc. | +| `Qwen3_5MoeForConditionalGeneration` | T + IE+ + VE+ | `NextTokenStrategy` | `generate` | N/A | `Qwen/Qwen3.5-35B-A3B`, etc. | +| `DiffusionGemmaForBlockDiffusion` | T + I+ + V+ | `DiffusionGemmaCanvasStrategy` | `generate` | N/A | `google/diffusiongemma-26B-A4B-it`, etc. | ## Online Serving diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 6ddfa5944443..32eb70d09f68 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -97,6 +97,12 @@ def test_strategy_selection(): ) lfm = "Lfm2VlForConditionalGeneration" assert select_read_strategy(model(lfm), generate) is LiquidStrategy + for gemma4 in ( + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "Gemma4UnifiedForConditionalGeneration", + ): + assert select_read_strategy(model(gemma4), generate) is NextTokenStrategy with pytest.raises(ValueError, match="does not support LlamaForCausalLM"): select_read_strategy(model("LlamaForCausalLM"), generate) with pytest.raises(ValueError, match="not raw_logits"): @@ -224,6 +230,26 @@ def test_labels_start_the_reply(qwen): reply_label_ids(tokenizer, tokenizer.encode("team:")) +@pytest.mark.parametrize( + "model_name,reply_start", + [ + # E2B and E4B: the reply starts right after the model turn. + ("google/gemma-4-E2B-it", "<|turn>model\n"), + # 12B, 26B and 31B: the template closes an empty thought first. + ("google/gemma-4-12B-it", "<|turn>model\n<|channel>thought\n"), + ], +) +def test_gemma4_labels_start_the_reply(model_name, reply_start): + from vllm.tokenizers import get_tokenizer + + tokenizer = get_tokenizer(model_name) + renderer = SimpleNamespace(get_tokenizer=lambda: tokenizer) + context = ReadContext(None, SimpleNamespace(renderer=renderer), None, "auto", {}) + # Construction fails unless every label is one token at the reply's start. + strategy = NextTokenStrategy(context) + assert tokenizer.decode(strategy._generation_prompt()).endswith(reply_start) + + @pytest.fixture def engine_and_renderer(qwen): tokenizer, _ = qwen diff --git a/vllm/entrypoints/openai/decisions/strategies.py b/vllm/entrypoints/openai/decisions/strategies.py index fd241ff34fe6..1baf8ce928ee 100644 --- a/vllm/entrypoints/openai/decisions/strategies.py +++ b/vllm/entrypoints/openai/decisions/strategies.py @@ -1161,6 +1161,9 @@ async def read_one(n: int, engine_input: EngineInput) -> list[float]: "ClefFlashForJointSchema": ClefFlashStrategy, "ClefOmniForJointSchema": JointSchemaStrategy, "D3ForSequenceClassification": D3Strategy, + "Gemma4ForCausalLM": NextTokenStrategy, + "Gemma4ForConditionalGeneration": NextTokenStrategy, + "Gemma4UnifiedForConditionalGeneration": NextTokenStrategy, "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, "Qwen3_5ForConditionalGeneration": NextTokenStrategy, From 2b50973f43b11b9e17ddef512413f5bde9e1b711 Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Mon, 19 Oct 2026 11:39:05 +0800 Subject: [PATCH 847/952] [Quantization][Autoround][ARK] Add ARK W4A16(moe) Support (#50396) Signed-off-by: Zhenzhong1 --- tests/quantization/test_auto_round.py | 193 +++++++- .../quantization/inc/schemes/inc_wna16_moe.py | 440 ++++++++++++++++++ .../inc/schemes/inc_wna16_scheme.py | 163 +++---- 3 files changed, 668 insertions(+), 128 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_moe.py diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 267e1a11ba86..8869b8d177ed 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -26,6 +26,7 @@ INCMxfp4Scheme, INCMxfp8Scheme, INCWna16Scheme, + inc_wna16_moe, resolve_scheme, ) from vllm.model_executor.layers.quantization.inc.schemes.inc_fp8_linear import ( @@ -45,10 +46,6 @@ INCWNA16LinearScheme, INCXPULinearMethod, ) -from vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_scheme import ( - _resolve_awq_moe, - _resolve_gptq_moe, -) from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.platforms import current_platform @@ -833,15 +830,19 @@ def test_qwen3_30b_a3b_w4a16_autoround_routes_to_gptq_moe( class DummyMoeConfig: pass - def fake_resolve_gptq_moe(layer, layer_config): - captured["layer"] = layer - captured["layer_config"] = layer_config - return expected_method + class DummyMoEScheme: + def __init__(self, layer_config) -> None: + captured["layer_config"] = layer_config + + def get_method(self, layer): + captured["layer"] = layer + return expected_method + monkeypatch.setattr(current_platform, "is_xpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes." - "inc_wna16_scheme._resolve_gptq_moe", - fake_resolve_gptq_moe, + inc_wna16_moe, + "INCWNA16MoEScheme", + DummyMoEScheme, ) config = make_qwen3_autoround_config("qwen3_30b_a3b_w4a16") @@ -1159,15 +1160,25 @@ def test_wna16_xpu_moe_routes_to_gptq_moe(monkeypatch) -> None: class DummyMoeConfig: pass + class DummyMoEScheme: + def __init__(self, layer_config) -> None: + captured["layer_config"] = layer_config + + def get_method(self, layer): + captured["layer"] = layer + return expected_method + + monkeypatch.delenv(_BACKEND_ENV, raising=False) monkeypatch.setattr(current_platform, "is_xpu", lambda: True) monkeypatch.setattr(current_platform, "is_cpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes." - "inc_wna16_scheme._resolve_gptq_moe", - lambda layer, layer_config: captured.update( - {"layer": layer, "layer_config": layer_config} - ) - or expected_method, + _ARK_STATE, + lambda: (False, "missing", None, None), + ) + monkeypatch.setattr( + inc_wna16_moe, + "INCWNA16MoEScheme", + DummyMoEScheme, ) layer = object.__new__(RoutedExperts) @@ -1184,6 +1195,146 @@ class DummyMoeConfig: assert captured["layer_config"].is_gptq is True +def test_wna16_xpu_moe_prefers_ark_when_moe_kernel_available(monkeypatch) -> None: + captured = {} + + class DummyMoeConfig: + pass + + class DummyArk: + xpu_lib = object() + + class MoeSymmetricGemm: + pass + + class DummyArkMethod: + def __init__(self, quant_config, moe_config) -> None: + captured["quant_config"] = quant_config + captured["moe_config"] = moe_config + + monkeypatch.delenv(_BACKEND_ENV, raising=False) + monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + monkeypatch.setattr(_ARK_STATE, lambda: (True, None, DummyArk(), object())) + monkeypatch.setattr( + inc_wna16_moe, + "INCARKWNA16MoEMethod", + DummyArkMethod, + ) + + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + method = INCWna16Scheme().get_moe_method( + make_config(), + layer, + "model.layers.0.mlp", + make_layer_config(group_size=32), + ) + + assert isinstance(method, DummyArkMethod) + assert captured["quant_config"].weight_bits == 4 + assert captured["quant_config"].group_size == 32 + assert captured["moe_config"] is layer.moe_config + + +def test_wna16_xpu_moe_falls_back_when_ark_lacks_moe_kernel(monkeypatch) -> None: + expected_method = object() + + class DummyMoeConfig: + pass + + class DummyArk: + xpu_lib = object() + + class DummyMoEScheme: + def __init__(self, layer_config) -> None: + self.layer_config = layer_config + + def get_method(self, layer): + return expected_method + + monkeypatch.delenv(_BACKEND_ENV, raising=False) + monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + monkeypatch.setattr(_ARK_STATE, lambda: (True, None, DummyArk(), object())) + monkeypatch.setattr( + inc_wna16_moe, + "INCWNA16MoEScheme", + DummyMoEScheme, + ) + + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + method = INCWna16Scheme().get_moe_method( + make_config(), + layer, + "model.layers.0.mlp", + make_layer_config(group_size=32), + ) + + assert method is expected_method + + +def test_wna16_xpu_moe_ark_backend_requested_but_unavailable_raises( + monkeypatch, +) -> None: + class DummyMoeConfig: + pass + + monkeypatch.setenv(_BACKEND_ENV, "ark") + monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + monkeypatch.setattr(_ARK_STATE, lambda: (False, "missing", None, None)) + + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + with pytest.raises(NotImplementedError, match="ARK MoE kernels are unavailable"): + INCWna16Scheme().get_moe_method( + make_config(), + layer, + "model.layers.0.mlp", + make_layer_config(group_size=32), + ) + + +def test_wna16_xpu_moe_w4a16_backend_overrides_available_ark(monkeypatch) -> None: + expected_method = object() + + class DummyMoeConfig: + pass + + class DummyMoEScheme: + def __init__(self, layer_config) -> None: + self.layer_config = layer_config + + def get_method(self, layer): + return expected_method + + monkeypatch.setenv(_BACKEND_ENV, "w4a16") + monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + monkeypatch.setattr( + inc_wna16_moe, + "INCWNA16MoEScheme", + DummyMoEScheme, + ) + + layer = object.__new__(RoutedExperts) + layer.moe_config = DummyMoeConfig() + + method = INCWna16Scheme().get_moe_method( + make_config(), + layer, + "model.layers.0.mlp", + make_layer_config(group_size=32), + ) + + assert method is expected_method + + def test_inc_resolve_scheme_selects_mxfp8() -> None: layer_config = INCLayerConfig( bits=8, @@ -1780,7 +1931,7 @@ def __init__(self, cfg, moe): quantized=True, ) - _resolve_gptq_moe(DummyLayer(), layer_config) + inc_wna16_moe.INCWNA16MoEScheme(layer_config)._build_gptq_method(DummyLayer()) assert captured["from_config"] == { "quant_method": "gptq", @@ -1817,7 +1968,9 @@ def __init__(self, cfg, moe): DummyMethod, ) - _resolve_gptq_moe(DummyLayer(), make_layer_config()) + inc_wna16_moe.INCWNA16MoEScheme(make_layer_config())._build_gptq_method( + DummyLayer() + ) assert isinstance(captured["cfg"], AutoGPTQConfig) assert captured["cfg"].weight_bits == 4 @@ -1863,7 +2016,7 @@ def __init__(self, cfg, moe): quantized=True, ) - _resolve_awq_moe(DummyLayer(), layer_config) + inc_wna16_moe.INCWNA16MoEScheme(layer_config)._build_awq_method(DummyLayer()) assert captured["cfg"].weight_bits == 4 assert captured["cfg"].zero_point is True diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_moe.py new file mode 100644 index 000000000000..8eca99daec89 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_moe.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import ( + ApplyMoEActivationConfig, + apply_moe_activation, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig +from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig +from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from ..config_parser import INCLayerConfig + +logger = init_logger(__name__) + + +class INCARKWNA16MoEMethod(MoeWNA16Method): + def __init__( + self, + quant_config: MoeWNA16Config, + moe: FusedMoEConfig, + ) -> None: + if quant_config.weight_bits != 4 or quant_config.has_zp: + raise NotImplementedError( + "ARK WNA16 MoE only supports symmetric int4 W4A16 for now." + ) + + super().__init__(quant_config, moe) + + from .inc_ark_ops import get_ark_state + + is_available, error_str, ark, _ = get_ark_state() + xpu_lib = getattr(ark, "xpu_lib", None) if ark is not None else None + has_moe_kernel = ( + is_available + and ark is not None + and hasattr(ark, "MoeSymmetricGemm") + and xpu_lib is not None + ) + if not has_moe_kernel: + reason = error_str or "ARK MoE kernels are unavailable." + raise ImportError(f"Failed to initialize ARK WNA16 MoE. {reason}") + + assert ark is not None + self.ark = ark + self.remap_hidden_states_op = None + self.moe_gather_op = None + self.local_num_experts: int | None = None + self.global_num_experts: int | None = None + self.expert_map: torch.Tensor | None = None + self.rows_per_expert: torch.Tensor | None = None + self.unpermuted_row_to_permuted_row: torch.Tensor | None = None + self.router_weight_ones: torch.Tensor | None = None + self.router_weights_fn = self._keep_router_weights + self.activation = None + self.activation_config = ApplyMoEActivationConfig.from_configs( + moe, + FusedMoEQuantConfig.make(), + ) + self.inter_size: int = 0 + self.inter_size_scale: int = 1 + self.w13_moe = None + self.w2_moe = None + + logger.info_once("Using ARK XPU WNA16 MoE kernel.") + + def process_weights_after_loading(self, layer) -> None: + # ARK int4 MoE with asym=False expects signed nibbles [-8, 7]. + # Loaded GPTQ weights are unsigned nibbles [0, 15], so flip the sign + # bit of each nibble, equivalent to q - 8 in 4-bit two's-complement. + w13_qweight = torch.bitwise_xor( + layer.w13_qweight.detach().contiguous(), + 0x88, + ) + w2_qweight = torch.bitwise_xor( + layer.w2_qweight.detach().contiguous(), + 0x88, + ) + + replace_parameter(layer, "w13_qweight", w13_qweight) + replace_parameter(layer, "w2_qweight", w2_qweight) + replace_parameter(layer, "w13_scales", layer.w13_scales.detach().contiguous()) + replace_parameter(layer, "w2_scales", layer.w2_scales.detach().contiguous()) + + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight + num_local_experts = layer.w13_weight.shape[0] + self.local_num_experts = num_local_experts + self.global_num_experts = layer.global_num_experts + self.expert_map = layer.expert_map + self.rows_per_expert = torch.empty( + (num_local_experts,), + dtype=torch.int32, + device=layer.w13_weight.device, + ) + self.unpermuted_row_to_permuted_row = None + self.router_weight_ones = None + if layer.apply_router_weight_on_input: + if layer.top_k != 1: + raise NotImplementedError( + "apply_router_weight_on_input is only supported for topk=1." + ) + self.router_weights_fn = self._apply_router_weight_to_input + else: + self.router_weights_fn = self._keep_router_weights + self.activation = layer.activation + self.activation_config = ApplyMoEActivationConfig.from_configs( + self.moe, + FusedMoEQuantConfig.make( + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + ), + ) + self.inter_size = layer.w13_weight.shape[-2] // 2 + self.inter_size_scale = 1 if layer.activation.is_gated else 2 + self.remap_hidden_states_op = torch.ops._moe_C.remap_hidden_states + self.moe_gather_op = torch.ops._moe_C.moe_gather + + self.w13_moe = self.ark.MoeSymmetricGemm.prepare( + layer.w13_weight, + layer.w13_scales, + weight_bits=4, + group_size=layer.group_size, + activation_dtype=layer.w13_scales.dtype, + ) + + self.w2_moe = self.ark.MoeSymmetricGemm.prepare( + layer.w2_weight, + layer.w2_scales, + weight_bits=4, + group_size=layer.group_size, + activation_dtype=layer.w2_scales.dtype, + ) + + def _get_rows_per_expert( + self, + local_num_experts: int, + device: torch.device, + ) -> torch.Tensor: + rows_per_expert = self.rows_per_expert + if rows_per_expert is None or rows_per_expert.device != device: + rows_per_expert = torch.empty( + (local_num_experts,), + dtype=torch.int32, + device=device, + ) + self.rows_per_expert = rows_per_expert + rows_per_expert.zero_() + return rows_per_expert + + def _get_unpermuted_row_to_permuted_row( + self, + num_rows: int, + topk: int, + device: torch.device, + ) -> torch.Tensor: + mapping = self.unpermuted_row_to_permuted_row + if ( + mapping is None + or mapping.device != device + or mapping.shape[0] < num_rows + or mapping.shape[1] != topk + ): + mapping = torch.empty( + (num_rows, topk), + dtype=torch.int32, + device=device, + ) + self.unpermuted_row_to_permuted_row = mapping + return mapping[:num_rows] + + def _get_router_weight_ones( + self, + num_rows: int, + topk: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + router_weight_ones = self.router_weight_ones + if ( + router_weight_ones is None + or router_weight_ones.device != device + or router_weight_ones.dtype != dtype + or router_weight_ones.shape[0] < num_rows + or router_weight_ones.shape[1] != topk + ): + router_weight_ones = torch.ones( + (num_rows, topk), + dtype=dtype, + device=device, + ) + self.router_weight_ones = router_weight_ones + return router_weight_ones[:num_rows] + + def _keep_router_weights( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + num_rows: int, + topk: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + del num_rows, topk + return x, topk_weights + + def _apply_router_weight_to_input( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + num_rows: int, + topk: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + x = x * topk_weights.to(x.dtype) + return x, self._get_router_weight_ones( + num_rows, + topk, + topk_weights.dtype, + topk_weights.device, + ) + + def apply( + self, + layer, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + del layer, shared_experts, shared_experts_input + + num_rows, hidden_size = x.shape + topk = topk_ids.shape[1] + + x, gather_topk_weights = self.router_weights_fn( + x, + topk_weights, + num_rows, + topk, + ) + + num_moe_inputs = num_rows * topk + output = torch.empty_like(x) + + if num_moe_inputs == 0: + return output + + assert self.local_num_experts is not None + assert self.global_num_experts is not None + assert self.remap_hidden_states_op is not None + assert self.moe_gather_op is not None + assert self.w13_moe is not None + assert self.w2_moe is not None + assert self.activation is not None + + remapped_hidden_states = torch.empty( + (num_moe_inputs, hidden_size), + dtype=x.dtype, + device=x.device, + ) + rows_per_expert = self._get_rows_per_expert(self.local_num_experts, x.device) + unpermuted_row_to_permuted_row = self._get_unpermuted_row_to_permuted_row( + num_rows, topk, x.device + ) + + self.remap_hidden_states_op( + hidden_states=x, + hidden_states_scales=None, + remapped_hidden_states=remapped_hidden_states, + remapped_hidden_states_scales=None, + expert_map=self.expert_map, + rows_per_expert=rows_per_expert, + unpermuted_row_to_permuted_row=unpermuted_row_to_permuted_row, + topk_ids=topk_ids, + total_experts_num=self.global_num_experts, + local_experts_num=self.local_num_experts, + ) + + gemm1_output = self.w13_moe.apply( + remapped_hidden_states, + rows_per_expert, + phase="auto", + ) + + act_output = gemm1_output.new_empty( + (gemm1_output.shape[0], self.inter_size * self.inter_size_scale) + ) + + apply_moe_activation( + self.activation, + act_output, + gemm1_output, + activation_config=self.activation_config, + ) + + gemm2_output = self.w2_moe.apply( + act_output, + rows_per_expert, + phase="auto", + ) + + self.moe_gather_op( + output, + gemm2_output, + gather_topk_weights, + unpermuted_row_to_permuted_row, + self.local_num_experts, + ) + return output + + +class INCWNA16MoEScheme: + def __init__(self, layer_config: "INCLayerConfig") -> None: + self.layer_config = layer_config + + def get_method(self, layer: torch.nn.Module): + if current_platform.is_cpu(): + return self._build_cpu_method(layer) + if self.layer_config.is_gptq: + return self._build_gptq_method(layer) + if self.layer_config.is_awq: + return self._build_awq_method(layer) + raise NotImplementedError( + f"WNA16 MoE does not support config {self.layer_config}" + ) + + def _build_cpu_method(self, layer: torch.nn.Module): + from vllm.model_executor.layers.fused_moe import ( + UnquantizedFusedMoEMethod, + ) + + return UnquantizedFusedMoEMethod(layer.moe_config) + + def _build_gptq_method(self, layer: torch.nn.Module): + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQMoEMethod, + ) + from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, + ) + from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_moe_marlin_supports_layer, + ) + + group_size = self.layer_config.group_size + assert isinstance(group_size, int), "WNA16 only supports integer group_size." + + use_marlin = (self.layer_config.bits, self.layer_config.sym) in { + (4, True), + (8, True), + } and check_moe_marlin_supports_layer( + layer, + group_size, + ) + + if use_marlin: + return AutoGPTQMoEMethod( + AutoGPTQConfig( + weight_bits=self.layer_config.bits, + group_size=group_size, + desc_act=False, + is_sym=self.layer_config.sym, + lm_head_quantized=False, + dynamic={}, + full_config={}, + ), + layer.moe_config, + ) + + moe_config = MoeWNA16Config.from_config( + { + "quant_method": "gptq", + "bits": self.layer_config.bits, + "group_size": group_size, + "sym": self.layer_config.sym, + "lm_head": False, + } + ) + return MoeWNA16Method(moe_config, layer.moe_config) + + def _build_awq_method(self, layer: torch.nn.Module): + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQMoEMethod + from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, + ) + from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_moe_marlin_supports_layer, + ) + + group_size = self.layer_config.group_size + assert isinstance(group_size, int), "WNA16 only supports integer group_size." + + use_marlin = self.layer_config.bits in ( + 4, + 8, + ) and check_moe_marlin_supports_layer( + layer, + group_size, + ) + + if use_marlin: + return AutoAWQMoEMethod( + AutoAWQConfig( + weight_bits=self.layer_config.bits, + group_size=group_size, + zero_point=not self.layer_config.sym, + lm_head_quantized=False, + modules_to_not_convert=[], + full_config={}, + ), + layer.moe_config, + ) + + moe_config = MoeWNA16Config.from_config( + { + "quant_method": "awq", + "bits": self.layer_config.bits, + "group_size": group_size, + "zero_point": not self.layer_config.sym, + "lm_head": False, + } + ) + return MoeWNA16Method(moe_config, layer.moe_config) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py index 82bb4bbb05b4..60916cd10e08 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py @@ -5,8 +5,6 @@ from vllm import envs from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig -from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.platforms import current_platform from ..inc_linear import INCLinearMethod @@ -100,7 +98,7 @@ def get_linear_method( ) if is_ark_available: return INCLinearMethod(INCARKLinearMethod(layer_config)) - elif layer_config.bits == 2: + if layer_config.bits == 2: raise NotImplementedError( "INC int2 on XPU requires the ARK backend. " f"Layer: {prefix}. " @@ -158,119 +156,68 @@ def get_moe_method( prefix: str, layer_config: "INCLayerConfig", ): - del config, prefix - # CPU does not support quantized MoE yet. - if current_platform.is_cpu(): - from vllm.model_executor.layers.fused_moe import ( - UnquantizedFusedMoEMethod, - ) - - return UnquantizedFusedMoEMethod(layer.moe_config) - # CUDA low-bit (2/3/5/6/7): route to the humming MoE kernel (see above). + del config + if ( - current_platform.is_cuda() - and layer_config.bits in CUDA_HUMMING_SUPPORTED_BITS + current_platform.is_xpu() + and layer_config.is_gptq + and layer_config.bits == 4 + and layer_config.sym ): - return _build_humming_moe_method(layer, layer_config) - if layer_config.is_gptq: - return _resolve_gptq_moe(layer, layer_config) - if layer_config.is_awq: - return _resolve_awq_moe(layer, layer_config) - raise NotImplementedError(f"WNA16 MoE does not support config {layer_config}") - - -def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"): - from vllm.model_executor.layers.quantization.auto_gptq import ( - AutoGPTQMoEMethod, - ) - from vllm.model_executor.layers.quantization.moe_wna16 import ( - MoeWNA16Config, - MoeWNA16Method, - ) - from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_moe_marlin_supports_layer, - ) - - assert isinstance(layer_config.group_size, int), ( - "WNA16 only supports integer group_size." - ) - - # AutoGPTQMoEMethod selects its fused-MoE backend through the WNA16 oracle - # (Marlin on CUDA, XPUExpertsWNA16 on XPU). Gate only on the layer-shape - # check like compressed-tensors does; the capability-based - # check_marlin_supported is skipped so the XPU path is reachable. - use_marlin = (layer_config.bits, layer_config.sym) in { - (4, True), - (8, True), - } and check_moe_marlin_supports_layer(layer, layer_config.group_size) - - if use_marlin: - return AutoGPTQMoEMethod( - AutoGPTQConfig( - weight_bits=layer_config.bits, - group_size=layer_config.group_size, - desc_act=False, - is_sym=layer_config.sym, - lm_head_quantized=False, - dynamic={}, - full_config={}, - ), - layer.moe_config, - ) - - moe_config = MoeWNA16Config.from_config( - { - "quant_method": "gptq", - "bits": layer_config.bits, - "group_size": layer_config.group_size, - "sym": layer_config.sym, - "lm_head": False, - } - ) - return MoeWNA16Method(moe_config, layer.moe_config) + from .inc_ark_ops import get_ark_state + from .inc_wna16_moe import INCARKWNA16MoEMethod + + backend = envs.VLLM_XPU_INC_WNA16_BACKEND + if backend not in XPU_ONEDNN_BACKENDS: + is_ark_available, ark_error, ark, _ = get_ark_state() + xpu_lib = getattr(ark, "xpu_lib", None) if ark is not None else None + is_ark_moe_available = ( + is_ark_available + and ark is not None + and hasattr(ark, "MoeSymmetricGemm") + and xpu_lib is not None + ) + ark_moe_error = ark_error or "ARK MoE kernels are unavailable" + if backend == "ark" and not is_ark_moe_available: + raise NotImplementedError( + "VLLM_XPU_INC_WNA16_BACKEND=ark was requested but " + f"ARK MoE kernels are unavailable: {ark_moe_error}. " + f"Layer: {prefix}." + ) + if is_ark_moe_available: + from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + ) -def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"): - from vllm.model_executor.layers.quantization.auto_awq import AutoAWQMoEMethod - from vllm.model_executor.layers.quantization.moe_wna16 import ( - MoeWNA16Config, - MoeWNA16Method, - ) - from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_moe_marlin_supports_layer, - ) + moe_config = MoeWNA16Config.from_config( + { + "quant_method": "gptq", + "bits": layer_config.bits, + "group_size": layer_config.group_size, + "sym": layer_config.sym, + "lm_head": False, + } + ) + return INCARKWNA16MoEMethod(moe_config, layer.moe_config) - assert isinstance(layer_config.group_size, int), ( - "WNA16 only supports integer group_size." - ) + logger.info( + "ARK backend is unavailable for MoE layer %s; " + "falling back to the default WNA16 MoE path. Error: %s", + prefix, + ark_moe_error, + ) - use_marlin = layer_config.bits in (4, 8) and check_moe_marlin_supports_layer( - layer, layer_config.group_size - ) + # CUDA low-bit (2/3/5/6/7): route to the humming MoE kernel (see above). + if ( + current_platform.is_cuda() + and layer_config.bits in CUDA_HUMMING_SUPPORTED_BITS + ): + return _build_humming_moe_method(layer, layer_config) - if use_marlin: - return AutoAWQMoEMethod( - AutoAWQConfig( - weight_bits=layer_config.bits, - group_size=layer_config.group_size, - zero_point=not layer_config.sym, - lm_head_quantized=False, - modules_to_not_convert=[], - full_config={}, - ), - layer.moe_config, - ) + from .inc_wna16_moe import INCWNA16MoEScheme - moe_config = MoeWNA16Config.from_config( - { - "quant_method": "awq", - "bits": layer_config.bits, - "group_size": layer_config.group_size, - "zero_point": not layer_config.sym, - "lm_head": False, - } - ) - return MoeWNA16Method(moe_config, layer.moe_config) + return INCWNA16MoEScheme(layer_config).get_method(layer) def _humming_weight_config(layer_config: "INCLayerConfig") -> dict: From ed871c91c749d23d1d9a0110ea7655f3b7d4e048 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Sun, 18 Oct 2026 20:55:54 -0700 Subject: [PATCH 848/952] [Bugfix][DSv4.1] Keep the Engram all-to-all buffers in one allocator (#62269) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 --- vllm/models/deepseek_v41/common/engram.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index 236235dfa82f..447feec9413d 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -1410,6 +1410,9 @@ def __init__(self, layers: list["Engram"], max_tokens: int) -> None: dtype=torch.uint8, device=device, ) + # Same allocator as the send buffer: NCCL hangs this all-to-all when its + # buffers mix the cumem pool and the caching allocator. + self.exchanged_rows = torch.empty_like(self.packed_rows) width = table.n_hash_cols * table.dim self.fp8_values = torch.empty( len(layers), max_tokens, width, dtype=torch.uint8, device=device @@ -1477,7 +1480,7 @@ def _stage(self, hashes: torch.Tensor, slot: int, peer: bool) -> None: ) group = get_engram_dp_group() assert group is not None - sent, rows = rows, torch.empty_like(rows) + sent, rows = rows, self.exchanged_rows[: rows.shape[0]] torch.distributed.all_to_all_single(rows, sent, group=group.device_group) table = self.layers[0].embed_tokens heads = table.n_hash_cols From 43caf5b3520b704fa766f00fc2b24da25c78e084 Mon Sep 17 00:00:00 2001 From: Giulio De Pasquale Date: Mon, 19 Oct 2026 05:01:49 +0100 Subject: [PATCH 849/952] [Bugfix][Core] Requeue encoder-stalled waiting requests (#59030) Signed-off-by: Giulio De Pasquale Signed-off-by: Giulio De Pasquale Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: vllm-agent --- tests/v1/core/test_scheduler.py | 65 +++++++++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 6 ++- 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7b2b5b13b667..b233d2b20653 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -6968,6 +6968,71 @@ def test_free_encoder_inputs_unchanged_without_spec_decode(): assert manager.get_cached_input_ids(request) == set() +def test_waiting_encoder_stall_does_not_starve_later_requests(): + """A waiting request stalled by the encoder cache must not end the pass. + + Regression test: the zero-token branch used to `break`, letting the first + encoder-stalled request starve every request queued behind it. + """ + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + scheduler.max_num_encoder_input_tokens = 0 + stalled, ready = create_requests( + num_requests=2, + num_tokens=100, + mm_positions=[[PlaceholderRange(offset=0, length=50)], []], + req_ids=["stalled", "ready"], + ) + scheduler.add_request(stalled) + scheduler.add_request(ready) + + output = scheduler.schedule() + + assert "ready" in output.num_scheduled_tokens + + +def test_waiting_encoder_stall_is_counted_as_deferred(): + """The requeue must go through skip_request's deferred_waiting accounting. + + A hand-rolled pop-and-prepend left the request out of `deferred_waiting`, + skewing the DP prefill-capacity signal. + """ + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + scheduler.max_num_encoder_input_tokens = 0 + stalled = create_requests( + num_requests=1, + num_tokens=100, + mm_positions=[[PlaceholderRange(offset=0, length=50)]], + req_ids=["stalled"], + )[0] + scheduler.add_request(stalled) + + scheduler.schedule() + + assert stalled in scheduler.deferred_waiting + + +def test_waiting_encoder_stall_kv_holder_keeps_drain_priority(): + """A stalled request holding KV blocks must stay on the drain-first queue. + + Async KV-load resumes hold blocks and are drained before plain waiting + requests; a hand-rolled requeue sent them to the plain queue. + """ + scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") + scheduler.max_num_encoder_input_tokens = 0 + stalled = create_requests( + num_requests=1, + num_tokens=100, + mm_positions=[[PlaceholderRange(offset=0, length=50)]], + req_ids=["stalled"], + )[0] + stalled.num_computed_tokens = 10 + scheduler.add_request(stalled) + + scheduler.schedule() + + assert stalled in scheduler.kv_holding_waiting + + def test_encoder_cache_retained_across_preemption_and_resume(): """Regression guard for issue #38551 (preemption path). diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ee9175304d19..6e07654d8f15 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1270,8 +1270,10 @@ def skip_request(from_queue: RequestQueue) -> None: ) if num_new_tokens == 0: - # The request cannot be scheduled. - break + # The request cannot be scheduled. Requeue it so it does + # not stall the requests behind it. + skip_request(request_queue) + continue # During async KV load, no forward pass is run yet. # Allocate speculative lookahead slots later to avoid From 66b431888dd8efdc9afbfb745d3e26be4e7a431a Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 19 Oct 2026 00:16:29 -0400 Subject: [PATCH 850/952] [PCP] Allow decode-only full CUDA graphs for DSv3.2/GLM-5.2/5.3 sparse MLA with DCP=1 (#62377) --- vllm/v1/worker/gpu/pcp_manager.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 16f2886358f7..5143d30c528e 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -189,10 +189,12 @@ def validate_config( is_sparse_mla and parallel_config.decode_context_parallel_size == 1 and cudagraph_mode != CUDAGraphMode.NONE + and getattr(model_config.hf_text_config, "model_type", None) + not in ("deepseek_v32", "glm_moe_dsa") ): raise NotImplementedError( - "MRV2 sparse MLA PCP does not support CUDA graphs yet. " - "Set -cc.cudagraph_mode=NONE." + "MRV2 sparse MLA PCP supports CUDA graphs only for DeepSeek-V3.2 " + "and GLM-5.2/5.3. Set -cc.cudagraph_mode=NONE." ) if ( cudagraph_mode.has_full_cudagraphs() From 6a0b146c1674bac69a783fda8bdbdcd0dfc5f998 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 23:23:43 -0500 Subject: [PATCH 851/952] [ROCm][Compile] Fuse AITER static FP8 quant into MLA attention output (#62371) Signed-off-by: Andreas Karatzas Co-authored-by: Claude Co-authored-by: vllm-agent --- .../passes/test_mla_attn_quant_fusion.py | 55 +++++++++++++++++-- .../passes/fusion/mla_attn_quant_fusion.py | 55 ++++++++++++++++--- .../layers/attention/mla_attention.py | 13 ++++- 3 files changed, 109 insertions(+), 14 deletions(-) diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index a945f000799e..05bfda483a98 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -8,6 +8,7 @@ from tests.compile.backend import LazyInitPass, TestBackend from tests.utils import TestFP8Layer, flat_product from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS from vllm.compilation.passes.fusion.mla_attn_quant_fusion import ( @@ -210,6 +211,18 @@ def forward( return self.fp8_linear(attn_output) +class TestMLAAttentionAiterFp8StaticQuantPatternModel( + TestMLAAttentionFp8StaticQuantPatternModel +): + aiter_quant = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.fp8_linear.kernel.quant_fp8.use_aiter = True + assert self.fp8_linear.input_scale is not None + self.fp8_linear.input_scale.fill_(0.125) + + class TestMLAAttentionNvfp4QuantPatternModel(MLAAttentionQuantPatternModel): """Test model for MLA Attention + NVFP4 quant fusion.""" @@ -550,12 +563,16 @@ def test_mla_attention_quant_pattern( # Check quantization ops in the graph is_per_group = quant_key.scale.group_shape.is_per_group() - quant_op = ( - torch.ops.aten.reciprocal - if "-quant_fp8" in custom_ops_list - else QUANT_OPS[quant_key] + aiter_quant = getattr(model_class, "aiter_quant", False) + if "-quant_fp8" in custom_ops_list: + quant_op = torch.ops.aten.reciprocal + elif aiter_quant or rocm_aiter_ops.is_linear_fp8_enabled(): + quant_op = torch.ops.vllm.rocm_aiter_per_tensor_quant.default + else: + quant_op = QUANT_OPS[quant_key] + test_backend.check_before_ops( + [quant_op], fully_replaced=is_per_group or aiter_quant ) - test_backend.check_before_ops([quant_op], fully_replaced=is_per_group) assert attn_pass.pass_.matched_count == sum(attn_fusion_supported) @@ -588,3 +605,31 @@ def test_mla_attention_quant_pattern( # Check numerical correctness torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not current_platform.is_rocm() or not is_aiter_found_and_supported(), + reason="Requires ROCm AITER MLA attention and quantization", +) +def test_rocm_aiter_static_quant_mla_attention( + dist_init, disable_vllm_compile_cache, workspace_init, monkeypatch +): + """Fuse AITER's static FP8 quant into the AITER MLA attention output.""" + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_ROCM_USE_AITER", "1") + rocm_aiter_ops.refresh_env_variables() + test_mla_attention_quant_pattern( + *MLA_DIMS[0], + batch_size=8, + dtype=torch.bfloat16, + custom_ops="+quant_fp8", + model_name="deepseek-ai/DeepSeek-V2-Lite", + model_class=TestMLAAttentionAiterFp8StaticQuantPatternModel, + backend=AttentionBackendEnum.ROCM_AITER_MLA, + dist_init=dist_init, + disable_vllm_compile_cache=disable_vllm_compile_cache, + workspace_init=workspace_init, + ) + finally: + rocm_aiter_ops.refresh_env_variables() diff --git a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py index 982d90726e06..48ebcb0d93b7 100644 --- a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py @@ -6,6 +6,7 @@ import torch from torch._higher_order_ops.auto_functionalize import auto_functionalized +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import MLAAttention @@ -32,7 +33,9 @@ MLA_ATTN_OP = torch.ops.vllm.unified_mla_attention_with_output.default -class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): +class MLAAttnFp8StaticQuantPattern( + VllmPatternReplacement[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]] +): """Fusion for MLA Attention+Fp8StaticQuant. Matches the pattern: MLA attention -> static FP8 quant, and replaces @@ -50,8 +53,20 @@ def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None: self._dtype = dtype self._quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) + def _quantize( + self, value: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._quant_matcher(value, scale)[0] + + def _output( + self, value: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return value + @property - def pattern(self) -> Callable[..., torch.Tensor]: + def pattern( + self, + ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]: _ln = _encode_layer_name(self._layer_name) if _USE_LAYERNAME: @@ -77,7 +92,7 @@ def _pattern_with_ln( # type: ignore[misc] kv_cache_dummy_dep=kv_cache_dummy_dep, ) # MLA output is already 2D (T, N*V), no reshape needed - return self._quant_matcher(at1[1], scale)[0] + return self._quantize(at1[1], scale) return _pattern_with_ln @@ -94,12 +109,14 @@ def _pattern(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep): kv_cache_dummy_dep=kv_cache_dummy_dep, ) # MLA output is already 2D (T, N*V), no reshape needed - return self._quant_matcher(at1[1], scale)[0] + return self._quantize(at1[1], scale) return _pattern @property - def replacement(self) -> Callable[..., torch.Tensor]: + def replacement( + self, + ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]: _ln = _encode_layer_name(self._layer_name) if _USE_LAYERNAME: @@ -130,7 +147,7 @@ def _replacement_with_ln( # type: ignore[misc] output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, ) - return at1[1] + return self._output(at1[1], scale) return _replacement_with_ln @@ -152,7 +169,7 @@ def _replacement(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep): output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, ) - return at1[1] + return self._output(at1[1], scale) return _replacement @@ -170,6 +187,28 @@ def get_inputs(self) -> list[torch.Tensor]: return inputs +class RocmMLAAttnFp8StaticQuantPattern(MLAAttnFp8StaticQuantPattern): + """Fuse AITER static quantization while preserving its scale output.""" + + def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None: + super().__init__(layer, dtype) + self._quant_matcher = MatcherQuantFP8( + kFp8StaticTensorSym, match_rocm_aiter=True + ) + + def _quantize( + self, value: torch.Tensor, scale: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + # The shared AITER op declares scale mutable for its dynamic variant. + # Static quant preserves it; keep this graph output when replacing it. + return self._quant_matcher(value, scale) + + def _output( + self, value: torch.Tensor, scale: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return value, scale + + class MLAAttnNvfp4QuantPattern( VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]] ): @@ -595,6 +634,8 @@ def __init__(self, config: VllmConfig) -> None: for layer in layers: if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym): self.register(MLAAttnFp8StaticQuantPattern(layer, dtype)) + if current_platform.is_rocm() and rocm_aiter_ops.is_enabled(): + self.register(RocmMLAAttnFp8StaticQuantPattern(layer, dtype)) if _USE_LAYERNAME: break diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 07130abd7f75..cf7ab7ad4595 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -737,6 +737,10 @@ def __init__( group_shape=GroupShape.PER_TENSOR, compile_native=True, ) + # Matches the graph quant replaced by fusion; writes into the output. + self._aiter_output_quant = ( + self._quant_fp8_op.use_aiter and self._quant_fp8_op.enabled() + ) def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: # [B, H=1, N, C] -> [B, N, C] @@ -1188,8 +1192,13 @@ def forward_impl( ) elif quant_key == kFp8StaticTensorSym: # Static FP8 quantization - fp8_data, _ = self._quant_fp8_op(actual, output_scale) - quant_output[:quant_idx].copy_(fp8_data) + if self._aiter_output_quant: + torch.ops.vllm.rocm_aiter_per_tensor_quant( + quant_output[:quant_idx], actual, output_scale, False + ) + else: + fp8_data, _ = self._quant_fp8_op(actual, output_scale) + quant_output[:quant_idx].copy_(fp8_data) else: raise ValueError(f"Unsupported quant_key: {quant_key}") return quant_output From e68efc685ab1e6f1eda004ced7157e3ca411da2f Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 18 Oct 2026 23:26:20 -0500 Subject: [PATCH 852/952] [CI][AMD] Split LM Eval Small by CDNA generation, unify AMD step titles, clone MoRIIO steps into test-amd.yaml (#62363) Signed-off-by: Andreas Karatzas Co-authored-by: Claude Co-authored-by: vllm-agent --- .buildkite/hardware_tests/amd.yaml | 6 +- .buildkite/test-amd.yaml | 103 ++++++++++++------ .buildkite/test_areas/basic_correctness.yaml | 2 +- .buildkite/test_areas/entrypoints.yaml | 2 +- .buildkite/test_areas/lm_eval.yaml | 14 +-- .buildkite/test_areas/lora.yaml | 2 +- .buildkite/test_areas/torch_abi.yaml | 2 +- tests/evals/gsm8k/README.md | 2 +- .../gsm8k/configs/models-gfx950-small.txt | 5 - .../gsm8k/configs/models-small-tp-mi355.txt | 2 - tests/evals/gsm8k/configs/models-small-tp.txt | 1 - .../evals/gsm8k/configs/models-small-tp.yaml | 8 ++ tests/evals/gsm8k/configs/models-small.txt | 8 -- tests/evals/gsm8k/configs/models-small.yaml | 17 +++ tests/evals/gsm8k/conftest.py | 47 ++++++-- tests/evals/gsm8k/test_gsm8k_correctness.py | 2 +- 16 files changed, 147 insertions(+), 76 deletions(-) delete mode 100644 tests/evals/gsm8k/configs/models-gfx950-small.txt delete mode 100644 tests/evals/gsm8k/configs/models-small-tp-mi355.txt delete mode 100644 tests/evals/gsm8k/configs/models-small-tp.txt create mode 100644 tests/evals/gsm8k/configs/models-small-tp.yaml delete mode 100644 tests/evals/gsm8k/configs/models-small.txt create mode 100644 tests/evals/gsm8k/configs/models-small.yaml diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index 66d8a589abc8..9649b948c83d 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -5,7 +5,7 @@ group: Hardware - AMD Build # 2. Build ci_base from either the stable base or the newly selected base. # 3. Build the build-scoped ROCm CI image and smoke-test it before GPU jobs run. steps: - - label: "AMD: :docker: refresh ROCm base" + - label: ":docker: Refresh ROCm base" key: refresh-rocm-base-amd depends_on: [] device: amd_cpu @@ -31,7 +31,7 @@ steps: # Compares a content hash of ci_base-affecting files against the remote # image label. If hashes match the build is skipped (< 30 s); if they # differ ci_base is rebuilt and pushed automatically. - - label: "AMD: :docker: ensure ci_base" + - label: ":docker: Ensure ROCm ci_base" key: ensure-ci-base-amd soft_fail: false depends_on: @@ -60,7 +60,7 @@ steps: - exit_status: -10 # Agent was lost limit: 1 - - label: "AMD: :docker: build test image and artifacts" + - label: ":docker: Build ROCm test image and artifacts" key: image-build-amd soft_fail: false depends_on: diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a28227689dc6..156b501257c0 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -111,7 +111,7 @@ steps: # # ######################################################################################################################################### -- label: ":computer: (CPU) Docker Build Metadata" # TBD +- label: ":amd: (MI250) Docker Build Metadata" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] dind: false @@ -134,7 +134,7 @@ steps: commands: - pytest -v -s tools/test_docker_build_metadata_args.py -- label: ":amd: (MI250) Torch Stable ABI Audit" # TBD +- label: ":amd: (MI300) Torch Stable ABI Audit" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1073,6 +1073,25 @@ steps: #----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# +- label: ":amd: (MI300) LM Eval Small Models" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + optional: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml + - label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1605,7 +1624,7 @@ steps: #----------------------------------------------------------- mi300 · lora ------------------------------------------------------------# -- label: ":amd: (MI300) LoRA TP (Distributed) %N" +- label: ":amd: (MI300) LoRA TP (Distributed) Shard %N" timeout_in_minutes: 65 parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2095,7 +2114,7 @@ steps: - pytest -v -s entrypoints/serve/instrumentator/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" -- label: ":amd: (MI300) Rust Frontend Serve-Admin Coverage" +- label: ":amd: (MI300) Rust Frontend Serve/Admin Coverage" timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -2830,7 +2849,7 @@ steps: #----------------------------------------------------- mi355 · basic_correctness -----------------------------------------------------# -- label: ":amd: (MI355 DPX) Basic Correctness Models %N" # TBD +- label: ":amd: (MI355 DPX) Basic Correctness Models Shard %N" # TBD timeout_in_minutes: 60 parallelism: 2 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3543,7 +3562,7 @@ steps: - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests -- label: ":amd: (MI355 DPX) Entrypoints Integration (API Server) %N" +- label: ":amd: (MI355 DPX) Entrypoints Integration (API Server) Shard %N" timeout_in_minutes: 65 parallelism: 4 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3745,7 +3764,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-gfx950-small.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml - label: ":amd: (MI355 DPX) LM Eval Watermarking" timeout_in_minutes: 45 @@ -3792,7 +3811,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt -- label: ":amd: (MI355) LM Eval Qwen3-5 Models" # TBD +- label: ":amd: (MI355) LM Eval Qwen3.5 Models" # TBD timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3861,7 +3880,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep-mi355.txt -- label: ":amd: (MI355) Qwen3-8-Flash-Next-FP8 Accuracy Eval" +- label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" timeout_in_minutes: 105 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -4042,11 +4061,9 @@ steps: - tests/utils.py - vllm/distributed/ - vllm/v1/attention/ops/ - - tests/evals/gsm8k/configs/models-small-tp-mi355.txt commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp-mi355.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.yaml - label: ":amd: (MI355) MoE Refactor Integration Shard %N" timeout_in_minutes: 105 @@ -4304,25 +4321,6 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-nemotron-3-super.txt -- label: ":amd: (MI355) LM Eval Small Models" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - optional: true - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt - - label: ":amd: (MI355 DPX) LM Eval KV-Offload" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -5600,6 +5598,49 @@ steps: - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB optional: true +- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/unit/test_moriio_ + - tests/v1/kv_connector/unit/utils.py + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py + +- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + parallelism: 6 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - vllm/model_executor/model_loader/dummy_loader.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/moriio_integration/ + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - label: ":amd: (MI355 DPX) V1 Metrics + LM Eval" timeout_in_minutes: 65 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index d1d595ad0a98..95471746e921 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -16,7 +16,7 @@ steps: - pytest -v -s basic_correctness/models --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: - label: ":amd: (MI355 DPX) Basic Correctness Models %N" + label: ":amd: (MI355 DPX) Basic Correctness Models Shard %N" dind: false device: mi355_dpx timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index ae59d3cb5376..f23a51577c59 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -90,7 +90,7 @@ steps: - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: - label: ":amd: (MI355 DPX) Entrypoints Integration (API Server) %N" + label: ":amd: (MI355 DPX) Entrypoints Integration (API Server) Shard %N" dind: false device: mi355_dpx timeout_in_minutes: 65 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index cb4e7eaedcd2..0b44b5d5877f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -12,12 +12,12 @@ steps: - vllm/model_executor/layers/quantization autorun_on_main: true commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml mirror: amd: - label: ":amd: (MI355) LM Eval Small Models" + label: ":amd: (MI300) LM Eval Small Models" dind: false - device: mi355_dpx + device: mi300_1 timeout_in_minutes: 45 depends_on: - image-build-amd @@ -173,7 +173,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-gfx950-small.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml - label: ":nvidia: (B200) LM Eval Small Models Distributed" key: lm-eval-small-models-distributed-2xb200 @@ -186,7 +186,7 @@ steps: - vllm/model_executor/layers/quantization autorun_on_main: true commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.yaml mirror: amd: label: ":amd: (MI355) LM Eval Small Models Distributed" @@ -214,9 +214,7 @@ steps: - vllm/v1/attention/ops/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt - # Dense MLA PCP-only accuracy; DP1/DCP1. - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp-mi355.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.yaml - label: ":nvidia: (B200) LM Eval HiSparse Nightly" key: lm-eval-hisparse-nightly-4xb200 diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 8321bc7a4f0a..1fa94247ccda 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -57,7 +57,7 @@ steps: - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp mirror: amd: - label: ":amd: (MI300) LoRA TP (Distributed) %N" + label: ":amd: (MI300) LoRA TP (Distributed) Shard %N" dind: false device: mi300_4 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml index ee95203a0143..576f9b8c9e10 100644 --- a/.buildkite/test_areas/torch_abi.yaml +++ b/.buildkite/test_areas/torch_abi.yaml @@ -14,7 +14,7 @@ steps: - python3 /vllm-workspace/.buildkite/check-torch-abi.py mirror: amd: - label: ":amd: (MI250) Torch Stable ABI Audit" + label: ":amd: (MI300) Torch Stable ABI Audit" dind: false device: mi300_1 timeout_in_minutes: 35 diff --git a/tests/evals/gsm8k/README.md b/tests/evals/gsm8k/README.md index ea61b6bf4562..d570259c0506 100644 --- a/tests/evals/gsm8k/README.md +++ b/tests/evals/gsm8k/README.md @@ -8,7 +8,7 @@ This directory contains a replacement for the lm-eval-harness GSM8K evaluation, ```bash pytest -s -v tests/evals/gsm8k/test_gsm8k_correctness.py \ - --config-list-file=configs/models-small.txt + --config-list-file=configs/models-small.yaml ``` ### Run standalone evaluation script diff --git a/tests/evals/gsm8k/configs/models-gfx950-small.txt b/tests/evals/gsm8k/configs/models-gfx950-small.txt deleted file mode 100644 index 05b78e934c6b..000000000000 --- a/tests/evals/gsm8k/configs/models-gfx950-small.txt +++ /dev/null @@ -1,5 +0,0 @@ -Qwen3-0.6B-FP8.yaml -Qwen3-1.7B-MXFP4.yaml -Qwen3-30B-A3B-NVFP4.yaml -Laguna-XS.2-NVFP4.yaml -Laguna-XS-2.1-INT4.yaml diff --git a/tests/evals/gsm8k/configs/models-small-tp-mi355.txt b/tests/evals/gsm8k/configs/models-small-tp-mi355.txt deleted file mode 100644 index 78283f92af6b..000000000000 --- a/tests/evals/gsm8k/configs/models-small-tp-mi355.txt +++ /dev/null @@ -1,2 +0,0 @@ -# Dense MLA PCP accuracy on MI355. -DeepSeek-V2-Lite-Instruct-FP8-TP1-PCP2-ROCm.yaml diff --git a/tests/evals/gsm8k/configs/models-small-tp.txt b/tests/evals/gsm8k/configs/models-small-tp.txt deleted file mode 100644 index 63bba5bcd1d6..000000000000 --- a/tests/evals/gsm8k/configs/models-small-tp.txt +++ /dev/null @@ -1 +0,0 @@ -DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml diff --git a/tests/evals/gsm8k/configs/models-small-tp.yaml b/tests/evals/gsm8k/configs/models-small-tp.yaml new file mode 100644 index 000000000000..70d3bcbaffbd --- /dev/null +++ b/tests/evals/gsm8k/configs/models-small-tp.yaml @@ -0,0 +1,8 @@ +# Every machine runs `common`, then the sections named after its +# current_platform.device_name (cuda, rocm) and, on ROCm, its CDNA generation +# (cdna3: gfx942/MI300, cdna4: gfx950/MI355). +common: +- DiffusionGemma-26B-A4B-it-FP8-dynamic.yaml +cdna4: +# Dense MLA PCP accuracy (DP1/DCP1). +- DeepSeek-V2-Lite-Instruct-FP8-TP1-PCP2-ROCm.yaml diff --git a/tests/evals/gsm8k/configs/models-small.txt b/tests/evals/gsm8k/configs/models-small.txt deleted file mode 100644 index 8ab287e31f6a..000000000000 --- a/tests/evals/gsm8k/configs/models-small.txt +++ /dev/null @@ -1,8 +0,0 @@ -Qwen3-0.6B-FP8.yaml -Llama-3.2-1B-Instruct-INT8-CT.yaml -Llama-3-8B-Instruct-nonuniform-CT.yaml -Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml -Qwen1.5-MoE-W4A16-CT.yaml -DeepSeek-V2-Lite-Instruct-FP8.yaml -Qwen3-30B-A3B-MXFP4A16.yaml -gemma-4-E4B-it-qat-mobile-ct.yaml diff --git a/tests/evals/gsm8k/configs/models-small.yaml b/tests/evals/gsm8k/configs/models-small.yaml new file mode 100644 index 000000000000..dc2bbbcb07c5 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-small.yaml @@ -0,0 +1,17 @@ +# Every machine runs `common`, then the sections named after its +# current_platform.device_name (cuda, rocm) and, on ROCm, its CDNA generation +# (cdna3: gfx942/MI300, cdna4: gfx950/MI355). +common: +- Qwen3-0.6B-FP8.yaml +- Llama-3.2-1B-Instruct-INT8-CT.yaml +- Llama-3-8B-Instruct-nonuniform-CT.yaml +- Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml +- Qwen1.5-MoE-W4A16-CT.yaml +- DeepSeek-V2-Lite-Instruct-FP8.yaml +- Qwen3-30B-A3B-MXFP4A16.yaml +- gemma-4-E4B-it-qat-mobile-ct.yaml +cdna4: +- Qwen3-1.7B-MXFP4.yaml +- Qwen3-30B-A3B-NVFP4.yaml +- Laguna-XS.2-NVFP4.yaml +- Laguna-XS-2.1-INT4.yaml diff --git a/tests/evals/gsm8k/conftest.py b/tests/evals/gsm8k/conftest.py index 6f25fe6414af..2df97e3010d7 100644 --- a/tests/evals/gsm8k/conftest.py +++ b/tests/evals/gsm8k/conftest.py @@ -3,16 +3,42 @@ from pathlib import Path +import yaml + +from vllm.platforms import current_platform + def pytest_addoption(parser): """Add custom command line options.""" parser.addoption( "--config-list-file", - default="configs/models-small.txt", + default="configs/models-small.yaml", help="File containing list of config files to test", ) +def _list_keys() -> list[str]: + """The .yaml list sections this machine runs: ``common``, its platform's + device name and, on ROCm, its CDNA generation (``cdna3``, ``cdna4``).""" + keys = ["common", current_platform.device_name] + if current_platform.is_rocm(): + from vllm.platforms.rocm import get_cdna_version + + if cdna := get_cdna_version(): + keys.append(f"cdna{cdna}") + return keys + + +def _config_names(config_list_path: Path) -> list[str]: + """Config file names from a .txt list, or from a .yaml list's sections + that this machine runs (see ``_list_keys``).""" + if config_list_path.suffix in (".yaml", ".yml"): + lists = yaml.safe_load(config_list_path.read_text()) or {} + return [name for key in _list_keys() for name in lists.get(key) or []] + lines = (line.strip() for line in config_list_path.read_text().splitlines()) + return [line for line in lines if line and not line.startswith("#")] + + def pytest_generate_tests(metafunc): """Generate test parameters from config files.""" if "config_filename" in metafunc.fixturenames: @@ -36,17 +62,14 @@ def pytest_generate_tests(metafunc): # Determine config directory (same directory as the list file) config_dir = config_list_path.parent - with open(config_list_path) as f: - for line in f: - line = line.strip() - if line and not line.startswith("#"): - config_path = config_dir / line - print(f"Checking config file: {config_path}") - if config_path.exists(): - config_files.append(config_path) - print(f" ✓ Found: {config_path}") - else: - print(f" ✗ Missing: {config_path}") + for name in _config_names(config_list_path): + config_path = config_dir / name + print(f"Checking config file: {config_path}") + if config_path.exists(): + config_files.append(config_path) + print(f" ✓ Found: {config_path}") + else: + print(f" ✗ Missing: {config_path}") else: print(f"Config list file not found: {config_list_path}") diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index 3b66de6ad022..589a2e87cc39 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -5,7 +5,7 @@ Usage: pytest -s -v tests/evals/gsm8k/test_gsm8k_correctness.py \ - --config-list-file=configs/models-small.txt + --config-list-file=configs/models-small.yaml """ import shlex From 38168d5b82f1c4015f8f37372f1670e11b6afdba Mon Sep 17 00:00:00 2001 From: Iryna Boiko Date: Mon, 19 Oct 2026 07:04:33 +0200 Subject: [PATCH 853/952] [XPU] [Feature] Adding triton/sycl GDN toggle (#56835) Signed-off-by: Iryna Boiko Signed-off-by: Iryna Boiko Co-authored-by: Kunshang Ji Co-authored-by: vllm-agent --- vllm/engine/arg_utils.py | 4 +- .../layers/mamba/gdn/olmo_gdn_linear_attn.py | 12 ++++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 66 ++++++++++++++++--- vllm/v1/attention/backends/gdn_attn.py | 2 +- 4 files changed, 71 insertions(+), 13 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ccf66f15b9dc..0aebd723f0b4 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -860,7 +860,7 @@ class EngineArgs: fail_on_environ_validation: bool = False gdn_prefill_backend: ( - Literal["flashinfer", "triton", "cutedsl", "aiter_flydsl"] | None + Literal["flashinfer", "triton", "cutedsl", "aiter_flydsl", "sycl"] | None ) = None kda_prefill_backend: ( Literal["auto", "triton", "flashkda", "flashinfer", "fused"] | None @@ -1930,7 +1930,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: parser.add_argument( "--gdn-prefill-backend", dest="gdn_prefill_backend", - choices=["flashinfer", "triton", "cutedsl", "aiter_flydsl"], + choices=["flashinfer", "triton", "cutedsl", "aiter_flydsl", "sycl"], default=None, help="Select GDN prefill backend.", ) diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index a7bbac74f66d..bd30103019e3 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -21,6 +21,10 @@ RowParallelLinear, ) from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _log_gdn_backend_decision, + _resolve_gdn_prefill_backend, +) from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateShapeCalculator, is_conv_state_dim_first, @@ -85,6 +89,14 @@ def __init__( self.value_dim = self.head_v_dim * self.num_v_heads self.allow_neg_eigval = getattr(config, "linear_allow_neg_eigval", False) + # OLMo always runs FLA; the resolver may still pick FlashInfer on CUDA. + requested_backend, _ = _resolve_gdn_prefill_backend(vllm_config) + if requested_backend == "sycl": + raise ValueError( + "GDN prefill backend 'sycl' is not supported for OLMo Hybrid." + ) + _log_gdn_backend_decision(vllm_config, requested_backend, "triton") + # Fused QKVG projection: 1 matmul instead of 4 self.in_proj_qkvg = MergedColumnParallelLinear( input_size=self.hidden_size, diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 774e6cde60a8..547a37593906 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -107,6 +107,11 @@ FUSED_GDN_STATE_DTYPES = (torch.float32, torch.bfloat16) +def xpu_gdn_kernel_available() -> bool: + """Whether vllm-xpu-kernels was built with the fused SYCL GDN op.""" + return hasattr(torch.ops._xpu_C, "gdn_attention") + + def _aiter_flydsl_unsupported_reasons( vllm_config: VllmConfig, head_k_dim: int | None, head_v_dim: int | None ) -> list[str]: @@ -126,9 +131,14 @@ def _aiter_flydsl_unsupported_reasons( def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, -) -> tuple[str, Literal["triton", "flashinfer", "cutedsl", "aiter_flydsl"]]: +) -> tuple[str, Literal["triton", "flashinfer", "cutedsl", "aiter_flydsl", "sycl"]]: """Resolve GDN prefill backend. + The SYCL kernel is chosen when ``requested in ["sycl", "auto"]``, the + platform is XPU and ``_xpu_C::gdn_attention`` is built. Unlike the other + backends it is a fused whole-layer op, so selecting it also replaces the + conv1d and decode kernels, not just prefill. + FlashInfer's GDN prefill kernel is chosen when: * ``requested in ["flashinfer", "auto"]``; * ``platform == cuda``; @@ -154,6 +164,20 @@ def _resolve_gdn_prefill_backend( ) backend = str(backend_cfg).strip().lower() + if backend == "sycl" and not current_platform.is_xpu(): + raise ValueError("GDN prefill backend 'sycl' is only supported on XPU.") + + if current_platform.is_xpu(): + if backend in ("auto", "sycl") and xpu_gdn_kernel_available(): + return backend, "sycl" + if backend == "sycl": + raise ValueError( + "GDN prefill backend 'sycl' was requested but " + "torch.ops._xpu_C.gdn_attention is not available. Rebuild " + "vllm-xpu-kernels with VLLM_GDN_ENABLED." + ) + return backend, "triton" + head_k_dim = getattr( vllm_config.model_config.hf_text_config, "linear_key_head_dim", None ) @@ -255,6 +279,7 @@ def _log_gdn_backend_decision( "cutedsl": "CuteDSL", "aiter_flydsl": "AITER FlyDSL", "triton": "Triton/FLA", + "sycl": "SYCL", }[active_backend] logger.info_once( "Using %s GDN prefill kernel (requested=%s, head_k_dim=%s).", @@ -359,6 +384,9 @@ def __init__(self) -> None: ) _log_gdn_backend_decision(vllm_config, backend, active_backend) + # On XPU "sycl" bypasses this module entirely (fused whole-layer op), + # so the chunked prefill path stays on Triton/FLA. + if active_backend == "flashinfer": self._forward_method = self.forward_cuda elif active_backend == "cutedsl": @@ -543,8 +571,12 @@ def __init__( self.key_dim = self.head_k_dim * self.num_k_heads self.value_dim = self.head_v_dim * self.num_v_heads self.gqa_interleaved_layout = gqa_interleaved_layout + self.gdn_xpu_backend: str | None = None self.qkvz_layout = "interleaved" if gqa_interleaved_layout else "flat" if current_platform.is_xpu(): + requested, gdn_xpu_backend = _resolve_gdn_prefill_backend(vllm_config) + _log_gdn_backend_decision(vllm_config, requested, gdn_xpu_backend) + self.gdn_xpu_backend = gdn_xpu_backend self._forward_method = self.forward_xpu elif current_platform.is_cpu(): from vllm.model_executor.layers.mamba.ops.cpu.gdn_attention import ( @@ -1177,7 +1209,7 @@ def forward_xpu( ) -> torch.Tensor: """Forward pass with three parts: 1. Input projection - 2. Core attention (custom op) + 2. Core attention (custom op, SYCL or Triton per --gdn-prefill-backend) 3. Output projection """ num_tokens = hidden_states.size(0) @@ -1196,15 +1228,29 @@ def forward_xpu( dtype=hidden_states.dtype, device=hidden_states.device, ) - z = torch.empty_like(core_attn_out) - torch.ops.vllm.gdn_attention_core_xpu( - core_attn_out, - z, - projected_states_qkvz, - projected_states_ba, - self.prefix, - ) + if self.gdn_xpu_backend == "triton": + # Triton consumes unpacked q/k/v/z/b/a; the SYCL op unpacks + # internally and takes the raw projections instead. + mixed_qkv, z, b, a = self.prepare_gdn_attention_core_inputs( + projected_states_qkvz, projected_states_ba, num_tokens + ) + torch.ops.vllm.qwen_gdn_attention_core( + mixed_qkv, + b.contiguous(), + a.contiguous(), + core_attn_out, + layer_name=_encode_layer_name(self.prefix), + ) + else: + z = torch.empty_like(core_attn_out) + torch.ops.vllm.gdn_attention_core_xpu( + core_attn_out, + z, + projected_states_qkvz, + projected_states_ba, + self.prefix, + ) # ============================================================ # Part 3: Output Projection diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 6b99ec5452f9..660ef5387b17 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -155,7 +155,7 @@ def __init__( ) self.gdn_prefill_backend: Literal[ - "triton", "flashinfer", "cutedsl", "aiter_flydsl" + "triton", "flashinfer", "cutedsl", "aiter_flydsl", "sycl" ] _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config) self._check_chunk_metadata_override(type(self), self.gdn_prefill_backend) From dec6ddadc08a7bae08ff4d2f8939faa308cf7f81 Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Sun, 18 Oct 2026 22:05:28 -0700 Subject: [PATCH 854/952] [CPU] Optimize routed INT8 W8A8 MoE GEMM dispatch (#60740) Signed-off-by: Li, Tianmu --- csrc/cpu/sgl-kernels/gemm.h | 6 ++ csrc/cpu/sgl-kernels/moe_int8.cpp | 71 +++++++------------ tests/kernels/moe/test_cpu_quant_fused_moe.py | 34 +++++++++ 3 files changed, 64 insertions(+), 47 deletions(-) diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h index dc5f9c585d5d..13f9760a7ce4 100644 --- a/csrc/cpu/sgl-kernels/gemm.h +++ b/csrc/cpu/sgl-kernels/gemm.h @@ -64,6 +64,12 @@ inline bool can_use_brgemm(int M, int N) { return brgemm_supported() && (can_use_brgemm(M) || N != block_size_n()); } +// Routed INT8 MoE uses a lower threshold than the M-only GEMM selector. +template <> +inline bool can_use_brgemm(int M, int N) { + return brgemm_supported() && (M > 3 || N != block_size_n()); +} + // work around compiler internal error #define BLOCK_K 128 // 4 * TILE_K diff --git a/csrc/cpu/sgl-kernels/moe_int8.cpp b/csrc/cpu/sgl-kernels/moe_int8.cpp index 4bdddb9c53e0..e0b81505b726 100644 --- a/csrc/cpu/sgl-kernels/moe_int8.cpp +++ b/csrc/cpu/sgl-kernels/moe_int8.cpp @@ -138,27 +138,7 @@ struct tinygemm_kernel_vnni { int64_t lda, int64_t ldb, int64_t ldc) { - TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); - } -}; - #if defined(CPU_CAPABILITY_AVX512) -template -struct tinygemm_kernel_vnni { - static inline void apply( - const uint8_t* __restrict__ A, - const int8_t* __restrict__ B0, - const int8_t* __restrict__ B1, - at::BFloat16* __restrict__ C, - const float* __restrict__ As, - const float* __restrict__ Bs0, - const float* __restrict__ Bs1, - const int32_t* __restrict__ Bcomp0, - const int32_t* __restrict__ Bcomp1, - int64_t K, - int64_t lda, - int64_t ldb, - int64_t ldc) { constexpr int ROWS = BLOCK_M; constexpr int COLS = BLOCK_N / 16; static_assert(COLS % 2 == 0); @@ -227,6 +207,8 @@ struct tinygemm_kernel_vnni { }; Unroll{}(scalec); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; auto storec = [&](auto i) { constexpr int row = i / COLS; constexpr int col = i % COLS; @@ -239,15 +221,16 @@ struct tinygemm_kernel_vnni { x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0); x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1); - _mm512_storeu_si512( - reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), - (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0)))); + bVec out_vec = convert_from_float_ext(fVec(x0), fVec(x1)); + out_vec.store(C + row * ldc + col * 16); } }; Unroll{}(storec); +#else + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); +#endif } }; -#endif #define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \ tinygemm_kernel_vnni::apply( \ @@ -329,24 +312,7 @@ struct tinygemm_kernel_vnni2 { int64_t lda, int64_t ldb, int64_t ldc) { - TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); - } -}; - #if defined(CPU_CAPABILITY_AVX512) -template -struct tinygemm_kernel_vnni2 { - static inline void apply( - const uint8_t* __restrict__ A, - const int8_t* __restrict__ B, - float* __restrict__ C, - const float* __restrict__ As, - const float* __restrict__ Bs, - const int32_t* __restrict__ Bcomp, - int64_t K, - int64_t lda, - int64_t ldb, - int64_t ldc) { constexpr int ROWS = BLOCK_M; constexpr int COLS = BLOCK_N / 16; static_assert(COLS % 2 == 0); @@ -406,9 +372,11 @@ struct tinygemm_kernel_vnni2 { _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x); }; Unroll{}(storec); +#else + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); +#endif } }; -#endif #define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE) \ tinygemm_kernel_vnni2::apply( \ @@ -522,9 +490,6 @@ void fused_experts_int8_kernel_impl( const int64_t stride_e = 2 * N * packed_K; const int64_t stride_n = packed_K; - int64_t avg_M = std::max(int64_t(1), M * topk / E); - const bool use_brgemm = can_use_brgemm(avg_M); - // here we only parallel on half of 2N to fuse silu_and_mul with gemm parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { // get local pointers @@ -532,6 +497,7 @@ void fused_experts_int8_kernel_impl( uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K; int32_t* __restrict__ C0 = reinterpret_cast(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N; int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N; + bool use_brgemm = false; alignas(64) float As[BLOCK_M]; @@ -548,6 +514,11 @@ void fused_experts_int8_kernel_impl( const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb_lower * BLOCK_N; int64_t m_size = offsets[mb + 1] - offsets[mb]; + if (m_size == 0) { + return; + } + const bool brg = can_use_brgemm(m_size, n_size); + use_brgemm |= brg; if (nb_offset == 0) { // 1.a load A @@ -559,7 +530,7 @@ void fused_experts_int8_kernel_impl( } } - if (use_brgemm) { + if (brg) { // 1.b gemm: C0 = A @ B0 at::native::cpublas::brgemm( /* M */ m_size, @@ -640,10 +611,16 @@ void fused_experts_int8_kernel_impl( int tid = get_thread_num(); float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N; int32_t* __restrict__ C32 = reinterpret_cast(C + BLOCK_M * BLOCK_N); + bool use_brgemm = false; loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) { int64_t m_size = offsets[mb + 1] - offsets[mb]; + if (m_size == 0) { + return; + } int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N); + const bool brg = can_use_brgemm(m_size, n_size); + use_brgemm |= brg; // A ptr from ic1 of [M * topk, N] in sorted order // so as to avoid copy A to tmp buffer again @@ -657,7 +634,7 @@ void fused_experts_int8_kernel_impl( const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N; // 2.a gemm: C = A @ B - if (use_brgemm) { + if (brg) { at::native::cpublas::brgemm( /* M */ m_size, /* N */ n_size, diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index e89fbb55a089..02e4120d2d6c 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -1114,6 +1114,40 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace): ) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_int8_w8a8_cpu_fused_moe_small_expert_blocks(dtype): + """Test INT8 dispatch at exact and multi-block expert boundaries.""" + set_random_seed(0) + block_sizes = (3, 4, 5, 33) + N, K, E = 128, 128, len(block_sizes) + M = sum(block_sizes) + + a = torch.randn(M, K, dtype=dtype) / (0.5 * K**0.5) + w1_q, w2_q, w1_s, w2_s = _make_int8_moe_weights(E, N, K) + topk_weight, topk_ids = _deterministic_expert_routes(block_sizes) + + ref_out = _ref_int8_moe(a, w1_q, w2_q, w1_s, w2_s, topk_weight, topk_ids) + pw1, pw2 = _prepack_experts(w1_q), _prepack_experts(w2_q) + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, + a, + pw1, + pw2, + topk_weight, + topk_ids, + ops.CPUQuantMethod.INT8_W8A8, + w1_s, + w2_s, + None, + None, + None, + is_vnni=True, + ) + + torch.testing.assert_close(ref_out, out, atol=2e-1, rtol=2e-1) + + # =========================================================================== # FP8 W8A8 MoE # =========================================================================== From 1aa5b6d9fbcc0984675d978195a31830b058d2aa Mon Sep 17 00:00:00 2001 From: xiaolinchen <2990624738@qq.com> Date: Mon, 19 Oct 2026 13:06:03 +0800 Subject: [PATCH 855/952] [BugFix][MyPy] Guard Linux-only mmap/fadvise calls with sys.platform (#61836) Signed-off-by: wentian-byte <2990624738@qq.com> Co-authored-by: vllm-agent Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/model_executor/model_loader/weight_utils.py | 3 +++ vllm/models/deepseek_v41/common/engram.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 3262e9408846..03d2a218af95 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -9,6 +9,7 @@ import hashlib import json import os +import sys import tempfile import threading import time @@ -898,6 +899,8 @@ def drop_checkpoint_cache( revision: str | None = None, cache_dir: str | None = None, ) -> None: + if sys.platform != "linux": + return checkpoint_dir: Path | None = None if os.path.isdir(model_name_or_path): checkpoint_dir = Path(model_name_or_path) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index 447feec9413d..e92b0a1371b7 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/vllm/models/deepseek_v41/common/engram.py @@ -37,6 +37,7 @@ import ctypes import mmap import os +import sys import tempfile import weakref from contextlib import ExitStack @@ -906,6 +907,8 @@ def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool: def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None: """Register prefaulted huge pages, or return None for pinned-memory fallback.""" + if sys.platform != "linux": + return None try: mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS) mapping.madvise(mmap.MADV_HUGEPAGE) From 93de77c5ed7ecde1b92649318e95a1e1c45edf5e Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Sun, 18 Oct 2026 22:12:19 -0700 Subject: [PATCH 856/952] [CPU] Fix underflow in FP8 AMX attention (#61329) Signed-off-by: Li, Tianmu Co-authored-by: Codex --- csrc/cpu/cpu_attn_amx.hpp | 22 +++----- csrc/cpu/cpu_attn_fp8.hpp | 69 +++++++++++------------- csrc/cpu/cpu_types_x86.hpp | 32 +++++------ tests/kernels/attention/test_cpu_attn.py | 43 +++++++++++++++ 4 files changed, 94 insertions(+), 72 deletions(-) diff --git a/csrc/cpu/cpu_attn_amx.hpp b/csrc/cpu/cpu_attn_amx.hpp index 6a0341085dce..0bd4fc090543 100644 --- a/csrc/cpu/cpu_attn_amx.hpp +++ b/csrc/cpu/cpu_attn_amx.hpp @@ -415,19 +415,18 @@ class AttentionImpl { void init_from_input(const AttentionInput* input) { if constexpr (fp8_kv) { - k_scale = input->k_scale_fp8; - v_scale = input->v_scale_fp8; + // Compensate for the E4M3/E5M2 decoder's 2^8/2^16 exponent lift. + constexpr float dequant_scale = + std::is_same_v ? 0x1.0p-8f + : 0x1.0p-16f; + k_scale = input->k_scale_fp8 * dequant_scale; + v_scale = input->v_scale_fp8 * dequant_scale; } } float get_output_v_scale() const noexcept { if constexpr (fp8_kv) { - // AMX dequant places FP8 payload into a BF16 field (exponent bias 127). - // Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 → - // 2^112. - constexpr float bias = - std::is_same_v ? 0x1p112f : 0x1p120f; - return v_scale * bias; + return v_scale; } return 1.0f; } @@ -435,12 +434,7 @@ class AttentionImpl { template